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,30 @@ 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
+        failed_benchmarks=""
+        for entry in configs/*; do
+          if [ -f "$entry" ]; then
+            echo "Running benchmark for $entry"
+            if ! python run_benchmarks.py --config_file="$entry" --output_directory=$GCS_BUCKET_PATH; then
+              echo "Benchmark $entry failed"
+              failed_benchmarks="$failed_benchmarks $entry"
+            fi
+          fi
+        done
         cd ../../../../..
+        if [ -n "$failed_benchmarks" ]; then
+          echo "The following benchmarks failed:$failed_benchmarks"
+          exit 1
+        fi
+    # 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 +294,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

@@ -1,9 +1,12 @@
 suite_name: "ArrayHandler Benchmark"
 
-mesh_config:
-  mesh_axes: ["data", "model"]
-  ici_parallelism: {"data": 2, "model": 2}
-  dcn_parallelism: {"data": 2, "model": 1}
+mesh_configs:
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 2, "model": 2}
+    dcn_parallelism: {"data": 2, "model": 1}
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 1, "model": 1}
+    dcn_parallelism: {"data": 4, "model": 1}
 
 checkpoint_config:
   spec:
@@ -17,5 +20,4 @@ benchmarks:
       use_replica_parallel: [True, False]
       enable_replica_parallel_separate_folder: [True, False]
       use_metadata_store: [True, False]
-      use_colocated_python: [True, False]
 

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

@@ -1,8 +1,11 @@
 suite_name: "CheckpointManager Benchmark"
 
-mesh_config:
-  mesh_axes: ["data", "model"]
-  ici_parallelism: {"data": 2, "model": 2}
+mesh_configs:
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 2, "model": 2}
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 1, "model": 1}
+    dcn_parallelism: {"data": 4, "model": 1}
 
 checkpoint_config:
   spec:

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

@@ -2,10 +2,13 @@
 suite_name: "CheckpointPolicy Benchmark"
 
 
-mesh_config:
-  mesh_axes: ["data", "model"]
-  ici_parallelism: {"data": 4, "model": 8}
-  dcn_parallelism: {"data": 1, "model": 1}
+mesh_configs:
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 4, "model": 8}
+    dcn_parallelism: {"data": 1, "model": 1}
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 1, "model": 1}
+    dcn_parallelism: {"data": 4, "model": 1}
 
 checkpoint_config:
   spec:

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

@@ -1,13 +1,16 @@
 # The name for the entire test suite run.
 suite_name: "EmergencyCheckpointManager Benchmark"
 
-mesh_config:
-  mesh_axes: ["data", "stage", "fsdp", "fsdp_transpose", "sequence", "tensor", "expert", "autoregressive"]
-  # ICI: Within a slice. Assuming 8 devices per slice.
-  # DCN: Across slices.
-  ici_parallelism: {"fsdp": 4, "tensor": 1, "data": 2}
-  dcn_parallelism: {"data": 2} # num_slices on the axis at replica_axis_index
-  allow_split_physical_axes: true
+mesh_configs:
+  - mesh_axes: ["data", "stage", "fsdp", "fsdp_transpose", "sequence", "tensor", "expert", "autoregressive"]
+    # ICI: Within a slice. Assuming 8 devices per slice.
+    # DCN: Across slices.
+    ici_parallelism: {"fsdp": 4, "tensor": 1, "data": 2}
+    dcn_parallelism: {"data": 2} # num_slices on the axis at replica_axis_index
+    allow_split_physical_axes: true
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 1, "model": 1}
+    dcn_parallelism: {"data": 4, "model": 1}
 
 checkpoint_config:
   spec:

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

@@ -1,9 +1,12 @@
 suite_name: "Multihost Dispatchers Benchmark"
 
-mesh_config:
-  mesh_axes: ["data", "model"]
-  ici_parallelism: {"data": 2, "model": 2}
-  dcn_parallelism: {"data": 2}
+mesh_configs:
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 2, "model": 2}
+    dcn_parallelism: {"data": 2}
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 1, "model": 1}
+    dcn_parallelism: {"data": 4, "model": 1}
 
 checkpoint_config:
   spec:

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

@@ -1,12 +1,15 @@
 # The name for the entire test suite run.
 suite_name: "Multislice Broadcast Benchmark"
 
-mesh_config:
-  mesh_axes: ["data", "stage", "fsdp", "fsdp_transpose", "sequence", "tensor", "expert", "autoregressive"]
-  # ICI: Within a slice.
-  ici_parallelism: {"fsdp": 4, "data": 2}
-  # DCN: Across slices.
-  dcn_parallelism: {"data": 2}
+mesh_configs:
+  - mesh_axes: ["data", "stage", "fsdp", "fsdp_transpose", "sequence", "tensor", "expert", "autoregressive"]
+    # ICI: Within a slice.
+    ici_parallelism: {"fsdp": 4, "data": 2}
+    # DCN: Across slices.
+    dcn_parallelism: {"data": 2}
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 1, "model": 1}
+    dcn_parallelism: {"data": 4, "model": 1}
 
 checkpoint_config:
   spec:

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

@@ -29,6 +29,9 @@ mesh_configs:
     ici_parallelism: {"fsdp": 32, "tensor": 2, "data": 1}
     dcn_parallelism: {"data": 2} # num_slices on the axis at replica_axis_index
     allow_split_physical_axes: True
+  - mesh_axes: ["data", "model"]
+    ici_parallelism: {"data": 1, "model": 1}
+    dcn_parallelism: {"data": 4, "model": 1}
 
 checkpoint_config:
   spec: