Remove manually synced arch versions in tools/nightly.py (#166616)

Discussed with @atalman offline. To reduce duplicate changes and reduce the number of files to change when updating arch versions.

------

Pull Request resolved: https://github.com/pytorch/pytorch/pull/166616
Approved by: https://github.com/ezyang

.github/scripts/generate_binary_build_matrix.py

@@ -11,11 +11,17 @@ architectures:
     * Latest XPU
 """
 
+import json
 import os
+import re
+from pathlib import Path
 from typing import Optional
 
 
-# NOTE: Please also update the CUDA sources in `PIP_SOURCES` in tools/nightly.py when changing this
+SCRIPT_DIR = Path(__file__).absolute().parent
+REPO_ROOT = SCRIPT_DIR.parent.parent
+
+
 CUDA_ARCHES = ["12.6", "12.8", "12.9", "13.0"]
 CUDA_STABLE = "12.8"
 CUDA_ARCHES_FULL_VERSION = {
@@ -31,7 +37,6 @@ CUDA_ARCHES_CUDNN_VERSION = {
     "13.0": "9",
 }
 
-# NOTE: Please also update the ROCm sources in `PIP_SOURCES` in tools/nightly.py when changing this
 ROCM_ARCHES = ["6.4", "7.0"]
 
 XPU_ARCHES = ["xpu"]
@@ -137,9 +142,48 @@ PYTORCH_EXTRA_INSTALL_REQUIREMENTS = {
 }
 
 
-def get_nccl_wheel_version(arch_version: str) -> str:
+# Used by tools/nightly.py
-    import re
+PYTORCH_NIGHTLY_PIP_INDEX_URL = "https://download.pytorch.org/whl/nightly"
+NIGHTLY_SOURCE_MATRIX = {
+    "cpu": dict(
+        name="cpu",
+        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/cpu",
+        supported_platforms=["Linux", "macOS", "Windows"],
+        accelerator="cpu",
+    )
+}
+CUDA_NIGHTLY_SOURCE_MATRIX = {
+    f"cuda-{major}.{minor}": dict(
+        name=f"cuda-{major}.{minor}",
+        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/cu{major}{minor}",
+        supported_platforms=["Linux", "Windows"],
+        accelerator="cuda",
+    )
+    for major, minor in (map(int, version.split(".")) for version in CUDA_ARCHES)
+}
+ROCM_NIGHTLY_SOURCE_MATRIX = {
+    f"rocm-{major}.{minor}": dict(
+        name=f"rocm-{major}.{minor}",
+        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/rocm{major}.{minor}",
+        supported_platforms=["Linux"],
+        accelerator="rocm",
+    )
+    for major, minor in (map(int, version.split(".")) for version in ROCM_ARCHES)
+}
+XPU_NIGHTLY_SOURCE_MATRIX = {
+    "xpu": dict(
+        name="xpu",
+        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/xpu",
+        supported_platforms=["Linux"],
+        accelerator="xpu",
+    )
+}
+NIGHTLY_SOURCE_MATRIX.update(CUDA_NIGHTLY_SOURCE_MATRIX)
+NIGHTLY_SOURCE_MATRIX.update(ROCM_NIGHTLY_SOURCE_MATRIX)
+NIGHTLY_SOURCE_MATRIX.update(XPU_NIGHTLY_SOURCE_MATRIX)
 
+
+def get_nccl_wheel_version(arch_version: str) -> str:
     requirements = map(
         str.strip, re.split("[;|]", PYTORCH_EXTRA_INSTALL_REQUIREMENTS[arch_version])
     )
@@ -147,17 +191,14 @@ def get_nccl_wheel_version(arch_version: str) -> str:
 
 
 def read_nccl_pin(arch_version: str) -> str:
-    from pathlib import Path
+    nccl_pin_path = (
-
+        REPO_ROOT
-    nccl_pin_path = os.path.join(
+        / ".ci"
-        Path(__file__).absolute().parents[2],
+        / "docker"
-        ".ci",
+        / "ci_commit_pins"
-        "docker",
+        / f"nccl-cu{arch_version[:2]}.txt"
-        "ci_commit_pins",
-        f"nccl-cu{arch_version[:2]}.txt",
     )
-    with open(nccl_pin_path) as f:
+    return nccl_pin_path.read_text().strip()
-        return f.read().strip()
 
 
 def validate_nccl_dep_consistency(arch_version: str) -> None:
@@ -165,7 +206,8 @@ def validate_nccl_dep_consistency(arch_version: str) -> None:
     wheel_ver = get_nccl_wheel_version(arch_version)
     if not nccl_release_tag.startswith(f"v{wheel_ver}"):
         raise RuntimeError(
-            f"{arch_version} NCCL release tag version {nccl_release_tag} does not correspond to wheel version {wheel_ver}"
+            f"{arch_version} NCCL release tag version {nccl_release_tag} "
+            f"does not correspond to wheel version {wheel_ver}"
         )
 
 
@@ -412,7 +454,14 @@ def generate_wheels_matrix(
     return ret
 
 
-validate_nccl_dep_consistency("13.0")
+arch_version = ""
-validate_nccl_dep_consistency("12.9")
+for arch_version in CUDA_ARCHES:
-validate_nccl_dep_consistency("12.8")
+    validate_nccl_dep_consistency(arch_version)
-validate_nccl_dep_consistency("12.6")
+del arch_version
+
+
+if __name__ == "__main__":
+    # Used by tools/nightly.py
+    (SCRIPT_DIR / "nightly_source_matrix.json").write_text(
+        json.dumps(NIGHTLY_SOURCE_MATRIX, indent=4) + "\n"
+    )

.gitignore

@@ -143,6 +143,7 @@ scripts/release_notes/*.json
 sccache-stats*.json
 lint.json
 merge_record.json
+.github/scripts/nightly_source_matrix.json
 
 # These files get copied over on invoking setup.py
 torchgen/packaged/*

tools/nightly.py

@@ -53,6 +53,7 @@ import atexit
 import contextlib
 import functools
 import itertools
+import json
 import logging
 import os
 import re
@@ -128,47 +129,19 @@ class PipSource(NamedTuple):
     accelerator: str
 
 
-PYTORCH_NIGHTLY_PIP_INDEX_URL = "https://download.pytorch.org/whl/nightly"
+# Generate: .github/scripts/nightly_source_matrix.json
-PIP_SOURCES = {
+GENERATE_MATRIX_SCRIPT = (
-    "cpu": PipSource(
+    REPO_ROOT / ".github" / "scripts" / "generate_binary_build_matrix.py"
-        name="cpu",
+)
-        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/cpu",
+subprocess.check_call(
-        supported_platforms={"Linux", "macOS", "Windows"},
+    [sys.executable, str(GENERATE_MATRIX_SCRIPT)],
-        accelerator="cpu",
+    cwd=GENERATE_MATRIX_SCRIPT.parent,
-    ),
+)
-    # NOTE: Sync with CUDA_ARCHES in .github/scripts/generate_binary_build_matrix.py
+
-    "cuda-12.6": PipSource(
+# See: .github/scripts/nightly_source_matrix.json
-        name="cuda-12.6",
+NIGHTLY_SOURCE_FILE = GENERATE_MATRIX_SCRIPT.with_name("nightly_source_matrix.json")
-        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/cu126",
+NIGHTLY_SOURCE_MATRIX = json.loads(NIGHTLY_SOURCE_FILE.read_text(encoding="utf-8"))
-        supported_platforms={"Linux", "Windows"},
+PIP_SOURCES = {name: PipSource(**data) for name, data in NIGHTLY_SOURCE_MATRIX.items()}
-        accelerator="cuda",
-    ),
-    "cuda-12.8": PipSource(
-        name="cuda-12.8",
-        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/cu128",
-        supported_platforms={"Linux", "Windows"},
-        accelerator="cuda",
-    ),
-    "cuda-13.0": PipSource(
-        name="cuda-13.0",
-        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/cu130",
-        supported_platforms={"Linux", "Windows"},
-        accelerator="cuda",
-    ),
-    # NOTE: Sync with ROCM_ARCHES in .github/scripts/generate_binary_build_matrix.py
-    "rocm-6.4": PipSource(
-        name="rocm-6.4",
-        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/rocm6.4",
-        supported_platforms={"Linux"},
-        accelerator="rocm",
-    ),
-    "rocm-7.0": PipSource(
-        name="rocm-7.0",
-        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/rocm7.0",
-        supported_platforms={"Linux"},
-        accelerator="rocm",
-    ),
-}
 
 
 class Formatter(logging.Formatter):