Add checkpoint_manager_test.py to multiprocess tests.

PiperOrigin-RevId: 828450426

Author: Orbax Authors

.github/workflows/build.yml

@@ -242,7 +242,7 @@ jobs:
         working-directory: checkpoint
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.10"]
         jax-version: ["0.6.0"]
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
@@ -262,14 +262,22 @@ jobs:
         else
           pip install "jax[k8s,cuda12]>=${{ matrix.jax-version }}" "jaxlib>=${{ matrix.jax-version }}"
         fi
+        pip install fsspec
         pip install gcsfs
         pip install portpicker
     - name: Run benchmarks
       env:
         GCS_BUCKET_PATH: gs://orbax-benchmarks/benchmark-results/${{ github.run_id }}
       run: |
-        cd orbax/checkpoint/_src/testing/benchmarks && python run_benchmarks.py --config_file=configs/pytree_checkpoint_benchmark.yaml --output_directory=$GCS_BUCKET_PATH
+        cd orbax/checkpoint/_src/testing/benchmarks
+        for entry in configs/*; do
+          if [ -f "$entry" ]; then
+            echo "Running benchmark for $entry"
+            python run_benchmarks.py --config_file="$entry" --output_directory=$GCS_BUCKET_PATH
+          fi
+        done
         cd ../../../../..
+    # cd orbax/checkpoint/_src/testing/benchmarks && python run_benchmarks.py --config_file=configs/pytree_checkpoint_benchmark.yaml --output_directory=$GCS_BUCKET_PATH
     # python -m pytest orbax/checkpoint/_src/handlers/array_checkpoint_handler_test.py
     # cd orbax/checkpoint/_src/testing/benchmarks && python run_benchmarks.py --config_file=configs/pytree_checkpoint_benchmark.yaml --output_directory=$GCS_BUCKET_PATH
     # The below step just reports the success or failure of tests as a "commit status".
@@ -278,7 +286,8 @@ jobs:
       env:
         TEST_TMPDIR: gs://orbax-benchmarks/unit-tests/${{ github.run_id }}
       run: |
-        python -c "import jax; jax.distributed.initialize(); print(jax.devices()); import pytest; test_files = [line.strip() for line in open('orbax/checkpoint/_src/testing/multiprocess_tests.txt') if line.strip()]; pytest.main(['-c', '/dev/null'] + test_files)"
+        python -c "import jax; jax.distributed.initialize(); print(jax.devices()); import pytest; pytest.main(['orbax/checkpoint/checkpoint_manager_test.py'])"
+    # python -c "import jax; jax.distributed.initialize(); print(jax.devices()); import pytest; test_files = [line.strip() for line in open('orbax/checkpoint/_src/testing/multiprocess_tests.txt') if line.strip()]; pytest.main(['-c', '/dev/null'] + test_files)"
     # python -m pytest orbax/checkpoint/_src/handlers/array_checkpoint_handler_test.py;"
     # cd orbax/checkpoint/_src/testing/benchmarks && python run_benchmarks.py --config_file=configs/pytree_checkpoint_benchmark.yaml --output_directory=$GCS_BUCKET_PATH
     # python -m pytest orbax/checkpoint/checkpoint_manager_test.py

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/array_handler_benchmark.yaml

@@ -3,7 +3,7 @@ suite_name: "ArrayHandler Benchmark"
 mesh_config:
   mesh_axes: ["data", "model"]
   ici_parallelism: {"data": 2, "model": 2}
-  dcn_parallelism: {"data": 2, "model": 1}
+  dcn_parallelism: {"data": 4, "model": 1}
 
 checkpoint_config:
   spec: