@cfasim-ui/docs 0.3.18 → 0.4.1

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 {