#!/usr/bin/env python # Copyright Advanced Micro Devices, Inc. # SPDX-License-Identifier: MIT """Download prerelease packages from S3 bucket for promotion to release. This script downloads release candidate packages from the therock-prerelease-python S3 bucket. It discovers all architectures that contain packages matching a specific version pattern and downloads them into a local directory structure. In addition to downloading packages to promote, the script also allows to download PyPI dependencies that are part of the same S3 bucket. Selection of those packages can be done by version pattern and architecture. Options to just list architectures found matching the version or list all packages per architecture without downloading are also available. Packages that are not known to be either a package to promote or a PyPI dependency are also listed. Previously downloaded packages are skipped. PREREQUISITES: - pip install -r ./build_tools/packaging/requirements.txt - AWS credentials configured - read and list bucket access for therock-prerelease-python bucket TYPICAL USAGE (Command Line): # Download all 7.10.0rc2 packages for all architectures: python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --output-dir=./downloads # Download only for a specific architecture: python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --arch=gfx950-dcgpu \ --output-dir=./downloads # Download for multiple specific architectures: python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --arch=gfx1151,gfx950-dcgpu \ --output-dir=./downloads # Download packages including tarballs: python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --arch=gfx1151,gfx950-dcgpu \ --output-dir=./downloads \ --include-tarballs # List available architectures without downloading: python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --list-archs # List all packages per architecture without downloading: python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --list-packages-per-arch # List all packages and tarballs per architecture with sizes: python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --list-packages-per-arch \ --include-tarballs # Download packages to promote and all known PyPI dependencies: python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --output-dir=./downloads \ --include-dependencies # Download multi-arch packages python download_python_packages.py --version=7.13.0rc1 --multi-arch --output-dir=./downloads --bucket-prefix=v4/whl/ # Download multi-arch packages including tarballs python download_python_packages.py --version=7.13.0rc1 --multi-arch --output-dir=./downloads --bucket-prefix=v4/whl/ --include-tarballs # List multi-arch packages python download_python_packages.py --version=7.13.0rc1 --multi-arch --list-multi-arch-packages --bucket-prefix=v4/whl/ DIRECTORY STRUCTURE: Output directory structure will be: / / package1.whl package2.whl ... / package1.whl ... tarball/ (single-arch tarballs, if --include-tarballs is specified) tarball-multi-arch/ (multi-arch tarballs, if --include-tarballs is specified) therock-dist-linux--.tar.gz therock-dist-windows--.tar.gz ... PACKAGE CATEGORIES: Packages to promote (ROCm and PyTorch): - rocm (sdist) - rocm-sdk-core - rocm-sdk-devel - rocm-sdk-libraries-* - torch - torchaudio - torchvision - triton Known dependencies: - filelock, fsspec, jinja2, markupsafe, mpmath, networkx - numpy, pillow, setuptools, sympy, typing-extensions Unknown packages: - everything else that is not a package to promote or a dependency package """ import argparse import fnmatch import re import sys from pathlib import Path from typing import List, Tuple, Union, Dict try: import boto3 from botocore.exceptions import ClientError, NoCredentialsError except ImportError: print("[ERROR]: boto3 not installed. Please run:") print(" pip install boto3") sys.exit(1) # Constants # Maximum number of S3 keys to check when validating architecture MAX_S3_KEYS_CHECK = 100 BYTES_TO_MB = 1024 * 1024 # Conversion factor from bytes to MB # The JAX plugin/pjrt wheels embed the ROCm major version in their package name # (jax_rocm7_plugin for ROCm 7, jax_rocm10_plugin for ROCm 10), so they are # matched by pattern instead of listed once per ROCm release. JAX_ROCM_PACKAGE_PATTERN = re.compile(r"^jax_rocm\d+_(plugin|pjrt)$") # Package categories # Note: replace - in package names with _ to match the filename patterns in S3 PACKAGES_TO_PROMOTE = { "apex", "jaxlib", "rocm", "rocm_profiler", "rocm_sdk_core", "rocm_sdk_devel", "rocm_sdk_libraries-*", "torch", "torchaudio", "torchvision", "triton", } PACKAGES_TO_PROMOTE_MULTI_ARCH = { "amd_torch_device", "amd_torchvision_device", "apex", "rocm", "rocm_profiler", "rocm_sdk_core", "rocm_sdk_devel", "rocm_sdk_device", "rocm_sdk_libraries", "torch", "torchaudio", "torchvision", "triton", # TODO: Enable once upstream support lands # Tracking: PyTorch/audio#4180 # "amd_torchaudio_device", } # copied from build_tools/third_party/s3_management/update_dependencies.py PACKAGES_PER_PROJECT # Note: replace - in package names with _ to match the filename patterns in S3 DEPENDENCY_PACKAGES = { # Jax "dbus_python", "flatbuffers", "ml_dtypes", "opt_einsum", "tomli", # # Torch "sympy", "mpmath", "pillow", "networkx", "numpy", "jinja2", "markupsafe", "filelock", "fsspec", "typing_extensions", # # ROCM "setuptools", } def filter_package_entries( packages: List[Tuple[str, int]], include_package_globs: List[str] | None = None, exclude_package_globs: List[str] | None = None, ) -> List[Tuple[str, int]]: """Filter package entries by filename glob patterns. Args: packages: List of tuples (s3_key, size). include_package_globs: Optional glob patterns. If set, only matching filenames are kept. exclude_package_globs: Optional glob patterns. Matching filenames are removed. Returns: Filtered list of package entries. """ include_package_globs = include_package_globs or [] exclude_package_globs = exclude_package_globs or [] filtered = [] for key, size in packages: filename = key.split("/")[-1] if include_package_globs and not any( fnmatch.fnmatch(filename, pattern) for pattern in include_package_globs ): print(f" FILTERED OUT by include filters: {filename}") continue if exclude_package_globs and any( fnmatch.fnmatch(filename, pattern) for pattern in exclude_package_globs ): print(f" FILTERED OUT by exclude filters: {filename}") continue filtered.append((key, size)) return filtered def is_allowed_multi_arch_package( filename: str, requested_arches: list[str] | None = None, ) -> bool: """ Check if filename belongs to allowed packages for multi-arch flow. Uses: - PACKAGES_TO_PROMOTE_MULTI_ARCH Rules: - Regular packages require exact match - Device packages require: - package prefix match - gfx target in filename - optional requested arch match """ # Extract package name from wheel/sdist filename by splitting at the # beginning of the version segment (e.g. "-7", "-2"). # This preserves hyphenated package names like: # amd-torch-device-gfx942 -> amd_torch_device_gfx942 # while also correctly handling regular packages like: # torch -> torch base = re.split(r"-\d", filename, maxsplit=1)[0] base = base.lower().replace("-", "_") if JAX_ROCM_PACKAGE_PATTERN.match(base): return True for pkg in PACKAGES_TO_PROMOTE_MULTI_ARCH: # Device packages contain gfx targets if "device" in pkg: if not (base.startswith(pkg) and "gfx" in base): continue # If user requested specific arches, # ensure one matches the filename if requested_arches: normalized_arches = { arch.lower().replace("-", "_").split("-")[0] for arch in requested_arches } if not any(arch in base for arch in normalized_arches): continue return True # Regular packages use prefix match if base.startswith(pkg): return True return False def categorize_package(filename: str) -> str: """Categorize a package file. Returns: "promote" - package that needs RC version promotion "dependency" - dependency package to copy as-is "unknown" - unrecognized package """ pkg_name = filename.split("-", 1)[0] # Check for rocm-sdk-libraries-* and jax-rocm-* patterns if ( pkg_name.startswith("rocm_sdk_libraries") or JAX_ROCM_PACKAGE_PATTERN.match(pkg_name) or pkg_name in PACKAGES_TO_PROMOTE ): return "promote" if pkg_name in DEPENDENCY_PACKAGES: return "dependency" return "unknown" def list_architectures( s3_client, bucket_name: str, bucket_prefix: str, version: str ) -> List[str]: """List all architectures in the bucket that have packages matching the version. Args: s3_client: boto3 S3 client bucket_name: S3 bucket name (e.g., "therock-prerelease-python") bucket_prefix: S3 bucket prefix (e.g., "v3/whl/") version: Version pattern to search for (e.g., "7.10.0rc2") Returns: List of architecture names (e.g., ["gfx950-dcgpu", "gfx94X-dcgpu"]) """ print(f"Discovering architectures with version {version}...") # List all "directories" in the bucket prefix try: response = s3_client.list_objects_v2( Bucket=bucket_name, Prefix=bucket_prefix, Delimiter="/", MaxKeys=MAX_S3_KEYS_CHECK, ) architectures = [] if "CommonPrefixes" in response.keys(): # "CommonPrefixes" is a list of prefixes that are the "directories" in the S3 bucket for prefix in response["CommonPrefixes"]: arch = prefix["Prefix"].replace(bucket_prefix, "").rstrip("/") # Check if the arch folder contains files with the specified version if exists_version_single_arch( s3_client, bucket_name, bucket_prefix, arch, version ): architectures.append(arch) print(f" Found: {arch}") if len(architectures) == 0: print( f"""[ERROR]: No architecture subdirectories found in bucket prefix '{bucket_prefix}' of bucket '{bucket_name}' that contain packages matching the version '{version}'""" ) sys.exit(1) return sorted(architectures) except ClientError as e: print(f"[ERROR]: Failed to list architectures: {e}") sys.exit(1) except NoCredentialsError: print("[ERROR]: AWS credentials not configured") print("Please configure credentials via the IAM role") sys.exit(1) def exists_version_single_arch( s3_client, bucket_name: str, bucket_prefix: str, arch: str, version: str ) -> bool: """Check if an architecture folder contains files with the specified version.""" try: paginator = s3_client.get_paginator("list_objects_v2") pages = paginator.paginate(Bucket=bucket_name, Prefix=f"{bucket_prefix}{arch}/") for page in pages: if "Contents" not in page: continue for obj in page["Contents"]: if version in obj["Key"]: return True return False except ClientError: return False def exists_version_multi_arch(s3_client, bucket, prefix, directory, version): """Check whether a directory contains files matching the requested version. Args: s3_client: boto3 S3 client bucket: S3 bucket name prefix: Base S3 prefix directory: Optional subdirectory to search under version: Version pattern to search for Returns: True if at least one matching file exists, otherwise False """ paginator = s3_client.get_paginator("list_objects_v2") if directory is None: prefix_to_use = prefix else: prefix_to_use = f"{prefix}{directory}/" pages = paginator.paginate( Bucket=bucket, Prefix=prefix_to_use, ) for page in pages: if "Contents" not in page: continue for obj in page["Contents"]: if version in obj["Key"]: return True return False def list_packages_multi_arch( s3_client, bucket, prefix, version, architectures=None, ): """List multi-arch packages matching the requested version and architectures. Args: s3_client: boto3 S3 client bucket: S3 bucket name prefix: S3 prefix containing multi-arch packages version: Version pattern to filter packages architectures: Optional list of architectures used for filtering Returns: List of tuples (key, size) for matching packages. """ paginator = s3_client.get_paginator("list_objects_v2") pages = paginator.paginate( Bucket=bucket, Prefix=prefix, ) packages_to_promote = [] for page in pages: if "Contents" not in page: continue for obj in page["Contents"]: key = obj["Key"] filename = key.split("/")[-1] # Skip directories and generated indexes if not filename or filename == "index.html": continue # Match either: # rocm_sdk_core-7.13.0rc0-... # or: # torch-2.10.0+rocm7.13.0rc0-... matches_version = version in filename or f"rocm{version}" in filename if matches_version and is_allowed_multi_arch_package( filename, architectures, ): size = obj["Size"] packages_to_promote.append((key, size)) if not packages_to_promote: print(f"[ERROR]: No packages found for version {version}") sys.exit(1) return packages_to_promote def list_packages_for_arch( s3_client, bucket_name: str, bucket_prefix: str, arch: str, version: str ) -> Tuple[List[Tuple[str, int]], List[Tuple[str, int]], List[Tuple[str, int]]]: """List all packages for an architecture matching the version. Args: s3_client: boto3 S3 client bucket_name: S3 bucket name bucket_prefix: S3 bucket prefix (e.g., "v3/whl/") arch: Architecture name (e.g., "gfx950-dcgpu") version: Version pattern to filter (e.g., "7.10.0rc2") Returns: Tuple of (packages_to_promote, dependencies, unknown_packages) Each is a list of tuples (full_filepath_in_s3, size) """ prefix = f"{bucket_prefix}{arch}/" packages_to_promote = [] dependencies = [] unknown = [] try: paginator = s3_client.get_paginator("list_objects_v2") pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix) for page in pages: if "Contents" not in page: continue for obj in page["Contents"]: key = obj["Key"] filename = key.split("/")[-1] # Skip directories and index files if not filename or filename == "index.html": continue # Skip files that don't match version (for packages to promote) # Dependencies don't need version matching category = categorize_package(filename) if category == "promote": if version in filename: packages_to_promote.append((key, obj["Size"])) elif category == "dependency": dependencies.append((key, obj["Size"])) else: unknown.append((key, obj["Size"])) return packages_to_promote, dependencies, unknown except ClientError as e: print(f"[ERROR]: Failed to list packages for {arch}: {e}") return [], [], [] def list_tarball_for_package( s3_client, bucket_name: str, bucket_prefix: str, package: str | None, version: str, ) -> List[Tuple[str, int]]: """List tarballs and their sizes matching the requested package/version. Args: s3_client: boto3 S3 client bucket_name: S3 bucket name bucket_prefix: S3 bucket prefix (e.g., "tarball/") package: Package identifier to match in tarball filename. Examples: - "gfx942" - "gfx950-dcgpu" - "multiarch" If None, all tarballs matching the version are returned. version: Version pattern to filter (e.g., "7.10.0rc2") Returns: List of tuples (tarball_key, size) for matching tarballs """ try: paginator = s3_client.get_paginator("list_objects_v2") pages = paginator.paginate( Bucket=bucket_name, Prefix=bucket_prefix, ) tarballs = [] for page in pages: if "Contents" not in page: continue for obj in page["Contents"]: key = obj["Key"] filename = key.split("/")[-1] # Skip directories and index files if not filename or filename == "index.html": continue # Only tarballs for requested version if ( filename.startswith("therock-dist") and filename.endswith(".tar.gz") and version in filename ): # If package filter specified, enforce it if package is not None and package + "-" + version not in filename: continue tarballs.append((key, obj["Size"])) return tarballs except ClientError as e: print(f"[ERROR]: Failed to list tarballs for {package}: {e}") return [] def download_file(s3_client, bucket_name: str, key: str, local_path: Path) -> bool: """Download a single file from S3. Args: s3_client: boto3 S3 client bucket_name: S3 bucket name key: S3 object key local_path: Local file path to save to Returns: True if successful, False otherwise """ try: local_path.parent.mkdir(parents=True, exist_ok=True) s3_client.download_file(bucket_name, key, str(local_path)) return True except ClientError as e: print(f" [ERROR]: Failed to download {key}: {e}") return False def handle_multi_arch_downloads( s3_client, version, output_dir, architectures, bucket_name, bucket_prefix, include_tarballs, tarball_bucket_name, tarball_bucket_prefix, list_multi_arch_packages, include_package_globs=None, exclude_package_globs=None, ): """Handle multi-arch package listing and downloads.""" # Validate that packages exist if not exists_version_multi_arch( s3_client, bucket_name, bucket_prefix, None, version ): print(f"[ERROR]: No packages found for version {version}") sys.exit(1) packages_to_promote = list_packages_multi_arch( s3_client, bucket_name, bucket_prefix, version, architectures, ) packages_to_promote = filter_package_entries( packages_to_promote, include_package_globs=include_package_globs, exclude_package_globs=exclude_package_globs, ) if not packages_to_promote: print("[ERROR]: No packages left after applying package filters") sys.exit(1) if list_multi_arch_packages: print("\nPackages") print("-" * 60) for key, size in packages_to_promote: filename = key.split("/")[-1] print(f" - {filename} ({size / BYTES_TO_MB:.2f} MB)") total_size = sum(size for _, size in packages_to_promote) print("\n" + "=" * 60) print(f"TOTAL SIZE: {total_size / BYTES_TO_MB:.2f} MB") return if include_tarballs: tarball_output_dir = output_dir / "tarball-multi-arch" tarball_output_dir.mkdir(parents=True, exist_ok=True) print(f"Tarball output directory: {tarball_output_dir.absolute()}") print(f"\nOutput directory: {output_dir.absolute()}") success, fail = download_multi_arch_packages( s3_client, bucket_name, packages_to_promote, output_dir, ) if include_tarballs: tar_success, tar_fail = download_tarball( s3_client, tarball_bucket_name, tarball_bucket_prefix, None, version, tarball_output_dir, ) success += tar_success fail += tar_fail print("\n" + "=" * 80) print("DOWNLOAD COMPLETE (MULTI-ARCH)") print("=" * 80) print(f"Total successful downloads: {success}") print(f"Total failed downloads: {fail}") if architectures: print(f"\nArchitectures downloaded: {', '.join(architectures)}") if fail > 0: print("\nWARNING: Some downloads failed!") sys.exit(1) def download_multi_arch_packages( s3_client, bucket, packages_to_promote, output_dir, ): """Download multi-arch packages. Args: s3_client: boto3 S3 client bucket: S3 bucket name packages_to_promote: List of tuples (key, size) for packages to download output_dir: Local output directory Returns: Tuple of (successful_downloads, failed_downloads) """ wheels_dir = output_dir / "wheels" wheels_dir.mkdir(parents=True, exist_ok=True) total_success = 0 total_fail = 0 print("\nDownloading packages") print("=" * 80) for key, size in packages_to_promote: filename = key.split("/")[-1] local_path = wheels_dir / filename if local_path.exists(): print(f" SKIP (exists): {filename}") total_success += 1 continue print(f" Downloading: {filename} " f"({size / BYTES_TO_MB:.2f} MB)") if download_file(s3_client, bucket, key, local_path): total_success += 1 else: total_fail += 1 return total_success, total_fail def download_packages( s3_client, bucket_name: str, bucket_prefix: str, arch: str, version: str, output_dir: Path, include_dependencies: bool = False, include_package_globs: List[str] | None = None, exclude_package_globs: List[str] | None = None, ) -> Tuple[int, int]: """Download packages for an architecture. By default, only packages to promote are downloaded. Unknown packages are always skipped. Args: s3_client: boto3 S3 client bucket_name: S3 bucket name bucket_prefix: S3 bucket prefix (e.g., "v3/whl/") arch: Architecture name version: Version pattern output_dir: Base output directory include_dependencies: Include dependency packages in download (default: False) Returns: Tuple of (successful_downloads, failed_downloads) """ print(f"\nProcessing architecture: {arch}") print("=" * 80) packages_to_promote, dependencies, unknown = list_packages_for_arch( s3_client, bucket_name, bucket_prefix, arch, version ) print(f" Packages to promote: {len(packages_to_promote)}") print(f" Dependencies found: {len(dependencies)}") if unknown: print(f" Unknown packages (skipped): {len(unknown)}") for key in unknown: print(f" - {key[0].split('/')[-1]}") print("") print("-" * 80) arch_dir = output_dir / arch success_count = 0 fail_count = 0 if include_dependencies: all_packages = packages_to_promote + dependencies print( f" Downloading {len(all_packages)} packages to promote and their dependencies for {arch} with version {version}..." ) else: all_packages = packages_to_promote print( f" Downloading {len(all_packages)} python packages to promote for {arch} with version {version}..." ) all_packages = filter_package_entries( all_packages, include_package_globs=include_package_globs, exclude_package_globs=exclude_package_globs, ) if not all_packages: print( f" [ERROR]: No packages left for {arch} with version {version} after applying package filters. Skipping!" ) return 0, 0 for idx, obj in enumerate(all_packages): key, size = obj filename = key.split("/")[-1] local_path = arch_dir / filename # Skip if already exists if local_path.exists(): print( f" ({idx+1}/{len(all_packages)}) \tSKIP (exists): {filename} ({size/BYTES_TO_MB:.2f} MB)" ) success_count += 1 continue print( f" ({idx+1}/{len(all_packages)}) \tDownloading: {filename} ({size/BYTES_TO_MB:.2f} MB)" ) if download_file(s3_client, bucket_name, key, local_path): success_count += 1 else: fail_count += 1 return success_count, fail_count def download_tarball( s3_client, bucket_name: str, bucket_prefix: str, package: str | None, version: str, output_dir: Path, ) -> Tuple[int, int]: """Download tarballs matching the requested package/version. Args: s3_client: boto3 S3 client bucket_name: S3 bucket name bucket_prefix: S3 bucket prefix package: Package identifier used for tarball filtering. If None, all tarballs matching the version are downloaded. version: Version pattern output_dir: Base output directory Returns: Tuple of (successful_download, failed_download) """ tarballs = list_tarball_for_package( s3_client, bucket_name, bucket_prefix, package, version ) if not tarballs: print( f" [ERROR]: No tarball found for {package} with version {version}. Skipping!" ) return 0, 1 success_count = 0 fail_count = 0 print(f" Downloading tarball...") for name, size in tarballs: filename = name.split("/")[-1] local_path = output_dir / filename # Skip if already exists if local_path.exists(): print(f" \tSKIP (exists): {filename} ({size/BYTES_TO_MB:.2f} MB)") success_count += 1 continue print(f" \tDownloading: {filename} ({size/BYTES_TO_MB:.2f} MB)") if download_file(s3_client, bucket_name, name, local_path): success_count += 1 else: fail_count += 1 return success_count, fail_count def parse_arguments(argv): parser = argparse.ArgumentParser( description="Download prerelease packages from S3 for promotion to release", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Download all architectures for version 7.10.0rc2 python download_python_packages.py --version=7.10.0rc2 --output-dir=./downloads # Download only specific architecture python download_python_packages.py --version=7.10.0rc2 --arch=gfx950-dcgpu --output-dir=./downloads # Download multiple specific architectures and including tarballs python download_python_packages.py --version=7.10.0rc2 --arch=gfx1151,gfx950-dcgpu --output-dir=./downloads --include-tarballs # Download packages to promote AND their dependencies python download_python_packages.py --version=7.10.0rc2 --output-dir=./downloads --include-dependencies # Use custom bucket prefix python download_python_packages.py --version=7.10.0rc2 --output-dir=./downloads --bucket-prefix=v3/whl/ # List available architectures python download_python_packages.py --version=7.10.0rc2 --list-archs # List all packages per architecture with sizes python download_python_packages.py --version=7.10.0rc2 --list-packages-per-arch # List packages and tarballs with sizes python download_python_packages.py --version=7.10.0rc2 --list-packages-per-arch --include-tarballs # Download multi-arch packages python download_python_packages.py --version=7.13.0rc1 --multi-arch --output-dir=./downloads --bucket-prefix=v4/whl/ # Download multi-arch packages including tarballs python download_python_packages.py --version=7.13.0rc1 --multi-arch --output-dir=./downloads --bucket-prefix=v4/whl/ --include-tarballs # List multi-arch packages python download_python_packages.py --version=7.13.0rc1 --multi-arch --list-multi-arch-packages --bucket-prefix=v4/whl/ # List multi-arch packages while excluding PyTorch 2.8 packages python download_python_packages.py --version=7.13.0rc1 --multi-arch --list-multi-arch-packages --bucket-prefix=v4/whl/ --exclude-package-glob "torch-2.8.*" # Download packages while excluding specific package variants python download_python_packages.py --version=7.13.0rc1 --output-dir=./downloads --exclude-package-glob "torch-2.8.*" --exclude-package-glob "torchaudio-2.8.*" --exclude-package-glob "torchvision-0.23.*" """, ) parser.add_argument( "--version", required=True, help="Version pattern to download (e.g., '7.10.0rc2')", ) parser.add_argument( "--output-dir", type=Path, help="Output directory for downloaded packages (required unless --list-archs or --list-packages-per-arch)", ) parser.add_argument( "--arch", help="Download only this specific architecture (e.g. 'gfx950-dcgpu,gfx1151') (default: all architectures)", ) parser.add_argument( "--bucket", default="therock-prerelease-python", help="S3 bucket name (default: therock-prerelease-python)", ) parser.add_argument( "--bucket-prefix", default="v3/whl/", help="S3 bucket prefix for packages (default: v3/whl/)", ) parser.add_argument( "--include-dependencies", action=argparse.BooleanOptionalAction, default=False, help="Include dependency packages in download (default: only packages to promote)", ) parser.add_argument( "--include-package-glob", action="append", default=[], help=( "Only include package filenames matching this glob pattern. " "Can be specified multiple times. Example: --include-package-glob 'torch-2.10.*'" ), ) parser.add_argument( "--exclude-package-glob", action="append", default=[], help=( "Exclude package filenames matching this glob pattern. " "Can be specified multiple times. Example: --exclude-package-glob 'torch-2.8.*'" ), ) parser.add_argument( "--include-tarballs", action=argparse.BooleanOptionalAction, default=False, help="Download also tarball packages (default: False)", ) parser.add_argument( "--tarball-bucket", default="therock-prerelease-tarball", help="S3 bucket name for tarball packages (default: therock-prerelease-tarball)", ) parser.add_argument( "--tarball-bucket-prefix", default="v3/tarball/", help="S3 bucket prefix for tarball packages (default: v3/tarball/)", ) parser.add_argument( "--tarball-output-dir", type=Path, help="Output directory for downloaded tarball packages (default: < --output-dir >/tarball)", ) parser.add_argument( "--list-archs", action=argparse.BooleanOptionalAction, default=False, help="Only list available architectures, do not download", ) parser.add_argument( "--list-packages-per-arch", action=argparse.BooleanOptionalAction, default=False, help="List all packages per architecture, do not download", ) parser.add_argument( "--multi-arch", action=argparse.BooleanOptionalAction, default=False, help="Enable multi-arch package handling. Requires bucket prefix compatible with multi-arch layout (e.g. v4/whl/)", ) parser.add_argument( "--list-multi-arch-packages", action=argparse.BooleanOptionalAction, default=False, help="List all multi-arch packages matching version", ) args = parser.parse_args(argv) if args.arch: args.arch = [arch.strip() for arch in args.arch.split(",")] # Validate arguments if args.include_tarballs: if args.tarball_output_dir is None and args.output_dir is not None: args.tarball_output_dir = args.output_dir / "tarball" else: args.tarball_output_dir = ( Path(args.tarball_output_dir) if args.tarball_output_dir is not None else None ) if ( not args.list_archs and not args.list_packages_per_arch and not args.list_multi_arch_packages and not args.output_dir ): parser.error( "--output-dir is required unless --list-archs, --list-packages-per-arch, or --list-multi-arch-packages is specified" ) return args def print_packages_per_arch( s3_client, bucket_name: str, bucket_prefix: str, architectures: List[str], version: str, include_tarballs: bool = False, tarball_bucket_name: str = None, tarball_bucket_prefix: str = None, include_package_globs: List[str] | None = None, exclude_package_globs: List[str] | None = None, ) -> Dict[str, Dict[str, List[str]]]: print( "\n--list-packages-per-arch specified, listing packages and their sizes without download" ) print("=" * 80) all_packages = {} total_size_promote = 0 total_size_deps = 0 total_size_unknown = 0 total_size_tarball = 0 for arch in architectures: print(f"\n{arch}:") print("-" * 80) packages_to_promote, dependencies, unknown = list_packages_for_arch( s3_client, bucket_name, bucket_prefix, arch, version ) packages_to_promote = filter_package_entries( packages_to_promote, include_package_globs=include_package_globs, exclude_package_globs=exclude_package_globs, ) if include_package_globs or exclude_package_globs: dependencies = filter_package_entries( dependencies, include_package_globs=include_package_globs, exclude_package_globs=exclude_package_globs, ) all_packages[arch] = { "packages_to_promote": packages_to_promote, "dependencies": dependencies, "unknown": unknown, } size_promote = 0 size_dependencies = 0 size_unknown = 0 size_tarball = 0 print(f" Packages to promote ({len(packages_to_promote)}):") for pkg in all_packages[arch]["packages_to_promote"]: name, size = pkg name = name.split("/")[-1] size_promote += size total_size_promote += size print(f" - {name} ({size/BYTES_TO_MB:.2f} MB)") print(f"\n Dependencies found ({len(dependencies)}):") for pkg in all_packages[arch]["dependencies"]: name, size = pkg name = name.split("/")[-1] size_dependencies += size total_size_deps += size print(f" - {name} ({size/BYTES_TO_MB:.2f} MB)") if unknown: print(f"\n Unknown packages ({len(unknown)}):") for pkg in all_packages[arch]["unknown"]: name, size = pkg name = name.split("/")[-1] size_unknown += size total_size_unknown += size print(f" - {name} ({size/BYTES_TO_MB:.2f} MB)") if include_tarballs: print(f"\n Tarball:") tarballs = list_tarball_for_package( s3_client, tarball_bucket_name, tarball_bucket_prefix, arch, version ) for tarball in tarballs: tarball_name, tarball_size = tarball size_tarball += tarball_size total_size_tarball += tarball_size size_mb = tarball_size / BYTES_TO_MB print(f" - {tarball_name.split('/')[-1]} ({size_mb:.2f} MB)") if not tarballs: print( f" [WARN]: No tarball found for {arch} with version {version}. Skipping!" ) print("") print("-" * 80) print("") print(f" Space needed for {arch}:") print(f" Packages to promote: {size_promote/BYTES_TO_MB:>12.2f} MB") print(f" Dependencies packages: {size_dependencies/BYTES_TO_MB:>12.2f} MB") print(f" Unknown packages: {size_unknown/BYTES_TO_MB:>12.2f} MB") print(f" Tarball: {size_tarball/BYTES_TO_MB:>12.2f} MB") print("") print("-" * 80) print("") print("") print("-" * 80) print("") print(" Total space needed for all architectures:") print(f" All packages to promote: {total_size_promote/BYTES_TO_MB:>12.2f} MB") print(f" All dependencies packages: {total_size_deps/BYTES_TO_MB:>12.2f} MB") print(f" All unknown packages: {total_size_unknown/BYTES_TO_MB:>12.2f} MB") print(f" All tarball: {total_size_tarball/BYTES_TO_MB:>12.2f} MB") print("") if include_tarballs: print( f" All packages to promote and tarballs: {(total_size_promote + total_size_tarball)/BYTES_TO_MB:>12.2f} MB" ) print("-" * 80) print("") return all_packages def download_prerelease_packages( version: str, output_dir: Path = None, architectures: List[str] = None, bucket_name: str = "therock-prerelease-python", bucket_prefix: str = "v3/whl/", include_dependencies: bool = False, multi_arch: bool = False, list_multi_arch_packages: bool = False, include_tarballs: bool = False, tarball_bucket_name: str = "therock-prerelease-tarball", tarball_bucket_prefix: str = "v3/tarball/", tarball_output_dir: Path = None, list_archs: bool = False, list_packages_per_arch: bool = False, include_package_globs: List[str] | None = None, exclude_package_globs: List[str] | None = None, ) -> Union[List[str], Dict[str, Dict[str, List[str]]], Tuple[int, int, List[str]]]: """Download prerelease packages from S3 bucket for promotion to release. Args: version: Version pattern to download (e.g., '7.10.0rc2') output_dir: Output directory for downloaded packages (required unless list_archs or list_packages_per_arch) architectures: Download only these specific architectures (default: all architectures) bucket_name: S3 bucket name (default: therock-prerelease-python) bucket_prefix: S3 bucket prefix for packages (default: v3/whl/) include_dependencies: Include dependency packages in download (default: False). Ignored if list_archs or list_packages_per_arch is True. include_tarballs: Include tarballs in download and listings (default: False). tarball_bucket_name: S3 bucket name for tarball packages (default: therock-prerelease-tarball) tarball_bucket_prefix: S3 bucket prefix for tarball packages (default: v3/tarball/) tarball_output_dir: Output directory for downloaded tarball packages (default: < --output-dir >/tarballs) list_archs: Only list available architectures, do not download (default: False). Set by CLI flag --list-archs list_packages_per_arch: List all packages per architecture and their sizes, do not download (default: False) Returns: If list_archs=True: List of architecture names If list_packages_per_arch=True: Dict mapping arch to dict of package categories Otherwise: Tuple of (total_success, total_fail, architectures) of downloaded packages Raises: SystemExit: If AWS credentials are not configured, no architectures found, or downloads fail """ if ( not list_archs and not list_packages_per_arch and not list_multi_arch_packages and output_dir is None ): print( "[ERROR]: output_dir is required unless list_archs=True, list_packages_per_arch=True, or list_multi_arch_packages=True" ) sys.exit(1) print("=" * 80) print("Download Prerelease Packages") print("=" * 80) print(f"Bucket: {bucket_name}") print(f"Version: {version}") if architectures: print(f"Architectures: {architectures} (user-specified)") else: print("Architecture: ALL") print("=" * 80) s3_client = boto3.client("s3") if multi_arch: return handle_multi_arch_downloads( s3_client=s3_client, version=version, output_dir=output_dir, architectures=architectures, bucket_name=bucket_name, bucket_prefix=bucket_prefix, include_tarballs=include_tarballs, tarball_bucket_name=tarball_bucket_name, tarball_bucket_prefix=tarball_bucket_prefix, list_multi_arch_packages=list_multi_arch_packages, include_package_globs=include_package_globs, exclude_package_globs=exclude_package_globs, ) # List architectures if architectures: # Validate they exist for arch in architectures: if not exists_version_single_arch( s3_client, bucket_name, bucket_prefix, arch, version ): print( f"[ERROR]: Architecture {arch} not found or has no packages with version {version}" ) sys.exit(1) print(f"\nVerified user-specified architectures exist: {architectures}") else: architectures = list_architectures( s3_client, bucket_name, bucket_prefix, version ) print(f"\nFound {len(architectures)} architecture(s)") if list_archs: print("\n--list-archs specified, exiting without download") return architectures if list_packages_per_arch: return print_packages_per_arch( s3_client=s3_client, bucket_name=bucket_name, bucket_prefix=bucket_prefix, architectures=architectures, version=version, include_tarballs=include_tarballs, tarball_bucket_name=tarball_bucket_name, tarball_bucket_prefix=tarball_bucket_prefix, include_package_globs=include_package_globs, exclude_package_globs=exclude_package_globs, ) # Create output directory output_dir.mkdir(parents=True, exist_ok=True) print(f"\nOutput directory: {output_dir.absolute()}") if include_tarballs: tarball_output_dir.mkdir(parents=True, exist_ok=True) print(f"Tarball output directory: {tarball_output_dir.absolute()}") # Download packages for each architecture total_success = 0 total_fail = 0 for arch in architectures: success_count = 0 fail_count = 0 success, fail = download_packages( s3_client, bucket_name, bucket_prefix, arch, version, output_dir, include_dependencies, include_package_globs, exclude_package_globs, ) success_count += success fail_count += fail total_success += success total_fail += fail if include_tarballs: success, fail = download_tarball( s3_client, tarball_bucket_name, tarball_bucket_prefix, arch, version, tarball_output_dir, ) success_count += success fail_count += fail total_success += success total_fail += fail print(f"\n Summary for {arch}:") print(f" Successful: {success_count}") print(f" Failed: {fail_count}") # Final summary print("\n" + "=" * 80) print("DOWNLOAD COMPLETE") print("=" * 80) print(f"Total architectures: {len(architectures)}") print(f"Total successful downloads: {total_success}") print(f"Total failed downloads: {total_fail}") print(f"\nOutput directory: {output_dir.absolute()}/") if include_tarballs: print(f"Tarball output directory: {tarball_output_dir.absolute()}") print(f"\nArchitectures downloaded: {architectures}") if total_fail > 0: print("\nWARNING: Some downloads failed!") sys.exit(1) print("\nFor next steps check: build_tools/packaging/how_to_do_release.md") print("=" * 80) return total_success, total_fail, architectures if __name__ == "__main__": args = parse_arguments(sys.argv[1:]) download_prerelease_packages( version=args.version, output_dir=args.output_dir, architectures=args.arch, bucket_name=args.bucket, multi_arch=args.multi_arch, list_multi_arch_packages=args.list_multi_arch_packages, bucket_prefix=args.bucket_prefix, include_dependencies=args.include_dependencies, include_tarballs=args.include_tarballs, tarball_bucket_name=args.tarball_bucket, tarball_bucket_prefix=args.tarball_bucket_prefix, tarball_output_dir=args.tarball_output_dir, list_archs=args.list_archs, list_packages_per_arch=args.list_packages_per_arch, include_package_globs=args.include_package_glob, exclude_package_globs=args.exclude_package_glob, )