Delete local image

AI-Hypercomputer · Jul 23, 2024 · 825ed6d · 825ed6d
1 parent f6c8df2
commit 825ed6d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
@@ -37,4 +37,4 @@ jobs:
         ls -lah /dev/shm/docker-context
     - name: build jax stable stack image
       run: |
-        bash docker_maxdiffusion_image_upload.sh PROJECT_ID=tpu-prod-env-multipod BASEIMAGE=us-docker.pkg.dev/tpu-prod-env-multipod/jax-ss/tpu:jax0.4.30-rev1 CLOUD_IMAGE_NAME=maxdiffusion-jax-ss IMAGE_TAG=jax0.4.30-rev1 MAXDIFFUSION_REQUIREMENTS_FILE=requirements_with_jax_ss.txt BUILD_CONTEXT=/dev/shm/docker-context
+        bash docker_maxdiffusion_image_upload.sh PROJECT_ID=tpu-prod-env-multipod BASEIMAGE=us-docker.pkg.dev/tpu-prod-env-multipod/jax-ss/tpu:jax0.4.30-rev1 CLOUD_IMAGE_NAME=maxdiffusion-jax-ss IMAGE_TAG=jax0.4.30-rev1 MAXDIFFUSION_REQUIREMENTS_FILE=requirements_with_jax_ss.txt BUILD_CONTEXT=/dev/shm/docker-context DELETE_LOCAL_IMAGE=true
diff --git a/docker_maxdiffusion_image_upload.sh b/docker_maxdiffusion_image_upload.sh
@@ -59,11 +59,13 @@ if [[ ! -v MAXDIFFUSION_REQUIREMENTS_FILE ]]; then
   exit 1
 fi
 
-# Optional: BUILD_CONTEXT, default to current directory if not set
+# BUILD_CONTEXT, default to current directory if not set
 BUILD_CONTEXT=${BUILD_CONTEXT:-.}
 
-gcloud auth configure-docker us-docker.pkg.dev --quiet
-gcloud auth configure-docker --quiet
+# Default: Don't delete local image
+DELETE_LOCAL_IMAGE="false"
+
+gcloud auth configure-docker us-docker.pkg.dev gcr.io --quiet
 
 COMMIT_HASH=$(git rev-parse --short HEAD)
 
@@ -78,13 +80,14 @@ docker build \
   --build-arg COMMIT_HASH=${COMMIT_HASH} \
   --build-arg MAXDIFFUSION_REQUIREMENTS_FILE=${MAXDIFFUSION_REQUIREMENTS_FILE} \
   --network=host \
-  -t ${LOCAL_IMAGE_NAME}_runner \
+  -t gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG} \
   -f ${BUILD_CONTEXT}/maxdiffusion_jax_ss_tpu.Dockerfile ${BUILD_CONTEXT}
 
-docker tag \
-  ${LOCAL_IMAGE_NAME}_runner \
-  gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG} \
-
 docker push gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG}
 
-echo "All done, check out your artifacts at: gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG}"
+echo "All done, check out your artifacts at: gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG}"
+
+if [ "$DELETE_LOCAL_IMAGE" == "true" ]; then
+  docker rmi gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG}
+  echo "Local image deleted."
+fi