@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 192 lines 5.91 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Run mypy on changed files.
4
5 This script is designed to be used as a pre-commit hook. It runs mypy
6 on files that have been changed. It groups files into different mypy calls
7 based on their directory to avoid import following issues.
8
9 Usage:
10 python tools/pre_commit/mypy.py <python_version> <changed_files...>
11
12 Args:
13 python_version: Python version to use (e.g., "3.10") or "local" to use
14 the local Python version.
15 changed_files: List of changed files to check.
16
17 """
18
19 import subprocess
20 import sys
21
22 import regex as re
23
24 # Paths verified clean under follow_imports="silent". Matched before
25 # SEPARATE_GROUPS, so these files join the default group and are checked at the
26 # stricter setting even while a parent directory remains in SEPARATE_GROUPS.
27 #
28 # Fixing a directory means moving it from SEPARATE_GROUPS to here. Without that
29 # move the fixes are not enforced because "tests" claims every file below it.
30 SILENT_GROUPS = [
31 "tests/benchmarks",
32 "tests/compile/correctness_e2e",
33 "tests/compile/fullgraph",
34 "tests/compile/fusions_e2e",
35 "tests/compile/passes",
36 "tests/config",
37 "tests/entrypoints/anthropic",
38 "tests/entrypoints/generate",
39 "tests/entrypoints/multimodal",
40 "tests/entrypoints/pooling",
41 "tests/entrypoints/serve",
42 "tests/entrypoints/tool_parsers",
43 "tests/entrypoints/unit_tests",
44 "tests/entrypoints/weight_transfer",
45 "tests/kernels/attention",
46 "tests/kernels/core",
47 "tests/kernels/helion",
48 "tests/kernels/mamba",
49 "tests/kernels/quantization",
50 "tests/model_executor/layers",
51 "tests/models/language",
52 "tests/models/quantization",
53 "tests/multimodal",
54 "tests/plugins/bge_m3_sparse_plugin",
55 "tests/plugins/prithvi_io_processor_plugin",
56 "tests/plugins/vllm_add_dummy_platform",
57 "tests/plugins/vllm_add_dummy_stat_logger",
58 "tests/plugins_tests",
59 "tests/quantization",
60 "tests/reasoning",
61 "tests/samplers",
62 "tests/spec_decode",
63 "tests/tokenizers_",
64 "tests/tool_use",
65 "tests/transformers_utils",
66 "tests/utils_",
67 "tests/v1/cudagraph",
68 "tests/v1/determinism",
69 "tests/v1/distributed",
70 "tests/v1/executor",
71 "tests/v1/logits_processors",
72 "tests/v1/sample",
73 "tests/v1/shutdown",
74 "tests/v1/simple_kv_offload",
75 "tests/v1/spec_decode",
76 "tests/v1/structured_output",
77 ]
78
79 # After fixing errors resulting from changing follow_imports
80 # from "skip" to "silent", move its directory to SILENT_GROUPS.
81 SEPARATE_GROUPS = [
82 "tests",
83 "tests/compile",
84 "tests/distributed",
85 "tests/entrypoints/llm",
86 "tests/entrypoints/openai",
87 "tests/entrypoints/speech_to_text",
88 "tests/kernels",
89 "tests/kernels/moe",
90 "tests/lora",
91 "tests/model_executor",
92 "tests/model_executor/model_loader",
93 "tests/models",
94 "tests/models/test_initialization.py",
95 "tests/models/multimodal",
96 "tests/parser",
97 "tests/renderers",
98 "tests/tool_parsers",
99 "tests/v1",
100 "tests/v1/attention",
101 "tests/v1/core",
102 "tests/v1/e2e",
103 "tests/v1/ec_connector",
104 "tests/v1/engine",
105 "tests/v1/kv_connector",
106 "tests/v1/kv_offload",
107 "tests/v1/metrics",
108 "tests/v1/streaming_input",
109 "tests/v1/worker",
110 ]
111
112
113 def group_files(changed_files: list[str]) -> dict[str, list[str]]:
114 """Group changed files into different mypy calls.
115
116 Args:
117 changed_files: List of changed files.
118
119 Returns:
120 A dictionary mapping file group names to lists of changed files.
121
122 """
123 silent_pattern = re.compile(f"^({'|'.join(SILENT_GROUPS)}).*")
124 file_groups: dict[str, list[str]] = {"": []}
125 file_groups.update({k: [] for k in SEPARATE_GROUPS})
126 # Longest path first so a sub-directory is not shadowed by its parent
127 separate_groups = sorted(SEPARATE_GROUPS, key=len, reverse=True)
128 for changed_file in changed_files:
129 # Already-fixed paths go in the default group, which runs at the
130 # stricter follow_imports setting from pyproject.toml
131 if silent_pattern.match(changed_file):
132 file_groups[""].append(changed_file)
133 continue
134 # Group files by mypy call
135 for directory in separate_groups:
136 if re.match(f"^{directory}.*", changed_file):
137 file_groups[directory].append(changed_file)
138 break
139 else:
140 if changed_file.startswith(("vllm/", "tests/")):
141 file_groups[""].append(changed_file)
142 return file_groups
143
144
145 def mypy(
146 targets: list[str],
147 python_version: str | None,
148 follow_imports: str | None,
149 file_group: str,
150 ) -> int:
151 """Run mypy on the given targets.
152
153 Args:
154 targets: List of files or directories to check.
155 python_version: Python version to use (e.g., "3.10") or None to use
156 the default mypy version.
157 follow_imports: Value for the --follow-imports option or None to use
158 the default mypy behavior.
159 file_group: The file group name for logging purposes.
160
161 Returns:
162 The return code from mypy.
163
164 """
165 args = ["mypy"]
166 if python_version is not None:
167 args += ["--python-version", python_version]
168 if follow_imports is not None:
169 args += ["--follow-imports", follow_imports]
170 print(f"$ {' '.join(args)} {file_group}")
171 return subprocess.run(args + targets, check=False).returncode
172
173
174 def main():
175 python_version = sys.argv[1]
176 file_groups = group_files(sys.argv[2:])
177
178 if python_version == "local":
179 python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
180
181 returncode = 0
182 for file_group, changed_files in file_groups.items():
183 follow_imports = None if file_group == "" else "skip"
184 if changed_files:
185 returncode |= mypy(
186 changed_files, python_version, follow_imports, file_group
187 )
188 return returncode
189
190
191 if __name__ == "__main__":
192 sys.exit(main())