Add testing to verify backwards compatibility when missing _CHECKPOINT_METADATA or saving pytree directly to checkpoint.

PiperOrigin-RevId: 840391322

Author: angel-core

checkpoint/CHANGELOG.md

@@ -9,10 +9,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [0.11.30] - 2025-11-26
 
+### Added
+
+- Added testing to verify backwards compatibility when a checkpoint is missing
+`_CHECKPOINT_METADATA` or saving pytree directly as a checkpoint with no
+checkpointable.
+
 ### Fixed
 
 - Roll back earlier change altering metadata format, which was observed to cause
 breakages.
+- Fix `CompositeHandler` extraction logic for handler type strings when missing `_CHECKPOINT_METADATA`
+
+### Changed
+
+- Split v0 checkpoint format/layout logic out from `OrbaxLayout` and into
+seperate `V0Layout`.
 
 ## [0.11.29] - 2025-11-25
 

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

@@ -25,7 +25,9 @@
 from orbax.checkpoint._src.metadata import step_metadata_serialization
 from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.path import async_path
+from orbax.checkpoint._src.path import format_utils
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
+from orbax.checkpoint.experimental.v1._src.handlers import pytree_handler
 from orbax.checkpoint.experimental.v1._src.handlers import registration
 from orbax.checkpoint.experimental.v1._src.handlers import types as handler_types
 import orbax.checkpoint.experimental.v1._src.handlers.global_registration  # pylint: disable=unused-import
@@ -262,7 +264,7 @@ def get_handlers_for_load(
 
   def _get_saved_handler_typestrs(
       self, directory: path_types.Path
-  ) -> dict[str, str]:
+  ) -> dict[str, str | None]:
     """Reads from the checkpoint metadata to get saved handler typestrs."""
     step_metadata_file_path = checkpoint_metadata.step_metadata_file_path(
         directory
@@ -291,28 +293,40 @@ def _get_saved_handler_typestrs(
         directory,
     )
 
-    saved_handler_typestrs: dict[str, str] = {}
+    # The following generically handles the case where the checkpoint is missing
+    # the checkpoint metadata file. This is a fallback for older v0 checkpoints.
+
+    # We check each subdirectory treating it as a checkpointable.
+    # The order of presedence for mapping checkpointable names to handlers is:
+    # 1. A pytree metadata file, indicating a pytree handler
+    # 2. A checkpointable with a handler registered in the handler registry.
+    # 3. If neither, we skip it and treat it as garbage since we don't know
+    #    how to handle it.
+
+    saved_handler_typestrs: dict[str, str | None] = {}
     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:
+      if not checkpointable_path.is_dir():
         continue
-      saved_metadata = step_metadata_serialization.deserialize(
-          serialized_metadata
+      checkpointable_name = checkpointable_path.name
+
+      pytree_metadata_path = (
+          checkpointable_path / format_utils.PYTREE_METADATA_FILE
       )
-      if isinstance(saved_metadata.item_handlers, dict):
+
+      if pytree_metadata_path.exists():
+        saved_handler_typestrs[checkpointable_name] = handler_types.typestr(
+            pytree_handler.PyTreeHandler
+        )
+      elif self._handler_registry.has(checkpointable_name):
+        # If the handler is registered we can assume it will be found in
+        # resolve_handler_for_load.
+        saved_handler_typestrs[checkpointable_name] = None
+      else:
         raise ValueError(
-            f'Path at {directory} contains subdirectories:'
-            f' {_subdirs(directory)}, which are expected to'
-            ' match the keys given by the _CHECKPOINT_METADATA file:'
-            f' {saved_metadata.item_handlers}. If you intended to load a pytree'
-            ' checkpoint from the given path, then please consider using'
-            ' `loading.load_pytree(..., checkpointable_name=None)` instead.'
-            f' {_V0_ERROR_MESSAGE}'
+            'Cannot determine handler for checkpointable'
+            f" '{checkpointable_name}' in directory {directory}. The top-level"
+            " '_CHECKPOINT_METADATA' is missing, and this item does not have a"
+            " '_METADATA' file to be loaded as a PyTree, nor is it registered"
+            " in the handler registry."
         )
-      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/layout/orbax_layout.py

@@ -25,7 +25,6 @@
 from orbax.checkpoint.experimental.v1._src.handlers import registration
 from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout
 from orbax.checkpoint.experimental.v1._src.loading import v0_compatibility
-from orbax.checkpoint.experimental.v1._src.metadata import serialization as metadata_serialization
 from orbax.checkpoint.experimental.v1._src.metadata import types as metadata_types
 from orbax.checkpoint.experimental.v1._src.path import types as path_types
 
@@ -148,24 +147,23 @@ async def _validate_pytree(self, checkpointable_name: str | None):
     """Validates a checkpoint path written by `ocp.save_pytree`.
 
     Args:
-      checkpointable_name: The name of the checkpointable to load. A
-        subdirectory with this name must exist in `directory`. If None then
-        `directory` is expected to contain the checkpoint directly. Defaults to
-        `pytree`.
+      checkpointable_name: The name of the checkpointable to load. For Orbax V1,
+        a subdirectory with this name must exist in `directory`.
 
     Raises:
       FileNotFoundError: If the path does not exist, or if `pytree` is not found
         in the directory
       ValueError: If the PyTree checkpoint is malformed.
     """
-    pytree_dir = (
-        self.path
-        if checkpointable_name is None
-        else self.path / checkpointable_name
-    )
-    if checkpointable_name is not None and not await async_path.exists(
-        pytree_dir
-    ):
+    if checkpointable_name is None:
+      raise ValueError(
+          "A V1 checkpoint was saved and user is attempting to load it,"
+          " treating it as a V0 Orbax checkpoint saved directly from"
+          f" {self.path}, this is not a characteristic of V1 saved checkpoints"
+      )
+
+    pytree_dir = self.path / checkpointable_name
+    if not await async_path.exists(pytree_dir):
       subdirs = [
           d.name
           for d in await _subpaths(self.path)
@@ -180,17 +178,13 @@ async def _validate_pytree(self, checkpointable_name: str | None):
           " using"
           " `ocp.load_checkpointables()`."
       )
-    if not await _has_pytree_metadata_file(pytree_dir):
-      # TODO(niketkb): Add following details to the error message:
-      # 1. we should check other available subdirectories and see if any of them
-      #   look like PyTree checkpoints, and instruct the user to consider
-      #   whether they meant to specify any of those.
-      # 2. we need to check the directory - if it contains PyTree files, suggest
-      #   loading with checkpointable_name=None
+
+    if not (pytree_dir / PYTREE_METADATA_FILE).exists():
       raise FileNotFoundError(
           f"Checkpoint path {self.path} does not contain a PyTree metadata"
           " file."
       )
+
     if not await _has_tensorstore_data_files(pytree_dir):
       logging.warning(
           "TensorStore data files not found in checkpoint path %s. This may be"
@@ -200,17 +194,13 @@ async def _validate_pytree(self, checkpointable_name: str | None):
       )
 
   async def _validate(self):
-    """Validates a checkpoint directory.
+    """Validates a checkpoint directory to be a V1 Orbax checkpoint.
 
-    Must be:
+    Must fulfill all of the following:
     - Existing
-    - A directory.
-    - Not a temporary path.
-    - OR
-      - Has orbax.checkpoint indicator file.
-      - OR
-        - Has _CHECKPOINT_METADATA file.
-        - A subdirectory has _METADATA file (PyTree checkpoint).
+    - A directory
+    - Not a temporary path
+    - Has orbax.checkpoint indicator file
 
     Raises:
       FileNotFoundError: If the path does not exist.
@@ -235,33 +225,10 @@ async def _validate(self):
     if ORBAX_CHECKPOINT_INDICATOR_FILE in [p.name for p in subpaths]:
       return
 
-    # Path points to a single step checkpoint with valid metadata.
-    if await async_path.exists(
-        metadata_serialization.checkpoint_metadata_file_path(self.path)
-    ):
-      return
-
-    # The path itself points to a PyTree checkpointable.
-    if await async_path.exists(self.path / PYTREE_METADATA_FILE):
-      return
-    # The path points to a directory containing at least one PyTree
-    # checkpointable.
-    for subpath in subpaths:
-      if await async_path.is_dir(subpath) and await async_path.exists(
-          subpath / PYTREE_METADATA_FILE
-      ):
-        return
-
     raise FileNotFoundError(
         f"Checkpoint path {self.path} could not be identified as a valid Orbax"
-        " checkpoint. The path must conform to one of the following"
-        " conditions:\n  - Contain the indicator file"
-        f" {ORBAX_CHECKPOINT_INDICATOR_FILE}. This should be true of all"
-        " checkpoints saved with the Orbax V1 API. If not present, the"
-        " checkpoint may have been saved with the V0 API.\n  - Contain the"
-        " _CHECKPOINT_METADATA file.\n  - Point directly to a PyTree"
-        " checkpointable (contain _METADATA file).\n  - Contain a subdirectory"
-        " which is a PyTree checkpointable (contain _METADATA file).\n"
+        " V1 checkpoint. It is missing the indicator file"
+        f" '{ORBAX_CHECKPOINT_INDICATOR_FILE}'."
     )
 
   async def validate(self):

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

@@ -18,9 +18,6 @@
 from etils import epath
 import numpy as np
 from orbax.checkpoint import test_utils
-from orbax.checkpoint._src.checkpointers import checkpointer
-from orbax.checkpoint._src.handlers import composite_checkpoint_handler
-from orbax.checkpoint._src.handlers import standard_checkpoint_handler
 from orbax.checkpoint._src.metadata import value as value_metadata
 from orbax.checkpoint._src.path import async_path
 from orbax.checkpoint.experimental.v1._src.handlers import composite_handler
@@ -102,7 +99,8 @@ async def test_validate_no_indicator_file(self):
         / composite_handler.ORBAX_CHECKPOINT_INDICATOR_FILE
     )
     indicator_path.rmtree()  # Remove the indicator file
-    await layout.validate()
+    with self.assertRaises(InvalidLayoutError):
+      await layout.validate()
 
   async def test_validate_no_metadata_file(self):
     layout = OrbaxLayout(self.orbax_path / '0')
@@ -176,93 +174,6 @@ async def test_metadata(self):
     self.assertGreater(result_metadata.commit_timestamp_nsecs, 0)
 
 
-class V0ValidationTest(
-    unittest.IsolatedAsyncioTestCase, parameterized.TestCase
-):
-
-  def setUp(self):
-    super().setUp()
-    self.directory = epath.Path(self.create_tempdir().full_path) / 'ckpt'
-    self.pytree, _ = array_test_utils.create_numpy_pytree()
-    # Save a checkpoint with a checkpointable name, `state`.
-    ckptr = checkpointer.Checkpointer(
-        composite_checkpoint_handler.CompositeCheckpointHandler()
-    )
-    ckptr.save(
-        self.directory,
-        composite_checkpoint_handler.CompositeArgs(
-            state=standard_checkpoint_handler.StandardSaveArgs(self.pytree)
-        ),
-    )
-
-  async def test_nonexistent_path(self):
-    with self.assertRaises(FileNotFoundError):
-      await OrbaxLayout(self.directory / 'foo')._validate()
-
-  async def test_not_a_directory(self):
-    await async_path.write_text(self.directory / 'foo', 'foo')
-    with self.assertRaises(NotADirectoryError):
-      await OrbaxLayout(self.directory / 'foo')._validate()
-
-  @parameterized.product(checkpointable_name=['state', None])
-  async def test_no_checkpoint_metadata(self, checkpointable_name: str | None):
-    directory = (
-        self.directory / checkpointable_name
-        if checkpointable_name is not None
-        else self.directory
-    )
-    await _unlink_checkpoint_metadata(directory)
-
-    await OrbaxLayout(directory)._validate()
-    if checkpointable_name is None:
-      await OrbaxLayout(directory)._validate_pytree('state')
-    else:
-      await OrbaxLayout(directory)._validate_pytree(None)
-
-  async def test_deleted_pytree(self):
-    directory = self.directory
-    (directory / 'state').rmtree()
-
-    await OrbaxLayout(directory)._validate()
-    with self.assertRaises(FileNotFoundError):
-      await OrbaxLayout(directory)._validate_pytree('state')
-
-  async def test_missing_checkpointable_matching_name(self):
-    with self.assertRaises(FileNotFoundError):
-      await OrbaxLayout(self.directory)._validate_pytree('foo')
-
-  @parameterized.product(checkpointable_name=['state', None])
-  async def test_no_pytree_metadata(self, checkpointable_name: str | None):
-    directory = (
-        self.directory / checkpointable_name
-        if checkpointable_name is not None
-        else self.directory
-    )
-    await _unlink_pytree_metadata(directory)
-
-    if checkpointable_name is None:
-      # Passes because we still have the checkpoint metadata.
-      await OrbaxLayout(directory)._validate()
-      with self.assertRaises(FileNotFoundError):
-        await OrbaxLayout(directory)._validate_pytree('state')
-    else:
-      with self.assertRaises(FileNotFoundError):
-        await OrbaxLayout(directory)._validate()
-        await OrbaxLayout(directory)._validate_pytree(None)
-
-  @parameterized.product(checkpointable_name=['state', None])
-  async def test_valid_pytree(self, checkpointable_name: str | None):
-    directory = (
-        self.directory / checkpointable_name
-        if checkpointable_name is not None
-        else self.directory
-    )
-    if checkpointable_name is None:
-      await OrbaxLayout(directory)._validate_pytree('state')
-    else:
-      await OrbaxLayout(directory)._validate_pytree(None)
-
-
 class V1ValidationTest(
     unittest.IsolatedAsyncioTestCase, parameterized.TestCase
 ):
@@ -290,7 +201,8 @@ async def test_missing_indicator_file(self, checkpointable_name: str | None):
         else self.directory
     )
     await _unlink_indicator(directory)
-    await OrbaxLayout(directory)._validate()
+    with self.assertRaises(FileNotFoundError):
+      await OrbaxLayout(directory)._validate()
 
   async def test_deleted_pytree(self):
     directory = self.directory
@@ -317,10 +229,12 @@ async def test_no_pytree_metadata(self, checkpointable_name: str | None):
 
     with self.assertRaises(FileNotFoundError):
       await OrbaxLayout(directory)._validate()
-    with self.assertRaises(FileNotFoundError):
-      if checkpointable_name is None:
+
+    if checkpointable_name is None:
+      with self.assertRaises(FileNotFoundError):
         await OrbaxLayout(directory)._validate_pytree('pytree')
-      else:
+    else:
+      with self.assertRaises(ValueError):
         await OrbaxLayout(directory)._validate_pytree(None)