@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 220 lines 7.5 KB
Raw
1 #!/usr/bin/env python3
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 """Pin vLLM dependencies to exact versions of custom ROCm wheels.
5
6 This script modifies vLLM's requirements files to replace version constraints
7 with exact versions of custom-built ROCm wheels (torch, triton, torchvision, amdsmi).
8
9 This ensures that 'pip install vllm' automatically installs the correct custom wheels
10 instead of allowing pip to download different versions from PyPI.
11 """
12
13 import sys
14 from pathlib import Path
15
16 import regex as re
17
18
19 def extract_version_from_wheel(wheel_name: str) -> str:
20 """Extract version from wheel filename.
21
22 Example:
23 torch-2.9.0a0+git1c57644-cp312-cp312-linux_x86_64.whl -> 2.9.0a0+git1c57644
24 triton-3.4.0-cp312-cp312-linux_x86_64.whl -> 3.4.0
25
26 """
27 # Wheel format:
28 # {distribution}-{version}(-{build tag})?-{python}-{abi}-{platform}.whl
29 parts = wheel_name.replace(".whl", "").split("-")
30
31 if len(parts) < 5:
32 raise ValueError(f"Invalid wheel filename format: {wheel_name}")
33
34 # Version is the second part
35 version = parts[1]
36 return version
37
38
39 def get_custom_wheel_versions(install_dir: str) -> dict[str, str]:
40 """Read /install directory and extract versions of custom wheels.
41
42 Returns:
43 Dict mapping package names to exact versions
44
45 """
46 install_path = Path(install_dir)
47 if not install_path.exists():
48 print(f"ERROR: Install directory not found: {install_dir}", file=sys.stderr)
49 sys.exit(1)
50
51 versions = {}
52
53 # Map wheel prefixes to package names
54 # IMPORTANT: Use dashes to avoid matching substrings
55 # (e.g., 'torch' would match 'torchvision')
56 # ORDER MATTERS: This order is preserved when pinning dependencies
57 # in requirements files
58 package_mapping = [
59 ("torch-", "torch"), # Match torch- (not torchvision)
60 ("triton-", "triton"), # Match triton- (not triton_kernels)
61 ("triton_kernels-", "triton-kernels"), # Match triton_kernels-
62 ("torchvision-", "torchvision"), # Match torchvision-
63 ("torchaudio-", "torchaudio"), # Match torchaudio-
64 ("amdsmi-", "amdsmi"), # Match amdsmi-
65 ("flash_attn-", "flash-attn"), # Match flash_attn-
66 ("amd_aiter-", "amd-aiter"), # Match amd_aiter-
67 ]
68
69 for wheel_file in install_path.glob("*.whl"):
70 wheel_name = wheel_file.name
71
72 for prefix, package_name in package_mapping:
73 if wheel_name.startswith(prefix):
74 try:
75 version = extract_version_from_wheel(wheel_name)
76 versions[package_name] = version
77 print(f"Found {package_name}=={version}", file=sys.stderr)
78 except Exception as e:
79 print(
80 f"WARNING: Could not extract version from {wheel_name}: {e}",
81 file=sys.stderr,
82 )
83 break
84
85 # Return versions in the order defined by package_mapping
86 ordered_versions = {}
87 for _, package_name in package_mapping:
88 if package_name in versions:
89 ordered_versions[package_name] = versions[package_name]
90 return ordered_versions
91
92
93 def pin_dependencies_in_requirements(requirements_path: str, versions: dict[str, str]):
94 """Insert custom wheel pins at the TOP of requirements file.
95
96 This ensures that when setup.py processes the file line-by-line,
97 custom wheels (torch, triton, etc.) are encountered FIRST, before
98 any `-r common.txt` includes that might pull in other dependencies.
99
100 Creates:
101 # Custom ROCm wheel pins (auto-generated)
102 torch==2.9.0a0+git1c57644
103 triton==3.4.0
104 torchvision==0.23.0a0+824e8c8
105 amdsmi==26.1.0+5df6c765
106
107 -r common.txt
108 ... rest of file ...
109 """
110 requirements_file = Path(requirements_path)
111
112 if not requirements_file.exists():
113 print(
114 f"ERROR: Requirements file not found: {requirements_path}", file=sys.stderr
115 )
116 sys.exit(1)
117
118 # Backup original file
119 backup_file = requirements_file.with_suffix(requirements_file.suffix + ".bak")
120 with open(requirements_file) as f:
121 original_lines = f.readlines()
122
123 # Write backup
124 with open(backup_file, "w") as f:
125 f.writelines(original_lines)
126
127 # Build header with pinned custom wheels
128 header_lines = [
129 "# Custom ROCm wheel pins (auto-generated by pin_rocm_dependencies.py)\n",
130 "# These must come FIRST to ensure correct dependency resolution\n",
131 ]
132
133 for package_name, exact_version in versions.items():
134 header_lines.append(f"{package_name}=={exact_version}\n")
135
136 header_lines.append("\n") # Blank line separator
137
138 # Filter out any existing entries for custom packages from original file
139 filtered_lines = []
140 removed_packages = []
141
142 for line in original_lines:
143 stripped = line.strip()
144 should_keep = True
145
146 # Check if this line is for one of our custom packages
147 if stripped and not stripped.startswith("#") and not stripped.startswith("-"):
148 for package_name in versions:
149 # Handle both hyphen and underscore variations
150 pattern_name = package_name.replace("-", "[-_]")
151 pattern = rf"^{pattern_name}\s*[=<>]=?\s*[\d.a-zA-Z+]+"
152
153 if re.match(pattern, stripped, re.IGNORECASE):
154 removed_packages.append(f"{package_name}: {stripped}")
155 should_keep = False
156 break
157
158 if should_keep:
159 filtered_lines.append(line)
160
161 # Combine: header + filtered original content
162 final_lines = header_lines + filtered_lines
163
164 # Write modified content
165 with open(requirements_file, "w") as f:
166 f.writelines(final_lines)
167
168 # Print summary
169 print("\n✓ Inserted custom wheel pins at TOP of requirements:", file=sys.stderr)
170 for package_name, exact_version in versions.items():
171 print(f" - {package_name}=={exact_version}", file=sys.stderr)
172
173 if removed_packages:
174 print("\n✓ Removed old package entries:", file=sys.stderr)
175 for pkg in removed_packages:
176 print(f" - {pkg}", file=sys.stderr)
177
178 print(f"\n✓ Patched requirements file: {requirements_path}", file=sys.stderr)
179 print(f" Backup saved: {backup_file}", file=sys.stderr)
180
181
182 def main():
183 if len(sys.argv) != 3:
184 print(
185 f"Usage: {sys.argv[0]} <install_dir> <requirements_file>", file=sys.stderr
186 )
187 print(
188 f"Example: {sys.argv[0]} /install /app/vllm/requirements/rocm.txt",
189 file=sys.stderr,
190 )
191 sys.exit(1)
192
193 install_dir = sys.argv[1]
194 requirements_path = sys.argv[2]
195
196 print("=" * 70, file=sys.stderr)
197 print("Pinning vLLM dependencies to custom ROCm wheel versions", file=sys.stderr)
198 print("=" * 70, file=sys.stderr)
199
200 # Get versions from custom wheels
201 print(f"\nScanning {install_dir} for custom wheels...", file=sys.stderr)
202 versions = get_custom_wheel_versions(install_dir)
203
204 if not versions:
205 print("\nERROR: No custom wheels found in /install!", file=sys.stderr)
206 sys.exit(1)
207
208 # Pin dependencies in requirements file
209 print(f"\nPatching {requirements_path}...", file=sys.stderr)
210 pin_dependencies_in_requirements(requirements_path, versions)
211
212 print("\n" + "=" * 70, file=sys.stderr)
213 print("✓ Dependency pinning complete!", file=sys.stderr)
214 print("=" * 70, file=sys.stderr)
215
216 sys.exit(0)
217
218
219 if __name__ == "__main__":
220 main()