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.

Author: vboussot

konfai/data/patching.py

@@ -3209,19 +3209,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``
@@ -3521,24 +3529,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,

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"