mlx-bun 0.0.13 → 0.2.0

package/src/serve/lab-routes.ts

@@ -0,0 +1,338 @@
+import type { JobStore } from "../jobs";
+
+export interface LabRouteDeps {
+  ensureJobs(): Promise<JobStore>;
+  serverPort(): number | undefined;
+  invalidateLibrary(): void;
+}
+
+export type LabRoute =
+  | "dataset-templates"
+  | "dataset-submit"
+  | "quantize-inspect"
+  | "model-resolve-folder"
+  | "quantize-submit"
+  | "finetune-inspect-dataset"
+  | "finetune-submit"
+  | "finetune-merge"
+  | "finetune-export";
+
+export function matchLabRoute(method: string, pathname: string): LabRoute | null {
+  switch (`${method} ${pathname}`) {
+    case "GET /api/dataset/templates": return "dataset-templates";
+    case "POST /api/dataset/submit": return "dataset-submit";
+    case "POST /api/quantize/inspect": return "quantize-inspect";
+    case "POST /api/quantize/resolve-folder":
+    case "POST /api/model/resolve-folder":
+      return "model-resolve-folder";
+    case "POST /api/quantize/submit": return "quantize-submit";
+    case "POST /api/finetune/inspect-dataset": return "finetune-inspect-dataset";
+    case "POST /api/finetune/submit": return "finetune-submit";
+    case "POST /api/finetune/merge": return "finetune-merge";
+    case "POST /api/finetune/export": return "finetune-export";
+    default: return null;
+  }
+}
+
+export async function handleLabRoute(
+  url: URL,
+  request: Request,
+  deps: LabRouteDeps,
+): Promise<Response | null> {
+  switch (matchLabRoute(request.method, url.pathname)) {
+    case "dataset-templates": {
+      const { TEMPLATES } = await import("../dataset");
+      return Response.json({ templates: TEMPLATES });
+    }
+    case "dataset-submit": {
+      const body = (await request.json().catch(() => ({}))) as {
+        template_id?: string;
+        inputs?: Record<string, unknown>;
+        model_name?: string;
+      };
+      const { getTemplate } = await import("../dataset");
+      if (!body.template_id || !getTemplate(body.template_id)) {
+        return Response.json(
+          { ok: false, error: `unknown template ${body.template_id}` },
+          { status: 400 },
+        );
+      }
+      const store = await deps.ensureJobs();
+      const { submitInProcess } = await import("../jobs");
+      const { homedir } = await import("node:os");
+      const safe = body.template_id.replace(/[^a-z0-9_-]/gi, "");
+      const outDir = `${homedir()}/.cache/mlx-bun/datasets/dataset-${safe}-${Date.now()}`;
+      const { jobId } = submitInProcess(store, "dataset", {
+        template_id: body.template_id,
+        inputs: body.inputs ?? {},
+        output_dir: outDir,
+        api_url: `http://127.0.0.1:${deps.serverPort()}`,
+        model_name: body.model_name ?? "local",
+      }, outDir);
+      return Response.json({ ok: true, job_id: jobId, output_dir: outDir });
+    }
+    case "quantize-inspect": {
+      const body = (await request.json().catch(() => ({}))) as { model_id?: string };
+      const { inspectModel } = await import("../quantize");
+      return Response.json(await inspectModel(body.model_id ?? ""));
+    }
+    case "model-resolve-folder":
+      return resolveModelFolder(request);
+    case "quantize-submit":
+      return submitQuantize(request, deps);
+    case "finetune-inspect-dataset": {
+      const body = (await request.json().catch(() => ({}))) as { path?: string };
+      const { inspectDataset } = await import("../train");
+      return Response.json(await inspectDataset(body.path ?? ""));
+    }
+    case "finetune-submit": {
+      const body = (await request.json().catch(() => ({}))) as Record<string, unknown>;
+      if (!body.model_dir || !body.data_dir) {
+        return Response.json(
+          { ok: false, error: "model_dir and data_dir required" },
+          { status: 400 },
+        );
+      }
+      const store = await deps.ensureJobs();
+      const { submitSubprocess } = await import("../jobs");
+      const { homedir } = await import("node:os");
+      const adapterPath = (body.adapter_path as string) ||
+        `${homedir()}/.cache/mlx-bun/adapters/adapter-${Date.now()}`;
+      const { jobId } = submitSubprocess(
+        store,
+        "finetune",
+        { ...body, adapter_path: adapterPath },
+        adapterPath,
+      );
+      return Response.json({ ok: true, job_id: jobId, adapter_path: adapterPath });
+    }
+    case "finetune-merge": {
+      const body = (await request.json().catch(() => ({}))) as {
+        adapter_a?: string;
+        adapter_b?: string;
+        scales?: number[];
+      };
+      if (!body.adapter_a || !body.adapter_b) {
+        return Response.json(
+          { ok: false, error: "adapter_a and adapter_b required" },
+          { status: 400 },
+        );
+      }
+      try {
+        const { mergeAdapters } = await import("../train");
+        const { homedir } = await import("node:os");
+        const mergedPath = `${homedir()}/.cache/mlx-bun/adapters/merged-${Date.now()}`;
+        const stats = await mergeAdapters(
+          [body.adapter_a, body.adapter_b],
+          mergedPath,
+          body.scales,
+        );
+        return Response.json({ ok: true, merged_path: mergedPath, stats });
+      } catch (error) {
+        return Response.json({ ok: false, error: (error as Error).message }, { status: 400 });
+      }
+    }
+    case "finetune-export": {
+      const body = (await request.json().catch(() => ({}))) as {
+        base_model?: string;
+        adapter_path?: string;
+        method?: string;
+      };
+      if (!body.base_model || !body.adapter_path) {
+        return Response.json(
+          { ok: false, error: "base_model and adapter_path required" },
+          { status: 400 },
+        );
+      }
+      try {
+        const { exportAdapter } = await import("../train");
+        const { homedir } = await import("node:os");
+        const exportPath = `${homedir()}/.cache/mlx-bun/exports/export-${Date.now()}`;
+        const manifest = await exportAdapter(
+          exportPath,
+          body.base_model,
+          body.adapter_path,
+          body.method,
+        );
+        return Response.json({ ok: true, export_path: exportPath, manifest });
+      } catch (error) {
+        return Response.json({ ok: false, error: (error as Error).message }, { status: 400 });
+      }
+    }
+    default:
+      return null;
+  }
+}
+
+async function resolveModelFolder(request: Request): Promise<Response> {
+  const body = (await request.json().catch(() => ({}))) as {
+    folder_name?: string;
+    rel_path?: string;
+  };
+  const { statSync, readdirSync, readFileSync } = await import("node:fs");
+  const { join } = await import("node:path");
+  const { homedir } = await import("node:os");
+  const hubRoot = process.env.HF_HUB_CACHE ??
+    (process.env.HF_HOME
+      ? join(process.env.HF_HOME, "hub")
+      : join(homedir(), ".cache/huggingface/hub"));
+  const roots = [hubRoot, join(homedir(), ".cache/mlx-bun")];
+  const hasConfig = (dir: string) => {
+    try {
+      return statSync(join(dir, "config.json")).isFile();
+    } catch {
+      return false;
+    }
+  };
+  const folder = (body.folder_name ?? "")
+    .replace(/\\/g, "/")
+    .split("/")
+    .filter(Boolean)
+    .pop() ?? "";
+  const relSegments = (body.rel_path ?? "")
+    .replace(/\\/g, "/")
+    .split("/")
+    .filter(Boolean);
+  const configDir = relSegments.length >= 2 ? relSegments[relSegments.length - 2] : "";
+  const repoIdOf = (modelsDir: string) =>
+    modelsDir.slice("models--".length).replaceAll("--", "/");
+  const pickSnapshot = (repoDir: string): string | null => {
+    const snapshots = join(repoDir, "snapshots");
+    let head = "";
+    try {
+      head = readFileSync(join(repoDir, "refs", "main"), "utf8").trim();
+    } catch {}
+    if (head && hasConfig(join(snapshots, head))) return join(snapshots, head);
+    try {
+      for (const hash of readdirSync(snapshots)) {
+        if (hasConfig(join(snapshots, hash))) return join(snapshots, hash);
+      }
+    } catch {}
+    return null;
+  };
+
+  if (folder.startsWith("models--")) {
+    const path = pickSnapshot(join(hubRoot, folder));
+    if (path) return Response.json({ ok: true, path, repo_id: repoIdOf(folder) });
+  }
+
+  const hashDir = configDir || folder;
+  if (hashDir) {
+    try {
+      for (const repo of readdirSync(hubRoot)) {
+        if (!repo.startsWith("models--")) continue;
+        const candidate = join(hubRoot, repo, "snapshots", hashDir);
+        if (hasConfig(candidate)) {
+          return Response.json({ ok: true, path: candidate, repo_id: repoIdOf(repo) });
+        }
+      }
+    } catch {}
+  }
+
+  for (const root of roots) {
+    if (folder && hasConfig(join(root, folder))) {
+      return Response.json({ ok: true, path: join(root, folder) });
+    }
+  }
+
+  const { Registry } = await import("../registry");
+  const registry = new Registry();
+  try {
+    await registry.scan();
+    const all = registry.list();
+    const record = (folder.startsWith("models--")
+      ? all.find((model) => model.repoId === repoIdOf(folder))
+      : undefined) ??
+      all.find((model) => model.path.split("/").pop() === hashDir) ??
+      all.find((model) => model.repoId.split("/").pop() === folder);
+    if (record) {
+      return Response.json({
+        ok: true,
+        path: record.path,
+        repo_id: record.repoId,
+        model_type: record.modelType,
+      });
+    }
+  } finally {
+    registry.close();
+  }
+  return Response.json({
+    ok: false,
+    error: "Couldn't locate this folder on disk — paste the path instead.",
+  });
+}
+
+async function submitQuantize(request: Request, deps: LabRouteDeps): Promise<Response> {
+  const body = (await request.json().catch(() => ({}))) as {
+    model_id?: string;
+    bits?: number;
+    group_size?: number;
+    target_bpw?: number;
+    candidate_bits?: number[];
+    reference?: string;
+    calibration_mix?: string;
+    n_calibration?: number;
+    rotate_weights?: boolean;
+    rotation_seed?: number;
+  };
+  if (!body.model_id)
+    return Response.json({ ok: false, error: "model_id required" }, { status: 400 });
+  const store = await deps.ensureJobs();
+  const { submitSubprocess } = await import("../jobs");
+  const { homedir } = await import("node:os");
+  const { join } = await import("node:path");
+  const { mkdirSync, writeFileSync } = await import("node:fs");
+  const { createHash } = await import("node:crypto");
+  const bits = body.bits ?? 4;
+  const groupSize = body.group_size ?? 64;
+  const snapshotMatch = body.model_id.match(/(models--[^/]+)\/snapshots\//);
+  let org = "local";
+  let name: string;
+  if (snapshotMatch) {
+    const parts = snapshotMatch[1]!.split("--");
+    org = parts[1] ?? "local";
+    name = parts.slice(2).join("--");
+  } else if (
+    body.model_id.includes("/") &&
+    !body.model_id.startsWith("/") &&
+    !body.model_id.startsWith("~")
+  ) {
+    const segments = body.model_id.split("/");
+    org = segments[0]!;
+    name = segments.slice(1).join("-");
+  } else {
+    name = body.model_id.split("/").filter(Boolean).at(-1) ?? "model";
+  }
+  name = (name || "model").replace(/[^a-z0-9_.-]/gi, "");
+  org = (org || "local").replace(/[^a-z0-9_.-]/gi, "");
+  const suffix = `${body.target_bpw ? `mixed-${body.target_bpw}bpw` : `${bits}bit`}` +
+    `${body.rotate_weights ? `-rot${body.rotation_seed ?? 42}` : ""}`;
+  const quantRepo = `${name}-OptiQ-${suffix}`;
+  const hubRoot = process.env.HF_HUB_CACHE ??
+    (process.env.HF_HOME
+      ? join(process.env.HF_HOME, "hub")
+      : join(homedir(), ".cache/huggingface/hub"));
+  const repoDir = join(hubRoot, `models--${org}--${quantRepo}`);
+  const snapshotHash = createHash("sha1").update(`${org}/${quantRepo}`).digest("hex");
+  const outDir = join(repoDir, "snapshots", snapshotHash);
+  try {
+    mkdirSync(join(repoDir, "refs"), { recursive: true });
+    writeFileSync(join(repoDir, "refs", "main"), snapshotHash);
+  } catch {}
+  const { jobId } = submitSubprocess(store, "quantize", {
+    model_id: body.model_id,
+    out_dir: outDir,
+    bits,
+    group_size: groupSize,
+    target_bpw: body.target_bpw,
+    candidate_bits: body.candidate_bits,
+    reference: body.reference,
+    calibration_mix: body.calibration_mix,
+    n_calibration: body.n_calibration,
+    rotate_weights: body.rotate_weights,
+    rotation_seed: body.rotation_seed,
+  }, outDir, {
+    onComplete: () => deps.invalidateLibrary(),
+  });
+  return Response.json({ ok: true, job_id: jobId, output_dir: outDir });
+}

package/src/serve/model-admin-routes.ts

@@ -0,0 +1,153 @@
+import type { ServerContext } from "../server";
+import type { GenerationGateway } from "./generation-gateway";
+import { embedMany, isEmbeddingModel } from "../embed";
+import { listAvailableAdapters } from "../lora";
+
+export type ModelAdminRoute =
+  | { kind: "embeddings" }
+  | { kind: "adapters-available" }
+  | { kind: "adapters-list" }
+  | { kind: "adapters-mount" }
+  | { kind: "adapters-unmount"; id: string };
+
+export function matchModelAdminRoute(method: string, pathname: string): ModelAdminRoute | null {
+  switch (`${method} ${pathname}`) {
+    case "POST /v1/embeddings": return { kind: "embeddings" };
+    case "GET /v1/adapters/available": return { kind: "adapters-available" };
+    case "GET /v1/adapters": return { kind: "adapters-list" };
+    case "POST /v1/adapters": return { kind: "adapters-mount" };
+  }
+  if (method === "DELETE" && pathname.startsWith("/v1/adapters/")) {
+    return {
+      kind: "adapters-unmount",
+      id: decodeURIComponent(pathname.slice("/v1/adapters/".length)),
+    };
+  }
+  return null;
+}
+
+export async function handleModelAdminRoute(
+  url: URL,
+  request: Request,
+  ctx: ServerContext,
+  gateway: Pick<GenerationGateway, "runExclusive">,
+): Promise<Response | null> {
+  const route = matchModelAdminRoute(request.method, url.pathname);
+  if (!route) return null;
+
+  switch (route.kind) {
+    case "embeddings": {
+      if (!isEmbeddingModel(ctx.model)) {
+        return Response.json({
+          error: {
+            message: `served model "${ctx.modelId}" is not an embedding model; ` +
+              "serve an embedding model (e.g. Qwen3-Embedding) to use /v1/embeddings",
+            type: "invalid_request_error",
+          },
+        }, { status: 400 });
+      }
+      let body: { input?: string | string[]; instruction?: string };
+      try {
+        body = (await request.json()) as typeof body;
+      } catch {
+        return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
+      }
+      const inputs = Array.isArray(body.input)
+        ? body.input
+        : body.input != null
+          ? [body.input]
+          : [];
+      if (inputs.length === 0 || !inputs.every((input) => typeof input === "string")) {
+        return Response.json({
+          error: {
+            message: "`input` must be a string or array of strings",
+            type: "invalid_request_error",
+          },
+        }, { status: 400 });
+      }
+      const instruction = typeof body.instruction === "string" ? body.instruction : undefined;
+      const results = embedMany(ctx.model, ctx.tokenizer, inputs, instruction);
+      let totalTokens = 0;
+      const data = results.map((result, index) => {
+        totalTokens += result.tokens;
+        return { object: "embedding", index, embedding: Array.from(result.vector) };
+      });
+      return Response.json({
+        object: "list",
+        data,
+        model: ctx.modelId,
+        usage: { prompt_tokens: totalTokens, total_tokens: totalTokens },
+      });
+    }
+
+    case "adapters-available": {
+      const { homedir } = await import("node:os");
+      const stores = [
+        `${homedir()}/.cache/mlx-bun-finetunes`,
+        `${homedir()}/.cache/mlx-bun/adapters`,
+      ];
+      const mounted = new Set(ctx.adapters.list().map((adapter) => adapter.id));
+      const bareName = (value: string) => value.split("/").pop()!.toLowerCase();
+      const servedName = bareName(ctx.modelId);
+      const adapters = (await listAvailableAdapters(stores)).map((adapter) => ({
+        id: adapter.id,
+        path: adapter.path,
+        rank: adapter.rank,
+        scale: adapter.scale,
+        base_model: adapter.baseModel,
+        mounted: mounted.has(adapter.id),
+        compatible: adapter.baseModel == null || bareName(adapter.baseModel) === servedName,
+      }));
+      return Response.json({ adapters });
+    }
+
+    case "adapters-list":
+      return Response.json({
+        adapters: ctx.adapters.list().map((adapter) => ({
+          id: adapter.id,
+          path: adapter.path,
+          rank: adapter.rank,
+          scale: adapter.scale,
+          size_bytes: adapter.sizeBytes,
+          mounted_layers: adapter.mountedLayers,
+          ram_bytes: adapter.ramBytes,
+        })),
+      });
+
+    case "adapters-mount": {
+      let body: { id?: string; path?: string };
+      try {
+        body = (await request.json()) as typeof body;
+      } catch {
+        return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
+      }
+      if (!body.id || !body.path) {
+        return Response.json({ error: { message: "id and path required" } }, { status: 400 });
+      }
+      try {
+        const info = await gateway.runExclusive(() =>
+          ctx.adapters.mount(body.id!, body.path!),
+        );
+        return Response.json({
+          id: info.id,
+          mounted_layers: info.mountedLayers,
+          rank: info.rank,
+          scale: info.scale,
+          ram_bytes: info.ramBytes,
+        });
+      } catch (error) {
+        return Response.json({ error: { message: (error as Error).message } }, { status: 400 });
+      }
+    }
+
+    case "adapters-unmount": {
+      const removed = await gateway.runExclusive(async () => ctx.adapters.unmount(route.id));
+      return removed > 0
+        ? Response.json({ id: route.id, removed_layers: removed })
+        : Response.json(
+            { error: { message: `adapter ${route.id} not mounted` } },
+            { status: 404 },
+          );
+    }
+  }
+}

package/src/serve/prompt-response-trace.ts

@@ -0,0 +1,147 @@
+import { runtimeValue } from "../runtime-config";
+
+export const P2R_TRACE_PREFIX = "[p2r] ";
+export const P2R_TRACE_VERSION = 1;
+
+export type P2RTraceOutcome = "success" | "error" | "abort";
+
+export type P2RTracePhase =
+  | "request.body_parse"
+  | "request.prompt_prepare"
+  | "completion.total"
+  | "completion.placement"
+  | "engine.admission_wait"
+  | "cache.lookup_restore"
+  | "prefill.total"
+  | "prefill.batch_setup"
+  | "prefill.chunk"
+  | "prefill.kv_maintenance"
+  | "token_zero.total"
+  | "token_zero.forward"
+  | "token_zero.head"
+  | "token_zero.sample"
+  | "response.first_write"
+  | "response.final_write";
+
+export type P2RTraceAttributes = Readonly<Record<
+  string,
+  string | number | boolean | null
+>>;
+
+export interface P2RTraceEvent {
+  phase: P2RTracePhase;
+  startMs: number;
+  durationMs: number;
+  attributes?: P2RTraceAttributes;
+}
+
+export interface P2RTraceRecord {
+  version: typeof P2R_TRACE_VERSION;
+  traceId: string;
+  requestId: string;
+  route: string;
+  clock: "monotonic-ms";
+  outcome: P2RTraceOutcome;
+  totalMs: number;
+  events: P2RTraceEvent[];
+  attributes?: P2RTraceAttributes;
+}
+
+type Clock = () => number;
+type Emit = (record: P2RTraceRecord) => void;
+
+export interface PromptResponseTraceOptions {
+  traceId: string;
+  requestId: string;
+  route: string;
+  clock?: Clock;
+  emit?: Emit;
+}
+
+/** Request-local diagnostic trace. Callers pass it explicitly because the
+ * continuous scheduler interleaves rows and cannot safely use ambient async
+ * context. No instance is created unless MLX_BUN_P2R_TRACE=1. */
+export class PromptResponseTrace {
+  readonly #startedAt: number;
+  readonly #events: P2RTraceEvent[] = [];
+  readonly #clock: Clock;
+  readonly #emit: Emit;
+  #finished = false;
+
+  readonly traceId: string;
+  readonly requestId: string;
+  readonly route: string;
+
+  constructor(options: PromptResponseTraceOptions) {
+    this.traceId = options.traceId;
+    this.requestId = options.requestId;
+    this.route = options.route;
+    this.#clock = options.clock ?? performance.now.bind(performance);
+    this.#emit = options.emit ?? ((record) => {
+      console.error(`${P2R_TRACE_PREFIX}${JSON.stringify(record)}`);
+    });
+    this.#startedAt = this.#clock();
+  }
+
+  get finished(): boolean {
+    return this.#finished;
+  }
+
+  /** Start one span. The returned closer is idempotent so finally blocks and
+   * early exits cannot double-record a phase. */
+  begin(phase: P2RTracePhase, attributes?: P2RTraceAttributes): () => void {
+    const startedAt = this.#clock();
+    let closed = false;
+    return () => {
+      if (closed || this.#finished) return;
+      closed = true;
+      const endedAt = this.#clock();
+      this.#events.push({
+        phase,
+        startMs: startedAt - this.#startedAt,
+        durationMs: endedAt - startedAt,
+        ...(attributes ? { attributes } : {}),
+      });
+    };
+  }
+
+  mark(phase: P2RTracePhase, attributes?: P2RTraceAttributes): void {
+    if (this.#finished) return;
+    this.#events.push({
+      phase,
+      startMs: this.#clock() - this.#startedAt,
+      durationMs: 0,
+      ...(attributes ? { attributes } : {}),
+    });
+  }
+
+  finish(
+    outcome: P2RTraceOutcome,
+    attributes?: P2RTraceAttributes,
+  ): P2RTraceRecord | null {
+    if (this.#finished) return null;
+    this.#finished = true;
+    const totalMs = this.#clock() - this.#startedAt;
+    const record: P2RTraceRecord = {
+      version: P2R_TRACE_VERSION,
+      traceId: this.traceId,
+      requestId: this.requestId,
+      route: this.route,
+      clock: "monotonic-ms",
+      outcome,
+      totalMs,
+      events: [...this.#events].sort((a, b) => a.startMs - b.startMs),
+      ...(attributes ? { attributes } : {}),
+    };
+    this.#emit(record);
+    return record;
+  }
+}
+
+export function createPromptResponseTrace(
+  options: PromptResponseTraceOptions,
+): PromptResponseTrace | undefined {
+  return runtimeValue("MLX_BUN_P2R_TRACE") === "1"
+    ? new PromptResponseTrace(options)
+    : undefined;
+}