mlx-bun 0.4.0 → 1.0.0

package/src/engine/telemetry.ts

@@ -0,0 +1,109 @@
+// The engine's side of the `events` service: a thin adapter at the app boundary
+// that publishes what the engine already measures. It wraps the completion
+// engine to time each request (arrival, first token, the scheduler's own
+// prefill and decode figures from the run's stats) and samples the gateway and
+// caches on an interval. It changes no scheduling or numerics: the wrapper only
+// counts tokens on their way through, publishing appends to bounded queues and
+// never waits on a subscriber, and the libraries below the app know nothing of it.
+import type { AppEvent, EventBus } from "@mlx-bun/app-core";
+import type { GenerateStats } from "@mlx-bun/inference/generation";
+import type { KvCacheService } from "@mlx-bun/inference/state";
+import type { CompletionEngine } from "./completion";
+import type { GenerationGateway } from "./generation-gateway";
+
+export interface EngineTelemetryInput {
+  events: Pick<EventBus, "publish">;
+  /** The model's exact id, as `/stats` reports it. */
+  model: string;
+  /** The batch capacity the gateway runs with. */
+  capacity: number;
+  /** Bytes of resident weights. */
+  weightsBytes: number;
+  gateway: Pick<GenerationGateway, "activeRows">;
+  /** The loaded model's cache service: its own report of the KV it admitted and the prefixes it holds. */
+  cache: Pick<KvCacheService, "report">;
+  /** Milliseconds between samples of the scheduler and caches. Default 1000. */
+  intervalMs?: number;
+  /** Test seams. */
+  now?: () => number;
+  performanceNow?: () => number;
+  timers?: { setInterval(run: () => void, ms: number): unknown; clearInterval(handle: unknown): void };
+}
+
+export interface EngineTelemetry {
+  /** The same engine with each `run` observed. Everything else passes through untouched. */
+  observe<T extends CompletionEngine>(engine: T): T;
+  /** Take one sample now, even if nothing changed since the last (the interval's own samples skip an idle, unchanged reading). */
+  sample(): void;
+  /** Stops sampling; runs still in flight publish their own finish. Idempotent. */
+  stop(): void;
+}
+
+const realTimers = {
+  setInterval: (run: () => void, ms: number) => { const handle = setInterval(run, ms); handle.unref(); return handle; },
+  clearInterval: (handle: unknown) => clearInterval(handle as ReturnType<typeof setInterval>),
+};
+
+export function createEngineTelemetry(input: EngineTelemetryInput): EngineTelemetry {
+  const publish = (event: AppEvent) => { try { input.events.publish(event); } catch { /* the bus never throws; a stand-in must not reach a request */ } };
+  const now = input.now ?? Date.now, clock = input.performanceNow ?? (() => performance.now()), timers = input.timers ?? realTimers;
+  const { model, capacity } = input;
+  let requestCount = 0, inflight = 0, tokensSinceSample = 0, lastSampleAt = clock(), lastSignature = "";
+  let stopped = false;
+
+  const sample = (quiet = false) => {
+    const at = now(), seconds = Math.max((clock() - lastSampleAt) / 1000, 1e-3);
+    const tokens = tokensSinceSample;
+    tokensSinceSample = 0; lastSampleAt = clock();
+    const active = input.gateway.activeRows, queued = Math.max(0, inflight - active);
+    const held = input.cache.report(), kv = held.kv, cache = held.prefixes;
+    const tokensPerSecond = tokens / seconds;
+    // Nothing moving and nothing changed since the last sample: say nothing, so an idle server keeps its subscribers quiet.
+    const signature = [active, queued, kv.projectedBytes, kv.budgetBytes, cache.bytes, cache.hits, cache.misses].join(":");
+    if (quiet && tokens === 0 && active === 0 && queued === 0 && signature === lastSignature) return;
+    lastSignature = signature;
+    publish({ type: "scheduler.sample", at, model, active, capacity, queued, tokensPerSecond });
+    publish({ type: "cache.sample", at, model, cache: "kv", bytes: kv.projectedBytes, capacityBytes: kv.budgetBytes, hits: null, misses: null });
+    publish({ type: "cache.sample", at, model, cache: "prefix", bytes: cache.bytes, capacityBytes: Number.isFinite(cache.maxBytes) ? cache.maxBytes : null,
+      hits: cache.hits, misses: cache.misses });
+    publish({ type: "model.memory", at, model, weightsBytes: input.weightsBytes, kvBytes: kv.projectedBytes, prefixCacheBytes: cache.bytes });
+  };
+
+  const timer = timers.setInterval(() => { if (!stopped) sample(true); }, input.intervalMs ?? 1000);
+
+  function finished(prompt: number, stats: GenerateStats | undefined, arrival: number, first: number | null, tokens: number, failure: "cancelled" | "error" | null) {
+    const total = clock() - arrival, ttft = first === null ? null : first - arrival;
+    // The run's own prefill span starts at admission, so what precedes it in the time to first token was queueing.
+    const queueMs = stats && ttft !== null ? Math.max(0, ttft - stats.prefillMs) : null;
+    publish({ type: "request.finished", at: now(), model, requestId: `req-${++requestCount}`,
+      finish: failure ?? (stats?.finishReason === "length" ? "length" : "stop"),
+      promptTokens: prompt, cachedPromptTokens: stats?.cachedTokens ?? 0, completionTokens: stats?.generatedTokens ?? tokens,
+      queueMs, ttftMs: ttft, prefillTokensPerSecond: stats && stats.prefillTps > 0 ? stats.prefillTps : null,
+      decodeTokensPerSecond: stats && stats.decodeTps > 0 ? stats.decodeTps : null, totalMs: total });
+  }
+
+  return {
+    sample: () => sample(),
+    stop() { if (stopped) return; stopped = true; timers.clearInterval(timer); },
+    observe<T extends CompletionEngine>(engine: T): T {
+      return { ...engine, async run(prompt: Parameters<T["run"]>[0], options: Parameters<T["run"]>[1], onToken: Parameters<T["run"]>[2], vision: Parameters<T["run"]>[3],
+        shape: Parameters<T["run"]>[4], placement: Parameters<T["run"]>[5], signal?: AbortSignal, trace?: Parameters<T["run"]>[7]) {
+        const arrival = clock();
+        let first: number | null = null, tokens = 0;
+        inflight++;
+        try {
+          const stats = await engine.run(prompt, options, (token, logprobs) => {
+            if (first === null) first = clock();
+            tokens++; tokensSinceSample++;
+            return onToken(token, logprobs);
+          }, vision, shape, placement, signal, trace);
+          finished(prompt.length, stats, arrival, first, tokens, null);
+          return stats;
+        } catch (error) {
+          finished(prompt.length, undefined, arrival, first, tokens, signal?.aborted ? "cancelled" : "error");
+          throw error;
+        } finally { inflight--; }
+      } } as T;
+    },
+  };
+}

package/src/hub/downloads.ts

@@ -0,0 +1,116 @@
+import { CatalogFailure, type CatalogHub } from "@mlx-bun/app-services/portable";
+import { downloadModel, downloadsSnapshot, type DownloadOptions, type DownloadStatus } from "@mlx-bun/hub/download";
+
+export class DuplicateDownloadError extends Error {
+  constructor(repoId: string) {
+    super(`a download for ${repoId} is already in progress`);
+    this.name = "DuplicateDownloadError";
+  }
+}
+
+export interface DownloadOwnerOptions {
+  /** Transfer implementation; the default is the hub package downloader. */
+  download?: (repoId: string, options: DownloadOptions) => Promise<string>;
+  /** Process-wide activity for transfers this owner did not start. */
+  active?: (repoId: string) => boolean;
+  /** Rows the hub package records once a listing exists; the default is its process tracker. */
+  tracker?: () => readonly DownloadStatus[];
+  /** Runs after a transfer completes; never after cancellation or failure. */
+  onComplete?: (repoId: string, snapshotPath: string) => void | Promise<void>;
+  /** Failures the hub tracker cannot show, such as a metadata request that
+   * fails before any tracker row exists. Never called for cancellation. */
+  onFailure?: (repoId: string, error: unknown) => void;
+  transfer?: Pick<DownloadOptions, "cacheDir" | "endpoint" | "token">;
+}
+
+export interface DownloadOwner {
+  /** Repos with a transfer this owner has admitted and not yet joined. */
+  readonly active: readonly string[];
+  /** Admits synchronously, before any network request, so a duplicate submit
+   * is refused even while the metadata request is still pending. */
+  start(repoId: string): void;
+  /** One row per transfer for its whole lifecycle: the admission row until the
+   * hub tracker publishes its live row (after listing and preflight), which
+   * then replaces it in place; a transfer that fails or is cancelled before
+   * that keeps its own terminal row. Tracker rows this owner did not start are
+   * included unchanged. */
+  snapshot(): DownloadStatus[];
+  /** Aborts every transfer and waits for each to settle its partial on disk.
+   * Idempotent; later `start` calls are refused. */
+  close(): Promise<void>;
+}
+
+/** Terminal rows retained for the progress surface, matching the hub tracker's window. */
+const RETAINED_ROWS = 5;
+
+/** Owns background Hugging Face transfers started from the web: admission,
+ * the progress surface before and after the hub tracker has a row, completion
+ * side effects, and cancel-and-join at shutdown. */
+export function createDownloadOwner(options: DownloadOwnerOptions = {}): DownloadOwner {
+  const transfers = new Map<string, { controller: AbortController; settled: Promise<void> }>();
+  const rows: DownloadStatus[] = [];
+  let closing: Promise<void> | undefined;
+  const message = (error: unknown) => error instanceof Error ? error.message : String(error);
+  /** The tracker settles its own row before the transfer resolves; only an
+   * admission row that never received one still needs a terminal state. */
+  const settle = (row: DownloadStatus, state: "done" | "error", error?: string) => {
+    if (row.state !== "active") return;
+    row.state = state; row.bytesPerSec = 0; row.finishedAt = Date.now();
+    if (error !== undefined) row.error = error;
+  };
+  const retain = () => {
+    while (rows.length > RETAINED_ROWS) {
+      const oldest = rows.findIndex(row => row.state !== "active");
+      if (oldest < 0) break;
+      rows.splice(oldest, 1);
+    }
+  };
+  return {
+    get active() { return [...transfers.keys()]; },
+    start(repoId) {
+      if (closing) throw new Error("downloads are closed");
+      if (transfers.has(repoId) || options.active?.(repoId)) throw new DuplicateDownloadError(repoId);
+      const controller = new AbortController();
+      let row: DownloadStatus = { repoId, state: "active", currentFile: null, receivedBytes: 0, totalBytes: 0,
+        filesDone: 0, filesTotal: 0, bytesPerSec: 0, startedAt: Date.now(), finishedAt: null };
+      rows.push(row); retain();
+      const entry = { controller, settled: Promise.resolve() };
+      transfers.set(repoId, entry);
+      entry.settled = (async () => {
+        try {
+          const path = await (options.download ?? downloadModel)(repoId, { ...options.transfer, signal: controller.signal,
+            onStatus: status => { rows[rows.indexOf(row)] = status; row = status; } });
+          if (controller.signal.aborted) { settle(row, "error", message(controller.signal.reason)); return; }
+          settle(row, "done");
+          await options.onComplete?.(repoId, path);
+        } catch (error) {
+          settle(row, "error", message(controller.signal.aborted ? controller.signal.reason : error));
+          if (!controller.signal.aborted) options.onFailure?.(repoId, error);
+        } finally { transfers.delete(repoId); }
+      })();
+    },
+    snapshot() {
+      const tracked = (options.tracker ?? downloadsSnapshot)();
+      const known = new Set<DownloadStatus>(tracked);
+      return [...tracked, ...rows.filter(row => !known.has(row))].sort((a, b) => a.startedAt - b.startedAt);
+    },
+    close() {
+      return closing ??= (async () => {
+        const pending = [...transfers.values()];
+        for (const { controller } of pending) controller.abort(new Error("server is shutting down"));
+        await Promise.all(pending.map(({ settled }) => settled));
+      })();
+    },
+  };
+}
+
+/** The owner as the catalog's transfers: a repo already downloading is the catalog's `duplicate-download` failure. */
+export function catalogTransfers(owner: Pick<DownloadOwner, "start" | "snapshot">): NonNullable<CatalogHub["transfers"]> {
+  return {
+    start(id) {
+      try { owner.start(id); }
+      catch (error) { throw error instanceof DuplicateDownloadError ? new CatalogFailure("duplicate-download", error.message) : error; }
+    },
+    snapshot: () => owner.snapshot(),
+  };
+}

package/src/jobs/db.ts

@@ -7,12 +7,10 @@
 // process) writes — without WAL the reader and writer fight over the lock.
 
 import { Database } from "bun:sqlite";
-import { mkdirSync } from "node:fs";
+import { mkdirSync, writeFileSync } from "node:fs";
 import { join } from "node:path";
-import type { JobRow, JobStatus } from "./types";
-
-export const DEFAULT_JOBS_DB = `${process.env.HOME}/.cache/mlx-bun/jobs.sqlite`;
-export const DEFAULT_JOBS_DIR = `${process.env.HOME}/.cache/mlx-bun/jobs`;
+import { storagePath } from "../storage/paths";
+import type { JobRow, JobStatus } from "./protocol";
 
 const SCHEMA = `
 CREATE TABLE IF NOT EXISTS jobs (
@@ -42,12 +40,23 @@ export function newJobId(): string {
   return `job_${hex}`;
 }
 
+/** The `ended_at` format: SQLite's datetime('now') ("YYYY-MM-DD HH:MM:SS",
+ * UTC), so the column is uniform whoever writes it (main's rows included). */
+export function nowIso(): string {
+  return new Date().toISOString().replace("T", " ").slice(0, 19);
+}
+
+/** A job row's `error`: the failure's name and message, as main recorded it. */
+export function jobError(e: unknown): string {
+  return e instanceof Error ? `${e.name}: ${e.message}` : String(e);
+}
+
 export class JobStore {
   readonly db: Database;
   readonly dbPath: string;
   readonly logsDir: string;
 
-  constructor(dbPath: string = DEFAULT_JOBS_DB, logsDir: string = DEFAULT_JOBS_DIR) {
+  constructor(dbPath: string = storagePath("jobsDb"), logsDir: string = storagePath("jobLogs")) {
     if (dbPath !== ":memory:") {
       try { mkdirSync(dbPath.slice(0, dbPath.lastIndexOf("/")), { recursive: true }); } catch {}
     }
@@ -82,7 +91,7 @@ export class JobStore {
     const id = newJobId();
     const logPath = join(this.logsDir, `${id}.log`);
     // touch the log so tailJob's first read finds an (empty) file
-    try { Bun.write(logPath, ""); } catch {}
+    try { writeFileSync(logPath, ""); } catch {}
     this.db.prepare(`
       INSERT INTO jobs (id, kind, status, config_json, log_path, output_path)
       VALUES (?, ?, 'queued', ?, ?, ?)

package/src/jobs/events.ts

@@ -0,0 +1,29 @@
+import { appendFileSync } from "node:fs";
+import type { Emit, JobEvent } from "./protocol";
+import type { JobStore } from "./db";
+
+export function makeEmit(store: JobStore, jobId: string, logPath: string): Emit {
+  return (e: JobEvent) => {
+    try {
+      appendFileSync(logPath, JSON.stringify(e) + "\n");
+    } catch {
+      // swallow: a job must survive a logging hiccup
+    }
+    if (e.type === "stage" || e.type === "metric") {
+      const progress = typeof e.progress === "number" ? e.progress : undefined;
+      const message = typeof e.message === "string" ? e.message : undefined;
+      if (progress !== undefined || message !== undefined) {
+        try {
+          if (progress !== undefined) store.setProgress(jobId, progress, message);
+          else if (message !== undefined) {
+            // message-only update preserves current progress
+            store.db.prepare("UPDATE jobs SET message = ? WHERE id = ?").run(message, jobId);
+          }
+        } catch {
+          // swallow: DB contention must not kill a job
+        }
+      }
+    }
+  };
+}
+

package/src/jobs/executable.ts

@@ -0,0 +1,5 @@
+// Bun resolves process.execPath lazily. The CLI imports this module before
+// dispatch so a later installer symlink swap cannot redirect child processes
+// into a different release. Source execution captures the Bun runtime instead.
+// Job and isolation workers must share this startup identity.
+export const executablePath = process.execPath;

package/src/jobs/host.ts

@@ -0,0 +1,51 @@
+import { createInProcessJobs } from "./in-process";
+import type { JobKind, JobRunner } from "./protocol";
+import { JobStore } from "./db";
+import { cancelSubprocessJob, closeSubprocessJobs, submitSubprocess, type SubprocessOpts } from "./runner";
+
+/** Owns the lazily opened store, in-process tasks, and managed subprocesses. Composition supplies
+ * an execution lease; the host never unloads the application's resident model. */
+export function createJobHost(options: SubprocessOpts & { createStore?: () => JobStore }) {
+  const tasks = createInProcessJobs();
+  let store: JobStore | undefined;
+  let closed = false;
+  let closing: Promise<void> | undefined;
+  const cancellation = new AbortController();
+  const ensureStore = () => {
+    if (closed) throw new Error("job host is closed");
+    if (!store) {
+      store = options.createStore?.() ?? new JobStore();
+      store.markZombies();
+    }
+    return store;
+  };
+  return {
+    signal: cancellation.signal,
+    ensureStore,
+    submit(kind: string, config: Record<string, unknown>, outputPath?: string, scratchDir?: string) {
+      return submitSubprocess(ensureStore(), kind, config, outputPath, options, scratchDir);
+    },
+    submitTask(kind: JobKind, config: Record<string, unknown>, runner: JobRunner, outputPath?: string) {
+      return tasks.submit(ensureStore(), kind, config, runner, outputPath);
+    },
+    /** Stop a running in-process task; it ends `failed` with "job cancelled". Resolves once it has ended. */
+    cancelTask(jobId: string) { return tasks.cancel(jobId); },
+    /** Stop a queued or running child process of this store and wait until its process group is gone. */
+    cancelProcess(jobId: string) { return cancelSubprocessJob(ensureStore(), jobId); },
+    close() {
+      if (closing) return closing;
+      closed = true;
+      cancellation.abort();
+      return closing = (async () => {
+        if (!store) return;
+        try {
+          const results = await Promise.allSettled([tasks.close(), closeSubprocessJobs(store)]);
+          const failures = results.filter((r): r is PromiseRejectedResult => r.status === "rejected").map(r => r.reason);
+          if (failures.length === 1) throw failures[0];
+          if (failures.length) throw new AggregateError(failures, "Job shutdown failed");
+        }
+        finally { store.close(); }
+      })();
+    },
+  };
+}

package/src/jobs/in-process.ts

@@ -0,0 +1,55 @@
+import { jobError, nowIso, type JobStore } from "./db";
+import { makeEmit } from "./events";
+import type { JobKind, JobRunner } from "./protocol";
+
+/** Owns tasks until they settle. Loopback inference borrows the server scheduler;
+ * these tasks must never hold its exclusive execution lease. */
+export function createInProcessJobs() {
+  const cancellation = new AbortController();
+  const tasks = new Set<Promise<void>>();
+  const running = new Map<string, { stop: AbortController; done: Promise<void> }>();
+  const failures: unknown[] = [];
+  const signal = cancellation.signal;
+  return {
+    submit(store: JobStore, kind: JobKind, config: Record<string, unknown>, runner: JobRunner, outputPath?: string) {
+      signal.throwIfAborted();
+      const row = store.create(kind, config, outputPath);
+      const emit = makeEmit(store, row.id, row.log_path);
+      // The host's shutdown and a single job's cancellation both stop the runner.
+      const stop = new AbortController(), own = AbortSignal.any([signal, stop.signal]);
+      const task = Promise.resolve().then(async () => {
+        try {
+          own.throwIfAborted();
+          store.setStatus(row.id, "running");
+          emit({ type: "started", ts: Date.now() });
+          const result = await runner(event => { own.throwIfAborted(); emit(event); }, config, own);
+          own.throwIfAborted();
+          if (result?.outputPath) store.setOutputPath(row.id, result.outputPath);
+          store.setProgress(row.id, 1);
+          store.setStatus(row.id, "done", { endedAt: nowIso() });
+          emit({ type: "done", ts: Date.now(), output_dir: result?.outputPath ?? outputPath });
+        } catch (error) {
+          // Main's row format for every job kind: the timestamp and the failure's name.
+          const message = jobError(error);
+          store.setStatus(row.id, "failed", { error: message, endedAt: nowIso() });
+          emit({ type: "failed", error: message, ts: Date.now() });
+        }
+      }).catch(error => { failures.push(error); }).finally(() => { tasks.delete(task); running.delete(row.id); });
+      tasks.add(task);
+      running.set(row.id, { stop, done: task });
+      return { jobId: row.id, outputPath };
+    },
+    /** Stop one task and wait until it has ended; a task that is not running is left alone. */
+    async cancel(jobId: string) {
+      const entry = running.get(jobId);
+      if (!entry) return;
+      entry.stop.abort(new Error("job cancelled"));
+      await entry.done;
+    },
+    async close() {
+      cancellation.abort(new Error("Server shutting down"));
+      await Promise.all([...tasks]);
+      if (failures.length) throw new AggregateError(failures, "Could not persist in-process job status");
+    },
+  };
+}

package/src/{contracts/jobs.ts → jobs/protocol.ts}

@@ -1,8 +1,3 @@
-// Shared job-system contract. Both the job runner (src/jobs/*) and every
-// job producer (src/dataset/*, src/quantize/*, src/train/*) speak this
-// vocabulary. Mirrors the optiq lab SSE event grammar so the web wizards
-// port over mechanically (started → log/stage/metric → done|failed).
-
 /** A single streamed job event, serialized one-per-line into the job's
  *  NDJSON log and re-emitted to the browser as SSE `data:` frames. */
 export type JobEvent =
@@ -35,18 +30,21 @@ export type JobEvent =
 
 /** Sink a job calls to report progress. Implementations append to the log
  *  file and (for stage/metric events carrying `progress`/`message`) update
- *  the SQLite row. Never throws — a logging failure must not kill a job. */
+ *  the SQLite row. Logging failures must not kill a job; the task owner may
+ * throw cancellation when a runner reports progress after shutdown. */
 export type Emit = (e: JobEvent) => void;
 
 /** The unit of work. Returns an optional output path recorded on the row.
- *  Receives the parsed `config` from the submit request. */
+ *  Receives the parsed `config` from the submit request. In-process runners
+ * receive a shutdown signal and must pass it to cancellable I/O. */
 export type JobRunner = (
   emit: Emit,
   config: Record<string, unknown>,
+  signal?: AbortSignal,
 ) => Promise<{ outputPath?: string } | void>;
 
-/** Job kinds in the system. */
-export type JobKind = "quantize" | "finetune" | "dataset" | "noop" | "crash";
+/** Job kinds in the system: each kind an installed module declares (quantize, finetune, dataset). */
+export type JobKind = string;
 
 export type JobStatus = "queued" | "running" | "done" | "failed" | "zombie";