Merge pull request #1132 from AI-Hypercomputer:small_gpu_build_optimm…

…izatioin PiperOrigin-RevId: 710749810
AI-Hypercomputer · Dec 30, 2024 · 21cb8ec · 21cb8ec
2 parents 5e41c61 + e954f33
commit 21cb8ec
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/maxtext_gpu_dependencies.Dockerfile b/maxtext_gpu_dependencies.Dockerfile
@@ -38,12 +38,13 @@ 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 constraints_gpu.txt requirements.txt requirements_with_jax_stable_stack.txt ./
 
+# Install dependencies - these steps are cached unless the copied files change
 RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"
 RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
 
-
-WORKDIR /deps
+# Now copy the remaining code (source files that may change frequently)
+COPY . .