IceChunkBackend

Import path: earth2studio.io.IceChunkBackend

View source on GitHub

Documentation

Bases: ZarrBackend

A backend that writes to an Icechunk repository.

Icechunk is a transactional storage engine for Zarr that adds version control (commits, branches, tags) on top of a regular object store. This backend behaves identically to earth2studio.io.ZarrBackend for add_array / write / read, using an Icechunk writable session's store in place of a plain Zarr store. Since Icechunk snapshots are immutable, call commit to persist accumulated writes; uncommitted writes are still visible to this backend but are lost if the process exits before committing (a warning is logged if the backend is destroyed with pending writes).

By default, write is non-blocking: the actual store write for a call runs in a background thread so the inference loop can move on to the next step while the previous step's write is still in flight. read, __getitem__ and commit all flush pending writes first, so they always observe the latest data; pass blocking=True to write synchronously instead, e.g. if a commit on every step makes the flush unconditional anyway.

Parameters:

• storage (Storage | str, default: None ) –

  Icechunk storage backend to open/create the repository with. If a string is provided, it is treated as a path for icechunk.local_filesystem_storage. If None, an in-memory Icechunk repository is created, by default None

• branch (str, default: 'main' ) –

  Branch to open a writable session on. Created (from the tip of "main") if it does not already exist, by default "main"

• repo_kwargs (dict[str, Any], default: {} ) –

  Key word arguments passed to icechunk.Repository.open_or_create, by default {}

• chunks (dict[str, int], default: {'ensemble': 1, 'time': 1, 'lead_time': 1, 'variable': 1} ) –

  An ordered dict of chunks to use with the data passed through data/coords, by default {}

• backend_kwargs (dict[str, Any], default: {'overwrite': False} ) –

  Key word arguments for zarr.Group root object, by default {"overwrite": False}

• zarr_codecs (CompressorsLike, default: None ) –

  Compression codec to use when creating any new arrays. Only effects Zarr 3.0. If None, will use no compressor, by default None

• blocking (bool, default: False ) –

  If False (default), write submits the store write to a background thread and returns immediately instead of waiting for it to complete. If True, write blocks until the store write finishes, by default False

• pool_size (int, default: 8 ) –

  Number of background threads used for non-blocking writes, ignored if blocking is True, by default 8

Note

For more information about Icechunk see: icechunk.io/en/latest/

Warning

In non-blocking mode, writing overlapping regions of the same array back to back before a flush point (read/__getitem__/commit) races: the background threads may apply in either order. This is not a concern for the typical inference-loop pattern of writing disjoint time/lead_time slices.

add_array

add_array(
    coords: CoordSystem,
    array_name: str | list[str],
    data: Tensor | list[Tensor] = None,
    **kwargs: Any
) -> None

Add an array to the existing zarr group.

Parameters:

• coords (CoordSystem) –

  Ordered dict of coordinate information.

• array_name (str) –

  Name to add to zarr group for the new array.

• data (Tensor | list[Tensor], default: None ) –

  Optional data to initialize the array with. If None, then the array is NaN initialized (zarr default). Can also pass a list of tensors, which must match in length to the list of array_names passed. If a list of tensors is passed, it is assumed that each tensor share coords.

• kwargs (Any, default: {} ) –

  Optional keyword arguments passed to zarr dataset constructor.

write

write(
    x: Tensor | list[Tensor],
    coords: CoordSystem,
    array_name: str | list[str],
) -> None

Write data to the current zarr group using the passed array_name.

Parameters:

• x (Tensor | list[Tensor]) –

  Tensor(s) to be written to zarr store.

• coords (OrderedDict) –

  Coordinates of the passed data.

• array_name (str | list[str]) –

  Name(s) of the array(s) that will be written to.

Examples using earth2studio.io.IceChunkBackend

IO Backend PerformanceLeverage different IO backends for storing inference results.