From 77624ebee19954fefab4ecc06a89114dded2c512 Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Wed, 24 Jul 2024 22:54:56 +0000
Subject: [PATCH] new runner

---
 .github/workflows/UploadDockerImages.yml | 9 ++++-----
 docker_maxdiffusion_image_upload.sh      | 2 --
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
index f9f8e6c0..dbf53261 100644
--- a/.github/workflows/UploadDockerImages.yml
+++ b/.github/workflows/UploadDockerImages.yml
@@ -22,11 +22,10 @@ on:
     branches: [ "pb-jax-ss" ]
 
 jobs:
-  Build-Image-On-Ubuntu:
-    name: Build Maxdiffusion JAX Stable Stack Image
-    runs-on: ubuntu-latest
+  build-image:
+    runs-on: ["self-hosted", "maxdiffusion-github-docker-image-build-runner"]
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v3
     - name: build jax stable stack image
       run: |
-        bash docker_maxdiffusion_image_upload.sh PROJECT_ID=tpu-prod-env-multipod BASEIMAGE=gcr.io/tpu-prod-env-multipod/jax-stable-stack/tpu:jax0.4.30-rev1 CLOUD_IMAGE_NAME=maxdiffusion-jax-stable-stack IMAGE_TAG=jax0.4.30-rev1 MAXDIFFUSION_REQUIREMENTS_FILE=requirements_with_jax_stable_stack.txt DELETE_LOCAL_IMAGE=true
\ No newline at end of file
+        bash docker_maxdiffusion_image_upload.sh PROJECT_ID=tpu-prod-env-multipod BASEIMAGE=us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/tpu:jax0.4.30-rev1 CLOUD_IMAGE_NAME=maxdiffusion-jax-stable-stack IMAGE_TAG=jax0.4.30-rev1 MAXDIFFUSION_REQUIREMENTS_FILE=requirements_with_jax_stable_stack.txt DELETE_LOCAL_IMAGE=true
\ No newline at end of file
diff --git a/docker_maxdiffusion_image_upload.sh b/docker_maxdiffusion_image_upload.sh
index 1355dcc8..7b358cc7 100644
--- a/docker_maxdiffusion_image_upload.sh
+++ b/docker_maxdiffusion_image_upload.sh
@@ -62,8 +62,6 @@ fi
 # Default: Don't delete local image
 DELETE_LOCAL_IMAGE="${DELETE_LOCAL_IMAGE:-false}"
 
-gcloud auth configure-docker us-docker.pkg.dev --quiet
-
 COMMIT_HASH=$(git rev-parse --short HEAD)
 
 echo "Building JAX Stable Stack MaxDiffusion at commit hash ${COMMIT_HASH} . . ."