main
py 467 lines 17.3 KB
Raw
1 #!/usr/bin/env python3
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
5 # do not complain about line length (for docstring)
6 # ruff: noqa: E501
7
8 import argparse
9 import json
10 import sys
11 from dataclasses import asdict, dataclass
12 from datetime import datetime
13 from pathlib import Path
14 from typing import Any
15 from urllib.parse import quote
16
17 import regex as re
18
19
20 def normalize_package_name(name: str) -> str:
21 """Normalize package name according to PEP 503.
22 https://peps.python.org/pep-0503/#normalized-names
23
24 Replace runs of underscores, hyphens, and periods with a single hyphen,
25 and lowercase the result.
26 """
27 return re.sub(r"[-_.]+", "-", name).lower()
28
29
30 if not sys.version_info >= (3, 12):
31 raise RuntimeError("This script requires Python 3.12 or higher.")
32
33 INDEX_HTML_TEMPLATE = """<!DOCTYPE html>
34 <html>
35 <!-- {comment} -->
36 <meta name="pypi:repository-version" content="1.0">
37 <body>
38 {items}
39 </body>
40 </html>
41 """
42
43
44 @dataclass
45 class WheelFileInfo:
46 package_name: str
47 version: str
48 build_tag: str | None
49 python_tag: str
50 abi_tag: str
51 platform_tag: str
52 variant: str | None
53 filename: str
54
55
56 def parse_from_filename(file: str) -> WheelFileInfo:
57 """Parse wheel file name to extract metadata.
58
59 The format of wheel names:
60 {package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
61 All versions could contain a variant like '+cu129' or '.cpu' or `.rocm` (or not).
62
63 Example:
64 vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl
65 vllm-0.10.2rc2+cu129-cp38-abi3-manylinux2014_aarch64.whl
66 vllm-0.11.1rc8.dev14+gaa384b3c0-cp38-abi3-manylinux2014_aarch64.whl
67 vllm-0.11.1rc8.dev14+gaa384b3c0.cu130-cp38-abi3-manylinux1_x86_64.whl
68
69 """
70 wheel_file_re = re.compile(
71 r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
72 )
73 match = wheel_file_re.match(file)
74 if not match:
75 raise ValueError(f"Invalid wheel file name: {file}")
76
77 package_name = match.group("package_name")
78 version = match.group("version")
79 build_tag = match.group("build_tag")
80 python_tag = match.group("python_tag")
81 abi_tag = match.group("abi_tag")
82 platform_tag = match.group("platform_tag")
83
84 # extract variant from version
85 variant = None
86 if "dev" in version:
87 ver_after_dev = version.split("dev")[-1]
88 if "." in ver_after_dev:
89 variant = ver_after_dev.split(".")[-1]
90 version = version.removesuffix("." + variant)
91 else:
92 if "+" in version:
93 version_part, suffix = version.split("+", 1)
94 # Only treat known patterns as variants (rocmXXX, cuXXX, cpu, xpu)
95 # Git hashes and other suffixes are NOT variants
96 if suffix.startswith(("rocm", "cu", "cpu", "xpu")):
97 variant = suffix
98 version = version_part
99 # Otherwise keep the full version string (variant stays None)
100
101 return WheelFileInfo(
102 package_name=package_name,
103 version=version,
104 build_tag=build_tag,
105 python_tag=python_tag,
106 abi_tag=abi_tag,
107 platform_tag=platform_tag,
108 variant=variant,
109 filename=file,
110 )
111
112
113 def generate_project_list(subdir_names: list[str], comment: str = "") -> str:
114 """Generate project list HTML content linking to each project & variant subdirectory."""
115 href_tags = []
116 for name in sorted(subdir_names):
117 name = name.strip("/").strip(".")
118 href_tags.append(f' <a href="{name}/">{name}/</a><br/>')
119 return INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
120
121
122 def generate_package_index_and_metadata(
123 wheel_files: list[WheelFileInfo],
124 wheel_base_dir: Path,
125 index_base_dir: Path,
126 comment: str = "",
127 ) -> tuple[str, str]:
128 """Generate package index HTML content for a specific package, linking to actual wheel files."""
129 href_tags = []
130 metadata = []
131 for file in sorted(wheel_files, key=lambda x: x.filename):
132 relative_path = (
133 wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename
134 )
135 # handle with '+' in URL, and avoid double-encoding '/' and already-encoded '%2B'
136 # NOTE: this is AWS S3 specific behavior!
137 file_path_quoted = quote(relative_path.as_posix(), safe=":%/")
138 href_tags.append(f' <a href="{file_path_quoted}">{file.filename}</a><br/>')
139 file_meta = asdict(file)
140 file_meta["path"] = file_path_quoted
141 metadata.append(file_meta)
142 index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
143 metadata_str = json.dumps(metadata, indent=2)
144 return index_str, metadata_str
145
146
147 def generate_index_and_metadata(
148 whl_files: list[str],
149 wheel_base_dir: Path,
150 index_base_dir: Path,
151 default_variant: str | None = None,
152 alias_to_default: str | None = None,
153 comment: str = "",
154 ):
155 """Generate index for all wheel files.
156
157 Args:
158 whl_files (list[str]): List of wheel files (must be directly under `wheel_base_dir`).
159 wheel_base_dir (Path): Base directory for wheel files.
160 index_base_dir (Path): Base directory to store index files.
161 default_variant (str | None): The default variant name, if any.
162 alias_to_default (str | None): Alias variant name for the default variant, if any.
163 comment (str | None): Optional comment to include in the generated HTML files.
164
165 First, parse all wheel files to extract metadata.
166 We need to collect all wheel files for each variant, and generate an index for it (in a subdirectory).
167 The index for the default variant (if any) is generated in the root index directory.
168
169 If `default_variant` is provided, all wheels must have variant suffixes, and the default variant index
170 is purely a copy of the corresponding variant index, with only the links adjusted.
171 Otherwise, all wheels without variant suffixes are treated as the default variant.
172
173 If `alias_to_default` is provided, an additional alias subdirectory is created, it has the same content
174 as the default variant index, but the links are adjusted accordingly.
175
176 Index directory structure:
177 index_base_dir/ (hosted at wheels.vllm.ai/{nightly,$commit,$version}/)
178 index.html # project list, linking to "vllm/" and other packages, and all variant subdirectories
179 vllm/
180 index.html # package index, pointing to actual files in wheel_base_dir (relative path)
181 metadata.json # machine-readable metadata for all wheels in this package
182 cpu/ # cpu variant subdirectory
183 index.html
184 vllm/
185 index.html
186 metadata.json
187 cu129/ # cu129 is actually the alias to default variant
188 index.html
189 vllm/
190 index.html
191 metadata.json
192 cu130/ # cu130 variant subdirectory
193 index.html
194 vllm/
195 index.html
196 metadata.json
197 ...
198
199 metadata.json stores a dump of all wheel files' metadata in a machine-readable format:
200 [
201 {
202 "package_name": "vllm",
203 "version": "0.10.2rc2",
204 "build_tag": null,
205 "python_tag": "cp38",
206 "abi_tag": "abi3",
207 "platform_tag": "manylinux2014_aarch64",
208 "variant": "cu129",
209 "filename": "vllm-0.10.2rc2+cu129-cp38-abi3-manylinux2014_aarch64.whl",
210 "path": "../vllm-0.10.2rc2%2Bcu129-cp38-abi3-manylinux2014_aarch64.whl" # to be concatenated with the directory URL and URL-encoded
211 },
212 ...
213 ]
214
215 """
216 parsed_files = [parse_from_filename(f) for f in whl_files]
217
218 if not parsed_files:
219 print("No wheel files found, skipping index generation.")
220 return
221
222 # For ROCm builds: inherit variant from vllm wheel
223 # All ROCm wheels should share the same variant as vllm
224 rocm_variant = None
225 for file in parsed_files:
226 if (
227 file.package_name == "vllm"
228 and file.variant
229 and file.variant.startswith("rocm")
230 ):
231 rocm_variant = file.variant
232 print(f"Detected ROCm variant from vllm: {rocm_variant}")
233 break
234
235 # Apply ROCm variant to all wheels without a variant
236 if rocm_variant:
237 for file in parsed_files:
238 if file.variant is None:
239 file.variant = rocm_variant
240 print(f"Inherited variant '{rocm_variant}' for {file.filename}")
241
242 # Group by variant
243 variant_to_files: dict[str, list[WheelFileInfo]] = {}
244 for file in parsed_files:
245 variant = file.variant or "default"
246 if variant not in variant_to_files:
247 variant_to_files[variant] = []
248 variant_to_files[variant].append(file)
249
250 print(f"Found variants: {list(variant_to_files.keys())}")
251
252 # sanity check for default variant
253 if default_variant:
254 if "default" in variant_to_files:
255 raise ValueError(
256 "All wheel files must have variant suffixes when `default_variant` is specified."
257 )
258 if default_variant not in variant_to_files:
259 raise ValueError(
260 f"Default variant '{default_variant}' not found among wheel files."
261 )
262
263 if alias_to_default:
264 if "default" not in variant_to_files:
265 # e.g. only some wheels are uploaded to S3 currently
266 print(
267 "[WARN] Alias to default variant specified, but no default variant found."
268 )
269 elif alias_to_default in variant_to_files:
270 raise ValueError(
271 f"Alias variant name '{alias_to_default}' already exists among wheel files."
272 )
273 else:
274 variant_to_files[alias_to_default] = variant_to_files["default"].copy()
275 print(f"Alias variant '{alias_to_default}' created for default variant.")
276
277 # Generate comment in HTML header
278 comment_str = f" ({comment})" if comment else ""
279 comment_tmpl = f"Generated on {datetime.now().isoformat()}{comment_str}"
280
281 # Generate index for each variant
282 subdir_names = set()
283 for variant, files in variant_to_files.items():
284 if variant == "default":
285 variant_dir = index_base_dir
286 else:
287 variant_dir = index_base_dir / variant
288 subdir_names.add(variant)
289
290 variant_dir.mkdir(parents=True, exist_ok=True)
291
292 # gather all package names in this variant (normalized per PEP 503)
293 packages = set(normalize_package_name(f.package_name) for f in files)
294 if variant == "default":
295 # these packages should also appear in the "project list"
296 # generate after all variants are processed
297 subdir_names = subdir_names.union(packages)
298 else:
299 # generate project list for this variant directly
300 project_list_str = generate_project_list(sorted(packages), comment_tmpl)
301 with open(variant_dir / "index.html", "w") as f:
302 f.write(project_list_str)
303
304 for package in packages:
305 # filter files belonging to this package only (compare normalized names)
306 package_files = [
307 f for f in files if normalize_package_name(f.package_name) == package
308 ]
309 package_dir = variant_dir / package
310 package_dir.mkdir(parents=True, exist_ok=True)
311 index_str, metadata_str = generate_package_index_and_metadata(
312 package_files, wheel_base_dir, package_dir, comment
313 )
314 with open(package_dir / "index.html", "w") as f:
315 f.write(index_str)
316 with open(package_dir / "metadata.json", "w") as f:
317 f.write(metadata_str)
318
319 # Generate top-level project list index
320 project_list_str = generate_project_list(sorted(subdir_names), comment_tmpl)
321 with open(index_base_dir / "index.html", "w") as f:
322 f.write(project_list_str)
323
324
325 if __name__ == "__main__":
326 """
327 Arguments:
328 --version <version> : version string for the current build (e.g., commit hash)
329 --wheel-dir <wheel_directory> : directory containing wheel files (default to be same as `version`)
330 --current-objects <path_to_json> : path to JSON file containing current S3 objects listing in this version directory
331 --output-dir <output_directory> : directory to store generated index files
332 --alias-to-default <alias_variant_name> : (optional) alias variant name for the default variant
333 --comment <comment_string> : (optional) comment string to include in generated HTML files
334 """
335
336 parser = argparse.ArgumentParser(
337 description="Process nightly build wheel files to generate indices."
338 )
339 parser.add_argument(
340 "--version",
341 type=str,
342 required=True,
343 help="Version string for the current build (e.g., commit hash)",
344 )
345 parser.add_argument(
346 "--current-objects",
347 type=str,
348 required=True,
349 help="Path to JSON file containing current S3 objects listing in this version directory",
350 )
351 parser.add_argument(
352 "--output-dir",
353 type=str,
354 required=True,
355 help="Directory to store generated index files",
356 )
357 parser.add_argument(
358 "--wheel-dir",
359 type=str,
360 default=None,
361 help="Directory containing wheel files (default to be same as `version`)",
362 )
363 parser.add_argument(
364 "--alias-to-default",
365 type=str,
366 default=None,
367 help="Alias variant name for the default variant",
368 )
369 parser.add_argument(
370 "--comment",
371 type=str,
372 default="",
373 help="Optional comment string to include in generated HTML files",
374 )
375
376 args = parser.parse_args()
377
378 version = args.version
379 # Allow rocm/ prefix, reject other slashes and all backslashes
380 if "\\" in version:
381 raise ValueError("Version string must not contain backslashes.")
382 if "/" in version and not version.startswith("rocm/"):
383 raise ValueError(
384 "Version string must not contain slashes (except for 'rocm/' prefix)."
385 )
386 current_objects_path = Path(args.current_objects)
387 output_dir = Path(args.output_dir)
388 if not output_dir.exists():
389 output_dir.mkdir(parents=True, exist_ok=True)
390
391 # Read current objects JSON
392 with open(current_objects_path) as f:
393 current_objects: dict[str, list[dict[str, Any]]] = json.load(f)
394
395 # current_objects looks like from list_objects_v2 S3 API:
396 """
397 "Contents": [
398 {
399 "Key": "e2f56c309d2a28899c68975a7e104502d56deb8f/vllm-0.11.2.dev363+ge2f56c309-cp38-abi3-manylinux1_x86_64.whl",
400 "LastModified": "2025-11-28T14:00:32+00:00",
401 "ETag": "\"37a38339c7cdb61ca737021b968075df-52\"",
402 "ChecksumAlgorithm": [
403 "CRC64NVME"
404 ],
405 "ChecksumType": "FULL_OBJECT",
406 "Size": 435649349,
407 "StorageClass": "STANDARD"
408 },
409 ...
410 ]
411 """
412
413 # Extract wheel file keys
414 wheel_files = []
415 for item in current_objects.get("Contents", []):
416 key: str = item["Key"]
417 if key.endswith(".whl"):
418 wheel_files.append(key.split("/")[-1]) # only the filename is used
419
420 print(f"Found {len(wheel_files)} wheel files for version {version}: {wheel_files}")
421
422 # keep only "official" files for a non-nightly version (specified by cli args)
423 PY_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+([a-zA-Z0-9.+-]*)?$")
424 if PY_VERSION_RE.match(version):
425 # upload-wheels.sh ensures no "dev" is in args.version
426 wheel_files = list(
427 filter(
428 lambda x: (version in x and "dev" not in x)
429 or (x.startswith("triton-") and "+xpu-" in x),
430 wheel_files,
431 )
432 )
433 print(f"Non-nightly version detected, wheel files used: {wheel_files}")
434 else:
435 print("Nightly version detected, keeping all wheel files.")
436
437 # Generate index and metadata, assuming wheels and indices are stored as:
438 # s3://vllm-wheels/{wheel_dir}/<wheel files>
439 # s3://vllm-wheels/<anything>/<index files>
440 #
441 # For ROCm builds, version is "rocm/{commit}" and indices are uploaded to:
442 # - rocm/{commit}/ (same as wheels)
443 # - rocm/nightly/
444 # - rocm/{version}/
445 # All these are under the "rocm/" prefix, so relative paths should be
446 # relative to "rocm/", not the bucket root.
447 if args.wheel_dir:
448 # Explicit wheel-dir provided (e.g., for version-specific indices pointing to commit dir)
449 wheel_dir = args.wheel_dir.strip().rstrip("/")
450 elif version.startswith("rocm/"):
451 # For rocm/commit, wheel_base_dir should be just the commit part
452 # so relative path from rocm/0.12.0/rocm710/vllm/ -> ../../../{commit}/
453 wheel_dir = version.split("/", 1)[1]
454 else:
455 wheel_dir = version
456 wheel_base_dir = Path(output_dir).parent / wheel_dir
457 index_base_dir = Path(output_dir)
458
459 generate_index_and_metadata(
460 whl_files=wheel_files,
461 wheel_base_dir=wheel_base_dir,
462 index_base_dir=index_base_dir,
463 default_variant=None,
464 alias_to_default=args.alias_to_default,
465 comment=args.comment.strip(),
466 )
467 print(f"Successfully generated index and metadata in {output_dir}")