Internal change

PiperOrigin-RevId: 876100881

Author: mxberlot

checkpoint/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.py

@@ -44,6 +44,7 @@
 from orbax.checkpoint._src.metadata import array_metadata_store as array_metadata_store_lib
 from orbax.checkpoint._src.metadata import empty_values
 from orbax.checkpoint._src.metadata import tree as tree_metadata
+from orbax.checkpoint._src.path import types as path_types
 from orbax.checkpoint._src.serialization import limits
 from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
 from orbax.checkpoint._src.serialization import type_handler_registry as handler_registry
@@ -470,7 +471,9 @@ def _concurrent_bytes(
     return concurrent_gb * 10**9
 
 
-class PyTreeCheckpointHandler(async_checkpoint_handler.AsyncCheckpointHandler):
+class PyTreeCheckpointHandler(
+    async_checkpoint_handler.DeferredPathAsyncCheckpointHandler
+):
   """A CheckpointHandler implementation for any PyTree structure.
 
   See JAX documentation for more information on what consistutes a "PyTree".
@@ -608,7 +611,7 @@ def __init__(
 
   async def async_save(
       self,
-      directory: epath.Path,
+      directory: epath.Path | path_types.PathAwaitingCreation,
       item: Optional[PyTree] = None,
       save_args: Optional[PyTreeSaveArgs] = None,
       args: Optional[PyTreeSaveArgs] = None,

checkpoint/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler_test.py

@@ -61,6 +61,7 @@
 from orbax.checkpoint._src.serialization import type_handlers
 from orbax.checkpoint._src.testing import multiprocess_test
 from orbax.checkpoint._src.tree import utils as tree_utils
+from orbax.checkpoint.google.path import tfhub_atomicity
 
 
 PyTree = Any
@@ -2948,6 +2949,51 @@ def test_partial_restore_with_omission_unexpected_keys(
       )
       test_utils.assert_tree_equal(self, expected, restored)
 
+  async def test_save_with_deferred_path(self):
+    """Tests that async_save works with deferred paths."""
+    deferred_path = tfhub_atomicity.DeferredPath()
+    save_dir = self.directory / 'deferred_path_ckpt'
+    await_creation_called = False
+    original_await = tfhub_atomicity.DeferredPath.await_creation
+
+    async def mock_await_creation(dp_self):
+      """Sets the path only once await_creation is called.
+
+      This ensures the path is not resolved before the handler awaits it, fully
+      exercising the deferred path resolution contract.
+
+      Args:
+        dp_self: The DeferredPath instance.
+
+      Returns:
+        The result of the original await_creation method.
+      """
+      nonlocal await_creation_called
+      if not dp_self._future_path.done():
+        save_dir.mkdir(parents=True, exist_ok=True)
+        dp_self.set_path(save_dir)
+      await_creation_called = True
+      return await original_await(dp_self)
+
+    with self.ocdbt_checkpoint_handler(use_ocdbt=False) as handler:
+      with mock.patch.object(
+          tfhub_atomicity.DeferredPath,
+          'await_creation',
+          mock_await_creation,
+      ):
+        await handler.async_save(
+            deferred_path,
+            args=PyTreeSaveArgs(self.pytree),
+        )
+
+      self.assertTrue(await_creation_called)
+      self.validate_save(
+          save_dir,
+          self.pytree,
+          handler,
+          restore_args=self.restore_args,
+      )
+
 
 if __name__ == '__main__':
   multiprocess_test.main()