#p2p Add Grain data iterator checkpointing to P2P

PiperOrigin-RevId: 863828755

Author: Orbax Authors

checkpoint/orbax/checkpoint/experimental/emergency/p2p/args.py

@@ -14,19 +14,41 @@
 
 """P2P composite checkpoint argument."""
 
-from typing import final
+from typing import Any, final
 from orbax.checkpoint import args as args_lib
 from orbax.checkpoint.experimental.emergency.p2p import constants
+from orbax.checkpoint.experimental.emergency.p2p import utils
+
+
+def _check_data_iter(value: Any):
+  """Checks if data_iter is valid."""
+  if utils.pygrain() is None:
+    raise ImportError(
+        'grain library is not available. Please install grain to use data_iter.'
+    )
+  if not isinstance(
+      value,
+      (
+          utils.pygrain().PyGrainCheckpointSave,
+          utils.pygrain().PyGrainCheckpointRestore,
+      ),
+  ):
+    raise TypeError(f'Unsupported type for data_iter: {type(value)}')
 
 
 @final
 class Composite(args_lib.Composite):
-  """Composite argument that only supports 'state' key."""
+  """Composite argument that supports 'state' and 'data_iter' keys."""
 
   def __init__(self, *args, **kwargs):
     super().__init__(*args, **kwargs)
-    if constants.STATE_SUBDIR not in self or len(self) > 1:
+    if constants.STATE_SUBDIR not in self:
       raise ValueError(
-          f'Composite must contain "{constants.STATE_SUBDIR}" key and no other'
-          f' keys: {list(self.keys())}'
+          f'Composite must contain "{constants.STATE_SUBDIR}" key:'
+          f' {list(self.keys())}'
       )
+    for key in self:
+      if key not in [constants.STATE_SUBDIR, constants.DATA_ITER_KEY]:
+        raise ValueError(f'Unsupported key in Composite: {key}')
+      if key == constants.DATA_ITER_KEY:
+        _check_data_iter(self[key])

checkpoint/orbax/checkpoint/experimental/emergency/p2p/constants.py

@@ -17,7 +17,9 @@
 # Directory names
 P2P_RESTORE_DIR_NAME = 'p2p_restore'
 STATE_SUBDIR = 'state'
+DATA_ITER_KEY = 'data_iter'
 PROCESS_SUBDIR_PREFIX = 'ocdbt.process_'
+PYGRAIN_STATES_FILENAME = 'pygrain_states.json'
 
 # Tuning for high-throughput networks (16MB buffers)
 SOCKET_BUFFER_SIZE = 16 * 1024 * 1024

checkpoint/orbax/checkpoint/experimental/emergency/p2p/local.py

@@ -14,21 +14,100 @@
 
 """Internal checkpoint manager for local P2P storage logic."""
 
+import dataclasses
+import json
 from typing import Any, Sequence, final
 
 from absl import logging
 from etils import epath
 import jax
 import orbax.checkpoint as ocp
+from orbax.checkpoint import args as args_lib
 from orbax.checkpoint import checkpoint_manager
 from orbax.checkpoint import type_handlers
 from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.serialization import type_handler_registry
 from orbax.checkpoint.experimental.emergency import checkpoint_manager as emergency_checkpoint_manager
+from orbax.checkpoint.experimental.emergency import path as emergency_path
 from orbax.checkpoint.experimental.emergency.p2p import args as p2p_args_lib
+from orbax.checkpoint.experimental.emergency.p2p import constants
 from orbax.checkpoint.experimental.emergency.p2p import utils
 
 
+if utils.pygrain() is not None:
+
+  class _LocalPyGrainHandlerMixin(utils.pygrain().PyGrainCheckpointHandler):
+    """Mixin for Local PyGrain handler."""
+
+    def __init__(self, process_index: int):
+      self._process_index = process_index
+
+    def save(self, directory: epath.Path, item: Any = None, args: Any = None):
+      """Saves the PyGrain iterator state to a JSON file."""
+      item = item or args.item
+      state = item.get_state()
+
+      if isinstance(item, utils.pygrain().DatasetIterator):
+        state_val = state
+      else:
+        # DataLoaderIterator state is bytes, decode to string for JSON
+        state_val = state.decode()
+
+      local_data = {
+          str(self._process_index): state_val,
+      }
+
+      all_data = emergency_path.sync_global_data(local_data)
+
+      combined_data = {}
+      for entry in all_data:
+        combined_data.update(entry)
+
+      (directory / constants.PYGRAIN_STATES_FILENAME).write_text(
+          json.dumps(combined_data, indent=2)
+      )
+
+    def restore(
+        self, directory: epath.Path, item: Any = None, args: Any = None
+    ):
+      """Restores the PyGrain iterator state from a JSON file."""
+      item = item or args.item
+      path = directory / constants.PYGRAIN_STATES_FILENAME
+
+      if not path.exists():
+        raise ValueError(f'PyGrain states not found at {path}')
+
+      combined_data = json.loads(path.read_text())
+      my_key = str(self._process_index)
+
+      if my_key not in combined_data:
+        raise ValueError(
+            f'Process index {self._process_index} not found in {path}'
+        )
+
+      state_val = combined_data[my_key]
+
+      if isinstance(item, utils.pygrain().DatasetIterator):
+        # DatasetIterator expects a dict
+        state = state_val
+      else:
+        # DataLoaderIterator expects bytes
+        state = state_val.encode()
+
+      item.set_state(state)
+      return item
+
+  @ocp.args.register_with_handler(_LocalPyGrainHandlerMixin, for_save=True)
+  @dataclasses.dataclass
+  class LocalPyGrainSave(ocp.args.CheckpointArgs):
+    item: Any
+
+  @ocp.args.register_with_handler(_LocalPyGrainHandlerMixin, for_restore=True)
+  @dataclasses.dataclass
+  class LocalPyGrainRestore(utils.pygrain().PyGrainCheckpointRestore):
+    item: Any
+
+
 @final
 class LocalCheckpointManager:
   """Wrapper around Orbax CheckpointManager for local P2P shards."""
@@ -77,10 +156,16 @@ def __init__(
         type_handler_registry=local_registry,
     )
 
+    item_handlers = dict(state=handler)
+    if utils.pygrain() is not None:
+      item_handlers[constants.DATA_ITER_KEY] = _LocalPyGrainHandlerMixin(
+          self._process_index
+      )
+
     self._manager = checkpoint_manager.CheckpointManager(
         self._directory,
         options=p2p_specific_options,
-        item_handlers=dict(state=handler),
+        item_handlers=item_handlers,
     )
 
   @property
@@ -122,6 +207,14 @@ def save(
       force: bool = False,
   ) -> bool:
     """Saves the checkpoint."""
+    if utils.pygrain() is not None and constants.DATA_ITER_KEY in args:
+      original_save = args[constants.DATA_ITER_KEY]
+      args_dict = dict(args.items())
+      args_dict[constants.DATA_ITER_KEY] = LocalPyGrainSave(
+          item=original_save.item
+      )
+      args = args_lib.Composite(**args_dict)
+
     return self._manager.save(step, args=args, force=force)
 
   def restore(
@@ -146,6 +239,14 @@ def restore(
         )
         raise ValueError(error_msg)
 
+    if utils.pygrain() is not None and args and constants.DATA_ITER_KEY in args:
+      original_restore = args[constants.DATA_ITER_KEY]
+      args_dict = dict(args.items())
+      args_dict[constants.DATA_ITER_KEY] = LocalPyGrainRestore(
+          original_restore.item
+      )
+      args = args_lib.Composite(**args_dict)
+
     # 2. Delegate to Orbax
     restored = self._manager.restore(
         step,

checkpoint/orbax/checkpoint/experimental/emergency/p2p/local_test.py

@@ -113,7 +113,9 @@ def test_save_restore(self, unused_process_index):
         ),
     )
 
-    jax.tree_util.tree_map(np.testing.assert_array_equal, state, restored.state)
+    jax.tree_util.tree_map(
+        np.testing.assert_array_equal, state, restored['state']
+    )
     manager.close()
 
 

checkpoint/orbax/checkpoint/experimental/emergency/p2p/persistent.py

@@ -29,6 +29,8 @@
 from orbax.checkpoint._src.serialization import type_handlers
 from orbax.checkpoint.experimental.emergency import checkpoint_manager as emergency_checkpoint_manager
 from orbax.checkpoint.experimental.emergency.p2p import args as p2p_args_lib
+from orbax.checkpoint.experimental.emergency.p2p import constants
+from orbax.checkpoint.experimental.emergency.p2p import utils
 
 _PRIMARY_REPLICA_ID = 0
 PyTree = Any
@@ -115,10 +117,14 @@ def __init__(
         enable_async_checkpointing=True,
     )
 
+    item_handlers = dict(state=_create_persistent_handler(mp_options))
+    if utils.pygrain() is not None:
+      item_handlers['data_iter'] = utils.pygrain().PyGrainCheckpointHandler()
+
     self._manager = checkpoint_manager.CheckpointManager(
         self._directory,
         options=internal_options,
-        item_handlers=dict(state=_create_persistent_handler(mp_options)),
+        item_handlers=item_handlers,
     )
 
   @property
@@ -169,8 +175,11 @@ def restore(
             abstract_state, sharding_tree
         ),
     )
+    restore_kwargs = {'state': restore_args_obj}
+    if constants.DATA_ITER_KEY in args:
+      restore_kwargs[constants.DATA_ITER_KEY] = args.data_iter
     return self._manager.restore(
-        step, args=p2p_args_lib.Composite(state=restore_args_obj)
+        step, args=p2p_args_lib.Composite(**restore_kwargs)
     )
 
   def delete(self, step: int):

checkpoint/orbax/checkpoint/experimental/emergency/p2p/persistent_test.py

@@ -16,6 +16,7 @@
 
 from absl.testing import absltest
 from etils import epath
+import grain.python as pygrain
 import jax
 import numpy as np
 from orbax.checkpoint import args as args_lib
@@ -220,6 +221,69 @@ def _to_abstract(x):
     test_utils.assert_tree_equal(self, state, restored_state)
     manager.close()
 
+  @mock.patch(
+      'orbax.checkpoint._src.multihost.multihost.process_index', return_value=0
+  )
+  def test_save_restore_with_grain_iterator(self, unused_process_index):
+    self._patch_process_index(process_index=0)
+    # persistent checkpoint manager with multiprocessing only works with a
+    # unified storage.
+    self.enter_context(mock.patch.object(jax, 'process_count', return_value=1))
+    devices = np.array([
+        [MockDevice(0, 0)],
+    ])
+    mesh = mock.Mock(
+        spec=jax.sharding.Mesh,
+        devices=devices,
+        axis_names=('replica', 'data'),
+        shape={'replica': 1, 'data': 1},
+        shape_tuple=devices.shape,
+        size=devices.size,
+    )
+
+    manager = persistent.PersistentCheckpointManager(
+        self.directory, mesh, replica_axis_index=0, options=self.options
+    )
+
+    ds = pygrain.MapDataset.source(list(range(10)))
+    dl = pygrain.DataLoader(
+        data_source=ds,
+        sampler=pygrain.SequentialSampler(10, pygrain.ShardOptions(0, 1)),
+        operations=[pygrain.Batch(1)],
+    )
+    data_iter = iter(dl)
+    for _ in range(3):
+      next(data_iter)
+
+    arr = jax.device_put(np.arange(self.mesh.size, dtype=np.int32))
+    state = {'a': arr}
+    save_args = p2p_args_lib.Composite(
+        state=args_lib.PyTreeSave(state),
+        data_iter=pygrain.PyGrainCheckpointSave(data_iter),
+    )
+    manager.save(1, args=save_args)
+    manager.wait_until_finished()
+
+    new_dl = pygrain.DataLoader(
+        data_source=ds,
+        sampler=pygrain.SequentialSampler(10, pygrain.ShardOptions(0, 1)),
+        operations=[pygrain.Batch(1)],
+    )
+    new_data_iter = iter(new_dl)
+    # PersistentCheckpointManager expects the state with sharding information
+    # in args.state.
+    restore_args = p2p_args_lib.Composite(
+        state=state,
+        data_iter=pygrain.PyGrainCheckpointRestore(new_data_iter),
+    )
+    restored = manager.restore(1, args=restore_args)
+
+    self.assertIn('state', restored)
+    self.assertIn('data_iter', restored)
+    test_utils.assert_tree_equal(self, state, restored['state'])
+    self.assertEqual(next(restored['data_iter']), 3)
+    manager.close()
+
   def test_delete_in_primary_slice_deletes(self):
     self._patch_process_index(process_index=0)
     manager = persistent.PersistentCheckpointManager(

checkpoint/orbax/checkpoint/experimental/emergency/p2p/utils.py

@@ -14,10 +14,27 @@
 
 """Utils for P2P checkpointing."""
 
+from typing import Any
+
 from absl import logging
 from etils import epath
+
 from orbax.checkpoint.experimental.emergency.p2p import constants
 
+# pytype:disable=import-error
+# pylint:disable=g-import-not-at-top
+try:
+  import grain.python as pygrain_module
+except ImportError:
+  pygrain_module = None
+# pytype:enable=import-error
+# pylint:enable=g-import-not-at-top
+
+
+def pygrain() -> Any | None:
+  """Returns the grain.python module if available, otherwise None."""
+  return pygrain_module
+
 
 def detect_process_index(directory: epath.Path, step: int) -> int | None:
   """Inspects the disk to find which process index created this step."""

checkpoint/pyproject.toml

@@ -72,6 +72,7 @@ testing = [
     'tensorflow',
     'fastapi',
     'httpx',
+    'grain',
 ]
 
 [tool.flit.sdist]