Synchronize OperationIdGenerator across processes in multihost tests.

PiperOrigin-RevId: 929817113

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,6 +133,26 @@ def main(argv: Sequence[str]) -> None:
 
     logging.info('Running test: %s (found from %s)', test_path, test_file_yaml)
     try:
+      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([test_path])
       if exit_code == 0:
         results[test_file_yaml] = 'PASSED'