Automate release branch preparation

.ci/scripts/test_wheel_package_qnn.sh

@@ -98,8 +98,7 @@ install_qnn
 export LD_LIBRARY_PATH="${QNN_SDK_ROOT}/lib/x86_64-linux-clang/:${LD_LIBRARY_PATH:-}"
 
 install_executorch
-EXECUTORCH_BUILDING_WHEEL=1 python setup.py bdist_wheel
-unset EXECUTORCH_BUILDING_WHEEL
+EXECUTORCH_BUILDING_WHEEL=1 EXECUTORCH_RELEASE_WHEEL_METADATA=1 EXECUTORCH_WHEEL_VARIANT=cpu python setup.py bdist_wheel
 
 WHEEL_FILE=$(ls dist/*.whl | head -n 1)
 echo "Found wheel: $WHEEL_FILE"

.ci/scripts/tests/test_cu134_dependencies.py

@@ -5,12 +5,14 @@
 # LICENSE file in the root directory of this source tree.
 
 import ast
+import functools
 import importlib.util
 import os
 import subprocess
 import sys
 import unittest
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import patch
 
 from packaging.requirements import Requirement
@@ -28,6 +30,7 @@ def load_module(name):
 class TestCu134Dependencies(unittest.TestCase):
     def setUp(self):
         self.utils = load_module("install_utils")
+        self.release_versions = load_module("scripts/release/release_versions")
         self.modules = patch.dict(sys.modules, {"install_utils": self.utils})
         self.modules.start()
         self.addCleanup(self.modules.stop)
@@ -57,12 +60,10 @@ def test_all_install_steps_preserve_exact_cu134_selection(self):
             with self.subTest(machine=machine):
                 commands = self.install_commands((13, 4), machine)
                 self.assertEqual(len(commands), 4)
-                expected = {
-                    "torch==2.14.0.dev20260810+cu134",
-                    "torchvision==0.29.0.dev20260811+cu134",
-                    "torchaudio==2.11.0.dev20260811+cu134",
-                    f"torchao==0.19.0.dev20260907+{ao_variant}",
-                }
+                expected = set(self.installer.CU134_TORCH_PACKAGES)
+                expected.add(
+                    f"torchao=={self.installer.CU134_TORCHAO_NIGHTLY_VERSION}+{ao_variant}"
+                )
                 for index, command in enumerate(commands):
                     required = (
                         expected
@@ -81,17 +82,29 @@ def test_all_install_steps_preserve_exact_cu134_selection(self):
                                 for arg in command
                             )
                         )
-                    self.assertIn(
-                        "https://download.pytorch.org/whl/nightly/cu134", command
-                    )
-                    self.assertNotIn(
-                        "https://download.pytorch.org/whl/test/cu134", command
+                    expected_base = (
+                        self.installer.TORCH_URL_BASE
+                        if self.installer.RELEASE_WHEEL
+                        else self.installer.TORCHAO_URL_BASE
                     )
+                    self.assertIn(f"{expected_base}/cu134", command)
                     self.assertNotIn("--no-deps", command)
                     if machine == "aarch64":
-                        self.assertIn(
-                            "https://download.pytorch.org/whl/nightly/cpu", command
-                        )
+                        self.assertIn(f"{self.installer.TORCHAO_URL_BASE}/cpu", command)
+
+        release_packages = [
+            "torch==2.15.0+cu134",
+            "torchvision==0.30.0+cu134",
+            "torchaudio==2.12.0+cu134",
+        ]
+        with (
+            patch.object(self.installer, "RELEASE_WHEEL", True),
+            patch.object(self.installer, "CU134_TORCH_PACKAGES", release_packages),
+        ):
+            commands = self.install_commands((13, 4))
+        for command in commands:
+            self.assertIn("https://download.pytorch.org/whl/test/cu134", command)
+        self.assertTrue(set(release_packages).issubset(commands[2]))
 
     def test_other_cuda_trains_keep_existing_pins(self):
         for cuda in ((12, 6), (13, 0), (13, 2)):
@@ -100,10 +113,16 @@ def test_other_cuda_trains_keep_existing_pins(self):
                     core, local, domains, examples = self.install_commands(
                         cuda, machine
                     )
-                    self.assertIn("torch==2.14.0", core)
-                    self.assertIn("torchao==0.19.0.dev20260907", core)
-                    self.assertIn("torchvision==0.29.0", domains)
-                    self.assertIn("torchaudio==2.11.0", domains)
+                    self.assertIn(f"torch=={self.installer.TORCH_VERSION}", core)
+                    self.assertIn(
+                        f"torchao=={self.installer.TORCHAO_NIGHTLY_VERSION}", core
+                    )
+                    self.assertIn(
+                        f"torchvision=={self.installer.TORCHVISION_VERSION}", domains
+                    )
+                    self.assertIn(
+                        f"torchaudio=={self.installer.TORCHAUDIO_VERSION}", domains
+                    )
                     self.assertFalse(any("==" in arg for arg in local))
                     self.assertFalse(any("==" in arg for arg in examples))
 
@@ -112,21 +131,21 @@ def test_source_pinned_torch_is_not_replaced(self):
             with self.subTest(cuda=cuda):
                 core, _, domains, _ = self.install_commands(cuda, nightly=False)
                 self.assertIn("torch", core)
-                self.assertNotIn("torch==2.14.0.dev20260810+cu134", core)
+                self.assertNotIn(self.installer.CU134_TORCH_PACKAGES[0], core)
                 self.assertIn("torchvision", domains)
                 self.assertIn("torchaudio", domains)
 
     def test_no_cuda_keeps_default_pins(self):
         core, _, domains, _ = self.install_commands(None)
-        self.assertIn("torch==2.14.0", core)
-        self.assertIn("torchao==0.19.0.dev20260907", core)
-        self.assertIn("torchvision==0.29.0", domains)
+        self.assertIn(f"torch=={self.installer.TORCH_VERSION}", core)
+        self.assertIn(f"torchao=={self.installer.TORCHAO_NIGHTLY_VERSION}", core)
+        self.assertIn(f"torchvision=={self.installer.TORCHVISION_VERSION}", domains)
         self.assertIn("https://download.pytorch.org/whl/test/cpu", core)
 
     def test_windows_does_not_select_cu134(self):
         core, _, domains, _ = self.install_commands((13, 4), system="Windows")
-        self.assertIn("torch==2.14.0", core)
-        self.assertIn("torchvision==0.29.0", domains)
+        self.assertIn(f"torch=={self.installer.TORCH_VERSION}", core)
+        self.assertIn(f"torchvision=={self.installer.TORCHVISION_VERSION}", domains)
         self.assertIn("https://download.pytorch.org/whl/test/cpu", core)
 
     def test_failure_is_not_retried_with_another_cuda_train(self):
@@ -143,26 +162,63 @@ def test_failure_is_not_retried_with_another_cuda_train(self):
                 self.installer.install_requirements(True)
         self.assertEqual(run.call_count, 1)
 
-    def torchao_requirement(self):
+    def setup_requirement(
+        self,
+        function_name,
+        *,
+        installed_torch="2.15.0",
+        building_wheel=True,
+        wheel_variant="cpu",
+    ):
         path = ROOT / "setup.py"
         tree = ast.parse(path.read_text())
-        function = next(
+        functions = [
             node
             for node in tree.body
-            if isinstance(node, ast.FunctionDef) and node.name == "_torchao_requirement"
-        )
+            if isinstance(node, ast.FunctionDef)
+            and node.name
+            in {
+                "_load_install_requirements",
+                "_torchao_requirement",
+                "_release_torch_requirement",
+            }
+        ]
         namespace = {
             "__file__": str(path),
             "Path": Path,
+            "List": list,
+            "functools": functools,
             "importlib": importlib,
+            "os": os,
             "sys": sys,
             "install_utils": self.utils,
+            "release_versions": self.release_versions,
+            "torch_pin": SimpleNamespace(RELEASE_WHEEL=True, TORCH_VERSION="2.15.0"),
         }
         exec(
-            compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"),
+            compile(ast.Module(body=functions, type_ignores=[]), str(path), "exec"),
             namespace,
         )
-        return namespace["_torchao_requirement"]()
+        environment = (
+            {
+                "EXECUTORCH_RELEASE_WHEEL_METADATA": "1",
+                "EXECUTORCH_WHEEL_VARIANT": wheel_variant,
+            }
+            if building_wheel
+            else {}
+        )
+        with (
+            patch.dict(os.environ, environment, clear=True),
+            patch("importlib.metadata.version", return_value=installed_torch),
+        ):
+            return namespace[function_name]()
+
+    def torchao_requirement(self):
+        return self.setup_requirement("_torchao_requirement")
+
+    def release_torch_requirement(self, **kwargs):
+        requirements = self.setup_requirement("_release_torch_requirement", **kwargs)
+        return requirements[0] if requirements else None
 
     def test_package_install_preserves_source_pinned_torchao(self):
         with patch.dict(sys.modules, {"install_requirements": self.installer}):
@@ -241,23 +297,30 @@ def test_cu134_keeps_explicit_torchao_source_build(self):
                 self.assertIn("torch==2.14.0.dev20260810+cu134", commands[-1])
                 self.assertIn("0.19.0+gitb7ac3aa", metadata.specifier)
 
-    def test_wheel_torchao_bound_matches_selected_train(self):
-        for cuda, expected in (
-            ((13, 4), "torchao>=0.19.0.dev20260907,<0.20"),
-            ((13, 2), "torchao>=0.19.0.dev20260907,<0.20"),
-            (None, "torchao>=0.19.0.dev20260907,<0.20"),
+    def test_wheel_bounds_match_selected_train(self):
+        for wheel_variant, installed_torch, expected_torch in (
+            ("cu132", "2.15.0+cu132", "torch==2.15.0+cu132"),
+            ("cpu", "2.15.0", "torch>=2.15.0,<2.16"),
         ):
-            self.utils.determine_torch_url.cache_clear()
-            with (
-                patch.object(
-                    self.utils,
-                    "_get_cuda_version",
-                    return_value=cuda,
-                    side_effect=RuntimeError("no nvcc") if cuda is None else None,
-                ),
-                patch.object(self.installer.platform, "system", return_value="Linux"),
-            ):
-                self.assertEqual(self.torchao_requirement(), expected)
+            with self.subTest(wheel_variant=wheel_variant):
+                self.assertEqual(
+                    self.torchao_requirement(),
+                    "torchao>=0.19.0.dev20260907,<0.20",
+                )
+                self.assertEqual(
+                    self.release_torch_requirement(
+                        installed_torch=installed_torch,
+                        wheel_variant=wheel_variant,
+                    ),
+                    expected_torch,
+                )
+
+        with self.assertRaisesRegex(RuntimeError, "for Torch 2.15.0"):
+            self.release_torch_requirement(
+                installed_torch="2.14.0.dev20260810+cu134",
+                wheel_variant="cu134",
+            )
+        self.assertIsNone(self.release_torch_requirement(building_wheel=False))
 
 
 if __name__ == "__main__":

.ci/scripts/tests/test_filter_cuda_matrix.py

@@ -16,6 +16,7 @@
 import json
 import os
 import re
+import runpy
 import subprocess
 import unittest
 from pathlib import Path
@@ -37,6 +38,7 @@ def _load_module(name, path):
     "filter_cuda_matrix", ROOT / ".github" / "scripts" / "filter_cuda_matrix.py"
 )
 INSTALL_UTILS = _load_module("install_utils", ROOT / "install_utils.py")
+RELEASE_WHEEL = runpy.run_path(str(ROOT / "torch_pin.py"))["RELEASE_WHEEL"]
 
 
 def _full_matrix():
@@ -284,7 +286,12 @@ class TestPublishedSets(unittest.TestCase):
     """
 
     def test_published_cuda_versions(self):
-        self.assertEqual(FILTER.SUPPORTED_CUDA_VERSIONS, ["cu130", "cu132", "cu134"])
+        expected = {"cu130", "cu132", "cu134"}
+        if RELEASE_WHEEL:
+            self.assertTrue(FILTER.SUPPORTED_CUDA_VERSIONS)
+            self.assertLessEqual(set(FILTER.SUPPORTED_CUDA_VERSIONS), expected)
+        else:
+            self.assertEqual(set(FILTER.SUPPORTED_CUDA_VERSIONS), expected)
 
     def test_published_cuda_versions_are_documented(self):
         # The install table on the getting started page is the only place a user is told
@@ -297,11 +304,15 @@ def test_published_cuda_versions_are_documented(self):
         self.assertIn(heading, text, f"no Installation section in {page.name}")
         section = text.split(heading, 1)[1].split("\n## ", 1)[0]
         documented = set(re.findall(r"cu\d+", section))
-        self.assertEqual(
-            documented,
-            set(FILTER.SUPPORTED_CUDA_VERSIONS),
-            f"the install table in {page.name} lists {sorted(documented)}",
-        )
+        published = set(FILTER.SUPPORTED_CUDA_VERSIONS)
+        if RELEASE_WHEEL:
+            self.assertLessEqual(published, documented)
+        else:
+            self.assertEqual(
+                documented,
+                published,
+                f"the install table in {page.name} lists {sorted(documented)}",
+            )
 
     def test_published_cuda_versions_are_supported_by_the_installer(self):
         supported = {

.ci/scripts/wheel/cuda_arch_list.sh

@@ -73,6 +73,17 @@ _executorch_unknown_train() {
   return 64
 }
 
+# Normalize and validate the CUDA wheel variant once for every wheel script.
+executorch_cuda_variant() {
+  local raw="${CU_VERSION:-${DESIRED_CUDA:-}}"
+  local digits="${raw//[!0-9]/}"
+  if [[ ! "${digits}" =~ ^[0-9]{3}$ ]]; then
+    echo "CU_VERSION or DESIRED_CUDA must identify a CUDA wheel variant, got '${raw}'" >&2
+    return 64
+  fi
+  printf 'cu%s' "${digits}"
+}
+
 # The architectures for the current row, space separated in the dotted form PyTorch expects.
 executorch_cuda_arch_list() {
   local machine
@@ -97,8 +108,10 @@ executorch_cuda_arch_list() {
       ;;
   esac
   # The value arrives as cu130, while some callers pass 13.0 instead.
+  if ! train="$(executorch_cuda_variant)"; then
+    return 64
+  fi
   train="${train#cu}"
-  train="${train//./}"
 
   case "${machine}" in
     aarch64 | arm64)

.ci/scripts/wheel/envvar_base.sh

@@ -11,3 +11,9 @@
 # Ensure that CMAKE_ARGS is defined before referencing it. Defaults to empty
 # if not defined.
 export CMAKE_ARGS="${CMAKE_ARGS:-}"
+
+# setup.py only adds release dependency metadata for artifacts built by the
+# binary wheel jobs. A source install on a release branch must preserve the
+# PyTorch build already installed by that checkout's CI job.
+export EXECUTORCH_RELEASE_WHEEL_METADATA=1
+export EXECUTORCH_WHEEL_VARIANT=cpu

.ci/scripts/wheel/envvar_cuda_linux.sh

@@ -9,6 +9,12 @@
 # any variables so that subprocesses will see them.
 
 source "${GITHUB_WORKSPACE}/${REPOSITORY}/.ci/scripts/wheel/envvar_base.sh"
+source "${GITHUB_WORKSPACE}/${REPOSITORY}/.ci/scripts/wheel/cuda_arch_list.sh"
+
+if ! _executorch_cuda_variant="$(executorch_cuda_variant)"; then
+  exit 1
+fi
+export EXECUTORCH_WHEEL_VARIANT="${_executorch_cuda_variant}"
 
 # Ask for the CUDA delegate explicitly rather than letting the build detect a toolkit. A detected
 # build is fine locally, but a release row states what it is producing, and a row that silently
@@ -31,7 +37,6 @@ fi
 # Compile device code for the GPUs this release row claims, rather than for whichever GPU the
 # builder happens to have. A wheel built by detection alone installs on every machine the row covers
 # and then fails when a model runs on a different generation.
-source "${GITHUB_WORKSPACE}/${REPOSITORY}/.ci/scripts/wheel/cuda_arch_list.sh"
 # The status is checked rather than only the output, so an unrecognised row reports why it stopped.
 # A bare assignment would end the build on the lookup's own exit status with no message, since this
 # file is sourced into a shell that exits on a failing command.