#v1 Refactor logic for handler resolution and loading checkpointables + additional fallback capabilities for non-standard checkpoint formats.

PiperOrigin-RevId: 869904948

Author: angel-core

checkpoint/CHANGELOG.md

@@ -26,6 +26,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 `load_checkpointables()` each with their own dedicated loading logic
 - Refactor v0 Pytree validation and metadata resolution and add `OrbaxV0Layout`
 tests
+- Refactor `CompositeHandler` logic into the orbax layout objects and handler
+resolution utility, deprecating and deleting the `CompositeHandler` class.
 
 ## [0.11.32] - 2026-01-20
 

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

@@ -444,7 +444,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.
 
@@ -492,15 +492,16 @@ 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]

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

@@ -284,6 +284,17 @@ def test_resolve_handler_for_load_checkpointable(self):
           handler_typestr='unused',
       )
 
+  def test_resolve_handler_for_load_no_handler_typestr(self):
+    local_registry = registration.local_registry()
+    local_registry.add(handler_utils.FooHandler)
+    with self.assertRaises(registration.NoEntryError):
+      registration.resolve_handler_for_load(
+          local_registry,
+          handler_utils.Foo(1, 'hi'),
+          name='unregistered_name',
+          handler_typestr=None,
+      )
+
 
 if __name__ == '__main__':
   absltest.main()