77 lines
2.5 KiB
Python
Executable File
77 lines
2.5 KiB
Python
Executable File
#!/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()
|