Compare commits

..
12 Commits
Author SHA1 Message Date
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
21 changed files with 1186 additions and 180 deletions
+17 -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
@@ -64,7 +64,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 +149,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
+19 -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:
@@ -230,7 +236,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 |
+231 -49
View File
@@ -16,6 +16,7 @@ import (
"os"
"os/signal"
"path/filepath"
"reflect"
"runtime"
"strings"
"sync"
@@ -39,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
@@ -59,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, monitor, 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 {
@@ -82,6 +84,8 @@ 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":
@@ -90,7 +94,7 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
rest = append(rest, args[i])
}
}
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, monitor, 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.
@@ -109,6 +113,9 @@ Usage:
(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
@@ -141,12 +148,12 @@ func fatalUsage(format string, args ...any) {
}
func main() {
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, monitor, 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 && !monitor && !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
@@ -171,7 +178,9 @@ 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 monitor && (install || remove || forceUpdate || updateNow):
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))
@@ -181,6 +190,8 @@ func main() {
os.Exit(forceUpdateCommand(configPath, elevatedChild))
case updateNow:
os.Exit(updateNowCommand())
case reloadEnv:
os.Exit(reloadEnvCommand())
case monitor:
os.Exit(monitorCommand())
}
@@ -202,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)
}
@@ -210,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)
}
@@ -451,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
@@ -531,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)
@@ -569,6 +592,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,
@@ -703,11 +727,14 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
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
@@ -747,6 +774,22 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
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 {
if p, err := os.Executable(); err != nil {
@@ -754,34 +797,24 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
} else {
exePath = p
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) {
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)
}
}
// The control channel (status for --monitor, update-now trigger) is
// served whenever running as a service, independent of AUTO_UPDATE.
// The control channel (status for --monitor, update-now and reload-env
// triggers) is served whenever running as a service, independent of
// AUTO_UPDATE.
if isService {
serveControl(ctx, log, u, exePath, applyStaged,
statusProvider(cfg, lk, comfySup, health, started))
statusProvider(cfg, lk, comfySup, health, started, gw, ollamaClient),
reloadHandler(cfg, configPath, restartWhenIdle))
}
select {
@@ -805,12 +838,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
@@ -824,6 +897,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
@@ -949,17 +1026,29 @@ 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 query status
// (--monitor) and trigger an update check (--force-update/--update-now)
// without admin rights. The update 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. u is nil when AUTO_UPDATE=false.
func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string), status func() string) {
// (--monitor), trigger an update check (--force-update/--update-now) and
// poke a config reload (--reload-env) without admin rights. The update
// payload is signature-verified regardless of who asks; update triggers
// are rate-limited to one per minute so the channel cannot be used to spam
// restarts. u is nil when AUTO_UPDATE=false.
func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string), status func() string, reload func() string) {
var mu sync.Mutex
var lastTrigger time.Time
var lastTrigger, lastReload time.Time
h := 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
@@ -1021,6 +1110,16 @@ type statusDownstream struct {
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 {
@@ -1038,22 +1137,105 @@ type statusSnapshot struct {
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"`
// MonitorNote is set client-side (never over the wire) when the
// monitor's own binary differs from the service's version.
MonitorNote string `json:"-"`
}
// statusProvider assembles the one-line JSON snapshot for CmdStatus.
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time) func() string {
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()),
}
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 != "" {
snap.Downstreams = append(snap.Downstreams, statusDownstream{
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")}
@@ -1062,15 +1244,15 @@ func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Proces
}
snap.Downstreams = append(snap.Downstreams, d)
}
st := lk.Status()
lst := lk.Status()
snap.Lock = statusLock{
State: string(st.State),
Detail: st.Detail,
LLMInflight: st.LLMInflight,
LLMWaiting: st.LLMWaiting,
ImageQueue: st.ImageQueue,
External: st.External,
SinceS: int64(time.Since(st.Since).Seconds()),
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 {
+165 -8
View File
@@ -21,17 +21,27 @@ const (
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.
// Ctrl+C quits.
// 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()
@@ -39,7 +49,38 @@ func monitorCommand() int {
if st, err := os.Stat(exe); err == nil {
exeStamp = st.ModTime()
}
for {
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 {
@@ -50,16 +91,64 @@ func monitorCommand() int {
fmt.Print("\x1b[2J\x1b[H")
fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version)
restartSelf(exe, "--monitor")
return 0
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
time.Sleep(time.Second)
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
}
}
}
@@ -82,7 +171,7 @@ func restartSelf(exe string, args ...string) {
}
func renderWaiting() string {
return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n"
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
@@ -109,7 +198,9 @@ func renderMonitor(snap statusSnapshot, width int) string {
b.WriteString(cDim + " " + strings.Repeat("─", width-2) + cReset + "\x1b[K\n")
for _, d := range snap.Downstreams {
b.WriteString(renderDownstream(d) + "\x1b[K\n")
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")
@@ -117,17 +208,83 @@ func renderMonitor(snap statusSnapshot, width int) string {
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.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) string {
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",
@@ -141,7 +298,7 @@ func renderDownstream(d statusDownstream) string {
suffix = " (external)"
}
if d.Up {
return fmt.Sprintf(" %s●%s %-8s %sUP%s%s %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, url)
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)
}
+63
View File
@@ -1,6 +1,8 @@
package main
import (
"gpu-turnstile/internal/config"
"reflect"
"strings"
"testing"
)
@@ -29,6 +31,41 @@ func TestRenderMonitor(t *testing.T) {
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)
}
}
func TestFmtDur(t *testing.T) {
@@ -39,3 +76,29 @@ func TestFmtDur(t *testing.T) {
}
}
}
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)
}
}
}
+18
View File
@@ -19,3 +19,21 @@ func termWidth() int {
}
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
}
+20
View File
@@ -27,3 +27,23 @@ func termWidth() int {
}
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
}
+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.7
image: git.rambossek.at/public/gpu-turnstile:v0.3.1
restart: unless-stopped
environment:
# Each consumer is enabled by setting its URL; leave one unset to
+55 -42
View File
@@ -13,43 +13,45 @@ 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"`
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 +61,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 +122,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,
@@ -241,6 +247,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 {
+2 -1
View File
@@ -34,7 +34,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},
+53
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,6 +24,7 @@ import (
"fmt"
"io"
"strings"
"time"
)
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
@@ -26,6 +34,10 @@ 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")
@@ -33,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")
}
}
+44 -12
View File
@@ -8,6 +8,7 @@ import (
"log/slog"
"os"
"syscall"
"time"
"unsafe"
"golang.org/x/sys/windows"
@@ -86,7 +87,10 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
windows.CloseHandle(pipe)
continue
}
go serveConn(&pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}, h)
conn := &pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}
if !serve(conn, h) {
conn.ForceClose()
}
}
}()
return nil
@@ -96,6 +100,11 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
// 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
@@ -114,22 +123,45 @@ func (c *pipeConn) Close() error {
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))
}
+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
}