diff --git a/build_tools/packaging/download_prerelease_packages.py b/build_tools/packaging/download_python_packages.py similarity index 76% rename from build_tools/packaging/download_prerelease_packages.py rename to build_tools/packaging/download_python_packages.py index 474c646139d..57148544112 100644 --- a/build_tools/packaging/download_prerelease_packages.py +++ b/build_tools/packaging/download_python_packages.py @@ -25,47 +25,47 @@ TYPICAL USAGE (Command Line): # Download all 7.10.0rc2 packages for all architectures: - python ./build_tools/packaging/download_prerelease_packages.py \ + 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_prerelease_packages.py \ + 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_prerelease_packages.py \ + 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_prerelease_packages.py \ + 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_prerelease_packages.py \ + 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_prerelease_packages.py \ + 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_prerelease_packages.py \ + 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_prerelease_packages.py \ + python ./build_tools/packaging/download_python_packages.py \ --version=7.10.0rc2 \ --output-dir=./downloads \ --include-dependencies @@ -130,6 +130,7 @@ "rocm_sdk_core", "rocm_sdk_devel", "rocm_sdk_libraries-*", + "rocm_profiler", "torch", "torchaudio", "torchvision", @@ -140,6 +141,27 @@ "jaxlib", } +PACKAGES_TO_PROMOTE_MULTI_ARCH = { + "torch", + "torchaudio", + "torchvision", + "triton", + "apex", + "rocm", + "rocm_sdk_core", + "rocm_sdk_devel", + "rocm_sdk_libraries", + "rocm_profiler", + "rocm_bootstrap", + # device packages (IMPORTANT) + "amd_torch_device", + "amd_torchvision_device", + "rocm_sdk_device", + # 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 = { @@ -167,6 +189,64 @@ } +def is_allowed_multi_arch_package(filename: str) -> bool: + """ + Check if filename belongs to allowed packages for multi-arch flow. + Uses BOTH: + - PACKAGES_TO_PROMOTE + - PACKAGES_TO_PROMOTE_MULTI_ARCH + + Supports: + - exact match + - wildcard '*' suffix (legacy) + - prefix match (for device packages) + """ + base = filename.split("-")[0] + + def matches(pkg: str) -> bool: + # Handle wildcard like "rocm_sdk_libraries-*" + if pkg.endswith("*"): + return base.startswith(pkg[:-1]) + + # Exact match OR prefix match (for device families) + return base == pkg or base.startswith(pkg) + + return any( + matches(pkg) for pkg in (PACKAGES_TO_PROMOTE | PACKAGES_TO_PROMOTE_MULTI_ARCH) + ) + + +def matches_requested_arches(filename: str, requested_arches: list[str]) -> bool: + """ + Check whether a multi-arch package matches one of the requested gfx arches. + + Rules: + - If no requested arches are provided -> allow + - If filename does not contain gfx -> allow (generic package) + - If filename contains gfx -> require one requested arch match + """ + + if not requested_arches: + return True + + lower_filename = filename.lower() + + # Generic package (not gfx-specific) + if "gfx" not in lower_filename: + return True + + normalized_arches = set() + + for arch in requested_arches: + lower_arch = arch.lower() + + normalized_arches.add(lower_arch) + normalized_arches.add(lower_arch.replace("-", "_")) + normalized_arches.add(lower_arch.split("-")[0]) + + return any(arch in lower_filename for arch in normalized_arches) + + def categorize_package(filename: str) -> str: """Categorize a package file. @@ -263,6 +343,78 @@ def has_version_in_arch( return False +def has_version_in_directory(s3_client, bucket, prefix, directory, version): + 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_verbose( + s3_client, + bucket, + prefix, + version, + architectures=None, +): + paginator = s3_client.get_paginator("list_objects_v2") + + total_size = 0 + found = False + + print("\nPackages") + print("-" * 60) + + pages = paginator.paginate( + Bucket=bucket, + Prefix=prefix, + ) + + for page in pages: + if "Contents" not in page: + continue + + for obj in page["Contents"]: + key = obj["Key"] + filename = key.split("/")[-1] + + if not filename or filename == "index.html": + continue + + if version in filename and is_allowed_multi_arch_package(filename): + if not matches_requested_arches(filename, architectures): + continue + found = True + size = obj["Size"] + + print(f" - {filename} ({size / BYTES_TO_MB:.2f} MB)") + total_size += size + + if not found: + print(f"[ERROR]: No packages found for version {version}") + sys.exit(1) + + print("\n" + "=" * 60) + print(f"TOTAL SIZE: {total_size / BYTES_TO_MB:.2f} MB") + + 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]]]: @@ -387,6 +539,69 @@ def download_file(s3_client, bucket_name: str, key: str, local_path: Path) -> bo return False +def download_multi_arch_packages( + s3_client, + bucket, + prefix, + version, + output_dir, + architectures=None, +): + paginator = s3_client.get_paginator("list_objects_v2") + + wheels_dir = output_dir / "wheels" + wheels_dir.mkdir(parents=True, exist_ok=True) + + total_success = 0 + total_fail = 0 + + pages = paginator.paginate( + Bucket=bucket, + Prefix=prefix, + ) + + print("\nDownloading packages") + print("=" * 80) + + for page in pages: + if "Contents" not in page: + continue + + for obj in page["Contents"]: + key = obj["Key"] + filename = key.split("/")[-1] + + if not filename or filename == "index.html": + continue + + if version not in filename: + continue + + if not is_allowed_multi_arch_package(filename): + continue + + if not matches_requested_arches(filename, architectures): + print(f" SKIP (arch filter): {filename}") + continue + + local_path = wheels_dir / filename + + if local_path.exists(): + print(f" SKIP (exists): {filename}") + total_success += 1 + continue + + size = obj["Size"] + print(f" Downloading: {filename} ({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, @@ -423,7 +638,7 @@ def download_packages( if unknown: print(f" Unknown packages (skipped): {len(unknown)}") for key in unknown: - print(f" - {key[0].split("/")[-1]}") + print(f" - {key[0].split('/')[-1]}") print("") print("-" * 80) @@ -535,28 +750,28 @@ def parse_arguments(argv): epilog=""" Examples: # Download all architectures for version 7.10.0rc2 - python download_prerelease_packages.py --version=7.10.0rc2 --output-dir=./downloads + python download_python_packages.py --version=7.10.0rc2 --output-dir=./downloads # Download only specific architecture - python download_prerelease_packages.py --version=7.10.0rc2 --arch=gfx950-dcgpu --output-dir=./downloads + python download_python_packages.py --version=7.10.0rc2 --arch=gfx950-dcgpu --output-dir=./downloads # Download multiple specific architectures and including tarballs - python download_prerelease_packages.py --version=7.10.0rc2 --arch=gfx1151,gfx950-dcgpu --output-dir=./downloads --include-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_prerelease_packages.py --version=7.10.0rc2 --output-dir=./downloads --include-dependencies + python download_python_packages.py --version=7.10.0rc2 --output-dir=./downloads --include-dependencies # Use custom bucket prefix - python download_prerelease_packages.py --version=7.10.0rc2 --output-dir=./downloads --bucket-prefix=v3/whl/ + python download_python_packages.py --version=7.10.0rc2 --output-dir=./downloads --bucket-prefix=v3/whl/ # List available architectures - python download_prerelease_packages.py --version=7.10.0rc2 --list-archs + python download_python_packages.py --version=7.10.0rc2 --list-archs # List all packages per architecture with sizes - python download_prerelease_packages.py --version=7.10.0rc2 --list-packages-per-arch + python download_python_packages.py --version=7.10.0rc2 --list-packages-per-arch # List packages and tarballs with sizes - python download_prerelease_packages.py --version=7.10.0rc2 --list-packages-per-arch --include-tarballs + python download_python_packages.py --version=7.10.0rc2 --list-packages-per-arch --include-tarballs """, ) @@ -632,6 +847,20 @@ def parse_arguments(argv): 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: @@ -647,9 +876,14 @@ def parse_arguments(argv): 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.output_dir: + 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 or --list-packages-per-arch is specified" + "--output-dir is required unless --list-archs, --list-packages-per-arch, or --list-multi-arch-packages is specified" ) return args @@ -727,7 +961,7 @@ def print_packages_per_arch( 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)") + 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!" @@ -771,6 +1005,8 @@ def download_prerelease_packages( 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/", @@ -804,10 +1040,14 @@ def download_prerelease_packages( Raises: SystemExit: If AWS credentials are not configured, no architectures found, or downloads fail """ - # Validate arguments - if not list_archs and not list_packages_per_arch and output_dir is None: + 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 or list_packages_per_arch=True" + "[ERROR]: output_dir is required unless list_archs=True, list_packages_per_arch=True, or list_multi_arch_packages=True" ) sys.exit(1) @@ -816,6 +1056,7 @@ def download_prerelease_packages( print("=" * 80) print(f"Bucket: {bucket_name}") print(f"Version: {version}") + if architectures: print(f"Architectures: {architectures} (user-specified)") else: @@ -823,6 +1064,47 @@ def download_prerelease_packages( print("=" * 80) s3_client = boto3.client("s3") + if multi_arch: + # Validate that packages exist + if not has_version_in_directory( + s3_client, bucket_name, bucket_prefix, None, version + ): + print(f"[ERROR]: No packages found for version {version}") + sys.exit(1) + + if list_multi_arch_packages: + list_packages_multi_arch_verbose( + s3_client, + bucket_name, + bucket_prefix, + version, + architectures, + ) + return + + output_dir.mkdir(parents=True, exist_ok=True) + print(f"\nOutput directory: {output_dir.absolute()}") + + success, fail = download_multi_arch_packages( + s3_client, + bucket_name, + bucket_prefix, + version, + output_dir, + architectures, + ) + + print("\n" + "=" * 80) + print("DOWNLOAD COMPLETE (MULTI-ARCH)") + print("=" * 80) + print(f"Total successful downloads: {success}") + print(f"Total failed downloads: {fail}") + + if fail > 0: + sys.exit(1) + + return + # List architectures if architectures: @@ -937,6 +1219,8 @@ def download_prerelease_packages( 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, diff --git a/build_tools/packaging/how_to_do_release.md b/build_tools/packaging/how_to_do_release.md index e56b2f794fa..211cc524668 100644 --- a/build_tools/packaging/how_to_do_release.md +++ b/build_tools/packaging/how_to_do_release.md @@ -13,21 +13,64 @@ The steps for the release are: Need: -- `build_tools/packaging/download_prerelease_packages.py` +- `build_tools/packaging/download_python_packages.py` - IAM role: read and list bucket access for therock-prerelease-python and therock-prerelease-tarball +By default, the script operates in single-architecture mode using per-arch +directories (e.g., `v3/whl//`). + Example: Download all prerelease candidates 7.10.0rc2 to ./promotion/download ```bash # 1. (Optional) Check which architectures are available -python build_tools/packaging/download_prerelease_packages.py --version=7.10.0rc2 --list-archs +python build_tools/packaging/download_python_packages.py --version=7.10.0rc2 --list-archs # 2. (Recommended) Check which packages are available and their sizes # Make sure you have enough disk space available for what you want to download! -python build_tools/packaging/download_prerelease_packages.py --version=7.10.0rc2 --list-packages-per-arch --include-tarballs +python build_tools/packaging/download_python_packages.py --version=7.10.0rc2 --list-packages-per-arch --include-tarballs # 3. Download all ROCm/PyTorch packages that need promotion (all architectures) -python build_tools/packaging/download_prerelease_packages.py --version=7.10.0rc2 --output-dir=./promotion/download/ --include-tarballs +python build_tools/packaging/download_python_packages.py --version=7.10.0rc2 --output-dir=./promotion/download/ --include-tarballs +``` + +### Multi-arch packages (flat layout) + +Use this mode when packages are stored without per-architecture subdirectories +(e.g., `v4/whl/` layout). + +Requirements for multi-arch mode: + +- Must use a compatible bucket prefix (e.g., `v4/whl/`) +- Output is a flat directory (no `/` folders) +- All downloaded wheels are placed under `/wheels/`. + +#### List multi-arch packages + +```bash +python build_tools/packaging/download_python_packages.py \ + --version=7.10.0rc2 \ + --bucket-prefix=v4/whl/ \ + --multi-arch \ + --list-multi-arch-packages +``` + +#### Download multi-arch packages + +```bash +python build_tools/packaging/download_python_packages.py \ + --version=7.10.0rc2 \ + --bucket-prefix=v4/whl/ \ + --multi-arch \ + --output-dir=./promotion/download/ + +``` + +Output structure: + +``` +/ + wheels/ + *.whl ``` ## 2. Promote prerelease candidates to release @@ -63,7 +106,7 @@ Need: - IAM role: - for testing: write access to therock-testing-bucket - for production: write access to therock-release-python and therock-release-tarball -- Same folder structure as created by `download_prerelease_packages.py`: +- Same folder structure as created by `download_python_packages.py`: ``` /