: GHA: pin torch-2.9.0 (#1742)

Summary:

yesterday's fix D86025607 is flaky. today's observed behavior shows the torch API still suffers from caching / runner variability. this fix is more authoritative: pin torch to 2.9.0 (stable):
replace `pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126` with `pip install torch==2.9.0+cu126 --index-url https://download.pytorch.org/whl/cu126`

pulling "nightly" exposes us to non-determinism. relying on a published version restores hermeticity. anyone can stamp/land and we should expect to recover green actions.

Differential Revision: D86052365

Author: shayne-fletcher

.github/workflows/build-cpu.yml

@@ -29,8 +29,9 @@ jobs:
         # Setup build environment (conda + system deps + rust + build deps)
         setup_build_environment
 
-        # Install torch nightly (CPU version)
-        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
+        # Install torch 2.9.0 (CPU version) (versioned wheel preferred
+        # for build hermeticity)
+        pip install torch==2.9.0+cpu --index-url https://download.pytorch.org/whl/cpu
         pip install -r build-requirements.txt
 
         # Build monarch (No tensor engine, CPU version)

.github/workflows/build-cuda.yml

@@ -17,7 +17,7 @@ jobs:
         include:
           - name: 4xlargegpu
             runs-on: linux.g5.4xlarge.nvidia.gpu
-            torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu126'
+            torch-spec: 'torch==2.9.0+cu126 --index-url https://download.pytorch.org/whl/cu126'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.6"
     with:

.github/workflows/doc_build.yml

@@ -39,7 +39,7 @@ jobs:
         pip install -r docs/requirements.txt
 
         # Install PyTorch with CUDA support (matching build-cuda.yml)
-        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126
+        pip install torch==2.9.0+cu126 --index-url https://download.pytorch.org/whl/cu126
 
         # Set environment variables for CUDA build
         export USE_CUDA=1

.github/workflows/test-cpu-python.yml

@@ -32,7 +32,7 @@ jobs:
         export USE_TENSOR_ENGINE=0
 
         # Install PyTorch nightly
-        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
+        pip install torch==2.9.0+cpu --index-url https://download.pytorch.org/whl/cpu
 
         # Install the built wheel from artifact
         install_wheel_from_artifact

.github/workflows/test-gpu-python.yml

@@ -22,7 +22,7 @@ jobs:
         include:
           - name: 4xlargegpu
             runs-on: linux.g5.4xlarge.nvidia.gpu
-            torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu126'
+            torch-spec: 'torch==2.9.0+cu126 --index-url https://download.pytorch.org/whl/cu126'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.6"
     with:

.github/workflows/test-gpu-rust.yml

@@ -22,7 +22,7 @@ jobs:
         include:
           - name: 4xlargegpu
             runs-on: linux.g5.4xlarge.nvidia.gpu
-            torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu126'
+            torch-spec: 'torch==2.9.0+cu126 --index-url https://download.pytorch.org/whl/cu126'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.6"
     with:

.github/workflows/wheels.yml

@@ -21,7 +21,7 @@ jobs:
         include:
           - name: 4xlarge
             runs-on: linux.g5.4xlarge.nvidia.gpu
-            install-args: '--pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126'
+            install-args: 'torch==2.9.0+cu126 --index-url https://download.pytorch.org/whl/cu126'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.6"
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
@@ -35,6 +35,7 @@ jobs:
       script: |
         source scripts/common-setup.sh
         setup_build_environment ${{ matrix.python-version }} "${{ matrix.install-args }}"
+        pip install -r build-requirements.txt
 
         # Setup Tensor Engine dependencies
         setup_tensor_engine

scripts/common-setup.sh

@@ -64,7 +64,7 @@ setup_tensor_engine() {
 # Install PyTorch with C++ development headers (libtorch) for Rust compilation
 setup_pytorch_with_headers() {
     local gpu_arch_version=${1:-"12.6"}
-    local torch_spec=${2:-"--pre torch --index-url https://download.pytorch.org/whl/nightly/cu126"}
+    local torch_spec=${2:-"torch==2.9.0+cu126 --index-url https://download.pytorch.org/whl/cu126"}
 
     echo "Setting up PyTorch with C++ headers (GPU arch: ${gpu_arch_version})..."
 
@@ -110,6 +110,10 @@ setup_build_environment() {
     setup_conda_environment "${python_version}"
     install_system_dependencies
     setup_rust_toolchain
+    if [ -n "$install_args" ]; then
+        echo "Installing PyTorch with: ${install_args}"
+        pip install ${install_args}
+    fi
 }
 
 # Detect and configure CUDA environment for linking