| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark the cold and warm startup time of vLLM models. |
| 4 | |
| 5 | This script measures total startup time (including model loading, compilation, |
| 6 | and cache operations) for both cold and warm scenarios: |
| 7 | - Cold startup: Fresh start with no caches (temporary cache directories) |
| 8 | - Warm startup: Using cached compilation and model info |
| 9 | """ |
| 10 | |
| 11 | import argparse |
| 12 | import json |
| 13 | import multiprocessing |
| 14 | import os |
| 15 | import shutil |
| 16 | import tempfile |
| 17 | import time |
| 18 | from contextlib import contextmanager |
| 19 | from typing import Any, NamedTuple |
| 20 | |
| 21 | import numpy as np |
| 22 | from tqdm import tqdm |
| 23 | |
| 24 | from vllm.benchmarks.lib.utils import ( |
| 25 | convert_to_pytorch_benchmark_format, |
| 26 | write_to_json, |
| 27 | ) |
| 28 | from vllm.engine.arg_utils import EngineArgs |
| 29 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 30 | |
| 31 | PERCENTAGES = [10, 25, 50, 75, 90, 99] |
| 32 | |
| 33 | |
| 34 | class MetricDesc(NamedTuple): |
| 35 | """Descriptor for a metric to collect from each iteration.""" |
| 36 | |
| 37 | iter_key: str # key in the iteration result dict |
| 38 | suffix: str # result key suffix, e.g. "startup", "compilation" |
| 39 | display_name: str |
| 40 | |
| 41 | |
| 42 | class MetricStats(NamedTuple): |
| 43 | """Aggregated statistics for a single benchmark metric.""" |
| 44 | |
| 45 | key: str # e.g. "cold_startup", "warm_encoder_compilation" |
| 46 | display_name: str |
| 47 | values: list[float] |
| 48 | avg: float |
| 49 | percentiles: dict[int, float] |
| 50 | |
| 51 | |
| 52 | _BASE_METRICS = [ |
| 53 | MetricDesc("total_startup_time", "startup", "Startup time"), |
| 54 | MetricDesc("compilation_time", "compilation", "Compilation time"), |
| 55 | ] |
| 56 | _ENCODER_METRIC = MetricDesc( |
| 57 | "encoder_compilation_time", |
| 58 | "encoder_compilation", |
| 59 | "Encoder compilation time", |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def _compute_metric( |
| 64 | phase: str, |
| 65 | desc: MetricDesc, |
| 66 | iterations: list[dict[str, float]], |
| 67 | ) -> MetricStats: |
| 68 | values = [m[desc.iter_key] for m in iterations] |
| 69 | arr = np.array(values) |
| 70 | return MetricStats( |
| 71 | key=f"{phase}_{desc.suffix}", |
| 72 | display_name=desc.display_name, |
| 73 | values=values, |
| 74 | avg=float(np.mean(arr)), |
| 75 | percentiles=dict(zip(PERCENTAGES, np.percentile(arr, PERCENTAGES).tolist())), |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | def _collect_phase_metrics( |
| 80 | phase: str, |
| 81 | iterations: list[dict[str, float]], |
| 82 | has_encoder: bool, |
| 83 | ) -> list[MetricStats]: |
| 84 | metrics = [_compute_metric(phase, desc, iterations) for desc in _BASE_METRICS] |
| 85 | if has_encoder: |
| 86 | metrics.append(_compute_metric(phase, _ENCODER_METRIC, iterations)) |
| 87 | return metrics |
| 88 | |
| 89 | |
| 90 | def _print_phase(phase_name: str, metrics: list[MetricStats]) -> None: |
| 91 | print(f"\n{phase_name}:") |
| 92 | for m in metrics: |
| 93 | print(f"Avg {m.display_name.lower()}: {m.avg:.2f} seconds") |
| 94 | for m in metrics: |
| 95 | print(f"{m.display_name} percentiles:") |
| 96 | for pct, val in m.percentiles.items(): |
| 97 | print(f" {pct}%: {val:.2f} seconds") |
| 98 | |
| 99 | |
| 100 | def _metric_to_json(m: MetricStats) -> dict[str, Any]: |
| 101 | return { |
| 102 | f"avg_{m.key}_time": m.avg, |
| 103 | f"{m.key}_times": m.values, |
| 104 | f"{m.key}_percentiles": m.percentiles, |
| 105 | } |
| 106 | |
| 107 | |
| 108 | @contextmanager |
| 109 | def cold_startup(): |
| 110 | """Context manager to measure cold startup time: |
| 111 | 1. Uses a temporary directory for vLLM cache to avoid any pollution |
| 112 | between cold startup iterations. |
| 113 | 2. Uses inductor's fresh_cache to clear torch.compile caches. |
| 114 | """ |
| 115 | from torch._inductor.utils import fresh_cache |
| 116 | |
| 117 | # Use temporary directory for caching to avoid any pollution between cold startups |
| 118 | original_cache_root = os.environ.get("VLLM_CACHE_ROOT") |
| 119 | temp_cache_dir = tempfile.mkdtemp(prefix="vllm_startup_bench_cold_") |
| 120 | try: |
| 121 | os.environ["VLLM_CACHE_ROOT"] = temp_cache_dir |
| 122 | with fresh_cache(): |
| 123 | yield |
| 124 | finally: |
| 125 | # Clean up temporary cache directory |
| 126 | shutil.rmtree(temp_cache_dir, ignore_errors=True) |
| 127 | if original_cache_root: |
| 128 | os.environ["VLLM_CACHE_ROOT"] = original_cache_root |
| 129 | else: |
| 130 | os.environ.pop("VLLM_CACHE_ROOT", None) |
| 131 | |
| 132 | |
| 133 | def run_startup_in_subprocess(engine_args, result_queue): |
| 134 | """Run LLM startup in a subprocess and return timing metrics via a queue. |
| 135 | This ensures complete isolation between iterations. |
| 136 | """ |
| 137 | try: |
| 138 | # Import inside the subprocess to avoid issues with forking |
| 139 | from vllm import LLM |
| 140 | |
| 141 | # Measure total startup time |
| 142 | start_time = time.perf_counter() |
| 143 | |
| 144 | llm = LLM.from_engine_args(engine_args) |
| 145 | |
| 146 | total_startup_time = time.perf_counter() - start_time |
| 147 | |
| 148 | # Extract compilation time if available |
| 149 | compilation_time = 0.0 |
| 150 | encoder_compilation_time = 0.0 |
| 151 | if hasattr(llm.llm_engine, "vllm_config"): |
| 152 | vllm_config = llm.llm_engine.vllm_config |
| 153 | if ( |
| 154 | hasattr(vllm_config, "compilation_config") |
| 155 | and vllm_config.compilation_config is not None |
| 156 | ): |
| 157 | compilation_time = vllm_config.compilation_config.compilation_time |
| 158 | encoder_compilation_time = ( |
| 159 | vllm_config.compilation_config.encoder_compilation_time |
| 160 | ) |
| 161 | |
| 162 | result_queue.put( |
| 163 | { |
| 164 | "total_startup_time": total_startup_time, |
| 165 | "compilation_time": compilation_time, |
| 166 | "encoder_compilation_time": encoder_compilation_time, |
| 167 | } |
| 168 | ) |
| 169 | |
| 170 | except Exception as e: |
| 171 | result_queue.put(None) |
| 172 | result_queue.put(str(e)) |
| 173 | |
| 174 | |
| 175 | def save_to_pytorch_benchmark_format( |
| 176 | args: argparse.Namespace, metrics: list[MetricStats] |
| 177 | ) -> None: |
| 178 | base_name = os.path.splitext(args.output_json)[0] |
| 179 | for m in metrics: |
| 180 | records = convert_to_pytorch_benchmark_format( |
| 181 | args=args, |
| 182 | metrics={f"avg_{m.key}_time": [m.avg]}, |
| 183 | extra_info={ |
| 184 | f"{m.key}_times": m.values, |
| 185 | f"{m.key}_percentiles": m.percentiles, |
| 186 | }, |
| 187 | ) |
| 188 | if records: |
| 189 | write_to_json(f"{base_name}.{m.key}.pytorch.json", records) |
| 190 | |
| 191 | |
| 192 | def add_cli_args(parser: FlexibleArgumentParser): |
| 193 | parser.add_argument( |
| 194 | "--num-iters-cold", |
| 195 | type=int, |
| 196 | default=3, |
| 197 | help="Number of cold startup iterations.", |
| 198 | ) |
| 199 | parser.add_argument( |
| 200 | "--num-iters-warmup", |
| 201 | type=int, |
| 202 | default=1, |
| 203 | help="Number of warmup iterations before benchmarking warm startups.", |
| 204 | ) |
| 205 | parser.add_argument( |
| 206 | "--num-iters-warm", |
| 207 | type=int, |
| 208 | default=3, |
| 209 | help="Number of warm startup iterations.", |
| 210 | ) |
| 211 | parser.add_argument( |
| 212 | "--output-json", |
| 213 | type=str, |
| 214 | default=None, |
| 215 | help="Path to save the startup time results in JSON format.", |
| 216 | ) |
| 217 | |
| 218 | parser = EngineArgs.add_cli_args(parser) |
| 219 | return parser |
| 220 | |
| 221 | |
| 222 | def main(args: argparse.Namespace): |
| 223 | # Set multiprocessing start method to 'spawn' for clean process isolation |
| 224 | # This ensures each subprocess starts fresh without inheriting state |
| 225 | multiprocessing.set_start_method("spawn", force=True) |
| 226 | |
| 227 | engine_args = EngineArgs.from_cli_args(args) |
| 228 | |
| 229 | def create_llm_and_measure_startup(): |
| 230 | """Create LLM instance in a subprocess and measure startup time. |
| 231 | Returns timing metrics, using subprocess for complete isolation. |
| 232 | """ |
| 233 | # Create a queue for inter-process communication |
| 234 | result_queue: multiprocessing.Queue[Any] = multiprocessing.Queue() |
| 235 | process = multiprocessing.Process( |
| 236 | target=run_startup_in_subprocess, |
| 237 | args=( |
| 238 | engine_args, |
| 239 | result_queue, |
| 240 | ), |
| 241 | ) |
| 242 | process.start() |
| 243 | process.join() |
| 244 | |
| 245 | if not result_queue.empty(): |
| 246 | result = result_queue.get() |
| 247 | if result is None: |
| 248 | if not result_queue.empty(): |
| 249 | error_msg = result_queue.get() |
| 250 | raise RuntimeError(f"Subprocess failed: {error_msg}") |
| 251 | else: |
| 252 | raise RuntimeError("Subprocess failed with unknown error") |
| 253 | return result |
| 254 | else: |
| 255 | raise RuntimeError("Subprocess did not return a result") |
| 256 | |
| 257 | os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" |
| 258 | print("Setting VLLM_ENABLE_V1_MULTIPROCESSING=0 to collect startup metrics.\n") |
| 259 | |
| 260 | # Collect cold startup iterations |
| 261 | print("Measuring cold startup time...\n") |
| 262 | cold_iterations = [] |
| 263 | for i in tqdm(range(args.num_iters_cold), desc="Cold startup iterations"): |
| 264 | with cold_startup(): |
| 265 | cold_iterations.append(create_llm_and_measure_startup()) |
| 266 | |
| 267 | # Warmup for warm startup |
| 268 | print("\nWarming up for warm startup measurement...\n") |
| 269 | for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"): |
| 270 | create_llm_and_measure_startup() |
| 271 | |
| 272 | # Collect warm startup iterations |
| 273 | print("\nMeasuring warm startup time...\n") |
| 274 | warm_iterations = [] |
| 275 | for i in tqdm(range(args.num_iters_warm), desc="Warm startup iterations"): |
| 276 | warm_iterations.append(create_llm_and_measure_startup()) |
| 277 | |
| 278 | # Determine if encoder compilation occurred in any iteration |
| 279 | has_encoder = any( |
| 280 | m["encoder_compilation_time"] > 0 for m in cold_iterations + warm_iterations |
| 281 | ) |
| 282 | |
| 283 | cold_metrics = _collect_phase_metrics("cold", cold_iterations, has_encoder) |
| 284 | warm_metrics = _collect_phase_metrics("warm", warm_iterations, has_encoder) |
| 285 | all_metrics = cold_metrics + warm_metrics |
| 286 | |
| 287 | # Print results |
| 288 | print("\n" + "=" * 60) |
| 289 | print("STARTUP TIME BENCHMARK RESULTS") |
| 290 | print("=" * 60) |
| 291 | _print_phase("COLD STARTUP", cold_metrics) |
| 292 | _print_phase("WARM STARTUP", warm_metrics) |
| 293 | print("=" * 60) |
| 294 | |
| 295 | # Output JSON results if specified |
| 296 | if args.output_json: |
| 297 | results: dict[str, Any] = {} |
| 298 | for m in all_metrics: |
| 299 | results.update(_metric_to_json(m)) |
| 300 | with open(args.output_json, "w") as f: |
| 301 | json.dump(results, f, indent=4) |
| 302 | save_to_pytorch_benchmark_format(args, all_metrics) |