@otto-code/brain 0.8.9 → 0.8.12

package/dist/service/scheduler.d.ts

@@ -1,25 +1,80 @@
 /**
- * Cooperative single-GPU scheduler with per-model concurrency.
+ * Cooperative single-GPU scheduler: one resident model, many concurrent chats.
  *
- * Only one model is resident at a time. Completion requests are queued rather
- * than refused for "wrong / no model loaded".
+ * There is exactly one place a job is ever started - `#dispatch`. Every event
+ * that could change the answer (a submit, a job settling, a model finishing its
+ * load, a slot poll) calls it again, and it re-derives the whole decision from
+ * scratch. There is no worker pool, no parked promise, no wake generation and
+ * no claim lock: the dispatcher is single-threaded by construction (`#busy`),
+ * so the state it reads cannot move underneath it.
  *
- *  - Requests for the resident model run concurrently, up to that model's
- *    `parallelSlots` (the same number of sequence slots llama-server was
- *    launched with - sending more would only queue inside llama-server). Extra
- *    same-model requests wait for a free slot; they never trigger a load.
- *  - Requests for a *different* model wait for a model switch. When the resident
- *    model's current batch drains, the scheduler switches and serves the other
- *    model's batch - so two clients wanting different models share the GPU by
- *    taking turns.
+ * The variables it arbitrates, in the order it applies them:
  *
- * Fairness: after a model finishes a turn, the next turn prefers a *different*
- * model when one is waiting, so a steady stream for model A cannot starve model
- * B. A turn is a snapshot: requests that arrive for a model mid-turn wait for
- * its next turn.
+ *  1. RESIDENCY. Only one model is resident, so `#running` never mixes models.
+ *     A switch happens only from a fully drained engine. Nothing that reaches
+ *     the scheduler is ever refused for "wrong model loaded" - it is queued
+ *     until its model is resident. Whether a switch is *allowed* at all is
+ *     decided upstream by `decideModelGate`: with `config.lockModel` on, the
+ *     host serves one pinned model and a request naming another is refused
+ *     there with a 409, so the scheduler never sees it.
+ *
+ *  2. EXCLUSIVITY. Calibrate, sweep and benchmark own the engine alone: the
+ *     engine drains before one starts, and nothing is admitted while it runs.
+ *
+ *  3. THE TURN, and what may join it. A turn takes every queued job of its
+ *     model up to the first exclusive operation. When that batch drains the
+ *     turn stays OPEN: a same-model job arriving a moment later joins it and
+ *     takes a free slot immediately. That is the difference between two chats
+ *     each holding their own slot and two chats trading one slot back and
+ *     forth - the latter is what a closed per-turn batch produces, because
+ *     agentic traffic arrives one request per chat per turn and never as a
+ *     burst.
+ *
+ *  4. FAIRNESS. A turn stops absorbing the moment the queue head is another
+ *     model's job or an exclusive operation. Since the queue is FIFO, "another
+ *     model is waiting" means it is at the head, so a steady stream for model
+ *     A cannot starve model B: B waits exactly as long as A's already-running
+ *     jobs take, then the turn retires and B is picked (preferring a model
+ *     other than the one that just ran).
+ *
+ *  5. CAPACITY. `parallelSlots` is the number of sequence slots llama-server
+ *     was launched with and therefore the KV pool we own; `#running.size`
+ *     against it is the hard bound. The measured free-slot count (llama-server's
+ *     own `/slots`) is the second bound, and it exists to notice slots taken by
+ *     traffic Otto did not schedule, or an engine mid-eviction. The two are
+ *     combined as `min(ceiling - running, measured)` and never both subtracted:
+ *     the engine's idle count already excludes our running jobs, so subtracting
+ *     them from it a second time leaves zero capacity with a single chat live
+ *     and silently serializes every other chat behind it.
+ *
+ * Ordering within a turn is session-affine: a job whose session already holds a
+ * slot goes first, because that chat's KV state is what llama-server's
+ * longest-common-prefix slot selection just filled, so it pays no re-prefill.
+ * Session identity is the standard `prompt_cache_key`; clients that omit it get
+ * plain FIFO.
+ *
+ *  6. OWNERSHIP. A slot's KV belongs to the chat that last ran on it.
+ *     llama-server never clears a released slot's prompt, so a slot handed to a
+ *     different chat would still hold the previous chat's KV - and that is
+ *     exactly the cross-chat bleed users see in thinking blocks. The scheduler
+ *     therefore tracks `slotId -> session` for every slot it names (the
+ *     `prompt_cache_key` that owned it, or `null` for keyless jobs) and erases
+ *     the slot before handoff to a different owner. Same session reusing its
+ *     own slot keeps its KV (that is the cache the whole point of `--cache-ram`
+ *     is to protect). The eraser is injected, because the engine endpoint is a
+ *     transport detail the scheduler stays agnostic of. The map is wiped when
+ *     the engine reloads a model: slot ids do not survive a switch, and a
+ *     fresh engine has no stale KV.
+ *
+ *     The erase is awaited BEFORE the job's run() reaches the engine: the
+ *     engine runs tasks in arrival order, so the erase must land in its queue
+ *     ahead of the completion that follows. Firing the erase and the
+ *     completion back-to-back as independent requests would let the engine see
+ *     the completion first, run it on the dirty slot, and only then erase the
+ *     KV the completion just produced - the fix inverting into the bug.
  *
  * The scheduler is transport-agnostic: a job is a resolved catalog model plus a
- * `run()` that does the proxying, which keeps the turn logic unit-testable.
+ * `run()` that does the proxying, which keeps the logic unit-testable.
  */
 import type { Model } from "../types.js";
 import type { Profile } from "../config/schema.js";
@@ -29,6 +84,37 @@ export interface SchedulerSupervisor {
     model: Model | null;
     profile: Profile | null;
 }
+/**
+ * The answer to a live slot measurement: how many sequence slots are free, and
+ * optionally WHICH ones. `idle` drives admission (see CAPACITY); `ids` lets
+ * `#pass` hand a distinct slot id to each job it admits, so the router can pin
+ * a completion to the exact slot this sample saw free. `ids` is absent when the
+ * engine reports a count but no per-slot rows - then admission still works, but
+ * no honest pin can be named and the affected requests run unpinned.
+ */
+export interface SlotMeasurement {
+    idle: number;
+    ids?: number[];
+}
+/**
+ * Erase one engine slot's retained KV state before the next task may land on
+ * it. This is the fix for the cross-chat KV bleed (see OWNERSHIP below):
+ * llama.cpp never clears a slot's prompt on release, so a slot that served one
+ * chat hands that chat's KV to whoever is pinned to the slot next. The caller
+ * (the router) performs the engine-side erase; the scheduler only decides WHEN
+ * a handoff happens and reports it, because only it knows which session was
+ * admitted to which slot.
+ *
+ * The promise resolves (never rejects) once the engine has ACKNOWLEDGED the
+ * erase. That acknowledgment matters: the engine runs tasks in arrival order,
+ * so the erase must sit in its queue ahead of the completion the scheduler is
+ * about to post. The scheduler awaits the promise before `run()` reaches the
+ * engine, which is what keeps the order honest across the HTTP boundary.
+ *
+ * Must never throw or reject: an erase failure degrades to the old (bleedy)
+ * behavior but must never fail the completion that is about to run.
+ */
+export type SlotEraser = (slotId: number) => Promise<void>;
 export interface SchedulerOptions {
     supervisor: SchedulerSupervisor;
     loadModel: (model: Model) => Promise<void>;
@@ -40,11 +126,89 @@ export interface SchedulerOptions {
      * the moment it becomes true rather than up to a poll later.
      */
     onChange?: (() => void) | null;
+    /**
+     * Live slot measurement: how many sequence slots llama-server actually has
+     * free right now, or null when unknown (server not ready, sample failed).
+     * See CAPACITY in the file header for how it combines with `parallelSlots`.
+     * Absent means "trust the profile count".
+     *
+     * May also answer with the free slots NAMED, not just counted:
+     * `{ idle, ids }` (the engine's idle slot ids from `/slots`). The count is
+     * what admission needs; the ids are what `onSlotFree` hands to admitted
+     * completions so the router can pin each one to the slot this very sample
+     * saw free. A plain number still works - it just yields no pin data.
+     */
+    freeSlots?: (() => Promise<SlotMeasurement | number | null> | SlotMeasurement | number | null) | null;
+    /**
+     * How long to wait before re-checking the engine when capacity is zero and
+     * no job of ours is running - the only state where nothing else will wake
+     * the dispatcher.
+     */
+    slotPollMs?: number;
+    /**
+     * Erase an engine slot's KV before it is handed to a different chat (see
+     * OWNERSHIP). Injected rather than built in: the engine endpoint is a
+     * transport detail, and its availability is runtime-dependent (the
+     * llama-server build must support `POST /slots?action=erase` - it does not
+     * on a server launched without `--slot-save-path`). Absent (null) means
+     * ownership is tracked and reported but nothing is erased - the old
+     * behavior, for runtimes where the engine cannot wipe a slot.
+     */
+    eraseSlot?: SlotEraser | null;
+}
+/** Work kinds sharing the single resident model. Operations own a full turn. */
+export type SchedulerJobKind = "completion" | "calibrate" | "sweep" | "benchmark";
+export interface SchedulerSubmitOptions {
+    kind?: SchedulerJobKind;
+    /** An operation never shares its model turn with inference. */
+    exclusive?: boolean;
+    /** Called exactly when this job starts running, before `run()` is awaited. */
+    onStart?: (() => void) | null;
+    /**
+     * The chat's stable identity (its `prompt_cache_key`). Jobs from a session
+     * that already holds a slot run next in line for their model, so a resident
+     * chat's KV cache is reused instead of evicted by a different chat. Absent
+     * (third-party clients) means no affinity - plain FIFO.
+     */
+    session?: string | null;
+    /**
+     * Called exactly once, when this job is admitted to a slot - i.e. the moment
+     * the scheduler's own free-slot measurement just counted it. The argument is
+     * the engine slot id the job may pin to, or null when no slot can be named
+     * honestly (sample failed, engine reports no per-slot rows, or the engine's
+     * slot pool is fully busy).
+     *
+     * Why the scheduler, not the job, names the slot: a job that samples `/slots`
+     * itself at dispatch time can race another job admitted in the same pass -
+     * neither has reached the engine yet, so both see the same slot free and both
+     * pin it, and the engine's busy-pinned-slot deferral turns the pair back into
+     * a serial queue. The pin must be drawn from the very measurement that
+     * admitted the job, one distinct id per job, so the ids it hands out are
+     * exactly the slots the engine still has free at the moment of admission.
+     *
+     * Exclusive operations are never called: they run alone on a drained engine
+     * and their attribution is not per-slot.
+     */
+    onSlotFree?: ((slotId: number | null) => void) | null;
 }
-/** A queued completion request bound to a resolved catalog model. */
+/** A queued request or host operation bound to a resolved catalog model. */
 export interface QueuedJob {
     modelId: string;
     model: Model;
+    kind: SchedulerJobKind;
+    exclusive: boolean;
+    session: string | null;
+    onStart: (() => void) | null;
+    onSlotFree: ((slotId: number | null) => void) | null;
+    /**
+     * The engine slot this job was pinned to at admission, or null when none was
+     * named. Kept on the job so a later pass can exclude it from the ids it hands
+     * out: a job admitted moments ago may not have reached llama-server yet, so
+     * the next `/slots` sample can still report its slot idle and offer it to a
+     * second job. Both would then pin the same slot, and the engine's defer-when-
+     * busy behavior would serialize the pair onto it while another slot sat empty.
+     */
+    slotId: number | null;
     run: () => Promise<unknown>;
     resolve: (value: unknown) => void;
     reject: (error: unknown) => void;
@@ -52,30 +216,69 @@ export interface QueuedJob {
 export interface SchedulerStats {
     queued: number;
     waiting: Record<string, number>;
+    /** Stable ids for host surfaces. Display names are not scheduler keys. */
+    waitingModelIds: Record<string, number>;
     lastTurn: string | null;
+    active: {
+        modelId: string;
+        kind: Exclude<SchedulerJobKind, "completion">;
+    } | null;
 }
 export declare class Scheduler {
     #private;
     supervisor: SchedulerSupervisor;
     loadModel: (model: Model) => Promise<void>;
     logger: ((message: string) => void) | null;
+    /** Submitted, not yet claimed by a turn. */
     queue: QueuedJob[];
     lastTurnId: string | null;
-    pumping: boolean;
+    /**
+     * Per model: the session whose jobs most recently filled its slots. That
+     * session's KV state is the one llama-server's LCP selection will match
+     * against, so its queued jobs go first on the next dispatch.
+     */
+    hotSessions: Map<string, string>;
+    /** The exclusive operation in flight, if any. Reported by `stats()`. */
+    activeJob: QueuedJob | null;
     onChange: (() => void) | null;
-    constructor({ supervisor, loadModel, logger, onChange }: SchedulerOptions);
+    freeSlots: SchedulerOptions["freeSlots"];
+    slotPollMs: number;
+    eraseSlot: SlotEraser | null;
+    /**
+     * Engine slot -> the session that last ran on it (see OWNERSHIP). The value
+     * is null for keyless jobs: a keyless job cannot prove ownership of anything
+     * after it settles, so the next keyed chat landing on that slot is treated as
+     * a handoff and the slot is erased for it. Wiped wholesale whenever the
+     * engine reloads a model - slot ids do not survive the relaunch.
+     */
+    slotOwners: Map<number, string | null>;
+    constructor({ supervisor, loadModel, logger, onChange, freeSlots, slotPollMs, eraseSlot, }: SchedulerOptions);
     /** Id of the model that is actually loaded and ready, or null. */
     get loadedId(): string | null;
-    /** How many requests may run at once against the resident model. */
+    /** How many requests may run at once against the resident model (static ceiling). */
     get concurrency(): number;
     /**
      * Queue a job for an already-resolved catalog model. `run` is invoked once
      * that model is the resident one; the returned promise settles when run does.
-     * Pumping is deferred a microtask so a burst of requests submitted together
-     * share one turn rather than the first one snapshotting a turn by itself.
+     * Dispatch is deferred a microtask so a burst of requests submitted together
+     * shares one turn rather than the first one taking a turn by itself.
+     */
+    submit(model: Model, run: () => Promise<unknown>, { kind, exclusive, onStart, session, onSlotFree, }?: SchedulerSubmitOptions): Promise<unknown>;
+    /**
+     * Drop every recorded slot owner. The engine's slots do not survive a model
+     * (re)launch, so their owners do not either - a stale entry would make the
+     * next admission think a FRESH slot still holds a previous chat's KV and
+     * erase it (an unnecessary re-prefill), or, worse, let a keyless job be
+     * mistaken for the owner of a slot it is not.
+     *
+     * Called from two places: `#pass` clears the owners the moment a turn begins
+     * on a different model than the one it last served (a model switch, where
+     * the scheduler itself sees the relaunch), and the router calls it on the
+     * supervisor's `starting` state - the relaunch that keeps the SAME model
+     * resident (a live profile edit), where the turn never changes and only the
+     * supervisor says the slots are gone. Both paths are idempotent.
      */
-    submit(model: Model, run: () => Promise<unknown>): Promise<unknown>;
-    pump(): Promise<void>;
+    forgetSlots(): void;
     /** Queue snapshot for the status endpoint / UI. */
     stats(): SchedulerStats;
 }