diff --git a/README.md b/README.md index f85d0cf..ec8807d 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,10 @@ Open WebUI / n8n ────► :8188 ───┘ - LLM endpoints (`/api/generate`, `/api/chat`, `/api/embed`, `/v1/*`) take the LLM lock: concurrent requests allowed, but blocked while an image job is - active or waiting (image priority). + active or waiting (image priority). Blocked requests either hang until the + lock is free (`LLM_BUSY_MODE=wait`, default) or fail immediately with 503 + (or 429) + `Retry-After` (`LLM_BUSY_MODE=reject`) — the latter lets routers + like LiteLLM cool down and retry instead of holding a hung connection. - `POST /prompt` on the ComfyUI listener takes the image lock: new LLM requests block, in-flight LLMs drain, Ollama models are unloaded, the prompt is forwarded, and the lock is held until the job finishes and ComfyUI frees @@ -42,7 +45,10 @@ override file values. Invalid values fail at startup. | `COMFY_URL` | `http://127.0.0.1:8189` | ComfyUI upstream | | `UNLOAD_TIMEOUT` | `60s` | Wait for Ollama to unload before an image job | | `JOB_TIMEOUT` | `15m` | Wait for a ComfyUI job to finish | -| `LLM_WAIT_TIMEOUT` | `10m` | Max lock wait for an LLM request before 503 | +| `LLM_WAIT_TIMEOUT` | `10m` | Max lock wait for an LLM request before 503 (wait mode) | +| `LLM_BUSY_MODE` | `wait` | `wait` = hold blocked LLM requests; `reject` = fail them immediately | +| `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400–599, e.g. 429) | +| `BUSY_RETRY_AFTER` | `30` | Seconds sent as `Retry-After` on busy responses (both modes) | | `WARM_MODEL` | _(empty)_ | Model to reload after an image job (off by default) | | `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` works as an alias | | `LOG_FORMAT` | `text` | `json` for structured JSON logs | @@ -96,6 +102,13 @@ gpu-turnstile.exe service remove The service uses the config file (services have no convenient environment); set `LOG_FILE` in it since there is no console. +Suggested layout: `C:\Program Files\gpu-turnstile\` for the exe and +`gpu-turnstile.env`, logs under `C:\ProgramData\gpu-turnstile\` via +`LOG_FILE`. The service runs as `LocalSystem` by default, which can write +the install directory for self-updates. For least privilege, run it as the +virtual account `NT SERVICE\gpu-turnstile` and grant write access to just +those two directories. + **Auto-update is on by default**: the binary checks the repo's latest release on startup and every `UPDATE_INTERVAL`, verifies the Ed25519 signature of the download against the public key embedded at build time, diff --git a/SPEC.md b/SPEC.md index 374c7c6..ed4a048 100644 --- a/SPEC.md +++ b/SPEC.md @@ -51,6 +51,11 @@ are LLM requests and the single "writer" is an image job): `image` **or while an image job is waiting**. Then state := `llm`, n++. On completion (response fully written, including streamed bodies, or client disconnect) n--; if n == 0 state := `idle`. + `LLM_BUSY_MODE` selects what a blocked LLM request sees: `wait` (default) + hangs until the lock is free or `LLM_WAIT_TIMEOUT` expires (then 503 + + `Retry-After`); `reject` answers immediately with `LLM_BUSY_STATUS` + (default 503; 429 works too) + `Retry-After: BUSY_RETRY_AFTER`, which + routers like LiteLLM honor for cooldowns/retries. - **Image job**: `AcquireImage()` marks "image pending" (so no new LLM requests start), waits until n == 0, sets state := `image`. Released after the ComfyUI job finished and models were freed. @@ -116,7 +121,10 @@ override file values. A missing file is fine; a malformed one is fatal. | `COMFY_URL` | `http://127.0.0.1:8189` | upstream | | `UNLOAD_TIMEOUT` | `60s` | wait for Ollama to unload | | `JOB_TIMEOUT` | `15m` | wait for ComfyUI job | -| `LLM_WAIT_TIMEOUT` | `10m` | max time an LLM request waits for the lock before 503 | +| `LLM_WAIT_TIMEOUT` | `10m` | max time an LLM request waits for the lock before 503 (wait mode) | +| `LLM_BUSY_MODE` | `wait` | `wait` = hold blocked LLM requests; `reject` = fail them immediately | +| `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400–599, e.g. 429) | +| `BUSY_RETRY_AFTER` | `30` | seconds sent as `Retry-After` on busy responses (both modes) | | `WARM_MODEL` | `` | optional model to reload after an image job | | `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` is accepted as an alias | | `LOG_FORMAT` | `text` | `json` for structured JSON logs | @@ -146,6 +154,15 @@ Docker. - `gpu-turnstile.exe service install [-config path]` registers an auto-start Windows service (needs an elevated shell). Recovery actions restart it 5 s after any failure. `service remove` uninstalls. +- **Layout**: install to `C:\Program Files\gpu-turnstile\` (exe plus + `gpu-turnstile.env`); logs belong in `C:\ProgramData\gpu-turnstile\` via + `LOG_FILE`. The service must be able to write its install directory for + self-updates — Program Files is writable by LocalSystem and admins, which + is why running as the default `LocalSystem` account is the simple choice. +- **Account**: the default `LocalSystem` works out of the box. For least + privilege, create the service with the virtual account + `NT SERVICE\gpu-turnstile` and grant it write access to the install and + log directories only (no network logon, no user profile). - Use a config file (above) for the service — Windows services have no convenient environment. Logs go to `LOG_FILE` since there is no console. - **Auto-update**: on startup and every `UPDATE_INTERVAL`, the binary diff --git a/cmd/gpu-turnstile/main.go b/cmd/gpu-turnstile/main.go index 6083496..dcd6f4f 100644 --- a/cmd/gpu-turnstile/main.go +++ b/cmd/gpu-turnstile/main.go @@ -192,6 +192,9 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri "unload_timeout", cfg.UnloadTimeout, "job_timeout", cfg.JobTimeout, "llm_wait_timeout", cfg.LLMWaitTimeout, + "llm_busy_mode", cfg.LLMBusyMode, + "llm_busy_status", cfg.LLMBusyStatus, + "busy_retry_after", cfg.BusyRetryAfter, "unload_poll_interval", cfg.UnloadPollInterval, "history_poll_interval", cfg.HistoryPollInterval, "probe_timeout", cfg.ProbeTimeout, @@ -238,6 +241,9 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri HistoryPollInterval: cfg.HistoryPollInterval, FreeTimeout: cfg.FreeTimeout, WarmTimeout: cfg.WarmTimeout, + LLMBusyMode: cfg.LLMBusyMode, + LLMBusyStatus: cfg.LLMBusyStatus, + BusyRetryAfter: cfg.BusyRetryAfter, BackoffInitial: cfg.BackoffInitial, BackoffMax: cfg.BackoffMax, PromptCaptureLimit: cfg.PromptCaptureLimit, diff --git a/compose.yaml.example b/compose.yaml.example index 2ba61d6..5323d75 100644 --- a/compose.yaml.example +++ b/compose.yaml.example @@ -18,6 +18,9 @@ services: # UNLOAD_TIMEOUT: 60s # JOB_TIMEOUT: 15m # LLM_WAIT_TIMEOUT: 10m + # LLM_BUSY_MODE: reject # wait (default) hangs; reject fails fast + # LLM_BUSY_STATUS: 429 # status in reject mode (default 503) + # BUSY_RETRY_AFTER: 30 # Retry-After seconds on busy responses # WARM_MODEL: qwen3:14b # LOG_LEVEL: info # LOG_FORMAT: json diff --git a/internal/config/config.go b/internal/config/config.go index add9fdd..5785dd1 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -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 == "" { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 2e081ba..fa51d5a 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -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 { diff --git a/internal/lock/lock.go b/internal/lock/lock.go index 1cf1eb5..c98dd51 100644 --- a/internal/lock/lock.go +++ b/internal/lock/lock.go @@ -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() diff --git a/internal/lock/try_test.go b/internal/lock/try_test.go new file mode 100644 index 0000000..d60dd93 --- /dev/null +++ b/internal/lock/try_test.go @@ -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() +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index f6bccd3..c4217c9 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -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 diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index c04a0f3..2b68702 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -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) + } +}