Create v1 ocp.pytree_metadata backwards compatibility tests against static v0 and v1 checkpoints.

PiperOrigin-RevId: 895570820

Author: angel-core

checkpoint/orbax/checkpoint/experimental/v1/_src/layout/orbax_v0_layout.py

@@ -112,16 +112,18 @@ async def metadata(
     Returns:
       The metadata describing the Orbax checkpoint.
     """
-    checkpoint_metadata = await orbax_layout.read_checkpoint_metadata(
-        path
-    )
-    # Delegate to OrbaxLayout if the checkpoint is a composite checkpoint.
-    if checkpoint_metadata and isinstance(
-        checkpoint_metadata.item_handlers, dict
+    checkpoint_metadata = await orbax_layout.read_checkpoint_metadata(path)
+    if (
+        checkpoint_metadata
+        and isinstance(checkpoint_metadata.item_handlers, str)
+        or await orbax_layout.has_pytree_metadata_file(path)
     ):
-      return await self._orbax_layout.metadata(path)
-    # Otherwise, load the metadata as a PyTree checkpoint.
-    return await self._load_pytree_metadata(path, checkpoint_metadata)
+      return await self._load_pytree_metadata(path, checkpoint_metadata)
+    # Delegate to OrbaxLayout if the checkpoint is a composite checkpoint.
+    # If there is no checkpoint metadata, we assume it is a composite
+    # checkpoint, and even if it is a direct pytree checkpoint it will load as
+    # a composite checkpoint as there is no metadata to indicate otherwise.
+    return await self._orbax_layout.metadata(path)
 
   async def _validate(self, path: Path) -> None:
     """Validates a V0 checkpoint directory.

checkpoint/orbax/checkpoint/experimental/v1/_src/testing/compatibility/load_pytree_compatibility_test.py

@@ -141,14 +141,13 @@ def _determine_expected_outcome(
       is_pytree: bool,
       handler_registered: bool,
       pytree_registered: bool,
-  ) -> Tuple[bool, Type[Exception] | None, str | None]:
+  ) -> Tuple[Type[Exception] | None, str | None]:
     """Encapsulates the complex boolean logic to determine load behavior."""
     # LAYOUT VALIDATION BEHAVIOR:
     if version == 'v1':
       # V1 strictly requires that checkpoint metadata is present.
       if not metadata_present:
         return (
-            True,
             InvalidLayoutError,
             (
                 r'Could not recognize the checkpoint at .* as a valid Orbax'
@@ -158,7 +157,6 @@ def _determine_expected_outcome(
       # V1 does not support loading a top-level pytree, redirects to V0.
       if checkpointable_name is None:
         return (
-            True,
             ValueError,
             (
                 r'Failed to interpret path .* as a .* Orbax PyTree'
@@ -173,7 +171,6 @@ def _determine_expected_outcome(
         or (not is_direct_checkpoint and checkpointable_name is None)
     ):
       return (
-          True,
           InvalidLayoutError,
           (
               r'Failed to interpret path .* as a .* Orbax PyTree'
@@ -197,15 +194,14 @@ def _determine_expected_outcome(
 
     if not can_resolve:
       return (
-          True,
           registration.NoEntryError,
           (
               r'Could not resolve a handler for .* and no \'pytree\' handler'
               r' found in .*'
           ),
       )
 
-    return False, None, None
+    return None, None
 
   @parameterized.product(
       version=['v0', 'v1'],
@@ -244,7 +240,7 @@ def test_load_pytree_compatibility(
         pytree_registered,
     )
 
-    should_fail, error_type, expected_error_msg = (
+    error_type, expected_error_msg = (
         self._determine_expected_outcome(
             version,
             checkpointable_name,
@@ -267,20 +263,20 @@ def test_load_pytree_compatibility(
             registry=registry
         )
     ):
-      if should_fail:
-        with self.assertRaisesRegex(error_type, expected_error_msg):
-          ocp.load_pytree(
-              path,
-              checkpointable_name=checkpointable_name,
-              abstract_pytree=actual_abstract_pytree,
-          )
-      else:
+      if error_type is None:
         loaded = ocp.load_pytree(
             path,
             checkpointable_name=checkpointable_name,
             abstract_pytree=actual_abstract_pytree,
         )
         test_utils.assert_tree_equal(self, loaded, self.expected_state)
+      else:
+        with self.assertRaisesRegex(error_type, expected_error_msg):
+          ocp.load_pytree(
+              path,
+              checkpointable_name=checkpointable_name,
+              abstract_pytree=actual_abstract_pytree,
+          )
 
   @parameterized.product(
       version=['v0', 'v1'],