diff --git a/public/index.ts b/public/index.ts index 0c6eb25a..f43afe93 100644 --- a/public/index.ts +++ b/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(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 = { diff --git a/public/test_worker.ts b/public/test_worker.ts new file mode 100644 index 00000000..962404a8 --- /dev/null +++ b/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]); diff --git a/src/data_manager/data_manager.ts b/src/data_manager/data_manager.ts index 52594a35..bce1d8af 100644 --- a/src/data_manager/data_manager.ts +++ b/src/data_manager/data_manager.ts @@ -13,6 +13,7 @@ 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"; @@ -20,27 +21,15 @@ import { ChunkSource, type VolumeMetadata, type ExtVolumeDims } from "./sources/ 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>(Symbol("DataManager.subscriberId")); export interface IDataSubscriber { - [SUBSCRIBER_ID]?: number; onChunkLoaded?: (id: ChunkId, chunk: Chunk) => 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): number => { - if (subscriber[SUBSCRIBER_ID] === undefined) { - subscriber[SUBSCRIBER_ID] = subscriberCount; - subscriberCount += 1; - } - return subscriber[SUBSCRIBER_ID]; -}; - type SourceEntry = { source: ChunkSource; subscribers: IDataSubscriber[]; @@ -51,18 +40,6 @@ type RequestEntry = { controller: AbortController; }; -const swapRemove = (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; class ChunkQueues { diff --git a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts new file mode 100644 index 00000000..f18aa1d6 --- /dev/null +++ b/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("zarrEncode", (chunk) => [chunk.data.buffer]); +const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); + +type WorkerConfig = { + kind: CodecKind; + codecs: zarr.CodecMetadata[]; +}; + +const getPoolId = defineGlobalId(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) => 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 => { + 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); + }, + }; +}); diff --git a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts new file mode 100644 index 00000000..6765b26f --- /dev/null +++ b/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 = { + encode: (data: In) => Out | Promise; + decode: (data: Out) => In | Promise; +}; +type ArrayToArrayCodec = Codec, Chunk>; +type ArrayToBytesCodec = Codec, Uint8Array>; +type BytesToBytesCodec = Codec; +type AnyCodec = ArrayToArrayCodec | ArrayToBytesCodec | BytesToBytesCodec; + +// hack the private `ChunkMetadata` type out of zarrita +type CodecEntry = typeof registry extends Map Promise> ? V : never; +type ChunkMetadata = Parameters[1]; + +type CodecDescriptor = { + id: number; + codecs: CodecMetadata[]; + meta: ChunkMetadata; +}; + +const codecs = new Map(); + +const joinCodecs = (codecs: { codec: AnyCodec; name: string }[]): ArrayToBytesCodec => { + // Simplified version of zarrita's `createCodecPipeline` + const runStep = async (direction: "encode" | "decode", codec: string, fn: () => Promise | T): Promise => { + 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> => { + let result: Uint8Array | Chunk = data; + for (const { name, codec } of codecs) { + // evil cast + result = await runStep("decode", name, () => codec.decode(result as Chunk & Uint8Array)); + } + return result as Chunk; + }, + encode: async (data: Chunk): Promise => { + let result: Uint8Array | Chunk = data; + for (const { name, codec } of codecs.toReversed()) { + // evil cast pt. 2 + result = await runStep("encode", name, () => codec.encode(result as Chunk & Uint8Array)); + } + return result as Uint8Array; + }, + }; +}; + +const loadCodec = async (descriptor: CodecDescriptor): Promise => { + 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 => { + 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, 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] }; +}); diff --git a/src/data_manager/sources/ome_zarr_source.ts b/src/data_manager/sources/ome_zarr_source.ts index d81545d2..830e2c9e 100644 --- a/src/data_manager/sources/ome_zarr_source.ts +++ b/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 { - const store = new zarr.FetchStore(url); + static async new(url: string, options?: { pool?: TaskPool }): Promise { + 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" }); diff --git a/src/data_manager/task_pool/borrow_guard.ts b/src/data_manager/task_pool/borrow_guard.ts index fa191541..4ca2004b 100644 --- a/src/data_manager/task_pool/borrow_guard.ts +++ b/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 extends EventDispatcher { private [guardMarker] = true as const; diff --git a/src/data_manager/task_pool/runner.ts b/src/data_manager/task_pool/runner.ts index f8b24fc2..a8336159 100644 --- a/src/data_manager/task_pool/runner.ts +++ b/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: Id, - handler: (...args: In) => { result: Out; transfer: Transferable[] } + handler: (...args: In) => TaskResult ): Task => { const task = handler as Task; task.taskId = id; @@ -45,7 +45,7 @@ export const task = ( 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,7 +67,7 @@ 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); @@ -75,8 +75,8 @@ export const initWorker = (tasks: Task[]) => { } }; - self.onmessage = ({ data }: MessageEvent) => { - const [result, transfer] = runTask(data); + self.onmessage = async ({ data }: MessageEvent) => { + const [result, transfer] = await runTask(data); self.postMessage(result, { transfer }); }; diff --git a/src/data_manager/task_pool/task_pool.ts b/src/data_manager/task_pool/task_pool.ts index da250702..06e50ff5 100644 --- a/src/data_manager/task_pool/task_pool.ts +++ b/src/data_manager/task_pool/task_pool.ts @@ -6,9 +6,7 @@ import { BorrowGuard, markBorrowed } from "./borrow_guard.js"; export const taskHandle = ( id: T["taskId"], transfer?: (...args: Parameters) => Transferable[] -): TaskHandle, ReturnType["result"]> => { - return { id, transfer }; -}; +): TaskHandle, Awaited>["result"]> => ({ id, transfer }); type StoredPromise = { resolve: (value: unknown) => void; diff --git a/src/data_manager/task_pool/types.ts b/src/data_manager/task_pool/types.ts index 7448226a..a35c7b94 100644 --- a/src/data_manager/task_pool/types.ts +++ b/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 = + | { result: Out; transfer: Transferable[] } + | Promise<{ result: Out; transfer: Transferable[] }>; + // eslint-disable-next-line @typescript-eslint/no-explicit-any export type Task = { taskId: Id; - (...args: In): { result: Out; transfer: Transferable[] }; + (...args: In): TaskResult; }; export type TaskHandle = { @@ -60,7 +64,7 @@ export type WorkerResponse = { } & ( | { error: false; - result: ReturnType["result"]; + result: Awaited>["result"]; } | { error: true; diff --git a/src/data_manager/utils.ts b/src/data_manager/utils.ts index 122931ae..69a48670 100644 --- a/src/data_manager/utils.ts +++ b/src/data_manager/utils.ts @@ -10,3 +10,28 @@ export const dataTypeToSize: { [T in NumberType]: number } = { float32: 4, float64: 8, }; + +/** Creates a function that returns a stable and globally unique number for any object passed to it. */ +export const defineGlobalId = (key: symbol): ((obj: T) => number) => { + let id = 0; + return (obj: T): number => { + if (obj[key] === undefined) { + obj[key] = id; + id += 1; + } + return obj[key]; + }; +}; + +/** Removes an element from an array by replacing it with the element at the end of the array. */ +export const swapRemove = (arr: T[], index: number) => { + if (index < 0 || index >= arr.length) { + return; + } + + const { length } = arr; + const replace = arr.pop(); + if (replace !== undefined && index < length) { + arr[index] = replace; + } +}; diff --git a/tsconfig.json b/tsconfig.json index abd0857c..fd26722f 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -4,7 +4,7 @@ "allowSyntheticDefaultImports": true, "declaration": true, "declarationDir": "es/types", - "lib": ["DOM", "ES2022", "ES6", "DOM.Iterable", "ScriptHost", "ES2016.Array.Include", "ES2024.ArrayBuffer"], + "lib": ["DOM", "ES2023", "ES6", "DOM.Iterable", "ScriptHost", "ES2016.Array.Include", "ES2024.ArrayBuffer"], "module": "Node16", "moduleResolution": "node16", "noImplicitAny": false,