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, 08:45 UTC
[Bench] Forward --seed to CustomDataset and CustomMMDataset shuffle (#40788)
Signed-off-by: akimtsvigun <akimtsvigun@gmail.com>
Akim Tsvigun committed
May 4, 2026 at 02:39 UTC
894a02500b16fceb77d92110d02c9b5127de0e5f
2 files changed
+83
-2
tests/benchmarks/test_custom_dataset_seed.py
new
+77
@@ -0,0 +1,77 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import argparse
4
+import json
5
+from pathlib import Path
6
+
7
+import pytest
8
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
9
+
10
+from vllm.benchmarks.datasets import get_samples
11
+
12
+
13
+@pytest.fixture(scope="session")
14
+def hf_tokenizer() -> PreTrainedTokenizerBase:
15
+ return AutoTokenizer.from_pretrained("gpt2")
16
+
17
+
18
+def _write_jsonl(path: Path, n_rows: int) -> None:
19
+ with path.open("w") as f:
20
+ for i in range(n_rows):
21
+ f.write(json.dumps({"prompt": f"row {i}: unique prompt content."}) + "\n")
22
+
23
+
24
+def _args_for_custom(dataset_path: str, seed: int) -> argparse.Namespace:
25
+ return argparse.Namespace(
26
+ dataset_name="custom",
27
+ dataset_path=dataset_path,
28
+ disable_shuffle=False,
29
+ num_prompts=30,
30
+ custom_output_len=32,
31
+ skip_chat_template=True,
32
+ no_oversample=False,
33
+ seed=seed,
34
+ request_id_prefix="",
35
+ )
36
+
37
+
38
+@pytest.mark.benchmark
39
+def test_custom_dataset_seed_propagates(
40
+ hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
41
+) -> None:
42
+ """--seed must control the CustomDataset shuffle used by get_samples.
43
+
44
+ Without the fix, CustomDataset was instantiated without random_seed,
45
+ so its load-time shuffle always used DEFAULT_SEED=0 regardless of
46
+ args.seed, causing every run with --dataset-name custom to pick the
47
+ same subset of rows from a larger file.
48
+ """
49
+ jsonl = tmp_path / "data.jsonl"
50
+ _write_jsonl(jsonl, n_rows=60)
51
+
52
+ samples_a = get_samples(_args_for_custom(str(jsonl), seed=0), hf_tokenizer)
53
+ samples_b = get_samples(_args_for_custom(str(jsonl), seed=42), hf_tokenizer)
54
+
55
+ prompts_a = {s.prompt for s in samples_a}
56
+ prompts_b = {s.prompt for s in samples_b}
57
+
58
+ assert len(prompts_a) == 30
59
+ assert len(prompts_b) == 30
60
+ assert prompts_a != prompts_b
61
+
62
+
63
+@pytest.mark.benchmark
64
+def test_custom_dataset_same_seed_is_deterministic(
65
+ hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
66
+) -> None:
67
+ """Same --seed must yield the same CustomDataset subset."""
68
+ jsonl = tmp_path / "data.jsonl"
69
+ _write_jsonl(jsonl, n_rows=60)
70
+
71
+ samples_a = get_samples(_args_for_custom(str(jsonl), seed=7), hf_tokenizer)
72
+ samples_b = get_samples(_args_for_custom(str(jsonl), seed=7), hf_tokenizer)
73
+
74
+ prompts_a = [s.prompt for s in samples_a]
75
+ prompts_b = [s.prompt for s in samples_b]
76
+
77
+ assert prompts_a == prompts_b
vllm/benchmarks/datasets/datasets.py
+6
-2
@@ -1803,7 +1803,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
1803
1804
if args.dataset_name == "custom":
1805
dataset = CustomDataset(
1806
- dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
1806
+ dataset_path=args.dataset_path,
1807
+ disable_shuffle=args.disable_shuffle,
1808
+ random_seed=args.seed,
1809
)
1810
input_requests = dataset.sample(
1811
num_requests=args.num_prompts,
@@ -1816,7 +1818,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
1818
1819
elif args.dataset_name == "custom_mm":
1820
dataset = CustomMMDataset(
1819
- dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
1821
+ dataset_path=args.dataset_path,
1822
+ disable_shuffle=args.disable_shuffle,
1823
+ random_seed=args.seed,
1824
)
1825
input_requests = dataset.sample(
1826
num_requests=args.num_prompts,