mlx-bun 0.0.13 → 0.3.0

package/src/serve/lab-routes.ts

@@ -0,0 +1,341 @@
+import type { JobStore } from "../jobs";
+
+export interface LabRouteDeps {
+  ensureJobs(): Promise<JobStore>;
+  acquireGpu?: (signal: AbortSignal) => Promise<import("../contracts/resources").DisposableResource>;
+  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,
+        { acquire: deps.acquireGpu },
+      );
+      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, {
+    acquire: deps.acquireGpu,
+    onComplete: () => deps.invalidateLibrary(),
+  });
+  return Response.json({ ok: true, job_id: jobId, output_dir: outDir });
+}

package/src/serve/lane-registry.ts

@@ -16,7 +16,8 @@
 // a same-process debugging/telemetry aid, not a durable store — entries for
 // long-dead requests must not leak memory over a long-running server.
 
-export type Lane = "serial" | "serial+spec" | "batched";
+import type { Lane } from "../contracts/pi";
+export type { Lane } from "../contracts/pi";
 
 const MAX_ENTRIES = 512;
 const lanes = new Map<string, Lane>();

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

@@ -0,0 +1,155 @@
+import { modelServingBinding } from "../backends/mlx/model-serving";
+import type { ServingContext } from "./model-host";
+import type { GenerationGateway } from "./generation-gateway";
+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: ServingContext,
+  gateway: Pick<GenerationGateway, "runExclusive">,
+): Promise<Response | null> {
+  const route = matchModelAdminRoute(request.method, url.pathname);
+  if (!route) return null;
+
+  switch (route.kind) {
+    case "embeddings": {
+      const binding = modelServingBinding(ctx);
+      const embed = binding.embed?.bind(binding);
+      if (!embed) {
+        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 = await gateway.runExclusive(async () => embed(inputs, instruction), undefined, request.signal);
+      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/model-binding.ts

@@ -0,0 +1,43 @@
+import type { ModelConfig } from "../config";
+import type { MlxArray } from "../mlx/array";
+import type { MlxGatewayBinding } from "../backends/mlx/gateway-binding";
+import type { MlxSerialServices } from "../backends/mlx/serial-executor";
+import type { SsdCacheStore, SsdIndexEntry } from "../ssd-cache";
+import type { ChatRequestParams } from "./chat-request";
+import type { SerialRun, Vision } from "./generation-gateway";
+import type { RequestOwnership } from "./request-plan";
+import type { RequestPrep } from "./request-prep";
+
+export interface ServedModelInfo {
+  readonly config: ModelConfig;
+  readonly weightsBytes: number;
+}
+
+export interface BuiltPrompt {
+  promptIds: number[];
+  vision: Vision | undefined;
+  startInThinking: boolean;
+  probeStableLen: boolean;
+  diffusionPixels: MlxArray | null;
+}
+
+export interface ModelPromptBuilder {
+  (body: ChatRequestParams, tools: ChatRequestParams["tools"] | null,
+    ownership: RequestOwnership, prep: RequestPrep): Promise<BuiltPrompt>;
+}
+
+/** Model-owned native services. A replacement supplies this one binding;
+ * routes and sessions consume it without requiring a concrete model class. */
+export interface ModelServingBinding {
+  /** Opaque identity of the backend numerics that produced persisted state. */
+  readonly stateCompatibility: string;
+  readonly gateway: MlxGatewayBinding;
+  createSerial(services: MlxSerialServices): SerialRun;
+  readonly buildPrompt: ModelPromptBuilder;
+  restore(store: SsdCacheStore, entry: SsdIndexEntry): ReturnType<SsdCacheStore["restore"]>;
+  signal(promptIds: number[], bins: number, minimum: number): Promise<{ bins: number[]; vocab: number }>;
+  diagnostics(): Record<string, unknown>;
+  readonly discovery: { readonly adapters: boolean; readonly training: boolean;
+    readonly dsa: unknown; readonly embeddings: boolean };
+  embed?(inputs: string[], instruction?: string): ReturnType<typeof import("../embed").embedMany>;
+}