Specific PYTORCH_VERSION when building RC for test channel (#4803)

With 2 candidate RC in the test channel (2.1.2 and 2.2.0), we need to
have a way to specify which PYTORCH_VERSION to use. The issue surfaces
when we are trying to cut RC3 for 2.1.2 today while there is already RC1
for 2.2.0 in place. In the current state, `PYTHON_VERSION=3.11
PACKAGE_TYPE=conda CU_VERSION=cpu ARCH_NAME=arm64 CHANNEL=test python -m
pytorch_pkg_helpers` will set `PYTORCH_VERSION` to 2.2.0 whenever
domains are built

For example,
https://github.com/pytorch/vision/actions/runs/7190086502/job/19582646141#step:6:743
build M1 conda for vision 0.16.2 and it's expected to use 2.1.2 there,
not 2.2.0.

The fix here is to:

1. Correctly set the channel to `test` when building RC
2. If the channel is test, use the current candidate version as
PYTORCH_VERSION. The value is exposed via the build matrix as
`stable_version`
3. When doing branch cut for test-infra, the `CURRENT_CANDIDATE_VERSION`
variable could be set accordingly, for example, use 2.1.2 for
`release/2.1` and 2.2.0 for `release/2.2`

### Testing plan

This PR #4803 tests `nightly`
while the cherry pick #4804
covers `test` channel

.github/actions/set-channel/action.yml

@@ -6,11 +6,6 @@ description: Add CHANNEL to GITHUB_ENV
 runs:
   using: composite
   steps:
-    - name: Set default CHANNEL
-      shell: bash
-      run: |
-        set -euxo pipefail
-        echo "CHANNEL=nightly" >> "${GITHUB_ENV}"
     - name: Set CHANNEL for tagged pushes
       if: ${{ github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/') }}
       shell: bash
@@ -26,3 +21,12 @@ runs:
       run: |
         set -euxo pipefail
         echo "CHANNEL=test" >> "$GITHUB_ENV"
+    - name: Set default CHANNEL
+      # Set this to nightly only when the CHANNEL hasn't been set yet. In GHA, once
+      # an env is set, it's fixed unless we can figure out a way to overwrite it in
+      # $GITHUB_ENV
+      if: ${{ env.CHANNEL == '' }}
+      shell: bash
+      run: |
+        set -euxo pipefail
+        echo "CHANNEL=nightly" >> "${GITHUB_ENV}"

.github/workflows/build_conda_linux.yml

@@ -96,6 +96,12 @@ jobs:
           ref: ${{ inputs.test-infra-ref }}
           path: test-infra
       - uses: ./test-infra/.github/actions/set-channel
+      - name: Set PYTORCH_VERSION
+        if: ${{ env.CHANNEL == 'test' }}
+        run: |
+          # When building RC, set the version to be the current candidate version,
+          # otherwise, leave it alone so nightly will pick up the latest
+          echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}"
       - uses: ./test-infra/.github/actions/setup-binary-builds
         with:
           repository: ${{ inputs.repository }}

.github/workflows/build_conda_macos.yml

@@ -106,6 +106,12 @@ jobs:
       - name: Set env variable for architecture name
         run: echo "ARCH_NAME=$(uname -m)" >> "${GITHUB_ENV}"
       - uses: ./test-infra/.github/actions/set-channel
+      - name: Set PYTORCH_VERSION
+        if: ${{ env.CHANNEL == 'test' }}
+        run: |
+          # When building RC, set the version to be the current candidate version,
+          # otherwise, leave it alone so nightly will pick up the latest
+          echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}"
       - uses: ./test-infra/.github/actions/setup-binary-builds
         with:
           repository: ${{ inputs.repository }}

.github/workflows/build_conda_windows.yml

@@ -96,6 +96,12 @@ jobs:
           ref: ${{ inputs.test-infra-ref }}
           path: test-infra
       - uses: ./test-infra/.github/actions/set-channel
+      - name: Set PYTORCH_VERSION
+        if: ${{ env.CHANNEL == 'test' }}
+        run: |
+          # When building RC, set the version to be the current candidate version,
+          # otherwise, leave it alone so nightly will pick up the latest
+          echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}"
       - uses: ./test-infra/.github/actions/setup-ssh
         name: Setup SSH
         with:

.github/workflows/build_wheels_linux.yml

@@ -134,6 +134,12 @@ jobs:
           echo "/opt/conda/bin" >> $GITHUB_PATH
           set -e
       - uses: ./test-infra/.github/actions/set-channel
+      - name: Set PYTORCH_VERSION
+        if: ${{ env.CHANNEL == 'test' }}
+        run: |
+          # When building RC, set the version to be the current candidate version,
+          # otherwise, leave it alone so nightly will pick up the latest
+          echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}"
       - uses: ./test-infra/.github/actions/setup-binary-builds
         env:
           PLATFORM: ${{ inputs.architecture == 'aarch64'  && 'linux-aarch64' || ''}}

.github/workflows/build_wheels_macos.yml

@@ -98,6 +98,12 @@ jobs:
           ref: ${{ inputs.test-infra-ref }}
           path: test-infra
       - uses: ./test-infra/.github/actions/set-channel
+      - name: Set PYTORCH_VERSION
+        if: ${{ env.CHANNEL == 'test' }}
+        run: |
+          # When building RC, set the version to be the current candidate version,
+          # otherwise, leave it alone so nightly will pick up the latest
+          echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}"
       - uses: ./test-infra/.github/actions/setup-binary-builds
         with:
           repository: ${{ inputs.repository }}

.github/workflows/build_wheels_windows.yml

@@ -105,6 +105,12 @@ jobs:
         run: |
           echo "C:/Jenkins/Miniconda3/Scripts" >> $GITHUB_PATH
       - uses: ./test-infra/.github/actions/set-channel
+      - name: Set PYTORCH_VERSION
+        if: ${{ env.CHANNEL == 'test' }}
+        run: |
+          # When building RC, set the version to be the current candidate version,
+          # otherwise, leave it alone so nightly will pick up the latest
+          echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}"
       - uses: ./test-infra/.github/actions/setup-binary-builds
         with:
           repository: ${{ inputs.repository }}

.github/workflows/generate_binary_build_matrix.yml

@@ -13,7 +13,7 @@ on:
         type: string
       channel:
         description: "Channel to use (nightly, test, release, all)"
-        default: "nightly"
+        default: ""
         type: string
       test-infra-repository:
         description: "Test infra repository to use"
@@ -60,7 +60,7 @@ jobs:
         env:
           PACKAGE_TYPE: ${{ inputs.package-type }}
           OS: ${{ inputs.os }}
-          CHANNEL: ${{ inputs.channel }}
+          CHANNEL: ${{ inputs.channel != '' && inputs.channel || env.CHANNEL }}
           WITH_CUDA: ${{ inputs.with-cuda }}
           WITH_ROCM: ${{ inputs.with-rocm }}
           WITH_CPU: ${{ inputs.with-cpu }}

tools/scripts/generate_binary_build_matrix.py

@@ -58,7 +58,7 @@
 ROCM = "rocm"
 
 
-CURRENT_CANDIDATE_VERSION = "2.1.2"
+CURRENT_CANDIDATE_VERSION = "2.2.0"
 CURRENT_STABLE_VERSION = "2.1.1"
 mod.CURRENT_VERSION = CURRENT_STABLE_VERSION