mlx-bun 0.4.0 → 1.0.0

package/src/cli/serve-host.ts

@@ -0,0 +1,423 @@
+// The model-scoped half of the serve composition. startModelHost is the CLI's
+// loader: runtime switches, expert offload, the model context and its startup
+// adapter. startContextHost serves one or more models from the model host
+// (residency/model-residency.ts): each loaded model is a serving unit (binding,
+// caches, engine, routes: serving-unit.ts), the Whisper companion, the model
+// router (server/model-routes.ts), and the listener that serves them with the
+// state's routes and the modules' sockets. Both borrow persistent services from the AppState they are
+// given, and one close releases everything they created in the app's order.
+import { requireChatTemplate, releaseContext, type LoadedModelContext } from "../engine/model-host";
+import { createKvBudget, type KvBudget } from "../engine/kv-budget";
+import { createResidencyHost, ResidencyError, type ResidencyEntry, type ResidencyHost, type ResidencySource } from "../residency/model-residency";
+import { leasedAdapters } from "../residency/leased-adapters";
+import { defaultBudgetBytes, estimateRecordBytes } from "../residency/resident-estimate";
+import { createRecordIndex } from "../residency/record-index";
+import { listLocalRecords } from "../residency/local-records";
+import type { DurabilityFlushResult, DurabilitySnapshotStats } from "@mlx-bun/inference/state";
+import type { ModelRecord } from "@mlx-bun/hub/registry";
+import pkgJson from "../../package.json" with { type: "json" };
+import { loadInstalledModules } from "./module-host";
+import { type RunningApp, type ServeOptions } from "./serve-options";
+import type { AppState, RouteGroup } from "./serve-state";
+import { createServingUnit, type ContextHost, type ContextHostOptions, type ServingUnit, type UnitShared } from "./serving-unit";
+
+export type { ContextHost, ContextHostOptions } from "./serving-unit";
+export { resolveServingLimits } from "./serve-options";
+
+export interface ModelHostHooks {
+  /** Runs inside the listener's drain step, after cache demotion stops and
+   * before chat sessions and HTTP responses drain: the app stops its
+   * persistent producers (jobs, downloads) here while the engine is alive. */
+  beforeDrain?(): void | Promise<void>;
+  /** Internal (worker mode, `cli/worker-entry.ts`): bind this Unix socket path
+   * instead of the TCP `port`/`hostname`; the host then reports no port, and
+   * the link it lends the state carries the socket for loopback clients. */
+  unix?: string;
+  /** Internal (worker mode): the worker's admin surface wraps the model routes
+   * and answers ahead of them (health, lease, the drain gate). */
+  routes?(model: RouteGroup): RouteGroup;
+  /** Internal (worker mode): serve only the model this host was started with; another is the parent's to load in another worker. */
+  oneModel?: true;
+}
+
+/** Worker mode: the listener is a Unix socket, so there is no port to report. */
+export interface RunningWorkerHost {
+  /** Drains the listener, releases the Whisper companion and every resident
+   * model (draining, flushing its saved state, releasing it), and restores
+   * process-wide settings. One drain however often it is called; it resolves
+   * with the cache persistence result of the models resident at close. */
+  close(): Promise<DurabilityFlushResult>;
+  /** Flush cache persistence of every resident model while serving (`POST /admin/cache/flush`'s result). */
+  flush(): Promise<DurabilityFlushResult>;
+  /** Pending cache persistence counters, readable at any time. */
+  stats(): DurabilitySnapshotStats;
+}
+export interface RunningModelHost extends RunningWorkerHost { port: number }
+
+/** The other local models a host may load beside (or instead of) its first, by exact id. */
+export interface ModelSource {
+  /** The local models this host can serve, as the registry knows them; never scans or downloads. */
+  records(): readonly ModelRecord[] | Promise<readonly ModelRecord[]>;
+  /** Load one model for serving; the host takes ownership of the context. */
+  load(record: ModelRecord): Promise<{ context: LoadedModelContext; defaultAdapter?: string; loadMs: number }>;
+  /** The model the host starts with, when it is not one `records` lists (a path given to `--model`). */
+  startup?: ModelRecord;
+}
+
+/** Internal: process-wide settings the loader applied, restored once with the
+ * host's own after the model is released, on close and on startup failure. */
+interface ContextHostHooks extends ModelHostHooks { restoreLoader?(): void }
+
+const idle = (): DurabilitySnapshotStats => ({ pendingSnapshots: 0, pendingSpills: 0, pendingSpillBytes: 0, droppedSpills: 0, failedSpills: 0 });
+function mergeStats(all: readonly DurabilitySnapshotStats[]): DurabilitySnapshotStats {
+  return all.reduce((sum, item) => ({ pendingSnapshots: sum.pendingSnapshots + item.pendingSnapshots, pendingSpills: sum.pendingSpills + item.pendingSpills,
+    pendingSpillBytes: sum.pendingSpillBytes + item.pendingSpillBytes, droppedSpills: sum.droppedSpills + item.droppedSpills,
+    failedSpills: sum.failedSpills + item.failedSpills }), idle());
+}
+function mergeDurability(all: readonly DurabilityFlushResult[]): DurabilityFlushResult {
+  return { ...mergeStats(all), durable: all.every(item => item.durable), flushedSnapshots: all.reduce((n, item) => n + item.flushedSnapshots, 0),
+    missingSnapshots: all.reduce((n, item) => n + item.missingSnapshots, 0), elapsedMs: all.reduce((n, item) => Math.max(n, item.elapsedMs), 0) };
+}
+
+/** Serve's KV flags as the composition request names them, `--fused-sdpa` (the
+ * affine read kernels) with them. */
+const kvFlags = ({ cache: { kvQuant, turboQuant, quantizedKvStart }, fusedSdpa }: ServeOptions) => ({
+  ...(kvQuant !== undefined ? { override: kvQuant } : {}), ...(turboQuant ? { turboQuant } : {}),
+  ...(quantizedKvStart !== undefined ? { quantizedKvStart } : {}), ...(fusedSdpa !== undefined ? { fusedSdpa } : {}) });
+
+/** Model composition owns resources until each explicit ownership transfer. */
+export function startModelHost(state: AppState, model: ModelRecord, options: ServeOptions, hooks?: ModelHostHooks & { unix?: undefined }): Promise<RunningModelHost>;
+export function startModelHost(state: AppState, model: ModelRecord, options: ServeOptions, hooks: ModelHostHooks & { unix: string }): Promise<RunningWorkerHost>;
+export function startModelHost(state: AppState, model: ModelRecord, options: ServeOptions, hooks?: ModelHostHooks): Promise<RunningModelHost | RunningWorkerHost>;
+export async function startModelHost(state: AppState, model: ModelRecord, options: ServeOptions, hooks: ModelHostHooks = {}): Promise<RunningModelHost | RunningWorkerHost> {
+  const [{ loadContext }, { configureRuntime }] = await Promise.all([import("../engine"), import("@mlx-bun/inference/runtime/config")]);
+  const restoreRuntime = configureRuntime({
+    ...(options.forceWire ? { MLX_BUN_FORCE_WIRE: "1" } : {}),
+    ...(options.allowPrivateMedia ? { MLX_BUN_ALLOW_PRIVATE_MEDIA: "1" } : {}) });
+  // Offload routing and the runtime switches are restored only after the
+  // engine has released the model, ahead of the context host's allocator limit.
+  let restoreOffload: (() => void) | undefined;
+  let restoreLoader: (() => void) | undefined = () => { try { restoreOffload?.(); } finally { restoreRuntime(); } };
+  let cleanup: (() => void) | undefined;
+  try {
+    if (options.expertOffload) {
+      // Main: dense models log and continue; MoE experts route through the
+      // page-aligned file built on first use, activated before construction.
+      if (model.expertsBytes === 0) console.warn("--expert-offload ignored: this model has no experts (dense)");
+      else {
+        const { ensureOffloadFile, activateExpertOffload } = await import("@mlx-bun/inference/artifacts");
+        restoreOffload = activateExpertOffload(await ensureOffloadFile(model.path, message => console.log(`[serve] expert offload: ${message}`)));
+      }
+    }
+    /** One model, loaded for serving. `--draft-*` and `--adapter` belong to the model named at startup; the others load plain. */
+    const loadServed = async (record: ModelRecord): Promise<{ context: LoadedModelContext; defaultAdapter?: string; loadMs: number }> => {
+      const startup = record.repoId === model.repoId;
+      const draft = startup ? options.draft ?? {} : {};
+      const loadStarted = performance.now();
+      const context = await loadContext(record.path, record.repoId, {
+        ...(options.memoryBudgetBytes !== undefined ? { memoryBudgetBytes: options.memoryBudgetBytes } : {}),
+        // The composition's facts: the KV flags and batch cap. Startup adapters mount below.
+        kv: kvFlags(options), maxRows: options.capacity,
+        ...(options.cache.pagedKv ? { pagedKv: options.cache.pagedKv } : {}),
+        // Resource plan inputs for runtimes that plan memory up front; other models ignore this block.
+        runtime: { batchSize: options.capacity, maxGenerationTokens: options.defaultGeneratedTokens ?? 128,
+          ...(options.memoryBudgetBytes !== undefined ? { memoryBudgetBytes: options.memoryBudgetBytes } : {}),
+          ...(options.contextTokens !== undefined ? { contextTokens: options.contextTokens } : {}),
+          ...(options.mtp !== undefined ? { nativeDraft: options.mtp } : {}) },
+        ...(draft.adaptiveDepth ? { adaptiveDraftDepth: true } : {}),
+        // Main's gate: a draft model, or the model-free ngram kind, or mtp alone
+        // (the host resolves the bundled <model>/mtp/ companion).
+        ...(draft.modelDir || draft.kind === "ngram" || draft.kind === "mtp" ? {
+          ...(draft.modelDir ? { draftModelDir: draft.modelDir } : {}),
+          ...(draft.numTokens !== undefined ? { numDraftTokens: draft.numTokens } : {}),
+          ...(draft.kind ? { draftKind: draft.kind } : {}),
+          ...(draft.ngramMax !== undefined ? { ngramMax: draft.ngramMax } : {}),
+          ...(draft.ngramMin !== undefined ? { ngramMin: draft.ngramMin } : {}) } : {}),
+      });
+      try {
+        if (draft.modelDir) console.log(`[serve] draft: ${draft.modelDir.split("/").filter(Boolean).at(-1)}`);
+        else if (draft.kind === "ngram") console.log("[serve] draft: ngram (prompt lookup)");
+        requireChatTemplate(context);
+        // Main: a startup adapter mounts before any request and becomes the default
+        // for requests without an adapter field (an explicit adapter, including
+        // "none", still wins); a bad adapter fails startup and releases the model.
+        if (startup && options.adapterDir) {
+          const directory = options.adapterDir.replace(/\/+$/, "");
+          try {
+            const info = await context.adapters.mount(directory.split("/").pop()!, directory);
+            console.log(`[serve] adapter ${info.id} mounted (${info.mountedLayers} layers) · default for requests (select others via \`adapter\`)`);
+            return { context, defaultAdapter: info.id, loadMs: performance.now() - loadStarted };
+          } catch (error) {
+            throw new Error(`adapter mount failed: ${error instanceof Error ? error.message : String(error)}`);
+          }
+        }
+        return { context, loadMs: performance.now() - loadStarted };
+      } catch (error) {
+        try { context.dispose(); } catch (failure) { throw new AggregateError([error, failure], "startup and cleanup failed"); }
+        throw error;
+      }
+    };
+    const first = await loadServed(model);
+    cleanup = () => first.context.dispose();
+    // The context host takes the model and the loader's process settings on entry, including a failed start.
+    const restore = restoreLoader;
+    cleanup = undefined; restoreLoader = undefined;
+    // Expert offload routes the process's loads through one file: that model is all this process serves.
+    const models: ModelSource | undefined = options.expertOffload || hooks.oneModel ? undefined : { startup: model, load: loadServed, records: localRecords };
+    return await startContextHost(state, first.context, options,
+      { ownership: "owned", artifact: model, owner: "serve", loadMs: first.loadMs, ...(first.defaultAdapter ? { defaultAdapter: first.defaultAdapter } : {}),
+        ...(models ? { models } : {}) },
+      { ...hooks, restoreLoader: restore });
+  } catch (error) {
+    try { cleanup?.(); }
+    catch (failure) { throw new AggregateError([error, failure], "startup and cleanup failed"); }
+    finally { restoreLoader?.(); }
+    throw error;
+  }
+}
+
+/** The models `/v1/models` may list and the host may load: supported generation checkpoints the registry knows. */
+const localRecords = () => listLocalRecords("generate");
+
+/** Serve one or more models with the app's routes, chat, and listener. The
+ * first context is resident from the start; it is released by
+ * `input.ownership`, on close and on a failed start alike. Without
+ * `input.models` it is the only model and is never evicted; with it, any other
+ * local model loads when a request names it and fits the memory budget (the
+ * least recently used unpinned, unleased model is drained, its saved state
+ * flushed, and released first), and the first can be evicted and reloaded like
+ * the rest. Everything else the host created is released either way. It
+ * applies no runtime switches and activates no expert offload: those belong to
+ * whoever loaded the context. A context without a chat template is refused
+ * unless `input.buildPrompt` builds its prompts. */
+export function startContextHost(state: AppState, context: LoadedModelContext, options: ContextHostOptions, input: ContextHost & { models?: ModelSource },
+  hooks?: ContextHostHooks & { unix?: undefined }): Promise<RunningModelHost>;
+export function startContextHost(state: AppState, context: LoadedModelContext, options: ContextHostOptions, input: ContextHost & { models?: ModelSource },
+  hooks: ContextHostHooks): Promise<RunningModelHost | RunningWorkerHost>;
+export async function startContextHost(state: AppState, context: LoadedModelContext, options: ContextHostOptions, input: ContextHost & { models?: ModelSource },
+  hooks: ContextHostHooks = {}): Promise<RunningModelHost | RunningWorkerHost> {
+  // Process-wide settings (the loader's, then this host's allocator limit) are
+  // restored only after the engine has released the model, on close and on
+  // startup failure alike, so a later host in the same process starts from
+  // the state it found. The restore runs once: a repeated close must not undo
+  // a later host's settings.
+  let restoreAllocator: (() => void) | undefined, restored = false;
+  const restoreProcess = () => {
+    if (restored) return;
+    restored = true;
+    try { hooks.restoreLoader?.(); } finally { restoreAllocator?.(); }
+  };
+  let cleanup: (() => void | Promise<unknown>) | undefined = () => releaseContext(context, input.ownership);
+  // The link is returned to the state once, like the process restore.
+  let detachLink = () => {};
+  const detach = () => { const release = detachLink; detachLink = () => {}; release(); };
+  try {
+    const [{ startServer }, { createModelRoutes },
+      { createHostServices, createModuleRoutes, declaredOperations }] = await Promise.all([
+      import("../server/start"), import("../server/model-routes"),
+      import("@mlx-bun/app-services"),
+    ]);
+    // Main: the plan's allocator limit, else the explicit budget, caps the
+    // allocator for the whole process and bounds optional cache residency.
+    const allocatorLimitBytes = context.memoryPlan?.allocatorLimitBytes ?? options.memoryBudgetBytes;
+    if (allocatorLimitBytes) {
+      const { setMemoryLimit } = await import("@mlx-bun/mlx/ffi");
+      const previous = setMemoryLimit(allocatorLimitBytes);
+      restoreAllocator = () => { setMemoryLimit(previous); };
+    }
+    // One saved-state budget over every model's store in this process.
+    const ssdBudget: KvBudget | undefined = options.cache.ssdCacheDir
+      ? createKvBudget(options.cache.ssdCacheDir, options.cache.ssdCacheMaxBytes ?? Infinity) : undefined;
+    let residency!: ResidencyHost<ServingUnit>;
+    // Each closing model's saved-state result; the ones at host close are the close's evidence.
+    const finalResults: DurabilityFlushResult[] = [];
+    let closingHost = false;
+    // Main's speech-to-text companion, served by the transcription module: an
+    // explicit --whisper-model, else the first downloaded Whisper checkpoint,
+    // resolved once on the first audio request (a later download needs a
+    // restart, as in main). The weights load per take and release per the
+    // --whisper-* policy; every decode runs with every resident model paused so
+    // it never overlaps generation, and a load first makes room in the budget.
+    const moduleHost = createHostServices({ whisper: options.whisper, events: state.events,
+      exclusive: async (fn, signal) => { const pause = await residency.pauseAll(signal); try { return await fn(); } finally { pause.dispose(); } },
+      admit: bytes => residency.makeRoom(bytes) });
+    const whisperInfo = async () => {
+      const id = await moduleHost.whisper.defaultFor("transcribe");
+      return id === undefined ? null : { id, resident: moduleHost.whisper.stats(id).resident };
+    };
+    const shared: UnitShared = { responses: state.responses, events: state.events, transcription: whisperInfo,
+      ...(ssdBudget ? { ssdBudget } : {}), ...(allocatorLimitBytes ? { allocatorLimitBytes } : {}),
+      onClosed: result => { if (closingHost) finalResults.push(result); } };
+    // The unit takes the context and releases it (by ownership) on close and on its own failed start.
+    cleanup = undefined;
+    const first = await createServingUnit(context, options, input, shared);
+    cleanup = () => first.close({ flush: true });
+    // The other local models, as the registry knows them: refreshed after a download or job, and at most every few seconds on a miss.
+    const records = input.models ? createRecordIndex(() => input.models!.records(), input.models.startup) : undefined;
+    const known = async (id: string): Promise<ModelRecord | undefined> => records?.find(id);
+    const machine = input.models ? (await import("@mlx-bun/inference/execution/fit")).thisMachine() : undefined;
+    const memory = input.models ? await import("@mlx-bun/mlx/ffi") : undefined;
+    // Every resident model together may use this much: by default a share of what the GPU can wire.
+    const budgetBytes = !input.models ? Infinity
+      : options.modelBudgetBytes ?? defaultBudgetBytes(memory!.maxRecommendedWorkingSetSize());
+    const estimate = (record: ModelRecord) => estimateRecordBytes(record, { budgetBytes, cache: options.cache, ...(machine ? { machine } : {}) });
+    const entryOf = (unit: ServingUnit): ResidencyEntry => ({ id: unit.id, bytes: unit.bytes(), operations: unit.operations });
+    const source: ResidencySource<ServingUnit> = {
+      async resolve(id) {
+        const record = await known(id);
+        if (!record) return undefined;
+        return { id: record.repoId, bytes: await estimate(record), operations: declaredOperations(record.modelType) };
+      },
+      async load(entry) {
+        const record = (await known(entry.id))!;
+        const loaded = await input.models!.load(record);
+        // A planned runtime's own allocator limit applies while it is resident.
+        const limit = loaded.context.memoryPlan?.allocatorLimitBytes;
+        const previous = limit ? (await import("@mlx-bun/mlx/ffi")).setMemoryLimit(limit) : undefined;
+        const unit = await createServingUnit(loaded.context, options,
+          { ownership: "owned", artifact: record, ...(input.owner ? { owner: input.owner } : {}), ...(loaded.defaultAdapter ? { defaultAdapter: loaded.defaultAdapter } : {}) },
+          { ...shared, ...(limit ? { allocatorLimitBytes: limit } : {}) });
+        if (previous === undefined) return unit;
+        const close = unit.close.bind(unit);
+        return Object.assign(unit, { close: async (options: { readonly flush: boolean }) => {
+          try { return await close(options); } finally { (await import("@mlx-bun/mlx/ffi")).setMemoryLimit(previous); }
+        } });
+      },
+    };
+    let current = first.id;
+    residency = createResidencyHost<ServingUnit>({ source, budgetBytes, events: state.events,
+      ...(memory ? { measured: () => memory.activeMemory() } : {}),
+      external: () => moduleHost.whisper.resident().reduce((sum, model) => sum + model.bytes, 0),
+      defaultFor: async operation => operation === "generate" ? current : undefined,
+      // Switching: the models this host may load; a caller's own context and an expert-offload process hold one model and never switch.
+      ...(input.models ? { async serve(id: string, signal: AbortSignal | undefined) {
+        if (id !== first.id && !await known(id)) throw new ResidencyError("not-found", `${id} is not a local model; download it first`);
+        (await residency.acquire(id, { ...(signal ? { signal } : {}), need: ["generate"] })).release();
+        current = id;
+      } } : {}),
+      uses: unit => [unit.artifactPath],
+      log: line => console.log(line) });
+    // A context the caller supplied cannot be reloaded, so it is never evicted.
+    residency.adopt(entryOf(first), first, { pin: !input.models, ...(input.loadMs !== undefined ? { loadMs: input.loadMs } : {}) });
+    let modules: Awaited<ReturnType<typeof loadInstalledModules>> | undefined;
+    // The modules stop first (admission stops, takes in flight are joined), then the models release.
+    const closeModules = async () => {
+      const errors: unknown[] = [];
+      try { await modules?.stop(); } catch (error) { errors.push(error); }
+      try { await moduleHost.whisper.close(); } catch (error) { errors.push(error); }
+      if (errors.length) throw new AggregateError(errors, "module cleanup failed");
+    };
+    const closeApp = async () => {
+      const errors: unknown[] = [];
+      // A request admitted before shutdown may lease the companion while
+      // responses drain. Close again here; each owner joins/releases only once.
+      try { await closeModules(); } catch (error) { errors.push(error); }
+      closingHost = true;
+      try { await residency.close(); } catch (error) { errors.push(error); }
+      if (errors.length) throw new AggregateError(errors, "application cleanup failed");
+    };
+    cleanup = closeApp;
+    // Attach before activation, so modules can resolve either operation while they start.
+    // A Unix listener has no port: the requested one stands in for Pi's TCP
+    // loopback, and for the link's URL placeholder (its clients use the socket).
+    let boundPort = options.port;
+    const invalidateLibrary = () => { records?.invalidate(); };
+    // Jobs and loopback clients reach this host from the first served request. A job pauses every resident model.
+    // Modules leasing the served model reach it as the current one: the id and weights follow a switch.
+    detachLink = state.attach({ get model() { return { id: current, bytes: (residency.peek(current) ?? first).context.model.weightsBytes }; },
+      get port() { return boundPort; }, ...(hooks.unix ? { unix: hooks.unix } : {}),
+      acquireExecutionLease: signal => residency.pauseAll(signal),
+      // Modules that generate through the model host (a dataset job) hold the model's residency lease until they release it.
+      hold: (id, signal) => residency.acquire(id, { need: ["generate"], ...(signal ? { signal } : {}) }),
+      invalidateLibrary,
+      resident: () => residency.resident(),
+      companions: moduleHost.whisper,
+      adapters: leasedAdapters(residency, () => current),
+      ...(input.models ? { serve: (id: string, signal: AbortSignal) => residency.serve(id, { signal }) } : {}) });
+    modules = await loadInstalledModules({ bindings: { ...moduleHost.bindings, modelHost: () => state.modelHost } });
+    const models = createModelRoutes({ host: residency, current: () => current, serves: async id => id === first.id || !!await known(id) });
+    const moduleRoutes = createModuleRoutes(modules.routes);
+    const residentUnits = () => residency.resident().flatMap(model => { const unit = residency.peek(model.id); return unit ? [unit] : []; });
+    const persistent = state.routes;
+    const modelRoutes = { handle: async (request: Request) => await models.handle(request) ?? await moduleRoutes.handle(request) ?? await persistent.jobs.handle(request) ??
+      await persistent.appModules.handle(request) ?? await persistent.publishing.handle(request) };
+    const routes = hooks.routes?.(modelRoutes) ?? modelRoutes;
+    // startServer owns engine cleanup on entry, including a bind failure.
+    cleanup = undefined;
+    const listener = await startServer({ routes, web: state.web, sockets: state.sockets,
+      beforeDrain: async () => {
+        const errors: unknown[] = [];
+        try { for (const unit of residentUnits()) unit.stopBackground(); } catch (error) { errors.push(error); }
+        // Whisper closes with the persistent owners before drain: admission stops,
+        // in-flight takes are joined, weights release ahead of the chat models.
+        for (const result of await Promise.allSettled([hooks.beforeDrain?.(), closeModules()]))
+          if (result.status === "rejected") errors.push(result.reason);
+        if (errors.length === 1) throw errors[0];
+        if (errors.length) throw new AggregateError(errors, "background shutdown failed");
+      },
+      closeEngine: closeApp }, hooks.unix ? { unix: hooks.unix } : { port: options.port, hostname: options.hostname });
+    if (listener.server.port !== undefined) boundPort = listener.server.port;
+    let closing: Promise<DurabilityFlushResult> | undefined;
+    const close = () => closing ??= (async () => {
+      try { await listener.close(); } finally { try { restoreProcess(); } finally { detach(); } }
+      return mergeDurability(finalResults);
+    })();
+    const host = { close, flush: async () => mergeDurability(await Promise.all(residentUnits().map(unit => unit.flush()))),
+      stats: () => mergeStats(residentUnits().map(unit => unit.stats())) };
+    return hooks.unix ? host : { port: boundPort, ...host };
+  } catch (error) {
+    try { await cleanup?.(); }
+    catch (failure) { throw new AggregateError([error, failure], "startup and cleanup failed"); }
+    finally { try { restoreProcess(); } finally { detach(); } }
+    throw error;
+  }
+}
+
+/** Transcription-only composition (main's `serve <whisper checkpoint>`): the
+ * audio routes plus `/v1`, `/v1/models`, `/health`, and `/stats` over the Whisper
+ * checkpoint alone; no chat model, prompt cache, jobs, or web app. `--preload`
+ * loads the weights before the listener binds; otherwise the first request
+ * pages them in. Every take runs one at a time inside the service. The
+ * internal hooks are the model host's: a Unix socket instead of TCP, a route
+ * wrapper, and a step ahead of the Whisper close (the worker app form). */
+export async function startTranscriptionHost(model: ModelRecord, options: ServeOptions,
+  hooks: Pick<ModelHostHooks, "unix" | "routes" | "beforeDrain"> = {}): Promise<RunningApp> {
+  const [{ createCompanionInfoRoutes, createHostServices, createModuleRoutes, createModuleSockets }, { startServer }] = await Promise.all([
+    import("@mlx-bun/app-services"), import("../server/start"),
+  ]);
+  const whisper = options.whisper ?? {};
+  const moduleHost = createHostServices({ whisper: { ...whisper, modelDir: model.path, modelId: model.repoId } });
+  let modules: Awaited<ReturnType<typeof loadInstalledModules>> | undefined;
+  // The modules stop first (takes in flight are joined), then the weights release.
+  const close = async () => {
+    const errors: unknown[] = [];
+    try { await modules?.stop(); } catch (error) { errors.push(error); }
+    try { await moduleHost.whisper.close(); } catch (error) { errors.push(error); }
+    moduleHost.events.close();
+    if (errors.length) throw new AggregateError(errors, "transcription cleanup failed");
+  };
+  let cleanup: (() => Promise<void>) | undefined = close;
+  try {
+    modules = await loadInstalledModules(moduleHost);
+    if (whisper.preload) await moduleHost.whisper.preload(model.repoId);
+    const audio = createModuleRoutes(modules.routes);
+    const loaded = modules;
+    const info = createCompanionInfoRoutes({ modelId: model.repoId, models: moduleHost.whisper, counters: () => loaded.status("transcription"),
+      name: "mlx-bun", version: pkgJson.version, startedAt: Date.now(),
+      endpoints: [...modules.routes.map(route => `${route.spec.method} ${route.path}`), "GET /v1/models", "GET /health", "GET /stats"] });
+    // startServer owns the cleanup on entry, including a bind failure.
+    cleanup = undefined;
+    const routes: RouteGroup = { handle: async request => await audio.handle(request) ?? await info.handle(request) };
+    const listener = await startServer({
+      routes: hooks.routes?.(routes) ?? routes,
+      web: () => null,
+      // No chat: a transcription-only server declares no sockets.
+      sockets: createModuleSockets([]),
+      // Whisper closes before drain: admission stops, in-flight takes are joined, weights release.
+      beforeDrain: async () => { try { await hooks.beforeDrain?.(); } finally { await close(); } },
+      closeEngine: async () => {},
+    }, hooks.unix ? { unix: hooks.unix } : { port: options.port, hostname: options.hostname });
+    return { ...(hooks.unix ? {} : { port: listener.server.port! }), close: listener.close,
+      downloads: { active: [], start() { throw new Error("transcription-only server owns no downloads"); } } };
+  } catch (error) { await cleanup?.(); throw error; }
+}

package/src/cli/serve-isolated.ts

@@ -0,0 +1,217 @@
+// The isolated serve composition, the default for `serve`: this process keeps the persistent
+// CPU state (serve-state.ts), the web app, the chat module, the Responses history, the managed jobs and
+// the modules that run beside them, and loads no model. Each resident model runs in its own
+// worker process (worker-unit.ts; worker-entry.ts composes the model host over a Unix socket).
+// This process holds their residency by memory fit with the same manager the in-process host
+// uses (residency/model-residency.ts), counting each worker at the MLX memory it measures and reports
+// (cli/worker-unit.ts; the estimate stands only before its first report): a model that fits the budget
+// gets its own worker beside the others; otherwise the least recently used unpinned, unleased model is drained, its worker
+// flushes its saved state durably and exits (the worker gets the CLI's shutdown budget, never a
+// kill first), and naming it again spawns a worker that resumes from that state. Requests route
+// by their `model` (server/model-routes.ts) and forward over the worker's socket; a crash
+// respawns that worker with its model while the app stays up. Whisper is a companion worker of
+// its own, which can be pinned. Each worker's events are relayed onto this process's bus.
+// Nothing here imports the engine or a native module.
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import type { AppModule } from "@mlx-bun/app-core";
+import type { ModelRecord } from "@mlx-bun/hub/registry";
+import { EngineUnavailableError, type WorkerRestartBudget } from "../jobs/worker-supervisor";
+import { createWorkerMemoryClient } from "../server/memory-completion-client";
+import { createModelRoutes } from "../server/model-routes";
+import { createProxyRoutes } from "../server/proxy-routes";
+import { createResponsesClient } from "../server/responses-client";
+import { startServer } from "../server/start";
+import { listLocalRecords } from "../residency/local-records";
+import { leasedAdapters } from "../residency/leased-adapters";
+import { createResidencyHost, ResidencyError, type ResidencyEntry, type ResidencyHost, type ResidencySource } from "../residency/model-residency";
+import { createRecordIndex } from "../residency/record-index";
+import { defaultBudgetBytes, estimateRecordBytes } from "../residency/resident-estimate";
+import type { RunningApp, ServeOptions } from "./serve-options";
+import { createAppState, type RouteGroup } from "./serve-state";
+import { spawnWorkerUnit, type WorkerUnit, type WorkerUnitContext } from "./worker-unit";
+
+/** What the composition root supplies (`modules`: the installed modules that run in the persistent state)
+ * and, internally for tests, stand-ins for the worker entry, the restart policy, and the model listing. */
+export interface IsolatedServeHooks {
+  modules?: Parameters<typeof createAppState>[2];
+  /** Selects the dedicated task-model snapshot, without loading it in the parent. */
+  taskSnapshot?(): Promise<string>;
+  entry?: string;
+  restarts?: Partial<WorkerRestartBudget>;
+  env?: Record<string, string | undefined>;
+  readyTimeoutMs?: number;
+  graceMs?: number;
+  notice?(line: string): void;
+  log?(line: string): void;
+  error?(line: string): void;
+  /** Internal (tests): sees the app's event bus, where each worker's relayed events and the host's load/unload events arrive. */
+  observe?(events: Pick<import("@mlx-bun/app-core").EventBus, "subscribe">): void;
+  /** The local chat models the host may load, instead of the registry's. */
+  records?(): readonly ModelRecord[];
+  /** The Whisper checkpoints it may load, instead of the registry's. */
+  companions?(): readonly ModelRecord[];
+}
+
+/** Isolated composition: the persistent state and the startup model's worker first, the listener
+ * once that worker serves, so startup fails the way the in-process composition does when the model cannot load. */
+export async function startIsolatedServer(model: ModelRecord, options: ServeOptions, hooks: IsolatedServeHooks = {}): Promise<RunningApp> {
+  let residency!: ResidencyHost<WorkerUnit>;
+  /** The model requests that name none are answered by, and Pi's `local`: the one the server started with, then the last switch. */
+  let current = model.repoId;
+  const workers = () => residency ? residency.resident().flatMap(({ id }) => { const unit = residency.peek(id); return unit ? [unit] : []; }) : [];
+  const currentUnit = () => residency?.peek(current);
+  // Snapshots a worker reads besides its own model: the memory task model, selected here and kept out of hub GC
+  // from the call that carries it for as long as that worker lives (a respawned worker holds none until asked).
+  const retained = new WeakMap<WorkerUnit, { restarts: number; paths: Set<string> }>();
+  const retention = (unit: WorkerUnit) => {
+    let held = retained.get(unit);
+    if (!held || held.restarts !== unit.supervisor.restarts) retained.set(unit, held = { restarts: unit.supervisor.restarts, paths: new Set() });
+    return held.paths;
+  };
+  // The parent loads no model: each synthesis stage call or batch runs on the current model's worker (its memory task
+  // model, over its private route). A run holds that model's residency lease from its first call to its last (the same
+  // `acquire` every consumer takes), so the worker it depends on is never drained under it; the worker takes its own
+  // execution lease per call. Each call selects the task model snapshot once, here (the worker never scans the cache):
+  // the call carries it, and a call that loads the task model loads exactly it.
+  const state = await createAppState({ ...options, memoryCompletions: async signal => {
+    let release: (() => void) | undefined;
+    let selected: Promise<WorkerUnit> | undefined;
+    const select = async () => {
+      if (!residency) throw new EngineUnavailableError("starting", null);
+      const lease = await residency.acquire(current, { signal, need: ["generate"] });
+      release = () => lease.release();
+      return lease.unit;
+    };
+    const client = createWorkerMemoryClient(async () => {
+      if (!hooks.taskSnapshot) throw new Error("memory synthesis has no task model snapshot selector in this composition");
+      const snapshot = await hooks.taskSnapshot();
+      // Planning and empty pipelines never acquire or reload a worker. The first real call captures one worker,
+      // whose residency stays held until the run (including every in-flight row) settles.
+      const unit = await (selected ??= select());
+      retention(unit).add(snapshot);
+      return { worker: unit.supervisor, snapshot };
+    }, signal);
+    return Object.assign(client, { release: () => release?.() });
+  } }, options.storagePaths ?? {}, hooks.modules ?? []);
+  hooks.observe?.(state.events);
+  // The sockets live in a private directory (0700) this process removes, one per worker.
+  const socketDir = mkdtempSync(join(tmpdir(), "mlx-worker-"));
+  const removeSocketDir = () => rmSync(socketDir, { recursive: true, force: true });
+  const notice = hooks.notice ?? (line => console.log(`[isolate] ${line}`));
+  // The startup worker until residency adopts it: a failure in between still stops it.
+  let unadopted: WorkerUnit | undefined;
+  const closeEngine = async () => {
+    try { await unadopted?.close({ flush: false }); } finally { try { await residency?.close(); } finally { removeSocketDir(); } }
+  };
+  let detachLink = () => {};
+  const detach = () => { const release = detachLink; detachLink = () => {}; release(); };
+  let cleanup: (() => Promise<void>) | undefined = closeEngine;
+  try {
+    const { thisMachine } = await import("@mlx-bun/inference/execution/fit");
+    const machine = thisMachine();
+    const context: WorkerUnitContext = { options, startup: model, socketDir, publish: event => state.events.publish(event), notice,
+      ...(hooks.entry ? { entry: hooks.entry } : {}), ...(hooks.restarts ? { restarts: hooks.restarts } : {}), ...(hooks.env ? { env: hooks.env } : {}),
+      ...(hooks.readyTimeoutMs !== undefined ? { readyTimeoutMs: hooks.readyTimeoutMs } : {}), ...(hooks.graceMs !== undefined ? { graceMs: hooks.graceMs } : {}),
+      ...(hooks.log ? { log: hooks.log } : {}), ...(hooks.error ? { error: hooks.error } : {}) };
+    // Startup fails fast, as the in-process composition does: the first model's worker serves before the listener binds. Its
+    // report is also how this process, which asks no Metal device itself, learns the device working set.
+    const first = unadopted = await spawnWorkerUnit(context, model, "primary", model.sizeBytes);
+    // What all resident workers may use together: `--model-budget`, else the share of the device working set the in-process
+    // host uses (residency/resident-estimate.ts), as the first worker reported it.
+    const budgetBytes = options.modelBudgetBytes ?? defaultBudgetBytes(first.measured()?.workingSetBytes);
+    const chatModels = createRecordIndex(() => hooks.records?.() ?? listLocalRecords("generate"), model);
+    const companionList = () => hooks.companions?.() ?? listLocalRecords("transcribe");
+    // Whisper: the checkpoint `--whisper-model` resolved to, else the first downloaded one, looked up once.
+    let whisper: Promise<ModelRecord | undefined> | undefined;
+    const whisperRecord = () => whisper ??= (async () => {
+      const explicit = options.whisper;
+      if (explicit?.modelDir && explicit.modelId) {
+        const listed = (await companionList()).find(record => record.repoId === explicit.modelId);
+        return listed ?? { repoId: explicit.modelId, path: explicit.modelDir, modelType: "whisper", sizeBytes: 0, expertsBytes: 0 } as unknown as ModelRecord;
+      }
+      return (await companionList())[0];
+    })();
+    const source: ResidencySource<WorkerUnit> = {
+      async resolve(id) {
+        const chat = await chatModels.find(id);
+        // A config the fit model cannot read leaves the checkpoint's own size as the estimate.
+        if (chat) return { id, bytes: await estimateRecordBytes(chat, { budgetBytes, cache: options.cache, machine }).catch(() => chat.sizeBytes), operations: ["generate"], role: "primary" };
+        const companion = await whisperRecord();
+        if (companion && companion.repoId === id) return { id, bytes: companion.sizeBytes, operations: ["transcribe"], role: "companion" };
+        return undefined;
+      },
+      async load(entry) {
+        const record = entry.role === "companion" ? (await whisperRecord())! : (await chatModels.find(entry.id))!;
+        return spawnWorkerUnit(context, record, entry.role ?? "primary", entry.bytes);
+      },
+    };
+    residency = createResidencyHost<WorkerUnit>({ source, budgetBytes, events: state.events,
+      defaultFor: async operation => operation === "generate" ? current : operation === "transcribe" ? (await whisperRecord())?.repoId : undefined,
+      async serve(id, signal) {
+        if (!await chatModels.find(id)) throw new ResidencyError("not-found", `${id} is not a local model; download it first`);
+        (await residency.acquire(id, { ...(signal ? { signal } : {}), need: ["generate"] })).release();
+        current = id;
+      },
+      // Hub cleanup keeps the snapshot of every resident worker's model, and the task model's on the worker that read it.
+      uses: unit => [unit.record.path, ...retention(unit)],
+      log: line => notice(line) });
+    const entryOf = (unit: WorkerUnit): ResidencyEntry => ({ id: unit.id, bytes: unit.bytes(), role: unit.role, operations: unit.operations });
+    residency.adopt(entryOf(first), first, { loadMs: first.readyMs });
+    unadopted = undefined;
+    // Whisper stays resident when asked to (`--whisper-resident`); otherwise it is the first to go when memory is short.
+    const companion = options.whisper?.resident ? await whisperRecord() : undefined;
+    if (companion) residency.pin(companion.repoId);
+
+    const models = createModelRoutes({ host: residency, current: () => current,
+      serves: async id => !!await chatModels.find(id),
+      transcription: async () => (await whisperRecord())?.repoId,
+      // Audio requests belong to the Whisper worker, which loads on first use and drains a chat model when memory is short.
+      companion: async pathname => (pathname.startsWith("/v1/audio/") || pathname.startsWith("/admin/transcription/")) ? (await whisperRecord())?.repoId : undefined });
+    const responses = createResponsesClient(state.responses);
+    const proxy = createProxyRoutes({ workers: () => workers(), current: () => current, models, responses,
+      modelId: model.repoId, startedAt: Date.now() });
+    const invalidateLibrary = () => { chatModels.invalidate(); };
+    const persistent = state.routes;
+    // The persistent groups answer first, in the direct host's order among themselves; the proxy takes every remaining path.
+    const routes: RouteGroup = { handle: async request => await persistent.jobs.handle(request) ??
+      await persistent.appModules.handle(request) ??
+      await persistent.publishing.handle(request) ?? await proxy.handle(request) };
+    let boundPort = options.port;
+    // The chat module lives in the state and reaches the model through the served model host (over this listener's proxy;
+    // its `local` model id is whatever is current), so web chat survives a worker restart and reports its failures as errors.
+    // Jobs pause every resident worker; loopback clients reach the workers through the proxy; a hub switch is this host's to make.
+    detachLink = state.attach({ get model() { return { id: current, bytes: currentUnit()?.bytes() ?? model.sizeBytes }; }, get port() { return boundPort; },
+      acquireExecutionLease: signal => residency.pauseAll(signal),
+      // Modules that generate through the model host (a dataset job) hold the model's residency lease until they release it.
+      hold: (id, signal) => residency.acquire(id, { need: ["generate"], ...(signal ? { signal } : {}) }),
+      invalidateLibrary,
+      resident: () => residency.resident(),
+      adapters: leasedAdapters(residency, () => current),
+      serve: (id, signal) => residency.serve(id, { signal }) });
+    // startServer owns engine cleanup on entry, including a bind failure.
+    cleanup = undefined;
+    const listener = await startServer({ routes, web: state.web, sockets: state.sockets,
+      // The app stops its producers (jobs, downloads) while the workers are alive, then chat and HTTP drain, then every
+      // worker is drained, flushes its saved state, and exits.
+      beforeDrain: () => state.close(),
+      closeEngine }, { port: options.port, hostname: options.hostname });
+    if (listener.server.port !== undefined) boundPort = listener.server.port;
+    let closing: Promise<void> | undefined;
+    return { port: boundPort, downloads: state.downloads, close: () => closing ??= (async () => {
+      const errors: unknown[] = [];
+      try { await listener.close(); } catch (error) { errors.push(error); } finally { detach(); }
+      // The listener's drain already closed the state; a repeated failure here is the same one.
+      try { await state.close(); } catch (error) { errors.push(error); }
+      if (errors.length) throw errors[0];
+    })() };
+  } catch (error) {
+    const failures: unknown[] = [];
+    try { await cleanup?.(); } catch (failure) { failures.push(failure); }
+    finally { detach(); }
+    try { await state.close(); } catch (failure) { failures.push(failure); }
+    if (failures.length) throw new AggregateError([error, ...failures], "startup and cleanup failed");
+    throw error;
+  }
+}