| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import dataclasses |
| 5 | from collections.abc import Callable |
| 6 | |
| 7 | from _typeshed import DataclassInstance |
| 8 | from torch._C._profiler import _EventType, _ProfilerEvent, _TensorMetadata |
| 9 | |
| 10 | # |
| 11 | # String / Print Manipulation |
| 12 | # |
| 13 | |
| 14 | |
| 15 | def trim_string_back(string: str, width: int) -> str: |
| 16 | if len(string) > width: |
| 17 | offset = len(string) - width + 3 |
| 18 | string = string[:-offset] |
| 19 | if len(string) > 3: |
| 20 | string = string + "..." |
| 21 | return string |
| 22 | |
| 23 | |
| 24 | class TablePrinter: |
| 25 | def __init__(self, row_cls: type[DataclassInstance], column_widths: dict[str, int]): |
| 26 | self.row_cls = row_cls |
| 27 | self.fieldnames = [x.name for x in dataclasses.fields(row_cls)] |
| 28 | self.column_widths = column_widths |
| 29 | assert set(self.column_widths.keys()) == set(self.fieldnames) |
| 30 | |
| 31 | def print_table(self, rows: list[DataclassInstance]): |
| 32 | self._print_header() |
| 33 | self._print_line() |
| 34 | for row in rows: |
| 35 | self._print_row(row) |
| 36 | |
| 37 | def _print_header(self): |
| 38 | for i, f in enumerate(self.fieldnames): |
| 39 | last = i == len(self.fieldnames) - 1 |
| 40 | col_width = self.column_widths[f] |
| 41 | print( |
| 42 | trim_string_back(f, col_width).ljust(col_width), |
| 43 | end=" | " if not last else "\n", |
| 44 | ) |
| 45 | |
| 46 | def _print_row(self, row): |
| 47 | assert isinstance(row, self.row_cls) |
| 48 | |
| 49 | for i, f in enumerate(self.fieldnames): |
| 50 | last = i == len(self.fieldnames) - 1 |
| 51 | col_width = self.column_widths[f] |
| 52 | val = getattr(row, f) |
| 53 | |
| 54 | val_str = "" |
| 55 | if isinstance(val, str): |
| 56 | val_str = trim_string_back(val, col_width).ljust(col_width) |
| 57 | elif type(val) in [float, int]: |
| 58 | val_str = f"{float(val):>.2f}".rjust(col_width) |
| 59 | else: |
| 60 | val_str = f"{val}".rjust(col_width) |
| 61 | print(val_str, end=" | " if not last else "\n") |
| 62 | |
| 63 | def _print_line(self): |
| 64 | total_col_width = 0 |
| 65 | for column_width in self.column_widths.values(): |
| 66 | total_col_width += column_width |
| 67 | print("=" * (total_col_width + 3 * (len(self.column_widths) - 1))) |
| 68 | |
| 69 | |
| 70 | def indent_string( |
| 71 | string: str, indent: int, indent_style: Callable[[int], str] | str = " " |
| 72 | ) -> str: |
| 73 | if indent: |
| 74 | if isinstance(indent_style, str): |
| 75 | return indent_style * indent + string |
| 76 | else: |
| 77 | return indent_style(indent) + string |
| 78 | else: |
| 79 | return string |
| 80 | |
| 81 | |
| 82 | # |
| 83 | # _ProfilerEvent utils |
| 84 | # |
| 85 | |
| 86 | |
| 87 | def event_has_module(event: _ProfilerEvent) -> bool: |
| 88 | event_type, typed_event = event.typed |
| 89 | if event_type == _EventType.PyCall: |
| 90 | return typed_event.module is not None |
| 91 | return False |
| 92 | |
| 93 | |
| 94 | def event_is_torch_op(event: _ProfilerEvent) -> bool: |
| 95 | return event.tag == _EventType.TorchOp |
| 96 | |
| 97 | |
| 98 | def event_arg_repr(arg) -> str: |
| 99 | if arg is None or type(arg) in [float, int, bool, str]: |
| 100 | return f"{arg}" |
| 101 | elif isinstance(arg, list): |
| 102 | return f"[{', '.join([event_arg_repr(x) for x in arg])}]" |
| 103 | elif isinstance(arg, tuple): |
| 104 | return f"({', '.join([event_arg_repr(x) for x in arg])})" |
| 105 | else: |
| 106 | assert isinstance(arg, _TensorMetadata), f"Unsupported type: {type(arg)}" |
| 107 | sizes_str = ", ".join([str(x) for x in arg.sizes]) |
| 108 | return f"{str(arg.dtype).replace('torch.', '')}[{sizes_str}]" |
| 109 | |
| 110 | |
| 111 | def event_torch_op_repr(event: _ProfilerEvent) -> str: |
| 112 | assert event.tag == _EventType.TorchOp |
| 113 | args_str = ", ".join([event_arg_repr(x) for x in event.typed[1].inputs]) |
| 114 | return f"{event.name}({args_str})".replace("aten::", "") |
| 115 | |
| 116 | |
| 117 | def event_module_repr(event: _ProfilerEvent) -> str: |
| 118 | assert event_has_module(event) |
| 119 | module = event.typed[1].module |
| 120 | if module.parameters and len(module.parameters) > 0: |
| 121 | args_str = ", ".join( |
| 122 | [f"{x[0]}={event_arg_repr(x[1])}" for x in module.parameters] |
| 123 | ) |
| 124 | return f"{module.cls_name}({args_str})" |
| 125 | else: |
| 126 | return module.cls_name |
| 127 | |
| 128 | |
| 129 | def event_torch_op_stack_trace( |
| 130 | curr_event: _ProfilerEvent, until: Callable[[_ProfilerEvent], bool] |
| 131 | ) -> str: |
| 132 | trace = "" |
| 133 | curr_event = curr_event.parent |
| 134 | while curr_event and not until(curr_event): |
| 135 | if event_is_torch_op(curr_event): |
| 136 | if len(trace) > 0: |
| 137 | trace += " <- " |
| 138 | trace += event_torch_op_repr(curr_event) |
| 139 | curr_event = curr_event.parent |
| 140 | |
| 141 | return trace |