@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1b77cc3d67

[Benchmark] Add sweep warmup and failure recovery (#57305)

Signed-off-by: louie-tsai <louie.tsai@intel.com>

Louie Tsai committed Sep 29, 2026 at 20:16 UTC 1b77cc3d67eddd98c7ecbad7024213a3eeaf7e93

4 files changed +300 -35

docs/benchmarking/sweeps.md

@@ -88,6 +88,12 @@ vllm bench sweep serve \

89 By default, each parameter combination is benchmarked 3 times to make the results more reliable. You can adjust the number of runs by setting `--num-runs`.

91 +You can optionally run an unmeasured warmup before each parameter combination by setting `--warmup-num-prompts`. The warmup result is saved as `warmup.json` for inspection, but is excluded from the measured summary.

92 +

93 +You can use `--continue-on-error` to record a failed warmup, benchmark run, or server configuration and continue processing the remaining sweep combinations. Failure details are saved in `.failure.json` files.

94 +

95 +Both options are also supported by `vllm bench sweep serve_workload`.

96 +

97 !!! important

98 If both `--serve-params` and `--bench-params` are passed, the script will iterate over the Cartesian product between them.

99 You can use `--dry-run` to preview the commands to be run.

tests/benchmarks/sweep/test_param_sweep.py

+126

@@ -6,6 +6,7 @@ from pathlib import Path

7 import pytest

9 +from vllm.benchmarks.sweep import serve as sweep_serve

10 from vllm.benchmarks.sweep.param_sweep import ParameterSweep, ParameterSweepItem

@@ -247,3 +248,128 @@ class TestParameterSweepItemKeyNormalization:

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}]

vllm/benchmarks/sweep/serve.py

+142 -31

@@ -221,29 +221,80 @@ def run_comb(

221 base_path: Path,

222 num_runs: int,

223 dry_run: bool,

224 + warmup_num_prompts: int = 0,

225 + continue_on_error: bool = False,

226 ):

227 if not _comb_is_valid(serve_comb, bench_comb, link_vars):

228 return None

229

230 comb_data = list[dict[str, object]]()

231

232 + def run_benchmark_with_error_handling(

233 + *,

234 + bench_overrides: ParameterSweepItem,

235 + run_number: int,

236 + output_path: Path,

237 + ):

238 + try:

239 + return run_benchmark(

240 + server,

241 + bench_cmd,

242 + serve_overrides=serve_comb,

243 + bench_overrides=bench_overrides,

244 + run_number=run_number,

245 + output_path=output_path,

246 + dry_run=dry_run,

247 + )

248 + except Exception as exc:

249 + if not continue_on_error:

250 + raise

251 + failure_path = output_path.with_name(output_path.stem + ".failure.json")

252 + failure_path.parent.mkdir(parents=True, exist_ok=True)

253 + failure_path.write_text(

254 + json.dumps(

255 + {

256 + "run_number": run_number,

257 + "error_type": type(exc).__name__,

258 + "error": str(exc),

259 + "serve_overrides": dict(serve_comb),

260 + "bench_overrides": dict(bench_overrides),

261 + },

262 + indent=2,

263 + )

264 + + "\n",

265 + encoding="utf-8",

266 + )

267 + print(f"[FAILED BENCHMARK] run={run_number}: {type(exc).__name__}: {exc}")

268 + print(f"Failure details: {failure_path}")

269 + if server is not None:

270 + with contextlib.suppress(Exception):

271 + server.after_bench()

272 + return None

273 +

274 + if warmup_num_prompts > 0:

275 + warmup_overrides = bench_comb | {"num_prompts": warmup_num_prompts}

276 + run_benchmark_with_error_handling(

277 + bench_overrides=warmup_overrides,

278 + run_number=-1,

279 + output_path=base_path / "warmup.json",

280 + )

281 +

282 for run_number in range(num_runs):

231 - run_data = run_benchmark(

232 - server,

233 - bench_cmd,

234 - serve_overrides=serve_comb,

283 + run_data = run_benchmark_with_error_handling(

284 bench_overrides=bench_comb,

285 run_number=run_number,

286 output_path=_get_comb_run_path(base_path, run_number),

238 - dry_run=dry_run,

287 )

240 -

288 if run_data is not None:

289 comb_data.append(run_data)

290

291 if dry_run:

292 return None

293

294 + if not comb_data:

295 + print("[NO SUCCESSFUL MEASURED RUNS] leaving combination retryable")

296 + return []

297 +

298 with _get_comb_run_path(base_path, run_number=None).open("w") as f:

299 json.dump(comb_data, f, indent=4)

300

@@ -262,36 +313,69 @@ def run_combs(

313 link_vars: list[tuple[str, str]],

314 experiment_dir: Path,

315 num_runs: int,

316 + warmup_num_prompts: int,

317 dry_run: bool,

318 + continue_on_error: bool = False,

319 ):

320 all_data = list[dict[str, object]]()

321 for serve_comb in serve_params:

269 - with server_ctx(

270 - serve_cmd,

271 - after_bench_cmd,

272 - show_stdout=show_stdout,

273 - serve_comb=serve_comb,

274 - bench_params=bench_params,

275 - experiment_dir=experiment_dir,

276 - dry_run=dry_run,

277 - server_ready_timeout=server_ready_timeout,

278 - ) as server:

279 - for bench_comb in bench_params:

280 - base_path = _get_comb_base_path(experiment_dir, serve_comb, bench_comb)

281 -

282 - comb_data = run_comb(

283 - server,

284 - bench_cmd,

285 - serve_comb=serve_comb,

286 - bench_comb=bench_comb,

287 - link_vars=link_vars,

288 - base_path=base_path,

289 - num_runs=num_runs,

290 - dry_run=dry_run,

322 + try:

323 + with server_ctx(

324 + serve_cmd,

325 + after_bench_cmd,

326 + show_stdout=show_stdout,

327 + serve_comb=serve_comb,

328 + bench_params=bench_params,

329 + experiment_dir=experiment_dir,

330 + dry_run=dry_run,

331 + server_ready_timeout=server_ready_timeout,

332 + ) as server:

333 + for bench_comb in bench_params:

334 + base_path = _get_comb_base_path(

335 + experiment_dir, serve_comb, bench_comb

336 + )

337 +

338 + comb_data = run_comb(

339 + server,

340 + bench_cmd,

341 + serve_comb=serve_comb,

342 + bench_comb=bench_comb,

343 + link_vars=link_vars,

344 + base_path=base_path,

345 + num_runs=num_runs,

346 + warmup_num_prompts=warmup_num_prompts,

347 + dry_run=dry_run,

348 + continue_on_error=continue_on_error,

349 + )

350 +

351 + if comb_data is not None:

352 + all_data.extend(comb_data)

353 + except Exception as exc:

354 + if not continue_on_error:

355 + raise

356 + failure_name = sanitize_filename(

357 + f"SERVE-{serve_comb.name or 'default'}-failure.json"

358 + )

359 + failure_path = experiment_dir / failure_name

360 + failure_path.parent.mkdir(parents=True, exist_ok=True)

361 + failure_path.write_text(

362 + json.dumps(

363 + {

364 + "error_type": type(exc).__name__,

365 + "error": str(exc),

366 + "serve_overrides": dict(serve_comb),

367 + },

368 + indent=2,

369 )

292 -

293 - if comb_data is not None:

294 - all_data.extend(comb_data)

370 + + "\n",

371 + encoding="utf-8",

372 + )

373 + print(

374 + "[FAILED SERVER COMBINATION] "

375 + f"{serve_comb.name or '(default)'}: "

376 + f"{type(exc).__name__}: {exc}"

377 + )

378 + print(f"Failure details: {failure_path}")

379

380 if dry_run:

381 return None

@@ -315,6 +399,8 @@ class SweepServeArgs:

399 output_dir: Path

400 experiment_name: str

401 num_runs: int

402 + warmup_num_prompts: int

403 + continue_on_error: bool

404 dry_run: bool

405 resume: bool

406

@@ -351,6 +437,9 @@ class SweepServeArgs:

437 num_runs = args.num_runs

438 if num_runs < 1:

439 raise ValueError("`num_runs` should be at least 1.")

440 + warmup_num_prompts = args.warmup_num_prompts

441 + if warmup_num_prompts < 0:

442 + raise ValueError("`warmup_num_prompts` should be at least 0.")

443

444 return cls(

445 serve_cmd=serve_cmd,

@@ -363,6 +452,8 @@ class SweepServeArgs:

452 output_dir=Path(args.output_dir),

453 experiment_name=experiment_name,

454 num_runs=num_runs,

455 + warmup_num_prompts=warmup_num_prompts,

456 + continue_on_error=args.continue_on_error,

457 dry_run=args.dry_run,

458 resume=args.resume,

459 server_ready_timeout=args.server_ready_timeout,

@@ -453,6 +544,24 @@ class SweepServeArgs:

544 default=3,

545 help="Number of runs per parameter combination.",

546 )

547 + parser.add_argument(

548 + "--warmup-num-prompts",

549 + type=int,

550 + default=0,

551 + help=(

552 + "Number of prompts in one unmeasured warmup run before each "

553 + "parameter combination. The warmup is saved as warmup.json but "

554 + "excluded from summary statistics."

555 + ),

556 + )

557 + parser.add_argument(

558 + "--continue-on-error",

559 + action="store_true",

560 + help=(

561 + "Record and skip a failed warmup or benchmark run instead of "

562 + "terminating the entire sweep."

563 + ),

564 + )

565 parser.add_argument(

566 "--dry-run",

567 action="store_true",

@@ -523,7 +632,9 @@ def run_main(args: SweepServeArgs):

632 bench_params=args.bench_params,

633 experiment_dir=experiment_dir,

634 num_runs=args.num_runs,

635 + warmup_num_prompts=args.warmup_num_prompts,

636 dry_run=args.dry_run,

637 + continue_on_error=args.continue_on_error,

638 )

639

640

vllm/benchmarks/sweep/serve_workload.py

+26 -4

@@ -63,6 +63,8 @@ def run_comb_workload(

63 experiment_dir: Path,

64 num_runs: int,

65 dry_run: bool,

66 + warmup_num_prompts: int,

67 + continue_on_error: bool,

68 workload_var: WorkloadVariable,

69 workload_value: int,

70 ) -> list[dict[str, object]] | None:

@@ -82,6 +84,8 @@ def run_comb_workload(

84 ),

85 num_runs=num_runs,

86 dry_run=dry_run,

87 + warmup_num_prompts=warmup_num_prompts,

88 + continue_on_error=continue_on_error,

89 )

@@ -97,6 +101,8 @@ def explore_comb_workloads(

101 experiment_dir: Path,

102 num_runs: int,

103 dry_run: bool,

104 + warmup_num_prompts: int,

105 + continue_on_error: bool,

106 ):

107 print("[WL START]")

108 print(f"Serve parameters: {serve_comb.as_text() or '(None)'}")

@@ -129,6 +135,8 @@ def explore_comb_workloads(

135 experiment_dir=experiment_dir,

136 num_runs=num_runs,

137 dry_run=dry_run,

138 + warmup_num_prompts=warmup_num_prompts,

139 + continue_on_error=continue_on_error,

140 workload_var=workload_var,

141 workload_value=1,

142 )

@@ -141,16 +149,22 @@ def explore_comb_workloads(

149 experiment_dir=experiment_dir,

150 num_runs=num_runs,

151 dry_run=dry_run,

152 + warmup_num_prompts=warmup_num_prompts,

153 + continue_on_error=continue_on_error,

154 workload_var=workload_var,

155 workload_value=dataset_size,

156 )

157

148 - if serial_workload_data is None or batch_workload_data is None:

158 + if not serial_workload_data or not batch_workload_data:

159 if dry_run:

160 print("Omitting intermediate Workload iterations.")

151 - print("[WL END]")

152 -

153 - return

161 + else:

162 + print(

163 + "Unable to establish both Workload Explorer endpoints; "

164 + "skipping intermediate workload levels."

165 + )

166 + print("[WL END]")

167 + return []

168

169 serial_workload_value = math.ceil(

170 _estimate_workload_avg(serial_workload_data, workload_var)

@@ -181,6 +195,8 @@ def explore_comb_workloads(

195 experiment_dir=experiment_dir,

196 num_runs=num_runs,

197 dry_run=dry_run,

198 + warmup_num_prompts=warmup_num_prompts,

199 + continue_on_error=continue_on_error,

200 workload_var=workload_var,

201 workload_value=inter_workload_value,

202 )

@@ -207,6 +223,8 @@ def explore_combs_workloads(

223 experiment_dir: Path,

224 num_runs: int,

225 dry_run: bool,

226 + warmup_num_prompts: int,

227 + continue_on_error: bool,

228 ):

229 if any(bench_comb.has_param(workload_var) for bench_comb in bench_params):

230 raise ValueError(

@@ -238,6 +256,8 @@ def explore_combs_workloads(

256 experiment_dir=experiment_dir,

257 num_runs=num_runs,

258 dry_run=dry_run,

259 + warmup_num_prompts=warmup_num_prompts,

260 + continue_on_error=continue_on_error,

261 )

262

263 if comb_data is not None:

@@ -315,6 +335,8 @@ def run_main(args: SweepServeWorkloadArgs):

335 experiment_dir=experiment_dir,

336 num_runs=args.num_runs,

337 dry_run=args.dry_run,

338 + warmup_num_prompts=args.warmup_num_prompts,

339 + continue_on_error=args.continue_on_error,

340 )

341

342