mlx-bun 0.2.0 → 0.4.0

package/src/jobs/runner.ts

@@ -1,3 +1,4 @@
+import { CancellationSource } from "../engine/cancellation";
 // Job runner: turns a registered `JobRunner` into a tracked, logged job.
 // Two execution modes — both create a row, both stream NDJSON events to the
 // log, both own terminal status:
@@ -14,6 +15,8 @@
 // Port of optiq lab jobs.py `submit` + `_job_main`, split across this file
 // and job-entry.ts.
 
+import type { DisposableResource } from "../contracts/resources";
+import { createJobTaskClient } from "./task-client";
 import { appendFileSync } from "node:fs";
 import { fileURLToPath } from "node:url";
 import type { Emit, JobEvent, JobRunner } from "./types";
@@ -104,18 +107,28 @@ export interface SubmitResult {
  *  on success sets 'done'/progress 1/ended_at + emits 'done', on throw sets
  *  'failed'/error/ended_at + emits 'failed'. Returns once the task is
  *  scheduled (does not await completion — poll the row / tail the log). */
+interface InProcessWork { cancellation: CancellationSource; done: Promise<void>; }
+const inProcess = new WeakMap<JobStore, Set<InProcessWork>>();
+
 export function submitInProcess(
   store: JobStore,
   kind: string,
   config: Record<string, unknown>,
   outputPath?: string,
 ): SubmitResult {
+  if (closedStores.has(store)) throw new Error("job host is closed");
   const row = store.create(kind, config, outputPath);
   const emit = makeEmit(store, row.id, row.log_path);
   store.setStatus(row.id, "running");
   emit({ type: "started", ts: Date.now() });
 
   const runner = getRunner(kind);
+  const cancellation = new CancellationSource();
+  const running = inProcess.get(store) ?? new Set<InProcessWork>();
+  inProcess.set(store, running);
+  let finish!: () => void;
+  const work = { cancellation, done: new Promise<void>((resolve) => { finish = resolve; }) };
+  running.add(work);
   void (async () => {
     if (!runner) {
       const error = `no runner registered for kind "${kind}"`;
@@ -124,7 +137,7 @@ export function submitInProcess(
       return;
     }
     try {
-      const result = await runner(emit, JSON.parse(row.config_json));
+      const result = await createJobTaskClient(runner).run(JSON.parse(row.config_json), emit, cancellation);
       const out = result?.outputPath;
       if (out) store.setOutputPath(row.id, out);
       store.setProgress(row.id, 1);
@@ -135,7 +148,7 @@ export function submitInProcess(
       store.setStatus(row.id, "failed", { error, endedAt: nowIso() });
       emit({ type: "failed", error, ts: Date.now() });
     }
-  })();
+  })().finally(() => { running.delete(work); finish(); });
 
   return { jobId: row.id, outputPath };
 }
@@ -145,16 +158,25 @@ export function submitInProcess(
 // ---------------------------------------------------------------------------
 
 interface QueuedSpawn {
+  acquire?: (signal: AbortSignal) => Promise<DisposableResource>;
+  lease?: DisposableResource;
+  abort: AbortController;
+  proc?: Bun.Subprocess;
+  finished: Promise<void>;
+  finish(): void;
   store: JobStore;
   jobId: string;
   entry: string;
   bin: string;
+  compiled: boolean;
   spawn: typeof Bun.spawn;
   onComplete?: (jobId: string, code: number) => void;
 }
 
 let gpuLeaseHolder: string | null = null;
 const spawnQueue: QueuedSpawn[] = [];
+let activeSpawn: QueuedSpawn | undefined;
+const closedStores = new WeakSet<JobStore>();
 
 /** True while a GPU (subprocess) job is running — the server gates token
  *  generation on this so quantize/finetune don't fight live inference. */
@@ -168,6 +190,8 @@ export function currentGpuJob(): string | null {
 }
 
 export interface SubprocessOpts {
+  /** Host execution lease, held until child exit (including crashes). */
+  acquire?: (signal: AbortSignal) => Promise<DisposableResource>;
   /** Override the entry script (tests). Defaults to JOB_ENTRY_PATH. */
   entry?: string;
   /** Override the runtime binary (tests). Defaults to "bun". */
@@ -189,14 +213,20 @@ export function submitSubprocess(
   outputPath?: string,
   opts: SubprocessOpts = {},
 ): SubmitResult {
+  if (closedStores.has(store)) throw new Error("job host is closed");
   const row = store.create(kind, config, outputPath);
+  let finish!: () => void;
+  const finished = new Promise<void>((resolve) => { finish = resolve; });
   const item: QueuedSpawn = {
+    abort: new AbortController(), finished, finish,
     store,
     jobId: row.id,
     entry: opts.entry ?? JOB_ENTRY_PATH,
+    compiled: opts.entry === undefined && opts.bin === undefined && JOB_ENTRY_PATH.includes("$bunfs"),
     bin: opts.bin ?? "bun",
     spawn: opts.spawn ?? Bun.spawn,
     onComplete: opts.onComplete,
+    acquire: opts.acquire,
   };
   spawnQueue.push(item);
   drainQueue();
@@ -215,7 +245,7 @@ function spawnNow(item: QueuedSpawn): void {
   // a file-backed DB.
   let proc: Bun.Subprocess<"ignore", "pipe", "pipe">;
   try {
-    proc = spawn([bin, entry, jobId], {
+    proc = spawn(item.compiled ? [process.execPath, "__job", jobId] : [bin, entry, jobId], {
       stdout: "pipe",
       stderr: "pipe",
       env: {
@@ -230,14 +260,16 @@ function spawnNow(item: QueuedSpawn): void {
         error: errString(e),
         endedAt: nowIso(),
       });
+    } catch (error) {
+      console.error(`[jobs] failed to record spawn failure: ${errString(error)}`);
     } finally {
-      releaseLease(jobId);
+      releaseLease(item);
     }
     return;
   }
 
-  const row = store.get(jobId);
-  const logPath = row?.log_path;
+  item.proc = proc;
+  const logPath = store.get(jobId)?.log_path;
   const logLine = (line: string) => {
     if (!line || !logPath) return;
     const ev: JobEvent = { type: "log", line };
@@ -252,7 +284,7 @@ function spawnNow(item: QueuedSpawn): void {
     // code 0 ⇒ trust the child's terminal status (it set done/failed itself).
     // non-zero ⇒ if the row never reached terminal (crash before the wrapper
     // could write), force it failed.
-    if (code !== 0) {
+    try { if (code !== 0) {
       const cur = store.get(jobId);
       if (cur && (cur.status === "queued" || cur.status === "running")) {
         store.setStatus(jobId, "failed", {
@@ -260,15 +292,25 @@ function spawnNow(item: QueuedSpawn): void {
           endedAt: nowIso(),
         });
       }
+    } } catch (error) {
+      console.error(`[jobs] failed to record child exit: ${errString(error)}`);
+    } finally {
+      releaseLease(item);
+      try { item.onComplete?.(jobId, code); } catch {}
+      drainQueue();
     }
-    releaseLease(jobId);
-    try { item.onComplete?.(jobId, code); } catch {}
-    drainQueue();
   })();
 }
 
-function releaseLease(jobId: string): void {
-  if (gpuLeaseHolder === jobId) gpuLeaseHolder = null;
+function releaseLease(item: QueuedSpawn): void {
+  try { item.lease?.dispose(); }
+  catch (error) { console.error(`[jobs] execution lease cleanup failed: ${errString(error)}`); }
+  finally {
+    item.lease = undefined;
+    if (gpuLeaseHolder === item.jobId) gpuLeaseHolder = null;
+    if (activeSpawn === item) activeSpawn = undefined;
+    item.finish();
+  }
 }
 
 /** Run the next queued subprocess job if the lease is free. */
@@ -279,8 +321,51 @@ export function drainQueue(): void {
     // A synchronous spawn failure releases the lease. Keep draining
     // iteratively so a run of bad queue entries cannot recurse or wedge the
     // first startable job behind them.
-    spawnNow(next);
+    activeSpawn = next;
+    if (next.acquire) {
+      gpuLeaseHolder = next.jobId; // reserve FIFO order while inference drains
+      void Promise.resolve().then(() => next.acquire!(next.abort.signal)).then((lease) => {
+        next.lease = lease;
+        if (next.abort.signal.aborted) throw next.abort.signal.reason;
+        spawnNow(next);
+      }).catch((error) => {
+        try {
+          next.store.setStatus(next.jobId, "failed", { error: errString(error), endedAt: nowIso() });
+        } catch (recordError) {
+          console.error(`[jobs] failed to record admission failure: ${errString(recordError)}`);
+        } finally { releaseLease(next); }
+      }).finally(drainQueue);
+    } else spawnNow(next);
+  }
+}
+
+/** Dataset work stays in-process, but belongs to the application host too. */
+export async function closeInProcessJobs(store: JobStore): Promise<void> {
+  closedStores.add(store);
+  const active = [...inProcess.get(store) ?? []];
+  for (const work of active) work.cancellation.cancel("engine_closed");
+  await Promise.all(active.map((work) => work.done));
+}
+
+/** Stop only this host's managed GPU jobs. Await process death before releasing
+ * the execution lease; queued and admission-waiting jobs never spawn. */
+export async function closeSubprocessJobs(store: JobStore): Promise<void> {
+  closedStores.add(store);
+  for (let i = spawnQueue.length - 1; i >= 0; i--) {
+    const item = spawnQueue[i]!;
+    if (item.store !== store) continue;
+    spawnQueue.splice(i, 1);
+    store.setStatus(item.jobId, "failed", { error: "job host closed", endedAt: nowIso() });
+    item.finish();
   }
+  const active = activeSpawn;
+  if (!active || active.store !== store) return;
+  active.abort.abort(new Error("job host closed"));
+  const proc = active.proc;
+  if (proc) proc.kill("SIGTERM");
+  const force = proc ? setTimeout(() => { if (proc.exitCode === null) proc.kill("SIGKILL"); }, 3000) : undefined;
+  try { await active.finished; }
+  finally { if (force) clearTimeout(force); }
 }
 
 /** Read a child stream line-by-line, buffering partial trailing lines, and

package/src/jobs/task-client.ts

@@ -0,0 +1,19 @@
+import type { TaskClient } from "../contracts/completion";
+import { throwIfCancelled } from "../engine/cancellation";
+import type { JobEvent, JobRunner } from "./types";
+
+/** Numerical runners retain their artifact policy; the client owns delivery
+ * and cooperative cancellation at progress boundaries. No event buffering. */
+export function createJobTaskClient(runner: JobRunner): TaskClient<
+  Record<string, unknown>, JobEvent, { outputPath?: string } | void
+> {
+  return {
+    async run(request, report, cancellation) {
+      const check = () => { if (cancellation) throwIfCancelled(cancellation); };
+      check();
+      const result = await runner((event) => { check(); report(event); check(); }, request);
+      check();
+      return result;
+    },
+  };
+}

package/src/jobs/types.ts

@@ -1,58 +1,2 @@
-// 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 =
-  | { type: "started"; ts: number }
-  | { type: "log"; line: string }
-  | {
-      type: "stage";
-      stage: string;
-      progress?: number; // 0..1
-      message?: string;
-      [k: string]: unknown;
-    }
-  | {
-      type: "metric";
-      kind: "train" | "val";
-      step: number;
-      loss: number;
-      [k: string]: unknown; // grad_norm, learning_rate, tokens_per_sec, accuracy, margin, progress, message
-    }
-  | { type: "done"; ts: number; output_dir?: string; summary?: unknown; [k: string]: unknown }
-  | { type: "failed"; error: string; ts: number };
-
-/** 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. */
-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. */
-export type JobRunner = (
-  emit: Emit,
-  config: Record<string, unknown>,
-) => Promise<{ outputPath?: string } | void>;
-
-/** Job kinds in the system. */
-export type JobKind = "quantize" | "finetune" | "dataset" | "noop" | "crash";
-
-export type JobStatus = "queued" | "running" | "done" | "failed" | "zombie";
-
-/** A persisted job row (shape returned by the HTTP `GET /api/jobs/:id`). */
-export interface JobRow {
-  id: string;
-  kind: string;
-  status: JobStatus;
-  config_json: string;
-  progress: number;
-  message: string | null;
-  log_path: string;
-  output_path: string | null;
-  error: string | null;
-  started_at: string;
-  ended_at: string | null;
-}
+/** Compatibility exports; producer and browser share the portable wire contract. */
+export type * from "../contracts/jobs";

package/src/kv-scheme.ts

@@ -86,6 +86,42 @@ export function kvBytesAt(
     value.linearStateBytes;
 }
 
+/** Head dims MLX 0.32.2's fused multi-query SDPA kernel serves for L > 8
+ *  queries on M1–M4 GPUs (mlx/backend/metal/scaled_dot_product_attention.cpp
+ *  `use_fallback`: 192 and 256 always fall back off-NAX; `has_fused_kernel`
+ *  lists 64/72/80/96/128 for the full kernel). Every other head dim runs
+ *  mlx/fast.cpp's composed fallback, which materializes the bf16 scores
+ *  tensor [heads, L, S] plus a bool mask per full-attention layer. The
+ *  decode kernel (L ≤ 8) fuses 256, so this term is a prefill-only cost. */
+export const FUSED_SDPA_HEAD_DIMS: ReadonlySet<number> = new Set([64, 72, 80, 96, 128]);
+const SDPA_SCORES_BYTES = 2; // bf16 scores in q's dtype
+const SDPA_FUSED_MAX_QUERIES = 8;
+
+/** Peak bytes of the SDPA-fallback scores tensor live during one prefill
+ *  chunk of `chunk` queries against a `ctx`-token context: the last chunk
+ *  sees the longest keys, and the lazy graph keeps one layer's scores live
+ *  at a time. Zero when every attention layer's head dim fuses. Sliding
+ *  layers see at most `window + L` keys. Admission-only (fit and the
+ *  in-flight memory guard); nothing here changes what is computed. */
+export function sdpaFallbackBytes(config: ModelConfig, chunk: number, ctx: number): number {
+  const text = config.text;
+  const heads = text.numAttentionHeads || 0;
+  const queries = Math.min(chunk, ctx);
+  if (heads <= 0 || queries <= SDPA_FUSED_MAX_QUERIES) return 0;
+  let peak = 0;
+  const layers = text.numHiddenLayers || text.layerTypes.length;
+  for (let layer = 0; layer < layers; layer++) {
+    const type = text.layerTypes[layer] ?? "full_attention";
+    if (type === "linear_attention") continue;
+    const sliding = type === "sliding_attention" && text.slidingWindow > 0;
+    const headDim = sliding ? text.headDim : text.globalHeadDim;
+    if (FUSED_SDPA_HEAD_DIMS.has(headDim)) continue;
+    const keys = sliding ? Math.min(ctx, text.slidingWindow + queries) : ctx;
+    peak = Math.max(peak, heads * queries * keys * SDPA_SCORES_BYTES);
+  }
+  return peak;
+}
+
 export class KvScheme {
   readonly kind: KvSchemeKind;
   readonly options: Readonly<ResolvedKvSchemeOptions>;
@@ -112,10 +148,17 @@ export class KvScheme {
   get cacheKey(): string {
     if (this.kind === "turbo") {
       const turbo = this.options.turboQuant!;
-      return `turbo-k${turbo.kBits}v${turbo.vBits}`;
+      const start = this.options.quantizedKvStart ?? 0;
+      return `turbo-k${turbo.kBits}v${turbo.vBits}${start > 0 ? `-start${start}` : ""}`;
+    }
+    if (this.kind === "affine-uniform") {
+      const start = this.options.quantizedKvStart ?? 5000, group = this.options.kvGroupSize ?? 64;
+      return `kv${this.options.kvBits}${group !== 64 ? `-g${group}` : ""}${start > 0 ? `-start${start}` : ""}`;
+    }
+    if (this.kind === "affine-config") {
+      const start = this.options.quantizedKvStart ?? 0;
+      return `config${start > 0 ? `-start${start}` : ""}`;
     }
-    if (this.kind === "affine-uniform") return `kv${this.options.kvBits}`;
-    if (this.kind === "affine-config") return "config";
     return "bf16";
   }
 
@@ -154,14 +197,24 @@ export class KvScheme {
     return this.kind === "turbo" ? {} : this.generationOptions;
   }
 
-  /** The current batch engine accepts only per-layer affine schemes whose
-   * named cache instances can perform the required conversion. The capability
-   * probe is mandatory for quantized schemes so config shape alone can never
-   * authorize placement. */
-  batchable(config: ModelConfig, canConvert?: (layerIdx: number) => boolean): boolean {
+  /** Affine layouts convert during solo prefill. TQ also has a transitional
+   * row layout that preserves each request's delayed conversion boundary. */
+  batchable(config: ModelConfig, canConvert?: (layerIdx: number) => boolean,
+    cacheLayerCount = config.text.numHiddenLayers, capabilities: { delayedAffine?: boolean } = {}): boolean {
     if (this.kind === "bf16") return true;
-    if (this.kind !== "affine-config") return false;
+    if (this.kind !== "affine-config" && this.kind !== "affine-uniform" && this.kind !== "turbo") return false;
     if (!canConvert) return false;
+    const start = this.options.quantizedKvStart ?? (this.kind === "affine-uniform" ? 5000 : 0);
+    if (start > 0 && this.kind !== "turbo" && !capabilities.delayedAffine) return false;
+    if (this.kind === "affine-uniform" || this.kind === "turbo") {
+      // Some models share the donor prefix's KV in later layers. Probe the
+      // actual cache list, just as uniform conversion does.
+      for (let layer = 0; layer < cacheLayerCount; layer++) {
+        if (config.text.layerTypes[layer] === "linear_attention") continue;
+        if (!canConvert(layer)) return false;
+      }
+      return true;
+    }
     return this.options.kvConfig!.every((entry) => {
       if (entry.layerIdx < 0 || entry.layerIdx >= config.text.numHiddenLayers) return false;
       return (config.text.layerTypes[entry.layerIdx] ?? "full_attention") !== "linear_attention" &&
@@ -175,11 +228,15 @@ export function resolveKvScheme(input: {
   turboQuant?: TurboQuantScheme;
   config?: readonly KvQuantSpec[] | null;
   missingConfig?: "bf16" | "error";
+  quantizedKvStart?: number;
 }): KvScheme {
+  const start = input.quantizedKvStart;
+  if (start !== undefined && (!Number.isSafeInteger(start) || start < 0))
+    throw new Error("quantizedKvStart must be a nonnegative integer");
   if (input.turboQuant) {
     return new KvScheme("turbo", {
       turboQuant: input.turboQuant,
-      quantizedKvStart: 0,
+      quantizedKvStart: start ?? 0,
     });
   }
   if (input.override === "off" || input.override === undefined)
@@ -190,10 +247,11 @@ export function resolveKvScheme(input: {
         throw new Error("model has no kv_config.json (--kv-quant config)");
       return new KvScheme("bf16", {});
     }
-    return new KvScheme("affine-config", { kvConfig: [...input.config] });
+    return new KvScheme("affine-config", { kvConfig: [...input.config],
+      ...(start === undefined ? {} : { quantizedKvStart: start }) });
   }
   return new KvScheme("affine-uniform", {
     kvBits: input.override,
-    quantizedKvStart: 0,
+    quantizedKvStart: start ?? 0,
   });
 }