Reland torch-family pin centralization

.ci/docker/build.sh

@@ -97,6 +97,34 @@ 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
+# 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" - <<'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')
+
 if [[ "${GCC_VERSION:-}" == "11" && -z "${SKIP_PYTORCH:-}" ]]; then
   PYTORCH_BUILD_MAX_JOBS=6
 fi
@@ -113,6 +141,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 "PYTORCH_BUILD_MAX_JOBS=${PYTORCH_BUILD_MAX_JOBS:-}" \
   --build-arg "BUCK2_VERSION=${BUCK2_VERSION}" \
   --build-arg "LINTRUNNER=${LINTRUNNER:-}" \

.ci/docker/common/install_pytorch.sh

@@ -17,6 +17,19 @@ 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.
+    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
+
   git clone https://github.com/pytorch/pytorch.git
 
   # Fetch the target commit
@@ -37,10 +50,10 @@ install_pytorch_and_domains() {
   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.27
+  TORCHVISION_VERSION=$(cat .github/ci_commit_pins/vision.txt)
   export TORCHVISION_VERSION
 
   install_domains

.ci/docker/ubuntu/Dockerfile

@@ -66,6 +66,11 @@ 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/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" --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} ==="
 
-#   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" --extra-index-url "$TORCH_INDEX"
   "$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,59 @@
+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 = "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_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]}"
+    )

.ci/scripts/utils.sh

@@ -89,6 +89,34 @@ 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
+    local torchvision_spec
+    local torchaudio_spec
+    local torch_index_url
+    torch_spec=$(python -c "from torch_pin import torch_spec; print(torch_spec())")
+    torchvision_spec=$(python -c "from torch_pin import torchvision_spec; print(torchvision_spec())")
+    torchaudio_spec=$(python -c "from torch_pin import torchaudio_spec; print(torchaudio_spec())")
+    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
@@ -175,10 +203,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.27
+  TORCHVISION_VERSION=$(cat .github/ci_commit_pins/vision.txt)
   export TORCHVISION_VERSION
 
   install_domains
@@ -253,17 +281,17 @@ 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')
+  # For CHANNEL=nightly, CI source-builds pytorch from the SHA in pytorch.txt,
+  # so the installed torch should include +git. For CHANNEL=test/release, we
+  # install published wheels by design, so this 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
 }