chunkmirage materialize: a dataset written to a store as the server serves it

CHANGELOG.md

@@ -55,6 +55,10 @@ The first release, an alpha: the plugin API can still change before 0.1.0.
   in another app.
 - `chunkmirage.serve` and `Server`: an app served from Python on a port bound first, with a
   callback once it accepts connections, in a background thread or this one.
+- `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.
 
 ### Sources
 

docs/concepts/formats.md

@@ -73,6 +73,31 @@ view as zarr v2 with OME-Zarr 0.4 (`virtual/<page>/zarr2/<view>/`), the same chu
 readers that predate zarr v3's final spec: GDAL 3.8, as the fires demo runs it in the page
 (gdal3.js), opens a level as a zarr array and translates it to a GeoTIFF.
 
+## A dataset written to a store (Python and the command line)
+
+```bash
+chunkmirage materialize SOURCE OUT [--op ...] [--format zarr3|zarr|n5] [--level N ...] [--box lo:hi]
+```
+
+computes SOURCE through its ops and writes it to OUT (a directory, or a `file://`, `gs://`
+or `s3://` URL, through tensorstore) as a frontend serves it: zarr v3 with OME-Zarr 0.5
+(the default), zarr v2 with OME-Zarr 0.4 and consolidated metadata, or N5. The metadata
+and every chunk are byte for byte what a client fetches, so the store opens wherever the
+served dataset did and reads back as a stored source of its own: a stage computed once (a
+registration's pass, a model's output over a region) is then read from disk instead of
+computed again. Every level's metadata is written first, so a store whose writing stopped
+part way is still the dataset, zero where no chunk was written yet; then the chunks of the
+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. From Python,
+`chunkmirage.materialize(pipeline, where, format=, levels=, box=, 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'
+requests as any work asked for from Python does (never dropped, in arrival order), and
+its chunks pass through the server's cache.
+
 ## A region written to a store (browser engine)
 
 The browser engine also writes a view, or a box of it, into a store as files

docs/contributing.md

@@ -114,7 +114,7 @@ Packages build on chunkmirage through its plugin API, which is stable:
 * **Schema.** A parameter model listed under `chunkmirage.models` appears in
   `chunkmirage schema` under the entry point's name, for the browser engine's types.
 * **Serving.** `PipelineSpec`, `Pipeline`, `create_app` and its parameters, `serve` and
-  `Server`, `gpu.set_device`, `demand.set_gather_threads`,
+  `Server`, `materialize`, `gpu.set_device`, `demand.set_gather_threads`,
   `DatasetRegistry` (`add`, `get`, `resolve`, `refresh`, `remove`, `subscribe`) and its
   resolver, the `chunkmirage.routes` entry point, the REST routes in the
   [API reference](reference/api.md), the served URL layout, and the CLI's flags.

docs/design.md

@@ -275,7 +275,7 @@ or astronomy? is a ratio of two bands geo?): if it has an algorithm in it, it go
 
 | package | holds |
 | ------- | ----- |
-| `chunkmirage` | the pipeline, cache, demand queue, server, REST API and CLI; the frontends; readers of stored formats (zarr, N5, precomputed, HDF5); `scene://` and the transforms (the pipeline's own resampler uses them), and that resampler on a GPU (the same code with CuPy for numpy, `gpu.py`, loaded only when a device is set: interpolation is what the core's own sources do, so a faster way to do it stays with them; CuPy rather than PyTorch because its calls are numpy's and scipy's, so one code path serves both and the CPU's tests cover it); the sources that compose other sources (`stack://`, `flip://`, `warp://`); meshes, a frontend's way of serving a surface; and two ops: `downsample`, which the pipeline runs to make the pyramid of an op bound to one resolution, and `python`, code carried in the spec |
+| `chunkmirage` | the pipeline, cache, demand queue, server, REST API and CLI; the frontends, and `materialize`, which writes a dataset as they serve it (their own metadata and encoded chunks, so a store reads back as the dataset did); readers of stored formats (zarr, N5, precomputed, HDF5); `scene://` and the transforms (the pipeline's own resampler uses them), and that resampler on a GPU (the same code with CuPy for numpy, `gpu.py`, loaded only when a device is set: interpolation is what the core's own sources do, so a faster way to do it stays with them; CuPy rather than PyTorch because its calls are numpy's and scipy's, so one code path serves both and the CPU's tests cover it); the sources that compose other sources (`stack://`, `flip://`, `warp://`); meshes, a frontend's way of serving a surface; and two ops: `downsample`, which the pipeline runs to make the pyramid of an op bound to one resolution, and `python`, code carried in the spec |
 | `chunkmirage_contrib` | every other op (filters, thresholds, morphology, labelling, detectors, terrain); the computed sources (`synthetic://`, `register://`, `stitch://`, GeoTIFF reading); the registration solver, stitching and tracking; and, in `web/src/contrib/`, the browser pages built on them |
 
 Contrib reaches the core only through the entry points any other package would use:

docs/getting-started.md

@@ -129,3 +129,10 @@ arr = ts.open({"driver": "zarr3",
                "kvstore": "http://localhost:8000/em-thresh/zarr3/s0/"}).result()
 arr[0:64, 0:64, 0:64].read().result()
 ```
+
+To keep a result instead, write it to disk as the server serves it, a zarr of its own that
+any tool reads (see [formats](concepts/formats.md#a-dataset-written-to-a-store-python-and-the-command-line)):
+
+```bash
+chunkmirage materialize /path/to/data.zarr/em/fibsem-uint8 em-thresh.zarr --op threshold:low=120
+```

docs/reference/api.md

@@ -143,6 +143,12 @@ server.stop()
 `Server(app, sock, ...)` takes a socket bound already (`netutil.bind_socket`), when the
 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
+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)).
+
 Two process-wide settings `chunkmirage serve` makes from its flags, for an application
 serving its own: `chunkmirage.gpu.set_device("cpu" | "auto" | "cuda" | "cuda:N")` (where
 `warp://` and `scene://` resample, `--device`), and `chunkmirage.demand.set_gather_threads(n)`

docs/reference/cli.md

@@ -49,6 +49,30 @@ the control page (`/ui`, with the token when there is one), the control API URL
 `--python-viewer`, the viewer URL; the port and thread count when they differ from what
 was asked for.
 
+## `materialize`
+
+```bash
+chunkmirage materialize SOURCE OUT [--op ...] [--format zarr3] [--level N ...] [--box lo:hi]
+```
+
+Computes SOURCE through its ops and writes it to OUT as a frontend serves it, a stored
+dataset of its own ([formats](../concepts/formats.md#a-dataset-written-to-a-store-python-and-the-command-line)).
+Prints the counts as JSON (`chunks`, `written`, `skipped`, `bytes`, `seconds`) and each
+batch's progress on stderr.
+
+| flag | default | meaning |
+| ---- | ------- | ------- |
+| `OUT` | | a directory, or a `file://`, `gs://` or `s3://` URL |
+| `--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 |
+| `--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) |
+| `--device` | `cpu` | where `warp://` and `scene://` resample, as for `serve` |
+| `--cache-gb` | `2.0` | in-process chunk cache |
+
 ## `ops`
 
 Lists registered ops with halo, cache flag and parameters.

docs/roadmap.md

@@ -36,9 +36,12 @@ describe shipped features as future work.
    viewer becomes the job scheduler. (Ordering the work by what clients ask for, and
    dropping what they stop waiting for, has shipped; see
    [caching](concepts/caching.md#order-of-work-and-requests-given-up-on). Inference ops
-   would join its queues. A first piece has shipped in the browser engine: an explicit
-   write of a box of a view into a store, a folder the visitor picks or memory, as an
-   OME-Zarr, its chunks computed in the background behind the viewer's
+   would join its queues. A first piece has shipped in both engines: an explicit write of
+   a dataset, or a box of it, into a store as the frontends serve it, from Python and the
+   command line (`chunkmirage materialize`,
+   [formats](concepts/formats.md#a-dataset-written-to-a-store-python-and-the-command-line)),
+   and in the browser into a folder the visitor picks or memory, its chunks computed in the
+   background behind the viewer's
    ([formats](concepts/formats.md#a-region-written-to-a-store-browser-engine)). The cache on
    disk itself, in either engine, its sharing across workers and the filler that grows it
    from what is browsed are still to do.)

src/chunkmirage/__init__.py

@@ -1,10 +1,11 @@
 """chunkmirage: spoof chunked array formats over HTTP with on-the-fly processing."""
 
 from chunkmirage.core import ArrayInfo, Box
+from chunkmirage.materialize import materialize
 from chunkmirage.pipeline import Pipeline, PipelineSpec
 from chunkmirage.server import create_app
 from chunkmirage.serving import Server, serve
 from chunkmirage.sources import open_source
 
-__all__ = ["ArrayInfo", "Box", "Pipeline", "PipelineSpec", "Server", "create_app", "open_source", "serve"]
+__all__ = ["ArrayInfo", "Box", "Pipeline", "PipelineSpec", "Server", "create_app", "materialize", "open_source", "serve"]
 __version__ = "0.1.0a1"

src/chunkmirage/cli.py

@@ -397,6 +397,52 @@ def schema(
         typer.echo(dumps(), nl=False)
 
 
+@app.command()
+def materialize(
+    source: str = typer.Argument(..., help="what to compute, as for serve"),
+    out: str = typer.Argument(..., help="where to write it: a directory, or a file://, gs:// or s3:// URL"),
+    op: list[str] = typer.Option([], "--op", "-o", help="op spec, as for serve (repeatable)"),
+    format: str = typer.Option("zarr3", help="zarr3 (OME-Zarr 0.5), zarr (zarr v2, OME-Zarr 0.4) or n5"),
+    level: list[int] = typer.Option([], "--level", "-l", help="a level to write (repeatable; default every level)"),
+    box: str | None = typer.Option(
+        None,
+        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)",
+    ),
+    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)"),
+    keep_empty: bool = typer.Option(False, "--keep-empty", help="also write chunks of nothing but zeros"),
+    device: str = typer.Option("cpu", help="where warp:// and scene:// resample, as for serve"),
+    cache_gb: float = typer.Option(2.0, help="in-process chunk cache size"),
+):
+    """Compute SOURCE through its ops and write it to OUT as FORMAT serves it, read back as a
+    stored dataset of its own."""
+    import json
+    import time
+
+    from chunkmirage import gpu
+    from chunkmirage.core import Box
+    from chunkmirage.materialize import materialize as write
+
+    gpu.set_device(device)
+    registry = build_registry(source, "processed", op, chunk, cache_gb=cache_gb, select=select)
+    region = None
+    if box:
+        lo, _, hi = box.partition(":")
+        region = Box(tuple(int(v) for v in lo.split(",")), tuple(int(v) for v in hi.split(",")))
+    started = time.perf_counter()
+
+    def report(c):
+        typer.echo(f"s{c['level']}: {c['written'] + c['skipped']}/{c['chunks']} chunks", err=True)
+
+    counts = write(
+        registry.get("processed"), out, format=format, levels=level or None, box=region,
+        compressor=compressor, skip_empty=not keep_empty, progress=report,
+    )
+    typer.echo(json.dumps({**counts, "seconds": round(time.perf_counter() - started, 1)}))
+
+
 @app.command()
 def inspect(source: str):
     """Print what chunkmirage sees when opening SOURCE."""

src/chunkmirage/frontends/base.py

@@ -61,6 +61,12 @@ def compute(self, pipeline: Pipeline, req: ChunkRequest) -> bytes:
     root_keys: ClassVar[tuple[str, ...]] = ()  # metadata keys of the group
     level_keys: ClassVar[tuple[str, ...]] = ()  # metadata keys of each level
 
+    def chunk_key(self, level: int, index: tuple[int, ...]) -> str | None:
+        """The path chunk ``index`` of ``level`` is served at, which ``resolve`` maps back to
+        it; None for a format whose chunks are not files of their own (what ``materialize``
+        cannot write)."""
+        return None
+
     def listing(self, pipeline: Pipeline, path: str) -> list[str] | None:
         """What a directory listing of ``path`` shows: the group's metadata keys and one
         directory per level, or a level's metadata keys (never its chunks); None if ``path``

src/chunkmirage/frontends/n5.py

@@ -88,6 +88,9 @@ def _level_attrs(self, pipeline: Pipeline, lvl: int) -> dict:
             },
         }
 
+    def chunk_key(self, level: int, index) -> str:
+        return f"s{level}/" + "/".join(str(i) for i in index[::-1])  # N5's axes are x first
+
     def encode(self, info: ArrayInfo, index, block: np.ndarray) -> bytes:
         # mode 0 = default; dims are F-order (x first) = reversed C shape; data big-endian.
         header = struct.pack(">HH", 0, block.ndim) + struct.pack(

src/chunkmirage/frontends/zarr2.py

@@ -78,6 +78,9 @@ def _consolidated(self, pipeline: Pipeline) -> dict:
             meta[f"s{lvl}/.zattrs"] = {}
         return {"zarr_consolidated_format": 1, "metadata": meta}
 
+    def chunk_key(self, level: int, index) -> str:
+        return f"s{level}/" + self.separator.join(str(i) for i in index)
+
     def encode(self, info: ArrayInfo, index, block: np.ndarray) -> bytes:
         full = self.pad_to_full(info, index, block)
         raw = np.ascontiguousarray(full).astype(full.dtype.newbyteorder("<"), copy=False).tobytes()

src/chunkmirage/frontends/zarr3.py

@@ -84,6 +84,9 @@ def _array_meta(self, info: ArrayInfo) -> dict:
             "dimension_names": list(info.axes),
         }
 
+    def chunk_key(self, level: int, index) -> str:
+        return f"s{level}/c/" + "/".join(str(i) for i in index)
+
     def encode(self, info: ArrayInfo, index, block: np.ndarray) -> bytes:
         full = self.pad_to_full(info, index, block)
         raw = np.ascontiguousarray(full).astype(full.dtype.newbyteorder("<"), copy=False).tobytes()