Skip to content

Read training data with the zarrs codec pipeline - #1422

Open
oliverwm1 wants to merge 1 commit into
feature/cache-zarr-array-handlesfrom
feature/zarrs-codec-pipeline
Open

Read training data with the zarrs codec pipeline#1422
oliverwm1 wants to merge 1 commit into
feature/cache-zarr-array-handlesfrom
feature/zarrs-codec-pipeline

Conversation

@oliverwm1

@oliverwm1 oliverwm1 commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Stacked on #1421; review that first. The base branch is feature/cache-zarr-array-handles, so this diff shows only the zarrs change.

zarrs-python provides a codec pipeline that decodes chunks in Rust and does its own chunk IO, rather than going through zarr-python's per-chunk buffer and store handling. zarr-python still parses metadata, resolves the selection and computes the per-chunk work list; only fetch, decode and scatter move to Rust.

Measured end-to-end on a local zarr store, 51 variables, 8 dataloader workers, warm page cache, medians of interleaved A/B pairs:

sample length zarr-python zarrs speedup
3 timesteps 2763 MB/s 4392 MB/s 1.59x
30 timesteps 1570 MB/s 3333 MB/s 2.12x

The pipeline is bound to an array when that array is opened, so setting it around the array open in _get_async_array applies it to the arrays read here and leaves zarr's global pipeline untouched. Everything else in the process, including the inference data writers in fme.ace.inference.data_writer.zarr, keeps zarr-python's default pipeline. There is a test asserting this.

Reads were verified bitwise identical to zarr-python's pipeline across 212 arrays spanning four sample offsets and all 51 variables.

Changes:

  • pyproject.toml, constraints.txt: add zarrs>=0.2.3. Wheels are published for manylinux_2_28_x86_64, which covers the CUDA base image.
  • fme.core.dataset.utils._get_async_array: open arrays under a scoped zarr.config setting the codec pipeline to zarrs. Because array handles are cached, this runs once per (file, variable) rather than per read.

Two things a reviewer should weigh:

These numbers are warm page cache on a local filesystem, so they measure decode and Python overhead rather than IO. On Weka or GCS, where IO is a larger share of the time, the gain will be smaller, and zarrs does its own IO so its behavior against those stores is not exercised here. This is worth confirming on Weka before relying on the speedup; there is no runtime switch, so turning it off means reverting.

zarrs falls back silently. For dtypes or selections it does not support (variable-length strings, datetimes, some fancy indexing) it defers to BatchedCodecPipeline rather than failing, so an unsupported dataset would quietly read at the old speed. Our data is float32 with basic slicing, and the fallback was confirmed not to trigger. zarr.config's codec_pipeline.strict can turn the fallback into an error if we later want that.

  • Tests added
  • If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated

zarrs decodes chunks in Rust and does its own chunk IO, rather than going
through zarr-python's per-chunk buffer and store handling. Measured 1.59x
end-to-end throughput on a local store with 3 timesteps per sample, and
2.12x with 30 timesteps per sample.

The pipeline is bound to an array when it is opened, so it is scoped to the
arrays opened for reading training data and leaves the pipeline used
elsewhere, such as by the inference data writers, unchanged.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant