| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | import math |
| 5 | from dataclasses import dataclass, fields |
| 6 | from pathlib import Path |
| 7 | from typing import ClassVar, Literal, get_args |
| 8 | |
| 9 | import numpy as np |
| 10 | from typing_extensions import assert_never |
| 11 | |
| 12 | from vllm.benchmarks.datasets import DEFAULT_NUM_PROMPTS |
| 13 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 14 | from vllm.utils.import_utils import PlaceholderModule |
| 15 | |
| 16 | from .param_sweep import ParameterSweep, ParameterSweepItem |
| 17 | from .serve import ( |
| 18 | SweepServeArgs, |
| 19 | _get_comb_base_path, |
| 20 | run_comb, |
| 21 | server_ctx, |
| 22 | ) |
| 23 | from .server import ServerProcess |
| 24 | |
| 25 | try: |
| 26 | import pandas as pd |
| 27 | except ImportError: |
| 28 | pd = PlaceholderModule("pandas") |
| 29 | |
| 30 | |
| 31 | WorkloadVariable = Literal["request_rate", "max_concurrency"] |
| 32 | |
| 33 | |
| 34 | def _estimate_workload_value( |
| 35 | run_data: dict[str, object], |
| 36 | workload_var: WorkloadVariable, |
| 37 | ): |
| 38 | request_throughput = float(run_data["request_throughput"]) # type: ignore |
| 39 | if workload_var == "request_rate": |
| 40 | return request_throughput |
| 41 | if workload_var == "max_concurrency": |
| 42 | mean_latency_ms = float(run_data["mean_e2el_ms"]) # type: ignore |
| 43 | return request_throughput * mean_latency_ms / 1000 |
| 44 | |
| 45 | assert_never(workload_var) |
| 46 | |
| 47 | |
| 48 | def _estimate_workload_avg( |
| 49 | runs: list[dict[str, object]], |
| 50 | workload_var: WorkloadVariable, |
| 51 | ): |
| 52 | total = sum(_estimate_workload_value(run, workload_var) for run in runs) |
| 53 | return total / len(runs) |
| 54 | |
| 55 | |
| 56 | def run_comb_workload( |
| 57 | server: ServerProcess | None, |
| 58 | bench_cmd: list[str], |
| 59 | *, |
| 60 | serve_comb: ParameterSweepItem, |
| 61 | bench_comb: ParameterSweepItem, |
| 62 | link_vars: list[tuple[str, str]], |
| 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: |
| 71 | bench_comb_workload = bench_comb | {workload_var: workload_value} |
| 72 | |
| 73 | return run_comb( |
| 74 | server, |
| 75 | bench_cmd, |
| 76 | serve_comb=serve_comb, |
| 77 | bench_comb=bench_comb_workload, |
| 78 | link_vars=link_vars, |
| 79 | base_path=_get_comb_base_path( |
| 80 | experiment_dir, |
| 81 | serve_comb, |
| 82 | bench_comb, |
| 83 | extra_parts=("WL-", f"{workload_var}={workload_value}"), |
| 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 | ) |
| 90 | |
| 91 | |
| 92 | def explore_comb_workloads( |
| 93 | server: ServerProcess | None, |
| 94 | bench_cmd: list[str], |
| 95 | *, |
| 96 | serve_comb: ParameterSweepItem, |
| 97 | bench_comb: ParameterSweepItem, |
| 98 | link_vars: list[tuple[str, str]], |
| 99 | workload_var: WorkloadVariable, |
| 100 | workload_iters: int, |
| 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)'}") |
| 109 | print(f"Bench parameters: {bench_comb.as_text() or '(None)'}") |
| 110 | print(f"Number of workload iterations: {workload_iters}") |
| 111 | |
| 112 | if workload_iters < 2: |
| 113 | raise ValueError("`workload_iters` should be at least 2") |
| 114 | |
| 115 | dataset_size = DEFAULT_NUM_PROMPTS |
| 116 | if "num_prompts" in bench_comb: |
| 117 | dataset_size = int(bench_comb["num_prompts"]) # type: ignore |
| 118 | else: |
| 119 | for i, arg in enumerate(bench_cmd): |
| 120 | if arg == "--num-prompts" and i + 1 < len(bench_cmd): |
| 121 | dataset_size = int(bench_cmd[i + 1]) |
| 122 | break |
| 123 | elif arg.startswith("--num-prompts="): |
| 124 | dataset_size = int(arg.split("=", 1)[1]) |
| 125 | break |
| 126 | |
| 127 | print(f"Dataset size: {dataset_size}") |
| 128 | |
| 129 | serial_workload_data = run_comb_workload( |
| 130 | server, |
| 131 | bench_cmd, |
| 132 | serve_comb=serve_comb, |
| 133 | bench_comb=bench_comb | {"max_concurrency": 1}, |
| 134 | link_vars=link_vars, |
| 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 | ) |
| 143 | batch_workload_data = run_comb_workload( |
| 144 | server, |
| 145 | bench_cmd, |
| 146 | serve_comb=serve_comb, |
| 147 | bench_comb=bench_comb | {"max_concurrency": dataset_size}, |
| 148 | link_vars=link_vars, |
| 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 | |
| 158 | if not serial_workload_data or not batch_workload_data: |
| 159 | if dry_run: |
| 160 | print("Omitting intermediate Workload iterations.") |
| 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) |
| 171 | ) |
| 172 | print(f"Serial inference: {workload_var}={serial_workload_value}") |
| 173 | |
| 174 | batch_workload_value = math.floor( |
| 175 | _estimate_workload_avg(batch_workload_data, workload_var) |
| 176 | ) |
| 177 | print(f"Batch inference: {workload_var}={batch_workload_value}") |
| 178 | |
| 179 | # Avoid duplicated runs for intermediate values if the range between |
| 180 | # `serial_workload_value` and `batch_workload_value` is small |
| 181 | inter_workload_values = np.linspace( |
| 182 | serial_workload_value, batch_workload_value, workload_iters |
| 183 | )[1:-1] |
| 184 | inter_workload_values = sorted(set(map(round, inter_workload_values))) |
| 185 | |
| 186 | inter_workloads_data: list[dict[str, object]] = [] |
| 187 | for inter_workload_value in inter_workload_values: |
| 188 | print(f"Exploring: {workload_var}={inter_workload_value}") |
| 189 | inter_workload_data = run_comb_workload( |
| 190 | server, |
| 191 | bench_cmd, |
| 192 | serve_comb=serve_comb, |
| 193 | bench_comb=bench_comb, |
| 194 | link_vars=link_vars, |
| 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 | ) |
| 203 | if inter_workload_data is not None: |
| 204 | inter_workloads_data.extend(inter_workload_data) |
| 205 | |
| 206 | print("[WL END]") |
| 207 | |
| 208 | return serial_workload_data + inter_workloads_data + batch_workload_data |
| 209 | |
| 210 | |
| 211 | def explore_combs_workloads( |
| 212 | serve_cmd: list[str], |
| 213 | bench_cmd: list[str], |
| 214 | after_bench_cmd: list[str], |
| 215 | *, |
| 216 | show_stdout: bool, |
| 217 | server_ready_timeout: int, |
| 218 | serve_params: ParameterSweep, |
| 219 | bench_params: ParameterSweep, |
| 220 | link_vars: list[tuple[str, str]], |
| 221 | workload_var: WorkloadVariable, |
| 222 | workload_iters: int, |
| 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( |
| 231 | f"You should not override `{workload_var}` in `bench_params` " |
| 232 | "since it is supposed to be explored automatically." |
| 233 | ) |
| 234 | |
| 235 | all_data = list[dict[str, object]]() |
| 236 | for serve_comb in serve_params: |
| 237 | with server_ctx( |
| 238 | serve_cmd, |
| 239 | after_bench_cmd, |
| 240 | show_stdout=show_stdout, |
| 241 | server_ready_timeout=server_ready_timeout, |
| 242 | serve_comb=serve_comb, |
| 243 | bench_params=bench_params, |
| 244 | experiment_dir=experiment_dir, |
| 245 | dry_run=dry_run, |
| 246 | ) as server: |
| 247 | for bench_comb in bench_params: |
| 248 | comb_data = explore_comb_workloads( |
| 249 | server, |
| 250 | bench_cmd, |
| 251 | serve_comb=serve_comb, |
| 252 | bench_comb=bench_comb, |
| 253 | link_vars=link_vars, |
| 254 | workload_var=workload_var, |
| 255 | workload_iters=workload_iters, |
| 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: |
| 264 | all_data.extend(comb_data) |
| 265 | |
| 266 | if dry_run: |
| 267 | return None |
| 268 | |
| 269 | combined_df = pd.DataFrame.from_records(all_data) |
| 270 | combined_df.to_csv(experiment_dir / "summary.csv") |
| 271 | |
| 272 | return combined_df |
| 273 | |
| 274 | |
| 275 | @dataclass |
| 276 | class SweepServeWorkloadArgs(SweepServeArgs): |
| 277 | workload_var: WorkloadVariable |
| 278 | workload_iters: int |
| 279 | |
| 280 | parser_name: ClassVar[str] = "serve_workload" |
| 281 | parser_help: ClassVar[str] = ( |
| 282 | "Explore the latency-throughput tradeoff for different workload levels." |
| 283 | ) |
| 284 | |
| 285 | @classmethod |
| 286 | def from_cli_args(cls, args: argparse.Namespace): |
| 287 | # NOTE: Don't use super() as `from_cli_args` calls `cls()` |
| 288 | base_args = SweepServeArgs.from_cli_args(args) |
| 289 | |
| 290 | return cls( |
| 291 | **{f.name: getattr(base_args, f.name) for f in fields(base_args)}, |
| 292 | workload_var=args.workload_var, |
| 293 | workload_iters=args.workload_iters, |
| 294 | ) |
| 295 | |
| 296 | @classmethod |
| 297 | def add_cli_args(cls, parser: FlexibleArgumentParser) -> FlexibleArgumentParser: |
| 298 | parser = super().add_cli_args(parser) |
| 299 | |
| 300 | workload_group = parser.add_argument_group("workload options") |
| 301 | workload_group.add_argument( |
| 302 | "--workload-var", |
| 303 | type=str, |
| 304 | choices=get_args(WorkloadVariable), |
| 305 | default="request_rate", |
| 306 | help="The variable to adjust in each iteration.", |
| 307 | ) |
| 308 | workload_group.add_argument( |
| 309 | "--workload-iters", |
| 310 | type=int, |
| 311 | default=10, |
| 312 | help="Number of workload levels to explore. " |
| 313 | "This includes the first two iterations used to interpolate the value of " |
| 314 | "`workload_var` for remaining iterations.", |
| 315 | ) |
| 316 | |
| 317 | return parser |
| 318 | |
| 319 | |
| 320 | def run_main(args: SweepServeWorkloadArgs): |
| 321 | experiment_dir = args.resolve_experiment_dir() |
| 322 | |
| 323 | with args.run_ctx(experiment_dir): |
| 324 | return explore_combs_workloads( |
| 325 | serve_cmd=args.serve_cmd, |
| 326 | bench_cmd=args.bench_cmd, |
| 327 | after_bench_cmd=args.after_bench_cmd, |
| 328 | show_stdout=args.show_stdout, |
| 329 | server_ready_timeout=args.server_ready_timeout, |
| 330 | serve_params=args.serve_params, |
| 331 | bench_params=args.bench_params, |
| 332 | link_vars=args.link_vars, |
| 333 | workload_var=args.workload_var, |
| 334 | workload_iters=args.workload_iters, |
| 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 | |
| 343 | def main(args: argparse.Namespace): |
| 344 | run_main(SweepServeWorkloadArgs.from_cli_args(args)) |
| 345 | |
| 346 | |
| 347 | if __name__ == "__main__": |
| 348 | parser = FlexibleArgumentParser(description=SweepServeWorkloadArgs.parser_help) |
| 349 | SweepServeWorkloadArgs.add_cli_args(parser) |
| 350 | |
| 351 | main(parser.parse_args()) |