diff --git a/.gitignore b/.gitignore index e4c1146..7b1905e 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ /gpu-turnstile.exe.old /gpu-turnstile.exe.new +/gpu-turnstile.exe~ diff --git a/README.md b/README.md index 8d8a827..19fae8c 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,10 @@ GPU arbitration proxy for Ollama + ComfyUI. One consumer GPU is shared by an LLM server (Ollama) and an image generator (ComfyUI); gpu-turnstile sits in -front of both and guarantees the GPU is always in exactly one of three states: -`idle`, `llm` (N ≥ 1 Ollama requests in flight), or `image` (exactly one -ComfyUI job, Ollama models unloaded). See [SPEC.md](SPEC.md) for the full +front of both and guarantees the GPU is always in exactly one of four states: +`idle`, `llm` (N ≥ 1 Ollama requests in flight), `image` (exactly one +ComfyUI job, Ollama models unloaded), or `external` (a foreign process such +as a game holds the GPU). See [SPEC.md](SPEC.md) for the full design. gpu-turnstile listens on the ports the services normally use; the actual @@ -32,8 +33,9 @@ Open WebUI / n8n ────► :8188 ───┘ Each consumer is enabled by setting its URL (`OLLAMA_URL`, `COMFY_URL`) and disabled by leaving it empty — at least one is required. With only Ollama the proxy is a pass-through (no image jobs can arrive); with only ComfyUI -the Ollama unload/warm steps are skipped. Future consumers (e.g. local game -detection) plug into the same lock the same way. +the Ollama unload/warm steps are skipped. A third, URL-less consumer — +detection of foreign GPU holders such as games — is enabled by `GAME_PROCS` +and/or `GPU_FOREIGN_VRAM_MB` (see below). ## Configuration @@ -60,6 +62,10 @@ override file values. Invalid values fail at startup. | `COMFY_DIR` | _(empty)_ | Working directory for `COMFY_CMD` | | `COMFY_IDLE_TIMEOUT` | `5m` | Stop the managed ComfyUI after this long idle | | `COMFY_START_TIMEOUT` | `2m` | Max wait for the managed ComfyUI to come up | +| `GAME_PROCS` | _(empty = disabled)_ | Process names (comma-separated); while any runs, the GPU counts as held: requests wait, Ollama unloads, managed ComfyUI stops | +| `GPU_FOREIGN_VRAM_MB` | `0` (disabled) | Also treat the GPU as held when a non-ignored process uses more VRAM than this (needs nvidia-smi) | +| `GPU_IGNORE_PROCS` | `ollama,ollama app,ollama_llama_server,python,pythonw` | Process names never counted as foreign GPU users | +| `GAME_POLL_INTERVAL` | `5s` | How often game/VRAM detection runs | | `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 | | `LOG_FILE` | _(empty)_ | Append logs to this file instead of stderr | @@ -81,7 +87,7 @@ override file values. Invalid values fail at startup. ## Observability -- `GET /healthz` (both listeners): `{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}` +- `GET /healthz` (both listeners): `{"state":"idle|llm|image|external","llm_inflight":N,"image_pending":B}` - `GET /metrics` (both listeners): Prometheus text format — `gpu_turnstile_state`, `gpu_turnstile_llm_inflight`, `gpu_turnstile_image_pending`, `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds` @@ -118,6 +124,27 @@ answers on the port, gpu-turnstile just uses it instead of spawning instance already holds the port when you open the desktop app, the desktop's server is the one that fails to bind. +## Game detection + +Want to game on the same GPU without Ollama/ComfyUI squatting on the VRAM? +gpu-turnstile can watch for foreign GPU holders and, while one is active, +make LLM/image requests wait (or 503, per `LLM_BUSY_MODE`), unload Ollama's +models and stop the managed ComfyUI so the game gets the memory. Two +detection paths, each optional, polled every `GAME_POLL_INTERVAL` (5s): + +``` +GAME_PROCS=cyberpunk2077.exe,bg3.exe # the reliable way on Windows +GPU_FOREIGN_VRAM_MB=1024 # catch-all via nvidia-smi +``` + +`GAME_PROCS` matches running process names (case-insensitive, `.exe` +optional). `GPU_FOREIGN_VRAM_MB` asks nvidia-smi which processes hold GPU +memory and treats anything not in `GPU_IGNORE_PROCS` above the threshold as +foreign — handy as a catch-all, but note that under Windows' WDDM driver +graphics-only games may not show up in nvidia-smi's per-process list, so +name your games in `GAME_PROCS` there; on Linux both paths work. When the +game exits, requests resume automatically. + ## Build and run ```sh diff --git a/SPEC.md b/SPEC.md index 2dffb38..51a60e1 100644 --- a/SPEC.md +++ b/SPEC.md @@ -10,11 +10,12 @@ becomes very slow; on Linux it would OOM instead. ## Goal A single Go binary that sits in front of **both** services and guarantees that -at any moment the GPU is in exactly one of three states: +at any moment the GPU is in exactly one of four states: - `idle` — nothing in flight - `llm` — N ≥ 1 Ollama inference requests in flight (concurrency allowed) - `image` — exactly one ComfyUI job in flight, Ollama models unloaded +- `external` — a foreign process (e.g. a game) holds the GPU; new work waits Clients (LiteLLM, Open WebUI, n8n) point at gpu-turnstile instead of at the services. gpu-turnstile is transparent for everything that does not touch the @@ -53,9 +54,10 @@ consumer gets no listener, no startup probe, and no lock participation: blocks since no image jobs can arrive. - **Only `COMFY_URL`**: image jobs are tracked and ComfyUI's VRAM is freed afterwards, but the Ollama unload and warm-reload steps are skipped. -- Future consumers (e.g. detecting a local game holding VRAM) plug into the - same lock the same way: enabled by their config knob, excluded when - absent. +- **Game detection** is a third, optional consumer without a URL: enabled by + `GAME_PROCS` and/or `GPU_FOREIGN_VRAM_MB` it watches for foreign processes + holding the GPU (see below) and plugs into the same lock the same way — + excluded when both knobs are unset. ### Lock semantics @@ -75,6 +77,11 @@ are LLM requests and the single "writer" is an image job): requests start), waits until n == 0, sets state := `image`. Released after the ComfyUI job finished and models were freed. - Concurrent image jobs queue FIFO behind each other. +- **External hold**: game detection calls `SetExternal(holder)` while a + foreign process holds the GPU. New LLM and image grants block (same busy + handling as above) until `ClearExternal()`; in-flight work is not + preempted, it drains. Snapshot reports state `external` once nothing else + is in flight. - All waits are context-aware: a client that disconnects while waiting is removed from the queue. @@ -148,6 +155,41 @@ process instead of expecting an always-on server: skips the intentionally-stopped/starting states; a failed probe while the process is alive and was previously ready is logged as DOWN. +## Game detection (foreign GPU holders) + +Games and other foreign GPU users sit outside the URL-based consumer model — +nothing proxies through gpu-turnstile for them. Two independent detection +paths, polled every `GAME_POLL_INTERVAL` (default 5 s); either one being +configured enables the feature: + +- **Process watch list** (`GAME_PROCS`, comma-separated, case-insensitive, + `.exe` optional): while any listed process runs, the GPU counts as held. + This is the reliable path on Windows. +- **Foreign VRAM threshold** (`GPU_FOREIGN_VRAM_MB`): `nvidia-smi + --query-compute-apps` lists per-process GPU memory; any process not in + `GPU_IGNORE_PROCS` (default: Ollama and python — ComfyUI runs under python) + holding more than the threshold counts as a foreign holder. Needs + nvidia-smi on the PATH (absent: logged once, path disabled) and works best + on Linux — under Windows' WDDM driver, graphics-only games may not appear + in the per-process list. + +While a holder is detected, gpu-turnstile: + +1. takes the external hold on the lock (`SetExternal`), so new LLM/image + requests wait or are rejected per `LLM_BUSY_MODE`; busy responses name the + holder, +2. once in-flight work has drained, **frees VRAM for the foreign process**: + the managed ComfyUI is stopped (an external server on its port is never + touched) and Ollama's resident models are unloaded, +3. refuses to spawn the managed ComfyUI; ComfyUI requests that would need a + spawn are answered 503 + `Retry-After` (a running desktop instance keeps + being proxied), +4. logs the transitions at WARN ("GPU held by an external process …" / + "released the GPU; resuming") and reports state `external` in `/healthz` + and `/metrics`. + +When the holder disappears, the hold is lifted and queued requests proceed. + ## Configuration (env) Configuration comes from environment variables and/or an `.env`-style @@ -173,6 +215,10 @@ override file values. A missing file is fine; a malformed one is fatal. | `COMFY_DIR` | `` | working directory for `COMFY_CMD` | | `COMFY_IDLE_TIMEOUT` | `5m` | stop the managed ComfyUI after this long without requests or jobs | | `COMFY_START_TIMEOUT` | `2m` | how long a request waits for the managed ComfyUI to come up | +| `GAME_PROCS` | _(empty = disabled)_ | comma-separated process names (case-insensitive, `.exe` optional); while any runs, the GPU counts as held by it: requests wait, Ollama unloads, the managed ComfyUI stops | +| `GPU_FOREIGN_VRAM_MB` | `0` (disabled) | also treat the GPU as held when a process not in `GPU_IGNORE_PROCS` uses more VRAM than this; needs nvidia-smi | +| `GPU_IGNORE_PROCS` | `ollama,ollama app,ollama_llama_server,python,pythonw` | process names never counted as foreign GPU users | +| `GAME_POLL_INTERVAL` | `5s` | how often game/VRAM detection runs | | `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 | | `LOG_FILE` | `` | append logs to this file instead of stderr (useful as a service) | diff --git a/cmd/gpu-turnstile/main.go b/cmd/gpu-turnstile/main.go index c6f8c67..7c27324 100644 --- a/cmd/gpu-turnstile/main.go +++ b/cmd/gpu-turnstile/main.go @@ -21,6 +21,7 @@ import ( "gpu-turnstile/internal/comfy" "gpu-turnstile/internal/config" + "gpu-turnstile/internal/game" "gpu-turnstile/internal/lock" "gpu-turnstile/internal/metrics" "gpu-turnstile/internal/ollama" @@ -441,6 +442,10 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri "comfy_dir", cfg.ComfyDir, "comfy_idle_timeout", cfg.ComfyIdleTimeout, "comfy_start_timeout", cfg.ComfyStartTimeout, + "game_procs", cfg.GameProcs, + "gpu_foreign_vram_mb", cfg.GPUForeignVRAMMB, + "gpu_ignore_procs", cfg.GPUIgnoreProcs, + "game_poll_interval", cfg.GamePollInterval, "auto_update", cfg.AutoUpdate, "update_interval", cfg.UpdateInterval, "update_repo", cfg.UpdateRepo, @@ -555,6 +560,13 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes) } + // Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS + // and/or GPU_FOREIGN_VRAM_MB — hold the lock externally while they run. + if len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0 { + det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUIgnoreProcs, log) + go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup) + } + // Bind the listeners up front so a port conflict fails fast and the // readiness notification below really means "accepting connections". var servers []*http.Server @@ -615,6 +627,80 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri // (idle); health checks skip it instead of logging an outage. var errManagedDown = errors.New("managed upstream intentionally stopped") +// gameLoop polls for foreign GPU holders (a game, another ML job). While one +// is detected it holds the lock externally so new LLM and image requests +// wait (or are rejected per LLM_BUSY_MODE), and — once in-flight work has +// drained — frees VRAM for it: the managed ComfyUI is stopped and Ollama's +// resident models are unloaded. +func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *game.Detector, lk *lock.Lock, ollamaClient *ollama.Client, comfySup *supervise.Process) { + ticker := time.NewTicker(cfg.GamePollInterval) + defer ticker.Stop() + held, freed := false, false + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + holders, err := det.Check(ctx) + if err != nil && ctx.Err() == nil { + log.Warn("game detection failed", "err", err) + } + switch { + case len(holders) > 0 && !held: + held = true + lk.SetExternal(summarizeHolders(holders)) + log.Warn("GPU held by an external process; new LLM/image requests wait", + "holders", summarizeHolders(holders)) + case len(holders) == 0 && held: + held, freed = false, false + lk.ClearExternal() + log.Warn("external process released the GPU; resuming") + } + if held && !freed { + if state, _, _ := lk.Snapshot(); state != lock.StateLLM && state != lock.StateImage { + freed = true + freeVRAM(ctx, cfg.UnloadTimeout, log, ollamaClient, comfySup) + } + } + } +} + +// summarizeHolders joins holder descriptions for logs and busy responses, +// capping the list so a process name matching dozens of PIDs (system +// services) does not flood the log. +func summarizeHolders(holders []string) string { + const max = 5 + if len(holders) > max { + return strings.Join(holders[:max], "; ") + fmt.Sprintf("; +%d more", len(holders)-max) + } + return strings.Join(holders, "; ") +} + +// freeVRAM stops the managed ComfyUI (never an external server on its port) +// and unloads Ollama's resident models so the foreign process gets the GPU +// memory. +func freeVRAM(ctx context.Context, unloadTimeout time.Duration, log *slog.Logger, ollamaClient *ollama.Client, comfySup *supervise.Process) { + if comfySup != nil && comfySup.Running() { + log.Warn("stopping the managed ComfyUI to free VRAM") + comfySup.Stop() + } + if ollamaClient == nil { + return + } + uctx, cancel := context.WithTimeout(ctx, unloadTimeout) + defer cancel() + if models, err := ollamaClient.LoadedModels(uctx); err != nil || len(models) == 0 { + return // nothing resident (or ollama unreachable); nothing to free + } + elapsed, err := ollamaClient.UnloadAll(uctx) + if err != nil { + log.Warn("ollama unload incomplete; continuing", "err", err) + return + } + log.Warn("ollama models unloaded to free VRAM", "seconds", elapsed.Seconds()) +} + // healthLoop probes the enabled upstreams every interval and logs status // transitions — "is DOWN" when a previously healthy upstream stops // answering, "recovered" when it comes back. The first round only diff --git a/internal/config/config.go b/internal/config/config.go index fb47d42..dbb498e 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -61,6 +61,17 @@ type Config struct { ComfyIdleTimeout time.Duration ComfyStartTimeout time.Duration + // GameProcs (GAME_PROCS) is a watch list of process names; while any of + // them runs, the GPU is treated as held by a foreign process. The + // nvidia-smi path (GPUForeignVRAMMB, GPU_FOREIGN_VRAM_MB) does the same + // when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than + // that many MiB of VRAM. GamePollInterval (GAME_POLL_INTERVAL) is how + // often both checks run. + GameProcs []string + GPUForeignVRAMMB int + GPUIgnoreProcs []string + GamePollInterval time.Duration + WarmModel string LogLevel slog.Level LogJSON bool @@ -102,6 +113,11 @@ func Defaults() Config { ComfyIdleTimeout: 5 * time.Minute, ComfyStartTimeout: 2 * time.Minute, + // ComfyUI runs under python; excluding it (and Ollama) by name keeps + // our own consumers from tripping the foreign-VRAM check. + GPUIgnoreProcs: []string{"ollama", "ollama app", "ollama_llama_server", "python", "pythonw"}, + GamePollInterval: 5 * time.Second, + LogLevel: slog.LevelWarn, } } @@ -137,6 +153,17 @@ func ParseEnvFile(r io.Reader) (map[string]string, error) { return values, scanner.Err() } +// splitList parses a comma-separated setting into trimmed, non-empty items. +func splitList(v string) []string { + var out []string + for _, item := range strings.Split(v, ",") { + if item = strings.TrimSpace(item); item != "" { + out = append(out, item) + } + } + return out +} + func envDuration(getenv func(string) string, name string, dst *time.Duration) error { v := getenv(name) if v == "" { @@ -192,11 +219,25 @@ func Load(getenv func(string) string) (Config, error) { {"COMFY_IDLE_TIMEOUT", &cfg.ComfyIdleTimeout}, {"COMFY_START_TIMEOUT", &cfg.ComfyStartTimeout}, {"UPDATE_INTERVAL", &cfg.UpdateInterval}, + {"GAME_POLL_INTERVAL", &cfg.GamePollInterval}, } { if err := envDuration(getenv, e.name, e.dst); err != nil { return cfg, err } } + if v := getenv("GAME_PROCS"); v != "" { + cfg.GameProcs = splitList(v) + } + if v := getenv("GPU_IGNORE_PROCS"); v != "" { + cfg.GPUIgnoreProcs = splitList(v) + } + if v := getenv("GPU_FOREIGN_VRAM_MB"); v != "" { + n, err := strconv.Atoi(v) + if err != nil || n < 0 { + return cfg, fmt.Errorf("GPU_FOREIGN_VRAM_MB: must be a non-negative integer (MiB, 0 = disabled)") + } + cfg.GPUForeignVRAMMB = n + } if v := getenv("PROMPT_CAPTURE_LIMIT"); v != "" { n, err := strconv.ParseInt(v, 10, 64) if err != nil || n < 0 { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 20fd499..5d3d4f0 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -157,6 +157,8 @@ func TestLoadErrors(t *testing.T) { {"LLM_BUSY_MODE", "bogus"}, {"LLM_BUSY_STATUS", "200"}, {"BUSY_RETRY_AFTER", "0"}, + {"GPU_FOREIGN_VRAM_MB", "-1"}, + {"GAME_POLL_INTERVAL", "bogus"}, } { _, err := Load(func(k string) string { if k == tc.key { @@ -172,3 +174,51 @@ func TestLoadErrors(t *testing.T) { } } } + +func TestGameDetectionSettings(t *testing.T) { + cfg, err := Load(func(k string) string { + switch k { + case "OLLAMA_URL": + return "http://127.0.0.1:11435" + case "GAME_PROCS": + return " cyberpunk2077.exe, hl2.exe ,, " + case "GPU_FOREIGN_VRAM_MB": + return "1024" + case "GPU_IGNORE_PROCS": + return "ollama, my-trainer" + } + return "" + }) + if err != nil { + t.Fatal(err) + } + if len(cfg.GameProcs) != 2 || cfg.GameProcs[0] != "cyberpunk2077.exe" || cfg.GameProcs[1] != "hl2.exe" { + t.Fatalf("GameProcs = %v", cfg.GameProcs) + } + if cfg.GPUForeignVRAMMB != 1024 { + t.Fatalf("GPUForeignVRAMMB = %d", cfg.GPUForeignVRAMMB) + } + if len(cfg.GPUIgnoreProcs) != 2 || cfg.GPUIgnoreProcs[1] != "my-trainer" { + t.Fatalf("GPUIgnoreProcs = %v", cfg.GPUIgnoreProcs) + } + if cfg.GamePollInterval != 5*time.Second { + t.Fatalf("GamePollInterval = %v, want 5s default", cfg.GamePollInterval) + } + + // Defaults: both detection paths off, ignore list covers our consumers. + cfg, err = Load(func(k string) string { + if k == "OLLAMA_URL" { + return "http://127.0.0.1:11435" + } + return "" + }) + if err != nil { + t.Fatal(err) + } + if len(cfg.GameProcs) != 0 || cfg.GPUForeignVRAMMB != 0 { + t.Fatalf("detection must be off by default: %v %d", cfg.GameProcs, cfg.GPUForeignVRAMMB) + } + if len(cfg.GPUIgnoreProcs) == 0 { + t.Fatal("GPUIgnoreProcs default must not be empty") + } +} diff --git a/internal/config/sample.go b/internal/config/sample.go index be1e096..92e0223 100644 --- a/internal/config/sample.go +++ b/internal/config/sample.go @@ -32,6 +32,10 @@ func sampleEntries(logFile string) []sampleEntry { {"COMFY_DIR", `C:\ComfyUI`, "Working directory for COMFY_CMD (default: empty = inherit)", false}, {"COMFY_IDLE_TIMEOUT", "5m", "Stop the managed ComfyUI after this long without requests or jobs (frees VRAM)", false}, {"COMFY_START_TIMEOUT", "2m", "How long a request waits for the managed ComfyUI to come up", false}, + {"GAME_PROCS", "cyberpunk2077.exe,hl2.exe", "While a listed process runs, the GPU counts as held by it: requests wait, Ollama unloads, managed ComfyUI stops (default: empty = disabled)", false}, + {"GPU_FOREIGN_VRAM_MB", "1024", "Also treat the GPU as held when a process not in GPU_IGNORE_PROCS uses more VRAM than this (needs nvidia-smi; 0/empty = disabled)", false}, + {"GPU_IGNORE_PROCS", "ollama,ollama app,ollama_llama_server,python,pythonw", "Process names never counted as foreign GPU users (ComfyUI runs under python)", false}, + {"GAME_POLL_INTERVAL", "5s", "How often game/VRAM detection runs", false}, {"UNLOAD_TIMEOUT", "60s", "How long to wait for Ollama to unload a model", false}, {"JOB_TIMEOUT", "15m", "Maximum time to wait for a ComfyUI job", false}, {"LLM_WAIT_TIMEOUT", "10m", "Max time an LLM request waits for the GPU before being answered 503 (wait mode)", false}, diff --git a/internal/game/game.go b/internal/game/game.go new file mode 100644 index 0000000..b425cbf --- /dev/null +++ b/internal/game/game.go @@ -0,0 +1,159 @@ +// Package game detects processes outside gpu-turnstile's control that hold +// the GPU — typically a game — so the proxy can block new GPU work and free +// VRAM while they run. Two detection paths: an explicit process watch list +// (GAME_PROCS) and a foreign-VRAM threshold via nvidia-smi +// (GPU_FOREIGN_VRAM_MB) that catches anything not on the ignore list. +package game + +import ( + "context" + "errors" + "fmt" + "log/slog" + "os/exec" + "strconv" + "strings" +) + +// Process is one running OS process. +type Process struct { + PID int + Name string +} + +// computeApp is one process holding GPU memory, as reported by nvidia-smi. +type computeApp struct { + PID int + UsedMB int +} + +// Detector checks whether a foreign process holds the GPU. The zero value +// (no watch list, no threshold) never detects anything; main only starts the +// poll loop when at least one path is configured. +type Detector struct { + procs map[string]bool // normalized names from GAME_PROCS + vramMB int // foreign VRAM threshold; 0 = disabled + ignore map[string]bool // normalized names never counted as foreign + log *slog.Logger + noNvidia bool // nvidia-smi was not found; VRAM path disabled for good +} + +// New builds a Detector from the configured watch list, VRAM threshold in +// MiB (0 disables the nvidia-smi path) and ignore list. Names are matched +// case-insensitively, with or without a trailing ".exe". +func New(procs []string, vramMB int, ignore []string, log *slog.Logger) *Detector { + if log == nil { + log = slog.Default() + } + return &Detector{ + procs: nameSet(procs), + vramMB: vramMB, + ignore: nameSet(ignore), + log: log, + } +} + +// normName lowercases a process name and strips a trailing ".exe" so the +// watch and ignore lists match on Windows and Linux spellings alike. +func normName(s string) string { + return strings.TrimSuffix(strings.ToLower(strings.TrimSpace(s)), ".exe") +} + +func nameSet(names []string) map[string]bool { + set := make(map[string]bool, len(names)) + for _, n := range names { + if n = normName(n); n != "" { + set[n] = true + } + } + return set +} + +// Check looks once for foreign GPU holders and returns a human-readable +// description of each (empty when the GPU is free for gpu-turnstile's +// consumers). A failing nvidia-smi call is returned as an error only when +// the process list found nothing; a missing nvidia-smi binary disables the +// VRAM path permanently (logged once). +func (d *Detector) Check(ctx context.Context) ([]string, error) { + ps, psErr := processes() + if d.vramMB <= 0 || d.noNvidia { + return d.detect(ps, nil), psErr + } + apps, err := queryComputeApps(ctx) + if errors.Is(err, exec.ErrNotFound) { + d.noNvidia = true + d.log.Warn("GPU_FOREIGN_VRAM_MB is set but nvidia-smi was not found; VRAM detection disabled") + return d.detect(ps, nil), nil + } + if err != nil { + return d.detect(ps, nil), err + } + return d.detect(ps, apps), nil +} + +// detect is the pure core of Check: given the process table and (optionally) +// the nvidia-smi compute-apps list, it returns the foreign holders. +func (d *Detector) detect(ps []Process, apps []computeApp) []string { + var holders []string + for _, p := range ps { + if d.procs[normName(p.Name)] { + holders = append(holders, fmt.Sprintf("%s (pid %d)", p.Name, p.PID)) + } + } + if d.vramMB > 0 && apps != nil { + names := make(map[int]string, len(ps)) + for _, p := range ps { + names[p.PID] = p.Name + } + for _, a := range apps { + name := names[a.PID] + if d.ignore[normName(name)] || a.UsedMB < d.vramMB { + continue + } + if name == "" { + name = "unknown process" + } + holders = append(holders, fmt.Sprintf("%s (pid %d) using %d MiB VRAM", name, a.PID, a.UsedMB)) + } + } + return holders +} + +// queryComputeApps runs nvidia-smi and parses the per-process VRAM list. +// Note: under Windows' WDDM driver, nvidia-smi only sees compute +// allocations, so graphics-only games may not appear there — GAME_PROCS is +// the reliable path on Windows; on Linux both work. +func queryComputeApps(ctx context.Context) ([]computeApp, error) { + out, err := exec.CommandContext(ctx, "nvidia-smi", + "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits").Output() + if err != nil { + return nil, err + } + return parseComputeApps(string(out)) +} + +// parseComputeApps parses "pid, used_memory" CSV lines (no header, MiB +// units). Unsupported rows ("N/A" on WDDM) are skipped. +func parseComputeApps(out string) ([]computeApp, error) { + var apps []computeApp + for _, line := range strings.Split(out, "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + pidStr, memStr, ok := strings.Cut(line, ",") + if !ok { + return nil, fmt.Errorf("nvidia-smi: unexpected line %q", line) + } + pid, err := strconv.Atoi(strings.TrimSpace(pidStr)) + if err != nil { + return nil, fmt.Errorf("nvidia-smi: unexpected pid in %q", line) + } + mem, err := strconv.Atoi(strings.TrimSpace(memStr)) + if err != nil { + continue // "N/A" and friends: unsupported under WDDM + } + apps = append(apps, computeApp{PID: pid, UsedMB: mem}) + } + return apps, nil +} diff --git a/internal/game/game_test.go b/internal/game/game_test.go new file mode 100644 index 0000000..3325c38 --- /dev/null +++ b/internal/game/game_test.go @@ -0,0 +1,96 @@ +package game + +import ( + "runtime" + "slices" + "testing" +) + +func TestNormName(t *testing.T) { + for in, want := range map[string]string{ + "Cyberpunk2077.exe": "cyberpunk2077", + "ollama": "ollama", + "OLLAMA APP.EXE": "ollama app", + " python.exe ": "python", + "hl2": "hl2", + } { + if got := normName(in); got != want { + t.Errorf("normName(%q) = %q, want %q", in, got, want) + } + } +} + +func TestParseComputeApps(t *testing.T) { + apps, err := parseComputeApps("1234, 512\n 42 , 8192 \n\n") + if err != nil { + t.Fatal(err) + } + want := []computeApp{{PID: 1234, UsedMB: 512}, {PID: 42, UsedMB: 8192}} + if !slices.Equal(apps, want) { + t.Errorf("got %+v, want %+v", apps, want) + } + + // WDDM reports "N/A" for memory; those rows are skipped, not fatal. + apps, err = parseComputeApps("1234, N/A\n42, 1024\n") + if err != nil { + t.Fatal(err) + } + if !slices.Equal(apps, []computeApp{{PID: 42, UsedMB: 1024}}) { + t.Errorf("got %+v", apps) + } + + if _, err := parseComputeApps("garbage\n"); err == nil { + t.Error("expected an error for a malformed line") + } +} + +func TestDetect(t *testing.T) { + d := New([]string{"Cyberpunk2077.exe", "hl2"}, 1024, + []string{"ollama", "python", "pythonw"}, nil) + ps := []Process{ + {PID: 10, Name: "ollama.exe"}, + {PID: 20, Name: "python.exe"}, + {PID: 30, Name: "cyberpunk2077.exe"}, + } + apps := []computeApp{ + {PID: 10, UsedMB: 8192}, // ignored: ollama + {PID: 20, UsedMB: 4096}, // ignored: python (ComfyUI) + {PID: 40, UsedMB: 2048}, // foreign, above threshold + {PID: 50, UsedMB: 100}, // foreign but below threshold + } + holders := d.detect(ps, apps) + if len(holders) != 2 { + t.Fatalf("got %v, want 2 holders", holders) + } + if holders[0] != "cyberpunk2077.exe (pid 30)" { + t.Errorf("holders[0] = %q", holders[0]) + } + if holders[1] != "unknown process (pid 40) using 2048 MiB VRAM" { + t.Errorf("holders[1] = %q", holders[1]) + } +} + +func TestDetectNothingConfigured(t *testing.T) { + d := New(nil, 0, nil, nil) + if got := d.detect([]Process{{PID: 1, Name: "game.exe"}}, nil); len(got) != 0 { + t.Errorf("got %v, want none", got) + } +} + +func TestProcessesLive(t *testing.T) { + if runtime.GOOS != "windows" && runtime.GOOS != "linux" { + t.Skip("no process listing on this platform") + } + ps, err := processes() + if err != nil { + t.Fatal(err) + } + if len(ps) == 0 { + t.Fatal("no processes listed") + } + for _, p := range ps { + if p.Name == "" { + t.Errorf("pid %d has an empty name", p.PID) + } + } +} diff --git a/internal/game/procs_linux.go b/internal/game/procs_linux.go new file mode 100644 index 0000000..7af5226 --- /dev/null +++ b/internal/game/procs_linux.go @@ -0,0 +1,30 @@ +//go:build linux + +package game + +import ( + "os" + "strconv" + "strings" +) + +// processes lists the running processes from /proc//comm. +func processes() ([]Process, error) { + entries, err := os.ReadDir("/proc") + if err != nil { + return nil, err + } + var ps []Process + for _, e := range entries { + pid, err := strconv.Atoi(e.Name()) + if err != nil { + continue + } + comm, err := os.ReadFile("/proc/" + e.Name() + "/comm") + if err != nil { + continue // process vanished mid-walk + } + ps = append(ps, Process{PID: pid, Name: strings.TrimSpace(string(comm))}) + } + return ps, nil +} diff --git a/internal/game/procs_other.go b/internal/game/procs_other.go new file mode 100644 index 0000000..d307894 --- /dev/null +++ b/internal/game/procs_other.go @@ -0,0 +1,7 @@ +//go:build !windows && !linux + +package game + +// processes is unsupported on this platform; the process watch list never +// matches and the nvidia-smi path reports PIDs without names. +func processes() ([]Process, error) { return nil, nil } diff --git a/internal/game/procs_windows.go b/internal/game/procs_windows.go new file mode 100644 index 0000000..d5a96cf --- /dev/null +++ b/internal/game/procs_windows.go @@ -0,0 +1,35 @@ +//go:build windows + +package game + +import ( + "unsafe" + + "golang.org/x/sys/windows" +) + +// processes lists the running processes via the Toolhelp32 snapshot API. +func processes() ([]Process, error) { + h, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0) + if err != nil { + return nil, err + } + defer windows.CloseHandle(h) //nolint:errcheck // best effort + + var entry windows.ProcessEntry32 + entry.Size = uint32(unsafe.Sizeof(entry)) + if err := windows.Process32First(h, &entry); err != nil { + return nil, err + } + var ps []Process + for { + ps = append(ps, Process{ + PID: int(entry.ProcessID), + Name: windows.UTF16ToString(entry.ExeFile[:]), + }) + if err := windows.Process32Next(h, &entry); err != nil { + break // ERROR_NO_MORE_FILES ends the walk + } + } + return ps, nil +} diff --git a/internal/lock/lock.go b/internal/lock/lock.go index c98dd51..792c22e 100644 --- a/internal/lock/lock.go +++ b/internal/lock/lock.go @@ -1,6 +1,8 @@ // Package lock implements the two-mode GPU arbitration lock: any number of // concurrent LLM requests ("readers") or exactly one image job ("writer"), -// with image jobs taking priority over newly arriving LLM requests. +// with image jobs taking priority over newly arriving LLM requests. An +// external hold (SetExternal) blocks new grants of both kinds while a +// foreign process — e.g. a game — holds the GPU; in-flight work drains. package lock import ( @@ -13,9 +15,10 @@ import ( type State string const ( - StateIdle State = "idle" - StateLLM State = "llm" - StateImage State = "image" + StateIdle State = "idle" + StateLLM State = "llm" + StateImage State = "image" + StateExternal State = "external" ) type imageWaiter struct{ id uint64 } @@ -30,6 +33,7 @@ type Lock struct { imageActive bool // an image job holds the GPU imageQ []imageWaiter nextID uint64 + external string // non-empty: a foreign process (e.g. a game) holds the GPU log *slog.Logger } @@ -52,12 +56,41 @@ func (l *Lock) logTransition(msg string, args ...any) { } } +// SetExternal records that a process outside gpu-turnstile's control (a +// game, another ML job) holds the GPU: new LLM and image grants block until +// ClearExternal. In-flight work is not preempted. holder describes the +// process for logs and busy responses. +func (l *Lock) SetExternal(holder string) { + l.mu.Lock() + l.external = holder + l.broadcast() + l.mu.Unlock() + l.logTransition("lock transition", "state", StateExternal, "holder", holder) +} + +// ClearExternal lifts the external hold; waiting LLM and image requests +// proceed. +func (l *Lock) ClearExternal() { + l.mu.Lock() + l.external = "" + l.broadcast() + l.mu.Unlock() + l.logTransition("lock transition", "state", StateIdle) +} + +// External returns the current external holder, or "" when none. +func (l *Lock) External() string { + l.mu.Lock() + defer l.mu.Unlock() + return l.external +} + // AcquireLLM blocks until no image job is active or pending, then registers // one in-flight LLM request. Returns ctx.Err() if the context is cancelled // while waiting; no state is changed in that case. func (l *Lock) AcquireLLM(ctx context.Context) error { l.mu.Lock() - for l.imageActive || len(l.imageQ) > 0 { + for l.imageActive || len(l.imageQ) > 0 || l.external != "" { ch := l.change l.mu.Unlock() select { @@ -76,10 +109,10 @@ func (l *Lock) AcquireLLM(ctx context.Context) error { // TryAcquireLLM acquires one in-flight LLM slot without waiting and // reports whether it succeeded. It fails when an image job is active or -// pending. +// pending or an external hold is set. func (l *Lock) TryAcquireLLM() bool { l.mu.Lock() - if l.imageActive || len(l.imageQ) > 0 { + if l.imageActive || len(l.imageQ) > 0 || l.external != "" { l.mu.Unlock() return false } @@ -120,7 +153,7 @@ func (l *Lock) AcquireImage(ctx context.Context) error { for { l.mu.Lock() - if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive { + if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive && l.external == "" { l.imageQ = l.imageQ[1:] l.imageActive = true l.mu.Unlock() @@ -166,6 +199,8 @@ func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) { state = StateImage case l.n > 0: state = StateLLM + case l.external != "": + state = StateExternal default: state = StateIdle } diff --git a/internal/lock/lock_test.go b/internal/lock/lock_test.go index 7854e8f..aea1e97 100644 --- a/internal/lock/lock_test.go +++ b/internal/lock/lock_test.go @@ -226,3 +226,70 @@ func TestRace(t *testing.T) { t.Fatalf("leaked lock state: state=%s n=%d pending=%v", state, n, pending) } } + +func TestExternalHoldBlocksBoth(t *testing.T) { + lk := New(nil) + ctx := context.Background() + + lk.SetExternal("game.exe (pid 42)") + if lk.TryAcquireLLM() { + t.Fatal("TryAcquireLLM succeeded during external hold") + } + if got := lk.External(); got != "game.exe (pid 42)" { + t.Fatalf("External() = %q", got) + } + if state, _, _ := lk.Snapshot(); state != StateExternal { + t.Fatalf("state=%s, want external", state) + } + + llmAcquired := make(chan struct{}) + go func() { + if err := lk.AcquireLLM(ctx); err != nil { + t.Error(err) + } + close(llmAcquired) + }() + assertBlocked(t, llmAcquired, "LLM acquire during external hold") + + imageAcquired := make(chan struct{}) + go func() { + if err := lk.AcquireImage(ctx); err != nil { + t.Error(err) + } + close(imageAcquired) + }() + assertBlocked(t, imageAcquired, "image acquire during external hold") + + lk.ClearExternal() + // The queued image job wins over the LLM waiter (image priority). + waitFor(t, imageAcquired, "image acquire after ClearExternal") + assertBlocked(t, llmAcquired, "LLM acquire while image active") + lk.ReleaseImage() + waitFor(t, llmAcquired, "LLM acquire after image release") + lk.ReleaseLLM() +} + +func TestExternalHoldDoesNotPreempt(t *testing.T) { + lk := New(nil) + ctx := context.Background() + if err := lk.AcquireLLM(ctx); err != nil { + t.Fatal(err) + } + + lk.SetExternal("game.exe") + // In-flight LLM work keeps the llm state; the hold blocks new grants. + if state, _, _ := lk.Snapshot(); state != StateLLM { + t.Fatalf("state=%s, want llm while work in flight", state) + } + if lk.TryAcquireLLM() { + t.Fatal("TryAcquireLLM succeeded during external hold") + } + lk.ReleaseLLM() + if state, _, _ := lk.Snapshot(); state != StateExternal { + t.Fatalf("state=%s, want external after drain", state) + } + lk.ClearExternal() + if state, _, _ := lk.Snapshot(); state != StateIdle { + t.Fatalf("state=%s, want idle", state) + } +} diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index 8305beb..389e6a0 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -89,7 +89,7 @@ func (m *Metrics) Render(w io.Writer, state string, llmInflight int, imagePendin fmt.Fprint(w, `# HELP gpu_turnstile_state Current GPU state (1 for the active state). # TYPE gpu_turnstile_state gauge `) - for _, s := range []string{"idle", "llm", "image"} { + for _, s := range []string{"idle", "llm", "image", "external"} { v := 0 if s == state { v = 1 diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 05eb807..f481921 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -464,10 +464,14 @@ func (s *Server) OllamaHandler() http.Handler { if s.busyMode == "reject" { if !s.cfg.Lock.TryAcquireLLM() { s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds()) + msg := "GPU busy: image job active or queued" + if holder := s.cfg.Lock.External(); holder != "" { + msg = "GPU busy: " + holder + } 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) + http.Error(w, msg, s.busyStatus) return } s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds()) @@ -508,7 +512,14 @@ func (s *Server) ComfyHandler() http.Handler { } if s.cfg.ComfySup != nil { // Any other ComfyUI request also wakes the managed server; the - // retry backoff bridges the time it needs to come up. + // retry backoff bridges the time it needs to come up. While a + // foreign process holds the GPU we refuse to spawn it — the + // request gets the busy answer instead of fighting for VRAM. + if holder := s.cfg.Lock.External(); holder != "" && !s.cfg.ComfySup.Running() { + w.Header().Set("Retry-After", strconv.Itoa(s.busyRetryAfter)) + http.Error(w, "GPU busy: "+holder, http.StatusServiceUnavailable) + return + } if err := s.cfg.ComfySup.EnsureRunning(); err != nil { http.Error(w, fmt.Sprintf("cannot start ComfyUI: %v", err), http.StatusBadGateway) return @@ -553,10 +564,13 @@ func (w *captureWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter } func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) { log := s.log.With("op", "image") - if s.cfg.ComfySup != nil { - // Bring the managed server up *before* taking the GPU lock: torch - // can take a minute to load, and LLM traffic should keep flowing - // in the meantime. + // Normally the managed server is brought up *before* taking the GPU + // lock: torch can take a minute to load, and LLM traffic should keep + // flowing in the meantime. While a foreign process holds the GPU, LLM + // traffic is blocked anyway and a fresh ComfyUI would fight it for + // VRAM — so the lock comes first in that case. + comfyFirst := s.cfg.ComfySup != nil && s.cfg.Lock.External() == "" + if comfyFirst { if err := s.cfg.ComfySup.EnsureRunning(); err != nil { http.Error(w, fmt.Sprintf("cannot start ComfyUI: %v", err), http.StatusBadGateway) return @@ -577,6 +591,19 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) { s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds()) log.Info("image lock acquired") + if s.cfg.ComfySup != nil && !comfyFirst { + if err := s.cfg.ComfySup.EnsureRunning(); err != nil { + s.cfg.Lock.ReleaseImage() + http.Error(w, fmt.Sprintf("cannot start ComfyUI: %v", err), http.StatusBadGateway) + return + } + if err := s.cfg.ComfySup.WaitReady(r.Context()); err != nil { + s.cfg.Lock.ReleaseImage() + http.Error(w, fmt.Sprintf("ComfyUI did not become ready: %v", err), http.StatusBadGateway) + return + } + } + if s.cfg.Ollama != nil { uctx, ucancel := context.WithTimeout(r.Context(), s.cfg.UnloadTimeout) elapsed, uerr := s.cfg.Ollama.UnloadAll(uctx)