#!/usr/bin/env python3
"""Send a small number of warmup requests to a running backend.

Uses sglang.bench_serving with a single request so that the same code path
(prefill / decode kernels, CUDA graphs, etc.) is exercised before the real
benchmark begins. Discards the output.
"""
import argparse
import subprocess
import sys
import tempfile
from pathlib import Path


def run_warmup(backend: str, host: str, port: int, input_len: int, output_len: int, num: int, env_python: Path) -> None:
    with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=True) as tmp:
        cmd = [
            str(env_python),
            "-m",
            "sglang.bench_serving",
            "--backend",
            backend,
            "--host",
            host,
            "--port",
            str(port),
            "--dataset-name",
            "random",
            "--random-input-len",
            str(input_len),
            "--random-output-len",
            str(output_len),
            "--num-prompts",
            str(num),
            "--max-concurrency",
            "1",
            "--request-rate",
            "10000",
            "--output-file",
            tmp.name,
            "--output-details",
        ]
        print(f"[warmup] {' '.join(cmd)}", flush=True)
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            print("[warmup] FAILED", file=sys.stderr)
            print(result.stdout, file=sys.stderr)
            print(result.stderr, file=sys.stderr)
            sys.exit(1)
        print(f"[warmup] OK: backend={backend} port={port} input={input_len} output={output_len} num={num}")


def main():
    parser = argparse.ArgumentParser(description="Warmup a serving backend.")
    parser.add_argument("--backend", required=True, choices=["sglang", "vllm"])
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument("--input-len", type=int, required=True)
    parser.add_argument("--output-len", type=int, required=True)
    parser.add_argument("--num", type=int, default=1)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--env-python", default="/data/user1/yy/envs/sglang/bin/python")
    args = parser.parse_args()

    run_warmup(
        backend=args.backend,
        host=args.host,
        port=args.port,
        input_len=args.input_len,
        output_len=args.output_len,
        num=args.num,
        env_python=Path(args.env_python),
    )


if __name__ == "__main__":
    main()