| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | import math |
| 5 | from concurrent.futures import ProcessPoolExecutor |
| 6 | from dataclasses import dataclass |
| 7 | from functools import partial |
| 8 | from pathlib import Path |
| 9 | from typing import TYPE_CHECKING, ClassVar |
| 10 | |
| 11 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 12 | from vllm.utils.collection_utils import full_groupby |
| 13 | from vllm.utils.import_utils import PlaceholderModule |
| 14 | |
| 15 | from .plot import DummyExecutor, _json_load_bytes |
| 16 | from .utils import sanitize_filename |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | import pandas as pd |
| 20 | |
| 21 | |
| 22 | def _first_present(run_data: dict[str, object], keys: list[str]): |
| 23 | for key in keys: |
| 24 | for candidate in {key, key.replace("_", "-"), key.replace("-", "_")}: |
| 25 | if candidate in run_data: |
| 26 | return run_data[candidate] |
| 27 | return None |
| 28 | |
| 29 | |
| 30 | def _get_numeric( |
| 31 | run_data: dict[str, object], |
| 32 | keys: list[str], |
| 33 | *, |
| 34 | allow_zero: bool = True, |
| 35 | ) -> float | None: |
| 36 | value = _first_present(run_data, keys) |
| 37 | if value is None: |
| 38 | return None |
| 39 | |
| 40 | try: |
| 41 | numeric = float(value) |
| 42 | except (TypeError, ValueError) as exc: |
| 43 | raise ValueError( |
| 44 | f"Expected numeric value for one of {keys}, " |
| 45 | f"but found {value!r} in {run_data=}" |
| 46 | ) from exc |
| 47 | |
| 48 | if not allow_zero and numeric == 0: |
| 49 | return None |
| 50 | |
| 51 | return numeric |
| 52 | |
| 53 | |
| 54 | def _infer_user_count( |
| 55 | run_data: dict[str, object], |
| 56 | user_count_var: str | None, |
| 57 | ) -> float | None: |
| 58 | candidates = [user_count_var] if user_count_var else [] |
| 59 | candidates.extend(["request_rate"]) |
| 60 | user_count = _get_numeric(run_data, candidates, allow_zero=False) |
| 61 | if user_count is not None: |
| 62 | return user_count |
| 63 | |
| 64 | # Fallback to the observed peak if configured value is missing. |
| 65 | return _get_numeric(run_data, ["max_concurrent_requests"], allow_zero=False) |
| 66 | |
| 67 | |
| 68 | def _infer_gpu_count( |
| 69 | run_data: dict[str, object], |
| 70 | gpu_count_var: str | None, |
| 71 | ) -> float: |
| 72 | direct_candidates = [gpu_count_var] if gpu_count_var else [] |
| 73 | direct_gpu_count = _get_numeric(run_data, direct_candidates, allow_zero=False) |
| 74 | if direct_gpu_count: |
| 75 | return direct_gpu_count |
| 76 | |
| 77 | tp_size = _get_numeric(run_data, ["tensor_parallel_size", "tp"]) |
| 78 | pp_size = _get_numeric(run_data, ["pipeline_parallel_size", "pp"]) |
| 79 | dp_size = _get_numeric(run_data, ["data_parallel_size", "dp"]) |
| 80 | world_size = 1.0 |
| 81 | if tp_size: |
| 82 | world_size *= tp_size |
| 83 | if pp_size: |
| 84 | world_size *= pp_size |
| 85 | if dp_size: |
| 86 | world_size *= dp_size |
| 87 | |
| 88 | return world_size |
| 89 | |
| 90 | |
| 91 | def _get_throughput( |
| 92 | run_data: dict[str, object], |
| 93 | throughput_var: str, |
| 94 | ) -> float: |
| 95 | throughput = _get_numeric(run_data, [throughput_var]) |
| 96 | if throughput is None: |
| 97 | raise ValueError( |
| 98 | f"Cannot find throughput metric {throughput_var!r} in run data. " |
| 99 | f"Available keys: {sorted(run_data)}" |
| 100 | ) |
| 101 | |
| 102 | return throughput |
| 103 | |
| 104 | |
| 105 | def _prepare_records( |
| 106 | all_data: list[dict[str, object]], |
| 107 | *, |
| 108 | user_count_var: str | None, |
| 109 | gpu_count_var: str | None, |
| 110 | ) -> tuple[list[dict[str, object]], int]: |
| 111 | prepared = [] |
| 112 | skipped_missing_users = 0 |
| 113 | |
| 114 | for record in all_data: |
| 115 | throughput = _get_throughput(record, "output_throughput") |
| 116 | user_count = _infer_user_count(record, user_count_var) |
| 117 | if user_count is None: |
| 118 | skipped_missing_users += 1 |
| 119 | continue |
| 120 | |
| 121 | gpu_count = _infer_gpu_count(record, gpu_count_var) |
| 122 | tokens_per_user = throughput / user_count |
| 123 | tokens_per_gpu = throughput / gpu_count |
| 124 | |
| 125 | prepared.append( |
| 126 | { |
| 127 | **record, |
| 128 | "tokens_per_user": tokens_per_user, |
| 129 | "tokens_per_gpu": tokens_per_gpu, |
| 130 | "user_count_estimate": user_count, |
| 131 | "gpu_count": gpu_count, |
| 132 | } |
| 133 | ) |
| 134 | |
| 135 | return prepared, skipped_missing_users |
| 136 | |
| 137 | |
| 138 | def _pareto_frontier( |
| 139 | df: "pd.DataFrame", |
| 140 | x_col: str, |
| 141 | y_col: str, |
| 142 | *, |
| 143 | epsilon: float = 1e-9, |
| 144 | ) -> "pd.DataFrame": |
| 145 | sorted_df = df.sort_values([x_col, y_col], ascending=[False, False]) |
| 146 | frontier_indices = [] |
| 147 | best_y = -math.inf |
| 148 | |
| 149 | for idx, row in sorted_df.iterrows(): |
| 150 | y_val = row[y_col] |
| 151 | if y_val >= best_y - epsilon: |
| 152 | frontier_indices.append(idx) |
| 153 | best_y = max(best_y, y_val) |
| 154 | |
| 155 | return df.loc[frontier_indices] |
| 156 | |
| 157 | |
| 158 | def _get_fig_path( |
| 159 | fig_dir: Path, |
| 160 | fig_group: tuple[tuple[str, str], ...], |
| 161 | ) -> Path: |
| 162 | parts = ["PARETO"] |
| 163 | if fig_group: |
| 164 | parts.extend(f"{k}={v}" for k, v in fig_group) |
| 165 | filename = sanitize_filename("-".join(parts) + ".png") |
| 166 | return fig_dir / filename |
| 167 | |
| 168 | |
| 169 | def _plot_fig( |
| 170 | fig_dir: Path, |
| 171 | fig_group_data: tuple[tuple[tuple[str, str], ...], list[dict[str, object]]], |
| 172 | label_by: list[str], |
| 173 | *, |
| 174 | dry_run: bool, |
| 175 | ): |
| 176 | fig_group, fig_data = fig_group_data |
| 177 | fig_path = _get_fig_path(fig_dir, fig_group) |
| 178 | |
| 179 | print("[BEGIN FIGURE]") |
| 180 | print(f"Group: {dict(fig_group)}") |
| 181 | print(f"Output file: {fig_path}") |
| 182 | |
| 183 | if dry_run: |
| 184 | print("[END FIGURE]") |
| 185 | return |
| 186 | |
| 187 | # Lazy-import matplotlib/pandas/seaborn |
| 188 | try: |
| 189 | import matplotlib.pyplot as plt |
| 190 | except ImportError: |
| 191 | plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot") |
| 192 | try: |
| 193 | import pandas as pd |
| 194 | except ImportError: |
| 195 | pd = PlaceholderModule("pandas") |
| 196 | try: |
| 197 | import seaborn as sns |
| 198 | except ImportError: |
| 199 | sns = PlaceholderModule("seaborn") |
| 200 | |
| 201 | df = pd.DataFrame.from_records(fig_data) |
| 202 | df = df.dropna(subset=["tokens_per_user", "tokens_per_gpu"]) |
| 203 | |
| 204 | if df.empty: |
| 205 | print("No data points available after filtering; skipping.") |
| 206 | print("[END FIGURE]") |
| 207 | return |
| 208 | |
| 209 | frontier = _pareto_frontier(df, "tokens_per_user", "tokens_per_gpu") |
| 210 | frontier = frontier.sort_values("tokens_per_user") |
| 211 | |
| 212 | fig, ax = plt.subplots() |
| 213 | sns.scatterplot( |
| 214 | data=df, |
| 215 | x="tokens_per_user", |
| 216 | y="tokens_per_gpu", |
| 217 | color="0.5", |
| 218 | alpha=0.6, |
| 219 | ax=ax, |
| 220 | label="All runs", |
| 221 | ) |
| 222 | sns.lineplot( |
| 223 | data=frontier, |
| 224 | x="tokens_per_user", |
| 225 | y="tokens_per_gpu", |
| 226 | marker="o", |
| 227 | ax=ax, |
| 228 | label="Pareto frontier", |
| 229 | ) |
| 230 | |
| 231 | if label_by: |
| 232 | for _, row in frontier.iterrows(): |
| 233 | label_parts = [] |
| 234 | for key in label_by: |
| 235 | if key in row: |
| 236 | label_parts.append(f"{key}={row[key]}") |
| 237 | if label_parts: |
| 238 | ax.text( |
| 239 | row["tokens_per_user"], |
| 240 | row["tokens_per_gpu"], |
| 241 | "\n".join(label_parts), |
| 242 | fontsize=8, |
| 243 | ) |
| 244 | |
| 245 | ax.set_xlabel("Tokens/s/user") |
| 246 | ax.set_ylabel("Tokens/s/GPU") |
| 247 | ax.grid(True, linestyle="--", linewidth=0.5, alpha=0.6) |
| 248 | fig.tight_layout() |
| 249 | fig.savefig(fig_path) |
| 250 | plt.close(fig) |
| 251 | |
| 252 | print( |
| 253 | f"Plotted {len(df)} points; Pareto frontier size: {len(frontier)}.", |
| 254 | ) |
| 255 | print("[END FIGURE]") |
| 256 | |
| 257 | |
| 258 | def plot_pareto( |
| 259 | output_dir: Path, |
| 260 | user_count_var: str | None, |
| 261 | gpu_count_var: str | None, |
| 262 | label_by: list[str], |
| 263 | *, |
| 264 | dry_run: bool, |
| 265 | ): |
| 266 | fig_dir = output_dir / "pareto" |
| 267 | raw_data = [ |
| 268 | run_data |
| 269 | for path in output_dir.rglob("**/summary.json") |
| 270 | for run_data in _json_load_bytes(path) |
| 271 | ] |
| 272 | |
| 273 | if not raw_data: |
| 274 | raise ValueError(f"Did not find any parameter sweep results under {output_dir}") |
| 275 | |
| 276 | fig_dir.mkdir(parents=True, exist_ok=True) |
| 277 | |
| 278 | prepared_data, skipped_missing_users = _prepare_records( |
| 279 | raw_data, |
| 280 | user_count_var=user_count_var, |
| 281 | gpu_count_var=gpu_count_var, |
| 282 | ) |
| 283 | |
| 284 | if skipped_missing_users: |
| 285 | print( |
| 286 | f"Skipped {skipped_missing_users} runs without a user count " |
| 287 | "(`max_concurrency` or `max_concurrent_requests`).", |
| 288 | ) |
| 289 | |
| 290 | if not prepared_data: |
| 291 | raise ValueError( |
| 292 | "No data points with both throughput and user count available " |
| 293 | "to plot Pareto frontier.", |
| 294 | ) |
| 295 | |
| 296 | fig_groups = full_groupby( |
| 297 | prepared_data, |
| 298 | key=lambda item: tuple(), |
| 299 | ) |
| 300 | |
| 301 | with DummyExecutor() if len(fig_groups) <= 1 else ProcessPoolExecutor() as executor: |
| 302 | all( |
| 303 | executor.map( |
| 304 | partial( |
| 305 | _plot_fig, |
| 306 | fig_dir, |
| 307 | label_by=label_by, |
| 308 | dry_run=dry_run, |
| 309 | ), |
| 310 | fig_groups, |
| 311 | ) |
| 312 | ) |
| 313 | |
| 314 | |
| 315 | @dataclass |
| 316 | class SweepPlotParetoArgs: |
| 317 | output_dir: Path |
| 318 | user_count_var: str | None |
| 319 | gpu_count_var: str | None |
| 320 | label_by: list[str] |
| 321 | dry_run: bool |
| 322 | |
| 323 | parser_name: ClassVar[str] = "plot_pareto" |
| 324 | parser_help: ClassVar[str] = ( |
| 325 | "Plot Pareto frontier between tokens/s/user and tokens/s/GPU " |
| 326 | "from parameter sweep results." |
| 327 | ) |
| 328 | |
| 329 | @classmethod |
| 330 | def from_cli_args(cls, args: argparse.Namespace): |
| 331 | output_dir = Path(args.EXPERIMENT_DIR) |
| 332 | if not output_dir.exists(): |
| 333 | raise ValueError(f"No parameter sweep results under {output_dir}") |
| 334 | |
| 335 | label_by = [] if not args.label_by else args.label_by.split(",") |
| 336 | |
| 337 | return cls( |
| 338 | output_dir=output_dir, |
| 339 | user_count_var=args.user_count_var, |
| 340 | gpu_count_var=args.gpu_count_var, |
| 341 | label_by=label_by, |
| 342 | dry_run=args.dry_run, |
| 343 | ) |
| 344 | |
| 345 | @classmethod |
| 346 | def add_cli_args(cls, parser: FlexibleArgumentParser): |
| 347 | parser.add_argument( |
| 348 | "EXPERIMENT_DIR", |
| 349 | type=str, |
| 350 | help="The directory containing the sweep results to plot.", |
| 351 | ) |
| 352 | parser.add_argument( |
| 353 | "--user-count-var", |
| 354 | type=str, |
| 355 | default="max_concurrency", |
| 356 | help="Result key that stores concurrent user count. " |
| 357 | "Falls back to max_concurrent_requests if missing.", |
| 358 | ) |
| 359 | parser.add_argument( |
| 360 | "--gpu-count-var", |
| 361 | type=str, |
| 362 | default=None, |
| 363 | help="Result key that stores GPU count. " |
| 364 | "If not provided, falls back to num_gpus/gpu_count " |
| 365 | "or tensor_parallel_size * pipeline_parallel_size.", |
| 366 | ) |
| 367 | parser.add_argument( |
| 368 | "--label-by", |
| 369 | type=str, |
| 370 | default="max_concurrency,gpu_count", |
| 371 | help="Comma-separated list of fields to annotate on Pareto frontier " |
| 372 | "points.", |
| 373 | ) |
| 374 | parser.add_argument( |
| 375 | "--dry-run", |
| 376 | action="store_true", |
| 377 | help="If set, prints the figures to plot without drawing them.", |
| 378 | ) |
| 379 | |
| 380 | return parser |
| 381 | |
| 382 | |
| 383 | def run_main(args: SweepPlotParetoArgs): |
| 384 | return plot_pareto( |
| 385 | output_dir=args.output_dir, |
| 386 | user_count_var=args.user_count_var, |
| 387 | gpu_count_var=args.gpu_count_var, |
| 388 | label_by=args.label_by, |
| 389 | dry_run=args.dry_run, |
| 390 | ) |
| 391 | |
| 392 | |
| 393 | def main(args: argparse.Namespace): |
| 394 | run_main(SweepPlotParetoArgs.from_cli_args(args)) |
| 395 | |
| 396 | |
| 397 | if __name__ == "__main__": |
| 398 | parser = FlexibleArgumentParser(description=SweepPlotParetoArgs.parser_help) |
| 399 | SweepPlotParetoArgs.add_cli_args(parser) |
| 400 | |
| 401 | main(parser.parse_args()) |