| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | 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 |
| 11 | from pathlib import Path |
| 12 | from typing import ClassVar |
| 13 | |
| 14 | from vllm.benchmarks.startup import add_cli_args as add_startup_cli_args |
| 15 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 16 | from vllm.utils.import_utils import PlaceholderModule |
| 17 | |
| 18 | from .param_sweep import ParameterSweep, ParameterSweepItem |
| 19 | from .utils import sanitize_filename |
| 20 | |
| 21 | try: |
| 22 | import pandas as pd |
| 23 | except ImportError: |
| 24 | pd = PlaceholderModule("pandas") |
| 25 | |
| 26 | |
| 27 | @lru_cache(maxsize=1) |
| 28 | def _get_supported_startup_keys() -> set[str]: |
| 29 | parser = FlexibleArgumentParser(add_help=False) |
| 30 | add_startup_cli_args(parser) |
| 31 | |
| 32 | supported: set[str] = {"config"} |
| 33 | for action in parser._actions: |
| 34 | if action.dest and action.dest is not argparse.SUPPRESS: |
| 35 | supported.add(action.dest) |
| 36 | for option in action.option_strings: |
| 37 | if option.startswith("--"): |
| 38 | supported.add(option.lstrip("-").replace("-", "_")) |
| 39 | |
| 40 | return supported |
| 41 | |
| 42 | |
| 43 | def _is_supported_param(param_key: str, supported: set[str]) -> bool: |
| 44 | if param_key == "_benchmark_name": |
| 45 | return True |
| 46 | prefix = param_key.split(".", 1)[0] |
| 47 | normalized = prefix.replace("-", "_") |
| 48 | return normalized in supported |
| 49 | |
| 50 | |
| 51 | def _filter_params( |
| 52 | params: ParameterSweep, *, supported: set[str], strict: bool |
| 53 | ) -> ParameterSweep: |
| 54 | filtered = [] |
| 55 | for item in params: |
| 56 | kept: dict[str, object] = {} |
| 57 | dropped: list[str] = [] |
| 58 | for key, value in item.items(): |
| 59 | if _is_supported_param(key, supported): |
| 60 | kept[key] = value |
| 61 | else: |
| 62 | dropped.append(key) |
| 63 | |
| 64 | if dropped: |
| 65 | label = item.get("_benchmark_name") or item.as_text() |
| 66 | message = ( |
| 67 | "Ignoring unsupported startup params" |
| 68 | f"{' for ' + str(label) if label else ''}: " |
| 69 | f"{', '.join(sorted(dropped))}" |
| 70 | ) |
| 71 | if strict: |
| 72 | raise ValueError(message) |
| 73 | print(message) |
| 74 | |
| 75 | filtered.append(ParameterSweepItem.from_record(kept)) |
| 76 | |
| 77 | return ParameterSweep(filtered) |
| 78 | |
| 79 | |
| 80 | def _update_run_data( |
| 81 | run_data: dict[str, object], |
| 82 | serve_overrides: ParameterSweepItem, |
| 83 | startup_overrides: ParameterSweepItem, |
| 84 | run_number: int, |
| 85 | ) -> dict[str, object]: |
| 86 | run_data["run_number"] = run_number |
| 87 | run_data.update(serve_overrides) |
| 88 | run_data.update(startup_overrides) |
| 89 | return run_data |
| 90 | |
| 91 | |
| 92 | def _strip_arg(cmd: list[str], keys: tuple[str, ...]) -> list[str]: |
| 93 | stripped: list[str] = [] |
| 94 | skip_next = False |
| 95 | for arg in cmd: |
| 96 | if skip_next: |
| 97 | skip_next = False |
| 98 | continue |
| 99 | if arg in keys: |
| 100 | skip_next = True |
| 101 | continue |
| 102 | if any(arg.startswith(f"{key}=") for key in keys): |
| 103 | continue |
| 104 | stripped.append(arg) |
| 105 | return stripped |
| 106 | |
| 107 | |
| 108 | def _apply_output_json(cmd: list[str], output_path: Path) -> list[str]: |
| 109 | keys = ("--output-json", "--output_json") |
| 110 | cmd = _strip_arg(cmd, keys) |
| 111 | return [*cmd, keys[0], str(output_path)] |
| 112 | |
| 113 | |
| 114 | def _get_comb_base_path( |
| 115 | experiment_dir: Path, |
| 116 | serve_comb: ParameterSweepItem, |
| 117 | startup_comb: ParameterSweepItem, |
| 118 | ) -> Path: |
| 119 | parts = list[str]() |
| 120 | if serve_comb: |
| 121 | parts.extend(("SERVE-", serve_comb.name)) |
| 122 | if startup_comb: |
| 123 | parts.extend(("STARTUP-", startup_comb.name)) |
| 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: |
| 129 | if run_number is None: |
| 130 | return base_path / "summary.json" |
| 131 | return base_path / f"run={run_number}.json" |
| 132 | |
| 133 | |
| 134 | def run_benchmark( |
| 135 | startup_cmd: list[str], |
| 136 | *, |
| 137 | serve_overrides: ParameterSweepItem, |
| 138 | startup_overrides: ParameterSweepItem, |
| 139 | run_number: int, |
| 140 | output_path: Path, |
| 141 | show_stdout: bool, |
| 142 | dry_run: bool, |
| 143 | ) -> dict[str, object] | None: |
| 144 | cmd = serve_overrides.apply_to_cmd(startup_cmd) |
| 145 | cmd = startup_overrides.apply_to_cmd(cmd) |
| 146 | cmd = _apply_output_json(cmd, output_path) |
| 147 | |
| 148 | print("[BEGIN BENCHMARK]") |
| 149 | print(f"Serve overrides: {serve_overrides}") |
| 150 | print(f"Startup overrides: {startup_overrides}") |
| 151 | print(f"Run Number: {run_number}") |
| 152 | print(f"Benchmark command: {cmd}") |
| 153 | print(f"Output file: {output_path}") |
| 154 | |
| 155 | if output_path.exists(): |
| 156 | print("Found existing results.") |
| 157 | print("[SKIPPED BENCHMARK]") |
| 158 | |
| 159 | with output_path.open("r", encoding="utf-8") as f: |
| 160 | run_data = json.load(f) |
| 161 | return _update_run_data( |
| 162 | run_data, serve_overrides, startup_overrides, run_number |
| 163 | ) |
| 164 | |
| 165 | if dry_run: |
| 166 | print("[END BENCHMARK]") |
| 167 | return None |
| 168 | |
| 169 | output_path.parent.mkdir(parents=True, exist_ok=True) |
| 170 | subprocess.run( |
| 171 | cmd, |
| 172 | stdout=None if show_stdout else subprocess.DEVNULL, |
| 173 | check=True, |
| 174 | ) |
| 175 | |
| 176 | with output_path.open("r", encoding="utf-8") as f: |
| 177 | run_data = json.load(f) |
| 178 | |
| 179 | run_data = _update_run_data( |
| 180 | run_data, serve_overrides, startup_overrides, run_number |
| 181 | ) |
| 182 | |
| 183 | with output_path.open("w", encoding="utf-8") as f: |
| 184 | json.dump(run_data, f, indent=4) |
| 185 | |
| 186 | print("[END BENCHMARK]") |
| 187 | return run_data |
| 188 | |
| 189 | |
| 190 | def run_comb( |
| 191 | startup_cmd: list[str], |
| 192 | *, |
| 193 | serve_comb: ParameterSweepItem, |
| 194 | startup_comb: ParameterSweepItem, |
| 195 | base_path: Path, |
| 196 | num_runs: int, |
| 197 | show_stdout: bool, |
| 198 | dry_run: bool, |
| 199 | ) -> list[dict[str, object]] | None: |
| 200 | comb_data = list[dict[str, object]]() |
| 201 | for run_number in range(num_runs): |
| 202 | run_data = run_benchmark( |
| 203 | startup_cmd, |
| 204 | serve_overrides=serve_comb, |
| 205 | startup_overrides=startup_comb, |
| 206 | run_number=run_number, |
| 207 | output_path=_get_comb_run_path(base_path, run_number), |
| 208 | show_stdout=show_stdout, |
| 209 | dry_run=dry_run, |
| 210 | ) |
| 211 | if run_data is not None: |
| 212 | comb_data.append(run_data) |
| 213 | |
| 214 | if dry_run: |
| 215 | return None |
| 216 | |
| 217 | with _get_comb_run_path(base_path, run_number=None).open( |
| 218 | "w", encoding="utf-8" |
| 219 | ) as f: |
| 220 | json.dump(comb_data, f, indent=4) |
| 221 | |
| 222 | return comb_data |
| 223 | |
| 224 | |
| 225 | def run_combs( |
| 226 | startup_cmd: list[str], |
| 227 | *, |
| 228 | serve_params: ParameterSweep, |
| 229 | startup_params: ParameterSweep, |
| 230 | experiment_dir: Path, |
| 231 | num_runs: int, |
| 232 | show_stdout: bool, |
| 233 | dry_run: bool, |
| 234 | ) -> "pd.DataFrame | None": |
| 235 | all_data = list[dict[str, object]]() |
| 236 | for serve_comb in serve_params: |
| 237 | for startup_comb in startup_params: |
| 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, |
| 242 | startup_comb=startup_comb, |
| 243 | base_path=base_path, |
| 244 | num_runs=num_runs, |
| 245 | show_stdout=show_stdout, |
| 246 | dry_run=dry_run, |
| 247 | ) |
| 248 | if comb_data is not None: |
| 249 | all_data.extend(comb_data) |
| 250 | |
| 251 | if dry_run: |
| 252 | return None |
| 253 | |
| 254 | combined_df = pd.DataFrame.from_records(all_data) |
| 255 | combined_df.to_csv(experiment_dir / "summary.csv") |
| 256 | return combined_df |
| 257 | |
| 258 | |
| 259 | @dataclass |
| 260 | class SweepStartupArgs: |
| 261 | startup_cmd: list[str] |
| 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 |
| 269 | resume: bool |
| 270 | |
| 271 | parser_name: ClassVar[str] = "startup" |
| 272 | parser_help: ClassVar[str] = ( |
| 273 | "Benchmark vLLM startup time over parameter combinations." |
| 274 | ) |
| 275 | |
| 276 | @classmethod |
| 277 | def from_cli_args(cls, args: argparse.Namespace): |
| 278 | startup_cmd = shlex.split(args.startup_cmd) |
| 279 | |
| 280 | if args.serve_params: |
| 281 | serve_params = ParameterSweep.read_json(args.serve_params) |
| 282 | else: |
| 283 | serve_params = ParameterSweep.from_records([{}]) |
| 284 | |
| 285 | if args.startup_params: |
| 286 | startup_params = ParameterSweep.read_json(args.startup_params) |
| 287 | else: |
| 288 | startup_params = ParameterSweep.from_records([{}]) |
| 289 | |
| 290 | supported = _get_supported_startup_keys() |
| 291 | strict_params = args.strict_params |
| 292 | serve_params = _filter_params( |
| 293 | serve_params, supported=supported, strict=strict_params |
| 294 | ) |
| 295 | startup_params = _filter_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 | |
| 307 | return cls( |
| 308 | startup_cmd=startup_cmd, |
| 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, |
| 317 | ) |
| 318 | |
| 319 | @classmethod |
| 320 | def add_cli_args(cls, parser: FlexibleArgumentParser) -> FlexibleArgumentParser: |
| 321 | parser.add_argument( |
| 322 | "--startup-cmd", |
| 323 | type=str, |
| 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 | default=None, |
| 332 | help="Path to JSON file containing parameter combinations " |
| 333 | "for the `vllm serve` command. Only parameters supported by " |
| 334 | "`vllm bench startup` will be applied.", |
| 335 | ) |
| 336 | parser.add_argument( |
| 337 | "--startup-params", |
| 338 | type=str, |
| 339 | default=None, |
| 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", |
| 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", |
| 367 | type=int, |
| 368 | default=1, |
| 369 | help="Number of runs per parameter combination.", |
| 370 | ) |
| 371 | parser.add_argument( |
| 372 | "--show-stdout", |
| 373 | action="store_true", |
| 374 | help="If set, logs the standard output of subcommands.", |
| 375 | ) |
| 376 | parser.add_argument( |
| 377 | "--dry-run", |
| 378 | action="store_true", |
| 379 | help="If set, prints the commands to run, " |
| 380 | "then exits without executing them.", |
| 381 | ) |
| 382 | parser.add_argument( |
| 383 | "--resume", |
| 384 | action="store_true", |
| 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 | |
| 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 | |
| 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 | |
| 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, |
| 429 | experiment_dir=experiment_dir, |
| 430 | num_runs=args.num_runs, |
| 431 | show_stdout=args.show_stdout, |
| 432 | dry_run=args.dry_run, |
| 433 | ) |
| 434 | |
| 435 | |
| 436 | def main(args: argparse.Namespace): |
| 437 | run_main(SweepStartupArgs.from_cli_args(args)) |
| 438 | |
| 439 | |
| 440 | if __name__ == "__main__": |
| 441 | parser = FlexibleArgumentParser(description=SweepStartupArgs.parser_help) |
| 442 | SweepStartupArgs.add_cli_args(parser) |
| 443 | main(parser.parse_args()) |