Merge pull request #2358 from AI-Hypercomputer:rbierneni-add-pinned-github-commits

PiperOrigin-RevId: 807883660

Author: Google-ML-Automation

maxtext_jax_ai_image.Dockerfile

@@ -55,7 +55,7 @@ RUN ls .
 ARG TEST_TYPE
 # Copy over test assets if building image for end-to-end tests or unit tests
 RUN if [ "$TEST_TYPE" = "xlml" ] || [ "$TEST_TYPE" = "unit_test" ]; then \
-      if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_DIR}"; then \
+      if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}"; then \
         echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."; \
       fi; \
     fi

requirements.txt

@@ -3,15 +3,16 @@ aqtp
 array-record
 cloud-accelerator-diagnostics
 cloud-tpu-diagnostics
+datasets
 flax
 gcsfs
 google-api-python-client
 google-cloud-aiplatform
 google-cloud-monitoring
 grain[parquet]
 huggingface_hub
-jax<0.7.1
-jaxlib<0.7.1
+jax!=0.7.1
+jaxlib!=0.7.1
 jaxtyping
 jsonlines
 ml-collections
@@ -39,5 +40,4 @@ transformers
 tunix @ https://github.com/google/tunix/archive/d770659621eb16ef6588268e26fa687fa068df20.zip
 google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
-# datasets 0.4.0 is not compatible with fsspec==2025.7.0. It will be updated in the next datasets release.
-datasets @ https://github.com/huggingface/datasets/archive/6790e138c00b87a1ddc72184f89e7814cf784360.zip
+qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip

requirements_with_jax_ai_image.txt

@@ -1,13 +1,13 @@
 # Requirements for Building the MaxText Docker Image
 # These requirements are additional to the dependencies present in the JAX AI base image.
-datasets
+datasets @ https://github.com/huggingface/datasets/archive/6790e138c00b87a1ddc72184f89e7814cf784360.zip
 flax>=0.11.0
 google-api-python-client
-google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
+google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
 grain[parquet]>=0.2.12
 jaxtyping
 jsonlines
-mlperf-logging@git+https://github.com/mlperf/logging.git
+mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
 omegaconf
 orbax-checkpoint>=0.11.22
 pathwaysutils>=0.1.1
@@ -23,4 +23,5 @@ tensorflow-datasets
 tensorflow-text>=2.17.0
 tiktoken
 transformers
-tunix@git+https://github.com/google/tunix.git
+tunix @ https://github.com/google/tunix/archive/d770659621eb16ef6588268e26fa687fa068df20.zip
+qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip

setup.sh

@@ -139,7 +139,27 @@ run_name_folder_path=$(pwd)
 
 # Install dependencies from requirements.txt
 cd "$run_name_folder_path" && python3 -m uv pip install --upgrade pip
-python3 -m uv pip install --no-cache-dir -U -r requirements.txt
+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|' \
+      -e 's|^([^ ]*) @ https?://github.com/([^/]*\/[^/]*)/archive/.*\.zip$|\1@git+https://github.com/\2.git|' \
+      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
+    rm requirements.txt.nightly-temp
+else
+    # stable or stable_stack mode: Install with pinned commits
+    echo "Installing requirements.txt with pinned commits."
+    python3 -m uv pip install --no-cache-dir -U -r requirements.txt
+fi
 
 # Install maxtext package
 if [ -f 'pyproject.toml' ]; then