ci: consolidate pinned PyTorch installation

.ci/docker/README.md

@@ -32,3 +32,42 @@ torch identity and schema-1.0 reporter after installing the domain libraries.
 These pins participate in the existing global Linux image hash. A pin change
 can require other Linux images to rebuild, even though only clang12 reads it.
 Validate the rebuilt image before enabling test-report uploads.
+
+## Reusing the pinned PyTorch installer
+
+From an ExecuTorch checkout, activate a venv or conda environment and run:
+
+```bash
+bash .ci/scripts/install_pinned_pytorch.sh
+bash .ci/scripts/install_pinned_pytorch.sh --verify-only
+```
+
+The installer reads `.ci/docker/ci_commit_pins/pytorch.txt`, resolves named refs,
+and reuses an installed build with the same commit. Set `PYTHON_EXECUTABLE` to
+select another interpreter, or `TORCH_VERSION` to override the source ref as the
+Docker build does. A mismatched build is replaced, including when downgrading.
+Install mode also installs the domain packages when the torch build is reused.
+`--verify-only` checks the installed PyTorch commit and LAPACK support, returning
+nonzero on a mismatch or failed check. It does not install packages or change the
+active environment. Named refs are resolved with `git ls-remote`, without cloning.
+This is not a check of every ExecuTorch build or test dependency.
+
+Use it inside a CI container to check that its prebuilt PyTorch matches the pin
+in the current checkout, or locally to check the selected venv/conda environment
+before running tests:
+
+```bash
+PYTHON_EXECUTABLE="$PWD/env/bin/python" \
+  bash .ci/scripts/install_pinned_pytorch.sh --verify-only
+```
+
+It can also run after dependency installation to detect replacement of the
+pinned build. Linux and macOS CI verify after installing ExecuTorch, before
+building the test runner. On Linux the check preserves the commit provided by
+the image or PyTorch CI, which may intentionally differ from the shared checkout
+pin.
+
+The shared implementation lives in `common/install_pinned_pytorch.sh` so it is
+also available in the Docker build context. Docker and macOS CI use this same
+implementation. System build dependencies (including BLAS/LAPACK) must already
+be installed. Windows is not supported.

.ci/docker/common/install_pinned_pytorch.sh

@@ -0,0 +1,194 @@
+#!/bin/bash
+# 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.
+
+# Called directly or through .ci/scripts/install_pinned_pytorch.sh.
+# Activate a venv/conda environment first, or set PYTHON_EXECUTABLE.
+set -euo pipefail
+
+dedupe_macos_loader_path_rpaths() {
+  if [[ "$(uname)" != "Darwin" ]]; then
+    return
+  fi
+
+  local torch_lib_dir
+  torch_lib_dir=$(cd / && "${python_executable}" -c "import importlib.util; print(importlib.util.find_spec('torch').submodule_search_locations[0])")/lib
+
+  if [[ -z "${torch_lib_dir}" || ! -d "${torch_lib_dir}" ]]; then
+    return
+  fi
+
+  local torch_libs=(
+    "libtorch_cpu.dylib"
+    "libtorch.dylib"
+    "libc10.dylib"
+  )
+
+  for lib_name in "${torch_libs[@]}"; do
+    local lib_path="${torch_lib_dir}/${lib_name}"
+    if [[ ! -f "${lib_path}" ]]; then
+      continue
+    fi
+
+    local removed=0
+    # Repeatedly remove the @loader_path rpath entries until none remain.
+    while install_name_tool -delete_rpath @loader_path "${lib_path}" 2>/dev/null; do
+      removed=1
+    done
+
+    if [[ "${removed}" == "1" ]]; then
+      install_name_tool -add_rpath @loader_path "${lib_path}" || true
+    fi
+  done
+}
+
+verify_only=false
+if [[ "${1:-}" == "--verify-only" ]]; then
+  verify_only=true
+  shift
+fi
+if [[ $# -gt 1 || "${1:-}" == -* ]]; then
+  echo "Usage: $0 [--verify-only] [EXECUTORCH_DIRECTORY]" >&2
+  exit 2
+fi
+executorch_dir=$(cd "${1:-.}" && pwd)
+python_executable=${PYTHON_EXECUTABLE:-python}
+# Resolve the interpreter before entering the temporary source directory.
+python_executable=$("${python_executable}" -c 'import sys; print(sys.executable)')
+system_name=$(uname)
+case "${system_name}" in
+  Linux | Darwin) ;;
+  *) echo "Unsupported system: ${system_name}" >&2; exit 1 ;;
+esac
+
+torch_ref=${TORCH_VERSION:-$(cat "${executorch_dir}/.ci/docker/ci_commit_pins/pytorch.txt")}
+torch_ref=$(printf '%s' "${torch_ref}" | tr -d '[:space:]')
+git check-ref-format --branch "${torch_ref}" >/dev/null
+installed_pin=$(cd / && "${python_executable}" -c \
+  'import torch; print(torch.version.git_version or "")' 2>/dev/null || true)
+expected_pin=${torch_ref}
+
+verify_pin() {
+  if [[ "${installed_pin}" != "${expected_pin}" ]]; then
+    echo "PyTorch pin mismatch: expected ${expected_pin}, found ${installed_pin:-none}" >&2
+    return 1
+  fi
+  (cd / && "${python_executable}" -c '
+import torch
+assert torch._C.has_lapack, "PyTorch was built without LAPACK"
+torch.linalg.qr(torch.randn(4, 4))
+')
+}
+
+install_domains() {
+  "${python_executable}" -m pip install numpy pillow
+  "${python_executable}" -m pip install --no-build-isolation --no-deps \
+    "git+https://github.com/pytorch/audio.git@${TORCHAUDIO_VERSION:-release/2.11}" \
+    "git+https://github.com/pytorch/vision.git@${TORCHVISION_VERSION:-release/0.29}"
+}
+
+if [[ ! "${torch_ref}" =~ ^[0-9a-f]{40}$ ]]; then
+  expected_pin=$(git ls-remote --exit-code https://github.com/pytorch/pytorch.git \
+    "${torch_ref}" "${torch_ref}^{}" | awk '
+      { sub(/\^\{\}$/, "", $2); pins[$2] = $1 }
+      END { for (ref in pins) print pins[ref] }
+    ' | sort -u) || expected_pin=""
+  if [[ ! "${expected_pin}" =~ ^[0-9a-f]{40}$ ]]; then
+    echo "Cannot resolve PyTorch ref to one commit: ${torch_ref}" >&2
+    exit 1
+  fi
+fi
+
+if [[ "${installed_pin}" == "${expected_pin}" || "${verify_only}" == true ]]; then
+  verify_pin
+  echo "PyTorch pin ${expected_pin} is already installed"
+  if [[ "${verify_only}" == false ]]; then
+    dedupe_macos_loader_path_rpaths
+    install_domains
+  fi
+  exit 0
+fi
+
+build_root=$(mktemp -d "${TMPDIR:-/tmp}/executorch-pytorch.XXXXXXXX")
+trap '
+  status=$?
+  if [[ ${status} -ne 0 ]]; then
+    echo "Failed to prepare the replacement PyTorch wheel." >&2
+    echo "Existing torch and domain packages have not been removed." >&2
+  fi
+  rm -rf "${build_root}"
+  exit "${status}"
+' EXIT
+git clone --filter=blob:none --no-checkout https://github.com/pytorch/pytorch.git "${build_root}/pytorch"
+cd "${build_root}/pytorch"
+git fetch --depth=1 origin "${expected_pin}"
+git checkout --detach FETCH_HEAD
+
+cached=false
+if [[ "${system_name}" == Darwin ]]; then
+  platform=$("${python_executable}" -c 'import sysconfig; import platform; v=platform.mac_ver()[0].split(".")[0]; platform=sysconfig.get_platform().split("-"); platform[1]=f"{v}_0"; print("_".join(platform))')
+  python_version=$("${python_executable}" -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
+  torch_release=$(cat version.txt)
+  torch_short_hash=$(git rev-parse --short=7 HEAD)
+  torch_wheel_path="cached_artifacts/pytorch/executorch/pytorch_wheels/${system_name}/${python_version}"
+  torch_wheel_name="torch-${torch_release}+git${torch_short_hash}-cp${python_version}-cp${python_version}-${platform}.whl"
+  cached_torch_wheel="https://gha-artifacts.s3.us-east-1.amazonaws.com/${torch_wheel_path}/${torch_wheel_name//+/%2B}"
+  if "${python_executable}" -m pip download --no-deps --dest dist "${cached_torch_wheel}"; then
+    cached=true
+  fi
+fi
+
+if [[ "${cached}" == false ]]; then
+  echo "Building PyTorch at ${expected_pin}"
+  git submodule update --init --recursive
+  if [[ "$(uname -m)" == aarch64 ]]; then
+    export BUILD_IGNORE_SVE_UNAVAILABLE=1
+  fi
+  if [[ "${system_name}" == Linux ]]; then
+    export _GLIBCXX_USE_CXX11_ABI=1
+  fi
+  if [[ -n "${PYTORCH_BUILD_MAX_JOBS:-}" ]]; then
+    export MAX_JOBS="${PYTORCH_BUILD_MAX_JOBS}"
+  fi
+  build_venv="${build_root}/build-env"
+  "${python_executable}" -m venv --system-site-packages "${build_venv}"
+  # Keep in sync with pytorch/pyproject.toml [build-system].requires.
+  # Do not install pip cmake: it hides the image's MKL/libomp search paths.
+  "${build_venv}/bin/python" -m pip install build "scikit-build-core>=1.0" \
+    ninja "packaging>=24.2" "typing-extensions>=4.10.0" pyyaml six numpy
+  CMAKE_CXX_SCAN_FOR_MODULES=OFF USE_DISTRIBUTED=1 \
+    "${build_venv}/bin/python" -m build --wheel --no-isolation
+fi
+
+shopt -s nullglob
+wheels=(dist/*.whl)
+shopt -u nullglob
+if [[ ${#wheels[@]} -ne 1 ]]; then
+  echo "Expected one PyTorch wheel, found ${#wheels[@]}" >&2
+  exit 1
+fi
+if [[ "${system_name}" == Darwin ]]; then
+  if [[ "$(basename "${wheels[0]}")" != "${torch_wheel_name}" ]]; then
+    echo "Built PyTorch wheel does not match cache key: ${torch_wheel_name}" >&2
+    exit 1
+  fi
+fi
+
+trap 'rm -rf "${build_root}"' EXIT
+# Remove the old builds explicitly so equal package versions also replace
+# different source commits, and domain dependencies cannot replace the pin.
+"${python_executable}" -m pip uninstall -y torch torchaudio torchvision
+"${python_executable}" -m pip install "${wheels[0]}"
+
+dedupe_macos_loader_path_rpaths
+install_domains
+installed_pin=$(cd / && "${python_executable}" -c 'import torch; print(torch.version.git_version or "")')
+verify_pin
+
+# Preserve the existing macOS wheel cache on AWS CI runners.
+if [[ "${system_name}" == Darwin && "${cached}" == false && "${GITHUB_ACTIONS:-}" == true && -z "${GITHUB_RUNNER:-}" ]] && command -v aws >/dev/null; then
+  aws s3 cp "${wheels[0]}" "s3://gha-artifacts/${torch_wheel_path}/${torch_wheel_name}"
+fi
+sccache --show-stats || true

.ci/docker/common/install_pytorch.sh

@@ -1,6 +1,7 @@
 #!/bin/bash
 # 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.
@@ -23,12 +24,6 @@ else
   export SCCACHE_DIR=/tmp/sccache
 fi
 
-install_domains() {
-  echo "Install torchvision and torchaudio"
-  pip_install --no-build-isolation --user "git+https://github.com/pytorch/audio.git@${TORCHAUDIO_VERSION}"
-  pip_install --no-build-isolation --user "git+https://github.com/pytorch/vision.git@${TORCHVISION_VERSION}"
-}
-
 configure_pytorch_compiler() {
   local cxx_version
   cxx_version=$(/usr/bin/c++ --version | head -n1)
@@ -44,69 +39,13 @@ configure_pytorch_compiler() {
   fi
 }
 
-install_pytorch_and_domains() {
-  git clone https://github.com/pytorch/pytorch.git
-
-  # Fetch the target commit
-  pushd pytorch || true
-  git checkout "${TORCH_VERSION}"
-  git submodule update --init --recursive
-
-  chown -R ci-user .
-
-  export _GLIBCXX_USE_CXX11_ABI=1
-  if [[ "$(uname -m)" == "aarch64" ]]; then
-    export BUILD_IGNORE_SVE_UNAVAILABLE=1
-  fi
-  if [[ -n "${PYTORCH_BUILD_MAX_JOBS:-}" ]]; then
-    export MAX_JOBS="${PYTORCH_BUILD_MAX_JOBS}"
-  fi
-  configure_pytorch_compiler
-  # PyTorch no longer supports "python setup.py bdist_wheel"; it now builds
-  # through scikit-build-core (PEP 517). Build the wheel with the standard
-  # frontend and keep build isolation off, so PyTorch builds against this
-  # environment's numpy and toolchain (avoids an ABI mismatch) and reuses
-  # sccache.
-  #
-  # With isolation off the frontend does not fetch the PEP 517 build
-  # requirements, so install them here. They go into a throwaway venv rather
-  # than into the image, because scikit-build-core registers a setuptools
-  # build_ext plugin: left in the image it hijacks every later
-  # "pip install --no-build-isolation" and turns on C++20 module scanning that
-  # the image compiler cannot satisfy. The venv inherits the image's
-  # site-packages, so PyTorch still builds against the same numpy.
-  #
-  # Keep in sync with pytorch/pyproject.toml [build-system].requires.
-  local build_venv=/tmp/pytorch-build-venv
-  rm -rf "${build_venv}"
-  conda_run python -m venv --system-site-packages "${build_venv}"
-  # No pip cmake: scikit-build-core would prefer it over the image's, and it
-  # searches site-packages, where MKL and libomp are not.
-  conda_run "${build_venv}/bin/pip" install build "scikit-build-core>=1.0" \
-    ninja "packaging>=24.2" "typing-extensions>=4.10.0" pyyaml six numpy
-  # These images have no module scanner, and nothing here uses modules.
-  conda_run env CMAKE_CXX_SCAN_FOR_MODULES=OFF \
-    "${build_venv}/bin/python" -m build --wheel --no-isolation
-  rm -rf "${build_venv}"
-  pip_install "$(echo dist/*.whl)"
-
-  # A build with no BLAS succeeds silently. Run from / to import the wheel.
-  (cd / && conda_run python -c "
-import torch
-assert torch._C.has_lapack, 'built without LAPACK'
-torch.linalg.qr(torch.randn(4, 4))
-")
+configure_pytorch_compiler
+export TORCH_VERSION
+export PYTHON_EXECUTABLE="/opt/conda/envs/py_${PYTHON_VERSION}/bin/python"
+conda_run bash "$(dirname "${BASH_SOURCE[0]}")/install_pinned_pytorch.sh"
 
-  # Grab the pinned audio and vision commits from PyTorch
-  TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION:-release/2.11}
-  export TORCHAUDIO_VERSION
-  TORCHVISION_VERSION=${TORCHVISION_VERSION:-release/0.29}
-  export TORCHVISION_VERSION
-
-  install_domains
-
-  if [[ "${TORCHCI_REPORTS:-}" == "1" ]]; then
-    (cd / && conda_run python -c '
+if [[ "${TORCHCI_REPORTS:-}" == "1" ]]; then
+  (cd / && conda_run python -c '
 import os
 import torch
 from torch.testing._internal.torchci import environment, plugin, report
@@ -117,14 +56,4 @@ if report.SCHEMA_VERSION != "1.0":
     raise RuntimeError("The image requires torchci report schema 1.0")
 environment.capture()
 ')
-  fi
-
-  popd || true
-  # Clean up the cloned PyTorch repo to reduce the Docker image size
-  rm -rf pytorch
-
-  # Print sccache stats for debugging
-  as_ci_user sccache --show-stats
-}
-
-install_pytorch_and_domains
+fi

.ci/docker/ubuntu/Dockerfile

@@ -72,13 +72,14 @@ ARG TORCHCI_REPORTS
 ARG SKIP_PYTORCH
 ARG PYTORCH_BUILD_MAX_JOBS
 COPY ./common/install_pytorch.sh install_pytorch.sh
+COPY ./common/install_pinned_pytorch.sh install_pinned_pytorch.sh
 COPY ./common/utils.sh utils.sh
 # A host docker build let sccache reach its S3 bucket with the EC2 instance
 # role; a BuildKit pod has no instance metadata, so docker-builds passes the
 # credentials in as this secret. World-readable because the sccache server runs
 # as ci-user; a secret mount never lands in a layer.
 RUN --mount=type=secret,id=aws-credentials,mode=0444 \
-    if [ -z "${SKIP_PYTORCH}" ]; then bash ./install_pytorch.sh; fi && rm install_pytorch.sh utils.sh
+    if [ -z "${SKIP_PYTORCH}" ]; then bash ./install_pytorch.sh; fi && rm install_pytorch.sh install_pinned_pytorch.sh utils.sh
 
 ARG LINTRUNNER
 # Install lintrunner if needed

.ci/scripts/install_pinned_pytorch.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+# 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.
+
+set -euo pipefail
+
+script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+exec bash "${script_dir}/../docker/common/install_pinned_pytorch.sh" "$@"