From e4f3df318d79b216720192020544a87fc87bd995 Mon Sep 17 00:00:00 2001 From: Zingo Andersen Date: Fri, 6 Nov 2026 08:55:52 +0100 Subject: [PATCH] Fix TorchAO source installs and automatic pin updates Skip the prebuilt TorchAO requirement when building from source. An unavailable wheel should not prevent installing the local source. Allow source revision lookup from partially failed wheel workflows. The nightly selection already checks that the required wheels exist across all supported channels, so unrelated build failures should not block pin updates. Add regression tests for source installs, nightly selection and revision lookup from partially failed wheel workflows. Authored with assistance from OpenAI Codex. Change-Id: Ia106b28ffeaf564b78d664c18e87f35e7fbea00e Signed-off-by: Zingo Andersen --- .ci/scripts/tests/test_cu134_dependencies.py | 28 ++++++++-- .ci/scripts/tests/test_update_pytorch_pin.py | 59 ++++++++++++++++++++ .github/scripts/update_pytorch_pin.py | 10 +++- install_requirements.py | 3 +- 4 files changed, 93 insertions(+), 7 deletions(-) create mode 100644 .ci/scripts/tests/test_update_pytorch_pin.py diff --git a/.ci/scripts/tests/test_cu134_dependencies.py b/.ci/scripts/tests/test_cu134_dependencies.py index 2812a34d5df..8655ba6aadf 100644 --- a/.ci/scripts/tests/test_cu134_dependencies.py +++ b/.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) diff --git a/.ci/scripts/tests/test_update_pytorch_pin.py b/.ci/scripts/tests/test_update_pytorch_pin.py new file mode 100644 index 00000000000..f422b3d3172 --- /dev/null +++ b/.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() diff --git a/.github/scripts/update_pytorch_pin.py b/.github/scripts/update_pytorch_pin.py index c6a6318e965..35efe66fc87 100644 --- a/.github/scripts/update_pytorch_pin.py +++ b/.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"] diff --git a/install_requirements.py b/install_requirements.py index 4d519747a5d..a01b845db09 100644 --- a/install_requirements.py +++ b/install_requirements.py @@ -122,8 +122,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