mlx-bun 0.0.13 → 0.3.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

@@ -0,0 +1,199 @@
+import type { KvQuantSpec, ModelConfig, TurboQuantScheme } from "./config";
+
+const BF16_BYTES = 2;
+
+export type KvQuantOverride = "off" | "config" | number | undefined;
+
+export interface KvSchemeOptions {
+  kvBits?: number;
+  kvGroupSize?: number;
+  kvConfig?: KvQuantSpec[];
+  quantizedKvStart?: number;
+  turboQuant?: TurboQuantScheme;
+}
+
+interface ResolvedKvSchemeOptions extends Omit<KvSchemeOptions, "kvConfig" | "turboQuant"> {
+  readonly kvConfig?: readonly Readonly<KvQuantSpec>[];
+  readonly turboQuant?: Readonly<TurboQuantScheme>;
+}
+
+export type KvSchemeKind = "bf16" | "affine-uniform" | "affine-config" | "turbo";
+
+export interface KvGeometry {
+  fullBytesPerToken: number;
+  slidingBytesPerToken: number;
+  linearStateBytes: number;
+  window: number;
+}
+
+/** Bytes per stored affine-quantized KV element, including scale and bias. */
+export function kvQuantBytesPerElement(bits: number, groupSize: number): number {
+  return bits / 8 + 4 / groupSize;
+}
+
+export function kvGeometry(config: ModelConfig, options: KvSchemeOptions = {}): KvGeometry {
+  const text = config.text;
+  const byLayer = options.kvConfig?.length
+    ? new Map(options.kvConfig.map((entry) => [entry.layerIdx, entry]))
+    : null;
+  const uniformBytes = options.kvBits
+    ? kvQuantBytesPerElement(options.kvBits, options.kvGroupSize ?? 64)
+    : BF16_BYTES;
+  const elementBytes = (layer: number): number => {
+    const entry = byLayer?.get(layer);
+    return entry
+      ? kvQuantBytesPerElement(entry.bits, entry.groupSize)
+      : byLayer ? BF16_BYTES : uniformBytes;
+  };
+
+  let linearLayers = 0;
+  let fullBytesPerToken = 0;
+  let slidingBytesPerToken = 0;
+  for (let layer = 0; layer < text.numHiddenLayers; layer++) {
+    const type = text.layerTypes[layer] ?? "full_attention";
+    if (type === "linear_attention") {
+      linearLayers++;
+    } else if (type === "sliding_attention") {
+      slidingBytesPerToken +=
+        2 * text.numKeyValueHeads * text.headDim * elementBytes(layer);
+    } else {
+      fullBytesPerToken +=
+        2 * text.numGlobalKeyValueHeads * text.globalHeadDim * elementBytes(layer);
+    }
+  }
+  const convDim = 2 * text.linearNumKeyHeads * text.linearKeyHeadDim +
+    text.linearNumValueHeads * text.linearValueHeadDim;
+  const linearStateBytes = linearLayers === 0 ? 0 : linearLayers * (
+    text.linearNumValueHeads * text.linearValueHeadDim * text.linearKeyHeadDim * 4 +
+    Math.max(0, text.linearConvKernelDim - 1) * convDim * BF16_BYTES
+  );
+  return {
+    fullBytesPerToken,
+    slidingBytesPerToken,
+    linearStateBytes,
+    window: text.slidingWindow,
+  };
+}
+
+export function kvBytesAt(
+  config: ModelConfig,
+  tokens: number,
+  options: KvSchemeOptions = {},
+): number {
+  const value = kvGeometry(config, options);
+  return value.fullBytesPerToken * tokens +
+    value.slidingBytesPerToken * Math.min(tokens, value.window) +
+    value.linearStateBytes;
+}
+
+export class KvScheme {
+  readonly kind: KvSchemeKind;
+  readonly options: Readonly<ResolvedKvSchemeOptions>;
+
+  constructor(kind: KvSchemeKind, options: KvSchemeOptions) {
+    this.kind = kind;
+    const kvConfig = options.kvConfig?.map((entry) =>
+      Object.freeze({ ...entry }),
+    );
+    const turboQuant = options.turboQuant
+      ? Object.freeze({ ...options.turboQuant })
+      : undefined;
+    this.options = Object.freeze({
+      ...options,
+      ...(kvConfig ? { kvConfig: Object.freeze(kvConfig) } : {}),
+      ...(turboQuant ? { turboQuant } : {}),
+    });
+  }
+
+  get quantized(): boolean {
+    return this.kind !== "bf16";
+  }
+
+  get cacheKey(): string {
+    if (this.kind === "turbo") {
+      const turbo = this.options.turboQuant!;
+      return `turbo-k${turbo.kBits}v${turbo.vBits}`;
+    }
+    if (this.kind === "affine-uniform") return `kv${this.options.kvBits}`;
+    if (this.kind === "affine-config") return "config";
+    return "bf16";
+  }
+
+  get label(): string {
+    if (this.kind === "turbo") {
+      const turbo = this.options.turboQuant!;
+      return `turbo k${turbo.kBits}v${turbo.vBits}`;
+    }
+    if (this.kind === "affine-uniform") return `uniform-kv${this.options.kvBits}`;
+    if (this.kind === "affine-config") return "mixed (kv_config.json)";
+    return "bf16";
+  }
+
+  bytesAt(config: ModelConfig, tokens: number): number {
+    // TurboQuant stays conservatively billed as bf16 until its packed layout
+    // exposes a stable projector.
+    return kvBytesAt(config, tokens, this.fitOptions);
+  }
+
+  /** Mutable compatibility value for the numerical generator. The scheme
+   *  keeps its declaration immutable and hands legacy consumers owned copies. */
+  get generationOptions(): KvSchemeOptions {
+    const { kvConfig, turboQuant, ...scalar } = this.options;
+    return {
+      ...scalar,
+      ...(kvConfig
+        ? { kvConfig: kvConfig.map((entry) => ({ ...entry })) }
+        : {}),
+      ...(turboQuant
+        ? { turboQuant: { ...turboQuant } }
+        : {}),
+    };
+  }
+
+  get fitOptions(): KvSchemeOptions {
+    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 {
+    if (this.kind === "bf16") return true;
+    if (this.kind !== "affine-config") return false;
+    if (!canConvert) return false;
+    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" &&
+        canConvert(entry.layerIdx);
+    });
+  }
+}
+
+export function resolveKvScheme(input: {
+  override?: KvQuantOverride;
+  turboQuant?: TurboQuantScheme;
+  config?: readonly KvQuantSpec[] | null;
+  missingConfig?: "bf16" | "error";
+}): KvScheme {
+  if (input.turboQuant) {
+    return new KvScheme("turbo", {
+      turboQuant: input.turboQuant,
+      quantizedKvStart: 0,
+    });
+  }
+  if (input.override === "off" || input.override === undefined)
+    return new KvScheme("bf16", {});
+  if (input.override === "config") {
+    if (!input.config?.length) {
+      if (input.missingConfig === "error")
+        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-uniform", {
+    kvBits: input.override,
+    quantizedKvStart: 0,
+  });
+}