Pipeline.plan: what a read would compute, computing nothing

CHANGELOG.md

@@ -114,6 +114,9 @@ The first release, an alpha: the plugin API can still change before 0.1.0.
 - `cache_token()` for ops whose output depends on state outside their parameters, and
   `DatasetRegistry.refresh` to read it again.
 - `slots` to bound how many calls of an op run at once.
+- `Pipeline.plan(level, box)`: what a read would compute, computing nothing: per stage, the
+  chunks it is asked for, those held and those to compute, following `dataset://` into
+  the datasets it reads and stopping at other sources.
 - Zero padding past the volume's edge (`padding`).
 - Channel axes read whole, so an op may change their length: select or combine channels.
 - A level shrunk by a whole factor for an op at one voxel size is the mean of each block,

docs/concepts/caching.md

@@ -247,6 +247,24 @@ level, saying whether that chunk is in the cache and whether every finest-level
 it is, read afresh at every request; `GET /api/datasets/{name}/coverage` gives the counts
 per level. Both read the cache as it is, so what was evicted counts as not computed.
 
+What a read would compute, before it does, is `Pipeline.plan(level, box)`: per stage, from
+the first to the output, the chunks the read takes from it (`chunks`), those the cache
+holds now (`held`) and those it would compute (`compute`), whose inputs, padded by the
+ops' halos, are what the stage before is asked for. An op with `overlap` has its blocks
+as a stage of their own (`blocks:...`). The walk stops at a stage holding all it is asked
+for; it goes on through `dataset://` into the named dataset's stages, and stops at a source
+that is not a stage (a stored array, a `warp://` whose reads depend on its field's values),
+whose `reads` are the boxes it would be read over. A stage that keeps no chunks computes
+one as often as it is read; `compute` counts it once.
+
+```python
+>>> pipeline.plan(0, Box((0, 0, 0), (64, 64, 64)))
+[{'source': ..., 'reads': 12},
+ {'stage': 'raw:...', 'keeps': True, 'chunks': 12, 'held': 0, 'compute': 12},
+ {'stage': 'blocks:match:...', 'keeps': True, 'chunks': 8, 'held': 0, 'compute': 8},
+ {'stage': 'match:...', 'keeps': False, 'chunks': 1, 'held': 0, 'compute': 1}]
+```
+
 ## Inspecting and clearing
 
 ```bash

src/chunkmirage/pipeline.py

@@ -218,15 +218,22 @@ def compute(out_box: Box, idx: tuple[int, ...], stage_key: str) -> np.ndarray:
         run = lambda: fused.run(ops, block, in_box, out_box, info, total_halo, scale)  # noqa: E731
         return run() if queue is None else queue.run((stage_key, idx), level, run)
 
+    def reads(box_of):  # what a chunk's computation reads, for Pipeline.plan
+        return prev, lambda idx: fused.input_box(prev_info, box_of(idx), lead, total_halo, scale)
+
     if not any(overlap):
-        return ChunkedSource(info, lambda idx: compute(info.chunk_box(idx), idx, key), cache, key)
+        stage = ChunkedSource(info, lambda idx: compute(info.chunk_box(idx), idx, key), cache, key)
+        stage.reads = reads(info.chunk_box)
+        return stage
     # the blocks, chunk plus overlap, each from its inputs alone; the chunk is their blend
     bkey = f"blocks:{key}"
     blocks = BlockSource(
         info, overlap, lambda idx: compute(blocks.chunk_box(idx), idx, bkey), cache if blocks_cache is None else blocks_cache, bkey
     )
+    blocks.reads = reads(blocks.chunk_box)
     stage = ChunkedSource(info, lambda idx: blocks.read(info.chunk_box(idx)), cache, key)
     stage.blocks = blocks
+    stage.reads = blocks, info.chunk_box
     return stage
 
 
@@ -402,6 +409,7 @@ def _raw(self, raw: Source, lvl_i: int, chunk_shape, cache_source: bool) -> tupl
             self.cache if cache_source else None,
             f"raw:{h}",
         )
+        stage.reads = raw, stage.chunk_box
         return stage, h
 
     def _chain(
@@ -492,6 +500,34 @@ def chunk(self, level: int, index: Sequence[int]) -> np.ndarray:
     def read(self, level: int, box: Box) -> np.ndarray:
         return self.levels[level].read(box)
 
+    def plan(self, level: int, box: Box) -> list[dict]:
+        """What reading ``box`` of level ``level`` would compute, computing nothing: per
+        stage, from the first to the output, the chunks the read takes from it (``chunks``),
+        those its cache holds now (``held``) and those it would compute (``compute``), whose
+        inputs are what the stage before is asked for; an overlapping op's blocks are a
+        stage of their own (``blocks:``). Through ``dataset://`` the walk goes on into the
+        named dataset's stages; it stops at a source that is not a stage (a stored array, a
+        ``warp://`` whose reads depend on its field's values), whose ``reads`` are the boxes
+        it would be read over. A stage that keeps no chunks computes each as often as it is
+        read; ``compute`` counts them once."""
+        steps: list[dict] = []
+        src: Source = self.levels[level]
+        boxes = [box.clip(src.info.shape)]
+        while isinstance(src, ChunkedSource):
+            needed: set[tuple[int, ...]] = set()
+            for b in boxes:
+                needed.update(src.covering(b))
+            held = {idx for idx in needed if src.in_cache(idx)}
+            todo = sorted(needed - held)
+            steps.append({"stage": src.cache_key(), "keeps": src.cached, "chunks": len(needed),
+                          "held": len(held), "compute": len(todo)})
+            if src.reads is None or not todo:
+                return steps[::-1]
+            src, box_of = src.reads
+            boxes = [b for b in (box_of(idx).clip(src.info.shape) for idx in todo) if not b.empty]
+        steps.append({"source": src.cache_key(), "reads": len(boxes)})
+        return steps[::-1]
+
     def digest(self) -> str:
         return hashlib.sha1("|".join(lvl.cache_key() for lvl in self.levels).encode()).hexdigest()[
             :12

src/chunkmirage/sources/base.py

@@ -87,6 +87,9 @@ def __init__(
         self._key = key
         #: the overlapping blocks this stage is blended from, if any (an op with ``overlap``)
         self.blocks: BlockSource | None = None
+        #: what computing a chunk reads, if said: the source and the box of it for a chunk's
+        #: index (``Pipeline.plan`` follows it)
+        self.reads: tuple[Source, Callable[[tuple[int, ...]], Box]] | None = None
         # In-flight computations keyed by chunk index: concurrent requests for the same chunk
         # (a viewer retrying, or two clients) wait for one computation instead of duplicating it.
         self._inflight: dict[tuple[int, ...], threading.Event] = {}
@@ -112,6 +115,10 @@ def chunk_box(self, index: Sequence[int]) -> Box:
         """What chunk ``index`` covers: ``compute_chunk`` returns an array of its shape."""
         return self._info.chunk_box(index)
 
+    def covering(self, box: Box) -> list[tuple[int, ...]]:
+        """The chunks a read of ``box`` (within the array) takes data from."""
+        return list(self._info.chunks_covering(box))
+
     def chunk(self, index: Sequence[int]) -> np.ndarray:
         """Chunk ``index`` as an array of shape ``chunk_box(index).shape`` (edge-clipped)."""
         index = tuple(int(i) for i in index)
@@ -221,19 +228,25 @@ def chunk_box(self, index: Sequence[int]) -> Box:
             whole.stop[:lead] + tuple(b + o for b, o in zip(core.stop[lead:], self.overlap)),
         )
 
-    def read(self, box: Box) -> np.ndarray:
+    def covering(self, box: Box) -> list[tuple[int, ...]]:
+        """The blocks whose extent reaches ``box``."""
         info, k = self._info, len(self.overlap)
         lead = info.ndim - k
         chunks, grid = info.chunk_shape[lead:], info.chunk_grid[lead:]
-        acc = np.zeros(box.shape, np.float64)
-        total = np.zeros(box.shape[lead:], np.float64)
-        # the blocks whose extent reaches the box
         first = [max((a - o) // c, 0) for a, o, c in zip(box.start[lead:], self.overlap, chunks)]
         last = [min((b - 1 + o) // c, g - 1) for b, o, c, g in zip(box.stop[lead:], self.overlap, chunks, grid)]
         ahead = tuple(a // c for a, c in zip(box.start[:lead], info.chunk_shape[:lead]))
+        return [ahead + tuple(int(lo + r) for lo, r in zip(first, rel))
+                for rel in np.ndindex(*[hi - lo + 1 for lo, hi in zip(first, last)])]
+
+    def read(self, box: Box) -> np.ndarray:
+        info, k = self._info, len(self.overlap)
+        lead = info.ndim - k
+        chunks = info.chunk_shape[lead:]
+        acc = np.zeros(box.shape, np.float64)
+        total = np.zeros(box.shape[lead:], np.float64)
         parts = []
-        for rel in np.ndindex(*[hi - lo + 1 for lo, hi in zip(first, last)]):
-            index = ahead + tuple(int(lo + r) for lo, r in zip(first, rel))
+        for index in self.covering(box):
             extent = self.chunk_box(index)
             inter = Box(
                 tuple(max(a, b) for a, b in zip(box.start, extent.start)),

tests/test_overlap.py

@@ -79,3 +79,19 @@ def output_dtype(self, in_dtype):
         assert (out == np.round(out)).all(), "rounded after the mean, as the op's uint8 output asks"
     finally:
         del base._REGISTRY[Rounded.name]
+
+
+def test_a_plan_says_what_a_read_would_compute_and_computes_nothing():
+    p = Pipeline(open_source(SRC), [{"op": "block_centre_test"}])
+    BlockCentre.calls = 0
+    box = Box((0, 0, 0), (8, 8, 8))  # one chunk, which 8 blocks reach
+    plan = p.plan(0, box)
+    assert BlockCentre.calls == 0
+    source, raw, blocks, out = plan
+    assert "source" in source and source["reads"] == raw["compute"] == 12, "the blocks' inputs: 2 x 2 x 3 chunks"
+    assert blocks["stage"].startswith("blocks:") and (blocks["chunks"], blocks["compute"]) == (8, 8)
+    assert out["stage"] == p.levels[0].cache_key() and not out["keeps"] and out["compute"] == 1
+    p.read(0, box)
+    blocks, out = p.plan(0, box)
+    assert (blocks["held"], blocks["compute"]) == (8, 0), "held now: the walk stops there"
+    assert out["compute"] == 1, "a stage that keeps nothing computes again"