#v1 Add checkpointables support for training.Checkpointer.

PiperOrigin-RevId: 752424341

Author: Orbax Authors

checkpoint/CHANGELOG.md

@@ -21,7 +21,6 @@ would want them to be preferred.
 
 - #v1 Add `JsonHandler`.
 - #v1 Add `training.Checkpointer`.
-- #v1 Add checkpointables support for `training.Checkpointer`.
 - `PartsOf` structure which holds a PyTree whose leaf nodes may be missing.
 - #v1 Add compatibility tests for save-by-v0-load-by-v1 and also fix code.
 

checkpoint/orbax/checkpoint/_src/handlers/BUILD

@@ -13,6 +13,7 @@ py_library(
         ":checkpoint_handler",
         ":handler_registration",
         ":handler_type_registry",
+        ":proto_checkpoint_handler",
         "//checkpoint/orbax/checkpoint:checkpoint_args",
         "//checkpoint/orbax/checkpoint:options",
         "//checkpoint/orbax/checkpoint/_src:asyncio_utils",

checkpoint/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py

@@ -48,7 +48,7 @@
 import asyncio
 import concurrent.futures
 import dataclasses
-from typing import Any, Dict, List, Mapping, MutableSet, Optional, Tuple, Type
+from typing import Any, Coroutine, Dict, List, Mapping, MutableSet, Optional, Tuple, Type
 
 from absl import logging
 from etils import epath
@@ -63,6 +63,7 @@
 from orbax.checkpoint._src.handlers import checkpoint_handler
 from orbax.checkpoint._src.handlers import handler_registration
 from orbax.checkpoint._src.handlers import handler_type_registry
+from orbax.checkpoint._src.handlers import proto_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint
 from orbax.checkpoint._src.metadata import step_metadata_serialization
 from orbax.checkpoint._src.path import atomicity
@@ -78,6 +79,10 @@
 CompositeItemMetadata = checkpoint.CompositeItemMetadata
 AsyncCheckpointHandler = async_checkpoint_handler.AsyncCheckpointHandler
 register_with_handler = checkpoint_args.register_with_handler
+ProtoCheckpointHandler = proto_checkpoint_handler.ProtoCheckpointHandler
+ProtoSaveArgs = proto_checkpoint_handler.ProtoSaveArgs
+ProtoRestoreArgs = proto_checkpoint_handler.ProtoRestoreArgs
+AsyncSaveCoroutine = Coroutine[Any, Any, Optional[List[Future]]]
 Composite = composite.Composite
 HandlerAwaitableSignal = synchronization.HandlerAwaitableSignal
 

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -92,9 +92,10 @@
 FileOptions = options_lib.FileOptions
 
 DEFAULT_ITEM_NAME = 'default'
+DESCRIPTOR_ITEM_NAME = 'descriptor'
 METRIC_ITEM_NAME = 'metrics'
 METADATA_ITEM_NAME = 'metadata'
-RESERVED_ITEM_NAMES = []
+RESERVED_ITEM_NAMES = [DESCRIPTOR_ITEM_NAME, METRIC_ITEM_NAME]
 
 _INIT_TIME = datetime.datetime.now(tz=datetime.timezone.utc)
 
@@ -318,7 +319,6 @@ class CheckpointManagerOptions:
     is the sole means of determining when a checkpoint should be saved. If not
     provided, these other options are used instead. Prefer to use this option
     over others.
-  prevent_write_metrics: False by default. If True, metrics will not be written.
   """
 
   save_interval_steps: int = 1
@@ -351,7 +351,6 @@ class CheckpointManagerOptions:
   save_decision_policy: Optional[
       save_decision_policy_lib.SaveDecisionPolicy
   ] = None
-  prevent_write_metrics: bool = False
 
   def __post_init__(self):
     step_name_format_single_host_load_and_broadcast = (
@@ -911,6 +910,7 @@ def _configure_checkpointer_legacy_init(
           f'Invalid type for `checkpointers`. Found {checkpointers}.'
       )
 
+    # if options.best_fn:
     item_handlers[METRIC_ITEM_NAME] = self._metrics_handler
     if options.async_options is None:
       options.async_options = (
@@ -1366,10 +1366,8 @@ def save(
           'Some provided items have prohibited reserved names:'
           f' {args_dict.keys()}. Reserved names: {RESERVED_ITEM_NAMES}.'
       )
-    if (
-        metrics is not None and self._track_best
-    ) and not self._options.prevent_write_metrics:
-      args_dict[METRIC_ITEM_NAME] = args_lib.JsonSave(metrics)
+    if metrics is not None and self._track_best:
+      args_dict['metrics'] = args_lib.JsonSave(metrics)
     args = args_lib.Composite(**args_dict)
 
     save_directory = self._get_write_step_directory(step, self.directory)
@@ -1640,18 +1638,20 @@ def item_metadata(
 
   # TODO(b/370812224): Deprecate in favor of StepMetadata.metrics
   def metrics(self, step: int) -> Optional[PyTree]:
-    try:
-      # Use handler directly, since this happens in a background thread and
-      # barriers cannot be used. This usage pattern is not
-      # recommended in other contexts.
-      metrics = self._metrics_handler.restore(
-          self._get_read_step_directory(step, self.directory)
-          / METRIC_ITEM_NAME
-      )
-      return metrics
-    except FileNotFoundError as e:
-      logging.warning('Missing metrics for step %d', step)
-      logging.error(e)
+    if self._track_best:
+      try:
+        # Use handler directly, since this happens in a background thread and
+        # barriers cannot be used. This usage pattern is not
+        # recommended in other contexts.
+        return self._metrics_handler.restore(
+            self._get_read_step_directory(step, self.directory)
+            / METRIC_ITEM_NAME
+        )
+      except FileNotFoundError as e:
+        logging.warning('Missing metrics for step %d', step)
+        logging.error(e)
+        return None
+    else:
       return None
 
   @property

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/compatibility.py

@@ -114,17 +114,11 @@ def typestr(cls) -> str:
     """A unique identifier for the CheckpointHandler type."""
     ...
 
-  def __repr__(self):
-    return f'CompatibilityCheckpointHandler({handler_types.typestr(type(self._handler))})'
-
 
 @dataclasses.dataclass
 class Args(checkpoint_args.CheckpointArgs):
   checkpointable: Any
 
-  def __repr__(self):
-    return f'CompatibilityArgs({type(self.checkpointable)})'
-
 
 def get_compatibility_handler(
     handler: handler_types.CheckpointableHandler,

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/global_registration.py

@@ -31,9 +31,5 @@
 
 registration.global_registry().add(json_handler.JsonHandler)
 
-registration.global_registry().add(
-    json_handler.MetricsHandler,
-    format_utils.METRICS_CHECKPOINTABLE_KEY,
-)
 
 registration.global_registry().add(pytree_handler.PyTreeHandler)

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/json_handler.py

@@ -33,8 +33,10 @@
 class JsonHandler(CheckpointableHandler[JsonType, None]):
   """An implementation of `CheckpointableHandler` for Json."""
 
-  def __init__(self, filename: str = 'metadata'):
-    self._filename = filename
+  def __init__(
+      self,
+  ):
+    self._filename = 'metadata'
 
   async def _background_save(
       self,
@@ -85,9 +87,3 @@ def is_handleable(self, checkpointable: Any) -> bool:
 
   def is_abstract_handleable(self, abstract_checkpointable: Any) -> bool | None:
     return None
-
-
-class MetricsHandler(JsonHandler):
-
-  def __init__(self):
-    super().__init__(filename='metrics')

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/pytree_handler.py

@@ -278,15 +278,7 @@ def is_abstract_handleable(self, abstract_checkpointable: Any) -> bool:
 
 @contextlib.contextmanager
 def pytree_handler_context():
-  """Creates a local context for PyTree handling.
-
-  `PYTREE_CHECKPOINTABLE_KEY` is explicitly registered linking to
-  `PyTreeHandler`. Note that all globally-registered handlers are still included
-  as backup options. Other options from the parent context are carried through.
-
-  Yields:
-    A new context.
-  """
+  """Creates a local context where only `PyTreeHandler` is registered."""
   # TODO(b/398310070): Verify behavior with nested Contexts.
   checkpointables_options = options_lib.CheckpointablesOptions(
       registry=registration.local_registry(include_global_registry=True).add(

checkpoint/orbax/checkpoint/experimental/v1/_src/metadata/loading.py

@@ -118,10 +118,6 @@ def checkpointables_metadata(
   )
   metadata = checkpointer.metadata(path)
   item_metadata = {k: v for k, v in metadata.item_metadata.items()}
-  # Exclude `metrics` if present. This is relevant only for
-  # `training.Checkpointer`, and is separately added to the
-  # `training.CheckpointMetadata` object.
-  item_metadata.pop('metrics', None)
   return CheckpointMetadata[dict[str, Any]](
       metadata=item_metadata,
       init_timestamp_nsecs=metadata.init_timestamp_nsecs,

checkpoint/orbax/checkpoint/experimental/v1/_src/metadata/types.py

@@ -107,18 +107,3 @@ def from_metadata(
       cls, metadata: CheckpointableMetadataT
   ) -> CheckpointMetadata[CheckpointableMetadataT]:
     return cls(metadata=metadata)
-
-  def _properties_strings(self) -> dict[str, str]:
-    return {
-        'metadata': str(self.metadata),
-        'init_timestamp_nsecs': str(self.init_timestamp_nsecs),
-        'commit_timestamp_nsecs': str(self.commit_timestamp_nsecs),
-        'custom_metadata': str(self.custom_metadata),
-    }
-
-  def __repr__(self):
-    s = 'CheckpointMetadata('
-    for k, v in self._properties_strings().items():
-      s += f' {k} = {v}, '
-    s += ')'
-    return s

checkpoint/orbax/checkpoint/experimental/v1/_src/path/format_utils.py

@@ -28,10 +28,7 @@
 
 PYTREE_CHECKPOINTABLE_KEY = 'pytree'
 
-METRICS_CHECKPOINTABLE_KEY = 'metrics'
-
 RESERVED_CHECKPOINTABLE_KEYS = frozenset({
-    METRICS_CHECKPOINTABLE_KEY,
 })
 
 

checkpoint/orbax/checkpoint/experimental/v1/_src/saving/saving.py

@@ -239,7 +239,6 @@ def save_checkpointables_async(
 def get_v0_checkpointer_and_args(
     checkpointables: dict[str, Any],
     *,
-    metrics: tree_types.JsonType | None = None,
     context: context_lib.Context,
 ) -> tuple[
     async_checkpointer.AsyncCheckpointer,
@@ -253,9 +252,6 @@ def get_v0_checkpointer_and_args(
     raise ValueError(
         f'Provided reserved checkpointable keys: {provided_reserved_keys}.'
     )
-  # Global registration ties metrics key to JsonHandler.
-  if metrics:
-    checkpointables[format_utils.METRICS_CHECKPOINTABLE_KEY] = metrics
 
 
   handlers = {

checkpoint/orbax/checkpoint/experimental/v1/_src/testing/BUILD

@@ -6,7 +6,6 @@ py_library(
     deps = [
         ":array_utils",
         ":handler_utils",
-        ":tree_utils",
         "//checkpoint/orbax/checkpoint:test_utils",
         "//checkpoint/orbax/checkpoint/_src/path:atomicity",
         "//checkpoint/orbax/checkpoint/_src/serialization",
@@ -45,16 +44,6 @@ py_library(
     deps = ["//orbax/checkpoint/experimental/v1/_src/path:types"],
 )
 
-py_library(
-    name = "tree_utils",
-    srcs = ["tree_utils.py"],
-    deps = [
-        "//orbax/checkpoint/experimental/v1/_src/path:format_utils",
-        "//orbax/checkpoint/experimental/v1/_src/path:types",
-        "//orbax/checkpoint/experimental/v1/_src/synchronization:multihost",
-    ],
-)
-
 py_library(
     name = "v0v1_compatibility_save_load_test_base",
     srcs = ["v0v1_compatibility_save_load_test_base.py"],

checkpoint/orbax/checkpoint/experimental/v1/_src/testing/save_load_test_base.py

@@ -42,7 +42,6 @@
 from orbax.checkpoint.experimental.v1._src.synchronization import multihost
 from orbax.checkpoint.experimental.v1._src.testing import array_utils as array_test_utils
 from orbax.checkpoint.experimental.v1._src.testing import handler_utils
-from orbax.checkpoint.experimental.v1._src.testing import tree_utils as tree_test_utils
 from orbax.checkpoint.experimental.v1._src.tree import types as tree_types
 
 
@@ -72,6 +71,33 @@ async def _sleep_and_create_paths(*args, **kwargs):
   return await _original_create_paths(*args, **kwargs)
 
 
+def get_d_files_mtimes(path: Path) -> list[int]:
+  mtimes = []
+  matching_dirs = list(path.parent.glob(f'{path.name}*'))
+  if not matching_dirs:
+    # Temp path not created yet.
+    return []
+  assert (
+      len(matching_dirs) == 1
+  ), f'Expected exactly one matching directory, got {matching_dirs}.'
+  tmpdir = matching_dirs[0]
+  matching_pytree_dirs = list(tmpdir.glob(f'{PYTREE_CHECKPOINTABLE_KEY}*'))
+  if not matching_pytree_dirs:
+    # Temp path not created yet.
+    return []
+  assert len(matching_pytree_dirs) == 1, (
+      'Expected exactly one matching pytree directory, got'
+      f' {matching_pytree_dirs}.'
+  )
+  pytree_dir = matching_pytree_dirs[0]
+  for idx in range(multihost.process_count()):
+    d_path = pytree_dir / f'ocdbt.process_{idx}' / 'd'
+    if not d_path.exists():
+      continue
+    mtimes.extend([f.stat().mtime for f in d_path.iterdir()])
+  return mtimes
+
+
 class SaveLoadTestBase:
 
   class Test(parameterized.TestCase):
@@ -121,6 +147,11 @@ def test_load_default(self, use_async):
       test_utils.assert_tree_equal(self, self.pytree, loaded)
 
     def test_save_pytree_async(self):
+      def is_save_complete(directory):
+        return (
+            directory / PYTREE_CHECKPOINTABLE_KEY / 'manifest.ocdbt'
+        ).exists()
+
       start_serialize = threading.Event()
       original_serialize = serialization.async_serialize_from_host
 
@@ -136,21 +167,15 @@ def mock_serialize(*args, **kwargs):
       )
 
       response = ocp.save_pytree_async(self.directory, self.pytree)
-      initial_d_files_mtimes = tree_test_utils.get_d_files_mtimes(
-          self.directory
-      )
-      self.assertFalse(
-          tree_test_utils.is_pytree_checkpoint_complete(self.directory)
-      )
+      initial_d_files_mtimes = get_d_files_mtimes(self.directory)
+      self.assertFalse(is_save_complete(self.directory))
       start_serialize.set()
 
       response.result()
-      final_d_files_mtimes = tree_test_utils.get_d_files_mtimes(self.directory)
+      final_d_files_mtimes = get_d_files_mtimes(self.directory)
       self.assertNotEmpty(final_d_files_mtimes)
       self.assertNotEqual(initial_d_files_mtimes, final_d_files_mtimes)
-      self.assertTrue(
-          tree_test_utils.is_pytree_checkpoint_complete(self.directory)
-      )
+      self.assertTrue(is_save_complete(self.directory))
 
       restored = ocp.load_pytree(
           self.directory,