grain.checkpoint module

APIs for saving and restoring pipeline state.

List of Members

class grain.checkpoint.CheckpointHandler

Orbax CheckpointHandler for PyGrain iterators.

This handler provides a bridge between Orbax checkpointing and PyGrain iterators. It manages the serialization and deserialization of iterator states across distributed processes, ensuring each process saves its unique shard state to a JSON file.

The handler supports both the legacy item parameter (for backward compatibility with older Orbax API prior to 0.5.0, see https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) and the newer Orbax V0 args API. It dynamically handles state formats: dataset.DatasetIterator types are serialized as JSON dictionaries, while other iterator types are handled as decoded strings.

Examples

Refer to CheckpointRestore and CheckpointSave for examples.

restore(directory, item=None, args=None)

Restores the given iterator from the checkpoint in directory.

Reads the process-specific JSON file from the checkpoint directory, deserializes the state, and applies it to the provided iterator using set_state(). The iterator will be modified in place. Critically, this method also triggers start_prefetch() after restoration to ensure the data pipeline resumes immediately.

Parameters:

• directory (Path) – The directory containing the checkpoint file.

• item (IteratorType | None) – A freshly created DataLoaderIterator or DatasetIterator. The state will be applied to this object. Should be provided if args is None. For backward compatibility with older Orbax API (prior to 0.5.0), see https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html.

• args (Any) – The Orbax V0 arguments container (typically a CheckpointRestore instance, since orbax-checkpoint-0.5.0) holding the target iterator. Should be provided if item is None. Note that the newest Orbax V1 API does not use args and instead operates on checkpointables directly.

Returns:

The restored DataLoaderIterator or DatasetIterator with restored state and active prefetching. This is the same object as the item argument (or args.item).

Raises:

ValueError – If the required process-specific checkpoint file does not exist.

Return type:

IteratorType

save(directory, item=None, args=None)

Saves the given iterator to the checkpoint in directory.

Retrieves the internal state of the iterator, formats it appropriately (using JSON serialization for DatasetIterator`s, and raw string decoding for other iterator types), and writes it to a process-specific file named `process_<index>-of-<count>.json in the specified checkpoint directory.

Parameters:

• directory (Path) – The directory where the checkpoint file will be written.

• item (IteratorType | None) – The DataLoaderIterator or DatasetIterator to be saved. Should be provided if args is None. For backward compatibility with older Orbax API (prior to 0.5.0), see https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html.

• args (Any) – For the Orbax V0 API (since orbax-checkpoint-0.5.0), args will contain the item to save (typically as args.item). Should be provided if item is None. Note that the newest Orbax V1 API does not use args and instead operates on checkpointables directly.

class grain.checkpoint.CheckpointSave(*args, **kwargs)

Arguments for saving a PyGrain iterator via Orbax.

This dataclass registers the CheckpointHandler for save operations within the Orbax API. It wraps the iterator instance, allowing it to be passed to a CheckpointManager using the newer Orbax V0 args interface.

Parameters:

item (Any)

item

The iterator instance to be saved.

Type:

Any

Example

Saving the iterator alongside model weights using Orbax Composite:

import orbax.checkpoint as ocp  # v0 API
import grain

from etils import epath
import tempfile

# Setup Data
model_weights = {'layer1': 1.0, 'layer2': 2.0}
ds = grain.MapDataset.range(100)
iterator = iter(ds)
for _ in range(10):
    next(iterator)

# Setup Orbax CheckpointManager
with tempfile.TemporaryDirectory() as temp_dir:
  path = epath.Path(temp_dir) / 'orbax_grain_checkpoint'
  path.mkdir(parents=True, exist_ok=True)

  with ocp.CheckpointManager(path) as mngr:
    mngr.save(
        step=0,
        args=ocp.args.Composite(
            model=ocp.args.StandardSave(model_weights),
            data_iter=grain.checkpoint.CheckpointSave(item=iterator)
        )
    )

class grain.checkpoint.CheckpointRestore(*args, **kwargs)

Arguments for restoring a PyGrain iterator via Orbax.

This dataclass registers the CheckpointHandler for restore operations within the Orbax API. It wraps the iterator instance, allowing it to be passed to a CheckpointManager using the newer Orbax V0 args interface.

Parameters:

item (Any)

item

The freshly created iterator instance to restore into.

Type:

Any

Example

Restoring the iterator alongside model weights using Orbax Composite:

import orbax.checkpoint as ocp  # v0 API
import grain

from etils import epath
import tempfile

# First, save a checkpoint to restore from.
model_weights = {'layer1': 1.0, 'layer2': 2.0}
ds = grain.MapDataset.range(100)
save_iterator = iter(ds)
for _ in range(10):
    next(save_iterator)

with tempfile.TemporaryDirectory() as temp_dir:
  path = epath.Path(temp_dir) / 'orbax_grain_checkpoint'
  path.mkdir(parents=True, exist_ok=True)

  with ocp.CheckpointManager(path) as mngr:
    mngr.save(
        step=0,
        args=ocp.args.Composite(
            model=ocp.args.StandardSave(model_weights),
            data_iter=grain.checkpoint.CheckpointSave(item=save_iterator)
        )
    )

  # Now, restore into fresh iterator and weights using a new manager
  # instance.
  restore_iterator = iter(ds)
  with ocp.CheckpointManager(path) as mngr:
    restored_data = mngr.restore(
        step=0,
        args=ocp.args.Composite(
            model=ocp.args.StandardRestore(),
            data_iter=grain.checkpoint.CheckpointRestore(item=restore_iterator)
        )
    )
  # Model weights are returned, iterator is restored in-place.
  print(restored_data['model'])
  # {'layer1': 1.0, 'layer2': 2.0}
  print(next(restore_iterator))
  # 10