@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/execution/row-sampling.ts

@@ -0,0 +1,49 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { TokenLogprobs } from "../contracts/portable/generation";
+import { readStepExtras, stepExtrasArrays, disposeStepExtras,
+  type DeviceStepSampler, type StepExtras } from "../sampling/index";
+import { disposeResources, cleanupFailure } from "../runtime/resources";
+
+type TokenSink = (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>;
+
+/** Request-owned sampling and output metadata. Execution groups receive token
+ * tensors and an output callback; they need no logprob or retention policy. */
+export function createRowSampling(
+  sampler: DeviceStepSampler, onToken: TokenSink,
+) {
+  if (!sampler.capturesLogprobs) return {
+    sample: (logits: MlxArray, step: number) => sampler.sample(logits, step).token,
+    plainGreedy: sampler.isPlainGreedy, onToken, dispose: () => sampler.dispose(),
+  };
+
+  // A pipelined group can sample the next token before emitting the previous
+  // one. Keep captures by request-local step, independently of row position.
+  const pending = new Map<number, StepExtras>();
+  let emitted = 0;
+  return {
+    sample(logits: MlxArray, step: number): MlxArray {
+      const result = sampler.sample(logits, step);
+      try {
+        ops.asyncEvalAll([result.token, ...stepExtrasArrays(result.extras)]);
+        if (result.extras) pending.set(step, result.extras);
+        return result.token;
+      } catch (error) {
+        return cleanupFailure(error, () => disposeResources([result.token,
+          { dispose: () => disposeStepExtras(result.extras) }]));
+      }
+    },
+    // The vectorized token-only shortcut bypasses the per-request sampler.
+    plainGreedy: false,
+    onToken(token: number) {
+      const extras = pending.get(emitted) ?? null;
+      pending.delete(emitted++);
+      return onToken(token, readStepExtras(extras));
+    },
+    dispose() {
+      const extras = [...pending.values()];
+      pending.clear();
+      disposeResources([sampler, ...extras.map(value => ({ dispose: () => disposeStepExtras(value) }))]);
+    },
+  };
+}

package/src/execution/scheduler.ts

@@ -0,0 +1,92 @@
+import type { ExecutionGroup, SchedulingClock } from "../contracts/portable/scheduling";
+import { disposeResources } from "../runtime/resources";
+
+/** Continuous scheduling policy. Each advance is a backend-defined safe unit;
+ * this driver neither handles native state nor forces pending work to evaluate. */
+export async function driveExecutionGroup(
+  group: ExecutionGroup, clock: SchedulingClock, yieldAfterPreparation = false,
+): Promise<void> {
+  let execution: (() => void) | undefined;
+  let residency: (() => void) | undefined;
+  let lastYield = 0;
+  let terminalError: unknown = new Error("scheduler closed");
+  const release = () => {
+    const releases = [execution, residency];
+    execution = residency = undefined;
+    disposeResources(releases.filter((fn): fn is () => void => !!fn).map((dispose) => ({ dispose })));
+  };
+  try {
+    while (!group.closed) {
+      group.pruneCancelled();
+      let held = group.admissionHeld;
+      if (!group.active && !group.preparing && (held || (!group.queued && !group.pendingTasks))) {
+        release();
+        await group.waitForWork();
+        continue;
+      }
+      residency ??= group.reserveResidency();
+      if (!execution && group.acquireExecution) execution = await group.acquireExecution();
+      if (group.closed) break;
+
+      // One preparation task per iteration. An exclusive model mutation
+      // keeps its existing drain priority; preparation cannot prolong it.
+      if (!held && group.pendingTasks && group.advanceTask) {
+        await group.advanceTask();
+        if (group.closed) break;
+        group.pruneCancelled();
+        held = group.admissionHeld;
+      }
+
+      if (!group.preparing && !held && group.queued && group.active < group.maxActive)
+        group.admitNext();
+      while (group.preparing && group.canPrepareMore && !held && group.queued &&
+          group.active + (group.preparingRows ?? 1) < group.maxActive && group.canBurst() &&
+          (group.preparingTokens ?? 0) + (group.nextPreparationTokens ?? 0) <= (group.maxPreparationTokens ?? Infinity)) {
+        if (!group.admitNext()) break;
+      }
+      const mixed = group.preparing && group.active ? group.mixedPreparation : undefined;
+      if (mixed && group.advanceMixed) {
+        // Reserve the method's running-token demand first; prompt work fills
+        // the rest. A budget smaller than one token per row still progresses.
+        const budget = Math.max(group.maxIterationTokens ?? Infinity,
+          mixed.runningTokens + mixed.minimumPreparationTokens);
+        await group.advanceMixed(budget);
+        lastYield = clock.now();
+        await clock.yield();
+        continue;
+      }
+      if (group.preparing) {
+        const activeBefore = group.active;
+        await group.advancePreparation();
+        // A bounded preparation may publish output before its state work ends.
+        // Flush it before the next work unit; queued short admissions still
+        // form a group across yields. Re-enter policy to observe new work,
+        // cancellation, admission holds and shutdown.
+        if (group.preparationPublishedOutput ||
+            (yieldAfterPreparation && !activeBefore && !group.preparing && group.active === 1 && !group.queued)) {
+          lastYield = clock.now();
+          await clock.yield();
+          continue;
+        }
+      }
+
+      // Finish short admissions before decode; long preparation interleaves
+      // with active work because the group still reports preparing=true.
+      if (!group.preparing && !held && group.queued && group.active < group.maxActive && group.canBurst())
+        continue;
+      if (group.active) {
+        try { await group.advance(); }
+        catch (error) { group.failActive(error); }
+      }
+      // Preserve the serial responsiveness budget for the single active run.
+      if (group.active === 1 && !group.queued && !group.preparing &&
+          !group.admissionHeld && clock.now() - lastYield < 25) continue;
+      lastYield = clock.now();
+      await clock.yield();
+    }
+  } catch (error) {
+    terminalError = error;
+  } finally {
+    disposeResources([{ dispose: () => group.failAll(terminalError) }, { dispose: release }]);
+  }
+}

package/src/execution/session.ts

@@ -0,0 +1,247 @@
+import type {
+  Cancellation, CancelReason, ExecutionPlan, GenerationEvent, GenerationOutcome,
+  GenerationOutput, GenerationSession, MethodResult, MethodRun, RunControl,
+  SessionState, Timer, TokenLogprobs,
+} from "../contracts/portable/generation";
+import { CancellationSource, throwIfCancelled } from "./cancellation";
+
+export interface SessionResources {
+  readonly timer: Timer;
+  readonly idleTimeoutMs: number;
+  readonly maxQueuedEvents: number;
+  readonly maxQueuedTokens: number;
+  readonly maxTopLogprobs: number;
+  reserveOutput(tokens: number): { buffer: Uint32Array; release(): void };
+}
+
+/** Lifecycle and bounded delivery only. No graph, cache, sampler, or method branches. */
+export class EngineSession<Metrics> implements GenerationSession<Metrics> {
+  readonly result: Promise<GenerationOutcome<Metrics>>;
+  readonly events: AsyncIterable<GenerationEvent>;
+  #state: SessionState = "created";
+  #cancel = new CancellationSource();
+  #settle!: (outcome: GenerationOutcome<Metrics>) => void;
+  #started = false;
+  #readerClaimed = false;
+  #reading = false;
+  #publishing = false;
+  #queue: GenerationEvent[] = [];
+  #queuedTokens = 0;
+  #committed = 0;
+  #sequence = 0;
+  #wakeReader: (() => void) | undefined;
+  #wakeProducer: (() => void) | undefined;
+  #stopTimer: (() => void) | undefined;
+  #unsubscribe: (() => void) | undefined;
+  #buffer: Uint32Array | undefined;
+  #limit = 0;
+  #outputStopped = false;
+
+  constructor(
+    readonly id: string,
+    private readonly prepare: (cancellation: Cancellation) => Promise<ExecutionPlan<Metrics>>,
+    private readonly control: RunControl,
+    private readonly resources: SessionResources,
+    private readonly onTerminal: () => void,
+  ) {
+    this.result = new Promise((resolve) => { this.#settle = resolve; });
+    this.events = { [Symbol.asyncIterator]: () => this.#iterator() };
+    const unsubscribe = control.cancellation?.subscribe((reason) => { void this.cancel(reason); });
+    if (this.#state === "terminal") unsubscribe?.();
+    else this.#unsubscribe = unsubscribe;
+    if (this.#state !== "terminal") {
+      if (control.output !== "stream") this.#start();
+      else this.#armIdle();
+    }
+  }
+
+  get state(): SessionState { return this.#state; }
+
+  async cancel(reason: CancelReason = "requested"): Promise<void> {
+    if (this.#state !== "terminal") {
+      this.#cancel.cancel(reason);
+      this.#queue = [];
+      this.#queuedTokens = 0;
+      this.#wakeReader?.();
+      this.#wakeProducer?.();
+      this.#stopTimer?.();
+      if (!this.#started) this.#finish({ status: "cancelled", reason, committedTokens: 0 });
+    }
+    await this.result;
+  }
+
+  close(): Promise<void> { return this.cancel("consumer_closed"); }
+
+  #armIdle(): void {
+    this.#stopTimer?.();
+    this.#stopTimer = this.resources.timer.after(this.resources.idleTimeoutMs,
+      () => { void this.cancel("consumer_idle"); });
+  }
+
+  #iterator(): AsyncIterator<GenerationEvent> {
+    if (this.control.output !== "stream") throw new Error("collect sessions return output through result");
+    if (this.#readerClaimed) throw new Error("generation events have one consumer");
+    this.#readerClaimed = true;
+    return {
+      next: async () => {
+        if (this.#reading) throw new Error("concurrent next() is not supported");
+        this.#reading = true;
+        this.#stopTimer?.();
+        try {
+          this.#start();
+          while (!this.#queue.length && this.#state !== "terminal") {
+            await new Promise<void>((resolve) => { this.#wakeReader = resolve; });
+            this.#wakeReader = undefined;
+          }
+          const event = this.#queue.shift();
+          if (!event) return { done: true, value: undefined };
+          if (event.type === "committed") this.#queuedTokens -= event.tokenIds.length;
+          this.#wakeProducer?.();
+          if (this.#state !== "terminal") this.#armIdle();
+          return { done: false, value: event };
+        } finally { this.#reading = false; }
+      },
+      return: async () => { await this.close(); return { done: true, value: undefined }; },
+      throw: async (error) => { await this.close(); throw error; },
+    };
+  }
+
+  #start(): void {
+    if (this.#started || this.#state === "terminal") return;
+    this.#started = true;
+    void this.#execute();
+  }
+
+  async #execute(): Promise<void> {
+    let run: MethodRun<Metrics> | undefined;
+    let reservation: ReturnType<SessionResources["reserveOutput"]> | undefined;
+    let result: MethodResult<Metrics> | undefined;
+    let failure: unknown;
+    let failed = false;
+    try {
+      this.#state = "preparing";
+      throwIfCancelled(this.#cancel);
+      const plan = await this.prepare(this.#cancel);
+      throwIfCancelled(this.#cancel);
+      if (!Number.isSafeInteger(plan.outputTokenLimit) || plan.outputTokenLimit < 0)
+        throw new Error("invalid planned output token limit");
+      this.#limit = plan.outputTokenLimit;
+      if (this.control.output === "collect") {
+        reservation = this.resources.reserveOutput(this.#limit);
+        this.#buffer = reservation.buffer;
+      }
+      run = await plan.method.createRun(this.#cancel);
+      throwIfCancelled(this.#cancel);
+      this.#state = "running";
+      const output: GenerationOutput = {
+        commit: (ids, logprobs) => this.#publish(ids, logprobs),
+        progress: (completed, total) => this.#progress(completed, total),
+      };
+      result = await run.execute(output, this.#cancel);
+    } catch (error) { failure = error; failed = true; }
+    this.#state = "settling";
+    // Terminal means resources are safe to release, not merely that JS stopped reading.
+    let cleanupFailed = false;
+    let cleanupError: string | undefined;
+    try { await run?.close(); }
+    catch (error) {
+      cleanupError = error instanceof Error ? error.message : String(error);
+      if (!failed) failure = error;
+      failed = true;
+      cleanupFailed = true;
+    }
+    reservation?.release(); // Completed output ownership transfers to the caller.
+    const common = { committedTokens: this.#committed,
+      ...(this.#buffer ? { output: this.#buffer.subarray(0, this.#committed) } : {}) };
+    if (this.#cancel.reason && !cleanupFailed) {
+      this.#finish({ status: "cancelled", reason: this.#cancel.reason, ...common });
+    } else if (failed) {
+      this.#finish({ status: "failed", error: { code: cleanupFailed ? "cleanup_failed" : "execution_failed",
+        message: failure instanceof Error ? failure.message : String(failure),
+        ...(cleanupError === undefined ? {} : { cleanupError }) }, ...common });
+    } else {
+      this.#finish({ status: "completed", result: result!, ...common });
+    }
+  }
+
+  async #publish(ids: readonly number[], logprobs?: readonly (TokenLogprobs | undefined)[]): Promise<void | false> {
+    if (this.#publishing) throw new Error("method must await each publication");
+    this.#publishing = true;
+    try {
+      throwIfCancelled(this.#cancel);
+      if (this.#state !== "running") throw new Error("publication outside method execution");
+      if (this.#outputStopped) throw new Error("method published after the consumer requested a stop");
+      if (ids.length > this.#limit - this.#committed) throw new Error("method exceeded planned output limit");
+      if (logprobs && logprobs.length !== ids.length) throw new Error("logprobs must align with committed tokens");
+      if (logprobs?.some((item) => (item?.top?.length ?? 0) > this.resources.maxTopLogprobs))
+        throw new Error("top logprobs exceed session delivery capacity");
+      for (const id of ids)
+        if (!Number.isInteger(id) || id < 0 || id > 0xffffffff) throw new Error("invalid committed token id");
+      if (this.control.output === "callback") {
+        this.#committed += ids.length;
+        const result = await this.control.onTokens(ids, logprobs);
+        if (result === false) this.#outputStopped = true;
+        return result;
+      }
+      if (this.#buffer) {
+        if (logprobs) throw new Error("token logprobs require stream consumption");
+        this.#buffer.set(ids, this.#committed);
+        this.#committed += ids.length;
+        return;
+      }
+      for (let offset = 0; offset < ids.length;) {
+        const length = Math.min(ids.length - offset, this.resources.maxQueuedTokens);
+        await this.#space(length);
+        // Snapshot producer-owned arrays; later scratch-buffer reuse cannot alter events.
+        const event: GenerationEvent = { type: "committed", sequence: this.#sequence++,
+          tokenIds: ids.slice(offset, offset + length),
+          ...(logprobs ? { logprobs: logprobs.slice(offset, offset + length).map((item) => item && ({
+            ...(item.logprob === undefined ? {} : { logprob: item.logprob }),
+            ...(item.top ? { top: item.top.map((entry) => ({ id: entry.id, logprob: entry.logprob })) } : {}),
+          })) } : {}) };
+        this.#queue.push(event);
+        this.#queuedTokens += length;
+        this.#committed += length;
+        offset += length;
+        this.#wakeReader?.();
+      }
+    } finally { this.#publishing = false; }
+  }
+
+  async #progress(completed: number, total?: number): Promise<void> {
+    if (this.#publishing) throw new Error("method must await each publication");
+    this.#publishing = true;
+    try {
+      throwIfCancelled(this.#cancel);
+      if (this.#state !== "running") throw new Error("publication outside method execution");
+      if (!Number.isFinite(completed) || completed < 0 ||
+          (total !== undefined && (!Number.isFinite(total) || total < completed)))
+        throw new Error("invalid progress");
+      if (this.control.output !== "stream") return;
+      await this.#space(0);
+      this.#queue.push({ type: "progress", completed, ...(total === undefined ? {} : { total }) });
+      this.#wakeReader?.();
+    } finally { this.#publishing = false; }
+  }
+
+  async #space(tokens: number): Promise<void> {
+    while (this.#queue.length >= this.resources.maxQueuedEvents ||
+           this.#queuedTokens + tokens > this.resources.maxQueuedTokens) {
+      throwIfCancelled(this.#cancel);
+      this.#armIdle();
+      await new Promise<void>((resolve) => { this.#wakeProducer = resolve; });
+      this.#wakeProducer = undefined;
+    }
+    throwIfCancelled(this.#cancel);
+  }
+
+  #finish(outcome: GenerationOutcome<Metrics>): void {
+    if (this.#state === "terminal") return;
+    this.#state = "terminal";
+    this.#stopTimer?.();
+    this.#unsubscribe?.();
+    this.#settle(outcome);
+    this.#wakeReader?.();
+    this.onTerminal();
+  }
+}