Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
31ee6013db | ||
|
|
5506163493 | ||
|
|
18a7336c02 | ||
|
|
0796092b80 | ||
|
|
897163042c | ||
|
|
e58ff33912 | ||
|
|
2040b30c94 | ||
|
|
a07b6726ea | ||
|
|
d2c49e52fa | ||
|
|
e98331bb9e | ||
|
|
54174787d5 | ||
|
|
d259b2e96c | ||
|
|
8fccd333aa | ||
|
|
4bd5f34ce7 | ||
|
|
b9d3f91403 | ||
|
|
98d2d714ca | ||
|
|
7a6409831c | ||
|
|
c7b8747d85 | ||
|
|
09d4e81eab | ||
|
|
1264afe40e | ||
|
|
c95d401a36 | ||
|
|
f7ea30a494 | ||
|
|
620be9d57c | ||
|
|
cffb4151a6 | ||
|
|
16dc7fe462 | ||
|
|
a74cd49fe5 | ||
|
|
cc2a2cad27 | ||
|
|
f707d07fd8 |
@@ -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
|
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 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 —
|
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`
|
detection of foreign GPU holders such as games — is enabled by `GAME_PROCS`,
|
||||||
and/or `GPU_FOREIGN_VRAM_MB` (see below).
|
`GPU_FOREIGN_VRAM_MB` and/or `GPU_FOREIGN_UTIL_PCT` (see below).
|
||||||
|
|
||||||
## Configuration
|
## Configuration
|
||||||
|
|
||||||
@@ -49,6 +49,8 @@ override file values. Invalid values fail at startup.
|
|||||||
|---|---|---|
|
|---|---|---|
|
||||||
| `LISTEN_OLLAMA` | `:11434` | Listener for Ollama-compatible clients |
|
| `LISTEN_OLLAMA` | `:11434` | Listener for Ollama-compatible clients |
|
||||||
| `LISTEN_COMFY` | `:8188` | Listener for ComfyUI clients |
|
| `LISTEN_COMFY` | `:8188` | Listener for ComfyUI clients |
|
||||||
|
| `LISTEN_UI` | _(empty = disabled)_ | Web UI listener (`host:port`) mirroring `--monitor` (live status, update/reload buttons). Keep it on `127.0.0.1` unless you set auth |
|
||||||
|
| `UI_USER` / `UI_PASS` | _(empty = no auth)_ | HTTP basic auth for the web UI; must be set together |
|
||||||
| `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer |
|
| `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer |
|
||||||
| `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer |
|
| `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer |
|
||||||
| `UNLOAD_TIMEOUT` | `60s` | Wait for Ollama to unload before an image job |
|
| `UNLOAD_TIMEOUT` | `60s` | Wait for Ollama to unload before an image job |
|
||||||
@@ -64,7 +66,8 @@ override file values. Invalid values fail at startup.
|
|||||||
| `COMFY_START_TIMEOUT` | `2m` | Max wait for the managed ComfyUI to come up |
|
| `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 |
|
| `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_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) |
|
| `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 |
|
| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` works as an alias |
|
||||||
| `LOG_FORMAT` | `text` | `json` for structured JSON logs |
|
| `LOG_FORMAT` | `text` | `json` for structured JSON logs |
|
||||||
@@ -148,20 +151,25 @@ adds a `BindPaths=` to the systemd unit on Linux. Re-run it after changing
|
|||||||
Want to game on the same GPU without Ollama/ComfyUI squatting on the VRAM?
|
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,
|
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
|
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):
|
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
|
GPU_FOREIGN_VRAM_MB=1024 # catch-all via nvidia-smi
|
||||||
```
|
```
|
||||||
|
|
||||||
`GAME_PROCS` matches running process names (case-insensitive, `.exe`
|
`GPU_FOREIGN_UTIL_PCT` reads the same per-process GPU engine counters as
|
||||||
optional). `GPU_FOREIGN_VRAM_MB` asks nvidia-smi which processes hold GPU
|
Task Manager (PDH), which cover graphics work under Windows' WDDM driver —
|
||||||
memory and treats anything not in `GPU_IGNORE_PROCS` above the threshold as
|
so games show up without naming them, and the offender is named in the log
|
||||||
foreign — handy as a catch-all, but note that under Windows' WDDM driver
|
and monitor. Windows-only. `GAME_PROCS` matches running process names
|
||||||
graphics-only games may not show up in nvidia-smi's per-process list, so
|
(case-insensitive, `.exe` optional). `GPU_FOREIGN_VRAM_MB` asks nvidia-smi
|
||||||
name your games in `GAME_PROCS` there; on Linux both paths work. When the
|
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.
|
game exits, requests resume automatically.
|
||||||
|
|
||||||
## Build and run
|
## 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
|
- **Only `COMFY_URL`**: image jobs are tracked and ComfyUI's VRAM is freed
|
||||||
afterwards, but the Ollama unload and warm-reload steps are skipped.
|
afterwards, but the Ollama unload and warm-reload steps are skipped.
|
||||||
- **Game detection** is a third, optional consumer without a URL: enabled by
|
- **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
|
`GAME_PROCS`, `GPU_FOREIGN_VRAM_MB` and/or `GPU_FOREIGN_UTIL_PCT` it
|
||||||
holding the GPU (see below) and plugs into the same lock the same way —
|
watches for foreign processes holding the GPU (see below) and plugs into
|
||||||
excluded when both knobs are unset.
|
the same lock the same way — excluded when all knobs are unset.
|
||||||
|
|
||||||
### Lock semantics
|
### Lock semantics
|
||||||
|
|
||||||
@@ -171,20 +171,26 @@ files are flagged in the startup log.
|
|||||||
## Game detection (foreign GPU holders)
|
## Game detection (foreign GPU holders)
|
||||||
|
|
||||||
Games and other foreign GPU users sit outside the URL-based consumer model —
|
Games and other foreign GPU users sit outside the URL-based consumer model —
|
||||||
nothing proxies through gpu-turnstile for them. Two independent detection
|
nothing proxies through gpu-turnstile for them. Three independent detection
|
||||||
paths, polled every `GAME_POLL_INTERVAL` (default 15 s); either one being
|
paths, polled every `GAME_POLL_INTERVAL` (default 15 s); any one being
|
||||||
configured enables the feature:
|
configured enables the feature:
|
||||||
|
|
||||||
- **Process watch list** (`GAME_PROCS`, comma-separated, case-insensitive,
|
- **Process watch list** (`GAME_PROCS`, comma-separated, case-insensitive,
|
||||||
`.exe` optional): while any listed process runs, the GPU counts as held.
|
`.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
|
- **Foreign VRAM threshold** (`GPU_FOREIGN_VRAM_MB`): `nvidia-smi
|
||||||
--query-compute-apps` lists per-process GPU memory; any process not in
|
--query-compute-apps` lists per-process GPU memory; any process not in
|
||||||
`GPU_IGNORE_PROCS` (default: Ollama and python — ComfyUI runs under python)
|
`GPU_IGNORE_PROCS` (default: Ollama and python — ComfyUI runs under python
|
||||||
holding more than the threshold counts as a foreign holder. Needs
|
— plus dwm, the desktop compositor) holding more than the threshold counts
|
||||||
nvidia-smi on the PATH (absent: logged once, path disabled) and works best
|
as a foreign holder. Needs nvidia-smi on the PATH (absent: logged once,
|
||||||
on Linux — under Windows' WDDM driver, graphics-only games may not appear
|
path disabled) and works best on Linux — under Windows' WDDM driver,
|
||||||
in the per-process list.
|
graphics-only games may not appear in the per-process list.
|
||||||
|
|
||||||
While a holder is detected, gpu-turnstile:
|
While a holder is detected, gpu-turnstile:
|
||||||
|
|
||||||
@@ -215,6 +221,8 @@ override file values. A missing file is fine; a malformed one is fatal.
|
|||||||
|---|---|---|
|
|---|---|---|
|
||||||
| `LISTEN_OLLAMA` | `:11434` | listener for Ollama-compatible clients |
|
| `LISTEN_OLLAMA` | `:11434` | listener for Ollama-compatible clients |
|
||||||
| `LISTEN_COMFY` | `:8188` | listener for ComfyUI clients |
|
| `LISTEN_COMFY` | `:8188` | listener for ComfyUI clients |
|
||||||
|
| `LISTEN_UI` | _(empty = disabled)_ | web UI listener (`host:port`) mirroring `--monitor` (live status, update/reload buttons); keep it on `127.0.0.1` unless auth is set |
|
||||||
|
| `UI_USER` / `UI_PASS` | _(empty = no auth)_ | HTTP basic auth for the web UI; must be set together |
|
||||||
| `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer |
|
| `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer |
|
||||||
| `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer |
|
| `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer |
|
||||||
| `UNLOAD_TIMEOUT` | `60s` | wait for Ollama to unload |
|
| `UNLOAD_TIMEOUT` | `60s` | wait for Ollama to unload |
|
||||||
@@ -230,7 +238,8 @@ override file values. A missing file is fine; a malformed one is fatal.
|
|||||||
| `COMFY_START_TIMEOUT` | `2m` | how long a request waits for the managed ComfyUI to come up |
|
| `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 |
|
| `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_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) |
|
| `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 |
|
| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` is accepted as an alias |
|
||||||
| `LOG_FORMAT` | `text` | `json` for structured JSON logs |
|
| `LOG_FORMAT` | `text` | `json` for structured JSON logs |
|
||||||
|
|||||||
+465
-43
@@ -5,6 +5,7 @@ package main
|
|||||||
import (
|
import (
|
||||||
"bufio"
|
"bufio"
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
@@ -15,13 +16,16 @@ import (
|
|||||||
"os"
|
"os"
|
||||||
"os/signal"
|
"os/signal"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"reflect"
|
||||||
"runtime"
|
"runtime"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"gpu-turnstile/internal/comfy"
|
"gpu-turnstile/internal/comfy"
|
||||||
"gpu-turnstile/internal/config"
|
"gpu-turnstile/internal/config"
|
||||||
|
"gpu-turnstile/internal/control"
|
||||||
"gpu-turnstile/internal/game"
|
"gpu-turnstile/internal/game"
|
||||||
"gpu-turnstile/internal/lock"
|
"gpu-turnstile/internal/lock"
|
||||||
"gpu-turnstile/internal/metrics"
|
"gpu-turnstile/internal/metrics"
|
||||||
@@ -36,7 +40,8 @@ import (
|
|||||||
var version = "dev"
|
var version = "dev"
|
||||||
|
|
||||||
// exitCodeUpdate tells the service recovery configuration to restart the
|
// 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
|
const exitCodeUpdate = 3
|
||||||
|
|
||||||
// exitCodeStaged is returned by an elevated --force-update child when it
|
// exitCodeStaged is returned by an elevated --force-update child when it
|
||||||
@@ -56,7 +61,7 @@ func stdoutIsTerminal() bool {
|
|||||||
// --install-service / --remove-service switches, --no-copy, -h/--help,
|
// --install-service / --remove-service switches, --no-copy, -h/--help,
|
||||||
// -v/--version, --force-update and the hidden --elevated-child marker from
|
// -v/--version, --force-update and the hidden --elevated-child marker from
|
||||||
// args.
|
// args.
|
||||||
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, 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]
|
rest = args[:0]
|
||||||
for i := 0; i < len(args); i++ {
|
for i := 0; i < len(args); i++ {
|
||||||
switch {
|
switch {
|
||||||
@@ -65,9 +70,9 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
|
|||||||
i++
|
i++
|
||||||
case strings.HasPrefix(args[i], "-config="):
|
case strings.HasPrefix(args[i], "-config="):
|
||||||
configPath = strings.TrimPrefix(args[i], "-config=")
|
configPath = strings.TrimPrefix(args[i], "-config=")
|
||||||
case args[i] == "--install-service" || args[i] == "-install-service":
|
case args[i] == "--install-service" || args[i] == "-install-service" || args[i] == "-i":
|
||||||
install = true
|
install = true
|
||||||
case args[i] == "--remove-service" || args[i] == "-remove-service":
|
case args[i] == "--remove-service" || args[i] == "-remove-service" || args[i] == "-r":
|
||||||
remove = true
|
remove = true
|
||||||
case args[i] == "--no-copy" || args[i] == "-no-copy":
|
case args[i] == "--no-copy" || args[i] == "-no-copy":
|
||||||
noCopy = true
|
noCopy = true
|
||||||
@@ -77,13 +82,19 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
|
|||||||
showVersion = true
|
showVersion = true
|
||||||
case args[i] == "--force-update" || args[i] == "-force-update":
|
case args[i] == "--force-update" || args[i] == "-force-update":
|
||||||
forceUpdate = true
|
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":
|
case args[i] == "--elevated-child":
|
||||||
elevatedChild = true
|
elevatedChild = true
|
||||||
default:
|
default:
|
||||||
rest = append(rest, args[i])
|
rest = append(rest, args[i])
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, 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.
|
// versionLine is printed at the top of every help and error screen.
|
||||||
@@ -93,11 +104,20 @@ const usageText = `GPU arbitration proxy for Ollama + ComfyUI
|
|||||||
|
|
||||||
Usage:
|
Usage:
|
||||||
gpu-turnstile -config <path> run the proxy
|
gpu-turnstile -config <path> run the proxy
|
||||||
gpu-turnstile --install-service [--no-copy] [-config path] install + start as a service
|
gpu-turnstile -i | --install-service [--no-copy] [-config path]
|
||||||
gpu-turnstile --remove-service stop + uninstall the service
|
install + start as a service
|
||||||
|
gpu-turnstile -r | --remove-service stop + uninstall the service
|
||||||
gpu-turnstile -v | --version print just the version
|
gpu-turnstile -v | --version print just the version
|
||||||
gpu-turnstile --force-update check for a signed update now,
|
gpu-turnstile --force-update check for a signed update now,
|
||||||
apply it and restart the service
|
apply it and restart the service
|
||||||
|
(no admin needed when the service runs)
|
||||||
|
gpu-turnstile --update-now like --force-update, but only
|
||||||
|
through the running service
|
||||||
|
gpu-turnstile --reload-env make the service re-read and
|
||||||
|
validate its config file, then
|
||||||
|
restart onto it if it changed
|
||||||
|
gpu-turnstile -m | --monitor live status view (downstreams,
|
||||||
|
GPU lock, queue); Ctrl+C quits
|
||||||
gpu-turnstile -h | --help this help
|
gpu-turnstile -h | --help this help
|
||||||
|
|
||||||
Options:
|
Options:
|
||||||
@@ -128,12 +148,12 @@ func fatalUsage(format string, args ...any) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
configPath, install, remove, noCopy, help, showVersion, forceUpdate, elevatedChild, args := parseFlags(os.Args[1:])
|
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, args := parseFlags(os.Args[1:])
|
||||||
if showVersion {
|
if showVersion {
|
||||||
fmt.Println(version)
|
fmt.Println(version)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
bare := configPath == "" && !install && !remove && !forceUpdate && !elevatedChild && len(args) == 0
|
bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !reloadEnv && !monitor && !elevatedChild && len(args) == 0
|
||||||
if help || (bare && stdoutIsTerminal()) {
|
if help || (bare && stdoutIsTerminal()) {
|
||||||
// Bare invocation in a terminal (e.g. double-clicked on Windows)
|
// Bare invocation in a terminal (e.g. double-clicked on Windows)
|
||||||
// shows the help instead of starting a proxy window with no visible
|
// shows the help instead of starting a proxy window with no visible
|
||||||
@@ -156,12 +176,24 @@ func main() {
|
|||||||
fatalUsage("error: --install-service and --remove-service are mutually exclusive")
|
fatalUsage("error: --install-service and --remove-service are mutually exclusive")
|
||||||
case forceUpdate && (install || remove):
|
case forceUpdate && (install || remove):
|
||||||
fatalUsage("error: --force-update cannot be combined with --install-service/--remove-service")
|
fatalUsage("error: --force-update cannot be combined with --install-service/--remove-service")
|
||||||
|
case updateNow && (install || remove || forceUpdate):
|
||||||
|
fatalUsage("error: --update-now cannot be combined with other commands")
|
||||||
|
case reloadEnv && (install || remove || forceUpdate || updateNow):
|
||||||
|
fatalUsage("error: --reload-env cannot be combined with other commands")
|
||||||
|
case monitor && (install || remove || forceUpdate || updateNow || reloadEnv):
|
||||||
|
fatalUsage("error: --monitor cannot be combined with other commands")
|
||||||
case install:
|
case install:
|
||||||
os.Exit(serviceCommand(configPath, true, noCopy, elevatedChild))
|
os.Exit(serviceCommand(configPath, true, noCopy, elevatedChild))
|
||||||
case remove:
|
case remove:
|
||||||
os.Exit(serviceCommand(configPath, false, noCopy, elevatedChild))
|
os.Exit(serviceCommand(configPath, false, noCopy, elevatedChild))
|
||||||
case forceUpdate:
|
case forceUpdate:
|
||||||
os.Exit(forceUpdateCommand(configPath, elevatedChild))
|
os.Exit(forceUpdateCommand(configPath, elevatedChild))
|
||||||
|
case updateNow:
|
||||||
|
os.Exit(updateNowCommand())
|
||||||
|
case reloadEnv:
|
||||||
|
os.Exit(reloadEnvCommand())
|
||||||
|
case monitor:
|
||||||
|
os.Exit(monitorCommand())
|
||||||
}
|
}
|
||||||
if len(args) > 0 {
|
if len(args) > 0 {
|
||||||
fatalUsage("error: unknown arguments: %s", strings.Join(args, " "))
|
fatalUsage("error: unknown arguments: %s", strings.Join(args, " "))
|
||||||
@@ -181,7 +213,7 @@ func main() {
|
|||||||
syncEnvFile(resolveConfigPath(configPath), cfg.LogFile, log)
|
syncEnvFile(resolveConfigPath(configPath), cfg.LogFile, log)
|
||||||
|
|
||||||
if service.IsService() {
|
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)
|
log.Error("service failed", "err", err)
|
||||||
os.Exit(1)
|
os.Exit(1)
|
||||||
}
|
}
|
||||||
@@ -189,7 +221,7 @@ func main() {
|
|||||||
}
|
}
|
||||||
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||||
defer stop()
|
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)
|
log.Error("listener failed", "err", err)
|
||||||
os.Exit(1)
|
os.Exit(1)
|
||||||
}
|
}
|
||||||
@@ -341,12 +373,23 @@ func forceUpdateCommand(configPath string, elevatedChild bool) int {
|
|||||||
fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot locate executable: %v\n", err)
|
fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot locate executable: %v\n", err)
|
||||||
return 1
|
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)
|
log, _, logCloser := newLogger(cfg)
|
||||||
defer logCloser.Close()
|
defer logCloser.Close()
|
||||||
if cfg.AppVersion == "dev" {
|
if cfg.AppVersion == "dev" {
|
||||||
fmt.Printf("%s: APP_VER=dev, updates disabled\n", versionLine())
|
fmt.Printf("%s: APP_VER=dev, updates disabled\n", versionLine())
|
||||||
return 0
|
return 0
|
||||||
}
|
}
|
||||||
|
// A running service can do the privileged work itself (its account owns
|
||||||
|
// the install dir and it knows when the GPU is idle): ask it over the
|
||||||
|
// local control channel first, no admin rights needed. Fails fast when
|
||||||
|
// no service is listening, in which case we do the direct check below.
|
||||||
|
if reply, err := control.Ask(control.CmdUpdateNow); err == nil {
|
||||||
|
return printControlReply(reply)
|
||||||
|
}
|
||||||
u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
|
u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
|
||||||
|
|
||||||
// Single-shot: one attempt, fail fast when the server is unreachable
|
// Single-shot: one attempt, fail fast when the server is unreachable
|
||||||
@@ -397,6 +440,40 @@ func forceUpdateCommand(configPath string, elevatedChild bool) int {
|
|||||||
return 0
|
return 0
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// printControlReply prints the service's answer to a control-channel
|
||||||
|
// request: "OK ..." on stdout (exit 0), "ERR ..." on stderr (exit 1).
|
||||||
|
func printControlReply(reply string) int {
|
||||||
|
if msg, ok := strings.CutPrefix(reply, "ERR "); ok {
|
||||||
|
fmt.Fprintf(os.Stderr, "gpu-turnstile: %s\n", msg)
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
fmt.Println("gpu-turnstile: " + strings.TrimPrefix(reply, "OK "))
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// updateNowCommand only goes through the running service's control channel
|
||||||
|
// (no direct check, no elevation): the unprivileged update trigger.
|
||||||
|
func updateNowCommand() int {
|
||||||
|
reply, err := control.Ask(control.CmdUpdateNow)
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: no running service to ask — use --force-update for a direct check\n", versionLine())
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
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
|
// reportElevatedUpdate prints the parent's summary of an elevated
|
||||||
// --force-update child: exitCodeStaged means the child staged a new binary,
|
// --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
|
// 0 means it found nothing to do. to is the tag the parent's own check
|
||||||
@@ -477,7 +554,7 @@ func managedComfyCommand(cfg config.Config) string {
|
|||||||
return ""
|
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
|
// ComfyUI can run as a managed child — COMFY_CMD verbatim, or the
|
||||||
// standard venv layout derived from COMFY_DIR alone: started on demand
|
// standard venv layout derived from COMFY_DIR alone: started on demand
|
||||||
// by the proxy, stopped after COMFY_IDLE_TIMEOUT idle (and on shutdown)
|
// by the proxy, stopped after COMFY_IDLE_TIMEOUT idle (and on shutdown)
|
||||||
@@ -490,6 +567,8 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
"version", version,
|
"version", version,
|
||||||
"listen_ollama", cfg.ListenOllama,
|
"listen_ollama", cfg.ListenOllama,
|
||||||
"listen_comfy", cfg.ListenComfy,
|
"listen_comfy", cfg.ListenComfy,
|
||||||
|
"listen_ui", orDisabled(cfg.ListenUI),
|
||||||
|
"ui_auth", cfg.UIUser != "",
|
||||||
"ollama_url", orDisabled(cfg.OllamaURL),
|
"ollama_url", orDisabled(cfg.OllamaURL),
|
||||||
"comfy_url", orDisabled(cfg.ComfyURL),
|
"comfy_url", orDisabled(cfg.ComfyURL),
|
||||||
"unload_timeout", cfg.UnloadTimeout,
|
"unload_timeout", cfg.UnloadTimeout,
|
||||||
@@ -515,6 +594,7 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
"comfy_start_timeout", cfg.ComfyStartTimeout,
|
"comfy_start_timeout", cfg.ComfyStartTimeout,
|
||||||
"game_procs", cfg.GameProcs,
|
"game_procs", cfg.GameProcs,
|
||||||
"gpu_foreign_vram_mb", cfg.GPUForeignVRAMMB,
|
"gpu_foreign_vram_mb", cfg.GPUForeignVRAMMB,
|
||||||
|
"gpu_foreign_util_pct", cfg.GPUForeignUtilPct,
|
||||||
"gpu_ignore_procs", cfg.GPUIgnoreProcs,
|
"gpu_ignore_procs", cfg.GPUIgnoreProcs,
|
||||||
"game_poll_interval", cfg.GamePollInterval,
|
"game_poll_interval", cfg.GamePollInterval,
|
||||||
"auto_update", cfg.AutoUpdate,
|
"auto_update", cfg.AutoUpdate,
|
||||||
@@ -527,6 +607,8 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
)
|
)
|
||||||
|
|
||||||
lk := lock.New(log)
|
lk := lock.New(log)
|
||||||
|
health := newHealthTracker()
|
||||||
|
started := time.Now()
|
||||||
// Each consumer is enabled by setting its URL; a disabled consumer gets
|
// Each consumer is enabled by setting its URL; a disabled consumer gets
|
||||||
// no client, no listener and no probe.
|
// no client, no listener and no probe.
|
||||||
var ollamaClient *ollama.Client
|
var ollamaClient *ollama.Client
|
||||||
@@ -636,20 +718,25 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
}
|
}
|
||||||
probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout)
|
probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout)
|
||||||
for name, probe := range probes {
|
for name, probe := range probes {
|
||||||
if err := probe(probeCtx); err != nil && !errors.Is(err, errManagedDown) {
|
err := probe(probeCtx)
|
||||||
|
health.set(name, err == nil)
|
||||||
|
if err != nil && !errors.Is(err, errManagedDown) {
|
||||||
log.Warn(name+" probe failed", "err", err)
|
log.Warn(name+" probe failed", "err", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
probeCancel()
|
probeCancel()
|
||||||
if cfg.HealthInterval > 0 {
|
if cfg.HealthInterval > 0 {
|
||||||
go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes)
|
go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes, health)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS
|
// Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS,
|
||||||
// and/or GPU_FOREIGN_VRAM_MB — hold the lock externally while they run.
|
// GPU_FOREIGN_VRAM_MB and/or GPU_FOREIGN_UTIL_PCT — hold the lock
|
||||||
if len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0 {
|
// externally while they run. gw collects the VRAM reading and the last
|
||||||
det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUIgnoreProcs, log)
|
// check result for the status channel.
|
||||||
go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup)
|
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
|
// Bind the listeners up front so a port conflict fails fast and the
|
||||||
@@ -687,8 +774,63 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
service.NotifyReady()
|
service.NotifyReady()
|
||||||
service.StartWatchdog(ctx)
|
service.StartWatchdog(ctx)
|
||||||
|
|
||||||
|
var u *update.Updater
|
||||||
|
var exePath string
|
||||||
|
// restartWhenIdle exits with exitCodeUpdate once the GPU lock is idle;
|
||||||
|
// the service recovery configuration brings the process back. Shared by
|
||||||
|
// staged updates and config reloads; the once guard makes repeated
|
||||||
|
// triggers idempotent.
|
||||||
|
var restartOnce sync.Once
|
||||||
|
restartWhenIdle := func(reason string) {
|
||||||
|
log.Warn("restarting once the GPU is idle", "reason", reason)
|
||||||
|
restartOnce.Do(func() {
|
||||||
|
go func() {
|
||||||
|
if waitForIdle(ctx, lk, 24*time.Hour) {
|
||||||
|
log.Warn("restarting now", "reason", reason)
|
||||||
|
os.Exit(exitCodeUpdate)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
applyStaged := func(to string) {}
|
||||||
if cfg.AutoUpdate {
|
if cfg.AutoUpdate {
|
||||||
go updateLoop(ctx, cfg, log, lk, isService)
|
if p, err := os.Executable(); err != nil {
|
||||||
|
log.Warn("auto-update disabled: cannot locate executable", "err", err)
|
||||||
|
} else {
|
||||||
|
exePath = p
|
||||||
|
u = &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
|
||||||
|
applyStaged = func(to string) {
|
||||||
|
if !isService {
|
||||||
|
log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply", "version", to)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
restartWhenIdle("update to " + to)
|
||||||
|
}
|
||||||
|
go updateLoop(ctx, cfg.UpdateInterval, log, u, exePath, applyStaged)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// The command handler backs both the local control channel (served
|
||||||
|
// whenever running as a service) and the optional web UI (LISTEN_UI),
|
||||||
|
// independent of AUTO_UPDATE.
|
||||||
|
handler := controlHandler(ctx, u, exePath, applyStaged,
|
||||||
|
statusProvider(cfg, lk, comfySup, health, started, gw, ollamaClient),
|
||||||
|
reloadHandler(cfg, configPath, restartWhenIdle))
|
||||||
|
if isService {
|
||||||
|
serveControl(ctx, log, handler)
|
||||||
|
}
|
||||||
|
if cfg.ListenUI != "" {
|
||||||
|
ln, err := net.Listen("tcp", cfg.ListenUI)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("listen ui on %s: %w", cfg.ListenUI, err)
|
||||||
|
}
|
||||||
|
if host, _, err := net.SplitHostPort(cfg.ListenUI); err == nil && !isLoopbackHost(host) && cfg.UIUser == "" {
|
||||||
|
log.Warn("web UI is reachable from other machines WITHOUT auth; set UI_USER/UI_PASS or bind to 127.0.0.1", "addr", cfg.ListenUI)
|
||||||
|
}
|
||||||
|
uiSrv := &http.Server{Addr: cfg.ListenUI, Handler: uiHandler(handler, cfg.UIUser, cfg.UIPass)}
|
||||||
|
servers = append(servers, uiSrv)
|
||||||
|
log.Warn("listening", "consumer", "ui", "addr", cfg.ListenUI)
|
||||||
|
go func() { errCh <- uiSrv.Serve(ln) }()
|
||||||
}
|
}
|
||||||
|
|
||||||
select {
|
select {
|
||||||
@@ -712,12 +854,52 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
// (idle); health checks skip it instead of logging an outage.
|
// (idle); health checks skip it instead of logging an outage.
|
||||||
var errManagedDown = errors.New("managed upstream intentionally stopped")
|
var errManagedDown = errors.New("managed upstream intentionally stopped")
|
||||||
|
|
||||||
|
// gpuWatch records the latest VRAM reading and game-detector finding from
|
||||||
|
// the game detector's poll loop, for the status channel. Enabled is false
|
||||||
|
// when game detection is not configured (no polling happens then); Known is
|
||||||
|
// false until the first successful nvidia-smi reading.
|
||||||
|
type gpuWatch struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
enabled bool
|
||||||
|
usedMB int
|
||||||
|
total int
|
||||||
|
tempC int
|
||||||
|
fanPct int
|
||||||
|
known bool
|
||||||
|
foreign string // last Check result: external holders, "" when none
|
||||||
|
at time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *gpuWatch) setVRAM(st game.GPUStats) {
|
||||||
|
g.mu.Lock()
|
||||||
|
g.usedMB, g.total = st.UsedMB, st.TotalMB
|
||||||
|
g.tempC, g.fanPct, g.known = st.TempC, st.FanPct, true
|
||||||
|
g.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *gpuWatch) setCheck(foreign string) {
|
||||||
|
g.mu.Lock()
|
||||||
|
g.foreign, g.at = foreign, time.Now()
|
||||||
|
g.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *gpuWatch) get() (st game.GPUStats, known bool, foreign string, ageS int64) {
|
||||||
|
g.mu.Lock()
|
||||||
|
defer g.mu.Unlock()
|
||||||
|
ageS = -1
|
||||||
|
if !g.at.IsZero() {
|
||||||
|
ageS = int64(time.Since(g.at).Seconds())
|
||||||
|
}
|
||||||
|
return game.GPUStats{UsedMB: g.usedMB, TotalMB: g.total, TempC: g.tempC, FanPct: g.fanPct},
|
||||||
|
g.known, g.foreign, ageS
|
||||||
|
}
|
||||||
|
|
||||||
// gameLoop polls for foreign GPU holders (a game, another ML job). While one
|
// 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
|
// 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
|
// 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
|
// drained — frees VRAM for it: the managed ComfyUI is stopped and Ollama's
|
||||||
// resident models are unloaded.
|
// 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)
|
ticker := time.NewTicker(cfg.GamePollInterval)
|
||||||
defer ticker.Stop()
|
defer ticker.Stop()
|
||||||
held, freed := false, false
|
held, freed := false, false
|
||||||
@@ -731,6 +913,10 @@ func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *gam
|
|||||||
if err != nil && ctx.Err() == nil {
|
if err != nil && ctx.Err() == nil {
|
||||||
log.Warn("game detection failed", "err", err)
|
log.Warn("game detection failed", "err", err)
|
||||||
}
|
}
|
||||||
|
gw.setCheck(summarizeHolders(holders))
|
||||||
|
if st, verr := game.QueryGPUStats(ctx); verr == nil {
|
||||||
|
gw.setVRAM(st)
|
||||||
|
}
|
||||||
switch {
|
switch {
|
||||||
case len(holders) > 0 && !held:
|
case len(holders) > 0 && !held:
|
||||||
held = true
|
held = true
|
||||||
@@ -790,7 +976,8 @@ func freeVRAM(ctx context.Context, unloadTimeout time.Duration, log *slog.Logger
|
|||||||
// transitions — "is DOWN" when a previously healthy upstream stops
|
// transitions — "is DOWN" when a previously healthy upstream stops
|
||||||
// answering, "recovered" when it comes back. The first round only
|
// answering, "recovered" when it comes back. The first round only
|
||||||
// establishes the baseline; the startup probe already reported that state.
|
// establishes the baseline; the startup probe already reported that state.
|
||||||
func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *slog.Logger, probes map[string]func(context.Context) error) {
|
// Every result goes into the tracker for the status channel.
|
||||||
|
func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *slog.Logger, probes map[string]func(context.Context) error, tracker *healthTracker) {
|
||||||
ticker := time.NewTicker(interval)
|
ticker := time.NewTicker(interval)
|
||||||
defer ticker.Stop()
|
defer ticker.Stop()
|
||||||
up := map[string]bool{}
|
up := map[string]bool{}
|
||||||
@@ -807,9 +994,11 @@ func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *
|
|||||||
cancel()
|
cancel()
|
||||||
if errors.Is(err, errManagedDown) {
|
if errors.Is(err, errManagedDown) {
|
||||||
managed[name] = true // intentionally stopped; not an outage
|
managed[name] = true // intentionally stopped; not an outage
|
||||||
|
tracker.set(name, false)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
now := err == nil
|
now := err == nil
|
||||||
|
tracker.set(name, now)
|
||||||
if managed[name] {
|
if managed[name] {
|
||||||
// First real probe after an idle stop only re-baselines —
|
// First real probe after an idle stop only re-baselines —
|
||||||
// an on-demand start is not a "recovery".
|
// an on-demand start is not a "recovery".
|
||||||
@@ -830,42 +1019,275 @@ func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// updateLoop checks for signed updates on startup and every UPDATE_INTERVAL.
|
// updateLoop checks for signed updates on startup and every interval;
|
||||||
// In service mode a staged update is applied by exiting with exitCodeUpdate
|
// applyStaged decides what a staged update means (restart when idle as a
|
||||||
// once the GPU lock is idle; the service recovery configuration restarts the
|
// service, log only interactively).
|
||||||
// process with the new binary. Interactively it only logs.
|
func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string)) {
|
||||||
func updateLoop(ctx context.Context, cfg config.Config, log *slog.Logger, lk *lock.Lock, isService bool) {
|
|
||||||
exePath, err := os.Executable()
|
|
||||||
if err != nil {
|
|
||||||
log.Warn("auto-update disabled: cannot locate executable", "err", err)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
|
|
||||||
for {
|
for {
|
||||||
staged, to, err := u.Check(ctx, exePath)
|
staged, to, err := u.Check(ctx, exePath)
|
||||||
if err != nil && ctx.Err() == nil {
|
if err != nil && ctx.Err() == nil {
|
||||||
log.Warn("auto-update check failed", "err", err)
|
log.Warn("auto-update check failed", "err", err)
|
||||||
}
|
}
|
||||||
if staged {
|
if staged {
|
||||||
if !isService {
|
applyStaged(to)
|
||||||
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)
|
|
||||||
if waitForIdle(ctx, lk, 24*time.Hour) {
|
|
||||||
log.Warn("auto-update: restarting to apply update")
|
|
||||||
os.Exit(exitCodeUpdate)
|
|
||||||
}
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
return
|
return
|
||||||
case <-time.After(cfg.UpdateInterval):
|
case <-time.After(interval):
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// controlHandler builds the command handler shared by the control channel
|
||||||
|
// (named pipe / unix socket) and the web UI's action endpoints: status for
|
||||||
|
// --monitor / the UI, an update check (--force-update/--update-now) and a
|
||||||
|
// config reload (--reload-env), all safe for unprivileged local users. The
|
||||||
|
// update payload is signature-verified regardless of who asks; update
|
||||||
|
// triggers are rate-limited to one per minute and reloads to one per two
|
||||||
|
// seconds so neither can be used to spam restarts or disk churn. u is nil
|
||||||
|
// when AUTO_UPDATE=false.
|
||||||
|
func controlHandler(ctx context.Context, u *update.Updater, exePath string, applyStaged func(to string), status func() string, reload func() string) control.Handler {
|
||||||
|
var mu sync.Mutex
|
||||||
|
var lastTrigger, lastReload time.Time
|
||||||
|
return func(cmd string) string {
|
||||||
|
switch cmd {
|
||||||
|
case control.CmdStatus:
|
||||||
|
return "OK " + status()
|
||||||
|
case control.CmdReloadEnv:
|
||||||
|
// reload re-reads the config file from disk; a short limiter
|
||||||
|
// keeps a local flood from turning into disk churn.
|
||||||
|
mu.Lock()
|
||||||
|
if wait := 2*time.Second - time.Since(lastReload); wait > 0 {
|
||||||
|
mu.Unlock()
|
||||||
|
return fmt.Sprintf("ERR rate limited: retry in %ds", int(wait.Seconds())+1)
|
||||||
|
}
|
||||||
|
lastReload = time.Now()
|
||||||
|
mu.Unlock()
|
||||||
|
return reload()
|
||||||
|
case control.CmdUpdateNow:
|
||||||
|
default:
|
||||||
|
return "ERR unknown command: " + cmd
|
||||||
|
}
|
||||||
|
if u == nil {
|
||||||
|
return "ERR auto-update is disabled on this instance"
|
||||||
|
}
|
||||||
|
mu.Lock()
|
||||||
|
if wait := time.Minute - time.Since(lastTrigger); wait > 0 {
|
||||||
|
mu.Unlock()
|
||||||
|
return fmt.Sprintf("ERR rate limited: retry in %ds", int(wait.Seconds())+1)
|
||||||
|
}
|
||||||
|
lastTrigger = time.Now()
|
||||||
|
mu.Unlock()
|
||||||
|
cctx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
staged, to, err := u.Check(cctx, exePath)
|
||||||
|
if err != nil {
|
||||||
|
return "ERR update check failed: " + err.Error()
|
||||||
|
}
|
||||||
|
if !staged {
|
||||||
|
return "OK " + version + " is up to date"
|
||||||
|
}
|
||||||
|
applyStaged(to)
|
||||||
|
return "OK updated from " + version + " to " + to + "; the service restarts once the GPU is idle"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// serveControl opens the local control channel (named pipe on Windows,
|
||||||
|
// unix socket on Linux) for unprivileged local users (--monitor,
|
||||||
|
// --force-update/--update-now, --reload-env).
|
||||||
|
func serveControl(ctx context.Context, log *slog.Logger, h control.Handler) {
|
||||||
|
if err := control.Serve(ctx, h, log); err != nil {
|
||||||
|
log.Warn("control channel disabled", "err", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// healthTracker records the latest probe result per upstream for the
|
||||||
|
// status channel.
|
||||||
|
type healthTracker struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
up map[string]bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func newHealthTracker() *healthTracker {
|
||||||
|
return &healthTracker{up: map[string]bool{}}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *healthTracker) set(name string, up bool) {
|
||||||
|
h.mu.Lock()
|
||||||
|
h.up[name] = up
|
||||||
|
h.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *healthTracker) get(name string) bool {
|
||||||
|
h.mu.Lock()
|
||||||
|
defer h.mu.Unlock()
|
||||||
|
return h.up[name]
|
||||||
|
}
|
||||||
|
|
||||||
|
// statusDownstream/statusLock/statusSnapshot are the JSON the control
|
||||||
|
// channel serves on CmdStatus; the monitor mode renders them.
|
||||||
|
type statusDownstream struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
URL string `json:"url"`
|
||||||
|
Up bool `json:"up"`
|
||||||
|
Managed string `json:"managed,omitempty"`
|
||||||
|
// Models lists Ollama's loaded models with their VRAM footprint. Null
|
||||||
|
// when unknown (query failed / not applicable); [] means none loaded —
|
||||||
|
// deliberately no omitempty so the two stay distinguishable.
|
||||||
|
Models []statusModel `json:"models"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// statusModel is one loaded Ollama model.
|
||||||
|
type statusModel struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
VRAMMB int64 `json:"vram_mb"` // 0 = resident in RAM, not VRAM
|
||||||
|
}
|
||||||
|
|
||||||
|
type statusLock struct {
|
||||||
|
State string `json:"state"`
|
||||||
|
Detail string `json:"detail,omitempty"`
|
||||||
|
LLMInflight int `json:"llm_inflight"`
|
||||||
|
LLMWaiting int `json:"llm_waiting"`
|
||||||
|
ImageQueue int `json:"image_queue"`
|
||||||
|
External string `json:"external,omitempty"`
|
||||||
|
SinceS int64 `json:"since_s"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type statusSnapshot struct {
|
||||||
|
Version string `json:"version"`
|
||||||
|
UptimeS int64 `json:"uptime_s"`
|
||||||
|
Downstreams []statusDownstream `json:"downstreams"`
|
||||||
|
Lock statusLock `json:"lock"`
|
||||||
|
// GPU carries the latest VRAM reading and detector finding; Enabled is
|
||||||
|
// false when game detection (and with it nvidia-smi polling) is not
|
||||||
|
// configured.
|
||||||
|
GPU statusGPU `json:"gpu"`
|
||||||
|
// UI is the web UI's URL when LISTEN_UI is active, empty otherwise.
|
||||||
|
UI string `json:"ui,omitempty"`
|
||||||
|
// MonitorNote is set client-side (never over the wire) when the
|
||||||
|
// monitor's own binary differs from the service's version.
|
||||||
|
MonitorNote string `json:"-"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type statusGPU struct {
|
||||||
|
// Enabled reports whether game detection is configured (and with it
|
||||||
|
// VRAM polling); when false the other fields carry no information.
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
UsedMB int `json:"used_mb"`
|
||||||
|
TotalMB int `json:"total_mb"`
|
||||||
|
// TempC/FanPct are -1 when unknown (never sampled or nvidia-smi
|
||||||
|
// reported N/A).
|
||||||
|
TempC int `json:"temp_c"`
|
||||||
|
FanPct int `json:"fan_pct"`
|
||||||
|
Known bool `json:"known"`
|
||||||
|
// Foreign is the last detector finding (external GPU holders), empty
|
||||||
|
// when the last check found none.
|
||||||
|
Foreign string `json:"foreign,omitempty"`
|
||||||
|
// AgeS is how long ago the last check ran; -1 before the first check.
|
||||||
|
AgeS int64 `json:"age_s"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// reloadHandler re-reads and validates the service's config file for
|
||||||
|
// CmdReloadEnv. An invalid config is reported and the service keeps running
|
||||||
|
// untouched; a valid, changed config triggers a GPU-idle-gated restart onto
|
||||||
|
// it (same mechanism as staged updates); unchanged is a no-op.
|
||||||
|
func reloadHandler(current config.Config, configPath string, restartWhenIdle func(reason string)) func() string {
|
||||||
|
return func() string {
|
||||||
|
ncfg, err := loadMergedConfig(configPath)
|
||||||
|
if err != nil {
|
||||||
|
return "ERR config invalid: " + err.Error()
|
||||||
|
}
|
||||||
|
changed := diffConfig(current, ncfg)
|
||||||
|
if len(changed) == 0 {
|
||||||
|
return "OK config unchanged"
|
||||||
|
}
|
||||||
|
restartWhenIdle("config reload (" + strings.Join(changed, ", ") + ")")
|
||||||
|
return "OK config valid; restarting once the GPU is idle (changed: " + strings.Join(changed, ", ") + ")"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// diffConfig lists the env names of settings whose values differ between
|
||||||
|
// two configs (from each field's env tag, so users recognize them).
|
||||||
|
func diffConfig(a, b config.Config) []string {
|
||||||
|
va, vb := reflect.ValueOf(a), reflect.ValueOf(b)
|
||||||
|
t := va.Type()
|
||||||
|
var out []string
|
||||||
|
for i := 0; i < t.NumField(); i++ {
|
||||||
|
if !reflect.DeepEqual(va.Field(i).Interface(), vb.Field(i).Interface()) {
|
||||||
|
name := t.Field(i).Tag.Get("env")
|
||||||
|
if name == "" {
|
||||||
|
name = t.Field(i).Name
|
||||||
|
}
|
||||||
|
out = append(out, name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// statusProvider assembles the one-line JSON snapshot for CmdStatus. The
|
||||||
|
// loaded-model query to Ollama gets a short timeout so a wedged upstream
|
||||||
|
// cannot stall the status channel for long.
|
||||||
|
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time, gw *gpuWatch, ollamaClient *ollama.Client) func() string {
|
||||||
|
return func() string {
|
||||||
|
snap := statusSnapshot{
|
||||||
|
Version: version,
|
||||||
|
UptimeS: int64(time.Since(started).Seconds()),
|
||||||
|
UI: uiURL(cfg.ListenUI),
|
||||||
|
}
|
||||||
|
st, known, foreign, ageS := gw.get()
|
||||||
|
snap.GPU = statusGPU{
|
||||||
|
Enabled: gw.enabled,
|
||||||
|
UsedMB: st.UsedMB, TotalMB: st.TotalMB, Known: known,
|
||||||
|
TempC: -1, FanPct: -1,
|
||||||
|
Foreign: foreign, AgeS: ageS,
|
||||||
|
}
|
||||||
|
if known {
|
||||||
|
snap.GPU.TempC, snap.GPU.FanPct = st.TempC, st.FanPct
|
||||||
|
}
|
||||||
|
if cfg.OllamaURL != "" {
|
||||||
|
d := statusDownstream{
|
||||||
|
Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
|
||||||
|
}
|
||||||
|
if ollamaClient != nil {
|
||||||
|
mctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||||
|
models, err := ollamaClient.LoadedModelDetails(mctx)
|
||||||
|
cancel()
|
||||||
|
if err == nil {
|
||||||
|
d.Models = make([]statusModel, 0, len(models))
|
||||||
|
for _, m := range models {
|
||||||
|
d.Models = append(d.Models, statusModel{Name: m.Name, VRAMMB: m.SizeVRAM / (1024 * 1024)})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
snap.Downstreams = append(snap.Downstreams, d)
|
||||||
|
}
|
||||||
|
if cfg.ComfyURL != "" {
|
||||||
|
d := statusDownstream{Name: "comfy", URL: cfg.ComfyURL, Up: health.get("comfy")}
|
||||||
|
if comfySup != nil {
|
||||||
|
d.Managed = comfySup.Status()
|
||||||
|
}
|
||||||
|
snap.Downstreams = append(snap.Downstreams, d)
|
||||||
|
}
|
||||||
|
lst := lk.Status()
|
||||||
|
snap.Lock = statusLock{
|
||||||
|
State: string(lst.State),
|
||||||
|
Detail: lst.Detail,
|
||||||
|
LLMInflight: lst.LLMInflight,
|
||||||
|
LLMWaiting: lst.LLMWaiting,
|
||||||
|
ImageQueue: lst.ImageQueue,
|
||||||
|
External: lst.External,
|
||||||
|
SinceS: int64(time.Since(lst.Since).Seconds()),
|
||||||
|
}
|
||||||
|
b, err := json.Marshal(snap)
|
||||||
|
if err != nil {
|
||||||
|
return `{"version":"` + version + `"}`
|
||||||
|
}
|
||||||
|
return string(b)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// waitForIdle polls the lock until no LLM or image work is active or
|
// waitForIdle polls the lock until no LLM or image work is active or
|
||||||
// pending, max at most. Returns false on timeout or cancellation.
|
// pending, max at most. Returns false on timeout or cancellation.
|
||||||
func waitForIdle(ctx context.Context, lk *lock.Lock, max time.Duration) bool {
|
func waitForIdle(ctx context.Context, lk *lock.Lock, max time.Duration) bool {
|
||||||
|
|||||||
@@ -0,0 +1,345 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"gpu-turnstile/internal/control"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ANSI colors for the monitor frame.
|
||||||
|
const (
|
||||||
|
cReset = "\x1b[0m"
|
||||||
|
cDim = "\x1b[2m"
|
||||||
|
cRed = "\x1b[31m"
|
||||||
|
cGreen = "\x1b[32m"
|
||||||
|
cYellow = "\x1b[33m"
|
||||||
|
cCyan = "\x1b[36m"
|
||||||
|
)
|
||||||
|
|
||||||
|
// hotkeysLine is the monitor's footer.
|
||||||
|
const hotkeysLine = " " + cDim + "q quit · u update now · r reload config" + cReset + "\x1b[K\n"
|
||||||
|
|
||||||
|
// monitorCommand renders a live status view of the running service,
|
||||||
|
// refreshed every second from the control channel. When the service
|
||||||
|
// reports a different version and the executable on disk changed (the
|
||||||
|
// updater replaced it), the monitor restarts itself onto the new binary.
|
||||||
|
// Hotkeys: q quits, u triggers an update check on the service, r asks the
|
||||||
|
// service to reload its config file.
|
||||||
|
func monitorCommand() int {
|
||||||
|
if !stdoutIsTerminal() {
|
||||||
|
fmt.Fprintln(os.Stderr, "gpu-turnstile: --monitor needs an interactive terminal")
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
enableVirtualTerminal()
|
||||||
|
restore := enableRawKeys()
|
||||||
|
defer func() {
|
||||||
|
if restore != nil {
|
||||||
|
restore()
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
fmt.Print("\x1b[2J") // clear once; frames then redraw in place
|
||||||
|
defer fmt.Print(cReset + "\n")
|
||||||
|
exe, _ := os.Executable()
|
||||||
|
var exeStamp time.Time
|
||||||
|
if st, err := os.Stat(exe); err == nil {
|
||||||
|
exeStamp = st.ModTime()
|
||||||
|
}
|
||||||
|
|
||||||
|
keys := make(chan byte, 8)
|
||||||
|
go readKeys(keys)
|
||||||
|
ticker := time.NewTicker(time.Second)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
var note string
|
||||||
|
var noteAt time.Time
|
||||||
|
noteCh := make(chan string, 1)
|
||||||
|
askPending := false
|
||||||
|
|
||||||
|
// ask sends a one-shot command to the service and reports the reply in
|
||||||
|
// the note line. Only one ask runs at a time.
|
||||||
|
ask := func(cmd, busy, label string) {
|
||||||
|
if askPending {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
askPending = true
|
||||||
|
note, noteAt = busy, time.Now()
|
||||||
|
go func() {
|
||||||
|
reply, err := control.Ask(cmd)
|
||||||
|
if err != nil {
|
||||||
|
noteCh <- label + ": no answer from the service"
|
||||||
|
return
|
||||||
|
}
|
||||||
|
msg := strings.TrimPrefix(reply, "OK ")
|
||||||
|
msg = strings.TrimPrefix(msg, "ERR ")
|
||||||
|
noteCh <- label + ": " + msg
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
poll := func() string {
|
||||||
|
frame := renderWaiting()
|
||||||
|
if reply, err := control.Ask(control.CmdStatus); err == nil {
|
||||||
|
if msg, ok := strings.CutPrefix(reply, "OK "); ok {
|
||||||
|
var snap statusSnapshot
|
||||||
|
if json.Unmarshal([]byte(msg), &snap) == nil {
|
||||||
|
if snap.Version != "" && snap.Version != version {
|
||||||
|
if exeChanged(exe, exeStamp) {
|
||||||
|
fmt.Print("\x1b[2J\x1b[H")
|
||||||
|
fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version)
|
||||||
|
restartSelf(exe, "--monitor")
|
||||||
|
return "" // re-execed; this process exits below
|
||||||
|
}
|
||||||
|
snap.MonitorNote = fmt.Sprintf("note: the service runs %s, this monitor is %s", snap.Version, version)
|
||||||
|
}
|
||||||
|
if note != "" {
|
||||||
|
snap.MonitorNote = note
|
||||||
|
}
|
||||||
|
frame = renderMonitor(snap, termWidth())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return frame
|
||||||
|
}
|
||||||
|
|
||||||
|
for {
|
||||||
|
frame := poll()
|
||||||
|
if frame == "" {
|
||||||
|
return 0 // restartSelf fired
|
||||||
|
}
|
||||||
|
fmt.Print("\x1b[H" + frame + "\x1b[J") // home, frame, clear below
|
||||||
|
select {
|
||||||
|
case <-ticker.C:
|
||||||
|
if note != "" && time.Since(noteAt) > 15*time.Second {
|
||||||
|
note = ""
|
||||||
|
}
|
||||||
|
case k, ok := <-keys:
|
||||||
|
if !ok {
|
||||||
|
keys = nil
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
switch k {
|
||||||
|
case 'q', 'Q', 3: // q or Ctrl+C (raw mode delivers it as a byte)
|
||||||
|
return 0
|
||||||
|
case 'u', 'U':
|
||||||
|
ask(control.CmdUpdateNow, "checking for updates…", "update")
|
||||||
|
case 'r', 'R':
|
||||||
|
ask(control.CmdReloadEnv, "reloading config…", "reload")
|
||||||
|
}
|
||||||
|
case n := <-noteCh:
|
||||||
|
askPending = false
|
||||||
|
note, noteAt = n, time.Now()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// readKeys reads single keypresses from stdin (raw mode was enabled by the
|
||||||
|
// caller) and delivers them until stdin fails.
|
||||||
|
func readKeys(keys chan<- byte) {
|
||||||
|
defer close(keys)
|
||||||
|
buf := make([]byte, 1)
|
||||||
|
for {
|
||||||
|
n, err := os.Stdin.Read(buf)
|
||||||
|
if n > 0 {
|
||||||
|
keys <- buf[0]
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// exeChanged reports whether the executable on disk was replaced since the
|
||||||
|
// recorded stamp (the updater swaps it via rename, which changes ModTime).
|
||||||
|
func exeChanged(exe string, stamp time.Time) bool {
|
||||||
|
if exe == "" || stamp.IsZero() {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
st, err := os.Stat(exe)
|
||||||
|
return err == nil && !st.ModTime().Equal(stamp)
|
||||||
|
}
|
||||||
|
|
||||||
|
// restartSelf starts a fresh copy of this executable with the given args on
|
||||||
|
// the same console; the caller exits right after.
|
||||||
|
func restartSelf(exe string, args ...string) {
|
||||||
|
cmd := exec.Command(exe, args...)
|
||||||
|
cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
|
||||||
|
cmd.Start() //nolint:errcheck // best effort: on failure we just exit
|
||||||
|
}
|
||||||
|
|
||||||
|
func renderWaiting() string {
|
||||||
|
return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n\x1b[K\n" + hotkeysLine
|
||||||
|
}
|
||||||
|
|
||||||
|
// renderMonitor draws one full frame. Each line ends with \x1b[K (clear to
|
||||||
|
// end of line) so shrinking content leaves no residue.
|
||||||
|
func renderMonitor(snap statusSnapshot, width int) string {
|
||||||
|
if width < 40 {
|
||||||
|
width = 80
|
||||||
|
}
|
||||||
|
var b strings.Builder
|
||||||
|
|
||||||
|
left := " gpu-turnstile"
|
||||||
|
if snap.UptimeS > 0 {
|
||||||
|
left += " " + cDim + "up " + fmtDur(snap.UptimeS) + cReset
|
||||||
|
}
|
||||||
|
right := snap.Version
|
||||||
|
pad := width - printableLen(" gpu-turnstile up "+fmtDur(snap.UptimeS)) - len(right) - 1
|
||||||
|
if snap.UptimeS == 0 {
|
||||||
|
pad = width - len(" gpu-turnstile") - len(right) - 1
|
||||||
|
}
|
||||||
|
if pad < 1 {
|
||||||
|
pad = 1
|
||||||
|
}
|
||||||
|
b.WriteString(cDim + left + strings.Repeat(" ", pad) + right + cReset + "\x1b[K\n")
|
||||||
|
b.WriteString(cDim + " " + strings.Repeat("─", width-2) + cReset + "\x1b[K\n")
|
||||||
|
|
||||||
|
for _, d := range snap.Downstreams {
|
||||||
|
busy := (d.Name == "ollama" && snap.Lock.State == "llm" && snap.Lock.LLMInflight > 0) ||
|
||||||
|
(d.Name == "comfy" && snap.Lock.State == "image")
|
||||||
|
b.WriteString(renderDownstream(d, busy) + "\x1b[K\n")
|
||||||
|
}
|
||||||
|
b.WriteString("\x1b[K\n")
|
||||||
|
b.WriteString(renderLock(snap.Lock) + "\x1b[K\n")
|
||||||
|
if snap.Lock.ImageQueue > 0 {
|
||||||
|
b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n",
|
||||||
|
cYellow, snap.Lock.ImageQueue, cReset))
|
||||||
|
}
|
||||||
|
if snap.GPU.Enabled {
|
||||||
|
b.WriteString(renderGPU(snap.GPU) + "\x1b[K\n")
|
||||||
|
}
|
||||||
|
if snap.UI != "" {
|
||||||
|
b.WriteString(" " + cDim + "UI: " + snap.UI + cReset + "\x1b[K\n")
|
||||||
|
}
|
||||||
|
if snap.MonitorNote != "" {
|
||||||
|
b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
|
||||||
|
}
|
||||||
|
b.WriteString("\x1b[K\n")
|
||||||
|
b.WriteString(hotkeysLine)
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// renderGPU renders the GPU line: VRAM usage (when nvidia-smi answered),
|
||||||
|
// the game detector's last finding, and how long ago it ran.
|
||||||
|
func renderGPU(g statusGPU) string {
|
||||||
|
s := " GPU: "
|
||||||
|
if g.Known {
|
||||||
|
s += renderVRAM(g.UsedMB, g.TotalMB)
|
||||||
|
if g.TempC >= 0 {
|
||||||
|
s += fmt.Sprintf(" · %d°C", g.TempC)
|
||||||
|
}
|
||||||
|
if g.FanPct >= 0 {
|
||||||
|
s += fmt.Sprintf(" · fan %d%%", g.FanPct)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
s += cDim + "VRAM unknown (nvidia-smi not answering)" + cReset
|
||||||
|
}
|
||||||
|
if g.Foreign != "" {
|
||||||
|
s += " · external: " + cRed + g.Foreign + cReset
|
||||||
|
} else {
|
||||||
|
s += cDim + " · no external process" + cReset
|
||||||
|
}
|
||||||
|
if g.AgeS >= 0 {
|
||||||
|
s += cDim + fmt.Sprintf(" (checked %s ago)", fmtDur(g.AgeS)) + cReset
|
||||||
|
}
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatMB renders MiB as "4.2 GiB" (or MiB below 1 GiB).
|
||||||
|
func formatMB(mb int64) string {
|
||||||
|
if mb >= 1024 {
|
||||||
|
return fmt.Sprintf("%.1f GiB", float64(mb)/1024)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%d MiB", mb)
|
||||||
|
}
|
||||||
|
|
||||||
|
// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
|
||||||
|
func renderVRAM(used, total int) string {
|
||||||
|
if total > 0 {
|
||||||
|
return formatMB(int64(used)) + " / " + formatMB(int64(total)) + " used"
|
||||||
|
}
|
||||||
|
return formatMB(int64(used)) + " used"
|
||||||
|
}
|
||||||
|
|
||||||
|
// printableLen counts characters without ANSI escapes (ASCII-only content).
|
||||||
|
func printableLen(s string) int { return len(s) }
|
||||||
|
|
||||||
|
func renderDownstream(d statusDownstream, busy bool) string {
|
||||||
|
url := cDim + d.URL + cReset
|
||||||
|
detail := ""
|
||||||
|
if d.Models != nil {
|
||||||
|
if len(d.Models) == 0 {
|
||||||
|
detail = cDim + " · no models loaded" + cReset
|
||||||
|
} else {
|
||||||
|
parts := make([]string, 0, len(d.Models))
|
||||||
|
for _, m := range d.Models {
|
||||||
|
if m.VRAMMB > 0 {
|
||||||
|
parts = append(parts, fmt.Sprintf("%s (%s VRAM)", m.Name, formatMB(m.VRAMMB)))
|
||||||
|
} else {
|
||||||
|
parts = append(parts, m.Name+" (in RAM)")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
detail = " · " + strings.Join(parts, ", ")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if busy {
|
||||||
|
detail += " · " + cCyan + "busy" + cReset
|
||||||
|
}
|
||||||
|
switch d.Managed {
|
||||||
|
case "stopped":
|
||||||
|
return fmt.Sprintf(" %s○%s %-8s %sstopped (managed — starts on demand)%s %s",
|
||||||
|
cDim, cReset, d.Name, cDim, cReset, url)
|
||||||
|
case "starting":
|
||||||
|
return fmt.Sprintf(" %s◌%s %-8s %sstarting…%s %s",
|
||||||
|
cYellow, cReset, d.Name, cYellow, cReset, url)
|
||||||
|
}
|
||||||
|
suffix := ""
|
||||||
|
if d.Managed == "external" {
|
||||||
|
suffix = " (external)"
|
||||||
|
}
|
||||||
|
if d.Up {
|
||||||
|
return fmt.Sprintf(" %s●%s %-8s %sUP%s%s%s %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, detail, url)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf(" %s●%s %-8s %sDOWN%s %s", cRed, cReset, d.Name, cRed, cReset, url)
|
||||||
|
}
|
||||||
|
|
||||||
|
func renderLock(l statusLock) string {
|
||||||
|
dur := cDim + "(" + fmtDur(l.SinceS) + ")" + cReset
|
||||||
|
switch l.State {
|
||||||
|
case "idle":
|
||||||
|
return fmt.Sprintf(" Lock: %sidle%s %s", cGreen, cReset, dur)
|
||||||
|
case "llm":
|
||||||
|
s := fmt.Sprintf(" Lock: %sLLM%s — %d in flight", cCyan, cReset, l.LLMInflight)
|
||||||
|
if l.LLMWaiting > 0 {
|
||||||
|
s += fmt.Sprintf(", %d waiting", l.LLMWaiting)
|
||||||
|
}
|
||||||
|
if l.Detail != "" {
|
||||||
|
s += " — " + l.Detail
|
||||||
|
}
|
||||||
|
return s + " " + dur
|
||||||
|
case "image":
|
||||||
|
s := fmt.Sprintf(" Lock: %sIMAGE%s", cYellow, cReset)
|
||||||
|
if l.Detail != "" {
|
||||||
|
s += " — " + l.Detail
|
||||||
|
}
|
||||||
|
return s + " " + dur
|
||||||
|
case "external":
|
||||||
|
return fmt.Sprintf(" Lock: %sEXTERNAL%s — %s %s", cRed, cReset, l.External, dur)
|
||||||
|
}
|
||||||
|
return " Lock: unknown"
|
||||||
|
}
|
||||||
|
|
||||||
|
// fmtDur renders seconds as a compact duration ("1m32s", "2h07m").
|
||||||
|
func fmtDur(s int64) string {
|
||||||
|
if s < 0 {
|
||||||
|
s = 0
|
||||||
|
}
|
||||||
|
d := time.Duration(s) * time.Second
|
||||||
|
if d >= time.Hour {
|
||||||
|
return fmt.Sprintf("%dh%02dm", int(d.Hours()), int(d.Minutes())%60)
|
||||||
|
}
|
||||||
|
return d.Round(time.Second).String()
|
||||||
|
}
|
||||||
@@ -0,0 +1,110 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gpu-turnstile/internal/config"
|
||||||
|
"reflect"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRenderMonitor(t *testing.T) {
|
||||||
|
snap := statusSnapshot{
|
||||||
|
Version: "v0.2.3",
|
||||||
|
UptimeS: 3725,
|
||||||
|
Downstreams: []statusDownstream{
|
||||||
|
{Name: "ollama", URL: "http://127.0.0.1:11435", Up: true},
|
||||||
|
{Name: "comfy", URL: "http://127.0.0.1:8189", Managed: "stopped"},
|
||||||
|
},
|
||||||
|
Lock: statusLock{State: "llm", LLMInflight: 2, LLMWaiting: 1, Detail: "ollama: POST /api/generate", SinceS: 95},
|
||||||
|
}
|
||||||
|
frame := renderMonitor(snap, 80)
|
||||||
|
for _, want := range []string{"v0.2.3", "1h02m", "ollama", "UP", "comfy", "stopped", "LLM", "2 in flight", "1 waiting", "1m35s"} {
|
||||||
|
if !strings.Contains(frame, want) {
|
||||||
|
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
snap.Lock = statusLock{State: "external", External: "cyberpunk2077.exe (pid 1234)", SinceS: 3, ImageQueue: 2}
|
||||||
|
frame = renderMonitor(snap, 40) // narrow: falls back to 80
|
||||||
|
for _, want := range []string{"EXTERNAL", "cyberpunk2077.exe", "2 image job(s) waiting"} {
|
||||||
|
if !strings.Contains(frame, want) {
|
||||||
|
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
snap.GPU = statusGPU{Enabled: true, Known: true, UsedMB: 4300, TotalMB: 16384, TempC: 55, FanPct: 42, Foreign: "cyberpunk2077.exe (pid 1234)", AgeS: 12}
|
||||||
|
frame = renderMonitor(snap, 80)
|
||||||
|
for _, want := range []string{"GPU:", "4.2 GiB / 16.0 GiB used", "55°C", "fan 42%", "external:", "checked 12s ago"} {
|
||||||
|
if !strings.Contains(frame, want) {
|
||||||
|
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
snap.GPU = statusGPU{Enabled: true, AgeS: 4}
|
||||||
|
frame = renderMonitor(snap, 80)
|
||||||
|
for _, want := range []string{"VRAM unknown", "no external process"} {
|
||||||
|
if !strings.Contains(frame, want) {
|
||||||
|
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
snap.Downstreams[0].Models = []statusModel{{Name: "llama3.1:8b", VRAMMB: 4900}, {Name: "embed", VRAMMB: 0}}
|
||||||
|
snap.Lock = statusLock{State: "llm", LLMInflight: 1}
|
||||||
|
frame = renderMonitor(snap, 80)
|
||||||
|
for _, want := range []string{"llama3.1:8b (4.8 GiB VRAM)", "embed (in RAM)", "busy"} {
|
||||||
|
if !strings.Contains(frame, want) {
|
||||||
|
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
snap.Downstreams[0].Models = []statusModel{}
|
||||||
|
snap.Lock = statusLock{State: "idle"}
|
||||||
|
frame = renderMonitor(snap, 80)
|
||||||
|
if !strings.Contains(frame, "no models loaded") {
|
||||||
|
t.Errorf("frame missing %q:\n%s", "no models loaded", frame)
|
||||||
|
}
|
||||||
|
if strings.Contains(frame, "busy") {
|
||||||
|
t.Errorf("idle lock still shows busy:\n%s", frame)
|
||||||
|
}
|
||||||
|
|
||||||
|
snap.UI = "http://127.0.0.1:7860"
|
||||||
|
frame = renderMonitor(snap, 80)
|
||||||
|
if !strings.Contains(frame, "UI: http://127.0.0.1:7860") {
|
||||||
|
t.Errorf("frame missing UI line:\n%s", frame)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFmtDur(t *testing.T) {
|
||||||
|
cases := map[int64]string{0: "0s", 5: "5s", 95: "1m35s", 3725: "1h02m", -3: "0s"}
|
||||||
|
for in, want := range cases {
|
||||||
|
if got := fmtDur(in); got != want {
|
||||||
|
t.Errorf("fmtDur(%d) = %q, want %q", in, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDiffConfig(t *testing.T) {
|
||||||
|
a := config.Defaults()
|
||||||
|
b := a
|
||||||
|
if got := diffConfig(a, b); len(got) != 0 {
|
||||||
|
t.Fatalf("identical configs: got %v", got)
|
||||||
|
}
|
||||||
|
b.LogLevel = -4
|
||||||
|
b.GameProcs = []string{"game.exe"}
|
||||||
|
got := diffConfig(a, b)
|
||||||
|
if len(got) != 2 || got[0] != "GAME_PROCS" || got[1] != "LOGLEVEL" {
|
||||||
|
t.Fatalf("got %v, want [GAME_PROCS LOGLEVEL]", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Every Config field must carry an env tag so user-facing output (the
|
||||||
|
// reload diff) can name the setting the user would actually change.
|
||||||
|
func TestConfigFieldsHaveEnvTags(t *testing.T) {
|
||||||
|
typ := reflect.TypeOf(config.Config{})
|
||||||
|
for i := 0; i < typ.NumField(); i++ {
|
||||||
|
f := typ.Field(i)
|
||||||
|
if f.Tag.Get("env") == "" {
|
||||||
|
t.Errorf("config.Config.%s has no env tag", f.Name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
//go:build !windows
|
||||||
|
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
|
||||||
|
"golang.org/x/sys/unix"
|
||||||
|
)
|
||||||
|
|
||||||
|
// enableVirtualTerminal is a no-op: unix terminals speak ANSI natively.
|
||||||
|
func enableVirtualTerminal() {}
|
||||||
|
|
||||||
|
// termWidth is the terminal width in columns, 80 when unknown.
|
||||||
|
func termWidth() int {
|
||||||
|
ws, err := unix.IoctlGetWinsize(int(os.Stdout.Fd()), unix.TIOCGWINSZ)
|
||||||
|
if err != nil || ws.Col == 0 {
|
||||||
|
return 80
|
||||||
|
}
|
||||||
|
return int(ws.Col)
|
||||||
|
}
|
||||||
|
|
||||||
|
// enableRawKeys switches the terminal to per-keypress mode (ICANON and ECHO
|
||||||
|
// off) and returns the restore function, nil when stdin is not a terminal.
|
||||||
|
func enableRawKeys() func() {
|
||||||
|
fd := int(os.Stdin.Fd())
|
||||||
|
term, err := unix.IoctlGetTermios(fd, unix.TCGETS)
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
raw := *term
|
||||||
|
raw.Lflag &^= unix.ICANON | unix.ECHO
|
||||||
|
raw.Cc[unix.VMIN] = 1
|
||||||
|
raw.Cc[unix.VTIME] = 0
|
||||||
|
if err := unix.IoctlSetTermios(fd, unix.TCSETS, &raw); err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return func() { unix.IoctlSetTermios(fd, unix.TCSETS, term) } //nolint:errcheck
|
||||||
|
}
|
||||||
@@ -0,0 +1,49 @@
|
|||||||
|
//go:build windows
|
||||||
|
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
|
||||||
|
"golang.org/x/sys/windows"
|
||||||
|
)
|
||||||
|
|
||||||
|
// enableVirtualTerminal asks the console to honor ANSI escapes (Windows 10+);
|
||||||
|
// mintty/Git Bash already does, so errors are ignored.
|
||||||
|
func enableVirtualTerminal() {
|
||||||
|
h := windows.Handle(os.Stdout.Fd())
|
||||||
|
var mode uint32
|
||||||
|
if err := windows.GetConsoleMode(h, &mode); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
windows.SetConsoleMode(h, mode|windows.ENABLE_VIRTUAL_TERMINAL_PROCESSING) //nolint:errcheck
|
||||||
|
}
|
||||||
|
|
||||||
|
// termWidth is the console window width in columns, 80 when unknown.
|
||||||
|
func termWidth() int {
|
||||||
|
var info windows.ConsoleScreenBufferInfo
|
||||||
|
if err := windows.GetConsoleScreenBufferInfo(windows.Handle(os.Stdout.Fd()), &info); err != nil {
|
||||||
|
return 80
|
||||||
|
}
|
||||||
|
return int(info.Window.Right-info.Window.Left) + 1
|
||||||
|
}
|
||||||
|
|
||||||
|
// enableRawKeys puts the console's stdin into per-keypress mode (no line
|
||||||
|
// buffering, no echo) and returns the restore function. When stdin is not a
|
||||||
|
// real console (mintty/Git Bash pipes) it returns nil: ptys already deliver
|
||||||
|
// keystrokes immediately.
|
||||||
|
func enableRawKeys() func() {
|
||||||
|
h := windows.Handle(os.Stdin.Fd())
|
||||||
|
var mode uint32
|
||||||
|
if err := windows.GetConsoleMode(h, &mode); err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
const (
|
||||||
|
enableLineInput = 0x0002
|
||||||
|
enableEchoInput = 0x0004
|
||||||
|
)
|
||||||
|
if err := windows.SetConsoleMode(h, mode&^(enableLineInput|enableEchoInput)); err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return func() { windows.SetConsoleMode(h, mode) } //nolint:errcheck
|
||||||
|
}
|
||||||
@@ -0,0 +1,242 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/subtle"
|
||||||
|
"fmt"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"gpu-turnstile/internal/control"
|
||||||
|
)
|
||||||
|
|
||||||
|
// uiHandler serves the web UI (LISTEN_UI): a single page mirroring
|
||||||
|
// --monitor, its status JSON (the same snapshot the control channel
|
||||||
|
// serves) and action endpoints for the update/reload buttons. user/pass
|
||||||
|
// (UI_USER/UI_PASS) enable HTTP basic auth; both empty means open access,
|
||||||
|
// which is meant for localhost only.
|
||||||
|
func uiHandler(h control.Handler, user, pass string) http.Handler {
|
||||||
|
mux := http.NewServeMux()
|
||||||
|
mux.HandleFunc("GET /", func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.URL.Path != "/" {
|
||||||
|
http.NotFound(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||||
|
fmt.Fprint(w, uiPage)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("GET /api/status", func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
msg := strings.TrimPrefix(h(control.CmdStatus), "OK ")
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
if strings.HasPrefix(msg, "ERR ") {
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
msg = strings.TrimPrefix(msg, "ERR ")
|
||||||
|
}
|
||||||
|
fmt.Fprint(w, msg)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("POST /api/action/{cmd}", func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
cmd := r.PathValue("cmd")
|
||||||
|
if cmd != control.CmdUpdateNow && cmd != control.CmdReloadEnv {
|
||||||
|
http.NotFound(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
msg := strings.TrimPrefix(h(cmd), "OK ")
|
||||||
|
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
||||||
|
if strings.HasPrefix(msg, "ERR ") {
|
||||||
|
w.WriteHeader(http.StatusConflict)
|
||||||
|
msg = strings.TrimPrefix(msg, "ERR ")
|
||||||
|
}
|
||||||
|
fmt.Fprint(w, msg)
|
||||||
|
})
|
||||||
|
if user != "" {
|
||||||
|
return basicAuth(mux, user, pass)
|
||||||
|
}
|
||||||
|
return mux
|
||||||
|
}
|
||||||
|
|
||||||
|
// basicAuth wraps next with HTTP basic auth (constant-time compares).
|
||||||
|
func basicAuth(next http.Handler, user, pass string) http.Handler {
|
||||||
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
u, p, ok := r.BasicAuth()
|
||||||
|
if !ok ||
|
||||||
|
subtle.ConstantTimeCompare([]byte(u), []byte(user)) != 1 ||
|
||||||
|
subtle.ConstantTimeCompare([]byte(p), []byte(pass)) != 1 {
|
||||||
|
w.Header().Set("WWW-Authenticate", `Basic realm="gpu-turnstile"`)
|
||||||
|
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
next.ServeHTTP(w, r)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// isLoopbackHost reports whether a listen host is localhost-only. An empty
|
||||||
|
// host (":7860") binds all interfaces and is therefore not loopback.
|
||||||
|
func isLoopbackHost(host string) bool {
|
||||||
|
if host == "localhost" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
ip := net.ParseIP(host)
|
||||||
|
return ip != nil && ip.IsLoopback()
|
||||||
|
}
|
||||||
|
|
||||||
|
// uiURL renders a LISTEN_UI address as a browsable URL ("" when disabled).
|
||||||
|
// A wildcard host reads as 127.0.0.1 — that's where the browser is usually
|
||||||
|
// running.
|
||||||
|
func uiURL(addr string) string {
|
||||||
|
if addr == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if strings.HasPrefix(addr, ":") {
|
||||||
|
addr = "127.0.0.1" + addr
|
||||||
|
}
|
||||||
|
return "http://" + addr
|
||||||
|
}
|
||||||
|
|
||||||
|
// uiPage is the whole web UI: a terminal-styled page that polls /api/status
|
||||||
|
// once a second and renders the same lines as --monitor. The buttons (and
|
||||||
|
// the u/r keys) hit the action endpoints, like the monitor's hotkeys.
|
||||||
|
const uiPage = `<!doctype html>
|
||||||
|
<html lang="en">
|
||||||
|
<head>
|
||||||
|
<meta charset="utf-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||||
|
<title>gpu-turnstile</title>
|
||||||
|
<style>
|
||||||
|
body { background:#111; color:#ddd; font:14px/1.6 ui-monospace,Consolas,"Cascadia Mono",monospace;
|
||||||
|
margin:0 auto; padding:14px 18px; max-width:860px; }
|
||||||
|
header { display:flex; justify-content:space-between; border-bottom:1px solid #333;
|
||||||
|
padding-bottom:4px; margin-bottom:10px; }
|
||||||
|
footer { border-top:1px solid #333; margin-top:10px; padding-top:8px;
|
||||||
|
display:flex; gap:8px; align-items:center; }
|
||||||
|
button { background:#1d1d1d; color:#bbb; border:1px solid #444; border-radius:4px;
|
||||||
|
padding:2px 10px; font:inherit; cursor:pointer; }
|
||||||
|
button:hover { background:#2c2c2c; color:#eee; }
|
||||||
|
.dim { color:#777; } .up { color:#4caf50; } .down { color:#ef5350; }
|
||||||
|
.llm { color:#26c6da; } .warn { color:#fbc02d; }
|
||||||
|
#note { color:#fbc02d; min-height:1.6em; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<header><span>gpu-turnstile <span id="up" class="dim"></span></span><span id="ver" class="dim"></span></header>
|
||||||
|
<div id="downstreams"></div>
|
||||||
|
<div id="lock"></div>
|
||||||
|
<div id="gpu"></div>
|
||||||
|
<div id="note"></div>
|
||||||
|
<footer>
|
||||||
|
<button id="bu">u update now</button>
|
||||||
|
<button id="br">r reload config</button>
|
||||||
|
<span class="dim">refreshes every second</span>
|
||||||
|
</footer>
|
||||||
|
<script>
|
||||||
|
const $ = id => document.getElementById(id);
|
||||||
|
const el = (cls, text) => { const e = document.createElement("span"); if (cls) e.className = cls; e.textContent = text; return e; };
|
||||||
|
const line = (...parts) => { const d = document.createElement("div"); d.append(...parts); return d; };
|
||||||
|
const fmtMB = mb => mb >= 1024 ? (mb / 1024).toFixed(1) + " GiB" : mb + " MiB";
|
||||||
|
const fmtDur = s => {
|
||||||
|
s = Math.max(0, Math.floor(s));
|
||||||
|
if (s >= 3600) return Math.floor(s / 3600) + "h" + String(Math.floor(s / 60) % 60).padStart(2, "0") + "m";
|
||||||
|
if (s >= 60) return Math.floor(s / 60) + "m" + String(s % 60).padStart(2, "0") + "s";
|
||||||
|
return s + "s";
|
||||||
|
};
|
||||||
|
|
||||||
|
function renderDownstream(d, lock) {
|
||||||
|
const busy = (d.name === "ollama" && lock.state === "llm" && lock.llm_inflight > 0) ||
|
||||||
|
(d.name === "comfy" && lock.state === "image");
|
||||||
|
if (d.managed === "stopped")
|
||||||
|
return line(el("dim", "○ " + d.name + " stopped (managed — starts on demand) " + d.url));
|
||||||
|
if (d.managed === "starting")
|
||||||
|
return line(el("warn", "◌ " + d.name + " starting… "), el("dim", d.url));
|
||||||
|
const parts = [el(d.up ? "up" : "down", "● " + d.name + " " + (d.up ? "UP" : "DOWN"))];
|
||||||
|
if (d.managed === "external") parts.push(el("", " (external)"));
|
||||||
|
if (d.models) {
|
||||||
|
if (d.models.length === 0) {
|
||||||
|
parts.push(el("dim", " · no models loaded"));
|
||||||
|
} else {
|
||||||
|
parts.push(el("", " · " + d.models.map(m =>
|
||||||
|
m.vram_mb > 0 ? m.name + " (" + fmtMB(m.vram_mb) + " VRAM)" : m.name + " (in RAM)").join(", ")));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (busy) parts.push(el("llm", " · busy"));
|
||||||
|
parts.push(el("dim", " " + d.url));
|
||||||
|
return line(...parts);
|
||||||
|
}
|
||||||
|
|
||||||
|
function renderLock(l) {
|
||||||
|
const dur = el("dim", " (" + fmtDur(l.since_s || 0) + ")");
|
||||||
|
switch (l.state) {
|
||||||
|
case "idle": return line(el("up", "Lock: idle"), dur);
|
||||||
|
case "llm": {
|
||||||
|
let s = "Lock: LLM — " + l.llm_inflight + " in flight";
|
||||||
|
if (l.llm_waiting > 0) s += ", " + l.llm_waiting + " waiting";
|
||||||
|
if (l.detail) s += " — " + l.detail;
|
||||||
|
return line(el("llm", s), dur);
|
||||||
|
}
|
||||||
|
case "image": return line(el("warn", "Lock: IMAGE" + (l.detail ? " — " + l.detail : "")), dur);
|
||||||
|
case "external": return line(el("down", "Lock: EXTERNAL — " + (l.external || "")), dur);
|
||||||
|
}
|
||||||
|
return line(el("", "Lock: unknown"));
|
||||||
|
}
|
||||||
|
|
||||||
|
function renderGPU(g) {
|
||||||
|
const parts = [el("", "GPU: ")];
|
||||||
|
if (g.known) {
|
||||||
|
let s = fmtMB(g.used_mb) + " / " + fmtMB(g.total_mb) + " used";
|
||||||
|
if (g.temp_c >= 0) s += " · " + g.temp_c + "°C";
|
||||||
|
if (g.fan_pct >= 0) s += " · fan " + g.fan_pct + "%";
|
||||||
|
parts.push(el("", s));
|
||||||
|
} else {
|
||||||
|
parts.push(el("dim", "VRAM unknown (nvidia-smi not answering)"));
|
||||||
|
}
|
||||||
|
parts.push(g.foreign ? el("down", " · external: " + g.foreign) : el("dim", " · no external process"));
|
||||||
|
if (g.age_s >= 0) parts.push(el("dim", " (checked " + fmtDur(g.age_s) + " ago)"));
|
||||||
|
return line(...parts);
|
||||||
|
}
|
||||||
|
|
||||||
|
async function poll() {
|
||||||
|
let snap;
|
||||||
|
try {
|
||||||
|
const r = await fetch("/api/status");
|
||||||
|
if (!r.ok) throw new Error();
|
||||||
|
snap = await r.json();
|
||||||
|
} catch {
|
||||||
|
$("downstreams").replaceChildren(line(el("dim", "gpu-turnstile — waiting for a running service…")));
|
||||||
|
$("lock").replaceChildren(); $("gpu").replaceChildren();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
$("ver").textContent = snap.version || "";
|
||||||
|
$("up").textContent = snap.uptime_s > 0 ? "up " + fmtDur(snap.uptime_s) : "";
|
||||||
|
const ds = $("downstreams"); ds.replaceChildren();
|
||||||
|
for (const d of snap.downstreams || []) ds.appendChild(renderDownstream(d, snap.lock || {}));
|
||||||
|
const lk = $("lock"); lk.replaceChildren(renderLock(snap.lock || {}));
|
||||||
|
if ((snap.lock || {}).image_queue > 0)
|
||||||
|
lk.appendChild(line(el("warn", "Queue: " + snap.lock.image_queue + " image job(s) waiting")));
|
||||||
|
const g = $("gpu"); g.replaceChildren();
|
||||||
|
if (snap.gpu && snap.gpu.enabled) g.appendChild(renderGPU(snap.gpu));
|
||||||
|
}
|
||||||
|
|
||||||
|
let noteTimer;
|
||||||
|
async function act(cmd) {
|
||||||
|
const note = $("note");
|
||||||
|
note.textContent = "…";
|
||||||
|
try {
|
||||||
|
const r = await fetch("/api/action/" + cmd, { method: "POST" });
|
||||||
|
note.textContent = await r.text();
|
||||||
|
} catch {
|
||||||
|
note.textContent = "no answer from the service";
|
||||||
|
}
|
||||||
|
clearTimeout(noteTimer);
|
||||||
|
noteTimer = setTimeout(() => { note.textContent = ""; }, 15000);
|
||||||
|
}
|
||||||
|
|
||||||
|
$("bu").onclick = () => act("update-now");
|
||||||
|
$("br").onclick = () => act("reload-env");
|
||||||
|
document.addEventListener("keydown", e => {
|
||||||
|
if (e.key === "u") act("update-now");
|
||||||
|
if (e.key === "r") act("reload-env");
|
||||||
|
});
|
||||||
|
poll();
|
||||||
|
setInterval(poll, 1000);
|
||||||
|
</script>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
|
`
|
||||||
@@ -0,0 +1,185 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"gpu-turnstile/internal/control"
|
||||||
|
)
|
||||||
|
|
||||||
|
func testUIHandler() *httptest.Server {
|
||||||
|
h := uiHandler(func(cmd string) string {
|
||||||
|
switch cmd {
|
||||||
|
case control.CmdStatus:
|
||||||
|
return `OK {"version":"v9.9.9","uptime_s":5}`
|
||||||
|
case control.CmdUpdateNow:
|
||||||
|
return "OK v9.9.9 is up to date"
|
||||||
|
default:
|
||||||
|
return "ERR unknown command: " + cmd
|
||||||
|
}
|
||||||
|
}, "", "")
|
||||||
|
return httptest.NewServer(h)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUIPage(t *testing.T) {
|
||||||
|
srv := testUIHandler()
|
||||||
|
defer srv.Close()
|
||||||
|
resp, err := srv.Client().Get(srv.URL + "/")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body := make([]byte, 1024*64)
|
||||||
|
n, _ := resp.Body.Read(body)
|
||||||
|
if resp.StatusCode != 200 {
|
||||||
|
t.Fatalf("status %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "text/html") {
|
||||||
|
t.Errorf("Content-Type %q", ct)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(body[:n]), "gpu-turnstile") {
|
||||||
|
t.Error("page does not mention gpu-turnstile")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUIStatus(t *testing.T) {
|
||||||
|
srv := testUIHandler()
|
||||||
|
defer srv.Close()
|
||||||
|
resp, err := srv.Client().Get(srv.URL + "/api/status")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body := make([]byte, 4096)
|
||||||
|
n, _ := resp.Body.Read(body)
|
||||||
|
if resp.StatusCode != 200 {
|
||||||
|
t.Fatalf("status %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "application/json") {
|
||||||
|
t.Errorf("Content-Type %q", ct)
|
||||||
|
}
|
||||||
|
if got := strings.TrimSpace(string(body[:n])); got != `{"version":"v9.9.9","uptime_s":5}` {
|
||||||
|
t.Errorf("body %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUIActions(t *testing.T) {
|
||||||
|
srv := testUIHandler()
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
resp, err := srv.Client().Post(srv.URL+"/api/action/update-now", "", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body := make([]byte, 4096)
|
||||||
|
n, _ := resp.Body.Read(body)
|
||||||
|
if resp.StatusCode != 200 || string(body[:n]) != "v9.9.9 is up to date" {
|
||||||
|
t.Errorf("update-now: %d %q", resp.StatusCode, body[:n])
|
||||||
|
}
|
||||||
|
|
||||||
|
// The handler answers ERR for reload-env: mapped to 409 with the
|
||||||
|
// message stripped.
|
||||||
|
resp2, err := srv.Client().Post(srv.URL+"/api/action/reload-env", "", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer resp2.Body.Close()
|
||||||
|
n, _ = resp2.Body.Read(body)
|
||||||
|
if resp2.StatusCode != 409 || !strings.HasPrefix(string(body[:n]), "unknown command") {
|
||||||
|
t.Errorf("reload-env: %d %q", resp2.StatusCode, body[:n])
|
||||||
|
}
|
||||||
|
|
||||||
|
// Unknown actions are not routed to the handler.
|
||||||
|
resp3, err := srv.Client().Post(srv.URL+"/api/action/reboot", "", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
resp3.Body.Close()
|
||||||
|
if resp3.StatusCode != 404 {
|
||||||
|
t.Errorf("bogus action: %d, want 404", resp3.StatusCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
// GET on an action path falls through to the root handler's guard,
|
||||||
|
// which only serves "/" itself.
|
||||||
|
resp4, err := srv.Client().Get(srv.URL + "/api/action/update-now")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
resp4.Body.Close()
|
||||||
|
if resp4.StatusCode != 404 {
|
||||||
|
t.Errorf("GET action: %d, want 404", resp4.StatusCode)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUIBasicAuth(t *testing.T) {
|
||||||
|
h := uiHandler(func(cmd string) string { return `OK {"version":"v9.9.9"}` }, "admin", "s3cret")
|
||||||
|
srv := httptest.NewServer(h)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
// No credentials: 401 with the auth challenge on page and API alike.
|
||||||
|
for _, path := range []string{"/", "/api/status"} {
|
||||||
|
resp, err := srv.Client().Get(srv.URL + path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
resp.Body.Close()
|
||||||
|
if resp.StatusCode != 401 {
|
||||||
|
t.Errorf("GET %s without auth: %d, want 401", path, resp.StatusCode)
|
||||||
|
}
|
||||||
|
if resp.Header.Get("WWW-Authenticate") == "" {
|
||||||
|
t.Errorf("GET %s: no WWW-Authenticate header", path)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wrong credentials: 401. Right credentials: 200.
|
||||||
|
req, _ := http.NewRequest("GET", srv.URL+"/api/status", nil)
|
||||||
|
req.SetBasicAuth("admin", "wrong")
|
||||||
|
resp, err := srv.Client().Do(req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
resp.Body.Close()
|
||||||
|
if resp.StatusCode != 401 {
|
||||||
|
t.Errorf("wrong password: %d, want 401", resp.StatusCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
req, _ = http.NewRequest("GET", srv.URL+"/api/status", nil)
|
||||||
|
req.SetBasicAuth("admin", "s3cret")
|
||||||
|
resp, err = srv.Client().Do(req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body := make([]byte, 4096)
|
||||||
|
n, _ := resp.Body.Read(body)
|
||||||
|
if resp.StatusCode != 200 || strings.TrimSpace(string(body[:n])) != `{"version":"v9.9.9"}` {
|
||||||
|
t.Errorf("with auth: %d %q", resp.StatusCode, body[:n])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIsLoopbackHost(t *testing.T) {
|
||||||
|
for host, want := range map[string]bool{
|
||||||
|
"127.0.0.1": true, "localhost": true, "::1": true,
|
||||||
|
"": false, "0.0.0.0": false, "192.168.1.10": false, "example.com": false,
|
||||||
|
} {
|
||||||
|
if got := isLoopbackHost(host); got != want {
|
||||||
|
t.Errorf("isLoopbackHost(%q) = %v, want %v", host, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUIURL(t *testing.T) {
|
||||||
|
for addr, want := range map[string]string{
|
||||||
|
"": "",
|
||||||
|
"127.0.0.1:7860": "http://127.0.0.1:7860",
|
||||||
|
":7860": "http://127.0.0.1:7860",
|
||||||
|
"192.168.1.5:9000": "http://192.168.1.5:9000",
|
||||||
|
} {
|
||||||
|
if got := uiURL(addr); got != want {
|
||||||
|
t.Errorf("uiURL(%q) = %q, want %q", addr, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -6,7 +6,7 @@
|
|||||||
# ComfyUI --listen 0.0.0.0 --port 8189).
|
# ComfyUI --listen 0.0.0.0 --port 8189).
|
||||||
services:
|
services:
|
||||||
gpu-turnstile:
|
gpu-turnstile:
|
||||||
image: git.rambossek.at/public/gpu-turnstile:v0.2.2
|
image: git.rambossek.at/public/gpu-turnstile:v0.3.2
|
||||||
restart: unless-stopped
|
restart: unless-stopped
|
||||||
environment:
|
environment:
|
||||||
# Each consumer is enabled by setting its URL; leave one unset to
|
# Each consumer is enabled by setting its URL; leave one unset to
|
||||||
|
|||||||
+67
-42
@@ -13,43 +13,48 @@ import (
|
|||||||
"time"
|
"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 {
|
type Config struct {
|
||||||
ListenOllama string
|
ListenOllama string `env:"LISTEN_OLLAMA"`
|
||||||
ListenComfy string
|
ListenComfy string `env:"LISTEN_COMFY"`
|
||||||
OllamaURL string
|
ListenUI string `env:"LISTEN_UI"`
|
||||||
ComfyURL string
|
UIUser string `env:"UI_USER"`
|
||||||
UnloadTimeout time.Duration
|
UIPass string `env:"UI_PASS"`
|
||||||
JobTimeout time.Duration
|
OllamaURL string `env:"OLLAMA_URL"`
|
||||||
LLMWaitTimeout time.Duration
|
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
|
UnloadPollInterval time.Duration `env:"UNLOAD_POLL_INTERVAL"`
|
||||||
HistoryPollInterval time.Duration
|
HistoryPollInterval time.Duration `env:"HISTORY_POLL_INTERVAL"`
|
||||||
ProbeTimeout time.Duration
|
ProbeTimeout time.Duration `env:"PROBE_TIMEOUT"`
|
||||||
HealthInterval time.Duration
|
HealthInterval time.Duration `env:"HEALTH_INTERVAL"`
|
||||||
FreeTimeout time.Duration
|
FreeTimeout time.Duration `env:"FREE_TIMEOUT"`
|
||||||
WarmTimeout time.Duration
|
WarmTimeout time.Duration `env:"WARM_TIMEOUT"`
|
||||||
ShutdownTimeout time.Duration
|
ShutdownTimeout time.Duration `env:"SHUTDOWN_TIMEOUT"`
|
||||||
BackoffInitial time.Duration
|
BackoffInitial time.Duration `env:"BACKOFF_INITIAL"`
|
||||||
BackoffMax time.Duration
|
BackoffMax time.Duration `env:"BACKOFF_MAX"`
|
||||||
PromptCaptureLimit int64
|
PromptCaptureLimit int64 `env:"PROMPT_CAPTURE_LIMIT"`
|
||||||
|
|
||||||
AutoUpdate bool
|
AutoUpdate bool `env:"AUTO_UPDATE"`
|
||||||
UpdateInterval time.Duration
|
UpdateInterval time.Duration `env:"UPDATE_INTERVAL"`
|
||||||
UpdateRepo string
|
UpdateRepo string `env:"UPDATE_REPO"`
|
||||||
UpdateAsset string
|
UpdateAsset string `env:"UPDATE_ASSET"`
|
||||||
|
|
||||||
// AppVersion is the version the user wants to run: "dev" disables
|
// AppVersion is the version the user wants to run: "dev" disables
|
||||||
// updates, "stable" tracks the latest release, anything else is an
|
// updates, "stable" tracks the latest release, anything else is an
|
||||||
// exact vX.Y.Z release to pin. From APP_VER; defaults to "stable".
|
// 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
|
// LLMBusyMode is "wait" (hold requests until the lock is free or
|
||||||
// LLMWaitTimeout expires) or "reject" (immediately answer with
|
// LLMWaitTimeout expires) or "reject" (immediately answer with
|
||||||
// LLMBusyStatus + Retry-After when an image job is active or pending).
|
// LLMBusyStatus + Retry-After when an image job is active or pending).
|
||||||
LLMBusyMode string
|
LLMBusyMode string `env:"LLM_BUSY_MODE"`
|
||||||
LLMBusyStatus int
|
LLMBusyStatus int `env:"LLM_BUSY_STATUS"`
|
||||||
BusyRetryAfter int
|
BusyRetryAfter int `env:"BUSY_RETRY_AFTER"`
|
||||||
|
|
||||||
// ComfyCmd spawns and supervises a ComfyUI server on demand. When
|
// ComfyCmd spawns and supervises a ComfyUI server on demand. When
|
||||||
// ComfyCmd is empty but ComfyDir is set, management is enabled with the
|
// ComfyCmd is empty but ComfyDir is set, management is enabled with the
|
||||||
@@ -59,26 +64,29 @@ type Config struct {
|
|||||||
// set explicitly. The managed server is stopped after ComfyIdleTimeout
|
// set explicitly. The managed server is stopped after ComfyIdleTimeout
|
||||||
// without requests, freeing its VRAM; ComfyStartTimeout bounds how long
|
// without requests, freeing its VRAM; ComfyStartTimeout bounds how long
|
||||||
// a request waits for it to come up.
|
// a request waits for it to come up.
|
||||||
ComfyCmd string
|
ComfyCmd string `env:"COMFY_CMD"`
|
||||||
ComfyDir string
|
ComfyDir string `env:"COMFY_DIR"`
|
||||||
ComfyIdleTimeout time.Duration
|
ComfyIdleTimeout time.Duration `env:"COMFY_IDLE_TIMEOUT"`
|
||||||
ComfyStartTimeout time.Duration
|
ComfyStartTimeout time.Duration `env:"COMFY_START_TIMEOUT"`
|
||||||
|
|
||||||
// GameProcs (GAME_PROCS) is a watch list of process names; while any of
|
// 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
|
// them runs, the GPU is treated as held by a foreign process. The
|
||||||
// nvidia-smi path (GPUForeignVRAMMB, GPU_FOREIGN_VRAM_MB) does the same
|
// nvidia-smi path (GPUForeignVRAMMB, GPU_FOREIGN_VRAM_MB) does the same
|
||||||
// when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than
|
// when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than
|
||||||
// that many MiB of VRAM. GamePollInterval (GAME_POLL_INTERVAL) is how
|
// that many MiB of VRAM, and the PDH path (GPUForeignUtilPct,
|
||||||
// often both checks run.
|
// GPU_FOREIGN_UTIL_PCT) when such a process uses more than that many
|
||||||
GameProcs []string
|
// percent of the GPU 3D engine (Windows only). GamePollInterval
|
||||||
GPUForeignVRAMMB int
|
// (GAME_POLL_INTERVAL) is how often all checks run.
|
||||||
GPUIgnoreProcs []string
|
GameProcs []string `env:"GAME_PROCS"`
|
||||||
GamePollInterval time.Duration
|
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
|
WarmModel string `env:"WARM_MODEL"`
|
||||||
LogLevel slog.Level
|
LogLevel slog.Level `env:"LOGLEVEL"`
|
||||||
LogJSON bool
|
LogJSON bool `env:"LOG_FORMAT"`
|
||||||
LogFile string
|
LogFile string `env:"LOG_FILE"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// Defaults returns the configuration used when neither the environment nor
|
// Defaults returns the configuration used when neither the environment nor
|
||||||
@@ -117,8 +125,9 @@ func Defaults() Config {
|
|||||||
ComfyStartTimeout: 2 * time.Minute,
|
ComfyStartTimeout: 2 * time.Minute,
|
||||||
|
|
||||||
// ComfyUI runs under python; excluding it (and Ollama) by name keeps
|
// ComfyUI runs under python; excluding it (and Ollama) by name keeps
|
||||||
// our own consumers from tripping the foreign-VRAM check.
|
// our own consumers from tripping the foreign-VRAM check. dwm is the
|
||||||
GPUIgnoreProcs: []string{"ollama", "ollama app", "ollama_llama_server", "python", "pythonw"},
|
// desktop compositor — it always shows some 3D-engine usage.
|
||||||
|
GPUIgnoreProcs: []string{"ollama", "ollama app", "ollama_llama_server", "python", "pythonw", "dwm"},
|
||||||
GamePollInterval: 15 * time.Second,
|
GamePollInterval: 15 * time.Second,
|
||||||
|
|
||||||
LogLevel: slog.LevelWarn,
|
LogLevel: slog.LevelWarn,
|
||||||
@@ -190,6 +199,9 @@ func Load(getenv func(string) string) (Config, error) {
|
|||||||
}{
|
}{
|
||||||
{"LISTEN_OLLAMA", &cfg.ListenOllama},
|
{"LISTEN_OLLAMA", &cfg.ListenOllama},
|
||||||
{"LISTEN_COMFY", &cfg.ListenComfy},
|
{"LISTEN_COMFY", &cfg.ListenComfy},
|
||||||
|
{"LISTEN_UI", &cfg.ListenUI},
|
||||||
|
{"UI_USER", &cfg.UIUser},
|
||||||
|
{"UI_PASS", &cfg.UIPass},
|
||||||
{"OLLAMA_URL", &cfg.OllamaURL},
|
{"OLLAMA_URL", &cfg.OllamaURL},
|
||||||
{"COMFY_URL", &cfg.ComfyURL},
|
{"COMFY_URL", &cfg.ComfyURL},
|
||||||
{"WARM_MODEL", &cfg.WarmModel},
|
{"WARM_MODEL", &cfg.WarmModel},
|
||||||
@@ -241,6 +253,13 @@ func Load(getenv func(string) string) (Config, error) {
|
|||||||
}
|
}
|
||||||
cfg.GPUForeignVRAMMB = n
|
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 != "" {
|
if v := getenv("PROMPT_CAPTURE_LIMIT"); v != "" {
|
||||||
n, err := strconv.ParseInt(v, 10, 64)
|
n, err := strconv.ParseInt(v, 10, 64)
|
||||||
if err != nil || n < 0 {
|
if err != nil || n < 0 {
|
||||||
@@ -305,6 +324,12 @@ func Load(getenv func(string) string) (Config, error) {
|
|||||||
default:
|
default:
|
||||||
return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"")
|
return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"")
|
||||||
}
|
}
|
||||||
|
if (cfg.UIUser == "") != (cfg.UIPass == "") {
|
||||||
|
return cfg, fmt.Errorf("UI_USER and UI_PASS must be set together (both empty = no auth)")
|
||||||
|
}
|
||||||
|
if cfg.UIUser != "" && cfg.ListenUI == "" {
|
||||||
|
return cfg, fmt.Errorf("UI_USER/UI_PASS have no effect without LISTEN_UI")
|
||||||
|
}
|
||||||
if cfg.ComfyCmd != "" && cfg.ComfyURL == "" {
|
if cfg.ComfyCmd != "" && cfg.ComfyURL == "" {
|
||||||
return cfg, fmt.Errorf("COMFY_CMD requires COMFY_URL to be set (the proxy needs somewhere to forward)")
|
return cfg, fmt.Errorf("COMFY_CMD requires COMFY_URL to be set (the proxy needs somewhere to forward)")
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -97,6 +97,30 @@ func TestComfyCmdRequiresURL(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestUIAuthPairing(t *testing.T) {
|
||||||
|
env := func(set map[string]string) func(string) string {
|
||||||
|
return func(k string) string { return set[k] }
|
||||||
|
}
|
||||||
|
// Only one of UI_USER/UI_PASS: error.
|
||||||
|
_, err := Load(env(map[string]string{"OLLAMA_URL": "http://x", "LISTEN_UI": "127.0.0.1:7860", "UI_USER": "admin"}))
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "UI_USER and UI_PASS must be set together") {
|
||||||
|
t.Fatalf("err = %v, want pairing error", err)
|
||||||
|
}
|
||||||
|
// Auth without LISTEN_UI: error.
|
||||||
|
_, err = Load(env(map[string]string{"OLLAMA_URL": "http://x", "UI_USER": "admin", "UI_PASS": "x"}))
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "no effect without LISTEN_UI") {
|
||||||
|
t.Fatalf("err = %v, want LISTEN_UI error", err)
|
||||||
|
}
|
||||||
|
// Both with LISTEN_UI: loads.
|
||||||
|
cfg, err := Load(env(map[string]string{"OLLAMA_URL": "http://x", "LISTEN_UI": "127.0.0.1:7860", "UI_USER": "admin", "UI_PASS": "x"}))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("auth pair with LISTEN_UI must load: %v", err)
|
||||||
|
}
|
||||||
|
if cfg.UIUser != "admin" || cfg.UIPass != "x" {
|
||||||
|
t.Errorf("got %q/%q", cfg.UIUser, cfg.UIPass)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestParseEnvFile(t *testing.T) {
|
func TestParseEnvFile(t *testing.T) {
|
||||||
input := `# comment
|
input := `# comment
|
||||||
OLLAMA_URL=http://host:11435
|
OLLAMA_URL=http://host:11435
|
||||||
|
|||||||
@@ -25,6 +25,9 @@ func sampleEntries(logFile string) []sampleEntry {
|
|||||||
return []sampleEntry{
|
return []sampleEntry{
|
||||||
{"LISTEN_OLLAMA", ":11434", "Listen address for Ollama-compatible clients (gpu-turnstile poses as Ollama here)", false},
|
{"LISTEN_OLLAMA", ":11434", "Listen address for Ollama-compatible clients (gpu-turnstile poses as Ollama here)", false},
|
||||||
{"LISTEN_COMFY", ":8188", "Listen address for ComfyUI clients (gpu-turnstile poses as ComfyUI here)", false},
|
{"LISTEN_COMFY", ":8188", "Listen address for ComfyUI clients (gpu-turnstile poses as ComfyUI here)", false},
|
||||||
|
{"LISTEN_UI", "127.0.0.1:7860", "Web UI listen address (host:port): live status like --monitor, with update/reload buttons (default: empty = disabled; bind 127.0.0.1 unless you set auth)", false},
|
||||||
|
{"UI_USER", "", "HTTP basic auth for the web UI (UI_USER and UI_PASS must be set together; both empty = no auth)", false},
|
||||||
|
{"UI_PASS", "", "See UI_USER", false},
|
||||||
{"OLLAMA_URL", "http://127.0.0.1:11434", "Ollama upstream URL; setting it enables the Ollama consumer (default: empty = disabled)", false},
|
{"OLLAMA_URL", "http://127.0.0.1:11434", "Ollama upstream URL; setting it enables the Ollama consumer (default: empty = disabled)", false},
|
||||||
{"COMFY_URL", "http://127.0.0.1:8188", "ComfyUI upstream URL; setting it enables the ComfyUI consumer (default: empty = disabled)", false},
|
{"COMFY_URL", "http://127.0.0.1:8188", "ComfyUI upstream URL; setting it enables the ComfyUI consumer (default: empty = disabled)", false},
|
||||||
{"WARM_MODEL", "", "Optional model to reload after an image job (default: empty = none)", false},
|
{"WARM_MODEL", "", "Optional model to reload after an image job (default: empty = none)", false},
|
||||||
@@ -34,7 +37,8 @@ func sampleEntries(logFile string) []sampleEntry {
|
|||||||
{"COMFY_START_TIMEOUT", "2m", "How long a request waits for the managed ComfyUI to come up", false},
|
{"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},
|
{"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_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},
|
{"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},
|
{"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},
|
{"JOB_TIMEOUT", "15m", "Maximum time to wait for a ComfyUI job", false},
|
||||||
|
|||||||
@@ -0,0 +1,117 @@
|
|||||||
|
// Package control exposes a local-only command channel into a running
|
||||||
|
// gpu-turnstile service: a named pipe on Windows, a unix socket on Linux.
|
||||||
|
// It lets unprivileged local users ask the service to do privileged work
|
||||||
|
// that is safe to offer — currently triggering an update check, whose
|
||||||
|
// payload is signature-verified regardless of who asks. The channel never
|
||||||
|
// accepts data beyond a one-word command, and the server rate-limits
|
||||||
|
// 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
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
|
||||||
|
// idle) restart onto a signed update immediately.
|
||||||
|
const CmdUpdateNow = "update-now"
|
||||||
|
|
||||||
|
// CmdStatus asks for a one-line JSON status snapshot (monitor mode).
|
||||||
|
const CmdStatus = "status"
|
||||||
|
|
||||||
|
// CmdReloadEnv asks the service to re-read and validate its config file,
|
||||||
|
// and to restart onto it (once the GPU is idle) when it changed.
|
||||||
|
const CmdReloadEnv = "reload-env"
|
||||||
|
|
||||||
|
// ErrUnavailable means no running service offers the control channel.
|
||||||
|
var ErrUnavailable = errors.New("control channel unavailable")
|
||||||
|
|
||||||
|
// Handler answers one command; the returned string is sent back as one
|
||||||
|
// 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 == "" {
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Fprintln(c, "ERR empty command")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Fprintln(c, h(cmd))
|
||||||
|
}
|
||||||
|
|
||||||
|
// readReply writes cmd and reads the server's one-line reply.
|
||||||
|
func readReply(c io.ReadWriteCloser, cmd string) (string, error) {
|
||||||
|
if _, err := fmt.Fprintln(c, cmd); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
// The server hangs up after its reply; a broken-pipe error after the
|
||||||
|
// last byte still leaves the reply in the buffer.
|
||||||
|
data, _ := io.ReadAll(io.LimitReader(c, 4096))
|
||||||
|
line := strings.TrimSpace(string(data))
|
||||||
|
if line == "" {
|
||||||
|
return "", ErrUnavailable
|
||||||
|
}
|
||||||
|
return line, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
//go:build linux
|
||||||
|
|
||||||
|
package control
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"log/slog"
|
||||||
|
"net"
|
||||||
|
"os"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// sockPath lives in the unit's RuntimeDirectory; mode 0666 lets every
|
||||||
|
// local user ask, nothing can reach it from off the machine.
|
||||||
|
const sockPath = "/run/gpu-turnstile/control.sock"
|
||||||
|
|
||||||
|
// Serve starts the socket listener in the background and returns; only a
|
||||||
|
// setup failure is reported. Each client connection is answered in its own
|
||||||
|
// goroutine.
|
||||||
|
func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
||||||
|
os.Remove(sockPath) // stale socket from a previous run
|
||||||
|
ln, err := net.Listen("unix", sockPath)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := os.Chmod(sockPath, 0o666); err != nil {
|
||||||
|
ln.Close()
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
go func() {
|
||||||
|
<-ctx.Done()
|
||||||
|
ln.Close()
|
||||||
|
}()
|
||||||
|
go func() {
|
||||||
|
for {
|
||||||
|
c, err := ln.Accept()
|
||||||
|
if err != nil {
|
||||||
|
return // shutting down
|
||||||
|
}
|
||||||
|
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)
|
||||||
|
if err != nil {
|
||||||
|
return "", ErrUnavailable
|
||||||
|
}
|
||||||
|
defer c.Close()
|
||||||
|
return readReply(c, cmd)
|
||||||
|
}
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
//go:build !windows && !linux
|
||||||
|
|
||||||
|
package control
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"log/slog"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Serve is a no-op on platforms without a control channel.
|
||||||
|
func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
||||||
|
return ErrUnavailable
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ask always reports the channel as unavailable.
|
||||||
|
func Ask(cmd string) (string, error) {
|
||||||
|
return "", ErrUnavailable
|
||||||
|
}
|
||||||
@@ -0,0 +1,103 @@
|
|||||||
|
package control
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"net"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRoundTrip(t *testing.T) {
|
||||||
|
server, client := net.Pipe()
|
||||||
|
go serveConn(server, func(cmd string) string {
|
||||||
|
if cmd != CmdUpdateNow {
|
||||||
|
return "ERR unknown command: " + cmd
|
||||||
|
}
|
||||||
|
return "OK v0.2.2 is up to date"
|
||||||
|
})
|
||||||
|
reply, err := readReply(client, CmdUpdateNow)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if reply != "OK v0.2.2 is up to date" {
|
||||||
|
t.Fatalf("reply = %q", reply)
|
||||||
|
}
|
||||||
|
|
||||||
|
server2, client2 := net.Pipe()
|
||||||
|
go serveConn(server2, func(cmd string) string { return "ERR unknown command: " + cmd })
|
||||||
|
reply, err = readReply(client2, "bogus")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.HasPrefix(reply, "ERR ") {
|
||||||
|
t.Fatalf("reply = %q, want ERR prefix", reply)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEmptyReplyIsUnavailable(t *testing.T) {
|
||||||
|
server, client := net.Pipe()
|
||||||
|
go serveConn(server, func(cmd string) string {
|
||||||
|
server.Close() // hang up without answering
|
||||||
|
return ""
|
||||||
|
})
|
||||||
|
if _, err := readReply(client, CmdUpdateNow); !errors.Is(err, ErrUnavailable) {
|
||||||
|
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")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,167 @@
|
|||||||
|
//go:build windows
|
||||||
|
|
||||||
|
package control
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"os"
|
||||||
|
"syscall"
|
||||||
|
"time"
|
||||||
|
"unsafe"
|
||||||
|
|
||||||
|
"golang.org/x/sys/windows"
|
||||||
|
)
|
||||||
|
|
||||||
|
// pipePath is a kernel-local named pipe: no TCP, no firewall prompt.
|
||||||
|
const pipePath = `\\.\pipe\gpu-turnstile`
|
||||||
|
|
||||||
|
// sddlPipe grants full access to Administrators, SYSTEM and the pipe owner,
|
||||||
|
// and read+write to authenticated users — except network logons, so the
|
||||||
|
// pipe cannot be reached from another machine over SMB.
|
||||||
|
const sddlPipe = "D:(D;;GRGW;;;NU)(A;;GA;;;BA)(A;;GA;;;SY)(A;;GA;;;OW)(A;;GRGW;;;AU)"
|
||||||
|
|
||||||
|
var (
|
||||||
|
procConvertSDDL = windows.NewLazySystemDLL("advapi32.dll").
|
||||||
|
NewProc("ConvertStringSecurityDescriptorToSecurityDescriptorW")
|
||||||
|
procWaitNamedPipe = windows.NewLazySystemDLL("kernel32.dll").
|
||||||
|
NewProc("WaitNamedPipeW")
|
||||||
|
)
|
||||||
|
|
||||||
|
func waitNamedPipe(name *uint16, timeout uint32) error {
|
||||||
|
r, _, err := procWaitNamedPipe.Call(uintptr(unsafe.Pointer(name)), uintptr(timeout))
|
||||||
|
if r == 0 {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func securityAttributesFromSDDL(sddl string) (*windows.SecurityAttributes, error) {
|
||||||
|
s, err := windows.UTF16PtrFromString(sddl)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
var sd *uint16 // SECURITY_DESCRIPTOR*, kept for the process lifetime
|
||||||
|
r, _, callErr := procConvertSDDL.Call(
|
||||||
|
uintptr(unsafe.Pointer(s)), 1, /* SDDL_REVISION_1 */
|
||||||
|
uintptr(unsafe.Pointer(&sd)), 0)
|
||||||
|
if r == 0 {
|
||||||
|
return nil, fmt.Errorf("invalid SDDL: %w", callErr)
|
||||||
|
}
|
||||||
|
sa := &windows.SecurityAttributes{
|
||||||
|
Length: uint32(unsafe.Sizeof(windows.SecurityAttributes{})),
|
||||||
|
SecurityDescriptor: (*windows.SECURITY_DESCRIPTOR)(unsafe.Pointer(sd)),
|
||||||
|
}
|
||||||
|
return sa, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Serve starts the pipe listener in the background and returns; only a
|
||||||
|
// setup failure is reported. Each client connection is answered in its own
|
||||||
|
// goroutine. On shutdown the process exit reaps everything.
|
||||||
|
func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
||||||
|
sa, err := securityAttributesFromSDDL(sddlPipe)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
name, err := windows.UTF16PtrFromString(pipePath)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
go func() {
|
||||||
|
for ctx.Err() == nil {
|
||||||
|
pipe, err := windows.CreateNamedPipe(name,
|
||||||
|
windows.PIPE_ACCESS_DUPLEX,
|
||||||
|
windows.PIPE_TYPE_BYTE|windows.PIPE_READMODE_BYTE|windows.PIPE_WAIT,
|
||||||
|
16, 4096, 4096, 0, sa)
|
||||||
|
if err != nil {
|
||||||
|
log.Warn("control channel stopped", "err", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// Blocks until a client connects — only then is the next
|
||||||
|
// instance created, so instances are not burned without
|
||||||
|
// clients. ERROR_PIPE_CONNECTED means the client raced us
|
||||||
|
// and connected before the call: that is a success. Process
|
||||||
|
// exit reaps the blocked call on shutdown.
|
||||||
|
if err := windows.ConnectNamedPipe(pipe, nil); err != nil && err != errnoPipeConnected {
|
||||||
|
windows.CloseHandle(pipe)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
conn := &pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}
|
||||||
|
if !serve(conn, h) {
|
||||||
|
conn.ForceClose()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// errnoPipeConnected is ConnectNamedPipe's "the client connected before we
|
||||||
|
// called" result, which means the connection is established.
|
||||||
|
var errnoPipeConnected = syscall.Errno(535) // ERROR_PIPE_CONNECTED
|
||||||
|
|
||||||
|
// errnoPipeBusy is CreateFile's "all pipe instances are busy" result. It
|
||||||
|
// loses the race against another client that grabbed the instance
|
||||||
|
// WaitNamedPipe just reported — the caller must wait and retry.
|
||||||
|
var errnoPipeBusy = syscall.Errno(231) // ERROR_PIPE_BUSY
|
||||||
|
|
||||||
|
// pipeConn adapts a pipe handle to io.ReadWriteCloser. Close flushes first
|
||||||
|
// (FlushFileBuffers blocks until the client has read the reply) and then
|
||||||
|
// disconnects — closing the bare handle right after writing can discard
|
||||||
|
// unread reply bytes, which clients see as an empty, failed request.
|
||||||
|
type pipeConn struct {
|
||||||
|
f *os.File
|
||||||
|
h windows.Handle
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *pipeConn) Read(p []byte) (int, error) { return c.f.Read(p) }
|
||||||
|
func (c *pipeConn) Write(p []byte) (int, error) { return c.f.Write(p) }
|
||||||
|
|
||||||
|
func (c *pipeConn) Close() error {
|
||||||
|
windows.FlushFileBuffers(c.h) //nolint:errcheck // best effort
|
||||||
|
windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort
|
||||||
|
return c.f.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ForceClose aborts the connection without flushing: disconnecting
|
||||||
|
// unblocks pending reads and writes at the cost of possibly discarding an
|
||||||
|
// unread reply. Used by the connection watchdog; normal closes flush.
|
||||||
|
func (c *pipeConn) ForceClose() error {
|
||||||
|
windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort
|
||||||
|
return c.f.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ask sends one command to the running service and returns its reply.
|
||||||
|
//
|
||||||
|
// The server keeps exactly one listening instance per connection, so
|
||||||
|
// concurrent clients race for it: WaitNamedPipe can report an instance
|
||||||
|
// that another client grabs before our CreateFile runs (ERROR_PIPE_BUSY).
|
||||||
|
// Retry on that — with the monitor polling status every second, a single
|
||||||
|
// attempt loses that race regularly.
|
||||||
|
func Ask(cmd string) (string, error) {
|
||||||
|
name, err := windows.UTF16PtrFromString(pipePath)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+134
-30
@@ -1,8 +1,10 @@
|
|||||||
// Package game detects processes outside gpu-turnstile's control that hold
|
// 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
|
// 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
|
// VRAM while they run. Three detection paths: an explicit process watch
|
||||||
// (GAME_PROCS) and a foreign-VRAM threshold via nvidia-smi
|
// list (GAME_PROCS), a foreign-VRAM threshold via nvidia-smi
|
||||||
// (GPU_FOREIGN_VRAM_MB) that catches anything not on the ignore list.
|
// (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
|
package game
|
||||||
|
|
||||||
import (
|
import (
|
||||||
@@ -11,6 +13,7 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
|
"slices"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
)
|
)
|
||||||
@@ -28,28 +31,33 @@ type computeApp struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Detector checks whether a foreign process holds the GPU. The zero value
|
// Detector checks whether a foreign process holds the GPU. The zero value
|
||||||
// (no watch list, no threshold) never detects anything; main only starts the
|
// (no watch list, no thresholds) never detects anything; main only starts
|
||||||
// poll loop when at least one path is configured.
|
// the poll loop when at least one path is configured.
|
||||||
type Detector struct {
|
type Detector struct {
|
||||||
procs map[string]bool // normalized names from GAME_PROCS
|
procs map[string]bool // normalized names from GAME_PROCS
|
||||||
vramMB int // foreign VRAM threshold; 0 = disabled
|
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
|
ignore map[string]bool // normalized names never counted as foreign
|
||||||
log *slog.Logger
|
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
|
// 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
|
// MiB, 3D-engine utilization threshold in percent (both 0 = disabled) and
|
||||||
// case-insensitively, with or without a trailing ".exe".
|
// ignore list. Names are matched case-insensitively, with or without a
|
||||||
func New(procs []string, vramMB int, ignore []string, log *slog.Logger) *Detector {
|
// trailing ".exe".
|
||||||
|
func New(procs []string, vramMB, utilPct int, ignore []string, log *slog.Logger) *Detector {
|
||||||
if log == nil {
|
if log == nil {
|
||||||
log = slog.Default()
|
log = slog.Default()
|
||||||
}
|
}
|
||||||
return &Detector{
|
return &Detector{
|
||||||
procs: nameSet(procs),
|
procs: nameSet(procs),
|
||||||
vramMB: vramMB,
|
vramMB: vramMB,
|
||||||
ignore: nameSet(ignore),
|
utilPct: utilPct,
|
||||||
log: log,
|
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
|
// 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
|
// 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
|
// 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) {
|
func (d *Detector) Check(ctx context.Context) ([]string, error) {
|
||||||
ps, psErr := processes()
|
ps, psErr := processes()
|
||||||
if d.vramMB <= 0 || d.noNvidia {
|
utils := d.engineUtil()
|
||||||
return d.detect(ps, nil), psErr
|
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)
|
return d.detect(ps, apps, utils), psErr
|
||||||
if errors.Is(err, exec.ErrNotFound) {
|
|
||||||
d.noNvidia = true
|
|
||||||
d.log.Warn("GPU_FOREIGN_VRAM_MB is set but nvidia-smi was not found; VRAM detection disabled")
|
|
||||||
return d.detect(ps, nil), nil
|
|
||||||
}
|
|
||||||
if err != nil {
|
|
||||||
return d.detect(ps, nil), err
|
|
||||||
}
|
|
||||||
return d.detect(ps, apps), nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// detect is the pure core of Check: given the process table and (optionally)
|
// engineUtil samples per-process 3D-engine utilization via PDH. The first
|
||||||
// the nvidia-smi compute-apps list, it returns the foreign holders.
|
// call only primes the rate counters and returns nil. A failing open or
|
||||||
func (d *Detector) detect(ps []Process, apps []computeApp) []string {
|
// 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
|
var holders []string
|
||||||
for _, p := range ps {
|
for _, p := range ps {
|
||||||
if d.procs[normName(p.Name)] {
|
if d.procs[normName(p.Name)] {
|
||||||
holders = append(holders, fmt.Sprintf("%s (pid %d)", p.Name, p.PID))
|
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))
|
names := make(map[int]string, len(ps))
|
||||||
for _, p := range ps {
|
for _, p := range ps {
|
||||||
names[p.PID] = p.Name
|
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))
|
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
|
return holders
|
||||||
}
|
}
|
||||||
@@ -132,6 +186,56 @@ func queryComputeApps(ctx context.Context) ([]computeApp, error) {
|
|||||||
return parseComputeApps(string(out))
|
return parseComputeApps(string(out))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// GPUStats is one nvidia-smi reading of the whole card.
|
||||||
|
type GPUStats struct {
|
||||||
|
UsedMB int
|
||||||
|
TotalMB int
|
||||||
|
TempC int // -1 when nvidia-smi reports N/A
|
||||||
|
FanPct int // -1 when N/A (some cards don't expose the fan)
|
||||||
|
}
|
||||||
|
|
||||||
|
// QueryGPUStats returns VRAM usage, temperature and fan speed via
|
||||||
|
// nvidia-smi. Unlike the per-process list this works under WDDM too.
|
||||||
|
func QueryGPUStats(ctx context.Context) (GPUStats, error) {
|
||||||
|
out, err := exec.CommandContext(ctx, "nvidia-smi",
|
||||||
|
"--query-gpu=memory.used,memory.total,temperature.gpu,fan.speed", "--format=csv,noheader,nounits").Output()
|
||||||
|
if err != nil {
|
||||||
|
return GPUStats{}, err
|
||||||
|
}
|
||||||
|
return parseGPUStats(string(out))
|
||||||
|
}
|
||||||
|
|
||||||
|
// parseGPUStats parses one "used, total, temp, fan" CSV line (MiB, °C,
|
||||||
|
// percent). The memory fields must be numeric; temperature and fan fall
|
||||||
|
// back to -1 on "N/A" and friends.
|
||||||
|
func parseGPUStats(out string) (GPUStats, error) {
|
||||||
|
fields := strings.Split(strings.TrimSpace(out), ",")
|
||||||
|
if len(fields) != 4 {
|
||||||
|
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(out))
|
||||||
|
}
|
||||||
|
num := func(s string) (int, error) {
|
||||||
|
return strconv.Atoi(strings.TrimSpace(s))
|
||||||
|
}
|
||||||
|
optional := func(s string) int {
|
||||||
|
n, err := num(s)
|
||||||
|
if err != nil {
|
||||||
|
return -1
|
||||||
|
}
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
var st GPUStats
|
||||||
|
var err error
|
||||||
|
if st.UsedMB, err = num(fields[0]); err != nil {
|
||||||
|
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(out))
|
||||||
|
}
|
||||||
|
if st.TotalMB, err = num(fields[1]); err != nil {
|
||||||
|
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(out))
|
||||||
|
}
|
||||||
|
st.TempC = optional(fields[2])
|
||||||
|
st.FanPct = optional(fields[3])
|
||||||
|
return st, nil
|
||||||
|
}
|
||||||
|
|
||||||
// parseComputeApps parses "pid, used_memory" CSV lines (no header, MiB
|
// parseComputeApps parses "pid, used_memory" CSV lines (no header, MiB
|
||||||
// units). Unsupported rows ("N/A" on WDDM) are skipped.
|
// units). Unsupported rows ("N/A" on WDDM) are skipped.
|
||||||
func parseComputeApps(out string) ([]computeApp, error) {
|
func parseComputeApps(out string) ([]computeApp, error) {
|
||||||
|
|||||||
@@ -45,7 +45,7 @@ func TestParseComputeApps(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func TestDetect(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)
|
[]string{"ollama", "python", "pythonw"}, nil)
|
||||||
ps := []Process{
|
ps := []Process{
|
||||||
{PID: 10, Name: "ollama.exe"},
|
{PID: 10, Name: "ollama.exe"},
|
||||||
@@ -58,7 +58,7 @@ func TestDetect(t *testing.T) {
|
|||||||
{PID: 40, UsedMB: 2048}, // foreign, above threshold
|
{PID: 40, UsedMB: 2048}, // foreign, above threshold
|
||||||
{PID: 50, UsedMB: 100}, // foreign but below threshold
|
{PID: 50, UsedMB: 100}, // foreign but below threshold
|
||||||
}
|
}
|
||||||
holders := d.detect(ps, apps)
|
holders := d.detect(ps, apps, nil)
|
||||||
if len(holders) != 2 {
|
if len(holders) != 2 {
|
||||||
t.Fatalf("got %v, want 2 holders", holders)
|
t.Fatalf("got %v, want 2 holders", holders)
|
||||||
}
|
}
|
||||||
@@ -70,13 +70,77 @@ func TestDetect(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestDetectEngineUtil(t *testing.T) {
|
||||||
|
d := New(nil, 0, 30, []string{"dwm", "python"}, nil)
|
||||||
|
ps := []Process{
|
||||||
|
{PID: 10, Name: "dwm.exe"},
|
||||||
|
{PID: 20, Name: "game.exe"},
|
||||||
|
{PID: 40, Name: "browser.exe"},
|
||||||
|
}
|
||||||
|
utils := map[int]float64{
|
||||||
|
10: 45, // ignored: dwm
|
||||||
|
20: 61, // foreign, above threshold
|
||||||
|
30: 82, // foreign, unknown name
|
||||||
|
40: 5, // below threshold
|
||||||
|
}
|
||||||
|
holders := d.detect(ps, nil, utils)
|
||||||
|
want := []string{
|
||||||
|
"game.exe (pid 20) using 61% GPU",
|
||||||
|
"unknown process (pid 30) using 82% GPU",
|
||||||
|
}
|
||||||
|
if !slices.Equal(holders, want) {
|
||||||
|
t.Errorf("got %v, want %v", holders, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestDetectNothingConfigured(t *testing.T) {
|
func TestDetectNothingConfigured(t *testing.T) {
|
||||||
d := New(nil, 0, nil, nil)
|
d := New(nil, 0, 0, nil, nil)
|
||||||
if got := d.detect([]Process{{PID: 1, Name: "game.exe"}}, nil); len(got) != 0 {
|
if got := d.detect([]Process{{PID: 1, Name: "game.exe"}}, nil, nil); len(got) != 0 {
|
||||||
t.Errorf("got %v, want none", got)
|
t.Errorf("got %v, want none", got)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestParseGPUEngineInstance(t *testing.T) {
|
||||||
|
pid, eng, ok := parseGPUEngineInstance("pid_1234_luid_0x00000000_0x00011A2B_phys_0_eng_0_engtype_3D")
|
||||||
|
if !ok || pid != 1234 || eng != "3D" {
|
||||||
|
t.Errorf("got %d, %q, %v", pid, eng, ok)
|
||||||
|
}
|
||||||
|
pid, eng, ok = parseGPUEngineInstance("pid_42_luid_0x0_0x0_phys_0_eng_1_engtype_Copy")
|
||||||
|
if !ok || pid != 42 || eng != "Copy" {
|
||||||
|
t.Errorf("got %d, %q, %v", pid, eng, ok)
|
||||||
|
}
|
||||||
|
for _, bad := range []string{"", "something", "pid_", "pid_x_luid", "pid_-1_luid_0"} {
|
||||||
|
if _, _, ok := parseGPUEngineInstance(bad); ok {
|
||||||
|
t.Errorf("%q parsed, want failure", bad)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseGPUStats(t *testing.T) {
|
||||||
|
st, err := parseGPUStats("4300, 16384, 55, 42\n")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if st.UsedMB != 4300 || st.TotalMB != 16384 || st.TempC != 55 || st.FanPct != 42 {
|
||||||
|
t.Errorf("got %+v", st)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cards that don't expose temperature/fan report N/A.
|
||||||
|
st, err = parseGPUStats("1024, 16384, N/A, N/A")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if st.TempC != -1 || st.FanPct != -1 {
|
||||||
|
t.Errorf("got %+v, want -1 for N/A fields", st)
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, bad := range []string{"", "1, 2", "x, 16384, 55, 42", "1024, x, 55, 42", "1, 2, 3, 4, 5"} {
|
||||||
|
if _, err := parseGPUStats(bad); err == nil {
|
||||||
|
t.Errorf("%q parsed, want failure", bad)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestProcessesLive(t *testing.T) {
|
func TestProcessesLive(t *testing.T) {
|
||||||
if runtime.GOOS != "windows" && runtime.GOOS != "linux" {
|
if runtime.GOOS != "windows" && runtime.GOOS != "linux" {
|
||||||
t.Skip("no process listing on this platform")
|
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))
|
||||||
|
}
|
||||||
+75
-2
@@ -9,6 +9,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
"sync"
|
"sync"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
// State is the current GPU occupancy state.
|
// State is the current GPU occupancy state.
|
||||||
@@ -30,10 +31,13 @@ type Lock struct {
|
|||||||
change chan struct{} // closed and replaced on every state change
|
change chan struct{} // closed and replaced on every state change
|
||||||
|
|
||||||
n int // LLM requests in flight
|
n int // LLM requests in flight
|
||||||
|
llmWaiting int // LLM requests blocked waiting for the GPU
|
||||||
imageActive bool // an image job holds the GPU
|
imageActive bool // an image job holds the GPU
|
||||||
imageQ []imageWaiter
|
imageQ []imageWaiter
|
||||||
nextID uint64
|
nextID uint64
|
||||||
external string // non-empty: a foreign process (e.g. a game) holds the GPU
|
external string // non-empty: a foreign process (e.g. a game) holds the GPU
|
||||||
|
detail string // what the current holder is doing (best effort)
|
||||||
|
since time.Time // when the current state began
|
||||||
|
|
||||||
log *slog.Logger
|
log *slog.Logger
|
||||||
}
|
}
|
||||||
@@ -41,7 +45,7 @@ type Lock struct {
|
|||||||
// New returns a ready-to-use Lock. log may be nil; if set, every state
|
// New returns a ready-to-use Lock. log may be nil; if set, every state
|
||||||
// transition is logged at debug level.
|
// transition is logged at debug level.
|
||||||
func New(log *slog.Logger) *Lock {
|
func New(log *slog.Logger) *Lock {
|
||||||
return &Lock{change: make(chan struct{}), log: log}
|
return &Lock{change: make(chan struct{}), log: log, since: time.Now()}
|
||||||
}
|
}
|
||||||
|
|
||||||
// broadcast wakes all waiters. Call with mu held.
|
// broadcast wakes all waiters. Call with mu held.
|
||||||
@@ -63,6 +67,7 @@ func (l *Lock) logTransition(msg string, args ...any) {
|
|||||||
func (l *Lock) SetExternal(holder string) {
|
func (l *Lock) SetExternal(holder string) {
|
||||||
l.mu.Lock()
|
l.mu.Lock()
|
||||||
l.external = holder
|
l.external = holder
|
||||||
|
l.since = time.Now()
|
||||||
l.broadcast()
|
l.broadcast()
|
||||||
l.mu.Unlock()
|
l.mu.Unlock()
|
||||||
l.logTransition("lock transition", "state", StateExternal, "holder", holder)
|
l.logTransition("lock transition", "state", StateExternal, "holder", holder)
|
||||||
@@ -73,6 +78,7 @@ func (l *Lock) SetExternal(holder string) {
|
|||||||
func (l *Lock) ClearExternal() {
|
func (l *Lock) ClearExternal() {
|
||||||
l.mu.Lock()
|
l.mu.Lock()
|
||||||
l.external = ""
|
l.external = ""
|
||||||
|
l.since = time.Now()
|
||||||
l.broadcast()
|
l.broadcast()
|
||||||
l.mu.Unlock()
|
l.mu.Unlock()
|
||||||
l.logTransition("lock transition", "state", StateIdle)
|
l.logTransition("lock transition", "state", StateIdle)
|
||||||
@@ -90,17 +96,31 @@ func (l *Lock) External() string {
|
|||||||
// while waiting; no state is changed in that case.
|
// while waiting; no state is changed in that case.
|
||||||
func (l *Lock) AcquireLLM(ctx context.Context) error {
|
func (l *Lock) AcquireLLM(ctx context.Context) error {
|
||||||
l.mu.Lock()
|
l.mu.Lock()
|
||||||
|
waiting := false
|
||||||
for l.imageActive || len(l.imageQ) > 0 || l.external != "" {
|
for l.imageActive || len(l.imageQ) > 0 || l.external != "" {
|
||||||
|
if !waiting {
|
||||||
|
l.llmWaiting++
|
||||||
|
waiting = true
|
||||||
|
}
|
||||||
ch := l.change
|
ch := l.change
|
||||||
l.mu.Unlock()
|
l.mu.Unlock()
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
|
l.mu.Lock()
|
||||||
|
l.llmWaiting--
|
||||||
|
l.mu.Unlock()
|
||||||
return ctx.Err()
|
return ctx.Err()
|
||||||
case <-ch:
|
case <-ch:
|
||||||
}
|
}
|
||||||
l.mu.Lock()
|
l.mu.Lock()
|
||||||
}
|
}
|
||||||
|
if waiting {
|
||||||
|
l.llmWaiting--
|
||||||
|
}
|
||||||
l.n++
|
l.n++
|
||||||
|
if l.n == 1 {
|
||||||
|
l.since = time.Now()
|
||||||
|
}
|
||||||
n := l.n
|
n := l.n
|
||||||
l.mu.Unlock()
|
l.mu.Unlock()
|
||||||
l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
|
l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
|
||||||
@@ -129,6 +149,8 @@ func (l *Lock) ReleaseLLM() {
|
|||||||
l.n--
|
l.n--
|
||||||
n := l.n
|
n := l.n
|
||||||
if l.n == 0 {
|
if l.n == 0 {
|
||||||
|
l.since = time.Now()
|
||||||
|
l.detail = ""
|
||||||
l.broadcast()
|
l.broadcast()
|
||||||
}
|
}
|
||||||
l.mu.Unlock()
|
l.mu.Unlock()
|
||||||
@@ -156,6 +178,7 @@ func (l *Lock) AcquireImage(ctx context.Context) error {
|
|||||||
if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive && l.external == "" {
|
if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive && l.external == "" {
|
||||||
l.imageQ = l.imageQ[1:]
|
l.imageQ = l.imageQ[1:]
|
||||||
l.imageActive = true
|
l.imageActive = true
|
||||||
|
l.since = time.Now()
|
||||||
l.mu.Unlock()
|
l.mu.Unlock()
|
||||||
l.logTransition("lock transition", "state", StateImage)
|
l.logTransition("lock transition", "state", StateImage)
|
||||||
return nil
|
return nil
|
||||||
@@ -184,6 +207,8 @@ func (l *Lock) AcquireImage(ctx context.Context) error {
|
|||||||
func (l *Lock) ReleaseImage() {
|
func (l *Lock) ReleaseImage() {
|
||||||
l.mu.Lock()
|
l.mu.Lock()
|
||||||
l.imageActive = false
|
l.imageActive = false
|
||||||
|
l.since = time.Now()
|
||||||
|
l.detail = ""
|
||||||
l.broadcast()
|
l.broadcast()
|
||||||
l.mu.Unlock()
|
l.mu.Unlock()
|
||||||
l.logTransition("lock transition", "state", StateIdle)
|
l.logTransition("lock transition", "state", StateIdle)
|
||||||
@@ -206,3 +231,51 @@ func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) {
|
|||||||
}
|
}
|
||||||
return state, l.n, l.imageActive || len(l.imageQ) > 0
|
return state, l.n, l.imageActive || len(l.imageQ) > 0
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetDetail records what the current holder is doing (e.g. the request
|
||||||
|
// path), for status displays. Best effort: overwritten by each new holder,
|
||||||
|
// cleared when the GPU goes idle.
|
||||||
|
func (l *Lock) SetDetail(detail string) {
|
||||||
|
l.mu.Lock()
|
||||||
|
l.detail = detail
|
||||||
|
l.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Status is a point-in-time view of the lock for monitoring.
|
||||||
|
type Status struct {
|
||||||
|
State State
|
||||||
|
Detail string
|
||||||
|
LLMInflight int
|
||||||
|
LLMWaiting int
|
||||||
|
ImageActive bool
|
||||||
|
ImageQueue int
|
||||||
|
External string
|
||||||
|
Since time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// Status reports the full lock state, including waiters and how long the
|
||||||
|
// current state has held.
|
||||||
|
func (l *Lock) Status() Status {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
s := Status{
|
||||||
|
Detail: l.detail,
|
||||||
|
LLMInflight: l.n,
|
||||||
|
LLMWaiting: l.llmWaiting,
|
||||||
|
ImageActive: l.imageActive,
|
||||||
|
ImageQueue: len(l.imageQ),
|
||||||
|
External: l.external,
|
||||||
|
Since: l.since,
|
||||||
|
}
|
||||||
|
switch {
|
||||||
|
case l.imageActive:
|
||||||
|
s.State = StateImage
|
||||||
|
case l.n > 0:
|
||||||
|
s.State = StateLLM
|
||||||
|
case l.external != "":
|
||||||
|
s.State = StateExternal
|
||||||
|
default:
|
||||||
|
s.State = StateIdle
|
||||||
|
}
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|||||||
@@ -60,13 +60,21 @@ func (c *Client) Probe(ctx context.Context) error {
|
|||||||
|
|
||||||
type psResponse struct {
|
type psResponse struct {
|
||||||
Models []struct {
|
Models []struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
|
SizeVRAM int64 `json:"size_vram"` // bytes resident in VRAM (0 = RAM-only)
|
||||||
} `json:"models"`
|
} `json:"models"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// LoadedModels returns the names of models currently held in memory.
|
// LoadedModel is one model currently held in memory.
|
||||||
func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
|
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)
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/api/ps", nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
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 {
|
if err := json.NewDecoder(resp.Body).Decode(&ps); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
models := make([]string, 0, len(ps.Models))
|
models := make([]LoadedModel, 0, len(ps.Models))
|
||||||
for _, m := range ps.Models {
|
for _, m := range ps.Models {
|
||||||
if m.Name != "" {
|
name := m.Name
|
||||||
models = append(models, m.Name)
|
if name == "" {
|
||||||
} else {
|
name = m.Model
|
||||||
models = append(models, 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
|
return models, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -491,6 +491,7 @@ func (s *Server) OllamaHandler() http.Handler {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
defer s.cfg.Lock.ReleaseLLM()
|
defer s.cfg.Lock.ReleaseLLM()
|
||||||
|
s.cfg.Lock.SetDetail("ollama: " + r.Method + " " + r.URL.Path)
|
||||||
s.ollamaProxy.ServeHTTP(w, r)
|
s.ollamaProxy.ServeHTTP(w, r)
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
@@ -590,6 +591,7 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
|
|||||||
}
|
}
|
||||||
s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds())
|
s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds())
|
||||||
log.Info("image lock acquired")
|
log.Info("image lock acquired")
|
||||||
|
s.cfg.Lock.SetDetail("comfy: POST /prompt")
|
||||||
|
|
||||||
if s.cfg.ComfySup != nil && !comfyFirst {
|
if s.cfg.ComfySup != nil && !comfyFirst {
|
||||||
if err := s.cfg.ComfySup.EnsureRunning(); err != nil {
|
if err := s.cfg.ComfySup.EnsureRunning(); err != nil {
|
||||||
|
|||||||
@@ -95,6 +95,8 @@ RestartSec=5s
|
|||||||
|
|
||||||
DynamicUser=yes
|
DynamicUser=yes
|
||||||
StateDirectory=%s
|
StateDirectory=%s
|
||||||
|
RuntimeDirectory=%s
|
||||||
|
RuntimeDirectoryMode=0755
|
||||||
%sProtectSystem=strict
|
%sProtectSystem=strict
|
||||||
ProtectHome=yes
|
ProtectHome=yes
|
||||||
PrivateTmp=yes
|
PrivateTmp=yes
|
||||||
@@ -117,7 +119,7 @@ SystemCallErrorNumber=EPERM
|
|||||||
|
|
||||||
[Install]
|
[Install]
|
||||||
WantedBy=multi-user.target
|
WantedBy=multi-user.target
|
||||||
`, exePath, configPath, Name, bind)
|
`, exePath, configPath, Name, Name, bind)
|
||||||
}
|
}
|
||||||
|
|
||||||
// copyFile copies src to dst, creating dst with the given mode.
|
// copyFile copies src to dst, creating dst with the given mode.
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ func TestRenderUnit(t *testing.T) {
|
|||||||
"WantedBy=multi-user.target",
|
"WantedBy=multi-user.target",
|
||||||
"DynamicUser=yes",
|
"DynamicUser=yes",
|
||||||
"StateDirectory=gpu-turnstile",
|
"StateDirectory=gpu-turnstile",
|
||||||
|
"RuntimeDirectory=gpu-turnstile",
|
||||||
"ProtectSystem=strict",
|
"ProtectSystem=strict",
|
||||||
"NoNewPrivileges=yes",
|
"NoNewPrivileges=yes",
|
||||||
"RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6",
|
"RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6",
|
||||||
|
|||||||
@@ -164,6 +164,24 @@ func (p *Process) Running() bool {
|
|||||||
return p.cmd != nil
|
return p.cmd != nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Status describes the child for status displays: "external" when something
|
||||||
|
// else serves the port, "running" once ready, "starting" while the child
|
||||||
|
// boots, "stopped" otherwise.
|
||||||
|
func (p *Process) Status() string {
|
||||||
|
p.mu.Lock()
|
||||||
|
defer p.mu.Unlock()
|
||||||
|
switch {
|
||||||
|
case p.external:
|
||||||
|
return "external"
|
||||||
|
case p.cmd == nil:
|
||||||
|
return "stopped"
|
||||||
|
case p.ready:
|
||||||
|
return "running"
|
||||||
|
default:
|
||||||
|
return "starting"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Ready reports whether the server has answered a probe since its last
|
// Ready reports whether the server has answered a probe since its last
|
||||||
// (re)start. Health checks use it to tell "starting up" from "outage".
|
// (re)start. Health checks use it to tell "starting up" from "outage".
|
||||||
func (p *Process) Ready() bool {
|
func (p *Process) Ready() bool {
|
||||||
|
|||||||
Reference in New Issue
Block a user