[Bugfix][Benchmarks] Restore --skip-tokenizer-init with custom dataset (#49180)
Signed-off-by: Michele Gazzetti <michele.gazzetti1@ibm.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
mgazz committed
Jul 24, 2026 at 12:30 UTC
a454a1dd25a181334f6f126b2a2bec3d7c1d9aab
3 files changed
+185
-9
tests/benchmarks/test_skip_tokenizer_init.py
new
+162
@@ -0,0 +1,162 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Regression test for --skip-tokenizer-init with --dataset-name custom.
4
+
5
+Before the fix (introduced by #39896), running:
6
+
7
+ vllm bench serve \
8
+ --backend vllm-pooling \
9
+ --dataset-name custom \
10
+ --dataset-path <path> \
11
+ --model ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11 \
12
+ --endpoint /pooling \
13
+ --skip-tokenizer-init \
14
+ ...
15
+
16
+raised immediately with:
17
+
18
+ AssertionError: Tokenizer must be initialized before loading dataset
19
+
20
+even though CustomDataset.sample() already handles tokenizer=None.
21
+This test exercises main_async() directly so it catches any regression
22
+re-introduced at the serve.py level, not just inside get_samples().
23
+"""
24
+
25
+import argparse
26
+import asyncio
27
+import json
28
+from pathlib import Path
29
+from unittest.mock import AsyncMock, patch
30
+
31
+import pytest
32
+
33
+import vllm.benchmarks.serve as serve_module
34
+
35
+# Exact prompt payload from the failing benchmark run against a
36
+# Prithvi-EO-2.0 pooling endpoint (URL-in / base64-out format).
37
+_PRITHVI_PROMPT = {
38
+ "data": {
39
+ "data": "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/India_900498_S2Hand.tif",
40
+ "data_format": "url",
41
+ "out_data_format": "b64_json",
42
+ "indices": [1, 2, 3, 8, 11, 12],
43
+ },
44
+ "priority": 0,
45
+ "softmax": False,
46
+}
47
+
48
+
49
+def _write_dataset(path: Path) -> None:
50
+ path.write_text(json.dumps({"prompt": _PRITHVI_PROMPT}) + "\n")
51
+
52
+
53
+def _args(dataset_path: str) -> argparse.Namespace:
54
+ """Reproduce the argparse.Namespace that serve.py builds from the
55
+ failing command, including skip_tokenizer_init=True."""
56
+ return argparse.Namespace(
57
+ # dataset
58
+ dataset_name="custom",
59
+ dataset_path=dataset_path,
60
+ disable_shuffle=False,
61
+ num_prompts=1,
62
+ custom_output_len=256,
63
+ skip_chat_template=True,
64
+ chat_template_kwargs=None,
65
+ no_oversample=False,
66
+ seed=0,
67
+ request_id_prefix="bench-",
68
+ # model / tokenizer
69
+ model="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
70
+ served_model_name=None,
71
+ tokenizer=None,
72
+ tokenizer_mode="auto",
73
+ trust_remote_code=False,
74
+ skip_tokenizer_init=True, # <-- the flag under test
75
+ # backend / endpoint
76
+ backend="vllm-pooling",
77
+ base_url="http://127.0.0.1:8000",
78
+ host="127.0.0.1",
79
+ port=8000,
80
+ endpoint="/pooling",
81
+ header=None,
82
+ insecure=False,
83
+ # traffic
84
+ request_rate=16.0,
85
+ burstiness=1.0,
86
+ max_concurrency=None,
87
+ # misc serve args that main_async reads before reaching get_samples
88
+ plot_timeline=False,
89
+ plot_dataset_stats=False,
90
+ self_timed=None,
91
+ metadata=None,
92
+ label=None,
93
+ logprobs=None,
94
+ use_beam_search=False,
95
+ ignore_eos=False,
96
+ goodput=None,
97
+ percentile_metrics="ttft,tpot,itl,e2el",
98
+ metric_percentiles="25,50,75,99",
99
+ save_result=False,
100
+ append_result=False,
101
+ result_dir=".",
102
+ result_filename=None,
103
+ num_warmups=0,
104
+ profile=False,
105
+ disable_tqdm=True,
106
+ lora_modules=None,
107
+ lora_assignment="random",
108
+ ramp_up_strategy=None,
109
+ ramp_up_start_rps=None,
110
+ ramp_up_end_rps=None,
111
+ ready_check_timeout_sec=0,
112
+ extra_body=None,
113
+ top_p=None,
114
+ top_k=None,
115
+ min_p=None,
116
+ temperature=None,
117
+ frequency_penalty=None,
118
+ presence_penalty=None,
119
+ repetition_penalty=None,
120
+ save_detailed=False,
121
+ input_len=None,
122
+ output_len=None,
123
+ )
124
+
125
+
126
+@pytest.mark.benchmark
127
+def test_main_async_skip_tokenizer_init_does_not_raise(tmp_path: Path) -> None:
128
+ """main_async must not raise AssertionError when skip_tokenizer_init=True
129
+ and dataset_name='custom'.
130
+
131
+ On main (before the fix) this test fails with:
132
+ AssertionError: Tokenizer must be initialized before loading dataset
133
+ """
134
+ dataset_path = tmp_path / "dataset_url_input_india.jsonl"
135
+ _write_dataset(dataset_path)
136
+
137
+ args = _args(str(dataset_path))
138
+
139
+ # Patch benchmark() so we never make real HTTP requests — the regression
140
+ # triggers before benchmark() is ever called, so this just keeps the test
141
+ # fast and hermetic.
142
+ mock_result = {
143
+ "completed": 1,
144
+ "failed": 0,
145
+ "total_input_tokens": 1,
146
+ "total_output_tokens": 1,
147
+ "request_throughput": 1.0,
148
+ "output_throughput": 1.0,
149
+ "total_token_throughput": 1.0,
150
+ "input_lens": [],
151
+ "output_lens": [],
152
+ "ttfts": [],
153
+ "itls": [],
154
+ "generated_texts": [],
155
+ "errors": [],
156
+ "duration": 1.0,
157
+ }
158
+ with patch.object(
159
+ serve_module, "benchmark", new=AsyncMock(return_value=mock_result)
160
+ ):
161
+ # Must NOT raise AssertionError
162
+ asyncio.run(serve_module.main_async(args))
vllm/benchmarks/datasets/datasets.py
+21
-6
@@ -2091,7 +2091,7 @@ def _parse_range_ratio(value: str) -> RangeRatio:
2091
return json.loads(value)
2092
2093
2094
-def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2094
+def get_samples(args, tokenizer: TokenizerLike | None) -> list[SampleRequest]:
2095
if not hasattr(args, "request_id_prefix"):
2096
args.request_id_prefix = ""
2097
@@ -2148,6 +2148,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2148
)
2149
2150
elif args.dataset_name == "sonnet":
2151
+ assert tokenizer is not None, (
2152
+ "Tokenizer must be initialized for the 'sonnet' dataset."
2153
+ )
2154
sonnet_dataset = SonnetDataset(
2155
dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
2156
)
@@ -2317,6 +2320,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2320
"Multi-modal content is only supported on 'openai-chat' and "
2321
"'openai-audio' backends."
2322
)
2323
+ assert tokenizer is not None, (
2324
+ "Tokenizer must be initialized for the 'hf' dataset."
2325
+ )
2326
input_requests = dataset_class(
2327
dataset_path=args.dataset_path,
2328
dataset_subset=args.hf_subset,
@@ -2338,6 +2344,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2344
)
2345
2346
elif args.dataset_name == "timed_trace":
2347
+ assert tokenizer is not None, (
2348
+ "Tokenizer must be initialized for the 'timed_trace' dataset."
2349
+ )
2350
dataloader = TimedTrace(**vars(args))
2351
input_requests = dataloader.sample(
2352
num_requests=args.num_prompts,
@@ -2347,6 +2356,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2356
2357
else:
2358
# For datasets that follow a similar structure, use a mapping.
2359
+ assert tokenizer is not None, (
2360
+ f"Tokenizer must be initialized for the '{args.dataset_name}' dataset."
2361
+ )
2362
dataset_mapping = {
2363
"spec_bench": lambda: SpecBench(
2364
dataset_path=args.dataset_path,
@@ -2532,7 +2544,7 @@ class CustomDataset(BenchmarkDataset):
2544
2545
def sample(
2546
self,
2535
- tokenizer: TokenizerLike,
2547
+ tokenizer: TokenizerLike | None,
2548
num_requests: int,
2549
request_id_prefix: str = "",
2550
no_oversample: bool = False,
@@ -2809,7 +2821,7 @@ class CustomImageDataset(CustomDataset):
2821
2822
def sample(
2823
self,
2812
- tokenizer: TokenizerLike,
2824
+ tokenizer: TokenizerLike | None,
2825
num_requests: int,
2826
request_id_prefix: str = "",
2827
no_oversample: bool = False,
@@ -2907,7 +2919,7 @@ class CustomAudioDataset(CustomDataset):
2919
2920
def sample(
2921
self,
2910
- tokenizer: TokenizerLike,
2922
+ tokenizer: TokenizerLike | None,
2923
num_requests: int,
2924
request_id_prefix: str = "",
2925
no_oversample: bool = False,
@@ -2929,7 +2941,9 @@ class CustomAudioDataset(CustomDataset):
2941
prompt = item.get("prompt", "")
2942
if tokenizer is None:
2943
prompt_len = 1
2932
- new_output_len = output_len if output_len not in (None, -1) else 256
2944
+ new_output_len = (
2945
+ output_len if (output_len is not None and output_len != -1) else 256
2946
+ )
2947
mm_content = None
2948
else:
2949
use_chat_template = (
@@ -2979,6 +2993,7 @@ class CustomAudioDataset(CustomDataset):
2993
)
2994
new_output_len = int(item["output_tokens"])
2995
else:
2996
+ assert output_len is not None
2997
new_output_len = output_len
2998
sampled_requests.append(
2999
SampleRequest(
@@ -3037,7 +3052,7 @@ class SpecBench(CustomDataset):
3052
3053
def sample(
3054
self,
3040
- tokenizer: TokenizerLike,
3055
+ tokenizer: TokenizerLike | None,
3056
num_requests: int,
3057
request_id_prefix: str = "",
3058
no_oversample: bool = False,
vllm/benchmarks/serve.py
+2
-3
@@ -557,7 +557,7 @@ def calculate_metrics(
557
input_requests: list[SampleRequest],
558
outputs: list[RequestFuncOutput],
559
dur_s: float,
560
- tokenizer: TokenizerLike,
560
+ tokenizer: TokenizerLike | None,
561
selected_percentiles: list[float],
562
goodput_config_dict: dict[str, float],
563
) -> tuple[BenchmarkMetrics, list[int]]:
@@ -772,7 +772,7 @@ async def benchmark(
772
base_url: str,
773
model_id: str,
774
model_name: str,
775
- tokenizer: TokenizerLike,
775
+ tokenizer: TokenizerLike | None,
776
input_requests: list[SampleRequest],
777
logprobs: int | None,
778
request_rate: float,
@@ -2036,7 +2036,6 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
2036
args.self_timed = False
2037
2038
# Load the dataset.
2039
- assert tokenizer is not None, "Tokenizer must be initialized before loading dataset"
2039
input_requests = get_samples(args, tokenizer)
2040
2041
if args.dataset_name in ("random", "prefix_repetition"):