Introduce registered handler to alternate handler typestrs mapping to v1 API.

PiperOrigin-RevId: 874792153

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/global_registration.py

@@ -20,7 +20,7 @@
 given checkpointable will be used.
 """
 
-from typing import Type
+from typing import Sequence, Type
 
 from orbax.checkpoint.experimental.v1._src.handlers import json_handler
 from orbax.checkpoint.experimental.v1._src.handlers import proto_handler
@@ -34,23 +34,45 @@
 def _try_register_handler(
     handler_type: Type[handler_types.CheckpointableHandler],
     name: str | None = None,
+    secondary_typestrs: Sequence[str] | None = None,
 ):
+  """Tries to register handler globally with name and recognized typestrs."""
   try:
-    registration.global_registry().add(handler_type, name)
+    registration.global_registry().add(
+        handler_type,
+        name,
+        secondary_typestrs=secondary_typestrs,
+    )
   except registration.AlreadyExistsError:
     pass
 
 
-_try_register_handler(proto_handler.ProtoHandler)
-_try_register_handler(json_handler.JsonHandler)
+_try_register_handler(
+    proto_handler.ProtoHandler,
+    secondary_typestrs=[
+        'orbax.checkpoint._src.handlers.proto_checkpoint_handler.ProtoCheckpointHandler',
+    ],
+)
+_try_register_handler(
+    json_handler.JsonHandler,
+    secondary_typestrs=[
+        'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler',
+    ],
+)
 _try_register_handler(
     stateful_checkpointable_handler.StatefulCheckpointableHandler
 )
 _try_register_handler(
     json_handler.MetricsHandler,
     checkpoint_layout.METRICS_CHECKPOINTABLE_KEY,
 )
-_try_register_handler(pytree_handler.PyTreeHandler)
+_try_register_handler(
+    pytree_handler.PyTreeHandler,
+    secondary_typestrs=[
+        'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler',
+        'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler',
+    ],
+)
 _try_register_handler(
     pytree_handler.PyTreeHandler, checkpoint_layout.PYTREE_CHECKPOINTABLE_KEY
 )

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

@@ -35,25 +35,31 @@
   registry.add(BarHandler)
   # Scope this handler specifically to checkpointables named 'baz'.
   registry.add(BazHandler, 'baz')
+  # Scope this handler to legacy typestr 'OldBazHandlerTypestr'. Secondary
+  # typestrs provide a way to map legacy identifiers to a new handler class.
+  registry.add(BazHandler, secondary_typestrs=['OldBazHandlerTypestr'])
 
   checkpointables_options = ocp.options.CheckpointablesOptions(
       registry=registry
   )
   with ocp.Context(checkpointables_options=checkpointables_options):
     ocp.save_checkpointables(...)
 
-If a registered handler is scoped to a specific name (e.g.
-`registry.add(BazHandler, 'baz')`), then this handler will always be
-prioritized for saving or loading the checkpointable with that name, even if
-the handler is not capable of saving/loading the checkpointable.
-
-In the most common case, where a handler is not scoped to a specific name,
-a given checkpointable (or abstract_checkpointable) will be resolved to a
-handler returning True for `is_handleable` (or `is_abstract_handleable`),
-respectively. If multiple handlers are usable, the first usable handler will be
-returned. When loading, the handler type used for saving will be recorded in
-the metadata, and will be used to resolve the handler, if a corresponding
-handler is present in the registry.
+Handler resolution for saving/loading follows this logic:
+
+    1. If a registered handler is scoped to a specific name
+      (e.g. `registry.add(BazHandler, 'baz')`), then this handler will always
+      be prioritized for saving or loading the checkpointable with that name,
+      even if the handler is not capable of saving/loading the checkpointable.
+    2. In the absence of an explicit name match, the registry filters for
+      handlers returning `True` for `is_handleable` (during save) or
+      `is_abstract_handleable` (during load).
+    3. [Pertains to loading only] The handler type used for saving will be
+      recorded in the metadata, and will be used to resolve the handler, if a
+      corresponding handler is present in the registry. If not, scan the
+      secondary typestrs of registered handlers for a match.
+    4. If no metadata match is found (or during saving), the most recently
+       registered capable handler is returned.
 """
 
 from __future__ import annotations
@@ -73,7 +79,13 @@ def add_all(
 ) -> CheckpointableHandlerRegistry:
   """Adds all entries from `other_registry` to `registry`."""
   for handler, checkpointable in other_registry.get_all_entries():
-    registry.add(handler, checkpointable)
+    registry.add(
+        handler,
+        checkpointable,
+        secondary_typestrs=other_registry.get_secondary_typestrs(
+            handler
+        ),
+    )
   return registry
 
 
@@ -87,6 +99,7 @@ def add(
       self,
       handler_type: Type[CheckpointableHandler],
       checkpointable: str | None = None,
+      secondary_typestrs: Sequence[str] | None = None,
   ) -> CheckpointableHandlerRegistry:
     """Adds an entry to the registry."""
     ...
@@ -110,6 +123,13 @@ def get_all_entries(
   ) -> Sequence[RegistryEntry]:
     ...
 
+  def get_secondary_typestrs(
+      self,
+      handler_type: Type[CheckpointableHandler],
+  ) -> Sequence[str]:
+    """Returns all secondary typestrs associated with the given handler type."""
+    ...
+
 
 class AlreadyExistsError(ValueError):
   """Raised when an entry already exists in the registry."""
@@ -126,6 +146,9 @@ def __init__(
       self, other_registry: CheckpointableHandlerRegistry | None = None
   ):
     self._registry: list[RegistryEntry] = []
+    self._secondary_typestrs: dict[
+        Type[CheckpointableHandler], Sequence[str]
+    ] = {}
 
     # Initialize the registry with entries from other registry.
     if other_registry:
@@ -135,14 +158,25 @@ def add(
       self,
       handler_type: Type[CheckpointableHandler],
       checkpointable: str | None = None,
+      secondary_typestrs: Sequence[str] | None = None,
   ) -> CheckpointableHandlerRegistry:
     """Adds an entry to the registry.
 
+    Adds a primary handler_type to the registry with an optional checkpointable
+    name and an optional sequence of secondary typestrs that can be used to
+    identify the handler.
+
+    Note: We only guarantee unique handler type entries in the registry and do
+    not explicitly prevent a primary handler type from being registered and its
+    typestr being used as a secondary_typestr for itself or another handler.
+
     Args:
       handler_type: The handler type.
       checkpointable: The checkpointable name. If not-None, the registered
         handler will be scoped to that specific name. Otherwise, the handler
         will be available for any checkpointable name.
+      secondary_typestrs: A sequence of alternate typestrs that serve as
+        secondary identifiers for the handler.
 
     Returns:
       The registry itself.
@@ -170,6 +204,8 @@ def add(
           f'Handler type {handler_type} already exists in the registry.'
       )
     self._registry.append((handler_type, checkpointable))
+    if secondary_typestrs is not None:
+      self._secondary_typestrs[handler_type] = secondary_typestrs
     return self
 
   def get(
@@ -220,6 +256,13 @@ def get_all_entries(
     """Returns all entries in the registry."""
     return self._registry
 
+  def get_secondary_typestrs(
+      self,
+      handler_type: Type[CheckpointableHandler],
+  ) -> Sequence[str]:
+    """Returns all secondary typestrs associated with the given handler type."""
+    return self._secondary_typestrs.get(handler_type, [])
+
   def __repr__(self):
     return f'_DefaultCheckpointableHandlerRegistry({self.get_all_entries()})'
 
@@ -237,6 +280,7 @@ def add(
       self,
       handler_type: Type[CheckpointableHandler],
       checkpointable: str | None = None,
+      secondary_typestrs: Sequence[str] | None = None,
   ) -> CheckpointableHandlerRegistry:
     raise NotImplementedError('Adding not implemented for read-only registry.')
 
@@ -257,6 +301,12 @@ def get_all_entries(
   ) -> Sequence[RegistryEntry]:
     return self._registry.get_all_entries()
 
+  def get_secondary_typestrs(
+      self,
+      handler_type: Type[CheckpointableHandler],
+  ) -> Sequence[str]:
+    return self._registry.get_secondary_typestrs(handler_type)
+
   def __repr__(self):
     return f'ReadOnlyCheckpointableHandlerRegistry({self.get_all_entries()})'
 
@@ -303,6 +353,8 @@ def local_registry(
 
 def register_handler(
     cls: CheckpointableHandlerType,
+    *,
+    secondary_typestrs: Sequence[str] | None = None,
 ) -> CheckpointableHandlerType:
   """Registers a :py:class:`~.v1.handlers.CheckpointableHandler` globally.
 
@@ -322,16 +374,20 @@ class FooHandler(ocp.handlers.CheckpointableHandler[Foo, AbstractFoo]):
 
   Args:
     cls: The handler class.
+    secondary_typestrs: A sequence of alternate handler typestrs that serve as
+      secondary identifiers for the handler.
 
   Returns:
     The handler class.
   """
-  _GLOBAL_REGISTRY.add(cls)
+  _GLOBAL_REGISTRY.add(
+      cls, secondary_typestrs=secondary_typestrs
+  )
   return cls
 
 
 def _construct_handler_instance(
-    name: str,
+    name: str | None,
     handler_type: Type[CheckpointableHandler],
 ) -> CheckpointableHandler:
   """Attempts to default-construct a handler type if possible."""
@@ -392,6 +448,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 +501,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 +510,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.
 
@@ -456,8 +522,9 @@ def resolve_handler_for_load(
     4. If multiple handlers are usable, return the handler with the matching
       typestr. If no matching typestr is found, then the handler used for saving
       may not be available now.
-    4. Return the *last* usable handler. This allows us to resolve the most
-       recently-registered handler.
+    5. Return the *last* usable handler. This allows us to resolve the most
+      recently-registered handler, unless abstract_checkpointable is None, in
+      which case raise a NoEntryError.
 
   Raises:
     NoEntryError: If no compatible
@@ -471,7 +538,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 +561,30 @@ 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:
+  if handler_typestr:
+    if handler_typestr in possible_handler_typestrs:
+      idx = possible_handler_typestrs.index(handler_typestr)
+      return possible_handlers[idx]
+    # Attempt to find a handler with a matching secondary typestr.
+    for i in reversed(range(len(possible_handlers))):
+      if handler_typestr in registry.get_secondary_typestrs(
+          type(possible_handlers[i])
+      ):
+        return possible_handlers[i]
     logging.warning(
-        'No handler found for typestr %s. The checkpointable may be restored'
-        ' with different handler logic than was used for saving.',
+        'No handler found for typestr %s (or its converted form). 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 abstract_checkpointable:
+    # Prefer the last handler in the absence of any other information.
+    return possible_handlers[-1]
+
+  raise NoEntryError(
+      f'No entry for checkpointable={name} in the registry, using'
+      f' handler_typestr={handler_typestr} and'
+      f' abstract_checkpointable={abstract_checkpointable}. Registry contents:'
+      f' {registry.get_all_entries()}'
+  )