io.AsyncZarrBackend

Import path: earth2studio.io.AsyncZarrBackend

View source on GitHub

Documentation

Async Zarr v3 IO Backend

An asynchronous Zarr backend for inference pipelines that produce data faster than a store can absorb it synchronously. Iteratively generated dimensions (time, lead_time, ensemble) go into parallel_coords with their complete value sets, each inference step writes one slice along them, and close() is called at the end to drain pending writes and write out any incomplete shard. Remote stores are supported through fs_factory, and Zarr v3 sharding via shard_coords keeps the file count of large campaigns low.

Warning

This IO backend presently does not support overwritting existing Zarr stores. Only creation of new arrays or writing to existing.

Warning

Enabling sharding via shard_coords buffers chunks in host memory until a shard is complete, trading host memory for a smaller file count. Budget roughly

max_inflight_shards * 4 * prod(shard_shape) * itemsize + pool_size * write_bytes

per process, since a flushing shard costs several times its own size once Zarr's encoded copy is counted. Sharded writes are also slower than unsharded ones, which is hidden as long as the model takes longer to produce a step than the store takes to absorb it. This latency hiding only applies in non-blocking mode; with blocking=True each shard flush runs synchronously.

Warning

When sharding, a shard must not contain data owned by more than one process. This backend keeps every shard object to a single write by buffering its chunks, but that only holds within a process, separate ranks have separate buffers. If two ranks each hold part of the same shard they will both write it in full and the later write wins, silently discarding the other's data. Shard along a coordinate that each rank owns entirely (typically lead_time, since a rank runs a whole forecast), not along the coordinate the work is distributed over.

Parameters:

• file_name (str) –

  Path location to place zarr store

• parallel_coords (CoordSystem) –

  Coordinates that enable parallel writes during inference. These coordinates specify which dimensions will be written in parallel via async operations, typically representing dimensions that are iteratively generated (such as time or lead_time). The chunk size for each of these dimensions will be set to 1. These coordinates should contain the complete set of values needed for the entire inference pipeline. The remaining coordinates of a given array will be populated upon the first write to the respective array.

• fs_factory (Callable[..., AbstractFileSystem], default: LocalFileSystem ) –

  FSSpec file system factory method. This is a callable object that should return an instance of the desired filesystem to use, by default LocalFileSystem

• blocking (bool, default: True ) –

  Blocking write calls in the synchronous API. When set to false, the IO backend will execute write calls in separate threads. Users should call the close() API to ensure all threads have finished / cleaned up, by default True

• pool_size (int, default: 8 ) –

  The thread / async loop pool used with the synchronous write API in non-blocking mode, by default 8

• async_timeout (int, default: 600 ) –

  Async operation timeout for a given write operation, by default 600. When sharding, the write that completes a shard carries the entire flush plus any wait for a concurrency slot, so this should be scaled with shard size and expected store throughput.

• zarr_kwargs (dict[str, Any], default: {'mode': 'a'} ) –

  Additional keyword arguments to provide to the zarr.api.asynchronous.open function, by default {"mode": "a"}

• zarr_codecs (CompressorsLike, default: None ) –

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

• chunked_coords (dict[str, int], default: {} ) –

  Chunk sizes for coordinates that are not in parallel_coords. By default any such coordinate is stored as a single chunk spanning its full length. Keys not present in a given array are ignored, by default {}

• shard_coords (dict[str, int], default: {} ) –

  Number of elements per shard along the given coordinates, enabling Zarr v3 sharding. Each value must be a multiple of that coordinate's chunk size, and any coordinate not listed uses a shard size equal to its chunk size. See the Sharding notes below. By default, {} (unsharded).

• max_inflight_shards (int, default: 4 ) –

  Maximum number of shard flushes allowed to run at once. Concurrent flushes are what keep sharded write throughput up, at the cost of holding that many shards in memory. Lower it if memory is tight, raise it if writes are the bottleneck and the store has bandwidth to spare, by default 4

Raises:

• ImportError –

  If Zarr 2.0 is installed. This io backend only supports Zarr 3.0

• TypeError –

  If fs_factory is not a callable, this should be a callable method not an object

• ValueError –

  If a shard_coords value is not positive

Notes

Relation to ZarrBackend

Exposes the same surface as ZarrBackend and can be used as a drop-in replacement, with a few behavioral differences:

• add_array takes a dtype instead of a template data tensor and is idempotent, so it is safe to call from every rank of a distributed job.
• In non-blocking mode a failed write raises at a later write, flush or close rather than at the failing call, and inputs are copied so callers may freely reuse their buffers.
• coords is read back from the store, and __getitem__ flushes pending writes first; intended for inspection, not reads in a write loop.
• Consolidated metadata is not maintained; consolidate at the end of a pipeline if desired, e.g. zarr.consolidate_metadata(io.store).

Sharding

Because every coordinate in parallel_coords is chunked with a size of 1, a large inference campaign can produce an enormous number of small files, which is a common way to exhaust an inode quota on a parallel filesystem. Sharding packs many chunks into a single storage object to avoid that. The chunk layout is unchanged, so readers still fetch one chunk at a time and only the file count changes.

A shard is one object, so writing part of one would force Zarr to read, modify and rewrite all of it. To keep every shard to a single write, this backend accumulates a shard's chunks in host memory and writes it once complete, hence the memory and throughput tradeoffs in the warnings above.

Shard sizes need not divide evenly into a coordinate. close() writes out any shard that never filled, using the array fill value where nothing was supplied, which reads back exactly as an unwritten chunk would. Writing into a shard already present in the store still works but falls back to a read-modify-write of the whole shard and logs a warning. That happens when close() or flush() is called mid run and the same shards are written again, or when restarting into a store left with incomplete shards, so aligning restart boundaries with the shard size keeps writes on the fast path.

Sharding composes with zarr_codecs, which compresses the inner chunks within a shard, and with chunked_coords, which sets the chunk size of coordinates outside parallel_coords.

Examples:

Write a forecast one lead time at a time, hiding the IO behind the model steps:

>>> times = np.array([np.datetime64("2024-01-01")])
>>> lead_times = np.array([np.timedelta64(6 * i, "h") for i in range(4)])
>>> io = AsyncZarrBackend(
...     "forecast.zarr",
...     parallel_coords={"time": times, "lead_time": lead_times},
...     blocking=False,
...     shard_coords={"lead_time": 4},  # optional: 4 chunks per storage object
... )
>>> total_coords = OrderedDict(
...     {
...         "time": times,
...         "lead_time": lead_times,
...         "lat": np.linspace(-90, 90, 721),
...         "lon": np.linspace(0, 360, 1440, endpoint=False),
...     }
... )
>>> io.add_array(total_coords, ["t2m", "z500"])
>>> for i in range(len(lead_times)):
...     x = torch.randn(1, 1, 721, 1440)  # model output for this step
...     step_coords = total_coords.copy()
...     step_coords["lead_time"] = lead_times[i : i + 1]
...     io.write([x, x], step_coords, ["t2m", "z500"])
>>> io.close()  # drain pending writes, write out any incomplete shard

add_array

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

Create arrays and their coordinate arrays in the store

Arrays are otherwise created lazily on first write, which is per process: several processes writing the same array for the first time all see it as absent and race on creation. Creating the schema up front, from one process, avoids that race. It also allows arrays with different dimension sets, which no single first write could imply.

Parameters:

• coords (CoordSystem) –

  Coordinate system of the array(s).

• array_name (str | list[str]) –

  Name(s) of the array(s) to create.

• dtype (dtype, default: float32 ) –

  Data type of the array(s), by default np.float32 (matching ZarrBackend.add_array without data). Arrays created lazily by write instead use the written tensor's dtype.

• kwargs (Any, default: {} ) –

  Accepted for IOBackend compatibility (e.g. ZarrBackend's data=) but not supported here, array content comes from write and creation options from the constructor. Ignored with a warning.

write

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

Write data

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.AsyncZarrBackend

IO Backend PerformanceLeverage different IO backends for storing inference results.