@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4292e3b807

Fork 0

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 7, 12:45 UTC

[Benchmark] Improve UX of sweep scripts (#35600)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Feb 28, 2026 at 16:36 UTC 4292e3b807a51507f60f43b3829b5e5e918f5b87

6 files changed +191 -128

docs/benchmarking/sweeps.md

+14 -10

@@ -72,7 +72,7 @@ Follow these steps to run the script:

72 ]

73 ```

75 -5. Determine where you want to save the results, and pass that to `--output-dir`.

75 +5. Set `--output-dir` and optionally `--experiment-name` to control where to save the results.

77 Example command:

@@ -82,7 +82,8 @@ vllm bench sweep serve \

82 --bench-cmd 'vllm bench serve --model meta-llama/Llama-2-7b-chat-hf --backend vllm --endpoint /v1/completions --dataset-name sharegpt --dataset-path benchmarks/ShareGPT_V3_unfiltered_cleaned_split.json' \

83 --serve-params benchmarks/serve_hparams.json \

84 --bench-params benchmarks/bench_hparams.json \

85 - -o benchmarks/results

85 + --output-dir benchmarks/results \

86 + --experiment-name demo

87 ```

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`.

@@ -118,7 +119,8 @@ vllm bench sweep serve_workload \

119 --serve-params benchmarks/serve_hparams.json \

120 --bench-params benchmarks/bench_hparams.json \

121 --num-runs 1 \

121 - -o benchmarks/results

122 + --output-dir benchmarks/results \

123 + --experiment-name demo

124 ```

125

126 The algorithm for exploring different workload levels can be summarized as follows:

@@ -186,7 +188,8 @@ vllm bench sweep startup \

188 --startup-cmd 'vllm bench startup --model Qwen/Qwen3-0.6B' \

189 --serve-params benchmarks/serve_hparams.json \

190 --startup-params benchmarks/startup_hparams.json \

189 - -o benchmarks/results

191 + --output-dir benchmarks/results \

192 + --experiment-name demo

193 ```

194

195 !!! important

@@ -204,11 +207,10 @@ Control the variables to plot via `--var-x` and `--var-y`, optionally applying `

207 Example commands for visualizing [Workload Explorer](#workload-explorer) results:

208

209 ```bash

207 -# Name of the directory that stores the results

208 -TIMESTAMP=$1

210 +EXPERIMENT_DIR=${1:-"benchmarks/results/demo"}

211

212 # Latency increases as the workload increases

211 -vllm bench sweep plot benchmarks/results/$TIMESTAMP \

213 +vllm bench sweep plot $EXPERIMENT_DIR \

214 --var-x max_concurrency \

215 --var-y median_ttft_ms \

216 --col-by _benchmark_name \

@@ -216,7 +218,7 @@ vllm bench sweep plot benchmarks/results/$TIMESTAMP \

218 --fig-name latency_curve

219

220 # Throughput saturates as workload increases

219 -vllm bench sweep plot benchmarks/results/$TIMESTAMP \

221 +vllm bench sweep plot $EXPERIMENT_DIR \

222 --var-x max_concurrency \

223 --var-y total_token_throughput \

224 --col-by _benchmark_name \

@@ -224,7 +226,7 @@ vllm bench sweep plot benchmarks/results/$TIMESTAMP \

226 --fig-name throughput_curve

227

228 # Tradeoff between latency and throughput

227 -vllm bench sweep plot benchmarks/results/$TIMESTAMP \

229 +vllm bench sweep plot $EXPERIMENT_DIR \

230 --var-x total_token_throughput \

231 --var-y median_ttft_ms \

232 --col-by _benchmark_name \

@@ -249,7 +251,9 @@ Higher concurrency or batch size can raise GPU efficiency (per-GPU), but can add

251 Example:

252

253 ```bash

252 -vllm bench sweep plot_pareto benchmarks/results/<timestamp> \

254 +EXPERIMENT_DIR=${1:-"benchmarks/results/demo"}

255 +

256 +vllm bench sweep plot_pareto $EXPERIMENT_DIR \

257 --label-by max_concurrency,tensor_parallel_size,pipeline_parallel_size

258 ```

259

vllm/benchmarks/sweep/plot.py

+3 -5

@@ -499,7 +499,7 @@ class SweepPlotArgs:

499

500 @classmethod

501 def from_cli_args(cls, args: argparse.Namespace):

502 - output_dir = Path(args.OUTPUT_DIR)

502 + output_dir = Path(args.EXPERIMENT_DIR)

503 if not output_dir.exists():

504 raise ValueError(f"No parameter sweep results under {output_dir}")

505

@@ -531,11 +531,9 @@ class SweepPlotArgs:

531 @classmethod

532 def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:

533 parser.add_argument(

534 - "OUTPUT_DIR",

534 + "EXPERIMENT_DIR",

535 type=str,

536 - default="results",

537 - help="The directory containing the results to plot, "

538 - "i.e., the `--output-dir` argument to the parameter sweep script.",

536 + help="The directory containing the sweep results to plot.",

537 )

538 parser.add_argument(

539 "--fig-dir",

vllm/benchmarks/sweep/plot_pareto.py

+2 -3

@@ -325,7 +325,7 @@ class SweepPlotParetoArgs:

325

326 @classmethod

327 def from_cli_args(cls, args: argparse.Namespace):

328 - output_dir = Path(args.OUTPUT_DIR)

328 + output_dir = Path(args.EXPERIMENT_DIR)

329 if not output_dir.exists():

330 raise ValueError(f"No parameter sweep results under {output_dir}")

331

@@ -342,9 +342,8 @@ class SweepPlotParetoArgs:

342 @classmethod

343 def add_cli_args(cls, parser: argparse.ArgumentParser):

344 parser.add_argument(

345 - "OUTPUT_DIR",

345 + "EXPERIMENT_DIR",

346 type=str,

347 - default="results",

347 help="The directory containing the sweep results to plot.",

348 )

349 parser.add_argument(

vllm/benchmarks/sweep/serve.py

+80 -45

@@ -4,6 +4,7 @@ import argparse

4 import contextlib

5 import json

6 import shlex

7 +from contextlib import contextmanager

8 from dataclasses import dataclass

9 from datetime import datetime

10 from pathlib import Path

@@ -135,7 +136,7 @@ def run_benchmark(

136

137

138 def _get_comb_base_path(

138 - output_dir: Path,

139 + experiment_dir: Path,

140 serve_comb: ParameterSweepItem,

141 bench_comb: ParameterSweepItem,

142 *,

@@ -149,7 +150,7 @@ def _get_comb_base_path(

150 if extra_parts:

151 parts.extend(extra_parts)

152

152 - return output_dir / sanitize_filename("-".join(parts))

153 + return experiment_dir / sanitize_filename("-".join(parts))

154

155

156 def _get_comb_run_path(base_path: Path, run_number: int | None):

@@ -162,10 +163,10 @@ def _get_comb_run_path(base_path: Path, run_number: int | None):

163 def _comb_needs_server(

164 serve_comb: ParameterSweepItem,

165 bench_combs: ParameterSweep,

165 - output_dir: Path,

166 + experiment_dir: Path,

167 ):

168 for bench_comb in bench_combs:

168 - base_path = _get_comb_base_path(output_dir, serve_comb, bench_comb)

169 + base_path = _get_comb_base_path(experiment_dir, serve_comb, bench_comb)

170 if not _get_comb_run_path(base_path, run_number=None).exists():

171 return True

172

@@ -179,11 +180,11 @@ def server_ctx(

180 show_stdout: bool,

181 serve_comb: ParameterSweepItem,

182 bench_params: ParameterSweep,

182 - output_dir: Path,

183 + experiment_dir: Path,

184 dry_run: bool,

185 server_ready_timeout: int = 300,

186 ):

186 - if not _comb_needs_server(serve_comb, bench_params, output_dir):

187 + if not _comb_needs_server(serve_comb, bench_params, experiment_dir):

188 return contextlib.nullcontext()

189

190 return run_server(

@@ -215,10 +216,10 @@ def run_comb(

216 *,

217 serve_comb: ParameterSweepItem,

218 bench_comb: ParameterSweepItem,

219 + link_vars: list[tuple[str, str]],

220 base_path: Path,

221 num_runs: int,

222 dry_run: bool,

221 - link_vars: list[tuple[str, str]],

223 ):

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

225 return None

@@ -257,10 +258,10 @@ def run_combs(

258 server_ready_timeout: int,

259 serve_params: ParameterSweep,

260 bench_params: ParameterSweep,

260 - output_dir: Path,

261 + link_vars: list[tuple[str, str]],

262 + experiment_dir: Path,

263 num_runs: int,

264 dry_run: bool,

263 - link_vars: list[tuple[str, str]],

265 ):

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

267 for serve_comb in serve_params:

@@ -270,22 +271,22 @@ def run_combs(

271 show_stdout=show_stdout,

272 serve_comb=serve_comb,

273 bench_params=bench_params,

273 - output_dir=output_dir,

274 + experiment_dir=experiment_dir,

275 dry_run=dry_run,

276 server_ready_timeout=server_ready_timeout,

277 ) as server:

278 for bench_comb in bench_params:

278 - base_path = _get_comb_base_path(output_dir, serve_comb, bench_comb)

279 + base_path = _get_comb_base_path(experiment_dir, serve_comb, bench_comb)

280

281 comb_data = run_comb(

282 server,

283 bench_cmd,

284 serve_comb=serve_comb,

285 bench_comb=bench_comb,

286 + link_vars=link_vars,

287 base_path=base_path,

288 num_runs=num_runs,

289 dry_run=dry_run,

288 - link_vars=link_vars,

290 )

291

292 if comb_data is not None:

@@ -295,7 +296,7 @@ def run_combs(

296 return None

297

298 combined_df = pd.DataFrame.from_records(all_data)

298 - combined_df.to_csv(output_dir / "summary.csv")

299 + combined_df.to_csv(experiment_dir / "summary.csv")

300

301 return combined_df

302

@@ -309,11 +310,12 @@ class SweepServeArgs:

310 server_ready_timeout: int

311 serve_params: ParameterSweep

312 bench_params: ParameterSweep

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

314 output_dir: Path

315 + experiment_name: str

316 num_runs: int

317 dry_run: bool

315 - resume: str | None

316 - link_vars: list[tuple[str, str]]

318 + resume: bool

319

320 parser_name: ClassVar[str] = "serve"

321 parser_help: ClassVar[str] = "Run vLLM server benchmark under multiple settings."

@@ -340,6 +342,11 @@ class SweepServeArgs:

342

343 link_vars = cls.parse_link_vars(args.link_vars)

344

345 + if args.experiment_name:

346 + experiment_name = args.experiment_name

347 + else:

348 + experiment_name = datetime.now().strftime("%Y%m%d_%H%M%S")

349 +

350 num_runs = args.num_runs

351 if num_runs < 1:

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

@@ -351,11 +358,12 @@ class SweepServeArgs:

358 show_stdout=args.show_stdout,

359 serve_params=serve_params,

360 bench_params=bench_params,

361 + link_vars=link_vars,

362 output_dir=Path(args.output_dir),

363 + experiment_name=experiment_name,

364 num_runs=num_runs,

365 dry_run=args.dry_run,

366 resume=args.resume,

358 - link_vars=link_vars,

367 server_ready_timeout=args.server_ready_timeout,

368 )

369

@@ -392,6 +400,7 @@ class SweepServeArgs:

400 default=300,

401 help="Timeout in seconds to wait for the server to become ready.",

402 )

403 +

404 parser.add_argument(

405 "--serve-params",

406 type=str,

@@ -402,6 +411,16 @@ class SweepServeArgs:

411 "If both `serve_params` and `bench_params` are given, "

412 "this script will iterate over their Cartesian product.",

413 )

414 + parser.add_argument(

415 + "--link-vars",

416 + type=str,

417 + default="",

418 + help=(

419 + "Comma-separated list of linked variables between serve and bench, "

420 + "e.g. max_num_seqs=max_concurrency,max_model_len=random_input_len"

421 + ),

422 + )

423 +

424 parser.add_argument(

425 "--bench-params",

426 type=str,

@@ -417,7 +436,15 @@ class SweepServeArgs:

436 "--output-dir",

437 type=str,

438 default="results",

420 - help="The directory to which results are written.",

439 + help="The main directory to which results are written.",

440 + )

441 + parser.add_argument(

442 + "-e",

443 + "--experiment-name",

444 + type=str,

445 + default=None,

446 + help="The name of this experiment (defaults to current timestamp). "

447 + "Results will be stored under `output_dir/experiment_name`.",

448 )

449 parser.add_argument(

450 "--num-runs",

@@ -433,21 +460,10 @@ class SweepServeArgs:

460 )

461 parser.add_argument(

462 "--resume",

436 - type=str,

437 - default=None,

438 - help="Set this to the name of a directory under `output_dir` (which is a "

439 - "timestamp) to resume a previous execution of this script, i.e., only run "

440 - "parameter combinations for which there are still no output files.",

441 - )

442 -

443 - parser.add_argument(

444 - "--link-vars",

445 - type=str,

446 - default="",

447 - help=(

448 - "Comma-separated list of linked variables between serve and bench, "

449 - "e.g. max_num_seqs=max_concurrency,max_model_len=random_input_len"

450 - ),

463 + action="store_true",

464 + help="Resume a previous execution of this script, i.e., only run "

465 + "parameter combinations for which there are still no output files "

466 + "under `output_dir/experiment_name`.",

467 )

468

469 return parser

@@ -462,33 +478,52 @@ class SweepServeArgs:

478 pairs.append((a.strip(), b.strip()))

479 return pairs

480

481 + def resolve_experiment_dir(self) -> Path:

482 + experiment_dir = self.output_dir / self.experiment_name

483

466 -def run_main(args: SweepServeArgs):

467 - timestamp = args.resume or datetime.now().strftime("%Y%m%d_%H%M%S")

468 - output_dir = args.output_dir / timestamp

484 + if self.resume:

485 + if not experiment_dir.exists():

486 + raise ValueError(f"Cannot resume from non-existent {experiment_dir=}")

487 + else:

488 + if experiment_dir.exists():

489 + raise ValueError(f"Cannot overwrite existing {experiment_dir=}")

490 +

491 + return experiment_dir

492

470 - if args.resume and not output_dir.exists():

471 - raise ValueError(f"Cannot resume from non-existent directory ({output_dir})")

493 + @contextmanager

494 + def run_ctx(self, experiment_dir: Path):

495 + if self.dry_run:

496 + yield

497 + print(f"Experiment will be saved at: {experiment_dir}")

498 + return

499

473 - try:

500 + try:

501 + yield

502 + print(f"Experiment has been saved at: {experiment_dir}")

503 + except BaseException as exc:

504 + raise RuntimeError(

505 + "The script was terminated early. Use `--resume` "

506 + "to continue the script from its last checkpoint."

507 + ) from exc

508 +

509 +

510 +def run_main(args: SweepServeArgs):

511 + experiment_dir = args.resolve_experiment_dir()

512 +

513 + with args.run_ctx(experiment_dir):

514 return run_combs(

515 serve_cmd=args.serve_cmd,

516 bench_cmd=args.bench_cmd,

517 + link_vars=args.link_vars,

518 after_bench_cmd=args.after_bench_cmd,

519 show_stdout=args.show_stdout,

520 server_ready_timeout=args.server_ready_timeout,

521 serve_params=args.serve_params,

522 bench_params=args.bench_params,

482 - output_dir=output_dir,

523 + experiment_dir=experiment_dir,

524 num_runs=args.num_runs,

525 dry_run=args.dry_run,

485 - link_vars=args.link_vars,

526 )

487 - except BaseException as exc:

488 - raise RuntimeError(

489 - f"The script was terminated early. Use `--resume {timestamp}` "

490 - f"to continue the script from its last checkpoint."

491 - ) from exc

527

528

529 def main(args: argparse.Namespace):

vllm/benchmarks/sweep/serve_workload.py

+22 -32

@@ -3,7 +3,6 @@

3 import argparse

4 import math

5 from dataclasses import asdict, dataclass

6 -from datetime import datetime

6 from pathlib import Path

7 from typing import ClassVar, Literal, get_args

@@ -59,10 +58,10 @@ def run_comb_workload(

58 *,

59 serve_comb: ParameterSweepItem,

60 bench_comb: ParameterSweepItem,

62 - output_dir: Path,

61 + link_vars: list[tuple[str, str]],

62 + experiment_dir: Path,

63 num_runs: int,

64 dry_run: bool,

65 - link_vars: list[tuple[str, str]],

65 workload_var: WorkloadVariable,

66 workload_value: int,

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

@@ -73,15 +72,15 @@ def run_comb_workload(

72 bench_cmd,

73 serve_comb=serve_comb,

74 bench_comb=bench_comb_workload,

75 + link_vars=link_vars,

76 base_path=_get_comb_base_path(

77 - output_dir,

77 + experiment_dir,

78 serve_comb,

79 bench_comb,

80 extra_parts=("WL-", f"{workload_var}={workload_value}"),

81 ),

82 num_runs=num_runs,

83 dry_run=dry_run,

84 - link_vars=link_vars,

84 )

@@ -91,12 +90,12 @@ def explore_comb_workloads(

90 *,

91 serve_comb: ParameterSweepItem,

92 bench_comb: ParameterSweepItem,

93 + link_vars: list[tuple[str, str]],

94 workload_var: WorkloadVariable,

95 workload_iters: int,

96 - output_dir: Path,

96 + experiment_dir: Path,

97 num_runs: int,

98 dry_run: bool,

99 - link_vars: list[tuple[str, str]],

99 ):

100 print("[WL START]")

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

@@ -125,10 +124,10 @@ def explore_comb_workloads(

124 bench_cmd,

125 serve_comb=serve_comb,

126 bench_comb=bench_comb | {"max_concurrency": 1},

128 - output_dir=output_dir,

127 + link_vars=link_vars,

128 + experiment_dir=experiment_dir,

129 num_runs=num_runs,

130 dry_run=dry_run,

131 - link_vars=link_vars,

131 workload_var=workload_var,

132 workload_value=1,

133 )

@@ -137,10 +136,10 @@ def explore_comb_workloads(

136 bench_cmd,

137 serve_comb=serve_comb,

138 bench_comb=bench_comb | {"max_concurrency": dataset_size},

140 - output_dir=output_dir,

139 + link_vars=link_vars,

140 + experiment_dir=experiment_dir,

141 num_runs=num_runs,

142 dry_run=dry_run,

143 - link_vars=link_vars,

143 workload_var=workload_var,

144 workload_value=dataset_size,

145 )

@@ -177,10 +176,10 @@ def explore_comb_workloads(

176 bench_cmd,

177 serve_comb=serve_comb,

178 bench_comb=bench_comb,

180 - output_dir=output_dir,

179 + link_vars=link_vars,

180 + experiment_dir=experiment_dir,

181 num_runs=num_runs,

182 dry_run=dry_run,

183 - link_vars=link_vars,

183 workload_var=workload_var,

184 workload_value=inter_workload_value,

185 )

@@ -201,12 +200,12 @@ def explore_combs_workloads(

200 server_ready_timeout: int,

201 serve_params: ParameterSweep,

202 bench_params: ParameterSweep,

203 + link_vars: list[tuple[str, str]],

204 workload_var: WorkloadVariable,

205 workload_iters: int,

206 - output_dir: Path,

206 + experiment_dir: Path,

207 num_runs: int,

208 dry_run: bool,

209 - link_vars: list[tuple[str, str]],

209 ):

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

211 raise ValueError(

@@ -223,7 +222,7 @@ def explore_combs_workloads(

222 server_ready_timeout=server_ready_timeout,

223 serve_comb=serve_comb,

224 bench_params=bench_params,

226 - output_dir=output_dir,

225 + experiment_dir=experiment_dir,

226 dry_run=dry_run,

227 ) as server:

228 for bench_comb in bench_params:

@@ -232,12 +231,12 @@ def explore_combs_workloads(

231 bench_cmd,

232 serve_comb=serve_comb,

233 bench_comb=bench_comb,

234 + link_vars=link_vars,

235 workload_var=workload_var,

236 workload_iters=workload_iters,

237 - output_dir=output_dir,

237 + experiment_dir=experiment_dir,

238 num_runs=num_runs,

239 dry_run=dry_run,

240 - link_vars=link_vars,

240 )

241

242 if comb_data is not None:

@@ -247,7 +246,7 @@ def explore_combs_workloads(

246 return None

247

248 combined_df = pd.DataFrame.from_records(all_data)

250 - combined_df.to_csv(output_dir / "summary.csv")

249 + combined_df.to_csv(experiment_dir / "summary.csv")

250

251 return combined_df

252

@@ -298,13 +297,9 @@ class SweepServeWorkloadArgs(SweepServeArgs):

297

298

299 def run_main(args: SweepServeWorkloadArgs):

301 - timestamp = args.resume or datetime.now().strftime("%Y%m%d_%H%M%S")

302 - output_dir = args.output_dir / timestamp

303 -

304 - if args.resume and not output_dir.exists():

305 - raise ValueError(f"Cannot resume from non-existent directory ({output_dir})")

300 + experiment_dir = args.resolve_experiment_dir()

301

307 - try:

302 + with args.run_ctx(experiment_dir):

303 return explore_combs_workloads(

304 serve_cmd=args.serve_cmd,

305 bench_cmd=args.bench_cmd,

@@ -313,18 +308,13 @@ def run_main(args: SweepServeWorkloadArgs):

308 server_ready_timeout=args.server_ready_timeout,

309 serve_params=args.serve_params,

310 bench_params=args.bench_params,

311 + link_vars=args.link_vars,

312 workload_var=args.workload_var,

313 workload_iters=args.workload_iters,

318 - output_dir=output_dir,

314 + experiment_dir=experiment_dir,

315 num_runs=args.num_runs,

316 dry_run=args.dry_run,

321 - link_vars=args.link_vars,

317 )

323 - except BaseException as exc:

324 - raise RuntimeError(

325 - f"The script was terminated early. Use `--resume {timestamp}` "

326 - f"to continue the script from its last checkpoint."

327 - ) from exc

318

319

320 def main(args: argparse.Namespace):

vllm/benchmarks/sweep/startup.py

+70 -33

@@ -4,6 +4,7 @@ import argparse

4 import json

5 import shlex

6 import subprocess

7 +from contextlib import contextmanager

8 from dataclasses import dataclass

9 from datetime import datetime

10 from functools import lru_cache

@@ -111,7 +112,7 @@ def _apply_output_json(cmd: list[str], output_path: Path) -> list[str]:

112

113

114 def _get_comb_base_path(

114 - output_dir: Path,

115 + experiment_dir: Path,

116 serve_comb: ParameterSweepItem,

117 startup_comb: ParameterSweepItem,

118 ) -> Path:

@@ -120,7 +121,8 @@ def _get_comb_base_path(

121 parts.extend(("SERVE-", serve_comb.name))

122 if startup_comb:

123 parts.extend(("STARTUP-", startup_comb.name))

123 - return output_dir / sanitize_filename("-".join(parts))

124 +

125 + return experiment_dir / sanitize_filename("-".join(parts))

126

127

128 def _get_comb_run_path(base_path: Path, run_number: int | None) -> Path:

@@ -225,7 +227,7 @@ def run_combs(

227 *,

228 serve_params: ParameterSweep,

229 startup_params: ParameterSweep,

228 - output_dir: Path,

230 + experiment_dir: Path,

231 num_runs: int,

232 show_stdout: bool,

233 dry_run: bool,

@@ -233,7 +235,7 @@ def run_combs(

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

236 for serve_comb in serve_params:

237 for startup_comb in startup_params:

236 - base_path = _get_comb_base_path(output_dir, serve_comb, startup_comb)

238 + base_path = _get_comb_base_path(experiment_dir, serve_comb, startup_comb)

239 comb_data = run_comb(

240 startup_cmd,

241 serve_comb=serve_comb,

@@ -250,7 +252,7 @@ def run_combs(

252 return None

253

254 combined_df = pd.DataFrame.from_records(all_data)

253 - combined_df.to_csv(output_dir / "summary.csv")

255 + combined_df.to_csv(experiment_dir / "summary.csv")

256 return combined_df

257

258

@@ -260,11 +262,11 @@ class SweepStartupArgs:

262 serve_params: ParameterSweep

263 startup_params: ParameterSweep

264 output_dir: Path

265 + experiment_name: str

266 num_runs: int

267 show_stdout: bool

268 dry_run: bool

266 - resume: str | None

267 - strict_params: bool

269 + resume: bool

270

271 parser_name: ClassVar[str] = "startup"

272 parser_help: ClassVar[str] = (

@@ -286,13 +288,19 @@ class SweepStartupArgs:

288 startup_params = ParameterSweep.from_records([{}])

289

290 supported = _get_supported_startup_keys()

291 + strict_params = args.strict_params

292 serve_params = _filter_params(

290 - serve_params, supported=supported, strict=args.strict_params

293 + serve_params, supported=supported, strict=strict_params

294 )

295 startup_params = _filter_params(

293 - startup_params, supported=supported, strict=args.strict_params

296 + startup_params, supported=supported, strict=strict_params

297 )

298

299 + if args.experiment_name:

300 + experiment_name = args.experiment_name

301 + else:

302 + experiment_name = datetime.now().strftime("%Y%m%d_%H%M%S")

303 +

304 if args.num_runs < 1:

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

306

@@ -301,11 +309,11 @@ class SweepStartupArgs:

309 serve_params=serve_params,

310 startup_params=startup_params,

311 output_dir=Path(args.output_dir),

312 + experiment_name=experiment_name,

313 num_runs=args.num_runs,

314 show_stdout=args.show_stdout,

315 dry_run=args.dry_run,

316 resume=args.resume,

308 - strict_params=args.strict_params,

317 )

318

319 @classmethod

@@ -316,6 +324,7 @@ class SweepStartupArgs:

324 default="vllm bench startup",

325 help="The command used to run the startup benchmark.",

326 )

327 +

328 parser.add_argument(

329 "--serve-params",

330 type=str,

@@ -331,12 +340,27 @@ class SweepStartupArgs:

340 help="Path to JSON file containing parameter combinations "

341 "for the `vllm bench startup` command.",

342 )

343 + parser.add_argument(

344 + "--strict-params",

345 + action="store_true",

346 + help="If set, unknown parameters in sweep files raise an error "

347 + "instead of being ignored.",

348 + )

349 +

350 parser.add_argument(

351 "-o",

352 "--output-dir",

353 type=str,

354 default="results",

339 - help="The directory to which results are written.",

355 + help="The main directory to which results are written.",

356 + )

357 + parser.add_argument(

358 + "-e",

359 + "--experiment-name",

360 + type=str,

361 + default=None,

362 + help="The name of this experiment (defaults to current timestamp). "

363 + "Results will be stored under `output_dir/experiment_name`.",

364 )

365 parser.add_argument(

366 "--num-runs",

@@ -357,43 +381,56 @@ class SweepStartupArgs:

381 )

382 parser.add_argument(

383 "--resume",

360 - type=str,

361 - default=None,

362 - help="Set this to the name of a directory under `output_dir` (which is a "

363 - "timestamp) to resume a previous execution of this script, i.e., only run "

364 - "parameter combinations for which there are still no output files.",

365 - )

366 - parser.add_argument(

367 - "--strict-params",

384 action="store_true",

369 - help="If set, unknown parameters in sweep files raise an error "

370 - "instead of being ignored.",

385 + help="Resume a previous execution of this script, i.e., only run "

386 + "parameter combinations for which there are still no output files "

387 + "under `output_dir/experiment_name`.",

388 )

389 +

390 return parser

391

392 + def resolve_experiment_dir(self) -> Path:

393 + experiment_dir = self.output_dir / self.experiment_name

394

375 -def run_main(args: SweepStartupArgs):

376 - timestamp = args.resume or datetime.now().strftime("%Y%m%d_%H%M%S")

377 - output_dir = args.output_dir / timestamp

395 + if self.resume:

396 + if not experiment_dir.exists():

397 + raise ValueError(f"Cannot resume from non-existent {experiment_dir=}")

398 + else:

399 + if experiment_dir.exists():

400 + raise ValueError(f"Cannot overwrite existing {experiment_dir=}")

401 +

402 + return experiment_dir

403 +

404 + @contextmanager

405 + def run_ctx(self, experiment_dir: Path):

406 + if self.dry_run:

407 + yield

408 + print(f"Experiment will be saved at: {experiment_dir}")

409 + return

410

379 - if args.resume and not output_dir.exists():

380 - raise ValueError(f"Cannot resume from non-existent directory ({output_dir})")

411 + try:

412 + yield

413 + print(f"Experiment has been saved at: {experiment_dir}")

414 + except BaseException as exc:

415 + raise RuntimeError(

416 + "The script was terminated early. Use `--resume` "

417 + "to continue the script from its last checkpoint."

418 + ) from exc

419 +

420 +

421 +def run_main(args: SweepStartupArgs):

422 + experiment_dir = args.resolve_experiment_dir()

423

382 - try:

424 + with args.run_ctx(experiment_dir):

425 return run_combs(

426 startup_cmd=args.startup_cmd,

427 serve_params=args.serve_params,

428 startup_params=args.startup_params,

387 - output_dir=output_dir,

429 + experiment_dir=experiment_dir,

430 num_runs=args.num_runs,

431 show_stdout=args.show_stdout,

432 dry_run=args.dry_run,

433 )

392 - except BaseException as exc:

393 - raise RuntimeError(

394 - f"The script was terminated early. Use `--resume {timestamp}` "

395 - f"to continue the script from its last checkpoint."

396 - ) from exc

434

435

436 def main(args: argparse.Namespace):