Skip to content

Commit 0be39b5

Browse files
Add ROCm pip index mirroring for rocm7.14 dependencies (#8436)
## Summary Mirror ROCm SDK packages from AMD's multi-arch pip index (`https://repo.amd.com/rocm/whl-multi-arch/`) into download.pytorch.org dependency indexes for the `rocm7.14` target. - Add `torch_rocm` workflow option to `update-s3-dependencies.yml` - Register 28 ROCm packages under `torch_rocm` / `rocm7.14` in `PACKAGES_PER_PROJECT` (`rocm`, `rocm-sdk-core`, `rocm-sdk-libraries`, and all 25 `rocm-sdk-device-*` packages currently on the AMD index) - Route AMD packages via `is_amd_package()` and `get_package_source_url()` - Extend `get_packages_for_target()` so ROCm packages are included only for `rocm*` targets and `create-target` works for ROCm ## Test plan - [x] Verified AMD URL routing: ```bash python3 -c "from s3_management.update_dependencies import get_package_source_url, is_amd_package; assert is_amd_package('rocm-sdk-core'); assert get_package_source_url('rocm-sdk-core') == 'https://repo.amd.com/rocm/whl-multi-arch/rocm-sdk-core/'" ``` - [x] Dry-run torch_rocm update (all 28 packages fetched from AMD index): ```bash python3 s3_management/update_dependencies.py --package torch_rocm --dry-run ``` - [ ] After merge: trigger `Update S3 HTML dependencies` workflow with `package=torch_rocm`, `dryrun=enabled`, then `disabled` - [ ] Confirm `https://download.pytorch.org/whl/nightly/rocm7.14/rocm-sdk-core/index.html` serves absolute links to `repo.amd.com` Authored with assistance from Cursor Made with [Cursor](https://cursor.com) --------- Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent fc982f6 commit 0be39b5

2 files changed

Lines changed: 53 additions & 8 deletions

File tree

.github/workflows/update-s3-dependencies.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ on:
3030
- triton
3131
- torchtune
3232
- torch_xpu
33+
- torch_rocm
3334
- vllm
3435
target:
3536
description: 'Target to create (e.g., rocm7.2, cu130) - only for create-target mode'

s3_management/update_dependencies.py

Lines changed: 52 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -592,6 +592,34 @@
592592
"target": "cu132",
593593
},
594594
],
595+
"rocm": [{"project": "torch_rocm", "target": "rocm7.14"}],
596+
"rocm-sdk-core": [{"project": "torch_rocm", "target": "rocm7.14"}],
597+
"rocm-sdk-libraries": [{"project": "torch_rocm", "target": "rocm7.14"}],
598+
"rocm-sdk-device-gfx1010": [{"project": "torch_rocm", "target": "rocm7.14"}],
599+
"rocm-sdk-device-gfx1011": [{"project": "torch_rocm", "target": "rocm7.14"}],
600+
"rocm-sdk-device-gfx1012": [{"project": "torch_rocm", "target": "rocm7.14"}],
601+
"rocm-sdk-device-gfx1030": [{"project": "torch_rocm", "target": "rocm7.14"}],
602+
"rocm-sdk-device-gfx1031": [{"project": "torch_rocm", "target": "rocm7.14"}],
603+
"rocm-sdk-device-gfx1032": [{"project": "torch_rocm", "target": "rocm7.14"}],
604+
"rocm-sdk-device-gfx1033": [{"project": "torch_rocm", "target": "rocm7.14"}],
605+
"rocm-sdk-device-gfx1034": [{"project": "torch_rocm", "target": "rocm7.14"}],
606+
"rocm-sdk-device-gfx1035": [{"project": "torch_rocm", "target": "rocm7.14"}],
607+
"rocm-sdk-device-gfx1036": [{"project": "torch_rocm", "target": "rocm7.14"}],
608+
"rocm-sdk-device-gfx1100": [{"project": "torch_rocm", "target": "rocm7.14"}],
609+
"rocm-sdk-device-gfx1101": [{"project": "torch_rocm", "target": "rocm7.14"}],
610+
"rocm-sdk-device-gfx1102": [{"project": "torch_rocm", "target": "rocm7.14"}],
611+
"rocm-sdk-device-gfx1103": [{"project": "torch_rocm", "target": "rocm7.14"}],
612+
"rocm-sdk-device-gfx1150": [{"project": "torch_rocm", "target": "rocm7.14"}],
613+
"rocm-sdk-device-gfx1151": [{"project": "torch_rocm", "target": "rocm7.14"}],
614+
"rocm-sdk-device-gfx1152": [{"project": "torch_rocm", "target": "rocm7.14"}],
615+
"rocm-sdk-device-gfx1153": [{"project": "torch_rocm", "target": "rocm7.14"}],
616+
"rocm-sdk-device-gfx1200": [{"project": "torch_rocm", "target": "rocm7.14"}],
617+
"rocm-sdk-device-gfx1201": [{"project": "torch_rocm", "target": "rocm7.14"}],
618+
"rocm-sdk-device-gfx1250": [{"project": "torch_rocm", "target": "rocm7.14"}],
619+
"rocm-sdk-device-gfx908": [{"project": "torch_rocm", "target": "rocm7.14"}],
620+
"rocm-sdk-device-gfx90a": [{"project": "torch_rocm", "target": "rocm7.14"}],
621+
"rocm-sdk-device-gfx942": [{"project": "torch_rocm", "target": "rocm7.14"}],
622+
"rocm-sdk-device-gfx950": [{"project": "torch_rocm", "target": "rocm7.14"}],
595623
"arpeggio": [{"project": "triton"}],
596624
"caliper-reader": [{"project": "triton"}],
597625
"contourpy": [{"project": "triton"}],
@@ -786,12 +814,19 @@ def is_nvidia_package(pkg_name: str) -> bool:
786814
return pkg_name.startswith("nvidia-") or pkg_name.startswith("cuda-")
787815

788816

817+
def is_amd_package(pkg_name: str) -> bool:
818+
"""Check if a package is from AMD and should use repo.amd.com"""
819+
name = pkg_name.lower()
820+
return "rocm" in name or "amd" in name
821+
822+
789823
def get_package_source_url(pkg_name: str) -> str:
790824
"""Get the source URL for a package based on its type"""
791825
if is_nvidia_package(pkg_name):
792826
return f"https://pypi.nvidia.com/{pkg_name}/"
793-
else:
794-
return f"https://pypi.org/simple/{pkg_name}/"
827+
if is_amd_package(pkg_name):
828+
return f"https://repo.amd.com/rocm/whl-multi-arch/{pkg_name}/"
829+
return f"https://pypi.org/simple/{pkg_name}/"
795830

796831

797832
def download(url: str) -> bytes:
@@ -989,11 +1024,14 @@ def upload_package_using_simple_index(
9891024
Works for both NVIDIA and non-NVIDIA packages.
9901025
"""
9911026
source_url = get_package_source_url(pkg_name)
992-
is_nvidia = is_nvidia_package(pkg_name)
1027+
if is_nvidia_package(pkg_name):
1028+
source_label = "NVIDIA"
1029+
elif is_amd_package(pkg_name):
1030+
source_label = "AMD"
1031+
else:
1032+
source_label = "PyPI"
9931033

994-
print(
995-
f"Processing {pkg_name} using {'NVIDIA' if is_nvidia else 'PyPI'} Simple Index: {source_url}"
996-
)
1034+
print(f"Processing {pkg_name} using {source_label} Simple Index: {source_url}")
9971035

9981036
# Parse the index and get raw HTML
9991037
try:
@@ -1025,20 +1063,26 @@ def get_packages_for_target(target: str) -> List[str]:
10251063
Get packages from PACKAGES_PER_PROJECT that should be initialized for a target.
10261064
10271065
Returns packages where:
1028-
- project is "torch" AND
1066+
- project is "torch" (or "torch_rocm" for ROCm targets) AND
10291067
- either no target is specified (universal packages like filelock, numpy)
10301068
- or the target matches the specified target
10311069
- nvidia/cuda packages are only included for CUDA targets (cu*)
1070+
- amd/rocm packages are only included for ROCm targets (rocm*)
10321071
"""
10331072
is_cuda_target = target.startswith("cu")
1073+
is_rocm_target = target.startswith("rocm")
1074+
allowed_projects = ("torch", "torch_rocm") if is_rocm_target else ("torch",)
10341075
packages = []
10351076
for pkg_name, pkg_configs in PACKAGES_PER_PROJECT.items():
10361077
# Skip nvidia/cuda packages for non-CUDA targets
10371078
if not is_cuda_target and is_nvidia_package(pkg_name):
10381079
continue
1080+
# Skip amd/rocm packages for non-ROCm targets
1081+
if not is_rocm_target and is_amd_package(pkg_name):
1082+
continue
10391083

10401084
for config in pkg_configs:
1041-
if config.get("project") != "torch":
1085+
if config.get("project") not in allowed_projects:
10421086
continue
10431087
pkg_target = config.get("target", "")
10441088
# Include if no target specified (universal) or target matches

0 commit comments

Comments
 (0)