Read training data with the zarrs codec pipeline - #1422
Open
oliverwm1 wants to merge 1 commit into
Open
Conversation
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.
oliverwm1
force-pushed
the
feature/zarrs-codec-pipeline
branch
from
August 7, 2026 00:00
26f0936 to
07c8ee5
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #1421; review that first. The base branch is
feature/cache-zarr-array-handles, so this diff shows only the zarrs change.zarrs-pythonprovides 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:
The pipeline is bound to an array when that array is opened, so setting it around the array open in
_get_async_arrayapplies it to the arrays read here and leaves zarr's global pipeline untouched. Everything else in the process, including the inference data writers infme.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: addzarrs>=0.2.3. Wheels are published formanylinux_2_28_x86_64, which covers the CUDA base image.fme.core.dataset.utils._get_async_array: open arrays under a scopedzarr.configsetting 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
BatchedCodecPipelinerather 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'scodec_pipeline.strictcan turn the fallback into an error if we later want that.