Refine resolution logic to ensure valid handler matching for abstract checkpointables.

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,24 @@ 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,
-    )
+  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,
+      )
 
-  # 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()}'
+  )

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

@@ -197,18 +197,31 @@ def test_resolve_handler_for_load(
 
   def test_resolve_handler_for_load_resolution_order(self):
 
-    class HandlerOne(handler_utils.DictHandler):
-      pass
+    class HandlerOne(handler_utils.BazHandler):
+      def is_abstract_handleable(
+          self, abstract_checkpointable: handler_utils.AbstractBaz
+      ) -> bool:
+        return isinstance(abstract_checkpointable, handler_utils.AbstractBaz)
 
-    class HandlerTwo(handler_utils.DictHandler):
-      pass
+    class HandlerTwo(handler_utils.BazHandler):
+      def is_abstract_handleable(
+          self, abstract_checkpointable: handler_utils.AbstractBaz
+      ) -> bool:
+        return isinstance(abstract_checkpointable, handler_utils.AbstractBaz)
 
     handlers_to_register = [HandlerOne, HandlerTwo]
 
     with self.subTest('globally_registered'):
+      with self.assertRaises(registration.NoEntryError):
+        registration.resolve_handler_for_load(
+            registration.local_registry(),
+            None,
+            name='checkpointable_name',
+            handler_typestr='unknown_class',
+        )
       resolved_handler = registration.resolve_handler_for_load(
           registration.local_registry(),
-          None,
+          handler_utils.AbstractBaz(),
           name='checkpointable_name',
           handler_typestr='unknown_class',
       )
@@ -219,9 +232,16 @@ class HandlerTwo(handler_utils.DictHandler):
       )
       for handler in handlers_to_register:
         local_registry.add(handler)
+      with self.assertRaises(registration.NoEntryError):
+        registration.resolve_handler_for_load(
+            local_registry,
+            None,
+            name='checkpointable_name',
+            handler_typestr='unknown_class',
+        )
       resolved_handler = registration.resolve_handler_for_load(
           local_registry,
-          None,
+          handler_utils.AbstractBaz(),
           name='checkpointable_name',
           handler_typestr='unknown_class',
       )
@@ -243,9 +263,16 @@ class HandlerTwo(handler_utils.DictHandler):
       )
       for handler in reversed(handlers_to_register):
         local_registry.add(handler)
+      with self.assertRaises(registration.NoEntryError):
+        registration.resolve_handler_for_load(
+            local_registry,
+            None,
+            name='checkpointable_name',
+            handler_typestr='unknown_class',
+        )
       resolved_handler = registration.resolve_handler_for_load(
           local_registry,
-          None,
+          handler_utils.AbstractBaz(),
           name='checkpointable_name',
           handler_typestr='unknown_class',
       )
@@ -273,15 +300,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='unregistered_name',
+        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_no_metadata(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,
       )