#!/usr/bin/env python3
"""Benchmark orchestration for DeepSeek-V4 inference comparison.

Compares:
- vllm-dspark + DeepSeek-V4-Flash-DSpark + DSpark (various spec-tokens)
- vllm-dspark + DeepSeek-V4-Flash-DSpark without spec decode
- vllm (0.24.0) + DeepSeek-V4-Flash without spec decode
- sglang + DeepSeek-V4-Flash + EAGLE (reuse existing results)
"""

import json
import os
import subprocess
import sys
import time
import urllib.request
from datetime import datetime
from pathlib import Path

ROOT = Path("/data/user1/yy")
RESULT_DIR = ROOT / "bench_results" / f"dsv4_comparison_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
RESULT_DIR.mkdir(parents=True, exist_ok=True)

DATASET = "/data/user1/yy/datasets/ShareGPT_V4.3_unfiltered_cleaned_split.json"
NUM_PROMPTS = 200
SEED = 42
OUTPUT_LEN = 256
HOST = "127.0.0.1"

# Service configs to benchmark
SERVICES = [
    {
        "name": "vllm-dspark-dspark-st5",
        "engine": "vllm-dspark",
        "model": "/data/models/DeepSeek-V4-Flash-DSpark",
        "port": 30004,
        "spec_method": "dspark",
        "spec_tokens": 5,
        "cmd": [
            "/data/user1/yy/envs/vllm-dspark/bin/vllm", "serve",
            "/data/models/DeepSeek-V4-Flash-DSpark",
            "--trust-remote-code",
            "--tensor-parallel-size", "8",
            "--kv-cache-dtype", "fp8",
            "--block-size", "256",
            "--max-model-len", "auto",
            "--max-num-seqs", "256",
            "--tokenizer-mode", "deepseek_v4",
            "--reasoning-parser", "deepseek_v4",
            "--spec-method", "dspark",
            "--spec-model", "/data/models/DeepSeek-V4-Flash-DSpark",
            "--spec-tokens", "5",
            "--no-disable-hybrid-kv-cache-manager",
            "--disable-uvicorn-access-log",
            "--port", "30004",
        ],
        "env": {"CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", "TMPDIR": str(ROOT / "tmp")},
        "bench_backend": "openai",
    },
    {
        "name": "vllm-dspark-nospec",
        "engine": "vllm-dspark",
        "model": "/data/models/DeepSeek-V4-Flash-DSpark",
        "port": 30004,
        "spec_method": None,
        "cmd": [
            "/data/user1/yy/envs/vllm-dspark/bin/vllm", "serve",
            "/data/models/DeepSeek-V4-Flash-DSpark",
            "--trust-remote-code",
            "--tensor-parallel-size", "8",
            "--kv-cache-dtype", "fp8",
            "--block-size", "256",
            "--max-model-len", "auto",
            "--max-num-seqs", "256",
            "--tokenizer-mode", "deepseek_v4",
            "--reasoning-parser", "deepseek_v4",
            "--no-disable-hybrid-kv-cache-manager",
            "--disable-uvicorn-access-log",
            "--port", "30004",
        ],
        "env": {"CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", "TMPDIR": str(ROOT / "tmp")},
        "bench_backend": "openai",
    },
    {
        "name": "vllm-dspark-dspark-st3",
        "engine": "vllm-dspark",
        "model": "/data/models/DeepSeek-V4-Flash-DSpark",
        "port": 30004,
        "spec_method": "dspark",
        "spec_tokens": 3,
        "cmd": [
            "/data/user1/yy/envs/vllm-dspark/bin/vllm", "serve",
            "/data/models/DeepSeek-V4-Flash-DSpark",
            "--trust-remote-code",
            "--tensor-parallel-size", "8",
            "--kv-cache-dtype", "fp8",
            "--block-size", "256",
            "--max-model-len", "auto",
            "--max-num-seqs", "256",
            "--tokenizer-mode", "deepseek_v4",
            "--reasoning-parser", "deepseek_v4",
            "--spec-method", "dspark",
            "--spec-model", "/data/models/DeepSeek-V4-Flash-DSpark",
            "--spec-tokens", "3",
            "--no-disable-hybrid-kv-cache-manager",
            "--disable-uvicorn-access-log",
            "--port", "30004",
        ],
        "env": {"CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", "TMPDIR": str(ROOT / "tmp")},
        "bench_backend": "openai",
    },
    {
        "name": "vllm-dspark-dspark-st7",
        "engine": "vllm-dspark",
        "model": "/data/models/DeepSeek-V4-Flash-DSpark",
        "port": 30004,
        "spec_method": "dspark",
        "spec_tokens": 7,
        "cmd": [
            "/data/user1/yy/envs/vllm-dspark/bin/vllm", "serve",
            "/data/models/DeepSeek-V4-Flash-DSpark",
            "--trust-remote-code",
            "--tensor-parallel-size", "8",
            "--kv-cache-dtype", "fp8",
            "--block-size", "256",
            "--max-model-len", "auto",
            "--max-num-seqs", "256",
            "--tokenizer-mode", "deepseek_v4",
            "--reasoning-parser", "deepseek_v4",
            "--spec-method", "dspark",
            "--spec-model", "/data/models/DeepSeek-V4-Flash-DSpark",
            "--spec-tokens", "7",
            "--no-disable-hybrid-kv-cache-manager",
            "--disable-uvicorn-access-log",
            "--port", "30004",
        ],
        "env": {"CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", "TMPDIR": str(ROOT / "tmp")},
        "bench_backend": "openai",
    },
    {
        "name": "vllm-main-nospec",
        "engine": "vllm",
        "model": "/data/models/DeepSeek-V4-Flash",
        "port": 30005,
        "spec_method": None,
        "cmd": [
            "/data/user1/yy/envs/vllm/bin/vllm", "serve",
            "/data/models/DeepSeek-V4-Flash",
            "--trust-remote-code",
            "--tensor-parallel-size", "8",
            "--kv-cache-dtype", "fp8",
            "--block-size", "256",
            "--max-model-len", "auto",
            "--max-num-seqs", "256",
            "--tokenizer-mode", "deepseek_v4",
            "--reasoning-parser", "deepseek_v4",
            "--disable-uvicorn-access-log",
            "--port", "30005",
        ],
        "env": {"CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", "TMPDIR": str(ROOT / "tmp")},
        "bench_backend": "openai",
    },
]

CONCURRENCIES = [1, 16, 64]


def log(msg):
    print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {msg}", flush=True)


def wait_for_health(port, timeout=300):
    url = f"http://{HOST}:{port}/health"
    start = time.time()
    while time.time() - start < timeout:
        try:
            with urllib.request.urlopen(url, timeout=5) as resp:
                if resp.status == 200:
                    return True
        except Exception:
            pass
        time.sleep(2)
    return False


def start_service(service):
    log(f"Starting {service['name']} on port {service['port']}...")
    env = os.environ.copy()
    env.update(service["env"])
    log_file = RESULT_DIR / f"{service['name']}_service.log"
    proc = subprocess.Popen(
        service["cmd"],
        stdout=open(log_file, "w"),
        stderr=subprocess.STDOUT,
        env=env,
    )
    if not wait_for_health(service["port"]):
        log(f"ERROR: {service['name']} failed to start")
        proc.terminate()
        return None
    log(f"{service['name']} is ready")
    return proc


def stop_service(proc, name):
    if proc is None:
        return
    log(f"Stopping {name} (pid {proc.pid})...")
    proc.terminate()
    try:
        proc.wait(timeout=60)
    except subprocess.TimeoutExpired:
        proc.kill()
        proc.wait()
    log(f"{name} stopped")


def run_benchmark(service, concurrency):
    name = service["name"]
    port = service["port"]
    backend = service["bench_backend"]
    result_file = RESULT_DIR / f"{name}_c{concurrency}.json"
    log_file = RESULT_DIR / f"{name}_c{concurrency}.log"

    if service["engine"] == "vllm-dspark":
        bench_cmd = [
            "/data/user1/yy/envs/vllm-dspark/bin/vllm", "bench", "serve",
            "--host", HOST,
            "--port", str(port),
            "--backend", backend,
            "--dataset-name", "sharegpt",
            "--dataset-path", DATASET,
            "--sharegpt-output-len", str(OUTPUT_LEN),
            "--num-prompts", str(NUM_PROMPTS),
            "--max-concurrency", str(concurrency),
            "--endpoint", "/v1/completions",
            "--model", service["model"],
            "--seed", str(SEED),
            "--save-result",
            "--result-dir", str(RESULT_DIR),
            "--result-filename", result_file.name,
        ]
    else:
        bench_cmd = [
            "/data/user1/yy/envs/vllm/bin/vllm", "bench", "serve",
            "--host", HOST,
            "--port", str(port),
            "--backend", backend,
            "--dataset-name", "sharegpt",
            "--dataset-path", DATASET,
            "--sharegpt-output-len", str(OUTPUT_LEN),
            "--num-prompts", str(NUM_PROMPTS),
            "--max-concurrency", str(concurrency),
            "--endpoint", "/v1/completions",
            "--model", service["model"],
            "--seed", str(SEED),
            "--save-result",
            "--result-dir", str(RESULT_DIR),
            "--result-filename", result_file.name,
        ]

    log(f"Running benchmark {name} concurrency={concurrency}...")
    start = time.time()
    with open(log_file, "w") as f:
        proc = subprocess.Popen(bench_cmd, stdout=f, stderr=subprocess.STDOUT)
        proc.wait()
    duration = time.time() - start
    log(f"Benchmark {name} c={concurrency} finished in {duration:.1f}s, exit={proc.returncode}")

    if result_file.exists():
        with open(result_file) as f:
            data = json.load(f)
        return {
            "service": name,
            "engine": service["engine"],
            "spec_method": service.get("spec_method"),
            "spec_tokens": service.get("spec_tokens"),
            "concurrency": concurrency,
            "request_throughput": data.get("request_throughput"),
            "output_throughput": data.get("output_throughput"),
            "total_input_tokens": data.get("total_input_tokens"),
            "total_output_tokens": data.get("total_output_tokens"),
            "duration_s": data.get("duration_s"),
            "result_file": str(result_file),
        }
    else:
        log(f"WARNING: result file {result_file} not found")
        return {
            "service": name,
            "engine": service["engine"],
            "spec_method": service.get("spec_method"),
            "spec_tokens": service.get("spec_tokens"),
            "concurrency": concurrency,
            "error": "result file missing",
            "log_file": str(log_file),
        }


def main():
    summary = []
    for service in SERVICES:
        proc = start_service(service)
        if proc is None:
            continue
        try:
            for concurrency in CONCURRENCIES:
                result = run_benchmark(service, concurrency)
                summary.append(result)
                # Save incremental summary
                with open(RESULT_DIR / "summary.json", "w") as f:
                    json.dump(summary, f, indent=2)
        finally:
            stop_service(proc, service["name"])
        # Small gap between services
        time.sleep(10)

    log(f"All benchmarks complete. Results in {RESULT_DIR}")
    with open(RESULT_DIR / "summary.json", "w") as f:
        json.dump(summary, f, indent=2)


if __name__ == "__main__":
    main()