Internal change

PiperOrigin-RevId: 872295675

Author: mxberlot

checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py

@@ -648,9 +648,11 @@ async def async_save(
         use_compression=self._use_compression,
         use_zarr3=self._use_zarr3,
     )
-    assert all(
-        leaf.parent_dir is directory for leaf in jax.tree.leaves(param_infos)
-    )
+    # TODO(b/425293362): Add validation for PathAwaitingCreation.
+    if isinstance(directory, epath.Path):
+      assert all(
+          leaf.parent_dir is directory for leaf in jax.tree.leaves(param_infos)
+      )
 
     serialize_ops = []  # List of (coros -> List of futures)
     batch_requests = batched_serialization_requests(
@@ -710,7 +712,6 @@ async def async_save(
           future.CommitFutureAwaitingContractedSignals(
               self._write_metadata_after_commits(
                   commit_futures,
-                  checkpoint_dir=directory,
                   param_infos=param_infos,
                   save_args=save_args,
                   custom_metadata=custom_metadata,
@@ -790,7 +791,7 @@ async def _maybe_deserialize(
     flat_metadata = tree_utils.to_flat_dict(metadata)
     byte_limiter = limits.get_byte_limiter(self._restore_concurrent_bytes)
     param_infos = jax.tree.map(
-        lambda info: dataclasses.replace(info, byte_limiter=byte_limiter),
+        lambda info: info.replace(byte_limiter=byte_limiter),
         param_infos,
     )
     batch_requests = batched_serialization_requests(
@@ -1156,9 +1157,8 @@ def update_param_info(param_info: types.ParamInfo) -> types.ParamInfo:
             f'No ArrayMetadata found for param_info: {param_info}, checkpoint'
             f' directory: {checkpoint_dir}, process_index={process_index}.'
         )
-      return dataclasses.replace(
-          param_info,
-          write_shape=array_metadatas_cache[param_info.name].write_shape,
+      return param_info.replace(
+          write_shape=array_metadatas_cache[param_info.name].write_shape
       )
 
     return jax.tree.map(update_param_info, param_infos)
@@ -1210,7 +1210,6 @@ async def _write_metadata_file(
   async def _write_metadata_after_commits(
       self,
       commit_futures: List[future.Future],
-      checkpoint_dir: epath.Path,
       *,
       param_infos: PyTree,
       save_args: PyTree,
@@ -1225,6 +1224,11 @@ async def _write_metadata_after_commits(
     for commit_future in commit_futures:
       await asyncio.to_thread(commit_future.result)
 
+    await asyncio.gather(
+        *[info.await_path_creation() for info in jax.tree.leaves(param_infos)]
+    )
+    checkpoint_dir = jax.tree.leaves(param_infos)[0].parent_dir
+
     commit_time = time.time()
     # `write_shape` is extracted from ArrayMetadata store saved during
     # materialization of commit_futures. Then it is written to the pytree

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

@@ -680,7 +680,7 @@ async def _maybe_deserialize(
     """Deserializes values or gets them from the aggregate file."""
     byte_limiter = limits.get_byte_limiter(self._restore_concurrent_bytes)
     param_infos = jax.tree.map(
-        lambda info: dataclasses.replace(info, byte_limiter=byte_limiter),
+        lambda info: info.replace(byte_limiter=byte_limiter),
         param_infos,
     )
 

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

@@ -155,6 +155,7 @@ async def _async_serialize_shardings(
   for sharding, info in zip(shardings, infos):
     if sharding is None:
       continue
+    await info.await_path_creation()
     if info.parent_dir is None:
       raise ValueError('parent_dir cannot be None')
     tspec_sharding = ts_utils.get_sharding_tensorstore_spec(
@@ -568,6 +569,7 @@ async def _async_serialize_replica_slices(
             'Replica_separate_folder is disabled as OCDBT is not enabled.',
             1,
         )
+    await info.await_path_creation()
     array_write_spec = ts_utils.build_array_write_spec(
         info=info,
         arg=arg,
@@ -1003,6 +1005,7 @@ async def metadata(
     open_ops = []
     sharding_open_ops = []
     shardings = []
+    await asyncio.gather(*[info.await_path_creation() for info in infos])
     if infos[0].parent_dir is None:
       raise ValueError('parent_dir cannot be None')
     sharding_file_path = infos[0].parent_dir / _SHARDING_FILE_NAME
@@ -1197,6 +1200,7 @@ async def deserialize(
     if args is None:
       raise ValueError('Must provide ArrayRestoreArgs to restore as jax.Array.')
     types.check_input_arguments(infos, args)
+    await asyncio.gather(*[info.await_path_creation() for info in infos])
     if infos[0].parent_dir is None:
       raise ValueError('parent_dir cannot be None')
     sharding_file_path = infos[0].parent_dir / _SHARDING_FILE_NAME

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

@@ -85,6 +85,7 @@ async def metadata(
   ) -> Sequence[value_metadata.ArrayMetadata]:
     open_ops = []
     for info in infos:
+      await info.await_path_creation()
       # Use OCDBT flag from the existing checkpoint.
       use_ocdbt = info.is_ocdbt_checkpoint
       array_read_spec = ts_utils.build_array_read_spec(
@@ -122,6 +123,7 @@ async def _background_serialize(
     """Serializes numpy arrays in a background thread."""
     write_coros = []
     for value, info, arg in zip(values, infos, args):
+      await info.await_path_creation()
       array_write_spec = ts_utils.build_array_write_spec(
           info=info,
           arg=arg,
@@ -175,6 +177,7 @@ async def deserialize(
     types.check_input_arguments(infos, args)
     open_futures = []
     for info, arg in zip(infos, args):
+      await info.await_path_creation()
       if not info.is_ocdbt_checkpoint:
         await ts_utils.assert_parameter_files_exist(
             info.parent_dir / info.name, self._metadata_key, info.use_zarr3
@@ -308,6 +311,8 @@ def typestr(self) -> str:
   async def metadata(
       self, infos: Sequence[types.ParamInfo]
   ) -> Sequence[StringMetadata]:
+    for info in infos:
+      await info.await_path_creation()
     return [
         StringMetadata(name=info.name, directory=info.parent_dir)
         for info in infos
@@ -330,6 +335,7 @@ async def _background_serialize(
         info,
         value,
     ) in zip(infos, values):
+      await info.await_path_creation()
       tspec = self._get_json_tspec(info)
       if multihost.process_index() == 0:
         t = await ts.open(
@@ -368,6 +374,7 @@ async def deserialize(
     open_futures = []
 
     for info in infos:
+      await info.await_path_creation()
       tspec = self._get_json_tspec(info)
       open_future = ts.open(
           tspec, open=True, read=True, context=self._ts_context

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

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import abc
+import copy
 import dataclasses
 from typing import Any, Callable, Optional, Protocol, Sequence, Tuple
 
@@ -30,6 +31,7 @@
 from orbax.checkpoint._src.metadata import empty_values
 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.path import types as path_types
 from orbax.checkpoint._src.serialization import limits
 import tensorstore as ts
 
@@ -60,81 +62,123 @@ def check_input_arguments(*args):
       raise ValueError('Found input args with mismatched lengths.')
 
 
-@dataclasses.dataclass(kw_only=True)
 class ParamInfo:
   """Information describing a parameter in a PyTree.
 
   Note that ParamInfo is distinct from SaveArgs and RestoreArgs in that in
   represents information not provided by a user, and should be computed
   internally.
 
-  name:
-    Name of the parameter.
-  parent_dir:
-    A path providing location where all files under the same checkpoint should
-    be saved under. All `ParamInfo` provided to a given TypeHandler should have
-    the same `parent_dir`. The parent_dir is assumed to be a directory.
-  path:
-    Do not provide directly. Automatically set to `parent_dir / name`.
-  skip_deserialize:
-    If specified, skips deserialization of the given parameter using the
-    TypeHandler. This may be for multiple different reasons, including that the
-    parameter may have been aggregated, or it will be unneeded after
-    transformations. Note: this parameter is handled by PyTreeCheckpointHandler,
-    so it is unnecessary for TypeHandler implementations to deal with it.
-  byte_limiter:
-    Object to limit the number of bytes that can be read or written in
-    parallel.
-  device_host_byte_limiter:
-    Object to limit the number of bytes that can be transferred from device to
-    host memory in parallel.
-  is_ocdbt_checkpoint:
-    Indicates whether the checkpoint path uses OCDBT format
-    or not. Only used for restoration.
-  use_compression:
-    When True, turn on zstd compression. Default is True.
-  use_zarr3:
-    If True, use Zarr ver3 otherwise ver2.
-  ocdbt_target_data_file_size:
-    Specifies the target size (in bytes) of each OCDBT data file. If set to 0,
-    data file size is not limited. If omitted (None), the TensorStore default
-    is used.
-  ts_context:
-    Tensorstore context to use for reading/writing.
-  value_typestr: stores the original value's typestr (from TypeHandler).
-    Only required when saving.
-  enable_pinned_host_transfer:
-    True by default. If False, disables transfer to pinned host when copying
-    from device to host, regardless of the presence of pinned host memory.
-  raise_array_data_missing_error:
-    Only used for restoring. See documentation in `tensorstore_utils.py`. Comes
-    from tree metadata and should be the same across all parameters.
-  write_shape:
-    Shape of the array shard. Used in the subchunking context.
-  is_prioritized_key_fn: See `IsPrioritizedKeyFn` definition.
+  Attributes:
+    name: Name of the parameter.
+    parent_dir: A path providing location where all files under the same
+      checkpoint should be saved under. All ParamInfo provided to a given
+      TypeHandler should have the same parent_dir. The parent_dir is assumed to
+      be a directory. Accessing this property raises ValueError if the
+      underlying path is a PathAwaitingCreation that has not been resolved via
+      await_path_creation().
+    path: Automatically set to parent_dir / name. Accessing this property raises
+      ValueError if the underlying path is a PathAwaitingCreation that has not
+      been resolved via await_path_creation().
+    skip_deserialize: If specified, skips deserialization of the given parameter
+      using the TypeHandler.
+    byte_limiter: Object to limit the number of bytes that can be read or
+      written in parallel.
+    device_host_byte_limiter: Object to limit the number of bytes that can be
+      transferred from device to host memory in parallel.
+    is_ocdbt_checkpoint: Indicates whether the checkpoint path uses OCDBT format
+      or not. Only used for restoration.
+    use_compression: When True, turn on zstd compression. Default is True.
+    use_zarr3: If True, use Zarr ver3 otherwise ver2.
+    ocdbt_target_data_file_size: Specifies the target size (in bytes) of each
+      OCDBT data file.
+    ts_context: Tensorstore context to use for reading/writing.
+    value_typestr: Stores the original value's typestr (from TypeHandler).
+    enable_pinned_host_transfer: If False, disables transfer to pinned host.
+    raise_array_data_missing_error: Only used for restoring.
+    write_shape: Shape of the array shard. Used in the subchunking context.
+    is_prioritized_key_fn: See ``IsPrioritizedKeyFn`` definition.
+    keypath: Tuple of keys identifying the parameter's position in the PyTree.
   """
 
-  name: str
-  parent_dir: epath.Path
-  path: Optional[epath.Path] = None
-  keypath: Optional[Tuple[Any, ...]] = None
-  skip_deserialize: Optional[bool] = None
-  byte_limiter: Optional[limits.ByteLimiter] = None
-  device_host_byte_limiter: Optional[limits.ByteLimiter] = None
-  is_ocdbt_checkpoint: Optional[bool] = None
-  use_compression: bool | None = True
-  use_zarr3: Optional[bool] = False
-  ocdbt_target_data_file_size: Optional[int] = None
-  ts_context: Optional[ts.Context] = None
-  value_typestr: Optional[str] = None
-  enable_pinned_host_transfer: bool = False
-  raise_array_data_missing_error: bool = True
-  write_shape: arrays_types.Shape | None = None
-  is_prioritized_key_fn: Optional[IsPrioritizedKeyFn] = None
-
-  def __post_init__(self):
-    if self.path is None:
-      self.path = self.parent_dir / self.name
+  def __init__(
+      self,
+      *,
+      name: str,
+      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,
+      device_host_byte_limiter: Optional[limits.ByteLimiter] = None,
+      is_ocdbt_checkpoint: Optional[bool] = None,
+      use_compression: bool | None = True,
+      use_zarr3: Optional[bool] = False,
+      ocdbt_target_data_file_size: Optional[int] = None,
+      ts_context: Optional[ts.Context] = None,
+      value_typestr: Optional[str] = None,
+      enable_pinned_host_transfer: bool = False,
+      raise_array_data_missing_error: bool = True,
+      write_shape: arrays_types.Shape | None = None,
+      is_prioritized_key_fn: Optional[IsPrioritizedKeyFn] = None,
+  ):
+    self.name = name
+    self._parent_dir = parent_dir
+    self._path = path if path is not None else parent_dir / name
+    self.keypath = keypath
+    self.skip_deserialize = skip_deserialize
+    self.byte_limiter = byte_limiter
+    self.device_host_byte_limiter = device_host_byte_limiter
+    self.is_ocdbt_checkpoint = is_ocdbt_checkpoint
+    self.use_compression = use_compression
+    self.use_zarr3 = use_zarr3
+    self.ocdbt_target_data_file_size = ocdbt_target_data_file_size
+    self.ts_context = ts_context
+    self.value_typestr = value_typestr
+    self.enable_pinned_host_transfer = enable_pinned_host_transfer
+    self.raise_array_data_missing_error = raise_array_data_missing_error
+    self.write_shape = write_shape
+    self.is_prioritized_key_fn = is_prioritized_key_fn
+
+  @property
+  def parent_dir(self) -> epath.Path:
+    if isinstance(self._parent_dir, path_types.PathAwaitingCreation):
+      raise ValueError(
+          'parent_dir is a PathAwaitingCreation and has not been resolved yet. '
+          'Call `await info.await_path_creation()` first.'
+      )
+    return self._parent_dir
+
+  @parent_dir.setter
+  def parent_dir(self, value: epath.Path | path_types.PathAwaitingCreation):
+    self._parent_dir = value
+
+  @property
+  def path(self) -> epath.Path:
+    if isinstance(self._path, path_types.PathAwaitingCreation):
+      raise ValueError(
+          'path is a PathAwaitingCreation and has not been resolved yet. '
+          'Call `await info.await_path_creation()` first.'
+      )
+    if self._path is None:
+      raise ValueError('path is None.')
+    return self._path
+
+  @path.setter
+  def path(self, value: epath.Path | path_types.PathAwaitingCreation | None):
+    self._path = value
+
+  def replace(self, **kwargs) -> ParamInfo:
+    new = copy.copy(self)
+    for k, v in kwargs.items():
+      setattr(new, k, v)
+    return new
+
+  async def await_path_creation(self) -> None:
+    if isinstance(self._parent_dir, path_types.PathAwaitingCreation):
+      resolved_path = await self._parent_dir.await_creation()
+      self._parent_dir = resolved_path
+      self._path = resolved_path / self.name
 
 
 @dataclasses.dataclass