adding missing changes

AI-Hypercomputer · Jul 24, 2024 · 65cf2c0 · 65cf2c0
1 parent 5c98f4e
commit 65cf2c0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 10 deletions.
diff --git a/docker_maxdiffusion_image_upload.sh b/docker_maxdiffusion_image_upload.sh
@@ -21,7 +21,7 @@
 # (minutes). However, if you are simply changing local code and not updating dependencies, uploading just takes a few seconds.
 
 # Example command:
-# bash docker_maxdiffusion_image_upload.sh PROJECT_ID=tpu-prod-env-multipod BASEIMAGE=gcr.io/tpu-prod-env-multipod/jax-ss/tpu:jax0.4.28-v1.0.0 CLOUD_IMAGE_NAME=maxdiffusion-jax-ss-0.4.28-v1.0.0 IMAGE_TAG=latest MAXDIFFUSION_REQUIREMENTS_FILE=requirements_with_jax_ss.txt
+# bash docker_maxdiffusion_image_upload.sh PROJECT_ID=tpu-prod-env-multipod BASEIMAGE=gcr.io/tpu-prod-env-multipod/jax-ss/tpu:jax0.4.28-v1.0.0 CLOUD_IMAGE_NAME=maxdiffusion-jax-stable-stack IMAGE_TAG=latest MAXDIFFUSION_REQUIREMENTS_FILE=requirements_with_jax_stable_stack.txt
 
 set -e
 
@@ -59,9 +59,6 @@ if [[ ! -v MAXDIFFUSION_REQUIREMENTS_FILE ]]; then
   exit 1
 fi
 
-# BUILD_CONTEXT, default to current directory if not set
-BUILD_CONTEXT=${BUILD_CONTEXT:-.}
-
 # Default: Don't delete local image
 DELETE_LOCAL_IMAGE="${DELETE_LOCAL_IMAGE:-false}"
 
@@ -70,19 +67,19 @@ gcloud auth configure-docker us-docker.pkg.dev --quiet
 
 COMMIT_HASH=$(git rev-parse --short HEAD)
 
-echo "Building JAX SS MaxDiffusion at commit hash ${COMMIT_HASH} . . ."  
+echo "Building JAX Stable Stack MaxDiffusion at commit hash ${COMMIT_HASH} . . ."  
 
 IMAGE_DATE=$(date +%Y-%m-%d)
 
 FULL_IMAGE_TAG=${IMAGE_TAG}-${IMAGE_DATE}
 
 docker build \
-  --build-arg JAX_SS_BASEIMAGE=${BASEIMAGE} \
+  --build-arg JAX_STABLE_STACK_BASEIMAGE=${BASEIMAGE} \
   --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} \
-  -f ${BUILD_CONTEXT}/maxdiffusion_jax_ss_tpu.Dockerfile ${BUILD_CONTEXT}
+  -f maxdiffusion_jax_stable_stack_tpu.Dockerfile .
 
 docker push gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG}
 

diff --git a/maxdiffusion_jax_stable_stack_tpu.Dockerfile b/maxdiffusion_jax_stable_stack_tpu.Dockerfile
@@ -1,7 +1,7 @@
-ARG JAX_SS_BASEIMAGE
+ARG JAX_STABLE_STACK_BASEIMAGE
 
 # JAX Stable Stack Base Image
-From $JAX_SS_BASEIMAGE
+From $JAX_STABLE_STACK_BASEIMAGE
 
 ARG COMMIT_HASH
 
@@ -27,5 +27,5 @@ RUN if [ ! -z "${MAXDIFFUSION_REQUIREMENTS_FILE}" ]; then \
 # Install MaxDiffusion
 RUN pip install .
 
-# Run the script available in JAX-SS base image to generate the manifest file
+# Run the script available in JAX-Stable-Stack base image to generate the manifest file
 RUN bash /generate_manifest.sh PREFIX=maxdiffusion COMMIT_HASH=$COMMIT_HASH