From 20be6c6a11d577831306c303cc86960984927a1e Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 16 Oct 2026 12:31:13 -0400 Subject: [PATCH 1/2] The browser engine writes a region of a view into a store as an OME-Zarr of its own: Engine.materialize(view, store, {level, lo, hi, format, compression, skipEmpty, concurrency, signal, onProgress, attributes}) computes one level of a served view over a voxel box (rounded out to whole chunks, clipped to the level; the whole level by default) and writes it as zarr v2 with OME-Zarr 0.4 (the default) or zarr v3 with OME-Zarr 0.5, gzipped through CompressionStream, chunks of nothing but zeros left out, into a WritableStore: directoryStore() over a FileSystemDirectoryHandle (a folder picked with showDirectoryPicker, or the origin's private file system), each file written whole through createWritable, or memoryStore(). The metadata is the layout the engine serves for the view, now built in one module for the routes and the writer (web/src/layout.ts), the level written as the group's one dataset at 0 with its own scale and translation and the level's whole shape, and it is written first, so a store stopped part way is still an OME-Zarr. Its chunks come from the producer the routes use, through the same queue and kept chunks, but in the background: the queue (demand.ts) takes background claims, which no client giving up releases, a job only they hold waiting behind every job a client waits for and running in all the slots but one, so browsing goes on while a region is written; a client's claim brings such a job forward. An abort gives up the chunks still waiting and resolves at once with aborted; a chunk that fails rejects naming its key and why, once the writes under way have ended. The tests cover the box's chunks, both formats' files, gzip, skipEmpty, progress, abort, failure, the stores and the queue's background rule, and read the written store back with zarrita; the formats, caching and design pages and the roadmap (a first piece of item 2) say so, and the CHANGELOG has a line for it --- CHANGELOG.md | 11 ++ docs/concepts/caching.md | 23 +++ docs/concepts/formats.md | 44 +++++ docs/design.md | 33 ++++ docs/roadmap.md | 7 +- web/src/demand.ts | 69 ++++++-- web/src/engine.ts | 111 +++++++------ web/src/layout.ts | 89 ++++++++++ web/src/materialize.ts | 278 +++++++++++++++++++++++++++++++ web/test/demand.test.ts | 91 +++++++++- web/test/materialize.test.ts | 313 +++++++++++++++++++++++++++++++++++ 11 files changed, 995 insertions(+), 74 deletions(-) create mode 100644 web/src/layout.ts create mode 100644 web/src/materialize.ts create mode 100644 web/test/materialize.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 07eebdb..56dca7b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -146,3 +146,14 @@ The first release, an alpha: the plugin API can still change before 0.1.0. every session in one WebAssembly instance that its WebGPU build suspends while it waits. Two models' passes, or a second model built while the first computed, could hang the page and crash the tab. A page's own passes on the same ORT take the turn through `exclusive`. +- `Engine.materialize(view, store, options)` writes a box of a served view (or a whole + level) into a store as an OME-Zarr of its own, zarr v2 with OME 0.4 or zarr v3 with OME + 0.5, gzipped, its empty chunks left out, the metadata the engine serves for the view + (`web/src/layout.ts`, shared with the routes) with the one level as its one dataset; + `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 + chunks share the queue and the kept chunks with the viewer's but run in the background: + behind every chunk a client waits for, in all the queue's slots but one, and never dropped + as abandoned (`submit`, `claim` and `release` take `background`; + [caching](docs/concepts/caching.md#order-of-work-and-requests-given-up-on)). diff --git a/docs/concepts/caching.md b/docs/concepts/caching.md index ffc4426..e7b2298 100644 --- a/docs/concepts/caching.md +++ b/docs/concepts/caching.md @@ -208,6 +208,29 @@ one thread, so the first view of a deep chain used one core. The browser page keeps the same rules (`web/src/demand.ts`); it learns of a given-up request from its streamed reply being cancelled, a service worker being told of nothing else. +The browser page has one class of work more: chunks it computes for itself, in the +background, to write a region of a view into a store (`Engine.materialize`; see +[formats](formats.md#a-region-written-to-a-store-browser-engine)). They go through the same +queue as the viewer's chunks, so a chunk the viewer has computed is taken from the kept +chunks, one on its way for the viewer is shared, and one written is kept for the viewer +after. Three rules keep browsing usable while a region is written: + +* A background claim is never released because a client gave up, so nothing materialize + asked for is dropped as abandoned; only its own abort gives its waiting chunks up (and + they are dropped if no client wants them either). +* A job only background claims hold waits behind every job a client waits for, whatever + their levels; among themselves, background jobs keep the usual order. +* Background jobs run in all the queue's slots but one (in the one, when there is only + one), so a viewer's request finds a slot free, or the next one to free, rather than + waiting for the region's chunks in line before it. + +A viewer asking for a chunk that waits in the background brings it forward, and it goes +back if the viewer gives up on it. Materialize asks for a few chunks at once (as many as +the queue has slots, by default: the next one is waiting while one is written), not the +whole region, so the queue holds a few of its chunks at a time however large the region. +Its chunks are kept as the viewer's are (256 MB, least recently used out), as are a staged +view's stage outputs, so writing a large region cycles through both. + ## What has been computed Zoomed out, a dataset shows its ops run on the coarse level, not a downsample of the diff --git a/docs/concepts/formats.md b/docs/concepts/formats.md index 98abec5..23c80d5 100644 --- a/docs/concepts/formats.md +++ b/docs/concepts/formats.md @@ -73,6 +73,50 @@ view as zarr v2 with OME-Zarr 0.4 (`virtual//zarr2//`), 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 region written to a store (browser engine) + +The browser engine also writes a view, or a box of it, into a store as files +(`Engine.materialize`; see [the design notes](../design.md#another-page-on-the-engine)): a +folder the visitor picks, the origin's private file system or memory. What it writes is +the layout it serves for the view, one level of it, so the result opens wherever the view +did and lands in the same place: the group lists the level written as its one dataset, +`0`, with that level's scale and translation; the level's array has the level's whole +shape (a box keeps its place in it), the chunk shape and dtype served, and chunks +compressed with gzip unless asked not to be; and the chunks of the box, rounded out to +whole chunks. 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 one chunk, as served. As zarr v2 with OME-Zarr 0.4 (the default, which +more tools read): + +```text +.zgroup {"zarr_format": 2} +.zattrs {...extra attributes, "multiscales": [{"name": , "axes": [...], + "datasets": [{"path": "0", "coordinateTransformations": [scale, translation]}], + "version": "0.4"}]} +0/.zarray {"zarr_format": 2, "shape": , "chunks": , "dtype": "// each chunk, gzipped +``` + +As zarr v3 with OME-Zarr 0.5: + +```text +zarr.json {"zarr_format": 3, "node_type": "group", "attributes": {...extra attributes, + "ome": {"version": "0.5", "multiscales": [{... "datasets": [{"path": "0", ...}]}]}}} +0/zarr.json the level's array as served, its codecs [{"name": "bytes", ...}, + {"name": "gzip", "configuration": {"level": 6}}] +0/c/// each chunk, gzipped +``` + +The metadata is written first, so a store whose writing was stopped part way is still a +valid OME-Zarr, zero where no chunk was written. Extra attributes (a consumer's +provenance) sit beside the OME ones in the group's attributes, which they cannot replace. +gzip is what a browser compresses with on its own (`CompressionStream`; level 6 is the +level it writes at, which no reader needs); the Python server's default, blosc, would need +an encoder of its own in the page. + ## Meshes, computed when fetched The `mesh` frontend serves a dataset's surface in Neuroglancer's precomputed mesh formats, diff --git a/docs/design.md b/docs/design.md index d6821d5..e27227e 100644 --- a/docs/design.md +++ b/docs/design.md @@ -516,6 +516,39 @@ as text). What such a page gets, in the order it uses it: so the stage is typed with the rest and run in `npm test` through the WebAssembly backend, which runs in Node as in a browser without WebGPU, on two models of a hundred bytes (`web/test/onnx.test.ts`, the models written by `web/test/fixtures/make_onnx.py`). +* `Engine.materialize(view, store, options)` writes a region of a served view into a store + as an OME-Zarr of its own (`web/src/materialize.ts`), to open again later or copy where + other tools read it: a model's output over a box, computed once in the page and kept. + The store is anything with `put(path, bytes)`; `directoryStore(dir, root)` makes one of + a `FileSystemDirectoryHandle` (a folder the visitor picks with + `showDirectoryPicker({mode: "readwrite"})`, in Chrome and Edge, or the origin's private + file system, `navigator.storage.getDirectory()`), its subdirectories made as needed and + each file written whole through `createWritable()`, and `memoryStore()` one of a `Map`, + for tests and for zipping a small result. The options: the view's `level` (the finest by + default), a voxel box `lo`, `hi` on it (the whole level by default; rounded out to whole + chunks and clipped to the level), `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`) 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 + the group's one dataset at `0` with its own scale and translation, so it lands where the + view does ([formats](concepts/formats.md#a-region-written-to-a-store-browser-engine)). 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 + browsing goes on while a region is written. An abort starts no more chunks, gives up + those waiting for their turn, lets the writes under way end and resolves at once with + `aborted`; a chunk that fails to compute or be written rejects it with the chunk's key + and why, once the writes under way have ended. This is a first piece of + [roadmap](roadmap.md) item 2: an explicit write of a box, asked for by the page, not yet + the cache on disk that browsing fills. The parts that need no engine (the chunks of a + box, the files, gzip, the loop over any producer, the stores) run in `npm test`, the + result read back by zarrita (`web/test/materialize.test.ts`). * `Engine.unserve(view)` stops serving a view: its infos, producer and kept chunks are forgotten and chunks of it waiting to be computed dropped, so a page that serves a model again under a new name (viewers cache by URL, so an edit is a new name) does not keep the diff --git a/docs/roadmap.md b/docs/roadmap.md index a7d0aa2..e123376 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -36,7 +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.) + 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 + ([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.) 3. **Lazy-compute story for non-viewer clients.** Demonstrate dask/tensorstore reading a served pipeline and running downstream analysis with no intermediate written. (The read itself is shown in [getting started](getting-started.md#read-the-result-from-python-instead-of-a-viewer); diff --git a/web/src/demand.ts b/web/src/demand.ts index c06aed7..d770b05 100644 --- a/web/src/demand.ts +++ b/web/src/demand.ts @@ -8,6 +8,14 @@ // placeholder while the fine chunk computes) wants the finer. Within a level the first asked // for goes first. A few run at once, so the GPU works on one while the next ones read, and the // rest cannot flood the network with reads. +// +// One rule more than chunkmirage.demand's: work asked for in the background (a page writing a +// region to a store, Engine.materialize), not for a client waiting on it. A background claim +// is never released by a client giving up, so the work is not dropped, but a job only +// background claims hold waits behind every other, whatever its level, and such jobs take at +// most all the slots but one (one, if there is only one): a viewer's request finds a slot free +// or the next one to free, while the background work goes on in the rest. A client asking for +// a job waiting in the background brings it forward, and it goes back if the client gives up. const SLOTS = 3; // jobs run at once @@ -26,16 +34,21 @@ export class Cancelled extends Error { constructor(why = "no request wants this interface Job { run: () => Promise; level: number; seq: number; // seq: when first asked for - demand: number; // requests holding a claim on it + demand: number; // requests holding a claim on it + background: number; // of them, background claims: a job only they hold waits behind the rest resolve: (v: T) => void; reject: (e: unknown) => void; onDrop: () => void; } /** Jobs waiting, in the order described at the top; `slots` of them run at once. */ export class Queue { - private slots: number; - constructor(slots = SLOTS) { this.slots = slots; } + /** Jobs run at once. */ + readonly slots: number; + /** Of them, jobs only background claims hold (`submit`'s `background`): all but one, so a + * client's request finds a slot free; one, if there is only one. */ + readonly backgroundSlots: number; + constructor(slots = SLOTS) { this.slots = slots; this.backgroundSlots = Math.max(1, slots - 1); } private waiting = new Map>(); - private active = new Map(); // running: key -> level + private active = new Map(); // running, by key dropped = 0; // waiting fits no request wanted any more, never run onChange: () => void = () => {}; // a fit waiting, started, done or dropped private seq = 0; @@ -44,26 +57,38 @@ export class Queue { /** Fits running and waiting, per level. */ byLevel(): { running: Map; waiting: Map } { const count = (levels: Iterable) => { const m = new Map(); for (const l of levels) m.set(l, (m.get(l) ?? 0) + 1); return m; }; - return { running: count(this.active.values()), waiting: count([...this.waiting.values()].map((j) => j.level)) }; + return { running: count([...this.active.values()].map((a) => a.level)), waiting: count([...this.waiting.values()].map((j) => j.level)) }; } - /** A new fit, wanted by one request; `onDrop` runs, synchronously, if it is dropped unrun. */ - submit(key: string, level: number, run: () => Promise, onDrop: () => void): Promise { + /** A new fit, wanted by one request (one in the background, with `background`); `onDrop` + * runs, synchronously, if it is dropped unrun. */ + submit(key: string, level: number, run: () => Promise, onDrop: () => void, background = false): Promise { return new Promise((resolve, reject) => { - const job: Job = { run, level, resolve, reject, onDrop, seq: ++this.seq, demand: 1 }; + const job: Job = { run, level, resolve, reject, onDrop, seq: ++this.seq, demand: 1, background: background ? 1 : 0 }; this.waiting.set(key, job); this.pump(); this.onChange(); }); } - /** Another request wants a fit: if it is still waiting, it counts one more. */ - claim(key: string) { const job = this.waiting.get(key); if (job) job.demand++; } + /** Another request wants a fit (one in the background, with `background`): if it is still + * waiting, it counts one more, and a client's claim brings a background job forward. */ + claim(key: string, background = false) { + const job = this.waiting.get(key); + if (!job) return; + job.demand++; + if (background) job.background++; + else this.pump(); // in the background before: it may take a slot kept for clients + } - /** A request no longer wants a fit: waiting and wanted by no one, it is dropped unrun. */ - release(key: string) { + /** A request no longer wants a fit (`background`: one that claimed it in the background): + * waiting and wanted by no one, it is dropped unrun; wanted in the background alone, it goes + * back behind the rest. */ + release(key: string, background = false) { const job = this.waiting.get(key); - if (!job || --job.demand > 0) return; + if (!job) return; + if (background && job.background > 0) job.background--; + if (--job.demand > 0) return; this.drop(key); } @@ -84,18 +109,26 @@ export class Queue { while (this.active.size < this.slots && this.waiting.size) { let best: [string, Job] | null = null; for (const e of this.waiting) if (!best || before(e[1], best[1])) best = e; - const [key, job] = best!; + const [key, job] = best!, background = inBackground(job); + // the best in the background: so is every job waiting, and they have slots of their own + if (background && [...this.active.values()].filter((a) => a.background).length >= this.backgroundSlots) return; this.waiting.delete(key); - this.active.set(key, job.level); + this.active.set(key, { level: job.level, background }); job.run().then(job.resolve, job.reject).finally(() => { this.active.delete(key); this.pump(); this.onChange(); }); } } } -/** Which of two waiting fits goes first: the finer level (a client asking for one place at - * two levels, as a viewer does to show a coarse placeholder while the fine chunk computes, - * wants the finer one; coarser fits run when nothing finer waits), then the first asked for. */ +/** Whether only background claims hold a job. */ +const inBackground = (job: Job) => job.demand === job.background; + +/** Which of two waiting fits goes first: one a client waits for before one wanted in the + * background alone; then the finer level (a client asking for one place at two levels, as a + * viewer does to show a coarse placeholder while the fine chunk computes, wants the finer + * one; coarser fits run when nothing finer waits), then the first asked for. */ function before(a: Job, b: Job): boolean { + const ba = inBackground(a), bb = inBackground(b); + if (ba !== bb) return bb; return a.level !== b.level ? a.level < b.level : a.seq < b.seq; } diff --git a/web/src/engine.ts b/web/src/engine.ts index 19856fb..af3517a 100644 --- a/web/src/engine.ts +++ b/web/src/engine.ts @@ -10,17 +10,20 @@ // Neuroglancer (virtual///...) and GeoZarr for map clients such as OpenLayers // (virtual//geo////...), the zarr-conventions multiscales, proj: // and spatial: attributes on the group. And zarr v2 with OME 0.4 (virtual//zarr2//), -// for readers that predate zarr v3's final spec, such as GDAL 3.8. +// for readers that predate zarr v3's final spec, such as GDAL 3.8. A region of a view is +// written into a store as the same OME-Zarr by `materialize` (materialize.ts). import { asRead, computedSource, needsPython, TYPED, wholeChunk } from "./asread"; 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 { 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, ViewAxis, ViewInfo } from "./types"; +import type { Answer, Ask, Later, MeshCall, PipelineView, Plan, Printed, Reply, SourceInfo, ToPyWorker, ToReader, ViewInfo } from "./types"; -export const TO_SECONDS: Record = { s: 1, second: 1, millisecond: 1e-3, ms: 1e-3, minute: 60, hour: 3600, day: 86400 }; -const OME_UNIT: Record = { nm: "nanometer", um: "micrometer", m: "meter", s: "second" }; +export { isTime, TO_SECONDS } from "./layout"; +export { directoryStore, memoryStore, type MaterializeOptions, 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:" }, @@ -28,9 +31,6 @@ const CONVENTIONS = [ ]; const KEPT_BYTES = 256 * 2 ** 20; // computed chunks kept for clients that refetch const STAGED_BYTES = 512 * 2 ** 20; // the outputs of a staged view's stages (stages.ts), kept for the stages after them -const V2_DTYPE: Record = { - uint8: "|u1", int8: "|i1", uint16: " = Object.fromEntries(Object.values(schema.$defs as Record) .flatMap((d) => (d.properties?.op?.const && d["x-packages"] ? [[d.properties.op.const, d["x-packages"]]] : []))); @@ -77,7 +77,6 @@ class Rpc { const asJson = (o: unknown) => ({ status: 200, body: JSON.stringify(o), type: "application/json" }); const notFound = { status: 404, body: "", type: "text/plain" }; -export const isTime = (a: ViewAxis) => a.unit in TO_SECONDS; /** `Engine.start`'s options. `packages`: what Python loads beyond numpy and the ops' own * imports (each op's schema says those). `python`: when the Pyodide workers start; "auto" @@ -389,20 +388,8 @@ export class Engine { get running(): number { return this.chunks?.running ?? 0; } get waiting(): number { return this.chunks ? [...this.chunks.byLevel().waiting.values()].reduce((a, b) => a + b, 0) : 0; } - private omeGroup(view: string) { - const v = this.infos[view], added = (v.added ?? []).map(() => 0); // channel axes the ops add - return { - zarr_format: 3, node_type: "group", - attributes: { ome: { version: "0.5", multiscales: [{ - name: view, - axes: [...added.map(() => ({ name: "c", type: "channel" })), - ...v.axes.map((a) => ({ name: a.name, type: isTime(a) ? "time" : "space", unit: OME_UNIT[a.unit] ?? (a.unit || undefined) }))], - datasets: v.levels.map((l, i) => ({ path: String(i), coordinateTransformations: [ - { type: "scale", scale: [...added.map(() => 1), ...l.voxel] }, { type: "translation", translation: [...added, ...l.origin] }, - ] })), - }] } }, - }; - } + /** A view's levels, each a dataset of its group at its own index. */ + private datasets(view: string) { return this.infos[view].levels.map((_, level) => ({ path: String(level), level })); } /** The GeoZarr group of a view: the conventions, its levels as the multiscales layout, its * projection (the source's own, which the client must know by `proj`) and bounding box. */ @@ -419,20 +406,6 @@ export class Engine { }; } - /** A level's array; in the GeoZarr layout (`map`) as the map's y, x band, the view's one - * plane along z (a map client reads 2-D bands: with a third axis it would need the - * group's consolidated metadata to know to select it). */ - private array(view: string, level: number, map = false) { - const v = this.infos[view], keep = map ? -2 : 0, added = map ? [] : v.added ?? []; // channels: chunked whole - return { - zarr_format: 3, node_type: "array", shape: [...added, ...v.levels[level].shape.slice(keep)], data_type: v.out, fill_value: 0, - chunk_grid: { name: "regular", configuration: { chunk_shape: [...added, ...this.views[view].chunk.slice(keep)] } }, - chunk_key_encoding: { name: "default", configuration: { separator: "/" } }, - codecs: [{ name: "bytes", configuration: { endian: "little" } }], - dimension_names: [...added.map(() => "c"), ...v.axes.map((a) => a.name).slice(keep)], attributes: {}, - }; - } - /** The level a view's mesh is made from (chunkmirage.meshes.mesh_level). */ meshLevel(view: string): number { const v = this.sources[view], spec = this.views[view].mesh ?? {}; @@ -490,14 +463,17 @@ export class Engine { }, [block.buffer]); } - private chunk(view: string, level: number, index: number[], mesh?: MeshCall): Answered { + /** Chunk `index` of a view's level (or with `mesh`, its mesh): kept, or on its way for + * another request, or queued. `background`: for `materialize`, not a client, so it waits + * behind the clients' work and no client giving up drops it (demand.ts). */ + private chunk(view: string, level: number, index: number[], mesh?: MeshCall, background = false): Answered { const key = `${view}/${mesh ? `mesh${level}` : level}/${index.join(".")}`, done = this.kept.get(key); if (done) return { status: 200, body: done.slice(0), type: "application/octet-stream" }; const chunks = this.chunks!, claim = new Claim(); let pending = this.inflight.get(key); - if (pending) chunks.claim(key); // another request waits for it too + if (pending) chunks.claim(key, background); // another request waits for it too else { - pending = chunks.submit(key, level, () => this.compute(view, level, index, mesh), () => { this.counts.dropped++; }) + pending = chunks.submit(key, level, () => this.compute(view, level, index, mesh), () => { this.counts.dropped++; }, background) .then((b) => { this.counts.computed++; if (this.views[view]) this.kept.put(key, b); return b; }, (e) => { // kept unless unserved meanwhile if (!(e instanceof Cancelled)) { this.counts.failed++; this.failure = { view, message: String((e as Error).message ?? e) }; console.error(`chunk ${key}: ${String((e as Error).message ?? e).trim().split("\n").slice(-3).join(" / ")}`); } throw e; @@ -505,13 +481,48 @@ export class Engine { .finally(() => { this.inflight.delete(key); this.onChange(); }); this.inflight.set(key, pending); } - claim.onCancel(() => chunks.release(key)); + claim.onCancel(() => chunks.release(key, background)); const mine = pending.then((b) => b.slice(0)); // each request its own copy: replies transfer it mine.catch(() => {}); this.onChange(); return { status: 200, type: "application/octet-stream", pending: mine, claim }; } + /** Write a region of `view` into `store` as an OME-Zarr of its own (materialize.ts): one + * level of the view (`options.level`, the finest by default) over a box of its voxels + * (`lo`, `hi`, rounded out to whole chunks; the whole level by default), as zarr v2 with + * OME-Zarr 0.4 or zarr v3 with OME-Zarr 0.5, gzipped, its chunks of nothing but zeros left + * out. The group and the array are written first, so the store reads as the view from the + * start, zero where no chunk is yet; then the chunks, each computed as a client's request + * for it is (kept, or shared with one on its way, and kept after) but in the background: + * behind every chunk a client waits for, in all the queue's slots but one, and never + * dropped because a client gave up on it (demand.ts). `concurrency` chunks are asked for + * at once, 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. */ + async materialize(view: string, store: WritableStore, 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 }); + } + + /** Chunk `index` of a view's level for `materialize`: its whole bytes, as a client's request + * gets them; `signal`'s abort gives it up (dropped unrun if no client wants it either). */ + private async backgroundChunk(view: string, level: number, index: number[], signal: AbortSignal): Promise { + if (!this.chunks) throw new Error("the engine was stopped"); + if (!this.infos[view]) throw new Error(`${view} is no longer served`); + signal.throwIfAborted(); + const a = this.chunk(view, level, index, undefined, true)!; + if (!("pending" in a)) return a.body as ArrayBuffer; // kept + const cancel = () => a.claim.cancel(); + signal.addEventListener("abort", cancel, { once: true }); + try { return await a.pending; } finally { signal.removeEventListener("abort", cancel); } // done: no claim left to give up + } + /** Projection codes of the GeoZarr groups, by view (set by the map page). */ proj: Record = {}; @@ -539,12 +550,12 @@ export class Engine { view = rest[0]; rest = rest.slice(1); if (!this.infos[view]) return notFound; - if (rest.length === 1 && rest[0] === "zarr.json") return asJson(this.omeGroup(view)); + if (rest.length === 1 && rest[0] === "zarr.json") return asJson(group3(view, this.infos[view], this.datasets(view))); if (rest[0] === "mesh" && this.views[view].mesh) return this.mesh(view, rest.slice(1).join("/"), range); } const level = Number(rest[0]); if (!this.infos[view].levels[level]) return notFound; - if (rest.length === 2 && rest[1] === "zarr.json") return asJson(this.array(view, level, map)); + if (rest.length === 2 && rest[1] === "zarr.json") return asJson(array3(this.infos[view], this.views[view].chunk, level, { map })); const added = map ? 0 : this.infos[view].added?.length ?? 0, key = rest.slice(2).map(Number); if (rest[1] !== "c" || key.length !== added + this.infos[view].levels[level].shape.length || key.slice(0, added).some((i) => i !== 0)) return notFound; return this.chunk(view, level, key.slice(added)); // a channel axis is one chunk: the volume's index @@ -555,21 +566,13 @@ export class Engine { private zarr2(view: string, rest: string[]): Answered { const v = this.infos[view]; if (!v) return notFound; - const added = v.added ?? [], names = [...added.map(() => "c"), ...v.axes.map((a) => a.name)]; + const added = v.added ?? []; if (rest.length === 1 && rest[0] === ".zgroup") return asJson({ zarr_format: 2 }); - if (rest.length === 1 && rest[0] === ".zattrs") { - const g = this.omeGroup(view).attributes.ome.multiscales[0]; - return asJson({ multiscales: [{ ...g, version: "0.4" }] }); - } + if (rest.length === 1 && rest[0] === ".zattrs") return asJson(attrs2(view, v, this.datasets(view))); const level = Number(rest[0]), l = v.levels[level]; if (!l) return notFound; - if (rest.length === 2 && rest[1] === ".zattrs") return asJson({ _ARRAY_DIMENSIONS: names }); - if (rest.length === 2 && rest[1] === ".zarray") { - return asJson({ - zarr_format: 2, shape: [...added, ...l.shape], chunks: [...added, ...this.views[view].chunk], dtype: V2_DTYPE[v.out] ?? v.out, - compressor: null, fill_value: 0, order: "C", filters: null, dimension_separator: "/", - }); - } + if (rest.length === 2 && rest[1] === ".zattrs") return asJson({ _ARRAY_DIMENSIONS: dimensionNames(v) }); + if (rest.length === 2 && rest[1] === ".zarray") return asJson(array2(v, this.views[view].chunk, level)); const key = rest.slice(1).map(Number); if (key.length !== added.length + l.shape.length || key.some((i) => !Number.isInteger(i)) || key.slice(0, added.length).some((i) => i !== 0)) return notFound; return this.chunk(view, level, key.slice(added.length)); diff --git a/web/src/layout.ts b/web/src/layout.ts new file mode 100644 index 0000000..ecf24d5 --- /dev/null +++ b/web/src/layout.ts @@ -0,0 +1,89 @@ +// A view's OME-Zarr layouts, as the engine serves them (engine.ts's routes) and writes them +// into a store (materialize.ts), from one place so the two cannot differ: zarr v3 with +// OME-Zarr 0.5, and zarr v2 with OME-Zarr 0.4. The group's metadata, a level's array +// metadata and a chunk's key, as objects ready for JSON. A channel axis the view's ops add +// (`added`) comes first and is one chunk, as the routes serve it. Nothing here imports a +// value, so Node's tests run it. +import type { ViewAxis, ViewInfo } from "./types.ts"; + +export const TO_SECONDS: Record = { s: 1, second: 1, millisecond: 1e-3, ms: 1e-3, minute: 60, hour: 3600, day: 86400 }; +const OME_UNIT: Record = { nm: "nanometer", um: "micrometer", m: "meter", s: "second" }; +/** zarr v2's dtypes, by numpy's name (the engine's dtypes); little endian. */ +export const V2_DTYPE: Record = { + uint8: "|u1", int8: "|i1", uint16: " a.unit in TO_SECONDS; + +/** A level of a view as one of a group's datasets: its path in the group, and which of the + * view's levels it is. */ +export interface Dataset { path: string; level: number } + +/** A zarr v3 codec after `bytes` (a compressor). */ +export interface Codec { name: string; configuration: Record } + +/** The OME multiscales entry of view `name`, whose datasets are the levels `datasets` name: + * the channel axes its ops add, then its own (space, or time by their unit); each level's + * scale its voxel size and its translation its origin, so it lands where the view does. */ +export function multiscale(name: string, v: ViewInfo, datasets: Dataset[]) { + const added = (v.added ?? []).map(() => 0); + return { + name, + axes: [...added.map(() => ({ name: "c", type: "channel" })), + ...v.axes.map((a) => ({ name: a.name, type: isTime(a) ? "time" : "space", unit: OME_UNIT[a.unit] ?? (a.unit || undefined) }))], + datasets: datasets.map(({ path, level }) => ({ path, coordinateTransformations: [ + { type: "scale", scale: [...added.map(() => 1), ...v.levels[level].voxel] }, + { type: "translation", translation: [...added, ...v.levels[level].origin] }, + ] })), + }; +} + +/** A view's zarr v3 group (zarr.json): OME-Zarr 0.5's multiscales in its attributes, beside + * any `extra` attributes (which cannot replace `ome`). */ +export function group3(name: string, v: ViewInfo, datasets: Dataset[], extra: Record = {}) { + return { zarr_format: 3, node_type: "group", attributes: { ...extra, ome: { version: "0.5", multiscales: [multiscale(name, v, datasets)] } } }; +} + +/** A view's zarr v2 group attributes (.zattrs): OME-Zarr 0.4's multiscales, beside any + * `extra` attributes (which cannot replace `multiscales`). The group itself is + * `{"zarr_format": 2}` (.zgroup). */ +export function attrs2(name: string, v: ViewInfo, datasets: Dataset[], extra: Record = {}) { + return { ...extra, multiscales: [{ ...multiscale(name, v, datasets), version: "0.4" }] }; +} + +/** The names of a level's axes: "c" for each channel axis the ops add, then the view's. */ +export const dimensionNames = (v: ViewInfo) => [...(v.added ?? []).map(() => "c"), ...v.axes.map((a) => a.name)]; + +/** A level's zarr v3 array (its zarr.json), chunked `chunk` (the level's own axes), the + * channel axes the ops add one chunk each; its bytes little endian, then `codecs`. In the + * GeoZarr layout (`map`), the map's y, x band alone, the view's one plane along z (a map + * client reads 2-D bands: with a third axis it would need the group's consolidated + * metadata to know to select it). */ +export function array3(v: ViewInfo, chunk: number[], level: number, { map = false, codecs = [] }: { map?: boolean; codecs?: Codec[] } = {}) { + const keep = map ? -2 : 0, added = map ? [] : v.added ?? []; + return { + zarr_format: 3, node_type: "array", shape: [...added, ...v.levels[level].shape.slice(keep)], data_type: v.out, fill_value: 0, + chunk_grid: { name: "regular", configuration: { chunk_shape: [...added, ...chunk.slice(keep)] } }, + chunk_key_encoding: { name: "default", configuration: { separator: "/" } }, + codecs: [{ name: "bytes", configuration: { endian: "little" } }, ...codecs], + dimension_names: map ? v.axes.map((a) => a.name).slice(keep) : dimensionNames(v), attributes: {}, + }; +} + +/** A level's zarr v2 array (.zarray), chunked `chunk`, the channel axes the ops add one + * chunk each, compressed by `compressor` (numcodecs' config; null: as they are). Its axes' + * names are its .zattrs, `{"_ARRAY_DIMENSIONS": dimensionNames(v)}`. */ +export function array2(v: ViewInfo, chunk: number[], level: number, compressor: Record | null = null) { + const added = v.added ?? []; + return { + zarr_format: 2, shape: [...added, ...v.levels[level].shape], chunks: [...added, ...chunk], dtype: V2_DTYPE[v.out] ?? v.out, + compressor, fill_value: 0, order: "C", filters: null, dimension_separator: "/", + }; +} + +/** The key of chunk `index` (on the level's own axes) of the array at `path`: zarr v3's + * default encoding (`c/`, then the indices "/" apart), or zarr v2's with "/" as its + * dimension separator; each of the `added` channel axes is one chunk, index 0. */ +export function chunkKey(format: "zarr3" | "zarr2", path: string, added: number, index: number[]): string { + const at = [...Array(added).fill(0), ...index].join("/"); + return format === "zarr3" ? `${path}/c/${at}` : `${path}/${at}`; +} diff --git a/web/src/materialize.ts b/web/src/materialize.ts new file mode 100644 index 0000000..b16dc8a --- /dev/null +++ b/web/src/materialize.ts @@ -0,0 +1,278 @@ +// A region of a view the engine serves, written into a store as an OME-Zarr of its own +// (Engine.materialize): a box of the view computed once and kept, to open again later or to +// copy where other tools read it. The store is a directory (one the visitor picks with the +// File System Access API's showDirectoryPicker, in Chrome and Edge, or the origin's private +// file system) or memory. What is written is what the engine serves for the view +// (layout.ts), one level of it, its chunks compressed. Here are the parts that need no +// engine, so Node's tests run them: the chunks a box covers, the files written, gzip, and +// the loop that has the chunks made and writes them, from any producer (the engine's queue +// on the page, a fake one in the tests). +import { array2, array3, attrs2, chunkKey, dimensionNames, group3 } from "./layout.ts"; // with its extension: Node's runner resolves no other +import type { ViewInfo } from "./types.ts"; + +/** Where materialize writes: files by "/"-separated path relative to the store's root. */ +export interface WritableStore { + put(path: string, bytes: Uint8Array): Promise; +} + +/** The names of `path`, "/" apart; empty ones (a leading, trailing or doubled "/") left out. + * A store's paths go down from its root, so "." and ".." are refused. */ +function names(path: string): string[] { + const parts = path.split("/").filter(Boolean); + if (parts.some((p) => p === "." || p === "..")) throw new Error(`${path}: a store's paths go down from its root, with no "." or ".."`); + return parts; +} + +/** A directory picked with showDirectoryPicker({mode: "readwrite"}) (or any + * FileSystemDirectoryHandle, e.g. OPFS's navigator.storage.getDirectory()) as a + * WritableStore: subdirectories made as needed (getDirectoryHandle(name, {create: true}), + * each once however many files go in it), each file written whole with createWritable(), + * which writes to a swap file that replaces the file when closed, so a file is never seen + * half written. `root` (optional) is a subdirectory path under it, made as needed (e.g. + * "mito.zarr"). */ +export function directoryStore(dir: FileSystemDirectoryHandle, root?: string): WritableStore { + const base = root ? names(root) : []; + const dirs = new Map>(); + const at = (parts: string[]): Promise => { + if (!parts.length) return Promise.resolve(dir); + const key = parts.join("/"); + let made = dirs.get(key); + if (!made) { + made = at(parts.slice(0, -1)).then((parent) => parent.getDirectoryHandle(parts[parts.length - 1], { create: true })); + dirs.set(key, made); + made.catch(() => dirs.delete(key)); // tried again by the next file in it + } + return made; + }; + return { + async put(path, bytes) { + const parts = [...base, ...names(path)]; + if (parts.length === base.length) throw new Error(`"${path}" is not a file's path`); + const file = await (await at(parts.slice(0, -1))).getFileHandle(parts[parts.length - 1], { create: true }); + const w = await file.createWritable(); + try { + await w.write(bytes as Uint8Array); + await w.close(); + } catch (e) { + await w.abort().catch(() => {}); // the swap file discarded, the file as it was + throw e; + } + }, + }; +} + +/** An in-memory store (Map, by path without a leading "/"), for tests + * and for zipping small results. Each file is a copy of the bytes put. */ +export function memoryStore(): WritableStore & { files: Map } { + const files = new Map(); + return { + files, + async put(path, bytes) { + const parts = names(path); + if (!parts.length) throw new Error(`"${path}" is not a file's path`); + files.set(parts.join("/"), bytes.slice()); + }, + }; +} + +/** How a materialize is going: chunks done (written or skipped as empty) of `total`, and of + * them `written`, in `bytes` as stored. */ +export interface MaterializeProgress { done: number; total: number; written: number; bytes: number } + +/** `Engine.materialize`'s options. */ +export interface MaterializeOptions { + /** The view's level to write (default 0, the finest). */ + level?: number; + /** Voxel box [lo, hi) on that level, in the level's voxel axes (z, y, x...); default the + * whole level. Rounded out to whole chunks, and clipped to the level. */ + lo?: number[]; hi?: number[]; + /** "zarr2" (OME-Zarr 0.4: .zgroup, .zattrs, .zarray, "/" as the dimension separator) or + * "zarr3" (OME-Zarr 0.5, zarr.json); default "zarr2", which more tools read. */ + format?: "zarr3" | "zarr2"; + /** gzip through CompressionStream (zarr v2: the compressor {"id": "gzip", "level": 6}; + * zarr v3: the "gzip" codec after "bytes"), or null for none; default "gzip". */ + compression?: "gzip" | null; + /** A chunk of nothing but zero bytes (the fill value) is not written, zarr reading a + * missing chunk as the fill value; default true. */ + skipEmpty?: boolean; + /** Chunks asked for at once (default: as many as the engine's queue has slots, of which + * all but one run materialize's chunks at most, so the next is waiting while one is + * written). */ + concurrency?: number; + /** Aborted: no chunk is started after it, chunks waiting for their turn are given up, the + * writes under way end, and materialize resolves with `aborted`; chunks written stay. */ + signal?: AbortSignal; + /** Told once the metadata is written (done 0), then after each chunk. */ + onProgress?: (p: MaterializeProgress) => void; + /** Attributes beside the OME multiscales in the group's attributes (the consumer's + * provenance); `ome` (zarr v3) or `multiscales` (zarr v2) among them is the view's. */ + attributes?: Record; +} + +/** What a materialize wrote: of the box's `chunks`, those `written` (in `bytes`, as + * stored) and those `skipped` as empty; `aborted` if its signal stopped it first. */ +export interface Materialized { chunks: number; written: number; skipped: number; bytes: number; aborted: boolean } + +/** The written level's path in the store: its one dataset, "0", as an OME-Zarr's finest + * level is, whichever level of the view it is (the scale and translation say where it lies). */ +export const LEVEL_PATH = "0"; +/** The gzip level written in the metadata: CompressionStream's, zlib's default (a reader + * needs none of it to decompress). */ +export const GZIP_LEVEL = 6; + +/** The chunk grid's indices [lo, hi) covering voxels [lo, hi) of a level of `shape` cut in + * chunks of `chunk`: the box rounded out to whole chunks and clipped to the level; the + * whole level where `lo` or `hi` is left out. A box with no voxel of the level is refused. */ +export function chunkRange(shape: number[], chunk: number[], lo?: number[], hi?: number[]): { lo: number[]; hi: number[] } { + for (const [what, box] of [["lo", lo], ["hi", hi]] as const) { + if (box && (box.length !== shape.length || box.some((x) => typeof x !== "number" || Number.isNaN(x)))) { + throw new Error(`${what} ${JSON.stringify(box)}: one number for each of the level's ${shape.length} axes`); + } + } + const inLo = shape.map((n, a) => Math.min(Math.max(lo?.[a] ?? 0, 0), n)); // the box clipped to the level + const inHi = shape.map((n, a) => Math.min(Math.max(hi?.[a] ?? n, 0), n)); + if (inHi.some((h, a) => h <= inLo[a])) { + throw new Error(`the box from [${(lo ?? shape.map(() => 0)).join(", ")}] to [${(hi ?? shape).join(", ")}] holds no voxel of the level, which is ${shape.join("×")}`); + } + return { lo: inLo.map((x, a) => Math.floor(x / chunk[a])), hi: inHi.map((x, a) => Math.ceil(x / chunk[a])) }; +} + +/** The chunk indices [lo, hi), in C order (the last axis fastest). */ +export function* chunkIndices(lo: number[], hi: number[]): Generator { + if (hi.some((h, a) => h <= lo[a])) return; + const at = [...lo]; + for (;;) { + yield [...at]; + let a = at.length - 1; + while (a >= 0 && ++at[a] >= hi[a]) { at[a] = lo[a]; a--; } + if (a < 0) return; + } +} + +/** `bytes` gzipped, through CompressionStream (the browser's, and Node's). */ +export async function gzip(bytes: Uint8Array): Promise { + const stream = new Blob([bytes as Uint8Array]).stream().pipeThrough(new CompressionStream("gzip")); + return new Uint8Array(await new Response(stream).arrayBuffer()); +} + +/** Whether every byte of `bytes` is zero: an empty chunk (zarr's fill value, 0, for every + * dtype; a float's -0 is not empty, its sign bit set). */ +export function allZero(bytes: Uint8Array): boolean { + let i = 0; + if (bytes.byteOffset % 4 === 0) { // four bytes at a time + const words = new Uint32Array(bytes.buffer, bytes.byteOffset, bytes.byteLength >> 2); + for (const w of words) if (w) return false; + i = words.length * 4; + } + for (; i < bytes.length; i++) if (bytes[i]) return false; + return true; +} + +/** What a materialize writes of view `name`: which `level`, in which `format` and + * `compression`; the metadata `files` (path and bytes, the group's first), the chunk grid's + * indices [lo, hi) the box covers, `total` chunks, and each chunk's `key` in the store. */ +export interface Region { + name: string; level: number; format: "zarr3" | "zarr2"; compression: "gzip" | null; + files: [string, Uint8Array][]; + lo: number[]; hi: number[]; total: number; + key: (index: number[]) => string; +} + +/** What `options` asks to write of view `name` (served as `v`, chunked `chunk`): the group + * listing the one level written as its one dataset, at LEVEL_PATH, with the level's scale and + * translation; the level's array with its whole shape (so the box keeps its place in it), + * the chunks and dtype served, compressed as asked; and the chunks of the box. */ +export function region(name: string, v: ViewInfo, chunk: number[], options: MaterializeOptions = {}): Region { + const { level = 0, format = "zarr2", compression = "gzip", attributes = {} } = options; + if (!Number.isInteger(level) || !v.levels[level]) throw new Error(`${name} has no level ${level}: it has ${v.levels.length}`); + if (format !== "zarr2" && format !== "zarr3") throw new Error(`format ${JSON.stringify(format)}: "zarr2" or "zarr3"`); + if (compression !== "gzip" && compression !== null) throw new Error(`compression ${JSON.stringify(compression)}: "gzip" or null`); + const { lo, hi } = chunkRange(v.levels[level].shape, chunk, options.lo, options.hi); + const datasets = [{ path: LEVEL_PATH, level }], added = (v.added ?? []).length; + const files: [string, unknown][] = format === "zarr3" + ? [ + ["zarr.json", group3(name, v, datasets, attributes)], + [`${LEVEL_PATH}/zarr.json`, array3(v, chunk, level, { codecs: compression ? [{ name: "gzip", configuration: { level: GZIP_LEVEL } }] : [] })], + ] + : [ + [".zgroup", { zarr_format: 2 }], + [".zattrs", attrs2(name, v, datasets, attributes)], + [`${LEVEL_PATH}/.zarray`, array2(v, chunk, level, compression ? { id: "gzip", level: GZIP_LEVEL } : null)], + [`${LEVEL_PATH}/.zattrs`, { _ARRAY_DIMENSIONS: dimensionNames(v) }], + ]; + const json = (o: unknown) => new TextEncoder().encode(`${JSON.stringify(o, null, 2)}\n`); + return { + name, level, format, compression, files: files.map(([path, o]) => [path, json(o)]), + lo, hi, total: lo.reduce((n, l, a) => n * (hi[a] - l), 1), + key: (index) => chunkKey(format, LEVEL_PATH, added, index), + }; +} + +/** Chunk `index` of the level as whole bytes, as served (zarr v3's: little endian, C order, + * the channels first); `signal` aborts when it is no longer wanted. */ +export type Produce = (index: number[], signal: AbortSignal) => Promise; + +const why = (e: unknown) => String((e as Error)?.message ?? e); + +/** Write region `r` into `store`: its metadata, then its chunks, `concurrency` of them asked + * of `produce` at once (one by default), each as it comes, gzipped if `r` says so, those of + * nothing but zeros skipped (`skipEmpty`, the default). Resolves with what was written, or + * on `signal`'s abort with `aborted` (no chunk started after it, the ones asked for given up + * through `produce`'s signal and not waited for, the writes under way ended). A chunk that + * fails to be made or written rejects it, naming the chunk and why, once the writes under + * way have ended; the rest are given up as on an abort. */ +export async function writeChunks(store: WritableStore, r: Region, produce: Produce, + options: Pick = {}): Promise { + const { skipEmpty = true, concurrency = 1, signal, onProgress } = options; + if (!Number.isInteger(concurrency) || concurrency < 1) throw new Error(`concurrency ${concurrency}: a whole number of chunks, 1 or more`); + const got: Materialized = { chunks: r.total, written: 0, skipped: 0, bytes: 0, aborted: false }; + if (signal?.aborted) return { ...got, aborted: true }; + for (const [path, bytes] of r.files) { + try { await store.put(path, bytes); } catch (e) { throw new Error(`writing ${path} of ${r.name}: ${why(e)}`); } + } + // stopped by the caller's abort, or by a chunk that failed + const stop = new AbortController(), onAbort = () => stop.abort(signal?.reason); + signal?.addEventListener("abort", onAbort, { once: true }); + const stopped = new Promise((_, reject) => stop.signal.addEventListener("abort", () => reject(stop.signal.reason), { once: true })); + stopped.catch(() => {}); + let failure: Error | null = null; + const fail = (e: Error) => { failure ??= e; stop.abort(e); }; + const report = () => { + try { onProgress?.({ done: got.written + got.skipped, total: r.total, written: got.written, bytes: got.bytes }); } + catch (e) { console.error(`materialize ${r.name}: onProgress threw: ${why(e)}`); } // the page's own trouble: the writing goes on + }; + report(); + const indices = chunkIndices(r.lo, r.hi); + const lane = async () => { + for (let next = indices.next(); !next.done && !stop.signal.aborted; next = indices.next()) { + const index = next.value, key = r.key(index); + let body: ArrayBuffer; + try { + body = await Promise.race([produce(index, stop.signal), stopped]); // given up at once when stopped + } catch (e) { + if (!stop.signal.aborted) fail(new Error(`chunk ${key} of ${r.name}: ${why(e)}`)); + return; + } + if (stop.signal.aborted) return; // made as it stopped: not written + const bytes = new Uint8Array(body); + if (skipEmpty && allZero(bytes)) { got.skipped++; report(); continue; } + try { + const stored = r.compression === "gzip" ? await gzip(bytes) : bytes; + await store.put(key, stored); + got.written++; got.bytes += stored.byteLength; + } catch (e) { + fail(new Error(`writing chunk ${key} of ${r.name}: ${why(e)}`)); + return; + } + report(); + } + }; + try { + await Promise.all(Array.from({ length: Math.min(concurrency, r.total) }, lane)); + } finally { + signal?.removeEventListener("abort", onAbort); + } + if (signal?.aborted) return { ...got, aborted: true }; + if (failure) throw failure; + return got; +} diff --git a/web/test/demand.test.ts b/web/test/demand.test.ts index 6e93b1f..12658a5 100644 --- a/web/test/demand.test.ts +++ b/web/test/demand.test.ts @@ -3,7 +3,7 @@ // stripped), so it needs nothing installed. import { test } from "node:test"; import assert from "node:assert/strict"; -import { Cancelled, Claim, queue } from "../src/demand.ts"; +import { Cancelled, Claim, Queue, queue } from "../src/demand.ts"; const done = () => Promise.resolve(new Float32Array()); @@ -90,3 +90,92 @@ test("a claim cancelled already runs what is hooked to it at once", () => { claim.onCancel(() => calls++); assert.equal(calls, 2); }); + +// ------------------------------------------------ background work (Engine.materialize's) +/** A job that runs until its gate opens, noting when it starts. */ +function gated(started: string[], key: string) { + let open!: () => void; + const gate = new Promise((resolve) => (open = resolve)); + return { open, run: async () => { started.push(key); await gate; return new Float32Array(); } }; +} +const tick = () => new Promise((r) => setTimeout(r, 0)); + +test("background work waits behind every job a client waits for, whatever its level, and takes all the slots but one", async () => { + const q = new Queue(3), started: string[] = []; + assert.deepEqual([q.slots, q.backgroundSlots], [3, 2]); + const bg = [0, 1, 2].map((i) => gated(started, `background ${i}`)); + const runs = bg.map((b, i) => q.submit(`bg ${i}`, 0, b.run, () => {}, true)); + await tick(); + assert.deepEqual(started, ["background 0", "background 1"]); // the third slot is kept for clients + const client = gated(started, "client, level 3"); + const asked = q.submit("client", 3, client.run, () => {}); + await tick(); + assert.deepEqual(started.slice(2), ["client, level 3"]); // at once, in the slot kept for it + const later = gated(started, "client, later, level 5"); + const asked2 = q.submit("client later", 5, later.run, () => {}); + bg[0].open(); + await tick(); + // a slot freed: the client's job goes before the background one waiting, coarser as it is + assert.deepEqual(started.slice(3), ["client, later, level 5"]); + client.open(); + await tick(); + assert.deepEqual(started.slice(4), ["background 2"]); + for (const g of [...bg, later]) g.open(); + await Promise.all([...runs, asked, asked2]); +}); + +test("a background claim is never released by a client giving up; given up in the background, it is dropped", async () => { + const q = new Queue(1), started: string[] = []; + const busy = gated(started, "busy"); + const running = q.submit("busy", 0, busy.run, () => {}); + let dropped = 0; + const wanted = q.submit("both", 0, async () => { started.push("both"); return new Float32Array(); }, () => dropped++, true); + q.claim("both"); // a client asks for it too + q.release("both"); // and gives up + busy.open(); + await wanted; + assert.deepEqual([started, dropped], [["busy", "both"], 0]); + await running; + const again = gated(started, "busy again"); + const r2 = q.submit("busy again", 0, again.run, () => {}); + const unwanted = q.submit("background alone", 0, async () => new Float32Array(), () => dropped++, true); + q.release("background alone", true); // materialize aborted before its turn + await assert.rejects(unwanted, Cancelled); + assert.equal(dropped, 1); + again.open(); + await r2; +}); + +test("a client asking for a background job brings it forward, and it goes back when the client gives up", async () => { + const q = new Queue(3), started: string[] = []; + const bg = [0, 1, 2, 3].map((i) => gated(started, `background ${i}`)); + const runs = bg.map((b, i) => q.submit(`bg ${i}`, 0, b.run, () => {}, true)); + await tick(); + assert.deepEqual(started, ["background 0", "background 1"]); + q.claim("bg 3"); // a viewer wants the last one: it takes the slot kept for clients at once + await tick(); + assert.deepEqual(started.slice(2), ["background 3"]); + // a fourth waiting in the background, a client's claim on it given up: back behind the rest + const fourth = gated(started, "background 4"); + const r4 = q.submit("bg 4", 0, fourth.run, () => {}, true); + q.claim("bg 4"); + q.release("bg 4"); + const client = gated(started, "client"); + const asked = q.submit("client", 1, client.run, () => {}); + bg[0].open(); + await tick(); + assert.deepEqual(started.slice(3), ["client"]); + for (const g of [...bg, fourth, client]) g.open(); + await Promise.all([...runs, r4, asked]); +}); + +test("with one slot, background work runs when no client waits", async () => { + const q = new Queue(1), started: string[] = []; + assert.equal(q.backgroundSlots, 1); + const bg = gated(started, "background"); + const run = q.submit("bg", 0, bg.run, () => {}, true); + await tick(); + assert.deepEqual(started, ["background"]); + bg.open(); + await run; +}); diff --git a/web/test/materialize.test.ts b/web/test/materialize.test.ts new file mode 100644 index 0000000..08ff27f --- /dev/null +++ b/web/test/materialize.test.ts @@ -0,0 +1,313 @@ +// A region of a view written into a store (src/materialize.ts, Engine.materialize's parts +// that need no engine): the chunks a box covers, rounded out and clipped; the metadata both +// formats write, the layout the engine serves (src/layout.ts); gzip; the loop over a fake +// producer, with empty chunks skipped, progress, abort and failure; the stores; and the +// result read back by zarrita, a zarr implementation of its own, as a reader would. Run by +// `npm test`. +import { test } from "node:test"; +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, writeChunks, type MaterializeProgress, type Produce, type WritableStore } from "../src/materialize.ts"; +import type { ViewInfo } from "../src/types.ts"; + +const nm = (name: string) => ({ name, unit: "nm" }); +/** A view of two levels, served as uint16 in chunks of 2 × 4 × 4. */ +const view: ViewInfo = { + dtype: "uint8", channels: 1, axes: [nm("z"), nm("y"), nm("x")], halo: [0, 0, 0], lead: 0, out: "uint16", added: [], + levels: [{ shape: [5, 6, 7], voxel: [8, 8, 8], origin: [0, 0, 0] }, { shape: [3, 3, 4], voxel: [16, 16, 16], origin: [4, 4, 4] }], +}; +const chunk = [2, 4, 4]; +/** The same view with two channels its ops add (a gradient's, say), as float32. */ +const channels: ViewInfo = { ...view, out: "float32", added: [2] }; + +const json = (b: Uint8Array | undefined) => JSON.parse(new TextDecoder().decode(b)); +const files = (s: { files: Map }) => [...s.files.keys()].sort(); +async function gunzip(b: Uint8Array): Promise { + const stream = new Blob([b as Uint8Array]).stream().pipeThrough(new DecompressionStream("gzip")); + return new Uint8Array(await new Response(stream).arrayBuffer()); +} +const tick = () => new Promise((r) => setTimeout(r, 0)); + +/** Level 0's voxel (z, y, x) as the fake producer makes it: never zero inside the level. */ +const value = (z: number, y: number, x: number) => 1 + 100 * z + 10 * y + x; +/** Chunk `index` of level 0 as the engine gives it: whole, zero past the level's edge; the + * chunks of `empty` all zero, as a model's output where there is nothing. */ +function fakeChunk(index: number[], empty: string[] = []): ArrayBuffer { + const out = new Uint16Array(chunk[0] * chunk[1] * chunk[2]), [Z, Y, X] = view.levels[0].shape; + if (empty.includes(index.join("."))) return out.buffer; + for (let z = 0; z < chunk[0]; z++) for (let y = 0; y < chunk[1]; y++) for (let x = 0; x < chunk[2]; x++) { + const g = [index[0] * chunk[0] + z, index[1] * chunk[1] + y, index[2] * chunk[2] + x]; + if (g[0] < Z && g[1] < Y && g[2] < X) out[(z * chunk[1] + y) * chunk[2] + x] = value(g[0], g[1], g[2]); + } + return out.buffer; +} +const producer = (empty: string[] = []): Produce => async (index) => fakeChunk(index, empty); + +/** `store`'s files as zarrita reads a store (keys with a leading "/"). */ +const readable = (s: { files: Map }) => ({ async get(key: string) { return s.files.get(key.replace(/^\//, "")); } }); + +test("a box covers the chunks it touches, rounded out and clipped to the level; the whole level by default", () => { + const shape = view.levels[0].shape; // 5 × 6 × 7 in chunks of 2 × 4 × 4: a grid of 3 × 2 × 2 + assert.deepEqual(chunkRange(shape, chunk), { lo: [0, 0, 0], hi: [3, 2, 2] }); + assert.deepEqual(chunkRange(shape, chunk, [1, 3, 5], [3, 5, 6]), { lo: [0, 0, 1], hi: [2, 2, 2] }); + assert.deepEqual(chunkRange(shape, chunk, [2, 4, 4], [4, 8, 8]), { lo: [1, 1, 1], hi: [2, 2, 2] }); // on chunk edges: no more + assert.deepEqual(chunkRange(shape, chunk, [-10, 2, 3], [100, 6, 100]), { lo: [0, 0, 0], hi: [3, 2, 2] }); // past the edges: clipped + assert.deepEqual(chunkRange(shape, chunk, [4, 0, 0]), { lo: [2, 0, 0], hi: [3, 2, 2] }); // lo alone: to the level's end + assert.deepEqual(chunkRange(shape, chunk, [0.5, 0, 0], [1.5, 1, 1]), { lo: [0, 0, 0], hi: [1, 1, 1] }); // fractions: rounded out too + assert.throws(() => chunkRange(shape, chunk, [5, 0, 0]), /holds no voxel of the level, which is 5×6×7/); // starts at the end + assert.throws(() => chunkRange(shape, chunk, [0, 0, 0], [0, 6, 7]), /holds no voxel/); + assert.throws(() => chunkRange(shape, chunk, [0, 0], [1, 1]), /lo \[0,0\]: one number for each of the level's 3 axes/); + assert.throws(() => chunkRange(shape, chunk, undefined, [1, Number.NaN, 1]), /hi /); +}); + +test("the chunks of a range in C order, the last axis fastest", () => { + assert.deepEqual([...chunkIndices([0, 1], [2, 3])], [[0, 1], [0, 2], [1, 1], [1, 2]]); + assert.deepEqual([...chunkIndices([1, 1, 1], [2, 2, 2])], [[1, 1, 1]]); + assert.deepEqual([...chunkIndices([0, 0], [0, 3])], []); +}); + +test("zarr v2: .zgroup, .zattrs with OME 0.4 multiscales listing the one level at 0, its .zarray gzipped and .zattrs", () => { + const r = region("mito", view, chunk, { level: 1, attributes: { provenance: { model: "mito-v3" } } }); + assert.deepEqual(r.files.map(([p]) => p), [".zgroup", ".zattrs", "0/.zarray", "0/.zattrs"]); + const [group, attrs, zarray, dims] = r.files.map(([, b]) => json(b)); + assert.deepEqual(group, { zarr_format: 2 }); + assert.deepEqual(attrs, { + provenance: { model: "mito-v3" }, + multiscales: [{ + name: "mito", + axes: [{ name: "z", type: "space", unit: "nanometer" }, { name: "y", type: "space", unit: "nanometer" }, { name: "x", type: "space", unit: "nanometer" }], + // level 1's own scale and translation, so it lands where the view's level 1 does + datasets: [{ path: "0", coordinateTransformations: [{ type: "scale", scale: [16, 16, 16] }, { type: "translation", translation: [4, 4, 4] }] }], + version: "0.4", + }], + }); + assert.deepEqual(zarray, { + zarr_format: 2, shape: [3, 3, 4], chunks: [2, 4, 4], dtype: " { + const r = region("mito", view, chunk, { format: "zarr3", lo: [2, 0, 4], hi: [5, 3, 7], attributes: { provenance: 1 } }); + assert.deepEqual(r.files.map(([p]) => p), ["zarr.json", "0/zarr.json"]); + const [group, array] = r.files.map(([, b]) => json(b)); + assert.deepEqual(group, { + zarr_format: 3, node_type: "group", attributes: { + provenance: 1, + ome: { version: "0.5", multiscales: [{ + name: "mito", + axes: [{ name: "z", type: "space", unit: "nanometer" }, { name: "y", type: "space", unit: "nanometer" }, { name: "x", type: "space", unit: "nanometer" }], + datasets: [{ path: "0", coordinateTransformations: [{ type: "scale", scale: [8, 8, 8] }, { type: "translation", translation: [0, 0, 0] }] }], + }] }, + }, + }); + assert.deepEqual(array, { + zarr_format: 3, node_type: "array", shape: [5, 6, 7], data_type: "uint16", fill_value: 0, // the whole level: the box keeps its place + chunk_grid: { name: "regular", configuration: { chunk_shape: [2, 4, 4] } }, + chunk_key_encoding: { name: "default", configuration: { separator: "/" } }, + codecs: [{ name: "bytes", configuration: { endian: "little" } }, { name: "gzip", configuration: { level: 6 } }], + dimension_names: ["z", "y", "x"], attributes: {}, + }); + assert.deepEqual([r.lo, r.hi, r.total, r.key([2, 0, 1])], [[1, 0, 1], [3, 1, 2], 2, "0/c/2/0/1"]); + // the metadata is the engine's as served, but for the one dataset and the codec + assert.deepEqual(array.codecs.slice(0, 1), array3(view, chunk, 0).codecs); + assert.deepEqual(group3("mito", view, [{ path: "0", level: 0 }, { path: "1", level: 1 }]).attributes.ome.multiscales[0].datasets[1].path, "1"); +}); + +test("uncompressed when asked; a channel axis the ops add is one chunk, index 0, as the engine serves it", () => { + const two = region("grad", channels, chunk, { compression: null }); + const zarray = json(two.files.find(([p]) => p === "0/.zarray")![1]); + assert.deepEqual([zarray.shape, zarray.chunks, zarray.dtype, zarray.compressor], [[2, 5, 6, 7], [2, 2, 4, 4], " p === "0/.zattrs")![1]), { _ARRAY_DIMENSIONS: ["c", "z", "y", "x"] }); + const ms = json(two.files.find(([p]) => p === ".zattrs")![1]).multiscales[0]; + assert.deepEqual(ms.axes[0], { name: "c", type: "channel" }); + assert.deepEqual(ms.datasets[0].coordinateTransformations, [{ type: "scale", scale: [1, 8, 8, 8] }, { type: "translation", translation: [0, 0, 0, 0] }]); + assert.equal(two.key([1, 0, 1]), "0/0/1/0/1"); + const three = region("grad", channels, chunk, { format: "zarr3", compression: null }); + const array = json(three.files[1][1]); + assert.deepEqual([array.shape, array.chunk_grid.configuration.chunk_shape, array.codecs.length], [[2, 5, 6, 7], [2, 2, 4, 4], 1]); + assert.equal(three.key([1, 0, 1]), "0/c/0/1/0/1"); + assert.equal(chunkKey("zarr3", "0", 0, [3, 2, 1]), "0/c/3/2/1"); + assert.throws(() => region("grad", channels, chunk, { level: 2 }), /grad has no level 2: it has 2/); + assert.throws(() => region("grad", channels, chunk, { format: "n5" as "zarr2" }), /format "n5"/); +}); + +test("the group's attributes keep the view's OME metadata whatever the extra ones are called", () => { + const two = region("v", view, chunk, { attributes: { multiscales: "mine", note: "kept" } }); + const attrs = json(two.files[1][1]); + assert.equal(attrs.note, "kept"); + assert.equal(attrs.multiscales[0].name, "v"); + const three = region("v", view, chunk, { format: "zarr3", attributes: { ome: "mine" } }); + assert.equal(json(three.files[0][1]).attributes.ome.version, "0.5"); +}); + +test("gzip round trips through CompressionStream and DecompressionStream; zeros are told from the rest", async () => { + const chunk0 = new Uint8Array(fakeChunk([0, 0, 0])); + assert.deepEqual(await gunzip(await gzip(chunk0)), chunk0); + const bytes = new Uint8Array(4096).fill(7, 100, 3000), packed = await gzip(bytes); // a model's output: much of it alike + assert.deepEqual([packed[0], packed[1]], [0x1f, 0x8b]); // gzip's magic + assert.ok(packed.length < bytes.length / 10, `${packed.length} bytes`); + assert.deepEqual(await gunzip(packed), bytes); + assert.equal(allZero(new Uint8Array(64)), true); + assert.equal(allZero(new Uint8Array(new ArrayBuffer(67), 1, 65)), true); // not on a word: byte by byte + const one = new Uint8Array(67); + one[66] = 1; + assert.equal(allZero(one), false); // past the last whole word + assert.equal(allZero(new Uint8Array(new Float32Array([0, -0]).buffer)), false); // -0 has its sign bit +}); + +for (const format of ["zarr2", "zarr3"] as const) { + test(`${format}: the box's chunks written gzipped, the empty one skipped, and read back by zarrita as the view`, async () => { + const store = memoryStore(), seen: MaterializeProgress[] = []; + // z voxels 1 to 3 of the whole y, x: chunks z 0 and 1, the second of them empty + const got = await writeChunks(store, region("v", view, chunk, { format, lo: [1, 0, 0], hi: [3, 6, 7] }), producer(["1.0.1"]), + { concurrency: 3, onProgress: (p) => seen.push(p) }); + const keys = files(store).filter((f) => !f.endsWith("json") && !f.includes(".z")); + const at = (i: string) => (format === "zarr3" ? `0/c/${i}` : `0/${i}`); + assert.deepEqual(keys, ["0.0.0", "0.0.1", "0.1.0", "0.1.1", "1.0.0", "1.1.0", "1.1.1"].map((i) => at(i.replaceAll(".", "/"))).sort()); + const bytes = keys.reduce((n, k) => n + store.files.get(k)!.length, 0); + assert.deepEqual(got, { chunks: 8, written: 7, skipped: 1, bytes, aborted: false }); + assert.deepEqual(await gunzip(store.files.get(at("1/1/1"))!), new Uint8Array(fakeChunk([1, 1, 1]))); + // progress: once with the metadata written, then after each chunk, done = written + skipped + assert.deepEqual(seen[0], { done: 0, total: 8, written: 0, bytes: 0 }); + assert.deepEqual(seen.at(-1), { done: 8, total: 8, written: 7, bytes }); + const steps = seen.slice(1).map((p, i) => [p.done - seen[i].done, p.written - seen[i].written, p.bytes - seen[i].bytes]); + assert.deepEqual(steps.map(([d]) => d), [1, 1, 1, 1, 1, 1, 1, 1]); + assert.deepEqual(steps.filter(([, w, b]) => w === 0 && b === 0).length, 1); // the one skipped + // a reader: the box's chunks as made, the skipped one and those outside the box zero + const root = zarr.root(readable(store)); + const group = format === "zarr3" ? await zarr.open.v3(root, { kind: "group" }) : await zarr.open.v2(root, { kind: "group" }); + assert.ok(group.attrs.ome ?? group.attrs.multiscales); + const arr = format === "zarr3" ? await zarr.open.v3(root.resolve("0"), { kind: "array" }) : await zarr.open.v2(root.resolve("0"), { kind: "array" }); + assert.deepEqual([arr.shape, arr.dtype], [[5, 6, 7], "uint16"]); + const data = (await zarr.get(arr)).data as Uint16Array; + const [, Y, X] = arr.shape; + for (let z = 0; z < 5; z++) for (let y = 0; y < Y; y++) for (let x = 0; x < X; x++) { + const written = z < 4 && !(z >= 2 && y < 4 && x >= 4); // chunks z 0 and 1, but for chunk (1, 0, 1) + assert.equal(data[(z * Y + y) * X + x], written ? value(z, y, x) : 0, `voxel ${z}, ${y}, ${x}`); + } + }); +} + +test("empty chunks are written too when asked, and uncompressed ones as they are", async () => { + const store = memoryStore(); + const got = await writeChunks(store, region("v", view, chunk, { compression: null, lo: [2, 0, 4], hi: [4, 4, 8] }), producer(["1.0.1"]), { skipEmpty: false }); + assert.deepEqual(got, { chunks: 1, written: 1, skipped: 0, bytes: 64, aborted: false }); + assert.deepEqual(store.files.get("0/1/0/1"), new Uint8Array(32 * 2)); +}); + +test("no more than `concurrency` chunks are asked for at once", async () => { + let asked = 0, most = 0; + const produce: Produce = async (index) => { + most = Math.max(most, ++asked); + await tick(); + asked--; + return fakeChunk(index); + }; + const got = await writeChunks(memoryStore(), region("v", view, chunk), produce, { concurrency: 2 }); + assert.deepEqual([got.written, most], [12, 2]); + await assert.rejects(writeChunks(memoryStore(), region("v", view, chunk), produce, { concurrency: 0 }), /concurrency 0/); +}); + +test("an abort stops starting chunks, gives up those asked for and resolves at once; the chunks written stay", { timeout: 5000 }, async () => { + const store = memoryStore(), stop = new AbortController(), asked: number[][] = [], given: AbortSignal[] = []; + let atAbort = -1; + // the first two chunks come at once, the rest never (a model still running, say) + const produce: Produce = (index, signal) => { + asked.push(index); + given.push(signal); + return asked.length <= 2 ? Promise.resolve(fakeChunk(index)) : new Promise(() => {}); + }; + const got = await writeChunks(store, region("v", view, chunk), produce, { + concurrency: 2, signal: stop.signal, onProgress: (p) => { if (p.written === 2) { stop.abort(); atAbort = asked.length; } }, + }); + assert.deepEqual(got, { chunks: 12, written: 2, skipped: 0, bytes: got.bytes, aborted: true }); + assert.equal(asked.length, atAbort); // none asked for after the abort + assert.ok(atAbort > 2); // and the one asked for when it came, never made, was not waited for + assert.ok(given.every((s) => s.aborted)); + assert.deepEqual(files(store), [".zattrs", ".zgroup", "0/.zarray", "0/.zattrs", "0/0/0/0", "0/0/0/1"]); + // aborted before it starts: nothing is written + const none = memoryStore(); + assert.deepEqual(await writeChunks(none, region("v", view, chunk), producer(), { signal: AbortSignal.abort() }), + { chunks: 12, written: 0, skipped: 0, bytes: 0, aborted: true }); + assert.equal(none.files.size, 0); +}); + +test("a chunk that fails rejects naming it and why, once the writes under way have ended", async () => { + let slowWritten = false; + const memory = memoryStore(); + const store: WritableStore = { + async put(path, bytes) { + if (path === "0/0/0/0") { await new Promise((r) => setTimeout(r, 30)); slowWritten = true; } + return memory.put(path, bytes); + }, + }; + const produce: Produce = async (index) => { + if (index.join(".") === "0.0.1") { await tick(); throw new Error("the model ran out of GPU memory"); } + return fakeChunk(index); + }; + await assert.rejects(writeChunks(store, region("v", view, chunk), produce, { concurrency: 2 }), (e: Error) => { + assert.equal(e.message, "chunk 0/0/0/1 of v: the model ran out of GPU memory"); + assert.ok(slowWritten, "the chunk being written when the other failed was written first"); + return true; + }); + assert.ok(memory.files.has("0/0/0/0") && !memory.files.has("0/0/1/0")); // nothing started after the failure + const full: WritableStore = { async put(path) { if (path.startsWith("0/c/1/")) throw new Error("the disk is full"); } }; + await assert.rejects(writeChunks(full, region("v", view, chunk, { format: "zarr3" }), producer()), /^Error: writing chunk 0\/c\/1\/0\/0 of v: the disk is full$/); + const refused: WritableStore = { async put() { throw new Error("permission revoked"); } }; + await assert.rejects(writeChunks(refused, region("v", view, chunk), producer()), /^Error: writing \.zgroup of v: permission revoked$/); +}); + +test("the memory store keeps a copy of each file under its path, without empty names", async () => { + const store = memoryStore(), bytes = new Uint8Array([1, 2, 3]); + await store.put("/a//b/c", bytes); + bytes[0] = 9; + assert.deepEqual([...store.files], [["a/b/c", new Uint8Array([1, 2, 3])]]); + await assert.rejects(store.put("a/../b", bytes), /with no "\." or "\.\."/); + await assert.rejects(store.put("/", bytes), /not a file's path/); +}); + +/** A FileSystemDirectoryHandle's few methods directoryStore calls, over a tree of Maps. */ +interface Dir { dirs: Map; files: Map } +function fakeDirectory(calls: string[], failWrite?: string) { + const make = (path: string, d: Dir): FileSystemDirectoryHandle => ({ + async getDirectoryHandle(name: string, { create = false } = {}) { + calls.push(`dir ${path}${name}`); + await tick(); // as a real handle is: concurrent puts must not make one directory twice + if (!d.dirs.has(name)) { assert.ok(create); d.dirs.set(name, { dirs: new Map(), files: new Map() }); } + return make(`${path}${name}/`, d.dirs.get(name)!); + }, + async getFileHandle(name: string) { + return { + async createWritable() { + const parts: Uint8Array[] = []; + return { + async write(b: Uint8Array) { if (name === failWrite) throw new Error("quota exceeded"); parts.push(b.slice()); }, + async close() { d.files.set(name, new Uint8Array(parts.flatMap((p) => [...p]))); calls.push(`closed ${path}${name}`); }, + async abort() { calls.push(`aborted ${path}${name}`); }, + }; + }, + }; + }, + }) as unknown as FileSystemDirectoryHandle; + const tree: Dir = { dirs: new Map(), files: new Map() }; + return { tree, handle: make("", tree) }; +} + +test("a directory as a store: the root and subdirectories made once each, each file written whole", async () => { + const calls: string[] = [], { tree, handle } = fakeDirectory(calls); + const store = directoryStore(handle, "runs/mito.zarr"); + await Promise.all([store.put("0/c/0/0/0", new Uint8Array([1])), store.put("0/c/0/0/1", new Uint8Array([2])), store.put("zarr.json", new Uint8Array([3]))]); + assert.deepEqual(calls.filter((c) => c.startsWith("dir")), ["dir runs", "dir runs/mito.zarr", "dir runs/mito.zarr/0", "dir runs/mito.zarr/0/c", "dir runs/mito.zarr/0/c/0", "dir runs/mito.zarr/0/c/0/0"]); + const zarrDir = tree.dirs.get("runs")!.dirs.get("mito.zarr")!; + assert.deepEqual(zarrDir.files.get("zarr.json"), new Uint8Array([3])); + assert.deepEqual([...zarrDir.dirs.get("0")!.dirs.get("c")!.dirs.get("0")!.dirs.get("0")!.files], [["0", new Uint8Array([1])], ["1", new Uint8Array([2])]]); + // a write that fails is aborted (the swap file dropped) and rejects + const failing = fakeDirectory(calls, "bad"); + await assert.rejects(directoryStore(failing.handle).put("bad", new Uint8Array([1])), /quota exceeded/); + assert.ok(calls.includes("aborted bad")); + assert.equal(failing.tree.files.size, 0); +}); From ac41da8e8020e4d89a7e0fb0065ff8701fd7e7b0 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 16 Oct 2026 13:21:08 -0400 Subject: [PATCH 2/2] A chunk only Engine.materialize asked for is not kept once it is written: the kept chunks (256 MB) are for clients that fetch a chunk again, and a region written chunk by chunk pushed what the viewer was looking at out of them for chunks no client might ask for; the engine notes the keys of chunks on their way that a client asked for (chunk() without background) and keeps only those, so a chunk a viewer asked for while materialize had it on its way is kept as before; caching.md says so (the staged stage outputs, 512 MB, are kept as before), as the chunkmirage session suggested in its review of the materialize proposal --- docs/concepts/caching.md | 11 +++++++---- web/src/engine.ts | 15 +++++++++++++-- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/docs/concepts/caching.md b/docs/concepts/caching.md index e7b2298..d44158c 100644 --- a/docs/concepts/caching.md +++ b/docs/concepts/caching.md @@ -212,8 +212,8 @@ The browser page has one class of work more: chunks it computes for itself, in t background, to write a region of a view into a store (`Engine.materialize`; see [formats](formats.md#a-region-written-to-a-store-browser-engine)). They go through the same queue as the viewer's chunks, so a chunk the viewer has computed is taken from the kept -chunks, one on its way for the viewer is shared, and one written is kept for the viewer -after. Three rules keep browsing usable while a region is written: +chunks and one on its way for the viewer is shared. Three rules keep browsing usable while +a region is written: * A background claim is never released because a client gave up, so nothing materialize asked for is dropped as abandoned; only its own abort gives its waiting chunks up (and @@ -228,8 +228,11 @@ A viewer asking for a chunk that waits in the background brings it forward, and back if the viewer gives up on it. Materialize asks for a few chunks at once (as many as the queue has slots, by default: the next one is waiting while one is written), not the whole region, so the queue holds a few of its chunks at a time however large the region. -Its chunks are kept as the viewer's are (256 MB, least recently used out), as are a staged -view's stage outputs, so writing a large region cycles through both. +A chunk only materialize asked for is not kept once written (the kept chunks are for +clients that fetch a chunk again), so a large region does not push what the viewer is +looking at out of the 256 MB; one a client asked for too, while it was on its way, is kept +as any other. 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/web/src/engine.ts b/web/src/engine.ts index af3517a..d6add92 100644 --- a/web/src/engine.ts +++ b/web/src/engine.ts @@ -106,6 +106,9 @@ export class Engine { private chunks: Queue | null = null; private inflight = new Map>(); private kept = new Kept(KEPT_BYTES); + /** The keys of chunks on their way that a client asked for (`chunk` without `background`): + * only those are kept once computed. */ + private clientWanted = new Set(); private staged = new Kept(STAGED_BYTES); // by stage key, not view: shared by views, kept past an unserve private edits = 0; private producers = new Map Promise>(); @@ -470,15 +473,23 @@ export class Engine { const key = `${view}/${mesh ? `mesh${level}` : level}/${index.join(".")}`, done = this.kept.get(key); if (done) return { status: 200, body: done.slice(0), type: "application/octet-stream" }; const chunks = this.chunks!, claim = new Claim(); + if (!background) this.clientWanted.add(key); let pending = this.inflight.get(key); if (pending) chunks.claim(key, background); // another request waits for it too else { pending = chunks.submit(key, level, () => this.compute(view, level, index, mesh), () => { this.counts.dropped++; }, background) - .then((b) => { this.counts.computed++; if (this.views[view]) this.kept.put(key, b); return b; }, (e) => { // kept unless unserved meanwhile + .then((b) => { + // Kept for a client that fetches it again, unless the view was unserved meanwhile or + // only materialize asked for it: a region written chunk by chunk would otherwise push + // what the viewer is looking at out of the 256 MB, for chunks no client may ask for. + this.counts.computed++; + if (this.views[view] && this.clientWanted.has(key)) this.kept.put(key, b); + return b; + }, (e) => { if (!(e instanceof Cancelled)) { this.counts.failed++; this.failure = { view, message: String((e as Error).message ?? e) }; console.error(`chunk ${key}: ${String((e as Error).message ?? e).trim().split("\n").slice(-3).join(" / ")}`); } throw e; }) - .finally(() => { this.inflight.delete(key); this.onChange(); }); + .finally(() => { this.inflight.delete(key); this.clientWanted.delete(key); this.onChange(); }); this.inflight.set(key, pending); } claim.onCancel(() => chunks.release(key, background));