Fix resolution logic to correctly find handler based on checkpointable_name when missing abstract_checkpointable or checkpoint metadata.

PiperOrigin-RevId: 874801701

Author: angel-core

checkpoint/CHANGELOG.md

@@ -41,6 +41,9 @@ devices anyway.
 `load_checkpointables()` each with their own dedicated loading logic
 - Refactor v0 Pytree validation and metadata resolution and add `OrbaxV0Layout`
 tests
+- Refactor logic for handler resolution and loading checkpointables for
+`OrbaxLayout` and `OrbaxV0Layout`, adding additional fallback capabilities for
+non-standard checkpoint formats.
 
 ### Fixed
 

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

@@ -392,6 +392,16 @@ def _get_possible_handlers(
   return possible_handlers
 
 
+def get_registered_handler_by_name(
+    registry: CheckpointableHandlerRegistry,
+    name: str,
+) -> CheckpointableHandler | None:
+  """Returns the handler for the given name if registered."""
+  if registry.has(name):
+    return _construct_handler_instance(name, registry.get(name))
+  return None
+
+
 def resolve_handler_for_save(
     registry: CheckpointableHandlerRegistry,
     checkpointable: Any,
@@ -435,7 +445,7 @@ def is_handleable_fn(handler: CheckpointableHandler, ckpt: Any) -> bool:
       registry, is_handleable_fn, checkpointable, name
   )
 
-  # Prefer the first handler in the absence of any other information.
+  # Prefer the last handler in the absence of any other information.
   return possible_handlers[-1]
 
 
@@ -444,7 +454,7 @@ def resolve_handler_for_load(
     abstract_checkpointable: Any | None,
     *,
     name: str,
-    handler_typestr: str,
+    handler_typestr: str | None = None,
 ) -> CheckpointableHandler:
   """Resolves a :py:class:`~.v1.handlers.CheckpointableHandler` for loading.
 
@@ -471,7 +481,9 @@ def resolve_handler_for_load(
     abstract_checkpointable: An abstract checkpointable to resolve.
     name: The name of the checkpointable.
     handler_typestr: A :py:class:`~.v1.handlers.CheckpointableHandler` typestr
-      to guide resolution.
+      to guide resolution. We allow a None value for handler_typestr as its
+      possible to find the last registered handler given a specified
+      abstract_checkpointable.
 
   Returns:
     A :py:class:`~.v1.handlers.CheckpointableHandler` instance.
@@ -492,15 +504,23 @@ def is_handleable_fn(
       handler_types.typestr(type(handler)) for handler in possible_handlers
   ]
 
-  try:
-    idx = possible_handler_typestrs.index(handler_typestr)
-    return possible_handlers[idx]
-  except ValueError:
-    logging.warning(
-        'No handler found for typestr %s. The checkpointable may be restored'
-        ' with different handler logic than was used for saving.',
-        handler_typestr,
-    )
-
-  # Prefer the first handler in the absence of any other information.
-  return possible_handlers[-1]
+  if handler_typestr:
+    try:
+      idx = possible_handler_typestrs.index(handler_typestr)
+      return possible_handlers[idx]
+    except ValueError:
+      logging.warning(
+          'No handler found for typestr %s. The checkpointable may be restored'
+          ' with different handler logic than was used for saving.',
+          handler_typestr,
+      )
+  # For non-None, specified abstract_checkpointable, prefer the last handler in
+  # the absence of any other information.
+  if abstract_checkpointable:
+    return possible_handlers[-1]
+
+  raise NoEntryError(
+      f'No entry for checkpointable={name} found in the registry: {registry},'
+      f' using abstract_checkpointable={abstract_checkpointable} and'
+      f' handler_typestr={handler_typestr}'
+  )

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

@@ -273,15 +273,33 @@ def test_resolve_handler_for_load_no_matching_name(self):
           handler_typestr='unused',
       )
 
-  def test_resolve_handler_for_load_checkpointable(self):
+  def test_resolve_handler_for_load_no_handler_typestr(self):
     local_registry = registration.local_registry()
     local_registry.add(handler_utils.FooHandler)
+    resolved = registration.resolve_handler_for_load(
+        local_registry,
+        handler_utils.AbstractFoo(),
+        name='dummy_unregistered_nameame',
+        handler_typestr=None,
+    )
+    self.assertIsInstance(resolved, handler_utils.FooHandler)
+
     with self.assertRaises(registration.NoEntryError):
       registration.resolve_handler_for_load(
           local_registry,
-          handler_utils.Foo(1, 'hi'),
-          name='foo',
-          handler_typestr='unused',
+          handler_utils.AbstractBar(),
+          name='unregistered_name',
+          handler_typestr=None,
+      )
+
+  def test_resolve_handler_for_load_no_checkpointable_or_handler_typestr(self):
+    local_registry = registration.local_registry()
+    with self.assertRaises(registration.NoEntryError):
+      registration.resolve_handler_for_load(
+          local_registry,
+          None,
+          name='unregistered_name',
+          handler_typestr=None,
       )
 
 

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

@@ -15,38 +15,19 @@
 """Logic for resolving handlers for saving and loading."""
 from __future__ import annotations
 
-import itertools
 from typing import Any
 
 from absl import logging
 from orbax.checkpoint._src.metadata import step_metadata_serialization
 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.path import types as path_types
 
 InternalCheckpointMetadata = (
     step_metadata_serialization.InternalCheckpointMetadata
 )
 
 
-def _subdirs(directory: path_types.Path, *, limit: int = 3) -> list[str]:
-  return list(
-      itertools.islice(
-          (subdir.name for subdir in directory.iterdir() if subdir.is_dir()),
-          limit,
-      )
-  )
-
-
-_V0_ERROR_MESSAGE = (
-    'If your checkpoint was saved with the Orbax V0 API, please follow the'
-    ' instructions at'
-    ' https://orbax.readthedocs.io/en/latest/guides/checkpoint/v1/orbax_v0_to_v1_migration.html'
-    ' to load it with the Orbax V1 API.'
-)
-
-
 def get_handlers_for_save(
     handler_registry: registration.CheckpointableHandlerRegistry,
     checkpointables: dict[str, Any],
@@ -60,60 +41,162 @@ def get_handlers_for_save(
   }
 
 
+def _resolve_single_handler_for_load(
+    checkpointable_name: str,
+    handler_registry: registration.CheckpointableHandlerRegistry,
+    abstract_checkpointable: Any,
+    metadata_handler_typestr: str | None,
+) -> handler_types.CheckpointableHandler:
+  """Logic to resolve a checkpointable's loading handler.
+
+  1. registration.resolve_handler_for_load performs handler discovery based on
+  abstract_checkpointable type and handler_typestr.
+  2. If this fails or if abstract_checkpointable and handler_typestr are not
+  available, we try to resolve using the default pytree handler if registered.
+
+  Args:
+    checkpointable_name: The checkpointable name to resolve the handler for.
+    handler_registry: The handler registry to use for resolution.
+    abstract_checkpointable: The abstract checkpointable to load.
+    metadata_handler_typestr: The handler typestr from the checkpoint metadata.
+
+  Returns:
+    The handler for the checkpointable.
+
+  Raises:
+    registration.NoEntryError: If no handler is resolved and 'pytree' name is
+    not registered.
+  """
+  # 1. Resolve the checkpointable's handler using handler discovery.
+  try:
+    return registration.resolve_handler_for_load(
+        handler_registry,
+        abstract_checkpointable,
+        name=checkpointable_name,
+        handler_typestr=metadata_handler_typestr,
+    )
+  except registration.NoEntryError as e:
+    logging.warning(
+        "Failed to resolve handler for checkpointable: '%s'. Attempting to"
+        " load using pytree handler. Error: %s",
+        checkpointable_name,
+        e,
+    )
+
+  # 2. If no handler is resolved yet, try to resolve using the default
+  # pytree handler.
+  pytree_handler = registration.get_registered_handler_by_name(
+      handler_registry, "pytree"
+  )
+  if not pytree_handler:
+    raise registration.NoEntryError(
+        f"Could not resolve a handler for '{checkpointable_name}' and no"
+        f" 'pytree' handler found in {handler_registry})."
+        "Please inspect the checkpoint contents via"
+        " `loading.checkpointables_metadata`. You may need to provide an"
+        " abstract_checkpointable or register a missing handler for this name"
+        " or for 'pytree' name which is used as a fallback."
+    )
+  return pytree_handler
+
+
 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, checkpoint_metadata)
-  )
-  abstract_checkpointables = abstract_checkpointables or {
-      name: None for name in existing_checkpointable_names_to_handler_typestrs
-  }
-
-  loadable_checkpointable_names_to_handlers = {}
-  for name, abstract_checkpointable in abstract_checkpointables.items():
-    if name not in existing_checkpointable_names_to_handler_typestrs:
-      raise KeyError(
-          f'Checkpointable "{name}" was not found in the checkpoint.'
-          ' Available names:'
-          f' {existing_checkpointable_names_to_handler_typestrs.keys()}'
-      )
-    handler_typestr = existing_checkpointable_names_to_handler_typestrs[name]
-    handler = registration.resolve_handler_for_load(
+  """Returns a mapping from checkpointable name to handler.
+
+  Gathers and returns a mapping from checkpointable name to handler by
+  checking the following in order:
+
+  1. Check for handler_typestr in checkpoint metadata item_handlers using
+  checkpointable_name as key.
+  2. Find the handler for each checkpointable using
+  _resolve_single_handler_for_load.
+  3. If no handler is resolved for a checkpointable, raise a NoEntryError.
+
+  Args:
+    handler_registry: The handler registry to use for resolution.
+    abstract_checkpointables: The abstract checkpointables to load.
+    checkpoint_metadata: InternalCheckpointMetadata to read handler_typestr(s)
+      from.
+
+  Returns:
+    A mapping from checkpointable name to handler.
+
+  Raises:
+    registration.NoEntryError: If no handler is resolved.
+  """
+  handlers_for_load: dict[str, handler_types.CheckpointableHandler] = {}
+  for (
+      checkpointable_name,
+      abstract_checkpointable,
+  ) in abstract_checkpointables.items():
+    metadata_handler_typestr = _get_saved_handler_typestr(
+        checkpointable_name, checkpoint_metadata
+    )
+    handlers_for_load[checkpointable_name] = _resolve_single_handler_for_load(
+        checkpointable_name,
         handler_registry,
         abstract_checkpointable,
-        name=name,
-        handler_typestr=handler_typestr,
+        metadata_handler_typestr,
     )
-    loadable_checkpointable_names_to_handlers[name] = handler
-  return loadable_checkpointable_names_to_handlers
+  return handlers_for_load
 
 
-async def _get_saved_handler_typestrs(
-    directory: path_types.Path,
+async def get_handler_for_load_direct_pytree(
+    checkpointable_name: str,
+    handler_registry: registration.CheckpointableHandlerRegistry,
+    abstract_checkpointable: Any,
     checkpoint_metadata: InternalCheckpointMetadata,
-) -> dict[str, str]:
+) -> handler_types.CheckpointableHandler:
+  """Returns a handler for direct load of a pytree checkpoint.
+
+  1. Check for checkpointable_name in checkpoint metadata item_handlers.
+  2. resolve_handler_for_load performs handler discovery based on
+  abstract_checkpointable type and handler_typestr.
+  2. Find the handler for each checkpointable using
+  _resolve_single_handler_for_load.
+  3. If no handler is resolved for a checkpointable, raise a NoEntryError.
+
+  Args:
+    checkpointable_name: The checkpointable name to resolve the handler for.
+    handler_registry: The handler registry to use for resolution.
+    abstract_checkpointable: The abstract checkpointable to load.
+    checkpoint_metadata: InternalCheckpointMetadata to read handler_typestr
+      from.
+
+  Returns:
+    The handler for direct load of a pytree checkpoint.
+  """
+  metadata_handler_typestr = _get_saved_handler_typestr_direct_pytree(
+      checkpoint_metadata
+  )
+  return _resolve_single_handler_for_load(
+      checkpointable_name,
+      handler_registry,
+      abstract_checkpointable,
+      metadata_handler_typestr,
+  )
+
+
+def _get_saved_handler_typestr(
+    checkpointable_name: str,
+    checkpoint_metadata: InternalCheckpointMetadata,
+) -> str | None:
   """Reads from the checkpoint metadata to get saved handler typestrs."""
-  if checkpoint_metadata.item_handlers:
-    if isinstance(checkpoint_metadata.item_handlers, dict):
-      return checkpoint_metadata.item_handlers  # found step level metadata.
-    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}'
-    )
+  if isinstance(checkpoint_metadata.item_handlers, dict) and (
+      checkpointable_name in checkpoint_metadata.item_handlers
+  ):
+    return checkpoint_metadata.item_handlers[checkpointable_name]
+  return None
 
-  logging.warning(
-      'Given dir does not contain checkpoint metadata file: %s. No handler'
-      ' typestrs found.',
-      directory,
-  )
-  return {}
+
+def _get_saved_handler_typestr_direct_pytree(
+    checkpoint_metadata: InternalCheckpointMetadata,
+) -> str | None:
+  """Reads from the checkpoint metadata to get saved handler typestrs."""
+  if isinstance(checkpoint_metadata.item_handlers, str):
+    return checkpoint_metadata.item_handlers
+  return None