@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 163 lines 4.93 KB
Raw
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 probe_request_rate=0.0,
88 # misc serve args that main_async reads before reaching get_samples
89 plot_timeline=False,
90 plot_dataset_stats=False,
91 self_timed=None,
92 metadata=None,
93 label=None,
94 logprobs=None,
95 use_beam_search=False,
96 ignore_eos=False,
97 goodput=None,
98 percentile_metrics="ttft,tpot,itl,e2el",
99 metric_percentiles="25,50,75,99",
100 save_result=False,
101 append_result=False,
102 result_dir=".",
103 result_filename=None,
104 num_warmups=0,
105 profile=False,
106 disable_tqdm=True,
107 lora_modules=None,
108 lora_assignment="random",
109 ramp_up_strategy=None,
110 ramp_up_start_rps=None,
111 ramp_up_end_rps=None,
112 ready_check_timeout_sec=0,
113 extra_body=None,
114 top_p=None,
115 top_k=None,
116 min_p=None,
117 temperature=None,
118 frequency_penalty=None,
119 presence_penalty=None,
120 repetition_penalty=None,
121 save_detailed=False,
122 input_len=None,
123 output_len=None,
124 )
125
126
127 @pytest.mark.benchmark
128 def test_main_async_skip_tokenizer_init_does_not_raise(tmp_path: Path) -> None:
129 """main_async must not raise AssertionError when skip_tokenizer_init=True
130 and dataset_name='custom'.
131
132 On main (before the fix) this test fails with:
133 AssertionError: Tokenizer must be initialized before loading dataset
134 """
135 dataset_path = tmp_path / "dataset_url_input_india.jsonl"
136 _write_dataset(dataset_path)
137
138 args = _args(str(dataset_path))
139
140 # Patch benchmark() so we never make real HTTP requests — the regression
141 # triggers before benchmark() is ever called, so this just keeps the test
142 # fast and hermetic.
143 mock_result = {
144 "completed": 1,
145 "failed": 0,
146 "total_input_tokens": 1,
147 "total_output_tokens": 1,
148 "request_throughput": 1.0,
149 "output_throughput": 1.0,
150 "total_token_throughput": 1.0,
151 "input_lens": [],
152 "output_lens": [],
153 "ttfts": [],
154 "itls": [],
155 "generated_texts": [],
156 "errors": [],
157 "duration": 1.0,
158 }
159 with patch.object(
160 serve_module, "benchmark", new=AsyncMock(return_value=mock_result)
161 ):
162 # Must NOT raise AssertionError
163 asyncio.run(serve_module.main_async(args))