#v1 Enable v1 users to load v0 checkpoints without using contexts.

checkpoint/CHANGELOG.md

@@ -24,6 +24,7 @@ would want them to be preferred.
 - #v1 Add checkpointables support for `training.Checkpointer`.
 - `PartsOf` structure which holds a PyTree whose leaf nodes may be missing.
 - #v1 Add compatibility tests for save-by-v0-load-by-v1 and also fix code.
+- #v1 Enable v1 users to load v0 checkpoints without using contexts.
 
 ## [0.11.12] - 2025-04-09
 

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

@@ -217,21 +217,20 @@ def _get_saved_handler_typestrs(
         ' the checkpointable subdirectories.',
         directory,
     )
+
     saved_handler_typestrs = {}
-    for _, checkpointable_name in self._handler_registry.get_all_entries():
-      if not checkpointable_name:
-        continue
-      checkpointable_path = directory / checkpointable_name
-      if not checkpointable_path.exists() or not checkpointable_path.is_dir():
-        continue
+    for checkpointable_path in directory.iterdir():
       serialized_metadata = self._metadata_store.read(
           checkpoint_metadata.step_metadata_file_path(checkpointable_path)
       )
+      if serialized_metadata is None:
+        continue
       saved_metadata = step_metadata_serialization.deserialize(
-          serialized_metadata or {}
+          serialized_metadata
       )
       assert not isinstance(saved_metadata.item_handlers, dict)
       item_handlers = saved_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

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

@@ -217,11 +217,11 @@ def get_v0_checkpointer_and_args(
       name: handler_compatibility.get_compatibility_handler(handler)
       for name, handler in handlers.items()
   }
-  handler_registry = (
+  legacy_handler_registry = (
       legacy_handler_registration.DefaultCheckpointHandlerRegistry()
   )
   for name, handler in compatibility_handlers.items():
-    handler_registry.add(name, handler_compatibility.Args, handler)
+    legacy_handler_registry.add(name, handler_compatibility.Args, handler)
   composite_options = composite_checkpoint_handler.CompositeOptions(
       async_options=context.async_options.v0(),
       file_options=context.file_options.v0(),
@@ -230,7 +230,7 @@ def get_v0_checkpointer_and_args(
   )
   ckptr = async_checkpointer.AsyncCheckpointer(
       composite_checkpoint_handler.CompositeCheckpointHandler(
-          handler_registry=handler_registry,
+          handler_registry=legacy_handler_registry,
           composite_options=composite_options,
       ),
       async_options=context.async_options.v0(),

checkpoint/orbax/checkpoint/experimental/v1/_src/testing/v0v1_compatibility_save_load_test_base.py

@@ -173,35 +173,20 @@ def test_load_v0_checkpoint_with_v1_load_checkpointables(
               self.directory, abstract_checkpointables
           )
       else:
-        checkpointables_options = (
-            ocp.options.CheckpointablesOptions.create_with_handlers(
-                **{checkpointable_name: ocp.handlers.PyTreeHandler}
-            )
-        )
-        with ocp.Context(checkpointables_options=checkpointables_options):
-          loaded = ocp.load_checkpointables(
-              self.directory, abstract_checkpointables
-          )
-          test_utils.assert_tree_equal(
-              self, self.pytree, loaded[checkpointable_name]
+        with self.subTest('with_context'):
+          checkpointables_options = (
+              ocp.options.CheckpointablesOptions.create_with_handlers(
+                  **{checkpointable_name: ocp.handlers.PyTreeHandler}
+              )
           )
-        if to_checkpointable_subdir:
-          if abstract_checkpointables:
-            with self.assertRaisesRegex(
-                KeyError,
-                f'Checkpointable "{checkpointable_name}" was not found in the'
-                ' checkpoint',
-            ):
-              ocp.load_checkpointables(self.directory, abstract_checkpointables)
-          else:
-            with self.assertRaisesRegex(
-                ValueError,
-                f'Item "{checkpointable_name}" was found in the checkpoint, but'
-                ' could not be restored',
-            ):
-              # TODO(b/401585537): Support in v1 and decouple from v0.
-              ocp.load_checkpointables(self.directory)
-        else:
+          with ocp.Context(checkpointables_options=checkpointables_options):
+            loaded = ocp.load_checkpointables(
+                self.directory, abstract_checkpointables
+            )
+            test_utils.assert_tree_equal(
+                self, self.pytree, loaded[checkpointable_name]
+            )
+        with self.subTest('without_context'):
           loaded = ocp.load_checkpointables(
               self.directory, abstract_checkpointables
           )