A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 08:45 UTC
[UX] Make `vllm bench serve` discover model by default and use --input-len (#30816)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Dec 17, 2025 at 04:55 UTC
519ef9a91111d2d6f8545c8a6b2c1a28d87309fa
2 files changed
+79
-13
tests/benchmarks/test_serve_cli.py
+3
-6
@@ -19,21 +19,18 @@ def server():
19
20
@pytest.mark.benchmark
21
def test_bench_serve(server):
22
+ # Test default model detection and input/output len
23
command = [
24
"vllm",
25
"bench",
26
"serve",
26
- "--model",
27
- MODEL_NAME,
27
"--host",
28
server.host,
29
"--port",
30
str(server.port),
32
- "--dataset-name",
33
- "random",
34
- "--random-input-len",
31
+ "--input-len",
32
"32",
36
- "--random-output-len",
33
+ "--output-len",
34
"4",
35
"--num-prompts",
36
"5",
vllm/benchmarks/serve.py
+76
-7
@@ -10,8 +10,10 @@ On the client side, run:
10
vllm bench serve \
11
--backend <backend or endpoint type. Default 'openai'> \
12
--label <benchmark result label. Default using backend> \
13
- --model <your_model> \
13
+ --model <your_model. Optional, defaults to first model from server> \
14
--dataset-name <dataset_name. Default 'random'> \
15
+ --input-len <general input length. Optional, maps to dataset-specific args> \
16
+ --output-len <general output length. Optional, maps to dataset-specific args> \
17
--request-rate <request_rate. Default inf> \
18
--num-prompts <num_prompts. Default 1000>
19
"""
@@ -57,6 +59,33 @@ TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) a
59
)
60
61
62
+async def get_first_model_from_server(
63
+ base_url: str, headers: dict | None = None
64
+) -> str:
65
+ """Fetch the first model from the server's /v1/models endpoint."""
66
+ models_url = f"{base_url}/v1/models"
67
+ async with aiohttp.ClientSession() as session:
68
+ try:
69
+ async with session.get(models_url, headers=headers) as response:
70
+ response.raise_for_status()
71
+ data = await response.json()
72
+ if "data" in data and len(data["data"]) > 0:
73
+ return data["data"][0]["id"]
74
+ else:
75
+ raise ValueError(
76
+ f"No models found on the server at {base_url}. "
77
+ "Make sure the server is running and has models loaded."
78
+ )
79
+ except (aiohttp.ClientError, json.JSONDecodeError) as e:
80
+ raise RuntimeError(
81
+ f"Failed to fetch models from server at {models_url}. "
82
+ "Check that:\n"
83
+ "1. The server is running\n"
84
+ "2. The server URL is correct\n"
85
+ f"Error: {e}"
86
+ ) from e
87
+
88
+
89
class TaskType(Enum):
90
GENERATION = "generation"
91
POOLING = "pooling"
@@ -1025,8 +1054,26 @@ def add_cli_args(parser: argparse.ArgumentParser):
1054
parser.add_argument(
1055
"--model",
1056
type=str,
1028
- required=True,
1029
- help="Name of the model.",
1057
+ required=False,
1058
+ default=None,
1059
+ help="Name of the model. If not specified, will fetch the first model "
1060
+ "from the server's /v1/models endpoint.",
1061
+ )
1062
+ parser.add_argument(
1063
+ "--input-len",
1064
+ type=int,
1065
+ default=None,
1066
+ help="General input length for datasets. Maps to dataset-specific "
1067
+ "input length arguments (e.g., --random-input-len, --sonnet-input-len). "
1068
+ "If not specified, uses dataset defaults.",
1069
+ )
1070
+ parser.add_argument(
1071
+ "--output-len",
1072
+ type=int,
1073
+ default=None,
1074
+ help="General output length for datasets. Maps to dataset-specific "
1075
+ "output length arguments (e.g., --random-output-len, --sonnet-output-len). "
1076
+ "If not specified, uses dataset defaults.",
1077
)
1078
parser.add_argument(
1079
"--tokenizer",
@@ -1332,10 +1379,6 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1379
raise ValueError("For exponential ramp-up, the start RPS cannot be 0.")
1380
1381
label = args.label
1335
- model_id = args.model
1336
- model_name = args.served_model_name
1337
- tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
1338
- tokenizer_mode = args.tokenizer_mode
1382
1383
if args.base_url is not None:
1384
api_url = f"{args.base_url}{args.endpoint}"
@@ -1356,6 +1399,18 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1399
else:
1400
raise ValueError("Invalid header format. Please use KEY=VALUE format.")
1401
1402
+ # Fetch model from server if not specified
1403
+ if args.model is None:
1404
+ print("Model not specified, fetching first model from server...")
1405
+ model_id = await get_first_model_from_server(base_url, headers)
1406
+ print(f"Using model: {model_id}")
1407
+ else:
1408
+ model_id = args.model
1409
+
1410
+ model_name = args.served_model_name
1411
+ tokenizer_id = args.tokenizer if args.tokenizer is not None else model_id
1412
+ tokenizer_mode = args.tokenizer_mode
1413
+
1414
tokenizer = get_tokenizer(
1415
tokenizer_id,
1416
tokenizer_mode=tokenizer_mode,
@@ -1368,6 +1423,20 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1423
"'--dataset-path' if required."
1424
)
1425
1426
+ # Map general --input-len and --output-len to all dataset-specific arguments
1427
+ if args.input_len is not None:
1428
+ args.random_input_len = args.input_len
1429
+ args.sonnet_input_len = args.input_len
1430
+
1431
+ if args.output_len is not None:
1432
+ args.random_output_len = args.output_len
1433
+ args.sonnet_output_len = args.output_len
1434
+ args.sharegpt_output_len = args.output_len
1435
+ args.custom_output_len = args.output_len
1436
+ args.hf_output_len = args.output_len
1437
+ args.spec_bench_output_len = args.output_len
1438
+ args.prefix_repetition_output_len = args.output_len
1439
+
1440
# when using random datasets, default to ignoring EOS
1441
# so generation runs to the requested length
1442
if (