@pie-project/server-web 0.5.2

package/index.d.ts

@@ -0,0 +1,65 @@
+import { PieClient } from "@pie-project/client";
+
+export { PieClient };
+
+export interface LoadInfo {
+  worker: boolean;
+}
+
+export interface LoadOptions {
+  worker?: boolean;
+  log?: string;
+  wasmUrl?: string;
+  fresh?: boolean;
+}
+
+export interface BootSummary {
+  model: string;
+  sku: string;
+  trace: string;
+  weight_bytes: number;
+  kv_pages: number;
+  [key: string]: unknown;
+}
+
+export function load(options?: LoadOptions): Promise<LoadInfo>;
+
+export function bootLazy(
+  url: string,
+  config?: string,
+  onProgress?: (served: number, total: number, from: "network" | "cache") => void,
+): Promise<BootSummary>;
+
+export function awaitCache(): Promise<void>;
+
+export function install(wasmBytes: Uint8Array | ArrayBuffer, manifestToml: string): Promise<string>;
+
+export function client(): PieClient;
+
+export function connect(): Promise<PieClient>;
+
+export function memoryBytes(): Promise<number>;
+
+export { registerLanguage, registeredLanguages, attachLanguages } from "./src/languages.js";
+
+export interface StartOptions extends LoadOptions {
+  /** The URL of the `.wgpu.zt` artifact, served with Range support. */
+  model: string;
+  /** The boot config: an object, or TOML text. */
+  config?: Record<string, unknown> | string;
+  onProgress?: (served: number, total: number, from: "network" | "cache") => void;
+}
+
+/** One running pie in this tab, with the surface of the Node host. */
+export class Server {
+  static start(options: StartOptions): Promise<Server>;
+  readonly summary: BootSummary;
+  readonly url: string;
+  readonly running: boolean;
+  installLanguage(language: "python" | "javascript" | string, source: Uint8Array | ArrayBuffer | URL | string): Promise<string>;
+  install(source: Uint8Array | ArrayBuffer | URL | string, manifest: string): Promise<string>;
+  connect(): Promise<PieClient>;
+  memoryBytes(): Promise<number>;
+  awaitCache(): Promise<void>;
+  shutdown(): Promise<void>;
+}

package/package.json

@@ -0,0 +1,38 @@
+{
+  "name": "@pie-project/server-web",
+  "version": "0.5.2",
+  "description": "The browser half of @pie-project/server: the runtime and the WebGPU engine as wasm, reached through `import { Server } from \"@pie-project/server\"` under a bundler.",
+  "type": "module",
+  "main": "dist/pie.mjs",
+  "types": "index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./index.d.ts",
+      "default": "./dist/pie.mjs"
+    },
+    "./languages": {
+      "types": "./src/languages.d.ts",
+      "default": "./src/languages.mjs"
+    }
+  },
+  "files": [
+    "dist",
+    "src",
+    "index.d.ts"
+  ],
+  "scripts": {
+    "build": "./build.sh release",
+    "bundle": "esbuild src/pie.mjs src/worker.mjs --bundle --format=esm --outdir=dist --out-extension:.js=.mjs --external:./worker.mjs --external:fs --define:__PIE_WASM__='\"./pie_browser_bg.wasm\"' --log-level=warning --log-override:duplicate-object-key=silent && cp pkg/pie_browser_bg.wasm dist/"
+  },
+  "dependencies": {
+    "@pie-project/client": "^0.5.0"
+  },
+  "devDependencies": {
+    "esbuild": "^0.28.2"
+  },
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/pie-project/pie",
+    "directory": "javascript/server-web"
+  }
+}

package/src/languages.d.ts

@@ -0,0 +1,5 @@
+export function registerLanguage(language: string, url: URL | string): void;
+export function registeredLanguages(): Map<string, URL | string>;
+export function attachLanguages(server: {
+  installLanguage(language: string, source: URL | string): Promise<string>;
+}): Promise<() => void>;

package/src/languages.mjs

@@ -0,0 +1,27 @@
+// The language components a server installs when it starts. A language
+// package registers itself on import (`import "@pie-project/language-python"`);
+// `Server.start()` installs every registered language, and a registration
+// after the start lands on the running servers too. The registry is one per
+// process, however many copies of this module are loaded.
+
+const registry = (globalThis[Symbol.for("pie.languages")] ??= { languages: new Map(), servers: new Set() });
+
+/** Register a language component by the URL of its wasm. */
+export function registerLanguage(language, url) {
+  registry.languages.set(language, url);
+  for (const server of registry.servers) {
+    server.installLanguage(language, url).catch((error) => console.error(`pie: ${language}: ${error.message ?? error}`));
+  }
+}
+
+/** The registered languages, language → URL. */
+export function registeredLanguages() {
+  return new Map(registry.languages);
+}
+
+/** Install every registered language on `server`, and every one registered later while it runs. */
+export async function attachLanguages(server) {
+  registry.servers.add(server);
+  for (const [language, url] of registry.languages) await server.installLanguage(language, url);
+  return () => registry.servers.delete(server);
+}

package/src/pie.mjs

@@ -0,0 +1,211 @@
+import { PieClient } from "@pie-project/client";
+
+import { socketClass } from "./transport.mjs";
+import { attachLanguages } from "./languages.mjs";
+
+export { PieClient };
+
+const WASM_PATH = typeof __PIE_WASM__ === "undefined" ? "../pkg/pie_browser_bg.wasm" : __PIE_WASM__;
+
+class WorkerBackend {
+  terminate() {
+    this.worker.terminate();
+  }
+
+  constructor(worker) {
+    this.worker = worker;
+    this.pending = new Map();
+    this.next = 0;
+    worker.onmessage = ({ data }) => {
+      const call = this.pending.get(data.id);
+      if (!call) return;
+      if (data.progress) {
+        call.onProgress?.(...data.progress);
+        return;
+      }
+      this.pending.delete(data.id);
+      if (data.ok) call.resolve(data.value);
+      else call.reject(new Error(data.error));
+    };
+    worker.onerror = (e) => {
+      console.error(`pie worker: ${e.message ?? e}`);
+      const error = new Error(`pie worker: ${e.message ?? e}`);
+      for (const call of this.pending.values()) call.reject(error);
+      this.pending.clear();
+    };
+    worker.onmessageerror = (e) => console.error("pie worker: message could not be deserialised", e);
+  }
+
+  call(op, args = [], { transfer = [], onProgress } = {}) {
+    const id = ++this.next;
+    return new Promise((resolve, reject) => {
+      this.pending.set(id, { resolve, reject, onProgress });
+      this.worker.postMessage({ id, op, args }, transfer);
+    });
+  }
+}
+
+class LocalBackend {
+  constructor(host) {
+    this.host = host;
+  }
+
+  async call(op, args = [], { onProgress } = {}) {
+    return await this.host[op](...args, onProgress);
+  }
+}
+
+let backend = null;
+let Socket = null;
+
+function ready() {
+  if (!backend) throw new Error("pie: call load() first");
+  return backend;
+}
+
+export async function load({ worker = true, log, wasmUrl, fresh = false } = {}) {
+  if (backend && !fresh) throw new Error("pie: load() was already called");
+  if (backend) {
+    backend.terminate?.();
+    backend = null;
+  }
+  const url = new URL(wasmUrl ?? WASM_PATH, wasmUrl ? location.href : import.meta.url).href;
+  if (worker) {
+    const w = new Worker(new URL("./worker.mjs", import.meta.url), { type: "module", name: "pie" });
+    backend = new WorkerBackend(w);
+  } else {
+    const { host } = await import("./worker.mjs");
+    backend = new LocalBackend(host);
+  }
+  const info = await backend.call("load", [url, log]);
+  Socket = socketClass(backend);
+  return info;
+}
+
+export function client() {
+  ready();
+  return new PieClient("pie://local", { WebSocket: Socket });
+}
+
+export async function connect() {
+  const c = client();
+  await c.connect();
+  return c;
+}
+
+export async function install(wasmBytes, manifestToml) {
+  const copy = new Uint8Array(wasmBytes).slice();
+  return await ready().call("install", [copy, manifestToml], { transfer: [copy.buffer] });
+}
+
+/** Hand the runtime a language component (`"python"`, `"javascript"`)
+ * from bytes; a script inferlet in that language can run once this has. */
+export async function installLanguage(language, wasmBytes) {
+  const copy = new Uint8Array(wasmBytes).slice();
+  return await ready().call("installLanguage", [language, copy], { transfer: [copy.buffer] });
+}
+
+export async function memoryBytes() {
+  return await ready().call("memoryBytes");
+}
+
+export async function bootLazy(url, config, onProgress) {
+  return await ready().call("bootLazy", [url, config, location.href], { onProgress });
+}
+
+export async function awaitCache() {
+  return await ready().call("awaitCache", []);
+}
+
+async function bytesOf(source) {
+  if (source instanceof Uint8Array) return source;
+  if (source instanceof ArrayBuffer) return new Uint8Array(source);
+  const response = await fetch(source);
+  if (!response.ok) throw new Error(`pie: ${source}: ${response.status} ${response.statusText}`);
+  return new Uint8Array(await response.arrayBuffer());
+}
+
+/** One running pie in this tab, with the surface of the Node host. */
+export class Server {
+  #summary;
+  #open = true;
+  #detach = null;
+
+  constructor(summary) {
+    this.#summary = summary;
+  }
+
+  /**
+   * Boot the runtime and the WebGPU engine. `model` is the URL of a
+   * `.wgpu.zt` artifact served with Range support; `config` the boot config
+   * (an object, or TOML text). Resolves once the weights are on the device,
+   * with every registered language component installed.
+   */
+  static async start({ model, config = {}, onProgress, worker = true, log, wasmUrl } = {}) {
+    if (typeof model !== "string") throw new Error("Server.start: `model` is the URL of the artifact to serve");
+    if (!backend) await load({ worker, log, wasmUrl });
+    const summary = await bootLazy(model, typeof config === "string" ? config : JSON.stringify(config), onProgress);
+    const server = new Server(summary);
+    server.#detach = await attachLanguages(server);
+    return server;
+  }
+
+  /** What the boot found: model, sku, weight bytes, kv pages. */
+  get summary() {
+    return this.#summary;
+  }
+
+  /** The in-page address clients connect to. */
+  get url() {
+    return "pie://local";
+  }
+
+  /** True until `shutdown()`. */
+  get running() {
+    return this.#open && backend !== null;
+  }
+
+  /** A language component (`python`, `javascript`) from bytes or a URL. */
+  async installLanguage(language, source) {
+    this.#alive();
+    return installLanguage(language, await bytesOf(source));
+  }
+
+  /** An inferlet from its component bytes (or URL) and manifest text, replacing an installed version. */
+  async install(source, manifest) {
+    this.#alive();
+    return install(await bytesOf(source), manifest);
+  }
+
+  /** A `PieClient` connected in-page. */
+  async connect() {
+    this.#alive();
+    return connect();
+  }
+
+  /** Bytes of wasm memory in use. */
+  async memoryBytes() {
+    this.#alive();
+    return memoryBytes();
+  }
+
+  /** Resolves once the artifact is copied into the origin's private file system. */
+  async awaitCache() {
+    this.#alive();
+    return awaitCache();
+  }
+
+  /** Stop the worker and release the runtime. Idempotent. */
+  async shutdown() {
+    if (!this.#open) return;
+    this.#open = false;
+    this.#detach?.();
+    backend?.terminate?.();
+    backend = null;
+    Socket = null;
+  }
+
+  #alive() {
+    if (!this.running) throw new Error("the server is shut down");
+  }
+}

package/src/platform.mjs

@@ -0,0 +1,60 @@
+// Fibers (wasmtime-fiber's `custom` contract, see crates/wasmtime-web):
+// every context, the page's tick and each fiber, is one promising export
+// invocation; a park returns a promise from the Suspending import. JSPI keeps
+// the engine stack; Rust's shadow stack (`__stack_pointer`) is saved and
+// restored here so the two stay paired.
+let wasm = null;
+let fiberEntry = null;
+const slots = new Map();
+let onWake = () => {};
+
+export function attach(exports, wakeHook) {
+  wasm = exports;
+  fiberEntry = WebAssembly.promising(exports.pie_fiber_entry);
+  onWake = wakeHook;
+}
+
+function finishFiber(top) {
+  const parked = slots.get(top);
+  slots.delete(top);
+  if (!parked || parked.init) {
+    console.error(`[fiber] ${top} finished with no parked resumer`, parked);
+    return;
+  }
+  wasm.__stack_pointer.value = parked.sp;
+  parked.resolve();
+}
+
+export function pie_fiber_init(top, entry, arg0) {
+  slots.set(top >>> 0, { init: { entry, arg0 } });
+}
+
+export const pie_fiber_switch = new WebAssembly.Suspending((top) => {
+  top = top >>> 0;
+  const parked = slots.get(top);
+  if (!parked) throw new Error(`fiber switch to unknown stack ${top}`);
+  const me = { sp: wasm.__stack_pointer.value, resolve: null };
+  const resumed = new Promise((r) => (me.resolve = r));
+  slots.set(top, me);
+  if (parked.init) {
+    const { entry, arg0 } = parked.init;
+    queueMicrotask(() => {
+      wasm.__stack_pointer.value = (top - 16) >>> 0;
+      fiberEntry(entry, arg0, top).then((ret) => finishFiber(ret >>> 0));
+    });
+  } else {
+    queueMicrotask(() => {
+      wasm.__stack_pointer.value = parked.sp;
+      parked.resolve();
+    });
+  }
+  return resumed;
+});
+
+export function pie_now() {
+  return performance.now();
+}
+
+export function pie_wake() {
+  onWake();
+}

package/src/transport.mjs

@@ -0,0 +1,74 @@
+function moved(data) {
+  const bytes = data instanceof Uint8Array ? data : new Uint8Array(data);
+  if (bytes.byteOffset === 0 && bytes.byteLength === bytes.buffer.byteLength && !bytes.buffer.detached) {
+    return { bytes, transfer: [bytes.buffer] };
+  }
+  const copy = bytes.slice();
+  return { bytes: copy, transfer: [copy.buffer] };
+}
+
+export function socketClass(backend) {
+  return class PieSocket extends EventTarget {
+  static CONNECTING = 0;
+  static OPEN = 1;
+  static CLOSING = 2;
+  static CLOSED = 3;
+
+  constructor(url) {
+    super();
+    this.url = url;
+    this.binaryType = "blob";
+    this.readyState = PieSocket.CONNECTING;
+    this.onopen = null;
+    this.onmessage = null;
+    this.onerror = null;
+    this.onclose = null;
+    this.session = null;
+    backend.call("openSession").then(
+      (id) => {
+        if (this.readyState !== PieSocket.CONNECTING) {
+          backend.call("closeSession", [id]);
+          return;
+        }
+        this.session = id;
+        this.readyState = PieSocket.OPEN;
+        this.onopen?.(new Event("open"));
+        this.pump();
+      },
+      (e) => {
+        this.readyState = PieSocket.CLOSED;
+        this.onerror?.(e);
+      },
+    );
+  }
+
+  send(data) {
+    if (this.readyState !== PieSocket.OPEN) throw new Error("socket is not open");
+    const { bytes, transfer } = moved(data);
+    backend.call("sendFrame", [this.session, bytes], { transfer }).catch((e) => this.onerror?.(e));
+  }
+
+  async pump() {
+    while (this.readyState === PieSocket.OPEN) {
+      let frames;
+      try {
+        frames = await backend.call("recvFrames", [this.session, 1000, 64]);
+      } catch (e) {
+        if (this.readyState === PieSocket.OPEN) this.onerror?.(e);
+        break;
+      }
+      for (const frame of frames) {
+        const data = this.binaryType === "blob" ? new Blob([frame]) : frame.buffer;
+        this.onmessage?.({ data });
+      }
+    }
+  }
+
+  close() {
+    if (this.readyState === PieSocket.CLOSED) return;
+    this.readyState = PieSocket.CLOSED;
+    if (this.session !== null) backend.call("closeSession", [this.session]);
+    this.onclose?.(new Event("close"));
+  }
+  };
+}