Add pytorch-version-override to the build_wheels_*.yml workflows

This lets repos override the torch pip version to install before
building the wheel.

Author: dbort

.github/workflows/test_infra_build_wheels_linux.yml

@@ -23,6 +23,10 @@ on:
         description: "Build matrix to utilize"
         default: ""
         type: string
+      pytorch-version-override:
+        description: "Manual override of the pytorch version to install"
+        default: ""
+        type: string
       pre-script:
         description: "Pre script to run prior to build"
         default: ""
@@ -136,11 +140,18 @@ jobs:
           set -e
       - uses: pytorch/test-infra/.github/actions/set-channel@release/2.3
       - name: Set PYTORCH_VERSION
-        if: ${{ env.CHANNEL == 'test' }}
+        if: ${{ inputs.pytorch-version-override != '' || 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}"
+          if [[ ! -z "${{ inputs.pytorch-version-override }}" ]]; then
+            # If there is an override, use it instead of the version associated
+            # with the channel.
+            echo "PYTORCH_VERSION=${{ inputs.pytorch-version-override }}" >> "${GITHUB_ENV}"
+          else
+            # When building RC (test channel), 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}"
+          fi
       - uses: pytorch/test-infra/.github/actions/setup-binary-builds@release/2.3
         env:
           PLATFORM: ${{ inputs.architecture == 'aarch64'  && 'linux-aarch64' || ''}}

.github/workflows/test_infra_build_wheels_macos.yml

@@ -23,6 +23,10 @@ on:
         description: "Build matrix to utilize"
         default: ""
         type: string
+      pytorch-version-override:
+        description: "Manual override of the pytorch version to install"
+        default: ""
+        type: string
       pre-script:
         description: "Pre script to run prior to build"
         default: ""
@@ -101,11 +105,18 @@ jobs:
           path: test-infra
       - uses: pytorch/test-infra/.github/actions/set-channel@release/2.3
       - name: Set PYTORCH_VERSION
-        if: ${{ env.CHANNEL == 'test' }}
+        if: ${{ inputs.pytorch-version-override != '' || 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}"
+          if [[ ! -z "${{ inputs.pytorch-version-override }}" ]]; then
+            # If there is an override, use it instead of the version associated
+            # with the channel.
+            echo "PYTORCH_VERSION=${{ inputs.pytorch-version-override }}" >> "${GITHUB_ENV}"
+          else
+            # When building RC (test channel), 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}"
+          fi
       - uses: pytorch/test-infra/.github/actions/setup-binary-builds@release/2.3
         with:
           repository: ${{ inputs.repository }}