From e96341321fc17da6fd7c4c1249001fb4a7a342d5 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Mon, 19 Oct 2026 11:19:54 -0400 Subject: [PATCH] Pipeline.plan(level, box) says what a read would compute, computing nothing, as registration-station asked after a live view of its last pass pulled in a 75-chunk cone it found only in its logs: per stage, from the first to the output, the chunks the read takes from it, those its cache holds now and those it would 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, goes on through dataset:// into the named dataset's stages, and stops at a source that is not a stage, such as a warp:// whose reads depend on its field's values, whose reads are the boxes it would be read over; ChunkedSource.reads says what a chunk's computation reads (set by the pipeline's stages) and covering which chunks or blocks a read takes, which BlockSource.read now uses too; on registration-station's chain the counts matched every chunk it measured; caching docs and the CHANGELOG say so --- CHANGELOG.md | 3 +++ docs/concepts/caching.md | 18 ++++++++++++++++ src/chunkmirage/pipeline.py | 38 ++++++++++++++++++++++++++++++++- src/chunkmirage/sources/base.py | 25 ++++++++++++++++------ tests/test_overlap.py | 16 ++++++++++++++ 5 files changed, 93 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f3b0a5a..206271d 100644 --- a/CHANGELOG.md +++ b/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, diff --git a/docs/concepts/caching.md b/docs/concepts/caching.md index 346bfc0..9882986 100644 --- a/docs/concepts/caching.md +++ b/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 diff --git a/src/chunkmirage/pipeline.py b/src/chunkmirage/pipeline.py index 9e76992..b6fcb35 100644 --- a/src/chunkmirage/pipeline.py +++ b/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 diff --git a/src/chunkmirage/sources/base.py b/src/chunkmirage/sources/base.py index 36caff9..c3b0109 100644 --- a/src/chunkmirage/sources/base.py +++ b/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)), diff --git a/tests/test_overlap.py b/tests/test_overlap.py index d64bff8..df430c3 100644 --- a/tests/test_overlap.py +++ b/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"