diff --git a/CHANGELOG.md b/CHANGELOG.md index ec930f8..f3b0a5a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -165,8 +165,9 @@ 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 --extent`) writes - either too, the box by default; + (`computed`, `needed`), and `Engine.materializePlan(view, options)` says beforehand what + it will write and the most bytes of the box's chunks it will hold (`peakBytes`); 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/caching.md b/docs/concepts/caching.md index 8f20202..346bfc0 100644 --- a/docs/concepts/caching.md +++ b/docs/concepts/caching.md @@ -234,8 +234,8 @@ looking at out of the 256 MB; one a client asked for too, while it was on its wa as any other. So materialize asks for each chunk once: a box written as exactly the box (`extent: "box"`, the default) that does not start on chunk edges has each of its chunks stitched from the served chunks it overlaps as they come, rather than asking a served -chunk again for each of the box's chunks it overlaps. A staged view's stage outputs are kept as usual -(512 MB), so writing a large region does cycle through those. +chunk again for each of the box's chunks it overlaps. A staged view's stage outputs are +kept as usual (512 MB), so writing a large region does cycle through those. ## What has been computed diff --git a/docs/concepts/formats.md b/docs/concepts/formats.md index 8119f00..1fbc7a0 100644 --- a/docs/concepts/formats.md +++ b/docs/concepts/formats.md @@ -145,7 +145,17 @@ zero takes no memory. So a 512 × 2048 × 2048 box of 64³ chunks off the chunk everywhere (8 × 32 × 32 chunks) holds 32 × 8 + 8 + 1 = 265 chunks, which for 3-channel float32 chunks (3 MiB) is about 800 MiB; with `lo` on a chunk edge along z, 32 + 1. No order of the axes holds fewer; putting `lo` on the chunk edges along an axis is what -shrinks it further. +shrinks it further. These counts are exact where the box has three chunks or more along each +split axis, and the most it can hold otherwise. `Engine.materializePlan(view, options)` +gives the exact count for a box before anything is computed (`plan(region(...))` in +`web/src/materialize.ts`): `peakChunks`, and `peakBytes`, those chunks' bytes (`chunkBytes` +each), besides the `concurrency` served chunks on their way, with the array's `shape` and +`chunk`, its chunks (`total`), the served chunks it computes (`needed`), the axes it splits +and the walk's order. A page can warn, or offer to put `lo` on the chunk grid along an axis, +before it starts. Should a smaller peak ever matter, the way to get one is to compute the +plane of served chunks that two slabs of the box share twice, along the outermost split +axis, so each slab is finished before the next is begun, rather than compressing chunks +while they wait or refusing the box. A chunk of nothing but zeros is left out (zarr reads a missing chunk as the fill value, 0), unless asked to be written. A channel axis the view's ops add is the first axis and diff --git a/docs/design.md b/docs/design.md index ae53a9b..902ba9f 100644 --- a/docs/design.md +++ b/docs/design.md @@ -530,11 +530,11 @@ as text). What such a page gets, in the order it uses it: voxel, its chunks the served chunk shape counted from `lo`, each stitched from the served chunks it overlaps, since a page passing a box expects an array of that box; `"level"`: the array is the level's whole shape and the box's chunks are written, rounded out to - whole chunks, for several boxes written into one store), `format` (`"zarr2"`, OME-Zarr 0.4, the - default, since more tools read it; or `"zarr3"`, OME-Zarr 0.5), `compression` (`"gzip"`, - the default, or `null`), `skipEmpty` (a chunk of nothing but zeros left out, by default), - `concurrency`, `signal` (an `AbortSignal`), `onProgress` (`done`, `total`, `written`, - `bytes` for the array's chunks; `computed`, `needed` for the served chunks) and + whole chunks, for several boxes written into one store), `format` (`"zarr2"`, OME-Zarr + 0.4, the default, since more tools read it; or `"zarr3"`, OME-Zarr 0.5), `compression` + (`"gzip"`, the default, or `null`), `skipEmpty` (a chunk of nothing but zeros left out, by + default), `concurrency`, `signal` (an `AbortSignal`), `onProgress` (`done`, `total`, + `written`, `bytes` for the array's chunks; `computed`, `needed` for the served chunks) and `attributes` (the consumer's provenance, beside the OME metadata). It resolves with `{chunks, written, skipped, bytes, aborted}`. The layout is the one the engine serves for the view, from one module (`web/src/layout.ts`) for both, with the one level written as @@ -549,9 +549,14 @@ as text). What such a page gets, in the order it uses it: walk puts the axes the box does not split outermost and the split axis with the most of the box's chunks next, so the slab is the one across the other split axes, the smallest any order of the axes gives (a wide, shallow box walked in C order would hold a slab - across its two wide axes). A single level rather than the pyramid: the engine downsamples nothing itself, - and every level would be the model run again on its own grid. The metadata is written first, so a - store whose writing stopped is still an OME-Zarr, zero where no chunk was written. The + across its two wide axes). `Engine.materializePlan(view, options)` gives that peak for a + box before it starts (`peakChunks`, `peakBytes`), so a page can warn, or offer to put the + box on the chunk grid along an axis; should a smaller peak ever matter, computing the + plane of served chunks two slabs share twice, along the outermost split axis, is the way + to it, not compressing chunks while they wait or refusing a box. A single level rather + than the pyramid: the engine downsamples nothing itself, and every level would be the + model run again on its own grid. The metadata is written first, so a store whose writing + stopped is still an OME-Zarr, zero where no chunk was written. The chunks come from the same producer the routes use, through the same queue, but in the background: they wait behind the viewer's, in all its slots but one, and are never dropped as abandoned ([caching](concepts/caching.md#order-of-work-and-requests-given-up-on)), so diff --git a/web/src/engine.ts b/web/src/engine.ts index df14a51..34de49d 100644 --- a/web/src/engine.ts +++ b/web/src/engine.ts @@ -17,13 +17,13 @@ import { Cancelled, Claim, Queue } from "./demand"; import schema from "./generated/chunkmirage.schema.json"; import { Kept } from "./kept"; import { array2, array3, attrs2, dimensionNames, group3 } from "./layout"; -import { region, writeChunks, type MaterializeOptions, type Materialized, type WritableStore } from "./materialize"; +import { plan, region, writeChunks, type MaterializeOptions, type MaterializePlan, type Materialized, type Region, type WritableStore } from "./materialize"; import { route } from "./routes"; import { bytes, fitted, inputBox, levelAt, padded, prod, type Block, type Padding, type Stage, type StageAt } from "./stages"; import type { Answer, Ask, Later, MeshCall, PipelineView, Plan, Printed, Reply, SourceInfo, ToPyWorker, ToReader, ViewInfo } from "./types"; export { isTime, TO_SECONDS } from "./layout"; -export { directoryStore, memoryStore, type MaterializeOptions, type Materialized, type MaterializeProgress, type WritableStore } from "./materialize"; +export { directoryStore, memoryStore, type MaterializeOptions, type MaterializePlan, type Materialized, type MaterializeProgress, type WritableStore } from "./materialize"; const CONVENTIONS = [ { uuid: "d35379db-88df-4056-af3a-620245f8e347", name: "multiscales" }, { uuid: "f17cb550-5864-4468-aeb7-f3180cfb622f", name: "proj:" }, @@ -514,14 +514,30 @@ export class Engine { * as many as the queue has slots by default. Resolves with what was written, or * on `signal`'s abort with `aborted` (chunks waiting for their turn are dropped, the ones * written stay); rejects naming the chunk that failed to compute or be written, once the - * writes under way have ended. */ + * writes under way have ended. `materializePlan` says beforehand what it will write and + * hold. */ async materialize(view: string, store: WritableStore, options: MaterializeOptions = {}): Promise { + const r = await this.regionOf(view, options); + return writeChunks(store, r, (index, signal) => this.backgroundChunk(view, r.level, index, signal), { ...options, concurrency: options.concurrency ?? this.chunks?.slots ?? 1 }); + } + + /** What `materialize(view, store, options)` would write and hold, computing nothing + * (materialize.ts's plan): the array's shape and chunks, the view's chunks it computes, the + * order it walks them in, and the most bytes of the array's chunks it holds at once while + * they are stitched (`peakBytes`), for a page to warn, or offer to put the box's lo on the + * view's chunk edges along an axis, before it starts. */ + async materializePlan(view: string, options: MaterializeOptions = {}): Promise { + return plan(await this.regionOf(view, options)); + } + + /** What `options` asks materialize to write of `view` (materialize.ts's region), once its + * ops are planned. */ + private async regionOf(view: string, options: MaterializeOptions): Promise { if (!this.chunks) throw new Error(this.stopped ? "the engine was stopped" : "the engine has not started"); if (!this.infos[view] && this.sources[view]) await this.ready; // its ops are being planned const v = this.infos[view]; if (!v) throw new Error(`${view} is not served`); - const r = region(view, v, this.views[view].chunk, options); - return writeChunks(store, r, (index, signal) => this.backgroundChunk(view, r.level, index, signal), { ...options, concurrency: options.concurrency ?? this.chunks?.slots ?? 1 }); + return region(view, v, this.views[view].chunk, options); } /** Chunk `index` of a view's level for `materialize`: its whole bytes, as a client's request diff --git a/web/src/materialize.ts b/web/src/materialize.ts index 5916a2f..3d622e1 100644 --- a/web/src/materialize.ts +++ b/web/src/materialize.ts @@ -203,16 +203,16 @@ export function allZero(bytes: Uint8Array): boolean { * them; and the array written: its voxel 0 the level's voxel `at` (0 for the level's extent, * the box's lo for the box's), `shape` voxels on the level's axes (the level's own, * `levelShape`, or the box's), cut in chunks of `chunk` (the view's), those of its grid's - * indices [first, last) written, `total` of them, each chunk the channels the ops add first, - * `channels` values for each voxel, of `item` bytes each; and each chunk's `key` in the - * store. With the level's extent the array's grid is the view's, [first, last) = [lo, hi). */ + * indices [first, last) written, `total` of them, each chunk the channel axes the ops add + * first (`added`, `channels` values for each voxel in all), of `item` bytes each; and each + * chunk's `key` in the store. With the level's extent the array's grid is the view's, [first, last) = [lo, hi). */ export interface Region { name: string; level: number; format: "zarr3" | "zarr2"; compression: "gzip" | null; extent: "level" | "box"; files: [string, Uint8Array][]; lo: number[]; hi: number[]; needed: number; at: number[]; shape: number[]; levelShape: number[]; chunk: number[]; first: number[]; last: number[]; total: number; - channels: number; item: number; + added: number[]; channels: number; item: number; key: (index: number[]) => string; } @@ -254,7 +254,7 @@ export function region(name: string, v: ViewInfo, chunk: number[], options: Mate name, level, format, compression, extent, files: files.map(([path, o]) => [path, json(o)]), lo, hi, needed: count(lo, hi), at, shape, levelShape, chunk, first, last, total: count(first, last), - channels: prod(added), item: item ?? 0, + added, channels: prod(added), item: item ?? 0, key: (index) => chunkKey(format, LEVEL_PATH, added.length, index), }; } @@ -284,6 +284,65 @@ export function walkOrder(r: Region): number[] { /** The bytes of a chunk of region `r`'s array (or of the view's: they are cut alike). */ export const chunkBytes = (r: Region) => r.channels * prod(r.chunk) * r.item; +/** The most of region `r`'s array chunks its stitcher holds at once, the view's chunks added + * one at a time in `walkOrder`. The axes the box does not split are walked outermost and + * part it into blocks that share no chunk, so it is the most one block holds, over the split + * axes alone, n₁ ≥ n₂ ≥ ... of the box's chunks along them in walk order: along an axis of + * three chunks or more (or of two, the last spanning a chunk edge of the view's too), a slab + * across the axes inside it, n₂ × n₃ × ..., and what those axes hold of their own; so + * n₂ × n₃ + n₃ + 1 along three split axes, n₂ + 1 along two, 1 along one, none along none. + * Along an axis of fewer, every axis inside it has two chunks at most, and they are counted + * one by one. With the view's chunks coming out of order (`concurrency` above one), a few + * more. */ +export function peakChunks(r: Region): number { + const split = splitAxes(r), axes = walkOrder(r).filter((a) => split.includes(a)), C = r.chunk; + const n = axes.map((a) => r.last[a] - r.first[a]); + const e = axes.map((a, i) => { // the box's last chunk along the axis spans a chunk edge of the view's + const o = r.first[a] + n[i] - 1, from = r.at[a] + o * C[a], to = r.at[a] + Math.min((o + 1) * C[a], r.shape[a]); + return Math.floor(from / C[a]) !== Math.floor((to - 1) / C[a]); + }); + return mostHeld(n, e); +} + +/** The most chunks held at once by a grid of `n[i]` chunks along each axis (outermost first), + * each spanning two of the view's chunks along every axis but where it is the last along it + * and `e[i]` is false, the view's chunks walked in C order over these axes (peakChunks). */ +function mostHeld(n: number[], e: boolean[]): number { + if (!n.length) return 0; + if (n[0] >= 3 || (n[0] === 2 && e[0])) return prod(n.slice(1)) + mostHeld(n.slice(1), e.slice(1)); + // two chunks at most along this axis and so along those inside it: counted at each of the view's chunks + const zero = n.map(() => 0), boxes = [...chunkIndices(zero, n)].map((o) => [o, o.map((x, i) => x + (x < n[i] - 1 || e[i] ? 1 : 0))]); + const upTo = (o: number[], p: number[]) => { for (let i = 0; i < o.length; i++) if (o[i] !== p[i]) return o[i] < p[i]; return true; }; + let most = 0; + for (const p of chunkIndices(zero, n.map((x, i) => x + (e[i] ? 1 : 0)))) { + most = Math.max(most, boxes.filter(([first, last]) => upTo(first, p) && !upTo(last, p)).length); + } + return most; +} + +/** What materializing region `r` writes and holds, before it starts (`Engine.materializePlan`): + * the array's `shape` and `chunk` shape (the channel axes the ops add first), its chunks + * written (`total`), the view's chunks computed for it (`needed`, once each), the axes the box + * splits (`split`) and the order the view's chunks are walked in (`walk`, outermost first); + * and the most of the array's chunks held at once while they are stitched, `peakChunks`, in + * `peakBytes` (`chunkBytes` each), besides the view's chunks on their way (`concurrency` of + * them), so a page can warn, or offer to put the box's lo on the view's chunk edges along an + * axis, before it starts. */ +export interface MaterializePlan { + shape: number[]; chunk: number[]; total: number; needed: number; + split: number[]; walk: number[]; + chunkBytes: number; peakChunks: number; peakBytes: number; +} + +/** Region `r`'s MaterializePlan. */ +export function plan(r: Region): MaterializePlan { + const peak = peakChunks(r), bytes = chunkBytes(r); + return { + shape: [...r.added, ...r.shape], chunk: [...r.added, ...r.chunk], total: r.total, needed: r.needed, + split: splitAxes(r), walk: walkOrder(r), chunkBytes: bytes, peakChunks: peak, peakBytes: peak * bytes, + }; +} + /** The array's chunk that `add` completes: its index in the array's grid and its bytes, null * for nothing but the fill value. */ export type Completed = [index: number[], bytes: Uint8Array | null]; diff --git a/web/test/materialize.test.ts b/web/test/materialize.test.ts index 003ebc2..b1c65b1 100644 --- a/web/test/materialize.test.ts +++ b/web/test/materialize.test.ts @@ -12,7 +12,7 @@ import assert from "node:assert/strict"; import * as zarr from "zarrita"; import { array3, chunkKey, group3 } from "../src/layout.ts"; import { - allZero, chunkIndices, chunkRange, directoryStore, gzip, memoryStore, region, splitAxes, stitcher, voxelBox, walkOrder, writeChunks, + allZero, chunkIndices, chunkRange, directoryStore, gzip, memoryStore, peakChunks, plan, region, splitAxes, stitcher, voxelBox, walkOrder, writeChunks, type MaterializeOptions, type MaterializeProgress, type Produce, type Region, type WritableStore, } from "../src/materialize.ts"; import type { ViewInfo } from "../src/types.ts"; @@ -540,12 +540,14 @@ test("a box's chunk is written once the last of the view's chunks it is made of const last = o.map((x, a) => Math.floor((r.at[a] + Math.min((x + 1) * chunk[a], r.shape[a]) - 1) / chunk[a])); assert.equal(events.slice(0, i).filter((x) => x.startsWith("make")).at(-1), `make ${last.join(".")}`, `${e} as soon as it is complete`); } - // held: split along all three (4 ≥ 3 ≥ 3 chunks), a slab across the inner two (3 × 3), a row (3) and one more + // held, and projected before it starts (plan): split along all three (4 ≥ 3 ≥ 3 chunks), a + // slab across the inner two (3 × 3), a row (3) and one more; on edges along z, split along y + // and x, a row and one; along z and y, one; along all, none, each of the view's passing through const box = (lo: number[]) => region("v", big, chunk, { lo }); - assert.equal(held(box([1, 1, 1]), fake.make), 3 * 3 + 3 + 1); - assert.equal(held(box([2, 1, 1]), fake.make), 3 + 1); // on edges along z: split along y and x, a row and one - assert.equal(held(box([2, 4, 1]), fake.make), 1); // along z and y: one - assert.equal(held(box([2, 4, 4]), fake.make), 0); // along all: none, each of the view's passing through + for (const [lo, most] of [[[1, 1, 1], 3 * 3 + 3 + 1], [[2, 1, 1], 3 + 1], [[2, 4, 1], 1], [[2, 4, 4], 0]] as [number[], number][]) { + assert.equal(held(box(lo), fake.make), most, `held from ${lo}`); + assert.deepEqual([plan(box(lo)).peakChunks, plan(box(lo)).peakBytes], [most, most * 64], `projected from ${lo}`); + } assert.deepEqual(splitAxes(box([2, 4, 4])), []); // the view's chunks all empty: nothing held in memory but the counts (no buffer made) const r3 = box([1, 1, 1]), s3 = stitcher(r3); @@ -560,12 +562,14 @@ test("the view's chunks are walked with the split axis of the most chunks outerm assert.deepEqual([r.last, splitAxes(r), walkOrder(r)], [[2, 10, 10], [0, 1, 2], [1, 2, 0]]); // y, then x, z fastest // held: a slab across x and z (10 × 2), a row along z (2) and one more; in C order, a slab across y and x assert.equal(held(r, fake.make), 10 * 2 + 2 + 1); + assert.equal(plan(r).peakChunks, 10 * 2 + 2 + 1); assert.equal(held(r, fake.make, [0, 1, 2]), 10 * 10 + 10 + 1); assert.equal(held(r, fake.make, [0, 2, 1]), 10 * 10 + 10 + 1); // the fewest chunks outermost: the most held // the axes it does not split go outermost: on chunk edges along y, then x (10 chunks), then z (2) const edge = region("v", wide, chunk, { lo: [1, 4, 1] }); assert.deepEqual([splitAxes(edge), walkOrder(edge)], [[0, 2], [1, 2, 0]]); assert.equal(held(edge, fake.make), 2 + 1); + assert.equal(plan(edge).peakChunks, 2 + 1); assert.ok(held(edge, fake.make, [0, 1, 2]) > 9 * 10); // C order: a slab across y and x // writeChunks asks in that order, each once, and the box reads back as the view const store = memoryStore(); @@ -576,6 +580,32 @@ test("the view's chunks are walked with the split axis of the most chunks outerm assertBox((await readArray(store, "zarr2")).data, [3, 39, 39], [1, 1, 1]); }); +test("a box's plan: the array, the view's chunks it computes, and the most bytes it holds, the walk's own for any box", () => { + // two float32 channels, 2 × 4 × 4 chunks of 256 bytes, off the edges on every axis + const p = plan(region("grad", bigChannels, chunk, { lo: [1, 3, 2], hi: [8, 9, 11] })); + assert.deepEqual(p, { + shape: [2, 7, 6, 9], chunk: [2, 2, 4, 4], total: 24, needed: 36, split: [0, 1, 2], walk: [0, 2, 1], + chunkBytes: 2 * 32 * 4, peakChunks: 3 * 2 + 2 + 1, peakBytes: (3 * 2 + 2 + 1) * 256, + }); + // the level's extent holds none + const level = plan(region("v", big, chunk, { extent: "level", lo: [1, 3, 2], hi: [8, 9, 11] })); + assert.deepEqual([level.shape, level.split, level.walk, level.peakChunks, level.peakBytes], [[9, 10, 11], [], [0, 1, 2], 0, 0]); + // for boxes of every kind (split along none to three axes, with one or two chunks along some, + // the last spanning a chunk edge or not), what the walk holds is what the plan said + let seed = 7; + const rnd = (k: number) => { seed = (seed * 1103515245 + 12345) % 2 ** 31; return seed % k; }; + const seen = new Set(); + for (let t = 0; t < 400; t++) { + const C = [1 + rnd(5), 1 + rnd(5), 1 + rnd(5)], L = C.map((c) => 1 + rnd(c * 6)); + const lo = L.map((l) => rnd(l)), hi = lo.map((x, a) => x + 1 + rnd(L[a] - x)); + const r = region("v", { ...view, levels: [{ shape: L, voxel: [1, 1, 1], origin: [0, 0, 0] }] }, C, { lo, hi }); + const size = 2 * C[0] * C[1] * C[2], most = held(r, () => new Uint8Array(size).fill(1).buffer); + assert.equal(peakChunks(r), most, `${L} in ${C}, from ${lo} to ${hi}`); + seen.add(splitAxes(r).length); + } + assert.deepEqual([...seen].sort(), [0, 1, 2, 3]); +}); + test("a box's abort: no chunk asked for or written after it, those asked for given up; the chunks written stay", { timeout: 5000 }, async () => { const store = memoryStore(), stop = new AbortController(), asked: string[] = [], given: AbortSignal[] = [], fake = served(big); let atAbort = -1;