Skip to content

utils.checkpoint.Checkpoint

Import path: earth2studio.utils.checkpoint.Checkpoint

View source on GitHub

Documentation

Catalog of restart checkpoints for a named inference run.

A checkpoint owns the on-disk catalog and commit directories for one logical inference run. Each committed catalog row stores generic workflow metadata and component state bound through bind_checkpoint_state.

Parameters:

  • name (str) –

    Name of the checkpoint store. Used in the default checkpoint path and in commit manifests.

  • path (str | Path | None, default: None ) –

    Directory where checkpoint files are stored. If None, checkpoints are written under the default Earth2Studio cache location, by default None.

  • mode (Literal['overwrite', 'append'], default: 'overwrite' ) –

    Catalog update mode. "overwrite" keeps only the latest row; "append" preserves row history for repeated writes, by default "overwrite".

  • flush_interval (int | None, default: 1 ) –

    Number of workflow write calls between durable commits. Use None to keep updates pending until flush is called explicitly, by default 1.

  • history_size (int | None, default: None ) –

    Maximum number of rows to keep when mode="append". In overwrite mode this is ignored and treated as one latest row, by default None.

  • level (CheckpointLevel, default: 2 ) –

    Requested component logging level. 0 records workflow progress and explicit metadata only, 1 allows component state needed to restart workflow items such as ensemble members, and 2 allows component state needed to resume inside a forecast rollout, by default 2.

  • rank (int | None, default: None ) –

    Distributed rank for this process. If None, Earth2Studio attempts to detect the rank from PhysicsNeMo or common distributed environment variables, by default None.

  • world_size (int | None, default: None ) –

    Distributed world size. If None, Earth2Studio attempts to detect the world size from PhysicsNeMo or WORLD_SIZE, by default None.

  • device (str | device, default: device('cpu') ) –

    Device where components should stage live tensor state before it is serialized into the checkpoint, by default torch.device("cpu").

rank_path property 

rank_path: Path

Directory for this process checkpoint writes.

catalog property 

catalog: tuple[CheckpointEntry, ...]

Committed checkpoint entries for the current rank.

active property 

active: CheckpointSession | None

Active checkpoint selected from this catalog, if one is in scope.

refresh 

refresh() -> None

Refresh the checkpoint catalog from disk.

select 

select(row: int) -> CheckpointSession

Select an existing checkpoint row by position.

The selected session can be used as a context manager. Bound component state and metadata are restored from the selected catalog row when the session becomes active.

Parameters:

  • row (int) –

    Positional row index. Negative indexing is supported.

Returns:

  • CheckpointSession –

    Session representing the selected catalog row.

Raises:

  • IndexError –

    If no catalog row matches the selection.

Examples using earth2studio.utils.checkpoint.Checkpoint

Checkpointing a Deterministic ForecastBasic inference workflow checkpointing.