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 8, 19:00 UTC
Consolidate and fix forbidden import `pre-commit` checks (#33982)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Feb 6, 2026 at 09:47 UTC
791a94bed03f4011c98b7179a6fd2f4fed6d1d9a
5 files changed
+145
-300
.pre-commit-config.yaml
+3
-18
@@ -121,24 +121,9 @@ repos:
121
name: Update Dockerfile dependency graph
122
entry: tools/pre_commit/update-dockerfile-graph.sh
123
language: script
124
- - id: enforce-import-regex-instead-of-re
125
- name: Enforce import regex as re
126
- entry: python tools/pre_commit/enforce_regex_import.py
127
- language: python
128
- types: [python]
129
- pass_filenames: false
130
- additional_dependencies: [regex]
131
- # forbid directly import triton
132
- - id: forbid-direct-triton-import
133
- name: "Forbid direct 'import triton'"
134
- entry: python tools/pre_commit/check_triton_import.py
135
- language: python
136
- types: [python]
137
- pass_filenames: false
138
- additional_dependencies: [regex]
139
- - id: check-pickle-imports
140
- name: Prevent new pickle/cloudpickle imports
141
- entry: python tools/pre_commit/check_pickle_imports.py
124
+ - id: check-forbidden-imports
125
+ name: Check for forbidden imports
126
+ entry: python tools/pre_commit/check_forbidden_imports.py
127
language: python
128
types: [python]
129
additional_dependencies: [regex]
tools/pre_commit/check_forbidden_imports.py
new
+142
@@ -0,0 +1,142 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+import sys
5
+from dataclasses import dataclass, field
6
+
7
+import regex as re
8
+
9
+
10
+@dataclass
11
+class ForbiddenImport:
12
+ pattern: str
13
+ tip: str
14
+ allowed_pattern: re.Pattern = re.compile(r"^$") # matches nothing by default
15
+ allowed_files: set[str] = field(default_factory=set)
16
+
17
+
18
+CHECK_IMPORTS = {
19
+ "pickle/cloudpickle": ForbiddenImport(
20
+ pattern=(
21
+ r"^\s*(import\s+(pickle|cloudpickle)(\s|$|\sas)"
22
+ r"|from\s+(pickle|cloudpickle)\s+import\b)"
23
+ ),
24
+ tip=(
25
+ "Avoid using pickle or cloudpickle or add this file to "
26
+ "tools/pre_commit/check_forbidden_imports.py."
27
+ ),
28
+ allowed_files={
29
+ # pickle
30
+ "vllm/multimodal/hasher.py",
31
+ "vllm/transformers_utils/config.py",
32
+ "vllm/model_executor/models/registry.py",
33
+ "vllm/compilation/caching.py",
34
+ "vllm/compilation/piecewise_backend.py",
35
+ "vllm/distributed/utils.py",
36
+ "vllm/distributed/parallel_state.py",
37
+ "vllm/distributed/device_communicators/all_reduce_utils.py",
38
+ "vllm/distributed/device_communicators/shm_broadcast.py",
39
+ "vllm/distributed/device_communicators/shm_object_storage.py",
40
+ "vllm/utils/hashing.py",
41
+ "tests/multimodal/media/test_base.py",
42
+ "tests/tokenizers_/test_hf.py",
43
+ "tests/utils_/test_hashing.py",
44
+ "tests/compile/test_aot_compile.py",
45
+ "benchmarks/kernels/graph_machete_bench.py",
46
+ "benchmarks/kernels/benchmark_lora.py",
47
+ "benchmarks/kernels/benchmark_machete.py",
48
+ "benchmarks/fused_kernels/layernorm_rms_benchmarks.py",
49
+ "benchmarks/cutlass_benchmarks/w8a8_benchmarks.py",
50
+ "benchmarks/cutlass_benchmarks/sparse_benchmarks.py",
51
+ # cloudpickle
52
+ "vllm/v1/executor/multiproc_executor.py",
53
+ "vllm/v1/executor/ray_executor.py",
54
+ "vllm/entrypoints/llm.py",
55
+ "tests/utils.py",
56
+ # pickle and cloudpickle
57
+ "vllm/v1/serial_utils.py",
58
+ },
59
+ ),
60
+ "re": ForbiddenImport(
61
+ pattern=r"^\s*(?:import\s+re(?:$|\s|,)|from\s+re\s+import)",
62
+ tip="Replace 'import re' with 'import regex as re' or 'import regex'.",
63
+ allowed_pattern=re.compile(r"^\s*import\s+regex(\s*|\s+as\s+re\s*)$"),
64
+ allowed_files={"setup.py"},
65
+ ),
66
+ "triton": ForbiddenImport(
67
+ pattern=r"^(from|import)\s+triton(\s|\.|$)",
68
+ tip="Use 'from vllm.triton_utils import triton' instead.",
69
+ allowed_pattern=re.compile(
70
+ "from vllm.triton_utils import (triton|tl|tl, triton)"
71
+ ),
72
+ allowed_files={"vllm/triton_utils/importing.py"},
73
+ ),
74
+}
75
+
76
+
77
+def check_file(path: str) -> int:
78
+ with open(path, encoding="utf-8") as f:
79
+ content = f.read()
80
+ return_code = 0
81
+ # Check all patterns in the whole file
82
+ for import_name, forbidden_import in CHECK_IMPORTS.items():
83
+ # Skip files that are allowed for this import
84
+ if path in forbidden_import.allowed_files:
85
+ continue
86
+ # Search for forbidden imports
87
+ for match in re.finditer(forbidden_import.pattern, content, re.MULTILINE):
88
+ # Check if it's allowed
89
+ if forbidden_import.allowed_pattern.match(match.group()):
90
+ continue
91
+ # Calculate line number from match position
92
+ line_num = content[: match.start() + 1].count("\n") + 1
93
+ print(
94
+ f"{path}:{line_num}: "
95
+ "\033[91merror:\033[0m " # red color
96
+ f"Found forbidden import: {import_name}. {forbidden_import.tip}"
97
+ )
98
+ return_code = 1
99
+ return return_code
100
+
101
+
102
+def main():
103
+ returncode = 0
104
+ for path in sys.argv[1:]:
105
+ returncode |= check_file(path)
106
+ return returncode
107
+
108
+
109
+def test_regex():
110
+ test_cases = [
111
+ # Should match
112
+ ("import pickle", True),
113
+ ("import cloudpickle", True),
114
+ ("import pickle as pkl", True),
115
+ ("import cloudpickle as cpkl", True),
116
+ ("from pickle import *", True),
117
+ ("from cloudpickle import dumps", True),
118
+ ("from pickle import dumps, loads", True),
119
+ ("from cloudpickle import (dumps, loads)", True),
120
+ (" import pickle", True),
121
+ ("\timport cloudpickle", True),
122
+ ("from pickle import loads", True),
123
+ # Should not match
124
+ ("import somethingelse", False),
125
+ ("from somethingelse import pickle", False),
126
+ ("# import pickle", False),
127
+ ("print('import pickle')", False),
128
+ ("import pickleas as asdf", False),
129
+ ]
130
+ for i, (line, should_match) in enumerate(test_cases):
131
+ result = bool(CHECK_IMPORTS["pickle/cloudpickle"].pattern.match(line))
132
+ assert result == should_match, (
133
+ f"Test case {i} failed: '{line}' (expected {should_match}, got {result})"
134
+ )
135
+ print("All regex tests passed.")
136
+
137
+
138
+if __name__ == "__main__":
139
+ if "--test-regex" in sys.argv:
140
+ test_regex()
141
+ else:
142
+ sys.exit(main())
tools/pre_commit/check_pickle_imports.py
deleted
-111
@@ -1,111 +0,0 @@
1
-#!/usr/bin/env python3
2
-# SPDX-License-Identifier: Apache-2.0
3
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
-import sys
5
-
6
-import regex as re
7
-
8
-# List of files (relative to repo root) that are allowed to import pickle or
9
-# cloudpickle
10
-#
11
-# STOP AND READ BEFORE YOU ADD ANYTHING ELSE TO THIS LIST:
12
-# The pickle and cloudpickle modules are known to be unsafe when deserializing
13
-# data from potentially untrusted parties. They have resulted in multiple CVEs
14
-# for vLLM and numerous vulnerabilities in the Python ecosystem more broadly.
15
-# Before adding new uses of pickle/cloudpickle, please consider safer
16
-# alternatives like msgpack or pydantic that are already in use in vLLM. Only
17
-# add to this list if absolutely necessary and after careful security review.
18
-ALLOWED_FILES = {
19
- # pickle
20
- "vllm/multimodal/hasher.py",
21
- "vllm/transformers_utils/config.py",
22
- "vllm/model_executor/models/registry.py",
23
- "vllm/compilation/caching.py",
24
- "vllm/compilation/piecewise_backend.py",
25
- "vllm/distributed/utils.py",
26
- "vllm/distributed/parallel_state.py",
27
- "vllm/distributed/device_communicators/all_reduce_utils.py",
28
- "vllm/distributed/device_communicators/shm_broadcast.py",
29
- "vllm/distributed/device_communicators/shm_object_storage.py",
30
- "vllm/utils/hashing.py",
31
- "tests/multimodal/media/test_base.py",
32
- "tests/tokenizers_/test_hf.py",
33
- "tests/utils_/test_hashing.py",
34
- "tests/compile/test_aot_compile.py",
35
- "benchmarks/kernels/graph_machete_bench.py",
36
- "benchmarks/kernels/benchmark_lora.py",
37
- "benchmarks/kernels/benchmark_machete.py",
38
- "benchmarks/fused_kernels/layernorm_rms_benchmarks.py",
39
- "benchmarks/cutlass_benchmarks/w8a8_benchmarks.py",
40
- "benchmarks/cutlass_benchmarks/sparse_benchmarks.py",
41
- # cloudpickle
42
- "vllm/v1/executor/multiproc_executor.py",
43
- "vllm/v1/executor/ray_executor.py",
44
- "vllm/entrypoints/llm.py",
45
- "tests/utils.py",
46
- # pickle and cloudpickle
47
- "vllm/v1/serial_utils.py",
48
-}
49
-
50
-PICKLE_RE = re.compile(
51
- r"^\s*(import\s+(pickle|cloudpickle)(\s|$|\sas)"
52
- r"|from\s+(pickle|cloudpickle)\s+import\b)"
53
-)
54
-
55
-
56
-def scan_file(path: str) -> int:
57
- with open(path, encoding="utf-8") as f:
58
- for i, line in enumerate(f, 1):
59
- if PICKLE_RE.match(line):
60
- print(
61
- f"{path}:{i}: "
62
- "\033[91merror:\033[0m " # red color
63
- "Found pickle/cloudpickle import"
64
- )
65
- return 1
66
- return 0
67
-
68
-
69
-def main():
70
- returncode = 0
71
- for filename in sys.argv[1:]:
72
- if filename in ALLOWED_FILES:
73
- continue
74
- returncode |= scan_file(filename)
75
- return returncode
76
-
77
-
78
-def test_regex():
79
- test_cases = [
80
- # Should match
81
- ("import pickle", True),
82
- ("import cloudpickle", True),
83
- ("import pickle as pkl", True),
84
- ("import cloudpickle as cpkl", True),
85
- ("from pickle import *", True),
86
- ("from cloudpickle import dumps", True),
87
- ("from pickle import dumps, loads", True),
88
- ("from cloudpickle import (dumps, loads)", True),
89
- (" import pickle", True),
90
- ("\timport cloudpickle", True),
91
- ("from pickle import loads", True),
92
- # Should not match
93
- ("import somethingelse", False),
94
- ("from somethingelse import pickle", False),
95
- ("# import pickle", False),
96
- ("print('import pickle')", False),
97
- ("import pickleas as asdf", False),
98
- ]
99
- for i, (line, should_match) in enumerate(test_cases):
100
- result = bool(PICKLE_RE.match(line))
101
- assert result == should_match, (
102
- f"Test case {i} failed: '{line}' (expected {should_match}, got {result})"
103
- )
104
- print("All regex tests passed.")
105
-
106
-
107
-if __name__ == "__main__":
108
- if "--test-regex" in sys.argv:
109
- test_regex()
110
- else:
111
- sys.exit(main())
tools/pre_commit/check_triton_import.py
deleted
-88
@@ -1,88 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import subprocess
4
-import sys
5
-
6
-import regex as re
7
-
8
-FORBIDDEN_IMPORT_RE = re.compile(r"^(from|import)\s+triton(\s|\.|$)")
9
-
10
-# the way allowed to import triton
11
-ALLOWED_LINES = {
12
- "from vllm.triton_utils import triton",
13
- "from vllm.triton_utils import tl",
14
- "from vllm.triton_utils import tl, triton",
15
-}
16
-
17
-ALLOWED_FILES = {"vllm/triton_utils/importing.py"}
18
-
19
-
20
-def is_allowed_file(current_file: str) -> bool:
21
- return current_file in ALLOWED_FILES
22
-
23
-
24
-def is_forbidden_import(line: str) -> bool:
25
- stripped = line.strip()
26
- return bool(FORBIDDEN_IMPORT_RE.match(stripped)) and stripped not in ALLOWED_LINES
27
-
28
-
29
-def parse_diff(diff: str) -> list[str]:
30
- violations = []
31
- current_file = None
32
- current_lineno = None
33
- skip_allowed_file = False
34
-
35
- for line in diff.splitlines():
36
- if line.startswith("+++ b/"):
37
- current_file = line[6:]
38
- skip_allowed_file = is_allowed_file(current_file)
39
- elif skip_allowed_file:
40
- continue
41
- elif line.startswith("@@"):
42
- match = re.search(r"\+(\d+)", line)
43
- if match:
44
- current_lineno = int(match.group(1)) - 1 # next "+ line" is here
45
- elif line.startswith("+") and not line.startswith("++"):
46
- current_lineno += 1
47
- code_line = line[1:]
48
- if is_forbidden_import(code_line):
49
- violations.append(
50
- f"{current_file}:{current_lineno}: {code_line.strip()}"
51
- )
52
- return violations
53
-
54
-
55
-def get_diff(diff_type: str) -> str:
56
- if diff_type == "staged":
57
- return subprocess.check_output(
58
- ["git", "diff", "--cached", "--unified=0"], text=True
59
- )
60
- elif diff_type == "unstaged":
61
- return subprocess.check_output(["git", "diff", "--unified=0"], text=True)
62
- else:
63
- raise ValueError(f"Unknown diff_type: {diff_type}")
64
-
65
-
66
-def main():
67
- all_violations = []
68
- for diff_type in ["staged", "unstaged"]:
69
- try:
70
- diff_output = get_diff(diff_type)
71
- violations = parse_diff(diff_output)
72
- all_violations.extend(violations)
73
- except subprocess.CalledProcessError as e:
74
- print(f"[{diff_type}] Git diff failed: {e}", file=sys.stderr)
75
-
76
- if all_violations:
77
- print(
78
- "❌ Forbidden direct `import triton` detected."
79
- " ➤ Use `from vllm.triton_utils import triton` instead.\n"
80
- )
81
- for v in all_violations:
82
- print(f"❌ {v}")
83
- return 1
84
- return 0
85
-
86
-
87
-if __name__ == "__main__":
88
- sys.exit(main())
tools/pre_commit/enforce_regex_import.py
deleted
-83
@@ -1,83 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import subprocess
4
-from pathlib import Path
5
-
6
-import regex as re
7
-
8
-FORBIDDEN_PATTERNS = re.compile(r"^\s*(?:import\s+re(?:$|\s|,)|from\s+re\s+import)")
9
-ALLOWED_PATTERNS = [
10
- re.compile(r"^\s*import\s+regex\s+as\s+re\s*$"),
11
- re.compile(r"^\s*import\s+regex\s*$"),
12
-]
13
-
14
-
15
-def get_staged_python_files() -> list[str]:
16
- try:
17
- result = subprocess.run(
18
- ["git", "diff", "--cached", "--name-only", "--diff-filter=AM"],
19
- capture_output=True,
20
- text=True,
21
- check=True,
22
- )
23
- files = result.stdout.strip().split("\n") if result.stdout.strip() else []
24
- return [f for f in files if f.endswith(".py")]
25
- except subprocess.CalledProcessError:
26
- return []
27
-
28
-
29
-def is_forbidden_import(line: str) -> bool:
30
- line = line.strip()
31
- return bool(
32
- FORBIDDEN_PATTERNS.match(line)
33
- and not any(pattern.match(line) for pattern in ALLOWED_PATTERNS)
34
- )
35
-
36
-
37
-def check_file(filepath: str) -> list[tuple[int, str]]:
38
- violations = []
39
- try:
40
- with open(filepath, encoding="utf-8") as f:
41
- for line_num, line in enumerate(f, 1):
42
- if is_forbidden_import(line):
43
- violations.append((line_num, line.strip()))
44
- except (OSError, UnicodeDecodeError):
45
- pass
46
- return violations
47
-
48
-
49
-def main() -> int:
50
- files = get_staged_python_files()
51
- if not files:
52
- return 0
53
-
54
- total_violations = 0
55
-
56
- for filepath in files:
57
- if not Path(filepath).exists():
58
- continue
59
-
60
- if filepath == "setup.py":
61
- continue
62
-
63
- violations = check_file(filepath)
64
- if violations:
65
- print(f"\n❌ {filepath}:")
66
- for line_num, line in violations:
67
- print(f" Line {line_num}: {line}")
68
- total_violations += 1
69
-
70
- if total_violations > 0:
71
- print(f"\n💡 Found {total_violations} violation(s).")
72
- print("❌ Please replace 'import re' with 'import regex as re'")
73
- print(" Also replace 'from re import ...' with 'from regex import ...'") # noqa: E501
74
- print("✅ Allowed imports:")
75
- print(" - import regex as re")
76
- print(" - import regex") # noqa: E501
77
- return 1
78
-
79
- return 0
80
-
81
-
82
-if __name__ == "__main__":
83
- raise SystemExit(main())