Use prebuilt PyTorch wheels in CI

.ci/docker/build.sh

@@ -94,12 +94,35 @@ case "${IMAGE_NAME}" in
     exit 1
 esac
 
-TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
 BUILD_DOCS=1
 
-if [[ "${GCC_VERSION:-}" == "11" && -z "${SKIP_PYTORCH:-}" ]]; then
-  PYTORCH_BUILD_MAX_JOBS=6
-fi
+# Pull channel + spec/url helpers out of torch_pin.py so install_pytorch.sh
+# can install the selected PyTorch wheel channel.
+# Self-hosted runners often have python3 but not the unversioned python alias.
+PYTHON_BIN=$(command -v python3 || command -v python)
+TORCH_PIN_HELPERS=$(
+  cd ../.. &&
+    "$PYTHON_BIN" - <<'PY'
+from torch_pin import (
+    CHANNEL,
+    torch_index_url_base,
+    torch_spec,
+    torchaudio_spec,
+    torchvision_spec,
+)
+
+print(CHANNEL)
+print(torch_spec())
+print(torchaudio_spec())
+print(torchvision_spec())
+print(torch_index_url_base())
+PY
+)
+TORCH_CHANNEL=$(echo "${TORCH_PIN_HELPERS}" | sed -n '1p')
+TORCH_SPEC=$(echo "${TORCH_PIN_HELPERS}" | sed -n '2p')
+TORCHAUDIO_SPEC=$(echo "${TORCH_PIN_HELPERS}" | sed -n '3p')
+TORCHVISION_SPEC=$(echo "${TORCH_PIN_HELPERS}" | sed -n '4p')
+TORCH_INDEX_URL=$(echo "${TORCH_PIN_HELPERS}" | sed -n '5p')
 
 # Copy requirements-lintrunner.txt from root to here
 cp ../../requirements-lintrunner.txt ./
@@ -112,8 +135,11 @@ docker build \
   --build-arg "GCC_VERSION=${GCC_VERSION}" \
   --build-arg "PYTHON_VERSION=${PYTHON_VERSION}" \
   --build-arg "MINICONDA_VERSION=${MINICONDA_VERSION}" \
-  --build-arg "TORCH_VERSION=${TORCH_VERSION}" \
-  --build-arg "PYTORCH_BUILD_MAX_JOBS=${PYTORCH_BUILD_MAX_JOBS:-}" \
+  --build-arg "TORCH_CHANNEL=${TORCH_CHANNEL}" \
+  --build-arg "TORCH_SPEC=${TORCH_SPEC}" \
+  --build-arg "TORCHAUDIO_SPEC=${TORCHAUDIO_SPEC}" \
+  --build-arg "TORCHVISION_SPEC=${TORCHVISION_SPEC}" \
+  --build-arg "TORCH_INDEX_URL=${TORCH_INDEX_URL}" \
   --build-arg "BUCK2_VERSION=${BUCK2_VERSION}" \
   --build-arg "LINTRUNNER=${LINTRUNNER:-}" \
   --build-arg "BUILD_DOCS=${BUILD_DOCS}" \

.ci/docker/common/install_pytorch.sh

@@ -10,47 +10,47 @@ set -ex
 # shellcheck source=/dev/null
 source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 
-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}"
-}
-
 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}"
+  local cache_args=()
+  if [ "${TORCH_CHANNEL}" = "test" ]; then
+    cache_args=("--no-cache-dir")
   fi
-  # Then build and install PyTorch
-  conda_run python setup.py bdist_wheel
-  pip_install "$(echo dist/*.whl)"
 
-  # Grab the pinned audio and vision commits from PyTorch
-  TORCHAUDIO_VERSION=release/2.11
-  export TORCHAUDIO_VERSION
-  TORCHVISION_VERSION=release/0.27
-  export TORCHVISION_VERSION
+  local wheelhouse
+  wheelhouse=$(mktemp -d)
+  chown ci-user:ci-user "${wheelhouse}"
+
+  local system_name
+  local system_arch
+  local python_version
+  system_name=$(uname)
+  system_arch=$(uname -m)
+  python_version=$(conda_run python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
+  local torch_wheel_cache_path="cached_artifacts/pytorch/executorch/pytorch_wheels/${system_name}/${system_arch}/${python_version}/cpu/${TORCH_CHANNEL}"
+  local torch_wheel_cache_uri="s3://gha-artifacts/${torch_wheel_cache_path}"
+
+  # Do not cache test-channel wheels in S3: RC artifacts may be re-uploaded
+  # under the same package version.
+  if [[ "${TORCH_CHANNEL}" != "test" ]] && command -v aws >/dev/null 2>&1; then
+    aws s3 sync "${torch_wheel_cache_uri}" "${wheelhouse}" || true
+  fi
 
-  install_domains
+  conda_run python -m pip download --progress-bar off --no-deps "${cache_args[@]}" \
+    --dest "${wheelhouse}" \
+    --find-links "${wheelhouse}" \
+    "${TORCH_SPEC}" "${TORCHVISION_SPEC}" "${TORCHAUDIO_SPEC}" \
+    --index-url "${TORCH_INDEX_URL}/cpu"
 
-  popd || true
-  # Clean up the cloned PyTorch repo to reduce the Docker image size
-  rm -rf pytorch
+  if [[ "${TORCH_CHANNEL}" != "test" && -z "${GITHUB_RUNNER:-}" ]] && command -v aws >/dev/null 2>&1; then
+    aws s3 sync "${wheelhouse}" "${torch_wheel_cache_uri}" \
+      --exclude "*" --include "*.whl" || true
+  fi
 
-  # Print sccache stats for debugging
-  as_ci_user sccache --show-stats
+  pip_install --force-reinstall "${cache_args[@]}" \
+    --find-links "${wheelhouse}" \
+    "${TORCH_SPEC}" "${TORCHVISION_SPEC}" "${TORCHAUDIO_SPEC}" \
+    --index-url "${TORCH_INDEX_URL}/cpu"
+  rm -rf "${wheelhouse}"
 }
 
 install_pytorch_and_domains

.ci/docker/ubuntu/Dockerfile

@@ -54,7 +54,7 @@ COPY ./common/install_conda.sh install_conda.sh
 COPY ./common/utils.sh utils.sh
 RUN bash ./install_conda.sh && rm install_conda.sh utils.sh /opt/conda/requirements-ci.txt /opt/conda/conda-env-ci.txt
 
-# Install sccache before building torch
+# Install sccache before building ExecuTorch.
 COPY ./common/install_cache.sh install_cache.sh
 ENV PATH /opt/cache/bin:$PATH
 COPY ./common/utils.sh utils.sh
@@ -65,7 +65,11 @@ ENV SCCACHE_REGION us-east-1
 ENV AWS_REGION us-east-1
 ENV AWS_DEFAULT_REGION us-east-1
 
-ARG TORCH_VERSION
+ARG TORCH_CHANNEL
+ARG TORCH_SPEC
+ARG TORCHAUDIO_SPEC
+ARG TORCHVISION_SPEC
+ARG TORCH_INDEX_URL
 ARG SKIP_PYTORCH
 ARG PYTORCH_BUILD_MAX_JOBS
 COPY ./common/install_pytorch.sh install_pytorch.sh

.ci/scripts/setup-linux.sh

@@ -13,16 +13,13 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 read -r BUILD_TOOL BUILD_MODE EDITABLE < <(parse_args "$@")
 echo "Build tool: $BUILD_TOOL, Mode: $BUILD_MODE"
 
-# As Linux job is running inside a Docker container, all of its dependencies
-# have already been installed, so we use PyTorch build from source here instead
-# of nightly. This allows CI to test against latest commits from PyTorch
 if [[ "${EDITABLE:-false}" == "true" ]]; then
-  install_executorch --use-pt-pinned-commit --editable
+  install_executorch --editable
 else
-  install_executorch --use-pt-pinned-commit
+  install_executorch
 fi
 build_executorch_runner "${BUILD_TOOL}" "${BUILD_MODE}"
 
 if [[ "${GITHUB_BASE_REF:-}" == *main* || "${GITHUB_BASE_REF:-}" == *gh* ]]; then
-  do_not_use_nightly_on_ci
+  verify_torch_matches_pin_on_ci
 fi

.ci/scripts/setup-macos.sh

@@ -126,22 +126,19 @@ fi
 
 # Install pinned torch before requirements-ci.txt so torchsr's transitive
 # torch dep is satisfied by the existing install and pip does not pull a
-# separate copy from PyPI. sccache is initialized above so source-build
-# cache misses still hit the cache.
+# separate copy from PyPI.
 print_cmake_info
 install_pytorch_and_domains
 
 install_pip_dependencies
 
-# install_executorch's --use-pt-pinned-commit skips re-installing torch since
-# install_pytorch_and_domains already installed the pinned build above.
 if [[ "$EDITABLE" == "true" ]]; then
-  install_executorch --use-pt-pinned-commit --editable
+  install_executorch --editable
 else
-  install_executorch --use-pt-pinned-commit
+  install_executorch
 fi
 build_executorch_runner "${BUILD_TOOL}" "${BUILD_MODE}"
 
 if [[ "${GITHUB_BASE_REF:-}" == *main* || "${GITHUB_BASE_REF:-}" == *gh* ]]; then
-  do_not_use_nightly_on_ci
+  verify_torch_matches_pin_on_ci
 fi

.ci/scripts/test_model_e2e.sh

@@ -284,7 +284,9 @@ elif [[ "$MODEL_NAME" == *whisper* ]] || [ "$MODEL_NAME" = "voxtral_realtime" ];
     fi
   fi
   pip install datasets soundfile
-  pip install torchcodec==0.11.0 --extra-index-url https://download.pytorch.org/whl/test/cpu
+  TORCHCODEC_PKG=$(python -c "from torch_pin import torchcodec_spec; print(torchcodec_spec())")
+  TORCHCODEC_INDEX=$(python -c "from torch_pin import torch_index_url_base; print(torch_index_url_base())")
+  pip install "$TORCHCODEC_PKG" --index-url "${TORCHCODEC_INDEX}/cpu"
   python -c "from datasets import load_dataset;import soundfile as sf;sample = load_dataset('distil-whisper/librispeech_long', 'clean', split='validation')[0]['audio'];sf.write('${MODEL_DIR}/$AUDIO_FILE', sample['array'][:sample['sampling_rate']*30], sample['sampling_rate'])"
 fi
 

.ci/scripts/test_wheel_package_qnn.sh

@@ -150,25 +150,37 @@ run_core_tests () {
   echo "=== [$LABEL] Installing wheel & deps ==="
   "$PIPBIN" install --upgrade pip
   "$PIPBIN" install "$WHEEL_FILE"
-  TORCH_VERSION=$(
+  # runpy.run_path uses a relative path, so the caller must run this script
+  # from the executorch repo root (where torch_pin.py lives).
+  TORCH_SPEC=$(
   "$PYBIN" - <<'PY'
 import runpy
 module_vars = runpy.run_path("torch_pin.py")
-print(module_vars["TORCH_VERSION"])
+print(module_vars["torch_spec"]())
 PY
 )
+  TORCH_INDEX=$(
+  "$PYBIN" - <<'PY'
+import runpy
+module_vars = runpy.run_path("torch_pin.py")
+print(module_vars["torch_index_url_base"]())
+PY
+)
+  TORCH_CACHE_ARGS_OUTPUT=$(
+  "$PYBIN" - <<'PY'
+import runpy
+module_vars = runpy.run_path("torch_pin.py")
+print(" ".join(module_vars["pip_cache_args"]()))
+PY
+)
+  TORCH_CACHE_ARGS=()
+  if [[ -n "${TORCH_CACHE_ARGS_OUTPUT}" ]]; then
+    read -r -a TORCH_CACHE_ARGS <<< "${TORCH_CACHE_ARGS_OUTPUT}"
+  fi
+  echo "=== [$LABEL] Install $TORCH_SPEC from ${TORCH_INDEX}/cpu ==="
 
-#   NIGHTLY_VERSION=$(
-#   "$PYBIN" - <<'PY'
-# import runpy
-# module_vars = runpy.run_path("torch_pin.py")
-# print(module_vars["NIGHTLY_VERSION"])
-# PY
-# )
-  echo "=== [$LABEL] Install torch==${TORCH_VERSION} ==="
-
-  # Install torch based on the pinned PyTorch version, preferring the PyTorch test index
-  "$PIPBIN" install torch=="${TORCH_VERSION}" --extra-index-url "https://download.pytorch.org/whl/test"
+  # Install torch based on the pinned PyTorch version from the channel index.
+  "$PIPBIN" install "${TORCH_CACHE_ARGS[@]}" "$TORCH_SPEC" --index-url "${TORCH_INDEX}/cpu"
   "$PIPBIN" install wheel
 
   # Install torchao based on the pinned commit from third-party/ao submodule

.ci/scripts/tests/test_torch_pin.py

@@ -0,0 +1,72 @@
+import importlib
+
+import pytest
+
+
+@pytest.fixture
+def pin():
+    import torch_pin
+
+    yield torch_pin
+    importlib.reload(torch_pin)
+
+
+@pytest.mark.parametrize(
+    "channel, expected_torch, expected_url",
+    [
+        (
+            "nightly",
+            "torch=={TORCH_VERSION}.{NIGHTLY_VERSION}",
+            "https://download.pytorch.org/whl/nightly",
+        ),
+        ("test", "torch=={TORCH_VERSION}", "https://download.pytorch.org/whl/test"),
+        ("release", "torch=={TORCH_VERSION}", "https://download.pytorch.org/whl"),
+    ],
+)
+def test_channel_resolution(pin, channel, expected_torch, expected_url):
+    pin.CHANNEL = channel
+    expected = expected_torch.format(
+        TORCH_VERSION=pin.TORCH_VERSION, NIGHTLY_VERSION=pin.NIGHTLY_VERSION
+    )
+    assert pin.torch_spec() == expected
+    assert pin.torch_index_url_base() == expected_url
+
+
+def test_all_specs_share_nightly_suffix(pin):
+    pin.CHANNEL = "nightly"
+    suffix = f".{pin.NIGHTLY_VERSION}"
+    assert pin.torch_spec().endswith(suffix)
+    assert pin.torchaudio_spec().endswith(suffix)
+    assert pin.torchcodec_spec().endswith(suffix)
+    assert pin.torchvision_spec().endswith(suffix)
+
+
+def test_specs_drop_suffix_off_nightly(pin):
+    pin.CHANNEL = "release"
+    assert pin.torch_spec() == f"torch=={pin.TORCH_VERSION}"
+    assert pin.torchaudio_spec() == f"torchaudio=={pin.TORCHAUDIO_VERSION}"
+    assert pin.torchcodec_spec() == f"torchcodec=={pin.TORCHCODEC_VERSION}"
+    assert pin.torchvision_spec() == f"torchvision=={pin.TORCHVISION_VERSION}"
+
+
+@pytest.mark.parametrize(
+    "channel, expected",
+    [
+        ("nightly", []),
+        ("test", ["--no-cache-dir"]),
+        ("release", []),
+    ],
+)
+def test_cache_args_only_disable_cache_for_test_channel(pin, channel, expected):
+    pin.CHANNEL = channel
+    assert pin.pip_cache_args() == expected
+
+
+def test_release_branches_derived_from_versions(pin):
+    assert pin.torch_branch() == f"release/{pin.TORCH_VERSION.rsplit('.', 1)[0]}"
+    assert pin.torchaudio_branch() == (
+        f"release/{pin.TORCHAUDIO_VERSION.rsplit('.', 1)[0]}"
+    )
+    assert pin.torchvision_branch() == (
+        f"release/{pin.TORCHVISION_VERSION.rsplit('.', 1)[0]}"
+    )