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 4, 21:15 UTC
[Bugfix] Fix dataset name and path argument validation bug in vllm bench serve (#40288)
Signed-off-by: talora <talora@nvidia.com> Signed-off-by: Talor Abramovich <talor19@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Talor Abramovich committed
Apr 21, 2026 at 16:14 UTC
7b1e0b07d0ea9fe11f0c4a35001baade2c3b1074
2 files changed
+34
-26
vllm/benchmarks/datasets/datasets.py
+21
-26
@@ -1373,26 +1373,6 @@ class ShareGPTDataset(BenchmarkDataset):
1373
return samples
1374
1375
1376
-class _ValidateDatasetArgs(argparse.Action):
1377
- """Argparse action to validate dataset name and path compatibility."""
1378
-
1379
- def __call__(self, parser, namespace, values, option_string=None):
1380
- setattr(namespace, self.dest, values)
1381
-
1382
- # Get current values of both dataset_name and dataset_path
1383
- dataset_name = getattr(namespace, "dataset_name", "random")
1384
- dataset_path = getattr(namespace, "dataset_path", None)
1385
-
1386
- # Validate the combination
1387
- if dataset_name == "random" and dataset_path is not None:
1388
- parser.error(
1389
- "Cannot use 'random' dataset with --dataset-path. "
1390
- "Please specify the appropriate --dataset-name (e.g., "
1391
- "'sharegpt', 'custom', 'sonnet') for your dataset file: "
1392
- f"{dataset_path}"
1393
- )
1394
-
1395
-
1376
def add_dataset_parser(parser: FlexibleArgumentParser):
1377
parser.add_argument(
1378
"--trust-remote-code",
@@ -1410,7 +1390,6 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1390
"--dataset-name",
1391
type=str,
1392
default="random",
1413
- action=_ValidateDatasetArgs,
1393
choices=[
1394
"sharegpt",
1395
"burstgpt",
@@ -1436,7 +1415,6 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1415
"--dataset-path",
1416
type=str,
1417
default=None,
1439
- action=_ValidateDatasetArgs,
1418
help="Path to the sharegpt/sonnet dataset or the HF dataset ID if "
1419
"using HF dataset.",
1420
)
@@ -1608,7 +1586,9 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1586
"repetition dataset.",
1587
)
1588
1611
- speed_bench_group = parser.add_argument_group("speed bench dataset options")
1589
+ speed_bench_group = parser.add_argument_group(
1590
+ "speed bench dataset options", description=SpeedBench.__doc__
1591
+ )
1592
speed_bench_group.add_argument(
1593
"--speed-bench-dataset-subset",
1594
type=str,
@@ -3603,11 +3583,15 @@ class MMStarDataset(HuggingFaceDataset):
3583
3584
class SpeedBench(CustomDataset):
3585
"""
3606
- Implements the SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench
3586
+ SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench
3587
+
3588
Download the dataset using:
3608
- curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 -
3589
+
3590
+ `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 -`
3591
""" # noqa: E501
3592
3593
+ DOWNLOAD_SCRIPT_URL = "https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py"
3594
+
3595
def __init__(self, **kwargs) -> None:
3596
self.dataset_subset = kwargs.pop("dataset_subset", "qualitative")
3597
self.category = kwargs.pop("category", None)
@@ -3618,6 +3602,13 @@ class SpeedBench(CustomDataset):
3602
if self.dataset_path is None:
3603
raise ValueError("dataset_path must be provided for loading data.")
3604
3605
+ if not Path(self.dataset_path).is_dir():
3606
+ raise ValueError(
3607
+ f"dataset_path {self.dataset_path} is not a directory. "
3608
+ f"Please make sure to download the dataset from HuggingFace using "
3609
+ f"`curl -LsSf {self.DOWNLOAD_SCRIPT_URL} | python3 -`"
3610
+ )
3611
+
3612
self.data = []
3613
3614
# Load the JSONL file
@@ -3628,7 +3619,11 @@ class SpeedBench(CustomDataset):
3619
3620
# check if the JSONL file has a 'turns' column
3621
if "messages" not in jsonl_data.columns:
3631
- raise ValueError("JSONL file must contain a 'messages' column.")
3622
+ raise ValueError(
3623
+ "JSONL file must contain a 'messages' column. "
3624
+ "Please make sure to download the dataset from HuggingFace using "
3625
+ f"`curl -LsSf {self.DOWNLOAD_SCRIPT_URL} | python3 -`"
3626
+ )
3627
3628
for _, row in jsonl_data.iterrows():
3629
# sample only from a specific category if specified
vllm/benchmarks/serve.py
+13
@@ -1711,12 +1711,25 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1711
trust_remote_code=args.trust_remote_code,
1712
)
1713
1714
+ # Validate dataset name/path
1715
if args.dataset_name is None:
1716
raise ValueError(
1717
"Please specify '--dataset-name' and the corresponding "
1718
"'--dataset-path' if required."
1719
)
1720
1721
+ if (
1722
+ args.dataset_name
1723
+ in ["random", "random-mm", "random-rerank", "prefix_repetition"]
1724
+ and args.dataset_path is not None
1725
+ ):
1726
+ raise ValueError(
1727
+ f"Cannot use '{args.dataset_name}' dataset with --dataset-path. "
1728
+ "Please specify the appropriate --dataset-name (e.g., "
1729
+ "'sharegpt', 'custom', 'sonnet') for your dataset file: "
1730
+ f"{args.dataset_path}"
1731
+ )
1732
+
1733
# Map general --input-len and --output-len to all dataset-specific arguments
1734
if args.input_len is not None:
1735
args.random_input_len = args.input_len