[Release only] Pin rocm docker images (pytorch#126452)

* Pin rocm docker images * test * test * test * test
ROCm · May 17, 2024 · ee68b41 · ee68b41
1 parent 0365423
commit ee68b41
Show file tree

Hide file tree

Showing 7 changed files with 29 additions and 4 deletions.
diff --git a/.github/workflows/_linux-build.yml b/.github/workflows/_linux-build.yml
@@ -11,6 +11,11 @@ on:
         required: true
         type: string
         description: Name of the base docker image to build with.
+      docker-image-tag:
+        required: false
+        type: string
+        description: Name of the base docker image tag
+        default: ""
       build-generates-artifacts:
         required: false
         type: boolean
@@ -69,7 +74,7 @@ jobs:
     runs-on: ${{ inputs.runner }}
     timeout-minutes: 240
     outputs:
-      docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
+      docker-image: ${{ steps.calculate-docker.outputs.docker-image }}
       test-matrix: ${{ steps.filter.outputs.test-matrix }}
     steps:
       - name: Setup SSH (Click me for login details)
@@ -93,10 +98,24 @@ jobs:
         with:
           docker-image-name: ${{ inputs.docker-image-name }}
 
+      - name: Override docker image tag if pinned
+        id: calculate-docker
+        env:
+          ECR_DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
+        shell: bash
+        run: |
+          export NEW_TAG=${{ inputs.docker-image-tag }}
+          if [[ ${NEW_TAG} != '' ]]; then
+            IMAGE=${ECR_DOCKER_IMAGE%:*}
+            echo "docker-image=${IMAGE}:${NEW_TAG}" >> "${GITHUB_OUTPUT}"
+          else
+            echo "docker-image=${ECR_DOCKER_IMAGE}" >> "${GITHUB_OUTPUT}"
+          fi
+
       - name: Use following to pull public copy of the image
         id: print-ghcr-mirror
         env:
-          ECR_DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
+          ECR_DOCKER_IMAGE: ${{ steps.calculate-docker.outputs.docker-image }}
         shell: bash
         run: |
           tag=${ECR_DOCKER_IMAGE##*/}
@@ -105,7 +124,7 @@ jobs:
       - name: Pull docker image
         uses: pytorch/test-infra/.github/actions/pull-docker-image@release/2.3
         with:
-          docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
+          docker-image: ${{ steps.calculate-docker.outputs.docker-image }}
 
       - name: Parse ref
         id: parse-ref
@@ -149,7 +168,7 @@ jobs:
           XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla
           PR_LABELS: ${{ toJson(github.event.pull_request.labels.*.name) }}
           TORCH_CUDA_ARCH_LIST: ${{ inputs.cuda-arch-list }}
-          DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
+          DOCKER_IMAGE: ${{ steps.calculate-docker.outputs.docker-image }}
           XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }}
           DEBUG: ${{ inputs.build-with-debug && '1' || '0' }}
           OUR_GITHUB_JOB_ID: ${{ steps.get-job-id.outputs.job-id }}

diff --git a/.github/workflows/inductor.yml b/.github/workflows/inductor.yml
@@ -22,6 +22,7 @@ jobs:
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3
+      docker-image-tag: cea4be730564c18dd285a12828c7c449490b10b9
       test-matrix: |
         { include: [
           { config: "inductor", shard: 1, num_shards: 1, runner: "linux.rocm.gpu.2" },

diff --git a/.github/workflows/periodic.yml b/.github/workflows/periodic.yml
@@ -215,6 +215,7 @@ jobs:
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3
+      docker-image-tag: cea4be730564c18dd285a12828c7c449490b10b9
       test-matrix: |
         { include: [
           { config: "distributed", shard: 1, num_shards: 2, runner: "linux.rocm.gpu" },

diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml
@@ -409,6 +409,7 @@ jobs:
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3
+      docker-image-tag: cea4be730564c18dd285a12828c7c449490b10b9
       sync-tag: rocm-build
       test-matrix: |
         { include: [

diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml
@@ -31,6 +31,7 @@ jobs:
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3
+      docker-image-tag: cea4be730564c18dd285a12828c7c449490b10b9
       sync-tag: rocm-build
       test-matrix: |
         { include: [

diff --git a/.github/workflows/slow.yml b/.github/workflows/slow.yml
@@ -107,6 +107,7 @@ jobs:
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3
+      docker-image-tag: cea4be730564c18dd285a12828c7c449490b10b9
       test-matrix: |
         { include: [
           { config: "slow", shard: 1, num_shards: 1, runner: "linux.rocm.gpu" },

diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml
@@ -196,6 +196,7 @@ jobs:
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3
+      docker-image-tag: cea4be730564c18dd285a12828c7c449490b10b9
       sync-tag: rocm-build
       test-matrix: |
         { include: [