@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ff334ca1cd

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 11:00 UTC

Update deprecated type hinting in `vllm/profiler` (#18057)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed May 13, 2025 at 12:34 UTC ff334ca1cd92c41cc79e9dead91de40b87601daf

3 files changed +23 -24

pyproject.toml

-1

@@ -84,7 +84,6 @@ exclude = [

84 "vllm/model_executor/models/**/*.py" = ["UP006", "UP035"]

85 "vllm/platforms/**/*.py" = ["UP006", "UP035"]

86 "vllm/plugins/**/*.py" = ["UP006", "UP035"]

87 -"vllm/profiler/**/*.py" = ["UP006", "UP035"]

87 "vllm/prompt_adapter/**/*.py" = ["UP006", "UP035"]

88 "vllm/spec_decode/**/*.py" = ["UP006", "UP035"]

89 "vllm/transformers_utils/**/*.py" = ["UP006", "UP035"]

vllm/profiler/layerwise_profile.py

+19 -19

@@ -3,7 +3,7 @@

3 import copy

4 from collections import defaultdict

5 from dataclasses import asdict, dataclass, field

6 -from typing import Any, Callable, Dict, List, Optional, Tuple, TypeAlias, Union

6 +from typing import Any, Callable, Optional, TypeAlias, Union

8 import pandas as pd

9 from torch._C._autograd import DeviceType, _KinetoEvent, _ProfilerResult

@@ -20,7 +20,7 @@ from vllm.profiler.utils import (TablePrinter, event_has_module,

20 class _ModuleTreeNode:

21 event: _ProfilerEvent

22 parent: Optional['_ModuleTreeNode'] = None

23 - children: List['_ModuleTreeNode'] = field(default_factory=list)

23 + children: list['_ModuleTreeNode'] = field(default_factory=list)

24 trace: str = ""

26 @property

@@ -60,19 +60,19 @@ StatsEntry: TypeAlias = Union[ModelStatsEntry, SummaryStatsEntry]

60 @dataclass

61 class _StatsTreeNode:

62 entry: StatsEntry

63 - children: List[StatsEntry]

63 + children: list[StatsEntry]

64 parent: Optional[StatsEntry]

67 @dataclass

68 class LayerwiseProfileResults(profile):

69 _kineto_results: _ProfilerResult

70 - _kineto_event_correlation_map: Dict[int,

71 - List[_KinetoEvent]] = field(init=False)

72 - _event_correlation_map: Dict[int, List[FunctionEvent]] = field(init=False)

73 - _module_tree: List[_ModuleTreeNode] = field(init=False)

74 - _model_stats_tree: List[_StatsTreeNode] = field(init=False)

75 - _summary_stats_tree: List[_StatsTreeNode] = field(init=False)

70 + _kineto_event_correlation_map: dict[int,

71 + list[_KinetoEvent]] = field(init=False)

72 + _event_correlation_map: dict[int, list[FunctionEvent]] = field(init=False)

73 + _module_tree: list[_ModuleTreeNode] = field(init=False)

74 + _model_stats_tree: list[_StatsTreeNode] = field(init=False)

75 + _summary_stats_tree: list[_StatsTreeNode] = field(init=False)

77 # profile metadata

78 num_running_seqs: Optional[int] = None

@@ -82,7 +82,7 @@ class LayerwiseProfileResults(profile):

82 self._build_module_tree()

83 self._build_stats_trees()

85 - def print_model_table(self, column_widths: Dict[str, int] = None):

85 + def print_model_table(self, column_widths: dict[str, int] = None):

86 _column_widths = dict(name=60,

87 cpu_time_us=12,

88 cuda_time_us=12,

@@ -100,7 +100,7 @@ class LayerwiseProfileResults(profile):

100 filtered_model_table,

101 indent_style=lambda indent: "|" + "-" * indent + " "))

102

103 - def print_summary_table(self, column_widths: Dict[str, int] = None):

103 + def print_summary_table(self, column_widths: dict[str, int] = None):

104 _column_widths = dict(name=80,

105 cuda_time_us=12,

106 pct_cuda_time=12,

@@ -142,7 +142,7 @@ class LayerwiseProfileResults(profile):

142 }

143

144 @staticmethod

145 - def _indent_row_names_based_on_depth(depths_rows: List[Tuple[int,

145 + def _indent_row_names_based_on_depth(depths_rows: list[tuple[int,

146 StatsEntry]],

147 indent_style: Union[Callable[[int],

148 str],

@@ -229,7 +229,7 @@ class LayerwiseProfileResults(profile):

229 [self._cumulative_cuda_time(root) for root in self._module_tree])

230

231 def _build_stats_trees(self):

232 - summary_dict: Dict[str, _StatsTreeNode] = {}

232 + summary_dict: dict[str, _StatsTreeNode] = {}

233 total_cuda_time = self._total_cuda_time()

234

235 def pct_cuda_time(cuda_time_us):

@@ -238,7 +238,7 @@ class LayerwiseProfileResults(profile):

238 def build_summary_stats_tree_df(

239 node: _ModuleTreeNode,

240 parent: Optional[_StatsTreeNode] = None,

241 - summary_trace: Tuple[str] = ()):

241 + summary_trace: tuple[str] = ()):

242

243 if event_has_module(node.event):

244 name = event_module_repr(node.event)

@@ -313,8 +313,8 @@ class LayerwiseProfileResults(profile):

313 self._model_stats_tree.append(build_model_stats_tree_df(root))

314

315 def _flatten_stats_tree(

316 - self, tree: List[_StatsTreeNode]) -> List[Tuple[int, StatsEntry]]:

317 - entries: List[Tuple[int, StatsEntry]] = []

316 + self, tree: list[_StatsTreeNode]) -> list[tuple[int, StatsEntry]]:

317 + entries: list[tuple[int, StatsEntry]] = []

318

319 def df_traversal(node: _StatsTreeNode, depth=0):

320 entries.append((depth, node.entry))

@@ -327,10 +327,10 @@ class LayerwiseProfileResults(profile):

327 return entries

328

329 def _convert_stats_tree_to_dict(self,

330 - tree: List[_StatsTreeNode]) -> List[Dict]:

331 - root_dicts: List[Dict] = []

330 + tree: list[_StatsTreeNode]) -> list[dict]:

331 + root_dicts: list[dict] = []

332

333 - def df_traversal(node: _StatsTreeNode, curr_json_list: List[Dict]):

333 + def df_traversal(node: _StatsTreeNode, curr_json_list: list[dict]):

334 curr_json_list.append({

335 "entry": asdict(node.entry),

336 "children": []

vllm/profiler/utils.py

+4 -4

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import dataclasses

4 -from typing import Callable, Dict, List, Type, Union

4 +from typing import Callable, Union

6 from torch._C._profiler import _EventType, _ProfilerEvent, _TensorMetadata

@@ -30,14 +30,14 @@ def trim_string_back(string, width):

31 class TablePrinter:

33 - def __init__(self, row_cls: Type[dataclasses.dataclass],

34 - column_widths: Dict[str, int]):

33 + def __init__(self, row_cls: type[dataclasses.dataclass],

34 + column_widths: dict[str, int]):

35 self.row_cls = row_cls

36 self.fieldnames = [x.name for x in dataclasses.fields(row_cls)]

37 self.column_widths = column_widths

38 assert set(self.column_widths.keys()) == set(self.fieldnames)

40 - def print_table(self, rows: List[dataclasses.dataclass]):

40 + def print_table(self, rows: list[dataclasses.dataclass]):

41 self._print_header()

42 self._print_line()

43 for row in rows: