[CI] Fix Linux aarch64 nightly wheel uploads

.github/workflows/build-wheels-aarch64-linux.yml

@@ -24,9 +24,6 @@ on:
         required: false
         type: boolean
         default: false
-    secrets:
-      PYPI_TOKEN:
-        required: false
   workflow_dispatch:
 
 permissions:
@@ -59,12 +56,12 @@ jobs:
       test-infra-ref: ${{ inputs.test-infra-ref || 'main' }}
       build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
       package-name: ${{ matrix.package-name }}
-      trigger-event: ${{ inputs.nightly && 'schedule' || github.event_name }}
+      # Not 'schedule' for nightly: that enables test-infra's S3+PyPI upload jobs,
+      # which need an AWS role tensordict cannot assume; nightly_build.yml uploads
+      # to PyPI instead. ('' is falsy and would fall through to github.event_name.)
+      trigger-event: ${{ inputs.nightly && 'skip-upload' || github.event_name }}
       env-var-script: .github/scripts/version_script.sh
       architecture: aarch64
       setup-miniconda: false
       build-command: ${{ inputs.nightly && 'bash .github/scripts/build-aarch64-nightly-wheel.sh' || 'pip wheel --no-deps . && mkdir -p dist && mv tensordict*.whl ./dist' }}
       build-platform: python-build-package
-      upload-to-pypi: ${{ inputs.nightly && 'cpu' || '' }}
-    secrets:
-      PYPI_API_TOKEN: ${{ secrets.PYPI_TOKEN }}

.github/workflows/nightly_build.yml

@@ -35,12 +35,38 @@ jobs:
     uses: ./.github/workflows/build-wheels-aarch64-linux.yml
     with:
       nightly: true
-    secrets:
-      PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
     permissions:
       id-token: write
       contents: read
 
+  upload-wheel-aarch64:
+    # Don't run on forked repos.
+    # !cancelled(): upload whatever built even if some matrix entries failed.
+    if: ${{ !cancelled() && github.repository_owner == 'pytorch' && (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule') }}
+    needs: build-wheel-aarch64
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Download built wheels
+        uses: actions/download-artifact@v4
+        with:
+          # Artifacts uploaded by test-infra, named pytorch_tensordict__<python>_cpu_aarch64
+          pattern: pytorch_tensordict_*_aarch64
+          path: /tmp/wheels
+          merge-multiple: true
+      - name: Push tensordict Binary to PYPI
+        env:
+          PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
+        run: |
+          # Fail if no wheels were downloaded (e.g. test-infra renamed its artifacts).
+          ls /tmp/wheels/tensordict_nightly-*.whl
+          python3 -mpip install twine
+          python3 -m twine upload \
+              --username __token__ \
+              --password "$PYPI_TOKEN" \
+              --skip-existing \
+              /tmp/wheels/tensordict_nightly-*.whl \
+              --verbose
+
   build-wheel-unix:
     # Don't run on forked repos.
     if: github.repository_owner == 'pytorch'