Separate Pathways tests in Orbax OSS build.

PiperOrigin-RevId: 876687909

Author: Orbax Authors

.github/workflows/build.yml

@@ -344,17 +344,42 @@ jobs:
         pip install -e .[testing,gcs] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
         pip uninstall -y orbax
         pip install gcsfs
-        pip install portpicker pytest chex pyyaml
+        pip install portpicker pytest chex pyyaml pathwaysutils
         if [ "${{ matrix.jax-version }}" = "newest" ]; then
           pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
         elif [ "${{ matrix.jax-version }}" = "nightly" ]; then
           pip install -U --pre "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/
         else
           pip install "jax[tpu]==${{ matrix.jax-version }}" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
         fi
-    - name: Run multiprocess tests
+    - name: Run pathways tests
+      env:
+        JAX_DEFAULT_BACKEND: pathways
+        JAX_PLATFORMS: tpu
+        # Configures JAX to target a subslice within the TPU allocation.
+        JAX_BACKEND_TARGET: subslice
+        # Enables IFRT in Pathways.
+        PATHWAYS_IFRT: true
+        # Allows JAX to run even if some TPUs are not utilized.
+        JAX_ALLOW_UNUSED_TPUS: true
+      run: |
+        python -c "import pathwaysutils; pathwaysutils.initialize(); print('Pathways initialized')" && python orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests.yaml --pathways=1
+    - name: Run colacated pathways tests
+      env:
+        JAX_PLATFORMS: pathways
+        JAX_BACKEND_TARGET: subslice
+        PATHWAYS_IFRT: true
+        JAX_ALLOW_UNUSED_TPUS: true
+        PATHWAYS_EXPECTED_INSTANCES: df=1x1,df=1x1,df=1x1,df=1x1
+        USE_COLOCATED_PYTHON: true
+      run: |
+        python orbax/checkpoint/_src/testing/oss/run_multihost.py --num_processes=1 --tpu_chips_per_process=8 orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests.yaml --pathways=1
+    - name: Run 2 multiprocess tests
+      run: |
+        python orbax/checkpoint/_src/testing/oss/run_multihost.py --num_processes=2 --tpu_chips_per_process=4 orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests.yaml --processes=2
+    - name: Run 4 multiprocess tests
       run: |
-        python orbax/checkpoint/_src/testing/oss/run_multihost.py orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests.yaml --processes=4
+        python orbax/checkpoint/_src/testing/oss/run_multihost.py --num_processes=4 --tpu_chips_per_process=2 orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests.yaml --processes=4
     - name: Run single process tests
       run: |
         python orbax/checkpoint/_src/testing/oss/run_multihost.py --num_processes=1 --tpu_chips_per_process=8 orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests.yaml --processes=1

checkpoint/orbax/checkpoint/_src/multihost/multihost_test.py

@@ -0,0 +1,212 @@
+# Copyright 2026 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+from unittest import mock
+
+from absl.testing import flagsaver
+from absl.testing import parameterized
+from etils import epath
+import jax
+import numpy as np
+from orbax.checkpoint import test_utils
+from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.testing import multiprocess_test
+
+
+class MultihostUtilsTestBase:
+
+  class Test(parameterized.TestCase):
+
+    def setUp(self):
+      super().setUp()
+      self.assertEqual(jax.device_count(), 8)
+      self.assertEqual(jax.process_count(), 4)
+      self.assertEqual(jax.local_device_count(), 2)
+
+      if not multihost.is_runtime_to_distributed_ids_initialized():
+        multihost.initialize_runtime_to_distributed_ids()
+
+      self.tmpdir = epath.Path(
+          self.create_tempdir(name='multihost_test').full_path
+      )
+      test_utils.sync_global_processes('setUp')
+
+    def tearDown(self):
+      test_utils.sync_global_processes('tearDown')
+      super().tearDown()
+
+    def test_process_errors(self):
+      if multihost.process_index() == 1:
+        with self.assertRaises(ValueError):
+          multihost.sync_global_processes(
+              'test_process_errors_1', processes={0}
+          )
+
+    def test_sync_global_processes(self):
+      if multihost.process_index() == 0:
+        time.sleep(2)
+        (self.tmpdir / 'dummy').mkdir(parents=False, exist_ok=False)
+      multihost.sync_global_processes('test_sync_global_processes')
+      self.assertTrue((self.tmpdir / 'dummy').exists())
+
+    def test_sync_global_processes_partial(self):
+      participating_processes = {0, 2}
+      primary_process = 0
+      non_primary_process = 1
+
+      directory = self.tmpdir / 'testdir'
+      if multihost.process_index() == primary_process:
+        directory.mkdir(parents=False, exist_ok=False)
+      test_utils.sync_global_processes(
+          'test_sync_global_processes_partial_setup'
+      )
+
+      if multihost.process_index() == primary_process:
+        time.sleep(2)
+        (directory / 'dummy').mkdir(parents=False, exist_ok=False)
+      if multihost.process_index() in participating_processes:
+        multihost.sync_global_processes(
+            'test_sync_global_processes_partial',
+            processes=participating_processes,
+        )
+      if multihost.process_index() in participating_processes:
+        self.assertTrue((directory / 'dummy').exists())
+      else:
+        self.assertFalse((directory / 'dummy').exists())
+
+      if multihost.process_index() == primary_process:
+        time.sleep(2)
+        (directory / 'foo').mkdir(parents=False, exist_ok=False)
+      if multihost.process_index() in participating_processes:
+        multihost.sync_global_processes(
+            'test_sync_global_processes_partial_second',
+            processes=participating_processes,
+        )
+      if multihost.process_index() in participating_processes:
+        self.assertTrue((directory / 'foo').exists())
+      else:
+        self.assertFalse((directory / 'foo').exists())
+
+      multihost.sync_global_processes('test_sync_global_processes_partial_all')
+      # If non-primary processes get past the above barrier without waiting for
+      # all, then an error would happen for the primary process when trying to
+      # create subdirectories.
+      if multihost.process_index() == non_primary_process:
+        directory.rmtree()
+
+    def test_different_barriers(self):
+      slice1 = {0, 2}
+      slice2 = {1, 3}
+      primary_processes = [0, 1]
+
+      if multihost.process_index() in primary_processes:
+        # Don't sleep for slice1, but do sleep for slice2, so that when slice1
+        # finishes waiting at the barrier, one file exists but the other does
+        # not.
+        time.sleep(3 * multihost.process_index())
+        (self.tmpdir / f'dummy_{multihost.process_index()}').mkdir(
+            parents=False, exist_ok=False
+        )
+
+      if multihost.process_index() in slice1:
+        multihost.sync_global_processes(
+            'test_different_barriers_slice1',
+            processes=slice1,
+        )
+      else:
+        multihost.sync_global_processes(
+            'test_different_barriers_slice2',
+            processes=slice2,
+        )
+      if multihost.process_index() in slice1:
+        self.assertTrue((self.tmpdir / 'dummy_0').exists())
+        self.assertFalse((self.tmpdir / 'dummy_1').exists())
+      else:
+        self.assertTrue((self.tmpdir / 'dummy_0').exists())
+        self.assertTrue((self.tmpdir / 'dummy_1').exists())
+
+    def test_broadcast_one_to_all(self):
+      if multihost.process_index() == 0:
+        tree = {'bar': [5, 12]}
+      else:
+        tree = {'bar': [0, 0]}
+      result = multihost.broadcast_one_to_all(tree)
+
+      expected = {
+          'bar': [np.asarray(5, dtype=np.int32), np.asarray(12, dtype=np.int32)]
+      }
+      test_utils.assert_tree_equal(self, expected, result)
+
+
+    def test_sync_global_processes_with_distributed_barrier(self):
+      with flagsaver.flagsaver(
+          experimental_orbax_use_distributed_barrier=True
+      ), mock.patch.object(
+          multihost.multihost_utils, 'sync_global_devices', autospec=True
+      ) as mock_sync_global_devices, mock.patch.object(
+          multihost, 'get_barrier_sync_fn', autospec=True
+      ) as mock_get_barrier_sync_fn, mock.patch.object(
+          multihost, 'should_skip_process_sync', return_value=False
+      ):
+        multihost.sync_global_processes('test_barrier')
+
+        mock_sync_global_devices.assert_not_called()
+        mock_get_barrier_sync_fn.assert_called_once_with(processes=None)
+        mock_get_barrier_sync_fn.return_value.assert_called_once_with(
+            key='test_barrier', timeout_ms=300000
+        )
+
+    def test_sync_global_processes_without_distributed_barrier(self):
+      with flagsaver.flagsaver(
+          experimental_orbax_use_distributed_barrier=False
+      ), mock.patch.object(
+          multihost.multihost_utils, 'sync_global_devices', autospec=True
+      ) as mock_sync_global_devices, mock.patch.object(
+          multihost, 'get_barrier_sync_fn', autospec=True
+      ) as mock_get_barrier_sync_fn, mock.patch.object(
+          multihost, 'should_skip_process_sync', return_value=False
+      ):
+        multihost.sync_global_processes('test_barrier')
+
+        mock_sync_global_devices.assert_called_once()
+        mock_get_barrier_sync_fn.assert_not_called()
+
+
+class MultihostUtilsTestStandard(MultihostUtilsTestBase.Test):
+
+  def setUp(self):
+    self.enter_context(
+        flagsaver.flagsaver(experimental_orbax_use_distributed_process_id=False)
+    )
+    super().setUp()
+
+  def test_sync_global_processes_partial(self):
+    self.skipTest('Fix this scenario.')
+
+  def test_different_barriers(self):
+    self.skipTest('Fix this scenario.')
+
+
+class MultihostUtilsTestDistributedId(MultihostUtilsTestBase.Test):
+
+  def setUp(self):
+    self.enter_context(
+        flagsaver.flagsaver(experimental_orbax_use_distributed_process_id=True)
+    )
+    super().setUp()
+
+
+if __name__ == '__main__':
+  multiprocess_test.main()

checkpoint/orbax/checkpoint/_src/serialization/pathways_local_type_handlers_test.py

@@ -0,0 +1,118 @@
+# Copyright 2026 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test for Pathways local type handlers."""
+
+from typing import Any, Sequence
+
+from absl import flags
+import jax
+from orbax.checkpoint import test_utils
+from orbax.checkpoint._src.metadata import array_metadata_store as array_metadata_store_lib
+from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.multihost import pathways as multihost_pathways
+from orbax.checkpoint._src.serialization import local_type_handlers_test
+from orbax.checkpoint._src.serialization import ocdbt_utils
+from orbax.checkpoint._src.serialization import pathways_handler_registry
+from orbax.checkpoint._src.serialization import type_handler_registry
+import tensorstore as ts
+
+from .learning.deepmind.jax.ocean.remote_python import rp
+from .pyglib.contrib.g3_multiprocessing import g3_multiprocessing
+from absl.testing import absltest
+
+
+USE_COLOCATED_PYTHON = flags.DEFINE_boolean(
+    'use_colocated_python',
+    False,
+    'Whether to use colocated Python.',
+)
+
+
+class FakeArrayMetadataStore(array_metadata_store_lib.Store):
+  """A fake in-memory store that mimics the real checkpoint store API."""
+
+  def __init__(self):
+    self._data = {}
+
+  async def write(
+      self,
+      checkpoint_dir: Any,
+      array_metadatas: Sequence[Any],
+      process_index: int,
+  ) -> None:
+    """Simulates writing metadata to storage."""
+    if checkpoint_dir not in self._data:
+      self._data[checkpoint_dir] = []
+    self._data[checkpoint_dir].extend(array_metadatas)
+
+  async def read(
+      self, checkpoint_dir: Any, process_index: int | None = None
+  ) -> Any:
+    """Simulates reading metadata from storage."""
+    return {0: self._data.get(checkpoint_dir, [])}
+
+
+class PathwaysLocalTypeHandlersTest(
+    local_type_handlers_test.LocalTypeHandlersTest,
+):
+
+  def setUp(self):
+    super().setUp()
+    self.assertTrue(multihost.is_pathways_backend())
+    self.assertTrue(rp.available())
+    print('jax.devices(): ', jax.devices())
+
+  def validate_topology(self):
+    self.assertEqual(jax.device_count(), 8)
+    self.assertGreater(multihost_pathways.worker_count(None), 1)
+
+  def get_array_handler(self):
+    pathways_handler_registry.register_pathways_handlers(
+        checkpointing_impl=pathways_handler_registry.CheckpointingImpl.from_options(
+            use_colocated_python=USE_COLOCATED_PYTHON.value,
+            use_remote_python=True,  # Fallback
+        ),
+        primary_host=None,
+        replica_id=None,
+        use_replica_parallel=False,
+        thinmint_testing=True,
+        array_metadata_store=FakeArrayMetadataStore(),
+    )
+    return type_handler_registry.get_type_handler(jax.Array)
+
+  def validate_paths(self):
+    # Array files should not exist at the global path level.
+    self.assertFalse((self.base_directory / 'manifest.ocdbt').exists())
+    for worker_id in range(multihost_pathways.worker_count(None)):
+      self.assertTrue((self.base_directory / f'local_{worker_id}').exists())
+
+  async def finalize_save(
+      self, *, ts_context: ts.Context, use_zarr3: bool, use_ocdbt: bool
+  ):
+    if use_ocdbt:
+      for worker_id in range(multihost_pathways.worker_count(None)):
+        await ocdbt_utils.merge_ocdbt_per_process_files(
+            self.base_directory / f'local_{worker_id}',
+            ts_context=ts_context,
+            use_zarr3=use_zarr3,
+        )
+      test_utils.sync_global_processes(
+          'local_serialization:merge_ocdbt_complete'
+      )
+
+
+if __name__ == '__main__':
+  jax.config.parse_flags_with_absl()
+  g3_multiprocessing.handle_test_main(googletest.main)