Refactor grain checkpoint handler functionality into grain_utility file.

PiperOrigin-RevId: 951631507

Author: angel-core

src/maxtext/common/checkpointing.py

@@ -15,9 +15,8 @@
 
 """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer."""
 
-import asyncio
 import time
-from typing import Any, Optional
+from typing import Any
 
 from absl import flags
 import datetime
@@ -29,14 +28,14 @@
 from maxtext.input_pipeline.multihost_dataloading import MultiHostDataLoadIterator
 from maxtext.input_pipeline.multihost_dataloading import RemoteIteratorWrapper
 from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator
+from maxtext.common import grain_utility
 from maxtext.common import train_state_nnx
 from maxtext.utils import exceptions
 from maxtext.utils import max_logging
 from maxtext.utils import gcs_utils
 from maxtext.utils import elastic_utils
 from maxtext.checkpoint_conversion.utils.load_dynamic import load_safetensors_dynamic_state
 
-import numpy as np
 import orbax.checkpoint as ocp
 from orbax.checkpoint import v1 as ocp_v1
 from orbax.checkpoint._src.arrays import sharding as sharding_utils
@@ -45,11 +44,7 @@
 import orbax.checkpoint.experimental.emergency.checkpoint_manager as emergency_checkpoint_manager
 import orbax.checkpoint.experimental.emergency.replicator_checkpoint_manager as emergency_replicator_checkpoint_manager
 # pylint: disable=too-many-positional-arguments
-import dataclasses
-import json
 
-import grain
-from grain.python import PyGrainCheckpointHandler
 from grain.experimental import ElasticIterator
 
 CheckpointManager = ocp.CheckpointManager
@@ -62,152 +57,6 @@
 EmergencyReplicatorCheckpointManager = emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager
 
 
-class GrainCheckpointHandler(PyGrainCheckpointHandler, ocp.CheckpointHandler):
-  """A CheckpointHandler that allows specifying process_index and process_count."""
-
-  def save(
-      self,
-      directory: epath.Path,
-      # `item` is for backwards compatibility with older Orbax API, see
-      # https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html.
-      item: Optional[Any] = None,
-      args: Any = None,
-  ):
-    """Saves the given iterator to the checkpoint in `directory`."""
-    item = item or args.item  # pytype:disable=attribute-error
-
-    # RemoteIteratorWrapper handles checkpointing via colocated python
-    if isinstance(item, RemoteIteratorWrapper):
-      step = int(directory.parent.name)
-      item.save_state(step)
-      return
-
-    # ElasticIterator state is a single global scalar shared by all shards,
-    # so we write one fixed `process_0.json` from process 0 only. This file
-    # layout survives changes in `jax.process_count()`.
-    if isinstance(item, ElasticIterator):
-      if jax.process_index() == 0:
-        directory.mkdir(parents=True, exist_ok=True)
-        filename = directory / "process_0.json"
-        filename.write_text(json.dumps(item.get_state(), indent=4))
-      return
-
-    def save_single_process(item, process_index, process_count):
-      filename = directory / f"process_{process_index}-of-{process_count}.json"
-      if isinstance(item, grain.DatasetIterator):
-        state = json.dumps(item.get_state(), indent=4)
-      else:
-        state = item.get_state().decode()
-      filename.write_text(state)
-
-    if isinstance(item, list):
-      for local_iterator, process_index, process_count in item:
-        save_single_process(local_iterator, process_index, process_count)
-    else:
-      process_index, process_count = jax.process_index(), jax.process_count()
-      save_single_process(item, process_index, process_count)
-
-  def restore(
-      self,
-      directory: epath.Path,
-      item: Optional[Any] = None,
-      args: Any = None,
-  ) -> Any:
-    """Restores the given iterator from the checkpoint in `directory`."""
-    item = item or args.item
-    process_index = getattr(args, "process_index", None)
-    process_count = getattr(args, "process_count", None)
-
-    # In Pathways + colocated_python environment, RemoteIteratorWrapper handles checkpointing
-    if isinstance(item, RemoteIteratorWrapper):
-      step = int(directory.parent.name)
-      item.restore_state(step)
-      return item
-
-    # McJax and Pathways through controller cases
-    # ElasticIterator: every process reads the same shared `process_0.json`.
-    if isinstance(item, ElasticIterator):
-      filename = directory / "process_0.json"
-      if not filename.exists():
-        raise ValueError(f"File {filename} does not exist.")
-      item.set_state(json.loads(filename.read_text()))
-      return item
-
-    def restore_single_process(item, process_index, process_count):
-      filename = directory / f"process_{process_index}-of-{process_count}.json"
-      if not filename.exists():
-        raise ValueError(f"File {filename} does not exist.")
-      state = filename.read_text()
-      if isinstance(item, grain.DatasetIterator):
-        state = json.loads(state)
-      else:
-        state = state.encode()
-      item.set_state(state)
-      return item
-
-    if isinstance(item, list):
-      restored_items = []
-      for data_iter, process_idx in zip(item, process_index):  # pyrefly: ignore[bad-argument-type]
-        restored_items.append(restore_single_process(data_iter, process_idx, process_count))
-      return restored_items
-    else:
-      if process_index is None or process_count is None:
-        process_index, process_count = jax.process_index(), jax.process_count()
-      return restore_single_process(item, process_index, process_count)
-
-
-@ocp.args.register_with_handler(GrainCheckpointHandler, for_save=True)
-@dataclasses.dataclass
-class GrainCheckpointSave(ocp.args.CheckpointArgs):
-  item: Any
-
-
-@ocp.args.register_with_handler(GrainCheckpointHandler, for_restore=True)
-@dataclasses.dataclass
-class GrainCheckpointRestore(ocp.args.CheckpointArgs):
-  item: Any
-  process_index: Optional[int | list[int]] = None
-  process_count: Optional[int] = None
-
-
-class GrainCheckpointable(ocp_v1.StatefulCheckpointable):
-  """Adapts `GrainCheckpointHandler` to Orbax v1's `StatefulCheckpointable`."""
-
-  def __init__(
-      self,
-      *,
-      save_args: GrainCheckpointSave | None = None,
-      restore_args: GrainCheckpointRestore | None = None,
-  ):
-    self._handler = GrainCheckpointHandler()
-    self._save_args = save_args
-    self._restore_args = restore_args
-
-  async def save(self, directory):
-    """Saves the Grain iterator state to the given directory."""
-    # `GrainCheckpointHandler.save` snapshots iterator state (`get_state`) AND
-    # writes it; both must happen in this (blocking) save phase, NOT in the
-    # returned background coroutine.
-    path = await directory.await_creation()
-    self._handler.save(path, args=self._save_args)
-
-    async def _committed():  # nothing left for the background commit phase
-      return None
-
-    return _committed()
-
-  async def load(self, directory: epath.Path):
-    """Loads the Grain iterator state from the given directory."""
-    handler, args = self._handler, self._restore_args
-
-    # This will be ran to completion so asynchronous portion is just for
-    # compatibility with Orbax v1 API.
-    async def _background_load():
-      await asyncio.to_thread(handler.restore, directory, args=args)
-
-    return _background_load()
-
-
 def _weight_mismatches(want, have, path=()):
   """Returns `(path, problem)` for each weight in `want` that `have` didn't restore faithfully.
 
@@ -462,7 +311,7 @@ def create_orbax_checkpoint_manager(
 
   if dataset_type is not None and dataset_type == "grain":
     item_names += ("iter",)
-    item_handlers["iter"] = GrainCheckpointHandler()  # pyrefly: ignore[bad-assignment]
+    item_handlers["iter"] = grain_utility.GrainCheckpointHandler()  # pyrefly: ignore[bad-assignment]
 
   # local storage checkpoint needs parent directory created
   p = gcs_utils.mkdir_and_check_permissions(checkpoint_dir)
@@ -623,137 +472,6 @@ def print_save_message(step, async_checkpointing):
     max_logging.log(f"Saved a checkpoint at step {step}.")
 
 
-def _find_idx(array: np.ndarray, replica_axis_idx: int):
-  """Returns the index along given dimension that the current host belongs to."""
-  idx = None
-  for idx, val in np.ndenumerate(array):
-    if val.process_index == jax.process_index():
-      break
-  return idx[replica_axis_idx]
-
-
-def _replica_devices(device_array: np.ndarray, replica_axis_idx: int):
-  """Returns the devices from the replica that current host belongs to.
-
-  Replicas are assumed to be restricted to the first axis.
-
-  Args:
-    device_array: devices of the mesh that can be obtained by mesh.devices()
-    replica_axis_idx: axis dimension along which replica is taken
-
-  Returns:
-    devices inside the replica that current host is in
-  """
-  idx = _find_idx(device_array, replica_axis_idx)
-  replica_result = np.take(device_array, idx, axis=replica_axis_idx)
-  return np.expand_dims(replica_result, axis=replica_axis_idx)
-
-
-def _prepare_scaled_down_grain_restore_args(
-    data_iterator: list, process_count_jax: int, process_count_stored: int, directory: epath.Path
-) -> GrainCheckpointRestore:
-  """
-  Prepares the restore arguments for a scaled-up (list) data iterator.
-
-  This is used when restoring a checkpoint saved with more processes than
-  the current run (e.g., 64 files onto 32 JAX processes).
-  """
-  # 1. Validation Assertions
-  assert isinstance(data_iterator, list), (
-      f"{process_count_stored} processes found in Grain checkpoint directory {directory}, but only "
-      f"{process_count_jax} jax processes in this run, please set expansion_factor_real_data accordingly."
-  )
-
-  scaling_factor = len(data_iterator)
-  expected_process_count = process_count_stored / process_count_jax
-  assert scaling_factor == expected_process_count, (
-      f"Found {process_count_stored} processes in checkpoint and {process_count_jax} "
-      f"JAX processes, implying a scaling factor of {expected_process_count}. "
-      f"However, the data_iterator list has {scaling_factor} items."
-  )
-
-  # 2. Prepare Arguments
-  local_iterator_list = [x.local_iterator for x in data_iterator]
-  # Each JAX process calculates the global indices it's responsible for.
-  # e.g., process 0 with scaling_factor=2 handles checkpoints from processes [0, 32]
-  # e.g., process 1 with scaling_factor=2 handles checkpoints from processes [1, 33]
-  process_index_list = [jax.process_index() + i * process_count_jax for i in range(scaling_factor)]
-
-  return GrainCheckpointRestore(local_iterator_list, process_index=process_index_list, process_count=process_count_stored)
-
-
-def _restore_grain_iterator(
-    checkpoint_manager,
-    step: int,
-    data_iterator,
-    checkpoint_args,
-    expansion_factor_real_data: int,  # This must be defined in the outer scope
-) -> tuple[Any, None]:
-  """
-  Handles the complex logic for restoring a Grain data iterator checkpoint.
-  This function dispatches to the correct restore strategy based on
-  the number of stored checkpoint files vs. current JAX processes.
-  """
-  if isinstance(data_iterator, RemoteIteratorWrapper):
-    grain_restore_args = GrainCheckpointRestore(item=data_iterator)
-    restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args))
-    return (restored_state, None)
-
-  # ElasticIterator: one shared `process_0.json` regardless of shard count.
-  if not isinstance(data_iterator, list) and isinstance(data_iterator.local_iterator, ElasticIterator):
-    grain_restore_args = GrainCheckpointRestore(item=data_iterator.local_iterator)
-    restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args))
-    return (restored_state, None)
-
-  directory = checkpoint_manager.directory / str(step) / "iter"
-  process_count_jax = jax.process_count()
-
-  # Count the number of checkpoint files
-  process_count_stored = len(list(directory.glob("process_*-of-*.json")))
-
-  grain_restore_args = None
-
-  if process_count_stored > process_count_jax:
-    # Scaling down from a larger number of hosts. (e.g., 128 files -> 64 processes)
-    # In this case, each host restores a list of data iterators.
-    grain_restore_args = _prepare_scaled_down_grain_restore_args(
-        data_iterator, process_count_jax, process_count_stored, directory
-    )
-
-  elif process_count_stored == process_count_jax:
-    # Normal case: number of hosts is the same. (e.g., 64 files -> 64 processes)
-    assert not isinstance(data_iterator, list), (
-        f"{process_count_stored} processes found in Grain checkpoint directory {directory}, matching the number of "
-        "jax process, please do not set expansion_factor_real_data."
-    )
-    grain_restore_args = GrainCheckpointRestore(data_iterator.local_iterator)
-
-  elif expansion_factor_real_data > 1 and process_count_stored == process_count_jax // expansion_factor_real_data:
-    # Scaling up to a larger number of hosts.(e.g., 32 files -> 64 processes)
-    # In this case, a subset of hosts restore the data iterator.
-    assert not isinstance(
-        data_iterator, list
-    ), "when expansion_factor_real_data > 1, the data iterator should not be a list."
-    grain_restore_args = GrainCheckpointRestore(
-        data_iterator.local_iterator, process_index=jax.process_index(), process_count=process_count_stored
-    )
-
-  else:
-    # Case 4: Mismatch
-    raise ValueError(
-        f"Error restoring Grain checkpoint in {directory}: "
-        f"The number of stored checkpoint files ({process_count_stored}) "
-        f"is incompatible with the number of JAX processes ({process_count_jax}). "
-        "If you are resuming training with a different number of chips, see instructions in "
-        "https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline/"
-        "data_input_grain.md#using-grain"
-    )
-
-  # Call restore once with the composed arguments
-  restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args))
-  return (restored_state, None)
-
-
 def load_state_if_possible(
     checkpoint_manager: CheckpointManager | None,
     data_iterator: MultiHostDataLoadIterator | list[MultiHostDataLoadIterator] | None,
@@ -811,7 +529,7 @@ def map_to_pspec(data):
         pspec = data.sharding.spec
         mesh = data.sharding.mesh
         replica_axis_index = 0
-        replica_devices = _replica_devices(mesh.devices, replica_axis_index)
+        replica_devices = grain_utility.replica_devices(mesh.devices, replica_axis_index)
         replica_mesh = jax.sharding.Mesh(replica_devices, mesh.axis_names)
         single_replica_sharding = jax.sharding.NamedSharding(replica_mesh, pspec)
 
@@ -846,7 +564,7 @@ def map_to_pspec(data):
             and not isinstance(data_iterator, PlaceHolderDataIterator)
             and (checkpoint_manager.directory / str(step) / "iter").exists()
         ):
-          restored, _ = _restore_grain_iterator(
+          restored, _ = grain_utility.restore_grain_iterator(
               checkpoint_manager, step, data_iterator, checkpoint_args, expansion_factor_real_data
           )
         else:
@@ -907,7 +625,7 @@ def map_to_pspec(data):
             and not isinstance(data_iterator, PlaceHolderDataIterator)
             and (checkpoint_manager.directory / str(step) / "iter").exists()
         ):
-          return _restore_grain_iterator(
+          return grain_utility.restore_grain_iterator(
               checkpoint_manager,
               step,
               data_iterator,
@@ -1189,12 +907,14 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
   if config and config.dataset_type == "grain" and not isinstance(data_iterator, PlaceHolderDataIterator):
     if isinstance(data_iterator, RemoteIteratorWrapper):
       # Pass the wrapper directly; GrainCheckpointHandler will call save_state with the step
-      save_args_composite["iter"] = GrainCheckpointSave(item=data_iterator)  # pyrefly: ignore[bad-assignment]
+      save_args_composite["iter"] = grain_utility.GrainCheckpointSave(
+          item=data_iterator
+      )  # pyrefly: ignore[bad-assignment]
     elif not isinstance(data_iterator, list) and isinstance(
         data_iterator.local_iterator, ElasticIterator
     ):  # pyrefly: ignore[missing-attribute]
       # ElasticIterator checkpoints a single global scalar shared by all shards.
-      save_args_composite["iter"] = GrainCheckpointSave(
+      save_args_composite["iter"] = grain_utility.GrainCheckpointSave(
           item=data_iterator.local_iterator
       )  # pyrefly: ignore[bad-assignment]
     else:
@@ -1209,7 +929,9 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
         grain_iters_to_save.append(
             (data_iter.local_iterator, process_index, process_count_total)
         )  # pyrefly: ignore[missing-attribute]
-      save_args_composite["iter"] = GrainCheckpointSave(item=grain_iters_to_save)  # pyrefly: ignore[bad-assignment]
+      save_args_composite["iter"] = grain_utility.GrainCheckpointSave(
+          item=grain_iters_to_save
+      )  # pyrefly: ignore[bad-assignment]
 
   custom_metadata = {}
   if config: