From c76e92b8b1c2cb60f7d4d0fe298178a821cf6d98 Mon Sep 17 00:00:00 2001
From: Jon Bolin <jonbolin@google.com>
Date: Wed, 5 Jun 2024 18:20:26 +0000
Subject: [PATCH] Use run_id instead of sha for docker tag

---
 .github/workflows/UnitTests.yml | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml
index 4c214567a..d80e3b20a 100644
--- a/.github/workflows/UnitTests.yml
+++ b/.github/workflows/UnitTests.yml
@@ -49,10 +49,10 @@ jobs:
         bash docker_build_dependency_image.sh MODE=${{ matrix.device.mode }} DEVICE=${{ matrix.device.type }} 
     - name: Tag the image
       run: |
-        docker tag maxtext_base_image gcr.io/tpu-prod-env-multipod/maxtext_${{ github.sha }}:${{ matrix.device.type }}
+        docker tag maxtext_base_image gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:${{ matrix.device.type }}
     - name: Upload the image
       run: |
-        docker push gcr.io/tpu-prod-env-multipod/maxtext_${{ github.sha }}:${{ matrix.device.type }}
+        docker push gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:${{ matrix.device.type }}
 
   common:
     needs: build_and_upload_image
@@ -80,7 +80,7 @@ jobs:
     name: Common test (${{ matrix.device.name }})
     runs-on: ["self-hosted", "${{ matrix.device.type }}", "${{ matrix.device.name }}"]
     container:
-      image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.sha }}:${{ matrix.device.type }}
+      image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:${{ matrix.device.type }}
       volumes: 
         - /home/runner/actions-runner/_work/maxtext/maxtext:/deps
       env:
@@ -127,7 +127,7 @@ jobs:
     name: "TPU test (${{ matrix.device-type }})"
     runs-on: ["self-hosted", "tpu", "${{ matrix.device-type }}"]
     container:
-      image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.sha }}:tpu
+      image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu
       volumes: 
       - /home/runner/actions-runner/_work/maxtext/maxtext:/deps
       options: "--privileged"
@@ -146,7 +146,7 @@ jobs:
     name: "GPU test (${{ matrix.device-type }}, ${{ matrix.build-mode }})"
     runs-on: ["self-hosted", "gpu", "${{ matrix.device-type }}"]
     container:
-      image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.sha }}:gpu
+      image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:gpu
       volumes: 
       - /home/runner/actions-runner/_work/maxtext/maxtext:/deps
       env:
@@ -167,7 +167,7 @@ jobs:
     runs-on: ["self-hosted"]
     steps:
     - name: Delete GPU image
-      run: gcloud container images delete gcr.io/tpu-prod-env-multipod/maxtext_${{ github.sha }}:gpu --force-delete-tags --quiet
+      run: gcloud container images delete gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:gpu --force-delete-tags --quiet
     - name: Delete TPU image
-      run: gcloud container images delete gcr.io/tpu-prod-env-multipod/maxtext_${{ github.sha }}:tpu --force-delete-tags --quiet
+      run: gcloud container images delete gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu --force-delete-tags --quiet