Merge pull request #2589 from SamuelMarks:extra-deps-location

PiperOrigin-RevId: 827963193

Author: Google-ML-Automation

dependencies/dockerfiles/maxtext_db_dependencies.Dockerfile

@@ -40,9 +40,9 @@ ENV MAXTEXT_REPO_ROOT=/deps
 WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
-COPY tools/setup /deps/tools/setup/
-COPY dependencies/requirements/ /deps/dependencies/requirements/
-COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt /deps/dependencies/requirements/
+COPY tools/setup tools/setup/
+COPY dependencies/requirements/ dependencies/requirements/
+COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxtext_extra_deps/
 
 # 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 LIBTPU_GCS_PATH=${ENV_LIBTPU_GCS_PATH} DEVICE=${ENV_DEVICE}"

dependencies/dockerfiles/maxtext_dependencies.Dockerfile

@@ -40,9 +40,9 @@ ENV MAXTEXT_REPO_ROOT=/deps
 WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
-COPY tools/setup /deps/tools/setup/
-COPY dependencies/requirements/ /deps/dependencies/requirements/
-COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt /deps/dependencies/requirements/
+COPY tools/setup tools/setup/
+COPY dependencies/requirements/ dependencies/requirements/
+COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxtext_extra_deps/
 
 # 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 LIBTPU_GCS_PATH=${ENV_LIBTPU_GCS_PATH} DEVICE=${ENV_DEVICE}"

dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile

@@ -42,9 +42,9 @@ ENV MAXTEXT_REPO_ROOT=/deps
 WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
-COPY tools/setup /deps/tools/setup/
-COPY dependencies/requirements/ /deps/dependencies/requirements/
-COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt /deps/dependencies/requirements/
+COPY tools/setup tools/setup/
+COPY dependencies/requirements/ dependencies/requirements/
+COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxtext_extra_deps/
 
 # 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}"

dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile

@@ -16,9 +16,9 @@ ENV MAXTEXT_REPO_ROOT=/deps
 WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
-COPY tools/setup /deps/tools/setup/
-COPY dependencies/requirements/ /deps/dependencies/requirements/
-COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt /deps/dependencies/requirements/
+COPY tools/setup tools/setup/
+COPY dependencies/requirements/ dependencies/requirements/
+COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxtext_extra_deps/
 
 # For JAX AI tpu training images 0.4.37 AND 0.4.35
 # Orbax checkpoint installs the latest version of JAX,