diff --git a/.ci/scripts/tests/test_cu134_dependencies.py b/.ci/scripts/tests/test_cu134_dependencies.py index 3134a3318f1..53dfaf411e7 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 ecc72553679..00a7b9bf9dc 100644 --- a/install_requirements.py +++ b/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