docker build process improvement

AI-Hypercomputer · Jul 23, 2024 · f5532c9 · f5532c9
1 parent dba217f
commit f5532c9
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/docker_maxdiffusion_image_upload.sh b/docker_maxdiffusion_image_upload.sh
@@ -25,6 +25,8 @@
 
 set -e
 
+export LOCAL_IMAGE_NAME=maxdiffusion_base_image
+
 # Set environment variables
 for ARGUMENT in "$@"; do
     IFS='=' read -r KEY VALUE <<< "$ARGUMENT"
@@ -75,9 +77,13 @@ docker build \
   --build-arg COMMIT_HASH=${COMMIT_HASH} \
   --build-arg MAXDIFFUSION_REQUIREMENTS_FILE=${MAXDIFFUSION_REQUIREMENTS_FILE} \
   --network=host \
-  -t gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG} \
+  -t ${LOCAL_IMAGE_NAME}_runner \
   -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}"