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"