diff --git a/CHANGELOG.md b/CHANGELOG.md index 4e11226..ec930f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -58,7 +58,10 @@ The first release, an alpha: the plugin API can still change before 0.1.0. - `chunkmirage materialize` and `chunkmirage.materialize`: a dataset, or a box of it, computed and written into a store (a directory, `gs://`, `s3://`) as a frontend serves it, zarr v3, zarr v2 or N5, so it reads back as a stored dataset of its own; the metadata - first, chunks several at once, empty ones left out. + first, chunks several at once, empty ones left out. A box is exactly the box by default + (`--extent box`: arrays of its shape, placed where it is, each of the dataset's chunks it + reaches computed once), or its chunks in arrays of the levels' whole shapes + (`--extent level`). ### Sources @@ -162,8 +165,8 @@ The first release, an alpha: the plugin API can still change before 0.1.0. level's whole shape, the box's chunks in it rounded out to whole chunks (`extent: "level"`, for several boxes in one store); a box of the whole level is the same store either way; progress counts the array's chunks done and the served chunks computed - (`computed`, `needed`); the Python side (`chunkmirage materialize`) writes the level's - extent only, for now; + (`computed`, `needed`); the Python side (`chunkmirage materialize --extent`) writes + either too, the box by default; `directoryStore()` writes into a folder picked with the File System Access API (or the origin's private file system), `memoryStore()` into memory ([formats](docs/concepts/formats.md#a-region-written-to-a-store-browser-engine)). Its diff --git a/docs/concepts/formats.md b/docs/concepts/formats.md index 52941a0..8119f00 100644 --- a/docs/concepts/formats.md +++ b/docs/concepts/formats.md @@ -90,10 +90,18 @@ part way is still the dataset, zero where no chunk was written yet; then the chu levels asked for (all by default), several at once, each of nothing but zeros left out (`--keep-empty` writes them). `--box` takes level 0's voxels along the last axes, `lo:hi` each comma separated in C order (`0,0,0:64,512,512`), and every level writes the same -region on its own voxels, rounded out to whole chunks, in an array of the level's whole -shape: what the browser engine writes with `extent: "level"`. Writing exactly the box -(the browser engine's default, `extent: "box"`) is not supported here yet. From Python, -`chunkmirage.materialize(pipeline, where, format=, levels=, box=, compressor=, +region on its own voxels, rounded out to whole voxels. As the browser engine does, +`--extent box` (the default) writes exactly the box: each level an array of the box's +shape, its translation moved by the box's corner so it lands where the box is, in chunks +of the dataset's shape counted from that corner, those at the far edge stored whole with +zeros past the box. Each of its chunks is read from the dataset's chunks it overlaps; when +the last stage keeps no chunks of its own (no `cache: true`), those pass through a cache +of up to 1 GiB, walked so that the fewest wait (the axes the box's corner is on chunk +edges along first, then the one with the most chunks), so each is computed once while a +slab of them fits. `--extent level` writes the box's chunks, rounded out to whole chunks, +into arrays of the levels' whole shapes, for several boxes in one store; a box of the +whole level is the same store either way. From Python, +`chunkmirage.materialize(pipeline, where, format=, levels=, box=, extent=, compressor=, skip_empty=, progress=)` does the same for any pipeline, another dataset's chain read through `dataset://` among them (`registry.get(name)`), and returns the counts (`{chunks, written, skipped, bytes}`). Inside a running server it competes with clients' diff --git a/docs/reference/api.md b/docs/reference/api.md index 29eee42..6841df4 100644 --- a/docs/reference/api.md +++ b/docs/reference/api.md @@ -144,7 +144,7 @@ server.stop() address must be known before the app is built. `chunkmirage.materialize(pipeline, where, *, format="zarr3", levels=None, box=None, -compressor=None, skip_empty=True, progress=None)` writes a pipeline (or a box of it) into +extent="box", compressor=None, skip_empty=True, progress=None)` writes a pipeline (or a box of it) into a store as a frontend serves it, which reads back as a stored dataset: the Python side of `chunkmirage materialize` ([formats](../concepts/formats.md#a-dataset-written-to-a-store-python-and-the-command-line)). diff --git a/docs/reference/cli.md b/docs/reference/cli.md index d9a3d7c..f9c436d 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -66,7 +66,8 @@ batch's progress on stderr. | `--op`, `-o` | none | op spec, as for `serve` (repeatable) | | `--format` | `zarr3` | `zarr3` (OME-Zarr 0.5), `zarr` (zarr v2, OME-Zarr 0.4, consolidated) or `n5` | | `--level`, `-l` | every level | a level whose chunks to write (repeatable); every level's metadata is written | -| `--box` | all of it | level 0's voxels along the last axes, `lo:hi` each comma separated in C order; every level writes the same region, rounded out to whole chunks, in an array of the level's whole shape (the browser engine's exact box, `extent: "box"`, is not supported here yet) | +| `--box` | all of it | level 0's voxels along the last axes, `lo:hi` each comma separated in C order; every level writes the same region on its own voxels, rounded out to whole voxels | +| `--extent` | `box` | `box`: each level an array of exactly the box, placed where it is (its translation moved by the box's corner), in chunks counted from that corner; `level`: the box's chunks, rounded out to whole chunks, in an array of the level's whole shape (several boxes in one store) | | `--chunk`, `--select` | as for `serve` | output chunk shape; non-spatial axes pinned | | `--compressor` | blosc (zarr), gzip (n5) | as the format takes it | | `--keep-empty` | off | also write chunks of nothing but zeros (left out otherwise: a reader takes a missing chunk as the fill value) | diff --git a/src/chunkmirage/cli.py b/src/chunkmirage/cli.py index 24a4f19..40be920 100644 --- a/src/chunkmirage/cli.py +++ b/src/chunkmirage/cli.py @@ -409,6 +409,10 @@ def materialize( help="the region, in level 0's voxels along the last axes: lo:hi, each comma separated in " "C order, e.g. 0,0,0:64,512,512 (every level writes the same region; default all of it)", ), + extent: str = typer.Option( + "box", help="box: an array of exactly the box, placed where it is; level: its chunks in an " + "array of the level's whole shape (to write several boxes into one store)", + ), chunk: str | None = typer.Option(None, help="output chunk shape, as for serve"), select: str | None = typer.Option(None, help="pin non-spatial axes, as for serve"), compressor: str | None = typer.Option(None, help="as the format takes it (default: blosc for zarr, gzip for n5)"), @@ -438,7 +442,7 @@ def report(c): counts = write( registry.get("processed"), out, format=format, levels=level or None, box=region, - compressor=compressor, skip_empty=not keep_empty, progress=report, + extent=extent, compressor=compressor, skip_empty=not keep_empty, progress=report, ) typer.echo(json.dumps({**counts, "seconds": round(time.perf_counter() - started, 1)})) diff --git a/src/chunkmirage/materialize.py b/src/chunkmirage/materialize.py index 55164ee..0bd1582 100644 --- a/src/chunkmirage/materialize.py +++ b/src/chunkmirage/materialize.py @@ -11,6 +11,12 @@ (``demand.gather``), and one of nothing but zeros (the fill value) is left out unless asked for. The store is anything tensorstore writes: a directory, or a ``file://``, ``gs://`` or ``s3://`` URL. + +A box is written as exactly the box (``extent="box"``, the default): each level an array +of the box's shape on it, placed where the box is (its translation moved by the box's +corner), in chunks of the dataset's shape counted from that corner, each read from the +dataset's chunks it overlaps. ``extent="level"`` writes the box's chunks into arrays of the +levels' whole shapes instead, so several boxes can go into one store. """ from __future__ import annotations @@ -22,13 +28,19 @@ import numpy as np from chunkmirage import demand -from chunkmirage.core import Box +from chunkmirage.cache import LRUCache +from chunkmirage.core import ArrayInfo, Box from chunkmirage.frontends import get_frontend from chunkmirage.frontends.base import Metadata from chunkmirage.pipeline import Pipeline +from chunkmirage.sources.base import ChunkedSource FORMATS = ("zarr3", "zarr", "n5") +EXTENTS = ("box", "level") BATCH = 256 # chunks asked for at once: the pool's work, and what one wait holds +#: most a box's write keeps of the dataset's chunks it stitches from, when its last stage +#: keeps none itself: enough for a slab of them, each computed once; past it, recomputed +BOX_BUFFER_BYTES = 1 << 30 def _kvstore_url(where: str) -> str: @@ -53,6 +65,50 @@ def level_box(pipeline: Pipeline, level: int, box: Box | None) -> Box: return Box((0,) * lead + tuple(lo), info.shape[:lead] + tuple(hi)).clip(info.shape) +class _Cropped: + """``pipeline`` cut to ``boxes``, one per level, as ``extent="box"`` writes it: each + level an array of its box's shape whose voxel 0 is the box's first (its translation + moved there), chunked as the level is, counted from the box's corner. A chunk is read + from the level's chunks it overlaps, through a cache of them when the level keeps none + (an uncached last stage), so each is computed once while a slab of them fits.""" + + def __init__(self, pipeline: Pipeline, boxes: list[Box]): + self.pipeline, self.source, self.boxes = pipeline, pipeline.source, boxes + self._infos: list[ArrayInfo] = [] + self._reads = [] + for level, b in enumerate(boxes): + info = pipeline.info(level) + moved = tuple(t + s * v for t, s, v in zip(info.translation, b.start, info.voxel_size)) + self._infos.append(info.with_(shape=b.shape, translation=moved)) + stage = pipeline.levels[level] + if not isinstance(stage, ChunkedSource) or stage.cached: + self._reads.append(stage) + else: + self._reads.append(ChunkedSource(stage.info, stage.chunk, LRUCache(BOX_BUFFER_BYTES), f"box:{stage.cache_key()}")) + + @property + def num_levels(self) -> int: + return self.pipeline.num_levels + + def info(self, level: int = 0) -> ArrayInfo: + return self._infos[level] + + def chunk(self, level: int, index) -> np.ndarray: + at, b = self._infos[level].chunk_box(index), self.boxes[level] + return self._reads[level].read(Box(tuple(s + a for s, a in zip(b.start, at.start)), + tuple(s + a for s, a in zip(b.start, at.stop)))) + + def walk(self, level: int) -> list[tuple[int, ...]]: + """The box's chunks in the order that keeps fewest of the level's open: the axes + its corner is on the level's chunk edges along first, then those it is not, the one + with the most chunks outermost (a chunk across an edge waits for the far side's).""" + info, b = self._infos[level], self.boxes[level] + grid = info.chunk_grid + split = [a for a in range(info.ndim) if b.start[a] % info.chunk_shape[a]] + order = [a for a in range(info.ndim) if a not in split] + sorted(split, key=lambda a: -grid[a]) + return sorted(np.ndindex(*grid), key=lambda idx: tuple(idx[a] for a in order)) + + def materialize( pipeline: Pipeline, where: str, @@ -60,6 +116,7 @@ def materialize( format: str = "zarr3", levels: Sequence[int] | None = None, box: Box | None = None, + extent: str = "box", compressor: str | None = None, skip_empty: bool = True, progress: Callable[[dict], None] | None = None, @@ -69,23 +126,34 @@ def materialize( frontend takes it (blosc by default for zarr, gzip for N5). ``levels``: which levels' chunks to write (all by default; every level's metadata is written whatever it is). ``box``: voxels of level 0 along its last axes, each level the same region on its own - voxels, rounded out to whole chunks. ``progress`` is told the counts after each batch - of chunks. Returns ``{chunks, written, skipped, bytes}``.""" + voxels (rounded out to whole voxels), written as ``extent`` says: ``box`` (the default) + an array of exactly the box, placed where it is; ``level`` its chunks (rounded out to + whole chunks) in an array of the level's whole shape. ``progress`` is told the counts + after each batch of chunks. Returns ``{chunks, written, skipped, bytes}``.""" import tensorstore as ts if format not in FORMATS: raise ValueError(f"format {format!r}: one of {', '.join(FORMATS)}") + if extent not in EXTENTS: + raise ValueError(f"extent {extent!r}: one of {', '.join(EXTENTS)}") fe = get_frontend(format, **({} if compressor is None else {"compressor": compressor})) kv = ts.KvStore.open(_kvstore_url(where)).result() levels = list(range(pipeline.num_levels)) if levels is None else [int(lv) for lv in levels] for lv in levels: if not 0 <= lv < pipeline.num_levels: raise ValueError(f"no level {lv}: the dataset has {pipeline.num_levels}") + source = pipeline + if box is not None and extent == "box": + boxes = [level_box(pipeline, lv, box) for lv in range(pipeline.num_levels)] + for lv, b in enumerate(boxes): + if b.empty: + raise ValueError(f"the box {box} holds no voxel of level {lv}, which is {pipeline.info(lv).shape}") + source = _Cropped(pipeline, boxes) # the metadata first: from then on the store is the dataset, zero where nothing is yet keys = [*fe.root_keys, *(f"s{i}/{k}" for i in range(pipeline.num_levels) for k in fe.level_keys)] for key in keys: - meta = fe.resolve(pipeline, key) + meta = fe.resolve(source, key) if isinstance(meta, Metadata): kv.write(key, meta.body).result() @@ -93,16 +161,18 @@ def materialize( def one(level: int, index: tuple[int, ...]) -> int | None: """Chunk ``index`` of ``level`` computed and written: its size, or None if empty.""" - block = pipeline.chunk(level, index) + block = source.chunk(level, index) if skip_empty and not block.any(): return None - data = fe.encode(pipeline.info(level), index, block) + data = fe.encode(source.info(level), index, block) kv.write(fe.chunk_key(level, index), data).result() return len(data) for level in levels: - info = pipeline.info(level) - todo = list(info.chunks_covering(level_box(pipeline, level, box))) + if isinstance(source, _Cropped): + todo = source.walk(level) + else: + todo = list(pipeline.info(level).chunks_covering(level_box(pipeline, level, box))) counts["chunks"] += len(todo) for start in range(0, len(todo), BATCH): for size in demand.gather([partial(one, level, idx) for idx in todo[start : start + BATCH]]): diff --git a/tests/test_materialize.py b/tests/test_materialize.py index fd5b267..b210ffe 100644 --- a/tests/test_materialize.py +++ b/tests/test_materialize.py @@ -1,6 +1,7 @@ """A dataset written to a store as the server serves it, read back as a stored dataset.""" import json +from typing import ClassVar import numpy as np import pytest @@ -9,6 +10,8 @@ from chunkmirage import Pipeline, PipelineSpec, materialize, open_source from chunkmirage.cli import app from chunkmirage.core import Box +from chunkmirage.ops import base +from chunkmirage.ops.base import Op IMG = "synthetic://julia?shape=32,48,48&chunk=16,16,16&levels=2&seed=2&voxel_size=1" # moved 24 voxels along x: the last third of every row reads past the image, zeros @@ -35,7 +38,7 @@ def test_a_dataset_written_reads_back_as_it_was_served(tmp_path, fmt): def test_a_box_of_one_level_from_the_command_line(tmp_path): out = tmp_path / "box.zarr" - result = CliRunner().invoke(app, ["materialize", IMG, str(out), "--level", "1", "--box", "0,0,0:16,16,16"]) + result = CliRunner().invoke(app, ["materialize", IMG, str(out), "--level", "1", "--box", "0,0,0:16,16,16", "--extent", "level"]) assert result.exit_code == 0, result.output assert json.loads(result.stdout.strip().splitlines()[-1])["chunks"] == 1 # level 1's voxels 0-8: one chunk back = open_source(str(out)).levels @@ -43,3 +46,41 @@ def test_a_box_of_one_level_from_the_command_line(tmp_path): got = back[1].read(Box((0, 0, 0), back[1].info.shape)) np.testing.assert_array_equal(got[:16, :16, :16], _whole(p, 1)[:16, :16, :16]) assert not got[16:].any() and not got[:, 16:].any() and not back[0].read(Box((0, 0, 0), (16, 16, 16))).any() + + +class Count(Op): + """Each block as it is, counted: how often the dataset's chunks are computed.""" + + name: ClassVar[str] = "count_test" + calls: ClassVar[int] = 0 + + def apply(self, block): + Count.calls += 1 + return block + + +@pytest.mark.parametrize("fmt", ["zarr3", "zarr"]) +def test_a_box_is_written_as_exactly_the_box_where_it_is(tmp_path, monkeypatch, fmt): + monkeypatch.setitem(base._REGISTRY, Count.name, Count) + p = Pipeline.from_spec(PipelineSpec(source=IMG, ops=[{"op": "count_test"}])) # its last stage keeps nothing + box = Box((3, 5, 7), (29, 40, 45)) # off the chunk edges along every axis + Count.calls = 0 + materialize(p, str(tmp_path / "box"), format=fmt, box=box) + # each of the dataset's chunks the box reaches computed once, though the box's own + # chunks each stitch up to eight of them + reached = sum(len(list(p.info(lv).chunks_covering(_level(p, lv, box)))) for lv in range(p.num_levels)) + assert Count.calls == reached + back = open_source(str(tmp_path / "box")).levels + got, want = back[0].info, p.info(0) + assert got.shape == box.shape and got.chunk_shape == want.chunk_shape + assert got.translation == tuple(t + s * v for t, s, v in zip(want.translation, box.start, want.voxel_size)) + np.testing.assert_array_equal(back[0].read(Box((0, 0, 0), box.shape)), p.read(0, box)) + lv1 = _level(p, 1, box) # level 1: the same region on its voxels, rounded out to whole ones + assert back[1].info.shape == lv1.shape + np.testing.assert_array_equal(back[1].read(Box((0, 0, 0), lv1.shape)), p.read(1, lv1)) + + +def _level(p, level, box): + from chunkmirage.materialize import level_box + + return level_box(p, level, box)