Internal change.

PiperOrigin-RevId: 733059068

Author: cpgaffney1

.github/workflows/build.yml

@@ -60,7 +60,7 @@ jobs:
     - name: Test with pytest
     # TODO(yaning): Move these to an exclude target within pytest.ini.
       run: |
-        python -m pytest --ignore=orbax/checkpoint/experimental/emergency/broadcast_multislice_test.py --ignore=orbax/checkpoint/experimental/emergency/checkpoint_manager_test.py --ignore=orbax/checkpoint/experimental/emergency/local_checkpoint_data_debugging_test.py --ignore=orbax/checkpoint/experimental/emergency/local_checkpoint_manager_test.py --ignore=orbax/checkpoint/experimental/emergency/multihost_test.py --ignore=orbax/checkpoint/experimental/emergency/replicator_checkpoint_manager_test.py --ignore=orbax/checkpoint/_src/testing/multiprocess_test.py
+        python -m pytest --ignore=orbax/checkpoint/experimental/emergency/broadcast_multislice_test.py --ignore=orbax/checkpoint/experimental/emergency/checkpoint_manager_test.py --ignore=orbax/checkpoint/experimental/emergency/single_slice_checkpoint_manager_test.py --ignore=orbax/checkpoint/experimental/emergency/local_checkpoint_data_debugging_test.py --ignore=orbax/checkpoint/experimental/emergency/local_checkpoint_manager_test.py --ignore=orbax/checkpoint/experimental/emergency/multihost_test.py --ignore=orbax/checkpoint/experimental/emergency/replicator_checkpoint_manager_test.py --ignore=orbax/checkpoint/_src/testing/multiprocess_test.py
     # The below step just reports the success or failure of tests as a "commit status".
     # This is needed for copybara integration.
     - name: Report success or failure as github status

checkpoint/CHANGELOG.md

@@ -7,13 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Fix RESOURCE_EXHAUSTED while writing array_metadatas.
+
 ### Changed
 
 - Improve `Cannot serialize host local jax.Array` error message.
 
 ### Added
 
-- support saving and restoring jax.random.key() in PyTree
+- support saving and restoring jax.random.key() in PyTree.
+- `CheckpointableHandler` for V1.
+- Support single-slice checkpointing in `emergency.CheckpointManager`.
 
 ## [0.11.6] - 2025-02-20
 

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

@@ -323,6 +323,8 @@ def __init__(
             type_handler_registry
         )
     )
+    if self._array_metadata_store:
+      self._array_metadata_store.set_primary_host(self._primary_host)
     self._array_metadata_validator = array_metadata_validator
 
 

checkpoint/orbax/checkpoint/_src/metadata/array_metadata_store.py

@@ -133,9 +133,45 @@ def __init__(
       self,
       path_resolver: PathResolver = PathResolver(),
       ser_deser: SerDeserializer = SerDeserializer(),
+      primary_host: int | None = 0,  # None means all hosts are primary hosts.
+      write_timeout_secs: int = 600,  # 10 minutes.
   ):
     self._path_resolver = path_resolver
     self._ser_deser = ser_deser
+    self._primary_host = primary_host
+    self._write_timeout_secs = write_timeout_secs
+
+  def set_primary_host(self, primary_host: int | None) -> None:
+    """Sets the primary host."""
+    self._primary_host = primary_host
+
+  async def _maybe_create_base_dir(self, base_dir: epath.Path) -> None:
+    """Primary host creates the base directory, rest of the hosts wait."""
+    if multihost.is_primary_host(self._primary_host):
+      # primary host creates, rest of the hosts wait.
+      return await asyncio.to_thread(
+          base_dir.mkdir, parents=True, exist_ok=True
+      )
+
+    # non-primary host waits for primary host to create the base dir/folder.
+    async def wait_for_base_dir_creation():
+      while not await asyncio.to_thread(base_dir.exists):
+        await asyncio.sleep(0.25)
+
+    try:
+      await asyncio.wait_for(
+          wait_for_base_dir_creation(), timeout=self._write_timeout_secs
+      )
+    except asyncio.TimeoutError as e:
+      primary_process = (
+          'LOCAL' if self._primary_host is None else self._primary_host
+      )
+      raise ValueError(
+          f'[process_index={multihost.process_index()}] Timed out waiting for'
+          f' array_metadatas base directory creation: {base_dir}.'
+          f' timeout={self._write_timeout_secs} seconds.'
+          f' primary_process={primary_process}'
+      ) from e
 
   async def write(
       self,
@@ -155,7 +191,7 @@ async def write(
     file_path = self._path_resolver.get_write_file_path(
         checkpoint_dir, process_index
     )
-    await asyncio.to_thread(file_path.parent.mkdir, parents=True, exist_ok=True)
+    await self._maybe_create_base_dir(file_path.parent)
     await asyncio.to_thread(
         file_path.write_text, self._ser_deser.serialize(array_metadatas)
     )

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

@@ -76,6 +76,8 @@ def slice_count(
     global_mesh: jax.sharding.Mesh, *, replica_axis_index: int = 0
 ) -> int:
   """Number of slices implied by the mesh's replica dimension."""
+  if len(global_mesh.shape_tuple) == 1:
+    return 1
   return global_mesh.devices.shape[replica_axis_index]
 
 

checkpoint/orbax/checkpoint/abstract_checkpoint_manager.py

@@ -85,7 +85,11 @@ def best_step(self) -> Optional[int]:
 
   @abc.abstractmethod
   def reload(self):
-    """Performs disk reads to ensure internal properties are up to date."""
+    """Reloads internal properties.
+
+    Resets internal cache of checkpoint steps, in case the directory managed
+    by this object has been updated externally.
+    """
 
   @abc.abstractmethod
   def reached_preemption(self, step: int) -> bool:
@@ -112,186 +116,25 @@ def delete(self, step: int):
   def save(
       self,
       step: int,
-      items: Optional[Union[Any, Mapping[str, Any]]] = None,
-      save_kwargs: Optional[Union[SaveParams, Mapping[str, SaveParams]]] = None,
-      metrics: Optional[PyTree] = None,
-      force: Optional[bool] = False,
-      args: Optional[args_lib.CheckpointArgs] = None,
-      custom_metadata: dict[str, Any] | None = None,
+      *args,
+      **kwargs,
   ) -> bool:
-    """Saves the provided items.
-
-    This method should be called by all hosts - process synchronization and
-    actions that need to be performed on only one host are managed internally.
-
-    NOTE: The `items` and `save_kwargs` arguments are deprecated, use `args`
-    instead. Make sure to configure `CheckpointManager` with `item_names`.
-
-    `args` should be a subclass of
-    `orbax.checkpoint.args.CheckpointArgs`, the specific type of which is used
-    to indicate what logic is used to save the object. For a typical, PyTree of
-    arrays, use `StandardSave`/`StandardRestore`.
-
-    When constructing the `CheckpointManager`, if no `item_names` were provided,
-    it is assumed that we are managing a single object. If `item_names` were
-    provided, it is assumed that we are managing multiple objects, and `args`
-    must be `orbax.checkpoint.args.CompositeArgs`. See below for details.
-
-    Example::
-
-      # Single item
-      mngr = ocp.CheckpointManager(directory)
-      mngr.save(step, args=ocp.args.StandardSave(my_train_state))
-
-      # Multiple items
-      mngr = ocp.CheckpointManager(directory, item_names=('state', 'meta'))
-      mngr.save(step, args=ocp.args.Composite(
-          state=ocp.args.StandardSave(my_train_state),
-          meta=ocp.args.JsonSave(my_metadata)
-      ))
-
-    Args:
-      step: current step, int
-      items: a savable object, or a dictionary of object name to savable object.
-      save_kwargs: save kwargs for a single Checkpointer, or a dictionary of
-        object name to kwargs needed by the Checkpointer implementation to save
-        the object.
-      metrics: a dictionary of metric name (string) to numeric value to be
-        tracked along with this checkpoint. Required if `options.best_fn` is
-        set. Allows users to specify a metric value to determine which
-        checkpoints are best and should be kept (in conjunction with
-        `options.max_to_keep`).
-      force: if `True`, this method will attempt to save a checkpoint regardless
-        of the result of `AbstractCheckpointManager.should_save(step)`. By
-        default, `save` will only write a checkpoint to disk when the options
-        permit, e.g. when `step` is in `options.save_interval_steps` or
-        `options.save_on_steps`. Setting `force=True` will not overwrite
-        existing checkpoints.
-      args: `CheckpointArgs` which is used to save checkpointable objects with
-        the appropriate logic.
-      custom_metadata: a dictionary of custom metadata to be written to the
-        checkpoint directory via StepMetadata.
-
-    Returns:
-      bool indicating whether a save operation was performed.
-    Raises:
-      ValueError: if `track_best` was indicated but `metrics` is not provided.
-      ValueError: directory creation failed.
-      ValueError: if an item is provided for which no `Checkpointer` is
-      found.
-      ValueError: if the checkpoint already exists.
-    """
+    """Saves the given step."""
 
   @abc.abstractmethod
   def restore(
       self,
       step: Optional[int],
-      items: Optional[Union[Any, Mapping[str, Any]]] = None,
-      restore_kwargs: Optional[
-          Union[RestoreParams, Mapping[str, RestoreParams]]
-      ] = None,
-      directory: Optional[epath.PathLike] = None,
-      args: Optional[args_lib.CheckpointArgs] = None,
+      *args,
+      **kwargs,
   ) -> Union[Any, Mapping[str, Any], args_lib.Composite]:
-    """Restores from the given step and provided items.
-
-    This method should be called by all hosts - process synchronization and
-    actions that need to be performed on only one host are managed internally.
-
-    NOTE: The `items` and `restore_kwargs` arguments are deprecated, use `args`
-    instead. Make sure to configure `CheckpointManager` with `item_names`.
-    See `save` docstring for additional details.
-
-    Example::
-
-      # Single item
-      mngr = ocp.CheckpointManager(directory)
-      mngr.restore(step, args=ocp.args.StandardRestore(abstract_train_state))
-
-      # Multiple items
-      mngr = ocp.CheckpointManager(directory, item_names=('state', 'meta'))
-      mngr.restore(step, args=ocp.args.Composite(
-          state=ocp.args.StandardRestore(abstract_train_state),
-          meta=ocp.args.JsonRestore(),
-      ))
-      # If it is acceptable to restore without providing additional arguments,
-      # and if a save has already been performed, it is ok to do the following:
-      mngr.restore(step, args=ocp.args.Composite(state=None, meta=None))
-      # If a save has not already been performed, there is no way for Orbax to
-      # know how to restore the objects. If a save has already been performed,
-      # it remembers the logic used to save the objects.
-
-    Args:
-      step: current step, int
-      items: a restoreable object, or a dictionary of object name to restorable
-        object.
-      restore_kwargs: restore kwargs for a single Checkpointer, or a dictionary
-        of object name to kwargs needed by the Checkpointer implementation to
-        restore the object.
-      directory: if provided, uses the given directory rather than the
-        `directory` property of this class. Can be used to restore checkpoints
-        from an independent location.
-      args: `CheckpointArgs` which is used to restore checkpointable objects
-        with the appropriate logic.
-
-    Returns:
-      If managing a single item, returns a single checkpointable object.
-      If managing multiple items, returns ocp.args.Composite, where the keys
-      are item names, and values are checkpointable objects.
-    """
+    """Restores the given step."""
 
   @abc.abstractmethod
   def item_metadata(
       self, step: int
   ) -> Union[Any, Mapping[str, Any], args_lib.Composite]:
-    """For all Checkpointers, returns any metadata associated with the item.
-
-    Calls the `metadata` method for each Checkpointer and returns a
-    mapping of each item name to the restored metadata. If the manager only
-    manages a single item, a single metadata will be returned instead.
-
-    To avoid errors due to missing CheckpointHandlers, concrete
-    CheckpointManager constructor must allow mapping from item names to
-    respective CheckpointHandlers to be input other than via save() and
-    restore(). Please note that save() and restore() calls automatically
-    map CheckpointHandlers to respective item names and retain it during the
-    lifetime of the CheckpointManager instance.
-
-    Example::
-
-      # Single item
-      mngr = ocp.CheckpointManager(directory)
-      # No calls to save() or restore() before calling item_metadata().
-      mngr.item_metadata(step)  # Raises error.
-
-      mngr = ocp.CheckpointManager(directory,
-          item_handlers=ocp.StandardCheckpointHandler)
-      # No calls to save() or restore() before calling item_metadata().
-      metadata = mngr.item_metadata(step)  # Successful.
-
-      # Multiple items
-      mngr = ocp.CheckpointManager(directory, item_names=('state', 'extra'))
-      # No calls to save() or restore() before calling item_metadata().
-      mngr.item_metadata(step)  # Raises error.
-
-      mngr = ocp.CheckpointManager(directory,
-        item_names=('state', 'extra'),
-        item_handlers={
-            'state': ocp.StandardCheckpointHandler,
-            'extra': ocp.PytreeCheckpointHandler,
-        }
-      )
-      # No calls to save() or restore() before calling item_metadata().
-      metadata = mngr.item_metadata(step)  # Successful.
-
-    Metadata may be None for an individual item.
-
-    Args:
-      step: Step for which to retrieve metadata.
-
-    Returns:
-      A dictionary mapping name to item metadata, or a single item metadata.
-    """
+    """Returns metadata for all known items."""
 
   @abc.abstractmethod
   def metadata(