Add flags for PyTorch & DCP support in the Orbax checkpoint benchmark launcher.

PiperOrigin-RevId: 865621844

Author: lyglst

checkpoint/CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - #v1 Add `use_load_and_broadcast` option.
+- Add PyTorch DCP (Distributed Checkpoint) to the benchmark suite.
 
 ### Removed
 

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/core.py

@@ -25,12 +25,12 @@
 from absl import logging
 from etils import epath
 import jax
-from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.testing.benchmarks.core import checkpoint_generation
 from orbax.checkpoint._src.testing.benchmarks.core import configs
 from orbax.checkpoint._src.testing.benchmarks.core import device_mesh
 from orbax.checkpoint._src.testing.benchmarks.core import directory_setup
 from orbax.checkpoint._src.testing.benchmarks.core import metric as metric_lib
+from orbax.checkpoint._src.testing.benchmarks.core import multihost
 
 
 @dataclasses.dataclass(frozen=True)
@@ -148,7 +148,7 @@ def run(self, repeat_index: int | None = None) -> TestResult:
       name += f"_repeat_{repeat_index}"
     logging.info(
         "[process_id=%s] Setting up test: %s",
-        multihost.process_index(),
+        multihost.get_process_index(),
         name,
     )
 
@@ -193,7 +193,7 @@ def run(self, repeat_index: int | None = None) -> TestResult:
 
     logging.info(
         "[process_id=%s] Executing test function: %s",
-        multihost.process_index(),
+        multihost.get_process_index(),
         name,
     )
     try:
@@ -203,13 +203,13 @@ def run(self, repeat_index: int | None = None) -> TestResult:
       # execution is recorded in the TestResult.
       if sys.version_info >= (3, 11):
         e.add_note(
-            f"[process_id={multihost.process_index()}],"
+            f"[process_id={multihost.get_process_index()}],"
             f" {test_context_summary[:100]}"
         )
       logging.error(
           "[process_id=%s] Test function '%s' context: %s, raised an"
           " exception: %s",
-          multihost.process_index(),
+          multihost.get_process_index(),
           name,
           test_context_summary[:100],
           e,
@@ -223,7 +223,7 @@ def run(self, repeat_index: int | None = None) -> TestResult:
 
     logging.info(
         "[process_id=%s] Test finished: %s",
-        multihost.process_index(),
+        multihost.get_process_index(),
         name,
     )
 
@@ -306,13 +306,13 @@ def _get_options_product(self) -> Sequence[BenchmarkOptions]:
         option_instances.append(option_instance)
         logging.info(
             "[process_id=%s] Generating valid option combination: %s",
-            multihost.process_index(),
+            multihost.get_process_index(),
             option_instance,
         )
       else:
         logging.info(
             "[process_id=%s] Skipping invalid option combination: %s",
-            multihost.process_index(),
+            multihost.get_process_index(),
             option_instance,
         )
     return option_instances

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/core_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import dataclasses
+import os
 from typing import List
 from unittest import mock
 
@@ -23,12 +24,14 @@
 import jax
 import numpy as np
 from orbax.checkpoint import test_utils
+from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.testing.benchmarks.core import checkpoint_generation
 from orbax.checkpoint._src.testing.benchmarks.core import configs
 from orbax.checkpoint._src.testing.benchmarks.core import core
 from orbax.checkpoint._src.testing.benchmarks.core import device_mesh
 from orbax.checkpoint._src.testing.benchmarks.core import directory_setup
 from orbax.checkpoint._src.testing.benchmarks.core import metric as metric_lib
+import torch.distributed as dist
 
 
 @dataclasses.dataclass(frozen=True)
@@ -472,8 +475,13 @@ def test_run_no_benchmarks_generated(
   @mock.patch.object(directory_setup, 'setup_test_directory')
   @mock.patch.object(checkpoint_generation, 'generate_checkpoint')
   @mock.patch.object(logging, 'info')
+  @mock.patch.object(
+      multihost,
+      'sync_global_processes',
+  )
   def test_run_generates_report_with_failures(
       self,
+      mock_sync_global_processes,
       mock_logging_info,
       mock_generate_checkpoint,
       mock_setup_test_directory,
@@ -499,6 +507,11 @@ def test_fn(self, test_context: core.TestContext) -> core.TestResult:
     suite = core.TestSuite(name='report_suite', benchmarks_generators=[gen])
     suite.run()
 
+    mock_sync_global_processes.assert_any_call('benchmark:run')
+    mock_sync_global_processes.assert_any_call('benchmark:setup_test_directory')
+    mock_sync_global_processes.assert_any_call('benchmark:setup_pytree')
+    mock_sync_global_processes.assert_any_call('test_suite:run_end')
+
     # 3 benchmarks generated: (1,a), (1,b), (2,b).
     # (1,a), (1,b) pass. (2,b) fails because (2,a) is invalid.
 
@@ -512,6 +525,24 @@ def test_fn(self, test_context: core.TestContext) -> core.TestResult:
     self.assertIn('--- Failed Runs ---', report_log)
     self.assertIn("Error: ValueError('opt1=2, opt2=b failed')", report_log)
 
+  @mock.patch.object(core.Benchmark, 'run')
+  @mock.patch.object(dist, 'barrier')
+  def test_run_with_torch(self, mock_dist_barrier, mock_benchmark_run):
+    # Initialize torch.distributed for testing.
+    os.environ.setdefault('MASTER_ADDR', 'localhost')
+    os.environ.setdefault('MASTER_PORT', '12355')
+    dist.init_process_group(backend='gloo', rank=0, world_size=1)
+    gen = MyGenerator(
+        checkpoint_configs=[configs.CheckpointConfig(spec={})],
+        options=MyBenchmarkOptions(opt1=[1, 2]),
+    )
+    suite = core.TestSuite(name='my_suite', benchmarks_generators=[gen])
+
+    suite.run()
+    mock_dist_barrier.assert_any_call()
+
+    self.assertEqual(mock_benchmark_run.call_count, 1)
+
 
 if __name__ == '__main__':
   absltest.main()

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/directory_setup.py

@@ -16,7 +16,7 @@
 
 from absl import logging
 from etils import epath
-import jax
+from orbax.checkpoint._src.testing.benchmarks.core import multihost
 
 
 def setup_test_directory(
@@ -39,7 +39,7 @@ def setup_test_directory(
   if repeat_index is not None:
     path = path / f"repeat_{repeat_index}"
   logging.info("Setting up test directory at: %s", path)
-  if jax.process_index() == 0:
+  if multihost.get_process_index() == 0:
     if path.exists():
       logging.warning("Test directory %s already exists. Deleting it.", path)
       path.rmtree()

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/directory_setup_test.py

@@ -16,8 +16,8 @@
 
 from absl.testing import absltest
 from etils import epath
-import jax
 from orbax.checkpoint._src.testing.benchmarks.core import directory_setup
+from orbax.checkpoint._src.testing.benchmarks.core import multihost
 
 
 class DirectorySetupTest(absltest.TestCase):
@@ -51,7 +51,7 @@ def test_setup_test_directory_already_exists(self):
     self.assertTrue(path.exists())
     self.assertFalse((path / 'some_file').exists())
 
-  @mock.patch.object(jax, 'process_index', return_value=1)
+  @mock.patch.object(multihost, 'get_process_index', return_value=1)
   def test_setup_test_directory_non_zero_process_index_does_not_exist(self, _):
     temp_dir = self.create_tempdir()
 

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/metric.py

@@ -30,7 +30,7 @@
 from clu import metric_writers
 from etils import epath
 import numpy as np
-from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.testing.benchmarks.core import multihost
 import psutil
 import tensorstore as ts
 
@@ -47,7 +47,7 @@ def start(self):
     self._start_time = time.perf_counter()
     logging.info(
         "[process_id=%s] Starting metric: '%s'...",
-        multihost.process_index(),
+        multihost.get_process_index(),
         self.name,
     )
 
@@ -56,7 +56,7 @@ def stop(self) -> dict[str, tuple[Any, str]]:
     duration = time.perf_counter() - self._start_time
     logging.info(
         "[process_id=%s] Finished metric: '%s' (took %.4fs)",
-        multihost.process_index(),
+        multihost.get_process_index(),
         self.name,
         duration,
     )
@@ -168,7 +168,7 @@ def stop(self) -> dict[str, tuple[Any, str]]:
 
     self._log_tracemalloc_snapshot_diff(
         self.name,
-        multihost.process_index(),
+        multihost.get_process_index(),
         self._start_snapshot,
         end_snapshot,
         top_n=15,
@@ -285,7 +285,7 @@ def stop(self) -> dict[str, tuple[Any, str]]:
     diff = self._diff_metrics(self._start_metrics, end_metrics)
     logging.info(
         "[process_id=%s] Finished metric: %s, num_diffs=%d",
-        multihost.process_index(),
+        multihost.get_process_index(),
         self.name,
         len(diff),
     )
@@ -423,12 +423,12 @@ def report(self):
     """Logs a formatted report of all collected metrics."""
     report_lines = []
     report_lines.append(
-        f"---[process_id={multihost.process_index()}] {self.name} Metrics"
+        f"---[process_id={multihost.get_process_index()}] {self.name} Metrics"
         " Report ---"
     )
     if not self.results:
       report_lines.append(
-          f"[process_id={multihost.process_index()}] No metrics recorded."
+          f"[process_id={multihost.get_process_index()}] No metrics recorded."
       )
     else:
       for name, (value, unit) in sorted(self.results.items()):
@@ -565,7 +565,7 @@ def add_result(
   def _get_writer(self, benchmark_name: str) -> Any:
     """Gets or creates a TensorBoard writer for the given benchmark."""
     if benchmark_name not in self._writers:
-      is_primary_host = multihost.process_index() == 0
+      is_primary_host = multihost.get_process_index() == 0
       self._writers[benchmark_name] = metric_writers.create_default_writer(
           self._tensorboard_dir,
           just_logging=not is_primary_host,

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/multihost.py

@@ -0,0 +1,55 @@
+# 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.
+
+"""Multihost utilities for benchmarks."""
+
+import threading
+
+from absl import logging
+from orbax.checkpoint._src.multihost import multihost
+
+
+def get_process_index() -> int:
+  """Returns process index from torch if available, else from multihost."""
+  try:
+    import torch.distributed as dist  # pylint: disable=g-import-not-at-top
+
+    if dist.is_initialized():
+      return dist.get_rank()
+  except ImportError:
+    pass
+  return multihost.process_index()
+
+
+def sync_global_processes(
+    name: str,
+) -> None:
+  """Syncs global processes using torch if available, else multihost."""
+  try:
+    import torch.distributed as dist  # pylint: disable=g-import-not-at-top
+
+    if dist.is_initialized():
+      logging.vlog(
+          1,
+          "[process=%s][thread=%s] sync_global_processes with torch"
+          " barrier: %s",
+          dist.get_rank(),
+          threading.current_thread().name,
+          name,
+      )
+      dist.barrier()
+      return
+  except ImportError:
+    pass
+  multihost.sync_global_processes(name)