Adds /checkpoint/experimental/v1/_src/synchronization/multihost_test.py to the presubmit tests.

PiperOrigin-RevId: 930299393

Author: abhishek002002

checkpoint/orbax/checkpoint/_src/futures/synchronization.py

@@ -65,6 +65,12 @@ def get_current_operation_id(cls) -> str:
     """Returns the current operation id."""
     return str(cls._operation_id)
 
+  @classmethod
+  def set_operation_id(cls, operation_id: int) -> None:
+    """Synchronizes the operation id counter to allow lagging ranks to catch up."""
+    cls._operation_id = operation_id
+    cls._operation_id_counter = itertools.count(operation_id + 1)
+
 
 class MultihostSynchronizedValue(Generic[_T]):
   """A thread-safe value that is synchronized across all processes."""

checkpoint/orbax/checkpoint/_src/testing/oss/multiprocess_test.py

@@ -32,6 +32,7 @@
 from absl.testing import absltest
 import jax
 from jax import config
+from orbax.checkpoint._src.futures import synchronization
 from orbax.checkpoint._src.multihost import multihost
 import portpicker
 
@@ -318,6 +319,15 @@ def setUp(self):
           f"multiprocess_test_ensure_all_processes_arrive_at_test_case_{self._testMethodName}",
           10000,
       )
+      sync_key = f"sync_op_id_{self._testMethodName}"
+      if jax.process_index() == 0:
+        client.key_value_set(
+            sync_key,
+            synchronization.OperationIdGenerator.get_current_operation_id(),
+            allow_overwrite=True,
+        )
+      max_op_id = int(client.blocking_key_value_get(sync_key, 10000))
+      synchronization.OperationIdGenerator.set_operation_id(max_op_id)
 
   def multiprocess_create_tempdir(self, name: str | None = None) -> str:
     """Creates a temporary directory for the test."""

checkpoint/orbax/checkpoint/_src/testing/oss/run_tests.py

@@ -22,6 +22,9 @@
 from absl import app
 from absl import flags
 from absl import logging
+import jax
+from orbax.checkpoint._src.futures import synchronization
+from orbax.checkpoint._src.multihost import multihost
 import pytest
 import yaml
 
@@ -130,7 +133,27 @@ def main(argv: Sequence[str]) -> None:
 
     logging.info('Running test: %s (found from %s)', test_path, test_file_yaml)
     try:
-      exit_code = pytest.main([test_path])
+      try:
+        client = multihost.get_jax_distributed_client()
+        if client is not None:
+          normalized_name = (
+              test_file_yaml.replace('/', '_').replace(':', '_')
+          )
+          sync_key = f'sync_op_id_file_{normalized_name}'
+          if jax.process_index() == 0:
+            client.key_value_set(
+                sync_key,
+                synchronization.OperationIdGenerator.get_current_operation_id(),
+                allow_overwrite=True,
+            )
+          max_op_id = int(client.blocking_key_value_get(sync_key, 10000))
+          synchronization.OperationIdGenerator.set_operation_id(max_op_id)
+      except Exception as sync_e:  # pylint: disable=broad-exception-caught
+        logging.warning(
+            'Could not synchronize OperationIdGenerator for file: %s', sync_e
+        )
+
+      exit_code = pytest.main(['--import-mode=importlib', test_path])
       if exit_code == 0:
         results[test_file_yaml] = 'PASSED'
         logging.info('%s: PASSED', test_path)

checkpoint/orbax/checkpoint/checkpoint_manager_slice_test.py

@@ -0,0 +1,150 @@
+# 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 absl import flags
+from absl.testing import flagsaver
+from absl.testing import parameterized
+from etils import epath
+import jax
+import numpy as np
+from orbax.checkpoint import args
+from orbax.checkpoint import checkpoint_manager
+from orbax.checkpoint import checkpoint_utils
+from orbax.checkpoint import test_utils
+from orbax.checkpoint import utils
+from orbax.checkpoint._src.handlers import handler_registration
+from orbax.checkpoint._src.handlers import pytree_checkpoint_handler
+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.serialization import type_handler_registry
+from orbax.checkpoint._src.serialization import type_handlers
+from orbax.checkpoint._src.testing import multiprocess_test
+
+
+FLAGS = flags.FLAGS
+PyTreeCheckpointHandler = pytree_checkpoint_handler.PyTreeCheckpointHandler
+CheckpointManager = checkpoint_manager.CheckpointManager
+CheckpointManagerOptions = checkpoint_manager.CheckpointManagerOptions
+
+
+@test_utils.barrier_compatible_test
+class CheckpointManagerSliceTest(
+    parameterized.TestCase, multiprocess_test.MultiProcessTest
+):
+  """Structure allows test to run as subclasses, not base class."""
+
+  def setUp(self):
+    super().setUp()
+
+    if not multihost.is_runtime_to_distributed_ids_initialized():
+      multihost.initialize_runtime_to_distributed_ids()
+
+    self.assertEqual(jax.device_count(), 8)
+    self.assertEqual(jax.process_count(), 4)
+    self.assertEqual(jax.local_device_count(), 2)
+
+    self.directory = epath.Path(
+        self.multiprocess_create_tempdir(name='checkpoint_manager_slice_test')
+    )
+    test_utils.set_tensorstore_driver_for_test()
+
+    test_utils.sync_global_processes(
+        'CheckpointManagerSliceTest:setup_complete'
+    )
+
+  def tearDown(self):
+    test_utils.sync_global_processes(
+        'CheckpointManagerSliceTest:tests_complete'
+    )
+    super().tearDown()
+
+  def wait_if_async(self, manager):
+    manager.wait_until_finished()  # no-op if no async checkpointers.
+
+  @parameterized.product(
+      enable_async_checkpointing=(False, True),
+      array_metadata_store=(None, array_metadata_store_lib.Store()),
+  )
+  def test_slice(
+      self,
+      enable_async_checkpointing: bool,
+      array_metadata_store: array_metadata_store_lib.Store | None,
+  ):
+    """Test slice."""
+    self.enter_context(
+        flagsaver.flagsaver(experimental_orbax_use_distributed_process_id=True)
+    )
+    global_mesh = test_utils.get_fake_global_mesh_for_slices([{0, 1}, {2, 3}])
+
+    mesh_axes = jax.sharding.PartitionSpec('data')
+    arrays = [
+        test_utils.create_sharded_array(arr, global_mesh, mesh_axes)
+        for arr in [np.arange(8), np.arange(16)]
+    ]
+    assert len(global_mesh.devices[0]) == 4
+    assert jax.process_count() == 4
+    active_processes = {0, 1}
+    primary_host = 0
+    if multihost.process_index() in active_processes:
+      single_slice_arrays = test_utils.select_single_replica(
+          arrays, global_mesh
+      )
+      options = CheckpointManagerOptions(
+          create=False,
+          enable_async_checkpointing=enable_async_checkpointing,
+          multiprocessing_options=checkpoint_manager.MultiprocessingOptions(
+              primary_host=primary_host,
+              active_processes=active_processes,
+          ),
+      )
+      registry = type_handler_registry.create_type_handler_registry(
+          (
+              jax.Array,
+              type_handlers.ArrayHandler(
+                  primary_host=None,
+                  replica_id=None,
+                  use_replica_parallel=False,
+                  array_metadata_store=array_metadata_store,
+              ),
+          ),
+      )
+      handler = PyTreeCheckpointHandler(
+          multiprocessing_options=options.multiprocessing_options,
+          type_handler_registry=registry,
+      )
+      registry = handler_registration.DefaultCheckpointHandlerRegistry()
+      registry.add(None, args.PyTreeSave, handler)
+      registry.add(None, args.PyTreeRestore, handler)
+      with CheckpointManager(
+          self.directory,
+          options=options,
+          handler_registry=registry,
+      ) as manager:
+        self.assertTrue(manager.save(0, args=args.PyTreeSave(arrays)))
+        time.sleep(10)
+        self.wait_if_async(manager)
+        abstract_target = jax.tree.map(
+            utils.to_shape_dtype_struct, single_slice_arrays
+        )
+        restore_args = checkpoint_utils.construct_restore_args(abstract_target)
+        restored = manager.restore(
+            0, args=args.PyTreeRestore(restore_args=restore_args)
+        )
+        test_utils.assert_tree_equal(self, single_slice_arrays, restored)
+
+
+if __name__ == '__main__':
+  multiprocess_test.main()