feat(patching): give the read-ahead up before refusing a sweep

konfai/data/patching.py

@@ -3194,19 +3194,27 @@ 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
 
     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,

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 {

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"