Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e98331bb9e | ||
|
|
54174787d5 | ||
|
|
d259b2e96c | ||
|
|
8fccd333aa | ||
|
|
4bd5f34ce7 | ||
|
|
b9d3f91403 | ||
|
|
98d2d714ca | ||
|
|
7a6409831c | ||
|
|
c7b8747d85 | ||
|
|
09d4e81eab | ||
|
|
1264afe40e |
@@ -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
|
||||
|
||||
@@ -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 |
|
||||
|
||||
+191
-38
@@ -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)
|
||||
}
|
||||
@@ -362,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" {
|
||||
@@ -447,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
|
||||
@@ -527,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)
|
||||
@@ -565,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,
|
||||
@@ -699,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
|
||||
@@ -743,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 {
|
||||
@@ -750,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),
|
||||
reloadHandler(cfg, configPath, restartWhenIdle))
|
||||
}
|
||||
|
||||
select {
|
||||
@@ -801,12 +838,48 @@ 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
|
||||
known bool
|
||||
foreign string // last Check result: external holders, "" when none
|
||||
at time.Time
|
||||
}
|
||||
|
||||
func (g *gpuWatch) setVRAM(used, total int) {
|
||||
g.mu.Lock()
|
||||
g.usedMB, g.total, g.known = used, total, 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() (used, total int, 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 g.usedMB, g.total, 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
|
||||
@@ -820,6 +893,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 used, total, verr := game.QueryVRAMMB(ctx); verr == nil {
|
||||
gw.setVRAM(used, total)
|
||||
}
|
||||
switch {
|
||||
case len(holders) > 0 && !held:
|
||||
held = true
|
||||
@@ -945,17 +1022,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
|
||||
@@ -1034,15 +1123,79 @@ 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:"-"`
|
||||
}
|
||||
|
||||
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"`
|
||||
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.
|
||||
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time) func() string {
|
||||
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time, gw *gpuWatch) func() string {
|
||||
return func() string {
|
||||
snap := statusSnapshot{
|
||||
Version: version,
|
||||
UptimeS: int64(time.Since(started).Seconds()),
|
||||
}
|
||||
used, total, known, foreign, ageS := gw.get()
|
||||
snap.GPU = statusGPU{
|
||||
Enabled: gw.enabled,
|
||||
UsedMB: used, TotalMB: total, Known: known,
|
||||
Foreign: foreign, AgeS: ageS,
|
||||
}
|
||||
if cfg.OllamaURL != "" {
|
||||
snap.Downstreams = append(snap.Downstreams, statusDownstream{
|
||||
Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -20,33 +21,147 @@ const (
|
||||
cCyan = "\x1b[36m"
|
||||
)
|
||||
|
||||
// hotkeysLine is the monitor's footer.
|
||||
const hotkeysLine = " " + cDim + "q quit · u update now" + cReset + "\x1b[K\n"
|
||||
|
||||
// monitorCommand renders a live status view of the running service,
|
||||
// refreshed every second from the control channel. Ctrl+C quits.
|
||||
// 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.
|
||||
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")
|
||||
for {
|
||||
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)
|
||||
updatePending := false
|
||||
|
||||
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
|
||||
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':
|
||||
if !updatePending {
|
||||
updatePending = true
|
||||
note, noteAt = "checking for updates…", time.Now()
|
||||
go func() {
|
||||
reply, err := control.Ask(control.CmdUpdateNow)
|
||||
if err != nil {
|
||||
noteCh <- "update: no answer from the service"
|
||||
return
|
||||
}
|
||||
msg := strings.TrimPrefix(reply, "OK ")
|
||||
msg = strings.TrimPrefix(msg, "ERR ")
|
||||
noteCh <- "update: " + msg
|
||||
}()
|
||||
}
|
||||
}
|
||||
case n := <-noteCh:
|
||||
updatePending = 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"
|
||||
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
|
||||
@@ -81,9 +196,51 @@ 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)
|
||||
} 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
|
||||
}
|
||||
|
||||
// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
|
||||
func renderVRAM(used, total int) string {
|
||||
format := func(mb int) string {
|
||||
if mb >= 1024 {
|
||||
return fmt.Sprintf("%.1f GiB", float64(mb)/1024)
|
||||
}
|
||||
return fmt.Sprintf("%d MiB", mb)
|
||||
}
|
||||
if total > 0 {
|
||||
return format(used) + " / " + format(total) + " used"
|
||||
}
|
||||
return format(used) + " used"
|
||||
}
|
||||
|
||||
// printableLen counts characters without ANSI escapes (ASCII-only content).
|
||||
func printableLen(s string) int { return len(s) }
|
||||
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"gpu-turnstile/internal/config"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
@@ -29,6 +31,22 @@ 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, Foreign: "cyberpunk2077.exe (pid 1234)", AgeS: 12}
|
||||
frame = renderMonitor(snap, 80)
|
||||
for _, want := range []string{"GPU:", "4.2 GiB / 16.0 GiB used", "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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFmtDur(t *testing.T) {
|
||||
@@ -39,3 +57,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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
# ComfyUI --listen 0.0.0.0 --port 8189).
|
||||
services:
|
||||
gpu-turnstile:
|
||||
image: git.rambossek.at/public/gpu-turnstile:v0.2.5
|
||||
image: git.rambossek.at/public/gpu-turnstile:v0.2.9
|
||||
restart: unless-stopped
|
||||
environment:
|
||||
# Each consumer is enabled by setting its URL; leave one unset to
|
||||
|
||||
+55
-42
@@ -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 {
|
||||
|
||||
@@ -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},
|
||||
|
||||
@@ -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 == "" {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"syscall"
|
||||
"time"
|
||||
"unsafe"
|
||||
|
||||
"golang.org/x/sys/windows"
|
||||
@@ -77,18 +78,19 @@ 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. A client that
|
||||
// raced us and connected between CreateNamedPipe and
|
||||
// ConnectNamedPipe reports ERROR_PIPE_CONNECTED — that is
|
||||
// a success, not a failure.
|
||||
if err := windows.ConnectNamedPipe(pipe, nil); err != nil && err != errnoPipeConnected {
|
||||
windows.CloseHandle(pipe)
|
||||
return
|
||||
}
|
||||
serveConn(&pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}, h)
|
||||
}()
|
||||
// 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
|
||||
@@ -98,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
|
||||
@@ -116,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)
|
||||
}
|
||||
|
||||
+107
-30
@@ -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,29 @@ func queryComputeApps(ctx context.Context) ([]computeApp, error) {
|
||||
return parseComputeApps(string(out))
|
||||
}
|
||||
|
||||
// QueryVRAMMB returns used and total GPU VRAM in MiB via nvidia-smi.
|
||||
// Unlike the per-process list this works under WDDM too.
|
||||
func QueryVRAMMB(ctx context.Context) (used, total int, err error) {
|
||||
out, err := exec.CommandContext(ctx, "nvidia-smi",
|
||||
"--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits").Output()
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
usedStr, totalStr, ok := strings.Cut(strings.TrimSpace(string(out)), ",")
|
||||
if !ok {
|
||||
return 0, 0, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(string(out)))
|
||||
}
|
||||
used, err = strconv.Atoi(strings.TrimSpace(usedStr))
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(string(out)))
|
||||
}
|
||||
total, err = strconv.Atoi(strings.TrimSpace(totalStr))
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(string(out)))
|
||||
}
|
||||
return used, total, 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) {
|
||||
|
||||
@@ -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,52 @@ 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 TestProcessesLive(t *testing.T) {
|
||||
if runtime.GOOS != "windows" && runtime.GOOS != "linux" {
|
||||
t.Skip("no process listing on this platform")
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 }
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
Reference in New Issue
Block a user