Move aws OIDC credentials steps into setup-rocm.yml (#164769)

The AWS ECR login step needs `id-token: write` permissions. We move the steps to get OIDC-based credentials from `_rocm-test.yml` to `setup-rocm.yml`. This lays the groundwork to enable access to AWS ECR in workflows in other repos such as torchtitan that use [linux_job_v2.yml](https://github.com/pytorch/test-infra/blob/main/.github/workflows/linux_job_v2.yml), which also uses [setup-rocm.yml](335f4f80a0/.github/workflows/linux_job_v2.yml (L168)). Any caller workflows that eventually execute `setup-rocm` action will thus need to provide the `id-token: write` permission. Pull Request resolved: https://github.com/pytorch/pytorch/pull/164769 Approved by: https://github.com/huydhn
2025-10-20 21:14:14 +08:00 · 2025-10-10 21:24:29 +00:00
parent 8360f34c36
commit 0ec0120b19
5 changed files with 64 additions and 13 deletions
--- a/.github/actions/setup-rocm/action.yml
+++ b/.github/actions/setup-rocm/action.yml
@ -111,3 +111,16 @@ runs:
        # This video group ID maps to subgid 1 inside the docker image due to the /etc/subgid entries.
        # The group name corresponding to group ID 1 can change depending on the OS, so both are necessary.
        echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd $DEVICE_FLAG --group-add video --group-add $render_gid --group-add daemon --group-add bin --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --network=host" >> "${GITHUB_ENV}"
+
+    - name: configure aws credentials
+      id: aws_creds
+      uses: aws-actions/configure-aws-credentials@ececac1a45f3b08a01d2dd070d28d111c5fe6722 # v4.1.0
+      with:
+        role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
+        aws-region: us-east-1
+        role-duration-seconds: 18000
+
+    - name: Login to Amazon ECR
+      id: login-ecr
+      continue-on-error: true
+      uses: aws-actions/amazon-ecr-login@062b18b96a7aff071d4dc91bc00c4c1a7945b076 # v2.0.1
--- a/.github/templates/linux_binary_build_workflow.yml.j2
+++ b/.github/templates/linux_binary_build_workflow.yml.j2
@ -177,6 +177,9 @@ jobs:
    runs-on: linux.rocm.gpu.mi250
    timeout-minutes: !{{ common.timeout_minutes }}
    !{{ upload.binary_env(config) }}
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
--- a/.github/workflows/_rocm-test.yml
+++ b/.github/workflows/_rocm-test.yml
@ -102,19 +102,6 @@ jobs:
            exit 1
          fi

-      - name: configure aws credentials
-        id: aws_creds
-        uses: aws-actions/configure-aws-credentials@ececac1a45f3b08a01d2dd070d28d111c5fe6722 # v4.1.0
-        with:
-          role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
-          aws-region: us-east-1
-          role-duration-seconds: 18000
-
-      - name: Login to Amazon ECR
-        id: login-ecr
-        continue-on-error: true
-        uses: aws-actions/amazon-ecr-login@062b18b96a7aff071d4dc91bc00c4c1a7945b076 # v2.0.1
-
      - name: Calculate docker image
        id: calculate-docker-image
        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
--- a/.github/workflows/generated-linux-binary-libtorch-nightly.yml
+++ b/.github/workflows/generated-linux-binary-libtorch-nightly.yml
@ -358,6 +358,9 @@ jobs:
      DOCKER_IMAGE_TAG_PREFIX: rocm6.4
      LIBTORCH_CONFIG: release
      LIBTORCH_VARIANT: shared-with-deps
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -473,6 +476,9 @@ jobs:
      DOCKER_IMAGE_TAG_PREFIX: rocm7.0
      LIBTORCH_CONFIG: release
      LIBTORCH_VARIANT: shared-with-deps
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
--- a/.github/workflows/generated-linux-binary-manywheel-nightly.yml
+++ b/.github/workflows/generated-linux-binary-manywheel-nightly.yml
@ -347,6 +347,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm6.4
      DESIRED_PYTHON: "3.10"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -459,6 +462,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm7.0
      DESIRED_PYTHON: "3.10"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -941,6 +947,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm6.4
      DESIRED_PYTHON: "3.11"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -1053,6 +1062,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm7.0
      DESIRED_PYTHON: "3.11"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -1535,6 +1547,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm6.4
      DESIRED_PYTHON: "3.12"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -1647,6 +1662,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm7.0
      DESIRED_PYTHON: "3.12"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -2129,6 +2147,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm6.4
      DESIRED_PYTHON: "3.13"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -2241,6 +2262,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm7.0
      DESIRED_PYTHON: "3.13"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -2723,6 +2747,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm6.4
      DESIRED_PYTHON: "3.13t"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -2835,6 +2862,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm7.0
      DESIRED_PYTHON: "3.13t"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -3317,6 +3347,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm6.4
      DESIRED_PYTHON: "3.14"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -3429,6 +3462,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm7.0
      DESIRED_PYTHON: "3.14"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -3911,6 +3947,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm6.4
      DESIRED_PYTHON: "3.14t"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm
@ -4023,6 +4062,9 @@ jobs:
      DOCKER_IMAGE: manylinux2_28-builder
      DOCKER_IMAGE_TAG_PREFIX: rocm7.0
      DESIRED_PYTHON: "3.14t"
+    permissions:
+      id-token: write
+      contents: read
    steps:
      - name: Setup ROCm
        uses: ./.github/actions/setup-rocm