Fix TorchAO source installs and automatic pin updates

.ci/scripts/tests/test_cu134_dependencies.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -113,6 +114,13 @@ def test_other_cuda_trains_keep_existing_pins(self):
                     self.assertIn(
                         f"torchao=={self.installer.TORCHAO_NIGHTLY_VERSION}", core
                     )
+                    torchao_channel = (
+                        "cpu" if machine == "aarch64" else f"cu{cuda[0]}{cuda[1]}"
+                    )
+                    self.assertIn(
+                        f"https://download.pytorch.org/whl/nightly/{torchao_channel}",
+                        core,
+                    )
                     self.assertIn("torchvision==0.29.0", domains)
                     self.assertIn("torchaudio==2.11.0", domains)
                     self.assertFalse(any("==" in arg for arg in local))
@@ -171,7 +179,7 @@ def test_cpu_release_option_keeps_cuda_pins(self):
                         ),
                     )
 
-    def test_cpu_source_build_keeps_nightly_torchao(self):
+    def test_cpu_source_build_does_not_require_torchao_wheels(self):
         for flag in (
             "EXECUTORCH_BUILD_KERNELS_TORCHAO",
             "TORCHAO_BUILD_EXPERIMENTAL_MPS",
@@ -181,14 +189,26 @@ def test_cpu_source_build_keeps_nightly_torchao(self):
                     None,
                     env={"EXECUTORCH_USE_RELEASE_CPU_DEPS": "1", flag: "1"},
                 )
-                self.assertIn(
-                    f"torchao=={self.installer.TORCHAO_NIGHTLY_VERSION}", core
-                )
+                self.assertFalse(any(arg.startswith("torchao==") for arg in core))
                 self.assertIn("https://download.pytorch.org/whl/nightly/cpu", core)
                 for command in (core, domains):
                     self.assertIn("https://download.pytorch.org/whl/test/cpu", command)
                 self.assertIn("third-party/ao", local)
 
+    def test_cuda_source_build_does_not_require_torchao_wheels(self):
+        for cuda in ((13, 0), (13, 2), (13, 4)):
+            for flag in (
+                "EXECUTORCH_BUILD_KERNELS_TORCHAO",
+                "TORCHAO_BUILD_EXPERIMENTAL_MPS",
+            ):
+                with self.subTest(cuda=cuda, flag=flag):
+                    commands = self.install_commands(cuda, env={flag: "1"})
+                    self.assertIn("third-party/ao", commands[1])
+                    for command in commands:
+                        self.assertFalse(
+                            any(arg.startswith("torchao==") for arg in command)
+                        )
+
     def test_windows_does_not_select_cu134(self):
         core, _, domains, _ = self.install_commands((13, 4), system="Windows")
         self.assertIn("torch==2.14.0", core)

.ci/scripts/tests/test_update_pytorch_pin.py

@@ -0,0 +1,59 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import importlib.util
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+ROOT = Path(__file__).resolve().parents[3]
+
+
+def load_updater():
+    path = ROOT / ".github/scripts/update_pytorch_pin.py"
+    spec = importlib.util.spec_from_file_location("update_pytorch_pin", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class TestUpdatePytorchPin(unittest.TestCase):
+    def test_torchao_pin_covers_supported_nightly_channels(self):
+        updater = load_updater()
+        older = "0.19.0.dev20261031"
+        newer = "0.19.0.dev20261101"
+        versions = {
+            "cpu": {older, newer},
+            "cu130": {older},
+            "cu132": {older, newer},
+            "cu134": {older, newer},
+        }
+
+        with patch.object(
+            updater, "get_torchao_versions", side_effect=versions.__getitem__
+        ):
+            self.assertEqual(
+                updater.get_latest_torchao_nightly("20261101"),
+                older,
+            )
+
+    def test_torchao_commit_accepts_a_partially_failed_wheel_run(self):
+        updater = load_updater()
+        run = {"head_sha": "b87db29e", "conclusion": "failure"}
+
+        with patch.object(
+            updater, "get_json", return_value={"workflow_runs": [run]}
+        ) as get_json:
+            self.assertEqual(
+                updater.get_torchao_commit_hash("0.19.0.dev20261101"),
+                "b87db29e",
+            )
+
+        self.assertNotIn("status=success", get_json.call_args.args[0])
+
+
+if __name__ == "__main__":
+    unittest.main()

.github/scripts/update_pytorch_pin.py

@@ -1,4 +1,8 @@
 #!/usr/bin/env python3
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
 
 import base64
 import json
@@ -165,15 +169,17 @@ def get_latest_torchao_nightly(max_date):
 def get_torchao_commit_hash(nightly_version):
     date = nightly_version.rsplit(".dev", 1)[-1]
     formatted_date = parse_nightly_version(f"dev{date}")
+    # Nightly selection already checks the required wheels. Other jobs in
+    # the same workflow may fail without invalidating those wheels.
     url = (
         "https://api.github.com/repos/pytorch/ao/actions/workflows/"  # @lint-ignore
-        "build_wheels_linux_x86.yml/runs?event=schedule&status=success&"
+        "build_wheels_linux_x86.yml/runs?event=schedule&"
         f"created={formatted_date}&per_page=100"
     )
     runs = get_json(url).get("workflow_runs", [])
     if not runs:
         raise ValueError(
-            f"Could not find the successful TorchAO wheel build for {nightly_version}"
+            f"Could not find the TorchAO wheel build for {nightly_version}"
         )
     return runs[0]["head_sha"]
 

install_requirements.py

@@ -127,8 +127,9 @@ def install_requirements(use_pytorch_nightly):
         # that we don't need to set any version number there because they have already
         # been installed on CI before this step, so pip won't reinstall them
         ("torch==2.14.0" if use_pytorch_nightly else "torch"),
-        f"torchao=={torchao_version}",
     ]
+    if not cu134_packages and not torchao_from_source():
+        TORCH_PACKAGE.append(f"torchao=={torchao_version}")
 
     # Install the requirements for core ExecuTorch package.
     # `--extra-index-url` tells pip to look for package