Enable linux_job_v2.yml workflow for ROCm

pytorch · Jan 12, 2025 · 3007720 · 3007720
1 parent 46cd41c
commit 3007720
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 6 deletions.
diff --git a/.github/workflows/linux_job.yml b/.github/workflows/linux_job.yml
@@ -201,7 +201,7 @@ jobs:
         if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' }}
 
       - name: Download artifacts (if any)
-        uses: actions/download-artifact@v4
+        uses: actions/download-artifact@v3
         if: ${{ inputs.download-artifact != '' }}
         with:
           name: ${{ inputs.download-artifact }}
@@ -323,7 +323,7 @@ jobs:
       # NB: Keep this for compatibility with existing jobs and also keep in mind that only
       # our AWS runners have access to S3
       - name: Upload artifacts to GitHub (if any)
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v3
         if: ${{ always() && inputs.upload-artifact != '' && !inputs.upload-artifact-to-s3 }}
         with:
           name: ${{ inputs.upload-artifact }}

diff --git a/.github/workflows/linux_job_v2.yml b/.github/workflows/linux_job_v2.yml
@@ -104,17 +104,24 @@ on:
         required: false
         default: false
         type: boolean
+permissions:
+  id-token: write
+  contents: read
+
 jobs:
   job:
     strategy:
       fail-fast: false
     name: ${{ inputs.job-name }}
     env:
       DOCKER_IMAGE: >-
-        ${{ inputs.docker-image == 'pytorch/almalinux-builder' && format('pytorch/almalinux-builder:{0}{1}',
+        ${{ inputs.gpu-arch-type == 'rocm' && format('pytorch/manylinux2_28-builder:{0}{1}',
+                                                                      inputs.gpu-arch-type,
+                                                                      inputs.gpu-arch-version)
+                                                            || inputs.docker-image == 'pytorch/almalinux-builder' && format('pytorch/almalinux-builder:{0}{1}',
                                                                       inputs.gpu-arch-type,
                                                                       inputs.gpu-arch-version)
-                                                           || inputs.docker-image }}
+                                                            || inputs.docker-image }}
       REPOSITORY: ${{ inputs.repository || github.repository }}
       # Will be blank outside of this
       PR_NUMBER: ${{ github.event.pull_request.number }}
@@ -157,6 +164,11 @@ jobs:
 
       - name: Setup Linux
         uses: ./test-infra/.github/actions/setup-linux
+        if: ${{ inputs.gpu-arch-type != 'rocm' }}
+
+      - name: Setup ROCM 
+        uses: pytorch/pytorch/.github/actions/setup-rocm@main 
+        if: ${{ inputs.gpu-arch-type == 'rocm' }}
 
       - name: Setup SSH
         uses: ./test-infra/.github/actions/setup-ssh
@@ -266,15 +278,30 @@ jobs:
           path: ${{ env.RUNNER_TEST_RESULTS_DIR }}
           fail-on-empty: false
 
+      - name: configure aws credentials
+        id: aws_creds
+        if: ${{ inputs.gpu-arch-type == 'rocm' }}
+        uses: aws-actions/configure-aws-credentials@v4
+        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
+        if: ${{ inputs.gpu-arch-type == 'rocm' }}
+        continue-on-error: true
+        uses: aws-actions/amazon-ecr-login@v2
+
       - name: Chown repository directory
-        if: always()
+        if: ${{ inputs.gpu-arch-type != 'rocm' }}
         uses: ./test-infra/.github/actions/chown-directory
         with:
           directory: ${{ github.workspace }}/${{ env.repository }}
           ALPINE_IMAGE: ${{ startsWith(inputs.runner, 'linux.arm64') && 'arm64v8/alpine' || '308535385114.dkr.ecr.us-east-1.amazonaws.com/tool/alpine' }}
 
       - name: Chown runner temp
-        if: always()
+        if: ${{ inputs.gpu-arch-type != 'rocm' }}
         uses: ./test-infra/.github/actions/chown-directory
         with:
           directory: ${{ runner.temp }}