Added caching optimization in building Maxtext + JAII Image

Author: Rohan-Bierneni

docker_build_dependency_image.sh

@@ -75,7 +75,7 @@ build_ai_image() {
     COMMIT_HASH=$(git rev-parse --short HEAD)
     echo "Building JAX AI MaxText Imageat commit hash ${COMMIT_HASH}..."
 
-    docker build --no-cache \
+    docker build \
         --build-arg JAX_AI_IMAGE_BASEIMAGE=${BASEIMAGE} \
         --build-arg COMMIT_HASH=${COMMIT_HASH} \
         --build-arg DEVICE="$DEVICE" \

maxtext_jax_ai_image.Dockerfile

@@ -13,9 +13,9 @@ RUN mkdir -p /deps
 # Set the working directory in the container
 WORKDIR /deps
 
-# Copy all files from local workspace into docker container
-COPY . .
-RUN ls .
+# Copy setup files and dependency files separately for better caching
+COPY setup.sh ./
+COPY requirements.txt requirements_with_jax_ai_image.txt ./
 
 
 # For JAX AI tpu training images 0.4.37 AND 0.4.35
@@ -37,5 +37,9 @@ RUN apt-get update && apt-get install --yes && apt-get install --yes dnsutils
 RUN pip install google-cloud-monitoring
 RUN python3 -m pip install -r /deps/requirements_with_jax_ai_image.txt
 
+# Now copy the remaining code (source files that may change frequently)
+COPY . .
+RUN ls .
+
 # Run the script available in JAX AI base image to generate the manifest file
 RUN bash /jax-stable-stack/generate_manifest.sh PREFIX=maxtext COMMIT_HASH=$COMMIT_HASH