From e954f33cc78f096fcbef91b10a241ed8573d2131 Mon Sep 17 00:00:00 2001
From: Alex Shraer <shralex@gmail.com>
Date: Mon, 30 Dec 2024 03:59:50 +0000
Subject: [PATCH] Improve build caching for GPU build.

---
 maxtext_gpu_dependencies.Dockerfile | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/maxtext_gpu_dependencies.Dockerfile b/maxtext_gpu_dependencies.Dockerfile
index ef718a44..b09cf520 100644
--- 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 . .