| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import json |
| 6 | |
| 7 | from vllm.profiler.layerwise_profile import ModelStatsEntry, SummaryStatsEntry |
| 8 | from vllm.profiler.utils import TablePrinter, indent_string |
| 9 | |
| 10 | |
| 11 | def flatten_entries(entry_cls, profile_dict: dict): |
| 12 | entries_and_depth = [] |
| 13 | |
| 14 | def get_entries(node, curr_depth=0): |
| 15 | entries_and_depth.append((entry_cls(**node["entry"]), curr_depth)) |
| 16 | |
| 17 | for child in node["children"]: |
| 18 | get_entries( |
| 19 | child, |
| 20 | curr_depth=curr_depth + 1, |
| 21 | ) |
| 22 | |
| 23 | for root in profile_dict: |
| 24 | get_entries(root) |
| 25 | |
| 26 | return entries_and_depth |
| 27 | |
| 28 | |
| 29 | if __name__ == "__main__": |
| 30 | parser = argparse.ArgumentParser() |
| 31 | |
| 32 | parser.add_argument( |
| 33 | "--json-trace", |
| 34 | type=str, |
| 35 | required=True, |
| 36 | help=( |
| 37 | "JSON trace file generated by scripts that use " |
| 38 | "vllm.profiler.layerwise_profile" |
| 39 | ), |
| 40 | ) |
| 41 | parser.add_argument( |
| 42 | "--phase", |
| 43 | type=str, |
| 44 | required=True, |
| 45 | help="The phase to print the table for. This is either" |
| 46 | "prefill or decode_n, where n is the decode step " |
| 47 | "number", |
| 48 | ) |
| 49 | parser.add_argument( |
| 50 | "--table", |
| 51 | type=str, |
| 52 | choices=["summary", "model"], |
| 53 | default="summary", |
| 54 | help="Which table to print, the summary table or the layerwise model table", |
| 55 | ) |
| 56 | |
| 57 | args = parser.parse_args() |
| 58 | |
| 59 | with open(args.json_trace) as f: |
| 60 | profile_data = json.load(f) |
| 61 | |
| 62 | assert args.phase in profile_data, ( |
| 63 | f"Cannot find phase {args.phase} in profile data. Choose one among" |
| 64 | f"{[x for x in profile_data if 'prefill' in x or 'decode' in x]}" |
| 65 | ) # noqa |
| 66 | |
| 67 | if args.table == "summary": |
| 68 | entries_and_depths = flatten_entries( |
| 69 | SummaryStatsEntry, profile_data[args.phase]["summary_stats"] |
| 70 | ) |
| 71 | column_widths = dict(name=80, cuda_time_us=12, pct_cuda_time=12, invocations=15) |
| 72 | elif args.table == "model": |
| 73 | entries_and_depths = flatten_entries( |
| 74 | ModelStatsEntry, profile_data[args.phase]["model_stats"] |
| 75 | ) |
| 76 | column_widths = dict( |
| 77 | name=60, cpu_time_us=12, cuda_time_us=12, pct_cuda_time=12, trace=60 |
| 78 | ) |
| 79 | |
| 80 | # indent entry names based on the depth |
| 81 | entries = [] |
| 82 | for entry, depth in entries_and_depths: |
| 83 | entry.name = indent_string( |
| 84 | entry.name, |
| 85 | indent=depth, |
| 86 | indent_style=lambda indent: "|" + "-" * indent + " ", |
| 87 | ) |
| 88 | entries.append(entry) |
| 89 | |
| 90 | TablePrinter(type(entries[0]), column_widths).print_table(entries) |