Correctly checkout test-infra on Nova _binary_upload job (#4882)

I made a mistake in #4877 when
not explicitly select `test-infra` as the repo to checkout for the GHA.
When running on domains, the default repos are domains, i.e.
`pytorch/vision`, and they don't have the GHA we need, namely
`setup-binary-upload`.

An example failure when testing on vision nightly
https://github.com/pytorch/vision/actions/runs/7533535440/job/20506343331

.github/actions/setup-binary-upload/action.yml

@@ -53,7 +53,15 @@ runs:
         # and "/" is not allowed in artifact names
         echo "ARTIFACT_NAME=${REPOSITORY/\//_}_${REF}_${PYTHON_VERSION}_${CU_VERSION}_${ARCH}" >> "${GITHUB_ENV}"
 
+    # Need to checkout the target repository to run pkg-helpers
+    - uses: actions/checkout@v3
+      with:
+        repository: ${{ inputs.repository }}
+        ref: ${{ inputs.ref }}
+        path: ${{ inputs.repository }}
+
     - name: Generate env variables from pytorch_pkg_helpers
+      working-directory: ${{ inputs.repository }}
       shell: bash
       env:
         REPOSITORY: ${{ inputs.repository }}
@@ -64,7 +72,7 @@ runs:
       run: |
         set -ex
 
-        python -m pip install tools/pkg-helpers
+        python -m pip install ${GITHUB_WORKSPACE}/test-infra/tools/pkg-helpers
 
         BUILD_ENV_FILE="${RUNNER_TEMP}/build_env_${GITHUB_RUN_ID}"
         python -m pytorch_pkg_helpers > "${BUILD_ENV_FILE}"

.github/workflows/_binary_upload.yml

@@ -11,6 +11,14 @@ on:
         description: 'Reference to checkout, defaults to "nightly"'
         default: 'nightly'
         type: string
+      test-infra-repository:
+        description: "Test infra repository to use"
+        default: "pytorch/test-infra"
+        type: string
+      test-infra-ref:
+        description: "Test infra reference to use"
+        default: ""
+        type: string
       build-matrix:
         description: "Build matrix to utilize"
         default: ''
@@ -36,9 +44,13 @@ jobs:
     name: ${{ matrix.build_name }}
     steps:
       - uses: actions/checkout@v3
+        with:
+          repository: ${{ inputs.test-infra-repository }}
+          ref: ${{ inputs.test-infra-ref }}
+          path: test-infra
 
       # For pytorch_pkg_helpers which we need to run to generate the artifact name and target S3 buckets
-      - uses: ./.github/actions/setup-binary-upload
+      - uses: ./test-infra/.github/actions/setup-binary-upload
         with:
           repository: ${{ inputs.repository }}
           ref: ${{ inputs.ref }}
@@ -47,7 +59,7 @@ jobs:
           arch: ${{ inputs.architecture }}
           upload-to-base-bucket: ${{ matrix.upload_to_base_bucket }}
 
-      - uses: ./.github/actions/set-channel
+      - uses: ./test-infra/.github/actions/set-channel
 
       - name: Download the artifact
         uses: actions/download-artifact@v3

.github/workflows/build_wheels_linux.yml

@@ -235,6 +235,8 @@ jobs:
     with:
       repository: ${{ inputs.repository }}
       ref: ${{ inputs.ref }}
+      test-infra-repository: ${{ inputs.test-infra-repository }}
+      test-infra-ref: ${{ inputs.test-infra-ref }}
       build-matrix: ${{ inputs.build-matrix }}
       architecture: ${{ inputs.architecture }}
       trigger-event: ${{ inputs.trigger-event }}

.github/workflows/build_wheels_macos.yml

@@ -214,6 +214,8 @@ jobs:
     with:
       repository: ${{ inputs.repository }}
       ref: ${{ inputs.ref }}
+      test-infra-repository: ${{ inputs.test-infra-repository }}
+      test-infra-ref: ${{ inputs.test-infra-ref }}
       build-matrix: ${{ inputs.build-matrix }}
       trigger-event: ${{ inputs.trigger-event }}
 

.github/workflows/build_wheels_windows.yml

@@ -202,6 +202,8 @@ jobs:
     with:
       repository: ${{ inputs.repository }}
       ref: ${{ inputs.ref }}
+      test-infra-repository: ${{ inputs.test-infra-repository }}
+      test-infra-ref: ${{ inputs.test-infra-ref }}
       build-matrix: ${{ inputs.build-matrix }}
       trigger-event: ${{ inputs.trigger-event }}