| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Generate optional vLLM benchmark sweep files from one initial suggestion.""" |
| 5 | |
| 6 | from __future__ import annotations |
| 7 | |
| 8 | import json |
| 9 | import os |
| 10 | import shlex |
| 11 | from pathlib import Path |
| 12 | from typing import Any |
| 13 | |
| 14 | from runtime_tuning import WorkloadHints |
| 15 | |
| 16 | |
| 17 | def _positive_int(config: dict[str, Any], key: str) -> int: |
| 18 | value = config.get(key) |
| 19 | if not isinstance(value, int) or isinstance(value, bool) or value <= 0: |
| 20 | raise ValueError(f"Sweep generation requires a positive initial {key!r} value.") |
| 21 | return value |
| 22 | |
| 23 | |
| 24 | def _strict_lower_power_of_two(value: int) -> int: |
| 25 | if value <= 1: |
| 26 | return 1 |
| 27 | return 1 << ((value - 1).bit_length() - 1) |
| 28 | |
| 29 | |
| 30 | def _strict_upper_power_of_two(value: int) -> int: |
| 31 | return 1 << value.bit_length() |
| 32 | |
| 33 | |
| 34 | def validate_sweep_workload(workload: WorkloadHints) -> None: |
| 35 | required = { |
| 36 | "--input-tokens": workload.input_tokens, |
| 37 | "--output-tokens": workload.output_tokens, |
| 38 | "--concurrency": workload.concurrency, |
| 39 | } |
| 40 | missing = [name for name, value in required.items() if value is None] |
| 41 | if missing: |
| 42 | raise ValueError("--generate-sweep requires " + ", ".join(missing) + ".") |
| 43 | |
| 44 | |
| 45 | def build_serve_params(config: dict[str, Any]) -> list[dict[str, Any]]: |
| 46 | """Build a diverse, bounded sweep around the initial suggestion. |
| 47 | |
| 48 | The initial five-point sweep changed only one parameter at a time and |
| 49 | skipped the useful middle sequence count. Eight directed points cover the |
| 50 | batch-budget curve at full concurrency plus interactions at 3/4 and 1/2 of |
| 51 | the initial scheduler concurrency without paying for a full Cartesian grid. |
| 52 | """ |
| 53 | initial_seqs = _positive_int(config, "max-num-seqs") |
| 54 | initial_batch = _positive_int(config, "max-num-batched-tokens") |
| 55 | |
| 56 | minimum_batch = initial_seqs |
| 57 | if config.get("enable-chunked-prefill") is False: |
| 58 | max_model_len = config.get("max-model-len") |
| 59 | if ( |
| 60 | isinstance(max_model_len, int) |
| 61 | and not isinstance(max_model_len, bool) |
| 62 | and max_model_len > 0 |
| 63 | ): |
| 64 | minimum_batch = max(minimum_batch, max_model_len) |
| 65 | |
| 66 | lower_seqs = max(1, (initial_seqs + 1) // 2) |
| 67 | middle_seqs = max(lower_seqs, (3 * initial_seqs + 3) // 4) |
| 68 | lower_batch = max( |
| 69 | minimum_batch, |
| 70 | _strict_lower_power_of_two(initial_batch), |
| 71 | ) |
| 72 | smaller_batch = max( |
| 73 | minimum_batch, |
| 74 | _strict_lower_power_of_two(lower_batch), |
| 75 | ) |
| 76 | higher_batch = max( |
| 77 | minimum_batch, |
| 78 | _strict_upper_power_of_two(initial_batch), |
| 79 | ) |
| 80 | |
| 81 | candidates: list[dict[str, Any]] = [] |
| 82 | seen: set[tuple[int, int]] = set() |
| 83 | |
| 84 | def add(name: str, max_num_seqs: int, max_num_batched_tokens: int) -> None: |
| 85 | max_num_batched_tokens = max( |
| 86 | max_num_batched_tokens, |
| 87 | max_num_seqs, |
| 88 | ) |
| 89 | signature = (max_num_seqs, max_num_batched_tokens) |
| 90 | if signature in seen: |
| 91 | return |
| 92 | seen.add(signature) |
| 93 | candidates.append( |
| 94 | { |
| 95 | "_benchmark_name": name, |
| 96 | "max_num_seqs": max_num_seqs, |
| 97 | "max_num_batched_tokens": max_num_batched_tokens, |
| 98 | } |
| 99 | ) |
| 100 | |
| 101 | # Keep the exact initial suggestion as the measured baseline. Additional |
| 102 | # values exist only in the optional sweep package. |
| 103 | add("initial", initial_seqs, initial_batch) |
| 104 | add("smaller_batch_budget", initial_seqs, smaller_batch) |
| 105 | add("lower_batch_budget", initial_seqs, lower_batch) |
| 106 | add("higher_batch_budget", initial_seqs, higher_batch) |
| 107 | add("middle_seqs_lower_batch", middle_seqs, lower_batch) |
| 108 | add("middle_seqs_higher_batch", middle_seqs, higher_batch) |
| 109 | add("lower_seqs_lower_batch", lower_seqs, lower_batch) |
| 110 | add("lower_seqs_higher_batch", lower_seqs, higher_batch) |
| 111 | |
| 112 | return candidates |
| 113 | |
| 114 | |
| 115 | def build_bench_params(workload: WorkloadHints) -> list[dict[str, Any]]: |
| 116 | validate_sweep_workload(workload) |
| 117 | assert workload.input_tokens is not None |
| 118 | assert workload.output_tokens is not None |
| 119 | assert workload.concurrency is not None |
| 120 | |
| 121 | return [ |
| 122 | { |
| 123 | "_benchmark_name": "user_workload", |
| 124 | "random_input_len": workload.input_tokens, |
| 125 | "random_output_len": workload.output_tokens, |
| 126 | "max_concurrency": workload.concurrency, |
| 127 | } |
| 128 | ] |
| 129 | |
| 130 | |
| 131 | def _benchmark_models(config: dict[str, Any]) -> tuple[str, str]: |
| 132 | model = config.get("model") |
| 133 | if not isinstance(model, str) or not model: |
| 134 | raise ValueError("Sweep generation requires a model in config.yml.") |
| 135 | |
| 136 | served_model_name = config.get("served-model-name") |
| 137 | if isinstance(served_model_name, str) and served_model_name: |
| 138 | return served_model_name, model |
| 139 | if ( |
| 140 | isinstance(served_model_name, list) |
| 141 | and served_model_name |
| 142 | and isinstance(served_model_name[0], str) |
| 143 | ): |
| 144 | return served_model_name[0], model |
| 145 | |
| 146 | return model, model |
| 147 | |
| 148 | |
| 149 | def _relative_to(directory: Path, target: str) -> str: |
| 150 | return os.path.relpath(Path(target).resolve(), directory.resolve()) |
| 151 | |
| 152 | |
| 153 | def _write_json(path: Path, value: object) -> None: |
| 154 | path.write_text( |
| 155 | json.dumps(value, indent=2) + "\n", |
| 156 | encoding="utf-8", |
| 157 | ) |
| 158 | |
| 159 | |
| 160 | def _write_run_script( |
| 161 | path: Path, |
| 162 | *, |
| 163 | config_rel: str, |
| 164 | env_rel: str, |
| 165 | request_model: str, |
| 166 | tokenizer: str, |
| 167 | workload: WorkloadHints, |
| 168 | ) -> None: |
| 169 | bench_parts = [ |
| 170 | "vllm bench serve", |
| 171 | "--backend vllm", |
| 172 | f"--model {shlex.quote(request_model)}", |
| 173 | f"--tokenizer {shlex.quote(tokenizer)}", |
| 174 | "--dataset-name random", |
| 175 | "--request-rate inf", |
| 176 | "--ignore-eos", |
| 177 | "--metric-percentiles 99", |
| 178 | ] |
| 179 | |
| 180 | goodput_pairs: list[str] = [] |
| 181 | if workload.ttft_sla_ms is not None: |
| 182 | goodput_pairs.append(f"ttft:{workload.ttft_sla_ms:g}") |
| 183 | if workload.tpot_sla_ms is not None: |
| 184 | goodput_pairs.append(f"tpot:{workload.tpot_sla_ms:g}") |
| 185 | if goodput_pairs: |
| 186 | bench_parts.append( |
| 187 | "--goodput " + " ".join(shlex.quote(pair) for pair in goodput_pairs) |
| 188 | ) |
| 189 | |
| 190 | bench_cmd = " ".join(bench_parts) |
| 191 | script = f"""#!/usr/bin/env bash |
| 192 | set -euo pipefail |
| 193 | |
| 194 | SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" |
| 195 | CONFIG_PATH="${{SCRIPT_DIR}}/{config_rel}" |
| 196 | ENV_PATH="${{SCRIPT_DIR}}/{env_rel}" |
| 197 | |
| 198 | source "${{ENV_PATH}}" |
| 199 | |
| 200 | vllm bench sweep serve \ |
| 201 | --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ |
| 202 | --bench-cmd "{bench_cmd}" \ |
| 203 | --serve-params "${{SCRIPT_DIR}}/serve_params.json" \ |
| 204 | --bench-params "${{SCRIPT_DIR}}/bench_params.json" \ |
| 205 | --output-dir "${{SCRIPT_DIR}}/results" \ |
| 206 | --experiment-name runtime-tuning \ |
| 207 | "$@" |
| 208 | """ |
| 209 | path.write_text(script, encoding="utf-8") |
| 210 | path.chmod(path.stat().st_mode | 0o111) |
| 211 | |
| 212 | |
| 213 | def _write_recommend_script( |
| 214 | path: Path, |
| 215 | *, |
| 216 | config_rel: str, |
| 217 | env_rel: str, |
| 218 | workload: WorkloadHints, |
| 219 | ) -> None: |
| 220 | template_path = Path(__file__).with_name("sweep_recommendation.py") |
| 221 | source = template_path.read_text(encoding="utf-8") |
| 222 | |
| 223 | replacements = { |
| 224 | "DEFAULT_CONFIG_PATH: str | None = None": ( |
| 225 | f"DEFAULT_CONFIG_PATH: str | None = {config_rel!r}" |
| 226 | ), |
| 227 | "DEFAULT_ENV_PATH: str | None = None": ( |
| 228 | f"DEFAULT_ENV_PATH: str | None = {env_rel!r}" |
| 229 | ), |
| 230 | "DEFAULT_TTFT_SLA_MS: float | None = None": ( |
| 231 | f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" |
| 232 | ), |
| 233 | "DEFAULT_TPOT_SLA_MS: float | None = None": ( |
| 234 | f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" |
| 235 | ), |
| 236 | } |
| 237 | for marker, replacement in replacements.items(): |
| 238 | if marker not in source: |
| 239 | raise ValueError(f"Recommender template marker not found: {marker}") |
| 240 | source = source.replace(marker, replacement, 1) |
| 241 | |
| 242 | path.write_text(source, encoding="utf-8") |
| 243 | path.chmod(path.stat().st_mode | 0o111) |
| 244 | |
| 245 | |
| 246 | def _write_guide(path: Path, workload: WorkloadHints) -> None: |
| 247 | sla_lines: list[str] = [] |
| 248 | if workload.ttft_sla_ms is not None: |
| 249 | sla_lines.append(f"- TTFT objective: `{workload.ttft_sla_ms:g} ms`") |
| 250 | if workload.tpot_sla_ms is not None: |
| 251 | sla_lines.append(f"- TPOT objective: `{workload.tpot_sla_ms:g} ms`") |
| 252 | |
| 253 | sla_text = "\n".join(sla_lines) |
| 254 | if sla_text: |
| 255 | sla_text = ( |
| 256 | "\n## Supplied latency objectives\n\n" |
| 257 | + sla_text |
| 258 | + "\n\nThe generated benchmark uses these values with vLLM " |
| 259 | "`--goodput`.\n" |
| 260 | ) |
| 261 | |
| 262 | content = f"""# Optional Runtime Tuning Sweep |
| 263 | |
| 264 | The generated `config.yml` is the **single initial suggestion** and can be |
| 265 | deployed directly. The sweep benchmarks nearby values for: |
| 266 | |
| 267 | - `max-num-seqs` |
| 268 | - `max-num-batched-tokens` |
| 269 | |
| 270 | ## Run |
| 271 | |
| 272 | ```bash |
| 273 | ./run_sweep.sh --dry-run |
| 274 | ./run_sweep.sh |
| 275 | ./recommend.py |
| 276 | ``` |
| 277 | |
| 278 | For a quick one-run experiment: |
| 279 | |
| 280 | ```bash |
| 281 | rm -rf results/runtime-tuning |
| 282 | ./run_sweep.sh --num-runs 1 |
| 283 | ./recommend.py |
| 284 | ``` |
| 285 | |
| 286 | Resume an interrupted sweep: |
| 287 | |
| 288 | ```bash |
| 289 | ./run_sweep.sh --resume |
| 290 | ``` |
| 291 | |
| 292 | By default, vLLM benchmarks each parameter combination three times. |
| 293 | {sla_text} |
| 294 | ## Outputs |
| 295 | |
| 296 | `recommend.py` writes: |
| 297 | |
| 298 | ```text |
| 299 | recommended-config.yml |
| 300 | recommendation.json |
| 301 | ``` |
| 302 | |
| 303 | With TTFT/TPOT objectives it requires duration-weighted combined compliance of |
| 304 | at least 99% plus median P99 TTFT/TPOT compliance, then selects highest mean |
| 305 | output-token throughput. If no candidate qualifies, it records a best-effort |
| 306 | candidate but does not write `recommended-config.yml`. Without latency |
| 307 | objectives it selects highest mean output-token throughput. Configurations with |
| 308 | failed requests are excluded. |
| 309 | |
| 310 | `recommended-config.yml` copies the initial configuration and changes only |
| 311 | `max-num-seqs` and `max-num-batched-tokens`. |
| 312 | """ |
| 313 | path.write_text(content, encoding="utf-8") |
| 314 | |
| 315 | |
| 316 | def write_sweep_files( |
| 317 | output_dir: str, |
| 318 | *, |
| 319 | config_path: str, |
| 320 | env_path: str, |
| 321 | config: dict[str, Any], |
| 322 | workload: WorkloadHints, |
| 323 | ) -> list[Path]: |
| 324 | """Write an optional sweep package around the initial config suggestion.""" |
| 325 | validate_sweep_workload(workload) |
| 326 | |
| 327 | directory = Path(output_dir) |
| 328 | directory.mkdir(parents=True, exist_ok=True) |
| 329 | |
| 330 | serve_params = directory / "serve_params.json" |
| 331 | bench_params = directory / "bench_params.json" |
| 332 | run_script = directory / "run_sweep.sh" |
| 333 | recommend_script = directory / "recommend.py" |
| 334 | guide = directory / "SWEEP.md" |
| 335 | |
| 336 | config_rel = _relative_to(directory, config_path) |
| 337 | env_rel = _relative_to(directory, env_path) |
| 338 | |
| 339 | _write_json(serve_params, build_serve_params(config)) |
| 340 | _write_json(bench_params, build_bench_params(workload)) |
| 341 | request_model, tokenizer = _benchmark_models(config) |
| 342 | _write_run_script( |
| 343 | run_script, |
| 344 | config_rel=config_rel, |
| 345 | env_rel=env_rel, |
| 346 | request_model=request_model, |
| 347 | tokenizer=tokenizer, |
| 348 | workload=workload, |
| 349 | ) |
| 350 | _write_recommend_script( |
| 351 | recommend_script, |
| 352 | config_rel=config_rel, |
| 353 | env_rel=env_rel, |
| 354 | workload=workload, |
| 355 | ) |
| 356 | _write_guide(guide, workload) |
| 357 | |
| 358 | return [serve_params, bench_params, run_script, recommend_script, guide] |