feat(transform): stream stored displacement fields and price what they read

konfai/data/case_reduction.py

@@ -604,6 +604,17 @@ def _member_region(self, manager: DatasetManager, region: tuple[slice, ...]) ->
     # keeps a reserve for the same reason (Predictor._ACCUMULATE_MARGIN): the measurement is of the
     # region that just ran, and the next one meets an allocator in a different state.
     _MEASURED_MARGIN = 0.9
+    #: The probe's share of the planned height. The probe is the one region that runs BEFORE any
+    #: measurement can bound it, so it is the one region that must not be able to kill the run on
+    #: its own. At the planned height it could: a fold over registration fields held 1.42x, 1.47x
+    #: and 1.50x what its first region was allowed at three budgets, and at an `auto` budget of 77
+    #: GiB that first region reached 90 GiB resident on a 122 GiB host, and the host went down
+    #: before the probe could read anything. The host gives no OutOfMemoryError to catch: the
+    #: kernel kills. A quarter-height probe overshooting by the same 1.5x holds 0.4 of the budget,
+    #: which is survivable, and the ratio it measures is the same one -- the halo does not shrink
+    #: with the region, so a short region over-holds by MORE than a tall one, and a refit from it
+    #: is conservative. Its price is one extra region: seconds, on a fold of minutes.
+    _PROBE_SHARE = 0.25
 
     def _folds(self, spatial: list[int], measure: bool = False):
         """Every region's fold, in order: the loop both passes share.
@@ -619,7 +630,11 @@ def _folds(self, spatial: list[int], measure: bool = False):
         """
         start, refitted = 0, not measure
         while start < int(spatial[0]):
-            stop = min(start + self.slab_rows, int(spatial[0]))
+            # The probe is SHORT. Every later region is cut against what it measured; the probe
+            # itself is cut against nothing, so it is sized so that its own overshoot cannot
+            # reach the host's limit (_PROBE_SHARE).
+            rows = self.slab_rows if refitted else max(1, int(self.slab_rows * self._PROBE_SHARE))
+            stop = min(start + rows, int(spatial[0]))
             region = (slice(start, stop), *(slice(0, extent) for extent in spatial[1:]))
             # Only around the region that is actually the probe. The host meter RESETS the
             # process's resident high-water mark to take its reading, and that mark is what the
@@ -643,26 +658,37 @@ def _refit_to_measurement(self, meter: HeldMeter | None, rows: int, spatial: lis
         bounds the next one from above, exactly as the predictor's gate reads a forward's transient
         from the batch that just ran (:meth:`Predictor._accumulate_device`).
 
-        Against the BUDGET, not the share the sizing aims at: a share is how a height is chosen,
-        and what must not be exceeded is the whole declaration. This exists to prevent a kill, not
-        to shave bytes -- and it cannot shave many, since the region that set the peak has already
-        run and cutting the rest never undoes it. What it catches is a LATER region holding more
-        than the first: a case with a wider halo, a region touching more chunks.
+        Against the whole declaration LESS the chunk cache's share, because that is what the
+        reading covers. A share is how a height is chosen and what must not be exceeded is the
+        declaration -- but the meter no longer counts the decoded-chunk cache (it outlives the
+        region, and charging the region for it cut every region after the probe), so the cache's
+        bytes have to come off the other side of the comparison too. Judged against the whole
+        budget, a reading that excludes the cache lets the cache be spent twice: once inside the
+        allowance, and again by the cache itself. This exists to prevent a kill, not to shave
+        bytes. The probe is a short region (_PROBE_SHARE), so what it held is scaled to
+        the planned height before it is judged: a probe that held its share of the budget says the
+        full region would hold the budget, and a probe that held more says the full region would
+        be the kill this exists to prevent. Only ever shorter: a probe that came in under its
+        share does not talk the fold into a taller region than the plan allowed.
         """
         del spatial
         held = meter.held() if meter is not None else None
         if held is None or not self._budget_bytes or self._budget_bytes <= 0 or rows <= 0 or held <= 0:
             return
-        allowed = float(self._budget_bytes) * self._MEASURED_MARGIN
-        if held <= allowed:
+        cache = budget_share("cache", self._budget_bytes) or 0.0
+        allowed = (float(self._budget_bytes) - cache) * self._MEASURED_MARGIN
+        # What the FULL region would hold, from what the probe held: the halo is a fixed cost the
+        # probe paid in full, so scaling by height over-estimates, which is the safe direction.
+        projected = held * (self.slab_rows / float(rows))
+        if projected <= allowed:
             return
-        fitted = max(1, int(rows * allowed / held))
+        fitted = max(1, int(self.slab_rows * allowed / projected))
         if fitted >= self.slab_rows:
             return
         print(
-            f"[Reduce] '{self.reduce.output}': first region held {format_bytes(held)} of the"
-            f" {format_bytes(allowed)} its {rows} row(s) may hold --"
-            f" the rest are cut to {fitted} row(s).",
+            f"[Reduce] '{self.reduce.output}': a {rows}-row probe held {format_bytes(held)}, so the planned"
+            f" {self.slab_rows} row(s) would hold {format_bytes(projected)} of the {format_bytes(allowed)}"
+            f" allowed -- the rest are cut to {fitted} row(s).",
             flush=True,
         )
         self.slab_rows = fitted

konfai/data/patching.py

@@ -406,7 +406,25 @@ def device_peak() -> int | None:
     if not reset_resident_peak():
         return None
     resident = resident_bytes()
-    return None if resident is None else HeldMeter(peak_resident_bytes, int(resident))
+    if resident is None:
+        return None
+    # THE CACHE IS NOT THE SCOPE'S. A host peak is the whole process's high-water mark, and the
+    # decoded-chunk cache sits inside it: a scope that reads from a store fills the cache on its
+    # way, and the cache keeps what it decoded past the scope, for the next one. Charging the scope
+    # for that is charging it for a budget line that has its own share (BUDGET_SHARES['cache']).
+    # Measured on a fold's probe over ten native members: 24.4 GiB read, 13.2 of it the cache
+    # filling from empty, and the fold cut to 78 % of the height its regions actually needed.
+    from konfai.utils.ome_zarr import chunk_cache_held_bytes
+
+    cache_at_start = chunk_cache_held_bytes()
+
+    def resident_peak_less_cache() -> int | None:
+        peak = peak_resident_bytes()
+        if peak is None:
+            return None
+        return peak - max(0, chunk_cache_held_bytes() - cache_at_start)
+
+    return HeldMeter(resident_peak_less_cache, int(resident))
 
 
 def save_destination(save: Save, default_dataset: Dataset, default_group: str) -> tuple[Dataset, str]:

konfai/data/reduction.py

@@ -191,17 +191,29 @@ class Median(Reduction):
     """
 
     voxel_local = True
-    # ``torch.stack`` copies the buffer and the sort along the case axis returns values and
-    # int64 indices over that: measured at 4x the stack it is handed (6 x 16 MiB float32 cases).
-    # A fold of three, four or five members takes the network instead and costs far less; the
-    # attribute is the worst case, :meth:`working_multiple_for` is what the plan asks.
-    working_multiple = 4.0
-
-    #: What the selection networks below hold beside the members they are handed, measured on a
-    #: 24 MiB member (float32): the sort's own 4.0 is what anything wider still costs.
+    # THE MIDDLE IS SELECTED, NEVER SORTED. A sort along the case axis copies the stack and returns
+    # int64 indices over it -- eight bytes an element whatever the members weigh -- so ten uint16
+    # regions of 33 x 1331 x 1775 (1.45 GiB) sorted at 6.0x their own size, and ten float32 ones
+    # at 4.0x (peak resident above the members, measured). A selection network of element-wise
+    # min/max holds a WINDOW of the k+1 smallest members seen so far, in the averaging dtype, and
+    # inserts each member into it: no stack, no indices, and the members stay in the dtype they
+    # arrived in. Ten uint16 members: 1.8x, against 6.0x. Twice the arithmetic of the sort (3.4 s
+    # against 1.6 on that region) on a fold whose clock is the disk by 40 to 1, and whose regions
+    # the planner may now cut two to three times taller.
+    #
+    # The attribute is the worst case the plan may see; :meth:`working_multiple_for` prices the
+    # network for the count it is handed.
+    working_multiple = 2.5
+    #: What the hand-written networks (three to five) hold beside the members they are handed,
+    #: measured on a 24 MiB float32 member.
     _NETWORK_MULTIPLE: ClassVar[dict[int, float]] = {1: 1.0, 2: 1.5, 3: 1.0, 4: 2.5, 5: 1.5}
+    #: Past five, the window: k+1 float32 buffers for k = count // 2, and the two it blends,
+    #: measured on a 293 MiB uint16 member at ten.
+    _WINDOW_MULTIPLE = 1.8
 
     def working_multiple_for(self, cases: int) -> float:
+        if cases > 5:
+            return self._WINDOW_MULTIPLE
         return self._NETWORK_MULTIPLE.get(cases, float(self.working_multiple))
 
     @staticmethod
@@ -210,23 +222,23 @@ def _median_of_three(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch
 
     def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor:
         dtype = _averaged_dtype(tensors[0].dtype)
-        members = [tensor.to(dtype) for tensor in tensors]
-        if len(members) == 1:
-            return members[0]
-        # Three to five members is what a fold has, and there the middle is SELECTED by a network of
-        # element-wise min/max rather than found by sorting the whole stack: same values to the bit,
-        # a fraction of the time (CUDA 7.20 -> 0.45 ms at three, 8.42 -> 1.10 at five; CPU 55 -> 26
-        # at three), and no stack to hold, which is what lets the planner cut taller slabs. Beyond
-        # five the sort is simpler and no slower: what torch.quantile computes without its
-        # interpolation machinery (1.5-2x on CPU, 3.5x on CUDA, measured).
-        low, high = self._middle_pair(members)
+        if len(tensors) == 1:
+            return tensors[0].to(dtype)
+        # The members are handed over in the dtype they arrived in and widened one at a time as the
+        # network takes them: torch has no integer min/max kernel on the CPU, and widening ten
+        # members up front is what put ten float32 copies beside ten uint16 regions.
+        low, high = self._middle_pair(tensors, dtype)
         return low if low is high else torch.lerp(low, high, 0.5)
 
-    def _middle_pair(self, members: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    def _middle_pair(self, members: list[torch.Tensor], dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
         """The one middle member of an odd fold (the same tensor twice), or the two an even fold
-        averages: by network up to five members, off a sorted stack past it."""
+        averages: by a hand-written network up to five members, by the insertion window past it.
+        ``dtype`` is what the network computes in; the members are widened to it as they enter."""
         minimum, maximum = torch.minimum, torch.maximum
         count = len(members)
+        if count > 5:
+            return self._middle_pair_by_window(members, dtype)
+        members = [member.to(dtype) for member in members]
         if count == 2:
             first, second = members
             return minimum(first, second), maximum(first, second)
@@ -239,19 +251,40 @@ def _middle_pair(self, members: list[torch.Tensor]) -> tuple[torch.Tensor, torch
             c, d = minimum(c, d), maximum(c, d)
             second, third = maximum(a, c), minimum(b, d)
             return minimum(second, third), maximum(second, third)
-        if count == 5:
-            a, b, c, d, e = members
-            a, b = minimum(a, b), maximum(a, b)
-            c, d = minimum(c, d), maximum(c, d)
-            a, c = minimum(a, c), maximum(a, c)  # a is the fold's smallest: out of the running
-            b, d = minimum(b, d), maximum(b, d)  # d is its largest: out too
-            middle = self._median_of_three(b, c, e)
-            return middle, middle
-        ranked = torch.stack(members, dim=0).sort(dim=0).values
+        a, b, c, d, e = members
+        a, b = minimum(a, b), maximum(a, b)
+        c, d = minimum(c, d), maximum(c, d)
+        a, c = minimum(a, c), maximum(a, c)  # a is the fold's smallest: out of the running
+        b, d = minimum(b, d), maximum(b, d)  # d is its largest: out too
+        middle = self._median_of_three(b, c, e)
+        return middle, middle
+
+    @staticmethod
+    def _middle_pair_by_window(members: list[torch.Tensor], dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        """The middle pair of any count, by an insertion window of the ``k + 1`` smallest seen.
+
+        The middle of ``count`` members is rank ``count // 2`` (0-based; the pair ``count // 2 - 1``
+        and ``count // 2`` on an even count). A window that keeps the ``k + 1`` smallest members
+        seen so far, ``k = count // 2``, holds those ranks exactly once every member has passed
+        through it: a member larger than the whole window can be no smaller than rank ``k + 1`` of
+        the members seen, so dropping it off the end loses nothing the answer needs. Each insertion
+        is a chain of element-wise min/max, which is what makes the selection exact -- the same
+        values a full sort returns, to the bit (pinned against ``torch.sort`` in the tests).
+        """
+        count = len(members)
+        keep = count // 2 + 1
+        window: list[torch.Tensor] = []
+        for member in members:
+            window.append(member.to(dtype))
+            for index in range(len(window) - 1, 0, -1):
+                lower, upper = window[index - 1], window[index]
+                window[index - 1], window[index] = torch.minimum(lower, upper), torch.maximum(lower, upper)
+            if len(window) > keep:
+                window.pop()
+        middle = window[count // 2]
         if count % 2:
-            middle = ranked[count // 2]
             return middle, middle
-        return ranked[count // 2 - 1], ranked[count // 2]
+        return window[count // 2 - 1], middle
 
 
 class Vote(Reduction):