From b6ba83d425762a502c3b3e863f071b9b855c461e Mon Sep 17 00:00:00 2001 From: Lucas Tucker <47258766+lucas-tucker@users.noreply.github.com> Date: Wed, 25 Dec 2024 23:04:07 -0600 Subject: [PATCH] Mypy checking for vllm/compilation (#11496) Signed-off-by: lucast2021 Co-authored-by: lucast2021 --- vllm/compilation/backends.py | 10 +++++----- vllm/compilation/multi_output_match.py | 3 ++- vllm/compilation/pass_manager.py | 4 ++-- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 0c7bbfe599b02..826d1744d88a5 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -141,14 +141,14 @@ def produce_guards_expression(self, *args, **kwargs): return "" -def wrap_inductor(graph, +def wrap_inductor(graph: fx.GraphModule, example_inputs, additional_inductor_config, compilation_config: CompilationConfig, graph_index: int = 0, num_graphs: int = 1, runtime_shape: Optional[int] = None, - use_inductor: bool = True): + use_inductor: bool = True) -> Any: if graph_index == 0: # before compiling the first graph, record the start time global compilation_start_time @@ -209,7 +209,7 @@ def wrap_inductor(graph, returns_tuple = graph_returns_tuple(graph) # this is the graph we return to Dynamo to run - def compiled_graph(*args): + def compiled_graph(*args) -> Optional[fx.CompiledFxGraph]: # convert args to list list_args = list(args) graph_output = inductor_compiled_graph(list_args) @@ -247,7 +247,7 @@ def _check_can_cache(*args, **kwargs): # see https://github.com/pytorch/pytorch/blob/9f5ebf3fc609105a74eab4ccc24932d6353ff566/torch/_inductor/codecache.py#L1221 # noqa return - def _get_shape_env(): + def _get_shape_env() -> AlwaysHitShapeEnv: return AlwaysHitShapeEnv() with patch(# for hijacking the hash of the compiled graph @@ -537,7 +537,7 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable: example_inputs[x].clone() for x in self.sym_tensor_indices ] - def copy_and_call(*args): + def copy_and_call(*args) -> fx.GraphModule: list_args = list(args) for i, index in enumerate(self.sym_tensor_indices): runtime_tensor = list_args[index] diff --git a/vllm/compilation/multi_output_match.py b/vllm/compilation/multi_output_match.py index 0ad648abfbb3a..b6bcecdc89e26 100644 --- a/vllm/compilation/multi_output_match.py +++ b/vllm/compilation/multi_output_match.py @@ -7,6 +7,7 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized from torch._inductor import pattern_matcher as pm from torch._ops import OpOverload +from torch.fx import Node from vllm.compilation.fx_utils import find_auto_fn @@ -97,7 +98,7 @@ def insert_getitems(self, tuple_node: fx.Node, self.graph.call_function(operator.getitem, (tuple_node, idx)) for idx in indices) - def insert_auto_fn(self, op: OpOverload, kwargs): + def insert_auto_fn(self, op: OpOverload, kwargs) -> Node: """ Insert an auto_functionalized node with the given op and kwargs. """ diff --git a/vllm/compilation/pass_manager.py b/vllm/compilation/pass_manager.py index fb522ae053e97..34f5f355798b2 100644 --- a/vllm/compilation/pass_manager.py +++ b/vllm/compilation/pass_manager.py @@ -1,4 +1,4 @@ -from typing import List +from typing import Any, Dict, List from torch import fx as fx @@ -53,7 +53,7 @@ def add(self, pass_: InductorPass): assert isinstance(pass_, InductorPass) self.passes.append(pass_) - def __getstate__(self): + def __getstate__(self) -> Dict[str, List[Any]]: """ Custom pickling for the pass manager, as some passes cannot be pickled. Pickling occurs because the pass manager is set as the value of