Use oidc for wheel upload

pytorch · Jul 23, 2024 · db7cd30 · db7cd30
1 parent b2a654a
commit db7cd30
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 33 deletions.
diff --git a/.github/workflows/_build_test_upload.yml b/.github/workflows/_build_test_upload.yml
@@ -17,10 +17,6 @@ on:
         default: true
         type: boolean
     secrets:
-      PYTORCH_BINARY_AWS_ACCESS_KEY_ID:
-        required: true
-      PYTORCH_BINARY_AWS_SECRET_ACCESS_KEY:
-        required: true
       PYPI_TOKEN:
         required: false
       CONDA_PYTORCHBOT_TOKEN:
@@ -30,6 +26,10 @@ on:
       CONDA_NIGHTLY_PYTORCHBOT_TOKEN:
         required: false
 
+permissions:
+  id-token: write
+  contents: read
+
 jobs:
   get_release_type:
     runs-on: ubuntu-latest
@@ -158,6 +158,19 @@ jobs:
     outputs:
       upload: ${{ steps.trigger_upload.outputs.value }}
     steps:
+      - name: Configure aws credentials (pytorch account)
+        if: ${{ needs.get_release_type.outputs.type == 'nightly' }}
+        uses: aws-actions/configure-aws-credentials@v3
+        with:
+          role-to-assume: arn:aws:iam::749337293305:role/gha_workflow_nightly_build_wheels
+          aws-region: us-east-1
+
+      - name: Configure aws credentials (pytorch account)
+        if: ${{ needs.get_release_type.outputs.type == 'test' }}
+        uses: aws-actions/configure-aws-credentials@v3
+        with:
+          role-to-assume: arn:aws:iam::749337293305:role/gha_workflow_test_build_wheels
+          aws-region: us-east-1
       - name: Download Artifacts from Github
         continue-on-error: true
         uses: actions/download-artifact@v3
@@ -177,9 +190,6 @@ jobs:
         run: ls -lh torchdata*.whl
       - name: Upload Wheels to S3 Storage
         if: steps.trigger_upload.outputs.value == 'true'
-        env:
-          AWS_ACCESS_KEY_ID: ${{ secrets.PYTORCH_BINARY_AWS_ACCESS_KEY_ID }}
-          AWS_SECRET_ACCESS_KEY: ${{ secrets.PYTORCH_BINARY_AWS_SECRET_ACCESS_KEY }}
         run: |
           if [[ ${{ inputs.branch }} == 'main' ]]; then
             S3_PATH=s3://pytorch/whl/nightly/

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
diff --git a/.github/workflows/test_release.yml b/.github/workflows/test_release.yml
@@ -22,6 +22,10 @@ on:
 # env:
 #   RELEASE_BRANCH: ""
 
+permissions:
+  id-token: write
+  contents: read
+
 jobs:
   build_test_upload:
     if: github.repository == 'pytorch/data' && startsWith(github.ref_name, 'release/')
@@ -31,6 +35,4 @@ jobs:
       pre_dev_release: true
       pytorch_version: "2.4.0"
     secrets:
-      PYTORCH_BINARY_AWS_ACCESS_KEY_ID: ${{ secrets.PYTORCH_BINARY_AWS_ACCESS_KEY_ID }}
-      PYTORCH_BINARY_AWS_SECRET_ACCESS_KEY: ${{ secrets.PYTORCH_BINARY_AWS_SECRET_ACCESS_KEY }}
       CONDA_TEST_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_TEST_PYTORCHBOT_TOKEN }}