Add utilities for working with local steps on Pathways.

PiperOrigin-RevId: 834320290

Author: cpgaffney1

checkpoint/orbax/checkpoint/_src/multihost/multihost.py

@@ -67,6 +67,10 @@ def is_pathways_backend() -> bool:
   )
 
 
+def is_pathways_controller() -> bool:
+  return jax.local_devices()[0].client.runtime_type == 'pathways'
+
+
 def is_runtime_to_distributed_ids_initialized() -> bool:
   return _RUNTIME_TO_DISTRIBUTED_ID is not None
 

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/testing/local_path.py

@@ -18,22 +18,44 @@
 
 import os
 import typing
-from typing import Iterator
+from typing import Iterator, Sequence
 
 from etils import epath
 from orbax.checkpoint._src.multihost import multihost
 
 
+_LOCAL_PATH_BASE_NAME = '_local_path_base'
+_LOCAL_PART_PREFIX = 'local'
+
+
 # The following is a hack to pass the type checker.
 if typing.TYPE_CHECKING:
   _BasePath = epath.Path
 else:
   _BasePath = object
 
 
+def create_local_path_base(testclass) -> epath.Path:
+  return epath.Path(
+      testclass.create_tempdir(name=_LOCAL_PATH_BASE_NAME).full_path
+  )
+
+
+def _get_local_part_index(parts: Sequence[str]) -> int:
+  for i, part in enumerate(parts):
+    if part.startswith(_LOCAL_PART_PREFIX):
+      return i
+  raise ValueError(
+      f'Did not find a local part ({_LOCAL_PART_PREFIX}) in parts: {parts}'
+  )
+
+
 class LocalPath(_BasePath):
   """A Path implementation for testing process-local paths.
 
+  IMPORTANT: Use `create_local_path_base` to create the base path for test
+  cases.
+
   In the future, this class may more completely provide all functions and
   properties of a pathlib Path, but for now, it only provides the minimum
   needed to support relevant tests.
@@ -54,9 +76,12 @@ class LocalPath(_BasePath):
 
   def __init__(self, *parts: epath.PathLike):
     self._path = epath.Path('/'.join(os.fspath(p) for p in parts))
+    # Assumes this class will always be constructed on the controller first
+    # (otherwise this check will return the wrong value).
+    self._is_pathways_backend = multihost.is_pathways_backend()
 
   def __repr__(self) -> str:
-    return f'{self.__class__.__name__}({self.path})'
+    return f'LocalPath({self.path})'
 
   def __str__(self) -> str:
     return str(self.path)
@@ -67,7 +92,40 @@ def base_path(self) -> epath.Path:
 
   @property
   def path(self) -> epath.Path:
-    return self.base_path / str(f'local_{multihost.process_index()}')
+    parts = list(self.base_path.parts)
+
+    # Fail if the path is not properly configured. The local part should be
+    # immediately following the base name.
+    try:
+      base_idx = parts.index(_LOCAL_PATH_BASE_NAME)
+    except ValueError as e:
+      raise ValueError(
+          f'Base path for LocalPath must contain {_LOCAL_PATH_BASE_NAME}. Got:'
+          f' {self.base_path}'
+      ) from e
+
+    if multihost.is_pathways_controller():
+      local_part = f'{_LOCAL_PART_PREFIX}_controller'
+    else:
+      local_part = f'{_LOCAL_PART_PREFIX}_{multihost.process_index()}'
+
+    try:
+      # If the local part is already present, potentially replace it with the
+      # correct local part (e.g. controller vs worker).
+      local_part_idx = _get_local_part_index(parts)
+      assert local_part_idx == base_idx + 1
+      parts[local_part_idx] = local_part
+      return epath.Path(*parts)
+    except ValueError:
+      pass
+
+    # Otherwise, insert following the base part.
+    parts.insert(base_idx + 1, local_part)
+    return epath.Path(*parts)
+
+  @property
+  def parts(self) -> tuple[str, ...]:
+    return self.path.parts
 
   def exists(self) -> bool:
     """Returns True if self exists."""
@@ -119,6 +177,14 @@ def unlink(self, missing_ok: bool = False) -> None:
     """Remove this file or symbolic link."""
     self.path.unlink(missing_ok=missing_ok)
 
+  def touch(self, mode: int = 0o666, exist_ok: bool = False) -> None:
+    """Creates the file at this path."""
+    self.path.touch(exist_ok=exist_ok)
+
+  def rename(self, new_path: epath.PathLike) -> None:
+    """Renames this file or directory to the given path."""
+    self.path.rename(new_path)
+
   def write_bytes(self, data: bytes) -> int:
     """Writes content as bytes."""
     return self.path.write_bytes(data)
@@ -135,16 +201,16 @@ def write_text(
   def as_posix(self) -> str:
     return self.path.as_posix()
 
-  def __truediv__(self, key: epath.PathLike) -> epath.Path:
-    return self.path / key
+  def __truediv__(self, key: epath.PathLike) -> LocalPath:
+    return LocalPath(self.path / key)
 
   @property
   def name(self) -> str:
     return self.path.name
 
   @property
-  def parent(self) -> epath.Path:
-    return self.path.parent
+  def parent(self) -> LocalPath:
+    return LocalPath(self.path.parent)
 
   def __fspath__(self) -> str:
     return os.fspath(self.path)