PUBLIC/gpu-turnstile
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

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

Browse Source 
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.
This commit is contained in:
mram
2026-09-20 21:43:03 +02:00
parent e363c9e4e4
commit bacb26772a
10 changed files with 291 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/config/config.go
+31
View File
@@ -37,6 +37,13 @@ type Config struct {
UpdateRepo string
UpdateAsset string
// LLMBusyMode is "wait" (hold requests until the lock is free or
// LLMWaitTimeout expires) or "reject" (immediately answer with
// LLMBusyStatus + Retry-After when an image job is active or pending).
LLMBusyMode string
LLMBusyStatus int
BusyRetryAfter int
WarmModel string
LogLevel slog.Level
LogJSON bool
@@ -70,6 +77,10 @@ func Defaults() Config {
UpdateRepo: "https://git.rambossek.at/PUBLIC/gpu-turnstile",
UpdateAsset: "gpu-turnstile.exe",
LLMBusyMode: "wait",
LLMBusyStatus: 503,
BusyRetryAfter: 30,
LogLevel: slog.LevelWarn,
}
}
@@ -169,6 +180,26 @@ func Load(getenv func(string) string) (Config, error) {
}
cfg.AutoUpdate = b
}
if v := getenv("LLM_BUSY_MODE"); v != "" {
if v != "wait" && v != "reject" {
return cfg, fmt.Errorf("LLM_BUSY_MODE: must be \"wait\" or \"reject\"")
}
cfg.LLMBusyMode = v
}
if v := getenv("LLM_BUSY_STATUS"); v != "" {
n, err := strconv.Atoi(v)
if err != nil || n < 400 || n > 599 {
return cfg, fmt.Errorf("LLM_BUSY_STATUS: must be an HTTP status in 400-599")
}
cfg.LLMBusyStatus = n
}
if v := getenv("BUSY_RETRY_AFTER"); v != "" {
n, err := strconv.Atoi(v)
if err != nil || n <= 0 {
return cfg, fmt.Errorf("BUSY_RETRY_AFTER: must be a positive integer (seconds)")
}
cfg.BusyRetryAfter = n
}
// LOGLEVEL is the canonical spelling; LOG_LEVEL is kept as an alias.
logLevelValue := getenv("LOGLEVEL")
if logLevelValue == "" {
internal/config/config_test.go
+3
View File
@@ -83,6 +83,9 @@ func TestLoadErrors(t *testing.T) {
{"AUTO_UPDATE", "maybe"},
{"LOGLEVEL", "shouty"},
{"LOG_FORMAT", "yaml"},
{"LLM_BUSY_MODE", "bogus"},
{"LLM_BUSY_STATUS", "200"},
{"BUSY_RETRY_AFTER", "0"},
} {
_, err := Load(func(k string) string {
if k == tc.key {
internal/lock/lock.go
+16
View File
@@ -74,6 +74,22 @@ func (l *Lock) AcquireLLM(ctx context.Context) error {
return nil
}
// TryAcquireLLM acquires one in-flight LLM slot without waiting and
// reports whether it succeeded. It fails when an image job is active or
// pending.
func (l *Lock) TryAcquireLLM() bool {
l.mu.Lock()
if l.imageActive || len(l.imageQ) > 0 {
l.mu.Unlock()
return false
}
l.n++
n := l.n
l.mu.Unlock()
l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
return true
}
// ReleaseLLM marks one LLM request as finished.
func (l *Lock) ReleaseLLM() {
l.mu.Lock()
internal/lock/try_test.go
+51
View File
@@ -0,0 +1,51 @@
package lock
import (
"context"
"testing"
"time"
)
func TestTryAcquireLLM(t *testing.T) {
lk := New(nil)
if !lk.TryAcquireLLM() {
t.Fatal("TryAcquireLLM on idle lock should succeed")
}
if _, n, _ := lk.Snapshot(); n != 1 {
t.Fatalf("n = %d, want 1", n)
}
// While an image job is pending, TryAcquireLLM must fail.
imageWaiting := make(chan struct{})
go func() {
lk.AcquireImage(context.Background())
close(imageWaiting)
}()
deadline := time.Now().Add(2 * time.Second)
for {
lk.mu.Lock()
queued := len(lk.imageQ)
lk.mu.Unlock()
if queued == 1 {
break
}
if time.Now().After(deadline) {
t.Fatal("image waiter never queued")
}
time.Sleep(time.Millisecond)
}
if lk.TryAcquireLLM() {
t.Fatal("TryAcquireLLM with image pending should fail")
}
lk.ReleaseLLM()
<-imageWaiting
if lk.TryAcquireLLM() {
t.Fatal("TryAcquireLLM with image active should fail")
}
lk.ReleaseImage()
if !lk.TryAcquireLLM() {
t.Fatal("TryAcquireLLM after image release should succeed")
}
lk.ReleaseLLM()
}
internal/proxy/proxy.go
+42
View File
@@ -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
internal/proxy/proxy_test.go
+106
View File
@@ -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)
}
}