Add LLM busy modes: wait (hang) or reject with Retry-After

LLM_BUSY_MODE=reject answers blocked LLM requests immediately with LLM_BUSY_STATUS (default 503, 429 works) and Retry-After, so routers like LiteLLM can cool down and retry instead of holding a hung connection. The default wait mode now also sends Retry-After when LLM_WAIT_TIMEOUT expires. Document the service account (LocalSystem default, NT SERVICE virtual-account hardening) and the Program Files / ProgramData install layout.
2026-09-20 21:43:03 +02:00
parent e363c9e4e4
commit bacb26772a
10 changed files with 291 additions and 3 deletions
@@ -54,6 +54,15 @@ type Config struct {
 	UnloadTimeout  time.Duration
 	JobTimeout     time.Duration

+	// LLMBusyMode is "wait" (default) or "reject". In reject mode an LLM
+	// request that arrives while an image job is active or pending is
+	// answered immediately with LLMBusyStatus and a Retry-After header
+	// (BusyRetryAfter seconds) instead of waiting for the lock. In wait
+	// mode the Retry-After header is sent when LLMWaitTimeout expires.
+	LLMBusyMode    string
+	LLMBusyStatus  int
+	BusyRetryAfter int
+
 	// BackoffInitial and BackoffMax control the exponential retry backoff
 	// when an upstream refuses a connection: the wait doubles from
 	// BackoffInitial up to BackoffMax between attempts. Zero selects the
@@ -85,6 +94,9 @@ type Server struct {
 	captureLimit   int64
 	backoffInitial time.Duration
 	backoffMax     time.Duration
+	busyMode       string
+	busyStatus     int
+	busyRetryAfter int

 	ollamaProxy *httputil.ReverseProxy
 	comfyProxy  *httputil.ReverseProxy
@@ -130,6 +142,18 @@ func New(cfg Config) (*Server, error) {
 	if backoffMax <= 0 {
 		backoffMax = time.Minute
 	}
+	busyMode := "wait"
+	if cfg.LLMBusyMode == "reject" {
+		busyMode = "reject"
+	}
+	busyStatus := cfg.LLMBusyStatus
+	if busyStatus == 0 {
+		busyStatus = http.StatusServiceUnavailable
+	}
+	busyRetryAfter := cfg.BusyRetryAfter
+	if busyRetryAfter <= 0 {
+		busyRetryAfter = 30
+	}
 	retry := &retryTransport{
 		base:    http.DefaultTransport,
 		initial: backoffInitial,
@@ -145,6 +169,9 @@ func New(cfg Config) (*Server, error) {
 		captureLimit:   captureLimit,
 		backoffInitial: backoffInitial,
 		backoffMax:     backoffMax,
+		busyMode:       busyMode,
+		busyStatus:     busyStatus,
+		busyRetryAfter: busyRetryAfter,
 		ollamaProxy:    newReverseProxy(ollamaURL, retry, log.With("upstream", "ollama")),
 		comfyProxy:     newReverseProxy(comfyURL, retry, log.With("upstream", "comfy")),
 	}, nil
@@ -410,12 +437,27 @@ func (s *Server) OllamaHandler() http.Handler {
 		}

 		start := time.Now()
+		if s.busyMode == "reject" {
+			if !s.cfg.Lock.TryAcquireLLM() {
+				s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds())
+				s.log.Info("llm request rejected; GPU busy",
+					"path", r.URL.Path, "status", s.busyStatus)
+				w.Header().Set("Retry-After", strconv.Itoa(s.busyRetryAfter))
+				http.Error(w, "GPU busy: image job active or queued", s.busyStatus)
+				return
+			}
+			s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds())
+			defer s.cfg.Lock.ReleaseLLM()
+			s.ollamaProxy.ServeHTTP(w, r)
+			return
+		}
 		wctx, cancel := context.WithTimeout(r.Context(), s.cfg.LLMWaitTimeout)
 		err := s.cfg.Lock.AcquireLLM(wctx)
 		cancel()
 		s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds())
 		if err != nil {
 			if errors.Is(err, context.DeadlineExceeded) && r.Context().Err() == nil {
+				w.Header().Set("Retry-After", strconv.Itoa(s.busyRetryAfter))
 				http.Error(w, "GPU busy: timed out waiting for the lock", http.StatusServiceUnavailable)
 			}
 			return
@@ -2,6 +2,7 @@ package proxy

 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -345,3 +346,108 @@ func TestPassThroughNoLock(t *testing.T) {
 		t.Fatalf("pass-through = %d %s", resp.StatusCode, body)
 	}
 }
+
+func TestLLMBusyReject(t *testing.T) {
+	f := newFakes(t)
+
+	lk := lock.New(nil)
+	if err := lk.AcquireImage(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	ollamaClient, err := ollama.New(f.ollama.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	comfyClient, err := comfy.New(f.comfy.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	srv, err := New(Config{
+		OllamaURL:      f.ollama.URL,
+		ComfyURL:       f.comfy.URL,
+		Lock:           lk,
+		Ollama:         ollamaClient,
+		Comfy:          comfyClient,
+		Metrics:        metrics.New(),
+		LLMWaitTimeout: 2 * time.Second,
+		LLMBusyMode:    "reject",
+		BusyRetryAfter: 17,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	front := httptest.NewServer(srv.OllamaHandler())
+	defer front.Close()
+
+	start := time.Now()
+	resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, _ := io.ReadAll(resp.Body)
+	resp.Body.Close()
+	if resp.StatusCode != http.StatusServiceUnavailable {
+		t.Fatalf("busy chat status = %d %s", resp.StatusCode, body)
+	}
+	if got := resp.Header.Get("Retry-After"); got != "17" {
+		t.Fatalf("Retry-After = %q", got)
+	}
+	if elapsed := time.Since(start); elapsed > time.Second {
+		t.Fatalf("reject was not immediate: %v", elapsed)
+	}
+
+	// After the image lock is released the next LLM request goes through.
+	lk.ReleaseImage()
+	resp, err = http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp.Body.Close()
+	if resp.StatusCode != 200 {
+		t.Fatalf("chat after release status = %d", resp.StatusCode)
+	}
+}
+
+func TestLLMBusyWaitTimeoutRetryAfter(t *testing.T) {
+	f := newFakes(t)
+
+	lk := lock.New(nil)
+	if err := lk.AcquireImage(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	defer lk.ReleaseImage()
+	ollamaClient, err := ollama.New(f.ollama.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	comfyClient, err := comfy.New(f.comfy.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	srv, err := New(Config{
+		OllamaURL:      f.ollama.URL,
+		ComfyURL:       f.comfy.URL,
+		Lock:           lk,
+		Ollama:         ollamaClient,
+		Comfy:          comfyClient,
+		Metrics:        metrics.New(),
+		LLMWaitTimeout: 50 * time.Millisecond,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	front := httptest.NewServer(srv.OllamaHandler())
+	defer front.Close()
+
+	resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp.Body.Close()
+	if resp.StatusCode != http.StatusServiceUnavailable {
+		t.Fatalf("timed-out chat status = %d", resp.StatusCode)
+	}
+	if got := resp.Header.Get("Retry-After"); got != "30" {
+		t.Fatalf("Retry-After = %q", got)
+	}
+}