Merge pull request #2515 from SamuelMarks:install_maxtext_extra_deps__CACHE

PiperOrigin-RevId: 822224535

Author: Google-ML-Automation

base_requirements/requirements.txt

@@ -38,5 +38,5 @@ tensorflow
 tiktoken
 transformers
 qwix
-google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
+google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip

maxtext_db_dependencies.Dockerfile

@@ -41,7 +41,7 @@ WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
 COPY setup.sh ./
-COPY requirements.txt requirements_with_jax_ai_image.txt generated_requirements ./
+COPY requirements.txt requirements_with_jax_ai_image.txt src/install_maxtext_extra_deps/extra_deps_from_github.txt generated_requirements ./
 
 # 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}"

maxtext_dependencies.Dockerfile

@@ -41,7 +41,7 @@ WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
 COPY setup.sh ./
-COPY requirements.txt requirements_with_jax_ai_image.txt generated_requirements ./
+COPY requirements.txt requirements_with_jax_ai_image.txt src/install_maxtext_extra_deps/extra_deps_from_github.txt generated_requirements ./
 
 # 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}"

maxtext_gpu_dependencies.Dockerfile

@@ -43,7 +43,7 @@ WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
 COPY setup.sh ./
-COPY requirements.txt requirements_with_jax_ai_image.txt generated_requirements ./
+COPY requirements.txt requirements_with_jax_ai_image.txt src/install_maxtext_extra_deps/extra_deps_from_github.txt generated_requirements ./
 
 # 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}"

maxtext_jax_ai_image.Dockerfile

@@ -17,7 +17,7 @@ WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
 COPY setup.sh ./
-COPY requirements.txt requirements_with_jax_ai_image.txt requirements_with_jax_stable_stack_0_6_1_pipreqs.txt generated_requirements ./
+COPY requirements.txt requirements_with_jax_ai_image.txt requirements_with_jax_stable_stack_0_6_1_pipreqs.txt src/install_maxtext_extra_deps/extra_deps_from_github.txt generated_requirements ./
 
 
 # For JAX AI tpu training images 0.4.37 AND 0.4.35

requirements.txt

@@ -38,5 +38,5 @@ tensorflow-text
 tensorflow
 tiktoken
 transformers
-google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
+google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip

requirements_with_jax_ai_image.txt

@@ -3,7 +3,7 @@
 datasets @ https://github.com/huggingface/datasets/archive/6790e138c00b87a1ddc72184f89e7814cf784360.zip
 flax>=0.11.0
 google-api-python-client
-google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
+google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
 grain[parquet]>=0.2.12
 jaxtyping
 jsonlines

setup.sh

@@ -137,32 +137,29 @@ if [[ $DEVICE == "tpu" ]]; then
     fi
 fi
 
-# Save the script folder path of maxtext
-run_name_folder_path=$(pwd)
-
-# Install dependencies from requirements.txt
-cd "$run_name_folder_path" && python3 -m uv pip install --upgrade pip
 if [[ "$MODE" == "nightly" ]]; then
     echo "Nightly mode: Installing requirements.txt, stripping commit pins from git+ repos."
-    cp requirements.txt requirements.txt.nightly-temp
+
     # Create a temp file, strip commit pins from git+ repos in requirements.txt
     # Remove/update this section based on the pinned github repo commit in requirements.txt
-    sed -i -E \
-      -e 's|^mlperf-logging @ https?://github.com/mlcommons/logging/archive/.*\.zip$|mlperf-logging@git+https://github.com/mlperf/logging.git|' \
+    sed -E \
       -e 's|^([^ ]*) @ https?://github.com/([^/]*\/[^/]*)/archive/.*\.zip$|\1@git+https://github.com/\2.git|' \
-      requirements.txt.nightly-temp
+      -e '/JetStream/d' \
+      -e '/mlperf-logging/d' \
+      requirements.txt > requirements.txt.nightly-temp
 
     echo "--- Installing modified nightly requirements: ---"
     cat requirements.txt.nightly-temp
     echo "-------------------------------------------------"
 
-    python3 -m uv pip install --no-cache-dir -U -r requirements.txt.nightly-temp
+    python3 -m uv pip install --no-cache-dir -U -r requirements.txt.nightly-temp \
+                                                -r "${MAXTEXT_REPO_ROOT?}"'/extra_deps_from_github.txt'
     rm requirements.txt.nightly-temp
 else
     # stable or stable_stack mode: Install with pinned commits
     echo "Installing tpu-requirements.txt with pinned commits."
     tpu_requirements_txt=
-    for candidate in 'generated_requirements' "${MAXTEXT_REPO_ROOT}"'/generated_requirements' "$PWD"; do
+    for candidate in 'generated_requirements' "${MAXTEXT_REPO_ROOT?}"'/generated_requirements' "$PWD"; do
       if [ -f "$candidate"'/tpu-requirements.txt' ]; then
         tpu_requirements_txt="$candidate"'/tpu-requirements.txt'
         break
@@ -174,7 +171,8 @@ else
       >&2 printf 'Could not find "tpu-requirements.txt", looked in: %s\n' "${searched%?}"
       exit 2
     else
-      python3 -m uv pip install --resolution=lowest -r "$tpu_requirements_txt"
+      python3 -m uv pip install --resolution=lowest -r "$tpu_requirements_txt" \
+                                                    -r "${MAXTEXT_REPO_ROOT?}"'/extra_deps_from_github.txt'
     fi
 fi
 

src/install_maxtext_extra_deps/extra_deps_from_github.txt

@@ -1,2 +1,2 @@
-google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
+google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip