Engine.materializePlan: what materialize will write, and its peak memory, before it starts

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

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
 

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

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

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<Materialized> {
+    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<MaterializePlan> {
+    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<Region> {
     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

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];