[Rust][Benchmark] Make `vllm bench serve` Rust delegation opt-in (#50081)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
Jul 29, 2026 at 00:21 UTC
01661cc57f48ce95c639efce7c88e6dd37349007
8 files changed
+61
-76
tests/entrypoints/openai/test_dp_supervisor.py
+1
@@ -201,6 +201,7 @@ def test_run_vllm_dp_server_uses_rust_frontend_when_enabled(monkeypatch):
201
monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
202
monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
203
monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
204
+ monkeypatch.setattr(dp_sup.envs, "VLLM_USE_RUST_FRONTEND", True, raising=False)
205
monkeypatch.setattr(
206
dp_sup.envs,
207
"VLLM_RUST_FRONTEND_PATH",
tests/test_envs.py
+9
@@ -145,6 +145,15 @@ def test_precompiled_install_flags_are_orthogonal() -> None:
145
assert environment_variables["VLLM_USE_PRECOMPILED_RUST"]() is True
146
147
148
+def test_rust_bench_auto_path_missing_fails_fast() -> None:
149
+ with (
150
+ patch.dict(os.environ, {"VLLM_USE_RUST_BENCH": "1"}, clear=True),
151
+ patch("vllm.envs.os.path.isfile", return_value=False),
152
+ pytest.raises(FileNotFoundError, match="vllm-rs binary was not found"),
153
+ ):
154
+ environment_variables["VLLM_RUST_FRONTEND_PATH"]()
155
+
156
+
157
class TestEnvWithChoices:
158
"""Test cases for env_with_choices function."""
159
vllm/entrypoints/cli/benchmark/main.py
+20
@@ -2,18 +2,38 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import argparse
5
+import os
6
import sys
7
import typing
8
9
+from vllm import envs
10
from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
11
from vllm.entrypoints.cli.types import CLISubcommand
12
from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
13
+from vllm.logger import init_logger
14
15
if typing.TYPE_CHECKING:
16
from vllm.utils.argparse_utils import FlexibleArgumentParser
17
else:
18
FlexibleArgumentParser = argparse.ArgumentParser
19
20
+logger = init_logger(__name__)
21
+
22
+
23
+def maybe_exec_rust_bench() -> None:
24
+ if sys.argv[1:3] != ["bench", "serve"] or not envs.VLLM_USE_RUST_BENCH:
25
+ return
26
+
27
+ rust_cli = envs.VLLM_RUST_FRONTEND_PATH
28
+ if rust_cli is None:
29
+ raise RuntimeError(
30
+ "VLLM_USE_RUST_BENCH=1 requires VLLM_RUST_FRONTEND_PATH "
31
+ "to resolve to the vllm-rs binary."
32
+ )
33
+
34
+ logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli)
35
+ os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]])
36
+
37
38
def _import_bench_subcommand_modules() -> None:
39
# Imported lazily so `BenchmarkSubcommandBase` subclasses register only
vllm/entrypoints/cli/benchmark/serve.py
+2
-60
@@ -1,68 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import argparse
4
-import os
5
-import sys
6
-from pathlib import Path
4
8
-from vllm.benchmarks.serve import add_cli_args
9
-from vllm.benchmarks.serve import main as python_main
5
+from vllm.benchmarks.serve import add_cli_args, main
6
from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
11
-from vllm.logger import init_logger
7
from vllm.utils.argparse_utils import FlexibleArgumentParser
8
14
-logger = init_logger(__name__)
15
-_RUST_CLI_PATH = Path(__file__).resolve().parents[3] / "vllm-rs"
16
-_RUST_SUPPORTED_DATASETS = frozenset(
17
- {
18
- "custom",
19
- "hf",
20
- "prefix_repetition",
21
- "random",
22
- "random-mm",
23
- "random-rerank",
24
- "sharegpt",
25
- "sonnet",
26
- "speed_bench",
27
- }
28
-)
29
-_RUST_SUPPORTED_BACKENDS = frozenset(
30
- {
31
- "openai",
32
- "openai-chat",
33
- "openai-embeddings",
34
- "openai-embeddings-chat",
35
- "vllm",
36
- "vllm-pooling",
37
- "vllm-rerank",
38
- }
39
-)
40
-
41
-
42
-def _rust_unsupported_reason(args: argparse.Namespace) -> str | None:
43
- if args.dataset_name not in _RUST_SUPPORTED_DATASETS:
44
- return f"dataset {args.dataset_name!r} is not supported by the Rust benchmark"
45
- if args.backend not in _RUST_SUPPORTED_BACKENDS:
46
- return f"backend {args.backend!r} is not supported by the Rust benchmark"
47
- return None
48
-
49
-
50
-def _maybe_exec_rust_bench(args: argparse.Namespace) -> None:
51
- if reason := _rust_unsupported_reason(args):
52
- logger.info("Using Python benchmark: %s.", reason)
53
- return
54
-
55
- if not _RUST_CLI_PATH.is_file():
56
- logger.warning(
57
- "Rust benchmark binary not found at %s; falling back to Python.",
58
- _RUST_CLI_PATH,
59
- )
60
- return
61
-
62
- rust_cli = str(_RUST_CLI_PATH)
63
- logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli)
64
- os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]])
65
-
9
10
class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
11
"""The `serve` subcommand for `vllm bench`."""
@@ -76,5 +19,4 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
19
20
@staticmethod
21
def cmd(args: argparse.Namespace) -> None:
79
- _maybe_exec_rust_bench(args)
80
- python_main(args)
22
+ main(args)
vllm/entrypoints/cli/main.py
+2
@@ -54,6 +54,8 @@ def main():
54
logger.info("Delegating entrypoint handling to vllm-omni")
55
omni_main()
56
else:
57
+ vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench()
58
+
59
# For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default
60
if len(sys.argv) > 1 and sys.argv[1] == "bench":
61
logger.debug(
vllm/entrypoints/cli/serve.py
+10
-4
@@ -58,6 +58,10 @@ class ServeSubcommand(CLISubcommand):
58
uvloop.run(serve_grpc(args))
59
return
60
61
+ rust_frontend_path = (
62
+ envs.VLLM_RUST_FRONTEND_PATH if envs.VLLM_USE_RUST_FRONTEND else None
63
+ )
64
+
65
if args.headless:
66
if args.api_server_count is not None and args.api_server_count > 0:
67
raise ValueError(
@@ -103,7 +107,7 @@ class ServeSubcommand(CLISubcommand):
107
# - Hybrid LB: Use local DP size (internal LB for local ranks only)
108
# - Internal LB: Use full DP size
109
if args.api_server_count is None:
106
- if is_multi_port or is_external_lb or envs.VLLM_RUST_FRONTEND_PATH:
110
+ if is_multi_port or is_external_lb or rust_frontend_path:
111
args.api_server_count = 1
112
elif is_hybrid_lb:
113
args.api_server_count = args.data_parallel_size_local or 1
@@ -120,7 +124,7 @@ class ServeSubcommand(CLISubcommand):
124
"Defaulting api_server_count to data_parallel_size (%d).",
125
args.api_server_count,
126
)
123
- elif envs.VLLM_RUST_FRONTEND_PATH and args.api_server_count > 1:
127
+ elif rust_frontend_path and args.api_server_count > 1:
128
logger.warning(
129
"Ignoring --api-server-count=%d when using rust front-end process",
130
args.api_server_count,
@@ -140,7 +144,7 @@ class ServeSubcommand(CLISubcommand):
144
run_dp_supervisor(args)
145
elif args.api_server_count < 1:
146
run_headless(args)
143
- elif args.api_server_count > 1 or envs.VLLM_RUST_FRONTEND_PATH:
147
+ elif args.api_server_count > 1 or rust_frontend_path:
148
run_multi_api_server(args)
149
else:
150
# Single API server (this process).
@@ -256,7 +260,9 @@ def run_headless(args: argparse.Namespace):
260
261
def run_multi_api_server(args: argparse.Namespace):
262
assert not args.headless
259
- rust_frontend_path = envs.VLLM_RUST_FRONTEND_PATH
263
+ rust_frontend_path = (
264
+ envs.VLLM_RUST_FRONTEND_PATH if envs.VLLM_USE_RUST_FRONTEND else None
265
+ )
266
num_api_servers: int = args.api_server_count
267
assert num_api_servers > 0
268
vllm/entrypoints/openai/dp_supervisor.py
+1
-1
@@ -257,7 +257,7 @@ def _run_vllm_dp_server(child_args: argparse.Namespace) -> None:
257
name = f"APIServer_DP{child_args.data_parallel_rank}"
258
set_process_title(name)
259
decorate_logs(name)
260
- if envs.VLLM_RUST_FRONTEND_PATH:
260
+ if envs.VLLM_USE_RUST_FRONTEND and envs.VLLM_RUST_FRONTEND_PATH:
261
_run_rust_vllm_dp_server(child_args)
262
else:
263
_run_python_vllm_dp_server(child_args)
vllm/envs.py
+16
-11
@@ -153,6 +153,7 @@ if TYPE_CHECKING:
153
K_SCALE_CONSTANT: int = 200
154
V_SCALE_CONSTANT: int = 100
155
VLLM_USE_RUST_FRONTEND: bool = False
156
+ VLLM_USE_RUST_BENCH: bool = False
157
VLLM_RUST_FRONTEND_PATH: str | None = "auto"
158
VLLM_SERVER_DEV_MODE: bool = False
159
VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128
@@ -548,22 +549,24 @@ def _deprecated_triton_attn_use_td() -> None:
549
return None
550
551
551
-def _resolve_rust_frontend_path() -> str | None:
552
- """Resolve the Rust frontend binary path.
552
+def _resolve_rust_cli_path() -> str | None:
553
+ """Resolve the vllm-rs binary path.
554
554
- Returns None if VLLM_USE_RUST_FRONTEND is not enabled.
555
+ Returns None unless VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH is enabled.
556
When enabled, resolves VLLM_RUST_FRONTEND_PATH ("auto" by default)
557
to the actual binary path.
558
"""
558
- use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0")))
559
+ use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0"))) or bool(
560
+ int(os.environ.get("VLLM_USE_RUST_BENCH", "0"))
561
+ )
562
raw = os.environ.get("VLLM_RUST_FRONTEND_PATH", "auto")
563
564
if not use_rust:
565
if os.environ.get("VLLM_RUST_FRONTEND_PATH") is not None:
566
logger.warning(
564
- "VLLM_RUST_FRONTEND_PATH is set but VLLM_USE_RUST_FRONTEND "
565
- "is not enabled. The Rust frontend will not be used. "
566
- "Set VLLM_USE_RUST_FRONTEND=1 to enable it."
567
+ "VLLM_RUST_FRONTEND_PATH is set without enabling "
568
+ "VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH. "
569
+ "Set one of them to 1 to use the vllm-rs binary."
570
)
571
return None
572
@@ -1340,10 +1343,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1343
"VLLM_USE_RUST_FRONTEND": lambda: bool(
1344
int(os.getenv("VLLM_USE_RUST_FRONTEND", "0"))
1345
),
1343
- # Path to the Rust frontend binary. Defaults to "auto" which discovers
1344
- # the binary installed with the vllm package. Only used when
1345
- # VLLM_USE_RUST_FRONTEND=1.
1346
- "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_frontend_path(),
1346
+ # If set, use the packaged Rust client for `vllm bench serve`.
1347
+ "VLLM_USE_RUST_BENCH": lambda: bool(int(os.getenv("VLLM_USE_RUST_BENCH", "0"))),
1348
+ # Path to the vllm-rs binary. Defaults to "auto" which discovers the
1349
+ # binary installed with the vllm package. Used when VLLM_USE_RUST_FRONTEND=1
1350
+ # or VLLM_USE_RUST_BENCH=1.
1351
+ "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_cli_path(),
1352
# If set, vllm will run in development mode, which will enable
1353
# some additional endpoints for developing and debugging,
1354
# e.g. `/reset_prefix_cache`