Move MTC files to multi_tier_checkpointing and use local checkpoint engine

PiperOrigin-RevId: 908412104

Author: SujeethJinesh

.github/workflows/build.yml

@@ -72,7 +72,7 @@ jobs:
         --ignore=orbax/checkpoint/experimental/emergency/local_checkpoint_data_debugging_test.py \
         --ignore=orbax/checkpoint/experimental/emergency/local_checkpoint_manager_test.py \
         --ignore=orbax/checkpoint/experimental/emergency/multihost_test.py \
-        --ignore=orbax/checkpoint/experimental/emergency/replicator_checkpoint_manager_test.py \
+        --ignore=orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/replicator_checkpoint_manager_test.py \
         --ignore=orbax/checkpoint/_src/testing/multiprocess_test.py \
         --ignore=orbax/checkpoint/_src/testing/oss/multiprocess_test.py \
         $(python3 -c "import yaml; d=yaml.safe_load(open('orbax/checkpoint/_src/testing/oss/tagged_tests.yaml')); print(' '.join(['--ignore=' + t.replace(':', '/') + '.py' for k,v in d.items() if k.startswith('processes') and v for t in v]))")

checkpoint/CHANGELOG.md

@@ -18,6 +18,7 @@ merging.
 - Add Colocated Python Orchestration handling
 - #v1 Add `LeafHandler` as a `CheckpointableHandler`, so that ordinary PyTree
 leaves can also be saved as individual checkpointables.
+- Move MTC files to multi_tier_checkpointing and use local checkpoint engine
 
 ## [0.11.36] - 2026-04-14
 

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

@@ -34,7 +34,7 @@
     '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/experimental/emergency/multi_tier_checkpointing/replicator_checkpoint_manager_test.py',
     'orbax/checkpoint/google',
 ]
 

checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/process_metadata_checkpoint_handler.py

@@ -0,0 +1,135 @@
+# 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.
+
+"""ProcessMetadataCheckpointHandler class.
+
+Implementation of CheckpointHandler interface.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Callable, List, Tuple
+
+from absl import logging
+from etils import epath
+import jax
+from orbax.checkpoint import checkpoint_args
+from orbax.checkpoint import options as options_lib
+from orbax.checkpoint._src import asyncio_utils
+from orbax.checkpoint._src.futures import future
+from orbax.checkpoint._src.handlers import async_checkpoint_handler
+from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint.experimental.emergency import mesh_consistency
+
+
+CheckpointArgs = checkpoint_args.CheckpointArgs
+PreviousDistributedToDeviceIds = List[List[int]]
+PreviousDeviceIds = List[int]
+register_with_handler = checkpoint_args.register_with_handler
+
+
+class ProcessMetadataCheckpointHandler(
+    async_checkpoint_handler.AsyncCheckpointHandler
+):
+  """Saves process metadata."""
+
+  def __init__(
+      self,
+      multiprocessing_options: options_lib.MultiprocessingOptions = (
+          options_lib.MultiprocessingOptions()
+      ),
+      distributed_to_device_ids_fn: Callable[[], List[List[int]]] | None = None,
+  ) -> None:
+    """Initializes ProcessMetadataCheckpointHandler."""
+    self._distributed_to_device_ids_fn = (
+        distributed_to_device_ids_fn or multihost.distributed_to_device_ids
+    )
+
+  async def async_save(
+      self,
+      directory: epath.Path,
+      args: ProcessMetadataSaveArgs,
+  ) -> List[future.Future] | None:
+    """Saves the given process metadata.
+
+    Args:
+      directory: save location directory.
+      args: ProcessMetadataSaveArgs (see below).
+
+    Returns:
+      A list of commit futures.
+    """
+    logging.info('Saving process metadata to %s', directory)
+    return [
+        future.CommitFutureAwaitingContractedSignals(
+            mesh_consistency.save_process_metadata(
+                directory,
+                args.global_mesh,
+                self._distributed_to_device_ids_fn(),
+            ),
+            name='process_metadata_ch_save',
+        )
+    ]
+
+  def save(
+      self,
+      directory: epath.Path,
+      args: ProcessMetadataSaveArgs,
+  ) -> None:
+    async def async_save(
+        directory: epath.Path,
+        args: ProcessMetadataSaveArgs,
+    ) -> None:
+      commit_futures = await self.async_save(directory, args)
+      if commit_futures:
+        for f in commit_futures:
+          f.result()
+
+    asyncio_utils.run_sync(async_save(directory, args))
+
+  def restore(
+      self,
+      directory: epath.Path,
+      args: ProcessMetadataRestoreArgs,
+  ) -> Tuple[PreviousDistributedToDeviceIds, PreviousDeviceIds]:
+    """Restores metadata from directory.
+
+    Args:
+      directory: restore location directory.
+      args: ProcessMetadataRestoreArgs (see below).
+
+    Returns:
+      A tuple of previous distributed to device ids and previous device ids.
+    """
+    logging.info('Restoring process metadata from %s', directory)
+    return mesh_consistency.read_process_metadata(directory)
+
+
+@register_with_handler(ProcessMetadataCheckpointHandler, for_save=True)
+@dataclasses.dataclass
+class ProcessMetadataSaveArgs(CheckpointArgs):
+  """Parameters for saving process metadata.
+
+  Attributes:
+    global_mesh: the global mesh.
+  """
+
+  global_mesh: jax.sharding.Mesh
+
+
+@register_with_handler(ProcessMetadataCheckpointHandler, for_restore=True)
+@dataclasses.dataclass
+class ProcessMetadataRestoreArgs(CheckpointArgs):
+  """Parameters for restoring process metadata."""