| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import json |
| 4 | import tempfile |
| 5 | from pathlib import Path |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.benchmarks.sweep import serve as sweep_serve |
| 10 | from vllm.benchmarks.sweep.param_sweep import ParameterSweep, ParameterSweepItem |
| 11 | |
| 12 | |
| 13 | class TestParameterSweepItem: |
| 14 | """Test ParameterSweepItem functionality.""" |
| 15 | |
| 16 | @pytest.mark.parametrize( |
| 17 | "input_dict,expected", |
| 18 | [ |
| 19 | ( |
| 20 | {"compilation_config.use_inductor_graph_partition": False}, |
| 21 | "--compilation-config.use_inductor_graph_partition=false", |
| 22 | ), |
| 23 | ( |
| 24 | {"compilation_config.use_inductor_graph_partition": True}, |
| 25 | "--compilation-config.use_inductor_graph_partition=true", |
| 26 | ), |
| 27 | ], |
| 28 | ) |
| 29 | def test_nested_boolean_params(self, input_dict, expected): |
| 30 | """Test that nested boolean params use =true/false syntax.""" |
| 31 | item = ParameterSweepItem.from_record(input_dict) |
| 32 | cmd = item.apply_to_cmd(["vllm", "serve", "model"]) |
| 33 | assert expected in cmd |
| 34 | |
| 35 | @pytest.mark.parametrize( |
| 36 | "input_dict,expected", |
| 37 | [ |
| 38 | ({"enable_prefix_caching": False}, "--no-enable-prefix-caching"), |
| 39 | ({"enable_prefix_caching": True}, "--enable-prefix-caching"), |
| 40 | ({"disable_log_stats": False}, "--no-disable-log-stats"), |
| 41 | ({"disable_log_stats": True}, "--disable-log-stats"), |
| 42 | ], |
| 43 | ) |
| 44 | def test_non_nested_boolean_params(self, input_dict, expected): |
| 45 | """Test that non-nested boolean params use --no- prefix.""" |
| 46 | item = ParameterSweepItem.from_record(input_dict) |
| 47 | cmd = item.apply_to_cmd(["vllm", "serve", "model"]) |
| 48 | assert expected in cmd |
| 49 | |
| 50 | @pytest.mark.parametrize( |
| 51 | "compilation_config", |
| 52 | [ |
| 53 | {"cudagraph_mode": "full", "mode": 2, "use_inductor_graph_partition": True}, |
| 54 | { |
| 55 | "cudagraph_mode": "piecewise", |
| 56 | "mode": 3, |
| 57 | "use_inductor_graph_partition": False, |
| 58 | }, |
| 59 | ], |
| 60 | ) |
| 61 | def test_nested_dict_value(self, compilation_config): |
| 62 | """Test that nested dict values are serialized as JSON.""" |
| 63 | item = ParameterSweepItem.from_record( |
| 64 | {"compilation_config": compilation_config} |
| 65 | ) |
| 66 | cmd = item.apply_to_cmd(["vllm", "serve", "model"]) |
| 67 | assert "--compilation-config" in cmd |
| 68 | # The dict should be JSON serialized |
| 69 | idx = cmd.index("--compilation-config") |
| 70 | assert json.loads(cmd[idx + 1]) == compilation_config |
| 71 | |
| 72 | @pytest.mark.parametrize( |
| 73 | "input_dict,expected_key,expected_value", |
| 74 | [ |
| 75 | ({"model": "test-model"}, "--model", "test-model"), |
| 76 | ({"max_tokens": 100}, "--max-tokens", "100"), |
| 77 | ({"temperature": 0.7}, "--temperature", "0.7"), |
| 78 | ], |
| 79 | ) |
| 80 | def test_string_and_numeric_values(self, input_dict, expected_key, expected_value): |
| 81 | """Test that string and numeric values are handled correctly.""" |
| 82 | item = ParameterSweepItem.from_record(input_dict) |
| 83 | cmd = item.apply_to_cmd(["vllm", "serve"]) |
| 84 | assert expected_key in cmd |
| 85 | assert expected_value in cmd |
| 86 | |
| 87 | @pytest.mark.parametrize( |
| 88 | "input_dict,expected_key,key_idx_offset", |
| 89 | [ |
| 90 | ({"max_tokens": 200}, "--max-tokens", 1), |
| 91 | ({"enable_prefix_caching": False}, "--no-enable-prefix-caching", 0), |
| 92 | ], |
| 93 | ) |
| 94 | def test_replace_existing_parameter(self, input_dict, expected_key, key_idx_offset): |
| 95 | """Test that existing parameters in cmd are replaced.""" |
| 96 | item = ParameterSweepItem.from_record(input_dict) |
| 97 | |
| 98 | if key_idx_offset == 1: |
| 99 | # Key-value pair |
| 100 | cmd = item.apply_to_cmd(["vllm", "serve", "--max-tokens", "100", "model"]) |
| 101 | assert expected_key in cmd |
| 102 | idx = cmd.index(expected_key) |
| 103 | assert cmd[idx + 1] == "200" |
| 104 | assert "100" not in cmd |
| 105 | else: |
| 106 | # Boolean flag |
| 107 | cmd = item.apply_to_cmd( |
| 108 | ["vllm", "serve", "--enable-prefix-caching", "model"] |
| 109 | ) |
| 110 | assert expected_key in cmd |
| 111 | assert "--enable-prefix-caching" not in cmd |
| 112 | |
| 113 | |
| 114 | class TestParameterSweep: |
| 115 | """Test ParameterSweep functionality.""" |
| 116 | |
| 117 | def test_from_records_list(self): |
| 118 | """Test creating ParameterSweep from a list of records.""" |
| 119 | records: list[dict[str, object]] = [ |
| 120 | {"max_tokens": 100, "temperature": 0.7}, |
| 121 | {"max_tokens": 200, "temperature": 0.9}, |
| 122 | ] |
| 123 | sweep = ParameterSweep.from_records(records) |
| 124 | assert len(sweep) == 2 |
| 125 | assert sweep[0]["max_tokens"] == 100 |
| 126 | assert sweep[1]["max_tokens"] == 200 |
| 127 | |
| 128 | def test_read_from_dict(self): |
| 129 | """Test creating ParameterSweep from a dict format.""" |
| 130 | data: dict[str, dict[str, object]] = { |
| 131 | "experiment1": {"max_tokens": 100, "temperature": 0.7}, |
| 132 | "experiment2": {"max_tokens": 200, "temperature": 0.9}, |
| 133 | } |
| 134 | sweep = ParameterSweep.read_from_dict(data) |
| 135 | assert len(sweep) == 2 |
| 136 | |
| 137 | # Check that items have the _benchmark_name field |
| 138 | names = {item["_benchmark_name"] for item in sweep} |
| 139 | assert names == {"experiment1", "experiment2"} |
| 140 | |
| 141 | # Check that parameters are preserved |
| 142 | for item in sweep: |
| 143 | if item["_benchmark_name"] == "experiment1": |
| 144 | assert item["max_tokens"] == 100 |
| 145 | assert item["temperature"] == 0.7 |
| 146 | elif item["_benchmark_name"] == "experiment2": |
| 147 | assert item["max_tokens"] == 200 |
| 148 | assert item["temperature"] == 0.9 |
| 149 | |
| 150 | def test_read_json_list_format(self): |
| 151 | """Test reading JSON file with list format.""" |
| 152 | records: list[dict[str, object]] = [ |
| 153 | {"max_tokens": 100, "temperature": 0.7}, |
| 154 | {"max_tokens": 200, "temperature": 0.9}, |
| 155 | ] |
| 156 | |
| 157 | with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: |
| 158 | json.dump(records, f) |
| 159 | temp_path = Path(f.name) |
| 160 | |
| 161 | try: |
| 162 | sweep = ParameterSweep.read_json(temp_path) |
| 163 | assert len(sweep) == 2 |
| 164 | assert sweep[0]["max_tokens"] == 100 |
| 165 | assert sweep[1]["max_tokens"] == 200 |
| 166 | finally: |
| 167 | temp_path.unlink() |
| 168 | |
| 169 | def test_read_json_dict_format(self): |
| 170 | """Test reading JSON file with dict format.""" |
| 171 | data: dict[str, dict[str, object]] = { |
| 172 | "experiment1": {"max_tokens": 100, "temperature": 0.7}, |
| 173 | "experiment2": {"max_tokens": 200, "temperature": 0.9}, |
| 174 | } |
| 175 | |
| 176 | with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: |
| 177 | json.dump(data, f) |
| 178 | temp_path = Path(f.name) |
| 179 | |
| 180 | try: |
| 181 | sweep = ParameterSweep.read_json(temp_path) |
| 182 | assert len(sweep) == 2 |
| 183 | |
| 184 | # Check that items have the _benchmark_name field |
| 185 | names = {item["_benchmark_name"] for item in sweep} |
| 186 | assert names == {"experiment1", "experiment2"} |
| 187 | finally: |
| 188 | temp_path.unlink() |
| 189 | |
| 190 | def test_unique_benchmark_names_validation(self): |
| 191 | """Test that duplicate _benchmark_name values raise an error.""" |
| 192 | # Test with duplicate names in list format |
| 193 | records: list[dict[str, object]] = [ |
| 194 | {"_benchmark_name": "exp1", "max_tokens": 100}, |
| 195 | {"_benchmark_name": "exp1", "max_tokens": 200}, |
| 196 | ] |
| 197 | |
| 198 | with pytest.raises(ValueError, match="Duplicate _benchmark_name values"): |
| 199 | ParameterSweep.from_records(records) |
| 200 | |
| 201 | def test_unique_benchmark_names_multiple_duplicates(self): |
| 202 | """Test validation with multiple duplicate names.""" |
| 203 | records: list[dict[str, object]] = [ |
| 204 | {"_benchmark_name": "exp1", "max_tokens": 100}, |
| 205 | {"_benchmark_name": "exp1", "max_tokens": 200}, |
| 206 | {"_benchmark_name": "exp2", "max_tokens": 300}, |
| 207 | {"_benchmark_name": "exp2", "max_tokens": 400}, |
| 208 | ] |
| 209 | |
| 210 | with pytest.raises(ValueError, match="Duplicate _benchmark_name values"): |
| 211 | ParameterSweep.from_records(records) |
| 212 | |
| 213 | def test_no_benchmark_names_allowed(self): |
| 214 | """Test that records without _benchmark_name are allowed.""" |
| 215 | records: list[dict[str, object]] = [ |
| 216 | {"max_tokens": 100, "temperature": 0.7}, |
| 217 | {"max_tokens": 200, "temperature": 0.9}, |
| 218 | ] |
| 219 | sweep = ParameterSweep.from_records(records) |
| 220 | assert len(sweep) == 2 |
| 221 | |
| 222 | def test_mixed_benchmark_names_allowed(self): |
| 223 | """Test that mixing records with and without _benchmark_name is allowed.""" |
| 224 | records: list[dict[str, object]] = [ |
| 225 | {"_benchmark_name": "exp1", "max_tokens": 100}, |
| 226 | {"max_tokens": 200, "temperature": 0.9}, |
| 227 | ] |
| 228 | sweep = ParameterSweep.from_records(records) |
| 229 | assert len(sweep) == 2 |
| 230 | |
| 231 | |
| 232 | class TestParameterSweepItemKeyNormalization: |
| 233 | """Test key normalization in ParameterSweepItem.""" |
| 234 | |
| 235 | def test_underscore_to_hyphen_conversion(self): |
| 236 | """Test that underscores are converted to hyphens in CLI.""" |
| 237 | item = ParameterSweepItem.from_record({"max_tokens": 100}) |
| 238 | cmd = item.apply_to_cmd(["vllm", "serve"]) |
| 239 | assert "--max-tokens" in cmd |
| 240 | |
| 241 | def test_nested_key_preserves_suffix(self): |
| 242 | """Test that nested keys preserve the suffix format.""" |
| 243 | # The suffix after the dot should preserve underscores |
| 244 | item = ParameterSweepItem.from_record( |
| 245 | {"compilation_config.some_nested_param": "value"} |
| 246 | ) |
| 247 | cmd = item.apply_to_cmd(["vllm", "serve"]) |
| 248 | # The prefix (compilation_config) gets converted to hyphens, |
| 249 | # but the suffix (some_nested_param) is preserved |
| 250 | assert any("compilation-config.some_nested_param" in arg for arg in cmd) |
| 251 | |
| 252 | |
| 253 | def test_run_comb_excludes_warmup_from_measured_results( |
| 254 | monkeypatch: pytest.MonkeyPatch, tmp_path: Path |
| 255 | ): |
| 256 | calls: list[tuple[int, int, str]] = [] |
| 257 | |
| 258 | def fake_run_benchmark( |
| 259 | server, |
| 260 | bench_cmd, |
| 261 | *, |
| 262 | serve_overrides, |
| 263 | bench_overrides, |
| 264 | run_number, |
| 265 | output_path, |
| 266 | dry_run, |
| 267 | ): |
| 268 | calls.append( |
| 269 | (run_number, int(bench_overrides["num_prompts"]), output_path.name) |
| 270 | ) |
| 271 | return {"run_number": run_number} |
| 272 | |
| 273 | monkeypatch.setattr(sweep_serve, "run_benchmark", fake_run_benchmark) |
| 274 | base_path = tmp_path / "combination" |
| 275 | base_path.mkdir() |
| 276 | |
| 277 | measured = sweep_serve.run_comb( |
| 278 | None, |
| 279 | [], |
| 280 | serve_comb=ParameterSweepItem(), |
| 281 | bench_comb=ParameterSweepItem({"num_prompts": 320}), |
| 282 | link_vars=[], |
| 283 | base_path=base_path, |
| 284 | num_runs=2, |
| 285 | warmup_num_prompts=32, |
| 286 | dry_run=False, |
| 287 | ) |
| 288 | |
| 289 | assert calls == [ |
| 290 | (-1, 32, "warmup.json"), |
| 291 | (0, 320, "run=0.json"), |
| 292 | (1, 320, "run=1.json"), |
| 293 | ] |
| 294 | assert measured == [{"run_number": 0}, {"run_number": 1}] |
| 295 | |
| 296 | |
| 297 | def test_run_comb_continue_on_error_keeps_later_runs( |
| 298 | monkeypatch: pytest.MonkeyPatch, tmp_path: Path |
| 299 | ): |
| 300 | calls: list[int] = [] |
| 301 | |
| 302 | def fake_run_benchmark( |
| 303 | server, |
| 304 | bench_cmd, |
| 305 | *, |
| 306 | serve_overrides, |
| 307 | bench_overrides, |
| 308 | run_number, |
| 309 | output_path, |
| 310 | dry_run, |
| 311 | ): |
| 312 | calls.append(run_number) |
| 313 | if run_number == 0: |
| 314 | raise RuntimeError("synthetic run failure") |
| 315 | return {"run_number": run_number} |
| 316 | |
| 317 | monkeypatch.setattr(sweep_serve, "run_benchmark", fake_run_benchmark) |
| 318 | base_path = tmp_path / "combination" |
| 319 | base_path.mkdir() |
| 320 | |
| 321 | measured = sweep_serve.run_comb( |
| 322 | None, |
| 323 | [], |
| 324 | serve_comb=ParameterSweepItem(), |
| 325 | bench_comb=ParameterSweepItem({"num_prompts": 100}), |
| 326 | link_vars=[], |
| 327 | base_path=base_path, |
| 328 | num_runs=2, |
| 329 | warmup_num_prompts=0, |
| 330 | dry_run=False, |
| 331 | continue_on_error=True, |
| 332 | ) |
| 333 | |
| 334 | assert calls == [0, 1] |
| 335 | assert measured == [{"run_number": 1}] |
| 336 | assert (base_path / "run=0.failure.json").exists() |
| 337 | assert (base_path / "summary.json").exists() |
| 338 | |
| 339 | |
| 340 | def test_run_comb_warmup_default_is_backward_compatible( |
| 341 | monkeypatch: pytest.MonkeyPatch, tmp_path: Path |
| 342 | ): |
| 343 | """Older callers can omit warmup_num_prompts.""" |
| 344 | calls: list[int] = [] |
| 345 | |
| 346 | def fake_run_benchmark( |
| 347 | server, |
| 348 | bench_cmd, |
| 349 | *, |
| 350 | serve_overrides, |
| 351 | bench_overrides, |
| 352 | run_number, |
| 353 | output_path, |
| 354 | dry_run, |
| 355 | ): |
| 356 | calls.append(run_number) |
| 357 | return {"run_number": run_number} |
| 358 | |
| 359 | monkeypatch.setattr(sweep_serve, "run_benchmark", fake_run_benchmark) |
| 360 | base_path = tmp_path / "combination" |
| 361 | base_path.mkdir() |
| 362 | |
| 363 | measured = sweep_serve.run_comb( |
| 364 | None, |
| 365 | [], |
| 366 | serve_comb=ParameterSweepItem(), |
| 367 | bench_comb=ParameterSweepItem({"num_prompts": 10}), |
| 368 | link_vars=[], |
| 369 | base_path=base_path, |
| 370 | num_runs=1, |
| 371 | dry_run=False, |
| 372 | ) |
| 373 | |
| 374 | assert calls == [0] |
| 375 | assert measured == [{"run_number": 0}] |