Eliminate usages of ParamInfo.path in favor of parent_dir / name.

PiperOrigin-RevId: 834378892

Author: cpgaffney1

checkpoint/orbax/checkpoint/_src/handlers/array_checkpoint_handler.py

@@ -153,10 +153,8 @@ def restore(
     # checkpoints lacking PYTREE_METADATA_FILE is no longer needed.
     restore_args = args.restore_args or type_handlers.RestoreArgs()
 
-    checkpoint_path = directory / self._checkpoint_name
     info = type_handlers.ParamInfo(
         name=self._checkpoint_name,
-        path=checkpoint_path,
         parent_dir=directory,
         skip_deserialize=False,
         is_ocdbt_checkpoint=type_handlers.is_ocdbt_checkpoint(directory),

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

@@ -472,7 +472,6 @@ def _param_info(keypath, name, value):
       return ParamInfo(
           name=name,
           keypath=keypath,
-          path=(directory / name),
           parent_dir=directory,
           skip_deserialize=skip_deserialize,
           is_ocdbt_checkpoint=use_ocdbt,

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

@@ -268,7 +268,6 @@ def _get_param_info(
       skip_deserialize = meta_or_value.skip_deserialize
     return ParamInfo(
         name=name,
-        path=directory / name,
         parent_dir=directory,
         skip_deserialize=skip_deserialize,
         is_ocdbt_checkpoint=is_ocdbt_checkpoint,
@@ -283,7 +282,9 @@ def _get_param_info(
   if partial_restore:
     for key, meta in flat_structure.items():
       if key not in flat_item:
-        flat_param_infos[key] = ParamInfo(skip_deserialize=True)
+        flat_param_infos[key] = ParamInfo(
+            name='', parent_dir=directory, skip_deserialize=True
+        )
         flat_input_restore_args[key] = RestoreArgs()
       else:
         flat_param_infos[key] = _get_param_info(flat_param_names[key], meta)
@@ -322,7 +323,9 @@ def _get_param_info(
             # Specified `use_fallback`, but key was also present in the
             # checkpoint. This means we should skip loading, since it will be
             # overridden with a new value.
-            flat_param_infos[input_key] = ParamInfo(skip_deserialize=True)
+            flat_param_infos[input_key] = ParamInfo(
+                name='', parent_dir=directory, skip_deserialize=True
+            )
             flat_input_restore_args[input_key] = RestoreArgs()
           else:
             # Specified `use_fallback`, but `transforms_default_to_original`
@@ -343,12 +346,16 @@ def _get_param_info(
           else:
             # Take the value from the user-provided `item`, ignoring any value
             # in the checkpoint.
-            flat_param_infos[input_key] = ParamInfo(skip_deserialize=True)
+            flat_param_infos[input_key] = ParamInfo(
+                name='', parent_dir=directory, skip_deserialize=True
+            )
             flat_input_restore_args[input_key] = RestoreArgs()
       else:
         # No match, restoration not required since it will be dropped from the
         # output.
-        flat_param_infos[input_key] = ParamInfo(skip_deserialize=True)
+        flat_param_infos[input_key] = ParamInfo(
+            name='', parent_dir=directory, skip_deserialize=True
+        )
         flat_input_restore_args[input_key] = RestoreArgs()
 
     restore_args = tree_utils.from_flat_dict(

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

@@ -485,7 +485,6 @@ def as_custom_metadata(
       param_name = '.'.join(keypath)
       flat_param_infos[keypath] = types.ParamInfo(
           name=param_name,
-          path=directory / param_name,
           parent_dir=directory,
           skip_deserialize=value_meta.skip_deserialize,
           is_ocdbt_checkpoint=use_ocdbt,

checkpoint/orbax/checkpoint/_src/metadata/tree_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import chex
+from etils import epath
 import jax
 import numpy as np
 from orbax.checkpoint._src.metadata import tree as tree_metadata_lib
@@ -43,6 +44,8 @@ def _to_param_infos(
   return jax.tree.map(
       # Other properties are not relevant.
       lambda x: types.ParamInfo(
+          name='',
+          parent_dir=epath.Path(''),
           value_typestr=type_handler_registry.get_param_typestr(
               x,
               type_handler_registry.GLOBAL_TYPE_HANDLER_REGISTRY,

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

@@ -681,7 +681,7 @@ async def _validate_non_ocdbt_files(
 ):
   await asyncio.gather(*[
       ts_utils.assert_parameter_files_exist(  # pylint: disable=protected-access
-          info.path, metadata_key, info.use_zarr3
+          info.parent_dir / info.name, metadata_key, info.use_zarr3
       )
       for info in infos
   ])

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

@@ -583,8 +583,8 @@ def _get_json_tspec(
     raise_array_data_missing_error: bool = True,
 ) -> dict[str, Any]:
   """Gets Tensorstore spec in JSON format."""
-  if info.path is None:
-    raise ValueError('Must construct serialization path.')
+  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
   directory = parent_dir.as_posix()
@@ -691,8 +691,8 @@ def build_array_write_spec(
     ext_metadata: dict[str, Any] | None = None,
 ) -> ArrayWriteSpec:
   """Gets ArrayWriteSpec for writing."""
-  if info.path is None:
-    raise ValueError('Must construct serialization path.')
+  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
   directory = parent_dir.as_posix()

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

@@ -22,7 +22,6 @@
 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
@@ -202,7 +201,7 @@ async def deserialize(
     for info, arg in zip(infos, args):
       if not info.is_ocdbt_checkpoint:
         await ts_utils.assert_parameter_files_exist(
-            info.path, self._metadata_key, info.use_zarr3
+            info.parent_dir / info.name, self._metadata_key, info.use_zarr3
         )
       # Use OCDBT flag from the existing checkpoint.
       use_ocdbt = info.is_ocdbt_checkpoint
@@ -310,8 +309,8 @@ def _get_json_tspec(
       info: types.ParamInfo,
   ) -> Dict[str, Any]:
     """Gets Tensorstore spec in JSON format."""
-    if info.path is None:
-      raise ValueError('Must construct serialization path.')
+    if info.parent_dir is None:
+      raise ValueError('Must provide info.parent_dir.')
     directory = (info.parent_dir / self._filename).as_posix()
     kvstore_tspec = ts_utils.build_kvstore_tspec(directory, use_ocdbt=False)
     tspec = {
@@ -384,11 +383,9 @@ async def deserialize(
     """See superclass documentation."""
     del args
     types.check_input_arguments(infos)
-    directory = epath.Path(infos[0].path).parent
     open_futures = []
 
     for info in infos:
-      info.path = epath.Path(directory / self._filename)
       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

@@ -60,7 +60,7 @@ def check_input_arguments(*args):
       raise ValueError('Found input args with mismatched lengths.')
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(kw_only=True)
 class ParamInfo:
   """Information describing a parameter in a PyTree.
 
@@ -70,13 +70,12 @@ class ParamInfo:
 
   name:
     Name of the parameter.
-  path:
-    A path providing a location where file(s) should be saved. The path is
-    assumed to be a directory.
   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
@@ -115,10 +114,10 @@ class ParamInfo:
   is_prioritized_key_fn: See `IsPrioritizedKeyFn` definition.
   """
 
-  name: Optional[str] = None
-  keypath: Optional[Tuple[Any, ...]] = None
+  name: str
+  parent_dir: epath.Path
   path: Optional[epath.Path] = None
-  parent_dir: 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
@@ -133,6 +132,10 @@ class ParamInfo:
   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
+
 
 @dataclasses.dataclass
 class SaveArgs:

checkpoint/orbax/checkpoint/_src/testing/benchmarks/array_handler_benchmark.py

@@ -151,7 +151,6 @@ def test_fn(self, test_context: core.TestContext) -> core.TestResult:
     )
     sharded_array = test_context.pytree['array']
     array_name = 'array'
-    array_path = test_context.path / array_name
 
     ts_context = ts_utils.get_ts_context(use_ocdbt=options.use_ocdbt)
     value_typestr = type_handler_registry.get_param_typestr(
@@ -162,7 +161,6 @@ def test_fn(self, test_context: core.TestContext) -> core.TestResult:
 
     param_info = type_handlers.ParamInfo(
         name=array_name,
-        path=array_path,
         parent_dir=test_context.path,
         use_zarr3=options.use_zarr3,
         is_ocdbt_checkpoint=options.use_ocdbt,