Update generated requirements for JAX 0.9.2

Author: SurbhiJainUSC

.github/workflows/build_and_test_maxtext.yml

@@ -37,8 +37,8 @@ concurrency:
 permissions:
   contents: read
 jobs:
-  doc_only_check:
-    name: Check for Documentation-Only Changes
+  analyze_changes:
+    name: Analyze Changes for Test Orchestration
     runs-on: ubuntu-latest
     outputs:
       run_tests: ${{ steps.check.outputs.run_tests }}
@@ -47,7 +47,7 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      - name: Check if only documentation changed
+      - name: Analyze Changes for Test Orchestration
         id: check
         run: |
           if [ "${{ github.event_name }}" != "pull_request" ]; then
@@ -70,6 +70,14 @@ jobs:
             exit 0
           fi
 
+          # Check if dependencies are changed
+          if echo "$CHANGED_FILES" | grep -E '(^|/)src/dependencies/' > /dev/null; then
+            echo "MaxText dependencies changed, enabling all tests and notebooks."
+            echo "run_tests=true" >> $GITHUB_OUTPUT
+            echo "run_notebooks=true" >> $GITHUB_OUTPUT
+            exit 0
+          fi
+
           # Check for source code changes (anything not .md and not .ipynb)
           if echo "$CHANGED_FILES" | grep -v -E '\.(md|ipynb)$' > /dev/null; then
             echo "Source code files changed, enabling unit tests."
@@ -91,11 +99,11 @@ jobs:
           exit 0
 
   build_and_upload_maxtext_package:
-    needs: doc_only_check
+    needs: analyze_changes
     # Run if either tests or notebooks need to run
     if: |
-      needs.doc_only_check.outputs.run_tests == 'true' || 
-      needs.doc_only_check.outputs.run_notebooks == 'true'
+      needs.analyze_changes.outputs.run_tests == 'true' || 
+      needs.analyze_changes.outputs.run_notebooks == 'true'
     uses: ./.github/workflows/build_package.yml
     with:
       device_type: tpu
@@ -104,7 +112,7 @@ jobs:
 
   maxtext_jupyter_notebooks:
     needs: build_and_upload_maxtext_package
-    if: needs.doc_only_check.outputs.run_notebooks == 'true'
+    if: needs.analyze_changes.outputs.run_notebooks == 'true'
     uses: ./.github/workflows/run_jupyter_notebooks.yml
     strategy:
         fail-fast: false
@@ -120,7 +128,7 @@ jobs:
   tpu-tests:
     name: ${{ matrix.flavor }} tests
     needs: [build_and_upload_maxtext_package]
-    if: needs.doc_only_check.outputs.run_tests == 'true'
+    if: needs.analyze_changes.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_tests_coordinator.yml
     strategy:
       fail-fast: false
@@ -135,7 +143,7 @@ jobs:
   gpu-tests:
     name: ${{ matrix.flavor }} tests
     needs: [build_and_upload_maxtext_package]
-    if: needs.doc_only_check.outputs.run_tests == 'true'
+    if: needs.analyze_changes.outputs.run_tests == 'true'
     strategy:
       fail-fast: false
       matrix:
@@ -150,7 +158,7 @@ jobs:
   cpu-tests:
     name: ${{ matrix.flavor }} tests
     needs: [build_and_upload_maxtext_package]
-    if: needs.doc_only_check.outputs.run_tests == 'true'
+    if: needs.analyze_changes.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_tests_coordinator.yml
     strategy:
       fail-fast: false
@@ -164,7 +172,7 @@ jobs:
 
   maxtext_tpu_pathways_unit_tests:
     needs: build_and_upload_maxtext_package
-    if: needs.doc_only_check.outputs.run_tests == 'true'
+    if: needs.analyze_changes.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
@@ -182,7 +190,7 @@ jobs:
 
   maxtext_tpu_pathways_integration_tests:
     needs: build_and_upload_maxtext_package
-    if: needs.doc_only_check.outputs.run_tests == 'true'
+    if: needs.analyze_changes.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
@@ -206,9 +214,8 @@ jobs:
     steps:
       - name: Check test results
         run: |
-          # If doc-only, all tests should be skipped
-          if [ "${NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_TESTS}" == "false" ]; then
-            echo "Documentation-only changes detected, tests were skipped"
+          if [ "${NEEDS_ANALYZE_CHANGES_OUTPUTS_RUN_TESTS}" == "false" ]; then
+            echo "Tests were skipped"
             exit 0
           fi
 
@@ -228,7 +235,7 @@ jobs:
 
           echo "All required tests passed successfully"
         env:
-          NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_TESTS: ${{ needs.doc_only_check.outputs.run_tests }}
+          NEEDS_ANALYZE_CHANGES_OUTPUTS_RUN_TESTS: ${{ needs.analyze_changes.outputs.run_tests }}
           NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }}
           NEEDS_CPU_TESTS_RESULT: ${{ needs.cpu-tests.result }}
           NEEDS_TPU_TESTS_RESULT: ${{ needs.tpu-tests.result }}
@@ -238,14 +245,14 @@ jobs:
 
   all_notebooks_passed:
     name: All Notebooks Passed
-    needs: [doc_only_check, build_and_upload_maxtext_package, maxtext_jupyter_notebooks]
+    needs: [analyze_changes, build_and_upload_maxtext_package, maxtext_jupyter_notebooks]
     if: always()
     runs-on: ubuntu-latest
     steps:
       - name: Check notebooks results
         run: |
-          if [ "${NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_NOTEBOOKS}" == "false" ]; then
-            echo "Non-notebook changes detected, runs were skipped"
+          if [ "${NEEDS_ANALYZE_CHANGES_OUTPUTS_RUN_NOTEBOOKS}" == "false" ]; then
+            echo "Notebooks were skipped"
             exit 0
           fi
 
@@ -261,7 +268,7 @@ jobs:
 
           echo "All required notebooks passed successfully"
         env:
-          NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_NOTEBOOKS: ${{ needs.doc_only_check.outputs.run_notebooks }}
+          NEEDS_ANALYZE_CHANGES_OUTPUTS_RUN_NOTEBOOKS: ${{ needs.analyze_changes.outputs.run_notebooks }}
           NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }}
           NEEDS_MAXTEXT_JUPYTER_NOTEBOOKS_RESULT: ${{ needs.maxtext_jupyter_notebooks.result }}
 

.github/workflows/run_tests_against_package.yml

@@ -152,6 +152,14 @@ jobs:
           # omit this libtpu init args for gpu tests
           if [ "${INPUTS_DEVICE_TYPE}" != "cuda12" ]; then
             export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
+          else
+            # For cuda12, explicitly point to the pip-installed CUDA libraries
+            # to avoid conflicts with system-level installations on the runner.
+            if [ -d ".venv/lib/python3.12/site-packages/nvidia" ]; then
+              export LD_LIBRARY_PATH=$(pwd)/.venv/lib/python3.12/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH}
+            else
+              echo "Warning: Could not find pinned nvidia libraries in .venv."
+            fi
           fi
           if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then
             $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist

docs/install_maxtext.md

@@ -118,11 +118,10 @@ Please keep dependencies updated throughout development. This will allow each co
 
 To update dependencies, you will follow these general steps:
 
-1. **Modify Base Requirements**: Update the desired dependencies in `base_requirements/requirements.txt` or the hardware-specific files (`base_requirements/tpu-base-requirements.txt`, `base_requirements/gpu-base-requirements.txt`).
+1. **Modify Base Requirements**: Update the desired dependencies in `src/dependencies/requirements/base_requirements/requirements.txt` or the hardware-specific pre-training files (`base_requirements/tpu-requirements.txt`, `base_requirements/gpu-requirements.txt`) or post-training requirements.
 2. **Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
 3. **Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
-4. **Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/dependencies/github_deps/pre_train_deps.txt`.
-5. **Verify**: Test the new dependencies to ensure the project installs and runs correctly.
+4. **Verify**: Test the new dependencies to ensure the project installs and runs correctly.
 
 The following sections provide detailed instructions for each step.
 
@@ -133,59 +132,70 @@ First, you need to install the `seed-env` command-line tool by running `pip inst
 
 ## Step 2: Find the JAX Build Commit Hash
 
-The dependency generation process is pinned to a specific nightly build of JAX. You need to find the commit hash for the desired JAX build.
-
-You can find the latest commit hashes in the [JAX `build/` folder](https://github.com/jax-ml/jax/commits/main/build). Choose a recent, successful build and copy its full commit hash.
+The dependency generation process is pinned to a specific nightly build of JAX. You need to find the commit hash for the desired JAX build from [JAX `build/` folder](https://github.com/jax-ml/jax/commits/main/build).
 
 ## Step 3: Generate the Requirements Files
 
 Next, run the `seed-env` CLI to generate the new requirements files. You will need to do this separately for the TPU and GPU environments. The generated files will be placed in a directory specified by `--output-dir`.
 
-### For TPU
+> **Note:** The current `src/dependencies/requirements/generated_requirements/` in the repository were generated using JAX build commit hash: [e0d2967b50abbefd651d563dbcd7afbcb963d08c](https://github.com/jax-ml/jax/commit/e0d2967b50abbefd651d563dbcd7afbcb963d08c).
+
+### TPU Pre-Training
 
-Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step.
+If you have made changes to TPU pre-training dependencies in `src/dependencies/requirements/base_requirements/tpu-requirements.txt`, you need to regenerate the pinned pre-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
 ```bash
 seed-env \
-  --local-requirements=src/dependencies/requirements/base_requirements/tpu-base-requirements.txt \
+  --local-requirements=src/dependencies/requirements/base_requirements/tpu-requirements.txt \
   --host-name=MaxText \
   --seed-commit=<jax-build-commit-hash> \
   --python-version=3.12 \
   --requirements-txt=tpu-requirements.txt \
   --output-dir=generated_tpu_artifacts
+
+# Copy generated requirements to src/dependencies/requirements/generated_requirements
+mv generated_tpu_artifacts/tpu-requirements.txt src/dependencies/requirements/generated_requirements/tpu-requirements.txt
 ```
 
-### For GPU
+#### TPU Post-Training
 
-Similarly, run the command for the GPU requirements.
+If you have made changes to the post-training dependencies in `src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt`, you need to regenerate the pinned post-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
 ```bash
 seed-env \
-  --local-requirements=src/dependencies/requirements/base_requirements/cuda12-base-requirements.txt \
+  --local-requirements=src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt \
   --host-name=MaxText \
   --seed-commit=<jax-build-commit-hash> \
   --python-version=3.12 \
-  --requirements-txt=cuda12-requirements.txt \
-  --hardware=cuda12 \
-  --output-dir=generated_gpu_artifacts
-```
+  --requirements-txt=tpu-post-train-requirements.txt \
+  --output-dir=generated_tpu_post_train_artifacts
 
-## Step 4: Update Project Files
+# Copy generated requirements to src/dependencies/requirements/generated_requirements
+mv generated_tpu_post_train_artifacts/tpu-post-train-requirements.txt src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt
+```
 
-After generating the new requirements, you need to update the files in the MaxText repository.
+### GPU Pre-Training
 
-1. **Copy the generated files:**
+If you have made changes to the GPU pre-training dependencies in `src/dependencies/requirements/base_requirements/gpu-requirements.txt`, you need to regenerate the pinned pre-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
-   - Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
-   - Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
+```bash
+seed-env \
+  --local-requirements=src/dependencies/requirements/base_requirements/cuda12-requirements.txt \
+  --host-name=MaxText \
+  --seed-commit=<jax-build-commit-hash> \
+  --python-version=3.12 \
+  --requirements-txt=cuda12-requirements.txt \
+  --hardware=cuda12 \
+  --output-dir=generated_gpu_artifacts
 
-2. **Update `pre_train_deps.txt` (if necessary):**
-   Currently, MaxText uses a few dependencies, such as `mlperf-logging` and `google-jetstream`, that are installed directly from GitHub source. These are defined in `base_requirements/requirements.txt`, and the `seed-env` tool will carry them over to the generated requirements files.
+# Copy generated requirements to src/dependencies/requirements/generated_requirements
+mv generated_gpu_artifacts/cuda12-requirements.txt.txt src/dependencies/requirements/generated_requirements/cuda12-requirements.txt.txt
+```
 
-## Step 5: Verify the New Dependencies
+## Step 4: Verify the New Dependencies
 
 Finally, test that the new dependencies install correctly and that MaxText runs as expected.
 
-1. **Install MaxText and dependencies**: For instructions on installing MaxText on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/maxtext-v0.2.0/install_maxtext.html#from-source).
+1. **Install MaxText and dependencies**: For instructions on installing MaxText on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source).
 
 2. **Verify the installation**: Run MaxText tests to ensure everything is working as expected with the newly installed dependencies and there are no regressions.

src/dependencies/extra_deps/post_train_overrides.txt

@@ -0,0 +1 @@
+google-metrax>=0.2.3

…e_requirements/gpu-base-requirements.txt …ase_requirements/cuda12-requirements.txtsrc/dependencies/requirements/base_requirements/gpu-base-requirements.txt renamed to src/dependencies/requirements/base_requirements/cuda12-requirements.txt src/dependencies/requirements/base_requirements/gpu-base-requirements.txt renamed to src/dependencies/requirements/base_requirements/cuda12-requirements.txt

@@ -1,8 +1,9 @@
 absl-py
 aqtp
 array-record
+chex
 cloud-accelerator-diagnostics
-cloud-tpu-diagnostics
+cloud-tpu-diagnostics!=1.1.14
 datasets
 drjax
 flax
@@ -24,6 +25,7 @@ numpy
 omegaconf
 optax
 orbax-checkpoint
+parameterized
 pathwaysutils
 pillow
 pre-commit
@@ -34,15 +36,14 @@ pylint
 pytest
 pytype
 sentencepiece
+seqio
 tensorboard-plugin-profile
 tensorboardx
 tensorflow-datasets
 tensorflow-text
 tensorflow
 tiktoken
-tokamax
+tokamax!=0.1.0
 transformers
 uvloop
 qwix
-google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
-mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip

src/dependencies/requirements/base_requirements/requirements.txt

@@ -0,0 +1,6 @@
+-r requirements.txt
+google-metrax
+ipykernel
+kagglehub
+papermill
+perfetto

src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt

@@ -1,2 +1 @@
 -r requirements.txt
-google-tunix