From 0736b94dee65cb0686867ccc70d2a4605a4b37e7 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Thu, 27 Aug 2026 22:04:13 +0200 Subject: [PATCH] feat(patching): give the read-ahead up before refusing a sweep The block a sweep holds is mostly what the chain must hold to run at all: the source window it pulls, the rows it lands, the buffers its widest stage allocates. The read-ahead queue is the one part the sizing CHOSE rather than met, and what it buys is wall clock -- half a second of a 6.7 s run, by _sweep_depth's own measurement. A sweep about to refuse has no clock to buy. So it stops buying it. Three resident source regions become one and the search runs once more; only if that does not fit either does it refuse, now naming both figures so a reader knows which one to clear. On a chain whose stage buffers dominate the band is narrow, 1.31x to 1.34x on the tiling fixture; on a pointwise one it is 2.5x. Inside the band the alternative was not a slower run but no run. _sweep_depth had to learn to go down before it goes up: a tile found without the queue, swept with one, would hold what the sizing was never told about -- the budget's promise lost to a default nobody revisited. And _keeps_the_block now asks the search rather than _sweep_tile, which falls back: a depth that cannot afford the block would otherwise come back holding it, and every depth would look affordable. The refusal test's own budget stopped refusing, which is the feature: half of one row's queued price is bought by the serial retry. It is now set under the serial price, and asserts both figures are named. --- konfai/data/patching.py | 78 ++++++++++++++++++++-------- konfai/utils/ome_zarr.py | 91 ++++++++++++++++++++++++++++++--- tests/unit/test_sweep_tiling.py | 37 +++++++++++++- 3 files changed, 177 insertions(+), 29 deletions(-) diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 6adc7084..634f9b5c 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -3194,6 +3194,11 @@ def _sweep_depth( second block in flight recovers 0.5 s of a 6.7 s run and a third recovers none. """ depth = _sweep_pipeline_depth() + # DOWN BEFORE UP. `tile` may be the one the sizing found only after giving the queue up + # (:meth:`_sweep_tile`), and a run that kept the queue anyway would hold what the sizing was + # never told about -- the budget's whole promise, lost to a default nobody revisited. + while depth and not self._keeps_the_block(spatial, channels, plans, tile, depth): + depth -= 1 while depth and depth < _SWEEP_MAX_DEPTH and self._keeps_the_block(spatial, channels, plans, tile, depth + 1): depth += 1 return depth @@ -3201,12 +3206,15 @@ def _sweep_depth( def _keeps_the_block( self, spatial: list[int], channels: int, plans: Sequence["_ReadStagePlan"], tile: list[int], depth: int ) -> bool: - """Whether a queue of ``depth`` still leaves the sweep exactly ``tile``: a deeper one the - budget cannot hold refuses, and a refusal here is the answer no, not the sweep's failure.""" - try: - return self._sweep_tile(spatial, channels, plans, depth) == tile - except DatasetManagerError: - return False + """Whether a queue of ``depth`` both affords ``tile`` and still picks it. + + Asked of the search and not of :meth:`_sweep_tile`, which falls back to no queue at all: a + depth that cannot hold the block would come back holding it, and every depth would look + affordable. + """ + budget = self._sweep_budget_bytes + found, held = self._tile_within(spatial, channels, plans, depth, budget) + return found == tile and (not budget or budget <= 0 or held <= budget) def read_granularity(self) -> tuple[int, ...] | None: """The stored block this case's source reads are served in, spatial axes only, or ``None`` @@ -3506,24 +3514,54 @@ def _sweep_tile( search is over the height, because that is the one free parameter of the decomposition. """ depth = _sweep_pipeline_depth() if depth is None else depth - cap = self._sweep_rows(spatial, channels, plans, depth) budget = self._sweep_budget_bytes + tile, held = self._tile_within(spatial, channels, plans, depth, budget) + if not budget or budget <= 0 or held <= budget: + return tile + # THE READ-AHEAD IS THE ONE PART OF THE PRICE THE SIZING CHOSE. Everything else in the block + # is what the chain must hold to run at all; the queue is bought, and what it buys is wall + # clock (_sweep_depth: half a second of a 6.7 s run). A sweep about to refuse has no clock to + # buy, so it gives the queue up and asks once more. Three source regions resident become one, + # which is a quarter to a third of the block on a chain whose stage buffers dominate -- a + # narrow band, and inside it the difference is running against not running. + serial = None + if depth > 0: + candidate, serial = self._tile_within(spatial, channels, plans, 0, budget) + if serial <= budget: + return candidate + raise DatasetManagerError( + f"'{self.name}': no region of '{self.group_src}' fits the per-rank memory budget" + f" ({format_bytes(budget)}): the smallest one this chain can sweep holds" + f" {format_bytes(held)}" + + (f", and {format_bytes(serial)} with the read-ahead given up" if serial is not None else "") + + ".", + "Raise 'memory_budget'.", + ) + + def _tile_within( + self, + spatial: list[int], + channels: int, + plans: Sequence["_ReadStagePlan"], + depth: int, + budget: float | None, + ) -> tuple[list[int], int]: + """The best block a sweep of ``depth`` can afford, and what it holds: the search alone. + + No refusal and no fallback, because two callers ask it two different questions -- whether a + deeper queue still buys the same block (:meth:`_keeps_the_block`) and what to do when none + of them fits (:meth:`_sweep_tile`) -- and a search that answered either for them would + answer the other one wrong. + """ + cap = self._sweep_rows(spatial, channels, plans, depth) if not budget or budget <= 0: - return self._sweep_shape(spatial, plans, cap) - # The bisection never takes one row as affordable: the refusal below answers for it. What it - # finds is then judged against the store's own heights, because the price steps rather than - # climbs and bisection lands somewhere affordable, not on the best region the budget buys. + return self._sweep_shape(spatial, plans, cap), 0 + # The bisection never takes one row as affordable: the caller answers for it. What it finds + # is then judged against the store's own heights, because the price steps rather than climbs + # and bisection lands somewhere affordable, not on the best region the budget buys. low = self._rows_within(spatial, channels, plans, depth, budget, cap) tile = self._best_tile(spatial, channels, plans, depth, budget, [low, *self._grid_rows(cap)]) - held = self.sweep_block_bytes(spatial, channels, plans, tile, depth) - if held > budget: - raise DatasetManagerError( - f"'{self.name}': no region of '{self.group_src}' fits the per-rank memory budget" - f" ({format_bytes(budget)}): the smallest one this chain can sweep holds" - f" {format_bytes(held)}.", - "Raise 'memory_budget'.", - ) - return tile + return tile, self.sweep_block_bytes(spatial, channels, plans, tile, depth) def _get_streamed_data( self, diff --git a/konfai/utils/ome_zarr.py b/konfai/utils/ome_zarr.py index 6d162bc5..94363cf4 100644 --- a/konfai/utils/ome_zarr.py +++ b/konfai/utils/ome_zarr.py @@ -39,6 +39,7 @@ import itertools import operator import shutil +import tempfile import threading from collections import OrderedDict from collections.abc import Sequence @@ -48,6 +49,7 @@ import numpy as np +from konfai.utils import uri from konfai.utils.errors import DatasetManagerError from konfai.utils.runtime import map_over_rank_pool @@ -177,6 +179,22 @@ def _konfai_attributes(store_path: str) -> dict[str, Any]: return {} +def _from_ngff_zarr(store_path: str | Path) -> Any: + """ngff-zarr's multiscales for ``store_path``. A remote root goes in as a key-to-bytes mapping + over its own filesystem (ngff-zarr >= 0.44 reads a remote string as a local path) and as its URL + when ngff-zarr refuses the mapping, which older releases resolve themselves.""" + if not uri.is_uri(store_path): + return ngff_zarr.from_ngff_zarr(str(store_path)) + # Through uri.filesystem, so a missing fsspec backend or configuration is the structured + # DatasetManagerError, never a raw dependency error. + filesystem = uri.filesystem(store_path) + _, target = uri.split_scheme(str(store_path)) + try: + return ngff_zarr.from_ngff_zarr(filesystem.get_mapper(target)) + except ValueError: + return ngff_zarr.from_ngff_zarr(str(store_path)) + + @lru_cache(maxsize=8) def _load_image(store_path: str, level: int) -> Any: """Return the ``NgffImage`` for ``level`` of an OME-Zarr store, memoised per (store, level). @@ -194,7 +212,7 @@ def _load_image(store_path: str, level: int) -> Any: """ _require_ngff_zarr() try: - multiscales = ngff_zarr.from_ngff_zarr(str(store_path)) + multiscales = _from_ngff_zarr(store_path) except (KeyError, IndexError, OSError, TypeError, ValueError) as exc: raise DatasetManagerError( f"Cannot open OME-Zarr store '{store_path}' (level {level}).", @@ -258,7 +276,7 @@ def is_displacement_field(store_path: str | Path) -> bool: if not _NGFF_ZARR_AVAILABLE: return False try: - metadata = ngff_zarr.from_ngff_zarr(str(store_path)).metadata + metadata = _from_ngff_zarr(store_path).metadata except Exception: # "Not a displacement field" is the only answer this owes: it is asked purely to decide HOW to # read an entry, and an absent or unreadable store is not one either. @@ -608,7 +626,7 @@ def place(coords: tuple) -> None: def _level_path(store_path: str, level: int) -> str | None: """The zarr path of one level, from the store's multiscales metadata, memoised beside the image.""" try: - datasets = ngff_zarr.from_ngff_zarr(store_path).metadata.datasets + datasets = _from_ngff_zarr(store_path).metadata.datasets return str(datasets[level if len(datasets) > 1 else 0].path) except Exception: return None @@ -624,7 +642,32 @@ def _read_level_window(store_path: str, level: int, image: Any, index: tuple) -> array = _level_array(store_path, level_path) if tuple(array.shape) == tuple(image.data.shape): return _read_chunked(store_path, level_path, array, index) - return np.asarray(image.data[index]) + return _lazy_window(image.data, index) + + +def _lazy_window(data: Any, index: tuple) -> np.ndarray: + """The plain lazy read. A lazy array that refuses a stepped slice (ngff-zarr >= 0.44 wraps the + level in an adapter that takes unit steps only) is read over the unit-step span and stepped here.""" + try: + return np.asarray(data[index]) + except NotImplementedError: + # Each slice normalized against the shape, then its ascending unit-step span; a negative + # step reads that span backwards from its own end, which lands on the indices the original + # slice named. + bounds = tuple( + slice(*item.indices(size)) if isinstance(item, slice) else item + for item, size in zip(index, data.shape, strict=True) + ) + span = tuple( + slice(item.stop + 1, item.start + 1) + if isinstance(item, slice) and item.step < 0 + else slice(item.start, item.stop) + if isinstance(item, slice) + else item + for item in bounds + ) + steps = tuple(slice(None, None, item.step) for item in bounds if isinstance(item, slice)) + return np.asarray(data[span])[steps] def _store_index( @@ -813,6 +856,11 @@ def write_ome_zarr( exist only from 0.6, so a caller passing both could only ever pass them consistently: an invariant worth removing rather than documenting. """ + if scale_factors and uri.is_uri(store_path): + raise DatasetManagerError( + f"Cannot append pyramid levels to the remote store '{store_path}'.", + "Levels are derived in place through local paths; write the store locally and upload it.", + ) array_data = np.asarray(data) # The one write path: the store described and created empty (ngff-zarr's metadata, the # caller's chunking), filled by zarr itself, its levels grafted beside level 0. Handing @@ -876,6 +924,30 @@ def _type_component_axis(multiscales: Any, axis_type: str) -> None: axis.type = axis_type +def _write_skeleton(store_path: str | Path, multiscales: Any, version: str) -> None: + """ngff-zarr's metadata for the store, written in place. ngff-zarr (>= 0.44) writes local + directories only, so a remote root gets the skeleton written locally and uploaded through the + root's own filesystem: a few bytes of metadata, before the array is created underneath it.""" + if not uri.is_uri(store_path): + ngff_zarr.to_ngff_zarr(str(store_path), multiscales, overwrite=True, version=version) + return + filesystem = uri.filesystem(store_path) + _, target = uri.split_scheme(str(store_path)) + if "/" not in target.strip("/"): + raise DatasetManagerError( + f"Refusing to create the store at the filesystem root '{store_path}'.", + "Name a key under the root (e.g. '.../dataset.ome.zarr'): creating a store replaces what its path holds.", + ) + if filesystem.exists(target): + filesystem.rm(target, recursive=True) + filesystem.makedirs(target, exist_ok=True) + with tempfile.TemporaryDirectory() as scratch: + local = Path(scratch) / "skeleton" + ngff_zarr.to_ngff_zarr(str(local), multiscales, overwrite=True, version=version) + for file in sorted(path for path in local.rglob("*") if path.is_file()): + filesystem.put_file(str(file), uri.join(target, file.relative_to(local).as_posix())) + + def create_ome_zarr_store( store_path: str | Path, shape: Sequence[int], @@ -934,7 +1006,7 @@ def create_ome_zarr_store( _type_component_axis(multiscales, _DISPLACEMENT_AXIS_TYPE) version = _RFC5_VERSION # version is explicit because to_ngff_zarr defaults to 0.5, which zarr-python 2 cannot write. - ngff_zarr.to_ngff_zarr(str(store_path), multiscales, overwrite=True, version=version) + _write_skeleton(store_path, multiscales, version) # The level-0 key comes from the metadata rather than a literal: ngff-zarr builds it from the # image name, so "scale0/image" is its convention to change, not ours to hardcode. @@ -1077,13 +1149,18 @@ def append_ome_zarr_levels( KonfAI attribute sidecar is untouched, being a key beside theirs; a displacement field keeps its typed component axis through the same call that types it at creation. """ + if uri.is_uri(store_path): + raise DatasetManagerError( + f"Cannot append pyramid levels to the remote store '{store_path}'.", + "Levels are derived in place through local paths; write the store locally and upload it.", + ) _require_ngff_zarr() if not scale_factors: return store = Path(store_path) clear_ome_zarr_cache(store) field = is_displacement_field(store) - base = ngff_zarr.from_ngff_zarr(str(store)).images[0] + base = _from_ngff_zarr(store).images[0] stored_chunks = tuple(int(size) for size in base.data.chunksize) if downsample_method in (None, "ITKWASM_BIN_SHRINK"): multiscales = _bin_shrink_multiscales(base, _level_zero_scale_factors(scale_factors), stored_chunks) @@ -1152,7 +1229,7 @@ def get_ome_zarr_info(store_path: str | Path, level: int = 0) -> dict[str, Any]: image = _load_image(str(store_path), level) dims = [str(axis).lower() for axis in image.dims] try: - n_levels = len(ngff_zarr.from_ngff_zarr(str(store_path)).images) + n_levels = len(_from_ngff_zarr(store_path).images) except (OSError, TypeError, ValueError): n_levels = 1 return { diff --git a/tests/unit/test_sweep_tiling.py b/tests/unit/test_sweep_tiling.py index 321c6390..910b5cdb 100644 --- a/tests/unit/test_sweep_tiling.py +++ b/tests/unit/test_sweep_tiling.py @@ -253,15 +253,48 @@ def test_a_regrid_pays_for_what_it_pulls_and_not_for_what_it_lands(tmp_path: Pat assert _block_voxels(regrid, plans) < _block_voxels(pointwise, ()) +def test_a_budget_that_only_fits_without_the_queue_gives_the_queue_up(tmp_path: Path) -> None: + """The read-ahead is bought, not owed. A sweep that cannot afford it stops buying it. + + Three source regions are resident with a queue and one without, so a chain whose stage buffers + do not dominate holds a quarter to a third less serially. That band is narrow -- 1.31x to 1.34x + on this fixture -- and inside it the alternative is not a slower run but no run at all. + """ + source, _volume = _sheared_fixture(tmp_path) + resample = Resample(reference="TARGET", reference_group="GRID", reference_dataset=f"{tmp_path / 'ref'}:h5") + manager = _manager(source, [resample, Save(f"{tmp_path / 'out'}:h5")]) + plans = _sweep_plans(manager) + base = _sweep_pipeline_depth() + assert base > 0, "a rank with one core queues nothing and has nothing to give up" + + tile = manager._sweep_shape(list(LANDING), plans, 1) + queued = manager.sweep_block_bytes(list(LANDING), 1, plans, tile, base) + serial = manager.sweep_block_bytes(list(LANDING), 1, plans, tile, 0) + assert serial < queued, "the queue is part of what one row costs" + + # A budget between the two: the queue is what makes it refuse, and nothing else does. + manager.set_memory_budget((queued + serial) / 2.0) + found = manager._sweep_tile(list(LANDING), 1, plans) + assert manager.sweep_block_bytes(list(LANDING), 1, plans, found, 0) <= manager._sweep_budget_bytes + + # And the run walks the depth the sizing solved for, or it holds what it was never priced for. + assert manager._sweep_depth(list(LANDING), 1, plans, found) == 0 + + def test_a_budget_no_region_fits_refuses_with_both_figures(tmp_path: Path) -> None: """A budget one row of the landing does not fit is not a one-row sweep: it is a refusal naming the budget and what the smallest region holds, so the reader knows what to raise it to.""" source, _volume = _sheared_fixture(tmp_path) manager = _manager(source, [Save(f"{tmp_path / 'out'}:h5")]) - manager.set_memory_budget(_priced(manager, (), 1) / 2.0) + # Under what one row holds WITHOUT the queue, since that is the last thing the sizing tries: + # half of the queued price is a budget the serial retry now buys, and buying it is the point. + serial = manager.sweep_block_bytes(list(LANDING), 1, (), manager._sweep_shape(list(LANDING), (), 1), 0) + manager.set_memory_budget(serial / 2.0) - with pytest.raises(DatasetManagerError, match=r"no region of 'CT' fits the per-rank memory budget"): + with pytest.raises(DatasetManagerError, match=r"no region of 'CT' fits the per-rank memory budget") as raised: manager._sweep_tile(list(LANDING), 1) + # Both attempts, so the reader raising the budget knows which figure to clear. + assert "with the read-ahead given up" in str(raised.value) assert manager.stream_refusal(0) is not None, "and the plan routes the case away from streaming"