Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a07b6726ea | ||
|
|
d2c49e52fa | ||
|
|
e98331bb9e | ||
|
|
54174787d5 | ||
|
|
d259b2e96c | ||
|
|
8fccd333aa |
@@ -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 |
|
||||
|
||||
+97
-30
@@ -592,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,
|
||||
@@ -726,12 +727,13 @@ 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.
|
||||
// gw collects the VRAM reading for the status channel.
|
||||
gw := &gpuWatch{}
|
||||
if len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0 {
|
||||
det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUIgnoreProcs, log)
|
||||
// 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)
|
||||
}
|
||||
|
||||
@@ -811,7 +813,7 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
||||
// AUTO_UPDATE.
|
||||
if isService {
|
||||
serveControl(ctx, log, u, exePath, applyStaged,
|
||||
statusProvider(cfg, lk, comfySup, health, started, gw),
|
||||
statusProvider(cfg, lk, comfySup, health, started, gw, ollamaClient),
|
||||
reloadHandler(cfg, configPath, restartWhenIdle))
|
||||
}
|
||||
|
||||
@@ -836,26 +838,40 @@ 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 from the game detector's poll
|
||||
// loop, for the status channel. Known stays false when game detection is
|
||||
// not configured (no nvidia-smi polling happens then).
|
||||
// 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
|
||||
usedMB int
|
||||
total int
|
||||
known bool
|
||||
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) set(used, total int) {
|
||||
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) get() (used, total int, known bool) {
|
||||
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()
|
||||
return g.usedMB, g.total, g.known
|
||||
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
|
||||
@@ -877,8 +893,9 @@ 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.set(used, total)
|
||||
gw.setVRAM(used, total)
|
||||
}
|
||||
switch {
|
||||
case len(holders) > 0 && !held:
|
||||
@@ -1012,12 +1029,21 @@ func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u
|
||||
// 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:
|
||||
@@ -1080,6 +1106,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 {
|
||||
@@ -1097,8 +1133,9 @@ type statusSnapshot struct {
|
||||
UptimeS int64 `json:"uptime_s"`
|
||||
Downstreams []statusDownstream `json:"downstreams"`
|
||||
Lock statusLock `json:"lock"`
|
||||
// GPU carries the latest VRAM reading; Known is false when game
|
||||
// detection (and with it nvidia-smi polling) is not configured.
|
||||
// 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.
|
||||
@@ -1106,9 +1143,17 @@ type statusSnapshot struct {
|
||||
}
|
||||
|
||||
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
|
||||
@@ -1130,33 +1175,55 @@ func reloadHandler(current config.Config, configPath string, restartWhenIdle fun
|
||||
}
|
||||
}
|
||||
|
||||
// diffConfig lists the names of fields whose values differ between two
|
||||
// configs.
|
||||
// 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()) {
|
||||
out = append(out, t.Field(i).Name)
|
||||
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, gw *gpuWatch) func() string {
|
||||
// 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()),
|
||||
}
|
||||
used, total, known := gw.get()
|
||||
snap.GPU = statusGPU{UsedMB: used, TotalMB: total, Known: known}
|
||||
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{
|
||||
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")}
|
||||
|
||||
@@ -188,7 +188,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")
|
||||
@@ -196,8 +198,8 @@ 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.Known {
|
||||
b.WriteString(" GPU: " + renderVRAM(snap.GPU.UsedMB, snap.GPU.TotalMB) + "\x1b[K\n")
|
||||
if snap.GPU.Enabled {
|
||||
b.WriteString(renderGPU(snap.GPU) + "\x1b[K\n")
|
||||
}
|
||||
if snap.MonitorNote != "" {
|
||||
b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
|
||||
@@ -207,25 +209,66 @@ func renderMonitor(snap statusSnapshot, width int) string {
|
||||
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
|
||||
}
|
||||
|
||||
// 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 {
|
||||
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 formatMB(int64(used)) + " / " + formatMB(int64(total)) + " used"
|
||||
}
|
||||
return format(used) + " 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",
|
||||
@@ -239,7 +282,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)
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package main
|
||||
|
||||
import (
|
||||
"gpu-turnstile/internal/config"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
@@ -30,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, 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)
|
||||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
@@ -50,7 +86,19 @@ func TestDiffConfig(t *testing.T) {
|
||||
b.LogLevel = -4
|
||||
b.GameProcs = []string{"game.exe"}
|
||||
got := diffConfig(a, b)
|
||||
if len(got) != 2 || got[0] != "GameProcs" || got[1] != "LogLevel" {
|
||||
t.Fatalf("got %v, want [GameProcs LogLevel]", got)
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
# ComfyUI --listen 0.0.0.0 --port 8189).
|
||||
services:
|
||||
gpu-turnstile:
|
||||
image: git.rambossek.at/public/gpu-turnstile:v0.2.8
|
||||
image: git.rambossek.at/public/gpu-turnstile:v0.3.0
|
||||
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
|
||||
@@ -37,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"
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
+84
-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
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user