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

PiperOrigin-RevId: 823187713

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

@@ -474,7 +474,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/multihost/pathways.py

@@ -16,38 +16,27 @@
 
 import functools
 import jax
-import numpy as np
-from .learning.deepmind.jax.ocean.remote_python import rp
 
 
 @functools.lru_cache(maxsize=1)
-def worker_count() -> int:
-  """Gets the number of Pathways workers."""
-  fully_replicated_sharding = jax.sharding.NamedSharding(
-      jax.sharding.Mesh(jax.devices(), 'x'),
-      jax.sharding.PartitionSpec(),
-  )
+def worker_count(global_mesh: jax.sharding.Mesh | None) -> int:
+  """Gets the number of Pathways workers.
 
-  @rp.stateless_fn
-  def _get_worker_count(_) -> jax.Array:
-    wc = np.asarray(jax.process_count(), dtype=np.int32)
-    return jax.make_array_from_callback(
-        (),
-        fully_replicated_sharding,
-        lambda _: wc,
-        dtype=np.int32,
-    )
+  Args:
+    global_mesh: The global mesh of active devices. If None is provided,
+      `jax.devices()` will be used.
 
-  dummy_input = jax.device_put(
-      np.asarray(0, dtype=np.int32),
-      device=fully_replicated_sharding,
-  )
-  _get_worker_count.register_shape_fn(
-      lambda _: jax.ShapeDtypeStruct(
-          (), dtype=np.int32, sharding=fully_replicated_sharding
-      )
-  )
-  result = _get_worker_count(rp.to_remote_python(dummy_input))
-  jax.block_until_ready(result)
-  result = rp.from_remote_python(result)
-  return result.item()
+  Returns:
+    The number of Pathways workers in the mesh.
+  """
+  global_mesh = global_mesh or jax.sharding.Mesh(jax.devices(), 'x')
+  devices = global_mesh.devices.flatten()
+  workers = set()
+  for d in devices:
+    attrs = []
+    if hasattr(d, 'virtual_task_index'):
+      attrs.append(d.virtual_task_index)
+    if hasattr(d, 'slice_index'):
+      attrs.append(d.slice_index)
+    workers.add(tuple(attrs))
+  return len(workers)

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

@@ -656,7 +656,7 @@ async def deserialize(
         )
       if not info.is_ocdbt_checkpoint:
         await ts_utils.assert_parameter_files_exist(
-            info.path,
+            info.parent_dir / info.name,
             self._metadata_key,
             info.use_zarr3,
         )
@@ -887,8 +887,8 @@ async def deserialize(
     for info, arg in zip(infos, args):
       arg = cast(SingleReplicaArrayRestoreArgs, arg)
       if not info.is_ocdbt_checkpoint:
-        await ts_utils.assert_parameter_files_exist(  # pylint: disable=protected-access
-            info.path, self._metadata_key
+        await ts_utils.assert_parameter_files_exist(
+            info.parent_dir / info.name, self._metadata_key
         )
       if not isinstance(arg, SingleReplicaArrayRestoreArgs):
         raise ValueError(

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

@@ -582,8 +582,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()
@@ -690,8 +690,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
@@ -198,7 +197,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
@@ -306,8 +305,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 = {
@@ -380,11 +379,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,9 +70,6 @@ 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
@@ -115,10 +112,9 @@ class ParamInfo:
   is_prioritized_key_fn: See `IsPrioritizedKeyFn` definition.
   """
 
-  name: Optional[str] = None
+  name: str
+  parent_dir: epath.Path
   keypath: Optional[Tuple[Any, ...]] = None
-  path: Optional[epath.Path] = None
-  parent_dir: Optional[epath.Path] = None
   skip_deserialize: Optional[bool] = None
   byte_limiter: Optional[limits.ByteLimiter] = None
   device_host_byte_limiter: Optional[limits.ByteLimiter] = None