#v1 Refactor OrbaxLayout metadata retrieval and v1 Checkpointer to use internal free function logic, removing dependency on v0 checkpointer + refactor to read checkpoint metadata externally from get_handlers_from_load.

PiperOrigin-RevId: 874781759

Author: angel-core

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

@@ -20,11 +20,9 @@
 
 from absl import logging
 from orbax.checkpoint._src.metadata import step_metadata_serialization
-from orbax.checkpoint._src.path import async_path
 from orbax.checkpoint.experimental.v1._src.handlers import registration
 from orbax.checkpoint.experimental.v1._src.handlers import types as handler_types
 import orbax.checkpoint.experimental.v1._src.handlers.global_registration  # pylint: disable=unused-import
-from orbax.checkpoint.experimental.v1._src.metadata import serialization as metadata_serialization
 from orbax.checkpoint.experimental.v1._src.path import types as path_types
 
 InternalCheckpointMetadata = (
@@ -49,21 +47,6 @@ def _subdirs(directory: path_types.Path, *, limit: int = 3) -> list[str]:
 )
 
 
-async def read_checkpoint_metadata(
-    directory: path_types.Path,
-) -> InternalCheckpointMetadata:
-  """Returns the step metadata for a given path, normalized for V1."""
-
-  serialized_metadata = (
-      await metadata_serialization.read(
-          metadata_serialization.checkpoint_metadata_file_path(directory)
-      )
-      or {}
-  )
-
-  return InternalCheckpointMetadata.deserialize(serialized_metadata)
-
-
 def get_handlers_for_save(
     handler_registry: registration.CheckpointableHandlerRegistry,
     checkpointables: dict[str, Any],
@@ -81,10 +64,11 @@ async def get_handlers_for_load(
     directory: path_types.Path,
     handler_registry: registration.CheckpointableHandlerRegistry,
     abstract_checkpointables: dict[str, Any],
+    checkpoint_metadata: InternalCheckpointMetadata,
 ) -> dict[str, handler_types.CheckpointableHandler]:
   """Returns a mapping from checkpointable name to handler."""
   existing_checkpointable_names_to_handler_typestrs = (
-      await _get_saved_handler_typestrs(directory)
+      await _get_saved_handler_typestrs(directory, checkpoint_metadata)
   )
   abstract_checkpointables = abstract_checkpointables or {
       name: None for name in existing_checkpointable_names_to_handler_typestrs
@@ -111,13 +95,10 @@ async def get_handlers_for_load(
 
 async def _get_saved_handler_typestrs(
     directory: path_types.Path,
+    checkpoint_metadata: InternalCheckpointMetadata,
 ) -> dict[str, str]:
   """Reads from the checkpoint metadata to get saved handler typestrs."""
-  checkpoint_metadata_file_path = (
-      metadata_serialization.checkpoint_metadata_file_path(directory)
-  )
-  if await async_path.exists(checkpoint_metadata_file_path):
-    checkpoint_metadata = await read_checkpoint_metadata(directory)
+  if checkpoint_metadata.item_handlers:
     if isinstance(checkpoint_metadata.item_handlers, dict):
       return checkpoint_metadata.item_handlers  # found step level metadata.
     raise ValueError(
@@ -131,34 +112,8 @@ async def _get_saved_handler_typestrs(
     )
 
   logging.warning(
-      'Given dir does not contain checkpoint metadata file: %s. Trying to get'
-      ' saved handlers from checkpoint metadata in each of the checkpointable'
-      ' subdirectory.',
+      'Given dir does not contain checkpoint metadata file: %s. No handler'
+      ' typestrs found.',
       directory,
   )
-
-  # TODO(b/475265289): Currently, we rely solely on CHECKPOINT_METADATA to
-  # find available checkpointables, ignoring valid subdirectories. We
-  # should update the composite handler to validate subdirectories to
-  # check if any either represents a valid pytree checkpointable or has a
-  # name that is registered in the handler registry.
-  saved_handler_typestrs: dict[str, str] = {}
-  for checkpointable_path in await async_path.iterdir(directory):
-    if not await async_path.is_dir(checkpointable_path):
-      continue
-    checkpoint_metadata = await read_checkpoint_metadata(checkpointable_path)
-    if isinstance(checkpoint_metadata.item_handlers, dict):
-      raise ValueError(
-          f'Path at {directory} contains subdirectories:'
-          f' {_subdirs(directory)}, which are expected to'
-          ' match the keys given by the _CHECKPOINT_METADATA file:'
-          f' {checkpoint_metadata.item_handlers}. If you intended to load a'
-          ' pytree checkpoint from the given path, then please consider using'
-          ' `loading.load_pytree(..., checkpointable_name=None)` instead.'
-          f' {_V0_ERROR_MESSAGE}'
-      )
-    item_handlers = checkpoint_metadata.item_handlers
-    if item_handlers is not None:
-      checkpointable_name = checkpointable_path.name
-      saved_handler_typestrs[checkpointable_name] = item_handlers
-  return saved_handler_typestrs
+  return {}

checkpoint/orbax/checkpoint/experimental/v1/_src/layout/orbax_layout.py

@@ -20,14 +20,15 @@
 
 from absl import logging
 from orbax.checkpoint._src import asyncio_utils
+from orbax.checkpoint._src.metadata import step_metadata_serialization
 from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.path import async_path
 from orbax.checkpoint._src.path import temporary_paths
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
 from orbax.checkpoint.experimental.v1._src.handlers import registration
 from orbax.checkpoint.experimental.v1._src.handlers import resolution as handler_resolution
 from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout
-from orbax.checkpoint.experimental.v1._src.loading import v0_compatibility
+from orbax.checkpoint.experimental.v1._src.metadata import serialization as metadata_serialization
 from orbax.checkpoint.experimental.v1._src.metadata import types as metadata_types
 from orbax.checkpoint.experimental.v1._src.path import types as path_types
 from orbax.checkpoint.experimental.v1._src.tree import types as tree_types
@@ -41,6 +42,9 @@ class CheckpointVersion(enum.Enum):
 InvalidLayoutError = checkpoint_layout.InvalidLayoutError
 Path = path_types.Path
 CheckpointLayout = checkpoint_layout.CheckpointLayout
+InternalCheckpointMetadata = (
+    step_metadata_serialization.InternalCheckpointMetadata
+)
 
 PYTREE_METADATA_FILE = "_METADATA"
 ORBAX_CHECKPOINT_INDICATOR_FILE = "orbax.checkpoint"
@@ -149,6 +153,19 @@ async def _create_orbax_identifier_file(
     )
 
 
+async def _read_checkpoint_metadata(
+    directory: path_types.Path,
+) -> InternalCheckpointMetadata:
+  """Returns the step metadata for a given path."""
+  serialized_metadata = (
+      await metadata_serialization.read(
+          metadata_serialization.checkpoint_metadata_file_path(directory)
+      )
+      or {}
+  )
+  return InternalCheckpointMetadata.deserialize(serialized_metadata)
+
+
 class OrbaxLayout(CheckpointLayout):
   """OrbaxLayout.
 
@@ -172,23 +189,47 @@ async def metadata(
       self, path: Path
   ) -> metadata_types.CheckpointMetadata[dict[str, Any]]:
     """Returns the metadata describing the Orbax checkpoint."""
-    # Uses the v0 checkpointer to get v0 StepMetadata
-    checkpointer, _ = v0_compatibility.get_v0_checkpointer_and_args(
-        path, None, context=context_lib.get_context()
+    checkpoint_metadata = await _read_checkpoint_metadata(
+        path
     )
-    step_metadata = checkpointer.metadata(path)
+    handlers_for_load = await handler_resolution.get_handlers_for_load(
+        path, self._handler_registry, {}, checkpoint_metadata
+    )
+    existing_checkpointable_names = await _existing_checkpointable_names(path)
+    abstract_checkpointables = {
+        name: None
+        for name in handlers_for_load.keys()
+        if name in existing_checkpointable_names
+    }
+    if any(
+        name not in existing_checkpointable_names
+        for name in abstract_checkpointables.keys()
+    ):
+      raise KeyError(
+          "Inferred checkpointables from metadata:"
+          f" {abstract_checkpointables.keys()} for loading were not found in"
+          " the checkpoint. Available checkpointables:"
+          f" {existing_checkpointable_names}"
+      )
 
-    item_metadata = {k: v for k, v in step_metadata.item_metadata.items()}
+    # Default to none for all existing checkpointable names, for
+    # subdirectories that we are unable to find a handler for and load.
+    item_metadata = {name: None for name in existing_checkpointable_names}
+    for checkpointable_name in abstract_checkpointables.keys():
+      handler = handlers_for_load[checkpointable_name]
+      item_metadata[checkpointable_name] = await handler.metadata(
+          path / checkpointable_name
+      )
     # 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 metadata_types.CheckpointMetadata[dict[str, Any]](
         metadata=item_metadata,
-        init_timestamp_nsecs=step_metadata.init_timestamp_nsecs,
-        commit_timestamp_nsecs=step_metadata.commit_timestamp_nsecs,
-        custom_metadata=step_metadata.custom_metadata,
+        init_timestamp_nsecs=checkpoint_metadata.init_timestamp_nsecs,
+        commit_timestamp_nsecs=checkpoint_metadata.commit_timestamp_nsecs,
+        custom_metadata=checkpoint_metadata.custom_metadata,
     )
 
   async def _validate_pytree(self, path: Path, checkpointable_name: str | None):
@@ -356,8 +397,14 @@ async def load_checkpointables(
       the checkpoint.
     """
     abstract_checkpointables = abstract_checkpointables or {}
+    checkpoint_metadata = await _read_checkpoint_metadata(
+        path
+    )
     handlers_for_load = await handler_resolution.get_handlers_for_load(
-        path, self._handler_registry, abstract_checkpointables
+        path,
+        self._handler_registry,
+        abstract_checkpointables,
+        checkpoint_metadata,
     )
     existing_checkpointable_names = await _existing_checkpointable_names(path)
     if not abstract_checkpointables:

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

@@ -25,7 +25,7 @@
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
 import orbax.checkpoint.experimental.v1._src.handlers.global_registration  # pylint: disable=unused-import
 from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout
-from orbax.checkpoint.experimental.v1._src.loading import v0_compatibility as v0_loading_utils
+from orbax.checkpoint.experimental.v1._src.loading import loading
 from orbax.checkpoint.experimental.v1._src.metadata import loading as metadata_loading
 from orbax.checkpoint.experimental.v1._src.metadata import types as metadata_types
 from orbax.checkpoint.experimental.v1._src.path import step as path_step_lib
@@ -444,17 +444,10 @@ def load_checkpointables(
   ) -> dict[str, Any]:
     """Loads a set of checkpointables at the given step."""
     step = self._resolve_existing_checkpoint(step).step
-    checkpointer, args = v0_loading_utils.get_v0_checkpointer_and_args(
+    return  loading.load_checkpointables(
         self.directory / self._step_name_format.build_name(step),
         abstract_checkpointables,
-        context=context_lib.get_context(),
     )
-    self._manager._checkpointer = checkpointer  # pylint: disable=protected-access
-    restored = self._manager.restore(
-        step,
-        args=args,
-    )
-    return {k: v for k, v in zip(restored.keys(), restored.values())}
 
   def load_pytree_async(
       self,