Centralize torch-family version pins in torch_pin.py (#19155)

Centralize torch-family version pins in torch_pin.py

torch_pin.py is now the single source of truth for torch + the three
domain libraries (vision/audio/codec). It exposes a CHANNEL field
(nightly/test/release), the four version constants, NIGHTLY_VERSION,
and helpers — torch_spec() / torchaudio_spec() / torchcodec_spec() /
torchvision_spec() emit the right pip spec, torch_index_url_base()
returns the right wheel index, and torch_branch() / torchaudio_branch()
/ torchvision_branch() derive the upstream release/M.N branch from each
package's version. Every consumer — install_requirements.py, the two
install_pytorch.sh / utils.sh shell helpers, test_model_e2e.sh,
test_wheel_package_qnn.sh, the moshi/mimi install_requirements.sh, the
update_pytorch_pin.py script, and the weekly bump workflow — reads
through these helpers instead of re-encoding the version strings.

Switching to a release candidate is now a one-line change (CHANNEL =
"test") plus bumping the four version constants. The header in
torch_pin.py walks through the procedure.

update_pytorch_pin.py imports CHANNEL / NIGHTLY_VERSION / torch_branch
directly (no more regex parsing of the file). For nightly it pins to
an immutable SHA looked up by date; for test/release it writes
torch_branch() (e.g. "release/2.12") into
.ci/docker/ci_commit_pins/pytorch.txt so git checkout follows
cherry-picks as they land.

The weekly-pytorch-pin-bump workflow is guarded on CHANNEL == "nightly"
and uses an in-place re.sub on NIGHTLY_VERSION (the previous
`printf '...' > torch_pin.py` would have clobbered the new constants
and helpers).

test/test_torch_pin.py covers all three channels, all four specs, and
the release/M.N branch derivation.

Co-authored-by: Claude <noreply@anthropic.com>

Co-authored-by: Claude <noreply@anthropic.com>

Author: mergennachin

.ci/docker/build.sh

@@ -92,6 +92,18 @@ esac
 TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
 BUILD_DOCS=1
 
+# Pull channel + spec/url helpers out of torch_pin.py so install_pytorch.sh
+# (which runs inside the docker build, where torch_pin.py isn't available)
+# can decide between wheel install (test/release) and source build (nightly).
+# 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" -c "from torch_pin import CHANNEL, torch_spec, torchaudio_spec, torchvision_spec, torch_index_url_base; print(CHANNEL); print(torch_spec()); print(torchaudio_spec()); print(torchvision_spec()); print(torch_index_url_base())")
+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 ./
 
@@ -104,6 +116,11 @@ docker build \
   --build-arg "PYTHON_VERSION=${PYTHON_VERSION}" \
   --build-arg "MINICONDA_VERSION=${MINICONDA_VERSION}" \
   --build-arg "TORCH_VERSION=${TORCH_VERSION}" \
+  --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/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-release/2.11
+release/2.11

.ci/docker/common/install_pytorch.sh

@@ -17,6 +17,24 @@ install_domains() {
 }
 
 install_pytorch_and_domains() {
+  if [ "${TORCH_CHANNEL}" != "nightly" ]; then
+    # Test/release: install the published wheels directly. The specs and URL
+    # are passed in as docker build args (computed from torch_pin.py by
+    # .ci/docker/build.sh). RC wheels at /whl/test/ get re-uploaded under the
+    # same version, so use --no-cache-dir there to avoid stale cache hits.
+    local cache_flag=""
+    if [ "${TORCH_CHANNEL}" = "test" ]; then
+      cache_flag="--no-cache-dir"
+    fi
+    pip_install --force-reinstall ${cache_flag} \
+      "${TORCH_SPEC}" "${TORCHVISION_SPEC}" "${TORCHAUDIO_SPEC}" \
+      --index-url "${TORCH_INDEX_URL}/cpu"
+    return
+  fi
+
+  # Nightly: build pytorch from source against the pinned SHA in pytorch.txt
+  # so we catch upstream regressions, then install audio/vision from the
+  # commits that pytorch itself pins.
   git clone https://github.com/pytorch/pytorch.git
 
   # Fetch the target commit
@@ -27,14 +45,19 @@ install_pytorch_and_domains() {
   chown -R ci-user .
 
   export _GLIBCXX_USE_CXX11_ABI=1
+  # PyTorch's FindARM.cmake hard-fails when the SVE+BF16 compile probe
+  # doesn't pass — gcc-11 in this image is too old to accept the combined
+  # NEON/SVE/bfloat16 intrinsics the probe exercises. Executorch's aarch64
+  # runtime targets (phones, embedded) don't use SVE, so bypass the check.
+  export BUILD_IGNORE_SVE_UNAVAILABLE=1
   # 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
+  # Defer to PyTorch's own pinned audio/vision commits.
+  TORCHAUDIO_VERSION=$(cat .github/ci_commit_pins/audio.txt)
   export TORCHAUDIO_VERSION
-  TORCHVISION_VERSION=release/0.26
+  TORCHVISION_VERSION=$(cat .github/ci_commit_pins/vision.txt)
   export TORCHVISION_VERSION
 
   install_domains

.ci/docker/ubuntu/Dockerfile

@@ -64,6 +64,11 @@ ENV SCCACHE_S3_KEY_PREFIX executorch
 ENV SCCACHE_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
 COPY ./common/install_pytorch.sh install_pytorch.sh
 COPY ./common/utils.sh utils.sh

.ci/scripts/test_model_e2e.sh

@@ -260,7 +260,10 @@ if [ "$AUDIO_URL" != "" ]; then
 elif [[ "$MODEL_NAME" == *whisper* ]] || [ "$MODEL_NAME" = "voxtral_realtime" ]; then
   conda install -y -c conda-forge "ffmpeg<8"
   pip install datasets soundfile
-  pip install torchcodec==0.11.0 --extra-index-url https://download.pytorch.org/whl/test/cpu
+  # We pushd'd into EXECUTORCH_ROOT above, so torch_pin is importable here.
+  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" --extra-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,26 @@ 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
+)
+  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_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,54 @@
+import importlib
+
+import pytest
+
+
+@pytest.fixture
+def pin():
+    """Yield a fresh import of torch_pin so tests can mutate CHANNEL safely."""
+    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 = "test"
+    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}"
+
+
+def test_torch_branch_derived_from_version(pin):
+    assert pin.torch_branch() == f"release/{pin.TORCH_VERSION.rsplit('.', 1)[0]}"

.ci/scripts/utils.sh

@@ -53,7 +53,7 @@ dedupe_macos_loader_path_rpaths() {
   pushd ..
   torch_lib_dir=$(python -c "import importlib.util; print(importlib.util.find_spec('torch').submodule_search_locations[0])")/lib
   popd
-
+  
   if [[ -z "${torch_lib_dir}" || ! -d "${torch_lib_dir}" ]]; then
     return
   fi
@@ -89,6 +89,30 @@ install_domains() {
 }
 
 install_pytorch_and_domains() {
+  # CWD is the executorch repo root, where torch_pin.py lives.
+  TORCH_CHANNEL=$(python -c "from torch_pin import CHANNEL; print(CHANNEL)")
+
+  if [ "${TORCH_CHANNEL}" != "nightly" ]; then
+    # Test/release: install the published wheels directly from torch_pin.py's
+    # channel index, skipping the source-build path entirely. RC wheels at
+    # /whl/test/ get re-uploaded under the same version, so use --no-cache-dir
+    # there to avoid stale cache hits.
+    local torch_spec=$(python -c "from torch_pin import torch_spec; print(torch_spec())")
+    local torchvision_spec=$(python -c "from torch_pin import torchvision_spec; print(torchvision_spec())")
+    local torchaudio_spec=$(python -c "from torch_pin import torchaudio_spec; print(torchaudio_spec())")
+    local torch_index_url=$(python -c "from torch_pin import torch_index_url_base; print(torch_index_url_base())")
+    local cache_flag=""
+    if [ "${TORCH_CHANNEL}" = "test" ]; then
+      cache_flag="--no-cache-dir"
+    fi
+    pip install --force-reinstall ${cache_flag} \
+      "${torch_spec}" "${torchvision_spec}" "${torchaudio_spec}" \
+      --index-url "${torch_index_url}/cpu"
+    return
+  fi
+
+  # Nightly: source-build pytorch from the pinned SHA so CI catches upstream
+  # regressions; pytorch's own audio/vision pins drive those installs.
   pushd .ci/docker || return
   TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
   popd || return
@@ -140,10 +164,10 @@ install_pytorch_and_domains() {
   fi
 
   dedupe_macos_loader_path_rpaths
-  # Grab the pinned audio and vision commits from PyTorch
-  TORCHAUDIO_VERSION=release/2.11
+  # We're on the nightly path here; defer to PyTorch's own pinned commits.
+  TORCHAUDIO_VERSION=$(cat .github/ci_commit_pins/audio.txt)
   export TORCHAUDIO_VERSION
-  TORCHVISION_VERSION=release/0.26
+  TORCHVISION_VERSION=$(cat .github/ci_commit_pins/vision.txt)
   export TORCHVISION_VERSION
 
   install_domains
@@ -218,17 +242,21 @@ download_stories_model_artifacts() {
 }
 
 do_not_use_nightly_on_ci() {
-  # An assert to make sure that we are not using PyTorch nightly on CI to prevent
-  # regression as documented in https://github.com/pytorch/executorch/pull/6564
-  TORCH_VERSION=$(pip list | grep -w 'torch ' | awk -F ' ' {'print $2'} | tr -d '\n')
+  # Sanity check that prevents accidentally landing a PR that pins to PyTorch
+  # nightly without exercising the source-build path (see #6564).
+  #
+  # For CHANNEL=nightly, CI source-builds pytorch from the SHA in pytorch.txt,
+  # so the installed torch shows up as e.g. 2.13.0a0+gitc8a648d — assert that.
+  # For CHANNEL=test/release, we install published wheels by design (e.g.
+  # 2.11.0), so the +git assertion doesn't apply.
+  TORCH_CHANNEL=$(python -c "from torch_pin import CHANNEL; print(CHANNEL)")
+  if [ "${TORCH_CHANNEL}" != "nightly" ]; then
+    return 0
+  fi
 
-  # The version of PyTorch building from source looks like 2.6.0a0+gitc8a648d that
-  # includes the commit while nightly (2.6.0.dev20241019+cpu) or release (2.6.0)
-  # won't have that. Note that we couldn't check for the exact commit from the pin
-  # ci_commit_pins/pytorch.txt here because the value will be different when running
-  # this on PyTorch CI
+  TORCH_VERSION=$(pip list | grep -w 'torch ' | awk -F ' ' {'print $2'} | tr -d '\n')
   if [[ "${TORCH_VERSION}" != *"+git"* ]]; then
-    echo "Unexpected torch version. Expected binary built from source, got ${TORCH_VERSION}"
+    echo "Unexpected torch version. Expected binary built from source for CHANNEL=nightly, got ${TORCH_VERSION}"
     exit 1
   fi
 }