@cfasim-ui/docs 0.3.17 → 0.4.0

package/pyodide/pyodide.worker.ts

@@ -5,13 +5,25 @@ import {
 } from "@cfasim-ui/shared";
 import type { ColumnDescriptor, ModelOutputsWire } from "@cfasim-ui/shared";
 
-interface WorkerMessage {
+interface RunMessage {
   id: number;
-  type?: "run" | "loadModule";
-  python?: string;
-  module?: string;
+  type?: "run";
+  python: string;
   context?: Record<string, unknown>;
 }
+interface CallMessage {
+  id: number;
+  type: "call";
+  module: string;
+  fn: string;
+  kwargs?: Record<string, unknown>;
+}
+interface LoadModuleMessage {
+  id: number;
+  type: "loadModule";
+  module: string;
+}
+type WorkerMessage = RunMessage | CallMessage | LoadModuleMessage;
 
 let wheelMap: Record<string, string> = {};
 
@@ -92,24 +104,29 @@ function installAllWheels(): Promise<void> {
   return installPromise;
 }
 
-const modulePromises = new Map<string, Promise<void>>();
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const modulePromises = new Map<string, Promise<any>>();
 
 function ensureModule(
   pyodide: Awaited<typeof pyodideReadyPromise>,
   moduleName: string,
-): Promise<void> {
-  if (!modulePromises.has(moduleName)) {
-    if (!wheelMap[moduleName]) throw new Error(`Unknown module: ${moduleName}`);
-    const promise = (async () => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+): Promise<any> {
+  let p = modulePromises.get(moduleName);
+  if (!p) {
+    if (!wheelMap[moduleName]) {
+      return Promise.reject(new Error(`Unknown module: ${moduleName}`));
+    }
+    p = (async () => {
       await installAllWheels();
-      pyodide.pyimport(moduleName);
+      return pyodide.pyimport(moduleName);
     })();
-    promise.catch(() => {
+    p.catch(() => {
       modulePromises.delete(moduleName);
     });
-    modulePromises.set(moduleName, promise);
+    modulePromises.set(moduleName, p);
   }
-  return modulePromises.get(moduleName)!;
+  return p;
 }
 
 // Map Python struct format characters to TypedArray constructors
@@ -127,18 +144,20 @@ const FORMAT_TO_TYPED_ARRAY: Record<
   d: Float64Array,
 };
 
+// Copy a Pyodide getBuffer() result into a JS-owned typed array. Releases the
+// underlying Python buffer; the returned view's ArrayBuffer is safe to transfer.
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-function extractNumpyBuffer(proxy: any): ArrayBuffer {
+function copyPyBuffer(proxy: any): ArrayBufferView {
   const pyBuffer = proxy.getBuffer();
   const Ctor = FORMAT_TO_TYPED_ARRAY[pyBuffer.format] ?? Float64Array;
-  const typed = new Ctor(
+  const view = new Ctor(
     pyBuffer.data.buffer.slice(
       pyBuffer.data.byteOffset,
       pyBuffer.data.byteOffset + pyBuffer.data.byteLength,
     ),
   );
   pyBuffer.release();
-  return typed.buffer as ArrayBuffer;
+  return view;
 }
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -161,7 +180,7 @@ function convertModelOutputs(jsResult: any): ModelOutputsWire | null {
     for (const buf of wire.buffers) {
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       if (buf && typeof buf === "object" && (buf as any).getBuffer) {
-        buffers.push(extractNumpyBuffer(buf));
+        buffers.push(copyPyBuffer(buf).buffer as ArrayBuffer);
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         if ((buf as any).destroy) (buf as any).destroy();
       } else if (buf instanceof ArrayBuffer) {
@@ -187,76 +206,94 @@ function convertModelOutputs(jsResult: any): ModelOutputsWire | null {
   return { __modelOutputs: true, outputs };
 }
 
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function convertResult(rawResult: any): unknown {
+  if (!rawResult || typeof rawResult.destroy !== "function") return rawResult;
+  try {
+    if (typeof rawResult.toJs === "function") {
+      return rawResult.toJs({ dict_converter: Object.fromEntries });
+    }
+    if (typeof rawResult.getBuffer === "function") {
+      return copyPyBuffer(rawResult);
+    }
+    return rawResult.toString();
+  } finally {
+    rawResult.destroy();
+  }
+}
+
+function send(id: number, result: unknown) {
+  const modelOutputs = convertModelOutputs(result);
+  if (modelOutputs) {
+    postModelOutputsWithTransfer(self, id, modelOutputs);
+  } else {
+    postWithTransfer(self, id, result);
+  }
+}
+
 self.onmessage = async (event: MessageEvent<WorkerMessage>) => {
   const pyodide = await pyodideReadyPromise;
-  const { id, type, python, module: moduleName, context } = event.data;
+  const msg = event.data;
+  const { id } = msg;
 
   try {
-    if (type === "loadModule" && moduleName) {
-      await ensureModule(pyodide, moduleName);
+    if (msg.type === "loadModule") {
+      await ensureModule(pyodide, msg.module);
       postWithTransfer(self, id, true);
       return;
     }
 
-    let globals;
-    if (context) {
-      const dict = pyodide.globals.get("dict");
-      globals = dict(Object.entries(context));
+    if (msg.type === "call") {
+      const mod = await ensureModule(pyodide, msg.module);
+      const t0 = performance.now();
+      // mod[fn] creates a fresh PyProxy on every access — release in finally.
+      const pyFn = mod[msg.fn];
+      let result: unknown;
+      try {
+        if (typeof pyFn !== "function") {
+          throw new Error(`Module ${msg.module} has no function ${msg.fn}`);
+        }
+        const rawResult = msg.kwargs ? pyFn.callKwargs(msg.kwargs) : pyFn();
+        result = convertResult(rawResult);
+      } finally {
+        if (pyFn && typeof pyFn.destroy === "function") pyFn.destroy();
+      }
+      const tEnd = performance.now();
+      console.log(
+        `[pyodide-worker] ${msg.module}.${msg.fn} ${
+          Math.round((tEnd - t0) * 10) / 10
+        }ms`,
+      );
+      send(id, result);
+      return;
     }
 
-    const t0 = performance.now();
-
-    // Use synchronous runPython (models don't use top-level await)
-    const rawResult = pyodide.runPython(
-      python!,
-      globals ? { globals } : undefined,
-    );
-
-    const tPython = performance.now();
-
-    // Destroy PyProxy if returned to prevent memory leaks
-    let result = rawResult;
-    if (rawResult && typeof rawResult === "object" && rawResult.destroy) {
-      if (rawResult.toJs) {
-        result = rawResult.toJs({ dict_converter: Object.fromEntries });
-      } else if (rawResult.getBuffer) {
-        // Single numpy array: use getBuffer() for direct typed array access
-        const pyBuffer = rawResult.getBuffer();
-        const Ctor = FORMAT_TO_TYPED_ARRAY[pyBuffer.format] ?? Float64Array;
-        result = new Ctor(
-          pyBuffer.data.buffer.slice(
-            pyBuffer.data.byteOffset,
-            pyBuffer.data.byteOffset + pyBuffer.data.byteLength,
-          ),
-        );
-        pyBuffer.release();
-      } else {
-        result = rawResult.toString();
+    // type === "run" or omitted
+    let globals;
+    if (msg.context) {
+      const dict = pyodide.globals.get("dict");
+      try {
+        globals = dict(Object.entries(msg.context));
+      } finally {
+        dict.destroy();
       }
-      rawResult.destroy();
     }
-
-    // Destroy globals proxy if created
-    if (globals && globals.destroy) {
-      globals.destroy();
+    const t0 = performance.now();
+    let result: unknown;
+    try {
+      const rawResult = pyodide.runPython(
+        msg.python,
+        globals ? { globals } : undefined,
+      );
+      result = convertResult(rawResult);
+    } finally {
+      if (globals && globals.destroy) globals.destroy();
     }
-
-    const tConvert = performance.now();
-    const bench = {
-      python_ms: Math.round((tPython - t0) * 10) / 10,
-      convert_ms: Math.round((tConvert - tPython) * 10) / 10,
-    };
+    const tEnd = performance.now();
     console.log(
-      `[pyodide-worker] python=${bench.python_ms}ms convert=${bench.convert_ms}ms`,
+      `[pyodide-worker] runPython ${Math.round((tEnd - t0) * 10) / 10}ms`,
     );
-
-    // Check for ModelOutputs wire format
-    const modelOutputs = convertModelOutputs(result);
-    if (modelOutputs) {
-      postModelOutputsWithTransfer(self, id, modelOutputs);
-    } else {
-      postWithTransfer(self, id, result);
-    }
+    send(id, result);
   } catch (error) {
     postErrorWithTransfer(self, id, error);
   }

package/pyodide/pyodideWorkerApi.ts

@@ -1,102 +1,196 @@
 import type { TransferableResponse } from "@cfasim-ui/shared";
 import { unwrapResponse } from "@cfasim-ui/shared";
 
-interface PromiseResolver<T> {
-  promise: Promise<T>;
-  resolve: (value: T) => void;
-}
-
-function getPromiseAndResolve<T>(): PromiseResolver<T> {
-  let resolve: (value: T) => void;
-  const promise = new Promise<T>((res) => {
-    resolve = res;
-  });
-  return { promise, resolve: resolve! };
-}
-
 let lastId = 0;
 function getId(): number {
   return ++lastId;
 }
 
-interface WorkerMessage {
+interface RunMessage {
   id: number;
-  type?: "run" | "loadModule";
-  python?: string;
-  module?: string;
+  type?: "run";
+  python: string;
   context?: Record<string, unknown>;
 }
+interface CallMessage {
+  id: number;
+  type: "call";
+  module: string;
+  fn: string;
+  kwargs?: Record<string, unknown>;
+}
+interface LoadModuleMessage {
+  id: number;
+  type: "loadModule";
+  module: string;
+}
+type WorkerMessage = RunMessage | CallMessage | LoadModuleMessage;
+type OutgoingMessage =
+  | Omit<RunMessage, "id">
+  | Omit<CallMessage, "id">
+  | Omit<LoadModuleMessage, "id">;
 
 function requestResponse(
   worker: Worker,
-  msg: Omit<WorkerMessage, "id">,
+  msg: OutgoingMessage,
 ): Promise<{ result?: unknown; error?: string }> {
-  const { promise, resolve } = getPromiseAndResolve<{
-    result?: unknown;
-    error?: string;
-  }>();
-  const idWorker = getId();
-
-  function listener(event: MessageEvent<TransferableResponse>) {
-    if (event.data?.id !== idWorker) {
-      return;
+  return new Promise((resolve) => {
+    const id = getId();
+    function listener(event: MessageEvent<TransferableResponse>) {
+      if (event.data?.id !== id) return;
+      worker.removeEventListener("message", listener);
+      if (event.data.error) {
+        resolve({ error: event.data.error });
+      } else {
+        resolve({ result: unwrapResponse(event.data) });
+      }
     }
-    worker.removeEventListener("message", listener);
-    if (event.data.error) {
-      resolve({ error: event.data.error });
-    } else {
-      resolve({ result: unwrapResponse(event.data) });
-    }
-  }
-
-  worker.addEventListener("message", listener);
-  worker.postMessage({ id: idWorker, ...msg } as WorkerMessage);
-  return promise;
+    worker.addEventListener("message", listener);
+    worker.postMessage({ id, ...msg } as WorkerMessage);
+  });
 }
 
-export type WorkerName = "baseline" | "intervention";
+/**
+ * Worker names are arbitrary string keys. Workers are spawned lazily on first
+ * use; each one is an independent Pyodide interpreter. Use distinct names when
+ * you want runs to execute in parallel without contention.
+ */
+export type WorkerName = string;
 
-function createWorker(): Worker {
-  return new Worker(new URL("./pyodide.worker.ts", import.meta.url), {
-    type: "module",
-  });
+const DEFAULT_WORKER = "default";
+
+interface WorkerEntry {
+  worker: Worker;
+  modules: Map<string, Promise<{ result?: unknown; error?: string }>>;
 }
 
-const workers: Record<WorkerName, Worker> = {
-  baseline: createWorker(),
-  intervention: createWorker(),
-};
+const workers = new Map<string, WorkerEntry>();
+// Modules registered via loadModule() — auto-installed on any newly spawned worker.
+const sharedModules = new Set<string>();
 
-export async function loadModuleOnWorker(
+function getWorker(name: string): WorkerEntry {
+  let entry = workers.get(name);
+  if (!entry) {
+    const worker = new Worker(new URL("./pyodide.worker.ts", import.meta.url), {
+      type: "module",
+    });
+    entry = { worker, modules: new Map() };
+    workers.set(name, entry);
+    for (const mod of sharedModules) {
+      ensureModuleOn(entry, mod);
+    }
+  }
+  return entry;
+}
+
+function ensureModuleOn(
+  entry: WorkerEntry,
   moduleName: string,
-  worker: WorkerName,
 ): Promise<{ result?: unknown; error?: string }> {
-  return requestResponse(workers[worker], {
-    type: "loadModule",
-    module: moduleName,
-  });
+  let p = entry.modules.get(moduleName);
+  if (!p) {
+    p = requestResponse(entry.worker, {
+      type: "loadModule",
+      module: moduleName,
+    });
+    entry.modules.set(moduleName, p);
+    p.then((r) => {
+      if (r.error) entry.modules.delete(moduleName);
+    });
+  }
+  return p;
 }
 
+/**
+ * Mark `moduleName` as a shared module: install it on every currently-spawned
+ * worker, and on any future worker the moment it spawns. Returns once the
+ * default worker has finished installing.
+ */
 export async function loadModule(
   moduleName: string,
 ): Promise<{ result?: unknown; error?: string }> {
-  // Load module on all workers so any worker can run code that depends on it
-  const results = await Promise.all(
-    Object.values(workers).map((w) =>
-      requestResponse(w, { type: "loadModule", module: moduleName }),
-    ),
-  );
-  return results[0];
+  sharedModules.add(moduleName);
+  // Always install on the default worker — that's what callers expect when
+  // the function name has no "OnWorker" suffix.
+  const defaultInstall = ensureModuleOn(getWorker(DEFAULT_WORKER), moduleName);
+  const others: Array<Promise<unknown>> = [];
+  for (const [name, entry] of workers) {
+    if (name !== DEFAULT_WORKER) others.push(ensureModuleOn(entry, moduleName));
+  }
+  await Promise.all(others);
+  return defaultInstall;
 }
 
+/**
+ * Install `moduleName` on a single worker. Spawns the worker if needed. Does
+ * not mark the module as shared — other workers won't auto-load it. Use this
+ * when one worker should diverge from others.
+ */
+export async function loadModuleOnWorker(
+  moduleName: string,
+  worker: WorkerName,
+): Promise<{ result?: unknown; error?: string }> {
+  return ensureModuleOn(getWorker(worker), moduleName);
+}
+
+/**
+ * Run an arbitrary Python script. `context` keys are exposed as Python globals.
+ * Prefer {@link callPython} when calling a function on a loaded module — it
+ * skips Python source parsing on every invocation.
+ */
 export async function asyncRunPython(
   script: string,
   context?: Record<string, unknown>,
-  worker?: WorkerName,
+  worker: WorkerName = DEFAULT_WORKER,
 ): Promise<{ result?: unknown; error?: string }> {
-  const target = worker ? workers[worker] : workers.baseline;
-  return requestResponse(target, {
+  return requestResponse(getWorker(worker).worker, {
+    type: "run",
     python: script,
     context,
   });
 }
+
+/**
+ * Call `module.fn(**kwargs)` on the named worker. Faster than asyncRunPython
+ * for repeated invocations: the worker caches the imported module and dispatches
+ * directly to the cached function rather than re-parsing source each call.
+ */
+export async function callPython(
+  module: string,
+  fn: string,
+  kwargs?: Record<string, unknown>,
+  worker: WorkerName = DEFAULT_WORKER,
+): Promise<{ result?: unknown; error?: string }> {
+  const entry = getWorker(worker);
+  const installResult = await ensureModuleOn(entry, module);
+  if (installResult.error) return installResult;
+  return requestResponse(entry.worker, {
+    type: "call",
+    module,
+    fn,
+    kwargs,
+  });
+}
+
+/**
+ * Pre-spawn the named workers and (optionally) pre-install modules on each.
+ * Pyodide takes a few seconds to boot, so warming workers up front lets the
+ * first call return immediately. Call this once near app startup when you
+ * know you'll need parallel interpreters (e.g. for side-by-side comparisons).
+ *
+ * Modules listed in `modules` are also registered as shared, so any worker
+ * spawned later will auto-install them.
+ */
+export async function warmWorkers(options: {
+  workers: WorkerName[];
+  modules?: string[];
+}): Promise<void> {
+  const modules = options.modules ?? [];
+  for (const mod of modules) sharedModules.add(mod);
+  const installs: Array<Promise<unknown>> = [];
+  for (const name of options.workers) {
+    const entry = getWorker(name);
+    for (const mod of modules) installs.push(ensureModuleOn(entry, mod));
+  }
+  await Promise.all(installs);
+}

package/pyodide/useModel.ts

@@ -1,7 +1,14 @@
 import { ref, toRaw, toValue, watch } from "vue";
 import type { MaybeRef } from "vue";
 import type { ModelOutput } from "@cfasim-ui/shared";
-import { asyncRunPython, loadModule } from "./pyodideWorkerApi.js";
+import { callPython, loadModule } from "./pyodideWorkerApi.js";
+
+function plainKwargs(
+  ctx: Record<string, unknown> | undefined,
+): Record<string, unknown> | undefined {
+  if (!ctx) return undefined;
+  return Object.fromEntries(Object.entries(ctx).map(([k, v]) => [k, toRaw(v)]));
+}
 
 export function useModel<T = unknown>(moduleName: string) {
   const result = ref<T>();
@@ -18,17 +25,7 @@ export function useModel<T = unknown>(moduleName: string) {
     error.value = undefined;
     try {
       await loaded;
-      const argNames = context ? Object.keys(context) : [];
-      const callArgs = argNames.join(", ");
-      const plainContext = context
-        ? Object.fromEntries(
-            Object.entries(context).map(([k, v]) => [k, toRaw(v)]),
-          )
-        : undefined;
-      const response = await asyncRunPython(
-        `import ${moduleName}\n${moduleName}.${fn}(${callArgs})`,
-        plainContext,
-      );
+      const response = await callPython(moduleName, fn, plainKwargs(context));
       if (response.error) {
         error.value = response.error;
       } else {
@@ -56,15 +53,7 @@ export function useModel<T = unknown>(moduleName: string) {
         outputsError.value = undefined;
         try {
           await loaded;
-          const plain = Object.fromEntries(
-            Object.entries(p).map(([k, v]) => [k, toRaw(v)]),
-          );
-          const argNames = Object.keys(plain);
-          const callArgs = argNames.join(", ");
-          const response = await asyncRunPython(
-            `import ${moduleName}\n${moduleName}.${fn}(${callArgs})`,
-            plain,
-          );
+          const response = await callPython(moduleName, fn, plainKwargs(p));
           if (response.error) {
             outputsError.value = response.error;
           } else {