Internal change

PiperOrigin-RevId: 872295675

Author: mxberlot

checkpoint/orbax/checkpoint/_src/metadata/tree.py

@@ -534,6 +534,7 @@ def as_custom_metadata(
               restore_type, self.pytree_metadata_options
           )
         else:
+          assert isinstance(param_info.parent_dir, epath.Path)
           flat_metadatas[keypath] = value_metadata.Metadata(
               name=param_info.name, directory=param_info.parent_dir
           )

checkpoint/orbax/checkpoint/_src/path/types.py

@@ -67,3 +67,24 @@ async def await_creation(self) -> Path:
       The path that was created.
     """
     ...
+
+
+async def await_and_resolve_as_posix(
+    path: Path | PathAwaitingCreation,
+) -> str:
+  """Resolves a path to a POSIX string, awaiting creation if necessary.
+
+  For `PathAwaitingCreation` inputs, this blocks until the path has been
+  created before resolving it to a string. For regular ``Path`` inputs, it
+  returns the POSIX string immediately.
+
+  Args:
+    path: An `epath.Path` or a `PathAwaitingCreation`.
+
+  Returns:
+    The POSIX string representation of the resolved path.
+  """
+  if isinstance(path, PathAwaitingCreation):
+    resolved = await path.await_creation()
+    return resolved.as_posix()
+  return path.as_posix()

checkpoint/orbax/checkpoint/_src/path/types_test.py

@@ -0,0 +1,78 @@
+# Copyright 2026 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+
+from absl.testing import absltest
+from etils import epath
+from orbax.checkpoint._src.path import types as path_types
+from orbax.checkpoint.google.path import tfhub_atomicity
+
+
+class AwaitAndResolveAsPosixTest(absltest.TestCase):
+
+  def test_resolves_epath_immediately(self):
+    async def _test():
+      path = epath.Path('/tmp/test/dir')
+      result = await path_types.await_and_resolve_as_posix(path)
+      self.assertEqual(result, '/tmp/test/dir')
+
+    asyncio.run(_test())
+
+  def test_resolves_deferred_path_after_set(self):
+    async def _test():
+      directory = self.create_tempdir().full_path
+      deferred = tfhub_atomicity.DeferredPath()
+      deferred.set_path(epath.Path(directory))
+      result = await path_types.await_and_resolve_as_posix(deferred)
+      self.assertEqual(result, directory)
+
+    asyncio.run(_test())
+
+  def test_blocks_until_deferred_path_set(self):
+    async def _test():
+      directory = self.create_tempdir().full_path
+      deferred = tfhub_atomicity.DeferredPath()
+      task = asyncio.create_task(
+          path_types.await_and_resolve_as_posix(deferred)
+      )
+
+      await asyncio.sleep(0.1)
+      self.assertFalse(task.done())
+
+      deferred.set_path(epath.Path(directory))
+      result = await task
+      self.assertEqual(result, directory)
+
+    asyncio.run(_test())
+
+  def test_resolves_child_deferred_path(self):
+    async def _test():
+      directory = self.create_tempdir().full_path
+      deferred = tfhub_atomicity.DeferredPath()
+      child = deferred / 'subdir'
+      task = asyncio.create_task(path_types.await_and_resolve_as_posix(child))
+
+      await asyncio.sleep(0.1)
+      self.assertFalse(task.done())
+
+      deferred.set_path(epath.Path(directory))
+      result = await task
+      self.assertEqual(result, directory + '/subdir')
+
+    asyncio.run(_test())
+
+
+if __name__ == '__main__':
+  absltest.main()

checkpoint/orbax/checkpoint/_src/serialization/jax_array_handlers.py

@@ -24,6 +24,7 @@
 import warnings
 
 from absl import logging
+from etils import epath
 import humanize
 import jax
 import jax.numpy as jnp
@@ -38,6 +39,7 @@
 from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.multihost import multislice
 from orbax.checkpoint._src.path import async_path
+from orbax.checkpoint._src.path import types as path_types
 from orbax.checkpoint._src.path import utils as path_utils
 from orbax.checkpoint._src.serialization import jax_array_restore_args
 from orbax.checkpoint._src.serialization import limits
@@ -157,9 +159,8 @@ async def _async_serialize_shardings(
       continue
     if info.parent_dir is None:
       raise ValueError('parent_dir cannot be None')
-    tspec_sharding = ts_utils.get_sharding_tensorstore_spec(
-        info.parent_dir.as_posix(), info.name
-    )
+    dir_str = await path_types.await_and_resolve_as_posix(info.parent_dir)
+    tspec_sharding = ts_utils.get_sharding_tensorstore_spec(dir_str, info.name)
     if multihost.is_primary_host(primary_host):
       # OCDBT is not used for sharding metadata.
       sharding_ts_context = info.ts_context
@@ -568,8 +569,14 @@ async def _async_serialize_replica_slices(
             'Replica_separate_folder is disabled as OCDBT is not enabled.',
             1,
         )
+    resolved_parent_dir = await path_types.await_and_resolve_as_posix(
+        info.parent_dir
+    )
+    resolved_info = dataclasses.replace(
+        info, parent_dir=epath.Path(resolved_parent_dir)
+    )
     array_write_spec = ts_utils.build_array_write_spec(
-        info=info,
+        info=resolved_info,
         arg=arg,
         global_shape=value.global_shape,
         local_shape=value.local_shape,
@@ -718,8 +725,9 @@ async def _deserialize_shardings(
       )
       assert info.parent_dir is not None
       if info.name:
+        dir_str = await path_types.await_and_resolve_as_posix(info.parent_dir)
         tspec_sharding = ts_utils.get_sharding_tensorstore_spec(
-            info.parent_dir.as_posix(), info.name
+            dir_str, info.name
         )
         t = await ts.open(
             tspec_sharding,
@@ -765,8 +773,14 @@ async def _async_deserialize(
       await _validate_non_ocdbt_files(infos, metadata_key)
     deserialize_ops = []
     for info, arg, sharding in zip(infos, args, shardings):
+      resolved_parent_dir = await path_types.await_and_resolve_as_posix(
+          info.parent_dir
+      )
+      resolved_info = dataclasses.replace(
+          info, parent_dir=epath.Path(resolved_parent_dir)
+      )
       array_read_spec = ts_utils.build_array_read_spec(
-          info,
+          resolved_info,
           use_ocdbt=use_ocdbt,
           metadata_key=metadata_key,
           raise_array_data_missing_error=info.raise_array_data_missing_error,
@@ -1010,8 +1024,14 @@ async def metadata(
     for info in infos:
       # Use OCDBT flag from the existing checkpoint.
       use_ocdbt = info.is_ocdbt_checkpoint
+      resolved_parent_dir = await path_types.await_and_resolve_as_posix(
+          info.parent_dir
+      )
+      resolved_info = dataclasses.replace(
+          info, parent_dir=epath.Path(resolved_parent_dir)
+      )
       array_read_spec = ts_utils.build_array_read_spec(
-          info,
+          resolved_info,
           use_ocdbt=use_ocdbt,
           metadata_key=self._metadata_key,
           raise_array_data_missing_error=info.raise_array_data_missing_error,
@@ -1024,8 +1044,9 @@ async def metadata(
       assert info.parent_dir is not None
       sharding_op = None
       if info.name:
+        dir_str = await path_types.await_and_resolve_as_posix(info.parent_dir)
         tspec_sharding = ts_utils.get_sharding_tensorstore_spec(
-            info.parent_dir.as_posix(), info.name
+            dir_str, info.name
         )
         if sharding_file_exists:
           sharding_op = ts.open(

checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils.py

@@ -648,7 +648,7 @@ def _get_json_tspec(
   if info.name is None or info.parent_dir is None:
     raise ValueError('Must provide info.name and info.parent_dir.')
   parent_dir = info.parent_dir
-  assert parent_dir is not None
+  assert isinstance(parent_dir, epath.Path)
   directory = parent_dir.as_posix()
   kvstore_tspec = build_kvstore_tspec(
       directory,
@@ -750,6 +750,7 @@ def build_array_read_spec(
   """Gets ArrayReadSpec for reading."""
   if info.name is None or info.parent_dir is None:
     raise ValueError('Must provide info.name and info.parent_dir.')
+  assert isinstance(info.parent_dir, epath.Path)
   return ArrayReadSpec(
       directory=info.parent_dir.as_posix(),
       relative_array_filename=info.name,
@@ -778,7 +779,7 @@ def build_array_write_spec(
   if info.name is None or info.parent_dir is None:
     raise ValueError('Must provide info.name and info.parent_dir.')
   parent_dir = info.parent_dir
-  assert parent_dir is not None
+  assert isinstance(parent_dir, epath.Path)
   directory = parent_dir.as_posix()
 
   return ArrayWriteSpec(

checkpoint/orbax/checkpoint/_src/serialization/type_handlers.py

@@ -18,16 +18,19 @@
 
 import asyncio
 import copy
+import dataclasses
 import sys
 from typing import Any, Dict, Optional, Sequence, Tuple, TypeAlias, Union
 
 from absl import logging
+from etils import epath
 import jax
 import numpy as np
 from orbax.checkpoint._src.futures import future
 from orbax.checkpoint._src.metadata import value as value_metadata
 from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.path import format_utils
+from orbax.checkpoint._src.path import types as path_types
 from orbax.checkpoint._src.serialization import jax_array_handlers
 from orbax.checkpoint._src.serialization import ocdbt_utils
 from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
@@ -85,10 +88,16 @@ async def metadata(
   ) -> Sequence[value_metadata.ArrayMetadata]:
     open_ops = []
     for info in infos:
+      resolved_parent_dir = await path_types.await_and_resolve_as_posix(
+          info.parent_dir
+      )
+      resolved_info = dataclasses.replace(
+          info, parent_dir=epath.Path(resolved_parent_dir)
+      )
       # Use OCDBT flag from the existing checkpoint.
       use_ocdbt = info.is_ocdbt_checkpoint
       array_read_spec = ts_utils.build_array_read_spec(
-          info,
+          resolved_info,
           use_ocdbt=use_ocdbt,
           metadata_key=self._metadata_key,
           raise_array_data_missing_error=info.raise_array_data_missing_error,
@@ -122,8 +131,14 @@ async def _background_serialize(
     """Serializes numpy arrays in a background thread."""
     write_coros = []
     for value, info, arg in zip(values, infos, args):
+      resolved_parent_dir = await path_types.await_and_resolve_as_posix(
+          info.parent_dir
+      )
+      resolved_info = dataclasses.replace(
+          info, parent_dir=epath.Path(resolved_parent_dir)
+      )
       array_write_spec = ts_utils.build_array_write_spec(
-          info=info,
+          info=resolved_info,
           arg=arg,
           global_shape=value.shape,
           local_shape=value.shape,
@@ -175,14 +190,22 @@ async def deserialize(
     types.check_input_arguments(infos, args)
     open_futures = []
     for info, arg in zip(infos, args):
+      resolved_parent_dir = await path_types.await_and_resolve_as_posix(
+          info.parent_dir
+      )
+      resolved_info = dataclasses.replace(
+          info, parent_dir=epath.Path(resolved_parent_dir)
+      )
       if not info.is_ocdbt_checkpoint:
         await ts_utils.assert_parameter_files_exist(
-            info.parent_dir / info.name, self._metadata_key, info.use_zarr3
+            resolved_info.parent_dir / resolved_info.name,
+            self._metadata_key,
+            resolved_info.use_zarr3,
         )
       # Use OCDBT flag from the existing checkpoint.
       use_ocdbt = info.is_ocdbt_checkpoint
       array_read_spec = ts_utils.build_array_read_spec(
-          info,
+          resolved_info,
           use_ocdbt=use_ocdbt,
           metadata_key=self._metadata_key,
           raise_array_data_missing_error=info.raise_array_data_missing_error,
@@ -293,6 +316,7 @@ def _get_json_tspec(
     """Gets Tensorstore spec in JSON format."""
     if info.parent_dir is None:
       raise ValueError('Must provide info.parent_dir.')
+    assert isinstance(info.parent_dir, epath.Path)
     directory = (info.parent_dir / self._filename).as_posix()
     kvstore_tspec = ts_utils.build_kvstore_tspec(directory, use_ocdbt=False)
     tspec = {
@@ -330,7 +354,13 @@ async def _background_serialize(
         info,
         value,
     ) in zip(infos, values):
-      tspec = self._get_json_tspec(info)
+      resolved_parent_dir = await path_types.await_and_resolve_as_posix(
+          info.parent_dir
+      )
+      resolved_info = dataclasses.replace(
+          info, parent_dir=epath.Path(resolved_parent_dir)
+      )
+      tspec = self._get_json_tspec(resolved_info)
       if multihost.process_index() == 0:
         t = await ts.open(
             tspec,
@@ -368,7 +398,13 @@ async def deserialize(
     open_futures = []
 
     for info in infos:
-      tspec = self._get_json_tspec(info)
+      resolved_parent_dir = await path_types.await_and_resolve_as_posix(
+          info.parent_dir
+      )
+      resolved_info = dataclasses.replace(
+          info, parent_dir=epath.Path(resolved_parent_dir)
+      )
+      tspec = self._get_json_tspec(resolved_info)
       open_future = ts.open(
           tspec, open=True, read=True, context=self._ts_context
       )

checkpoint/orbax/checkpoint/_src/serialization/types.py

@@ -31,8 +31,10 @@
 from orbax.checkpoint._src.metadata import pytree_metadata_options as pytree_metadata_options_lib
 from orbax.checkpoint._src.metadata import value as value_metadata
 from orbax.checkpoint._src.serialization import limits
+from orbax.checkpoint.experimental.v1._src.path import types as path_types
 import tensorstore as ts
 
+
 PyTreeMetadataOptions = pytree_metadata_options_lib.PyTreeMetadataOptions
 
 
@@ -115,8 +117,8 @@ class ParamInfo:
   """
 
   name: str
-  parent_dir: epath.Path
-  path: Optional[epath.Path] = None
+  parent_dir: epath.Path | path_types.PathAwaitingCreation
+  path: epath.Path | path_types.PathAwaitingCreation | None = None
   keypath: Optional[Tuple[Any, ...]] = None
   skip_deserialize: Optional[bool] = None
   byte_limiter: Optional[limits.ByteLimiter] = None

checkpoint/orbax/checkpoint/experimental/v1/_src/serialization/compatibility.py

@@ -204,6 +204,7 @@ def _construct_deserialization_param(
 def _construct_deserialization_context(
     info: types_v0.ParamInfo,
 ) -> types.DeserializationContext:
+  assert isinstance(info.parent_dir, path_types.Path)
   return types.DeserializationContext(
       parent_dir=info.parent_dir,
       ocdbt_checkpoint=info.is_ocdbt_checkpoint,
@@ -377,6 +378,7 @@ async def metadata(
 
     ret = []
     for info, metadata in zip(infos, metadatas):
+      assert isinstance(info.parent_dir, path_types.Path)
       ret.append(
           _convert_v1_metadata_to_v0(
               info.name, info.parent_dir / info.name, metadata