diff --git a/docs/architecture.mdx b/docs/architecture.mdx index f7e7ffed..2a8a929e 100644 --- a/docs/architecture.mdx +++ b/docs/architecture.mdx @@ -418,6 +418,15 @@ the same node. are the same unit of pending work, so "fewest jobs" is not "least busy." A node running one enormous request looks more idle than a node running two trivial ones. +**A request commits to a node when it arrives.** The proxy picks the least loaded +eligible node at arrival and forwards immediately, and that choice is final: the +request queues on that node however the fleet moves afterwards, so a burst that +lands on a node which then turns out to be busy waits behind it while another node +goes idle. `ollama-proxy` carries an opt-in `--late-binding` mode that instead +waits for a node with a free generation slot, described in +[its component reference](../services/ollama-proxy/README.md). +It is off by default and does not change the behavior described here. + **Model load state is not considered.** Eligibility asks whether a node *has* the model, not whether it is already loaded in memory. PAIR knows which models are loaded, and the interface shows it, but routing does not use it, so a request can diff --git a/services/ollama-proxy/README.md b/services/ollama-proxy/README.md index 35f959b5..7016bc2c 100644 --- a/services/ollama-proxy/README.md +++ b/services/ollama-proxy/README.md @@ -28,6 +28,8 @@ ollama-proxy [flags] | `--ignore-persisted-port` | `false` | Use `--port` even when `proxy-port.json` contains a saved port (used by broker-managed startup) | | `--ipc` | *(empty — use stdio)* | Path to a Unix domain socket or Windows named pipe for IPC | | `--cluster-dir` | *(empty)* | Cluster trust directory (`node.crt`/`node.key` plus trusted pins). Enables the LAN mTLS inference ingress while this node is a cluster member; empty means no ingress and no peer candidates. | +| `--late-binding` | `false` | Hold an inference request until a node has a free generation slot instead of committing it to a node the moment it arrives. See [Late binding](#late-binding-opt-in) for the precondition it depends on. | +| `--node-parallel` | `1` | Concurrent generation slots per node, consulted only with `--late-binding`: `N` for every node, or `=N` to override one. Repeat the flag to mix them. Match each node's `OLLAMA_NUM_PARALLEL`. | | `--log-level` | *(`$NVPAIR_LOG_LEVEL`, else `info`)* | Initial log level: `debug`, `info`, `warn`, or `error`. Changeable at runtime with `log/set-level`. | | `--version` | | Print version and exit | @@ -52,6 +54,24 @@ Node selection: - **Manual**: Use the `node/select` JSON-RPC method to pin traffic to a specific node. A manual pin **overrides the priority list only when that node is eligible** for the requested model. - **Failover**: If the selected node disappears from the discovery set, the proxy falls back to auto-select and emits a `node/selection-changed` notification. A transport error or retryable status, including a model `404` from an advertised owner with stale inventory, steps to the next eligible owner. +### Late binding (opt-in) + +By default the proxy commits a request to a node the moment the request arrives and forwards it immediately. That decision is final: the request then queues on that node however the fleet moves afterwards, so a burst that lands on a node which turns out to be busy waits behind it while another node goes idle. + +`--late-binding` moves the commitment to the moment a node can actually start generating. The same least-estimated-loaded rule decides, but only among eligible nodes below their concurrent-generation ceiling (`--node-parallel`, matching the node's `OLLAMA_NUM_PARALLEL`), and a request that finds every eligible owner at capacity waits for a slot instead of joining a queue. The node chosen is the one the default path picks whenever that node is free; otherwise a free node wins over one that is generating. + +The scheduler contract is untouched: `node/set-priority`, its ranks, and the ordering they imply are consumed exactly as before, and an eligible `node/select` pin still bypasses the whole mechanism. + +**Occupancy.** The gate counts one thing: the generations this proxy has bound to a node and not yet seen finish. One entry is added when a request is bound and removed at its terminal workload transition, so each in-flight request is counted exactly once, for exactly as long as the node is generating for it. + +The scheduler's `pending` count is deliberately not part of it. `pending` still *orders* the choice, as it always has, but it is those same requests seen through a full round trip — proxy → broker → workload-manager → scheduler → `node/set-priority` — so a gate built on it counts each request twice and, worse, keeps counting one after it has finished: the node would stay at capacity until some later snapshot happened to report the lower count, leaving a slot dead on every completion. A gate must be prompt and local; a ranking may be lagging and remote. + +**Precondition.** Because the ledger counts this proxy's own dispatches, it only describes the node while the proxy is its **only client** — a local `ollama run`, a second router, or an application pointed straight at the engine port is invisible here, and the ceiling is then fiction. PAIR's own layout satisfies this, since the proxy is the front door and the engine is moved aside onto a loopback port. + +Bounded by design: a request waits at most two minutes for a slot before committing to the least loaded node anyway, so an occupancy model that has gone stale — a node that died mid-generation, a workload whose terminal transition never arrived — degrades routing to the default behavior rather than stalling it. + +A waiting request stays visible. It has not been forwarded, so it emits no `workload:started`; it emits `workload:submitted` instead (`state` `queued`, no `scheduledOn`), and gets its job card while it waits. Unplaced work counts toward no node's pending, so announcing it does not disturb the ranking the scheduler sends back. If the client hangs up before a slot frees, the queued card is retired with `workload:errored` and no node is contacted. + ### IPC Transport By default the proxy communicates over **stdin/stdout** using newline-delimited JSON-RPC 2.0 (one message per line). All diagnostic logging goes to **stderr**. @@ -176,11 +196,14 @@ A proxied request finished, or was rejected before forwarding. `duration_ms` cov {"jsonrpc":"2.0","method":"proxy/request","params":{"id":"17","node_id":"22222222-2222-2222-2222-222222222222","method":"POST","path":"/api/chat","target":"192.168.1.50:11434","status":200,"duration_ms":6120,"ttfb_ms":95}} ``` -#### `workload:started` / `workload:completed` / `workload:errored` +#### `workload:submitted` / `workload:started` / `workload:completed` / `workload:errored` + +One lifecycle transition per inference request, carrying a single `workloadInfo`. `engine` is always `ollama`; `originatedFrom` is left empty for the broker to stamp, and `scheduledOn` names the node that actually served (re-pointed if failover moved the request). The broker relays these to `nvpair-workload-manager`. -One lifecycle transition per forwarded inference request, carrying a single `workloadInfo`. `engine` is always `ollama`; `originatedFrom` is left empty for the broker to stamp, and `scheduledOn` names the node that actually served (re-pointed if failover moved the request). The broker relays these to `nvpair-workload-manager`. The proxy never emits `workload:submitted` — it forwards immediately rather than queueing. +By default the proxy emits no `workload:submitted`: it forwards on arrival, so it queues nothing and a request's first event is `workload:started`. **Under [`--late-binding`](#late-binding-opt-in) only**, a request that has to wait for a generation slot emits `workload:submitted` first — `state` `queued`, and no `scheduledOn`, because no node has been chosen yet. `workload:started` then re-points the same id at the node that runs it. A request that gets a slot immediately still emits nothing but `started`, so the announcement means what it says: this request queued. ```json +{"jsonrpc":"2.0","method":"workload:submitted","params":{"workloadInfo":{"id":"18","model":"llama3:latest","engine":"ollama","runId":"3ce8a1740b62df95","state":"queued","originatedFrom":"","createdAt":1716998400000,"startedAt":null,"completedAt":null,"error":null,"requesterId":null}}} {"jsonrpc":"2.0","method":"workload:started","params":{"workloadInfo":{"id":"17","model":"llama3:latest","engine":"ollama","runId":"3ce8a1740b62df95","state":"running","originatedFrom":"","scheduledOn":"22222222-2222-2222-2222-222222222222","createdAt":1716998400000,"startedAt":1716998400000,"completedAt":null,"error":null,"requesterId":null}}} ``` @@ -300,6 +323,16 @@ Semantics: - **Snapshot reset.** Each new snapshot replaces the pending and GPU-pressure baselines and clears the reservations. GPU pressure is clamped to the scheduler's 0–3 range. +- **Late binding (opt-in) gates the same choice on a free slot.** With + `--late-binding` the auto ordering above is applied only to eligible nodes + below their concurrent-generation ceiling, and a request that finds none waits + for a slot rather than committing to a queue. The snapshot the scheduler sends + and the rule applied to it are unchanged; only the set it is applied to, and + therefore when the request is committed, differ. Whether a node is below its + ceiling is answered from this proxy's own in-flight generations, not from the + snapshot — a snapshot reset must not hand back a slot that is still in use, + and a finished request must not keep one until the next snapshot arrives. See + [Late binding](#late-binding-opt-in). - **Eligible manual pin wins.** An active `node/select` pin takes precedence when it is in the request's owner set. An ineligible pin is ignored for that request, so automatic reservations still apply among eligible owners. Clearing the pin diff --git a/services/ollama-proxy/latebind.go b/services/ollama-proxy/latebind.go new file mode 100644 index 00000000..8eada6f8 --- /dev/null +++ b/services/ollama-proxy/latebind.go @@ -0,0 +1,263 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +// Late binding (--late-binding, off by default) changes *when* a request is +// committed to a node, and nothing else. +// +// Today the proxy picks the least estimated loaded eligible node the moment the +// request arrives and forwards immediately. That decision is irrevocable: the +// request queues on that node however the fleet moves afterwards, so a burst +// that lands on a node which then turns out to be busy waits behind it while +// another node goes idle. With late binding the same choice is made only among +// nodes that still have a free generation slot, and a request that finds every +// eligible node at capacity waits for one to free instead of joining a queue. +// +// The scheduler contract is untouched. node/set-priority, its ranks, and the +// ordering they imply are consumed exactly as before; the same least-estimated- +// loaded rule decides, applied to the nodes that have a slot. So the node chosen +// is the one the proxy picks today whenever that node is free, and otherwise a +// free node wins over one that is at capacity. +// +// OCCUPANCY. The gate counts one thing: the generations this proxy has bound to +// a node and not yet seen finish, measured against a per-node ceiling +// (--node-parallel, the node's OLLAMA_NUM_PARALLEL). That ledger is exact by +// construction — one entry is added when a request is bound and removed at its +// terminal workload transition, so every in-flight request is counted exactly +// once, for exactly as long as the node is generating for it. +// +// It deliberately does NOT include the scheduler's pending count, which the load +// estimate still uses to *order* the choice. Pending is the same requests seen +// through a full round trip — proxy -> broker -> workload-manager -> scheduler +// catalog -> recompute -> schedule:priority -> broker -> node/set-priority — so +// folding it into the gate counts a request twice over: once locally and again +// as it comes back around. Worse, it is asymmetric. SetPrioritySnapshot clears +// this proxy's reservations wholesale on every snapshot, so by the time a +// request finishes, its own reservation is usually already gone and giving the +// slot back is a no-op: the node stays at capacity until a later snapshot +// happens to report the lower count. The slot is then dead for a whole +// scheduler round trip on every completion, and if anything downstream fails to +// retire the workload it is dead until the wait budget expires. A gate must be +// prompt and local; a ranking may be lagging and remote. +// +// PRECONDITION. Because the ledger counts this proxy's own dispatches, it only +// describes the node while the proxy is that node's only client: anything else +// generating on that engine — a local `ollama run`, a second router, an +// application pointed straight at the engine port — is invisible here, and the +// ceiling is then fiction. PAIR's own design satisfies this, since the proxy is +// the front door and the engine is moved aside onto a loopback port, but a +// hand-assembled setup need not. + +import ( + "context" + "fmt" + "sort" + "strconv" + "strings" + "sync" + "time" +) + +// defaultNodeParallel is the concurrent-generation ceiling assumed for a node +// with no --node-parallel override. The value to match is the node's own +// OLLAMA_NUM_PARALLEL, which current Ollama releases default to a single +// generation at a time — that default is what makes a second request sent to a +// busy node a queued request rather than a concurrent one. +const defaultNodeParallel = 1 + +// lateBindWaitTimeout bounds how long one request waits for a generation slot +// before giving up and committing to the least loaded node anyway — exactly +// what the proxy does today. It is the safety valve for an occupancy model that +// has gone stale (a node that died mid-generation, a workload whose terminal +// transition never arrived, an engine someone else is also driving), so a wedged +// node degrades routing to the current behavior instead of stalling it. Sized in +// the same range as proxyResponseTimeout: a generation that is ever going to +// free its slot does so well inside this window. +const lateBindWaitTimeout = 2 * time.Minute + +// nodeParallelFlags collects repeated --node-parallel values and answers the +// per-node ceiling. A bare "N" sets the count used for every node; "=N" +// overrides one node, for a fleet whose machines are configured differently. +type nodeParallelFlags struct { + slots int + overrides map[string]int +} + +func (f *nodeParallelFlags) String() string { + if f == nil { + return strconv.Itoa(defaultNodeParallel) + } + parts := []string{strconv.Itoa(f.slotsFor(""))} + ids := make([]string, 0, len(f.overrides)) + for id := range f.overrides { + ids = append(ids, id) + } + sort.Strings(ids) + for _, id := range ids { + parts = append(parts, id+"="+strconv.Itoa(f.overrides[id])) + } + return strings.Join(parts, ",") +} + +func (f *nodeParallelFlags) Set(value string) error { + id, count, hasID := strings.Cut(value, "=") + if !hasID { + id, count = "", value + } + slots, err := strconv.Atoi(strings.TrimSpace(count)) + if err != nil || slots < 1 { + return fmt.Errorf("want a positive slot count as N or =N, got %q", value) + } + id = strings.TrimSpace(id) + if id == "" { + f.slots = slots + return nil + } + if f.overrides == nil { + f.overrides = make(map[string]int) + } + f.overrides[id] = slots + return nil +} + +// slotsFor returns the concurrent-generation ceiling for one node: its override +// when it has one, otherwise the fleet-wide value. +func (f nodeParallelFlags) slotsFor(id string) int { + if slots, ok := f.overrides[id]; ok { + return slots + } + if f.slots > 0 { + return f.slots + } + return defaultNodeParallel +} + +// lateBindConfig is the whole of the late-binding configuration. A nil +// *lateBindConfig on the proxy means the feature is off, which is the default. +type lateBindConfig struct { + capacity nodeParallelFlags + // wait bounds a single request's wait for a free slot (lateBindWaitTimeout + // in production; tests shorten it to keep the give-up path fast). + wait time.Duration +} + +// EnableLateBinding switches the proxy from committing a request to a node on +// arrival to committing it when that node has a free generation slot. It is +// called once from main before Run — the configuration and the condition +// variable are read-only from then on, so the request path needs no lock to +// find out whether the feature is on. +func (p *Proxy) EnableLateBinding(capacity nodeParallelFlags) { + p.lateBind = &lateBindConfig{capacity: capacity, wait: lateBindWaitTimeout} + p.slotFree = sync.NewCond(&p.priorityMu) + p.lateBindInFlight = make(map[string]int) +} + +// noRelease is the release side of a reservation that was never taken. +func noRelease() {} + +// releaseSlot returns the release side of one generation slot: it retires this +// request's entry from the node's in-flight ledger and wakes every parked +// request so one of them can take the freed slot. sync.Once makes the returned +// function idempotent — the handler calls it at the terminal workload transition +// and defers it as a backstop — so the ledger entry added when the request was +// bound is removed exactly once. Nothing else may retire it: a scheduler +// snapshot resets the reservations it owns, and must leave this ledger alone, +// or a node would keep a slot it is no longer using. +func (p *Proxy) releaseSlot(id string) func() { + var once sync.Once + return func() { + once.Do(func() { + p.priorityMu.Lock() + defer p.priorityMu.Unlock() + if held := p.lateBindInFlight[id]; held > 1 { + p.lateBindInFlight[id] = held - 1 + } else { + delete(p.lateBindInFlight, id) + } + p.slotFree.Broadcast() + }) + } +} + +// hasFreeSlotLocked reports whether a node is below its concurrent-generation +// ceiling. Occupancy is this proxy's in-flight ledger for that node and nothing +// else: each request it bound contributes one, from the moment it is bound to +// its terminal transition. The scheduler's pending count and GPU pressure are +// deliberately excluded — pending is the same requests arriving back around a +// round trip later (see the file header), and pressure is a smoothed 0-3 score +// rather than a count of jobs, so a node whose pressure alone reached the +// ceiling would look permanently full. Both still order the choice; neither +// gates it. Callers hold priorityMu. +func (p *Proxy) hasFreeSlotLocked(id string) bool { + return p.lateBindInFlight[id] < p.lateBind.capacity.slotsFor(id) +} + +// waitForFreeSlotLocked parks a request until an eligible candidate has a free +// generation slot, the wait budget expires, or the request is cancelled. It is +// called with priorityMu held and returns holding it: sync.Cond releases the +// mutex for the duration of each wait, so reservations, releases, scheduler +// snapshots and other requests all keep making progress while a request is +// parked here. +// +// best and free are the current pick pair from pickCandidateLocked. Returning +// free (>= 0) hands back a candidate with a slot; returning best is the give-up +// path, which commits to the least loaded node exactly as the proxy does with +// the feature off, so a wedged node degrades routing rather than stopping it. +// +// onQueued, when non-nil, is called once at the moment this request is about to +// park — and only then, so a request that binds straight away emits exactly the +// events it emits today. It is the request's cluster visibility: a parked +// request has not been forwarded, so nothing else would announce it. +func (p *Proxy) waitForFreeSlotLocked(ctx context.Context, candidateIndex map[string]int, onQueued func(), best, free int) int { + if free >= 0 || best < 0 { + return free + } + // This request is going to queue. Say so before parking, with priorityMu + // dropped: the announcement writes to the orchestrator channel, and no + // routing decision should ever queue behind that write. Re-deriving the pick + // after re-acquiring is what makes dropping the mutex safe — a slot freed + // during the window is seen here rather than missed, so the broadcast we + // were not holding the mutex to receive costs nothing. + if onQueued != nil { + p.priorityMu.Unlock() + onQueued() + p.priorityMu.Lock() + if best, free = p.pickCandidateLocked(candidateIndex); free >= 0 || best < 0 { + return free + } + } + // Every eligible node is at capacity. Park until something changes. + // + // No wakeup can be lost: each waker takes priorityMu before broadcasting, so + // a broadcast cannot land in the window between the checks below and + // Cond.Wait releasing the mutex, and every state change that could free a + // slot (releaseSlot, SetPrioritySnapshot) happens under the same mutex. The + // two wakers below fire independently of any node activity, so a parked + // request always leaves this loop: the timer bounds the wait, and because + // r.Context() descends from the proxy's root context (serveHTTP's + // BaseContext), the cancel path covers both a client hanging up and shutdown. + wake := func() { + p.priorityMu.Lock() + defer p.priorityMu.Unlock() + p.slotFree.Broadcast() + } + deadline := time.Now().Add(p.lateBind.wait) + timer := time.AfterFunc(p.lateBind.wait, wake) + defer timer.Stop() + stopWakeOnCancel := context.AfterFunc(ctx, wake) + defer stopWakeOnCancel() + + for { + if ctx.Err() != nil || !time.Now().Before(deadline) { + return best + } + p.slotFree.Wait() + // Broadcast wakes every parked request, so re-derive the pick under the + // mutex: another of them may have taken the slot that woke us, and a new + // snapshot may have changed the eligible set entirely. + if best, free = p.pickCandidateLocked(candidateIndex); free >= 0 || best < 0 { + return free + } + } +} diff --git a/services/ollama-proxy/latebind_test.go b/services/ollama-proxy/latebind_test.go new file mode 100644 index 00000000..265c56bd --- /dev/null +++ b/services/ollama-proxy/latebind_test.go @@ -0,0 +1,697 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync" + "testing" + "time" + + "nvpair-shared/schedulerwire" +) + +// lbProxy returns a proxy with late binding enabled at the given per-node slot +// ceiling. The wait budget is shortened from lateBindWaitTimeout so a test that +// exercises the give-up path fails with a wrong answer rather than a stalled +// suite; it stays long enough that a missed wakeup is still a visible failure. +func lbProxy(t *testing.T, slots int) *Proxy { + t.Helper() + p := prProxy(t) + var capacity nodeParallelFlags + if err := capacity.Set(strconv.Itoa(slots)); err != nil { + t.Fatalf("node-parallel %d: %v", slots, err) + } + p.EnableLateBinding(capacity) + p.lateBind.wait = 5 * time.Second + return p +} + +// reserveAsync runs one reservation on its own goroutine (reserveCandidate +// reorders the slice in place, so each caller gets a copy) and reports the +// reserved id plus its release through the returned channels. +func reserveAsync(ctx context.Context, p *Proxy, candidates []candidate) <-chan struct { + id string + release func() +} { + done := make(chan struct { + id string + release func() + }, 1) + go func() { + reserved, release := p.reserveCandidate(ctx, append([]candidate(nil), candidates...), nil) + done <- struct { + id string + release func() + }{reserved[0].id, release} + }() + return done +} + +// oneBusyOneIdle ranks "idle" as the more loaded node by the scheduler's +// estimate (GPU pressure) while "busy" is the one a caller then binds a request +// to. It is the fixture that separates the load estimate, which orders the +// choice, from occupancy, which gates it — and it deliberately leaves occupancy +// alone, because a snapshot cannot tell anyone what is generating right now. +func oneBusyOneIdle(p *Proxy) []candidate { + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"idle", "busy"}, + Ranks: []schedulerwire.NodeRank{ + {ID: "idle", Pending: 0, GPUPressure: 3, Rank: 0}, + {ID: "busy", Pending: 1, GPUPressure: 0, Rank: 1}, + }, + }) + return reservationCandidates("idle", "busy") +} + +// TestLateBindingOff_CommitsToTheLeastLoadedNodeAndNeverWaits is the +// default-behavior guard: with the flag off a node at capacity is still a +// candidate, the least loaded node wins exactly as before, and the proxy holds +// no waiting apparatus at all. +func TestLateBindingOff_CommitsToTheLeastLoadedNodeAndNeverWaits(t *testing.T) { + p := prProxy(t) + candidates := oneBusyOneIdle(p) + + got := reserveAsync(context.Background(), p, candidates) + select { + case reserved := <-got: + if reserved.id != "busy" { + t.Fatalf("reserved %q, want busy (the least estimated loaded node)", reserved.id) + } + case <-time.After(2 * time.Second): + t.Fatal("reserveCandidate blocked with late binding off") + } + if p.lateBind != nil || p.slotFree != nil { + t.Fatal("late binding state exists without --late-binding") + } +} + +// TestLateBinding_PrefersAFreeNodeOverALowerLoadedFullOne: the only routing +// difference late binding makes while a slot is available is that occupancy +// gates the choice. GPU pressure biases the order but is not a job count, so a +// node whose pressure is high yet has no generation running is still free. +func TestLateBinding_PrefersAFreeNodeOverALowerLoadedFullOne(t *testing.T) { + p := lbProxy(t, 1) + candidates := oneBusyOneIdle(p) + + // Make "busy" genuinely busy, and keep its slot. Even carrying that + // request the scheduler's estimate still ranks it below "idle", whose GPU + // pressure is high but whose engine is generating nothing. + if first, _ := p.reserveCandidate(context.Background(), append([]candidate(nil), candidates...), nil); first[0].id != "busy" { + t.Fatalf("first reservation went to %q, want busy (the least estimated loaded node)", first[0].id) + } + + reserved, _ := p.reserveCandidate(context.Background(), candidates, nil) + if reserved[0].id != "idle" { + t.Fatalf("reserved %q, want idle (busy is at its generation ceiling)", reserved[0].id) + } + if got := candidateIDsFrom(reserved); got[1] != "busy" { + t.Fatalf("failover list = %v, want busy retained behind idle", got) + } +} + +// TestLateBinding_WaitsForAReleasedSlot is the whole point of the change: with +// every eligible node generating, the request waits rather than committing to a +// queue, and takes the slot the moment one is released. +func TestLateBinding_WaitsForAReleasedSlot(t *testing.T) { + p := lbProxy(t, 1) + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a", "b"}, + Ranks: []schedulerwire.NodeRank{{ID: "a"}, {ID: "b"}}, + }) + candidates := reservationCandidates("a", "b") + + first, releaseFirst := p.reserveCandidate(context.Background(), append([]candidate(nil), candidates...), nil) + second, _ := p.reserveCandidate(context.Background(), append([]candidate(nil), candidates...), nil) + if first[0].id == second[0].id { + t.Fatalf("two reservations landed on %q; both nodes should be filled", first[0].id) + } + + third := reserveAsync(context.Background(), p, candidates) + select { + case reserved := <-third: + t.Fatalf("reserved %q with every node at capacity", reserved.id) + case <-time.After(100 * time.Millisecond): + } + + releaseFirst() + select { + case reserved := <-third: + if reserved.id != first[0].id { + t.Fatalf("reserved %q, want the released node %q", reserved.id, first[0].id) + } + case <-time.After(2 * time.Second): + t.Fatal("releasing a slot did not wake the waiting request") + } +} + +// TestLateBinding_SnapshotWakesAWaitingRequest: a snapshot replaces the eligible +// set and the ordering, so a request parked when every listed owner was +// generating has to be woken by one — otherwise it waits out its whole budget +// for a node that became routable in the meantime. +func TestLateBinding_SnapshotWakesAWaitingRequest(t *testing.T) { + p := lbProxy(t, 1) + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a"}, + Ranks: []schedulerwire.NodeRank{{ID: "a"}}, + }) + // "b" is an advertised owner the scheduler has not listed yet. + candidates := reservationCandidates("a", "b") + + // Take a's only slot and keep it: the release is deliberately never called. + if reserved, _ := p.reserveCandidate(context.Background(), append([]candidate(nil), candidates...), nil); reserved[0].id != "a" { + t.Fatalf("first reservation went to %q, want the only listed node a", reserved[0].id) + } + + waiting := reserveAsync(context.Background(), p, candidates) + select { + case <-waiting: + t.Fatal("reserved a node that is already generating") + case <-time.After(100 * time.Millisecond): + } + + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a", "b"}, + Ranks: []schedulerwire.NodeRank{{ID: "a", Pending: 1}, {ID: "b", Pending: 0}}, + }) + select { + case reserved := <-waiting: + if reserved.id != "b" { + t.Fatalf("reserved %q, want the newly listed free node b", reserved.id) + } + case <-time.After(2 * time.Second): + t.Fatal("a new scheduler snapshot did not wake the waiting request") + } +} + +// TestLateBinding_AReleasedSlotIsFreeWithoutWaitingForTheScheduler is the +// over-blocking regression. The scheduler's pending count is this proxy's own +// dispatches arriving back a full round trip later, and every snapshot clears +// the reservations, so an occupancy model built from pending + reservations +// keeps a finished request occupying its node until some later snapshot happens +// to report the lower count — a late-binding router leaving nodes idle. The +// ledger the gate actually consults is retired by the request that finished, so +// the slot is free the moment it finishes, with no further snapshot. +func TestLateBinding_AReleasedSlotIsFreeWithoutWaitingForTheScheduler(t *testing.T) { + p := lbProxy(t, 1) + // Far longer than this test's own patience, so the second request can only + // be reserved because the first freed its slot — never because the give-up + // path rescued it. + p.lateBind.wait = 10 * time.Second + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a"}, + Ranks: []schedulerwire.NodeRank{{ID: "a", Pending: 0}}, + }) + candidates := reservationCandidates("a") + + _, releaseFirst := p.reserveCandidate(context.Background(), append([]candidate(nil), candidates...), nil) + // The scheduler catches up: its next snapshot counts the request just + // dispatched and, as every snapshot does, clears this proxy's reservations. + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a"}, + Ranks: []schedulerwire.NodeRank{{ID: "a", Pending: 1}}, + }) + + waiting := reserveAsync(context.Background(), p, candidates) + select { + case <-waiting: + t.Fatal("reserved a node that is still generating") + case <-time.After(100 * time.Millisecond): + } + + releaseFirst() + select { + case reserved := <-waiting: + if reserved.id != "a" { + t.Fatalf("reserved %q, want a", reserved.id) + } + case <-time.After(2 * time.Second): + t.Fatal("the node stayed at capacity after its only request finished: " + + "occupancy is still waiting for a scheduler snapshot to retire it") + } +} + +// TestLateBinding_ASnapshotDoesNotFreeASlotThatIsStillGenerating is the other +// half of counting each in-flight request exactly once. A snapshot clears the +// reservations wholesale, so a snapshot taken before the scheduler heard about a +// request would hand its node's ceiling back while it is still generating. The +// request is counted by the ledger for its whole life, whatever the snapshot +// says. +func TestLateBinding_ASnapshotDoesNotFreeASlotThatIsStillGenerating(t *testing.T) { + p := lbProxy(t, 1) + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a"}, + Ranks: []schedulerwire.NodeRank{{ID: "a", Pending: 0}}, + }) + candidates := reservationCandidates("a") + + _, release := p.reserveCandidate(context.Background(), append([]candidate(nil), candidates...), nil) + // A snapshot the scheduler computed before it saw that request: node idle. + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a"}, + Ranks: []schedulerwire.NodeRank{{ID: "a", Pending: 0}}, + }) + + waiting := reserveAsync(context.Background(), p, candidates) + select { + case reserved := <-waiting: + t.Fatalf("reserved %q while its only slot is still generating: "+ + "the snapshot gave back capacity that is in use", reserved.id) + case <-time.After(200 * time.Millisecond): + } + + release() + select { + case <-waiting: + case <-time.After(2 * time.Second): + t.Fatal("releasing the slot did not wake the waiting request") + } +} + +// TestLateBinding_CancelledRequestStopsWaiting: a client that hangs up (and, +// through the same root context, proxy shutdown) must not leave a request parked +// for the rest of its budget. +func TestLateBinding_CancelledRequestStopsWaiting(t *testing.T) { + p := lbProxy(t, 1) + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a"}, + Ranks: []schedulerwire.NodeRank{{ID: "a"}}, + }) + // Occupy the only slot and keep it, so the next request has to wait. + p.reserveCandidate(context.Background(), reservationCandidates("a"), nil) + + ctx, cancel := context.WithCancel(context.Background()) + waiting := reserveAsync(ctx, p, reservationCandidates("a")) + select { + case <-waiting: + t.Fatal("reserved a node that is already generating") + case <-time.After(100 * time.Millisecond): + } + + cancel() + select { + case <-waiting: + case <-time.After(2 * time.Second): + t.Fatal("cancelling the request did not wake it") + } +} + +// TestLateBinding_WaitBudgetFallsBackToImmediateCommit: an occupancy model that +// has gone stale must degrade to the current behavior rather than stall. Nothing +// here ever frees the slot, so the request has to give up and commit. +func TestLateBinding_WaitBudgetFallsBackToImmediateCommit(t *testing.T) { + p := lbProxy(t, 1) + p.lateBind.wait = 25 * time.Millisecond + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"a", "b"}, + Ranks: []schedulerwire.NodeRank{ + {ID: "a", Pending: 4}, + {ID: "b", Pending: 2}, + }, + }) + + // Fill both nodes and never release, standing in for an occupancy model that + // has gone stale — a node that died mid-generation, a workload whose + // terminal transition never arrived. Nothing here can free a slot. + for range 2 { + p.reserveCandidate(context.Background(), reservationCandidates("a", "b"), nil) + } + + waiting := reserveAsync(context.Background(), p, reservationCandidates("a", "b")) + select { + case reserved := <-waiting: + if reserved.id != "b" { + t.Fatalf("gave up onto %q, want the least loaded node b", reserved.id) + } + case <-time.After(2 * time.Second): + t.Fatal("a wedged occupancy model stalled the request past its wait budget") + } +} + +// TestLateBinding_ConcurrentBurstNeverExceedsCapacity: the reservation is a +// semaphore, so a burst may hold at most the configured number of slots on any +// node at any instant. +func TestLateBinding_ConcurrentBurstNeverExceedsCapacity(t *testing.T) { + const slots = 2 + p := lbProxy(t, slots) + ids := []string{"a", "b", "c"} + ranks := make([]schedulerwire.NodeRank, 0, len(ids)) + for i, id := range ids { + ranks = append(ranks, schedulerwire.NodeRank{ID: id, Rank: i}) + } + p.SetPrioritySnapshot(schedulerwire.Priority{Nodes: ids, Ranks: ranks}) + candidates := reservationCandidates(ids...) + + var ( + mu sync.Mutex + held = map[string]int{} + worst = map[string]int{} + ) + var wg sync.WaitGroup + for range 60 { + wg.Add(1) + go func() { + defer wg.Done() + reserved, release := p.reserveCandidate(context.Background(), append([]candidate(nil), candidates...), nil) + id := reserved[0].id + mu.Lock() + held[id]++ + if held[id] > worst[id] { + worst[id] = held[id] + } + mu.Unlock() + + time.Sleep(time.Millisecond) + + mu.Lock() + held[id]-- + mu.Unlock() + release() + }() + } + wg.Wait() + + for _, id := range ids { + if worst[id] > slots { + t.Fatalf("node %q held %d concurrent generations, ceiling is %d (all: %v)", id, worst[id], slots, worst) + } + } +} + +// TestLateBinding_HandleHTTPHoldsTheSecondRequestUntilTheFirstFinishes drives +// the real handler: the second inference request must not reach an engine that +// is still generating, and must be forwarded as soon as the first request +// reaches its terminal workload transition. +func TestLateBinding_HandleHTTPHoldsTheSecondRequestUntilTheFirstFinishes(t *testing.T) { + var ( + mu sync.Mutex + received int + ) + release := make(chan struct{}) + var releaseOnce sync.Once + doRelease := func() { releaseOnce.Do(func() { close(release) }) } + inflight := func() int { + mu.Lock() + defer mu.Unlock() + return received + } + + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + received++ + first := received == 1 + mu.Unlock() + if first { + <-release // hold the single generation slot until the test lets go + } + w.WriteHeader(http.StatusOK) + io.WriteString(w, `{"done":true}`) + })) + // Deferred order matters (LIFO): unblock the held handler before the server + // is torn down, including on t.Fatal. + defer engine.Close() + defer doRelease() + + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "solo-node", engine.URL, "llama")) + p := testProxy(disc, 11435) + p.EnableLateBinding(nodeParallelFlags{slots: 1}) + // A budget far longer than the test's own patience, so the second request + // can only be forwarded because the first released its slot — never because + // the give-up path rescued it. + p.lateBind.wait = 10 * time.Second + p.SetPriority([]string{"solo-node"}) + + post := func() *httptest.ResponseRecorder { + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/api/chat", + strings.NewReader(`{"model":"llama"}`))) + return rec + } + + firstDone := make(chan struct{}) + go func() { + defer close(firstDone) + post() + }() + deadline := time.Now().Add(2 * time.Second) + for inflight() == 0 && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if inflight() != 1 { + t.Fatal("the first request never reached the engine") + } + + secondDone := make(chan struct{}) + go func() { + defer close(secondDone) + post() + }() + time.Sleep(200 * time.Millisecond) + if got := inflight(); got != 1 { + t.Fatalf("%d requests reached the engine while its only slot was busy, want 1", got) + } + + doRelease() + <-firstDone + select { + case <-secondDone: + case <-time.After(3 * time.Second): + t.Fatal("the second request was not forwarded when the first released its slot") + } + if got := inflight(); got != 2 { + t.Fatalf("engine saw %d requests, want 2", got) + } +} + +func TestNodeParallelFlags_ParsesDefaultsAndPerNodeOverrides(t *testing.T) { + var capacity nodeParallelFlags + if got := capacity.slotsFor("anything"); got != defaultNodeParallel { + t.Fatalf("unset --node-parallel = %d, want %d", got, defaultNodeParallel) + } + for _, value := range []string{"4", "big-rig=8", " small-rig = 1 "} { + if err := capacity.Set(value); err != nil { + t.Fatalf("--node-parallel %q: %v", value, err) + } + } + for id, want := range map[string]int{"unlisted": 4, "big-rig": 8, "small-rig": 1} { + if got := capacity.slotsFor(id); got != want { + t.Fatalf("slots for %q = %d, want %d", id, got, want) + } + } + if got := capacity.String(); got != "4,big-rig=8,small-rig=1" { + t.Fatalf("--node-parallel String() = %q", got) + } + for _, value := range []string{"0", "-1", "", "node=", "node=zero"} { + if err := capacity.Set(value); err == nil { + t.Fatalf("--node-parallel %q was accepted", value) + } + } +} + +// workloadSequence returns the workload:* lifecycle methods the codec wrote, in +// the order they were written, so a test can assert a lifecycle rather than a +// set of events that happened to occur. +func workloadSequence(rec *recRW) []string { + rec.mu.Lock() + defer rec.mu.Unlock() + var out []string + for _, frame := range strings.Split(string(rec.b), "\n") { + for _, method := range []string{ + workloadSubmittedMethod, workloadStartedMethod, + workloadCompletedMethod, workloadErroredMethod, + } { + if strings.Contains(frame, `"method":"`+method+`"`) { + out = append(out, method) + } + } + } + return out +} + +// workloadFrame returns the first frame the codec wrote for one workload method, +// so a test can assert what that event actually carried. +func workloadFrame(rec *recRW, method string) string { + rec.mu.Lock() + defer rec.mu.Unlock() + for _, frame := range strings.Split(string(rec.b), "\n") { + if strings.Contains(frame, `"method":"`+method+`"`) { + return frame + } + } + return "" +} + +// oneSlotEngine is a single-slot engine: the first request it receives is held +// until the returned release is called, standing in for a node that is +// generating. It reports how many requests have reached it. +func oneSlotEngine(t *testing.T) (server *httptest.Server, received func() int, release func()) { + t.Helper() + var ( + mu sync.Mutex + count int + ) + gate := make(chan struct{}) + var once sync.Once + release = func() { once.Do(func() { close(gate) }) } + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + first := count == 0 + count++ + mu.Unlock() + if first { + <-gate + } + w.WriteHeader(http.StatusOK) + io.WriteString(w, `{"done":true}`) + })) + // Deferred order matters (LIFO): unblock the held handler before the server + // is torn down, including on t.Fatal. + t.Cleanup(server.Close) + t.Cleanup(release) + return server, func() int { + mu.Lock() + defer mu.Unlock() + return count + }, release +} + +// TestLateBinding_AQueuedRequestIsAnnouncedThenRepointedAtItsNode is the +// visibility regression. Late binding parks a request instead of forwarding it, +// and an unforwarded request emits nothing at all — so a burst larger than the +// fleet showed job cards only for the requests that won a slot, and the rest +// looked like nothing was happening. The protocol already has the state: this +// asserts the parked request is announced as queued work with no node yet, and +// that the same workload id is re-pointed at the node once one is bound. +func TestLateBinding_AQueuedRequestIsAnnouncedThenRepointedAtItsNode(t *testing.T) { + engine, received, release := oneSlotEngine(t) + + rec := &recRW{} + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "solo-node", engine.URL, "llama")) + p := NewProxy(NewCodec(rec), disc, 11435) + p.EnableLateBinding(nodeParallelFlags{slots: 1}) + p.lateBind.wait = 10 * time.Second + p.SetPriority([]string{"solo-node"}) + + post := func() { + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest( + http.MethodPost, "/api/chat", strings.NewReader(`{"model":"llama"}`))) + } + + first := make(chan struct{}) + go func() { defer close(first); post() }() + deadline := time.Now().Add(2 * time.Second) + for received() == 0 && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if received() != 1 { + t.Fatal("the first request never reached the engine") + } + + second := make(chan struct{}) + go func() { defer close(second); post() }() + deadline = time.Now().Add(2 * time.Second) + for !rec.has(workloadSubmittedMethod) && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + + // The second request is parked: it has not been forwarded, and the only + // thing that can make it visible is the queued announcement. + if got := received(); got != 1 { + t.Fatalf("%d requests reached the engine while its only slot was busy, want 1", got) + } + queued := workloadFrame(rec, workloadSubmittedMethod) + if queued == "" { + t.Fatal("a request waiting for a generation slot emitted no workload:submitted, so it is invisible while it waits") + } + if !strings.Contains(queued, `"state":"queued"`) { + t.Fatalf("workload:submitted did not carry state queued: %s", queued) + } + if strings.Contains(queued, `"scheduledOn"`) { + t.Fatalf("workload:submitted named a node before one was chosen: %s", queued) + } + + release() + <-first + select { + case <-second: + case <-time.After(3 * time.Second): + t.Fatal("the queued request was never forwarded") + } + + // submitted names no node; started re-points the same id at the one that + // ran it. Both requests then reach a terminal, so neither card is left open. + started := workloadFrame(rec, workloadStartedMethod) + if !strings.Contains(started, `"scheduledOn":"solo-node"`) { + t.Fatalf("workload:started did not name the node that ran it: %s", started) + } + sequence := workloadSequence(rec) + var submitted, terminal int + for _, method := range sequence { + switch method { + case workloadSubmittedMethod: + submitted++ + case workloadCompletedMethod, workloadErroredMethod: + terminal++ + } + } + if submitted != 1 { + t.Fatalf("workload:submitted emitted %d times, want 1 (only the request that queued): %v", submitted, sequence) + } + if terminal != 2 { + t.Fatalf("%d terminal transitions for 2 requests: %v", terminal, sequence) + } + if sequence[0] != workloadStartedMethod || sequence[1] != workloadSubmittedMethod { + t.Fatalf("lifecycle order = %v, want the forwarded request started before the queued one is submitted", sequence) + } +} + +// TestLateBindingOff_NeverAnnouncesAQueuedWorkload is the default-behavior +// guard for the event stream. With the flag off the proxy still forwards on +// arrival and queues nothing, so workload:submitted must not appear: the +// component README's claim about the events it emits stays true for every +// installation that has not opted in. +func TestLateBindingOff_NeverAnnouncesAQueuedWorkload(t *testing.T) { + engine, received, release := oneSlotEngine(t) + + rec := &recRW{} + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "solo-node", engine.URL, "llama")) + p := NewProxy(NewCodec(rec), disc, 11435) + p.SetPriority([]string{"solo-node"}) + + post := func() { + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest( + http.MethodPost, "/api/chat", strings.NewReader(`{"model":"llama"}`))) + } + done := make(chan struct{}, 2) + for range 2 { + go func() { post(); done <- struct{}{} }() + } + // Both requests are forwarded on arrival; the engine serializes them. + deadline := time.Now().Add(2 * time.Second) + for received() < 2 && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if got := received(); got != 2 { + t.Fatalf("%d requests reached the engine, want both forwarded immediately", got) + } + release() + for range 2 { + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("a request never finished") + } + } + if rec.has(workloadSubmittedMethod) { + t.Fatalf("workload:submitted emitted without --late-binding: %v", workloadSequence(rec)) + } +} diff --git a/services/ollama-proxy/main.go b/services/ollama-proxy/main.go index afd147dd..93315b29 100644 --- a/services/ollama-proxy/main.go +++ b/services/ollama-proxy/main.go @@ -34,6 +34,9 @@ func main() { ignorePersistedPort := flag.Bool("ignore-persisted-port", false, "use --port even when a persisted port exists") ipcPath := flag.String("ipc", "", "IPC endpoint: Unix domain socket path or Windows named pipe (default: stdin/stdout)") clusterDir := flag.String("cluster-dir", "", "cluster trust directory (node.crt/key + trusted pins); enables the LAN mTLS inference ingress when this node is clustered") + lateBinding := flag.Bool("late-binding", false, "hold an inference request until a node has a free generation slot instead of committing it to a node on arrival; requires this proxy to be the node's only client") + var nodeParallel nodeParallelFlags + flag.Var(&nodeParallel, "node-parallel", "concurrent generation slots per node, consulted only with --late-binding: N for every node, or =N to override one (default 1, matching Ollama's OLLAMA_NUM_PARALLEL)") showVersion := flag.Bool("version", false, "print version and exit") resolveLevel := applog.RegisterFlag(nil, slog.LevelInfo) flag.Parse() @@ -84,6 +87,10 @@ func main() { codec := NewCodec(transport) disc := NewDiscovery() proxy := NewProxy(codec, disc, effectivePort) + if *lateBinding { + proxy.EnableLateBinding(nodeParallel) + log.Printf("late binding enabled (generation slots per node: %s)", nodeParallel.String()) + } for _, aliasAddress := range aliasAddresses { if err := proxy.setLoopbackAlias(aliasAddress); err != nil { log.Fatalf("invalid alias address: %v", err) diff --git a/services/ollama-proxy/proxy.go b/services/ollama-proxy/proxy.go index ad4ac7a2..d7882e99 100644 --- a/services/ollama-proxy/proxy.go +++ b/services/ollama-proxy/proxy.go @@ -129,12 +129,17 @@ type RequestEvent struct { } // Workload lifecycle method names (workload-manager spec 7). The proxy is -// a workload *producer*: it emits one of these per forwarded inference -// request so the broker can stamp the origin (originatedFrom) and forward it to the -// workload-manager, which broadcasts it cluster-wide. We don't emit -// workload:submitted (the proxy never queues — it forwards immediately) or -// workloads:remove (retirement is a broker concern). +// a workload *producer*: it emits these per inference request so the broker can +// stamp the origin (originatedFrom) and forward them to the workload-manager, +// which broadcasts them cluster-wide. +// +// workload:submitted is emitted only under --late-binding, and only for a +// request that actually waits for a generation slot: the default path still +// forwards on arrival, so it queues nothing and goes straight to +// workload:started. workloads:remove is never emitted (retirement is a broker +// concern). const ( + workloadSubmittedMethod = "workload:submitted" workloadStartedMethod = "workload:started" workloadCompletedMethod = "workload:completed" workloadErroredMethod = "workload:errored" @@ -142,6 +147,11 @@ const ( // workloadEngine is the opaque engine identifier carried in every // workload this proxy produces. The proxy only ever fronts Ollama. workloadEngine = "ollama" + + // statusClientClosedRequest is nginx's 499, reported in proxy/request for a + // request the client abandoned before any node was contacted. Nothing is + // written to the client — it is gone — so this only labels the event. + statusClientClosedRequest = 499 ) // inferenceEndpoints is the set of request paths that count as cluster @@ -369,6 +379,24 @@ type Proxy struct { priorityGPUPressure map[string]int priorityReservations map[string]int + // lateBind is nil — the default — unless --late-binding is set. When it is, + // reserveCandidate takes a free generation slot or waits for one, and the + // handler gives it back at the terminal workload transition. slotFree is the + // wait, created alongside lateBind and signalled under priorityMu whenever a + // slot is released or a snapshot changes the eligible set. Both are + // read-only after startup. + // + // lateBindInFlight is the occupancy ledger the gate consults, guarded by + // priorityMu like the maps above but emphatically NOT part of the scheduler + // baseline: SetPrioritySnapshot replaces those and must leave this alone. + // It holds one entry per request this proxy has bound to a node and not yet + // seen finish, which is what makes the count exact — the scheduler's pending + // is the same requests a round trip later, so a gate that added the two + // would count each of them twice. See latebind.go. + lateBind *lateBindConfig + slotFree *sync.Cond + lateBindInFlight map[string]int + // targets remembers, per node, which of its published addresses accepted a // connection, so a repeated forward costs no confirmation. An entry is // re-confirmed when the node's candidate list changes and forgotten on an @@ -1147,8 +1175,69 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { routingModel = model } candidates := p.resolveCandidates(routingModel) + // releaseSlot gives back the generation slot the reservation took under + // --late-binding. It is a no-op with late binding off, or when no + // reservation was made (an explicit pin, an empty priority list). The slot + // is released at the terminal workload transition below, which is where the + // node actually stops generating; the defer is the backstop for a path that + // never reaches one, such as a panic unwinding the handler. + // + // queued records that this request was announced to the cluster as waiting + // for a slot. It can only become true under --late-binding, and only for a + // request that actually parked. Both it and announceQueued are touched from + // this goroutine alone: reserveCandidate calls the hook synchronously. + releaseSlot := noRelease + queued := false if isInf && model != "" { - candidates = p.reserveCandidate(candidates) + // A parked request has not been forwarded, so nothing downstream would + // otherwise know it exists — it would sit invisible for as long as it + // waits. workload:submitted is the cluster's word for that: state + // queued, and no scheduledOn, because no node has been chosen yet. + // Unplaced work counts toward no node's pending, so announcing it does + // not disturb the ranking the scheduler sends back. The workload:started + // emitted below re-points the same id at the node once one is bound. + announceQueued := func() { + queued = true + p.emitWorkload(workloadSubmittedMethod, Workload{ + ID: reqID, + Model: model, + Engine: workloadEngine, + RunID: p.runID, + State: "queued", + CreatedAt: start.UnixMilli(), + }) + } + candidates, releaseSlot = p.reserveCandidate(r.Context(), candidates, announceQueued) + defer releaseSlot() + if queued && r.Context().Err() != nil { + // The client gave up while this request was still waiting for a + // slot, so it never reached a node. Retire the queued card with its + // own terminal instead of forwarding a request nobody is waiting + // for and reporting it started on a node that never saw it. Only + // late binding can reach this: without it a request is committed on + // arrival and there is no window to be abandoned in. + completedMs := time.Now().UnixMilli() + errText := "client disconnected while queued for a generation slot" + p.emitWorkload(workloadErroredMethod, Workload{ + ID: reqID, + Model: model, + Engine: workloadEngine, + RunID: p.runID, + State: "failed", + CreatedAt: start.UnixMilli(), + CompletedAt: &completedMs, + Error: &errText, + }) + p.codec.Notify("proxy/request", RequestEvent{ + ID: reqID, + Method: r.Method, + Path: r.URL.Path, + Status: statusClientClosedRequest, + Duration: time.Since(start).Milliseconds(), + Error: errText, + }) + return + } } if r.Method == http.MethodGet && (r.URL.Path == "/api/tags" || r.URL.Path == "/v1/models") { if len(candidates) > 0 { @@ -1263,6 +1352,12 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { return } terminalOnce.Do(func() { + // The workload is over — completed, failed, or cancelled — so the + // node is no longer generating for it. Give the generation slot back + // first, so a request parked in reserveCandidate can take it without + // waiting for this handler to unwind (the disconnect watcher reaches + // this before the stream copy returns). No-op with late binding off. + releaseSlot() now := time.Now().UnixMilli() wlMu.Lock() terminated = true @@ -1672,14 +1767,22 @@ func (p *Proxy) resolveCandidates(model string) []candidate { // Model eligibility was enforced before this function receives the list. An // explicit node/select pin bypasses reservations, and unlisted/manual owners // retain their existing fallback position. -func (p *Proxy) reserveCandidate(candidates []candidate) []candidate { +// +// With --late-binding the choice is restricted to candidates that still have a +// free generation slot, and a request that finds none parks here until one frees +// rather than committing to a queue (see latebind.go). onQueued, when non-nil, +// is called once if this request has to park, so the caller can announce the +// queued state; it is never called for a request that binds immediately. The +// second return value releases the slot this request took; it is a no-op when +// late binding is off or no reservation was made. +func (p *Proxy) reserveCandidate(ctx context.Context, candidates []candidate, onQueued func()) ([]candidate, func()) { if len(candidates) == 0 { - return candidates + return candidates, noRelease } if selectedID := p.SelectedID(); selectedID != "" { for _, cand := range candidates { if cand.id == selectedID { - return candidates + return candidates, noRelease } } } @@ -1692,15 +1795,51 @@ func (p *Proxy) reserveCandidate(candidates []candidate) []candidate { p.priorityMu.Lock() defer p.priorityMu.Unlock() if len(p.priority) == 0 { - return candidates + return candidates, noRelease } if p.priorityReservations == nil { p.priorityReservations = make(map[string]int) } - bestIndex := -1 - bestOrder := len(p.priority) - var bestLoad uint64 + bestIndex, freeIndex := p.pickCandidateLocked(candidateIndex) + if p.lateBind != nil { + bestIndex = p.waitForFreeSlotLocked(ctx, candidateIndex, onQueued, bestIndex, freeIndex) + } + if bestIndex < 0 { + return candidates, noRelease + } + + chosen := candidates[bestIndex] + p.priorityReservations[chosen.id]++ + if bestIndex > 0 { + copy(candidates[1:bestIndex+1], candidates[:bestIndex]) + candidates[0] = chosen + } + release := noRelease + if p.lateBind != nil { + // The reservation above keeps its existing meaning and lifetime — an + // optimistic tiebreak the next snapshot clears. The ledger entry is + // separate and is retired only by release, at the terminal transition, + // so this request occupies the node's ceiling exactly once and for + // exactly as long as the node is generating for it. + p.lateBindInFlight[chosen.id]++ + release = p.releaseSlot(chosen.id) + } + return candidates, release +} + +// pickCandidateLocked scans the scheduler's list for the eligible candidate +// carrying the least estimated load and returns its index in the failover list, +// or -1 when the list holds none of them. With late binding on it also returns +// the least loaded candidate that still has a free generation slot, again -1 for +// none. The two agree whenever the least loaded candidate is itself free; they +// differ when it is at capacity, which is the whole of the routing change (GPU +// pressure can rank an idle node above a node that is generating). Callers hold +// priorityMu. +func (p *Proxy) pickCandidateLocked(candidateIndex map[string]int) (best, free int) { + best, free = -1, -1 + bestOrder, freeOrder := len(p.priority), len(p.priority) + var bestLoad, freeLoad uint64 for order, id := range p.priority { index, ok := candidateIndex[id] if !ok { @@ -1709,23 +1848,21 @@ func (p *Proxy) reserveCandidate(candidates []candidate) []candidate { load := uint64(p.priorityPending[id]) + uint64(p.priorityGPUPressure[id]) + uint64(p.priorityReservations[id]) - if bestIndex < 0 || load < bestLoad || (load == bestLoad && order < bestOrder) { - bestIndex = index + if best < 0 || load < bestLoad || (load == bestLoad && order < bestOrder) { + best = index bestOrder = order bestLoad = load } + if p.lateBind == nil || !p.hasFreeSlotLocked(id) { + continue + } + if free < 0 || load < freeLoad || (load == freeLoad && order < freeOrder) { + free = index + freeOrder = order + freeLoad = load + } } - if bestIndex < 0 { - return candidates - } - - chosen := candidates[bestIndex] - p.priorityReservations[chosen.id]++ - if bestIndex > 0 { - copy(candidates[1:bestIndex+1], candidates[:bestIndex]) - candidates[0] = chosen - } - return candidates + return best, free } func nodeAdvertisesModel(n Node, model string) bool { @@ -2052,6 +2189,20 @@ func (p *Proxy) SetPrioritySnapshot(priority schedulerwire.Priority) int { p.priorityPending = pending p.priorityGPUPressure = gpuPressure p.priorityReservations = make(map[string]int) + // A snapshot replaces the scheduler baseline and, with it, the eligible set + // and the ordering — so a request parked in reserveCandidate has to + // re-derive its pick against the new one, which may have made a node + // routable that was not before. Broadcasting under the mutex is what makes + // that wakeup impossible to miss; nil until --late-binding creates it. + // + // lateBindInFlight is untouched on purpose. It records generations this + // proxy started and has not seen finish, which a snapshot knows nothing + // about: clearing it here would hand out a slot that is still in use, and + // deriving occupancy from the pending counts above instead would keep a + // finished request occupying the node until some later snapshot retires it. + if p.slotFree != nil { + p.slotFree.Broadcast() + } p.priorityMu.Unlock() return len(cleaned) } diff --git a/services/ollama-proxy/reservation_test.go b/services/ollama-proxy/reservation_test.go index 97049ebf..1abc28e5 100644 --- a/services/ollama-proxy/reservation_test.go +++ b/services/ollama-proxy/reservation_test.go @@ -4,6 +4,7 @@ package main import ( + "context" "sync" "testing" @@ -20,7 +21,8 @@ func reservationCandidates(ids ...string) []candidate { func reservedID(p *Proxy, candidates []candidate) string { candidates = append([]candidate(nil), candidates...) - return p.reserveCandidate(candidates)[0].id + reserved, _ := p.reserveCandidate(context.Background(), candidates, nil) + return reserved[0].id } func TestReserveCandidate_ConcurrentEqualLoadHasAtMostOneSkew(t *testing.T) { @@ -209,7 +211,7 @@ func TestReserveCandidate_PreservesFailoverAndSnapshotReset(t *testing.T) { {ID: "c", Pending: 6}, }, }) - got := p.reserveCandidate(reservationCandidates("a", "b", "c")) + got, _ := p.reserveCandidate(context.Background(), reservationCandidates("a", "b", "c"), nil) want := []string{"b", "a", "c"} for i, id := range want { if got[i].id != id { diff --git a/services/versions.json b/services/versions.json index 29d8c230..5a0f4337 100644 --- a/services/versions.json +++ b/services/versions.json @@ -1,9 +1,9 @@ { "$comment": "Single source of truth for all version numbers. See VERSIONING.md for bump rules.", - "product": "0.91.7", - "installer": "0.91.7", + "product": "0.92.0", + "installer": "0.92.0", "components": { - "ollama-proxy": "0.26.2", + "ollama-proxy": "0.27.0", "lmstudio-proxy": "0.16.2", "nvpair-node-info": "0.13.3", "nvpair-node-scanner": "0.20.3",