From d15a87935d2dc15832a4744b47ab45582ccee3b6 Mon Sep 17 00:00:00 2001 From: Jared Ping <13906915+JaredP94@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:49:15 +0200 Subject: [PATCH 1/2] fix(proxies): retry silent attempts and demote stalled nodes A peer can accept a connection without ever answering. Each request then spends the entire response-header deadline on that peer before failing over, because candidate ordering remembers nothing about the stall. Record full silent header and inference first-body deadlines in each proxy process and move the candidate behind its eligible peers for one configured deadline. Expiry and a successful response restore normal ordering. Keep cooling candidates routable, including in a single-node cluster, and stop reservations or prefix affinity from promoting them back over ready peers. Hold successful inference responses until initial body bytes or valid empty EOF arrive, so a silent or dropped pre-commit attempt can use the existing failover loop. Preserve the request and workload identity and never replay a committed stream. Apply the same implementation and frozen-upstream regressions to Ollama, LM Studio, and oMLX. The default remains 120 seconds. Adapt the design from NVIDIA/Personal-AI-Router#144 and the pre-commit retry in NVIDIA/Personal-AI-Router@e0f75dd, addressing NVIDIA/Personal-AI-Router#6. Retain the existing unreleased proxy patch versions relative to v0.8.0, following the fork's versioning practice. Co-authored-by: canja006 <82221082+canja006@users.noreply.github.com> Co-authored-by: Noah (Terve) Tervalon Signed-off-by: Jared Ping <13906915+JaredP94@users.noreply.github.com> --- docs/architecture.mdx | 20 ++ docs/prefix-aware-routing-spec.mdx | 3 +- services/lmstudio-proxy/README.md | 14 +- services/lmstudio-proxy/cooldown_test.go | 341 +++++++++++++++++++++++ services/lmstudio-proxy/proxy.go | 36 ++- services/lmstudio-proxy/stall.go | 140 ++++++++++ services/ollama-proxy/README.md | 14 +- services/ollama-proxy/cooldown_test.go | 341 +++++++++++++++++++++++ services/ollama-proxy/proxy.go | 36 ++- services/ollama-proxy/stall.go | 140 ++++++++++ services/omlx-proxy/README.md | 14 +- services/omlx-proxy/cooldown_test.go | 341 +++++++++++++++++++++++ services/omlx-proxy/proxy.go | 36 ++- services/omlx-proxy/stall.go | 140 ++++++++++ 14 files changed, 1594 insertions(+), 22 deletions(-) create mode 100644 services/lmstudio-proxy/cooldown_test.go create mode 100644 services/lmstudio-proxy/stall.go create mode 100644 services/ollama-proxy/cooldown_test.go create mode 100644 services/ollama-proxy/stall.go create mode 100644 services/omlx-proxy/cooldown_test.go create mode 100644 services/omlx-proxy/stall.go diff --git a/docs/architecture.mdx b/docs/architecture.mdx index 7e60218c5..5ef3e24a4 100644 --- a/docs/architecture.mdx +++ b/docs/architecture.mdx @@ -476,6 +476,15 @@ see the same idle node and pile onto it. To prevent that, each proxy adds the requests it has just dispatched itself to its own estimate and reserves its choice before forwarding, so a burst spreads without waiting for those reports. +After those preferences, each proxy moves candidates that just spent a full +silent response deadline behind their eligible peers, preserving relative order +within both groups. This measured, process-local cooldown overrides pins, +reservations, and prefix hints; it is not a liveness verdict. It never removes a +candidate, so a lone node is still tried and a cooling node can serve after the +others fail. It expires after one configured response-header duration, or clears +early on a successful response. Fast errors and caller cancellation do not set +or clear it. Nothing is persisted or shared between proxies. + ### Models Do Not Need to Be on Every Node **You do not have to put the same models on every node.** A cluster where one @@ -522,6 +531,17 @@ retryable and moves to the next advertised owner. It never fails over to an unknown or known-missing node. Genuine client errors such as `400` or `422` are not retried, because they would fail identically everywhere. +Successful inference responses do not commit on headers alone. The proxy waits +for initial body bytes (or a valid empty EOF), bounded by a separate wait equal +to the configured response-header timeout. Silence or a body read failure before +commit advances to the next existing candidate while retaining the same request +and workload identity. It does not re-resolve candidates or add retry rounds. +Once the response commits, errors only terminate that stream; it is never +replayed after bytes have reached the client. A full silent header or first-body +deadline also records the cooldown described above. The `120s` default and +`--response-header-timeout` / `NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT` configuration +remain unchanged; operators must allow for cold loads and long queues. + #### Planning Model Placement Where you put each model decides what routing can do with it: diff --git a/docs/prefix-aware-routing-spec.mdx b/docs/prefix-aware-routing-spec.mdx index 932a72df7..82439d792 100644 --- a/docs/prefix-aware-routing-spec.mdx +++ b/docs/prefix-aware-routing-spec.mdx @@ -26,7 +26,8 @@ design below is refined by these implementation decisions: KV-cache residence. Failover records the serving node. - Affinity load includes local optimistic reservations. The eligibility check, promotion, and reservation use the existing scheduler lock and retain manual - pin precedence and failover order. + pin precedence and failover order within each deadline-cooldown group. A warm + node in cooldown cannot be promoted over a non-cooling eligible peer. - Plain completions require at least 1,024 bytes before prefix affinity applies. Ollama supplied context arrays can identify repeated context only; updated response token arrays are not parsed or retained. diff --git a/services/lmstudio-proxy/README.md b/services/lmstudio-proxy/README.md index 24fb0dddb..f72f96e0b 100644 --- a/services/lmstudio-proxy/README.md +++ b/services/lmstudio-proxy/README.md @@ -35,7 +35,7 @@ lmstudio-proxy [flags] | `--port` | `1234` | HTTP listen port for request forwarding | | `--ignore-persisted-port` | `false` | Use `--port` even when a prior runtime port was saved | | `--ipc` | *(empty — use stdio)* | Path to a Unix domain socket or Windows named pipe for IPC | -| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | How long to wait for an upstream engine to start answering (Go duration, e.g. `10m`). An engine sends no headers until generation starts, so raise it when requests queue behind long jobs or wait on slow model loads. An invalid or non-positive value keeps the default | +| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline and, for successful inference, a separate wait for initial body bytes (Go duration, e.g. `10m`). Raise it for long queues or cold model loads. An invalid or non-positive value keeps the default | | `--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. | | `--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 | @@ -53,7 +53,8 @@ Node selection: - **Auto**: When no eligible node is explicitly selected, the proxy follows `node/set-priority` (see below), then discovered nodes in stable ID order. - **Priority (scheduler-driven)**: The Job Scheduler ranks the cluster least-loaded-first by pending workload plus smoothed GPU pressure and, via `nvpair-ui-broker`, pushes the ordered node list with those per-node counts to this proxy with `node/set-priority`. Auto routing sends the request to the listed node carrying the least estimated load. See [`nvpair-job-scheduler`](../nvpair-job-scheduler/README.md). - **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. +- **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing; a silent body times out using the configured response-header duration and fails over within the same candidate list. Once committed, a stream is never retried. Failover keeps one request/workload identity. +- **Deadline cooldown**: A full response-header wait with no answer, or an inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured deadline. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. The default deadline remains `120s`. ### IPC Transport @@ -338,8 +339,9 @@ 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. -- **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, +- **Eligible manual pin wins unless cooling.** An active `node/select` pin + takes precedence when it is in the request's owner set and is not behind + peers in deadline cooldown. An ineligible pin is ignored for that request, so automatic reservations still apply among eligible owners. Clearing the pin (`node/select` with an empty `id`) activates the most-recently-set list. Setting a priority list does **not** emit `node/selection-changed` (the manual selection @@ -347,8 +349,8 @@ Semantics: - **Unknown IDs are ignored.** IDs not currently in discovery are kept in the stored list (a node may appear later) but contribute nothing until discovered. - **Eligible unlisted nodes are a lowest-priority fallback.** An advertised owner - absent from the list stays routable, but only after every listed owner — - ordered among themselves by the default stable ID sort. This ensures an + absent from the list stays routable, but only after every listed owner in + the same cooldown group — ordered among themselves by the default stable ID sort. This ensures an eligible manually-added node the scheduler never saw is never stranded. - **Empty list reverts to default.** `{"nodes":[]}` clears the scheduler's influence and returns the proxy to its default auto ordering (eligible diff --git a/services/lmstudio-proxy/cooldown_test.go b/services/lmstudio-proxy/cooldown_test.go new file mode 100644 index 000000000..8663932eb --- /dev/null +++ b/services/lmstudio-proxy/cooldown_test.go @@ -0,0 +1,341 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "nvpair-shared/schedulerwire" +) + +type stallWire struct { + mu sync.Mutex + b []byte +} + +func (w *stallWire) Read([]byte) (int, error) { return 0, io.EOF } +func (w *stallWire) Write(b []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + w.b = append(w.b, b...) + return len(b), nil +} + +// Removing deadline demotion must send the second request into the frozen +// upstream again. Reset scheduler reservations so they cannot hide that bug. +func TestStalledCandidateSecondRequest(t *testing.T) { + old := proxyResponseTimeout + proxyResponseTimeout = time.Second + defer func() { proxyResponseTimeout = old }() + + for _, headers := range []bool{false, true} { + name := "no-headers" + if headers { + name = "headers-without-body" + } + t.Run(name, func(t *testing.T) { + var stalledHits, goodHits atomic.Int32 + release := make(chan struct{}) + stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + stalledHits.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + if headers { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + } + <-release + })) + defer stalled.Close() + defer close(release) + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "stalled", stalled.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + wire := &stallWire{} + p := NewProxy(NewCodec(wire), disc, 1235) + defer p.closeIdleTransports() + request := func() time.Duration { + t.Helper() + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"stalled", "good"}, + Ranks: []schedulerwire.NodeRank{{ID: "stalled"}, {ID: "good"}}, + }) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx) + rec := httptest.NewRecorder() + start := time.Now() + p.handleHTTP(rec, req) + if rec.Code != http.StatusOK || rec.Body.String() != "healthy" { + t.Fatalf("response = %d %q, want healthy failover", rec.Code, rec.Body.String()) + } + return time.Since(start) + } + + first := request() + if stalledHits.Load() != 1 { + t.Fatalf("first request frozen hits = %d, want 1", stalledHits.Load()) + } + second := request() + if stalledHits.Load() != 1 { + t.Fatalf("second request paid the deadline again: frozen hits=%d; durations first=%s second=%s", + stalledHits.Load(), first, second) + } + if goodHits.Load() != 2 { + t.Fatalf("healthy hits = %d, want 2", goodHits.Load()) + } + t.Logf("first=%s second=%s (configured deadline %s)", first, second, proxyResponseTimeout) + assertStallWorkloads(t, wire, 2) + }) + } +} + +func TestCooldownOrderingRecovery(t *testing.T) { + p := prProxy(t, "a", "b", "c", "d") + p.SetPriority([]string{"a", "b", "c", "d"}) + p.SetSelected("a") + now := time.Now() + p.noteStalled("a", now) + p.noteStalled("c", now) + assertOrder(t, candidateIDs(p), []string{"b", "d", "a", "c"}) + p.noteSuccess("a") + assertOrder(t, candidateIDs(p), []string{"a", "b", "d", "c"}) + if len(p.coolingSet(now.Add(proxyResponseTimeout))) != 0 { + t.Fatal("cooldown did not expire without traffic or success") + } + assertOrder(t, candidateIDs(p), []string{"a", "b", "c", "d"}) + other := prProxy(t, "a", "b") + p.noteStalled("a", now) + assertOrder(t, candidateIDs(other), []string{"a", "b"}) +} + +func TestCooldownReservationAndAffinity(t *testing.T) { + for _, tc := range []struct { + name, warm, pin string + ids, ranked, cooled, want []string + }{ + {name: "ranked", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "affinity", warm: "a", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "unlisted-peer", warm: "a", ids: []string{"a", "b"}, ranked: []string{"a"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "manual-selection", pin: "a", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "all-cooling", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a", "b"}, want: []string{"a", "b"}}, + {name: "sole-cooling", ids: []string{"a"}, ranked: []string{"a"}, cooled: []string{"a"}, want: []string{"a"}}, + } { + t.Run(tc.name, func(t *testing.T) { + p := prProxy(t, tc.ids...) + p.SetPriority(tc.ranked) + p.SetSelected(tc.pin) + for _, id := range tc.cooled { + p.noteStalled(id, time.Now()) + } + got, affinity := p.reserveCandidateWithAffinity(reservationCandidates(tc.ids...), tc.warm) + var ids []string + for _, cand := range got { + ids = append(ids, cand.id) + } + assertOrder(t, ids, tc.want) + if affinity { + t.Fatal("affinity promoted a cooling candidate over a ready peer") + } + }) + } +} + +func TestPrecommitDropRetriesWithoutCooldown(t *testing.T) { + dropped := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // A 200 followed by EOF before the declared body starts. + w.Header().Set("Content-Length", "999") + w.WriteHeader(http.StatusOK) + })) + defer dropped.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "dropped", dropped.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + wire := &stallWire{} + p := NewProxy(NewCodec(wire), disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"dropped", "good"}) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`))) + if rec.Code != 200 || rec.Body.String() != "healthy" { + t.Fatalf("precommit drop did not fail over: %d %q", rec.Code, rec.Body.String()) + } + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("fast body failure was incorrectly classified as a stall") + } + assertStallWorkloads(t, wire, 1) +} + +func TestCommittedStreamNeverRetries(t *testing.T) { + var goodHits atomic.Int32 + broken := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", "999") + _, _ = io.WriteString(w, "partial") + })) + defer broken.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "broken", broken.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"broken", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer server.Close() + resp, err := server.Client().Post(server.URL+"/v1/chat/completions", "application/json", + strings.NewReader(`{"model":"model"}`)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err == nil || string(data) != "partial" || goodHits.Load() != 0 { + t.Fatalf("committed stream was retried or not truncated: body=%q error=%v goodHits=%d", data, err, goodHits.Load()) + } +} + +func TestCooldownSuccessAndFastStatus(t *testing.T) { + for _, code := range []int{http.StatusOK, http.StatusInternalServerError} { + t.Run(http.StatusText(code), func(t *testing.T) { + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(code) // empty 200 is a valid complete response + })) + defer engine.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.noteStalled("only", time.Now()) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`))) + if rec.Code != code { + t.Fatalf("sole cooling candidate was not tried: status %d, want %d", rec.Code, code) + } + if cooling := p.coolingSet(time.Now())["only"]; cooling != (code != http.StatusOK) { + t.Fatalf("cooling=%v for status %d: only success should clear", cooling, code) + } + }) + } +} + +func TestCallerCancellationDoesNotCoolCandidate(t *testing.T) { + for _, headers := range []bool{false, true} { + t.Run(fmt.Sprint(headers), func(t *testing.T) { + release := make(chan struct{}) + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if headers { + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + } + <-release + })) + defer engine.Close() + defer close(release) + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("a caller deadline demoted the candidate") + } + }) + } +} + +func TestFinalBodyStallIsBounded(t *testing.T) { + old := proxyResponseTimeout + proxyResponseTimeout = 100 * time.Millisecond + defer func() { proxyResponseTimeout = old }() + release := make(chan struct{}) + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + <-release + })) + defer engine.Close() + defer close(release) + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if rec.Code != http.StatusBadGateway || ctx.Err() != nil { + t.Fatalf("final body stall was not bounded before client deadline: status=%d context=%v", rec.Code, ctx.Err()) + } + if !p.coolingSet(time.Now())["only"] { + t.Fatal("final body stall was not remembered") + } +} + +func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { + t.Helper() + wire.mu.Lock() + data := append([]byte(nil), wire.b...) + wire.mu.Unlock() + started := make(map[string]bool) + completed := make(map[string]bool) + dec := json.NewDecoder(strings.NewReader(string(data))) + for dec.More() { + var frame struct { + Method string `json:"method"` + Params struct { + WorkloadInfo Workload `json:"workloadInfo"` + } `json:"params"` + } + if err := dec.Decode(&frame); err != nil { + t.Fatal(err) + } + wl := frame.Params.WorkloadInfo + switch frame.Method { + case workloadStartedMethod: + started[wl.ID] = true + case workloadCompletedMethod: + if completed[wl.ID] || wl.ScheduledOn != "good" || !started[wl.ID] { + t.Fatalf("invalid or duplicate completion: %+v", wl) + } + completed[wl.ID] = true + case workloadErroredMethod: + t.Fatalf("failover emitted a failed workload: %+v", wl) + } + } + if len(started) != count || len(completed) != count { + t.Fatalf("workloads started=%v completed=%v, want %d identities", started, completed, count) + } +} diff --git a/services/lmstudio-proxy/proxy.go b/services/lmstudio-proxy/proxy.go index 95d4a9f4d..7cc50588e 100644 --- a/services/lmstudio-proxy/proxy.go +++ b/services/lmstudio-proxy/proxy.go @@ -388,6 +388,10 @@ type Proxy struct { priorityGPUPressure map[string]int priorityReservations map[string]int + // coolMu guards process-local deadline measurements, never persisted. + coolMu sync.Mutex + coolUntil map[string]time.Time + // 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 @@ -1353,7 +1357,7 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { Director: p.forwardDirector(cand, reqID), // A remote cluster peer is dialed over mTLS (per-peer pinned config); // self/manual candidates use the plain transport. See candidateTransport. - Transport: p.candidateTransport(cand), + Transport: stallTransport{base: p.candidateTransport(cand), proxy: p, id: cand.id, timeout: proxyResponseTimeout}, BufferPool: p.bufPool, // ModifyResponse fires when the upstream's status line + headers // have arrived but before the body streams. That's both the retry @@ -1365,6 +1369,21 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { retry = true return retrySignal{} } + // Only successful inference waits for content. All commit effects + // below still happen once, for the node that actually serves. + if isInf && resp.StatusCode >= 200 && resp.StatusCode < 300 { + body, err := awaitFirstBody(r.Context(), resp.Body, proxyResponseTimeout) + if err != nil { + if r.Context().Err() == nil && stderrors.Is(err, errFirstBodyTimeout) { + p.noteStalled(cand.id, time.Now()) + } + return err + } + resp.Body = body + } + if resp.StatusCode >= 200 && resp.StatusCode < 300 { + p.noteSuccess(cand.id) + } // Prefer an engine-declared preflight policy so an exact origin plus // Allow-Credentials can pass a credentialed browser fetch. Engines // that publish no policy retain the proxy's permissive 204 fallback. @@ -1791,6 +1810,8 @@ func (p *Proxy) resolveCandidates(model string) []candidate { } } + out = demoteCooling(out, p.coolingSet(time.Now())) + slog.Debug("resolveCandidates resolved", "selected", id, "priority", len(priority), "candidates", len(out), "eligible", len(nodes), "known", known) @@ -1815,6 +1836,15 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID if len(candidates) == 0 { return candidates, false } + cooling := p.coolingSet(time.Now()) + candidates = demoteCooling(candidates, cooling) + hasReady := false + for _, cand := range candidates { + if !cooling[cand.id] { + hasReady = true + break + } + } if selectedID := p.SelectedID(); selectedID != "" { for _, cand := range candidates { if cand.id == selectedID { @@ -1837,6 +1867,7 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID // Check affinity under the same lock as optimistic reservations. An eligible // manual pin was handled above. Missing or overloaded hints leave ordering alone. if index, ok := candidateIndex[warmNodeID]; ok && warmNodeID != "" && + (!hasReady || !cooling[warmNodeID]) && p.priorityPending[warmNodeID]+p.priorityReservations[warmNodeID] <= 1 && p.priorityGPUPressure[warmNodeID] < schedulerwire.MaxGPUPressure { chosen := candidates[index] @@ -1850,6 +1881,9 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID bestOrder := len(p.priority) var bestLoad uint64 for order, id := range p.priority { + if hasReady && cooling[id] { + continue + } index, ok := candidateIndex[id] if !ok { continue diff --git a/services/lmstudio-proxy/stall.go b/services/lmstudio-proxy/stall.go new file mode 100644 index 000000000..449a1311d --- /dev/null +++ b/services/lmstudio-proxy/stall.go @@ -0,0 +1,140 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptrace" + "sync" + "time" +) + +var errFirstBodyTimeout = errors.New("upstream timed out before first response body byte") + +// A cooldown is evidence of a full silent deadline, not a liveness verdict. +// State is local to this proxy; success or time alone restores normal ordering. +func (p *Proxy) noteStalled(id string, now time.Time) { + p.coolMu.Lock() + defer p.coolMu.Unlock() + if p.coolUntil == nil { + p.coolUntil = make(map[string]time.Time) + } + p.coolUntil[id] = now.Add(proxyResponseTimeout) +} + +func (p *Proxy) noteSuccess(id string) { + p.coolMu.Lock() + defer p.coolMu.Unlock() + delete(p.coolUntil, id) +} + +func (p *Proxy) coolingSet(now time.Time) map[string]bool { + p.coolMu.Lock() + defer p.coolMu.Unlock() + if len(p.coolUntil) == 0 { + return nil + } + cooling := make(map[string]bool, len(p.coolUntil)) + for id, until := range p.coolUntil { + if now.Before(until) { + cooling[id] = true + } else { + delete(p.coolUntil, id) + } + } + return cooling +} + +// Stable partition: keep every candidate, including a sole cooling candidate. +func demoteCooling(candidates []candidate, cooling map[string]bool) []candidate { + if len(cooling) == 0 || len(candidates) < 2 { + return candidates + } + out := make([]candidate, 0, len(candidates)) + for _, cand := range candidates { + if !cooling[cand.id] { + out = append(out, cand) + } + } + for _, cand := range candidates { + if cooling[cand.id] { + out = append(out, cand) + } + } + return out +} + +// Observe the response-header wait after the request was written, not dial or +// TLS time. A shorter caller deadline, cancellation, or fast transport failure +// is not evidence that this candidate spent a full response deadline silent. +type stallTransport struct { + base http.RoundTripper + proxy *Proxy + id string + timeout time.Duration +} + +func (t stallTransport) RoundTrip(req *http.Request) (*http.Response, error) { + var mu sync.Mutex + var written time.Time + trace := &httptrace.ClientTrace{WroteRequest: func(info httptrace.WroteRequestInfo) { + if info.Err == nil { + mu.Lock() + written = time.Now() + mu.Unlock() + } + }} + resp, err := t.base.RoundTrip(req.WithContext(httptrace.WithClientTrace(req.Context(), trace))) + mu.Lock() + waitStart := written + mu.Unlock() + var timeout net.Error + if err != nil && req.Context().Err() == nil && errors.As(err, &timeout) && + timeout.Timeout() && !waitStart.IsZero() && time.Since(waitStart) >= t.timeout { + t.proxy.noteStalled(t.id, time.Now()) + } + return resp, err +} + +type bodyWithPeek struct { + io.Reader + io.Closer +} + +// Hold inference headers until content (or a valid empty EOF) arrives. On an +// error the caller returns from ModifyResponse, so ReverseProxy closes the +// original body and unblocks the pending read. The buffered channel lets that +// reader exit even if the deadline or caller cancellation won. +func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duration) (io.ReadCloser, error) { + type result struct { + data []byte + err error + } + ready := make(chan result, 1) + go func() { + // Match ReverseProxy's copy buffer. Reading only one byte would split + // the first engine frame and publish progress before it can be parsed. + first := make([]byte, 32*1024) + n, err := body.Read(first) + ready <- result{data: first[:n], err: err} + }() + timer := time.NewTimer(budget) + defer timer.Stop() + select { + case got := <-ready: + if len(got.data) == 0 && got.err != nil && !errors.Is(got.err, io.EOF) { + return nil, got.err + } + return bodyWithPeek{Reader: io.MultiReader(bytes.NewReader(got.data), body), Closer: body}, nil + case <-ctx.Done(): + return nil, ctx.Err() + case <-timer.C: + return nil, errFirstBodyTimeout + } +} diff --git a/services/ollama-proxy/README.md b/services/ollama-proxy/README.md index 998e2f17b..8441b0e5d 100644 --- a/services/ollama-proxy/README.md +++ b/services/ollama-proxy/README.md @@ -34,7 +34,7 @@ ollama-proxy [flags] | `--alias-address` | *(empty)* | Optional secondary `host:port` for the same routing handler. Repeat the flag to reserve both loopback families for `localhost`. Only literal loopback addresses are accepted; the broker uses this for a safe inherited local `OLLAMA_HOST`, and the aliases are not advertised to peers. | | `--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 | -| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | How long to wait for an upstream engine to start answering (Go duration, e.g. `10m`). An engine sends no headers until generation starts, so raise it when requests queue behind long jobs or wait on slow model loads. An invalid or non-positive value keeps the default | +| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline and, for successful inference, a separate wait for initial body bytes (Go duration, e.g. `10m`). Raise it for long queues or cold model loads. An invalid or non-positive value keeps the default | | `--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. | | `--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 | @@ -58,7 +58,8 @@ Node selection: - **Auto**: When no eligible node is explicitly selected, the proxy follows `node/set-priority` (see below), then discovered nodes in stable ID order. - **Priority (scheduler-driven)**: The Job Scheduler ranks the cluster least-loaded-first by pending workload plus smoothed GPU pressure and, via `nvpair-ui-broker`, pushes the ordered node list with those per-node counts to this proxy with `node/set-priority`. Auto routing sends the request to the listed node carrying the least estimated load. See [`nvpair-job-scheduler`](../nvpair-job-scheduler/README.md). - **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. +- **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing; a silent body times out using the configured response-header duration and fails over within the same candidate list. Once committed, a stream is never retried. Failover keeps one request/workload identity. +- **Deadline cooldown**: A full response-header wait with no answer, or an inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured deadline. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. The default deadline remains `120s`. ### IPC Transport @@ -342,8 +343,9 @@ 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. -- **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, +- **Eligible manual pin wins unless cooling.** An active `node/select` pin + takes precedence when it is in the request's owner set and is not behind + peers in deadline cooldown. An ineligible pin is ignored for that request, so automatic reservations still apply among eligible owners. Clearing the pin (`node/select` with an empty `id`) activates the most-recently-set list. Setting a priority list does **not** emit `node/selection-changed` (the manual selection @@ -351,8 +353,8 @@ Semantics: - **Unknown IDs are ignored.** IDs not currently in discovery are kept in the stored list (a node may appear later) but contribute nothing until discovered. - **Eligible unlisted nodes are a lowest-priority fallback.** An advertised owner - absent from the list stays routable, but only after every listed owner — - ordered among themselves by the default stable ID sort. This ensures an + absent from the list stays routable, but only after every listed owner in + the same cooldown group — ordered among themselves by the default stable ID sort. This ensures an eligible manually-added node the scheduler never saw is never stranded. - **Empty list reverts to default.** `{"nodes":[]}` clears the scheduler's influence and returns the proxy to its default auto ordering (eligible diff --git a/services/ollama-proxy/cooldown_test.go b/services/ollama-proxy/cooldown_test.go new file mode 100644 index 000000000..8663932eb --- /dev/null +++ b/services/ollama-proxy/cooldown_test.go @@ -0,0 +1,341 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "nvpair-shared/schedulerwire" +) + +type stallWire struct { + mu sync.Mutex + b []byte +} + +func (w *stallWire) Read([]byte) (int, error) { return 0, io.EOF } +func (w *stallWire) Write(b []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + w.b = append(w.b, b...) + return len(b), nil +} + +// Removing deadline demotion must send the second request into the frozen +// upstream again. Reset scheduler reservations so they cannot hide that bug. +func TestStalledCandidateSecondRequest(t *testing.T) { + old := proxyResponseTimeout + proxyResponseTimeout = time.Second + defer func() { proxyResponseTimeout = old }() + + for _, headers := range []bool{false, true} { + name := "no-headers" + if headers { + name = "headers-without-body" + } + t.Run(name, func(t *testing.T) { + var stalledHits, goodHits atomic.Int32 + release := make(chan struct{}) + stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + stalledHits.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + if headers { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + } + <-release + })) + defer stalled.Close() + defer close(release) + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "stalled", stalled.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + wire := &stallWire{} + p := NewProxy(NewCodec(wire), disc, 1235) + defer p.closeIdleTransports() + request := func() time.Duration { + t.Helper() + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"stalled", "good"}, + Ranks: []schedulerwire.NodeRank{{ID: "stalled"}, {ID: "good"}}, + }) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx) + rec := httptest.NewRecorder() + start := time.Now() + p.handleHTTP(rec, req) + if rec.Code != http.StatusOK || rec.Body.String() != "healthy" { + t.Fatalf("response = %d %q, want healthy failover", rec.Code, rec.Body.String()) + } + return time.Since(start) + } + + first := request() + if stalledHits.Load() != 1 { + t.Fatalf("first request frozen hits = %d, want 1", stalledHits.Load()) + } + second := request() + if stalledHits.Load() != 1 { + t.Fatalf("second request paid the deadline again: frozen hits=%d; durations first=%s second=%s", + stalledHits.Load(), first, second) + } + if goodHits.Load() != 2 { + t.Fatalf("healthy hits = %d, want 2", goodHits.Load()) + } + t.Logf("first=%s second=%s (configured deadline %s)", first, second, proxyResponseTimeout) + assertStallWorkloads(t, wire, 2) + }) + } +} + +func TestCooldownOrderingRecovery(t *testing.T) { + p := prProxy(t, "a", "b", "c", "d") + p.SetPriority([]string{"a", "b", "c", "d"}) + p.SetSelected("a") + now := time.Now() + p.noteStalled("a", now) + p.noteStalled("c", now) + assertOrder(t, candidateIDs(p), []string{"b", "d", "a", "c"}) + p.noteSuccess("a") + assertOrder(t, candidateIDs(p), []string{"a", "b", "d", "c"}) + if len(p.coolingSet(now.Add(proxyResponseTimeout))) != 0 { + t.Fatal("cooldown did not expire without traffic or success") + } + assertOrder(t, candidateIDs(p), []string{"a", "b", "c", "d"}) + other := prProxy(t, "a", "b") + p.noteStalled("a", now) + assertOrder(t, candidateIDs(other), []string{"a", "b"}) +} + +func TestCooldownReservationAndAffinity(t *testing.T) { + for _, tc := range []struct { + name, warm, pin string + ids, ranked, cooled, want []string + }{ + {name: "ranked", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "affinity", warm: "a", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "unlisted-peer", warm: "a", ids: []string{"a", "b"}, ranked: []string{"a"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "manual-selection", pin: "a", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "all-cooling", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a", "b"}, want: []string{"a", "b"}}, + {name: "sole-cooling", ids: []string{"a"}, ranked: []string{"a"}, cooled: []string{"a"}, want: []string{"a"}}, + } { + t.Run(tc.name, func(t *testing.T) { + p := prProxy(t, tc.ids...) + p.SetPriority(tc.ranked) + p.SetSelected(tc.pin) + for _, id := range tc.cooled { + p.noteStalled(id, time.Now()) + } + got, affinity := p.reserveCandidateWithAffinity(reservationCandidates(tc.ids...), tc.warm) + var ids []string + for _, cand := range got { + ids = append(ids, cand.id) + } + assertOrder(t, ids, tc.want) + if affinity { + t.Fatal("affinity promoted a cooling candidate over a ready peer") + } + }) + } +} + +func TestPrecommitDropRetriesWithoutCooldown(t *testing.T) { + dropped := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // A 200 followed by EOF before the declared body starts. + w.Header().Set("Content-Length", "999") + w.WriteHeader(http.StatusOK) + })) + defer dropped.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "dropped", dropped.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + wire := &stallWire{} + p := NewProxy(NewCodec(wire), disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"dropped", "good"}) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`))) + if rec.Code != 200 || rec.Body.String() != "healthy" { + t.Fatalf("precommit drop did not fail over: %d %q", rec.Code, rec.Body.String()) + } + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("fast body failure was incorrectly classified as a stall") + } + assertStallWorkloads(t, wire, 1) +} + +func TestCommittedStreamNeverRetries(t *testing.T) { + var goodHits atomic.Int32 + broken := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", "999") + _, _ = io.WriteString(w, "partial") + })) + defer broken.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "broken", broken.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"broken", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer server.Close() + resp, err := server.Client().Post(server.URL+"/v1/chat/completions", "application/json", + strings.NewReader(`{"model":"model"}`)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err == nil || string(data) != "partial" || goodHits.Load() != 0 { + t.Fatalf("committed stream was retried or not truncated: body=%q error=%v goodHits=%d", data, err, goodHits.Load()) + } +} + +func TestCooldownSuccessAndFastStatus(t *testing.T) { + for _, code := range []int{http.StatusOK, http.StatusInternalServerError} { + t.Run(http.StatusText(code), func(t *testing.T) { + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(code) // empty 200 is a valid complete response + })) + defer engine.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.noteStalled("only", time.Now()) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`))) + if rec.Code != code { + t.Fatalf("sole cooling candidate was not tried: status %d, want %d", rec.Code, code) + } + if cooling := p.coolingSet(time.Now())["only"]; cooling != (code != http.StatusOK) { + t.Fatalf("cooling=%v for status %d: only success should clear", cooling, code) + } + }) + } +} + +func TestCallerCancellationDoesNotCoolCandidate(t *testing.T) { + for _, headers := range []bool{false, true} { + t.Run(fmt.Sprint(headers), func(t *testing.T) { + release := make(chan struct{}) + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if headers { + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + } + <-release + })) + defer engine.Close() + defer close(release) + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("a caller deadline demoted the candidate") + } + }) + } +} + +func TestFinalBodyStallIsBounded(t *testing.T) { + old := proxyResponseTimeout + proxyResponseTimeout = 100 * time.Millisecond + defer func() { proxyResponseTimeout = old }() + release := make(chan struct{}) + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + <-release + })) + defer engine.Close() + defer close(release) + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if rec.Code != http.StatusBadGateway || ctx.Err() != nil { + t.Fatalf("final body stall was not bounded before client deadline: status=%d context=%v", rec.Code, ctx.Err()) + } + if !p.coolingSet(time.Now())["only"] { + t.Fatal("final body stall was not remembered") + } +} + +func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { + t.Helper() + wire.mu.Lock() + data := append([]byte(nil), wire.b...) + wire.mu.Unlock() + started := make(map[string]bool) + completed := make(map[string]bool) + dec := json.NewDecoder(strings.NewReader(string(data))) + for dec.More() { + var frame struct { + Method string `json:"method"` + Params struct { + WorkloadInfo Workload `json:"workloadInfo"` + } `json:"params"` + } + if err := dec.Decode(&frame); err != nil { + t.Fatal(err) + } + wl := frame.Params.WorkloadInfo + switch frame.Method { + case workloadStartedMethod: + started[wl.ID] = true + case workloadCompletedMethod: + if completed[wl.ID] || wl.ScheduledOn != "good" || !started[wl.ID] { + t.Fatalf("invalid or duplicate completion: %+v", wl) + } + completed[wl.ID] = true + case workloadErroredMethod: + t.Fatalf("failover emitted a failed workload: %+v", wl) + } + } + if len(started) != count || len(completed) != count { + t.Fatalf("workloads started=%v completed=%v, want %d identities", started, completed, count) + } +} diff --git a/services/ollama-proxy/proxy.go b/services/ollama-proxy/proxy.go index 11f5f4be5..196b49c0e 100644 --- a/services/ollama-proxy/proxy.go +++ b/services/ollama-proxy/proxy.go @@ -398,6 +398,10 @@ type Proxy struct { priorityGPUPressure map[string]int priorityReservations map[string]int + // coolMu guards process-local deadline measurements, never persisted. + coolMu sync.Mutex + coolUntil map[string]time.Time + // 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 @@ -1539,7 +1543,7 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { Director: p.forwardDirector(cand, reqID), // A remote cluster peer is dialed over mTLS (per-peer pinned config); // self/manual candidates use the plain transport. See candidateTransport. - Transport: p.candidateTransport(cand), + Transport: stallTransport{base: p.candidateTransport(cand), proxy: p, id: cand.id, timeout: proxyResponseTimeout}, BufferPool: p.bufPool, // ModifyResponse fires when the upstream's status line + headers // have arrived but before the body streams. That's both the retry @@ -1551,6 +1555,21 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { retry = true return retrySignal{} } + // Only successful inference waits for content. All commit effects + // below still happen once, for the node that actually serves. + if isInf && resp.StatusCode >= 200 && resp.StatusCode < 300 { + body, err := awaitFirstBody(r.Context(), resp.Body, proxyResponseTimeout) + if err != nil { + if r.Context().Err() == nil && stderrors.Is(err, errFirstBodyTimeout) { + p.noteStalled(cand.id, time.Now()) + } + return err + } + resp.Body = body + } + if resp.StatusCode >= 200 && resp.StatusCode < 300 { + p.noteSuccess(cand.id) + } // Prefer an engine-declared preflight policy so an exact origin plus // Allow-Credentials can pass a credentialed browser fetch. Engines // that publish no policy retain the proxy's permissive 204 fallback. @@ -1986,6 +2005,8 @@ func (p *Proxy) resolveCandidates(model string) []candidate { } } + out = demoteCooling(out, p.coolingSet(time.Now())) + slog.Debug("resolveCandidates resolved", "selected", id, "priority", len(priority), "candidates", len(out), "eligible", len(nodes), "known", known) @@ -2010,6 +2031,15 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID if len(candidates) == 0 { return candidates, false } + cooling := p.coolingSet(time.Now()) + candidates = demoteCooling(candidates, cooling) + hasReady := false + for _, cand := range candidates { + if !cooling[cand.id] { + hasReady = true + break + } + } if selectedID := p.SelectedID(); selectedID != "" { for _, cand := range candidates { if cand.id == selectedID { @@ -2032,6 +2062,7 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID // Check affinity under the same lock as optimistic reservations. An eligible // manual pin was handled above. Missing or overloaded hints leave ordering alone. if index, ok := candidateIndex[warmNodeID]; ok && warmNodeID != "" && + (!hasReady || !cooling[warmNodeID]) && p.priorityPending[warmNodeID]+p.priorityReservations[warmNodeID] <= 1 && p.priorityGPUPressure[warmNodeID] < schedulerwire.MaxGPUPressure { chosen := candidates[index] @@ -2045,6 +2076,9 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID bestOrder := len(p.priority) var bestLoad uint64 for order, id := range p.priority { + if hasReady && cooling[id] { + continue + } index, ok := candidateIndex[id] if !ok { continue diff --git a/services/ollama-proxy/stall.go b/services/ollama-proxy/stall.go new file mode 100644 index 000000000..449a1311d --- /dev/null +++ b/services/ollama-proxy/stall.go @@ -0,0 +1,140 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptrace" + "sync" + "time" +) + +var errFirstBodyTimeout = errors.New("upstream timed out before first response body byte") + +// A cooldown is evidence of a full silent deadline, not a liveness verdict. +// State is local to this proxy; success or time alone restores normal ordering. +func (p *Proxy) noteStalled(id string, now time.Time) { + p.coolMu.Lock() + defer p.coolMu.Unlock() + if p.coolUntil == nil { + p.coolUntil = make(map[string]time.Time) + } + p.coolUntil[id] = now.Add(proxyResponseTimeout) +} + +func (p *Proxy) noteSuccess(id string) { + p.coolMu.Lock() + defer p.coolMu.Unlock() + delete(p.coolUntil, id) +} + +func (p *Proxy) coolingSet(now time.Time) map[string]bool { + p.coolMu.Lock() + defer p.coolMu.Unlock() + if len(p.coolUntil) == 0 { + return nil + } + cooling := make(map[string]bool, len(p.coolUntil)) + for id, until := range p.coolUntil { + if now.Before(until) { + cooling[id] = true + } else { + delete(p.coolUntil, id) + } + } + return cooling +} + +// Stable partition: keep every candidate, including a sole cooling candidate. +func demoteCooling(candidates []candidate, cooling map[string]bool) []candidate { + if len(cooling) == 0 || len(candidates) < 2 { + return candidates + } + out := make([]candidate, 0, len(candidates)) + for _, cand := range candidates { + if !cooling[cand.id] { + out = append(out, cand) + } + } + for _, cand := range candidates { + if cooling[cand.id] { + out = append(out, cand) + } + } + return out +} + +// Observe the response-header wait after the request was written, not dial or +// TLS time. A shorter caller deadline, cancellation, or fast transport failure +// is not evidence that this candidate spent a full response deadline silent. +type stallTransport struct { + base http.RoundTripper + proxy *Proxy + id string + timeout time.Duration +} + +func (t stallTransport) RoundTrip(req *http.Request) (*http.Response, error) { + var mu sync.Mutex + var written time.Time + trace := &httptrace.ClientTrace{WroteRequest: func(info httptrace.WroteRequestInfo) { + if info.Err == nil { + mu.Lock() + written = time.Now() + mu.Unlock() + } + }} + resp, err := t.base.RoundTrip(req.WithContext(httptrace.WithClientTrace(req.Context(), trace))) + mu.Lock() + waitStart := written + mu.Unlock() + var timeout net.Error + if err != nil && req.Context().Err() == nil && errors.As(err, &timeout) && + timeout.Timeout() && !waitStart.IsZero() && time.Since(waitStart) >= t.timeout { + t.proxy.noteStalled(t.id, time.Now()) + } + return resp, err +} + +type bodyWithPeek struct { + io.Reader + io.Closer +} + +// Hold inference headers until content (or a valid empty EOF) arrives. On an +// error the caller returns from ModifyResponse, so ReverseProxy closes the +// original body and unblocks the pending read. The buffered channel lets that +// reader exit even if the deadline or caller cancellation won. +func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duration) (io.ReadCloser, error) { + type result struct { + data []byte + err error + } + ready := make(chan result, 1) + go func() { + // Match ReverseProxy's copy buffer. Reading only one byte would split + // the first engine frame and publish progress before it can be parsed. + first := make([]byte, 32*1024) + n, err := body.Read(first) + ready <- result{data: first[:n], err: err} + }() + timer := time.NewTimer(budget) + defer timer.Stop() + select { + case got := <-ready: + if len(got.data) == 0 && got.err != nil && !errors.Is(got.err, io.EOF) { + return nil, got.err + } + return bodyWithPeek{Reader: io.MultiReader(bytes.NewReader(got.data), body), Closer: body}, nil + case <-ctx.Done(): + return nil, ctx.Err() + case <-timer.C: + return nil, errFirstBodyTimeout + } +} diff --git a/services/omlx-proxy/README.md b/services/omlx-proxy/README.md index ca1684745..c2917c92b 100644 --- a/services/omlx-proxy/README.md +++ b/services/omlx-proxy/README.md @@ -35,7 +35,7 @@ omlx-proxy [flags] | `--port` | `1234` | HTTP listen port for request forwarding | | `--ignore-persisted-port` | `false` | Use `--port` even when a prior runtime port was saved | | `--ipc` | *(empty — use stdio)* | Path to a Unix domain socket or Windows named pipe for IPC | -| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | How long to wait for an upstream engine to start answering (Go duration, e.g. `10m`). An engine sends no headers until generation starts, so raise it when requests queue behind long jobs or wait on slow model loads. An invalid or non-positive value keeps the default | +| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline and, for successful inference, a separate wait for initial body bytes (Go duration, e.g. `10m`). Raise it for long queues or cold model loads. An invalid or non-positive value keeps the default | | `--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. | | `--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 | @@ -53,7 +53,8 @@ Node selection: - **Auto**: When no eligible node is explicitly selected, the proxy follows `node/set-priority` (see below), then discovered nodes in stable ID order. - **Priority (scheduler-driven)**: The Job Scheduler ranks the cluster least-loaded-first by pending workload plus smoothed GPU pressure and, via `nvpair-ui-broker`, pushes the ordered node list with those per-node counts to this proxy with `node/set-priority`. Auto routing sends the request to the listed node carrying the least estimated load. See [`nvpair-job-scheduler`](../nvpair-job-scheduler/README.md). - **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. +- **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing; a silent body times out using the configured response-header duration and fails over within the same candidate list. Once committed, a stream is never retried. Failover keeps one request/workload identity. +- **Deadline cooldown**: A full response-header wait with no answer, or an inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured deadline. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. The default deadline remains `120s`. ### IPC Transport @@ -353,8 +354,9 @@ 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. -- **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, +- **Eligible manual pin wins unless cooling.** An active `node/select` pin + takes precedence when it is in the request's owner set and is not behind + peers in deadline cooldown. An ineligible pin is ignored for that request, so automatic reservations still apply among eligible owners. Clearing the pin (`node/select` with an empty `id`) activates the most-recently-set list. Setting a priority list does **not** emit `node/selection-changed` (the manual selection @@ -362,8 +364,8 @@ Semantics: - **Unknown IDs are ignored.** IDs not currently in discovery are kept in the stored list (a node may appear later) but contribute nothing until discovered. - **Eligible unlisted nodes are a lowest-priority fallback.** An advertised owner - absent from the list stays routable, but only after every listed owner — - ordered among themselves by the default stable ID sort. This ensures an + absent from the list stays routable, but only after every listed owner in + the same cooldown group — ordered among themselves by the default stable ID sort. This ensures an eligible manually-added node the scheduler never saw is never stranded. - **Empty list reverts to default.** `{"nodes":[]}` clears the scheduler's influence and returns the proxy to its default auto ordering (eligible diff --git a/services/omlx-proxy/cooldown_test.go b/services/omlx-proxy/cooldown_test.go new file mode 100644 index 000000000..8663932eb --- /dev/null +++ b/services/omlx-proxy/cooldown_test.go @@ -0,0 +1,341 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "nvpair-shared/schedulerwire" +) + +type stallWire struct { + mu sync.Mutex + b []byte +} + +func (w *stallWire) Read([]byte) (int, error) { return 0, io.EOF } +func (w *stallWire) Write(b []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + w.b = append(w.b, b...) + return len(b), nil +} + +// Removing deadline demotion must send the second request into the frozen +// upstream again. Reset scheduler reservations so they cannot hide that bug. +func TestStalledCandidateSecondRequest(t *testing.T) { + old := proxyResponseTimeout + proxyResponseTimeout = time.Second + defer func() { proxyResponseTimeout = old }() + + for _, headers := range []bool{false, true} { + name := "no-headers" + if headers { + name = "headers-without-body" + } + t.Run(name, func(t *testing.T) { + var stalledHits, goodHits atomic.Int32 + release := make(chan struct{}) + stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + stalledHits.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + if headers { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + } + <-release + })) + defer stalled.Close() + defer close(release) + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "stalled", stalled.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + wire := &stallWire{} + p := NewProxy(NewCodec(wire), disc, 1235) + defer p.closeIdleTransports() + request := func() time.Duration { + t.Helper() + p.SetPrioritySnapshot(schedulerwire.Priority{ + Nodes: []string{"stalled", "good"}, + Ranks: []schedulerwire.NodeRank{{ID: "stalled"}, {ID: "good"}}, + }) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx) + rec := httptest.NewRecorder() + start := time.Now() + p.handleHTTP(rec, req) + if rec.Code != http.StatusOK || rec.Body.String() != "healthy" { + t.Fatalf("response = %d %q, want healthy failover", rec.Code, rec.Body.String()) + } + return time.Since(start) + } + + first := request() + if stalledHits.Load() != 1 { + t.Fatalf("first request frozen hits = %d, want 1", stalledHits.Load()) + } + second := request() + if stalledHits.Load() != 1 { + t.Fatalf("second request paid the deadline again: frozen hits=%d; durations first=%s second=%s", + stalledHits.Load(), first, second) + } + if goodHits.Load() != 2 { + t.Fatalf("healthy hits = %d, want 2", goodHits.Load()) + } + t.Logf("first=%s second=%s (configured deadline %s)", first, second, proxyResponseTimeout) + assertStallWorkloads(t, wire, 2) + }) + } +} + +func TestCooldownOrderingRecovery(t *testing.T) { + p := prProxy(t, "a", "b", "c", "d") + p.SetPriority([]string{"a", "b", "c", "d"}) + p.SetSelected("a") + now := time.Now() + p.noteStalled("a", now) + p.noteStalled("c", now) + assertOrder(t, candidateIDs(p), []string{"b", "d", "a", "c"}) + p.noteSuccess("a") + assertOrder(t, candidateIDs(p), []string{"a", "b", "d", "c"}) + if len(p.coolingSet(now.Add(proxyResponseTimeout))) != 0 { + t.Fatal("cooldown did not expire without traffic or success") + } + assertOrder(t, candidateIDs(p), []string{"a", "b", "c", "d"}) + other := prProxy(t, "a", "b") + p.noteStalled("a", now) + assertOrder(t, candidateIDs(other), []string{"a", "b"}) +} + +func TestCooldownReservationAndAffinity(t *testing.T) { + for _, tc := range []struct { + name, warm, pin string + ids, ranked, cooled, want []string + }{ + {name: "ranked", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "affinity", warm: "a", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "unlisted-peer", warm: "a", ids: []string{"a", "b"}, ranked: []string{"a"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "manual-selection", pin: "a", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a"}, want: []string{"b", "a"}}, + {name: "all-cooling", ids: []string{"a", "b"}, ranked: []string{"a", "b"}, cooled: []string{"a", "b"}, want: []string{"a", "b"}}, + {name: "sole-cooling", ids: []string{"a"}, ranked: []string{"a"}, cooled: []string{"a"}, want: []string{"a"}}, + } { + t.Run(tc.name, func(t *testing.T) { + p := prProxy(t, tc.ids...) + p.SetPriority(tc.ranked) + p.SetSelected(tc.pin) + for _, id := range tc.cooled { + p.noteStalled(id, time.Now()) + } + got, affinity := p.reserveCandidateWithAffinity(reservationCandidates(tc.ids...), tc.warm) + var ids []string + for _, cand := range got { + ids = append(ids, cand.id) + } + assertOrder(t, ids, tc.want) + if affinity { + t.Fatal("affinity promoted a cooling candidate over a ready peer") + } + }) + } +} + +func TestPrecommitDropRetriesWithoutCooldown(t *testing.T) { + dropped := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // A 200 followed by EOF before the declared body starts. + w.Header().Set("Content-Length", "999") + w.WriteHeader(http.StatusOK) + })) + defer dropped.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "dropped", dropped.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + wire := &stallWire{} + p := NewProxy(NewCodec(wire), disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"dropped", "good"}) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`))) + if rec.Code != 200 || rec.Body.String() != "healthy" { + t.Fatalf("precommit drop did not fail over: %d %q", rec.Code, rec.Body.String()) + } + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("fast body failure was incorrectly classified as a stall") + } + assertStallWorkloads(t, wire, 1) +} + +func TestCommittedStreamNeverRetries(t *testing.T) { + var goodHits atomic.Int32 + broken := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", "999") + _, _ = io.WriteString(w, "partial") + })) + defer broken.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + _, _ = io.WriteString(w, "healthy") + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "broken", broken.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"broken", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer server.Close() + resp, err := server.Client().Post(server.URL+"/v1/chat/completions", "application/json", + strings.NewReader(`{"model":"model"}`)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err == nil || string(data) != "partial" || goodHits.Load() != 0 { + t.Fatalf("committed stream was retried or not truncated: body=%q error=%v goodHits=%d", data, err, goodHits.Load()) + } +} + +func TestCooldownSuccessAndFastStatus(t *testing.T) { + for _, code := range []int{http.StatusOK, http.StatusInternalServerError} { + t.Run(http.StatusText(code), func(t *testing.T) { + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(code) // empty 200 is a valid complete response + })) + defer engine.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.noteStalled("only", time.Now()) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`))) + if rec.Code != code { + t.Fatalf("sole cooling candidate was not tried: status %d, want %d", rec.Code, code) + } + if cooling := p.coolingSet(time.Now())["only"]; cooling != (code != http.StatusOK) { + t.Fatalf("cooling=%v for status %d: only success should clear", cooling, code) + } + }) + } +} + +func TestCallerCancellationDoesNotCoolCandidate(t *testing.T) { + for _, headers := range []bool{false, true} { + t.Run(fmt.Sprint(headers), func(t *testing.T) { + release := make(chan struct{}) + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if headers { + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + } + <-release + })) + defer engine.Close() + defer close(release) + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("a caller deadline demoted the candidate") + } + }) + } +} + +func TestFinalBodyStallIsBounded(t *testing.T) { + old := proxyResponseTimeout + proxyResponseTimeout = 100 * time.Millisecond + defer func() { proxyResponseTimeout = old }() + release := make(chan struct{}) + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + <-release + })) + defer engine.Close() + defer close(release) + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "only", engine.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if rec.Code != http.StatusBadGateway || ctx.Err() != nil { + t.Fatalf("final body stall was not bounded before client deadline: status=%d context=%v", rec.Code, ctx.Err()) + } + if !p.coolingSet(time.Now())["only"] { + t.Fatal("final body stall was not remembered") + } +} + +func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { + t.Helper() + wire.mu.Lock() + data := append([]byte(nil), wire.b...) + wire.mu.Unlock() + started := make(map[string]bool) + completed := make(map[string]bool) + dec := json.NewDecoder(strings.NewReader(string(data))) + for dec.More() { + var frame struct { + Method string `json:"method"` + Params struct { + WorkloadInfo Workload `json:"workloadInfo"` + } `json:"params"` + } + if err := dec.Decode(&frame); err != nil { + t.Fatal(err) + } + wl := frame.Params.WorkloadInfo + switch frame.Method { + case workloadStartedMethod: + started[wl.ID] = true + case workloadCompletedMethod: + if completed[wl.ID] || wl.ScheduledOn != "good" || !started[wl.ID] { + t.Fatalf("invalid or duplicate completion: %+v", wl) + } + completed[wl.ID] = true + case workloadErroredMethod: + t.Fatalf("failover emitted a failed workload: %+v", wl) + } + } + if len(started) != count || len(completed) != count { + t.Fatalf("workloads started=%v completed=%v, want %d identities", started, completed, count) + } +} diff --git a/services/omlx-proxy/proxy.go b/services/omlx-proxy/proxy.go index 306dd4954..15d83064d 100644 --- a/services/omlx-proxy/proxy.go +++ b/services/omlx-proxy/proxy.go @@ -430,6 +430,10 @@ type Proxy struct { priorityGPUPressure map[string]int priorityReservations map[string]int + // coolMu guards process-local deadline measurements, never persisted. + coolMu sync.Mutex + coolUntil map[string]time.Time + // 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 @@ -1435,7 +1439,7 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { Director: p.forwardDirector(cand, reqID), // A remote cluster peer is dialed over mTLS (per-peer pinned config); // self/manual candidates use the plain transport. See candidateTransport. - Transport: p.candidateTransport(cand), + Transport: stallTransport{base: p.candidateTransport(cand), proxy: p, id: cand.id, timeout: proxyResponseTimeout}, BufferPool: p.bufPool, // ModifyResponse fires when the upstream's status line + headers // have arrived but before the body streams. That's both the retry @@ -1447,6 +1451,21 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { retry = true return retrySignal{} } + // Only successful inference waits for content. All commit effects + // below still happen once, for the node that actually serves. + if isInf && resp.StatusCode >= 200 && resp.StatusCode < 300 { + body, err := awaitFirstBody(r.Context(), resp.Body, proxyResponseTimeout) + if err != nil { + if r.Context().Err() == nil && stderrors.Is(err, errFirstBodyTimeout) { + p.noteStalled(cand.id, time.Now()) + } + return err + } + resp.Body = body + } + if resp.StatusCode >= 200 && resp.StatusCode < 300 { + p.noteSuccess(cand.id) + } // Prefer an engine-declared preflight policy so an exact origin plus // Allow-Credentials can pass a credentialed browser fetch. Engines // that publish no policy retain the proxy's permissive 204 fallback. @@ -1905,6 +1924,8 @@ func (p *Proxy) resolveCandidates(model string) []candidate { } } + out = demoteCooling(out, p.coolingSet(time.Now())) + slog.Debug("resolveCandidates resolved", "selected", id, "priority", len(priority), "candidates", len(out), "eligible", len(nodes), "known", known) @@ -1929,6 +1950,15 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID if len(candidates) == 0 { return candidates, false } + cooling := p.coolingSet(time.Now()) + candidates = demoteCooling(candidates, cooling) + hasReady := false + for _, cand := range candidates { + if !cooling[cand.id] { + hasReady = true + break + } + } if selectedID := p.SelectedID(); selectedID != "" { for _, cand := range candidates { if cand.id == selectedID { @@ -1951,6 +1981,7 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID // Check affinity under the same lock as optimistic reservations. An eligible // manual pin was handled above. Missing or overloaded hints leave ordering alone. if index, ok := candidateIndex[warmNodeID]; ok && warmNodeID != "" && + (!hasReady || !cooling[warmNodeID]) && p.priorityPending[warmNodeID]+p.priorityReservations[warmNodeID] <= 1 && p.priorityGPUPressure[warmNodeID] < schedulerwire.MaxGPUPressure { chosen := candidates[index] @@ -1964,6 +1995,9 @@ func (p *Proxy) reserveCandidateWithAffinity(candidates []candidate, warmNodeID bestOrder := len(p.priority) var bestLoad uint64 for order, id := range p.priority { + if hasReady && cooling[id] { + continue + } index, ok := candidateIndex[id] if !ok { continue diff --git a/services/omlx-proxy/stall.go b/services/omlx-proxy/stall.go new file mode 100644 index 000000000..449a1311d --- /dev/null +++ b/services/omlx-proxy/stall.go @@ -0,0 +1,140 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptrace" + "sync" + "time" +) + +var errFirstBodyTimeout = errors.New("upstream timed out before first response body byte") + +// A cooldown is evidence of a full silent deadline, not a liveness verdict. +// State is local to this proxy; success or time alone restores normal ordering. +func (p *Proxy) noteStalled(id string, now time.Time) { + p.coolMu.Lock() + defer p.coolMu.Unlock() + if p.coolUntil == nil { + p.coolUntil = make(map[string]time.Time) + } + p.coolUntil[id] = now.Add(proxyResponseTimeout) +} + +func (p *Proxy) noteSuccess(id string) { + p.coolMu.Lock() + defer p.coolMu.Unlock() + delete(p.coolUntil, id) +} + +func (p *Proxy) coolingSet(now time.Time) map[string]bool { + p.coolMu.Lock() + defer p.coolMu.Unlock() + if len(p.coolUntil) == 0 { + return nil + } + cooling := make(map[string]bool, len(p.coolUntil)) + for id, until := range p.coolUntil { + if now.Before(until) { + cooling[id] = true + } else { + delete(p.coolUntil, id) + } + } + return cooling +} + +// Stable partition: keep every candidate, including a sole cooling candidate. +func demoteCooling(candidates []candidate, cooling map[string]bool) []candidate { + if len(cooling) == 0 || len(candidates) < 2 { + return candidates + } + out := make([]candidate, 0, len(candidates)) + for _, cand := range candidates { + if !cooling[cand.id] { + out = append(out, cand) + } + } + for _, cand := range candidates { + if cooling[cand.id] { + out = append(out, cand) + } + } + return out +} + +// Observe the response-header wait after the request was written, not dial or +// TLS time. A shorter caller deadline, cancellation, or fast transport failure +// is not evidence that this candidate spent a full response deadline silent. +type stallTransport struct { + base http.RoundTripper + proxy *Proxy + id string + timeout time.Duration +} + +func (t stallTransport) RoundTrip(req *http.Request) (*http.Response, error) { + var mu sync.Mutex + var written time.Time + trace := &httptrace.ClientTrace{WroteRequest: func(info httptrace.WroteRequestInfo) { + if info.Err == nil { + mu.Lock() + written = time.Now() + mu.Unlock() + } + }} + resp, err := t.base.RoundTrip(req.WithContext(httptrace.WithClientTrace(req.Context(), trace))) + mu.Lock() + waitStart := written + mu.Unlock() + var timeout net.Error + if err != nil && req.Context().Err() == nil && errors.As(err, &timeout) && + timeout.Timeout() && !waitStart.IsZero() && time.Since(waitStart) >= t.timeout { + t.proxy.noteStalled(t.id, time.Now()) + } + return resp, err +} + +type bodyWithPeek struct { + io.Reader + io.Closer +} + +// Hold inference headers until content (or a valid empty EOF) arrives. On an +// error the caller returns from ModifyResponse, so ReverseProxy closes the +// original body and unblocks the pending read. The buffered channel lets that +// reader exit even if the deadline or caller cancellation won. +func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duration) (io.ReadCloser, error) { + type result struct { + data []byte + err error + } + ready := make(chan result, 1) + go func() { + // Match ReverseProxy's copy buffer. Reading only one byte would split + // the first engine frame and publish progress before it can be parsed. + first := make([]byte, 32*1024) + n, err := body.Read(first) + ready <- result{data: first[:n], err: err} + }() + timer := time.NewTimer(budget) + defer timer.Stop() + select { + case got := <-ready: + if len(got.data) == 0 && got.err != nil && !errors.Is(got.err, io.EOF) { + return nil, got.err + } + return bodyWithPeek{Reader: io.MultiReader(bytes.NewReader(got.data), body), Closer: body}, nil + case <-ctx.Done(): + return nil, ctx.Err() + case <-timer.C: + return nil, errFirstBodyTimeout + } +} From fef5b772737406398ffd6485c5f0dc2101b0b847 Mon Sep 17 00:00:00 2001 From: Jared Ping <13906915+JaredP94@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:09:42 +0200 Subject: [PATCH 2/2] fix(proxies): make the first-body deadline opt-in A fixed post-header deadline can abandon healthy long-context prefill and duplicate work on another node. LM Studio can send headers immediately, and large MLX models on Macs can remain silent for several minutes before their first response body bytes. Default the independent first-body timeout to disabled in all three proxies. Enable it only through --first-body-timeout or NVPAIR_PROXY_FIRST_BODY_TIMEOUT, with the flag taking precedence and zero explicitly disabling an environment opt-in. Keep the 120-second header deadline, header-stall cooldown, and pre-commit error retry unchanged. Add a slow-after-headers regression that fails under the previous default, configuration coverage, and flag/env opt-in stall cases. Require healthy failover peers to receive the full original body and require running workload notifications to name the serving node before completion. Update the READMEs, architecture, and troubleshooting with the distinct settings and the duplicate-work trade-off of enabling a first-body deadline. Signed-off-by: Jared Ping <13906915+JaredP94@users.noreply.github.com> --- README.md | 7 ++ docs/architecture.mdx | 30 ++++-- docs/troubleshooting.mdx | 45 ++++++++- services/lmstudio-proxy/README.md | 7 +- services/lmstudio-proxy/cooldown_test.go | 99 +++++++++++++++---- .../lmstudio-proxy/first_body_timeout_test.go | 43 ++++++++ services/lmstudio-proxy/main.go | 2 + services/lmstudio-proxy/proxy.go | 3 +- services/lmstudio-proxy/stall.go | 35 ++++++- services/ollama-proxy/README.md | 7 +- services/ollama-proxy/cooldown_test.go | 99 +++++++++++++++---- .../ollama-proxy/first_body_timeout_test.go | 43 ++++++++ services/ollama-proxy/main.go | 2 + services/ollama-proxy/proxy.go | 3 +- services/ollama-proxy/stall.go | 35 ++++++- services/omlx-proxy/README.md | 7 +- services/omlx-proxy/cooldown_test.go | 99 +++++++++++++++---- .../omlx-proxy/first_body_timeout_test.go | 43 ++++++++ services/omlx-proxy/main.go | 2 + services/omlx-proxy/proxy.go | 3 +- services/omlx-proxy/stall.go | 35 ++++++- 21 files changed, 561 insertions(+), 88 deletions(-) create mode 100644 services/lmstudio-proxy/first_body_timeout_test.go create mode 100644 services/ollama-proxy/first_body_timeout_test.go create mode 100644 services/omlx-proxy/first_body_timeout_test.go diff --git a/README.md b/README.md index 31572e7ea..9c3c9eaed 100644 --- a/README.md +++ b/README.md @@ -240,6 +240,13 @@ The reply is ordinary OpenAI-shaped JSON, abbreviated here: If you changed a port, or you are using LM Studio or oMLX rather than Ollama, copy the URL from **Endpoints → API endpoints** instead of assuming the one above. +The proxies keep the `120s` response-header timeout, but impose **no first-body +deadline by default** after headers arrive. Long-context prefill can legitimately +be silent for several minutes. `--first-body-timeout` / +`NVPAIR_PROXY_FIRST_BODY_TIMEOUT` is a separate opt-in; see +[Troubleshooting](docs/troubleshooting.mdx#requests-are-silent-after-headers) +before enabling it. + That is a single machine working. To route across machines, pair a second one from **Settings → Cluster** and repeat the engine and model steps there. The inviting machine shows a six-digit PIN, and you enter that PIN on the machine you diff --git a/docs/architecture.mdx b/docs/architecture.mdx index 5ef3e24a4..b8574658b 100644 --- a/docs/architecture.mdx +++ b/docs/architecture.mdx @@ -477,8 +477,8 @@ requests it has just dispatched itself to its own estimate and reserves its choice before forwarding, so a burst spreads without waiting for those reports. After those preferences, each proxy moves candidates that just spent a full -silent response deadline behind their eligible peers, preserving relative order -within both groups. This measured, process-local cooldown overrides pins, +silent response-header deadline (or an explicitly enabled first-body deadline) +behind their eligible peers, preserving relative order within both groups. This measured, process-local cooldown overrides pins, reservations, and prefix hints; it is not a liveness verdict. It never removes a candidate, so a lone node is still tried and a cooling node can serve after the others fail. It expires after one configured response-header duration, or clears @@ -532,15 +532,27 @@ unknown or known-missing node. Genuine client errors such as `400` or `422` are not retried, because they would fail identically everywhere. Successful inference responses do not commit on headers alone. The proxy waits -for initial body bytes (or a valid empty EOF), bounded by a separate wait equal -to the configured response-header timeout. Silence or a body read failure before -commit advances to the next existing candidate while retaining the same request +for initial body bytes (or a valid empty EOF), and pre-commit body read failures +advance to the next existing candidate by default while retaining one request and workload identity. It does not re-resolve candidates or add retry rounds. Once the response commits, errors only terminate that stream; it is never -replayed after bytes have reached the client. A full silent header or first-body -deadline also records the cooldown described above. The `120s` default and -`--response-header-timeout` / `NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT` configuration -remain unchanged; operators must allow for cold loads and long queues. +replayed after bytes have reached the client. + +**There is no first-body deadline by default.** LM Studio can send headers before +long-context prefill finishes, and oMLX with large MLX models on Macs can also be +silent for several minutes. The unchanged `120s` response-header timeout +(`--response-header-timeout` / `NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`) does not +limit this post-header wait. + +Operators may separately opt in with `--first-body-timeout` or +`NVPAIR_PROXY_FIRST_BODY_TIMEOUT`, a positive Go duration such as `10m`. The flag +wins over the environment; `0` explicitly disables it, including an environment +opt-in. Invalid or negative values keep it disabled. Only an enabled first-body +deadline can turn post-header silence into timeout-based retry and cooldown. +The cooldown still lasts one configured response-header duration. This option +can abandon legitimate prefill and duplicate engine work if set too short, so +leave it off unless that trade-off is intended; see +[Troubleshooting](troubleshooting.mdx#requests-are-silent-after-headers). #### Planning Model Placement diff --git a/docs/troubleshooting.mdx b/docs/troubleshooting.mdx index c073784ab..dc3dc0613 100644 --- a/docs/troubleshooting.mdx +++ b/docs/troubleshooting.mdx @@ -154,16 +154,53 @@ runs on one selected node. A request that fails with `502` and `upstream error: net/http: timeout awaiting response headers` after about 120 seconds waited longer than PAIR allows for an -engine to start answering. An engine sends nothing until generation starts, so a -request queued behind other work, or waiting for a large model to load, can hit -this while the engine is healthy. +engine to send response headers. An engine such as Ollama can withhold headers +until generation starts, so a request queued behind other work, or waiting for +a large model to load, can hit this while the engine is healthy. This is the +response-header deadline only, not a limit on generation or prefill after +headers arrive. Raise the limit with the `NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT` environment variable, a Go duration such as `10m`, and restart PAIR. On macOS, set it for apps opened from the Finder with `launchctl setenv NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT 10m`. Every node that serves requests applies its own limit, so raise it on each. A longer limit also -means a node that has stopped responding is given up on later. +means a node that has stopped responding is given up on later. Standalone proxies +also accept `--response-header-timeout 10m`. Header-stall cooldown remains on by +default: a node that spent that full silent deadline is tried behind its peers +for one header-timeout duration, never excluded. + +## Requests Are Silent After Headers + +**By default there is no first-body deadline.** LM Studio can send headers +immediately and then spend several minutes on long-context prefill. Large MLX +models in oMLX on Macs can also legitimately remain silent before their first +body bytes. The `120s` response-header timeout does not cut off this wait. +Pre-commit connection/body read failures still retry on another candidate by +default, but silence after headers alone does not. + +An optional, separate first-body limit is available in all three proxies: +`--first-body-timeout 10m` for standalone use, or +`NVPAIR_PROXY_FIRST_BODY_TIMEOUT=10m` in PAIR's inherited environment, followed +by a restart. On macOS, an app opened from the Finder can inherit it via +`launchctl setenv NVPAIR_PROXY_FIRST_BODY_TIMEOUT 10m`. The flag overrides the +environment. `0` disables it explicitly; invalid or negative values also leave +it disabled. It does not change `--response-header-timeout` or its environment +variable. + +Only enable this if you accept the trade-off: a legitimate prefill longer than +the chosen first-body budget will be abandoned and retried, potentially +duplicating work on another engine. An enabled first-body deadline also records +cooldown for one configured response-header duration. Once any response body is +committed, the stream is never retried and this first-body timer no longer +applies. + +If the error says `upstream timed out before first response body byte`, this +optional limit was enabled. Increase it or disable it rather than raising the +unrelated response-header timeout. For a Finder-launched app, remove the opt-in +with `launchctl unsetenv NVPAIR_PROXY_FIRST_BODY_TIMEOUT` and restart PAIR; +standalone proxies can use `--first-body-timeout 0` even if the environment sets +a positive value. Each node applies its own settings. ## Where the Logs Are diff --git a/services/lmstudio-proxy/README.md b/services/lmstudio-proxy/README.md index f72f96e0b..160ee4227 100644 --- a/services/lmstudio-proxy/README.md +++ b/services/lmstudio-proxy/README.md @@ -35,7 +35,8 @@ lmstudio-proxy [flags] | `--port` | `1234` | HTTP listen port for request forwarding | | `--ignore-persisted-port` | `false` | Use `--port` even when a prior runtime port was saved | | `--ipc` | *(empty — use stdio)* | Path to a Unix domain socket or Windows named pipe for IPC | -| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline and, for successful inference, a separate wait for initial body bytes (Go duration, e.g. `10m`). Raise it for long queues or cold model loads. An invalid or non-positive value keeps the default | +| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline only (Go duration, e.g. `10m`). Does not limit prefill after headers. An invalid or non-positive value keeps the unchanged `120s` default | +| `--first-body-timeout` | `$NVPAIR_PROXY_FIRST_BODY_TIMEOUT`, else `0` (disabled) | Opt-in, independent wait for initial successful inference body bytes after headers (Go duration, e.g. `10m`). No post-header deadline by default. The flag overrides the environment; `0` explicitly disables it even if the environment opts in. Invalid or negative values leave it disabled | | `--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. | | `--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 | @@ -53,8 +54,8 @@ Node selection: - **Auto**: When no eligible node is explicitly selected, the proxy follows `node/set-priority` (see below), then discovered nodes in stable ID order. - **Priority (scheduler-driven)**: The Job Scheduler ranks the cluster least-loaded-first by pending workload plus smoothed GPU pressure and, via `nvpair-ui-broker`, pushes the ordered node list with those per-node counts to this proxy with `node/set-priority`. Auto routing sends the request to the listed node carrying the least estimated load. See [`nvpair-job-scheduler`](../nvpair-job-scheduler/README.md). - **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing; a silent body times out using the configured response-header duration and fails over within the same candidate list. Once committed, a stream is never retried. Failover keeps one request/workload identity. -- **Deadline cooldown**: A full response-header wait with no answer, or an inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured deadline. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. The default deadline remains `120s`. +- **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing. Pre-commit body read failures retry by default within the same candidate list. There is no first-body deadline unless `--first-body-timeout` / `NVPAIR_PROXY_FIRST_BODY_TIMEOUT` explicitly opts in; only then can post-header silence trigger timeout-based retry. Once committed, a stream is never retried. Failover keeps one request/workload identity. +- **Deadline cooldown**: A full response-header wait with no answer, or an explicitly enabled inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured response-header duration. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. Header-stall cooldown remains on by default with the unchanged `120s` header deadline; first-body timing is off by default to avoid interrupting legitimate long-context prefill. See [troubleshooting](../../docs/troubleshooting.mdx#requests-are-silent-after-headers) before enabling it. ### IPC Transport diff --git a/services/lmstudio-proxy/cooldown_test.go b/services/lmstudio-proxy/cooldown_test.go index 8663932eb..572455631 100644 --- a/services/lmstudio-proxy/cooldown_test.go +++ b/services/lmstudio-proxy/cooldown_test.go @@ -36,21 +36,31 @@ func (w *stallWire) Write(b []byte) (int, error) { // upstream again. Reset scheduler reservations so they cannot hide that bug. func TestStalledCandidateSecondRequest(t *testing.T) { old := proxyResponseTimeout - proxyResponseTimeout = time.Second + // Leave scheduling slack between the measured stall and the next request. + // First-body opt-ins below remain shorter and independent of this window. + proxyResponseTimeout = 3 * time.Second defer func() { proxyResponseTimeout = old }() - for _, headers := range []bool{false, true} { - name := "no-headers" - if headers { - name = "headers-without-body" - } - t.Run(name, func(t *testing.T) { + const originalBody = `{"model":"model","messages":[{"role":"user","content":"replay regression"}],"stream":false,"max_tokens":7}` + for _, tc := range []struct { + name, flag, env string + headers bool + }{ + {name: "no-headers-default"}, + {name: "headers-without-body-flag", headers: true, flag: "1s"}, + {name: "headers-without-body-env", headers: true, env: "1s"}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, tc.env) + oldFirstBody := proxyFirstBodyTimeout + proxyFirstBodyTimeout = resolveFirstBodyTimeout(tc.flag) + t.Cleanup(func() { proxyFirstBodyTimeout = oldFirstBody }) var stalledHits, goodHits atomic.Int32 release := make(chan struct{}) stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { stalledHits.Add(1) _, _ = io.Copy(io.Discard, r.Body) - if headers { + if tc.headers { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) w.(http.Flusher).Flush() @@ -61,7 +71,10 @@ func TestStalledCandidateSecondRequest(t *testing.T) { defer close(release) good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { goodHits.Add(1) - _, _ = io.Copy(io.Discard, r.Body) + body, err := io.ReadAll(r.Body) + if err != nil || string(body) != originalBody { + t.Error("healthy peer did not receive the complete original request body") + } _, _ = io.WriteString(w, "healthy") })) defer good.Close() @@ -78,10 +91,10 @@ func TestStalledCandidateSecondRequest(t *testing.T) { Nodes: []string{"stalled", "good"}, Ranks: []schedulerwire.NodeRank{{ID: "stalled"}, {ID: "good"}}, }) - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", - strings.NewReader(`{"model":"model"}`)).WithContext(ctx) + strings.NewReader(originalBody)).WithContext(ctx) rec := httptest.NewRecorder() start := time.Now() p.handleHTTP(rec, req) @@ -103,12 +116,56 @@ func TestStalledCandidateSecondRequest(t *testing.T) { if goodHits.Load() != 2 { t.Fatalf("healthy hits = %d, want 2", goodHits.Load()) } - t.Logf("first=%s second=%s (configured deadline %s)", first, second, proxyResponseTimeout) + t.Logf("first=%s second=%s (header deadline %s, first-body deadline %s)", first, second, proxyResponseTimeout, proxyFirstBodyTimeout) assertStallWorkloads(t, wire, 2) }) } } +// A response-header deadline must not become a post-header generation limit. +// Long-context prefill can remain silent well beyond the header timeout. +func TestDefaultFirstBodyWaitIsUnbounded(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, "") + oldFirstBody := proxyFirstBodyTimeout + proxyFirstBodyTimeout = resolveFirstBodyTimeout("") + defer func() { proxyFirstBodyTimeout = oldFirstBody }() + old := proxyResponseTimeout + proxyResponseTimeout = 100 * time.Millisecond + defer func() { proxyResponseTimeout = old }() + var fallbackHits atomic.Int32 + slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + time.Sleep(3 * proxyResponseTimeout) + _, _ = io.WriteString(w, "prefill-complete") + })) + defer slow.Close() + fallback := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fallbackHits.Add(1) + _, _ = io.WriteString(w, "duplicate") + })) + defer fallback.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "slow", slow.URL, "model")) + disc.AddManual(nodeForModel(t, "fallback", fallback.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"slow", "fallback"}) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if rec.Code != http.StatusOK || rec.Body.String() != "prefill-complete" || fallbackHits.Load() != 0 { + t.Fatalf("default cut off healthy post-header prefill: status=%d body=%q fallback hits=%d", + rec.Code, rec.Body.String(), fallbackHits.Load()) + } + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("healthy prefill created a default cooldown") + } +} + func TestCooldownOrderingRecovery(t *testing.T) { p := prProxy(t, "a", "b", "c", "d") p.SetPriority([]string{"a", "b", "c", "d"}) @@ -276,9 +333,9 @@ func TestCallerCancellationDoesNotCoolCandidate(t *testing.T) { } func TestFinalBodyStallIsBounded(t *testing.T) { - old := proxyResponseTimeout - proxyResponseTimeout = 100 * time.Millisecond - defer func() { proxyResponseTimeout = old }() + old := proxyFirstBodyTimeout + proxyFirstBodyTimeout = 100 * time.Millisecond + defer func() { proxyFirstBodyTimeout = old }() release := make(chan struct{}) engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) @@ -309,7 +366,7 @@ func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { wire.mu.Lock() data := append([]byte(nil), wire.b...) wire.mu.Unlock() - started := make(map[string]bool) + started := make(map[string]string) completed := make(map[string]bool) dec := json.NewDecoder(strings.NewReader(string(data))) for dec.More() { @@ -325,9 +382,15 @@ func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { wl := frame.Params.WorkloadInfo switch frame.Method { case workloadStartedMethod: - started[wl.ID] = true + if wl.State != "running" || completed[wl.ID] { + t.Fatalf("invalid running workload event: %+v", wl) + } + started[wl.ID] = wl.ScheduledOn case workloadCompletedMethod: - if completed[wl.ID] || wl.ScheduledOn != "good" || !started[wl.ID] { + if started[wl.ID] != "good" { + t.Fatalf("no running workload event on serving node before completion: id=%s runningOn=%q", wl.ID, started[wl.ID]) + } + if completed[wl.ID] || wl.ScheduledOn != "good" { t.Fatalf("invalid or duplicate completion: %+v", wl) } completed[wl.ID] = true diff --git a/services/lmstudio-proxy/first_body_timeout_test.go b/services/lmstudio-proxy/first_body_timeout_test.go new file mode 100644 index 000000000..9ae5d8f77 --- /dev/null +++ b/services/lmstudio-proxy/first_body_timeout_test.go @@ -0,0 +1,43 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "testing" + "time" +) + +func TestResolveFirstBodyTimeout(t *testing.T) { + for _, tc := range []struct { + name, flag, env string + want time.Duration + }{ + {"default disabled", "", "", 0}, + {"environment opt-in", "", "10m", 10 * time.Minute}, + {"flag opt-in", "7m", "", 7 * time.Minute}, + {"flag beats environment", "5m", "10m", 5 * time.Minute}, + {"explicit zero disables environment", "0s", "10m", 0}, + {"environment zero disables", "", "0", 0}, + {"invalid flag disables", "soon", "10m", 0}, + {"invalid environment disables", "", "soon", 0}, + {"negative flag disables", "-1s", "10m", 0}, + {"negative environment disables", "", "-1m", 0}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, tc.env) + if got := resolveFirstBodyTimeout(tc.flag); got != tc.want { + t.Fatalf("first-body timeout = %s, want %s", got, tc.want) + } + }) + } +} + +func TestFirstBodyTimeoutDoesNotChangeHeaderDeadline(t *testing.T) { + old := proxyFirstBodyTimeout + t.Cleanup(func() { proxyFirstBodyTimeout = old }) + proxyFirstBodyTimeout = 7 * time.Minute + if got := newProxyTransport(nil).ResponseHeaderTimeout; got != proxyResponseTimeout { + t.Fatalf("header timeout = %s, want unchanged %s", got, proxyResponseTimeout) + } +} diff --git a/services/lmstudio-proxy/main.go b/services/lmstudio-proxy/main.go index 06fee8873..a516ef598 100644 --- a/services/lmstudio-proxy/main.go +++ b/services/lmstudio-proxy/main.go @@ -26,12 +26,14 @@ func main() { port := flag.Int("port", defaultProxyPort, "HTTP listen port") ignorePersistedPort := flag.Bool("ignore-persisted-port", false, "use --port even when a persisted port exists") responseHeaderTimeout := flag.String("response-header-timeout", "", "upstream response-header timeout (Go duration, e.g. 10m); default: $NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT or 120s") + firstBodyTimeout := flag.String("first-body-timeout", "", "opt-in inference first-body timeout (Go duration, e.g. 10m; 0 disables); default: $NVPAIR_PROXY_FIRST_BODY_TIMEOUT or disabled") 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") showVersion := flag.Bool("version", false, "print version and exit") resolveLevel := applog.RegisterFlag(nil, slog.LevelInfo) flag.Parse() proxyResponseTimeout = resolveResponseHeaderTimeout(*responseHeaderTimeout) + proxyFirstBodyTimeout = resolveFirstBodyTimeout(*firstBodyTimeout) if *showVersion { fmt.Println(Version) diff --git a/services/lmstudio-proxy/proxy.go b/services/lmstudio-proxy/proxy.go index 7cc50588e..b41ddab99 100644 --- a/services/lmstudio-proxy/proxy.go +++ b/services/lmstudio-proxy/proxy.go @@ -680,6 +680,7 @@ func (p *Proxy) serveHTTP(ctx context.Context, ln net.Listener) { "dial_timeout", proxyDialTimeout, "keep_alive", proxyKeepAlive, "response_header_timeout", proxyResponseTimeout, + "first_body_timeout", proxyFirstBodyTimeout, "max_idle_conns", proxyMaxIdleConns, "idle_conn_timeout", proxyIdleConnTimeout, ) @@ -1372,7 +1373,7 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { // Only successful inference waits for content. All commit effects // below still happen once, for the node that actually serves. if isInf && resp.StatusCode >= 200 && resp.StatusCode < 300 { - body, err := awaitFirstBody(r.Context(), resp.Body, proxyResponseTimeout) + body, err := awaitFirstBody(r.Context(), resp.Body, proxyFirstBodyTimeout) if err != nil { if r.Context().Err() == nil && stderrors.Is(err, errFirstBodyTimeout) { p.noteStalled(cand.id, time.Now()) diff --git a/services/lmstudio-proxy/stall.go b/services/lmstudio-proxy/stall.go index 449a1311d..c88d20bde 100644 --- a/services/lmstudio-proxy/stall.go +++ b/services/lmstudio-proxy/stall.go @@ -8,13 +8,38 @@ import ( "context" "errors" "io" + "log" "net" "net/http" "net/http/httptrace" + "os" "sync" "time" ) +// The response-header deadline does not bound inference prefill after headers. +// First-body timing is opt-in and independent; zero means wait until content, +// an upstream error, or caller cancellation. +const firstBodyTimeoutEnv = "NVPAIR_PROXY_FIRST_BODY_TIMEOUT" + +var proxyFirstBodyTimeout time.Duration + +func resolveFirstBodyTimeout(flagVal string) time.Duration { + raw := flagVal + if raw == "" { + raw = os.Getenv(firstBodyTimeoutEnv) + } + if raw == "" { + return 0 + } + d, err := time.ParseDuration(raw) + if err != nil || d < 0 { + log.Printf("invalid first-body timeout %q, using default (disabled)", raw) + return 0 + } + return d +} + var errFirstBodyTimeout = errors.New("upstream timed out before first response body byte") // A cooldown is evidence of a full silent deadline, not a liveness verdict. @@ -124,8 +149,12 @@ func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duratio n, err := body.Read(first) ready <- result{data: first[:n], err: err} }() - timer := time.NewTimer(budget) - defer timer.Stop() + var deadline <-chan time.Time + if budget > 0 { + timer := time.NewTimer(budget) + defer timer.Stop() + deadline = timer.C + } select { case got := <-ready: if len(got.data) == 0 && got.err != nil && !errors.Is(got.err, io.EOF) { @@ -134,7 +163,7 @@ func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duratio return bodyWithPeek{Reader: io.MultiReader(bytes.NewReader(got.data), body), Closer: body}, nil case <-ctx.Done(): return nil, ctx.Err() - case <-timer.C: + case <-deadline: return nil, errFirstBodyTimeout } } diff --git a/services/ollama-proxy/README.md b/services/ollama-proxy/README.md index 8441b0e5d..aa1c1e2f3 100644 --- a/services/ollama-proxy/README.md +++ b/services/ollama-proxy/README.md @@ -34,7 +34,8 @@ ollama-proxy [flags] | `--alias-address` | *(empty)* | Optional secondary `host:port` for the same routing handler. Repeat the flag to reserve both loopback families for `localhost`. Only literal loopback addresses are accepted; the broker uses this for a safe inherited local `OLLAMA_HOST`, and the aliases are not advertised to peers. | | `--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 | -| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline and, for successful inference, a separate wait for initial body bytes (Go duration, e.g. `10m`). Raise it for long queues or cold model loads. An invalid or non-positive value keeps the default | +| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline only (Go duration, e.g. `10m`). Does not limit prefill after headers. An invalid or non-positive value keeps the unchanged `120s` default | +| `--first-body-timeout` | `$NVPAIR_PROXY_FIRST_BODY_TIMEOUT`, else `0` (disabled) | Opt-in, independent wait for initial successful inference body bytes after headers (Go duration, e.g. `10m`). No post-header deadline by default. The flag overrides the environment; `0` explicitly disables it even if the environment opts in. Invalid or negative values leave it disabled | | `--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. | | `--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 | @@ -58,8 +59,8 @@ Node selection: - **Auto**: When no eligible node is explicitly selected, the proxy follows `node/set-priority` (see below), then discovered nodes in stable ID order. - **Priority (scheduler-driven)**: The Job Scheduler ranks the cluster least-loaded-first by pending workload plus smoothed GPU pressure and, via `nvpair-ui-broker`, pushes the ordered node list with those per-node counts to this proxy with `node/set-priority`. Auto routing sends the request to the listed node carrying the least estimated load. See [`nvpair-job-scheduler`](../nvpair-job-scheduler/README.md). - **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing; a silent body times out using the configured response-header duration and fails over within the same candidate list. Once committed, a stream is never retried. Failover keeps one request/workload identity. -- **Deadline cooldown**: A full response-header wait with no answer, or an inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured deadline. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. The default deadline remains `120s`. +- **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing. Pre-commit body read failures retry by default within the same candidate list. There is no first-body deadline unless `--first-body-timeout` / `NVPAIR_PROXY_FIRST_BODY_TIMEOUT` explicitly opts in; only then can post-header silence trigger timeout-based retry. Once committed, a stream is never retried. Failover keeps one request/workload identity. +- **Deadline cooldown**: A full response-header wait with no answer, or an explicitly enabled inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured response-header duration. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. Header-stall cooldown remains on by default with the unchanged `120s` header deadline; first-body timing is off by default to avoid interrupting legitimate long-context prefill. See [troubleshooting](../../docs/troubleshooting.mdx#requests-are-silent-after-headers) before enabling it. ### IPC Transport diff --git a/services/ollama-proxy/cooldown_test.go b/services/ollama-proxy/cooldown_test.go index 8663932eb..572455631 100644 --- a/services/ollama-proxy/cooldown_test.go +++ b/services/ollama-proxy/cooldown_test.go @@ -36,21 +36,31 @@ func (w *stallWire) Write(b []byte) (int, error) { // upstream again. Reset scheduler reservations so they cannot hide that bug. func TestStalledCandidateSecondRequest(t *testing.T) { old := proxyResponseTimeout - proxyResponseTimeout = time.Second + // Leave scheduling slack between the measured stall and the next request. + // First-body opt-ins below remain shorter and independent of this window. + proxyResponseTimeout = 3 * time.Second defer func() { proxyResponseTimeout = old }() - for _, headers := range []bool{false, true} { - name := "no-headers" - if headers { - name = "headers-without-body" - } - t.Run(name, func(t *testing.T) { + const originalBody = `{"model":"model","messages":[{"role":"user","content":"replay regression"}],"stream":false,"max_tokens":7}` + for _, tc := range []struct { + name, flag, env string + headers bool + }{ + {name: "no-headers-default"}, + {name: "headers-without-body-flag", headers: true, flag: "1s"}, + {name: "headers-without-body-env", headers: true, env: "1s"}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, tc.env) + oldFirstBody := proxyFirstBodyTimeout + proxyFirstBodyTimeout = resolveFirstBodyTimeout(tc.flag) + t.Cleanup(func() { proxyFirstBodyTimeout = oldFirstBody }) var stalledHits, goodHits atomic.Int32 release := make(chan struct{}) stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { stalledHits.Add(1) _, _ = io.Copy(io.Discard, r.Body) - if headers { + if tc.headers { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) w.(http.Flusher).Flush() @@ -61,7 +71,10 @@ func TestStalledCandidateSecondRequest(t *testing.T) { defer close(release) good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { goodHits.Add(1) - _, _ = io.Copy(io.Discard, r.Body) + body, err := io.ReadAll(r.Body) + if err != nil || string(body) != originalBody { + t.Error("healthy peer did not receive the complete original request body") + } _, _ = io.WriteString(w, "healthy") })) defer good.Close() @@ -78,10 +91,10 @@ func TestStalledCandidateSecondRequest(t *testing.T) { Nodes: []string{"stalled", "good"}, Ranks: []schedulerwire.NodeRank{{ID: "stalled"}, {ID: "good"}}, }) - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", - strings.NewReader(`{"model":"model"}`)).WithContext(ctx) + strings.NewReader(originalBody)).WithContext(ctx) rec := httptest.NewRecorder() start := time.Now() p.handleHTTP(rec, req) @@ -103,12 +116,56 @@ func TestStalledCandidateSecondRequest(t *testing.T) { if goodHits.Load() != 2 { t.Fatalf("healthy hits = %d, want 2", goodHits.Load()) } - t.Logf("first=%s second=%s (configured deadline %s)", first, second, proxyResponseTimeout) + t.Logf("first=%s second=%s (header deadline %s, first-body deadline %s)", first, second, proxyResponseTimeout, proxyFirstBodyTimeout) assertStallWorkloads(t, wire, 2) }) } } +// A response-header deadline must not become a post-header generation limit. +// Long-context prefill can remain silent well beyond the header timeout. +func TestDefaultFirstBodyWaitIsUnbounded(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, "") + oldFirstBody := proxyFirstBodyTimeout + proxyFirstBodyTimeout = resolveFirstBodyTimeout("") + defer func() { proxyFirstBodyTimeout = oldFirstBody }() + old := proxyResponseTimeout + proxyResponseTimeout = 100 * time.Millisecond + defer func() { proxyResponseTimeout = old }() + var fallbackHits atomic.Int32 + slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + time.Sleep(3 * proxyResponseTimeout) + _, _ = io.WriteString(w, "prefill-complete") + })) + defer slow.Close() + fallback := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fallbackHits.Add(1) + _, _ = io.WriteString(w, "duplicate") + })) + defer fallback.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "slow", slow.URL, "model")) + disc.AddManual(nodeForModel(t, "fallback", fallback.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"slow", "fallback"}) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if rec.Code != http.StatusOK || rec.Body.String() != "prefill-complete" || fallbackHits.Load() != 0 { + t.Fatalf("default cut off healthy post-header prefill: status=%d body=%q fallback hits=%d", + rec.Code, rec.Body.String(), fallbackHits.Load()) + } + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("healthy prefill created a default cooldown") + } +} + func TestCooldownOrderingRecovery(t *testing.T) { p := prProxy(t, "a", "b", "c", "d") p.SetPriority([]string{"a", "b", "c", "d"}) @@ -276,9 +333,9 @@ func TestCallerCancellationDoesNotCoolCandidate(t *testing.T) { } func TestFinalBodyStallIsBounded(t *testing.T) { - old := proxyResponseTimeout - proxyResponseTimeout = 100 * time.Millisecond - defer func() { proxyResponseTimeout = old }() + old := proxyFirstBodyTimeout + proxyFirstBodyTimeout = 100 * time.Millisecond + defer func() { proxyFirstBodyTimeout = old }() release := make(chan struct{}) engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) @@ -309,7 +366,7 @@ func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { wire.mu.Lock() data := append([]byte(nil), wire.b...) wire.mu.Unlock() - started := make(map[string]bool) + started := make(map[string]string) completed := make(map[string]bool) dec := json.NewDecoder(strings.NewReader(string(data))) for dec.More() { @@ -325,9 +382,15 @@ func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { wl := frame.Params.WorkloadInfo switch frame.Method { case workloadStartedMethod: - started[wl.ID] = true + if wl.State != "running" || completed[wl.ID] { + t.Fatalf("invalid running workload event: %+v", wl) + } + started[wl.ID] = wl.ScheduledOn case workloadCompletedMethod: - if completed[wl.ID] || wl.ScheduledOn != "good" || !started[wl.ID] { + if started[wl.ID] != "good" { + t.Fatalf("no running workload event on serving node before completion: id=%s runningOn=%q", wl.ID, started[wl.ID]) + } + if completed[wl.ID] || wl.ScheduledOn != "good" { t.Fatalf("invalid or duplicate completion: %+v", wl) } completed[wl.ID] = true diff --git a/services/ollama-proxy/first_body_timeout_test.go b/services/ollama-proxy/first_body_timeout_test.go new file mode 100644 index 000000000..9ae5d8f77 --- /dev/null +++ b/services/ollama-proxy/first_body_timeout_test.go @@ -0,0 +1,43 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "testing" + "time" +) + +func TestResolveFirstBodyTimeout(t *testing.T) { + for _, tc := range []struct { + name, flag, env string + want time.Duration + }{ + {"default disabled", "", "", 0}, + {"environment opt-in", "", "10m", 10 * time.Minute}, + {"flag opt-in", "7m", "", 7 * time.Minute}, + {"flag beats environment", "5m", "10m", 5 * time.Minute}, + {"explicit zero disables environment", "0s", "10m", 0}, + {"environment zero disables", "", "0", 0}, + {"invalid flag disables", "soon", "10m", 0}, + {"invalid environment disables", "", "soon", 0}, + {"negative flag disables", "-1s", "10m", 0}, + {"negative environment disables", "", "-1m", 0}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, tc.env) + if got := resolveFirstBodyTimeout(tc.flag); got != tc.want { + t.Fatalf("first-body timeout = %s, want %s", got, tc.want) + } + }) + } +} + +func TestFirstBodyTimeoutDoesNotChangeHeaderDeadline(t *testing.T) { + old := proxyFirstBodyTimeout + t.Cleanup(func() { proxyFirstBodyTimeout = old }) + proxyFirstBodyTimeout = 7 * time.Minute + if got := newProxyTransport(nil).ResponseHeaderTimeout; got != proxyResponseTimeout { + t.Fatalf("header timeout = %s, want unchanged %s", got, proxyResponseTimeout) + } +} diff --git a/services/ollama-proxy/main.go b/services/ollama-proxy/main.go index 3415ac7c2..d4c4ad163 100644 --- a/services/ollama-proxy/main.go +++ b/services/ollama-proxy/main.go @@ -37,12 +37,14 @@ func main() { flag.Var(&aliasAddresses, "alias-address", "optional loopback-only HTTP alias address (repeat for another loopback family)") ignorePersistedPort := flag.Bool("ignore-persisted-port", false, "use --port even when a persisted port exists") responseHeaderTimeout := flag.String("response-header-timeout", "", "upstream response-header timeout (Go duration, e.g. 10m); default: $NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT or 120s") + firstBodyTimeout := flag.String("first-body-timeout", "", "opt-in inference first-body timeout (Go duration, e.g. 10m; 0 disables); default: $NVPAIR_PROXY_FIRST_BODY_TIMEOUT or disabled") 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") showVersion := flag.Bool("version", false, "print version and exit") resolveLevel := applog.RegisterFlag(nil, slog.LevelInfo) flag.Parse() proxyResponseTimeout = resolveResponseHeaderTimeout(*responseHeaderTimeout) + proxyFirstBodyTimeout = resolveFirstBodyTimeout(*firstBodyTimeout) if *showVersion { fmt.Println(Version) diff --git a/services/ollama-proxy/proxy.go b/services/ollama-proxy/proxy.go index 196b49c0e..2773a9a9f 100644 --- a/services/ollama-proxy/proxy.go +++ b/services/ollama-proxy/proxy.go @@ -697,6 +697,7 @@ func (p *Proxy) serveHTTP(ctx context.Context, ln net.Listener) { "dial_timeout", proxyDialTimeout, "keep_alive", proxyKeepAlive, "response_header_timeout", proxyResponseTimeout, + "first_body_timeout", proxyFirstBodyTimeout, "max_idle_conns", proxyMaxIdleConns, "idle_conn_timeout", proxyIdleConnTimeout, ) @@ -1558,7 +1559,7 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { // Only successful inference waits for content. All commit effects // below still happen once, for the node that actually serves. if isInf && resp.StatusCode >= 200 && resp.StatusCode < 300 { - body, err := awaitFirstBody(r.Context(), resp.Body, proxyResponseTimeout) + body, err := awaitFirstBody(r.Context(), resp.Body, proxyFirstBodyTimeout) if err != nil { if r.Context().Err() == nil && stderrors.Is(err, errFirstBodyTimeout) { p.noteStalled(cand.id, time.Now()) diff --git a/services/ollama-proxy/stall.go b/services/ollama-proxy/stall.go index 449a1311d..c88d20bde 100644 --- a/services/ollama-proxy/stall.go +++ b/services/ollama-proxy/stall.go @@ -8,13 +8,38 @@ import ( "context" "errors" "io" + "log" "net" "net/http" "net/http/httptrace" + "os" "sync" "time" ) +// The response-header deadline does not bound inference prefill after headers. +// First-body timing is opt-in and independent; zero means wait until content, +// an upstream error, or caller cancellation. +const firstBodyTimeoutEnv = "NVPAIR_PROXY_FIRST_BODY_TIMEOUT" + +var proxyFirstBodyTimeout time.Duration + +func resolveFirstBodyTimeout(flagVal string) time.Duration { + raw := flagVal + if raw == "" { + raw = os.Getenv(firstBodyTimeoutEnv) + } + if raw == "" { + return 0 + } + d, err := time.ParseDuration(raw) + if err != nil || d < 0 { + log.Printf("invalid first-body timeout %q, using default (disabled)", raw) + return 0 + } + return d +} + var errFirstBodyTimeout = errors.New("upstream timed out before first response body byte") // A cooldown is evidence of a full silent deadline, not a liveness verdict. @@ -124,8 +149,12 @@ func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duratio n, err := body.Read(first) ready <- result{data: first[:n], err: err} }() - timer := time.NewTimer(budget) - defer timer.Stop() + var deadline <-chan time.Time + if budget > 0 { + timer := time.NewTimer(budget) + defer timer.Stop() + deadline = timer.C + } select { case got := <-ready: if len(got.data) == 0 && got.err != nil && !errors.Is(got.err, io.EOF) { @@ -134,7 +163,7 @@ func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duratio return bodyWithPeek{Reader: io.MultiReader(bytes.NewReader(got.data), body), Closer: body}, nil case <-ctx.Done(): return nil, ctx.Err() - case <-timer.C: + case <-deadline: return nil, errFirstBodyTimeout } } diff --git a/services/omlx-proxy/README.md b/services/omlx-proxy/README.md index c2917c92b..409496503 100644 --- a/services/omlx-proxy/README.md +++ b/services/omlx-proxy/README.md @@ -35,7 +35,8 @@ omlx-proxy [flags] | `--port` | `1234` | HTTP listen port for request forwarding | | `--ignore-persisted-port` | `false` | Use `--port` even when a prior runtime port was saved | | `--ipc` | *(empty — use stdio)* | Path to a Unix domain socket or Windows named pipe for IPC | -| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline and, for successful inference, a separate wait for initial body bytes (Go duration, e.g. `10m`). Raise it for long queues or cold model loads. An invalid or non-positive value keeps the default | +| `--response-header-timeout` | `$NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT`, else `120s` | Upstream response-header deadline only (Go duration, e.g. `10m`). Does not limit prefill after headers. An invalid or non-positive value keeps the unchanged `120s` default | +| `--first-body-timeout` | `$NVPAIR_PROXY_FIRST_BODY_TIMEOUT`, else `0` (disabled) | Opt-in, independent wait for initial successful inference body bytes after headers (Go duration, e.g. `10m`). No post-header deadline by default. The flag overrides the environment; `0` explicitly disables it even if the environment opts in. Invalid or negative values leave it disabled | | `--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. | | `--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 | @@ -53,8 +54,8 @@ Node selection: - **Auto**: When no eligible node is explicitly selected, the proxy follows `node/set-priority` (see below), then discovered nodes in stable ID order. - **Priority (scheduler-driven)**: The Job Scheduler ranks the cluster least-loaded-first by pending workload plus smoothed GPU pressure and, via `nvpair-ui-broker`, pushes the ordered node list with those per-node counts to this proxy with `node/set-priority`. Auto routing sends the request to the listed node carrying the least estimated load. See [`nvpair-job-scheduler`](../nvpair-job-scheduler/README.md). - **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing; a silent body times out using the configured response-header duration and fails over within the same candidate list. Once committed, a stream is never retried. Failover keeps one request/workload identity. -- **Deadline cooldown**: A full response-header wait with no answer, or an inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured deadline. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. The default deadline remains `120s`. +- **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. Successful inference responses wait for initial body bytes (or a valid empty EOF) before committing. Pre-commit body read failures retry by default within the same candidate list. There is no first-body deadline unless `--first-body-timeout` / `NVPAIR_PROXY_FIRST_BODY_TIMEOUT` explicitly opts in; only then can post-header silence trigger timeout-based retry. Once committed, a stream is never retried. Failover keeps one request/workload identity. +- **Deadline cooldown**: A full response-header wait with no answer, or an explicitly enabled inference first-body deadline with no content, moves that candidate behind its eligible peers for one configured response-header duration. This stable demotion overrides selection, reservations, and prefix affinity but never excludes a node: a sole candidate, or a candidate after its peers fail, is still tried. State is per proxy process, in memory only; time expiry or a successful response clears it. Fast errors and caller cancellation neither set nor clear it. Header-stall cooldown remains on by default with the unchanged `120s` header deadline; first-body timing is off by default to avoid interrupting legitimate long-context prefill. See [troubleshooting](../../docs/troubleshooting.mdx#requests-are-silent-after-headers) before enabling it. ### IPC Transport diff --git a/services/omlx-proxy/cooldown_test.go b/services/omlx-proxy/cooldown_test.go index 8663932eb..572455631 100644 --- a/services/omlx-proxy/cooldown_test.go +++ b/services/omlx-proxy/cooldown_test.go @@ -36,21 +36,31 @@ func (w *stallWire) Write(b []byte) (int, error) { // upstream again. Reset scheduler reservations so they cannot hide that bug. func TestStalledCandidateSecondRequest(t *testing.T) { old := proxyResponseTimeout - proxyResponseTimeout = time.Second + // Leave scheduling slack between the measured stall and the next request. + // First-body opt-ins below remain shorter and independent of this window. + proxyResponseTimeout = 3 * time.Second defer func() { proxyResponseTimeout = old }() - for _, headers := range []bool{false, true} { - name := "no-headers" - if headers { - name = "headers-without-body" - } - t.Run(name, func(t *testing.T) { + const originalBody = `{"model":"model","messages":[{"role":"user","content":"replay regression"}],"stream":false,"max_tokens":7}` + for _, tc := range []struct { + name, flag, env string + headers bool + }{ + {name: "no-headers-default"}, + {name: "headers-without-body-flag", headers: true, flag: "1s"}, + {name: "headers-without-body-env", headers: true, env: "1s"}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, tc.env) + oldFirstBody := proxyFirstBodyTimeout + proxyFirstBodyTimeout = resolveFirstBodyTimeout(tc.flag) + t.Cleanup(func() { proxyFirstBodyTimeout = oldFirstBody }) var stalledHits, goodHits atomic.Int32 release := make(chan struct{}) stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { stalledHits.Add(1) _, _ = io.Copy(io.Discard, r.Body) - if headers { + if tc.headers { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) w.(http.Flusher).Flush() @@ -61,7 +71,10 @@ func TestStalledCandidateSecondRequest(t *testing.T) { defer close(release) good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { goodHits.Add(1) - _, _ = io.Copy(io.Discard, r.Body) + body, err := io.ReadAll(r.Body) + if err != nil || string(body) != originalBody { + t.Error("healthy peer did not receive the complete original request body") + } _, _ = io.WriteString(w, "healthy") })) defer good.Close() @@ -78,10 +91,10 @@ func TestStalledCandidateSecondRequest(t *testing.T) { Nodes: []string{"stalled", "good"}, Ranks: []schedulerwire.NodeRank{{ID: "stalled"}, {ID: "good"}}, }) - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", - strings.NewReader(`{"model":"model"}`)).WithContext(ctx) + strings.NewReader(originalBody)).WithContext(ctx) rec := httptest.NewRecorder() start := time.Now() p.handleHTTP(rec, req) @@ -103,12 +116,56 @@ func TestStalledCandidateSecondRequest(t *testing.T) { if goodHits.Load() != 2 { t.Fatalf("healthy hits = %d, want 2", goodHits.Load()) } - t.Logf("first=%s second=%s (configured deadline %s)", first, second, proxyResponseTimeout) + t.Logf("first=%s second=%s (header deadline %s, first-body deadline %s)", first, second, proxyResponseTimeout, proxyFirstBodyTimeout) assertStallWorkloads(t, wire, 2) }) } } +// A response-header deadline must not become a post-header generation limit. +// Long-context prefill can remain silent well beyond the header timeout. +func TestDefaultFirstBodyWaitIsUnbounded(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, "") + oldFirstBody := proxyFirstBodyTimeout + proxyFirstBodyTimeout = resolveFirstBodyTimeout("") + defer func() { proxyFirstBodyTimeout = oldFirstBody }() + old := proxyResponseTimeout + proxyResponseTimeout = 100 * time.Millisecond + defer func() { proxyResponseTimeout = old }() + var fallbackHits atomic.Int32 + slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + time.Sleep(3 * proxyResponseTimeout) + _, _ = io.WriteString(w, "prefill-complete") + })) + defer slow.Close() + fallback := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fallbackHits.Add(1) + _, _ = io.WriteString(w, "duplicate") + })) + defer fallback.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "slow", slow.URL, "model")) + disc.AddManual(nodeForModel(t, "fallback", fallback.URL, "model")) + p := testProxy(disc, 1235) + defer p.closeIdleTransports() + p.SetPriority([]string{"slow", "fallback"}) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"model"}`)).WithContext(ctx)) + if rec.Code != http.StatusOK || rec.Body.String() != "prefill-complete" || fallbackHits.Load() != 0 { + t.Fatalf("default cut off healthy post-header prefill: status=%d body=%q fallback hits=%d", + rec.Code, rec.Body.String(), fallbackHits.Load()) + } + if len(p.coolingSet(time.Now())) != 0 { + t.Fatal("healthy prefill created a default cooldown") + } +} + func TestCooldownOrderingRecovery(t *testing.T) { p := prProxy(t, "a", "b", "c", "d") p.SetPriority([]string{"a", "b", "c", "d"}) @@ -276,9 +333,9 @@ func TestCallerCancellationDoesNotCoolCandidate(t *testing.T) { } func TestFinalBodyStallIsBounded(t *testing.T) { - old := proxyResponseTimeout - proxyResponseTimeout = 100 * time.Millisecond - defer func() { proxyResponseTimeout = old }() + old := proxyFirstBodyTimeout + proxyFirstBodyTimeout = 100 * time.Millisecond + defer func() { proxyFirstBodyTimeout = old }() release := make(chan struct{}) engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) @@ -309,7 +366,7 @@ func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { wire.mu.Lock() data := append([]byte(nil), wire.b...) wire.mu.Unlock() - started := make(map[string]bool) + started := make(map[string]string) completed := make(map[string]bool) dec := json.NewDecoder(strings.NewReader(string(data))) for dec.More() { @@ -325,9 +382,15 @@ func assertStallWorkloads(t *testing.T, wire *stallWire, count int) { wl := frame.Params.WorkloadInfo switch frame.Method { case workloadStartedMethod: - started[wl.ID] = true + if wl.State != "running" || completed[wl.ID] { + t.Fatalf("invalid running workload event: %+v", wl) + } + started[wl.ID] = wl.ScheduledOn case workloadCompletedMethod: - if completed[wl.ID] || wl.ScheduledOn != "good" || !started[wl.ID] { + if started[wl.ID] != "good" { + t.Fatalf("no running workload event on serving node before completion: id=%s runningOn=%q", wl.ID, started[wl.ID]) + } + if completed[wl.ID] || wl.ScheduledOn != "good" { t.Fatalf("invalid or duplicate completion: %+v", wl) } completed[wl.ID] = true diff --git a/services/omlx-proxy/first_body_timeout_test.go b/services/omlx-proxy/first_body_timeout_test.go new file mode 100644 index 000000000..9ae5d8f77 --- /dev/null +++ b/services/omlx-proxy/first_body_timeout_test.go @@ -0,0 +1,43 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "testing" + "time" +) + +func TestResolveFirstBodyTimeout(t *testing.T) { + for _, tc := range []struct { + name, flag, env string + want time.Duration + }{ + {"default disabled", "", "", 0}, + {"environment opt-in", "", "10m", 10 * time.Minute}, + {"flag opt-in", "7m", "", 7 * time.Minute}, + {"flag beats environment", "5m", "10m", 5 * time.Minute}, + {"explicit zero disables environment", "0s", "10m", 0}, + {"environment zero disables", "", "0", 0}, + {"invalid flag disables", "soon", "10m", 0}, + {"invalid environment disables", "", "soon", 0}, + {"negative flag disables", "-1s", "10m", 0}, + {"negative environment disables", "", "-1m", 0}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(firstBodyTimeoutEnv, tc.env) + if got := resolveFirstBodyTimeout(tc.flag); got != tc.want { + t.Fatalf("first-body timeout = %s, want %s", got, tc.want) + } + }) + } +} + +func TestFirstBodyTimeoutDoesNotChangeHeaderDeadline(t *testing.T) { + old := proxyFirstBodyTimeout + t.Cleanup(func() { proxyFirstBodyTimeout = old }) + proxyFirstBodyTimeout = 7 * time.Minute + if got := newProxyTransport(nil).ResponseHeaderTimeout; got != proxyResponseTimeout { + t.Fatalf("header timeout = %s, want unchanged %s", got, proxyResponseTimeout) + } +} diff --git a/services/omlx-proxy/main.go b/services/omlx-proxy/main.go index 9554f20d2..672844fe7 100644 --- a/services/omlx-proxy/main.go +++ b/services/omlx-proxy/main.go @@ -26,12 +26,14 @@ func main() { port := flag.Int("port", defaultProxyPort, "HTTP listen port") ignorePersistedPort := flag.Bool("ignore-persisted-port", false, "use --port even when a persisted port exists") responseHeaderTimeout := flag.String("response-header-timeout", "", "upstream response-header timeout (Go duration, e.g. 10m); default: $NVPAIR_PROXY_RESPONSE_HEADER_TIMEOUT or 120s") + firstBodyTimeout := flag.String("first-body-timeout", "", "opt-in inference first-body timeout (Go duration, e.g. 10m; 0 disables); default: $NVPAIR_PROXY_FIRST_BODY_TIMEOUT or disabled") 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") showVersion := flag.Bool("version", false, "print version and exit") resolveLevel := applog.RegisterFlag(nil, slog.LevelInfo) flag.Parse() proxyResponseTimeout = resolveResponseHeaderTimeout(*responseHeaderTimeout) + proxyFirstBodyTimeout = resolveFirstBodyTimeout(*firstBodyTimeout) if *showVersion { fmt.Println(Version) diff --git a/services/omlx-proxy/proxy.go b/services/omlx-proxy/proxy.go index 15d83064d..52de603cf 100644 --- a/services/omlx-proxy/proxy.go +++ b/services/omlx-proxy/proxy.go @@ -726,6 +726,7 @@ func (p *Proxy) serveHTTP(ctx context.Context, ln net.Listener) { "dial_timeout", proxyDialTimeout, "keep_alive", proxyKeepAlive, "response_header_timeout", proxyResponseTimeout, + "first_body_timeout", proxyFirstBodyTimeout, "max_idle_conns", proxyMaxIdleConns, "idle_conn_timeout", proxyIdleConnTimeout, ) @@ -1454,7 +1455,7 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { // Only successful inference waits for content. All commit effects // below still happen once, for the node that actually serves. if isInf && resp.StatusCode >= 200 && resp.StatusCode < 300 { - body, err := awaitFirstBody(r.Context(), resp.Body, proxyResponseTimeout) + body, err := awaitFirstBody(r.Context(), resp.Body, proxyFirstBodyTimeout) if err != nil { if r.Context().Err() == nil && stderrors.Is(err, errFirstBodyTimeout) { p.noteStalled(cand.id, time.Now()) diff --git a/services/omlx-proxy/stall.go b/services/omlx-proxy/stall.go index 449a1311d..c88d20bde 100644 --- a/services/omlx-proxy/stall.go +++ b/services/omlx-proxy/stall.go @@ -8,13 +8,38 @@ import ( "context" "errors" "io" + "log" "net" "net/http" "net/http/httptrace" + "os" "sync" "time" ) +// The response-header deadline does not bound inference prefill after headers. +// First-body timing is opt-in and independent; zero means wait until content, +// an upstream error, or caller cancellation. +const firstBodyTimeoutEnv = "NVPAIR_PROXY_FIRST_BODY_TIMEOUT" + +var proxyFirstBodyTimeout time.Duration + +func resolveFirstBodyTimeout(flagVal string) time.Duration { + raw := flagVal + if raw == "" { + raw = os.Getenv(firstBodyTimeoutEnv) + } + if raw == "" { + return 0 + } + d, err := time.ParseDuration(raw) + if err != nil || d < 0 { + log.Printf("invalid first-body timeout %q, using default (disabled)", raw) + return 0 + } + return d +} + var errFirstBodyTimeout = errors.New("upstream timed out before first response body byte") // A cooldown is evidence of a full silent deadline, not a liveness verdict. @@ -124,8 +149,12 @@ func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duratio n, err := body.Read(first) ready <- result{data: first[:n], err: err} }() - timer := time.NewTimer(budget) - defer timer.Stop() + var deadline <-chan time.Time + if budget > 0 { + timer := time.NewTimer(budget) + defer timer.Stop() + deadline = timer.C + } select { case got := <-ready: if len(got.data) == 0 && got.err != nil && !errors.Is(got.err, io.EOF) { @@ -134,7 +163,7 @@ func awaitFirstBody(ctx context.Context, body io.ReadCloser, budget time.Duratio return bodyWithPeek{Reader: io.MultiReader(bytes.NewReader(got.data), body), Closer: body}, nil case <-ctx.Done(): return nil, ctx.Err() - case <-timer.C: + case <-deadline: return nil, errFirstBodyTimeout } }