Workerize zarr chunk decoding (Data Manager 8/9)

public/index.ts

@@ -30,6 +30,7 @@ import DataManager, { IDataSubscriber } from "../src/data_manager/data_manager.j
 import { OMEZarrSource } from "../src/data_manager/sources/index.js";
 import { ChunkId, ChunkPriority, ChunkPriorityLevel } from "../src/data_manager/types.js";
 import { ThreeInterface } from "../src/data_manager/device_interface.js";
+import { TaskPool } from "../src/data_manager/task_pool/task_pool.js";
 
 const CACHE_MAX_SIZE = 1_000_000_000;
 const CONCURRENCY_LIMIT = 8;
@@ -925,10 +926,12 @@ function shuffle<T>(arr: T[]): T[] {
 }
 
 async function testDataManager() {
+  const pool = new TaskPool(new URL("./test_worker.js", import.meta.url));
   const dataManager = new DataManager(new ThreeInterface());
   const zarrSource = await OMEZarrSource.new(
     // "https://animatedcell-test-data.s3.us-west-2.amazonaws.com/variance/1.zarr"
-    "https://uk1s3.embassy.ebi.ac.uk/idr/zarr/v0.4/idr0048A/9846152.zarr/"
+    "https://vast-files.int.allencell.org/users/alex.khang/repos/bioio_conversion/H2B_640_CAAX_488_20250904_flat_colony_1.ome.zarr/",
+    { pool }
   );
   const sourceId = dataManager.addSource(zarrSource);
   const mockSubscriber: IDataSubscriber<Data3DTexture> = {

public/test_worker.ts

@@ -0,0 +1,5 @@
+import { copyChunkTask, reorderChunkTask } from "../src/data_manager/indexing_task.js";
+import { decodeTask, encodeTask } from "../src/data_manager/sources/ome_zarr/zarr_codec_worker.js";
+import { initWorker } from "../src/data_manager/task_pool/runner.js";
+
+initWorker([encodeTask, decodeTask, copyChunkTask, reorderChunkTask]);

src/data_manager/data_manager.ts

@@ -13,34 +13,23 @@ import {
   deviceSizeLimitForPriority,
 } from "./types.js";
 import type { NumberType, TypedArray } from "../types.js";
+import { defineGlobalId, swapRemove } from "./utils.js";
 import PriorityQueue from "./priority_queue.js";
 import type { DeviceInterface } from "./device_interface.js";
 import { dataTypeToSize } from "./utils.js";
 import { ChunkSource, type VolumeMetadata, type ExtVolumeDims } from "./sources/chunk_source.js";
 import SlotMap from "./slot_map.js";
 import { BorrowGuard } from "./task_pool/borrow_guard.js";
 
-const SUBSCRIBER_ID = Symbol("DataManager.subscriberId");
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const getIdForSubscriber = defineGlobalId<IDataSubscriber<any>>(Symbol("DataManager.subscriberId"));
 
 export interface IDataSubscriber<Tex> {
-  [SUBSCRIBER_ID]?: number;
   onChunkLoaded?: (id: ChunkId, chunk: Chunk<NumberType>) => void;
   onChunkOnGpu?: (id: ChunkId, texture: Tex) => void;
   // TODO events for when chunks are evicted?
 }
 
-// Subscriber IDs increment globally, like timeout IDs
-let subscriberCount = 0;
-
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-const getIdForSubscriber = (subscriber: IDataSubscriber<any>): number => {
-  if (subscriber[SUBSCRIBER_ID] === undefined) {
-    subscriber[SUBSCRIBER_ID] = subscriberCount;
-    subscriberCount += 1;
-  }
-  return subscriber[SUBSCRIBER_ID];
-};
-
 type SourceEntry<Tex> = {
   source: ChunkSource;
   subscribers: IDataSubscriber<Tex>[];
@@ -51,18 +40,6 @@ type RequestEntry = {
   controller: AbortController;
 };
 
-const swapRemove = <T>(arr: T[], index: number) => {
-  if (index < 0) {
-    return;
-  }
-
-  const { length } = arr;
-  const replace = arr.pop();
-  if (replace !== undefined && index < length) {
-    arr[index] = replace;
-  }
-};
-
 type ChunkQueue = PriorityQueue<ChunkPriority, string>;
 
 class ChunkQueues {

src/data_manager/sources/ome_zarr/with_threaded_codecs.ts

@@ -0,0 +1,98 @@
+import * as zarr from "zarrita";
+
+import { taskHandle, TaskPool } from "../../task_pool/index.js";
+import { defineGlobalId } from "../../utils.js";
+import type { encodeTask as encodeTaskType, decodeTask as decodeTaskType } from "./zarr_codec_worker.js";
+import type { NumberType } from "../../../types.js";
+
+type CodecKind = "array_to_array" | "array_to_bytes" | "bytes_to_bytes";
+
+// MARK: Codec wrapper
+
+const encodeTask = taskHandle<typeof encodeTaskType>("zarrEncode", (chunk) => [chunk.data.buffer]);
+const decodeTask = taskHandle<typeof decodeTaskType>("zarrDecode", (data) => [data.buffer]);
+
+type WorkerConfig = {
+  kind: CodecKind;
+  codecs: zarr.CodecMetadata[];
+};
+
+const getPoolId = defineGlobalId<TaskPool>(Symbol("OMEZarrSource.codecPoolId"));
+
+/**
+ * Returns the name of the `zarrita` wrapper codec that runs codec pipelines in the given `pool`.
+ *
+ * If no such codec exists, this function creates one and registers it with `zarrita`.
+ */
+const getPoolCodecName = (pool: TaskPool): string => {
+  const id = getPoolId(pool);
+  const name = `vole-worker-${id}`;
+
+  if (zarr.registry.get(name) === undefined) {
+    // Assign every codec pipeline created in this pool a serial number for easy lookup worker-side
+    let codecCounter = 0;
+
+    zarr.registry.set(name, async () => ({
+      fromConfig: (config, meta) => {
+        const { codecs, kind } = config as WorkerConfig;
+        const descriptor = { codecs, meta, id: codecCounter };
+        codecCounter++;
+
+        return {
+          encode: (data: zarr.Chunk<NumberType>) => pool.runTask(encodeTask, data, descriptor),
+          decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor),
+          kind,
+        };
+      },
+    }));
+  }
+
+  return name;
+};
+
+// MARK: Metadata parsing
+
+const wrapCodecMetadata = (name: string, meta: zarr.CodecMetadata[]): zarr.CodecMetadata[] => {
+  const shardIndex = meta.findIndex(({ name }) => name === "sharding_indexed");
+  if (shardIndex > -1) {
+    const codecs = meta[shardIndex].configuration.codecs as zarr.CodecMetadata[];
+    meta[shardIndex].configuration.codecs = wrapCodecMetadata(name, codecs);
+    return meta;
+  }
+  return [{ name, configuration: { kind: "array_to_bytes", codecs: meta } satisfies WorkerConfig }];
+};
+
+/** zarrita `Store` extension that intercepts metadata to replace as many codecs as possible with workerized variants. */
+export const withWorkerCodec = zarr.defineStoreExtension((store, opts: { pool?: TaskPool }) => {
+  if (opts?.pool === undefined) {
+    return {};
+  }
+  const { pool } = opts;
+  const name = getPoolCodecName(pool);
+
+  return {
+    get: async (key: zarr.AbsolutePath, options?: zarr.GetOptions): Promise<Uint8Array | undefined> => {
+      const bytes = await store.get(key, options);
+
+      if (!key.endsWith("zarr.json")) {
+        return bytes;
+      }
+
+      const str = new TextDecoder().decode(bytes);
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      let meta: any;
+      try {
+        meta = JSON.parse(str);
+      } catch (cause) {
+        throw new zarr.InvalidMetadataError("Failed to decode JSON", { cause });
+      }
+
+      if (meta.zarr_format === 3 && Array.isArray(meta.codecs)) {
+        meta.codecs = wrapCodecMetadata(name, meta.codecs);
+      }
+
+      const result = JSON.stringify(meta);
+      return new TextEncoder().encode(result);
+    },
+  };
+});

src/data_manager/sources/ome_zarr/zarr_codec_worker.ts

@@ -0,0 +1,97 @@
+import { type Chunk, CodecPipelineError, registry, UnknownCodecError, type CodecMetadata } from "zarrita";
+import { task } from "../../task_pool/index.js";
+import type { NumberType } from "../../../types.js";
+
+type Codec<In, Out> = {
+  encode: (data: In) => Out | Promise<Out>;
+  decode: (data: Out) => In | Promise<In>;
+};
+type ArrayToArrayCodec = Codec<Chunk<NumberType>, Chunk<NumberType>>;
+type ArrayToBytesCodec = Codec<Chunk<NumberType>, Uint8Array>;
+type BytesToBytesCodec = Codec<Uint8Array, Uint8Array>;
+type AnyCodec = ArrayToArrayCodec | ArrayToBytesCodec | BytesToBytesCodec;
+
+// hack the private `ChunkMetadata` type out of zarrita
+type CodecEntry = typeof registry extends Map<string, () => Promise<infer V>> ? V : never;
+type ChunkMetadata = Parameters<CodecEntry["fromConfig"]>[1];
+
+type CodecDescriptor = {
+  id: number;
+  codecs: CodecMetadata[];
+  meta: ChunkMetadata;
+};
+
+const codecs = new Map<number, ArrayToBytesCodec>();
+
+const joinCodecs = (codecs: { codec: AnyCodec; name: string }[]): ArrayToBytesCodec => {
+  // Simplified version of zarrita's `createCodecPipeline`
+  const runStep = async <T>(direction: "encode" | "decode", codec: string, fn: () => Promise<T> | T): Promise<T> => {
+    try {
+      return await fn();
+    } catch (cause) {
+      throw new CodecPipelineError({ direction, codec, cause });
+    }
+  };
+  // Put codecs in bytes -> array order to prioritize decoding over encoding
+  codecs.reverse();
+
+  return {
+    decode: async (data: Uint8Array): Promise<Chunk<NumberType>> => {
+      let result: Uint8Array | Chunk<NumberType> = data;
+      for (const { name, codec } of codecs) {
+        // evil cast
+        result = await runStep("decode", name, () => codec.decode(result as Chunk<NumberType> & Uint8Array));
+      }
+      return result as Chunk<NumberType>;
+    },
+    encode: async (data: Chunk<NumberType>): Promise<Uint8Array> => {
+      let result: Uint8Array | Chunk<NumberType> = data;
+      for (const { name, codec } of codecs.toReversed()) {
+        // evil cast pt. 2
+        result = await runStep("encode", name, () => codec.encode(result as Chunk<NumberType> & Uint8Array));
+      }
+      return result as Uint8Array;
+    },
+  };
+};
+
+const loadCodec = async (descriptor: CodecDescriptor): Promise<ArrayToBytesCodec> => {
+  const codecPromises = descriptor.codecs.map(async ({ name, configuration }) => {
+    const entry = await registry.get(name)?.();
+    if (entry === undefined) {
+      throw new UnknownCodecError(name);
+    }
+    const codec = entry.fromConfig(configuration, descriptor.meta);
+    return { codec, name };
+  });
+
+  const codecs = await Promise.all(codecPromises);
+
+  if (codecs.length === 1) {
+    return codecs[0].codec;
+  }
+  return joinCodecs(codecs);
+};
+
+const getCodec = async (descriptor: CodecDescriptor): Promise<ArrayToBytesCodec> => {
+  const savedCodec = codecs.get(descriptor.id);
+  if (savedCodec !== undefined) {
+    return savedCodec;
+  }
+
+  const newCodec = await loadCodec(descriptor);
+  codecs.set(descriptor.id, newCodec);
+  return newCodec;
+};
+
+export const encodeTask = task("zarrEncode", async (chunk: Chunk<NumberType>, descriptor: CodecDescriptor) => {
+  const codec = await getCodec(descriptor);
+  const result = await codec.encode(chunk);
+  return { result, transfer: [result.buffer] };
+});
+
+export const decodeTask = task("zarrDecode", async (data: Uint8Array, descriptor: CodecDescriptor) => {
+  const codec = await getCodec(descriptor);
+  const result = await codec.decode(data);
+  return { result, transfer: [result.data.buffer] };
+});

src/data_manager/sources/ome_zarr_source.ts

@@ -17,6 +17,8 @@ import {
 import { unitNameToSymbol } from "../../loaders/VolumeLoaderUtils.js";
 import type { Chunk, LocalChunkId } from "../types.js";
 import type { NumberType, TypedArray } from "../../types.js";
+import type { TaskPool } from "../task_pool/task_pool.js";
+import { withWorkerCodec } from "./ome_zarr/with_threaded_codecs.js";
 
 const PLACEHOLDER_NAME = "zarr source";
 const PLACEHOLDER_SCENE_INDEX = 0;
@@ -34,8 +36,8 @@ export class OMEZarrSource implements ChunkSource {
     public readonly axesTCZYX: [number, number, number, number, number]
   ) {}
 
-  static async new(url: string): Promise<OMEZarrSource> {
-    const store = new zarr.FetchStore(url);
+  static async new(url: string, options?: { pool?: TaskPool }): Promise<OMEZarrSource> {
+    const store = zarr.extendStore(new zarr.FetchStore(url), (s) => withWorkerCodec(s, { pool: options?.pool }));
     const root = zarr.root(store);
 
     const group = await zarr.open(root, { kind: "group" });

src/data_manager/task_pool/borrow_guard.ts

@@ -10,7 +10,7 @@ type BorrowGuardEvents = { restored: void };
  * Wraps a `TypedArray` that may be temporarily transferred to a worker for processing.
  *
  * A `TypedArray` guarded by this class must be accessed through the `get` method, which returns `undefined` when the
- * buffer is on a worker. The guard will trigger the `"restored"` event when it is returned to the main thread.
+ * buffer is on a worker. The guard will trigger the `restored` event when it is returned to the main thread.
  */
 export class BorrowGuard<T extends TypedArray = TypedArray> extends EventDispatcher<BorrowGuardEvents> {
   private [guardMarker] = true as const;

src/data_manager/task_pool/runner.ts

@@ -2,7 +2,7 @@ import { serializeError } from "serialize-error";
 
 import type { TypedArray, NumberType } from "../../types.js";
 import { getBorrowed } from "./borrow_guard.js";
-import type { Task, WorkerRequest, WorkerResponse } from "./types.js";
+import type { Task, TaskResult, WorkerRequest, WorkerResponse } from "./types.js";
 
 /**
  * Combines a `string` task `id` and a `handler` into a `Task` that can be registered and called with a `TaskPool`.
@@ -35,7 +35,7 @@ import type { Task, WorkerRequest, WorkerResponse } from "./types.js";
  */
 export const task = <Id extends string, In extends unknown[], Out>(
   id: Id,
-  handler: (...args: In) => { result: Out; transfer: Transferable[] }
+  handler: (...args: In) => TaskResult<Out>
 ): Task<Id, In, Out> => {
   const task = handler as Task<Id, In, Out>;
   task.taskId = id;
@@ -45,7 +45,7 @@ export const task = <Id extends string, In extends unknown[], Out>(
 export const initWorker = (tasks: Task[]) => {
   const handlers = new Map(tasks.map((task) => [task.taskId, task]));
 
-  const runTask = ({ id, task, args }: WorkerRequest): [WorkerResponse, Transferable[]] => {
+  const runTask = async ({ id, task, args }: WorkerRequest): Promise<[WorkerResponse, Transferable[]]> => {
     const handler = handlers.get(task);
 
     // Extract borrowed `TypedArray`s, which must be returned to the main thread on completion
@@ -67,16 +67,16 @@ export const initWorker = (tasks: Task[]) => {
     }
 
     try {
-      const { result, transfer } = handler(...processedArgs);
+      const { result, transfer } = await handler(...processedArgs);
       return [{ id, task, borrows, result, error: false }, [...borrowedBuffers, ...transfer]];
     } catch (e) {
       const result = serializeError(e);
       return [{ id, task, borrows, result, error: true }, borrowedBuffers];
     }
   };
 
-  self.onmessage = ({ data }: MessageEvent<WorkerRequest>) => {
-    const [result, transfer] = runTask(data);
+  self.onmessage = async ({ data }: MessageEvent<WorkerRequest>) => {
+    const [result, transfer] = await runTask(data);
     self.postMessage(result, { transfer });
   };
 

src/data_manager/task_pool/task_pool.ts

@@ -6,9 +6,7 @@ import { BorrowGuard, markBorrowed } from "./borrow_guard.js";
 export const taskHandle = <T extends Task>(
   id: T["taskId"],
   transfer?: (...args: Parameters<T>) => Transferable[]
-): TaskHandle<Parameters<T>, ReturnType<T>["result"]> => {
-  return { id, transfer };
-};
+): TaskHandle<Parameters<T>, Awaited<ReturnType<T>>["result"]> => ({ id, transfer });
 
 type StoredPromise = {
   resolve: (value: unknown) => void;

src/data_manager/task_pool/types.ts

@@ -3,10 +3,14 @@ import type { ErrorObject } from "serialize-error";
 import type { TypedArray, NumberType } from "../../types.js";
 import type { BorrowGuard } from "./borrow_guard.js";
 
+export type TaskResult<Out> =
+  | { result: Out; transfer: Transferable[] }
+  | Promise<{ result: Out; transfer: Transferable[] }>;
+
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export type Task<Id extends string = string, In extends any[] = any[], Out = any> = {
   taskId: Id;
-  (...args: In): { result: Out; transfer: Transferable[] };
+  (...args: In): TaskResult<Out>;
 };
 
 export type TaskHandle<In extends unknown[], Out> = {
@@ -60,7 +64,7 @@ export type WorkerResponse<T extends Task = Task> = {
 } & (
   | {
       error: false;
-      result: ReturnType<T>["result"];
+      result: Awaited<ReturnType<T>>["result"];
     }
   | {
       error: true;