Include Orbax Checkpoint experimental/emergency tests in OSS.

PiperOrigin-RevId: 874406518

Author: Orbax Authors

.github/workflows/build.yml

@@ -352,9 +352,12 @@ jobs:
         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 2 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=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 --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/testing/oss/generate_multiprocess_test.py

@@ -15,6 +15,7 @@
 """Script to generate YAML file with test targets based on tags."""
 
 import ast
+import collections
 import os
 import sys
 
@@ -29,7 +30,11 @@
     'pytype_strict_contrib_test',
 ]
 EXCLUDED_PATHS = [
-    'orbax/checkpoint/experimental',
+    'orbax/checkpoint/experimental/model_surgery',
+    'orbax/checkpoint/experimental/v1',
+    'orbax/checkpoint/experimental/emergency/p2p',
+    'orbax/checkpoint/experimental/emergency/checkpoint_manager_test.py',
+    'orbax/checkpoint/experimental/emergency/replicator_checkpoint_manager_test.py',
     'orbax/checkpoint/google',
 ]
 
@@ -60,53 +65,15 @@ def get_str_val(node):
   return None
 
 
-def inherits_from_multiprocess_test(test_file_path):
-  """Checks if test file inherits from MultiProcessTest."""
-  try:
-    with open(test_file_path, 'r') as f:
-      content = f.read()
-  except FileNotFoundError:
-    return False
-  try:
-    tree = ast.parse(content, filename=test_file_path)
-  except SyntaxError:
-    return False
-
-  imported_as_name = None  # if imported as `from ... import MultiProcessTest`
-  imported_as_module = []  # if imported as `from ... import multiprocess_test`
-
-  for node in tree.body:
-    if isinstance(node, ast.ImportFrom):
-      if node.module == 'orbax.checkpoint._src.testing.multiprocess_test':
-        for alias in node.names:
-          if alias.name == 'MultiProcessTest':
-            imported_as_name = alias.asname or alias.name
-      elif node.module == 'orbax.checkpoint._src.testing':
-        for alias in node.names:
-          if alias.name == 'multiprocess_test':
-            imported_as_module.append(alias.asname or alias.name)
-
-  if not imported_as_name and not imported_as_module:
-    return False
-
-  for node in tree.body:
-    if isinstance(node, ast.ClassDef):
-      for base in node.bases:
-        if (
-            imported_as_name
-            and isinstance(base, ast.Name)
-            and base.id == imported_as_name
-        ):
-          return True
-        if (
-            imported_as_module
-            and isinstance(base, ast.Attribute)
-            and isinstance(base.value, ast.Name)
-            and base.value.id in imported_as_module
-            and base.attr == 'MultiProcessTest'
-        ):
-          return True
-  return False
+def get_num_processes(args):
+  """Returns num_processes from args."""
+  for arg in args:
+    if arg.startswith('--num_processes='):
+      try:
+        return int(arg.split('=', 1)[1])
+      except ValueError:
+        return None
+  return None
 
 
 def get_build_targets(build_file_path):
@@ -135,18 +102,18 @@ def get_build_targets(build_file_path):
 
     if rule_name in TEST_RULES:
       kwargs = get_kwargs(call)
-      if 'name' in kwargs and 'tags' in kwargs:
+      if 'name' in kwargs:
         name = get_str_val(kwargs['name'])
-        tags = get_list_val(kwargs['tags'])
+        tags = get_list_val(kwargs['tags']) if 'tags' in kwargs else []
         srcs = get_list_val(kwargs['srcs']) if 'srcs' in kwargs else []
-        if name and tags:
-          yield name, tags, srcs
+        args = get_list_val(kwargs['args']) if 'args' in kwargs else []
+        if name:
+          yield name, tags, srcs, args
 
 
 def run(root_dir, output_file):
   """Runs the script to generate tagged tests file."""
-  tests_by_tag = {tag: [] for tag in TAG_MAPPING.values()}
-  tests_by_tag['processes:1'] = []
+  tests_by_tag = collections.defaultdict(list)
 
   count = 0
   for dirpath, dirnames, filenames in os.walk(root_dir):
@@ -166,35 +133,33 @@ def run(root_dir, output_file):
       count += 1
       build_file = os.path.join(dirpath, 'BUILD')
       package_path = dirpath.removeprefix('third_party/py/')
-      for name, tags, srcs in get_build_targets(build_file):
+      for name, tags, srcs, args in get_build_targets(build_file):
+        if not any(tag in TAG_MAPPING for tag in tags):
+          continue
         if srcs and any(
             os.path.join(dirpath, srcs[0]).startswith(p) for p in EXCLUDED_PATHS
         ):
           continue
-        is_multiprocess = False
-        if srcs:
-          is_multiprocess = inherits_from_multiprocess_test(
-              os.path.join(dirpath, srcs[0])
-          )
         target_path = f'{package_path}:{name}'
-        if not is_multiprocess:
-          tests_by_tag['processes:1'].append(target_path)
+        num_processes = get_num_processes(args)
+        if num_processes and num_processes > 1:
+          tag = f'processes:{num_processes}'
+          tests_by_tag[tag].append(target_path)
         else:
-          for tag in tags:
-            if tag in TAG_MAPPING:
-              tests_by_tag[TAG_MAPPING[tag]].append(target_path)
+          tests_by_tag['processes:1'].append(target_path)
 
   print(f'Processed {count} BUILD files.')
 
+  result_dict = {}
   for tag in tests_by_tag:
-    tests_by_tag[tag] = sorted(list(set(tests_by_tag[tag])))
+    result_dict[tag] = sorted(list(set(tests_by_tag[tag])))
 
   header = """# DO NOT EDIT!
 """
   os.makedirs(os.path.dirname(output_file), exist_ok=True)
   with open(output_file, 'w') as f:
     f.write(header)
-    yaml.dump(tests_by_tag, f, default_flow_style=False)
+    yaml.dump(result_dict, f, default_flow_style=False)
   print(f'Output written to {output_file}')