Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3214352454 | ||
|
|
2a7e7e7b47 | ||
|
|
31ee6013db | ||
|
|
5506163493 | ||
|
|
18a7336c02 | ||
|
|
0796092b80 | ||
|
|
897163042c | ||
|
|
e58ff33912 | ||
|
|
2040b30c94 | ||
|
|
a07b6726ea | ||
|
|
d2c49e52fa | ||
|
|
e98331bb9e | ||
|
|
54174787d5 | ||
|
|
d259b2e96c | ||
|
|
8fccd333aa | ||
|
|
4bd5f34ce7 | ||
|
|
b9d3f91403 | ||
|
|
98d2d714ca |
@@ -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 |
|
||||||
|
|||||||
+261
-53
@@ -16,6 +16,7 @@ import (
|
|||||||
"os"
|
"os"
|
||||||
"os/signal"
|
"os/signal"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"reflect"
|
||||||
"runtime"
|
"runtime"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -39,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
|
||||||
@@ -59,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, updateNow, monitor, elevatedChild bool, rest []string) {
|
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild bool, rest []string) {
|
||||||
rest = args[:0]
|
rest = args[:0]
|
||||||
for i := 0; i < len(args); i++ {
|
for i := 0; i < len(args); i++ {
|
||||||
switch {
|
switch {
|
||||||
@@ -82,6 +84,8 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
|
|||||||
forceUpdate = true
|
forceUpdate = true
|
||||||
case args[i] == "--update-now" || args[i] == "-update-now":
|
case args[i] == "--update-now" || args[i] == "-update-now":
|
||||||
updateNow = true
|
updateNow = true
|
||||||
|
case args[i] == "--reload-env" || args[i] == "-reload-env":
|
||||||
|
reloadEnv = true
|
||||||
case args[i] == "--monitor" || args[i] == "-monitor" || args[i] == "-m":
|
case args[i] == "--monitor" || args[i] == "-monitor" || args[i] == "-m":
|
||||||
monitor = true
|
monitor = true
|
||||||
case args[i] == "--elevated-child":
|
case args[i] == "--elevated-child":
|
||||||
@@ -90,7 +94,7 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
|
|||||||
rest = append(rest, args[i])
|
rest = append(rest, args[i])
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, monitor, elevatedChild, rest
|
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, rest
|
||||||
}
|
}
|
||||||
|
|
||||||
// versionLine is printed at the top of every help and error screen.
|
// versionLine is printed at the top of every help and error screen.
|
||||||
@@ -109,6 +113,9 @@ Usage:
|
|||||||
(no admin needed when the service runs)
|
(no admin needed when the service runs)
|
||||||
gpu-turnstile --update-now like --force-update, but only
|
gpu-turnstile --update-now like --force-update, but only
|
||||||
through the running service
|
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-turnstile -m | --monitor live status view (downstreams,
|
||||||
GPU lock, queue); Ctrl+C quits
|
GPU lock, queue); Ctrl+C quits
|
||||||
gpu-turnstile -h | --help this help
|
gpu-turnstile -h | --help this help
|
||||||
@@ -141,12 +148,12 @@ func fatalUsage(format string, args ...any) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, monitor, elevatedChild, args := parseFlags(os.Args[1:])
|
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, args := parseFlags(os.Args[1:])
|
||||||
if showVersion {
|
if showVersion {
|
||||||
fmt.Println(version)
|
fmt.Println(version)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !monitor && !elevatedChild && len(args) == 0
|
bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !reloadEnv && !monitor && !elevatedChild && len(args) == 0
|
||||||
if help || (bare && stdoutIsTerminal()) {
|
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
|
||||||
@@ -171,7 +178,9 @@ func main() {
|
|||||||
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):
|
case updateNow && (install || remove || forceUpdate):
|
||||||
fatalUsage("error: --update-now cannot be combined with other commands")
|
fatalUsage("error: --update-now cannot be combined with other commands")
|
||||||
case monitor && (install || remove || forceUpdate || updateNow):
|
case reloadEnv && (install || remove || forceUpdate || updateNow):
|
||||||
|
fatalUsage("error: --reload-env cannot be combined with other commands")
|
||||||
|
case monitor && (install || remove || forceUpdate || updateNow || reloadEnv):
|
||||||
fatalUsage("error: --monitor cannot be combined with other commands")
|
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))
|
||||||
@@ -181,6 +190,8 @@ func main() {
|
|||||||
os.Exit(forceUpdateCommand(configPath, elevatedChild))
|
os.Exit(forceUpdateCommand(configPath, elevatedChild))
|
||||||
case updateNow:
|
case updateNow:
|
||||||
os.Exit(updateNowCommand())
|
os.Exit(updateNowCommand())
|
||||||
|
case reloadEnv:
|
||||||
|
os.Exit(reloadEnvCommand())
|
||||||
case monitor:
|
case monitor:
|
||||||
os.Exit(monitorCommand())
|
os.Exit(monitorCommand())
|
||||||
}
|
}
|
||||||
@@ -202,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)
|
||||||
}
|
}
|
||||||
@@ -210,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)
|
||||||
}
|
}
|
||||||
@@ -451,6 +462,18 @@ func updateNowCommand() int {
|
|||||||
return printControlReply(reply)
|
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
|
||||||
@@ -531,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)
|
||||||
@@ -544,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,
|
||||||
@@ -569,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,
|
||||||
@@ -703,11 +729,14 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes, health)
|
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
|
||||||
@@ -747,6 +776,22 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
|
|
||||||
var u *update.Updater
|
var u *update.Updater
|
||||||
var exePath string
|
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) {}
|
applyStaged := func(to string) {}
|
||||||
if cfg.AutoUpdate {
|
if cfg.AutoUpdate {
|
||||||
if p, err := os.Executable(); err != nil {
|
if p, err := os.Executable(); err != nil {
|
||||||
@@ -754,34 +799,38 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
} else {
|
} else {
|
||||||
exePath = p
|
exePath = p
|
||||||
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}
|
||||||
// applyStaged is shared by the hourly loop and the control
|
|
||||||
// channel; the once guard keeps a second trigger from
|
|
||||||
// double-waiting on the GPU lock.
|
|
||||||
var once sync.Once
|
|
||||||
applyStaged = func(to string) {
|
applyStaged = func(to string) {
|
||||||
if !isService {
|
if !isService {
|
||||||
log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply", "version", to)
|
log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply", "version", to)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
log.Warn("auto-update: staged; restarting once the GPU is idle", "version", to)
|
restartWhenIdle("update to " + to)
|
||||||
once.Do(func() {
|
|
||||||
go func() {
|
|
||||||
if waitForIdle(ctx, lk, 24*time.Hour) {
|
|
||||||
log.Warn("auto-update: restarting to apply update")
|
|
||||||
os.Exit(exitCodeUpdate)
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
go updateLoop(ctx, cfg.UpdateInterval, log, u, exePath, applyStaged)
|
go updateLoop(ctx, cfg.UpdateInterval, log, u, exePath, applyStaged)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// The control channel (status for --monitor, update-now trigger) is
|
// The command handler backs both the local control channel (served
|
||||||
// served whenever running as a service, independent of AUTO_UPDATE.
|
// 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 {
|
if isService {
|
||||||
serveControl(ctx, log, u, exePath, applyStaged,
|
serveControl(ctx, log, handler)
|
||||||
statusProvider(cfg, lk, comfySup, health, started))
|
}
|
||||||
|
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 {
|
||||||
@@ -805,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
|
||||||
@@ -824,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
|
||||||
@@ -947,19 +1040,32 @@ func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// serveControl opens the local control channel (named pipe on Windows,
|
// controlHandler builds the command handler shared by the control channel
|
||||||
// unix socket on Linux) so unprivileged local users can query status
|
// (named pipe / unix socket) and the web UI's action endpoints: status for
|
||||||
// (--monitor) and trigger an update check (--force-update/--update-now)
|
// --monitor / the UI, an update check (--force-update/--update-now) and a
|
||||||
// without admin rights. The update payload is signature-verified regardless
|
// config reload (--reload-env), all safe for unprivileged local users. The
|
||||||
// of who asks; triggers are rate-limited to one per minute so the channel
|
// update payload is signature-verified regardless of who asks; update
|
||||||
// cannot be used to spam restarts. u is nil when AUTO_UPDATE=false.
|
// triggers are rate-limited to one per minute and reloads to one per two
|
||||||
func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string), status func() string) {
|
// 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 mu sync.Mutex
|
||||||
var lastTrigger time.Time
|
var lastTrigger, lastReload time.Time
|
||||||
h := func(cmd string) string {
|
return func(cmd string) string {
|
||||||
switch cmd {
|
switch cmd {
|
||||||
case control.CmdStatus:
|
case control.CmdStatus:
|
||||||
return "OK " + status()
|
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:
|
case control.CmdUpdateNow:
|
||||||
default:
|
default:
|
||||||
return "ERR unknown command: " + cmd
|
return "ERR unknown command: " + cmd
|
||||||
@@ -986,6 +1092,12 @@ func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exeP
|
|||||||
applyStaged(to)
|
applyStaged(to)
|
||||||
return "OK updated from " + version + " to " + to + "; the service restarts once the GPU is idle"
|
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 {
|
if err := control.Serve(ctx, h, log); err != nil {
|
||||||
log.Warn("control channel disabled", "err", err)
|
log.Warn("control channel disabled", "err", err)
|
||||||
}
|
}
|
||||||
@@ -1021,6 +1133,16 @@ type statusDownstream struct {
|
|||||||
URL string `json:"url"`
|
URL string `json:"url"`
|
||||||
Up bool `json:"up"`
|
Up bool `json:"up"`
|
||||||
Managed string `json:"managed,omitempty"`
|
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 {
|
type statusLock struct {
|
||||||
@@ -1038,22 +1160,108 @@ type statusSnapshot struct {
|
|||||||
UptimeS int64 `json:"uptime_s"`
|
UptimeS int64 `json:"uptime_s"`
|
||||||
Downstreams []statusDownstream `json:"downstreams"`
|
Downstreams []statusDownstream `json:"downstreams"`
|
||||||
Lock statusLock `json:"lock"`
|
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
|
// MonitorNote is set client-side (never over the wire) when the
|
||||||
// monitor's own binary differs from the service's version.
|
// monitor's own binary differs from the service's version.
|
||||||
MonitorNote string `json:"-"`
|
MonitorNote string `json:"-"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// statusProvider assembles the one-line JSON snapshot for CmdStatus.
|
type statusGPU struct {
|
||||||
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time) func() string {
|
// 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 {
|
return func() string {
|
||||||
snap := statusSnapshot{
|
snap := statusSnapshot{
|
||||||
Version: version,
|
Version: version,
|
||||||
UptimeS: int64(time.Since(started).Seconds()),
|
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 != "" {
|
if cfg.OllamaURL != "" {
|
||||||
snap.Downstreams = append(snap.Downstreams, statusDownstream{
|
d := statusDownstream{
|
||||||
Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
|
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 != "" {
|
if cfg.ComfyURL != "" {
|
||||||
d := statusDownstream{Name: "comfy", URL: cfg.ComfyURL, Up: health.get("comfy")}
|
d := statusDownstream{Name: "comfy", URL: cfg.ComfyURL, Up: health.get("comfy")}
|
||||||
@@ -1062,15 +1270,15 @@ func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Proces
|
|||||||
}
|
}
|
||||||
snap.Downstreams = append(snap.Downstreams, d)
|
snap.Downstreams = append(snap.Downstreams, d)
|
||||||
}
|
}
|
||||||
st := lk.Status()
|
lst := lk.Status()
|
||||||
snap.Lock = statusLock{
|
snap.Lock = statusLock{
|
||||||
State: string(st.State),
|
State: string(lst.State),
|
||||||
Detail: st.Detail,
|
Detail: lst.Detail,
|
||||||
LLMInflight: st.LLMInflight,
|
LLMInflight: lst.LLMInflight,
|
||||||
LLMWaiting: st.LLMWaiting,
|
LLMWaiting: lst.LLMWaiting,
|
||||||
ImageQueue: st.ImageQueue,
|
ImageQueue: lst.ImageQueue,
|
||||||
External: st.External,
|
External: lst.External,
|
||||||
SinceS: int64(time.Since(st.Since).Seconds()),
|
SinceS: int64(time.Since(lst.Since).Seconds()),
|
||||||
}
|
}
|
||||||
b, err := json.Marshal(snap)
|
b, err := json.Marshal(snap)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -21,17 +21,27 @@ const (
|
|||||||
cCyan = "\x1b[36m"
|
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,
|
// monitorCommand renders a live status view of the running service,
|
||||||
// refreshed every second from the control channel. When the service
|
// refreshed every second from the control channel. When the service
|
||||||
// reports a different version and the executable on disk changed (the
|
// reports a different version and the executable on disk changed (the
|
||||||
// updater replaced it), the monitor restarts itself onto the new binary.
|
// updater replaced it), the monitor restarts itself onto the new binary.
|
||||||
// Ctrl+C quits.
|
// Hotkeys: q quits, u triggers an update check on the service, r asks the
|
||||||
|
// service to reload its config file.
|
||||||
func monitorCommand() int {
|
func monitorCommand() int {
|
||||||
if !stdoutIsTerminal() {
|
if !stdoutIsTerminal() {
|
||||||
fmt.Fprintln(os.Stderr, "gpu-turnstile: --monitor needs an interactive terminal")
|
fmt.Fprintln(os.Stderr, "gpu-turnstile: --monitor needs an interactive terminal")
|
||||||
return 1
|
return 1
|
||||||
}
|
}
|
||||||
enableVirtualTerminal()
|
enableVirtualTerminal()
|
||||||
|
restore := enableRawKeys()
|
||||||
|
defer func() {
|
||||||
|
if restore != nil {
|
||||||
|
restore()
|
||||||
|
}
|
||||||
|
}()
|
||||||
fmt.Print("\x1b[2J") // clear once; frames then redraw in place
|
fmt.Print("\x1b[2J") // clear once; frames then redraw in place
|
||||||
defer fmt.Print(cReset + "\n")
|
defer fmt.Print(cReset + "\n")
|
||||||
exe, _ := os.Executable()
|
exe, _ := os.Executable()
|
||||||
@@ -39,7 +49,38 @@ func monitorCommand() int {
|
|||||||
if st, err := os.Stat(exe); err == nil {
|
if st, err := os.Stat(exe); err == nil {
|
||||||
exeStamp = st.ModTime()
|
exeStamp = st.ModTime()
|
||||||
}
|
}
|
||||||
for {
|
|
||||||
|
keys := make(chan byte, 8)
|
||||||
|
go readKeys(keys)
|
||||||
|
ticker := time.NewTicker(time.Second)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
var note string
|
||||||
|
var noteAt time.Time
|
||||||
|
noteCh := make(chan string, 1)
|
||||||
|
askPending := false
|
||||||
|
|
||||||
|
// ask sends a one-shot command to the service and reports the reply in
|
||||||
|
// the note line. Only one ask runs at a time.
|
||||||
|
ask := func(cmd, busy, label string) {
|
||||||
|
if askPending {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
askPending = true
|
||||||
|
note, noteAt = busy, time.Now()
|
||||||
|
go func() {
|
||||||
|
reply, err := control.Ask(cmd)
|
||||||
|
if err != nil {
|
||||||
|
noteCh <- label + ": no answer from the service"
|
||||||
|
return
|
||||||
|
}
|
||||||
|
msg := strings.TrimPrefix(reply, "OK ")
|
||||||
|
msg = strings.TrimPrefix(msg, "ERR ")
|
||||||
|
noteCh <- label + ": " + msg
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
poll := func() string {
|
||||||
frame := renderWaiting()
|
frame := renderWaiting()
|
||||||
if reply, err := control.Ask(control.CmdStatus); err == nil {
|
if reply, err := control.Ask(control.CmdStatus); err == nil {
|
||||||
if msg, ok := strings.CutPrefix(reply, "OK "); ok {
|
if msg, ok := strings.CutPrefix(reply, "OK "); ok {
|
||||||
@@ -50,16 +91,64 @@ func monitorCommand() int {
|
|||||||
fmt.Print("\x1b[2J\x1b[H")
|
fmt.Print("\x1b[2J\x1b[H")
|
||||||
fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version)
|
fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version)
|
||||||
restartSelf(exe, "--monitor")
|
restartSelf(exe, "--monitor")
|
||||||
return 0
|
return "" // re-execed; this process exits below
|
||||||
}
|
}
|
||||||
snap.MonitorNote = fmt.Sprintf("note: the service runs %s, this monitor is %s", snap.Version, version)
|
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())
|
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
|
fmt.Print("\x1b[H" + frame + "\x1b[J") // home, frame, clear below
|
||||||
time.Sleep(time.Second)
|
select {
|
||||||
|
case <-ticker.C:
|
||||||
|
if note != "" && time.Since(noteAt) > 15*time.Second {
|
||||||
|
note = ""
|
||||||
|
}
|
||||||
|
case k, ok := <-keys:
|
||||||
|
if !ok {
|
||||||
|
keys = nil
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
switch k {
|
||||||
|
case 'q', 'Q', 3: // q or Ctrl+C (raw mode delivers it as a byte)
|
||||||
|
return 0
|
||||||
|
case 'u', 'U':
|
||||||
|
ask(control.CmdUpdateNow, "checking for updates…", "update")
|
||||||
|
case 'r', 'R':
|
||||||
|
ask(control.CmdReloadEnv, "reloading config…", "reload")
|
||||||
|
}
|
||||||
|
case n := <-noteCh:
|
||||||
|
askPending = false
|
||||||
|
note, noteAt = n, time.Now()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// readKeys reads single keypresses from stdin (raw mode was enabled by the
|
||||||
|
// caller) and delivers them until stdin fails.
|
||||||
|
func readKeys(keys chan<- byte) {
|
||||||
|
defer close(keys)
|
||||||
|
buf := make([]byte, 1)
|
||||||
|
for {
|
||||||
|
n, err := os.Stdin.Read(buf)
|
||||||
|
if n > 0 {
|
||||||
|
keys <- buf[0]
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -82,7 +171,7 @@ func restartSelf(exe string, args ...string) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func renderWaiting() string {
|
func renderWaiting() string {
|
||||||
return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n"
|
return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n\x1b[K\n" + hotkeysLine
|
||||||
}
|
}
|
||||||
|
|
||||||
// renderMonitor draws one full frame. Each line ends with \x1b[K (clear to
|
// renderMonitor draws one full frame. Each line ends with \x1b[K (clear to
|
||||||
@@ -109,7 +198,9 @@ func renderMonitor(snap statusSnapshot, width int) string {
|
|||||||
b.WriteString(cDim + " " + strings.Repeat("─", width-2) + cReset + "\x1b[K\n")
|
b.WriteString(cDim + " " + strings.Repeat("─", width-2) + cReset + "\x1b[K\n")
|
||||||
|
|
||||||
for _, d := range snap.Downstreams {
|
for _, d := range snap.Downstreams {
|
||||||
b.WriteString(renderDownstream(d) + "\x1b[K\n")
|
busy := (d.Name == "ollama" && snap.Lock.State == "llm" && snap.Lock.LLMInflight > 0) ||
|
||||||
|
(d.Name == "comfy" && snap.Lock.State == "image")
|
||||||
|
b.WriteString(renderDownstream(d, busy) + "\x1b[K\n")
|
||||||
}
|
}
|
||||||
b.WriteString("\x1b[K\n")
|
b.WriteString("\x1b[K\n")
|
||||||
b.WriteString(renderLock(snap.Lock) + "\x1b[K\n")
|
b.WriteString(renderLock(snap.Lock) + "\x1b[K\n")
|
||||||
@@ -117,17 +208,86 @@ func renderMonitor(snap statusSnapshot, width int) string {
|
|||||||
b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n",
|
b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n",
|
||||||
cYellow, snap.Lock.ImageQueue, cReset))
|
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 != "" {
|
if snap.MonitorNote != "" {
|
||||||
b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
|
b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
|
||||||
}
|
}
|
||||||
|
b.WriteString("\x1b[K\n")
|
||||||
|
b.WriteString(hotkeysLine)
|
||||||
return b.String()
|
return b.String()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// renderGPU renders the GPU line: VRAM usage (when nvidia-smi answered),
|
||||||
|
// the game detector's last finding, and how long ago it ran.
|
||||||
|
func renderGPU(g statusGPU) string {
|
||||||
|
s := " GPU: "
|
||||||
|
if g.Known {
|
||||||
|
s += renderVRAM(g.UsedMB, g.TotalMB)
|
||||||
|
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).
|
// printableLen counts characters without ANSI escapes (ASCII-only content).
|
||||||
func printableLen(s string) int { return len(s) }
|
func printableLen(s string) int { return len(s) }
|
||||||
|
|
||||||
func renderDownstream(d statusDownstream) string {
|
func renderDownstream(d statusDownstream, busy bool) string {
|
||||||
url := cDim + d.URL + cReset
|
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 {
|
switch d.Managed {
|
||||||
case "stopped":
|
case "stopped":
|
||||||
return fmt.Sprintf(" %s○%s %-8s %sstopped (managed — starts on demand)%s %s",
|
return fmt.Sprintf(" %s○%s %-8s %sstopped (managed — starts on demand)%s %s",
|
||||||
@@ -141,7 +301,7 @@ func renderDownstream(d statusDownstream) string {
|
|||||||
suffix = " (external)"
|
suffix = " (external)"
|
||||||
}
|
}
|
||||||
if d.Up {
|
if d.Up {
|
||||||
return fmt.Sprintf(" %s●%s %-8s %sUP%s%s %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, url)
|
return fmt.Sprintf(" %s●%s %-8s %sUP%s%s%s %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, detail, url)
|
||||||
}
|
}
|
||||||
return fmt.Sprintf(" %s●%s %-8s %sDOWN%s %s", cRed, cReset, d.Name, cRed, cReset, url)
|
return fmt.Sprintf(" %s●%s %-8s %sDOWN%s %s", cRed, cReset, d.Name, cRed, cReset, url)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
package main
|
package main
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"gpu-turnstile/internal/config"
|
||||||
|
"reflect"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
)
|
)
|
||||||
@@ -29,6 +31,47 @@ func TestRenderMonitor(t *testing.T) {
|
|||||||
t.Errorf("frame missing %q:\n%s", want, frame)
|
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) {
|
func TestFmtDur(t *testing.T) {
|
||||||
@@ -39,3 +82,29 @@ func TestFmtDur(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestDiffConfig(t *testing.T) {
|
||||||
|
a := config.Defaults()
|
||||||
|
b := a
|
||||||
|
if got := diffConfig(a, b); len(got) != 0 {
|
||||||
|
t.Fatalf("identical configs: got %v", got)
|
||||||
|
}
|
||||||
|
b.LogLevel = -4
|
||||||
|
b.GameProcs = []string{"game.exe"}
|
||||||
|
got := diffConfig(a, b)
|
||||||
|
if len(got) != 2 || got[0] != "GAME_PROCS" || got[1] != "LOGLEVEL" {
|
||||||
|
t.Fatalf("got %v, want [GAME_PROCS LOGLEVEL]", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Every Config field must carry an env tag so user-facing output (the
|
||||||
|
// reload diff) can name the setting the user would actually change.
|
||||||
|
func TestConfigFieldsHaveEnvTags(t *testing.T) {
|
||||||
|
typ := reflect.TypeOf(config.Config{})
|
||||||
|
for i := 0; i < typ.NumField(); i++ {
|
||||||
|
f := typ.Field(i)
|
||||||
|
if f.Tag.Get("env") == "" {
|
||||||
|
t.Errorf("config.Config.%s has no env tag", f.Name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -19,3 +19,21 @@ func termWidth() int {
|
|||||||
}
|
}
|
||||||
return int(ws.Col)
|
return int(ws.Col)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// enableRawKeys switches the terminal to per-keypress mode (ICANON and ECHO
|
||||||
|
// off) and returns the restore function, nil when stdin is not a terminal.
|
||||||
|
func enableRawKeys() func() {
|
||||||
|
fd := int(os.Stdin.Fd())
|
||||||
|
term, err := unix.IoctlGetTermios(fd, unix.TCGETS)
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
raw := *term
|
||||||
|
raw.Lflag &^= unix.ICANON | unix.ECHO
|
||||||
|
raw.Cc[unix.VMIN] = 1
|
||||||
|
raw.Cc[unix.VTIME] = 0
|
||||||
|
if err := unix.IoctlSetTermios(fd, unix.TCSETS, &raw); err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return func() { unix.IoctlSetTermios(fd, unix.TCSETS, term) } //nolint:errcheck
|
||||||
|
}
|
||||||
|
|||||||
@@ -27,3 +27,23 @@ func termWidth() int {
|
|||||||
}
|
}
|
||||||
return int(info.Window.Right-info.Window.Left) + 1
|
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,243 @@
|
|||||||
|
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; min-width:14ch; text-align:center; }
|
||||||
|
button:hover { background:#2c2c2c; color:#eee; }
|
||||||
|
button .key { color:#26c6da; }
|
||||||
|
.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"><span class="key">u</span>pdate now</button>
|
||||||
|
<button id="br"><span class="key">r</span>eload 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.7
|
image: git.rambossek.at/public/gpu-turnstile:v0.3.3
|
||||||
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},
|
||||||
|
|||||||
@@ -7,6 +7,13 @@
|
|||||||
// triggers, so the worst a local user can cause is a cheap, throttled
|
// triggers, so the worst a local user can cause is a cheap, throttled
|
||||||
// check and a GPU-idle-gated restart onto a signed binary.
|
// 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
|
// Protocol: the client writes one command line, the server answers with
|
||||||
// one reply line ("OK ..." or "ERR ...") and hangs up.
|
// one reply line ("OK ..." or "ERR ...") and hangs up.
|
||||||
package control
|
package control
|
||||||
@@ -17,6 +24,7 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"strings"
|
"strings"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
|
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
|
||||||
@@ -26,6 +34,10 @@ const CmdUpdateNow = "update-now"
|
|||||||
// CmdStatus asks for a one-line JSON status snapshot (monitor mode).
|
// CmdStatus asks for a one-line JSON status snapshot (monitor mode).
|
||||||
const CmdStatus = "status"
|
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.
|
// ErrUnavailable means no running service offers the control channel.
|
||||||
var ErrUnavailable = errors.New("control channel unavailable")
|
var ErrUnavailable = errors.New("control channel unavailable")
|
||||||
|
|
||||||
@@ -33,9 +45,50 @@ var ErrUnavailable = errors.New("control channel unavailable")
|
|||||||
// line. It must start with "OK " or "ERR ".
|
// line. It must start with "OK " or "ERR ".
|
||||||
type Handler func(cmd string) string
|
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.
|
// serveConn runs the line protocol on one accepted connection.
|
||||||
func serveConn(c io.ReadWriteCloser, h Handler) {
|
func serveConn(c io.ReadWriteCloser, h Handler) {
|
||||||
defer c.Close()
|
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')
|
line, err := bufio.NewReader(io.LimitReader(c, 4096)).ReadString('\n')
|
||||||
cmd := strings.TrimSpace(line)
|
cmd := strings.TrimSpace(line)
|
||||||
if cmd == "" {
|
if cmd == "" {
|
||||||
|
|||||||
@@ -37,12 +37,24 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return // shutting down
|
return // shutting down
|
||||||
}
|
}
|
||||||
go serveConn(c, h)
|
uc := unixConn{c}
|
||||||
|
if !serve(uc, h) {
|
||||||
|
uc.ForceClose()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
return nil
|
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.
|
// Ask sends one command to the running service and returns its reply.
|
||||||
func Ask(cmd string) (string, error) {
|
func Ask(cmd string) (string, error) {
|
||||||
c, err := net.DialTimeout("unix", sockPath, 2*time.Second)
|
c, err := net.DialTimeout("unix", sockPath, 2*time.Second)
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestRoundTrip(t *testing.T) {
|
func TestRoundTrip(t *testing.T) {
|
||||||
@@ -44,3 +45,59 @@ func TestEmptyReplyIsUnavailable(t *testing.T) {
|
|||||||
t.Fatalf("err = %v, want ErrUnavailable", err)
|
t.Fatalf("err = %v, want ErrUnavailable", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestServeCap(t *testing.T) {
|
||||||
|
for i := 0; i < maxConns; i++ {
|
||||||
|
connSem <- struct{}{}
|
||||||
|
}
|
||||||
|
defer func() {
|
||||||
|
for i := 0; i < maxConns; i++ {
|
||||||
|
<-connSem
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
server, client := net.Pipe()
|
||||||
|
defer server.Close()
|
||||||
|
defer client.Close()
|
||||||
|
if serve(server, func(string) string { return "OK" }) {
|
||||||
|
t.Fatal("serve accepted a connection beyond the cap")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// forcePipe records ForceClose calls for the watchdog test.
|
||||||
|
type forcePipe struct {
|
||||||
|
net.Conn
|
||||||
|
forced chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c forcePipe) ForceClose() error {
|
||||||
|
err := c.Conn.Close()
|
||||||
|
close(c.forced)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConnWatchdog(t *testing.T) {
|
||||||
|
old := connTimeout
|
||||||
|
connTimeout = 50 * time.Millisecond
|
||||||
|
defer func() { connTimeout = old }()
|
||||||
|
|
||||||
|
server, client := net.Pipe()
|
||||||
|
defer client.Close()
|
||||||
|
fc := forcePipe{Conn: server, forced: make(chan struct{})}
|
||||||
|
done := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
serveConn(fc, func(string) string { return "OK" })
|
||||||
|
close(done)
|
||||||
|
}()
|
||||||
|
// The client never sends anything; the watchdog must tear the
|
||||||
|
// connection down instead of blocking forever.
|
||||||
|
select {
|
||||||
|
case <-fc.forced:
|
||||||
|
case <-time.After(5 * time.Second):
|
||||||
|
t.Fatal("watchdog did not force-close the stalled connection")
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
case <-time.After(5 * time.Second):
|
||||||
|
t.Fatal("serveConn still blocked after the force close")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import (
|
|||||||
"log/slog"
|
"log/slog"
|
||||||
"os"
|
"os"
|
||||||
"syscall"
|
"syscall"
|
||||||
|
"time"
|
||||||
"unsafe"
|
"unsafe"
|
||||||
|
|
||||||
"golang.org/x/sys/windows"
|
"golang.org/x/sys/windows"
|
||||||
@@ -86,7 +87,10 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
|||||||
windows.CloseHandle(pipe)
|
windows.CloseHandle(pipe)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
go serveConn(&pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}, h)
|
conn := &pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}
|
||||||
|
if !serve(conn, h) {
|
||||||
|
conn.ForceClose()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
return nil
|
return nil
|
||||||
@@ -96,6 +100,11 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
|||||||
// called" result, which means the connection is established.
|
// called" result, which means the connection is established.
|
||||||
var errnoPipeConnected = syscall.Errno(535) // ERROR_PIPE_CONNECTED
|
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
|
// pipeConn adapts a pipe handle to io.ReadWriteCloser. Close flushes first
|
||||||
// (FlushFileBuffers blocks until the client has read the reply) and then
|
// (FlushFileBuffers blocks until the client has read the reply) and then
|
||||||
// disconnects — closing the bare handle right after writing can discard
|
// disconnects — closing the bare handle right after writing can discard
|
||||||
@@ -114,22 +123,45 @@ func (c *pipeConn) Close() error {
|
|||||||
return c.f.Close()
|
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.
|
// 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) {
|
func Ask(cmd string) (string, error) {
|
||||||
name, err := windows.UTF16PtrFromString(pipePath)
|
name, err := windows.UTF16PtrFromString(pipePath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
if err := waitNamedPipe(name, 2000); err != nil {
|
deadline := time.Now().Add(5 * time.Second)
|
||||||
return "", ErrUnavailable
|
for {
|
||||||
|
if err := waitNamedPipe(name, 2000); err != nil {
|
||||||
|
return "", ErrUnavailable
|
||||||
|
}
|
||||||
|
handle, err := windows.CreateFile(name,
|
||||||
|
windows.GENERIC_READ|windows.GENERIC_WRITE, 0, nil,
|
||||||
|
windows.OPEN_EXISTING, 0, 0)
|
||||||
|
if err == errnoPipeBusy {
|
||||||
|
if time.Now().After(deadline) {
|
||||||
|
return "", ErrUnavailable
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return "", ErrUnavailable
|
||||||
|
}
|
||||||
|
f := os.NewFile(uintptr(handle), pipePath)
|
||||||
|
defer f.Close()
|
||||||
|
return readReply(f, cmd)
|
||||||
}
|
}
|
||||||
handle, err := windows.CreateFile(name,
|
|
||||||
windows.GENERIC_READ|windows.GENERIC_WRITE, 0, nil,
|
|
||||||
windows.OPEN_EXISTING, 0, 0)
|
|
||||||
if err != nil {
|
|
||||||
return "", ErrUnavailable
|
|
||||||
}
|
|
||||||
f := os.NewFile(uintptr(handle), pipePath)
|
|
||||||
defer f.Close()
|
|
||||||
return readReply(f, cmd)
|
|
||||||
}
|
}
|
||||||
|
|||||||
+134
-30
@@ -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))
|
||||||
|
}
|
||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user