main
py 61 lines 2.28 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Keep the `vllm bench serve` flag snapshot in sync for the Rust CLI parity test.
4
5 `rust/src/bench/tests/cli_parity.rs` asserts that every flag in the snapshot is
6 either a known `vllm-bench` clap flag/alias or explicitly allowlisted as
7 Python-only, so `VLLM_USE_RUST_BENCH=1` delegation cannot silently reject
8 documented flags. This test keeps the snapshot itself current.
9
10 Run this file directly to regenerate the snapshot:
11
12 .venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py
13 """
14
15 from pathlib import Path
16
17 REPO_ROOT = Path(__file__).resolve().parents[2]
18 SNAPSHOT = REPO_ROOT / "rust" / "src" / "bench" / "tests" / "python_serve_flags.txt"
19 HEADER = """\
20 # Long CLI flags of the Python `vllm bench serve` parser
21 # (vllm.benchmarks.serve.add_cli_args), sorted; --help excluded.
22 # Regenerate: .venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py
23 # Consumed by rust/src/bench/tests/cli_parity.rs, which requires every flag
24 # below to be a known vllm-bench flag/alias or a PYTHON_ONLY allowlist entry.
25 """
26
27
28 def _current_flags() -> list[str]:
29 from vllm.benchmarks.serve import add_cli_args
30 from vllm.utils.argparse_utils import FlexibleArgumentParser
31
32 parser = FlexibleArgumentParser()
33 add_cli_args(parser)
34 return sorted(
35 {
36 opt
37 for action in parser._actions
38 for opt in action.option_strings
39 if opt.startswith("--") and opt != "--help"
40 }
41 )
42
43
44 def _snapshot_flags() -> list[str]:
45 lines = SNAPSHOT.read_text().splitlines()
46 return [line.strip() for line in lines if line.strip() and not line.startswith("#")]
47
48
49 def test_serve_flag_snapshot_is_current():
50 assert SNAPSHOT.is_file(), f"missing snapshot {SNAPSHOT}"
51 assert _current_flags() == _snapshot_flags(), (
52 "`vllm bench serve` flags changed. Regenerate the snapshot with "
53 "`.venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py`, "
54 "then for any added flag either support it in rust/src/bench/src/cli.rs "
55 "or add it to PYTHON_ONLY in rust/src/bench/tests/cli_parity.rs"
56 )
57
58
59 if __name__ == "__main__":
60 SNAPSHOT.write_text(HEADER + "\n".join(_current_flags()) + "\n")
61 print(f"wrote {SNAPSHOT}")