mlx-bun 0.0.13 → 0.3.0

package/src/serve/pipeline.ts

@@ -0,0 +1,33 @@
+// Shared contracts of the request pipeline. Each stage is one configured
+// program — `run(input) → output` — and the server composes them in order,
+// like a shell pipe:
+//
+//   new ChatRequest(body) → ChatStage.run → InferenceRequest
+//                         → InferenceStage.run → InferenceResult → wire JSON/SSE
+//
+// A stage knows only the interface it reads and the one it writes.
+import type { CompletionEvent } from "./completion-sink";
+import type { CompletionUsage } from "./completion-executor";
+import type { PromptResponseTrace } from "./prompt-response-trace";
+
+/** A request a stage refuses. `status` is the HTTP status the wire adapter
+ *  answers with; `body` the protocol-neutral error object (adapters reshape
+ *  it: OpenAI `{error}`, Anthropic `{type:"error"}`, Responses `{error}`). */
+export class RequestError extends Error {
+  constructor(
+    readonly status: number,
+    message: string,
+    readonly body: Record<string, unknown> = { message },
+  ) {
+    super(message);
+  }
+}
+
+/** What flows back to the caller while a request is in flight: cancellation,
+ *  tracing, and (for streaming surfaces) events + live usage. */
+export interface RunControl {
+  signal?: AbortSignal;
+  trace?: PromptResponseTrace;
+  onEvents?(events: readonly CompletionEvent[]): void | false | Promise<void | false>;
+  onUsageProgress?(usage: Readonly<CompletionUsage>): void;
+}

package/src/serve/preparation.ts

@@ -0,0 +1,33 @@
+/** Host-side admission before grammar/media preparation can allocate native
+ * resources. The implementation shares the generation execution domain. */
+export interface PreparationExecutor {
+  reserve?(kind: "media" | "constraint", signal?: AbortSignal): Promise<import("../contracts/resources").DisposableResource>;
+  run<T>(prepare: () => Promise<T>, signal?: AbortSignal): Promise<T>;
+}
+
+import { AdmissionPool } from "../engine/admission";
+import { CancellationSource } from "../engine/cancellation";
+
+export async function acquireReservation(pool: Pick<AdmissionPool, "acquire">, signal?: AbortSignal) {
+  signal?.throwIfAborted();
+  const cancellation = new CancellationSource();
+  const abort = () => cancellation.cancel("requested");
+  signal?.addEventListener("abort", abort, { once: true });
+  try { return await pool.acquire(cancellation); }
+  finally { signal?.removeEventListener("abort", abort); }
+}
+
+/** One retained media preparation at a time; grammar rows may form a batch. */
+export function createPreparationExecutor(
+  run: PreparationExecutor["run"], grammarCapacity: number,
+): PreparationExecutor & { close(): void } {
+  const media = new AdmissionPool(1);
+  const constraints = new AdmissionPool(Math.max(1, grammarCapacity));
+  return {
+    run,
+    async reserve(kind, signal) {
+      return acquireReservation(kind === "media" ? media : constraints, signal);
+    },
+    close() { media.close(); constraints.close(); },
+  };
+}

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;
+}

package/src/serve/request-plan.ts

@@ -0,0 +1,162 @@
+import { snapshotGenerationPolicy } from "../backends/mlx/request-policy";
+import type { GenerateOptions } from "../generate";
+import type { RequestShape } from "./generation-gateway";
+import type { DisposableResource } from "../contracts/resources";
+import { disposeResources, ownResource } from "../engine/resources";
+export type { DisposableResource } from "../contracts/resources";
+
+export class RequestOwnership {
+  #owner = ownResource(new Set<DisposableResource>(), disposeResources);
+  #reservations = ownResource(new Set<DisposableResource>(), disposeResources);
+
+  /** Reservations live until completion, including after native ownership transfers. */
+  retain<T extends DisposableResource>(reservation: T): T {
+    this.#reservations.borrow().add(reservation);
+    return reservation;
+  }
+
+  own<T extends DisposableResource | null | undefined>(resource: T): T {
+    if (resource) this.#owner.borrow().add(resource);
+    return resource;
+  }
+
+  transfer(): void {
+    this.#owner.transfer();
+  }
+
+  dispose(): void {
+    disposeResources([{ dispose: () => this.#owner.close() }, { dispose: () => this.#reservations.close() }]);
+  }
+}
+
+export interface RequestPlanInput {
+  promptIds: number[];
+  options: GenerateOptions & { stopSequences: string[] };
+  requestedMaxTokens: number;
+  maxSafeContext: number;
+  stream: boolean;
+  wantLogprobs: boolean;
+  topLogprobs: number;
+  adapterIds: string[];
+  hasVision: boolean;
+  userSeed: boolean;
+  hasGrammar: boolean;
+  hasDraft: boolean;
+  ownership: RequestOwnership;
+}
+
+export class RequestRejection {
+  readonly ok = false;
+  readonly status = 400;
+  readonly error: {
+    message: string;
+    type: "memory_admission";
+    code: "context_over_budget";
+  };
+
+  readonly #ownership: RequestOwnership;
+
+  constructor(message: string, ownership: RequestOwnership) {
+    this.error = {
+      message,
+      type: "memory_admission",
+      code: "context_over_budget",
+    };
+    this.#ownership = ownership;
+  }
+
+  dispose(): void {
+    this.#ownership.dispose();
+  }
+}
+
+export class RequestPlan {
+  readonly ok = true;
+  readonly promptIds: number[];
+  readonly options: GenerateOptions & { stopSequences: string[] };
+  readonly shape: RequestShape;
+  readonly captureLogprobs: boolean;
+  readonly wantLogprobs: boolean;
+  readonly topLogprobs: number;
+  readonly #ownership: RequestOwnership;
+
+  constructor(input: {
+    promptIds: number[];
+    options: GenerateOptions & { stopSequences: string[] };
+    shape: RequestShape;
+    captureLogprobs: boolean;
+    wantLogprobs: boolean;
+    topLogprobs: number;
+    ownership: RequestOwnership;
+  }) {
+    this.promptIds = input.promptIds;
+    this.options = input.options;
+    this.shape = input.shape;
+    this.captureLogprobs = input.captureLogprobs;
+    this.wantLogprobs = input.wantLogprobs;
+    this.topLogprobs = input.topLogprobs;
+    this.#ownership = input.ownership;
+  }
+
+  transferOwnership(): void {
+    this.#ownership.transfer();
+  }
+
+  dispose(): void {
+    this.#ownership.dispose();
+  }
+}
+
+export type PlanRequestResult = RequestPlan | RequestRejection;
+
+/** Compile the model-independent part of request execution into one value.
+ * Prompt/media construction happens before this boundary; admission, capture
+ * options, adapter selection, and lane shape are derived here once for chat
+ * and raw completions. */
+export function planRequest(input: RequestPlanInput): PlanRequestResult {
+  const available = input.maxSafeContext - input.promptIds.length;
+  if (available < 1) {
+    return new RequestRejection(
+      `prompt is ${input.promptIds.length} tokens but the memory budget caps ` +
+        `safe context at ${input.maxSafeContext} — no room to generate; ` +
+        "shorten the prompt or raise --memory-budget",
+      input.ownership,
+    );
+  }
+
+  const options = { ...input.options };
+  options.maxTokens = Math.min(input.requestedMaxTokens, available);
+  options.seedWasExplicit = input.userSeed;
+  if (input.adapterIds.length) options.adapters = [...input.adapterIds];
+  const captureLogprobs = !input.stream && (input.wantLogprobs || input.topLogprobs > 0);
+  if (captureLogprobs) {
+    options.logprobs = input.wantLogprobs;
+    options.topLogprobs = input.topLogprobs;
+  }
+
+  const shape: RequestShape = {
+    hasVision: input.hasVision,
+    hasAdapters: !!options.adapters?.length,
+    hasRepetitionPenalty: !!options.repetitionPenalty,
+    hasLogitsExtras: !!(
+      options.minP || options.xtcProbability || options.logitBias ||
+      options.presencePenalty || options.frequencyPenalty
+    ),
+    wantsLogprobs: captureLogprobs,
+    userSeed: input.userSeed,
+    kvQuant: !!(options.kvConfig?.length || options.kvBits),
+    turboQuant: !!options.turboQuant,
+    hasGrammar: input.hasGrammar,
+    hasDraft: input.hasDraft,
+  };
+
+  return new RequestPlan({
+    promptIds: Object.freeze([...input.promptIds]) as number[],
+    options: snapshotGenerationPolicy(options),
+    shape: Object.freeze(shape),
+    captureLogprobs,
+    wantLogprobs: input.wantLogprobs,
+    topLogprobs: input.topLogprobs,
+    ownership: input.ownership,
+  });
+}