Compare commits

..
25 Commits
Author SHA1 Message Date
mram 31ee6013db Pin compose example to v0.3.2
ci / test (push) Successful in 16s
ci / docker (push) Successful in 1m8s
ci / release (push) Successful in 16s
2026-09-22 18:55:10 +02:00
mram 5506163493 Monitor: show the web UI URL when LISTEN_UI is active
The status snapshot carries the UI address; the monitor renders it as a
dim 'UI: http://…' line below the GPU line.
2026-09-22 18:54:52 +02:00
mram 18a7336c02 Web UI: optional HTTP basic auth via UI_USER/UI_PASS
Both must be set together and require LISTEN_UI; constant-time compares,
challenge on every endpoint. Startup warns when LISTEN_UI binds a
non-loopback address without auth. Bind address and port were already
covered by LISTEN_UI itself (host:port).
2026-09-22 18:53:00 +02:00
mram 0796092b80 Add optional web UI mirroring --monitor (LISTEN_UI, off by default)
A single self-contained page polls /api/status once a second and renders
the same lines as the terminal monitor (downstreams with models/VRAM and
busy markers, lock, queue, GPU stats); the u/r buttons and keys trigger
update-now and reload-env through the same rate-limited command handler
the control pipe uses. Unauthenticated by design — the sample config
tells users to keep it on localhost.
2026-09-22 18:48:28 +02:00
mram 897163042c Pin compose example to v0.3.1
ci / test (push) Successful in 15s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 15s
2026-09-22 17:08:40 +02:00
mram e58ff33912 Monitor: show GPU temperature and fan speed
Extends the existing per-tick nvidia-smi query with temperature.gpu and
fan.speed (no extra call); N/A values (cards without fan telemetry) are
simply omitted from the GPU line.
2026-09-22 16:56:01 +02:00
mram 2040b30c94 Monitor: add r hotkey for config reload
Generalizes the one-shot ask behind u into a shared helper; r sends
reload-env and reports the service's reply (unchanged / restarting with
the changed setting names / invalid config) in the note line. Footer now
lists q, u, r.
2026-09-22 13:33:46 +02:00
mram a07b6726ea Pin compose example to v0.3.0
ci / test (push) Successful in 15s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 15s
2026-09-22 13:31:05 +02:00
mram d2c49e52fa Monitor: show Ollama's loaded models, their VRAM footprint, and busy state
The status snapshot now queries /api/ps (2s timeout so a wedged Ollama
cannot stall the channel) and the ollama line renders e.g.
'UP · llama3.1:8b (4.8 GiB VRAM) · busy'; comfy gets the busy marker
too while an image job runs.
2026-09-22 12:44:04 +02:00
mram e98331bb9e Pin compose example to v0.2.9
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m8s
ci / release (push) Successful in 16s
2026-09-22 11:23:17 +02:00
mram 54174787d5 Fix control-pipe client race: retry CreateFile on ERROR_PIPE_BUSY
WaitNamedPipe can report an instance that a concurrent client (the
monitor polls status every second) grabs before our CreateFile runs;
the single-attempt Ask then failed with 'no answer from the service'.
Retry the wait+open until a 5s overall deadline.
2026-09-22 11:22:27 +02:00
mram d259b2e96c Add GPU_FOREIGN_UTIL_PCT: game detection via per-process GPU 3D-engine usage
Reads the same PDH counters as Task Manager (\GPU Engine(*)\Utilization
Percentage, locale-independent via PdhAddEnglishCounterW), which cover
graphics work under WDDM — games are caught without an exe list and the
offender is named. Windows-only; a missing or failing counter disables
the path with one log line. dwm (the desktop compositor) joins the
default ignore list.
2026-09-22 10:53:29 +02:00
mram 8fccd333aa Env names in reload diff; monitor shows last VRAM check result; harden control channel against floods
- diffConfig reports the env var names (from new struct tags) instead of
  Go field names, so the reload-env reply names what the user can change
- the monitor's GPU line now shows the game detector's last finding
  (external holders or none) and how long ago the check ran
- control channel: 10s per-connection watchdog (abortive force-close),
  cap of 32 concurrent connections, reload-env rate-limited; command
  read was already capped at 4 KiB
2026-09-22 09:24:27 +02:00
mram 4bd5f34ce7 Pin compose example to v0.2.8
ci / test (push) Successful in 15s
ci / docker (push) Successful in 1m6s
ci / release (push) Successful in 15s
2026-09-22 09:10:14 +02:00
mram b9d3f91403 Monitor: hotkeys (q quit, u update now) with footer line; show GPU VRAM usage 2026-09-22 09:10:14 +02:00
mram 98d2d714ca Add --reload-env: service re-reads and validates its config, restarts when idle if it changed 2026-09-22 09:02:25 +02:00
mram 7a6409831c Pin compose example to v0.2.7
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 16s
2026-09-22 08:47:47 +02:00
mram c7b8747d85 Fix control pipe accept loop: create instances only after a client connects; CLI force-update never touches the service log file 2026-09-22 08:43:58 +02:00
mram 09d4e81eab Pin compose example to v0.2.6
ci / test (push) Successful in 15s
ci / docker (push) Successful in 1m6s
ci / release (push) Successful in 15s
2026-09-22 08:36:36 +02:00
mram 1264afe40e Monitor restarts itself when the service updated its binary on disk 2026-09-22 08:34:47 +02:00
mram c95d401a36 Pin compose example to v0.2.5
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m6s
ci / release (push) Successful in 16s
2026-09-22 08:32:06 +02:00
mram f7ea30a494 Fix control pipe races: treat ERROR_PIPE_CONNECTED as success, flush+disconnect before close so replies are never discarded 2026-09-22 08:31:13 +02:00
mram 620be9d57c Add short flags: -i (install), -r (remove), -m (monitor) 2026-09-22 08:23:40 +02:00
mram cffb4151a6 Pin compose example to v0.2.4
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m9s
ci / release (push) Successful in 15s
2026-09-22 08:19:13 +02:00
mram 16dc7fe462 Add --monitor: live status view (downstream health, GPU lock, queue) via the control channel 2026-09-22 08:17:05 +02:00
27 changed files with 2215 additions and 184 deletions
+19 -11
View File
@@ -34,8 +34,8 @@ 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. 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).
detection of foreign GPU holders such as games — is enabled by `GAME_PROCS`,
`GPU_FOREIGN_VRAM_MB` and/or `GPU_FOREIGN_UTIL_PCT` (see below).
## Configuration
@@ -49,6 +49,8 @@ override file values. Invalid values fail at startup.
|---|---|---|
| `LISTEN_OLLAMA` | `:11434` | Listener for Ollama-compatible clients |
| `LISTEN_COMFY` | `:8188` | Listener for ComfyUI clients |
| `LISTEN_UI` | _(empty = disabled)_ | Web UI listener (`host:port`) mirroring `--monitor` (live status, update/reload buttons). Keep it on `127.0.0.1` unless you set auth |
| `UI_USER` / `UI_PASS` | _(empty = no auth)_ | HTTP basic auth for the web UI; must be set together |
| `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer |
| `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer |
| `UNLOAD_TIMEOUT` | `60s` | Wait for Ollama to unload before an image job |
@@ -64,7 +66,8 @@ override file values. Invalid values fail at startup.
| `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 |
| `GPU_FOREIGN_UTIL_PCT` | `0` (disabled) | Also treat the GPU as held when a non-ignored process uses more than this percent of the GPU 3D engine (Windows PDH counters — what Task Manager shows; catches games without an exe list) |
| `GPU_IGNORE_PROCS` | `ollama,ollama app,ollama_llama_server,python,pythonw,dwm` | Process names never counted as foreign GPU users |
| `GAME_POLL_INTERVAL` | `15s` | How often game/VRAM detection runs (don't go below ~10s — nvidia-smi polls keep the GPU awake) |
| `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 |
@@ -148,20 +151,25 @@ adds a `BindPaths=` to the systemd unit on Linux. Re-run it after changing
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
models and stop the managed ComfyUI so the game gets the memory. Three
detection paths, each optional, polled every `GAME_POLL_INTERVAL` (15s):
```
GAME_PROCS=cyberpunk2077.exe,bg3.exe # the reliable way on Windows
GPU_FOREIGN_UTIL_PCT=30 # zero-config on Windows: any non-ignored
# process using >30% of the GPU 3D engine
GAME_PROCS=cyberpunk2077.exe,bg3.exe # explicit exe watch list
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
`GPU_FOREIGN_UTIL_PCT` reads the same per-process GPU engine counters as
Task Manager (PDH), which cover graphics work under Windows' WDDM driver —
so games show up without naming them, and the offender is named in the log
and monitor. Windows-only. `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 on
Linux, but under WDDM graphics-only games may not show up in nvidia-smi's
per-process list, so on Windows prefer `GPU_FOREIGN_UTIL_PCT`. When the
game exits, requests resume automatically.
## Build and run
+21 -12
View File
@@ -55,9 +55,9 @@ consumer gets no listener, no startup probe, and no lock participation:
- **Only `COMFY_URL`**: image jobs are tracked and ComfyUI's VRAM is freed
afterwards, but the Ollama unload and warm-reload steps are skipped.
- **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.
`GAME_PROCS`, `GPU_FOREIGN_VRAM_MB` and/or `GPU_FOREIGN_UTIL_PCT` it
watches for foreign processes holding the GPU (see below) and plugs into
the same lock the same way — excluded when all knobs are unset.
### Lock semantics
@@ -171,20 +171,26 @@ files are flagged in the startup log.
## 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 15 s); either one being
nothing proxies through gpu-turnstile for them. Three independent detection
paths, polled every `GAME_POLL_INTERVAL` (default 15 s); any 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 3D-engine utilization** (`GPU_FOREIGN_UTIL_PCT`, Windows only):
per-process GPU engine counters from PDH — the same data Task Manager
shows — cover graphics work under WDDM, so any process not in
`GPU_IGNORE_PROCS` using more than the threshold percent of the 3D engine
counts as a foreign holder, with no exe list needed. Values are summed per
process across engines; a missing/failing PDH counter disables the path
(logged once).
- **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.
`GPU_IGNORE_PROCS` (default: Ollama and python — ComfyUI runs under python
— plus dwm, the desktop compositor) 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:
@@ -215,6 +221,8 @@ override file values. A missing file is fine; a malformed one is fatal.
|---|---|---|
| `LISTEN_OLLAMA` | `:11434` | listener for Ollama-compatible clients |
| `LISTEN_COMFY` | `:8188` | listener for ComfyUI clients |
| `LISTEN_UI` | _(empty = disabled)_ | web UI listener (`host:port`) mirroring `--monitor` (live status, update/reload buttons); keep it on `127.0.0.1` unless auth is set |
| `UI_USER` / `UI_PASS` | _(empty = no auth)_ | HTTP basic auth for the web UI; must be set together |
| `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer |
| `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer |
| `UNLOAD_TIMEOUT` | `60s` | wait for Ollama to unload |
@@ -230,7 +238,8 @@ override file values. A missing file is fine; a malformed one is fatal.
| `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 |
| `GPU_FOREIGN_UTIL_PCT` | `0` (disabled) | also treat the GPU as held when a process not in `GPU_IGNORE_PROCS` uses more than this percent of the GPU 3D engine (Windows PDH counters, as shown by Task Manager; catches games without an exe list) |
| `GPU_IGNORE_PROCS` | `ollama,ollama app,ollama_llama_server,python,pythonw,dwm` | process names never counted as foreign GPU users |
| `GAME_POLL_INTERVAL` | `15s` | how often game/VRAM detection runs (nvidia-smi polls keep the GPU awake; don't go below ~10s) |
| `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 |
+383 -50
View File
@@ -5,6 +5,7 @@ package main
import (
"bufio"
"context"
"encoding/json"
"errors"
"fmt"
"io"
@@ -15,6 +16,7 @@ import (
"os"
"os/signal"
"path/filepath"
"reflect"
"runtime"
"strings"
"sync"
@@ -38,7 +40,8 @@ import (
var version = "dev"
// exitCodeUpdate tells the service recovery configuration to restart the
// process: a signed update has been staged and the GPU lock is idle.
// process: a signed update has been staged or a changed config reload was
// requested, and the GPU lock is idle.
const exitCodeUpdate = 3
// exitCodeStaged is returned by an elevated --force-update child when it
@@ -58,7 +61,7 @@ func stdoutIsTerminal() bool {
// --install-service / --remove-service switches, --no-copy, -h/--help,
// -v/--version, --force-update and the hidden --elevated-child marker from
// args.
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, elevatedChild bool, rest []string) {
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild bool, rest []string) {
rest = args[:0]
for i := 0; i < len(args); i++ {
switch {
@@ -67,9 +70,9 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
i++
case strings.HasPrefix(args[i], "-config="):
configPath = strings.TrimPrefix(args[i], "-config=")
case args[i] == "--install-service" || args[i] == "-install-service":
case args[i] == "--install-service" || args[i] == "-install-service" || args[i] == "-i":
install = true
case args[i] == "--remove-service" || args[i] == "-remove-service":
case args[i] == "--remove-service" || args[i] == "-remove-service" || args[i] == "-r":
remove = true
case args[i] == "--no-copy" || args[i] == "-no-copy":
noCopy = true
@@ -81,13 +84,17 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
forceUpdate = true
case args[i] == "--update-now" || args[i] == "-update-now":
updateNow = true
case args[i] == "--reload-env" || args[i] == "-reload-env":
reloadEnv = true
case args[i] == "--monitor" || args[i] == "-monitor" || args[i] == "-m":
monitor = true
case args[i] == "--elevated-child":
elevatedChild = true
default:
rest = append(rest, args[i])
}
}
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, elevatedChild, rest
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, rest
}
// versionLine is printed at the top of every help and error screen.
@@ -97,14 +104,20 @@ const usageText = `GPU arbitration proxy for Ollama + ComfyUI
Usage:
gpu-turnstile -config <path> run the proxy
gpu-turnstile --install-service [--no-copy] [-config path] install + start as a service
gpu-turnstile --remove-service stop + uninstall the service
gpu-turnstile -i | --install-service [--no-copy] [-config path]
install + start as a service
gpu-turnstile -r | --remove-service stop + uninstall the service
gpu-turnstile -v | --version print just the version
gpu-turnstile --force-update check for a signed update now,
apply it and restart the service
(no admin needed when the service runs)
gpu-turnstile --update-now like --force-update, but only
through the running service
gpu-turnstile --reload-env make the service re-read and
validate its config file, then
restart onto it if it changed
gpu-turnstile -m | --monitor live status view (downstreams,
GPU lock, queue); Ctrl+C quits
gpu-turnstile -h | --help this help
Options:
@@ -135,12 +148,12 @@ func fatalUsage(format string, args ...any) {
}
func main() {
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, elevatedChild, args := parseFlags(os.Args[1:])
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, args := parseFlags(os.Args[1:])
if showVersion {
fmt.Println(version)
return
}
bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !elevatedChild && len(args) == 0
bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !reloadEnv && !monitor && !elevatedChild && len(args) == 0
if help || (bare && stdoutIsTerminal()) {
// Bare invocation in a terminal (e.g. double-clicked on Windows)
// shows the help instead of starting a proxy window with no visible
@@ -165,6 +178,10 @@ func main() {
fatalUsage("error: --force-update cannot be combined with --install-service/--remove-service")
case updateNow && (install || remove || forceUpdate):
fatalUsage("error: --update-now cannot be combined with other commands")
case reloadEnv && (install || remove || forceUpdate || updateNow):
fatalUsage("error: --reload-env cannot be combined with other commands")
case monitor && (install || remove || forceUpdate || updateNow || reloadEnv):
fatalUsage("error: --monitor cannot be combined with other commands")
case install:
os.Exit(serviceCommand(configPath, true, noCopy, elevatedChild))
case remove:
@@ -173,6 +190,10 @@ func main() {
os.Exit(forceUpdateCommand(configPath, elevatedChild))
case updateNow:
os.Exit(updateNowCommand())
case reloadEnv:
os.Exit(reloadEnvCommand())
case monitor:
os.Exit(monitorCommand())
}
if len(args) > 0 {
fatalUsage("error: unknown arguments: %s", strings.Join(args, " "))
@@ -192,7 +213,7 @@ func main() {
syncEnvFile(resolveConfigPath(configPath), cfg.LogFile, log)
if service.IsService() {
if err := service.Run(func(ctx context.Context) error { return run(ctx, cfg, log, logOut, true) }); err != nil {
if err := service.Run(func(ctx context.Context) error { return run(ctx, cfg, log, logOut, true, configPath) }); err != nil {
log.Error("service failed", "err", err)
os.Exit(1)
}
@@ -200,7 +221,7 @@ func main() {
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
if err := run(ctx, cfg, log, logOut, false); err != nil {
if err := run(ctx, cfg, log, logOut, false, configPath); err != nil {
log.Error("listener failed", "err", err)
os.Exit(1)
}
@@ -352,6 +373,10 @@ func forceUpdateCommand(configPath string, elevatedChild bool) int {
fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot locate executable: %v\n", err)
return 1
}
// One-shot CLI: the updater logs to stderr, never to the service's
// LOG_FILE — that file is ACL'd to the service account, and a CLI run
// has nothing worth persisting there.
cfg.LogFile = ""
log, _, logCloser := newLogger(cfg)
defer logCloser.Close()
if cfg.AppVersion == "dev" {
@@ -437,6 +462,18 @@ func updateNowCommand() int {
return printControlReply(reply)
}
// reloadEnvCommand asks the running service, over the control channel, to
// re-read and validate its config file. The client sends no path — the
// service only ever re-reads its own configured file.
func reloadEnvCommand() int {
reply, err := control.Ask(control.CmdReloadEnv)
if err != nil {
fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: no running service to ask\n", versionLine())
return 1
}
return printControlReply(reply)
}
// reportElevatedUpdate prints the parent's summary of an elevated
// --force-update child: exitCodeStaged means the child staged a new binary,
// 0 means it found nothing to do. to is the tag the parent's own check
@@ -517,7 +554,7 @@ func managedComfyCommand(cfg config.Config) string {
return ""
}
func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Writer, isService bool) error {
func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Writer, isService bool, configPath string) error {
// ComfyUI can run as a managed child — COMFY_CMD verbatim, or the
// standard venv layout derived from COMFY_DIR alone: started on demand
// by the proxy, stopped after COMFY_IDLE_TIMEOUT idle (and on shutdown)
@@ -530,6 +567,8 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
"version", version,
"listen_ollama", cfg.ListenOllama,
"listen_comfy", cfg.ListenComfy,
"listen_ui", orDisabled(cfg.ListenUI),
"ui_auth", cfg.UIUser != "",
"ollama_url", orDisabled(cfg.OllamaURL),
"comfy_url", orDisabled(cfg.ComfyURL),
"unload_timeout", cfg.UnloadTimeout,
@@ -555,6 +594,7 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
"comfy_start_timeout", cfg.ComfyStartTimeout,
"game_procs", cfg.GameProcs,
"gpu_foreign_vram_mb", cfg.GPUForeignVRAMMB,
"gpu_foreign_util_pct", cfg.GPUForeignUtilPct,
"gpu_ignore_procs", cfg.GPUIgnoreProcs,
"game_poll_interval", cfg.GamePollInterval,
"auto_update", cfg.AutoUpdate,
@@ -567,6 +607,8 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
)
lk := lock.New(log)
health := newHealthTracker()
started := time.Now()
// Each consumer is enabled by setting its URL; a disabled consumer gets
// no client, no listener and no probe.
var ollamaClient *ollama.Client
@@ -676,20 +718,25 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
}
probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout)
for name, probe := range probes {
if err := probe(probeCtx); err != nil && !errors.Is(err, errManagedDown) {
err := probe(probeCtx)
health.set(name, err == nil)
if err != nil && !errors.Is(err, errManagedDown) {
log.Warn(name+" probe failed", "err", err)
}
}
probeCancel()
if cfg.HealthInterval > 0 {
go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes)
go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes, health)
}
// 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)
// Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS,
// GPU_FOREIGN_VRAM_MB and/or GPU_FOREIGN_UTIL_PCT — hold the lock
// externally while they run. gw collects the VRAM reading and the last
// check result for the status channel.
gw := &gpuWatch{enabled: len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0 || cfg.GPUForeignUtilPct > 0}
if gw.enabled {
det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUForeignUtilPct, cfg.GPUIgnoreProcs, log)
go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup, gw)
}
// Bind the listeners up front so a port conflict fails fast and the
@@ -727,38 +774,65 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
service.NotifyReady()
service.StartWatchdog(ctx)
var u *update.Updater
var exePath string
// restartWhenIdle exits with exitCodeUpdate once the GPU lock is idle;
// the service recovery configuration brings the process back. Shared by
// staged updates and config reloads; the once guard makes repeated
// triggers idempotent.
var restartOnce sync.Once
restartWhenIdle := func(reason string) {
log.Warn("restarting once the GPU is idle", "reason", reason)
restartOnce.Do(func() {
go func() {
if waitForIdle(ctx, lk, 24*time.Hour) {
log.Warn("restarting now", "reason", reason)
os.Exit(exitCodeUpdate)
}
}()
})
}
applyStaged := func(to string) {}
if cfg.AutoUpdate {
exePath, err := os.Executable()
if err != nil {
if p, err := os.Executable(); err != nil {
log.Warn("auto-update disabled: cannot locate executable", "err", err)
} else {
u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
// applyStaged is shared by the hourly loop and the control
// channel; the once guard keeps a second trigger from
// double-waiting on the GPU lock.
var once sync.Once
applyStaged := func(to string) {
exePath = p
u = &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
applyStaged = func(to string) {
if !isService {
log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply", "version", to)
return
}
log.Warn("auto-update: staged; restarting once the GPU is idle", "version", to)
once.Do(func() {
go func() {
if waitForIdle(ctx, lk, 24*time.Hour) {
log.Warn("auto-update: restarting to apply update")
os.Exit(exitCodeUpdate)
}
}()
})
restartWhenIdle("update to " + to)
}
go updateLoop(ctx, cfg.UpdateInterval, log, u, exePath, applyStaged)
if isService {
serveControl(ctx, log, u, exePath, applyStaged)
}
}
}
// The command handler backs both the local control channel (served
// whenever running as a service) and the optional web UI (LISTEN_UI),
// independent of AUTO_UPDATE.
handler := controlHandler(ctx, u, exePath, applyStaged,
statusProvider(cfg, lk, comfySup, health, started, gw, ollamaClient),
reloadHandler(cfg, configPath, restartWhenIdle))
if isService {
serveControl(ctx, log, handler)
}
if cfg.ListenUI != "" {
ln, err := net.Listen("tcp", cfg.ListenUI)
if err != nil {
return fmt.Errorf("listen ui on %s: %w", cfg.ListenUI, err)
}
if host, _, err := net.SplitHostPort(cfg.ListenUI); err == nil && !isLoopbackHost(host) && cfg.UIUser == "" {
log.Warn("web UI is reachable from other machines WITHOUT auth; set UI_USER/UI_PASS or bind to 127.0.0.1", "addr", cfg.ListenUI)
}
uiSrv := &http.Server{Addr: cfg.ListenUI, Handler: uiHandler(handler, cfg.UIUser, cfg.UIPass)}
servers = append(servers, uiSrv)
log.Warn("listening", "consumer", "ui", "addr", cfg.ListenUI)
go func() { errCh <- uiSrv.Serve(ln) }()
}
select {
case err := <-errCh:
if err != nil && !errors.Is(err, http.ErrServerClosed) {
@@ -780,12 +854,52 @@ 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")
// gpuWatch records the latest VRAM reading and game-detector finding from
// the game detector's poll loop, for the status channel. Enabled is false
// when game detection is not configured (no polling happens then); Known is
// false until the first successful nvidia-smi reading.
type gpuWatch struct {
mu sync.Mutex
enabled bool
usedMB int
total int
tempC int
fanPct int
known bool
foreign string // last Check result: external holders, "" when none
at time.Time
}
func (g *gpuWatch) setVRAM(st game.GPUStats) {
g.mu.Lock()
g.usedMB, g.total = st.UsedMB, st.TotalMB
g.tempC, g.fanPct, g.known = st.TempC, st.FanPct, true
g.mu.Unlock()
}
func (g *gpuWatch) setCheck(foreign string) {
g.mu.Lock()
g.foreign, g.at = foreign, time.Now()
g.mu.Unlock()
}
func (g *gpuWatch) get() (st game.GPUStats, known bool, foreign string, ageS int64) {
g.mu.Lock()
defer g.mu.Unlock()
ageS = -1
if !g.at.IsZero() {
ageS = int64(time.Since(g.at).Seconds())
}
return game.GPUStats{UsedMB: g.usedMB, TotalMB: g.total, TempC: g.tempC, FanPct: g.fanPct},
g.known, g.foreign, ageS
}
// 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) {
func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *game.Detector, lk *lock.Lock, ollamaClient *ollama.Client, comfySup *supervise.Process, gw *gpuWatch) {
ticker := time.NewTicker(cfg.GamePollInterval)
defer ticker.Stop()
held, freed := false, false
@@ -799,6 +913,10 @@ func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *gam
if err != nil && ctx.Err() == nil {
log.Warn("game detection failed", "err", err)
}
gw.setCheck(summarizeHolders(holders))
if st, verr := game.QueryGPUStats(ctx); verr == nil {
gw.setVRAM(st)
}
switch {
case len(holders) > 0 && !held:
held = true
@@ -858,7 +976,8 @@ func freeVRAM(ctx context.Context, unloadTimeout time.Duration, log *slog.Logger
// transitions — "is DOWN" when a previously healthy upstream stops
// answering, "recovered" when it comes back. The first round only
// establishes the baseline; the startup probe already reported that state.
func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *slog.Logger, probes map[string]func(context.Context) error) {
// Every result goes into the tracker for the status channel.
func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *slog.Logger, probes map[string]func(context.Context) error, tracker *healthTracker) {
ticker := time.NewTicker(interval)
defer ticker.Stop()
up := map[string]bool{}
@@ -875,9 +994,11 @@ func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *
cancel()
if errors.Is(err, errManagedDown) {
managed[name] = true // intentionally stopped; not an outage
tracker.set(name, false)
continue
}
now := err == nil
tracker.set(name, now)
if managed[name] {
// First real probe after an idle stop only re-baselines —
// an on-demand start is not a "recovery".
@@ -919,18 +1040,39 @@ func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u
}
}
// serveControl opens the local control channel (named pipe on Windows,
// unix socket on Linux) so unprivileged local users can trigger an update
// check via --force-update without admin rights. The check payload is
// signature-verified regardless of who asks; triggers are rate-limited to
// one per minute so the channel cannot be used to spam restarts.
func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string)) {
// controlHandler builds the command handler shared by the control channel
// (named pipe / unix socket) and the web UI's action endpoints: status for
// --monitor / the UI, an update check (--force-update/--update-now) and a
// config reload (--reload-env), all safe for unprivileged local users. The
// update payload is signature-verified regardless of who asks; update
// triggers are rate-limited to one per minute and reloads to one per two
// seconds so neither can be used to spam restarts or disk churn. u is nil
// when AUTO_UPDATE=false.
func controlHandler(ctx context.Context, u *update.Updater, exePath string, applyStaged func(to string), status func() string, reload func() string) control.Handler {
var mu sync.Mutex
var lastTrigger time.Time
h := func(cmd string) string {
if cmd != control.CmdUpdateNow {
var lastTrigger, lastReload time.Time
return func(cmd string) string {
switch cmd {
case control.CmdStatus:
return "OK " + status()
case control.CmdReloadEnv:
// reload re-reads the config file from disk; a short limiter
// keeps a local flood from turning into disk churn.
mu.Lock()
if wait := 2*time.Second - time.Since(lastReload); wait > 0 {
mu.Unlock()
return fmt.Sprintf("ERR rate limited: retry in %ds", int(wait.Seconds())+1)
}
lastReload = time.Now()
mu.Unlock()
return reload()
case control.CmdUpdateNow:
default:
return "ERR unknown command: " + cmd
}
if u == nil {
return "ERR auto-update is disabled on this instance"
}
mu.Lock()
if wait := time.Minute - time.Since(lastTrigger); wait > 0 {
mu.Unlock()
@@ -950,11 +1092,202 @@ func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exeP
applyStaged(to)
return "OK updated from " + version + " to " + to + "; the service restarts once the GPU is idle"
}
}
// serveControl opens the local control channel (named pipe on Windows,
// unix socket on Linux) for unprivileged local users (--monitor,
// --force-update/--update-now, --reload-env).
func serveControl(ctx context.Context, log *slog.Logger, h control.Handler) {
if err := control.Serve(ctx, h, log); err != nil {
log.Warn("control channel disabled", "err", err)
}
}
// healthTracker records the latest probe result per upstream for the
// status channel.
type healthTracker struct {
mu sync.Mutex
up map[string]bool
}
func newHealthTracker() *healthTracker {
return &healthTracker{up: map[string]bool{}}
}
func (h *healthTracker) set(name string, up bool) {
h.mu.Lock()
h.up[name] = up
h.mu.Unlock()
}
func (h *healthTracker) get(name string) bool {
h.mu.Lock()
defer h.mu.Unlock()
return h.up[name]
}
// statusDownstream/statusLock/statusSnapshot are the JSON the control
// channel serves on CmdStatus; the monitor mode renders them.
type statusDownstream struct {
Name string `json:"name"`
URL string `json:"url"`
Up bool `json:"up"`
Managed string `json:"managed,omitempty"`
// Models lists Ollama's loaded models with their VRAM footprint. Null
// when unknown (query failed / not applicable); [] means none loaded —
// deliberately no omitempty so the two stay distinguishable.
Models []statusModel `json:"models"`
}
// statusModel is one loaded Ollama model.
type statusModel struct {
Name string `json:"name"`
VRAMMB int64 `json:"vram_mb"` // 0 = resident in RAM, not VRAM
}
type statusLock struct {
State string `json:"state"`
Detail string `json:"detail,omitempty"`
LLMInflight int `json:"llm_inflight"`
LLMWaiting int `json:"llm_waiting"`
ImageQueue int `json:"image_queue"`
External string `json:"external,omitempty"`
SinceS int64 `json:"since_s"`
}
type statusSnapshot struct {
Version string `json:"version"`
UptimeS int64 `json:"uptime_s"`
Downstreams []statusDownstream `json:"downstreams"`
Lock statusLock `json:"lock"`
// GPU carries the latest VRAM reading and detector finding; Enabled is
// false when game detection (and with it nvidia-smi polling) is not
// configured.
GPU statusGPU `json:"gpu"`
// UI is the web UI's URL when LISTEN_UI is active, empty otherwise.
UI string `json:"ui,omitempty"`
// MonitorNote is set client-side (never over the wire) when the
// monitor's own binary differs from the service's version.
MonitorNote string `json:"-"`
}
type statusGPU struct {
// Enabled reports whether game detection is configured (and with it
// VRAM polling); when false the other fields carry no information.
Enabled bool `json:"enabled"`
UsedMB int `json:"used_mb"`
TotalMB int `json:"total_mb"`
// TempC/FanPct are -1 when unknown (never sampled or nvidia-smi
// reported N/A).
TempC int `json:"temp_c"`
FanPct int `json:"fan_pct"`
Known bool `json:"known"`
// Foreign is the last detector finding (external GPU holders), empty
// when the last check found none.
Foreign string `json:"foreign,omitempty"`
// AgeS is how long ago the last check ran; -1 before the first check.
AgeS int64 `json:"age_s"`
}
// reloadHandler re-reads and validates the service's config file for
// CmdReloadEnv. An invalid config is reported and the service keeps running
// untouched; a valid, changed config triggers a GPU-idle-gated restart onto
// it (same mechanism as staged updates); unchanged is a no-op.
func reloadHandler(current config.Config, configPath string, restartWhenIdle func(reason string)) func() string {
return func() string {
ncfg, err := loadMergedConfig(configPath)
if err != nil {
return "ERR config invalid: " + err.Error()
}
changed := diffConfig(current, ncfg)
if len(changed) == 0 {
return "OK config unchanged"
}
restartWhenIdle("config reload (" + strings.Join(changed, ", ") + ")")
return "OK config valid; restarting once the GPU is idle (changed: " + strings.Join(changed, ", ") + ")"
}
}
// diffConfig lists the env names of settings whose values differ between
// two configs (from each field's env tag, so users recognize them).
func diffConfig(a, b config.Config) []string {
va, vb := reflect.ValueOf(a), reflect.ValueOf(b)
t := va.Type()
var out []string
for i := 0; i < t.NumField(); i++ {
if !reflect.DeepEqual(va.Field(i).Interface(), vb.Field(i).Interface()) {
name := t.Field(i).Tag.Get("env")
if name == "" {
name = t.Field(i).Name
}
out = append(out, name)
}
}
return out
}
// statusProvider assembles the one-line JSON snapshot for CmdStatus. The
// loaded-model query to Ollama gets a short timeout so a wedged upstream
// cannot stall the status channel for long.
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time, gw *gpuWatch, ollamaClient *ollama.Client) func() string {
return func() string {
snap := statusSnapshot{
Version: version,
UptimeS: int64(time.Since(started).Seconds()),
UI: uiURL(cfg.ListenUI),
}
st, known, foreign, ageS := gw.get()
snap.GPU = statusGPU{
Enabled: gw.enabled,
UsedMB: st.UsedMB, TotalMB: st.TotalMB, Known: known,
TempC: -1, FanPct: -1,
Foreign: foreign, AgeS: ageS,
}
if known {
snap.GPU.TempC, snap.GPU.FanPct = st.TempC, st.FanPct
}
if cfg.OllamaURL != "" {
d := statusDownstream{
Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
}
if ollamaClient != nil {
mctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
models, err := ollamaClient.LoadedModelDetails(mctx)
cancel()
if err == nil {
d.Models = make([]statusModel, 0, len(models))
for _, m := range models {
d.Models = append(d.Models, statusModel{Name: m.Name, VRAMMB: m.SizeVRAM / (1024 * 1024)})
}
}
}
snap.Downstreams = append(snap.Downstreams, d)
}
if cfg.ComfyURL != "" {
d := statusDownstream{Name: "comfy", URL: cfg.ComfyURL, Up: health.get("comfy")}
if comfySup != nil {
d.Managed = comfySup.Status()
}
snap.Downstreams = append(snap.Downstreams, d)
}
lst := lk.Status()
snap.Lock = statusLock{
State: string(lst.State),
Detail: lst.Detail,
LLMInflight: lst.LLMInflight,
LLMWaiting: lst.LLMWaiting,
ImageQueue: lst.ImageQueue,
External: lst.External,
SinceS: int64(time.Since(lst.Since).Seconds()),
}
b, err := json.Marshal(snap)
if err != nil {
return `{"version":"` + version + `"}`
}
return string(b)
}
}
// waitForIdle polls the lock until no LLM or image work is active or
// pending, max at most. Returns false on timeout or cancellation.
func waitForIdle(ctx context.Context, lk *lock.Lock, max time.Duration) bool {
+345
View File
@@ -0,0 +1,345 @@
package main
import (
"encoding/json"
"fmt"
"os"
"os/exec"
"strings"
"time"
"gpu-turnstile/internal/control"
)
// ANSI colors for the monitor frame.
const (
cReset = "\x1b[0m"
cDim = "\x1b[2m"
cRed = "\x1b[31m"
cGreen = "\x1b[32m"
cYellow = "\x1b[33m"
cCyan = "\x1b[36m"
)
// hotkeysLine is the monitor's footer.
const hotkeysLine = " " + cDim + "q quit · u update now · r reload config" + cReset + "\x1b[K\n"
// monitorCommand renders a live status view of the running service,
// refreshed every second from the control channel. When the service
// reports a different version and the executable on disk changed (the
// updater replaced it), the monitor restarts itself onto the new binary.
// Hotkeys: q quits, u triggers an update check on the service, r asks the
// service to reload its config file.
func monitorCommand() int {
if !stdoutIsTerminal() {
fmt.Fprintln(os.Stderr, "gpu-turnstile: --monitor needs an interactive terminal")
return 1
}
enableVirtualTerminal()
restore := enableRawKeys()
defer func() {
if restore != nil {
restore()
}
}()
fmt.Print("\x1b[2J") // clear once; frames then redraw in place
defer fmt.Print(cReset + "\n")
exe, _ := os.Executable()
var exeStamp time.Time
if st, err := os.Stat(exe); err == nil {
exeStamp = st.ModTime()
}
keys := make(chan byte, 8)
go readKeys(keys)
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
var note string
var noteAt time.Time
noteCh := make(chan string, 1)
askPending := false
// ask sends a one-shot command to the service and reports the reply in
// the note line. Only one ask runs at a time.
ask := func(cmd, busy, label string) {
if askPending {
return
}
askPending = true
note, noteAt = busy, time.Now()
go func() {
reply, err := control.Ask(cmd)
if err != nil {
noteCh <- label + ": no answer from the service"
return
}
msg := strings.TrimPrefix(reply, "OK ")
msg = strings.TrimPrefix(msg, "ERR ")
noteCh <- label + ": " + msg
}()
}
poll := func() string {
frame := renderWaiting()
if reply, err := control.Ask(control.CmdStatus); err == nil {
if msg, ok := strings.CutPrefix(reply, "OK "); ok {
var snap statusSnapshot
if json.Unmarshal([]byte(msg), &snap) == nil {
if snap.Version != "" && snap.Version != version {
if exeChanged(exe, exeStamp) {
fmt.Print("\x1b[2J\x1b[H")
fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version)
restartSelf(exe, "--monitor")
return "" // re-execed; this process exits below
}
snap.MonitorNote = fmt.Sprintf("note: the service runs %s, this monitor is %s", snap.Version, version)
}
if note != "" {
snap.MonitorNote = note
}
frame = renderMonitor(snap, termWidth())
}
}
}
return frame
}
for {
frame := poll()
if frame == "" {
return 0 // restartSelf fired
}
fmt.Print("\x1b[H" + frame + "\x1b[J") // home, frame, clear below
select {
case <-ticker.C:
if note != "" && time.Since(noteAt) > 15*time.Second {
note = ""
}
case k, ok := <-keys:
if !ok {
keys = nil
continue
}
switch k {
case 'q', 'Q', 3: // q or Ctrl+C (raw mode delivers it as a byte)
return 0
case 'u', 'U':
ask(control.CmdUpdateNow, "checking for updates…", "update")
case 'r', 'R':
ask(control.CmdReloadEnv, "reloading config…", "reload")
}
case n := <-noteCh:
askPending = false
note, noteAt = n, time.Now()
}
}
}
// readKeys reads single keypresses from stdin (raw mode was enabled by the
// caller) and delivers them until stdin fails.
func readKeys(keys chan<- byte) {
defer close(keys)
buf := make([]byte, 1)
for {
n, err := os.Stdin.Read(buf)
if n > 0 {
keys <- buf[0]
}
if err != nil {
return
}
}
}
// exeChanged reports whether the executable on disk was replaced since the
// recorded stamp (the updater swaps it via rename, which changes ModTime).
func exeChanged(exe string, stamp time.Time) bool {
if exe == "" || stamp.IsZero() {
return false
}
st, err := os.Stat(exe)
return err == nil && !st.ModTime().Equal(stamp)
}
// restartSelf starts a fresh copy of this executable with the given args on
// the same console; the caller exits right after.
func restartSelf(exe string, args ...string) {
cmd := exec.Command(exe, args...)
cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
cmd.Start() //nolint:errcheck // best effort: on failure we just exit
}
func renderWaiting() string {
return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n\x1b[K\n" + hotkeysLine
}
// renderMonitor draws one full frame. Each line ends with \x1b[K (clear to
// end of line) so shrinking content leaves no residue.
func renderMonitor(snap statusSnapshot, width int) string {
if width < 40 {
width = 80
}
var b strings.Builder
left := " gpu-turnstile"
if snap.UptimeS > 0 {
left += " " + cDim + "up " + fmtDur(snap.UptimeS) + cReset
}
right := snap.Version
pad := width - printableLen(" gpu-turnstile up "+fmtDur(snap.UptimeS)) - len(right) - 1
if snap.UptimeS == 0 {
pad = width - len(" gpu-turnstile") - len(right) - 1
}
if pad < 1 {
pad = 1
}
b.WriteString(cDim + left + strings.Repeat(" ", pad) + right + cReset + "\x1b[K\n")
b.WriteString(cDim + " " + strings.Repeat("─", width-2) + cReset + "\x1b[K\n")
for _, d := range snap.Downstreams {
busy := (d.Name == "ollama" && snap.Lock.State == "llm" && snap.Lock.LLMInflight > 0) ||
(d.Name == "comfy" && snap.Lock.State == "image")
b.WriteString(renderDownstream(d, busy) + "\x1b[K\n")
}
b.WriteString("\x1b[K\n")
b.WriteString(renderLock(snap.Lock) + "\x1b[K\n")
if snap.Lock.ImageQueue > 0 {
b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n",
cYellow, snap.Lock.ImageQueue, cReset))
}
if snap.GPU.Enabled {
b.WriteString(renderGPU(snap.GPU) + "\x1b[K\n")
}
if snap.UI != "" {
b.WriteString(" " + cDim + "UI: " + snap.UI + cReset + "\x1b[K\n")
}
if snap.MonitorNote != "" {
b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
}
b.WriteString("\x1b[K\n")
b.WriteString(hotkeysLine)
return b.String()
}
// renderGPU renders the GPU line: VRAM usage (when nvidia-smi answered),
// the game detector's last finding, and how long ago it ran.
func renderGPU(g statusGPU) string {
s := " GPU: "
if g.Known {
s += renderVRAM(g.UsedMB, g.TotalMB)
if g.TempC >= 0 {
s += fmt.Sprintf(" · %d°C", g.TempC)
}
if g.FanPct >= 0 {
s += fmt.Sprintf(" · fan %d%%", g.FanPct)
}
} else {
s += cDim + "VRAM unknown (nvidia-smi not answering)" + cReset
}
if g.Foreign != "" {
s += " · external: " + cRed + g.Foreign + cReset
} else {
s += cDim + " · no external process" + cReset
}
if g.AgeS >= 0 {
s += cDim + fmt.Sprintf(" (checked %s ago)", fmtDur(g.AgeS)) + cReset
}
return s
}
// formatMB renders MiB as "4.2 GiB" (or MiB below 1 GiB).
func formatMB(mb int64) string {
if mb >= 1024 {
return fmt.Sprintf("%.1f GiB", float64(mb)/1024)
}
return fmt.Sprintf("%d MiB", mb)
}
// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
func renderVRAM(used, total int) string {
if total > 0 {
return formatMB(int64(used)) + " / " + formatMB(int64(total)) + " used"
}
return formatMB(int64(used)) + " used"
}
// printableLen counts characters without ANSI escapes (ASCII-only content).
func printableLen(s string) int { return len(s) }
func renderDownstream(d statusDownstream, busy bool) string {
url := cDim + d.URL + cReset
detail := ""
if d.Models != nil {
if len(d.Models) == 0 {
detail = cDim + " · no models loaded" + cReset
} else {
parts := make([]string, 0, len(d.Models))
for _, m := range d.Models {
if m.VRAMMB > 0 {
parts = append(parts, fmt.Sprintf("%s (%s VRAM)", m.Name, formatMB(m.VRAMMB)))
} else {
parts = append(parts, m.Name+" (in RAM)")
}
}
detail = " · " + strings.Join(parts, ", ")
}
}
if busy {
detail += " · " + cCyan + "busy" + cReset
}
switch d.Managed {
case "stopped":
return fmt.Sprintf(" %s○%s %-8s %sstopped (managed — starts on demand)%s %s",
cDim, cReset, d.Name, cDim, cReset, url)
case "starting":
return fmt.Sprintf(" %s◌%s %-8s %sstarting…%s %s",
cYellow, cReset, d.Name, cYellow, cReset, url)
}
suffix := ""
if d.Managed == "external" {
suffix = " (external)"
}
if d.Up {
return fmt.Sprintf(" %s●%s %-8s %sUP%s%s%s %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, detail, url)
}
return fmt.Sprintf(" %s●%s %-8s %sDOWN%s %s", cRed, cReset, d.Name, cRed, cReset, url)
}
func renderLock(l statusLock) string {
dur := cDim + "(" + fmtDur(l.SinceS) + ")" + cReset
switch l.State {
case "idle":
return fmt.Sprintf(" Lock: %sidle%s %s", cGreen, cReset, dur)
case "llm":
s := fmt.Sprintf(" Lock: %sLLM%s — %d in flight", cCyan, cReset, l.LLMInflight)
if l.LLMWaiting > 0 {
s += fmt.Sprintf(", %d waiting", l.LLMWaiting)
}
if l.Detail != "" {
s += " — " + l.Detail
}
return s + " " + dur
case "image":
s := fmt.Sprintf(" Lock: %sIMAGE%s", cYellow, cReset)
if l.Detail != "" {
s += " — " + l.Detail
}
return s + " " + dur
case "external":
return fmt.Sprintf(" Lock: %sEXTERNAL%s — %s %s", cRed, cReset, l.External, dur)
}
return " Lock: unknown"
}
// fmtDur renders seconds as a compact duration ("1m32s", "2h07m").
func fmtDur(s int64) string {
if s < 0 {
s = 0
}
d := time.Duration(s) * time.Second
if d >= time.Hour {
return fmt.Sprintf("%dh%02dm", int(d.Hours()), int(d.Minutes())%60)
}
return d.Round(time.Second).String()
}
+110
View File
@@ -0,0 +1,110 @@
package main
import (
"gpu-turnstile/internal/config"
"reflect"
"strings"
"testing"
)
func TestRenderMonitor(t *testing.T) {
snap := statusSnapshot{
Version: "v0.2.3",
UptimeS: 3725,
Downstreams: []statusDownstream{
{Name: "ollama", URL: "http://127.0.0.1:11435", Up: true},
{Name: "comfy", URL: "http://127.0.0.1:8189", Managed: "stopped"},
},
Lock: statusLock{State: "llm", LLMInflight: 2, LLMWaiting: 1, Detail: "ollama: POST /api/generate", SinceS: 95},
}
frame := renderMonitor(snap, 80)
for _, want := range []string{"v0.2.3", "1h02m", "ollama", "UP", "comfy", "stopped", "LLM", "2 in flight", "1 waiting", "1m35s"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.Lock = statusLock{State: "external", External: "cyberpunk2077.exe (pid 1234)", SinceS: 3, ImageQueue: 2}
frame = renderMonitor(snap, 40) // narrow: falls back to 80
for _, want := range []string{"EXTERNAL", "cyberpunk2077.exe", "2 image job(s) waiting"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.GPU = statusGPU{Enabled: true, Known: true, UsedMB: 4300, TotalMB: 16384, TempC: 55, FanPct: 42, Foreign: "cyberpunk2077.exe (pid 1234)", AgeS: 12}
frame = renderMonitor(snap, 80)
for _, want := range []string{"GPU:", "4.2 GiB / 16.0 GiB used", "55°C", "fan 42%", "external:", "checked 12s ago"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.GPU = statusGPU{Enabled: true, AgeS: 4}
frame = renderMonitor(snap, 80)
for _, want := range []string{"VRAM unknown", "no external process"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.Downstreams[0].Models = []statusModel{{Name: "llama3.1:8b", VRAMMB: 4900}, {Name: "embed", VRAMMB: 0}}
snap.Lock = statusLock{State: "llm", LLMInflight: 1}
frame = renderMonitor(snap, 80)
for _, want := range []string{"llama3.1:8b (4.8 GiB VRAM)", "embed (in RAM)", "busy"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.Downstreams[0].Models = []statusModel{}
snap.Lock = statusLock{State: "idle"}
frame = renderMonitor(snap, 80)
if !strings.Contains(frame, "no models loaded") {
t.Errorf("frame missing %q:\n%s", "no models loaded", frame)
}
if strings.Contains(frame, "busy") {
t.Errorf("idle lock still shows busy:\n%s", frame)
}
snap.UI = "http://127.0.0.1:7860"
frame = renderMonitor(snap, 80)
if !strings.Contains(frame, "UI: http://127.0.0.1:7860") {
t.Errorf("frame missing UI line:\n%s", frame)
}
}
func TestFmtDur(t *testing.T) {
cases := map[int64]string{0: "0s", 5: "5s", 95: "1m35s", 3725: "1h02m", -3: "0s"}
for in, want := range cases {
if got := fmtDur(in); got != want {
t.Errorf("fmtDur(%d) = %q, want %q", in, got, want)
}
}
}
func TestDiffConfig(t *testing.T) {
a := config.Defaults()
b := a
if got := diffConfig(a, b); len(got) != 0 {
t.Fatalf("identical configs: got %v", got)
}
b.LogLevel = -4
b.GameProcs = []string{"game.exe"}
got := diffConfig(a, b)
if len(got) != 2 || got[0] != "GAME_PROCS" || got[1] != "LOGLEVEL" {
t.Fatalf("got %v, want [GAME_PROCS LOGLEVEL]", got)
}
}
// Every Config field must carry an env tag so user-facing output (the
// reload diff) can name the setting the user would actually change.
func TestConfigFieldsHaveEnvTags(t *testing.T) {
typ := reflect.TypeOf(config.Config{})
for i := 0; i < typ.NumField(); i++ {
f := typ.Field(i)
if f.Tag.Get("env") == "" {
t.Errorf("config.Config.%s has no env tag", f.Name)
}
}
}
+39
View File
@@ -0,0 +1,39 @@
//go:build !windows
package main
import (
"os"
"golang.org/x/sys/unix"
)
// enableVirtualTerminal is a no-op: unix terminals speak ANSI natively.
func enableVirtualTerminal() {}
// termWidth is the terminal width in columns, 80 when unknown.
func termWidth() int {
ws, err := unix.IoctlGetWinsize(int(os.Stdout.Fd()), unix.TIOCGWINSZ)
if err != nil || ws.Col == 0 {
return 80
}
return int(ws.Col)
}
// enableRawKeys switches the terminal to per-keypress mode (ICANON and ECHO
// off) and returns the restore function, nil when stdin is not a terminal.
func enableRawKeys() func() {
fd := int(os.Stdin.Fd())
term, err := unix.IoctlGetTermios(fd, unix.TCGETS)
if err != nil {
return nil
}
raw := *term
raw.Lflag &^= unix.ICANON | unix.ECHO
raw.Cc[unix.VMIN] = 1
raw.Cc[unix.VTIME] = 0
if err := unix.IoctlSetTermios(fd, unix.TCSETS, &raw); err != nil {
return nil
}
return func() { unix.IoctlSetTermios(fd, unix.TCSETS, term) } //nolint:errcheck
}
+49
View File
@@ -0,0 +1,49 @@
//go:build windows
package main
import (
"os"
"golang.org/x/sys/windows"
)
// enableVirtualTerminal asks the console to honor ANSI escapes (Windows 10+);
// mintty/Git Bash already does, so errors are ignored.
func enableVirtualTerminal() {
h := windows.Handle(os.Stdout.Fd())
var mode uint32
if err := windows.GetConsoleMode(h, &mode); err != nil {
return
}
windows.SetConsoleMode(h, mode|windows.ENABLE_VIRTUAL_TERMINAL_PROCESSING) //nolint:errcheck
}
// termWidth is the console window width in columns, 80 when unknown.
func termWidth() int {
var info windows.ConsoleScreenBufferInfo
if err := windows.GetConsoleScreenBufferInfo(windows.Handle(os.Stdout.Fd()), &info); err != nil {
return 80
}
return int(info.Window.Right-info.Window.Left) + 1
}
// enableRawKeys puts the console's stdin into per-keypress mode (no line
// buffering, no echo) and returns the restore function. When stdin is not a
// real console (mintty/Git Bash pipes) it returns nil: ptys already deliver
// keystrokes immediately.
func enableRawKeys() func() {
h := windows.Handle(os.Stdin.Fd())
var mode uint32
if err := windows.GetConsoleMode(h, &mode); err != nil {
return nil
}
const (
enableLineInput = 0x0002
enableEchoInput = 0x0004
)
if err := windows.SetConsoleMode(h, mode&^(enableLineInput|enableEchoInput)); err != nil {
return nil
}
return func() { windows.SetConsoleMode(h, mode) } //nolint:errcheck
}
+242
View File
@@ -0,0 +1,242 @@
package main
import (
"crypto/subtle"
"fmt"
"net"
"net/http"
"strings"
"gpu-turnstile/internal/control"
)
// uiHandler serves the web UI (LISTEN_UI): a single page mirroring
// --monitor, its status JSON (the same snapshot the control channel
// serves) and action endpoints for the update/reload buttons. user/pass
// (UI_USER/UI_PASS) enable HTTP basic auth; both empty means open access,
// which is meant for localhost only.
func uiHandler(h control.Handler, user, pass string) http.Handler {
mux := http.NewServeMux()
mux.HandleFunc("GET /", func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "text/html; charset=utf-8")
fmt.Fprint(w, uiPage)
})
mux.HandleFunc("GET /api/status", func(w http.ResponseWriter, r *http.Request) {
msg := strings.TrimPrefix(h(control.CmdStatus), "OK ")
w.Header().Set("Content-Type", "application/json")
if strings.HasPrefix(msg, "ERR ") {
w.WriteHeader(http.StatusInternalServerError)
msg = strings.TrimPrefix(msg, "ERR ")
}
fmt.Fprint(w, msg)
})
mux.HandleFunc("POST /api/action/{cmd}", func(w http.ResponseWriter, r *http.Request) {
cmd := r.PathValue("cmd")
if cmd != control.CmdUpdateNow && cmd != control.CmdReloadEnv {
http.NotFound(w, r)
return
}
msg := strings.TrimPrefix(h(cmd), "OK ")
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
if strings.HasPrefix(msg, "ERR ") {
w.WriteHeader(http.StatusConflict)
msg = strings.TrimPrefix(msg, "ERR ")
}
fmt.Fprint(w, msg)
})
if user != "" {
return basicAuth(mux, user, pass)
}
return mux
}
// basicAuth wraps next with HTTP basic auth (constant-time compares).
func basicAuth(next http.Handler, user, pass string) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
u, p, ok := r.BasicAuth()
if !ok ||
subtle.ConstantTimeCompare([]byte(u), []byte(user)) != 1 ||
subtle.ConstantTimeCompare([]byte(p), []byte(pass)) != 1 {
w.Header().Set("WWW-Authenticate", `Basic realm="gpu-turnstile"`)
http.Error(w, "unauthorized", http.StatusUnauthorized)
return
}
next.ServeHTTP(w, r)
})
}
// isLoopbackHost reports whether a listen host is localhost-only. An empty
// host (":7860") binds all interfaces and is therefore not loopback.
func isLoopbackHost(host string) bool {
if host == "localhost" {
return true
}
ip := net.ParseIP(host)
return ip != nil && ip.IsLoopback()
}
// uiURL renders a LISTEN_UI address as a browsable URL ("" when disabled).
// A wildcard host reads as 127.0.0.1 — that's where the browser is usually
// running.
func uiURL(addr string) string {
if addr == "" {
return ""
}
if strings.HasPrefix(addr, ":") {
addr = "127.0.0.1" + addr
}
return "http://" + addr
}
// uiPage is the whole web UI: a terminal-styled page that polls /api/status
// once a second and renders the same lines as --monitor. The buttons (and
// the u/r keys) hit the action endpoints, like the monitor's hotkeys.
const uiPage = `<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>gpu-turnstile</title>
<style>
body { background:#111; color:#ddd; font:14px/1.6 ui-monospace,Consolas,"Cascadia Mono",monospace;
margin:0 auto; padding:14px 18px; max-width:860px; }
header { display:flex; justify-content:space-between; border-bottom:1px solid #333;
padding-bottom:4px; margin-bottom:10px; }
footer { border-top:1px solid #333; margin-top:10px; padding-top:8px;
display:flex; gap:8px; align-items:center; }
button { background:#1d1d1d; color:#bbb; border:1px solid #444; border-radius:4px;
padding:2px 10px; font:inherit; cursor:pointer; }
button:hover { background:#2c2c2c; color:#eee; }
.dim { color:#777; } .up { color:#4caf50; } .down { color:#ef5350; }
.llm { color:#26c6da; } .warn { color:#fbc02d; }
#note { color:#fbc02d; min-height:1.6em; }
</style>
</head>
<body>
<header><span>gpu-turnstile <span id="up" class="dim"></span></span><span id="ver" class="dim"></span></header>
<div id="downstreams"></div>
<div id="lock"></div>
<div id="gpu"></div>
<div id="note"></div>
<footer>
<button id="bu">u update now</button>
<button id="br">r reload config</button>
<span class="dim">refreshes every second</span>
</footer>
<script>
const $ = id => document.getElementById(id);
const el = (cls, text) => { const e = document.createElement("span"); if (cls) e.className = cls; e.textContent = text; return e; };
const line = (...parts) => { const d = document.createElement("div"); d.append(...parts); return d; };
const fmtMB = mb => mb >= 1024 ? (mb / 1024).toFixed(1) + " GiB" : mb + " MiB";
const fmtDur = s => {
s = Math.max(0, Math.floor(s));
if (s >= 3600) return Math.floor(s / 3600) + "h" + String(Math.floor(s / 60) % 60).padStart(2, "0") + "m";
if (s >= 60) return Math.floor(s / 60) + "m" + String(s % 60).padStart(2, "0") + "s";
return s + "s";
};
function renderDownstream(d, lock) {
const busy = (d.name === "ollama" && lock.state === "llm" && lock.llm_inflight > 0) ||
(d.name === "comfy" && lock.state === "image");
if (d.managed === "stopped")
return line(el("dim", "○ " + d.name + " stopped (managed — starts on demand) " + d.url));
if (d.managed === "starting")
return line(el("warn", "◌ " + d.name + " starting… "), el("dim", d.url));
const parts = [el(d.up ? "up" : "down", "● " + d.name + " " + (d.up ? "UP" : "DOWN"))];
if (d.managed === "external") parts.push(el("", " (external)"));
if (d.models) {
if (d.models.length === 0) {
parts.push(el("dim", " · no models loaded"));
} else {
parts.push(el("", " · " + d.models.map(m =>
m.vram_mb > 0 ? m.name + " (" + fmtMB(m.vram_mb) + " VRAM)" : m.name + " (in RAM)").join(", ")));
}
}
if (busy) parts.push(el("llm", " · busy"));
parts.push(el("dim", " " + d.url));
return line(...parts);
}
function renderLock(l) {
const dur = el("dim", " (" + fmtDur(l.since_s || 0) + ")");
switch (l.state) {
case "idle": return line(el("up", "Lock: idle"), dur);
case "llm": {
let s = "Lock: LLM — " + l.llm_inflight + " in flight";
if (l.llm_waiting > 0) s += ", " + l.llm_waiting + " waiting";
if (l.detail) s += " — " + l.detail;
return line(el("llm", s), dur);
}
case "image": return line(el("warn", "Lock: IMAGE" + (l.detail ? " — " + l.detail : "")), dur);
case "external": return line(el("down", "Lock: EXTERNAL — " + (l.external || "")), dur);
}
return line(el("", "Lock: unknown"));
}
function renderGPU(g) {
const parts = [el("", "GPU: ")];
if (g.known) {
let s = fmtMB(g.used_mb) + " / " + fmtMB(g.total_mb) + " used";
if (g.temp_c >= 0) s += " · " + g.temp_c + "°C";
if (g.fan_pct >= 0) s += " · fan " + g.fan_pct + "%";
parts.push(el("", s));
} else {
parts.push(el("dim", "VRAM unknown (nvidia-smi not answering)"));
}
parts.push(g.foreign ? el("down", " · external: " + g.foreign) : el("dim", " · no external process"));
if (g.age_s >= 0) parts.push(el("dim", " (checked " + fmtDur(g.age_s) + " ago)"));
return line(...parts);
}
async function poll() {
let snap;
try {
const r = await fetch("/api/status");
if (!r.ok) throw new Error();
snap = await r.json();
} catch {
$("downstreams").replaceChildren(line(el("dim", "gpu-turnstile — waiting for a running service…")));
$("lock").replaceChildren(); $("gpu").replaceChildren();
return;
}
$("ver").textContent = snap.version || "";
$("up").textContent = snap.uptime_s > 0 ? "up " + fmtDur(snap.uptime_s) : "";
const ds = $("downstreams"); ds.replaceChildren();
for (const d of snap.downstreams || []) ds.appendChild(renderDownstream(d, snap.lock || {}));
const lk = $("lock"); lk.replaceChildren(renderLock(snap.lock || {}));
if ((snap.lock || {}).image_queue > 0)
lk.appendChild(line(el("warn", "Queue: " + snap.lock.image_queue + " image job(s) waiting")));
const g = $("gpu"); g.replaceChildren();
if (snap.gpu && snap.gpu.enabled) g.appendChild(renderGPU(snap.gpu));
}
let noteTimer;
async function act(cmd) {
const note = $("note");
note.textContent = "…";
try {
const r = await fetch("/api/action/" + cmd, { method: "POST" });
note.textContent = await r.text();
} catch {
note.textContent = "no answer from the service";
}
clearTimeout(noteTimer);
noteTimer = setTimeout(() => { note.textContent = ""; }, 15000);
}
$("bu").onclick = () => act("update-now");
$("br").onclick = () => act("reload-env");
document.addEventListener("keydown", e => {
if (e.key === "u") act("update-now");
if (e.key === "r") act("reload-env");
});
poll();
setInterval(poll, 1000);
</script>
</body>
</html>
`
+185
View File
@@ -0,0 +1,185 @@
package main
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"gpu-turnstile/internal/control"
)
func testUIHandler() *httptest.Server {
h := uiHandler(func(cmd string) string {
switch cmd {
case control.CmdStatus:
return `OK {"version":"v9.9.9","uptime_s":5}`
case control.CmdUpdateNow:
return "OK v9.9.9 is up to date"
default:
return "ERR unknown command: " + cmd
}
}, "", "")
return httptest.NewServer(h)
}
func TestUIPage(t *testing.T) {
srv := testUIHandler()
defer srv.Close()
resp, err := srv.Client().Get(srv.URL + "/")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body := make([]byte, 1024*64)
n, _ := resp.Body.Read(body)
if resp.StatusCode != 200 {
t.Fatalf("status %d", resp.StatusCode)
}
if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "text/html") {
t.Errorf("Content-Type %q", ct)
}
if !strings.Contains(string(body[:n]), "gpu-turnstile") {
t.Error("page does not mention gpu-turnstile")
}
}
func TestUIStatus(t *testing.T) {
srv := testUIHandler()
defer srv.Close()
resp, err := srv.Client().Get(srv.URL + "/api/status")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body := make([]byte, 4096)
n, _ := resp.Body.Read(body)
if resp.StatusCode != 200 {
t.Fatalf("status %d", resp.StatusCode)
}
if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "application/json") {
t.Errorf("Content-Type %q", ct)
}
if got := strings.TrimSpace(string(body[:n])); got != `{"version":"v9.9.9","uptime_s":5}` {
t.Errorf("body %q", got)
}
}
func TestUIActions(t *testing.T) {
srv := testUIHandler()
defer srv.Close()
resp, err := srv.Client().Post(srv.URL+"/api/action/update-now", "", nil)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body := make([]byte, 4096)
n, _ := resp.Body.Read(body)
if resp.StatusCode != 200 || string(body[:n]) != "v9.9.9 is up to date" {
t.Errorf("update-now: %d %q", resp.StatusCode, body[:n])
}
// The handler answers ERR for reload-env: mapped to 409 with the
// message stripped.
resp2, err := srv.Client().Post(srv.URL+"/api/action/reload-env", "", nil)
if err != nil {
t.Fatal(err)
}
defer resp2.Body.Close()
n, _ = resp2.Body.Read(body)
if resp2.StatusCode != 409 || !strings.HasPrefix(string(body[:n]), "unknown command") {
t.Errorf("reload-env: %d %q", resp2.StatusCode, body[:n])
}
// Unknown actions are not routed to the handler.
resp3, err := srv.Client().Post(srv.URL+"/api/action/reboot", "", nil)
if err != nil {
t.Fatal(err)
}
resp3.Body.Close()
if resp3.StatusCode != 404 {
t.Errorf("bogus action: %d, want 404", resp3.StatusCode)
}
// GET on an action path falls through to the root handler's guard,
// which only serves "/" itself.
resp4, err := srv.Client().Get(srv.URL + "/api/action/update-now")
if err != nil {
t.Fatal(err)
}
resp4.Body.Close()
if resp4.StatusCode != 404 {
t.Errorf("GET action: %d, want 404", resp4.StatusCode)
}
}
func TestUIBasicAuth(t *testing.T) {
h := uiHandler(func(cmd string) string { return `OK {"version":"v9.9.9"}` }, "admin", "s3cret")
srv := httptest.NewServer(h)
defer srv.Close()
// No credentials: 401 with the auth challenge on page and API alike.
for _, path := range []string{"/", "/api/status"} {
resp, err := srv.Client().Get(srv.URL + path)
if err != nil {
t.Fatal(err)
}
resp.Body.Close()
if resp.StatusCode != 401 {
t.Errorf("GET %s without auth: %d, want 401", path, resp.StatusCode)
}
if resp.Header.Get("WWW-Authenticate") == "" {
t.Errorf("GET %s: no WWW-Authenticate header", path)
}
}
// Wrong credentials: 401. Right credentials: 200.
req, _ := http.NewRequest("GET", srv.URL+"/api/status", nil)
req.SetBasicAuth("admin", "wrong")
resp, err := srv.Client().Do(req)
if err != nil {
t.Fatal(err)
}
resp.Body.Close()
if resp.StatusCode != 401 {
t.Errorf("wrong password: %d, want 401", resp.StatusCode)
}
req, _ = http.NewRequest("GET", srv.URL+"/api/status", nil)
req.SetBasicAuth("admin", "s3cret")
resp, err = srv.Client().Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body := make([]byte, 4096)
n, _ := resp.Body.Read(body)
if resp.StatusCode != 200 || strings.TrimSpace(string(body[:n])) != `{"version":"v9.9.9"}` {
t.Errorf("with auth: %d %q", resp.StatusCode, body[:n])
}
}
func TestIsLoopbackHost(t *testing.T) {
for host, want := range map[string]bool{
"127.0.0.1": true, "localhost": true, "::1": true,
"": false, "0.0.0.0": false, "192.168.1.10": false, "example.com": false,
} {
if got := isLoopbackHost(host); got != want {
t.Errorf("isLoopbackHost(%q) = %v, want %v", host, got, want)
}
}
}
func TestUIURL(t *testing.T) {
for addr, want := range map[string]string{
"": "",
"127.0.0.1:7860": "http://127.0.0.1:7860",
":7860": "http://127.0.0.1:7860",
"192.168.1.5:9000": "http://192.168.1.5:9000",
} {
if got := uiURL(addr); got != want {
t.Errorf("uiURL(%q) = %q, want %q", addr, got, want)
}
}
}
+1 -1
View File
@@ -6,7 +6,7 @@
# ComfyUI --listen 0.0.0.0 --port 8189).
services:
gpu-turnstile:
image: git.rambossek.at/public/gpu-turnstile:v0.2.3
image: git.rambossek.at/public/gpu-turnstile:v0.3.2
restart: unless-stopped
environment:
# Each consumer is enabled by setting its URL; leave one unset to
+67 -42
View File
@@ -13,43 +13,48 @@ import (
"time"
)
// Config holds every gpu-turnstile setting.
// Config holds every gpu-turnstile setting. Each field's env tag names the
// environment variable / config-file key that sets it; user-facing output
// (e.g. the reload diff) uses those names, never the Go field names.
type Config struct {
ListenOllama string
ListenComfy string
OllamaURL string
ComfyURL string
UnloadTimeout time.Duration
JobTimeout time.Duration
LLMWaitTimeout time.Duration
ListenOllama string `env:"LISTEN_OLLAMA"`
ListenComfy string `env:"LISTEN_COMFY"`
ListenUI string `env:"LISTEN_UI"`
UIUser string `env:"UI_USER"`
UIPass string `env:"UI_PASS"`
OllamaURL string `env:"OLLAMA_URL"`
ComfyURL string `env:"COMFY_URL"`
UnloadTimeout time.Duration `env:"UNLOAD_TIMEOUT"`
JobTimeout time.Duration `env:"JOB_TIMEOUT"`
LLMWaitTimeout time.Duration `env:"LLM_WAIT_TIMEOUT"`
UnloadPollInterval time.Duration
HistoryPollInterval time.Duration
ProbeTimeout time.Duration
HealthInterval time.Duration
FreeTimeout time.Duration
WarmTimeout time.Duration
ShutdownTimeout time.Duration
BackoffInitial time.Duration
BackoffMax time.Duration
PromptCaptureLimit int64
UnloadPollInterval time.Duration `env:"UNLOAD_POLL_INTERVAL"`
HistoryPollInterval time.Duration `env:"HISTORY_POLL_INTERVAL"`
ProbeTimeout time.Duration `env:"PROBE_TIMEOUT"`
HealthInterval time.Duration `env:"HEALTH_INTERVAL"`
FreeTimeout time.Duration `env:"FREE_TIMEOUT"`
WarmTimeout time.Duration `env:"WARM_TIMEOUT"`
ShutdownTimeout time.Duration `env:"SHUTDOWN_TIMEOUT"`
BackoffInitial time.Duration `env:"BACKOFF_INITIAL"`
BackoffMax time.Duration `env:"BACKOFF_MAX"`
PromptCaptureLimit int64 `env:"PROMPT_CAPTURE_LIMIT"`
AutoUpdate bool
UpdateInterval time.Duration
UpdateRepo string
UpdateAsset string
AutoUpdate bool `env:"AUTO_UPDATE"`
UpdateInterval time.Duration `env:"UPDATE_INTERVAL"`
UpdateRepo string `env:"UPDATE_REPO"`
UpdateAsset string `env:"UPDATE_ASSET"`
// AppVersion is the version the user wants to run: "dev" disables
// updates, "stable" tracks the latest release, anything else is an
// exact vX.Y.Z release to pin. From APP_VER; defaults to "stable".
AppVersion string
AppVersion string `env:"APP_VER"`
// 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
LLMBusyMode string `env:"LLM_BUSY_MODE"`
LLMBusyStatus int `env:"LLM_BUSY_STATUS"`
BusyRetryAfter int `env:"BUSY_RETRY_AFTER"`
// ComfyCmd spawns and supervises a ComfyUI server on demand. When
// ComfyCmd is empty but ComfyDir is set, management is enabled with the
@@ -59,26 +64,29 @@ type Config struct {
// set explicitly. The managed server is stopped after ComfyIdleTimeout
// without requests, freeing its VRAM; ComfyStartTimeout bounds how long
// a request waits for it to come up.
ComfyCmd string
ComfyDir string
ComfyIdleTimeout time.Duration
ComfyStartTimeout time.Duration
ComfyCmd string `env:"COMFY_CMD"`
ComfyDir string `env:"COMFY_DIR"`
ComfyIdleTimeout time.Duration `env:"COMFY_IDLE_TIMEOUT"`
ComfyStartTimeout time.Duration `env:"COMFY_START_TIMEOUT"`
// 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
// that many MiB of VRAM, and the PDH path (GPUForeignUtilPct,
// GPU_FOREIGN_UTIL_PCT) when such a process uses more than that many
// percent of the GPU 3D engine (Windows only). GamePollInterval
// (GAME_POLL_INTERVAL) is how often all checks run.
GameProcs []string `env:"GAME_PROCS"`
GPUForeignVRAMMB int `env:"GPU_FOREIGN_VRAM_MB"`
GPUForeignUtilPct int `env:"GPU_FOREIGN_UTIL_PCT"`
GPUIgnoreProcs []string `env:"GPU_IGNORE_PROCS"`
GamePollInterval time.Duration `env:"GAME_POLL_INTERVAL"`
WarmModel string
LogLevel slog.Level
LogJSON bool
LogFile string
WarmModel string `env:"WARM_MODEL"`
LogLevel slog.Level `env:"LOGLEVEL"`
LogJSON bool `env:"LOG_FORMAT"`
LogFile string `env:"LOG_FILE"`
}
// Defaults returns the configuration used when neither the environment nor
@@ -117,8 +125,9 @@ func Defaults() Config {
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"},
// our own consumers from tripping the foreign-VRAM check. dwm is the
// desktop compositor — it always shows some 3D-engine usage.
GPUIgnoreProcs: []string{"ollama", "ollama app", "ollama_llama_server", "python", "pythonw", "dwm"},
GamePollInterval: 15 * time.Second,
LogLevel: slog.LevelWarn,
@@ -190,6 +199,9 @@ func Load(getenv func(string) string) (Config, error) {
}{
{"LISTEN_OLLAMA", &cfg.ListenOllama},
{"LISTEN_COMFY", &cfg.ListenComfy},
{"LISTEN_UI", &cfg.ListenUI},
{"UI_USER", &cfg.UIUser},
{"UI_PASS", &cfg.UIPass},
{"OLLAMA_URL", &cfg.OllamaURL},
{"COMFY_URL", &cfg.ComfyURL},
{"WARM_MODEL", &cfg.WarmModel},
@@ -241,6 +253,13 @@ func Load(getenv func(string) string) (Config, error) {
}
cfg.GPUForeignVRAMMB = n
}
if v := getenv("GPU_FOREIGN_UTIL_PCT"); v != "" {
n, err := strconv.Atoi(v)
if err != nil || n < 0 || n > 100 {
return cfg, fmt.Errorf("GPU_FOREIGN_UTIL_PCT: must be an integer in 0-100 (percent of the GPU 3D engine, 0 = disabled)")
}
cfg.GPUForeignUtilPct = n
}
if v := getenv("PROMPT_CAPTURE_LIMIT"); v != "" {
n, err := strconv.ParseInt(v, 10, 64)
if err != nil || n < 0 {
@@ -305,6 +324,12 @@ func Load(getenv func(string) string) (Config, error) {
default:
return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"")
}
if (cfg.UIUser == "") != (cfg.UIPass == "") {
return cfg, fmt.Errorf("UI_USER and UI_PASS must be set together (both empty = no auth)")
}
if cfg.UIUser != "" && cfg.ListenUI == "" {
return cfg, fmt.Errorf("UI_USER/UI_PASS have no effect without LISTEN_UI")
}
if cfg.ComfyCmd != "" && cfg.ComfyURL == "" {
return cfg, fmt.Errorf("COMFY_CMD requires COMFY_URL to be set (the proxy needs somewhere to forward)")
}
+24
View File
@@ -97,6 +97,30 @@ func TestComfyCmdRequiresURL(t *testing.T) {
}
}
func TestUIAuthPairing(t *testing.T) {
env := func(set map[string]string) func(string) string {
return func(k string) string { return set[k] }
}
// Only one of UI_USER/UI_PASS: error.
_, err := Load(env(map[string]string{"OLLAMA_URL": "http://x", "LISTEN_UI": "127.0.0.1:7860", "UI_USER": "admin"}))
if err == nil || !strings.Contains(err.Error(), "UI_USER and UI_PASS must be set together") {
t.Fatalf("err = %v, want pairing error", err)
}
// Auth without LISTEN_UI: error.
_, err = Load(env(map[string]string{"OLLAMA_URL": "http://x", "UI_USER": "admin", "UI_PASS": "x"}))
if err == nil || !strings.Contains(err.Error(), "no effect without LISTEN_UI") {
t.Fatalf("err = %v, want LISTEN_UI error", err)
}
// Both with LISTEN_UI: loads.
cfg, err := Load(env(map[string]string{"OLLAMA_URL": "http://x", "LISTEN_UI": "127.0.0.1:7860", "UI_USER": "admin", "UI_PASS": "x"}))
if err != nil {
t.Fatalf("auth pair with LISTEN_UI must load: %v", err)
}
if cfg.UIUser != "admin" || cfg.UIPass != "x" {
t.Errorf("got %q/%q", cfg.UIUser, cfg.UIPass)
}
}
func TestParseEnvFile(t *testing.T) {
input := `# comment
OLLAMA_URL=http://host:11435
+5 -1
View File
@@ -25,6 +25,9 @@ func sampleEntries(logFile string) []sampleEntry {
return []sampleEntry{
{"LISTEN_OLLAMA", ":11434", "Listen address for Ollama-compatible clients (gpu-turnstile poses as Ollama here)", false},
{"LISTEN_COMFY", ":8188", "Listen address for ComfyUI clients (gpu-turnstile poses as ComfyUI here)", false},
{"LISTEN_UI", "127.0.0.1:7860", "Web UI listen address (host:port): live status like --monitor, with update/reload buttons (default: empty = disabled; bind 127.0.0.1 unless you set auth)", false},
{"UI_USER", "", "HTTP basic auth for the web UI (UI_USER and UI_PASS must be set together; both empty = no auth)", false},
{"UI_PASS", "", "See UI_USER", false},
{"OLLAMA_URL", "http://127.0.0.1:11434", "Ollama upstream URL; setting it enables the Ollama consumer (default: empty = disabled)", false},
{"COMFY_URL", "http://127.0.0.1:8188", "ComfyUI upstream URL; setting it enables the ComfyUI consumer (default: empty = disabled)", false},
{"WARM_MODEL", "", "Optional model to reload after an image job (default: empty = none)", false},
@@ -34,7 +37,8 @@ func sampleEntries(logFile string) []sampleEntry {
{"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},
{"GPU_FOREIGN_UTIL_PCT", "30", "Also treat the GPU as held when a process not in GPU_IGNORE_PROCS uses more than this percent of the GPU 3D engine (Windows Task-Manager counters; catches games without an exe list; 0/empty = disabled)", false},
{"GPU_IGNORE_PROCS", "ollama,ollama app,ollama_llama_server,python,pythonw,dwm", "Process names never counted as foreign GPU users (ComfyUI runs under python; dwm is the desktop compositor)", false},
{"GAME_POLL_INTERVAL", "15s", "How often game/VRAM detection runs (nvidia-smi polls keep the GPU awake; don't go below ~10s)", 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},
+56
View File
@@ -7,6 +7,13 @@
// triggers, so the worst a local user can cause is a cheap, throttled
// check and a GPU-idle-gated restart onto a signed binary.
//
// Abuse hardening: the command read is capped (4 KiB), each connection is
// force-closed after connTimeout so a stalled client cannot pin a goroutine
// (or a Windows pipe instance) forever, and concurrently served connections
// are capped at maxConns — beyond that, connections are closed on arrival.
// On Windows the pipe's ACL additionally denies network logons, so the
// channel cannot be reached from another machine.
//
// Protocol: the client writes one command line, the server answers with
// one reply line ("OK ..." or "ERR ...") and hangs up.
package control
@@ -17,12 +24,20 @@ import (
"fmt"
"io"
"strings"
"time"
)
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
// idle) restart onto a signed update immediately.
const CmdUpdateNow = "update-now"
// CmdStatus asks for a one-line JSON status snapshot (monitor mode).
const CmdStatus = "status"
// CmdReloadEnv asks the service to re-read and validate its config file,
// and to restart onto it (once the GPU is idle) when it changed.
const CmdReloadEnv = "reload-env"
// ErrUnavailable means no running service offers the control channel.
var ErrUnavailable = errors.New("control channel unavailable")
@@ -30,9 +45,50 @@ var ErrUnavailable = errors.New("control channel unavailable")
// line. It must start with "OK " or "ERR ".
type Handler func(cmd string) string
// connTimeout bounds one connection's lifetime: a client that stops
// mid-command or never reads the reply would otherwise pin its goroutine
// (and on Windows one of the pipe instances) indefinitely. A var so tests
// can shrink it.
var connTimeout = 10 * time.Second
// maxConns caps concurrently served connections; beyond it, new
// connections are closed on arrival. Bound on the goroutines a local
// flood can pile up.
const maxConns = 32
var connSem = make(chan struct{}, maxConns)
// serve dispatches connection handling under the concurrency cap. It
// returns false when the cap is reached — the caller must then close the
// connection itself.
func serve(c io.ReadWriteCloser, h Handler) bool {
select {
case connSem <- struct{}{}:
go func() {
defer func() { <-connSem }()
serveConn(c, h)
}()
return true
default:
return false
}
}
// forceCloser is implemented by connections that can be torn down
// abortively, unblocking pending reads and writes (Windows pipe:
// DisconnectNamedPipe; unix socket: a deadline in the past). The
// connection watchdog uses it; normal closes still flush the reply.
type forceCloser interface {
ForceClose() error
}
// serveConn runs the line protocol on one accepted connection.
func serveConn(c io.ReadWriteCloser, h Handler) {
defer c.Close()
if fc, ok := c.(forceCloser); ok {
timer := time.AfterFunc(connTimeout, func() { fc.ForceClose() })
defer timer.Stop()
}
line, err := bufio.NewReader(io.LimitReader(c, 4096)).ReadString('\n')
cmd := strings.TrimSpace(line)
if cmd == "" {
+13 -1
View File
@@ -37,12 +37,24 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
if err != nil {
return // shutting down
}
go serveConn(c, h)
uc := unixConn{c}
if !serve(uc, h) {
uc.ForceClose()
}
}
}()
return nil
}
// unixConn adds an abortive ForceClose to net.Conn: a deadline in the
// past fails pending and future I/O immediately.
type unixConn struct{ net.Conn }
func (c unixConn) ForceClose() error {
c.SetDeadline(time.Now().Add(-time.Second)) //nolint:errcheck // best effort
return c.Conn.Close()
}
// Ask sends one command to the running service and returns its reply.
func Ask(cmd string) (string, error) {
c, err := net.DialTimeout("unix", sockPath, 2*time.Second)
+57
View File
@@ -5,6 +5,7 @@ import (
"net"
"strings"
"testing"
"time"
)
func TestRoundTrip(t *testing.T) {
@@ -44,3 +45,59 @@ func TestEmptyReplyIsUnavailable(t *testing.T) {
t.Fatalf("err = %v, want ErrUnavailable", err)
}
}
func TestServeCap(t *testing.T) {
for i := 0; i < maxConns; i++ {
connSem <- struct{}{}
}
defer func() {
for i := 0; i < maxConns; i++ {
<-connSem
}
}()
server, client := net.Pipe()
defer server.Close()
defer client.Close()
if serve(server, func(string) string { return "OK" }) {
t.Fatal("serve accepted a connection beyond the cap")
}
}
// forcePipe records ForceClose calls for the watchdog test.
type forcePipe struct {
net.Conn
forced chan struct{}
}
func (c forcePipe) ForceClose() error {
err := c.Conn.Close()
close(c.forced)
return err
}
func TestConnWatchdog(t *testing.T) {
old := connTimeout
connTimeout = 50 * time.Millisecond
defer func() { connTimeout = old }()
server, client := net.Pipe()
defer client.Close()
fc := forcePipe{Conn: server, forced: make(chan struct{})}
done := make(chan struct{})
go func() {
serveConn(fc, func(string) string { return "OK" })
close(done)
}()
// The client never sends anything; the watchdog must tear the
// connection down instead of blocking forever.
select {
case <-fc.forced:
case <-time.After(5 * time.Second):
t.Fatal("watchdog did not force-close the stalled connection")
}
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("serveConn still blocked after the force close")
}
}
+76 -21
View File
@@ -7,6 +7,8 @@ import (
"fmt"
"log/slog"
"os"
"syscall"
"time"
"unsafe"
"golang.org/x/sys/windows"
@@ -76,37 +78,90 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
log.Warn("control channel stopped", "err", err)
return
}
go func() {
// Blocks until a client connects; on process exit the
// handle goes away with everything else.
if err := windows.ConnectNamedPipe(pipe, nil); err != nil {
windows.CloseHandle(pipe)
return
}
f := os.NewFile(uintptr(pipe), pipePath)
serveConn(f, h) // closes f, and with it the pipe handle
}()
// Blocks until a client connects — only then is the next
// instance created, so instances are not burned without
// clients. ERROR_PIPE_CONNECTED means the client raced us
// and connected before the call: that is a success. Process
// exit reaps the blocked call on shutdown.
if err := windows.ConnectNamedPipe(pipe, nil); err != nil && err != errnoPipeConnected {
windows.CloseHandle(pipe)
continue
}
conn := &pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}
if !serve(conn, h) {
conn.ForceClose()
}
}
}()
return nil
}
// errnoPipeConnected is ConnectNamedPipe's "the client connected before we
// called" result, which means the connection is established.
var errnoPipeConnected = syscall.Errno(535) // ERROR_PIPE_CONNECTED
// errnoPipeBusy is CreateFile's "all pipe instances are busy" result. It
// loses the race against another client that grabbed the instance
// WaitNamedPipe just reported — the caller must wait and retry.
var errnoPipeBusy = syscall.Errno(231) // ERROR_PIPE_BUSY
// pipeConn adapts a pipe handle to io.ReadWriteCloser. Close flushes first
// (FlushFileBuffers blocks until the client has read the reply) and then
// disconnects — closing the bare handle right after writing can discard
// unread reply bytes, which clients see as an empty, failed request.
type pipeConn struct {
f *os.File
h windows.Handle
}
func (c *pipeConn) Read(p []byte) (int, error) { return c.f.Read(p) }
func (c *pipeConn) Write(p []byte) (int, error) { return c.f.Write(p) }
func (c *pipeConn) Close() error {
windows.FlushFileBuffers(c.h) //nolint:errcheck // best effort
windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort
return c.f.Close()
}
// ForceClose aborts the connection without flushing: disconnecting
// unblocks pending reads and writes at the cost of possibly discarding an
// unread reply. Used by the connection watchdog; normal closes flush.
func (c *pipeConn) ForceClose() error {
windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort
return c.f.Close()
}
// Ask sends one command to the running service and returns its reply.
//
// The server keeps exactly one listening instance per connection, so
// concurrent clients race for it: WaitNamedPipe can report an instance
// that another client grabs before our CreateFile runs (ERROR_PIPE_BUSY).
// Retry on that — with the monitor polling status every second, a single
// attempt loses that race regularly.
func Ask(cmd string) (string, error) {
name, err := windows.UTF16PtrFromString(pipePath)
if err != nil {
return "", err
}
if err := waitNamedPipe(name, 2000); err != nil {
return "", ErrUnavailable
deadline := time.Now().Add(5 * time.Second)
for {
if err := waitNamedPipe(name, 2000); err != nil {
return "", ErrUnavailable
}
handle, err := windows.CreateFile(name,
windows.GENERIC_READ|windows.GENERIC_WRITE, 0, nil,
windows.OPEN_EXISTING, 0, 0)
if err == errnoPipeBusy {
if time.Now().After(deadline) {
return "", ErrUnavailable
}
continue
}
if err != nil {
return "", ErrUnavailable
}
f := os.NewFile(uintptr(handle), pipePath)
defer f.Close()
return readReply(f, cmd)
}
handle, err := windows.CreateFile(name,
windows.GENERIC_READ|windows.GENERIC_WRITE, 0, nil,
windows.OPEN_EXISTING, 0, 0)
if err != nil {
return "", ErrUnavailable
}
f := os.NewFile(uintptr(handle), pipePath)
defer f.Close()
return readReply(f, cmd)
}
+134 -30
View File
@@ -1,8 +1,10 @@
// 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.
// VRAM while they run. Three detection paths: an explicit process watch
// list (GAME_PROCS), a foreign-VRAM threshold via nvidia-smi
// (GPU_FOREIGN_VRAM_MB), and a per-process GPU 3D-engine utilization
// threshold via Windows PDH counters (GPU_FOREIGN_UTIL_PCT). The latter two
// catch anything not on the ignore list without naming individual games.
package game
import (
@@ -11,6 +13,7 @@ import (
"fmt"
"log/slog"
"os/exec"
"slices"
"strconv"
"strings"
)
@@ -28,28 +31,33 @@ type computeApp struct {
}
// 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.
// (no watch list, no thresholds) 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
utilPct int // foreign 3D-engine utilization 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
noNvidia bool // nvidia-smi was not found; VRAM path disabled for good
sampler *gpuEngineSampler // open PDH query, opened lazily on first Check
noPDH bool // engine counters unavailable; util 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 {
// MiB, 3D-engine utilization threshold in percent (both 0 = disabled) and
// ignore list. Names are matched case-insensitively, with or without a
// trailing ".exe".
func New(procs []string, vramMB, utilPct 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,
procs: nameSet(procs),
vramMB: vramMB,
utilPct: utilPct,
ignore: nameSet(ignore),
log: log,
}
}
@@ -73,34 +81,63 @@ func nameSet(names []string) map[string]bool {
// 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).
// VRAM path permanently (logged once), as do missing engine counters.
func (d *Detector) Check(ctx context.Context) ([]string, error) {
ps, psErr := processes()
if d.vramMB <= 0 || d.noNvidia {
return d.detect(ps, nil), psErr
utils := d.engineUtil()
var apps []computeApp
if d.vramMB > 0 && !d.noNvidia {
var err error
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")
} else if err != nil {
return d.detect(ps, nil, utils), err
}
}
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
return d.detect(ps, apps, utils), psErr
}
// 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 {
// engineUtil samples per-process 3D-engine utilization via PDH. The first
// call only primes the rate counters and returns nil. A failing open or
// sample disables the path permanently (logged once).
func (d *Detector) engineUtil() map[int]float64 {
if d.utilPct <= 0 || d.noPDH {
return nil
}
if d.sampler == nil {
s, err := openGPUEngineSampler()
if err != nil {
d.noPDH = true
d.log.Warn("GPU_FOREIGN_UTIL_PCT is set but per-process GPU counters are unavailable; engine detection disabled", "err", err)
return nil
}
d.sampler = s
}
utils, err := d.sampler.sample()
if err != nil {
if errors.Is(err, errNotPrimed) {
return nil
}
d.noPDH = true
d.log.Warn("per-process GPU counters failed; engine detection disabled", "err", err)
return nil
}
return utils
}
// detect is the pure core of Check: given the process table and
// (optionally) the nvidia-smi compute-apps list and the PDH engine
// utilization, it returns the foreign holders.
func (d *Detector) detect(ps []Process, apps []computeApp, utils map[int]float64) []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 {
if (d.vramMB > 0 && apps != nil) || (d.utilPct > 0 && utils != nil) {
names := make(map[int]string, len(ps))
for _, p := range ps {
names[p.PID] = p.Name
@@ -115,6 +152,23 @@ func (d *Detector) detect(ps []Process, apps []computeApp) []string {
}
holders = append(holders, fmt.Sprintf("%s (pid %d) using %d MiB VRAM", name, a.PID, a.UsedMB))
}
// Sorted for stable output (map iteration order is random).
pids := make([]int, 0, len(utils))
for pid := range utils {
pids = append(pids, pid)
}
slices.Sort(pids)
for _, pid := range pids {
util := utils[pid]
name := names[pid]
if d.ignore[normName(name)] || util < float64(d.utilPct) {
continue
}
if name == "" {
name = "unknown process"
}
holders = append(holders, fmt.Sprintf("%s (pid %d) using %.0f%% GPU", name, pid, util))
}
}
return holders
}
@@ -132,6 +186,56 @@ func queryComputeApps(ctx context.Context) ([]computeApp, error) {
return parseComputeApps(string(out))
}
// GPUStats is one nvidia-smi reading of the whole card.
type GPUStats struct {
UsedMB int
TotalMB int
TempC int // -1 when nvidia-smi reports N/A
FanPct int // -1 when N/A (some cards don't expose the fan)
}
// QueryGPUStats returns VRAM usage, temperature and fan speed via
// nvidia-smi. Unlike the per-process list this works under WDDM too.
func QueryGPUStats(ctx context.Context) (GPUStats, error) {
out, err := exec.CommandContext(ctx, "nvidia-smi",
"--query-gpu=memory.used,memory.total,temperature.gpu,fan.speed", "--format=csv,noheader,nounits").Output()
if err != nil {
return GPUStats{}, err
}
return parseGPUStats(string(out))
}
// parseGPUStats parses one "used, total, temp, fan" CSV line (MiB, °C,
// percent). The memory fields must be numeric; temperature and fan fall
// back to -1 on "N/A" and friends.
func parseGPUStats(out string) (GPUStats, error) {
fields := strings.Split(strings.TrimSpace(out), ",")
if len(fields) != 4 {
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(out))
}
num := func(s string) (int, error) {
return strconv.Atoi(strings.TrimSpace(s))
}
optional := func(s string) int {
n, err := num(s)
if err != nil {
return -1
}
return n
}
var st GPUStats
var err error
if st.UsedMB, err = num(fields[0]); err != nil {
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(out))
}
if st.TotalMB, err = num(fields[1]); err != nil {
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(out))
}
st.TempC = optional(fields[2])
st.FanPct = optional(fields[3])
return st, nil
}
// 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) {
+68 -4
View File
@@ -45,7 +45,7 @@ func TestParseComputeApps(t *testing.T) {
}
func TestDetect(t *testing.T) {
d := New([]string{"Cyberpunk2077.exe", "hl2"}, 1024,
d := New([]string{"Cyberpunk2077.exe", "hl2"}, 1024, 0,
[]string{"ollama", "python", "pythonw"}, nil)
ps := []Process{
{PID: 10, Name: "ollama.exe"},
@@ -58,7 +58,7 @@ func TestDetect(t *testing.T) {
{PID: 40, UsedMB: 2048}, // foreign, above threshold
{PID: 50, UsedMB: 100}, // foreign but below threshold
}
holders := d.detect(ps, apps)
holders := d.detect(ps, apps, nil)
if len(holders) != 2 {
t.Fatalf("got %v, want 2 holders", holders)
}
@@ -70,13 +70,77 @@ func TestDetect(t *testing.T) {
}
}
func TestDetectEngineUtil(t *testing.T) {
d := New(nil, 0, 30, []string{"dwm", "python"}, nil)
ps := []Process{
{PID: 10, Name: "dwm.exe"},
{PID: 20, Name: "game.exe"},
{PID: 40, Name: "browser.exe"},
}
utils := map[int]float64{
10: 45, // ignored: dwm
20: 61, // foreign, above threshold
30: 82, // foreign, unknown name
40: 5, // below threshold
}
holders := d.detect(ps, nil, utils)
want := []string{
"game.exe (pid 20) using 61% GPU",
"unknown process (pid 30) using 82% GPU",
}
if !slices.Equal(holders, want) {
t.Errorf("got %v, want %v", holders, want)
}
}
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 {
d := New(nil, 0, 0, nil, nil)
if got := d.detect([]Process{{PID: 1, Name: "game.exe"}}, nil, nil); len(got) != 0 {
t.Errorf("got %v, want none", got)
}
}
func TestParseGPUEngineInstance(t *testing.T) {
pid, eng, ok := parseGPUEngineInstance("pid_1234_luid_0x00000000_0x00011A2B_phys_0_eng_0_engtype_3D")
if !ok || pid != 1234 || eng != "3D" {
t.Errorf("got %d, %q, %v", pid, eng, ok)
}
pid, eng, ok = parseGPUEngineInstance("pid_42_luid_0x0_0x0_phys_0_eng_1_engtype_Copy")
if !ok || pid != 42 || eng != "Copy" {
t.Errorf("got %d, %q, %v", pid, eng, ok)
}
for _, bad := range []string{"", "something", "pid_", "pid_x_luid", "pid_-1_luid_0"} {
if _, _, ok := parseGPUEngineInstance(bad); ok {
t.Errorf("%q parsed, want failure", bad)
}
}
}
func TestParseGPUStats(t *testing.T) {
st, err := parseGPUStats("4300, 16384, 55, 42\n")
if err != nil {
t.Fatal(err)
}
if st.UsedMB != 4300 || st.TotalMB != 16384 || st.TempC != 55 || st.FanPct != 42 {
t.Errorf("got %+v", st)
}
// Cards that don't expose temperature/fan report N/A.
st, err = parseGPUStats("1024, 16384, N/A, N/A")
if err != nil {
t.Fatal(err)
}
if st.TempC != -1 || st.FanPct != -1 {
t.Errorf("got %+v, want -1 for N/A fields", st)
}
for _, bad := range []string{"", "1, 2", "x, 16384, 55, 42", "1024, x, 55, 42", "1, 2, 3, 4, 5"} {
if _, err := parseGPUStats(bad); err == nil {
t.Errorf("%q parsed, want failure", bad)
}
}
}
func TestProcessesLive(t *testing.T) {
if runtime.GOOS != "windows" && runtime.GOOS != "linux" {
t.Skip("no process listing on this platform")
+36
View File
@@ -0,0 +1,36 @@
package game
import (
"errors"
"strconv"
"strings"
)
// errNotPrimed marks the first PDH sample after opening a query: rate-based
// counters (like engine utilization) need two collections before they
// return meaningful values.
var errNotPrimed = errors.New("GPU engine counter needs a second sample")
// parseGPUEngineInstance splits a PDH "GPU Engine" instance name —
// "pid_1234_luid_0x00000000_0x00011A2B_phys_0_eng_0_engtype_3D" — into PID
// and engine type ("3D", "Copy", "VideoDecode", ...). engType is empty when
// the name carries no engtype marker.
func parseGPUEngineInstance(name string) (pid int, engType string, ok bool) {
rest, found := strings.CutPrefix(name, "pid_")
if !found {
return 0, "", false
}
digits, rest, found := strings.Cut(rest, "_")
if !found {
return 0, "", false
}
pid, err := strconv.Atoi(digits)
if err != nil || pid < 0 {
return 0, "", false
}
const marker = "engtype_"
if i := strings.LastIndex(rest, marker); i >= 0 {
engType = rest[i+len(marker):]
}
return pid, engType, true
}
+16
View File
@@ -0,0 +1,16 @@
//go:build !windows
package game
import "errors"
// errNoEngineCounters marks platforms without per-process GPU engine
// counters (the PDH path is Windows-only).
var errNoEngineCounters = errors.New("per-process GPU engine counters are only available on Windows")
// gpuEngineSampler is a stub on non-Windows platforms.
type gpuEngineSampler struct{}
func openGPUEngineSampler() (*gpuEngineSampler, error) { return nil, errNoEngineCounters }
func (s *gpuEngineSampler) sample() (map[int]float64, error) { return nil, errNoEngineCounters }
+111
View File
@@ -0,0 +1,111 @@
//go:build windows
package game
import (
"fmt"
"unsafe"
"golang.org/x/sys/windows"
)
// Per-process GPU engine utilization via PDH — the same counters Task
// Manager's "GPU engine" columns read. Unlike nvidia-smi's compute-apps
// this covers graphics work under WDDM, so games show up. The counter is
// added with PdhAddEnglishCounterW, which is independent of the Windows
// display language.
var (
pdhDLL = windows.NewLazySystemDLL("pdh.dll")
procPdhOpenQuery = pdhDLL.NewProc("PdhOpenQueryW")
procPdhAddEnglishCounter = pdhDLL.NewProc("PdhAddEnglishCounterW")
procPdhCollectQueryData = pdhDLL.NewProc("PdhCollectQueryData")
procPdhGetFormattedCounterArray = pdhDLL.NewProc("PdhGetFormattedCounterArrayW")
procPdhCloseQuery = pdhDLL.NewProc("PdhCloseQuery")
)
const (
pdhFmtDouble = 0x00000200 // PDH_FMT_DOUBLE
pdhMoreData = 0x800007D2 // PDH_MORE_DATA
)
// pdhCountervalueItem mirrors PDH_FMT_COUNTERVALUE_ITEM (64-bit, double).
type pdhCountervalueItem struct {
name *uint16
cStatus uint32
_ uint32 // alignment padding
value float64
}
// gpuEngineSampler holds an open PDH query on the wildcard GPU Engine
// utilization counter. Keeping the query open across polls is what makes
// the rate-based values meaningful; the sampler lives as long as the
// process (PdhCloseQuery would only matter on unload).
type gpuEngineSampler struct {
query uintptr // PDH_HQUERY
counter uintptr // PDH_HCOUNTER
primed bool
}
// openGPUEngineSampler opens a query on the per-process GPU engine
// utilization counter (all instances).
func openGPUEngineSampler() (*gpuEngineSampler, error) {
var q uintptr
if r, _, _ := procPdhOpenQuery.Call(0, 0, uintptr(unsafe.Pointer(&q))); r != 0 {
return nil, fmt.Errorf("PdhOpenQuery: status %#x", r)
}
path, err := windows.UTF16PtrFromString(`\GPU Engine(*)\Utilization Percentage`)
if err != nil {
procPdhCloseQuery.Call(q)
return nil, err
}
var c uintptr
if r, _, _ := procPdhAddEnglishCounter.Call(q, uintptr(unsafe.Pointer(path)), 0, uintptr(unsafe.Pointer(&c))); r != 0 {
procPdhCloseQuery.Call(q)
return nil, fmt.Errorf("PdhAddEnglishCounter: status %#x", r)
}
return &gpuEngineSampler{query: q, counter: c}, nil
}
// sample collects the counter once and returns per-PID 3D-engine
// utilization in percent. The first call after open only primes the rate
// calculation and returns errNotPrimed. Processes can drive several 3D
// engines; their values are summed.
func (s *gpuEngineSampler) sample() (map[int]float64, error) {
if r, _, _ := procPdhCollectQueryData.Call(s.query); r != 0 {
return nil, fmt.Errorf("PdhCollectQueryData: status %#x", r)
}
if !s.primed {
s.primed = true
return nil, errNotPrimed
}
var size, count uint32
r, _, _ := procPdhGetFormattedCounterArray.Call(s.counter, pdhFmtDouble,
uintptr(unsafe.Pointer(&size)), uintptr(unsafe.Pointer(&count)), 0)
if r == pdhMoreData && size == 0 {
return nil, nil // no GPU engine instances at all
}
if r != pdhMoreData {
return nil, fmt.Errorf("PdhGetFormattedCounterArray(size): status %#x", r)
}
buf := make([]byte, size)
r, _, _ = procPdhGetFormattedCounterArray.Call(s.counter, pdhFmtDouble,
uintptr(unsafe.Pointer(&size)), uintptr(unsafe.Pointer(&count)),
uintptr(unsafe.Pointer(&buf[0])))
if r != 0 {
return nil, fmt.Errorf("PdhGetFormattedCounterArray: status %#x", r)
}
items := unsafe.Slice((*pdhCountervalueItem)(unsafe.Pointer(&buf[0])), int(count))
out := make(map[int]float64)
for i := range items {
if items[i].cStatus != 0 || items[i].name == nil {
continue
}
pid, engType, ok := parseGPUEngineInstance(windows.UTF16PtrToString(items[i].name))
if !ok || engType != "3D" {
continue // only the 3D engine marks game-like work
}
out[pid] += items[i].value
}
return out, nil
}
+33
View File
@@ -0,0 +1,33 @@
//go:build windows
package game
import (
"errors"
"testing"
"time"
)
// TestGPUEngineSamplerLive opens the real PDH query and takes two samples;
// the first only primes the rate counters. Skipped (not failed) when the
// machine has no GPU counters.
func TestGPUEngineSamplerLive(t *testing.T) {
s, err := openGPUEngineSampler()
if err != nil {
t.Skipf("no GPU engine counters: %v", err)
}
if _, err := s.sample(); !errors.Is(err, errNotPrimed) {
t.Fatalf("first sample: err = %v, want errNotPrimed", err)
}
time.Sleep(200 * time.Millisecond)
utils, err := s.sample()
if err != nil {
t.Fatalf("second sample: %v", err)
}
for pid, util := range utils {
if pid < 0 || util < 0 {
t.Errorf("pid %d: util %.2f", pid, util)
}
}
t.Logf("%d processes with 3D-engine usage", len(utils))
}
+75 -2
View File
@@ -9,6 +9,7 @@ import (
"context"
"log/slog"
"sync"
"time"
)
// State is the current GPU occupancy state.
@@ -30,10 +31,13 @@ type Lock struct {
change chan struct{} // closed and replaced on every state change
n int // LLM requests in flight
llmWaiting int // LLM requests blocked waiting for the GPU
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
external string // non-empty: a foreign process (e.g. a game) holds the GPU
detail string // what the current holder is doing (best effort)
since time.Time // when the current state began
log *slog.Logger
}
@@ -41,7 +45,7 @@ type Lock struct {
// New returns a ready-to-use Lock. log may be nil; if set, every state
// transition is logged at debug level.
func New(log *slog.Logger) *Lock {
return &Lock{change: make(chan struct{}), log: log}
return &Lock{change: make(chan struct{}), log: log, since: time.Now()}
}
// broadcast wakes all waiters. Call with mu held.
@@ -63,6 +67,7 @@ func (l *Lock) logTransition(msg string, args ...any) {
func (l *Lock) SetExternal(holder string) {
l.mu.Lock()
l.external = holder
l.since = time.Now()
l.broadcast()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateExternal, "holder", holder)
@@ -73,6 +78,7 @@ func (l *Lock) SetExternal(holder string) {
func (l *Lock) ClearExternal() {
l.mu.Lock()
l.external = ""
l.since = time.Now()
l.broadcast()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateIdle)
@@ -90,17 +96,31 @@ func (l *Lock) External() string {
// while waiting; no state is changed in that case.
func (l *Lock) AcquireLLM(ctx context.Context) error {
l.mu.Lock()
waiting := false
for l.imageActive || len(l.imageQ) > 0 || l.external != "" {
if !waiting {
l.llmWaiting++
waiting = true
}
ch := l.change
l.mu.Unlock()
select {
case <-ctx.Done():
l.mu.Lock()
l.llmWaiting--
l.mu.Unlock()
return ctx.Err()
case <-ch:
}
l.mu.Lock()
}
if waiting {
l.llmWaiting--
}
l.n++
if l.n == 1 {
l.since = time.Now()
}
n := l.n
l.mu.Unlock()
l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
@@ -129,6 +149,8 @@ func (l *Lock) ReleaseLLM() {
l.n--
n := l.n
if l.n == 0 {
l.since = time.Now()
l.detail = ""
l.broadcast()
}
l.mu.Unlock()
@@ -156,6 +178,7 @@ func (l *Lock) AcquireImage(ctx context.Context) error {
if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive && l.external == "" {
l.imageQ = l.imageQ[1:]
l.imageActive = true
l.since = time.Now()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateImage)
return nil
@@ -184,6 +207,8 @@ func (l *Lock) AcquireImage(ctx context.Context) error {
func (l *Lock) ReleaseImage() {
l.mu.Lock()
l.imageActive = false
l.since = time.Now()
l.detail = ""
l.broadcast()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateIdle)
@@ -206,3 +231,51 @@ func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) {
}
return state, l.n, l.imageActive || len(l.imageQ) > 0
}
// SetDetail records what the current holder is doing (e.g. the request
// path), for status displays. Best effort: overwritten by each new holder,
// cleared when the GPU goes idle.
func (l *Lock) SetDetail(detail string) {
l.mu.Lock()
l.detail = detail
l.mu.Unlock()
}
// Status is a point-in-time view of the lock for monitoring.
type Status struct {
State State
Detail string
LLMInflight int
LLMWaiting int
ImageActive bool
ImageQueue int
External string
Since time.Time
}
// Status reports the full lock state, including waiters and how long the
// current state has held.
func (l *Lock) Status() Status {
l.mu.Lock()
defer l.mu.Unlock()
s := Status{
Detail: l.detail,
LLMInflight: l.n,
LLMWaiting: l.llmWaiting,
ImageActive: l.imageActive,
ImageQueue: len(l.imageQ),
External: l.external,
Since: l.since,
}
switch {
case l.imageActive:
s.State = StateImage
case l.n > 0:
s.State = StateLLM
case l.external != "":
s.State = StateExternal
default:
s.State = StateIdle
}
return s
}
+30 -9
View File
@@ -60,13 +60,21 @@ func (c *Client) Probe(ctx context.Context) error {
type psResponse struct {
Models []struct {
Name string `json:"name"`
Model string `json:"model"`
Name string `json:"name"`
Model string `json:"model"`
SizeVRAM int64 `json:"size_vram"` // bytes resident in VRAM (0 = RAM-only)
} `json:"models"`
}
// LoadedModels returns the names of models currently held in memory.
func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
// LoadedModel is one model currently held in memory.
type LoadedModel struct {
Name string
SizeVRAM int64 // bytes resident in VRAM; 0 when the model sits in RAM
}
// LoadedModelDetails returns the models currently held in memory with
// their VRAM footprint.
func (c *Client) LoadedModelDetails(ctx context.Context) ([]LoadedModel, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/api/ps", nil)
if err != nil {
return nil, err
@@ -84,13 +92,26 @@ func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
if err := json.NewDecoder(resp.Body).Decode(&ps); err != nil {
return nil, err
}
models := make([]string, 0, len(ps.Models))
models := make([]LoadedModel, 0, len(ps.Models))
for _, m := range ps.Models {
if m.Name != "" {
models = append(models, m.Name)
} else {
models = append(models, m.Model)
name := m.Name
if name == "" {
name = m.Model
}
models = append(models, LoadedModel{Name: name, SizeVRAM: m.SizeVRAM})
}
return models, nil
}
// LoadedModels returns the names of models currently held in memory.
func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
details, err := c.LoadedModelDetails(ctx)
if err != nil {
return nil, err
}
models := make([]string, 0, len(details))
for _, m := range details {
models = append(models, m.Name)
}
return models, nil
}
+2
View File
@@ -491,6 +491,7 @@ func (s *Server) OllamaHandler() http.Handler {
return
}
defer s.cfg.Lock.ReleaseLLM()
s.cfg.Lock.SetDetail("ollama: " + r.Method + " " + r.URL.Path)
s.ollamaProxy.ServeHTTP(w, r)
}))
}
@@ -590,6 +591,7 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
}
s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds())
log.Info("image lock acquired")
s.cfg.Lock.SetDetail("comfy: POST /prompt")
if s.cfg.ComfySup != nil && !comfyFirst {
if err := s.cfg.ComfySup.EnsureRunning(); err != nil {
+18
View File
@@ -164,6 +164,24 @@ func (p *Process) Running() bool {
return p.cmd != nil
}
// Status describes the child for status displays: "external" when something
// else serves the port, "running" once ready, "starting" while the child
// boots, "stopped" otherwise.
func (p *Process) Status() string {
p.mu.Lock()
defer p.mu.Unlock()
switch {
case p.external:
return "external"
case p.cmd == nil:
return "stopped"
case p.ready:
return "running"
default:
return "starting"
}
}
// Ready reports whether the server has answered a probe since its last
// (re)start. Health checks use it to tell "starting up" from "outage".
func (p *Process) Ready() bool {