diff --git a/.gitignore b/.gitignore index 638dd49f..924049f6 100644 --- a/.gitignore +++ b/.gitignore @@ -93,5 +93,8 @@ tests/_v01 # Benchmarks benchmark/* +# Local planning notes (maintainer checklists, kept out of the repo on purpose) +/V1_*.md + # Claude settings .claude diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e0cf69da..9957cbb5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ ci: # Vendored third-party files, kept byte-identical to what upstream ships. The SIL OFL # text in particular should not be rewritten by the whitespace hooks. -exclude: ^brand/fonts/ +exclude: ^docs/assets/fonts/ repos: - repo: https://github.com/pre-commit/pre-commit-hooks @@ -47,11 +47,6 @@ repos: args: [--fix] # may also add '--unsafe-fixes' - id: ruff-format - - repo: https://github.com/kynan/nbstripout - rev: 0.9.1 - hooks: - - id: nbstripout - # `ty` is deliberately not a hook here: it needs the locked pixi environment to # resolve ngio's dependencies, which neither the `lint` job nor pre-commit.ci # provides. It runs in the `typecheck` job in ci.yml instead, which `deploy` diff --git a/CHANGELOG.md b/CHANGELOG.md index 57e50bc0..a4c5dd7c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,13 +24,13 @@ First stable release. Everything deprecated in `v0.5.0` (each warned "will be re Same call, different result: -- `derive_image` now inherits `dtype`, `dimension_separator` and `compressors` from the reference image instead of forcing `uint16`, `"/"` and `"auto"` — deriving from a `float32` image no longer silently downcasts it. +- `derive_image` inherits `dtype`, `dimension_separator` and `compressors` from the reference image instead of forcing `uint16` — deriving from a `float32` image no longer silently downcasts it. - `add_table` and `write_table` keep the source table's backend instead of rewriting it as `anndata_v1` ([#207](https://github.com/BioVisionCenter/ngio/issues/207)). Pass `backend=` to convert. -- Opening a container no longer reads every pyramid level (`validate_arrays=False` by default), so a missing or malformed array fails on first access rather than at open. Pass `validate_arrays=True` for the old eager check. +- Opening a container no longer reads every pyramid level (`validate_arrays=False` by default), so a bad array fails on first access rather than at open. - `open_image` and `open_label` default to `strict=False`, matching every other getter. - `list_roi_tables` returns `[]` instead of raising when there are no tables. - `get_masked_label(path=...)` resolves the masking label at the label's own pixel size, matching `get_masked_image`. -- `PixelSize`s with different `time_unit`s now compare unequal, and `==` against a non-`PixelSize` returns `NotImplemented` instead of raising `TypeError`. +- `PixelSize`s with different `time_unit`s now compare unequal, and `==` against a non-`PixelSize` returns `NotImplemented`. ### Deprecated, removal in `ngio=1.1` @@ -52,34 +52,32 @@ plate.get_images(max_workers=8) # was get_images_ ### Features - The iterators are stable API: `from ngio import SegmentationIterator`. -- Configurable IO retries: `NgioConfig.io_retry` (`max_retries`, constant/linear/exponential backoff, error matching) plus the `ngio.utils.retry_io` decorator. ngio's own `NgioError`s are never retried. See the Configuration page. -- `ngio.utils.NgioStore` wraps every zarr store ngio opens and applies that retry policy to all IO — metadata, pixel data, and lazy dask reads on workers. `ZipStore` is now supported. +- Configurable IO retries: `NgioConfig.io_retry` plus the `ngio.utils.retry_io` decorator. ngio's own `NgioError`s are never retried. See the Configuration page. +- `ngio.utils.NgioStore` wraps every zarr store ngio opens and applies that retry policy to all IO. `ZipStore` is now supported. - `max_workers=` on the sync plate and table APIs replaces the separate async surface; `None` keeps the serial behaviour. - A larger public namespace, including `MaskedImage`, `MaskedLabel`, `Channel`, `S3FSConfig`, `derive_ome_zarr_plate`, `__version__`, the `get_ngio_*_meta` readers and every error class. `AbstractBaseTable`, `ImplementedTables` and `write_table` are exported from `ngio.tables`, so a custom table type can be registered without private imports. - `NgioTableValidationError` now subclasses `NgioValidationError`, so `except ValueError` catches it like its siblings; new `NgioKeyError`. ### Fixes -- Windows: concurrent access to a store no longer fails with `PermissionError: [WinError 5]`/`[WinError 32]`. Windows refuses to replace or remove a file while another handle to it is open, so a concurrent *reader* of `zarr.json` could break a writer's atomic rename — including between parallel `atomic_add_image` workers, which ngio's lock cannot prevent since opening a group reads metadata before any lock exists. Store operations now absorb these transient conflicts with a short bounded retry, always on and independent of `io_retry`; the original error is raised once the bound is reached. No behaviour change on Linux or macOS. +- Dask writes could silently drop data. `da.store(..., lock=False)` let two blocks read-modify-write the same chunk — or shard, when the target is sharded — concurrently, losing one update. This hit every region write that was not chunk-aligned and every sharded target, including pyramid consolidation. All `da.store` calls now share a lock; block compute stays parallel. +- Windows: concurrent access to a store no longer fails with `PermissionError: [WinError 5]`/`[WinError 32]`. A concurrent reader of `zarr.json` could break a writer's atomic rename; store operations now absorb these transient conflicts with a short bounded retry. No behaviour change on Linux or macOS. - `import ngio` no longer raises `AttributeError` when an s3fs older than 2026.2.0 is installed. - `concatenate_image_tables` built a wrong index: unnamed, and duplicated under `mode="lazy"`. - `Roi.union`/`intersection` dropped ROI name `""` and label `0`; `Roi.from_values` now validates its inputs. - Plate and well metadata `add_*`/`remove_*` mutated the receiver instead of returning a copy. - `AxesSetup.from_ordered_list` silently dropped a non-canonical axis in some orders. - Grid iterator ROIs now get unique names, and `by_chunks` with overlap ≥ chunk size raises `NgioValueError`. -- Also: empty `RoiTable` usability, duplicate ROI names across a roundtrip, labels on read-only images with no `labels` group, `OmeZarrPlate.get_well` caching, and negative indices inside a slicing sequence. ### Packaging -- Ship `src/ngio/py.typed`. The `Typing :: Typed` classifier was declared since 0.x but the PEP 561 marker was missing, so downstream type checkers ignored ngio's annotations. -- Real lower bounds on every dependency, installed and exercised by a `test-min-deps` CI leg: `zarr>=3.1.6`, `numpy>=2.0`, `fsspec>=2025.3`, `anndata>=0.12.5`, `ome-zarr-models>=1.4` and the rest. -- `pandas` 3.x and `anndata` 0.13 are now allowed, the `requires-python` upper cap is gone, and unused `requests`/`distributed` are dropped. -- New `s3` extra: `pip install ngio[s3]`. The README advertised S3 streaming but `s3fs` was never declared. +- Ship `src/ngio/py.typed`. The PEP 561 marker was missing, so downstream type checkers ignored ngio's annotations. +- Real lower bounds on every dependency, exercised by a `test-min-deps` CI leg: `zarr>=3.1.6`, `numpy>=2.0`, `fsspec>=2025.3`, `anndata>=0.12.5`, `ome-zarr-models>=1.4` and the rest. `pandas` 3.x and `anndata` 0.13 are now allowed, the `requires-python` upper cap is gone, and unused `requests`/`distributed` are dropped. +- New `s3` extra: `pip install ngio[s3]`. -### Docs and internal +### Docs -- Docs rebuilt on [Zensical](https://zensical.org) with every code block executed at build time, plus new landing, glossary and Configuration pages. -- `ty` and the docs build now run in CI, coverage is up from 91% to 95%, and concrete-store dispatch is centralized behind `NgioStore`. +- Rebuilt on [Zensical](https://zensical.org) with every code block executed at build time, plus new landing, glossary and Configuration pages. ## [v0.5.14] diff --git a/V1_TODO.md b/V1_TODO.md deleted file mode 100644 index b11e944a..00000000 --- a/V1_TODO.md +++ /dev/null @@ -1,196 +0,0 @@ -# ngio 1.0 release checklist - -Working list of what remains before tagging `v1.0.0`. Grouped by whether 1.0 freezes the -decision — the semver promise in the README means anything in "Blockers" or "API surface" -becomes expensive to change afterwards. - -Not a public document: this is a maintainer checklist, kept at the repo root rather than -under `docs/` because every `.md` in `docs/` is published as a site page. - -## Done - -- [x] Remove every API deprecated in v0.5 that warned "will be removed in `ngio=0.6`", - with an `### API Breaking Changes` table and a migration guide in the changelog. -- [x] Promote the iterators out of `ngio.experimental` into `ngio.iterators`, re-export them - at top level, and leave `ngio.experimental.iterators` as a deprecation shim - (scheduled for removal in 1.1). - ---- - -## Blockers - -Wrong or misleading at release time. - -- [ ] **Add `src/ngio/py.typed`.** `pyproject.toml` declares the `Typing :: Typed` - classifier but the marker file does not exist, so every downstream type checker - silently ignores ngio's annotations. One empty file; hatchling picks it up. -- [ ] **`Development Status :: 3 - Alpha` → `5 - Production/Stable`** (`pyproject.toml`). -- [ ] **Set `major_version_zero = false`** (`[tool.commitizen]`). While it is `true`, - commitizen treats the project as 0.x and will not propose a major bump. -- [ ] **Decide the config-loading story** and make the code, the changelog and - `docs/getting_started/7_configuration.md` agree. `get_config()` is lazy, but - `ngio.utils._zarr_utils` applies the s3fs config at module scope, so `import ngio` - materializes the singleton and `NGIO_CONFIG_PATH` set afterwards is ignored. The - docs and docstring now describe the real behaviour; the open choice is whether to - make the s3fs refresh lazy so the env var can be set after import. - -- [ ] **Fix dependency bounds** (`pyproject.toml`): - - `zarr>3` excludes zarr 3.0.0 itself; the code needs `WrapperStore` and the - `is_group_listable` fix relies on zarr ≥ 3.1.6 behaviour. Declare a real floor. - - No lower bound at all on `numpy`, `filelock`, `scipy`, `fsspec`, `pydantic`, - `aiohttp`, `dask`, `ome-zarr-models`, `pooch`, `polars`, `pyarrow`, `pillow`. At - minimum `pydantic>=2` (the code is v2-only), `polars>=1.0`, `pyarrow>=15`, - `numpy>=1.26`, plus a floor on the pre-1.0 `ome-zarr-models`. - - `anndata<0.13.0` and `pandas<3.0.0` will both block installs soon after release. - -- [ ] **Drop unused hard dependencies**: `requests` and `distributed` are never imported - anywhere in `src/`, `tests/` or `docs/`. Collapse `dask[array]` + `dask[distributed]` - into a single entry. -- [ ] **Add an `s3` extra.** The README advertises S3 streaming and there is an `s3fs` - config section, but `s3fs` is not in `[project.optional-dependencies]` — it only - arrives transitively via the `test` extra, so following the README raises - `ImportError`. -- [ ] **Rename the misspelled public export `conctatenate_tables`** → - `concatenate_tables` (`ngio.images`). 1.0 freezes the typo otherwise. - -## API surface — one-way doors - -Freezing these is the point of 1.0. - -- [ ] **`derive_image` ignores its own documented contract.** It hardcodes - `dtype="uint16"` (and `dimension_separator`, `compressors`) while documenting - "the value from the reference image will be used", so deriving from a `float32` - image silently downcasts. `derive_label` and the underlying - `derive_image_container` correctly use `None` sentinels. - -- [ ] **`get_masked_label` resolves the masking label at the wrong pixel size** — it - passes the raw `pixel_size` argument where `get_masked_image` passes the resolved - `image.pixel_size`, so with `path=` and no explicit pixel size the two disagree. -- [ ] **`list_tables` returns `[]` but `list_roi_tables` raises**, on both - `OmeZarrContainer` and `OmeZarrPlate`. The same empty state also raises two - different types (`NgioValidationError` vs `NgioValueError`), and `list_roi_tables` - returns the same tables in reversed order on the two classes. -- [ ] **`strict` default flips** between the free functions (`open_image`/`open_label`, - `True`) and every method form (`False`). -- [ ] **Pick one `pixel_size` story.** `pixel_size` is deprecated on the `derive*` paths - in favour of `pixelsize`, but remains live on the getters (`get_image`, `get_label`, - `open_image`, …) with a different type. - -- [ ] **Decide what is public.** Currently importable only from private modules: - `MaskedImage`/`MaskedLabel` (the *return types* of `get_masked_image`/ - `get_masked_label`), `AbstractImage`, `AbstractBaseTable`, `write_table`, - `Channel`, `S3FSConfig`, `derive_ome_zarr_plate`, `MapperProtocol`/`BasicMapper`, - the `get_ngio_*_meta` getters (all six `update_*` setters are exported), - `get_sample_info`. Also: no error classes and no table types are exported at top - level, and `__version__` is missing from `__all__`. - -- [ ] **Settle naming inconsistencies**: `set_axes_unit` vs `set_axes_units`; - `open_ome_zarr_container` vs `create_empty_ome_zarr`; `validate_arrays` vs - `validate_paths` (same flag, different name *and* default); `mode` meaning three - unrelated things; `derive_image(ref_path=)` vs `derive_label(ref_image=)`; - `ImagePyramidBuilder` still taking the deprecated `levels_paths` spelling. - -- [ ] **Settle the exception hierarchy.** `NgioTableValidationError` is the only ngio - error that does not also subclass a builtin, so `except ValueError` catches its - siblings but not it. `NgioValidationError` and `NgioValueError` have identical bases - and undocumented boundaries. There is no `NgioKeyError`. -- [ ] **`TableBackend` alias is inert**: the `| str` member makes the `Literal` useless - for checking, and the default `"anndata_v1"` is not in the `Literal` at all. - -- [ ] **Decide the async surface.** The only async entry points are 6 methods on - `OmeZarrPlate`, all `asyncio.to_thread` + unbounded `asyncio.gather` (a 384-well - plate fans out 384 threads — likely the cause of #172). Either complete the surface - or drop them in favour of `max_workers=` on the sync methods. Adding `max_workers` - later is non-breaking; having `_async` methods at all is not. - - -## Correctness / robustness - -- [ ] Replace the 66 raw `ValueError`/`TypeError`/`KeyError`/`RuntimeError` raises in - `src/` with typed ngio errors, so `except NgioError` works. Worst cluster: - `images/_table_ops.py` (10, all in exported functions). -- [ ] Convert the 19 `assert` statements in data paths to real errors — they vanish - under `python -O`. - -- [ ] Resolve the `MaskedImage`/`MaskedLabel` override incompatibility: `ty` reports 4 - `invalid-method-override` errors and 4 `# type: ignore`s already acknowledge it, - which contradicts the docs' claim that every `Image` method works on them. -- [ ] Narrow the `except Exception` in the metadata version-probing loop - (`ome_zarr_meta/_meta_handlers.py`) — it currently converts any decoder bug into - "Failed to decode metadata". -- [ ] Document or close the `ZipStore` hole: the pyarrow/parquet backend raises - `NotImplementedError` for both read and write, though `ZipStore` is advertised as - supported. - -## Enforcement — do these early - -These are what stop the lists above from regrowing. - -- [ ] **Run `ty` in CI** (and/or pre-commit). It is currently run nowhere, so its - ~120 diagnostics are ungated despite the project claiming to be typed. -- [ ] **Build the docs on pull requests.** `docs.yml` has no `pull_request` trigger and - `ci.yml` runs neither `test_snippets` nor `build_docs`, so a broken snippet or - cross-reference only fails after merge to `main` — where it also blocks the - dev-docs deploy. -- [ ] **Gate `deploy` on `lint` and `check-manifest`**, not just `test`. Consider - `environment: pypi`, `twine check dist/*`, and asserting the tag matches the built - version. -- [ ] **Make the docstring rules actually fire.** Ruff's `D1xx` treats every symbol in a - `_`-prefixed module as private, and every implementation module here is - `_`-prefixed — so missing-docstring is silently disabled package-wide. Public API - docstring coverage is ~81%, and only ~69% style-complete. -- [ ] Add a minimum-dependency-version CI leg once the floors above exist, otherwise the - declared bounds stay untested. - -## Docs - -- [ ] Document `open_ome_zarr_container` properly — the first function in every tutorial - has a one-line docstring and no parameter documentation. Same for `open_table`, - `open_table_as`, `open_tables_container`. -- [ ] Add docstrings to `NgioStore` (17 undocumented public methods), `NgioCache`, - `ZoomTransform`, `FeatureTable`, `GenericRoiTable`, and `ngio.io_pipes` — the - `io_pipes` API page currently renders ~40 bare signatures. -- [ ] Add API reference pages for `ngio.config` (the backoff strategies are used in the - configuration guide but documented nowhere), `ngio.ome_zarr_meta`, `ngio.resources`, - `ngio.tables.backends`, `ngio.tables.v1`. -- [ ] Fix the landing-page example: `get_table(...).get(...)` does not type-check because - `get_table` returns the base `Table`. Either narrow/overload it or route the docs - through `get_roi_table`. -- [ ] Sweep the copy-paste docstrings that will render verbatim: three classes documented - as "Placeholder class for a label", `Image.consolidate` saying "the label", the - `OmeZarrPlate` table methods saying "in the image", `v05/__init__.py` saying "v0.4", - `rows`/`columns` saying "the number of" while returning lists. -- [ ] Add a code example to the README, and confirm the advertised conda-forge - availability. -- [ ] Verify `Label`'s eight aliased data-access methods actually render in the API - reference — they are class-body assignments to underscore-named functions and - mkdocstrings filters `!^_`. - -## Release mechanics - -- [ ] Merge the duplicate `### Fix` and `### Chores` headings in `[Unreleased]` and - retitle it to `## [v1.0.0]`. -- [ ] Populate `version:` and `date-released:` in `CITATION.cff` (the README points users - there to cite ngio); consider a Zenodo DOI. -- [ ] Reconcile the copyright holder: `LICENSE` says "2023, Lorenzo Cerrone", - `mkdocs.yml` says "2024-2026, BioVisionCenter UZH". -- [ ] Metadata polish: weak `description`, no `keywords`, `[project.urls]` missing - Documentation/Issues/Changelog (so PyPI has no docs link), PEP 639 `license` form, - missing science/bio-imaging classifiers, `requires-python` upper cap `<3.15`. -- [ ] Confirm every `/stable/` docs link resolves after the first `v*` tag deploys, and - note that the `^v[0-9.]+$` stable regex means a release candidate deploys no docs - version at all. -- [ ] Decide whether the published `dev` extra should stay — it pulls - `napari`/`pyqt5`/`notebook`, so `pip install ngio[dev]` fails on 3.13/3.14. -- [ ] Add an sdist include/exclude; the sdist currently ships `brand/`, `tests/data/` - and `docs/`. - -## Deferred past 1.0 - -- Optional table-backend dependencies (#173) — 1.0 ships with the current dependency set. -- `PointsTable` (#177). Worth confirming first that the table-extension mechanism - (`AbstractBaseTable`, `ImplementedTables`, `from_handler`) is exported and documented, - since that part *is* frozen by 1.0. -- Object-detection iterators (#111). -- Splitting the large modules (`_ome_zarr_container.py`, `_plate.py`, `_image.py`, - `_abstract_image.py`) — deprecation removal already shrank the first two. diff --git a/brand/logo-dark.svg b/brand/logo-dark.svg deleted file mode 100644 index cc0a9acf..00000000 --- a/brand/logo-dark.svg +++ /dev/null @@ -1,16 +0,0 @@ - - ngio - - - - - - - - - - - - - - diff --git a/brand/logo-lockup-dark.svg b/brand/logo-lockup-dark.svg deleted file mode 100644 index 66d5f786..00000000 --- a/brand/logo-lockup-dark.svg +++ /dev/null @@ -1,23 +0,0 @@ - - ngio - ngio horizontal lockup. Wordmark is Space Grotesk SemiBold, outlined, tracking -0.03em. Groups: #mark, #wordmark. - - - - - - - - - - - - - - - - - - - - diff --git a/brand/logo-lockup-stacked-dark.svg b/brand/logo-lockup-stacked-dark.svg deleted file mode 100644 index 22a6502b..00000000 --- a/brand/logo-lockup-stacked-dark.svg +++ /dev/null @@ -1,23 +0,0 @@ - - ngio - ngio stacked lockup. Mark centred over the wordmark, 0.28em of separation. - - - - - - - - - - - - - - - - - - - - diff --git a/brand/logo-lockup-stacked.svg b/brand/logo-lockup-stacked.svg deleted file mode 100644 index 259ebc75..00000000 --- a/brand/logo-lockup-stacked.svg +++ /dev/null @@ -1,23 +0,0 @@ - - ngio - ngio stacked lockup. Mark centred over the wordmark, 0.28em of separation. - - - - - - - - - - - - - - - - - - - - diff --git a/brand/logo-lockup.svg b/brand/logo-lockup.svg deleted file mode 100644 index c29a513e..00000000 --- a/brand/logo-lockup.svg +++ /dev/null @@ -1,23 +0,0 @@ - - ngio - ngio horizontal lockup. Wordmark is Space Grotesk SemiBold, outlined, tracking -0.03em. Groups: #mark, #wordmark. - - - - - - - - - - - - - - - - - - - - diff --git a/brand/logo-mono.svg b/brand/logo-mono.svg deleted file mode 100644 index 121af638..00000000 --- a/brand/logo-mono.svg +++ /dev/null @@ -1,7 +0,0 @@ - - ngio - - - - - diff --git a/brand/logo.svg b/brand/logo.svg deleted file mode 100644 index f0748369..00000000 --- a/brand/logo.svg +++ /dev/null @@ -1,16 +0,0 @@ - - ngio - - - - - - - - - - - - - - diff --git a/brand/fonts/OFL.txt b/docs/assets/fonts/OFL.txt similarity index 100% rename from brand/fonts/OFL.txt rename to docs/assets/fonts/OFL.txt diff --git a/brand/fonts/SpaceGrotesk[wght].ttf b/docs/assets/fonts/SpaceGrotesk[wght].ttf similarity index 100% rename from brand/fonts/SpaceGrotesk[wght].ttf rename to docs/assets/fonts/SpaceGrotesk[wght].ttf diff --git a/src/ngio/common/_locks.py b/src/ngio/common/_locks.py new file mode 100644 index 00000000..e6e06551 --- /dev/null +++ b/src/ngio/common/_locks.py @@ -0,0 +1,16 @@ +"""Locks shared across ngio's write paths.""" + +from dask.utils import SerializableLock + +# Serialises the store flushes of every `da.store` call in ngio. +# +# `da.store(..., lock=False)` gives each dask block its own write task. A block +# whose footprint does not align with the target's write unit — the chunk, or +# the shard when the array is sharded — makes zarr read-modify-write that unit, +# so two such blocks racing on the same key silently lose one update. Holding +# this lock around the flush serialises those writes; block compute stays +# parallel. +# +# `SerializableLock` rather than `threading.Lock` so the graph still pickles +# under a distributed client. +DASK_STORE_LOCK = SerializableLock() diff --git a/src/ngio/common/_pyramid.py b/src/ngio/common/_pyramid.py index 806a35c3..921771a1 100644 --- a/src/ngio/common/_pyramid.py +++ b/src/ngio/common/_pyramid.py @@ -8,6 +8,7 @@ import zarr from pydantic import BaseModel, ConfigDict, model_validator +from ngio.common._locks import DASK_STORE_LOCK from ngio.common._zoom import ( InterpolationOrder, _zoom_inputs_check, @@ -45,7 +46,10 @@ def _on_disk_dask_zoom( # re-derives chunks via normalize_chunks(chunks="auto", ...) and warns # (treated as error by our filterwarnings) when the result isn't a # multiple of the zarr target's chunks. da.store writes blocks 1:1. - da.store(target_array, target, lock=False) # type: ignore + # The shared lock serialises the flushes: blocks that only partially cover + # a chunk (or, for a sharded target, a shard) make zarr read-modify-write + # it, and two of them racing on the same key lose an update. + da.store(target_array, target, lock=DASK_STORE_LOCK) # type: ignore def _on_disk_coarsen( @@ -97,8 +101,8 @@ def _on_disk_coarsen( aggregation_function, source_array, coarsening_setup, trim_excess=True ) out_target = out_target.rechunk(target.chunks) - # See _on_disk_dask_zoom for rationale. - da.store(out_target, target, lock=False) # type: ignore + # See _on_disk_dask_zoom for rationale, including the lock. + da.store(out_target, target, lock=DASK_STORE_LOCK) # type: ignore def on_disk_zoom( diff --git a/src/ngio/io_pipes/_ops_slices.py b/src/ngio/io_pipes/_ops_slices.py index aa2fa776..1c874cdc 100644 --- a/src/ngio/io_pipes/_ops_slices.py +++ b/src/ngio/io_pipes/_ops_slices.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict from ngio.common._dimensions import Dimensions +from ngio.common._locks import DASK_STORE_LOCK from ngio.io_pipes._ops_slices_utils import compute_slice_chunks from ngio.ome_zarr_meta.ngio_specs import Axis from ngio.utils import NgioUserWarning, NgioValueError @@ -255,7 +256,11 @@ def set_slice_as_dask( # da.store instead of da.to_zarr: see ngio.common._pyramid for the # dask>=2025.11 PerformanceWarning regression that to_zarr triggers # when the input chunks aren't a multiple of the target's chunks. - da.store(patch, zarr_array, regions=slice_tuple, lock=False) # type: ignore + # The shared lock serialises the flushes: a region write whose blocks + # only partially cover a chunk (or, for a sharded target, a shard) + # makes zarr read-modify-write it, and two of them racing on the same + # key lose an update. + da.store(patch, zarr_array, regions=slice_tuple, lock=DASK_STORE_LOCK) # type: ignore return # Complex case, we have exactly one tuple in the slicing tuple @@ -264,7 +269,7 @@ def set_slice_as_dask( _sub_slice = (*slice_tuple[:ax], slice(idx, idx + 1), *slice_tuple[ax + 1 :]) sub_patch = da.take(patch, indices=i, axis=ax) sub_patch = da.expand_dims(sub_patch, axis=ax) - da.store(sub_patch, zarr_array, regions=_sub_slice, lock=False) # type: ignore + da.store(sub_patch, zarr_array, regions=_sub_slice, lock=DASK_STORE_LOCK) # type: ignore ############################################################## diff --git a/tests/unit/common/test_pyramid.py b/tests/unit/common/test_pyramid.py index 9329be17..bbe06866 100644 --- a/tests/unit/common/test_pyramid.py +++ b/tests/unit/common/test_pyramid.py @@ -1,6 +1,7 @@ from pathlib import Path from typing import Literal +import numpy as np import pytest import zarr @@ -28,3 +29,35 @@ def test_on_disk_zooms( target_array = zarr.create_array(target, shape=(16, 64, 64), dtype="uint8") on_disk_zoom(source_array, target_array, order=order, mode=mode) + + +@pytest.mark.parametrize("mode", ["dask", "coarsen"]) +def test_on_disk_zoom_sharded_matches_unsharded( + tmp_path: Path, mode: Literal["dask", "coarsen"] +): + """A zoom onto a sharded target matches the same zoom onto an unsharded one. + + Both dask paths rechunk to `target.chunks`, which for a sharded array is the + *inner* chunk shape, while writes are atomic per shard object. Several blocks + therefore read-modify-write one shard and race without the shared `da.store` + lock. The unsharded layout is race-free by alignment, so it is the reference. + """ + rng = np.random.default_rng(0) + source_array = zarr.create_array( + tmp_path / "source.zarr", shape=(16, 128, 128), dtype="uint8" + ) + source_array[...] = rng.integers(0, 255, size=(16, 128, 128), dtype="uint8") + + results = {} + for name, shards in (("unsharded", None), ("sharded", (8, 32, 32))): + target_array = zarr.create_array( + tmp_path / f"target_{mode}_{name}.zarr", + shape=(16, 64, 64), + chunks=(4, 16, 16), + shards=shards, + dtype="uint8", + ) + on_disk_zoom(source_array, target_array, order="nearest", mode=mode) + results[name] = target_array[...] + + np.testing.assert_array_equal(results["sharded"], results["unsharded"]) diff --git a/tests/unit/io_pipes/test_dask_write_race.py b/tests/unit/io_pipes/test_dask_write_race.py new file mode 100644 index 00000000..e9fe68dd --- /dev/null +++ b/tests/unit/io_pipes/test_dask_write_race.py @@ -0,0 +1,101 @@ +"""Deterministic repro for the `da.store` chunk write race. + +A region write whose start is not chunk-aligned makes two dask blocks +read-modify-write the same zarr chunk. `GatedStore` below holds the first +reader of the contested chunk until the second one arrives, so the lost update +happens on every run instead of depending on scheduling luck. With the store +flushes serialised the second reader cannot arrive while the first holds the +chunk, so the gate simply times out and the write is correct. +""" + +import asyncio +import contextlib +from typing import Any + +import dask +import dask.array as da +import numpy as np +import pytest +import zarr +import zarr.storage + +from ngio.io_pipes._ops_slices import ( + SlicingOps, + set_slice_as_dask, + set_slice_as_numpy, +) + +# Long enough that a genuinely concurrent second reader always arrives inside +# it, short enough not to slow the suite down when nothing does. +_GATE_TIMEOUT_S = 1.0 + +# The single chunk both blocks partially cover, in zarr v3 default chunk key +# encoding. +_CONTESTED_KEY = "c/0" + + +class GatedStore(zarr.storage.WrapperStore): + """Hold the first read of the contested chunk until a second one arrives.""" + + def __init__(self, store: zarr.storage.MemoryStore) -> None: + super().__init__(store) + self.armed = False + self.arrivals = 0 + # Created lazily: it must belong to zarr's I/O event loop. + self._second_reader: asyncio.Event | None = None + + async def get(self, key: str, prototype: Any, byte_range: Any = None) -> Any: + if self.armed and key == _CONTESTED_KEY: + if self._second_reader is None: + self._second_reader = asyncio.Event() + self.arrivals += 1 + if self.arrivals == 1: + with contextlib.suppress(TimeoutError): + await asyncio.wait_for( + self._second_reader.wait(), timeout=_GATE_TIMEOUT_S + ) + else: + self._second_reader.set() + return await super().get(key, prototype, byte_range) + + +@pytest.fixture +def slicing_ops() -> SlicingOps: + # Chunk 0 spans [0, 16) and chunk 1 spans [16, 32). Writing [4, 20) with + # 8-element dask blocks puts block 0 at [4, 12) and block 1 at [12, 20): + # both partially cover chunk 0. + return SlicingOps( + on_disk_axes=("x",), + on_disk_shape=(64,), + on_disk_chunks=(16,), + slicing_tuple=(slice(4, 20),), + ) + + +def test_misaligned_dask_write_matches_numpy(slicing_ops: SlicingOps) -> None: + baseline = np.arange(1000, 1064, dtype="uint16") + patch = np.arange(1, 17, dtype="uint16") + + reference = zarr.create_array( + store=zarr.storage.MemoryStore(), shape=(64,), chunks=(16,), dtype="uint16" + ) + reference[:] = baseline + set_slice_as_numpy(zarr_array=reference, patch=patch, slicing_ops=slicing_ops) + + gated = GatedStore(zarr.storage.MemoryStore()) + array = zarr.create_array(store=gated, shape=(64,), chunks=(16,), dtype="uint16") + array[:] = baseline + + gated.armed = True + with dask.config.set(scheduler="threads", num_workers=4): + set_slice_as_dask( + zarr_array=array, + patch=da.from_array(patch, chunks=(8,)), + slicing_ops=slicing_ops, + ) + gated.armed = False + + # Locked or not, the contested chunk is read once per block. Fewer means + # the chunk key drifted and the gate never fired, making the test vacuous. + assert gated.arrivals >= 2 + np.testing.assert_array_equal(array[:], reference[:])