fix(omezarr): self-describing streamed stores, typed fields, no spill - #73
Conversation
create_ome_zarr_store composed its own multiscales entry by hand, frozen at version 0.4 with the component axis typed "channel". write_ome_zarr had meanwhile moved to ngff-zarr and can type a displacement field (NGFF RFC-5); the streaming path could not. So the same DVF came out self-describing when it fit in memory, and as an anonymous 3-channel image when it did not -- and streaming is precisely the path taken when it does not fit. Hand-written spec metadata drifts, and this is what the drift cost. Both paths now take their metadata from ngff-zarr, described from a one-voxel stand-in rather than from an array of the target shape. With a single resolution level the metadata does not depend on the extent at all, and the attributes come out byte-identical either way (verified for 0.4 and 0.6). Handing over the real shape costs a pass over every chunk of an array that is entirely zeros: 33.4 s for a 13.6 GiB field, to write no bytes. From the stand-in, 0.027 s. The real array is created underneath that metadata, which is also what makes its chunking exactly the caller's -- the region grid is the one thing a writer given only a shape cannot infer, and a store whose chunks straddle it turns every region write into a read-modify-write. Three failure modes here are silent rather than loud, so each is pinned by a test: - to_ngff_zarr defaults to version 0.5, which zarr-python 2 cannot write. 0.4 is now explicit; without it the Python 3.10 CI leg fails at run time, mid-prediction, with no fallback. - to_ngff_zarr(overwrite=True) reopens the root in mode="w" and drops every attribute it finds, so the KonfAI sidecar is written after it. Written before, Direction disappears, the reader substitutes identity, and a plausibly oriented volume comes back. - the store carries a consolidated index that readers trust over the arrays themselves, and ngff-zarr writes it while the array is still the stand-in. It is rebuilt last, after the sidecar. The level-0 key now comes from the metadata rather than a literal, having moved from "0" to "scale0/image", and _OmeZarrDataStream._close clears the path-keyed read memo after the rename: a replaced store can now differ from its predecessor down to that key, so a stale entry would point at a component that is no longer there. dask is declared explicitly, being imported directly rather than only reached through ngff-zarr. 1134 unit tests pass, plus 7 new ones covering guarantees previously reached only end to end.
The goal is a registration driven by config alone, ending in an RFC-5 OME-Zarr displacement field.
Two things stood in the way, and both were on the write side.
data_to_file typed a field only when handed a sitk.Transform. The predictor emits arrays, so its
output was never recognised -- and a field too large to hold in memory cannot be wrapped in a
transform at all, since building one means assembling the volume that streaming exists to avoid. The
attributes now carry the declaration instead, which is the same marker a store reports when read
back, so the write and read sides finally speak of it the same way.
create_ome_zarr_store could not type a component axis at all. It can now, sharing the capability
check with write_ome_zarr rather than repeating it: RFC-5 needs a zarr v3 store, and asking for a
field where that is impossible raises instead of quietly writing an untyped 0.4 one. A store that
silently is not a field gets found out much later, by a reader that took its three channels for an
image.
The predictor learns nothing about displacement fields. It gains a generic `attributes` mapping that
an output stamps onto what it writes, applied over what it inherited from its input -- inheritance
carries the geometry, which is right, but also whatever the SOURCE entry said it was, which describes
the source and not this. Declared last so config wins; an empty value drops an inherited key.
Verified end to end: a field written region by region comes back as a DisplacementFieldTransform,
axes ('c', 'displacement'), NGFF 0.6, pixels unchanged.
write_ome_zarr is handed an array that is in memory. to_multiscales then decides for itself whether to serialise it to ~/.cache/ngff-zarr first, "to limit memory consumption" -- which cannot lower a peak that has already happened. It buys a full write and a full read, and nothing else. The estimate that triggers it is inflated on top of that: ngff-zarr's memory_usage multiplies itemsize once per dimension, so it overstates by itemsize^(ndim-1) -- 64x for a 4-D float32 array. A 13.6 GiB displacement field is reported as 868 GiB and takes the cache path every time, on exactly the writes that can least afford it.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 14 minutes Your organization has reached its usage spending cap. Adjust your spending cap in the billing tab. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (6)
Comment |
|
@coderabbitai review |
✅ Action performedReview finished.
|
1 participant
Three commits closing the gap between the whole-array and streaming OME-Zarr write paths:
to_ngff_zarr(overwrite=True)(which drops attributes — losing Direction is silent), consolidated index rebuilt last.data_to_filenever recognised a field; the declaration now travels in the attributes (the same marker a store reports on read-back), and a genericattributesmapping lets a config stamp its outputs, applied over inheritance so config wins. Asking for a field where RFC-5 is impossible (zarr 2) raises instead of quietly writing an untyped store.cache=False, since a peak that already happened cannot be lowered.daskis declared explicitly (imaging/omezarr extras, dev, pixi) — it was imported directly but only reached through ngff-zarr.Verified: 118 tests green across store creation, displacement-field round-trip and dataset streaming; end to end a region-written field reads back as a
DisplacementFieldTransform, axes('c','displacement'), NGFF 0.6, pixels unchanged.