Merge pull request #1595 from AI-Hypercomputer:bvandermoon-xpk-path

PiperOrigin-RevId: 748341979

Author: maxtext authors

maxtext_jax_stable_stack.Dockerfile

@@ -33,7 +33,14 @@ RUN if [ "$DEVICE" = "tpu" ] && ([ "$JAX_STABLE_STACK_BASEIMAGE" = "us-docker.pk
 
 # Install Maxtext requirements with Jax Stable Stack
 RUN apt-get update && apt-get install --yes google-cloud-cli && apt-get install --yes dnsutils
-RUN python3 -m pip install -r /deps/requirements_with_jax_stable_stack.txt
+
+# Install requirements file generated with pipreqs for JSS 0.5.2. 
+# Othewise use general requirements_with_jax_stable_stack.txt
+RUN if [ "$DEVICE" = "tpu" ] && [ "$JAX_STABLE_STACK_BASEIMAGE" = "us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.5.2-rev1" ]; then \
+        python3 -m pip install -r /deps/requirements_with_jax_stable_stack_0_5_2_pipreqs.txt; \
+  else \
+        python3 -m pip install -r /deps/requirements_with_jax_stable_stack.txt; \
+  fi
 
 # Run the script available in JAX Stable Stack base image to generate the manifest file
 RUN bash /jax-stable-stack/generate_manifest.sh PREFIX=maxtext COMMIT_HASH=$COMMIT_HASH

requirements_with_jax_stable_stack_0_5_2_pipreqs.txt

@@ -0,0 +1,55 @@
+absl_py==2.1.0
+# Replacing with aqtp due to conflicts during build
+# aqt==25.2
+aqtp==0.8.2 # Added manually
+cloud_accelerator_diagnostics==0.1.1
+cloud_tpu_diagnostics==0.1.5
+datasets==3.5.0
+etils==1.12.2
+evaluate==0.4.3
+flax==0.10.4
+# Replacing with specific git commit due to conflicts during build
+# google_jetstream==0.3.0
+# Adding manually
+google-jetstream @ git+https://github.com/AI-Hypercomputer/JetStream.git@082c0ac526e50d8f732a083ed43920590d7ffd22
+grain_nightly==0.0.10
+jax==0.5.2
+jaxlib==0.5.1 # Manually adding to ensure consistency in future
+jaxtyping==0.3.1
+jsonlines==4.0.0
+libtpu==0.0.10.1 # Manually adding to ensure consistency in future
+ml_collections==1.0.0
+ml_goodput_measurement==0.0.8
+nltk==3.9.1
+# Removing due to conflicts during build
+# numpy==2.2.4
+omegaconf==2.3.0
+optax==0.2.4
+orbax==0.1.9
+pandas==2.2.3
+pathwaysutils==0.1.0
+# Removing due to conflicts during build
+# protobuf==3.20.3
+protobuf
+psutil==7.0.0
+pytest==8.3.5
+PyYAML==6.0.2
+PyYAML==6.0.2
+Requests==2.32.3
+safetensors==0.5.3
+sentencepiece==0.1.97
+tensorboard_plugin_profile==2.17.0
+tensorboardX==2.6.2.2
+tensorboardX==2.6.2.2
+tensorflow==2.19.0
+tensorflow_datasets==4.9.8
+tensorflow_text==2.19.0
+tensorstore==0.1.72
+tfds_nightly==4.9.2.dev202308090034
+tiktoken==0.9.0
+torch==2.6.0
+tqdm==4.67.1
+transformer_engine==2.1.0
+transformers==4.51.3
+trl==0.16.1
+urllib3==2.4.0