main
py 690 lines 19.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import argparse
4 import json
5 from abc import ABC, abstractmethod
6 from concurrent.futures import ProcessPoolExecutor
7 from dataclasses import dataclass
8 from functools import partial
9 from pathlib import Path
10 from types import TracebackType
11 from typing import TYPE_CHECKING, ClassVar
12
13 from typing_extensions import Self, override
14
15 from vllm.utils.argparse_utils import FlexibleArgumentParser
16 from vllm.utils.collection_utils import full_groupby
17 from vllm.utils.import_utils import PlaceholderModule
18
19 from .utils import sanitize_filename
20
21 if TYPE_CHECKING:
22 import pandas as pd
23
24
25 @dataclass
26 class PlotFilterBase(ABC):
27 var: str
28 target: str
29
30 @classmethod
31 def parse_str(cls, s: str):
32 for op_key in PLOT_FILTERS:
33 if op_key in s:
34 key, value = s.split(op_key)
35 return PLOT_FILTERS[op_key](
36 key,
37 value.removeprefix(op_key).strip("'").strip('"'),
38 )
39 else:
40 raise ValueError(
41 f"Invalid operator for plot filter '{s}'. "
42 f"Valid operators are: {sorted(PLOT_FILTERS)}",
43 )
44
45 @abstractmethod
46 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
47 """Applies this filter to a DataFrame."""
48 raise NotImplementedError
49
50
51 @dataclass
52 class PlotEqualTo(PlotFilterBase):
53 @override
54 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
55 target: float | str
56 try:
57 target = float(self.target)
58 except ValueError:
59 target = self.target
60
61 return df[df[self.var] == target]
62
63
64 @dataclass
65 class PlotNotEqualTo(PlotFilterBase):
66 @override
67 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
68 target: float | str
69 try:
70 target = float(self.target)
71 except ValueError:
72 target = self.target
73
74 return df[df[self.var] != target]
75
76
77 @dataclass
78 class PlotLessThan(PlotFilterBase):
79 @override
80 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
81 return df[df[self.var] < float(self.target)]
82
83
84 @dataclass
85 class PlotLessThanOrEqualTo(PlotFilterBase):
86 @override
87 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
88 return df[df[self.var] <= float(self.target)]
89
90
91 @dataclass
92 class PlotGreaterThan(PlotFilterBase):
93 @override
94 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
95 return df[df[self.var] > float(self.target)]
96
97
98 @dataclass
99 class PlotGreaterThanOrEqualTo(PlotFilterBase):
100 @override
101 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
102 return df[df[self.var] >= float(self.target)]
103
104
105 # NOTE: The ordering is important! Match longer op_keys first
106 PLOT_FILTERS: dict[str, type[PlotFilterBase]] = {
107 "==": PlotEqualTo,
108 "!=": PlotNotEqualTo,
109 "<=": PlotLessThanOrEqualTo,
110 ">=": PlotGreaterThanOrEqualTo,
111 "<": PlotLessThan,
112 ">": PlotGreaterThan,
113 }
114
115
116 class PlotFilters(list[PlotFilterBase]):
117 @classmethod
118 def parse_str(cls, s: str):
119 if not s:
120 return cls()
121
122 return cls(PlotFilterBase.parse_str(e) for e in s.split(","))
123
124 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
125 for item in self:
126 df = item.apply(df)
127
128 return df
129
130
131 @dataclass
132 class PlotBinner:
133 var: str
134 bin_size: float
135
136 @classmethod
137 def parse_str(cls, s: str):
138 for op_key in PLOT_BINNERS:
139 if op_key in s:
140 key, value = s.split(op_key)
141 return PLOT_BINNERS[op_key](key, float(value.removeprefix(op_key)))
142 else:
143 raise ValueError(
144 f"Invalid operator for plot binner '{s}'. "
145 f"Valid operators are: {sorted(PLOT_BINNERS)}",
146 )
147
148 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
149 """Applies this binner to a DataFrame."""
150 df = df.copy()
151 df[self.var] = df[self.var] // self.bin_size * self.bin_size
152 return df
153
154
155 PLOT_BINNERS: dict[str, type[PlotBinner]] = {
156 "%": PlotBinner,
157 }
158
159
160 class PlotBinners(list[PlotBinner]):
161 @classmethod
162 def parse_str(cls, s: str):
163 if not s:
164 return cls()
165
166 return cls(PlotBinner.parse_str(e) for e in s.split(","))
167
168 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":
169 for item in self:
170 df = item.apply(df)
171
172 return df
173
174
175 def _json_load_bytes(path: Path) -> list[dict[str, object]]:
176 with path.open("rb") as f:
177 return json.load(f)
178
179
180 def _convert_inf_nan_strings(data: list[dict[str, object]]) -> list[dict[str, object]]:
181 """Convert string values "inf", "-inf", and "nan" to their float equivalents.
182
183 This handles the case where JSON serialization represents inf/nan as strings.
184 """
185 converted_data = []
186 for record in data:
187 converted_record: dict[str, object] = {}
188 for key, value in record.items():
189 if isinstance(value, str):
190 if value in ["inf", "-inf", "nan"]:
191 converted_record[key] = float(value)
192 else:
193 converted_record[key] = value
194 else:
195 converted_record[key] = value
196 converted_data.append(converted_record)
197 return converted_data
198
199
200 def _get_metric(run_data: dict[str, object], metric_key: str):
201 try:
202 return run_data[metric_key]
203 except KeyError as exc:
204 raise ValueError(f"Cannot find metric {metric_key!r} in {run_data=}") from exc
205
206
207 def _get_group(run_data: dict[str, object], group_keys: list[str]):
208 return tuple((k, str(_get_metric(run_data, k))) for k in group_keys)
209
210
211 def _get_fig_path(fig_dir: Path, group: tuple[tuple[str, str], ...], fig_name: str):
212 parts = list[str]()
213
214 # Start with figure name (always provided, defaults to "FIGURE")
215 parts.append(fig_name)
216
217 # Always append group data if present
218 if group:
219 parts.extend(f"{k}={v}" for k, v in group)
220
221 return fig_dir / sanitize_filename("-".join(parts) + ".png")
222
223
224 class DummyExecutor:
225 map = map
226
227 def __enter__(self) -> Self:
228 return self
229
230 def __exit__(
231 self,
232 exc_type: type[BaseException] | None,
233 exc_value: BaseException | None,
234 exc_traceback: TracebackType | None,
235 ) -> None:
236 return None
237
238
239 def _plot_fig(
240 fig_dir: Path,
241 fig_group_data: tuple[tuple[tuple[str, str], ...], list[dict[str, object]]],
242 row_by: list[str],
243 col_by: list[str],
244 curve_by: list[str],
245 *,
246 var_x: str,
247 var_y: str,
248 filter_by: PlotFilters,
249 bin_by: PlotBinners,
250 scale_x: str | None,
251 scale_y: str | None,
252 dry_run: bool,
253 fig_name: str,
254 error_bars: bool,
255 fig_height: float,
256 fig_dpi: int,
257 ):
258 # Lazy-import matplotlib/pandas/seaborn
259 try:
260 import matplotlib.pyplot as plt
261 except ImportError:
262 plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
263 try:
264 import pandas as pd
265 except ImportError:
266 pd = PlaceholderModule("pandas")
267 try:
268 import seaborn as sns
269 except ImportError:
270 sns = PlaceholderModule("seaborn")
271
272 fig_group, fig_data = fig_group_data
273
274 row_groups = full_groupby(
275 fig_data,
276 key=lambda item: _get_group(item, row_by),
277 )
278 num_rows = len(row_groups)
279 num_cols = max(
280 len(full_groupby(row_data, key=lambda item: _get_group(item, col_by)))
281 for _, row_data in row_groups
282 )
283
284 fig_path = _get_fig_path(fig_dir, fig_group, fig_name)
285
286 print("[BEGIN FIGURE]")
287 print(f"Group: {dict(fig_group)}")
288 print(f"Grid: {num_rows} rows x {num_cols} cols")
289 print(f"Output file: {fig_path}")
290
291 if dry_run:
292 print("[END FIGURE]")
293 return
294
295 # Convert string "inf", "-inf", and "nan" to their float equivalents
296 fig_data = _convert_inf_nan_strings(fig_data)
297 df = pd.DataFrame.from_records(fig_data)
298
299 if var_x not in df.columns:
300 raise ValueError(
301 f"Cannot find {var_x=!r} in parameter sweep results. "
302 f"Available variables: {df.columns.tolist()}"
303 )
304 if var_y not in df.columns:
305 raise ValueError(
306 f"Cannot find {var_y=!r} in parameter sweep results. "
307 f"Available variables: {df.columns.tolist()}"
308 )
309 for k in row_by:
310 if k not in df.columns:
311 raise ValueError(
312 f"Cannot find row_by={k!r} in parameter sweep results. "
313 f"Available variables: {df.columns.tolist()}"
314 )
315 for k in col_by:
316 if k not in df.columns:
317 raise ValueError(
318 f"Cannot find col_by={k!r} in parameter sweep results. "
319 f"Available variables: {df.columns.tolist()}"
320 )
321 for k in curve_by:
322 if k not in df.columns:
323 raise ValueError(
324 f"Cannot find curve_by={k!r} in parameter sweep results. "
325 f"Available variables: {df.columns.tolist()}"
326 )
327
328 df = filter_by.apply(df)
329 df = bin_by.apply(df)
330
331 if len(df) == 0:
332 print(f"No data to plot. Filters: {filter_by}")
333 print("[END FIGURE]")
334 return
335
336 # Sort by curve_by columns alphabetically for consistent legend ordering
337 if curve_by:
338 df = df.sort_values(by=curve_by)
339
340 df["row_group"] = (
341 pd.concat(
342 [k + "=" + df[k].astype(str) for k in row_by],
343 axis=1,
344 ).agg("\n".join, axis=1)
345 if row_by
346 else "(All)"
347 )
348
349 df["col_group"] = (
350 pd.concat(
351 [k + "=" + df[k].astype(str) for k in col_by],
352 axis=1,
353 ).agg("\n".join, axis=1)
354 if col_by
355 else "(All)"
356 )
357
358 if len(curve_by) <= 3:
359 hue, style, size, *_ = (*curve_by, None, None, None)
360
361 g = sns.relplot(
362 df,
363 x=var_x,
364 y=var_y,
365 hue=hue,
366 style=style,
367 size=size,
368 markers=True,
369 errorbar="sd" if error_bars else None,
370 kind="line",
371 row="row_group",
372 col="col_group",
373 height=fig_height,
374 )
375 else:
376 df["curve_group"] = (
377 pd.concat(
378 [k + "=" + df[k].astype(str) for k in curve_by],
379 axis=1,
380 ).agg("\n".join, axis=1)
381 if curve_by
382 else "(All)"
383 )
384
385 g = sns.relplot(
386 df,
387 x=var_x,
388 y=var_y,
389 hue="curve_group",
390 markers=True,
391 errorbar="sd" if error_bars else None,
392 kind="line",
393 row="row_group",
394 col="col_group",
395 height=fig_height,
396 )
397
398 if row_by and col_by:
399 g.set_titles("{row_name}\n{col_name}")
400 elif row_by:
401 g.set_titles("{row_name}")
402 elif col_by:
403 g.set_titles("{col_name}")
404 else:
405 g.set_titles("")
406
407 if scale_x:
408 g.set(xscale=scale_x)
409 if scale_y:
410 g.set(yscale=scale_y)
411
412 g.savefig(fig_path, dpi=fig_dpi)
413 plt.close(g.figure)
414
415 print("[END FIGURE]")
416
417
418 def plot(
419 output_dir: Path,
420 fig_dir: Path,
421 fig_by: list[str],
422 row_by: list[str],
423 col_by: list[str],
424 curve_by: list[str],
425 *,
426 var_x: str,
427 var_y: str,
428 filter_by: PlotFilters,
429 bin_by: PlotBinners,
430 scale_x: str | None,
431 scale_y: str | None,
432 dry_run: bool,
433 fig_name: str = "FIGURE",
434 error_bars: bool = True,
435 fig_height: float = 6.4,
436 fig_dpi: int = 300,
437 ):
438 all_data = [
439 run_data
440 for path in output_dir.rglob("**/summary.json")
441 for run_data in _json_load_bytes(path)
442 ]
443
444 if not all_data:
445 raise ValueError(f"Did not find any parameter sweep results under {output_dir}")
446
447 fig_dir.mkdir(parents=True, exist_ok=True)
448
449 fig_groups = full_groupby(
450 all_data,
451 key=lambda item: _get_group(item, fig_by),
452 )
453
454 with DummyExecutor() if len(fig_groups) <= 1 else ProcessPoolExecutor() as executor:
455 # Resolve the iterable to ensure that the workers are run
456 all(
457 executor.map(
458 partial(
459 _plot_fig,
460 fig_dir,
461 row_by=row_by,
462 col_by=col_by,
463 curve_by=curve_by,
464 var_x=var_x,
465 var_y=var_y,
466 filter_by=filter_by,
467 bin_by=bin_by,
468 scale_x=scale_x,
469 scale_y=scale_y,
470 dry_run=dry_run,
471 fig_name=fig_name,
472 error_bars=error_bars,
473 fig_height=fig_height,
474 fig_dpi=fig_dpi,
475 ),
476 fig_groups,
477 )
478 )
479
480
481 @dataclass
482 class SweepPlotArgs:
483 output_dir: Path
484 fig_dir: Path
485 fig_by: list[str]
486 row_by: list[str]
487 col_by: list[str]
488 curve_by: list[str]
489 var_x: str
490 var_y: str
491 filter_by: PlotFilters
492 bin_by: PlotBinners
493 scale_x: str | None
494 scale_y: str | None
495 dry_run: bool
496 fig_name: str = "FIGURE"
497 error_bars: bool = True
498 fig_height: float = 6.4
499 fig_dpi: int = 300
500
501 parser_name: ClassVar[str] = "plot"
502 parser_help: ClassVar[str] = "Plot performance curves from parameter sweep results."
503
504 @classmethod
505 def from_cli_args(cls, args: argparse.Namespace):
506 output_dir = Path(args.EXPERIMENT_DIR)
507 if not output_dir.exists():
508 raise ValueError(f"No parameter sweep results under {output_dir}")
509
510 curve_by = [] if not args.curve_by else args.curve_by.split(",")
511 row_by = [] if not args.row_by else args.row_by.split(",")
512 col_by = [] if not args.col_by else args.col_by.split(",")
513 fig_by = [] if not args.fig_by else args.fig_by.split(",")
514
515 return cls(
516 output_dir=output_dir,
517 fig_dir=output_dir / args.fig_dir,
518 fig_by=fig_by,
519 row_by=row_by,
520 col_by=col_by,
521 curve_by=curve_by,
522 var_x=args.var_x,
523 var_y=args.var_y,
524 filter_by=PlotFilters.parse_str(args.filter_by),
525 bin_by=PlotBinners.parse_str(args.bin_by),
526 scale_x=args.scale_x,
527 scale_y=args.scale_y,
528 dry_run=args.dry_run,
529 fig_name=args.fig_name,
530 error_bars=not args.no_error_bars,
531 fig_height=args.fig_height,
532 fig_dpi=args.fig_dpi,
533 )
534
535 @classmethod
536 def add_cli_args(cls, parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
537 parser.add_argument(
538 "EXPERIMENT_DIR",
539 type=str,
540 help="The directory containing the sweep results to plot.",
541 )
542 parser.add_argument(
543 "--fig-dir",
544 type=str,
545 default="",
546 help="The directory to save the figures, relative to `OUTPUT_DIR`. "
547 "By default, the same directory is used.",
548 )
549 parser.add_argument(
550 "--fig-by",
551 type=str,
552 default="",
553 help="A comma-separated list of variables, such that a separate figure "
554 "is created for each combination of these variables.",
555 )
556 parser.add_argument(
557 "--row-by",
558 type=str,
559 default="",
560 help="A comma-separated list of variables, such that a separate row "
561 "is created for each combination of these variables.",
562 )
563 parser.add_argument(
564 "--col-by",
565 type=str,
566 default="",
567 help="A comma-separated list of variables, such that a separate column "
568 "is created for each combination of these variables.",
569 )
570 parser.add_argument(
571 "--curve-by",
572 type=str,
573 default=None,
574 help="A comma-separated list of variables, such that a separate curve "
575 "is created for each combination of these variables.",
576 )
577 parser.add_argument(
578 "--var-x",
579 type=str,
580 default="total_token_throughput",
581 help="The variable for the x-axis.",
582 )
583 parser.add_argument(
584 "--var-y",
585 type=str,
586 default="median_ttft_ms",
587 help="The variable for the y-axis",
588 )
589 parser.add_argument(
590 "--filter-by",
591 type=str,
592 default="",
593 help="A comma-separated list of statements indicating values to filter by. "
594 "This is useful to remove outliers. "
595 "Example: `max_concurrency<1000,max_num_batched_tokens<=4096` means "
596 "plot only the points where `max_concurrency` is less than 1000 and "
597 "`max_num_batched_tokens` is no greater than 4096.",
598 )
599 parser.add_argument(
600 "--bin-by",
601 type=str,
602 default="",
603 help="A comma-separated list of statements indicating values to bin by. "
604 "This is useful to avoid plotting points that are too close together. "
605 "Example: `request_throughput%%1` means "
606 "use a bin size of 1 for the `request_throughput` variable.",
607 )
608 parser.add_argument(
609 "--scale-x",
610 type=str,
611 default=None,
612 help="The scale to use for the x-axis. "
613 "Currently only accepts string values such as 'log' and 'sqrt'. "
614 "See also: https://seaborn.pydata.org/generated/seaborn.objects.Plot.scale.html",
615 )
616 parser.add_argument(
617 "--scale-y",
618 type=str,
619 default=None,
620 help="The scale to use for the y-axis. "
621 "Currently only accepts string values such as 'log' and 'sqrt'. "
622 "See also: https://seaborn.pydata.org/generated/seaborn.objects.Plot.scale.html",
623 )
624 parser.add_argument(
625 "--fig-name",
626 type=str,
627 default="FIGURE",
628 help="Name prefix for the output figure file. "
629 "Group data is always appended when present. "
630 "Default: 'FIGURE'. Example: --fig-name my_performance_plot",
631 )
632 parser.add_argument(
633 "--no-error-bars",
634 action="store_true",
635 help="If set, disables error bars on the plot. "
636 "By default, error bars are shown.",
637 )
638 parser.add_argument(
639 "--fig-height",
640 type=float,
641 default=6.4,
642 help="Height of each subplot in inches. Default: 6.4",
643 )
644 parser.add_argument(
645 "--fig-dpi",
646 type=int,
647 default=300,
648 help="Resolution of the output figure in dots per inch. Default: 300",
649 )
650 parser.add_argument(
651 "--dry-run",
652 action="store_true",
653 help="If set, prints the information about each figure to plot, "
654 "then exits without drawing them.",
655 )
656
657 return parser
658
659
660 def run_main(args: SweepPlotArgs):
661 return plot(
662 output_dir=args.output_dir,
663 fig_dir=args.fig_dir,
664 fig_by=args.fig_by,
665 row_by=args.row_by,
666 col_by=args.col_by,
667 curve_by=args.curve_by,
668 var_x=args.var_x,
669 var_y=args.var_y,
670 filter_by=args.filter_by,
671 bin_by=args.bin_by,
672 scale_x=args.scale_x,
673 scale_y=args.scale_y,
674 dry_run=args.dry_run,
675 fig_name=args.fig_name,
676 error_bars=args.error_bars,
677 fig_height=args.fig_height,
678 fig_dpi=args.fig_dpi,
679 )
680
681
682 def main(args: argparse.Namespace):
683 run_main(SweepPlotArgs.from_cli_args(args))
684
685
686 if __name__ == "__main__":
687 parser = FlexibleArgumentParser(description=SweepPlotArgs.parser_help)
688 SweepPlotArgs.add_cli_args(parser)
689
690 main(parser.parse_args())