Compare commits

..
65 Commits
Author SHA1 Message Date
mram 897163042c Pin compose example to v0.3.1
ci / test (push) Successful in 15s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 15s
2026-09-22 17:08:40 +02:00
mram e58ff33912 Monitor: show GPU temperature and fan speed
Extends the existing per-tick nvidia-smi query with temperature.gpu and
fan.speed (no extra call); N/A values (cards without fan telemetry) are
simply omitted from the GPU line.
2026-09-22 16:56:01 +02:00
mram 2040b30c94 Monitor: add r hotkey for config reload
Generalizes the one-shot ask behind u into a shared helper; r sends
reload-env and reports the service's reply (unchanged / restarting with
the changed setting names / invalid config) in the note line. Footer now
lists q, u, r.
2026-09-22 13:33:46 +02:00
mram a07b6726ea Pin compose example to v0.3.0
ci / test (push) Successful in 15s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 15s
2026-09-22 13:31:05 +02:00
mram d2c49e52fa Monitor: show Ollama's loaded models, their VRAM footprint, and busy state
The status snapshot now queries /api/ps (2s timeout so a wedged Ollama
cannot stall the channel) and the ollama line renders e.g.
'UP · llama3.1:8b (4.8 GiB VRAM) · busy'; comfy gets the busy marker
too while an image job runs.
2026-09-22 12:44:04 +02:00
mram e98331bb9e Pin compose example to v0.2.9
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m8s
ci / release (push) Successful in 16s
2026-09-22 11:23:17 +02:00
mram 54174787d5 Fix control-pipe client race: retry CreateFile on ERROR_PIPE_BUSY
WaitNamedPipe can report an instance that a concurrent client (the
monitor polls status every second) grabs before our CreateFile runs;
the single-attempt Ask then failed with 'no answer from the service'.
Retry the wait+open until a 5s overall deadline.
2026-09-22 11:22:27 +02:00
mram d259b2e96c Add GPU_FOREIGN_UTIL_PCT: game detection via per-process GPU 3D-engine usage
Reads the same PDH counters as Task Manager (\GPU Engine(*)\Utilization
Percentage, locale-independent via PdhAddEnglishCounterW), which cover
graphics work under WDDM — games are caught without an exe list and the
offender is named. Windows-only; a missing or failing counter disables
the path with one log line. dwm (the desktop compositor) joins the
default ignore list.
2026-09-22 10:53:29 +02:00
mram 8fccd333aa Env names in reload diff; monitor shows last VRAM check result; harden control channel against floods
- diffConfig reports the env var names (from new struct tags) instead of
  Go field names, so the reload-env reply names what the user can change
- the monitor's GPU line now shows the game detector's last finding
  (external holders or none) and how long ago the check ran
- control channel: 10s per-connection watchdog (abortive force-close),
  cap of 32 concurrent connections, reload-env rate-limited; command
  read was already capped at 4 KiB
2026-09-22 09:24:27 +02:00
mram 4bd5f34ce7 Pin compose example to v0.2.8
ci / test (push) Successful in 15s
ci / docker (push) Successful in 1m6s
ci / release (push) Successful in 15s
2026-09-22 09:10:14 +02:00
mram b9d3f91403 Monitor: hotkeys (q quit, u update now) with footer line; show GPU VRAM usage 2026-09-22 09:10:14 +02:00
mram 98d2d714ca Add --reload-env: service re-reads and validates its config, restarts when idle if it changed 2026-09-22 09:02:25 +02:00
mram 7a6409831c Pin compose example to v0.2.7
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 16s
2026-09-22 08:47:47 +02:00
mram c7b8747d85 Fix control pipe accept loop: create instances only after a client connects; CLI force-update never touches the service log file 2026-09-22 08:43:58 +02:00
mram 09d4e81eab Pin compose example to v0.2.6
ci / test (push) Successful in 15s
ci / docker (push) Successful in 1m6s
ci / release (push) Successful in 15s
2026-09-22 08:36:36 +02:00
mram 1264afe40e Monitor restarts itself when the service updated its binary on disk 2026-09-22 08:34:47 +02:00
mram c95d401a36 Pin compose example to v0.2.5
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m6s
ci / release (push) Successful in 16s
2026-09-22 08:32:06 +02:00
mram f7ea30a494 Fix control pipe races: treat ERROR_PIPE_CONNECTED as success, flush+disconnect before close so replies are never discarded 2026-09-22 08:31:13 +02:00
mram 620be9d57c Add short flags: -i (install), -r (remove), -m (monitor) 2026-09-22 08:23:40 +02:00
mram cffb4151a6 Pin compose example to v0.2.4
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m9s
ci / release (push) Successful in 15s
2026-09-22 08:19:13 +02:00
mram 16dc7fe462 Add --monitor: live status view (downstream health, GPU lock, queue) via the control channel 2026-09-22 08:17:05 +02:00
mram a74cd49fe5 Pin compose example to v0.2.3
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m9s
ci / release (push) Successful in 15s
2026-09-21 23:03:45 +02:00
mram cc2a2cad27 Add --update-now: update trigger that only goes through the service's control channel 2026-09-21 23:03:45 +02:00
mram f707d07fd8 Local control channel: unprivileged users can trigger --force-update via the running service 2026-09-21 23:00:38 +02:00
mram ca33a3db82 Pin compose example to v0.2.2
ci / test (push) Successful in 16s
ci / docker (push) Successful in 1m8s
ci / release (push) Successful in 15s
2026-09-21 22:34:56 +02:00
mram 507af1dbff Managed ComfyUI picks up Comfy-Desktop shared models/input/output automatically 2026-09-21 22:34:10 +02:00
mram 45bc9fe27c Strip ANSI escapes from managed ComfyUI output in the log; spawn child with NO_COLOR/TERM=dumb 2026-09-21 22:29:41 +02:00
mram 5b752d5637 Pin compose example to v0.2.1
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m16s
ci / release (push) Successful in 15s
2026-09-21 22:19:49 +02:00
mram b55d548c5f install-service: grant venv base interpreter (pyvenv.cfg home) outside COMFY_DIR 2026-09-21 22:19:08 +02:00
mram 2d5be72436 --force-update reports from/to versions; elevated parent no longer claims "update applied" when nothing changed 2026-09-21 21:58:50 +02:00
mram 21e40a1774 Pin compose example to v0.2.0
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 15s
2026-09-21 21:47:48 +02:00
mram bdc844872d Install narrates itself: stop/copy/start steps and the potentially long icacls tree grant 2026-09-21 21:36:45 +02:00
mram 299b6dc0bb Distinguish EACCES from ENOENT in the COMFY_DIR startup check 2026-09-21 21:32:45 +02:00
mram 9589e58ce6 Install opens up COMFY_DIR for the sandboxed service: ACL grant on Windows, BindPaths on Linux 2026-09-21 21:30:35 +02:00
mram be0bb36317 Pin compose example to v0.1.10
ci / test (push) Successful in 28s
ci / docker (push) Successful in 1m15s
ci / release (push) Successful in 17s
2026-09-21 21:20:49 +02:00
mram 9b267533c9 COMFY_DIR alone manages ComfyUI: derive the launch command from the standard venv layout 2026-09-21 21:18:33 +02:00
mram 6f092ddc12 Default GAME_POLL_INTERVAL to 15s: nvidia-smi polls keep the GPU awake 2026-09-21 21:10:21 +02:00
mram 0228ccc296 Pin compose example to v0.1.9
ci / test (push) Successful in 17s
ci / docker (push) Successful in 1m17s
ci / release (push) Successful in 19s
2026-09-21 20:57:34 +02:00
mram 9997913929 Sync the env file at startup, not just at install: updates append new settings 2026-09-21 20:56:32 +02:00
mram 482731ed9e Pin compose example to v0.1.8
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m11s
ci / release (push) Successful in 15s
2026-09-21 17:19:16 +02:00
mram e4bdc92ece Game detection: foreign GPU holders take an external lock hold (GAME_PROCS, GPU_FOREIGN_VRAM_MB) 2026-09-21 17:16:56 +02:00
mram 14120bf4a4 Supervisor probes before spawning: an external server on the port is used, never fought or killed 2026-09-21 16:06:28 +02:00
mram d7566329ae Managed ComfyUI: COMFY_CMD starts it on demand, idle stop frees VRAM (internal/supervise) 2026-09-21 13:48:07 +02:00
mram e43ad02fc4 Ignore staged-update leftovers next to the dev exe 2026-09-21 13:14:33 +02:00
mram 232f5b61f2 Periodic upstream health checks (HEALTH_INTERVAL, default 30s); log down/recovered transitions 2026-09-21 13:14:18 +02:00
mram 5e7a042cad CFG_VER is always a concrete version: dev builds stamp v0.0.0, never "dev" 2026-09-21 12:42:32 +02:00
mram 14c2e30478 Clarify listener wording: Ollama-compatible clients, not Ollama-facing 2026-09-21 12:38:03 +02:00
mram dcdb2cc72e Pin compose example to v0.1.7
ci / test (push) Successful in 13s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 16s
2026-09-21 10:04:29 +02:00
mram 44a8e9fbdc CFG_VER/APP_VER in the env file: invalid configs replaced, updates follow APP_VER (dev/stable/pin) 2026-09-21 09:28:29 +02:00
mram 272a3a467d force-update: single-shot with a 30s fail-fast timeout 2026-09-21 09:09:27 +02:00
mram cedf28a809 force-update does not require consumer URLs (ErrNoConsumer sentinel) 2026-09-21 09:05:46 +02:00
mram e46e92bee8 Sample env carries a version marker; install appends settings missing since the writing version 2026-09-21 08:54:47 +02:00
mram cf48796183 Install writes a fully commented sample env file when no config exists 2026-09-21 08:48:51 +02:00
mram dd3187f951 Install sets a default LOG_FILE in the env file on Windows (no console as a service) 2026-09-21 08:37:25 +02:00
mram 1394a76eea Make service install converge: stop running service, refresh binary/config only on change, restart only if it was running 2026-09-21 08:30:55 +02:00
mram da02457fd5 Pin compose example to v0.1.6
ci / test (push) Successful in 14s
ci / docker (push) Successful in 1m7s
ci / release (push) Successful in 14s
2026-09-21 08:22:42 +02:00
mram 6a6359a630 Add --force-update: immediate signed-update check, stage + service restart 2026-09-21 08:22:11 +02:00
mram e4e4348e8d Start the Windows service right after install (parity with enable --now) 2026-09-21 08:12:13 +02:00
mram b421bb7bfb Add --version and print the version atop every help and error screen 2026-09-21 08:10:46 +02:00
mram 909918657f Fix install/remove success message (installd -> installed) 2026-09-21 08:08:24 +02:00
mram a0435c858e Bare run in a terminal prints the help screen; add -h/--help
A zero-argument invocation now shows the usage text when stdout is a
console (double-clicked exe, interactive shell). Without a terminal —
Docker entrypoint, services, pipes — a bare invocation still starts the
proxy, so the container image and service behavior are unchanged.
2026-09-21 08:06:35 +02:00
mram fbab0bba33 Relaunch through UAC when (un)installing the service unprivileged
--install-service/--remove-service on Windows no longer fail with
'Access is denied' from a normal shell: the process re-runs itself via
ShellExecuteEx 'runas', waits for the elevated child and mirrors its
exit code. The child gets --elevated-child and pauses for a keypress so
its console output stays readable. Declining the prompt reports
'UAC prompt declined'.
2026-09-21 08:00:24 +02:00
mram 802a64280f Self-install into canonical layout on both platforms, --no-copy to opt out
Windows: --install-service creates %ProgramFiles%\gpu-turnstile and
%ProgramData%\gpu-turnstile, copies the exe and (if absent) the env
file in, and registers the copy. Linux: binary goes to
/var/lib/gpu-turnstile (not /usr/local/sbin: replacing a running binary
needs directory write, which must not be granted on a shared system dir
to a sandboxed service). --no-copy registers the current location
as-is on both platforms.
2026-09-21 07:52:33 +02:00
mram a88955e35c Sandbox the systemd unit: DynamicUser, read-only FS, no capabilities
The Linux install now mirrors the Windows virtual-account hardening: the
unit runs with DynamicUser=yes (transient per-service UID, no login),
ProtectSystem=strict with only StateDirectory writable (the install dir,
so self-update can rewrite the binary), NoNewPrivileges, empty
capability sets, restricted address families and a @system-service
syscall filter. Install copies the binary to /var/lib/gpu-turnstile and
the config to /etc/gpu-turnstile.env; Remove cleans up the unit and
binary but keeps the config.
2026-09-21 00:11:25 +02:00
mram 97624470eb Install the Windows service as the NT SERVICE virtual account only
--install-service now registers the service under
NT SERVICE\gpu-turnstile (low-privilege, per-service, no password) and
grants it modify access to the install dir (for self-updates) and the
LOG_FILE dir, plus read access to an external config file. Grants run
after CreateService because the virtual account's SID does not exist
before registration; a failed grant rolls back the registration.
2026-09-20 23:53:54 +02:00
45 changed files with 5504 additions and 284 deletions
+3
View File
@@ -4,3 +4,6 @@
/compose.yml /compose.yml
/signing/ /signing/
/gpu-turnstile.exe.old
/gpu-turnstile.exe.new
/gpu-turnstile.exe~
+130 -27
View File
@@ -2,9 +2,10 @@
GPU arbitration proxy for Ollama + ComfyUI. One consumer GPU is shared by an GPU arbitration proxy for Ollama + ComfyUI. One consumer GPU is shared by an
LLM server (Ollama) and an image generator (ComfyUI); gpu-turnstile sits in LLM server (Ollama) and an image generator (ComfyUI); gpu-turnstile sits in
front of both and guarantees the GPU is always in exactly one of three states: front of both and guarantees the GPU is always in exactly one of four states:
`idle`, `llm` (N ≥ 1 Ollama requests in flight), or `image` (exactly one `idle`, `llm` (N ≥ 1 Ollama requests in flight), `image` (exactly one
ComfyUI job, Ollama models unloaded). See [SPEC.md](SPEC.md) for the full ComfyUI job, Ollama models unloaded), or `external` (a foreign process such
as a game holds the GPU). See [SPEC.md](SPEC.md) for the full
design. design.
gpu-turnstile listens on the ports the services normally use; the actual gpu-turnstile listens on the ports the services normally use; the actual
@@ -32,8 +33,9 @@ Open WebUI / n8n ────► :8188 ───┘
Each consumer is enabled by setting its URL (`OLLAMA_URL`, `COMFY_URL`) and 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. Future consumers (e.g. local game the Ollama unload/warm steps are skipped. A third, URL-less consumer —
detection) plug into the same lock the same way. detection of foreign GPU holders such as games — is enabled by `GAME_PROCS`,
`GPU_FOREIGN_VRAM_MB` and/or `GPU_FOREIGN_UTIL_PCT` (see below).
## Configuration ## Configuration
@@ -45,8 +47,8 @@ override file values. Invalid values fail at startup.
| Var | Default | Meaning | | Var | Default | Meaning |
|---|---|---| |---|---|---|
| `LISTEN_OLLAMA` | `:11434` | Ollama-facing listener | | `LISTEN_OLLAMA` | `:11434` | Listener for Ollama-compatible clients |
| `LISTEN_COMFY` | `:8188` | ComfyUI-facing listener | | `LISTEN_COMFY` | `:8188` | Listener for ComfyUI clients |
| `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 |
@@ -56,12 +58,22 @@ override file values. Invalid values fail at startup.
| `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400599, e.g. 429) | | `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400599, e.g. 429) |
| `BUSY_RETRY_AFTER` | `30` | Seconds sent as `Retry-After` on busy responses (both modes) | | `BUSY_RETRY_AFTER` | `30` | Seconds sent as `Retry-After` on busy responses (both modes) |
| `WARM_MODEL` | _(empty)_ | Model to reload after an image job (off by default) | | `WARM_MODEL` | _(empty)_ | Model to reload after an image job (off by default) |
| `COMFY_CMD` | _(derived from `COMFY_DIR`; both empty = unmanaged)_ | Supervise ComfyUI: start on demand, stop when idle to free VRAM. Requires `COMFY_URL` |
| `COMFY_DIR` | _(empty)_ | Standard venv install root: set alone to supervise ComfyUI with the derived command (`.venv` + `main.py`); also the working directory for `COMFY_CMD` |
| `COMFY_IDLE_TIMEOUT` | `5m` | Stop the managed ComfyUI after this long idle |
| `COMFY_START_TIMEOUT` | `2m` | Max wait for the managed ComfyUI to come up |
| `GAME_PROCS` | _(empty = disabled)_ | Process names (comma-separated); while any runs, the GPU counts as held: requests wait, Ollama unloads, managed ComfyUI stops |
| `GPU_FOREIGN_VRAM_MB` | `0` (disabled) | Also treat the GPU as held when a non-ignored process uses more VRAM than this (needs nvidia-smi) |
| `GPU_FOREIGN_UTIL_PCT` | `0` (disabled) | Also treat the GPU as held when a non-ignored process uses more than this percent of the GPU 3D engine (Windows PDH counters — what Task Manager shows; catches games without an exe list) |
| `GPU_IGNORE_PROCS` | `ollama,ollama app,ollama_llama_server,python,pythonw,dwm` | Process names never counted as foreign GPU users |
| `GAME_POLL_INTERVAL` | `15s` | How often game/VRAM detection runs (don't go below ~10s — nvidia-smi polls keep the GPU awake) |
| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` works as an alias | | `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 |
| `LOG_FILE` | _(empty)_ | Append logs to this file instead of stderr | | `LOG_FILE` | _(empty)_ | Append logs to this file instead of stderr |
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading | | `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs | | `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
| `PROBE_TIMEOUT` | `5s` | Startup probe of both upstreams | | `PROBE_TIMEOUT` | `5s` | Startup probe of both upstreams (also per-probe health check timeout) |
| `HEALTH_INTERVAL` | `30s` | Periodic upstream probe; down/recovered changes are logged |
| `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job | | `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job |
| `WARM_TIMEOUT` | `2m` | Warm-model reload after an image job | | `WARM_TIMEOUT` | `2m` | Warm-model reload after an image job |
| `SHUTDOWN_TIMEOUT` | `10s` | Graceful shutdown on SIGINT/SIGTERM | | `SHUTDOWN_TIMEOUT` | `10s` | Graceful shutdown on SIGINT/SIGTERM |
@@ -72,10 +84,11 @@ override file values. Invalid values fail at startup.
| `UPDATE_INTERVAL` | `6h` | Auto-update check interval | | `UPDATE_INTERVAL` | `6h` | Auto-update check interval |
| `UPDATE_REPO` | `https://git.rambossek.at/PUBLIC/gpu-turnstile` | Repository checked for releases | | `UPDATE_REPO` | `https://git.rambossek.at/PUBLIC/gpu-turnstile` | Repository checked for releases |
| `UPDATE_ASSET` | `gpu-turnstile.exe` | Release asset to download | | `UPDATE_ASSET` | `gpu-turnstile.exe` | Release asset to download |
| `APP_VER` | `stable` | `dev` disables updates, `stable` tracks latest, or pin an exact `vX.Y.Z` |
## Observability ## Observability
- `GET /healthz` (both listeners): `{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}` - `GET /healthz` (both listeners): `{"state":"idle|llm|image|external","llm_inflight":N,"image_pending":B}`
- `GET /metrics` (both listeners): Prometheus text format — `gpu_turnstile_state`, - `GET /metrics` (both listeners): Prometheus text format — `gpu_turnstile_state`,
`gpu_turnstile_llm_inflight`, `gpu_turnstile_image_pending`, `gpu_turnstile_llm_inflight`, `gpu_turnstile_image_pending`,
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds` `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
@@ -87,60 +100,150 @@ override file values. Invalid values fail at startup.
class) in text mode, which renders in `docker compose logs` on Windows class) in text mode, which renders in `docker compose logs` on Windows
Terminal. Set `NO_COLOR` to disable colors. Terminal. Set `NO_COLOR` to disable colors.
## Managed ComfyUI (`COMFY_CMD` / `COMFY_DIR`)
Don't want ComfyUI running 24/7 (it holds VRAM even when idle — and the
Desktop app kills its server when you close it)? gpu-turnstile can supervise
it: the first request starts it, it stops again after `COMFY_IDLE_TIMEOUT`
(default 5m) without work, freeing the GPU for games or the LLM.
The easy way — point `COMFY_DIR` at a standard venv install (a folder with
`.venv` and `main.py`, or `.venv` and `ComfyUI\main.py`) and the launch
command is derived from it, including `--port` from `COMFY_URL`:
```
COMFY_URL=http://127.0.0.1:8189
COMFY_DIR=C:\ComfyUI
```
For other layouts, spell the command out yourself (run it once manually to
confirm it works):
```
COMFY_URL=http://127.0.0.1:8189
COMFY_CMD="C:\ComfyUI\.venv\Scripts\python.exe" main.py --port 8189
COMFY_DIR=C:\ComfyUI
```
`POST /prompt` waits for the server to answer before taking the GPU lock
(LLM traffic keeps flowing while torch loads); crashes are logged and the
next request respawns. Shutting gpu-turnstile down stops the child too.
Running the ComfyUI Desktop app alongside is safe: if something already
answers on the port, gpu-turnstile just uses it instead of spawning
(and never kills it — it only ever stops its own child). If the managed
instance already holds the port when you open the desktop app, the
desktop's server is the one that fails to bind.
One catch when gpu-turnstile runs as a service: the sandboxed service
account may not enter your user profile, so a ComfyUI install under
`C:\Users\...` (or `/home/...`) fails with "Access is denied".
`--install-service` fixes that automatically — it grants
`NT SERVICE\gpu-turnstile` recursive access to `COMFY_DIR` on Windows and
adds a `BindPaths=` to the systemd unit on Linux. Re-run it after changing
`COMFY_DIR`; or grant by hand from an admin shell:
`icacls "<COMFY_DIR>" /grant "NT SERVICE\gpu-turnstile:(OI)(CI)M" /T`.
## Game detection
Want to game on the same GPU without Ollama/ComfyUI squatting on the VRAM?
gpu-turnstile can watch for foreign GPU holders and, while one is active,
make LLM/image requests wait (or 503, per `LLM_BUSY_MODE`), unload Ollama's
models and stop the managed ComfyUI so the game gets the memory. Three
detection paths, each optional, polled every `GAME_POLL_INTERVAL` (15s):
```
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_UTIL_PCT` reads the same per-process GPU engine counters as
Task Manager (PDH), which cover graphics work under Windows' WDDM driver —
so games show up without naming them, and the offender is named in the log
and monitor. Windows-only. `GAME_PROCS` matches running process names
(case-insensitive, `.exe` optional). `GPU_FOREIGN_VRAM_MB` asks nvidia-smi
which processes hold GPU memory and treats anything not in
`GPU_IGNORE_PROCS` above the threshold as foreign — handy as a catch-all on
Linux, but under WDDM graphics-only games may not show up in nvidia-smi's
per-process list, so on Windows prefer `GPU_FOREIGN_UTIL_PCT`. When the
game exits, requests resume automatically.
## Build and run ## Build and run
```sh ```sh
go build ./cmd/gpu-turnstile go build ./cmd/gpu-turnstile
./gpu-turnstile OLLAMA_URL=http://127.0.0.1:11435 COMFY_URL=http://127.0.0.1:8189 ./gpu-turnstile
``` ```
Running the binary with no arguments in a terminal prints the help screen
(same as `-h`/`--help`); without a terminal (services, containers) a bare
invocation starts the proxy.
### Run natively on Windows (current primary deployment) ### Run natively on Windows (current primary deployment)
Download `gpu-turnstile.exe` from a release, put a `gpu-turnstile.env` Download `gpu-turnstile.exe` from a release and install it as a Windows
next to it, and run it — or install it as a Windows service from an service — no admin shell needed, a UAC prompt appears automatically and
elevated shell: the elevated child does the work (its window waits for Enter so you can
read the result):
```sh ```sh
gpu-turnstile.exe --install-service # auto-start service, recovery = restart gpu-turnstile.exe --install-service # installs into Program Files, auto-start
gpu-turnstile.exe --install-service --no-copy # register in place instead
gpu-turnstile.exe --remove-service gpu-turnstile.exe --remove-service
``` ```
The service uses the config file (services have no convenient Layout: `C:\Program Files\gpu-turnstile\` holds the exe and
environment); set `LOG_FILE` in it since there is no console. `gpu-turnstile.env`, logs go to `C:\ProgramData\gpu-turnstile\`. If you
install without a config, the installer writes a sample env file with every
setting commented and explained — only `LOG_FILE` is active (a service has
no console). Your own `LOG_FILE` setting is always kept. The service always
runs
as the virtual account `NT SERVICE\gpu-turnstile` (low-privilege,
per-service, no password); the installer automatically grants it write
access to the install and data directories — nothing else to do.
Suggested layout: `C:\Program Files\gpu-turnstile\` for the exe and Re-running `--install-service` is safe: it stops a running service,
`gpu-turnstile.env`, logs under `C:\ProgramData\gpu-turnstile\` via replaces the installed binary only if it changed, fixes the registration
`LOG_FILE`. The service runs as `LocalSystem` by default, which can write only where it drifted, and restarts the service only if it was running.
the install directory for self-updates. For least privilege, run it as the
virtual account `NT SERVICE\gpu-turnstile` and grant write access to just
those two directories.
### Run natively on Linux (systemd) ### Run natively on Linux (systemd)
The same binary works on Linux. Install it as a systemd service as root: The same binary works on Linux. Install it as a systemd service as root:
```sh ```sh
gpu-turnstile --install-service # writes + enables + starts the unit gpu-turnstile --install-service # installs into /var/lib/gpu-turnstile, enables + starts
gpu-turnstile --install-service --no-copy # register in place instead
gpu-turnstile --remove-service gpu-turnstile --remove-service
``` ```
The unit (`/etc/systemd/system/gpu-turnstile.service`) is `Type=notify`: The unit (`/etc/systemd/system/gpu-turnstile.service`) is `Type=notify`:
`systemctl start` blocks until the listeners are actually bound, a 30 s `systemctl start` blocks until the listeners are actually bound, a 30 s
watchdog restarts the process if it wedges, and logs land in the journal watchdog restarts the process if it wedges, and logs land in the journal
(`journalctl -u gpu-turnstile -f`) unless `LOG_FILE` is set. Put the (`journalctl -u gpu-turnstile -f`) unless `LOG_FILE` is set. Install
config in a `gpu-turnstile.env` next to the binary (or pass copies the binary to `/var/lib/gpu-turnstile/` and the config to
`-config /path` during install). The notify integration is a no-op in `/etc/gpu-turnstile.env` (edit that one after installing). The service
runs sandboxed with `DynamicUser=yes` — a transient low-privilege UID,
read-only filesystem except its install dir (so self-update keeps
working), no capabilities, syscall-filtered: same least-privilege idea as
the Windows virtual account. The notify integration is a no-op in
containers and interactive shells. containers and interactive shells.
**Auto-update is on by default**: the binary checks the repo's latest **Auto-update is on by default**: the binary checks the repo's latest
release on startup and every `UPDATE_INTERVAL`, verifies the Ed25519 release on startup and every `UPDATE_INTERVAL`, verifies the Ed25519
signature of the download against the public key embedded at build time, signature of the download against the public key embedded at build time,
and — once the GPU lock is idle — restarts the service onto the new and — once the GPU lock is idle — restarts the service onto the new
version. Disable with `AUTO_UPDATE=false`. Releases are signed by CI with version. `APP_VER` controls the target: `dev` disables updates, `stable`
(the default) tracks the latest release, and an exact `vX.Y.Z` pins that
release (even as a downgrade or to replace a dev build). Disable entirely
with `AUTO_UPDATE=false`. Releases are signed by CI with
OpenSSL; the matching public key lives in `internal/update/pubkey.go` OpenSSL; the matching public key lives in `internal/update/pubkey.go`
(one-time setup: `openssl genpkey -algorithm ed25519 -out private.pem`, (one-time setup: `openssl genpkey -algorithm ed25519 -out private.pem`,
`openssl pkey -in private.pem -pubout -out public.pem`; private key goes `openssl pkey -in private.pem -pubout -out public.pem`; private key goes
to the `RELEASE_SIGNING_KEY` repo secret, public key is committed). to the `RELEASE_SIGNING_KEY` repo secret, public key is committed).
`gpu-turnstile --force-update` checks immediately, stages the new binary
and restarts the running service (elevating via UAC only if needed).
### Docker ### Docker
+190 -35
View File
@@ -10,11 +10,12 @@ becomes very slow; on Linux it would OOM instead.
## Goal ## Goal
A single Go binary that sits in front of **both** services and guarantees that A single Go binary that sits in front of **both** services and guarantees that
at any moment the GPU is in exactly one of three states: at any moment the GPU is in exactly one of four states:
- `idle` — nothing in flight - `idle` — nothing in flight
- `llm` — N ≥ 1 Ollama inference requests in flight (concurrency allowed) - `llm` — N ≥ 1 Ollama inference requests in flight (concurrency allowed)
- `image` — exactly one ComfyUI job in flight, Ollama models unloaded - `image` — exactly one ComfyUI job in flight, Ollama models unloaded
- `external` — a foreign process (e.g. a game) holds the GPU; new work waits
Clients (LiteLLM, Open WebUI, n8n) point at gpu-turnstile instead of at the Clients (LiteLLM, Open WebUI, n8n) point at gpu-turnstile instead of at the
services. gpu-turnstile is transparent for everything that does not touch the services. gpu-turnstile is transparent for everything that does not touch the
@@ -53,9 +54,10 @@ consumer gets no listener, no startup probe, and no lock participation:
blocks since no image jobs can arrive. blocks since no image jobs can arrive.
- **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.
- Future consumers (e.g. detecting a local game holding VRAM) plug into the - **Game detection** is a third, optional consumer without a URL: enabled by
same lock the same way: enabled by their config knob, excluded when `GAME_PROCS`, `GPU_FOREIGN_VRAM_MB` and/or `GPU_FOREIGN_UTIL_PCT` it
absent. watches for foreign processes holding the GPU (see below) and plugs into
the same lock the same way — excluded when all knobs are unset.
### Lock semantics ### Lock semantics
@@ -75,6 +77,11 @@ are LLM requests and the single "writer" is an image job):
requests start), waits until n == 0, sets state := `image`. Released after requests start), waits until n == 0, sets state := `image`. Released after
the ComfyUI job finished and models were freed. the ComfyUI job finished and models were freed.
- Concurrent image jobs queue FIFO behind each other. - Concurrent image jobs queue FIFO behind each other.
- **External hold**: game detection calls `SetExternal(holder)` while a
foreign process holds the GPU. New LLM and image grants block (same busy
handling as above) until `ClearExternal()`; in-flight work is not
preempted, it drains. Snapshot reports state `external` once nothing else
is in flight.
- All waits are context-aware: a client that disconnects while waiting is - All waits are context-aware: a client that disconnects while waiting is
removed from the queue. removed from the queue.
@@ -121,6 +128,87 @@ state is `idle`, send `POST /api/generate {"model":WARM_MODEL,"keep_alive":-1}`
with empty prompt to reload the chat model so the next chat doesn't pay the with empty prompt to reload the chat model so the next chat doesn't pay the
load time. Off by default. load time. Off by default.
### Managed ComfyUI (`COMFY_CMD` / `COMFY_DIR`)
When `COMFY_CMD` is set, gpu-turnstile runs ComfyUI as a supervised child
process instead of expecting an always-on server. Setting only `COMFY_DIR`
enables the same management with the launch command derived from the
standard venv layout under it (`.venv\Scripts\python.exe` on Windows,
`.venv/bin/python` on Linux; `ComfyUI\main.py`, or a flat `main.py` when
that is what exists; `--port` from the `COMFY_URL` port). Missing layout
files are flagged in the startup log.
- **Start on demand**: any ComfyUI request spawns it (double quotes in the
command line group arguments with spaces; `COMFY_DIR` sets the working
directory). `POST /prompt` additionally waits for readiness
(`/system_stats`) for up to `COMFY_START_TIMEOUT` *before* taking the GPU
lock, so LLM traffic flows while torch loads. Other requests are bridged
by the normal retry backoff. Spawn failure or a readiness timeout
answers 502.
- **Idle stop**: after `COMFY_IDLE_TIMEOUT` without requests or finished
jobs — and only while the GPU lock is idle — the process tree is killed,
freeing the VRAM ComfyUI holds. The next request restarts it.
- **Crash**: an unexpected exit is logged; the next request respawns.
gpu-turnstile's own shutdown stops the child too.
- **Coexistence**: before spawning, the URL is probed — if another server
already answers (e.g. the ComfyUI desktop app), it is used as-is and no
child is spawned; the idle watcher and shutdown only ever stop the
supervisor's own process, never the external one. The other direction —
starting the desktop app while the managed instance holds the port —
makes the *desktop* server fail to bind; gpu-turnstile is unaffected.
- Its stdout/stderr is forwarded to the log at INFO. The health check
skips the intentionally-stopped/starting states; a failed probe while
the process is alive and was previously ready is logged as DOWN.
- **Permissions**: the service account is sandboxed (Windows virtual
account, systemd `DynamicUser`), so a ComfyUI install inside a user
profile is off-limits by default. `--install-service` opens it up —
a recursive ACL grant for `NT SERVICE\gpu-turnstile` on Windows, a
`BindPaths=` in the unit on Linux — reading `COMFY_DIR` from the config
it installs. Re-run `--install-service` after changing `COMFY_DIR`, or
grant by hand (admin shell):
`icacls "<COMFY_DIR>" /grant "NT SERVICE\gpu-turnstile:(OI)(CI)M" /T`.
## Game detection (foreign GPU holders)
Games and other foreign GPU users sit outside the URL-based consumer model —
nothing proxies through gpu-turnstile for them. Three independent detection
paths, polled every `GAME_POLL_INTERVAL` (default 15 s); any one being
configured enables the feature:
- **Process watch list** (`GAME_PROCS`, comma-separated, case-insensitive,
`.exe` optional): while any listed process runs, the GPU counts as held.
- **Foreign 3D-engine utilization** (`GPU_FOREIGN_UTIL_PCT`, Windows only):
per-process GPU engine counters from PDH — the same data Task Manager
shows — cover graphics work under WDDM, so any process not in
`GPU_IGNORE_PROCS` using more than the threshold percent of the 3D engine
counts as a foreign holder, with no exe list needed. Values are summed per
process across engines; a missing/failing PDH counter disables the path
(logged once).
- **Foreign VRAM threshold** (`GPU_FOREIGN_VRAM_MB`): `nvidia-smi
--query-compute-apps` lists per-process GPU memory; any process not in
`GPU_IGNORE_PROCS` (default: Ollama and python — ComfyUI runs under python
— plus dwm, the desktop compositor) holding more than the threshold counts
as a foreign holder. Needs nvidia-smi on the PATH (absent: logged once,
path disabled) and works best on Linux — under Windows' WDDM driver,
graphics-only games may not appear in the per-process list.
While a holder is detected, gpu-turnstile:
1. takes the external hold on the lock (`SetExternal`), so new LLM/image
requests wait or are rejected per `LLM_BUSY_MODE`; busy responses name the
holder,
2. once in-flight work has drained, **frees VRAM for the foreign process**:
the managed ComfyUI is stopped (an external server on its port is never
touched) and Ollama's resident models are unloaded,
3. refuses to spawn the managed ComfyUI; ComfyUI requests that would need a
spawn are answered 503 + `Retry-After` (a running desktop instance keeps
being proxied),
4. logs the transitions at WARN ("GPU held by an external process …" /
"released the GPU; resuming") and reports state `external` in `/healthz`
and `/metrics`.
When the holder disappears, the hold is lifted and queued requests proceed.
## Configuration (env) ## Configuration (env)
Configuration comes from environment variables and/or an `.env`-style Configuration comes from environment variables and/or an `.env`-style
@@ -131,8 +219,8 @@ override file values. A missing file is fine; a malformed one is fatal.
| Var | Default | Meaning | | Var | Default | Meaning |
|---|---|---| |---|---|---|
| `LISTEN_OLLAMA` | `:11434` | Ollama-facing listener | | `LISTEN_OLLAMA` | `:11434` | listener for Ollama-compatible clients |
| `LISTEN_COMFY` | `:8188` | ComfyUI-facing listener | | `LISTEN_COMFY` | `:8188` | listener for ComfyUI clients |
| `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 |
@@ -142,12 +230,22 @@ override file values. A missing file is fine; a malformed one is fatal.
| `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400599, e.g. 429) | | `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400599, e.g. 429) |
| `BUSY_RETRY_AFTER` | `30` | seconds sent as `Retry-After` on busy responses (both modes) | | `BUSY_RETRY_AFTER` | `30` | seconds sent as `Retry-After` on busy responses (both modes) |
| `WARM_MODEL` | `` | optional model to reload after an image job | | `WARM_MODEL` | `` | optional model to reload after an image job |
| `COMFY_CMD` | _(derived from `COMFY_DIR`; both empty = unmanaged)_ | spawn and supervise ComfyUI on demand: first request starts it, idle stop after `COMFY_IDLE_TIMEOUT` frees its VRAM. Requires `COMFY_URL` |
| `COMFY_DIR` | `` | standard venv install root: set alone to supervise ComfyUI with the derived launch command (`.venv` + `main.py`, `--port` from `COMFY_URL`); also the working directory for `COMFY_CMD` |
| `COMFY_IDLE_TIMEOUT` | `5m` | stop the managed ComfyUI after this long without requests or jobs |
| `COMFY_START_TIMEOUT` | `2m` | how long a request waits for the managed ComfyUI to come up |
| `GAME_PROCS` | _(empty = disabled)_ | comma-separated process names (case-insensitive, `.exe` optional); while any runs, the GPU counts as held by it: requests wait, Ollama unloads, the managed ComfyUI stops |
| `GPU_FOREIGN_VRAM_MB` | `0` (disabled) | also treat the GPU as held when a process not in `GPU_IGNORE_PROCS` uses more VRAM than this; needs nvidia-smi |
| `GPU_FOREIGN_UTIL_PCT` | `0` (disabled) | also treat the GPU as held when a process not in `GPU_IGNORE_PROCS` uses more than this percent of the GPU 3D engine (Windows PDH counters, as shown by Task Manager; catches games without an exe list) |
| `GPU_IGNORE_PROCS` | `ollama,ollama app,ollama_llama_server,python,pythonw,dwm` | process names never counted as foreign GPU users |
| `GAME_POLL_INTERVAL` | `15s` | how often game/VRAM detection runs (nvidia-smi polls keep the GPU awake; don't go below ~10s) |
| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` is accepted as an alias | | `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 |
| `LOG_FILE` | `` | append logs to this file instead of stderr (useful as a service) | | `LOG_FILE` | `` | append logs to this file instead of stderr (useful as a service) |
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading | | `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs | | `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
| `PROBE_TIMEOUT` | `5s` | startup probe of both upstreams | | `PROBE_TIMEOUT` | `5s` | startup probe of both upstreams (also the per-probe health check timeout) |
| `HEALTH_INTERVAL` | `30s` | periodic probe of enabled upstreams; status changes (down/recovered) are logged |
| `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job | | `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job |
| `WARM_TIMEOUT` | `2m` | warm-model reload after an image job | | `WARM_TIMEOUT` | `2m` | warm-model reload after an image job |
| `SHUTDOWN_TIMEOUT` | `10s` | graceful shutdown on SIGINT/SIGTERM | | `SHUTDOWN_TIMEOUT` | `10s` | graceful shutdown on SIGINT/SIGTERM |
@@ -158,6 +256,8 @@ override file values. A missing file is fine; a malformed one is fatal.
| `UPDATE_INTERVAL` | `6h` | auto-update check interval | | `UPDATE_INTERVAL` | `6h` | auto-update check interval |
| `UPDATE_REPO` | `https://git.rambossek.at/PUBLIC/gpu-turnstile` | repository to check for releases | | `UPDATE_REPO` | `https://git.rambossek.at/PUBLIC/gpu-turnstile` | repository to check for releases |
| `UPDATE_ASSET` | `gpu-turnstile.exe` | release asset to download | | `UPDATE_ASSET` | `gpu-turnstile.exe` | release asset to download |
| `APP_VER` | `stable` | version to run: `dev` disables updates, `stable` tracks the latest release, or an exact `vX.Y.Z` pin (up- or downgraded to) |
| `CFG_VER` | _(installer-managed)_ | config format reference written by `--install-service` (always a concrete `vX.Y.Z`; a dev build stamps `v0.0.0`); missing = the file is replaced with a fresh sample (backup `.bak`). New settings are appended (commented out) at install and at every startup after an update changed the version |
Startup fails fast on unparsable values and when neither consumer URL is Startup fails fast on unparsable values and when neither consumer URL is
set. Enabled upstreams are probed once at start (`/api/version`, set. Enabled upstreams are probed once at start (`/api/version`,
@@ -169,35 +269,78 @@ The binary runs natively on Windows (the current primary deployment) and on
Linux with systemd (the future GPU server), as well as in Docker. Linux with systemd (the future GPU server), as well as in Docker.
Service management is the same on both platforms: Service management is the same on both platforms:
`gpu-turnstile --install-service [-config path]` registers and starts an `gpu-turnstile --install-service [-config path]` installs, registers and
auto-start service; `--remove-service` stops and unregisters it (both need starts an auto-start service; `--remove-service` stops and uninstalls it.
an elevated/root shell). The legacy form `gpu-turnstile service Both need admin/root; on Windows a non-elevated shell triggers a UAC
install|remove` does the same thing. prompt instead of failing — the command relaunches itself elevated, waits
for the child, and mirrors its exit code. The legacy form
`gpu-turnstile service install|remove` does the same thing.
Re-running install on an already-registered service converges instead of
failing: a running service is stopped first, the installed binary copy is
refreshed only when the content differs, the registration (Windows service
config / systemd unit) is updated only where it drifted, and the service is
started again only if it was running before.
By default install creates the canonical layout and copies the binary into
it (Windows: `%ProgramFiles%\gpu-turnstile\`, plus
`%ProgramData%\gpu-turnstile\` for logs; Linux: `/var/lib/gpu-turnstile/`
with the config at `/etc/gpu-turnstile.env`). If there is no config at all,
install writes a sample env file covering every setting — each with a
comment line, everything commented out — except the `CFG_VER`/`APP_VER`
header and `LOG_FILE`, which is active on Windows
(`%ProgramData%\gpu-turnstile\gpu-turnstile.log`) since a service has no
console; on Linux it stays commented because stderr goes to the journal.
The first line is `CFG_VER=vX.Y.Z`, recording the installer version. When a
later version's install finds an older `CFG_VER`, it appends every setting
the file does not mention (commented or not) at the end and updates
`CFG_VER`; a file without `CFG_VER` is invalid and gets replaced by a fresh
sample, with the old content kept as `<file>.bak`. `--no-copy` registers
the current executable location as-is and leaves the config untouched.
### Windows ### Windows
- `--install-service` registers a Windows service; recovery actions restart - `--install-service` creates `%ProgramFiles%\gpu-turnstile\` and
it 5 s after any failure. `%ProgramData%\gpu-turnstile\`, copies the exe and (if none exists there
- **Layout**: install to `C:\Program Files\gpu-turnstile\` (exe plus yet) the `gpu-turnstile.env` into the Program Files directory, and
`gpu-turnstile.env`); logs belong in `C:\ProgramData\gpu-turnstile\` via registers that copy as a Windows service; recovery actions restart it
`LOG_FILE`. The service must be able to write its install directory for 5 s after any failure. The install ensures the env file sets `LOG_FILE`
self-updates — Program Files is writable by LocalSystem and admins, which to `%ProgramData%\gpu-turnstile\gpu-turnstile.log` since there is no
is why running as the default `LocalSystem` account is the simple choice. console — an existing `LOG_FILE` setting is kept.
- **Account**: the default `LocalSystem` works out of the box. For least - **Account**: the service always runs as the virtual account
privilege, create the service with the virtual account `NT SERVICE\gpu-turnstile` — a per-service low-privilege identity the
`NT SERVICE\gpu-turnstile` and grant it write access to the install and SCM manages (no password, automatic logon-as-a-service right, no admin
log directories only (no network logon, no user profile). rights, gone when the service is removed). The installer grants it
- Use a config file (above) for the service — Windows services have no modify access to the install and data directories (self-updates rewrite
convenient environment. Logs go to `LOG_FILE` since there is no console. the exe) and the `LOG_FILE` directory (created if missing), plus read
access to the config file when it lives elsewhere. The grants happen
after service registration because the virtual account's SID only exists
from that point on; if a grant fails the service registration is rolled
back.
### Linux (systemd) ### Linux (systemd)
- `--install-service` writes `/etc/systemd/system/gpu-turnstile.service` - `--install-service` copies the binary to `/var/lib/gpu-turnstile/`,
with `ExecStart` pointing at the current executable and the `-config` copies the config to `/etc/gpu-turnstile.env` if none exists there yet,
file, then runs `systemctl daemon-reload` and `enable --now`. The unit writes `/etc/systemd/system/gpu-turnstile.service`, then runs `systemctl
runs as root (it must be able to overwrite its own binary for daemon-reload` and `enable --now`. `--remove-service` removes the unit
self-updates); harden with `ProtectSystem=strict` plus a writable and the installed binary; the `/etc` config stays. The binary does not
`ReadWritePaths` if desired. go to `/usr/local/sbin` on purpose: replacing a running binary needs
write access to its *directory*, and granting the sandboxed service
write access to a shared system directory would let a compromised
service overwrite other binaries — `/var/lib/gpu-turnstile` is
exclusively ours.
- **Sandboxing** mirrors the Windows virtual account: the unit runs with
`DynamicUser=yes` — a transient per-service UID with no login, no home
and no password, managed entirely by systemd. `ProtectSystem=strict`
makes the filesystem read-only except `StateDirectory=gpu-turnstile`
(the install dir, so self-updates can rewrite the binary), plus
`NoNewPrivileges`, `ProtectHome`, `PrivateTmp`, `ProtectKernel*`,
`ProtectControlGroups`, `RestrictNamespaces`, `RestrictSUIDSGID`,
`RestrictRealtime`, `LockPersonality`, `MemoryDenyWriteExecute`, empty
capability sets, `RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6` and
`SystemCallFilter=@system-service`. The proxy needs only outbound
TCP/UDP and the notify socket, so it loses nothing.
- The unit is `Type=notify`: the binary sends `READY=1` via - The unit is `Type=notify`: the binary sends `READY=1` via
`github.com/coreos/go-systemd` only after the listeners are bound, so `github.com/coreos/go-systemd` only after the listeners are bound, so
`systemctl start` blocks until the proxy accepts connections. A 30 s `systemctl start` blocks until the proxy accepts connections. A 30 s
@@ -211,14 +354,25 @@ install|remove` does the same thing.
`RestartSec=5s` brings up the staged binary after the updater exits with `RestartSec=5s` brings up the staged binary after the updater exits with
code 3. code 3.
- **Auto-update**: on startup and every `UPDATE_INTERVAL`, the binary - **Auto-update**: on startup and every `UPDATE_INTERVAL`, the binary
checks `UPDATE_REPO`'s latest release; if its tag is a newer `vX.Y.Z`, consults `APP_VER`: `dev` disables updates; `stable` (the default)
it downloads `UPDATE_ASSET` plus its `.sig` (and `.sha256` when present) fetches `UPDATE_REPO`'s latest release and applies it when its tag is a
and verifies an Ed25519 signature against the public key embedded in newer `vX.Y.Z` (a `dev` binary cannot be compared and is replaced by the
latest release); a `vX.Y.Z` pin fetches that exact tag and stages it on
any difference, including downgrades. Applying means downloading
`UPDATE_ASSET` plus its `.sig` (and `.sha256` when present) and verifying
an Ed25519 signature against the public key embedded in
`internal/update/pubkey.go`. A verified binary is swapped in next to the `internal/update/pubkey.go`. A verified binary is swapped in next to the
running exe (rename-aside, allowed on Windows), and once the GPU lock is running exe (rename-aside, allowed on Windows), and once the GPU lock is
idle the process exits with code 3 so the service recovery restarts it idle the process exits with code 3 so the service recovery restarts it
on the new version. Interactive runs only log "restart to apply". on the new version. Interactive runs only log "restart to apply".
`dev` builds and builds without an embedded public key never update. Builds without an embedded public key never update.
- **`--force-update`** runs the same check immediately, single-shot: one
attempt with a 30 s timeout, then exit — "up to date" (exit 0) or the
error (exit 1), no retries. When a newer release is found it downloads,
verifies and stages it, and if the service is running it restarts it
right away (otherwise the new version applies on next start). On Windows
it elevates via UAC only when the stage or restart needs permissions the
caller does not have.
- **Signing setup (one time)**: `openssl genpkey -algorithm ed25519 -out - **Signing setup (one time)**: `openssl genpkey -algorithm ed25519 -out
private.pem`; `openssl pkey -in private.pem -pubout -out public.pem`. private.pem`; `openssl pkey -in private.pem -pubout -out public.pem`.
Private key → repo secret `RELEASE_SIGNING_KEY`; public key → committed into Private key → repo secret `RELEASE_SIGNING_KEY`; public key → committed into
@@ -302,7 +456,8 @@ are new.
receives the first chunk before the last is sent (no buffering). receives the first chunk before the last is sent (no buffering).
- `internal/config`: env-file parsing, precedence, fail-fast values. - `internal/config`: env-file parsing, precedence, fail-fast values.
- `internal/update`: fake Gitea releases API; staged update happy path, - `internal/update`: fake Gitea releases API; staged update happy path,
tampered signature rejected, older versions and dev builds skipped. tampered signature rejected, older versions skipped, APP_VER=dev and
pinned releases honored.
## Build and CI ## Build and CI
+962 -89
View File
File diff suppressed because it is too large Load Diff
+342
View File
@@ -0,0 +1,342 @@
package main
import (
"encoding/json"
"fmt"
"os"
"os/exec"
"strings"
"time"
"gpu-turnstile/internal/control"
)
// ANSI colors for the monitor frame.
const (
cReset = "\x1b[0m"
cDim = "\x1b[2m"
cRed = "\x1b[31m"
cGreen = "\x1b[32m"
cYellow = "\x1b[33m"
cCyan = "\x1b[36m"
)
// hotkeysLine is the monitor's footer.
const hotkeysLine = " " + cDim + "q quit · u update now · r reload config" + cReset + "\x1b[K\n"
// monitorCommand renders a live status view of the running service,
// refreshed every second from the control channel. When the service
// reports a different version and the executable on disk changed (the
// updater replaced it), the monitor restarts itself onto the new binary.
// Hotkeys: q quits, u triggers an update check on the service, r asks the
// service to reload its config file.
func monitorCommand() int {
if !stdoutIsTerminal() {
fmt.Fprintln(os.Stderr, "gpu-turnstile: --monitor needs an interactive terminal")
return 1
}
enableVirtualTerminal()
restore := enableRawKeys()
defer func() {
if restore != nil {
restore()
}
}()
fmt.Print("\x1b[2J") // clear once; frames then redraw in place
defer fmt.Print(cReset + "\n")
exe, _ := os.Executable()
var exeStamp time.Time
if st, err := os.Stat(exe); err == nil {
exeStamp = st.ModTime()
}
keys := make(chan byte, 8)
go readKeys(keys)
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
var note string
var noteAt time.Time
noteCh := make(chan string, 1)
askPending := false
// ask sends a one-shot command to the service and reports the reply in
// the note line. Only one ask runs at a time.
ask := func(cmd, busy, label string) {
if askPending {
return
}
askPending = true
note, noteAt = busy, time.Now()
go func() {
reply, err := control.Ask(cmd)
if err != nil {
noteCh <- label + ": no answer from the service"
return
}
msg := strings.TrimPrefix(reply, "OK ")
msg = strings.TrimPrefix(msg, "ERR ")
noteCh <- label + ": " + msg
}()
}
poll := func() string {
frame := renderWaiting()
if reply, err := control.Ask(control.CmdStatus); err == nil {
if msg, ok := strings.CutPrefix(reply, "OK "); ok {
var snap statusSnapshot
if json.Unmarshal([]byte(msg), &snap) == nil {
if snap.Version != "" && snap.Version != version {
if exeChanged(exe, exeStamp) {
fmt.Print("\x1b[2J\x1b[H")
fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version)
restartSelf(exe, "--monitor")
return "" // re-execed; this process exits below
}
snap.MonitorNote = fmt.Sprintf("note: the service runs %s, this monitor is %s", snap.Version, version)
}
if note != "" {
snap.MonitorNote = note
}
frame = renderMonitor(snap, termWidth())
}
}
}
return frame
}
for {
frame := poll()
if frame == "" {
return 0 // restartSelf fired
}
fmt.Print("\x1b[H" + frame + "\x1b[J") // home, frame, clear below
select {
case <-ticker.C:
if note != "" && time.Since(noteAt) > 15*time.Second {
note = ""
}
case k, ok := <-keys:
if !ok {
keys = nil
continue
}
switch k {
case 'q', 'Q', 3: // q or Ctrl+C (raw mode delivers it as a byte)
return 0
case 'u', 'U':
ask(control.CmdUpdateNow, "checking for updates…", "update")
case 'r', 'R':
ask(control.CmdReloadEnv, "reloading config…", "reload")
}
case n := <-noteCh:
askPending = false
note, noteAt = n, time.Now()
}
}
}
// readKeys reads single keypresses from stdin (raw mode was enabled by the
// caller) and delivers them until stdin fails.
func readKeys(keys chan<- byte) {
defer close(keys)
buf := make([]byte, 1)
for {
n, err := os.Stdin.Read(buf)
if n > 0 {
keys <- buf[0]
}
if err != nil {
return
}
}
}
// exeChanged reports whether the executable on disk was replaced since the
// recorded stamp (the updater swaps it via rename, which changes ModTime).
func exeChanged(exe string, stamp time.Time) bool {
if exe == "" || stamp.IsZero() {
return false
}
st, err := os.Stat(exe)
return err == nil && !st.ModTime().Equal(stamp)
}
// restartSelf starts a fresh copy of this executable with the given args on
// the same console; the caller exits right after.
func restartSelf(exe string, args ...string) {
cmd := exec.Command(exe, args...)
cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
cmd.Start() //nolint:errcheck // best effort: on failure we just exit
}
func renderWaiting() string {
return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n\x1b[K\n" + hotkeysLine
}
// renderMonitor draws one full frame. Each line ends with \x1b[K (clear to
// end of line) so shrinking content leaves no residue.
func renderMonitor(snap statusSnapshot, width int) string {
if width < 40 {
width = 80
}
var b strings.Builder
left := " gpu-turnstile"
if snap.UptimeS > 0 {
left += " " + cDim + "up " + fmtDur(snap.UptimeS) + cReset
}
right := snap.Version
pad := width - printableLen(" gpu-turnstile up "+fmtDur(snap.UptimeS)) - len(right) - 1
if snap.UptimeS == 0 {
pad = width - len(" gpu-turnstile") - len(right) - 1
}
if pad < 1 {
pad = 1
}
b.WriteString(cDim + left + strings.Repeat(" ", pad) + right + cReset + "\x1b[K\n")
b.WriteString(cDim + " " + strings.Repeat("─", width-2) + cReset + "\x1b[K\n")
for _, d := range snap.Downstreams {
busy := (d.Name == "ollama" && snap.Lock.State == "llm" && snap.Lock.LLMInflight > 0) ||
(d.Name == "comfy" && snap.Lock.State == "image")
b.WriteString(renderDownstream(d, busy) + "\x1b[K\n")
}
b.WriteString("\x1b[K\n")
b.WriteString(renderLock(snap.Lock) + "\x1b[K\n")
if snap.Lock.ImageQueue > 0 {
b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n",
cYellow, snap.Lock.ImageQueue, cReset))
}
if snap.GPU.Enabled {
b.WriteString(renderGPU(snap.GPU) + "\x1b[K\n")
}
if snap.MonitorNote != "" {
b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
}
b.WriteString("\x1b[K\n")
b.WriteString(hotkeysLine)
return b.String()
}
// renderGPU renders the GPU line: VRAM usage (when nvidia-smi answered),
// the game detector's last finding, and how long ago it ran.
func renderGPU(g statusGPU) string {
s := " GPU: "
if g.Known {
s += renderVRAM(g.UsedMB, g.TotalMB)
if g.TempC >= 0 {
s += fmt.Sprintf(" · %d°C", g.TempC)
}
if g.FanPct >= 0 {
s += fmt.Sprintf(" · fan %d%%", g.FanPct)
}
} else {
s += cDim + "VRAM unknown (nvidia-smi not answering)" + cReset
}
if g.Foreign != "" {
s += " · external: " + cRed + g.Foreign + cReset
} else {
s += cDim + " · no external process" + cReset
}
if g.AgeS >= 0 {
s += cDim + fmt.Sprintf(" (checked %s ago)", fmtDur(g.AgeS)) + cReset
}
return s
}
// formatMB renders MiB as "4.2 GiB" (or MiB below 1 GiB).
func formatMB(mb int64) string {
if mb >= 1024 {
return fmt.Sprintf("%.1f GiB", float64(mb)/1024)
}
return fmt.Sprintf("%d MiB", mb)
}
// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
func renderVRAM(used, total int) string {
if total > 0 {
return formatMB(int64(used)) + " / " + formatMB(int64(total)) + " used"
}
return formatMB(int64(used)) + " used"
}
// printableLen counts characters without ANSI escapes (ASCII-only content).
func printableLen(s string) int { return len(s) }
func renderDownstream(d statusDownstream, busy bool) string {
url := cDim + d.URL + cReset
detail := ""
if d.Models != nil {
if len(d.Models) == 0 {
detail = cDim + " · no models loaded" + cReset
} else {
parts := make([]string, 0, len(d.Models))
for _, m := range d.Models {
if m.VRAMMB > 0 {
parts = append(parts, fmt.Sprintf("%s (%s VRAM)", m.Name, formatMB(m.VRAMMB)))
} else {
parts = append(parts, m.Name+" (in RAM)")
}
}
detail = " · " + strings.Join(parts, ", ")
}
}
if busy {
detail += " · " + cCyan + "busy" + cReset
}
switch d.Managed {
case "stopped":
return fmt.Sprintf(" %s○%s %-8s %sstopped (managed — starts on demand)%s %s",
cDim, cReset, d.Name, cDim, cReset, url)
case "starting":
return fmt.Sprintf(" %s◌%s %-8s %sstarting…%s %s",
cYellow, cReset, d.Name, cYellow, cReset, url)
}
suffix := ""
if d.Managed == "external" {
suffix = " (external)"
}
if d.Up {
return fmt.Sprintf(" %s●%s %-8s %sUP%s%s%s %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, detail, url)
}
return fmt.Sprintf(" %s●%s %-8s %sDOWN%s %s", cRed, cReset, d.Name, cRed, cReset, url)
}
func renderLock(l statusLock) string {
dur := cDim + "(" + fmtDur(l.SinceS) + ")" + cReset
switch l.State {
case "idle":
return fmt.Sprintf(" Lock: %sidle%s %s", cGreen, cReset, dur)
case "llm":
s := fmt.Sprintf(" Lock: %sLLM%s — %d in flight", cCyan, cReset, l.LLMInflight)
if l.LLMWaiting > 0 {
s += fmt.Sprintf(", %d waiting", l.LLMWaiting)
}
if l.Detail != "" {
s += " — " + l.Detail
}
return s + " " + dur
case "image":
s := fmt.Sprintf(" Lock: %sIMAGE%s", cYellow, cReset)
if l.Detail != "" {
s += " — " + l.Detail
}
return s + " " + dur
case "external":
return fmt.Sprintf(" Lock: %sEXTERNAL%s — %s %s", cRed, cReset, l.External, dur)
}
return " Lock: unknown"
}
// fmtDur renders seconds as a compact duration ("1m32s", "2h07m").
func fmtDur(s int64) string {
if s < 0 {
s = 0
}
d := time.Duration(s) * time.Second
if d >= time.Hour {
return fmt.Sprintf("%dh%02dm", int(d.Hours()), int(d.Minutes())%60)
}
return d.Round(time.Second).String()
}
+104
View File
@@ -0,0 +1,104 @@
package main
import (
"gpu-turnstile/internal/config"
"reflect"
"strings"
"testing"
)
func TestRenderMonitor(t *testing.T) {
snap := statusSnapshot{
Version: "v0.2.3",
UptimeS: 3725,
Downstreams: []statusDownstream{
{Name: "ollama", URL: "http://127.0.0.1:11435", Up: true},
{Name: "comfy", URL: "http://127.0.0.1:8189", Managed: "stopped"},
},
Lock: statusLock{State: "llm", LLMInflight: 2, LLMWaiting: 1, Detail: "ollama: POST /api/generate", SinceS: 95},
}
frame := renderMonitor(snap, 80)
for _, want := range []string{"v0.2.3", "1h02m", "ollama", "UP", "comfy", "stopped", "LLM", "2 in flight", "1 waiting", "1m35s"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.Lock = statusLock{State: "external", External: "cyberpunk2077.exe (pid 1234)", SinceS: 3, ImageQueue: 2}
frame = renderMonitor(snap, 40) // narrow: falls back to 80
for _, want := range []string{"EXTERNAL", "cyberpunk2077.exe", "2 image job(s) waiting"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.GPU = statusGPU{Enabled: true, Known: true, UsedMB: 4300, TotalMB: 16384, TempC: 55, FanPct: 42, Foreign: "cyberpunk2077.exe (pid 1234)", AgeS: 12}
frame = renderMonitor(snap, 80)
for _, want := range []string{"GPU:", "4.2 GiB / 16.0 GiB used", "55°C", "fan 42%", "external:", "checked 12s ago"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.GPU = statusGPU{Enabled: true, AgeS: 4}
frame = renderMonitor(snap, 80)
for _, want := range []string{"VRAM unknown", "no external process"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.Downstreams[0].Models = []statusModel{{Name: "llama3.1:8b", VRAMMB: 4900}, {Name: "embed", VRAMMB: 0}}
snap.Lock = statusLock{State: "llm", LLMInflight: 1}
frame = renderMonitor(snap, 80)
for _, want := range []string{"llama3.1:8b (4.8 GiB VRAM)", "embed (in RAM)", "busy"} {
if !strings.Contains(frame, want) {
t.Errorf("frame missing %q:\n%s", want, frame)
}
}
snap.Downstreams[0].Models = []statusModel{}
snap.Lock = statusLock{State: "idle"}
frame = renderMonitor(snap, 80)
if !strings.Contains(frame, "no models loaded") {
t.Errorf("frame missing %q:\n%s", "no models loaded", frame)
}
if strings.Contains(frame, "busy") {
t.Errorf("idle lock still shows busy:\n%s", frame)
}
}
func TestFmtDur(t *testing.T) {
cases := map[int64]string{0: "0s", 5: "5s", 95: "1m35s", 3725: "1h02m", -3: "0s"}
for in, want := range cases {
if got := fmtDur(in); got != want {
t.Errorf("fmtDur(%d) = %q, want %q", in, got, want)
}
}
}
func TestDiffConfig(t *testing.T) {
a := config.Defaults()
b := a
if got := diffConfig(a, b); len(got) != 0 {
t.Fatalf("identical configs: got %v", got)
}
b.LogLevel = -4
b.GameProcs = []string{"game.exe"}
got := diffConfig(a, b)
if len(got) != 2 || got[0] != "GAME_PROCS" || got[1] != "LOGLEVEL" {
t.Fatalf("got %v, want [GAME_PROCS LOGLEVEL]", got)
}
}
// Every Config field must carry an env tag so user-facing output (the
// reload diff) can name the setting the user would actually change.
func TestConfigFieldsHaveEnvTags(t *testing.T) {
typ := reflect.TypeOf(config.Config{})
for i := 0; i < typ.NumField(); i++ {
f := typ.Field(i)
if f.Tag.Get("env") == "" {
t.Errorf("config.Config.%s has no env tag", f.Name)
}
}
}
+39
View File
@@ -0,0 +1,39 @@
//go:build !windows
package main
import (
"os"
"golang.org/x/sys/unix"
)
// enableVirtualTerminal is a no-op: unix terminals speak ANSI natively.
func enableVirtualTerminal() {}
// termWidth is the terminal width in columns, 80 when unknown.
func termWidth() int {
ws, err := unix.IoctlGetWinsize(int(os.Stdout.Fd()), unix.TIOCGWINSZ)
if err != nil || ws.Col == 0 {
return 80
}
return int(ws.Col)
}
// enableRawKeys switches the terminal to per-keypress mode (ICANON and ECHO
// off) and returns the restore function, nil when stdin is not a terminal.
func enableRawKeys() func() {
fd := int(os.Stdin.Fd())
term, err := unix.IoctlGetTermios(fd, unix.TCGETS)
if err != nil {
return nil
}
raw := *term
raw.Lflag &^= unix.ICANON | unix.ECHO
raw.Cc[unix.VMIN] = 1
raw.Cc[unix.VTIME] = 0
if err := unix.IoctlSetTermios(fd, unix.TCSETS, &raw); err != nil {
return nil
}
return func() { unix.IoctlSetTermios(fd, unix.TCSETS, term) } //nolint:errcheck
}
+49
View File
@@ -0,0 +1,49 @@
//go:build windows
package main
import (
"os"
"golang.org/x/sys/windows"
)
// enableVirtualTerminal asks the console to honor ANSI escapes (Windows 10+);
// mintty/Git Bash already does, so errors are ignored.
func enableVirtualTerminal() {
h := windows.Handle(os.Stdout.Fd())
var mode uint32
if err := windows.GetConsoleMode(h, &mode); err != nil {
return
}
windows.SetConsoleMode(h, mode|windows.ENABLE_VIRTUAL_TERMINAL_PROCESSING) //nolint:errcheck
}
// termWidth is the console window width in columns, 80 when unknown.
func termWidth() int {
var info windows.ConsoleScreenBufferInfo
if err := windows.GetConsoleScreenBufferInfo(windows.Handle(os.Stdout.Fd()), &info); err != nil {
return 80
}
return int(info.Window.Right-info.Window.Left) + 1
}
// enableRawKeys puts the console's stdin into per-keypress mode (no line
// buffering, no echo) and returns the restore function. When stdin is not a
// real console (mintty/Git Bash pipes) it returns nil: ptys already deliver
// keystrokes immediately.
func enableRawKeys() func() {
h := windows.Handle(os.Stdin.Fd())
var mode uint32
if err := windows.GetConsoleMode(h, &mode); err != nil {
return nil
}
const (
enableLineInput = 0x0002
enableEchoInput = 0x0004
)
if err := windows.SetConsoleMode(h, mode&^(enableLineInput|enableEchoInput)); err != nil {
return nil
}
return func() { windows.SetConsoleMode(h, mode) } //nolint:errcheck
}
+1 -1
View File
@@ -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.1.5 image: git.rambossek.at/public/gpu-turnstile:v0.3.1
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
+135 -29
View File
@@ -4,6 +4,7 @@ package config
import ( import (
"bufio" "bufio"
"errors"
"fmt" "fmt"
"io" "io"
"log/slog" "log/slog"
@@ -12,42 +13,77 @@ 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 OllamaURL string `env:"OLLAMA_URL"`
ComfyURL string ComfyURL string `env:"COMFY_URL"`
UnloadTimeout time.Duration UnloadTimeout time.Duration `env:"UNLOAD_TIMEOUT"`
JobTimeout time.Duration JobTimeout time.Duration `env:"JOB_TIMEOUT"`
LLMWaitTimeout time.Duration 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"`
FreeTimeout time.Duration HealthInterval time.Duration `env:"HEALTH_INTERVAL"`
WarmTimeout time.Duration FreeTimeout time.Duration `env:"FREE_TIMEOUT"`
ShutdownTimeout time.Duration WarmTimeout time.Duration `env:"WARM_TIMEOUT"`
BackoffInitial time.Duration ShutdownTimeout time.Duration `env:"SHUTDOWN_TIMEOUT"`
BackoffMax time.Duration BackoffInitial time.Duration `env:"BACKOFF_INITIAL"`
PromptCaptureLimit int64 BackoffMax time.Duration `env:"BACKOFF_MAX"`
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
// updates, "stable" tracks the latest release, anything else is an
// exact vX.Y.Z release to pin. From APP_VER; defaults to "stable".
AppVersion string `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"`
WarmModel string // ComfyCmd spawns and supervises a ComfyUI server on demand. When
LogLevel slog.Level // ComfyCmd is empty but ComfyDir is set, management is enabled with the
LogJSON bool // standard venv layout under ComfyDir (.venv + main.py or
LogFile string // ComfyUI/main.py; --port from the COMFY_URL port) — ComfyCmd is the
// override for other layouts and doubles as the working directory when
// set explicitly. The managed server is stopped after ComfyIdleTimeout
// without requests, freeing its VRAM; ComfyStartTimeout bounds how long
// a request waits for it to come up.
ComfyCmd string `env:"COMFY_CMD"`
ComfyDir string `env:"COMFY_DIR"`
ComfyIdleTimeout time.Duration `env:"COMFY_IDLE_TIMEOUT"`
ComfyStartTimeout time.Duration `env:"COMFY_START_TIMEOUT"`
// GameProcs (GAME_PROCS) is a watch list of process names; while any of
// them runs, the GPU is treated as held by a foreign process. The
// nvidia-smi path (GPUForeignVRAMMB, GPU_FOREIGN_VRAM_MB) does the same
// when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than
// that many MiB of VRAM, and the PDH path (GPUForeignUtilPct,
// GPU_FOREIGN_UTIL_PCT) when such a process uses more than that many
// percent of the GPU 3D engine (Windows only). GamePollInterval
// (GAME_POLL_INTERVAL) is how often all checks run.
GameProcs []string `env:"GAME_PROCS"`
GPUForeignVRAMMB int `env:"GPU_FOREIGN_VRAM_MB"`
GPUForeignUtilPct int `env:"GPU_FOREIGN_UTIL_PCT"`
GPUIgnoreProcs []string `env:"GPU_IGNORE_PROCS"`
GamePollInterval time.Duration `env:"GAME_POLL_INTERVAL"`
WarmModel string `env:"WARM_MODEL"`
LogLevel slog.Level `env:"LOGLEVEL"`
LogJSON bool `env:"LOG_FORMAT"`
LogFile string `env:"LOG_FILE"`
} }
// Defaults returns the configuration used when neither the environment nor // Defaults returns the configuration used when neither the environment nor
@@ -64,6 +100,7 @@ func Defaults() Config {
UnloadPollInterval: 500 * time.Millisecond, UnloadPollInterval: 500 * time.Millisecond,
HistoryPollInterval: time.Second, HistoryPollInterval: time.Second,
ProbeTimeout: 5 * time.Second, ProbeTimeout: 5 * time.Second,
HealthInterval: 30 * time.Second,
FreeTimeout: 30 * time.Second, FreeTimeout: 30 * time.Second,
WarmTimeout: 2 * time.Minute, WarmTimeout: 2 * time.Minute,
ShutdownTimeout: 10 * time.Second, ShutdownTimeout: 10 * time.Second,
@@ -75,15 +112,30 @@ func Defaults() Config {
UpdateInterval: 6 * time.Hour, UpdateInterval: 6 * time.Hour,
UpdateRepo: "https://git.rambossek.at/PUBLIC/gpu-turnstile", UpdateRepo: "https://git.rambossek.at/PUBLIC/gpu-turnstile",
UpdateAsset: "gpu-turnstile.exe", UpdateAsset: "gpu-turnstile.exe",
AppVersion: "stable",
LLMBusyMode: "wait", LLMBusyMode: "wait",
LLMBusyStatus: 503, LLMBusyStatus: 503,
BusyRetryAfter: 30, BusyRetryAfter: 30,
ComfyIdleTimeout: 5 * time.Minute,
ComfyStartTimeout: 2 * time.Minute,
// ComfyUI runs under python; excluding it (and Ollama) by name keeps
// our own consumers from tripping the foreign-VRAM check. dwm is the
// desktop compositor — it always shows some 3D-engine usage.
GPUIgnoreProcs: []string{"ollama", "ollama app", "ollama_llama_server", "python", "pythonw", "dwm"},
GamePollInterval: 15 * time.Second,
LogLevel: slog.LevelWarn, LogLevel: slog.LevelWarn,
} }
} }
// ErrNoConsumer is returned by Load when neither OLLAMA_URL nor COMFY_URL
// is set. Commands that never talk to an upstream (--force-update) may
// ignore it and proceed with the rest of the configuration.
var ErrNoConsumer = errors.New("at least one of OLLAMA_URL or COMFY_URL must be set (each URL enables its consumer)")
// ParseEnvFile parses a .env-style file: KEY=VALUE lines, blank lines and // ParseEnvFile parses a .env-style file: KEY=VALUE lines, blank lines and
// #-comments are ignored, no quoting. A line without '=' is an error. // #-comments are ignored, no quoting. A line without '=' is an error.
func ParseEnvFile(r io.Reader) (map[string]string, error) { func ParseEnvFile(r io.Reader) (map[string]string, error) {
@@ -110,6 +162,17 @@ func ParseEnvFile(r io.Reader) (map[string]string, error) {
return values, scanner.Err() return values, scanner.Err()
} }
// splitList parses a comma-separated setting into trimmed, non-empty items.
func splitList(v string) []string {
var out []string
for _, item := range strings.Split(v, ",") {
if item = strings.TrimSpace(item); item != "" {
out = append(out, item)
}
}
return out
}
func envDuration(getenv func(string) string, name string, dst *time.Duration) error { func envDuration(getenv func(string) string, name string, dst *time.Duration) error {
v := getenv(name) v := getenv(name)
if v == "" { if v == "" {
@@ -136,6 +199,8 @@ func Load(getenv func(string) string) (Config, error) {
{"OLLAMA_URL", &cfg.OllamaURL}, {"OLLAMA_URL", &cfg.OllamaURL},
{"COMFY_URL", &cfg.ComfyURL}, {"COMFY_URL", &cfg.ComfyURL},
{"WARM_MODEL", &cfg.WarmModel}, {"WARM_MODEL", &cfg.WarmModel},
{"COMFY_CMD", &cfg.ComfyCmd},
{"COMFY_DIR", &cfg.ComfyDir},
{"UPDATE_REPO", &cfg.UpdateRepo}, {"UPDATE_REPO", &cfg.UpdateRepo},
{"UPDATE_ASSET", &cfg.UpdateAsset}, {"UPDATE_ASSET", &cfg.UpdateAsset},
{"LOG_FILE", &cfg.LogFile}, {"LOG_FILE", &cfg.LogFile},
@@ -154,17 +219,41 @@ func Load(getenv func(string) string) (Config, error) {
{"UNLOAD_POLL_INTERVAL", &cfg.UnloadPollInterval}, {"UNLOAD_POLL_INTERVAL", &cfg.UnloadPollInterval},
{"HISTORY_POLL_INTERVAL", &cfg.HistoryPollInterval}, {"HISTORY_POLL_INTERVAL", &cfg.HistoryPollInterval},
{"PROBE_TIMEOUT", &cfg.ProbeTimeout}, {"PROBE_TIMEOUT", &cfg.ProbeTimeout},
{"HEALTH_INTERVAL", &cfg.HealthInterval},
{"FREE_TIMEOUT", &cfg.FreeTimeout}, {"FREE_TIMEOUT", &cfg.FreeTimeout},
{"WARM_TIMEOUT", &cfg.WarmTimeout}, {"WARM_TIMEOUT", &cfg.WarmTimeout},
{"SHUTDOWN_TIMEOUT", &cfg.ShutdownTimeout}, {"SHUTDOWN_TIMEOUT", &cfg.ShutdownTimeout},
{"BACKOFF_INITIAL", &cfg.BackoffInitial}, {"BACKOFF_INITIAL", &cfg.BackoffInitial},
{"BACKOFF_MAX", &cfg.BackoffMax}, {"BACKOFF_MAX", &cfg.BackoffMax},
{"COMFY_IDLE_TIMEOUT", &cfg.ComfyIdleTimeout},
{"COMFY_START_TIMEOUT", &cfg.ComfyStartTimeout},
{"UPDATE_INTERVAL", &cfg.UpdateInterval}, {"UPDATE_INTERVAL", &cfg.UpdateInterval},
{"GAME_POLL_INTERVAL", &cfg.GamePollInterval},
} { } {
if err := envDuration(getenv, e.name, e.dst); err != nil { if err := envDuration(getenv, e.name, e.dst); err != nil {
return cfg, err return cfg, err
} }
} }
if v := getenv("GAME_PROCS"); v != "" {
cfg.GameProcs = splitList(v)
}
if v := getenv("GPU_IGNORE_PROCS"); v != "" {
cfg.GPUIgnoreProcs = splitList(v)
}
if v := getenv("GPU_FOREIGN_VRAM_MB"); v != "" {
n, err := strconv.Atoi(v)
if err != nil || n < 0 {
return cfg, fmt.Errorf("GPU_FOREIGN_VRAM_MB: must be a non-negative integer (MiB, 0 = disabled)")
}
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 {
@@ -199,6 +288,17 @@ func Load(getenv func(string) string) (Config, error) {
} }
cfg.BusyRetryAfter = n cfg.BusyRetryAfter = n
} }
if v := getenv("APP_VER"); v != "" {
switch {
case v == "dev" || v == "stable":
cfg.AppVersion = v
default:
if _, ok := parseVersion(v); !ok {
return cfg, fmt.Errorf("APP_VER: must be \"dev\", \"stable\" or a vX.Y.Z version")
}
cfg.AppVersion = "v" + strings.TrimPrefix(v, "v")
}
}
// LOGLEVEL is the canonical spelling; LOG_LEVEL is kept as an alias. // LOGLEVEL is the canonical spelling; LOG_LEVEL is kept as an alias.
logLevelValue := getenv("LOGLEVEL") logLevelValue := getenv("LOGLEVEL")
if logLevelValue == "" { if logLevelValue == "" {
@@ -218,8 +318,14 @@ 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.ComfyCmd != "" && cfg.ComfyURL == "" {
return cfg, fmt.Errorf("COMFY_CMD requires COMFY_URL to be set (the proxy needs somewhere to forward)")
}
if cfg.ComfyCmd == "" && cfg.ComfyDir != "" && cfg.ComfyURL == "" {
return cfg, fmt.Errorf("COMFY_DIR without COMFY_CMD requires COMFY_URL to be set (it enables the managed ComfyUI)")
}
if cfg.OllamaURL == "" && cfg.ComfyURL == "" { if cfg.OllamaURL == "" && cfg.ComfyURL == "" {
return cfg, fmt.Errorf("at least one of OLLAMA_URL or COMFY_URL must be set (each URL enables its consumer)") return cfg, ErrNoConsumer
} }
return cfg, nil return cfg, nil
} }
+131
View File
@@ -1,6 +1,7 @@
package config package config
import ( import (
"errors"
"log/slog" "log/slog"
"strings" "strings"
"testing" "testing"
@@ -39,6 +40,61 @@ func TestLoadRequiresConsumer(t *testing.T) {
if err == nil || !strings.Contains(err.Error(), "OLLAMA_URL") { if err == nil || !strings.Contains(err.Error(), "OLLAMA_URL") {
t.Fatalf("err = %v, want missing-consumer error", err) t.Fatalf("err = %v, want missing-consumer error", err)
} }
if !errors.Is(err, ErrNoConsumer) {
t.Fatalf("err = %v, want errors.Is(err, ErrNoConsumer)", err)
}
}
func TestAppVersion(t *testing.T) {
load := func(appVer string) (Config, error) {
return Load(func(k string) string {
switch k {
case "OLLAMA_URL":
return "http://127.0.0.1:11435"
case "APP_VER":
return appVer
}
return ""
})
}
cfg, err := load("")
if err != nil || cfg.AppVersion != "stable" {
t.Fatalf("default AppVersion = %q, err %v; want stable", cfg.AppVersion, err)
}
for _, v := range []string{"dev", "stable"} {
if cfg, err := load(v); err != nil || cfg.AppVersion != v {
t.Fatalf("APP_VER=%s: got %q, err %v", v, cfg.AppVersion, err)
}
}
if cfg, err := load("1.2.3"); err != nil || cfg.AppVersion != "v1.2.3" {
t.Fatalf("APP_VER=1.2.3: got %q, err %v; want normalized v1.2.3", cfg.AppVersion, err)
}
if _, err := load("nightly"); err == nil {
t.Fatal("APP_VER=nightly: want validation error")
}
}
func TestComfyCmdRequiresURL(t *testing.T) {
_, err := Load(func(k string) string {
if k == "COMFY_CMD" {
return "python main.py"
}
return ""
})
if err == nil || !strings.Contains(err.Error(), "COMFY_CMD requires COMFY_URL") {
t.Fatalf("err = %v, want COMFY_CMD/COMFY_URL validation error", err)
}
if _, err := Load(func(k string) string {
switch k {
case "COMFY_CMD":
return "python main.py"
case "COMFY_URL":
return "http://127.0.0.1:8188"
}
return ""
}); err != nil {
t.Fatalf("COMFY_CMD with COMFY_URL must load: %v", err)
}
} }
func TestParseEnvFile(t *testing.T) { func TestParseEnvFile(t *testing.T) {
@@ -101,6 +157,8 @@ func TestLoadErrors(t *testing.T) {
{"LLM_BUSY_MODE", "bogus"}, {"LLM_BUSY_MODE", "bogus"},
{"LLM_BUSY_STATUS", "200"}, {"LLM_BUSY_STATUS", "200"},
{"BUSY_RETRY_AFTER", "0"}, {"BUSY_RETRY_AFTER", "0"},
{"GPU_FOREIGN_VRAM_MB", "-1"},
{"GAME_POLL_INTERVAL", "bogus"},
} { } {
_, err := Load(func(k string) string { _, err := Load(func(k string) string {
if k == tc.key { if k == tc.key {
@@ -116,3 +174,76 @@ func TestLoadErrors(t *testing.T) {
} }
} }
} }
func TestGameDetectionSettings(t *testing.T) {
cfg, err := Load(func(k string) string {
switch k {
case "OLLAMA_URL":
return "http://127.0.0.1:11435"
case "GAME_PROCS":
return " cyberpunk2077.exe, hl2.exe ,, "
case "GPU_FOREIGN_VRAM_MB":
return "1024"
case "GPU_IGNORE_PROCS":
return "ollama, my-trainer"
}
return ""
})
if err != nil {
t.Fatal(err)
}
if len(cfg.GameProcs) != 2 || cfg.GameProcs[0] != "cyberpunk2077.exe" || cfg.GameProcs[1] != "hl2.exe" {
t.Fatalf("GameProcs = %v", cfg.GameProcs)
}
if cfg.GPUForeignVRAMMB != 1024 {
t.Fatalf("GPUForeignVRAMMB = %d", cfg.GPUForeignVRAMMB)
}
if len(cfg.GPUIgnoreProcs) != 2 || cfg.GPUIgnoreProcs[1] != "my-trainer" {
t.Fatalf("GPUIgnoreProcs = %v", cfg.GPUIgnoreProcs)
}
if cfg.GamePollInterval != 15*time.Second {
t.Fatalf("GamePollInterval = %v, want 15s default", cfg.GamePollInterval)
}
// Defaults: both detection paths off, ignore list covers our consumers.
cfg, err = Load(func(k string) string {
if k == "OLLAMA_URL" {
return "http://127.0.0.1:11435"
}
return ""
})
if err != nil {
t.Fatal(err)
}
if len(cfg.GameProcs) != 0 || cfg.GPUForeignVRAMMB != 0 {
t.Fatalf("detection must be off by default: %v %d", cfg.GameProcs, cfg.GPUForeignVRAMMB)
}
if len(cfg.GPUIgnoreProcs) == 0 {
t.Fatal("GPUIgnoreProcs default must not be empty")
}
}
func TestComfyDirOnlyEnablesManaged(t *testing.T) {
// COMFY_DIR without COMFY_CMD and without COMFY_URL is a mistake.
_, err := Load(func(k string) string {
if k == "COMFY_DIR" {
return `C:\ComfyUI`
}
return ""
})
if err == nil || !strings.Contains(err.Error(), "COMFY_DIR") {
t.Fatalf("err = %v, want COMFY_DIR/COMFY_URL validation error", err)
}
// With COMFY_URL it loads — the launch command is derived from the dir.
if _, err := Load(func(k string) string {
switch k {
case "COMFY_DIR":
return `C:\ComfyUI`
case "COMFY_URL":
return "http://127.0.0.1:8189"
}
return ""
}); err != nil {
t.Fatalf("COMFY_DIR with COMFY_URL must load: %v", err)
}
}
+194
View File
@@ -0,0 +1,194 @@
package config
import (
"fmt"
"strconv"
"strings"
)
// sampleEntry is one line pair in the sample env file: a comment and the
// (usually commented-out) KEY=VALUE line.
type sampleEntry struct {
name string
value string
comment string
active bool // rendered uncommented
}
// appVerComment documents APP_VER wherever it is rendered.
const appVerComment = `Version to run: "dev" disables updates, "stable" tracks the latest release, or pin an exact release like v0.1.7`
// sampleEntries lists every setting in sample order. logFile activates the
// LOG_FILE line (Windows install); empty keeps it commented like the rest.
// CFG_VER and APP_VER are not entries — they head the file, always active.
func sampleEntries(logFile string) []sampleEntry {
return []sampleEntry{
{"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},
{"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},
{"WARM_MODEL", "", "Optional model to reload after an image job (default: empty = none)", false},
{"COMFY_CMD", `"C:\ComfyUI\.venv\Scripts\python.exe" ComfyUI\main.py --port 8188`, "Spawn and supervise ComfyUI on demand: the first request starts it, it stops after COMFY_IDLE_TIMEOUT to free VRAM (default: derived from COMFY_DIR; both empty = unmanaged)", false},
{"COMFY_DIR", `C:\ComfyUI`, "Root of a standard ComfyUI venv install (.venv + main.py): setting it alone supervises ComfyUI with the derived launch command; also the working directory for COMFY_CMD", false},
{"COMFY_IDLE_TIMEOUT", "5m", "Stop the managed ComfyUI after this long without requests or jobs (frees VRAM)", 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},
{"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_UTIL_PCT", "30", "Also treat the GPU as held when a process not in GPU_IGNORE_PROCS uses more than this percent of the GPU 3D engine (Windows Task-Manager counters; catches games without an exe list; 0/empty = disabled)", false},
{"GPU_IGNORE_PROCS", "ollama,ollama app,ollama_llama_server,python,pythonw,dwm", "Process names never counted as foreign GPU users (ComfyUI runs under python; dwm is the desktop compositor)", false},
{"GAME_POLL_INTERVAL", "15s", "How often game/VRAM detection runs (nvidia-smi polls keep the GPU awake; don't go below ~10s)", false},
{"UNLOAD_TIMEOUT", "60s", "How long to wait for Ollama to unload a model", false},
{"JOB_TIMEOUT", "15m", "Maximum time to wait for a ComfyUI job", false},
{"LLM_WAIT_TIMEOUT", "10m", "Max time an LLM request waits for the GPU before being answered 503 (wait mode)", false},
{"LLM_BUSY_MODE", "wait", `How blocked LLM requests are handled: "wait" holds them, "reject" fails them immediately`, false},
{"LLM_BUSY_STATUS", "503", "HTTP status for rejected LLM requests in reject mode (400-599, e.g. 429)", false},
{"BUSY_RETRY_AFTER", "30", "Seconds sent as Retry-After on busy responses (both modes)", false},
{"LOGLEVEL", "warn", `Log verbosity: debug, info, warn, error; "info" logs every request (LOG_LEVEL works too)`, false},
{"LOG_FORMAT", "text", `Log format: "text" or "json"`, false},
{"LOG_FILE", logFile, "Append logs to this file instead of stderr (a Windows service has no console)", logFile != ""},
{"UNLOAD_POLL_INTERVAL", "500ms", "/api/ps poll interval while unloading", false},
{"HISTORY_POLL_INTERVAL", "1s", "/history/<id> poll interval while a job runs", false},
{"PROBE_TIMEOUT", "5s", "Probe timeout for the startup probe and the periodic health check", false},
{"HEALTH_INTERVAL", "30s", "How often enabled upstreams are probed; status changes are logged", false},
{"FREE_TIMEOUT", "30s", "Timeout for the POST /free call after an image job", false},
{"WARM_TIMEOUT", "2m", "Timeout for the warm-model reload after an image job", false},
{"SHUTDOWN_TIMEOUT", "10s", "Graceful shutdown timeout on SIGINT/SIGTERM", false},
{"BACKOFF_INITIAL", "1s", "First retry wait when an upstream connection fails", false},
{"BACKOFF_MAX", "60s", "Cap for the exponential retry backoff", false},
{"PROMPT_CAPTURE_LIMIT", "65536", "Bytes of the /prompt response buffered to find prompt_id (pass-through unaffected)", false},
{"AUTO_UPDATE", "true", "Poll the releases API for signed updates", false},
{"UPDATE_INTERVAL", "6h", "Auto-update check interval", false},
{"UPDATE_REPO", "https://git.rambossek.at/PUBLIC/gpu-turnstile", "Repository to check for releases", false},
{"UPDATE_ASSET", "gpu-turnstile.exe", "Release asset to download", false},
}
}
// SampleEnv renders a sample .env file covering every setting, each with a
// comment line. Everything is commented out — so all defaults apply —
// except the CFG_VER/APP_VER header and LOG_FILE when logFile is non-empty
// (a Windows service has no console). CFG_VER records the version that
// wrote the file so installs — and startups after an update — can upgrade
// it.
func SampleEnv(version, logFile string) string {
// CFG_VER is always a concrete vX.Y.Z — never "dev". A dev build
// stamps v0.0.0, which sorts older than any release, so the next
// release install upgrades the file and stamps a proper version.
cfgVer := version
if cfgVer == "" || cfgVer == "dev" {
cfgVer = "v0.0.0"
}
var b strings.Builder
fmt.Fprintf(&b, "CFG_VER=%s\n", cfgVer)
b.WriteString("# Config format reference, written by the installer — do not edit.\n")
b.WriteString("# The installer uses it to append newly added settings on updates.\n\n")
b.WriteString("# gpu-turnstile configuration\n")
b.WriteString("# KEY=VALUE lines; \"#\" starts a comment. Every setting below is at its\n")
b.WriteString("# default and commented out — remove the \"#\" to change it.\n")
b.WriteString("# At least one of OLLAMA_URL / COMFY_URL must be set for the proxy to start.\n\n")
writeEntry(&b, sampleEntry{"APP_VER", "stable", appVerComment, true})
for _, e := range sampleEntries(logFile) {
writeEntry(&b, e)
}
return b.String()
}
func writeEntry(b *strings.Builder, e sampleEntry) {
fmt.Fprintf(b, "# %s\n", e.comment)
if e.active {
fmt.Fprintf(b, "%s=%s\n\n", e.name, e.value)
} else {
fmt.Fprintf(b, "#%s=%s\n\n", e.name, e.value)
}
}
// SyncSample upgrades an installer-written env file: when its CFG_VER says
// it was written by an older gpu-turnstile, every setting the file does not
// mention — commented or not — is appended at the end, and CFG_VER is
// updated to version. Files without CFG_VER (hand-written or foreign),
// up-to-date files and "dev" builds are returned unchanged; changed reports
// whether the returned content differs.
func SyncSample(data, version, logFile string) (string, bool) {
if version == "" || version == "dev" {
return data, false
}
values, err := ParseEnvFile(strings.NewReader(data))
if err != nil || values["CFG_VER"] == "" {
return data, false // not installer-written; the caller decides
}
if compareVersions(values["CFG_VER"], version) >= 0 {
return data, false // same or newer
}
present := make(map[string]bool)
for _, line := range strings.Split(data, "\n") {
line = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "#"))
if name, _, ok := strings.Cut(line, "="); ok {
present[strings.TrimSpace(name)] = true
}
}
lines := strings.Split(data, "\n")
for i, l := range lines {
if strings.HasPrefix(strings.TrimSpace(l), "CFG_VER=") {
lines[i] = "CFG_VER=" + version
break
}
}
out := strings.Join(lines, "\n")
if !strings.HasSuffix(out, "\n") {
out += "\n"
}
var b strings.Builder
if !present["APP_VER"] {
fmt.Fprintf(&b, "\n# Added by gpu-turnstile %s:\n", version)
writeEntry(&b, sampleEntry{"APP_VER", "stable", appVerComment, true})
}
for _, e := range sampleEntries(logFile) {
if !present[e.name] {
fmt.Fprintf(&b, "\n# Added by gpu-turnstile %s:\n", version)
writeEntry(&b, e)
}
}
out += b.String()
return out, out != data
}
// compareVersions orders two vX.Y.Z version strings. Unparseable versions
// (including "dev") sort before any release.
func compareVersions(a, b string) int {
pa, oka := parseVersion(a)
pb, okb := parseVersion(b)
if oka != okb {
if oka {
return 1
}
return -1
}
for i := range pa {
if pa[i] != pb[i] {
if pa[i] > pb[i] {
return 1
}
return -1
}
}
return 0
}
func parseVersion(v string) ([3]int, bool) {
var out [3]int
v = strings.TrimPrefix(strings.TrimSpace(v), "v")
parts := strings.Split(v, ".")
if len(parts) != 3 {
return out, false
}
for i, p := range parts {
n, err := strconv.Atoi(p)
if err != nil || n < 0 {
return out, false
}
out[i] = n
}
return out, true
}
+156
View File
@@ -0,0 +1,156 @@
package config
import (
"strings"
"testing"
)
const sampleLogPath = `C:\ProgramData\gpu-turnstile\gpu-turnstile.log`
var allSettingNames = []string{
"LISTEN_OLLAMA", "LISTEN_COMFY", "OLLAMA_URL", "COMFY_URL",
"WARM_MODEL", "UNLOAD_TIMEOUT", "JOB_TIMEOUT", "LLM_WAIT_TIMEOUT",
"LLM_BUSY_MODE", "LLM_BUSY_STATUS", "BUSY_RETRY_AFTER",
"LOGLEVEL", "LOG_FORMAT", "LOG_FILE",
"UNLOAD_POLL_INTERVAL", "HISTORY_POLL_INTERVAL", "PROBE_TIMEOUT",
"FREE_TIMEOUT", "WARM_TIMEOUT", "SHUTDOWN_TIMEOUT",
"BACKOFF_INITIAL", "BACKOFF_MAX", "PROMPT_CAPTURE_LIMIT",
"AUTO_UPDATE", "UPDATE_INTERVAL", "UPDATE_REPO", "UPDATE_ASSET",
}
func TestSampleEnv(t *testing.T) {
sample := SampleEnv("v0.1.7", sampleLogPath)
if !strings.HasPrefix(sample, "CFG_VER=v0.1.7\n") {
t.Errorf("first line does not carry CFG_VER: %q", strings.SplitN(sample, "\n", 2)[0])
}
// Every setting known to Load must appear.
for _, name := range allSettingNames {
if !strings.Contains(sample, name+"=") {
t.Errorf("sample is missing %s", name)
}
}
// The sample must parse cleanly; active values are CFG_VER, APP_VER
// and LOG_FILE.
values, err := ParseEnvFile(strings.NewReader(sample))
if err != nil {
t.Fatalf("sample does not parse: %v", err)
}
want := map[string]string{"CFG_VER": "v0.1.7", "APP_VER": "stable", "LOG_FILE": sampleLogPath}
if len(values) != len(want) {
t.Fatalf("active values = %v, want %v", values, want)
}
for k, v := range want {
if values[k] != v {
t.Errorf("%s = %q, want %q", k, values[k], v)
}
}
// Without a log path LOG_FILE stays commented out.
values, err = ParseEnvFile(strings.NewReader(SampleEnv("v0.1.7", "")))
if err != nil {
t.Fatalf("sample without log path does not parse: %v", err)
}
if len(values) != 2 || values["LOG_FILE"] != "" {
t.Fatalf("active values = %v, want only CFG_VER and APP_VER", values)
}
}
func TestSyncSample(t *testing.T) {
old := SampleEnv("v0.1.6", sampleLogPath)
// Simulate a setting that did not exist yet when the file was written.
old = strings.Replace(old, "#UPDATE_ASSET=gpu-turnstile.exe\n", "", 1)
out, changed := SyncSample(old, "v0.1.7", sampleLogPath)
if !changed {
t.Fatal("older installer file was not upgraded")
}
if !strings.Contains(out, "\nCFG_VER=v0.1.7\n") && !strings.HasPrefix(out, "CFG_VER=v0.1.7\n") {
t.Error("CFG_VER was not updated to the new version")
}
if !strings.Contains(out, "#UPDATE_ASSET=gpu-turnstile.exe") {
t.Error("missing setting was not appended")
}
if !strings.Contains(out, "LOG_FILE="+sampleLogPath) {
t.Error("existing active LOG_FILE was lost")
}
if !strings.Contains(out, "Added by gpu-turnstile v0.1.7") {
t.Error("appended section is not attributed")
}
if _, err := ParseEnvFile(strings.NewReader(out)); err != nil {
t.Fatalf("upgraded file does not parse: %v", err)
}
// A file written by the same or a newer version is left alone.
if _, changed := SyncSample(SampleEnv("v0.1.7", sampleLogPath), "v0.1.7", sampleLogPath); changed {
t.Error("same-version file was modified")
}
if _, changed := SyncSample(SampleEnv("v0.2.0", sampleLogPath), "v0.1.7", sampleLogPath); changed {
t.Error("newer-version file was modified")
}
// Files without CFG_VER are not installer-written; the installer
// replaces them, SyncSample leaves them alone.
user := "OLLAMA_URL=http://host:11434\n"
if out, changed := SyncSample(user, "v0.1.7", sampleLogPath); changed || out != user {
t.Error("file without CFG_VER was modified")
}
// A dev build never upgrades.
if _, changed := SyncSample(old, "dev", sampleLogPath); changed {
t.Error("dev build modified the file")
}
}
func TestSyncSampleAppendsMissingAppVer(t *testing.T) {
old := SampleEnv("v0.1.6", sampleLogPath)
old = strings.Replace(old, "# "+appVerComment+"\n", "", 1)
old = strings.Replace(old, "APP_VER=stable\n", "", 1)
out, changed := SyncSample(old, "v0.1.7", sampleLogPath)
if !changed {
t.Fatal("file without APP_VER was not upgraded")
}
values, err := ParseEnvFile(strings.NewReader(out))
if err != nil {
t.Fatalf("upgraded file does not parse: %v", err)
}
if values["APP_VER"] != "stable" {
t.Fatalf("APP_VER = %q, want appended default \"stable\"", values["APP_VER"])
}
}
func TestSampleEnvDevStampsConcreteVersion(t *testing.T) {
// A dev build must never write CFG_VER=dev: it stamps v0.0.0, and the
// next release install upgrades the file to a proper version.
dev := SampleEnv("dev", sampleLogPath)
if !strings.HasPrefix(dev, "CFG_VER=v0.0.0\n") {
t.Errorf("dev sample first line: %q", strings.SplitN(dev, "\n", 2)[0])
}
out, changed := SyncSample(dev, "v0.1.7", sampleLogPath)
if !changed || !strings.HasPrefix(out, "CFG_VER=v0.1.7\n") {
t.Errorf("dev-stamped file was not upgraded to v0.1.7 (changed=%v)", changed)
}
}
func TestCompareVersions(t *testing.T) {
cases := []struct {
a, b string
want int
}{
{"v0.1.6", "v0.1.7", -1},
{"v0.1.7", "v0.1.7", 0},
{"v1.0.0", "v0.9.9", 1},
{"0.1.7", "v0.1.6", 1},
{"dev", "v0.1.7", -1},
{"v0.1.7", "dev", 1},
{"dev", "dev", 0},
}
for _, c := range cases {
if got := compareVersions(c.a, c.b); got != c.want {
t.Errorf("compareVersions(%q, %q) = %d, want %d", c.a, c.b, got, c.want)
}
}
}
+117
View File
@@ -0,0 +1,117 @@
// Package control exposes a local-only command channel into a running
// gpu-turnstile service: a named pipe on Windows, a unix socket on Linux.
// It lets unprivileged local users ask the service to do privileged work
// that is safe to offer — currently triggering an update check, whose
// payload is signature-verified regardless of who asks. The channel never
// accepts data beyond a one-word command, and the server rate-limits
// triggers, so the worst a local user can cause is a cheap, throttled
// check and a GPU-idle-gated restart onto a signed binary.
//
// Abuse hardening: the command read is capped (4 KiB), each connection is
// force-closed after connTimeout so a stalled client cannot pin a goroutine
// (or a Windows pipe instance) forever, and concurrently served connections
// are capped at maxConns — beyond that, connections are closed on arrival.
// On Windows the pipe's ACL additionally denies network logons, so the
// channel cannot be reached from another machine.
//
// Protocol: the client writes one command line, the server answers with
// one reply line ("OK ..." or "ERR ...") and hangs up.
package control
import (
"bufio"
"errors"
"fmt"
"io"
"strings"
"time"
)
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
// idle) restart onto a signed update immediately.
const CmdUpdateNow = "update-now"
// CmdStatus asks for a one-line JSON status snapshot (monitor mode).
const CmdStatus = "status"
// CmdReloadEnv asks the service to re-read and validate its config file,
// and to restart onto it (once the GPU is idle) when it changed.
const CmdReloadEnv = "reload-env"
// ErrUnavailable means no running service offers the control channel.
var ErrUnavailable = errors.New("control channel unavailable")
// Handler answers one command; the returned string is sent back as one
// line. It must start with "OK " or "ERR ".
type Handler func(cmd string) string
// connTimeout bounds one connection's lifetime: a client that stops
// mid-command or never reads the reply would otherwise pin its goroutine
// (and on Windows one of the pipe instances) indefinitely. A var so tests
// can shrink it.
var connTimeout = 10 * time.Second
// maxConns caps concurrently served connections; beyond it, new
// connections are closed on arrival. Bound on the goroutines a local
// flood can pile up.
const maxConns = 32
var connSem = make(chan struct{}, maxConns)
// serve dispatches connection handling under the concurrency cap. It
// returns false when the cap is reached — the caller must then close the
// connection itself.
func serve(c io.ReadWriteCloser, h Handler) bool {
select {
case connSem <- struct{}{}:
go func() {
defer func() { <-connSem }()
serveConn(c, h)
}()
return true
default:
return false
}
}
// forceCloser is implemented by connections that can be torn down
// abortively, unblocking pending reads and writes (Windows pipe:
// DisconnectNamedPipe; unix socket: a deadline in the past). The
// connection watchdog uses it; normal closes still flush the reply.
type forceCloser interface {
ForceClose() error
}
// serveConn runs the line protocol on one accepted connection.
func serveConn(c io.ReadWriteCloser, h Handler) {
defer c.Close()
if fc, ok := c.(forceCloser); ok {
timer := time.AfterFunc(connTimeout, func() { fc.ForceClose() })
defer timer.Stop()
}
line, err := bufio.NewReader(io.LimitReader(c, 4096)).ReadString('\n')
cmd := strings.TrimSpace(line)
if cmd == "" {
if err != nil {
return
}
fmt.Fprintln(c, "ERR empty command")
return
}
fmt.Fprintln(c, h(cmd))
}
// readReply writes cmd and reads the server's one-line reply.
func readReply(c io.ReadWriteCloser, cmd string) (string, error) {
if _, err := fmt.Fprintln(c, cmd); err != nil {
return "", err
}
// The server hangs up after its reply; a broken-pipe error after the
// last byte still leaves the reply in the buffer.
data, _ := io.ReadAll(io.LimitReader(c, 4096))
line := strings.TrimSpace(string(data))
if line == "" {
return "", ErrUnavailable
}
return line, nil
}
+66
View File
@@ -0,0 +1,66 @@
//go:build linux
package control
import (
"context"
"log/slog"
"net"
"os"
"time"
)
// sockPath lives in the unit's RuntimeDirectory; mode 0666 lets every
// local user ask, nothing can reach it from off the machine.
const sockPath = "/run/gpu-turnstile/control.sock"
// Serve starts the socket listener in the background and returns; only a
// setup failure is reported. Each client connection is answered in its own
// goroutine.
func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
os.Remove(sockPath) // stale socket from a previous run
ln, err := net.Listen("unix", sockPath)
if err != nil {
return err
}
if err := os.Chmod(sockPath, 0o666); err != nil {
ln.Close()
return err
}
go func() {
<-ctx.Done()
ln.Close()
}()
go func() {
for {
c, err := ln.Accept()
if err != nil {
return // shutting down
}
uc := unixConn{c}
if !serve(uc, h) {
uc.ForceClose()
}
}
}()
return nil
}
// unixConn adds an abortive ForceClose to net.Conn: a deadline in the
// past fails pending and future I/O immediately.
type unixConn struct{ net.Conn }
func (c unixConn) ForceClose() error {
c.SetDeadline(time.Now().Add(-time.Second)) //nolint:errcheck // best effort
return c.Conn.Close()
}
// Ask sends one command to the running service and returns its reply.
func Ask(cmd string) (string, error) {
c, err := net.DialTimeout("unix", sockPath, 2*time.Second)
if err != nil {
return "", ErrUnavailable
}
defer c.Close()
return readReply(c, cmd)
}
+18
View File
@@ -0,0 +1,18 @@
//go:build !windows && !linux
package control
import (
"context"
"log/slog"
)
// Serve is a no-op on platforms without a control channel.
func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
return ErrUnavailable
}
// Ask always reports the channel as unavailable.
func Ask(cmd string) (string, error) {
return "", ErrUnavailable
}
+103
View File
@@ -0,0 +1,103 @@
package control
import (
"errors"
"net"
"strings"
"testing"
"time"
)
func TestRoundTrip(t *testing.T) {
server, client := net.Pipe()
go serveConn(server, func(cmd string) string {
if cmd != CmdUpdateNow {
return "ERR unknown command: " + cmd
}
return "OK v0.2.2 is up to date"
})
reply, err := readReply(client, CmdUpdateNow)
if err != nil {
t.Fatal(err)
}
if reply != "OK v0.2.2 is up to date" {
t.Fatalf("reply = %q", reply)
}
server2, client2 := net.Pipe()
go serveConn(server2, func(cmd string) string { return "ERR unknown command: " + cmd })
reply, err = readReply(client2, "bogus")
if err != nil {
t.Fatal(err)
}
if !strings.HasPrefix(reply, "ERR ") {
t.Fatalf("reply = %q, want ERR prefix", reply)
}
}
func TestEmptyReplyIsUnavailable(t *testing.T) {
server, client := net.Pipe()
go serveConn(server, func(cmd string) string {
server.Close() // hang up without answering
return ""
})
if _, err := readReply(client, CmdUpdateNow); !errors.Is(err, ErrUnavailable) {
t.Fatalf("err = %v, want ErrUnavailable", err)
}
}
func TestServeCap(t *testing.T) {
for i := 0; i < maxConns; i++ {
connSem <- struct{}{}
}
defer func() {
for i := 0; i < maxConns; i++ {
<-connSem
}
}()
server, client := net.Pipe()
defer server.Close()
defer client.Close()
if serve(server, func(string) string { return "OK" }) {
t.Fatal("serve accepted a connection beyond the cap")
}
}
// forcePipe records ForceClose calls for the watchdog test.
type forcePipe struct {
net.Conn
forced chan struct{}
}
func (c forcePipe) ForceClose() error {
err := c.Conn.Close()
close(c.forced)
return err
}
func TestConnWatchdog(t *testing.T) {
old := connTimeout
connTimeout = 50 * time.Millisecond
defer func() { connTimeout = old }()
server, client := net.Pipe()
defer client.Close()
fc := forcePipe{Conn: server, forced: make(chan struct{})}
done := make(chan struct{})
go func() {
serveConn(fc, func(string) string { return "OK" })
close(done)
}()
// The client never sends anything; the watchdog must tear the
// connection down instead of blocking forever.
select {
case <-fc.forced:
case <-time.After(5 * time.Second):
t.Fatal("watchdog did not force-close the stalled connection")
}
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("serveConn still blocked after the force close")
}
}
+167
View File
@@ -0,0 +1,167 @@
//go:build windows
package control
import (
"context"
"fmt"
"log/slog"
"os"
"syscall"
"time"
"unsafe"
"golang.org/x/sys/windows"
)
// pipePath is a kernel-local named pipe: no TCP, no firewall prompt.
const pipePath = `\\.\pipe\gpu-turnstile`
// sddlPipe grants full access to Administrators, SYSTEM and the pipe owner,
// and read+write to authenticated users — except network logons, so the
// pipe cannot be reached from another machine over SMB.
const sddlPipe = "D:(D;;GRGW;;;NU)(A;;GA;;;BA)(A;;GA;;;SY)(A;;GA;;;OW)(A;;GRGW;;;AU)"
var (
procConvertSDDL = windows.NewLazySystemDLL("advapi32.dll").
NewProc("ConvertStringSecurityDescriptorToSecurityDescriptorW")
procWaitNamedPipe = windows.NewLazySystemDLL("kernel32.dll").
NewProc("WaitNamedPipeW")
)
func waitNamedPipe(name *uint16, timeout uint32) error {
r, _, err := procWaitNamedPipe.Call(uintptr(unsafe.Pointer(name)), uintptr(timeout))
if r == 0 {
return err
}
return nil
}
func securityAttributesFromSDDL(sddl string) (*windows.SecurityAttributes, error) {
s, err := windows.UTF16PtrFromString(sddl)
if err != nil {
return nil, err
}
var sd *uint16 // SECURITY_DESCRIPTOR*, kept for the process lifetime
r, _, callErr := procConvertSDDL.Call(
uintptr(unsafe.Pointer(s)), 1, /* SDDL_REVISION_1 */
uintptr(unsafe.Pointer(&sd)), 0)
if r == 0 {
return nil, fmt.Errorf("invalid SDDL: %w", callErr)
}
sa := &windows.SecurityAttributes{
Length: uint32(unsafe.Sizeof(windows.SecurityAttributes{})),
SecurityDescriptor: (*windows.SECURITY_DESCRIPTOR)(unsafe.Pointer(sd)),
}
return sa, nil
}
// Serve starts the pipe listener in the background and returns; only a
// setup failure is reported. Each client connection is answered in its own
// goroutine. On shutdown the process exit reaps everything.
func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
sa, err := securityAttributesFromSDDL(sddlPipe)
if err != nil {
return err
}
name, err := windows.UTF16PtrFromString(pipePath)
if err != nil {
return err
}
go func() {
for ctx.Err() == nil {
pipe, err := windows.CreateNamedPipe(name,
windows.PIPE_ACCESS_DUPLEX,
windows.PIPE_TYPE_BYTE|windows.PIPE_READMODE_BYTE|windows.PIPE_WAIT,
16, 4096, 4096, 0, sa)
if err != nil {
log.Warn("control channel stopped", "err", err)
return
}
// Blocks until a client connects — only then is the next
// instance created, so instances are not burned without
// clients. ERROR_PIPE_CONNECTED means the client raced us
// and connected before the call: that is a success. Process
// exit reaps the blocked call on shutdown.
if err := windows.ConnectNamedPipe(pipe, nil); err != nil && err != errnoPipeConnected {
windows.CloseHandle(pipe)
continue
}
conn := &pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}
if !serve(conn, h) {
conn.ForceClose()
}
}
}()
return nil
}
// errnoPipeConnected is ConnectNamedPipe's "the client connected before we
// called" result, which means the connection is established.
var errnoPipeConnected = syscall.Errno(535) // ERROR_PIPE_CONNECTED
// errnoPipeBusy is CreateFile's "all pipe instances are busy" result. It
// loses the race against another client that grabbed the instance
// WaitNamedPipe just reported — the caller must wait and retry.
var errnoPipeBusy = syscall.Errno(231) // ERROR_PIPE_BUSY
// pipeConn adapts a pipe handle to io.ReadWriteCloser. Close flushes first
// (FlushFileBuffers blocks until the client has read the reply) and then
// disconnects — closing the bare handle right after writing can discard
// unread reply bytes, which clients see as an empty, failed request.
type pipeConn struct {
f *os.File
h windows.Handle
}
func (c *pipeConn) Read(p []byte) (int, error) { return c.f.Read(p) }
func (c *pipeConn) Write(p []byte) (int, error) { return c.f.Write(p) }
func (c *pipeConn) Close() error {
windows.FlushFileBuffers(c.h) //nolint:errcheck // best effort
windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort
return c.f.Close()
}
// ForceClose aborts the connection without flushing: disconnecting
// unblocks pending reads and writes at the cost of possibly discarding an
// unread reply. Used by the connection watchdog; normal closes flush.
func (c *pipeConn) ForceClose() error {
windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort
return c.f.Close()
}
// Ask sends one command to the running service and returns its reply.
//
// The server keeps exactly one listening instance per connection, so
// concurrent clients race for it: WaitNamedPipe can report an instance
// that another client grabs before our CreateFile runs (ERROR_PIPE_BUSY).
// Retry on that — with the monitor polling status every second, a single
// attempt loses that race regularly.
func Ask(cmd string) (string, error) {
name, err := windows.UTF16PtrFromString(pipePath)
if err != nil {
return "", err
}
deadline := time.Now().Add(5 * time.Second)
for {
if err := waitNamedPipe(name, 2000); err != nil {
return "", ErrUnavailable
}
handle, err := windows.CreateFile(name,
windows.GENERIC_READ|windows.GENERIC_WRITE, 0, nil,
windows.OPEN_EXISTING, 0, 0)
if err == errnoPipeBusy {
if time.Now().After(deadline) {
return "", ErrUnavailable
}
continue
}
if err != nil {
return "", ErrUnavailable
}
f := os.NewFile(uintptr(handle), pipePath)
defer f.Close()
return readReply(f, cmd)
}
}
+263
View File
@@ -0,0 +1,263 @@
// Package game detects processes outside gpu-turnstile's control that hold
// the GPU — typically a game — so the proxy can block new GPU work and free
// VRAM while they run. Three detection paths: an explicit process watch
// list (GAME_PROCS), a foreign-VRAM threshold via nvidia-smi
// (GPU_FOREIGN_VRAM_MB), and a per-process GPU 3D-engine utilization
// threshold via Windows PDH counters (GPU_FOREIGN_UTIL_PCT). The latter two
// catch anything not on the ignore list without naming individual games.
package game
import (
"context"
"errors"
"fmt"
"log/slog"
"os/exec"
"slices"
"strconv"
"strings"
)
// Process is one running OS process.
type Process struct {
PID int
Name string
}
// computeApp is one process holding GPU memory, as reported by nvidia-smi.
type computeApp struct {
PID int
UsedMB int
}
// Detector checks whether a foreign process holds the GPU. The zero value
// (no watch list, no thresholds) never detects anything; main only starts
// the poll loop when at least one path is configured.
type Detector struct {
procs map[string]bool // normalized names from GAME_PROCS
vramMB int // foreign VRAM threshold; 0 = disabled
utilPct int // foreign 3D-engine utilization threshold; 0 = disabled
ignore map[string]bool // normalized names never counted as foreign
log *slog.Logger
noNvidia bool // nvidia-smi was not found; VRAM path disabled for good
sampler *gpuEngineSampler // open PDH query, opened lazily on first Check
noPDH bool // engine counters unavailable; util path disabled for good
}
// New builds a Detector from the configured watch list, VRAM threshold in
// MiB, 3D-engine utilization threshold in percent (both 0 = disabled) and
// ignore list. Names are matched case-insensitively, with or without a
// trailing ".exe".
func New(procs []string, vramMB, utilPct int, ignore []string, log *slog.Logger) *Detector {
if log == nil {
log = slog.Default()
}
return &Detector{
procs: nameSet(procs),
vramMB: vramMB,
utilPct: utilPct,
ignore: nameSet(ignore),
log: log,
}
}
// normName lowercases a process name and strips a trailing ".exe" so the
// watch and ignore lists match on Windows and Linux spellings alike.
func normName(s string) string {
return strings.TrimSuffix(strings.ToLower(strings.TrimSpace(s)), ".exe")
}
func nameSet(names []string) map[string]bool {
set := make(map[string]bool, len(names))
for _, n := range names {
if n = normName(n); n != "" {
set[n] = true
}
}
return set
}
// Check looks once for foreign GPU holders and returns a human-readable
// description of each (empty when the GPU is free for gpu-turnstile's
// consumers). A failing nvidia-smi call is returned as an error only when
// the process list found nothing; a missing nvidia-smi binary disables the
// VRAM path permanently (logged once), as do missing engine counters.
func (d *Detector) Check(ctx context.Context) ([]string, error) {
ps, psErr := processes()
utils := d.engineUtil()
var apps []computeApp
if d.vramMB > 0 && !d.noNvidia {
var err error
apps, err = queryComputeApps(ctx)
if errors.Is(err, exec.ErrNotFound) {
d.noNvidia = true
d.log.Warn("GPU_FOREIGN_VRAM_MB is set but nvidia-smi was not found; VRAM detection disabled")
} else if err != nil {
return d.detect(ps, nil, utils), err
}
}
return d.detect(ps, apps, utils), psErr
}
// engineUtil samples per-process 3D-engine utilization via PDH. The first
// call only primes the rate counters and returns nil. A failing open or
// sample disables the path permanently (logged once).
func (d *Detector) engineUtil() map[int]float64 {
if d.utilPct <= 0 || d.noPDH {
return nil
}
if d.sampler == nil {
s, err := openGPUEngineSampler()
if err != nil {
d.noPDH = true
d.log.Warn("GPU_FOREIGN_UTIL_PCT is set but per-process GPU counters are unavailable; engine detection disabled", "err", err)
return nil
}
d.sampler = s
}
utils, err := d.sampler.sample()
if err != nil {
if errors.Is(err, errNotPrimed) {
return nil
}
d.noPDH = true
d.log.Warn("per-process GPU counters failed; engine detection disabled", "err", err)
return nil
}
return utils
}
// detect is the pure core of Check: given the process table and
// (optionally) the nvidia-smi compute-apps list and the PDH engine
// utilization, it returns the foreign holders.
func (d *Detector) detect(ps []Process, apps []computeApp, utils map[int]float64) []string {
var holders []string
for _, p := range ps {
if d.procs[normName(p.Name)] {
holders = append(holders, fmt.Sprintf("%s (pid %d)", p.Name, p.PID))
}
}
if (d.vramMB > 0 && apps != nil) || (d.utilPct > 0 && utils != nil) {
names := make(map[int]string, len(ps))
for _, p := range ps {
names[p.PID] = p.Name
}
for _, a := range apps {
name := names[a.PID]
if d.ignore[normName(name)] || a.UsedMB < d.vramMB {
continue
}
if name == "" {
name = "unknown process"
}
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
}
// queryComputeApps runs nvidia-smi and parses the per-process VRAM list.
// Note: under Windows' WDDM driver, nvidia-smi only sees compute
// allocations, so graphics-only games may not appear there — GAME_PROCS is
// the reliable path on Windows; on Linux both work.
func queryComputeApps(ctx context.Context) ([]computeApp, error) {
out, err := exec.CommandContext(ctx, "nvidia-smi",
"--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits").Output()
if err != nil {
return nil, err
}
return parseComputeApps(string(out))
}
// GPUStats is one nvidia-smi reading of the whole card.
type GPUStats struct {
UsedMB int
TotalMB int
TempC int // -1 when nvidia-smi reports N/A
FanPct int // -1 when N/A (some cards don't expose the fan)
}
// QueryGPUStats returns VRAM usage, temperature and fan speed via
// nvidia-smi. Unlike the per-process list this works under WDDM too.
func QueryGPUStats(ctx context.Context) (GPUStats, error) {
out, err := exec.CommandContext(ctx, "nvidia-smi",
"--query-gpu=memory.used,memory.total,temperature.gpu,fan.speed", "--format=csv,noheader,nounits").Output()
if err != nil {
return GPUStats{}, err
}
return parseGPUStats(string(out))
}
// parseGPUStats parses one "used, total, temp, fan" CSV line (MiB, °C,
// percent). The memory fields must be numeric; temperature and fan fall
// back to -1 on "N/A" and friends.
func parseGPUStats(out string) (GPUStats, error) {
fields := strings.Split(strings.TrimSpace(out), ",")
if len(fields) != 4 {
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(out))
}
num := func(s string) (int, error) {
return strconv.Atoi(strings.TrimSpace(s))
}
optional := func(s string) int {
n, err := num(s)
if err != nil {
return -1
}
return n
}
var st GPUStats
var err error
if st.UsedMB, err = num(fields[0]); err != nil {
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(out))
}
if st.TotalMB, err = num(fields[1]); err != nil {
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(out))
}
st.TempC = optional(fields[2])
st.FanPct = optional(fields[3])
return st, nil
}
// parseComputeApps parses "pid, used_memory" CSV lines (no header, MiB
// units). Unsupported rows ("N/A" on WDDM) are skipped.
func parseComputeApps(out string) ([]computeApp, error) {
var apps []computeApp
for _, line := range strings.Split(out, "\n") {
line = strings.TrimSpace(line)
if line == "" {
continue
}
pidStr, memStr, ok := strings.Cut(line, ",")
if !ok {
return nil, fmt.Errorf("nvidia-smi: unexpected line %q", line)
}
pid, err := strconv.Atoi(strings.TrimSpace(pidStr))
if err != nil {
return nil, fmt.Errorf("nvidia-smi: unexpected pid in %q", line)
}
mem, err := strconv.Atoi(strings.TrimSpace(memStr))
if err != nil {
continue // "N/A" and friends: unsupported under WDDM
}
apps = append(apps, computeApp{PID: pid, UsedMB: mem})
}
return apps, nil
}
+160
View File
@@ -0,0 +1,160 @@
package game
import (
"runtime"
"slices"
"testing"
)
func TestNormName(t *testing.T) {
for in, want := range map[string]string{
"Cyberpunk2077.exe": "cyberpunk2077",
"ollama": "ollama",
"OLLAMA APP.EXE": "ollama app",
" python.exe ": "python",
"hl2": "hl2",
} {
if got := normName(in); got != want {
t.Errorf("normName(%q) = %q, want %q", in, got, want)
}
}
}
func TestParseComputeApps(t *testing.T) {
apps, err := parseComputeApps("1234, 512\n 42 , 8192 \n\n")
if err != nil {
t.Fatal(err)
}
want := []computeApp{{PID: 1234, UsedMB: 512}, {PID: 42, UsedMB: 8192}}
if !slices.Equal(apps, want) {
t.Errorf("got %+v, want %+v", apps, want)
}
// WDDM reports "N/A" for memory; those rows are skipped, not fatal.
apps, err = parseComputeApps("1234, N/A\n42, 1024\n")
if err != nil {
t.Fatal(err)
}
if !slices.Equal(apps, []computeApp{{PID: 42, UsedMB: 1024}}) {
t.Errorf("got %+v", apps)
}
if _, err := parseComputeApps("garbage\n"); err == nil {
t.Error("expected an error for a malformed line")
}
}
func TestDetect(t *testing.T) {
d := New([]string{"Cyberpunk2077.exe", "hl2"}, 1024, 0,
[]string{"ollama", "python", "pythonw"}, nil)
ps := []Process{
{PID: 10, Name: "ollama.exe"},
{PID: 20, Name: "python.exe"},
{PID: 30, Name: "cyberpunk2077.exe"},
}
apps := []computeApp{
{PID: 10, UsedMB: 8192}, // ignored: ollama
{PID: 20, UsedMB: 4096}, // ignored: python (ComfyUI)
{PID: 40, UsedMB: 2048}, // foreign, above threshold
{PID: 50, UsedMB: 100}, // foreign but below threshold
}
holders := d.detect(ps, apps, nil)
if len(holders) != 2 {
t.Fatalf("got %v, want 2 holders", holders)
}
if holders[0] != "cyberpunk2077.exe (pid 30)" {
t.Errorf("holders[0] = %q", holders[0])
}
if holders[1] != "unknown process (pid 40) using 2048 MiB VRAM" {
t.Errorf("holders[1] = %q", holders[1])
}
}
func TestDetectEngineUtil(t *testing.T) {
d := New(nil, 0, 30, []string{"dwm", "python"}, nil)
ps := []Process{
{PID: 10, Name: "dwm.exe"},
{PID: 20, Name: "game.exe"},
{PID: 40, Name: "browser.exe"},
}
utils := map[int]float64{
10: 45, // ignored: dwm
20: 61, // foreign, above threshold
30: 82, // foreign, unknown name
40: 5, // below threshold
}
holders := d.detect(ps, nil, utils)
want := []string{
"game.exe (pid 20) using 61% GPU",
"unknown process (pid 30) using 82% GPU",
}
if !slices.Equal(holders, want) {
t.Errorf("got %v, want %v", holders, want)
}
}
func TestDetectNothingConfigured(t *testing.T) {
d := New(nil, 0, 0, nil, nil)
if got := d.detect([]Process{{PID: 1, Name: "game.exe"}}, nil, nil); len(got) != 0 {
t.Errorf("got %v, want none", got)
}
}
func TestParseGPUEngineInstance(t *testing.T) {
pid, eng, ok := parseGPUEngineInstance("pid_1234_luid_0x00000000_0x00011A2B_phys_0_eng_0_engtype_3D")
if !ok || pid != 1234 || eng != "3D" {
t.Errorf("got %d, %q, %v", pid, eng, ok)
}
pid, eng, ok = parseGPUEngineInstance("pid_42_luid_0x0_0x0_phys_0_eng_1_engtype_Copy")
if !ok || pid != 42 || eng != "Copy" {
t.Errorf("got %d, %q, %v", pid, eng, ok)
}
for _, bad := range []string{"", "something", "pid_", "pid_x_luid", "pid_-1_luid_0"} {
if _, _, ok := parseGPUEngineInstance(bad); ok {
t.Errorf("%q parsed, want failure", bad)
}
}
}
func TestParseGPUStats(t *testing.T) {
st, err := parseGPUStats("4300, 16384, 55, 42\n")
if err != nil {
t.Fatal(err)
}
if st.UsedMB != 4300 || st.TotalMB != 16384 || st.TempC != 55 || st.FanPct != 42 {
t.Errorf("got %+v", st)
}
// Cards that don't expose temperature/fan report N/A.
st, err = parseGPUStats("1024, 16384, N/A, N/A")
if err != nil {
t.Fatal(err)
}
if st.TempC != -1 || st.FanPct != -1 {
t.Errorf("got %+v, want -1 for N/A fields", st)
}
for _, bad := range []string{"", "1, 2", "x, 16384, 55, 42", "1024, x, 55, 42", "1, 2, 3, 4, 5"} {
if _, err := parseGPUStats(bad); err == nil {
t.Errorf("%q parsed, want failure", bad)
}
}
}
func TestProcessesLive(t *testing.T) {
if runtime.GOOS != "windows" && runtime.GOOS != "linux" {
t.Skip("no process listing on this platform")
}
ps, err := processes()
if err != nil {
t.Fatal(err)
}
if len(ps) == 0 {
t.Fatal("no processes listed")
}
for _, p := range ps {
if p.Name == "" {
t.Errorf("pid %d has an empty name", p.PID)
}
}
}
+36
View File
@@ -0,0 +1,36 @@
package game
import (
"errors"
"strconv"
"strings"
)
// errNotPrimed marks the first PDH sample after opening a query: rate-based
// counters (like engine utilization) need two collections before they
// return meaningful values.
var errNotPrimed = errors.New("GPU engine counter needs a second sample")
// parseGPUEngineInstance splits a PDH "GPU Engine" instance name —
// "pid_1234_luid_0x00000000_0x00011A2B_phys_0_eng_0_engtype_3D" — into PID
// and engine type ("3D", "Copy", "VideoDecode", ...). engType is empty when
// the name carries no engtype marker.
func parseGPUEngineInstance(name string) (pid int, engType string, ok bool) {
rest, found := strings.CutPrefix(name, "pid_")
if !found {
return 0, "", false
}
digits, rest, found := strings.Cut(rest, "_")
if !found {
return 0, "", false
}
pid, err := strconv.Atoi(digits)
if err != nil || pid < 0 {
return 0, "", false
}
const marker = "engtype_"
if i := strings.LastIndex(rest, marker); i >= 0 {
engType = rest[i+len(marker):]
}
return pid, engType, true
}
+16
View File
@@ -0,0 +1,16 @@
//go:build !windows
package game
import "errors"
// errNoEngineCounters marks platforms without per-process GPU engine
// counters (the PDH path is Windows-only).
var errNoEngineCounters = errors.New("per-process GPU engine counters are only available on Windows")
// gpuEngineSampler is a stub on non-Windows platforms.
type gpuEngineSampler struct{}
func openGPUEngineSampler() (*gpuEngineSampler, error) { return nil, errNoEngineCounters }
func (s *gpuEngineSampler) sample() (map[int]float64, error) { return nil, errNoEngineCounters }
+111
View File
@@ -0,0 +1,111 @@
//go:build windows
package game
import (
"fmt"
"unsafe"
"golang.org/x/sys/windows"
)
// Per-process GPU engine utilization via PDH — the same counters Task
// Manager's "GPU engine" columns read. Unlike nvidia-smi's compute-apps
// this covers graphics work under WDDM, so games show up. The counter is
// added with PdhAddEnglishCounterW, which is independent of the Windows
// display language.
var (
pdhDLL = windows.NewLazySystemDLL("pdh.dll")
procPdhOpenQuery = pdhDLL.NewProc("PdhOpenQueryW")
procPdhAddEnglishCounter = pdhDLL.NewProc("PdhAddEnglishCounterW")
procPdhCollectQueryData = pdhDLL.NewProc("PdhCollectQueryData")
procPdhGetFormattedCounterArray = pdhDLL.NewProc("PdhGetFormattedCounterArrayW")
procPdhCloseQuery = pdhDLL.NewProc("PdhCloseQuery")
)
const (
pdhFmtDouble = 0x00000200 // PDH_FMT_DOUBLE
pdhMoreData = 0x800007D2 // PDH_MORE_DATA
)
// pdhCountervalueItem mirrors PDH_FMT_COUNTERVALUE_ITEM (64-bit, double).
type pdhCountervalueItem struct {
name *uint16
cStatus uint32
_ uint32 // alignment padding
value float64
}
// gpuEngineSampler holds an open PDH query on the wildcard GPU Engine
// utilization counter. Keeping the query open across polls is what makes
// the rate-based values meaningful; the sampler lives as long as the
// process (PdhCloseQuery would only matter on unload).
type gpuEngineSampler struct {
query uintptr // PDH_HQUERY
counter uintptr // PDH_HCOUNTER
primed bool
}
// openGPUEngineSampler opens a query on the per-process GPU engine
// utilization counter (all instances).
func openGPUEngineSampler() (*gpuEngineSampler, error) {
var q uintptr
if r, _, _ := procPdhOpenQuery.Call(0, 0, uintptr(unsafe.Pointer(&q))); r != 0 {
return nil, fmt.Errorf("PdhOpenQuery: status %#x", r)
}
path, err := windows.UTF16PtrFromString(`\GPU Engine(*)\Utilization Percentage`)
if err != nil {
procPdhCloseQuery.Call(q)
return nil, err
}
var c uintptr
if r, _, _ := procPdhAddEnglishCounter.Call(q, uintptr(unsafe.Pointer(path)), 0, uintptr(unsafe.Pointer(&c))); r != 0 {
procPdhCloseQuery.Call(q)
return nil, fmt.Errorf("PdhAddEnglishCounter: status %#x", r)
}
return &gpuEngineSampler{query: q, counter: c}, nil
}
// sample collects the counter once and returns per-PID 3D-engine
// utilization in percent. The first call after open only primes the rate
// calculation and returns errNotPrimed. Processes can drive several 3D
// engines; their values are summed.
func (s *gpuEngineSampler) sample() (map[int]float64, error) {
if r, _, _ := procPdhCollectQueryData.Call(s.query); r != 0 {
return nil, fmt.Errorf("PdhCollectQueryData: status %#x", r)
}
if !s.primed {
s.primed = true
return nil, errNotPrimed
}
var size, count uint32
r, _, _ := procPdhGetFormattedCounterArray.Call(s.counter, pdhFmtDouble,
uintptr(unsafe.Pointer(&size)), uintptr(unsafe.Pointer(&count)), 0)
if r == pdhMoreData && size == 0 {
return nil, nil // no GPU engine instances at all
}
if r != pdhMoreData {
return nil, fmt.Errorf("PdhGetFormattedCounterArray(size): status %#x", r)
}
buf := make([]byte, size)
r, _, _ = procPdhGetFormattedCounterArray.Call(s.counter, pdhFmtDouble,
uintptr(unsafe.Pointer(&size)), uintptr(unsafe.Pointer(&count)),
uintptr(unsafe.Pointer(&buf[0])))
if r != 0 {
return nil, fmt.Errorf("PdhGetFormattedCounterArray: status %#x", r)
}
items := unsafe.Slice((*pdhCountervalueItem)(unsafe.Pointer(&buf[0])), int(count))
out := make(map[int]float64)
for i := range items {
if items[i].cStatus != 0 || items[i].name == nil {
continue
}
pid, engType, ok := parseGPUEngineInstance(windows.UTF16PtrToString(items[i].name))
if !ok || engType != "3D" {
continue // only the 3D engine marks game-like work
}
out[pid] += items[i].value
}
return out, nil
}
+33
View File
@@ -0,0 +1,33 @@
//go:build windows
package game
import (
"errors"
"testing"
"time"
)
// TestGPUEngineSamplerLive opens the real PDH query and takes two samples;
// the first only primes the rate counters. Skipped (not failed) when the
// machine has no GPU counters.
func TestGPUEngineSamplerLive(t *testing.T) {
s, err := openGPUEngineSampler()
if err != nil {
t.Skipf("no GPU engine counters: %v", err)
}
if _, err := s.sample(); !errors.Is(err, errNotPrimed) {
t.Fatalf("first sample: err = %v, want errNotPrimed", err)
}
time.Sleep(200 * time.Millisecond)
utils, err := s.sample()
if err != nil {
t.Fatalf("second sample: %v", err)
}
for pid, util := range utils {
if pid < 0 || util < 0 {
t.Errorf("pid %d: util %.2f", pid, util)
}
}
t.Logf("%d processes with 3D-engine usage", len(utils))
}
+30
View File
@@ -0,0 +1,30 @@
//go:build linux
package game
import (
"os"
"strconv"
"strings"
)
// processes lists the running processes from /proc/<pid>/comm.
func processes() ([]Process, error) {
entries, err := os.ReadDir("/proc")
if err != nil {
return nil, err
}
var ps []Process
for _, e := range entries {
pid, err := strconv.Atoi(e.Name())
if err != nil {
continue
}
comm, err := os.ReadFile("/proc/" + e.Name() + "/comm")
if err != nil {
continue // process vanished mid-walk
}
ps = append(ps, Process{PID: pid, Name: strings.TrimSpace(string(comm))})
}
return ps, nil
}
+7
View File
@@ -0,0 +1,7 @@
//go:build !windows && !linux
package game
// processes is unsupported on this platform; the process watch list never
// matches and the nvidia-smi path reports PIDs without names.
func processes() ([]Process, error) { return nil, nil }
+35
View File
@@ -0,0 +1,35 @@
//go:build windows
package game
import (
"unsafe"
"golang.org/x/sys/windows"
)
// processes lists the running processes via the Toolhelp32 snapshot API.
func processes() ([]Process, error) {
h, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
if err != nil {
return nil, err
}
defer windows.CloseHandle(h) //nolint:errcheck // best effort
var entry windows.ProcessEntry32
entry.Size = uint32(unsafe.Sizeof(entry))
if err := windows.Process32First(h, &entry); err != nil {
return nil, err
}
var ps []Process
for {
ps = append(ps, Process{
PID: int(entry.ProcessID),
Name: windows.UTF16ToString(entry.ExeFile[:]),
})
if err := windows.Process32Next(h, &entry); err != nil {
break // ERROR_NO_MORE_FILES ends the walk
}
}
return ps, nil
}
+114 -6
View File
@@ -1,12 +1,15 @@
// Package lock implements the two-mode GPU arbitration lock: any number of // Package lock implements the two-mode GPU arbitration lock: any number of
// concurrent LLM requests ("readers") or exactly one image job ("writer"), // concurrent LLM requests ("readers") or exactly one image job ("writer"),
// with image jobs taking priority over newly arriving LLM requests. // with image jobs taking priority over newly arriving LLM requests. An
// external hold (SetExternal) blocks new grants of both kinds while a
// foreign process — e.g. a game — holds the GPU; in-flight work drains.
package lock package lock
import ( import (
"context" "context"
"log/slog" "log/slog"
"sync" "sync"
"time"
) )
// State is the current GPU occupancy state. // State is the current GPU occupancy state.
@@ -16,6 +19,7 @@ const (
StateIdle State = "idle" StateIdle State = "idle"
StateLLM State = "llm" StateLLM State = "llm"
StateImage State = "image" StateImage State = "image"
StateExternal State = "external"
) )
type imageWaiter struct{ id uint64 } type imageWaiter struct{ id uint64 }
@@ -27,9 +31,13 @@ type Lock struct {
change chan struct{} // closed and replaced on every state change change chan struct{} // closed and replaced on every state change
n int // LLM requests in flight n int // LLM requests in flight
llmWaiting int // LLM requests blocked waiting for the GPU
imageActive bool // an image job holds the GPU imageActive bool // an image job holds the GPU
imageQ []imageWaiter imageQ []imageWaiter
nextID uint64 nextID uint64
external string // non-empty: a foreign process (e.g. a game) holds the GPU
detail string // what the current holder is doing (best effort)
since time.Time // when the current state began
log *slog.Logger log *slog.Logger
} }
@@ -37,7 +45,7 @@ type Lock struct {
// New returns a ready-to-use Lock. log may be nil; if set, every state // New returns a ready-to-use Lock. log may be nil; if set, every state
// transition is logged at debug level. // transition is logged at debug level.
func New(log *slog.Logger) *Lock { func New(log *slog.Logger) *Lock {
return &Lock{change: make(chan struct{}), log: log} return &Lock{change: make(chan struct{}), log: log, since: time.Now()}
} }
// broadcast wakes all waiters. Call with mu held. // broadcast wakes all waiters. Call with mu held.
@@ -52,22 +60,67 @@ func (l *Lock) logTransition(msg string, args ...any) {
} }
} }
// SetExternal records that a process outside gpu-turnstile's control (a
// game, another ML job) holds the GPU: new LLM and image grants block until
// ClearExternal. In-flight work is not preempted. holder describes the
// process for logs and busy responses.
func (l *Lock) SetExternal(holder string) {
l.mu.Lock()
l.external = holder
l.since = time.Now()
l.broadcast()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateExternal, "holder", holder)
}
// ClearExternal lifts the external hold; waiting LLM and image requests
// proceed.
func (l *Lock) ClearExternal() {
l.mu.Lock()
l.external = ""
l.since = time.Now()
l.broadcast()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateIdle)
}
// External returns the current external holder, or "" when none.
func (l *Lock) External() string {
l.mu.Lock()
defer l.mu.Unlock()
return l.external
}
// AcquireLLM blocks until no image job is active or pending, then registers // AcquireLLM blocks until no image job is active or pending, then registers
// one in-flight LLM request. Returns ctx.Err() if the context is cancelled // one in-flight LLM request. Returns ctx.Err() if the context is cancelled
// while waiting; no state is changed in that case. // while waiting; no state is changed in that case.
func (l *Lock) AcquireLLM(ctx context.Context) error { func (l *Lock) AcquireLLM(ctx context.Context) error {
l.mu.Lock() l.mu.Lock()
for l.imageActive || len(l.imageQ) > 0 { waiting := false
for l.imageActive || len(l.imageQ) > 0 || l.external != "" {
if !waiting {
l.llmWaiting++
waiting = true
}
ch := l.change ch := l.change
l.mu.Unlock() l.mu.Unlock()
select { select {
case <-ctx.Done(): case <-ctx.Done():
l.mu.Lock()
l.llmWaiting--
l.mu.Unlock()
return ctx.Err() return ctx.Err()
case <-ch: case <-ch:
} }
l.mu.Lock() l.mu.Lock()
} }
if waiting {
l.llmWaiting--
}
l.n++ l.n++
if l.n == 1 {
l.since = time.Now()
}
n := l.n n := l.n
l.mu.Unlock() l.mu.Unlock()
l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n) l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
@@ -76,10 +129,10 @@ func (l *Lock) AcquireLLM(ctx context.Context) error {
// TryAcquireLLM acquires one in-flight LLM slot without waiting and // TryAcquireLLM acquires one in-flight LLM slot without waiting and
// reports whether it succeeded. It fails when an image job is active or // reports whether it succeeded. It fails when an image job is active or
// pending. // pending or an external hold is set.
func (l *Lock) TryAcquireLLM() bool { func (l *Lock) TryAcquireLLM() bool {
l.mu.Lock() l.mu.Lock()
if l.imageActive || len(l.imageQ) > 0 { if l.imageActive || len(l.imageQ) > 0 || l.external != "" {
l.mu.Unlock() l.mu.Unlock()
return false return false
} }
@@ -96,6 +149,8 @@ func (l *Lock) ReleaseLLM() {
l.n-- l.n--
n := l.n n := l.n
if l.n == 0 { if l.n == 0 {
l.since = time.Now()
l.detail = ""
l.broadcast() l.broadcast()
} }
l.mu.Unlock() l.mu.Unlock()
@@ -120,9 +175,10 @@ func (l *Lock) AcquireImage(ctx context.Context) error {
for { for {
l.mu.Lock() l.mu.Lock()
if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive { if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive && l.external == "" {
l.imageQ = l.imageQ[1:] l.imageQ = l.imageQ[1:]
l.imageActive = true l.imageActive = true
l.since = time.Now()
l.mu.Unlock() l.mu.Unlock()
l.logTransition("lock transition", "state", StateImage) l.logTransition("lock transition", "state", StateImage)
return nil return nil
@@ -151,6 +207,8 @@ func (l *Lock) AcquireImage(ctx context.Context) error {
func (l *Lock) ReleaseImage() { func (l *Lock) ReleaseImage() {
l.mu.Lock() l.mu.Lock()
l.imageActive = false l.imageActive = false
l.since = time.Now()
l.detail = ""
l.broadcast() l.broadcast()
l.mu.Unlock() l.mu.Unlock()
l.logTransition("lock transition", "state", StateIdle) l.logTransition("lock transition", "state", StateIdle)
@@ -166,8 +224,58 @@ func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) {
state = StateImage state = StateImage
case l.n > 0: case l.n > 0:
state = StateLLM state = StateLLM
case l.external != "":
state = StateExternal
default: default:
state = StateIdle state = StateIdle
} }
return state, l.n, l.imageActive || len(l.imageQ) > 0 return state, l.n, l.imageActive || len(l.imageQ) > 0
} }
// SetDetail records what the current holder is doing (e.g. the request
// path), for status displays. Best effort: overwritten by each new holder,
// cleared when the GPU goes idle.
func (l *Lock) SetDetail(detail string) {
l.mu.Lock()
l.detail = detail
l.mu.Unlock()
}
// Status is a point-in-time view of the lock for monitoring.
type Status struct {
State State
Detail string
LLMInflight int
LLMWaiting int
ImageActive bool
ImageQueue int
External string
Since time.Time
}
// Status reports the full lock state, including waiters and how long the
// current state has held.
func (l *Lock) Status() Status {
l.mu.Lock()
defer l.mu.Unlock()
s := Status{
Detail: l.detail,
LLMInflight: l.n,
LLMWaiting: l.llmWaiting,
ImageActive: l.imageActive,
ImageQueue: len(l.imageQ),
External: l.external,
Since: l.since,
}
switch {
case l.imageActive:
s.State = StateImage
case l.n > 0:
s.State = StateLLM
case l.external != "":
s.State = StateExternal
default:
s.State = StateIdle
}
return s
}
+67
View File
@@ -226,3 +226,70 @@ func TestRace(t *testing.T) {
t.Fatalf("leaked lock state: state=%s n=%d pending=%v", state, n, pending) t.Fatalf("leaked lock state: state=%s n=%d pending=%v", state, n, pending)
} }
} }
func TestExternalHoldBlocksBoth(t *testing.T) {
lk := New(nil)
ctx := context.Background()
lk.SetExternal("game.exe (pid 42)")
if lk.TryAcquireLLM() {
t.Fatal("TryAcquireLLM succeeded during external hold")
}
if got := lk.External(); got != "game.exe (pid 42)" {
t.Fatalf("External() = %q", got)
}
if state, _, _ := lk.Snapshot(); state != StateExternal {
t.Fatalf("state=%s, want external", state)
}
llmAcquired := make(chan struct{})
go func() {
if err := lk.AcquireLLM(ctx); err != nil {
t.Error(err)
}
close(llmAcquired)
}()
assertBlocked(t, llmAcquired, "LLM acquire during external hold")
imageAcquired := make(chan struct{})
go func() {
if err := lk.AcquireImage(ctx); err != nil {
t.Error(err)
}
close(imageAcquired)
}()
assertBlocked(t, imageAcquired, "image acquire during external hold")
lk.ClearExternal()
// The queued image job wins over the LLM waiter (image priority).
waitFor(t, imageAcquired, "image acquire after ClearExternal")
assertBlocked(t, llmAcquired, "LLM acquire while image active")
lk.ReleaseImage()
waitFor(t, llmAcquired, "LLM acquire after image release")
lk.ReleaseLLM()
}
func TestExternalHoldDoesNotPreempt(t *testing.T) {
lk := New(nil)
ctx := context.Background()
if err := lk.AcquireLLM(ctx); err != nil {
t.Fatal(err)
}
lk.SetExternal("game.exe")
// In-flight LLM work keeps the llm state; the hold blocks new grants.
if state, _, _ := lk.Snapshot(); state != StateLLM {
t.Fatalf("state=%s, want llm while work in flight", state)
}
if lk.TryAcquireLLM() {
t.Fatal("TryAcquireLLM succeeded during external hold")
}
lk.ReleaseLLM()
if state, _, _ := lk.Snapshot(); state != StateExternal {
t.Fatalf("state=%s, want external after drain", state)
}
lk.ClearExternal()
if state, _, _ := lk.Snapshot(); state != StateIdle {
t.Fatalf("state=%s, want idle", state)
}
}
+1 -1
View File
@@ -89,7 +89,7 @@ func (m *Metrics) Render(w io.Writer, state string, llmInflight int, imagePendin
fmt.Fprint(w, `# HELP gpu_turnstile_state Current GPU state (1 for the active state). fmt.Fprint(w, `# HELP gpu_turnstile_state Current GPU state (1 for the active state).
# TYPE gpu_turnstile_state gauge # TYPE gpu_turnstile_state gauge
`) `)
for _, s := range []string{"idle", "llm", "image"} { for _, s := range []string{"idle", "llm", "image", "external"} {
v := 0 v := 0
if s == state { if s == state {
v = 1 v = 1
+28 -7
View File
@@ -62,11 +62,19 @@ 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
} }
+64 -3
View File
@@ -26,6 +26,7 @@ import (
"gpu-turnstile/internal/lock" "gpu-turnstile/internal/lock"
"gpu-turnstile/internal/metrics" "gpu-turnstile/internal/metrics"
"gpu-turnstile/internal/ollama" "gpu-turnstile/internal/ollama"
"gpu-turnstile/internal/supervise"
) )
// defaultCaptureLimit bounds how much of a /prompt response body is // defaultCaptureLimit bounds how much of a /prompt response body is
@@ -44,6 +45,11 @@ type Config struct {
Metrics *metrics.Metrics Metrics *metrics.Metrics
Log *slog.Logger Log *slog.Logger
// ComfySup, when non-nil, is the managed ComfyUI process: any ComfyUI
// request starts it on demand, /prompt additionally waits for
// readiness before taking the GPU lock. nil = unmanaged upstream.
ComfySup *supervise.Process
// LogColor enables ANSI colors in per-request log lines. Ignored when // LogColor enables ANSI colors in per-request log lines. Ignored when
// the log level is above INFO (request lines are not emitted at all). // the log level is above INFO (request lines are not emitted at all).
LogColor bool LogColor bool
@@ -438,7 +444,7 @@ func isLLMRequest(r *http.Request) bool {
return r.Method == http.MethodPost && llmPaths[r.URL.Path] return r.Method == http.MethodPost && llmPaths[r.URL.Path]
} }
// OllamaHandler serves the Ollama-facing listener. // OllamaHandler serves the listener for Ollama-compatible clients.
func (s *Server) OllamaHandler() http.Handler { func (s *Server) OllamaHandler() http.Handler {
return s.logRequests("ollama", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { return s.logRequests("ollama", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path { switch r.URL.Path {
@@ -458,10 +464,14 @@ func (s *Server) OllamaHandler() http.Handler {
if s.busyMode == "reject" { if s.busyMode == "reject" {
if !s.cfg.Lock.TryAcquireLLM() { if !s.cfg.Lock.TryAcquireLLM() {
s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds()) s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds())
msg := "GPU busy: image job active or queued"
if holder := s.cfg.Lock.External(); holder != "" {
msg = "GPU busy: " + holder
}
s.log.Info("llm request rejected; GPU busy", s.log.Info("llm request rejected; GPU busy",
"path", r.URL.Path, "status", s.busyStatus) "path", r.URL.Path, "status", s.busyStatus)
w.Header().Set("Retry-After", strconv.Itoa(s.busyRetryAfter)) w.Header().Set("Retry-After", strconv.Itoa(s.busyRetryAfter))
http.Error(w, "GPU busy: image job active or queued", s.busyStatus) http.Error(w, msg, s.busyStatus)
return return
} }
s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds()) s.cfg.Metrics.ObserveLockWait("llm", time.Since(start).Seconds())
@@ -481,11 +491,12 @@ func (s *Server) OllamaHandler() http.Handler {
return return
} }
defer s.cfg.Lock.ReleaseLLM() defer s.cfg.Lock.ReleaseLLM()
s.cfg.Lock.SetDetail("ollama: " + r.Method + " " + r.URL.Path)
s.ollamaProxy.ServeHTTP(w, r) s.ollamaProxy.ServeHTTP(w, r)
})) }))
} }
// ComfyHandler serves the ComfyUI-facing listener. // ComfyHandler serves the listener for ComfyUI clients.
func (s *Server) ComfyHandler() http.Handler { func (s *Server) ComfyHandler() http.Handler {
return s.logRequests("comfy", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { return s.logRequests("comfy", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path { switch r.URL.Path {
@@ -500,6 +511,21 @@ func (s *Server) ComfyHandler() http.Handler {
s.handlePrompt(w, r) s.handlePrompt(w, r)
return return
} }
if s.cfg.ComfySup != nil {
// Any other ComfyUI request also wakes the managed server; the
// retry backoff bridges the time it needs to come up. While a
// foreign process holds the GPU we refuse to spawn it — the
// request gets the busy answer instead of fighting for VRAM.
if holder := s.cfg.Lock.External(); holder != "" && !s.cfg.ComfySup.Running() {
w.Header().Set("Retry-After", strconv.Itoa(s.busyRetryAfter))
http.Error(w, "GPU busy: "+holder, http.StatusServiceUnavailable)
return
}
if err := s.cfg.ComfySup.EnsureRunning(); err != nil {
http.Error(w, fmt.Sprintf("cannot start ComfyUI: %v", err), http.StatusBadGateway)
return
}
}
s.comfyProxy.ServeHTTP(w, r) s.comfyProxy.ServeHTTP(w, r)
})) }))
} }
@@ -539,6 +565,23 @@ func (w *captureWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) { func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
log := s.log.With("op", "image") log := s.log.With("op", "image")
// Normally the managed server is brought up *before* taking the GPU
// lock: torch can take a minute to load, and LLM traffic should keep
// flowing in the meantime. While a foreign process holds the GPU, LLM
// traffic is blocked anyway and a fresh ComfyUI would fight it for
// VRAM — so the lock comes first in that case.
comfyFirst := s.cfg.ComfySup != nil && s.cfg.Lock.External() == ""
if comfyFirst {
if err := s.cfg.ComfySup.EnsureRunning(); err != nil {
http.Error(w, fmt.Sprintf("cannot start ComfyUI: %v", err), http.StatusBadGateway)
return
}
if err := s.cfg.ComfySup.WaitReady(r.Context()); err != nil {
http.Error(w, fmt.Sprintf("ComfyUI did not become ready: %v", err), http.StatusBadGateway)
return
}
}
start := time.Now() start := time.Now()
if err := s.cfg.Lock.AcquireImage(r.Context()); err != nil { if err := s.cfg.Lock.AcquireImage(r.Context()); err != nil {
if errors.Is(err, context.DeadlineExceeded) { if errors.Is(err, context.DeadlineExceeded) {
@@ -548,6 +591,20 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
} }
s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds()) s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds())
log.Info("image lock acquired") log.Info("image lock acquired")
s.cfg.Lock.SetDetail("comfy: POST /prompt")
if s.cfg.ComfySup != nil && !comfyFirst {
if err := s.cfg.ComfySup.EnsureRunning(); err != nil {
s.cfg.Lock.ReleaseImage()
http.Error(w, fmt.Sprintf("cannot start ComfyUI: %v", err), http.StatusBadGateway)
return
}
if err := s.cfg.ComfySup.WaitReady(r.Context()); err != nil {
s.cfg.Lock.ReleaseImage()
http.Error(w, fmt.Sprintf("ComfyUI did not become ready: %v", err), http.StatusBadGateway)
return
}
}
if s.cfg.Ollama != nil { if s.cfg.Ollama != nil {
uctx, ucancel := context.WithTimeout(r.Context(), s.cfg.UnloadTimeout) uctx, ucancel := context.WithTimeout(r.Context(), s.cfg.UnloadTimeout)
@@ -608,6 +665,10 @@ func (s *Server) finishImageJob(promptID string) {
s.cfg.Lock.ReleaseImage() s.cfg.Lock.ReleaseImage()
log.Info("image lock released") log.Info("image lock released")
if s.cfg.ComfySup != nil {
// The idle clock starts when the job ends, not when it began.
s.cfg.ComfySup.NoteActivity()
}
if s.cfg.Ollama != nil && s.cfg.WarmModel != "" { if s.cfg.Ollama != nil && s.cfg.WarmModel != "" {
if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle { if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle {
+2 -2
View File
@@ -46,8 +46,8 @@ type fakes struct {
rec *recorder rec *recorder
ollama *httptest.Server ollama *httptest.Server
comfy *httptest.Server comfy *httptest.Server
server *httptest.Server // Ollama-facing gpu-turnstile listener server *httptest.Server // gpu-turnstile listener for Ollama-compatible clients
comfySrv *httptest.Server // ComfyUI-facing gpu-turnstile listener comfySrv *httptest.Server // gpu-turnstile listener for ComfyUI clients
freeCh chan struct{} freeCh chan struct{}
chatCh chan struct{} chatCh chan struct{}
historyMu sync.Mutex historyMu sync.Mutex
+63
View File
@@ -0,0 +1,63 @@
// Shared env-file handling for the installers (Windows and Linux).
package service
import (
"fmt"
"os"
"strings"
"gpu-turnstile/internal/config"
)
// syncEnvFile ensures the installed env file at path is a current
// installer-written sample. A missing file is created; a file without a
// CFG_VER line (hand-written or from before versioning) is invalid and
// replaced by a fresh sample after a .bak backup; a file written by an
// older version gets newly added settings appended via config.SyncSample.
// logPath activates the LOG_FILE line (Windows); empty leaves it commented.
func syncEnvFile(path, version, logPath string) error {
data, err := os.ReadFile(path)
switch {
case os.IsNotExist(err):
return writeEnvFile(path, config.SampleEnv(version, logPath))
case err != nil:
return fmt.Errorf("read %s: %w", path, err)
}
values, perr := config.ParseEnvFile(strings.NewReader(string(data)))
if perr == nil && values["CFG_VER"] != "" {
synced, changed := config.SyncSample(string(data), version, logPath)
if !changed {
return nil
}
return writeEnvFile(path, synced)
}
// No readable CFG_VER: the file is invalid — keep a backup and start
// from a fresh sample.
if err := os.WriteFile(path+".bak", data, 0o644); err != nil {
return fmt.Errorf("back up %s: %w", path, err)
}
return writeEnvFile(path, config.SampleEnv(version, logPath))
}
func writeEnvFile(path, content string) error {
if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
return fmt.Errorf("write %s: %w", path, err)
}
return nil
}
// configuredValue reads one key from the env file the service will load, so
// the installers can adapt the sandbox to it (ACL grants, unit directives).
// "" when unset or unreadable.
func configuredValue(configPath, key string) string {
f, err := os.Open(configPath)
if err != nil {
return ""
}
defer f.Close()
values, err := config.ParseEnvFile(f)
if err != nil {
return ""
}
return values[key]
}
+8
View File
@@ -0,0 +1,8 @@
package service
import "errors"
// ErrUserCancelled is returned by RelaunchElevated when the user declines
// the UAC prompt. Windows-only in practice; defined here so cross-platform
// callers can compare against it.
var ErrUserCancelled = errors.New("UAC prompt declined")
+203 -13
View File
@@ -2,17 +2,21 @@
// Package service integrates gpu-turnstile with systemd on Linux: running // Package service integrates gpu-turnstile with systemd on Linux: running
// under a unit with readiness notification and watchdog, plus // under a unit with readiness notification and watchdog, plus
// install/remove helpers that manage a system unit. // install/remove helpers that manage a hardened system unit.
package service package service
import ( import (
"bytes"
"context" "context"
"fmt" "fmt"
"io"
"os" "os"
"os/exec" "os/exec"
"os/signal" "os/signal"
"path/filepath" "path/filepath"
"syscall" "syscall"
"gpu-turnstile/internal/supervise"
) )
// Name matches the Windows service name; the systemd unit is Name + ".service". // Name matches the Windows service name; the systemd unit is Name + ".service".
@@ -21,9 +25,26 @@ const Name = "gpu-turnstile"
// unitPath is where Install writes the unit file. // unitPath is where Install writes the unit file.
const unitPath = "/etc/systemd/system/" + Name + ".service" const unitPath = "/etc/systemd/system/" + Name + ".service"
// stateDir holds the installed binary (and staged updates); the unit's
// StateDirectory= directive makes systemd own it and grant the dynamic
// user write access. etcConfig is the config file the unit loads.
const (
stateDir = "/var/lib/" + Name
etcConfig = "/etc/" + Name + ".env"
)
// IsService reports whether the process was started by systemd. // IsService reports whether the process was started by systemd.
func IsService() bool { return os.Getenv("INVOCATION_ID") != "" } func IsService() bool { return os.Getenv("INVOCATION_ID") != "" }
// Elevated is always true on Linux: there is no UAC equivalent; privilege
// errors surface from the failing operation with a "run as root" hint.
func Elevated() bool { return true }
// RelaunchElevated is a Windows-only concept (UAC).
func RelaunchElevated([]string) (int, error) {
return 0, fmt.Errorf("elevated relaunch is only supported on Windows")
}
// Run executes run with SIGINT/SIGTERM cancellation (which is how systemctl // Run executes run with SIGINT/SIGTERM cancellation (which is how systemctl
// stop signals the process) and tells systemd when the shutdown begins. // stop signals the process) and tells systemd when the shutdown begins.
func Run(run func(ctx context.Context) error) error { func Run(run func(ctx context.Context) error) error {
@@ -33,11 +54,33 @@ func Run(run func(ctx context.Context) error) error {
return run(ctx) return run(ctx)
} }
// renderUnit builds the systemd unit: Type=notify so systemctl start blocks // renderUnit builds the hardened systemd unit: Type=notify so systemctl
// until the listeners are bound, a 30s watchdog, and restart-on-failure // start blocks until the listeners are bound, a 30s watchdog, and
// with a 5s delay — which is also what brings up a staged update after the // restart-on-failure with a 5s delay — which is also what brings up a
// updater exits with a non-zero code. // staged update after the updater exits with a non-zero code.
func renderUnit(exePath, configPath string) string { //
// Sandboxing mirrors the Windows virtual account: DynamicUser=yes gives
// the service a transient per-service UID with no login and no home, the
// filesystem is read-only except StateDirectory (the install dir, so
// self-updates can rewrite the binary), and the usual no-privilege-escalation
// directives apply. The proxy needs nothing but outbound TCP/UDP and the
// notify socket, so it loses nothing. A managed ComfyUI (comfyDir) gets a
// BindPaths hole through ProtectHome/ProtectSystem: it reads its venv and
// writes output/temp/user data under COMFY_DIR. A venv whose base
// interpreter (pyvenv.cfg home) lives outside COMFY_DIR gets an additional
// read-only bind, and a Comfy-Desktop shared data dir (models, input,
// output) a read-write one.
func renderUnit(exePath, configPath, comfyDir string) string {
bind := ""
if comfyDir != "" {
bind = "BindPaths=" + comfyDir + "\n"
if home := comfyVenvHome(comfyDir); home != "" {
bind += "BindReadOnlyPaths=" + home + "\n"
}
if shared := supervise.DesktopSharedDir(comfyDir); shared != "" {
bind += "BindPaths=" + shared + "\n"
}
}
return fmt.Sprintf(`[Unit] return fmt.Sprintf(`[Unit]
Description=gpu-turnstile GPU arbitration proxy for Ollama and ComfyUI Description=gpu-turnstile GPU arbitration proxy for Ollama and ComfyUI
After=network-online.target After=network-online.target
@@ -50,14 +93,73 @@ ExecStart=%q -config %q
Restart=on-failure Restart=on-failure
RestartSec=5s RestartSec=5s
DynamicUser=yes
StateDirectory=%s
RuntimeDirectory=%s
RuntimeDirectoryMode=0755
%sProtectSystem=strict
ProtectHome=yes
PrivateTmp=yes
NoNewPrivileges=yes
ProtectKernelTunables=yes
ProtectKernelModules=yes
ProtectKernelLogs=yes
ProtectControlGroups=yes
ProtectClock=yes
RestrictNamespaces=yes
RestrictSUIDSGID=yes
RestrictRealtime=yes
LockPersonality=yes
MemoryDenyWriteExecute=yes
CapabilityBoundingSet=
AmbientCapabilities=
RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6
SystemCallFilter=@system-service
SystemCallErrorNumber=EPERM
[Install] [Install]
WantedBy=multi-user.target WantedBy=multi-user.target
`, exePath, configPath) `, exePath, configPath, Name, Name, bind)
} }
// Install writes the unit for the current executable and the given config // copyFile copies src to dst, creating dst with the given mode.
// file, then enables and starts it. Needs root. func copyFile(src, dst string, mode os.FileMode) error {
func Install(configPath string) error { in, err := os.Open(src)
if err != nil {
return err
}
defer in.Close()
out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
if err != nil {
return err
}
defer out.Close()
if _, err := io.Copy(out, in); err != nil {
return err
}
return out.Close()
}
// Install copies the current executable into /var/lib/gpu-turnstile, makes
// sure /etc/gpu-turnstile.env exists (copied from the given config file if
// provided), writes the hardened unit, then enables and starts it. With
// copyBin=false the current executable location and config path are
// registered as-is instead. When the config sets COMFY_DIR, the unit gets a
// BindPaths= for it so the sandboxed service can reach the managed ComfyUI
// even under /home. Needs root.
//
// Re-running install converges an existing unit instead of failing: it is
// stopped first if active, the installed binary is replaced only when the
// content differs, the unit file is rewritten (followed by daemon-reload)
// only when it changed, and the service is started again only if it was
// active before.
//
// The binary lives in the StateDirectory rather than /usr/local/sbin on
// purpose: replacing a running binary needs write access to its
// *directory*, and granting the sandboxed service user write access to a
// shared system directory would let a compromised service overwrite other
// binaries. /var/lib/gpu-turnstile is exclusively ours.
func Install(configPath string, copyBin bool, version string) error {
exe, err := os.Executable() exe, err := os.Executable()
if err != nil { if err != nil {
return err return err
@@ -65,26 +167,114 @@ func Install(configPath string) error {
if abs, absErr := filepath.Abs(exe); absErr == nil { if abs, absErr := filepath.Abs(exe); absErr == nil {
exe = abs exe = abs
} }
if err := os.WriteFile(unitPath, []byte(renderUnit(exe, configPath)), 0o644); err != nil { unit := Name + ".service"
_, statErr := os.Stat(unitPath)
fresh := os.IsNotExist(statErr)
wasRunning := exec.Command("systemctl", "is-active", "--quiet", unit).Run() == nil
if wasRunning {
if out, err := exec.Command("systemctl", "stop", unit).CombinedOutput(); err != nil {
return fmt.Errorf("systemctl stop (run as root): %w (%s)", err, out)
}
}
cfg := etcConfig
if copyBin {
if err := os.MkdirAll(stateDir, 0o755); err != nil {
return fmt.Errorf("create %s (run as root): %w", stateDir, err)
}
installedExe := filepath.Join(stateDir, Name)
if exe != installedExe {
if same, _ := sameFileContent(exe, installedExe); !same {
if err := copyFile(exe, installedExe, 0o755); err != nil {
return fmt.Errorf("install binary to %s: %w", installedExe, err)
}
}
}
exe = installedExe
if _, err := os.Stat(etcConfig); os.IsNotExist(err) && configPath != "" {
copyFile(configPath, etcConfig, 0o644) //nolint:errcheck // best effort
}
// Missing configs get a fully commented sample (LOG_FILE stays
// commented — stderr goes to the journal on Linux);
// installer-written ones from older versions get new settings
// appended; anything without CFG_VER is invalid and gets replaced
// (backup kept as .bak).
if err := syncEnvFile(etcConfig, version, ""); err != nil {
return err
}
} else if configPath != "" {
if abs, absErr := filepath.Abs(configPath); absErr == nil {
cfg = abs
}
}
rendered := renderUnit(exe, cfg, configuredValue(cfg, "COMFY_DIR"))
if old, _ := os.ReadFile(unitPath); string(old) != rendered {
if err := os.WriteFile(unitPath, []byte(rendered), 0o644); err != nil {
return fmt.Errorf("write %s (run as root): %w", unitPath, err) return fmt.Errorf("write %s (run as root): %w", unitPath, err)
} }
if out, err := exec.Command("systemctl", "daemon-reload").CombinedOutput(); err != nil { if out, err := exec.Command("systemctl", "daemon-reload").CombinedOutput(); err != nil {
return fmt.Errorf("systemctl daemon-reload: %w (%s)", err, out) return fmt.Errorf("systemctl daemon-reload: %w (%s)", err, out)
} }
if out, err := exec.Command("systemctl", "enable", "--now", Name+".service").CombinedOutput(); err != nil { }
if fresh {
if out, err := exec.Command("systemctl", "enable", "--now", unit).CombinedOutput(); err != nil {
return fmt.Errorf("systemctl enable --now: %w (%s)", err, out) return fmt.Errorf("systemctl enable --now: %w (%s)", err, out)
} }
return nil return nil
} }
if exec.Command("systemctl", "is-enabled", "--quiet", unit).Run() != nil {
if out, err := exec.Command("systemctl", "enable", unit).CombinedOutput(); err != nil {
return fmt.Errorf("systemctl enable: %w (%s)", err, out)
}
}
if wasRunning {
if out, err := exec.Command("systemctl", "start", unit).CombinedOutput(); err != nil {
return fmt.Errorf("systemctl start: %w (%s)", err, out)
}
}
return nil
}
// Remove stops and disables the service and deletes the unit file. // sameFileContent reports whether two files hold identical bytes. A missing
// destination is simply "different".
func sameFileContent(a, b string) (bool, error) {
ba, err := os.ReadFile(a)
if err != nil {
return false, err
}
bb, err := os.ReadFile(b)
if os.IsNotExist(err) {
return false, nil
}
if err != nil {
return false, err
}
return bytes.Equal(ba, bb), nil
}
// Remove stops and disables the service and deletes the unit file and the
// installed binary. The config file in /etc is left in place (user data).
func Remove() error { func Remove() error {
exec.Command("systemctl", "disable", "--now", Name+".service").Run() // ignore: may not exist exec.Command("systemctl", "disable", "--now", Name+".service").Run() // ignore: may not exist
if err := os.Remove(unitPath); err != nil && !os.IsNotExist(err) { if err := os.Remove(unitPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("remove %s: %w", unitPath, err) return fmt.Errorf("remove %s: %w", unitPath, err)
} }
os.RemoveAll(stateDir) // installed binary + staged updates; ignore error
if out, err := exec.Command("systemctl", "daemon-reload").CombinedOutput(); err != nil { if out, err := exec.Command("systemctl", "daemon-reload").CombinedOutput(); err != nil {
return fmt.Errorf("systemctl daemon-reload: %w (%s)", err, out) return fmt.Errorf("systemctl daemon-reload: %w (%s)", err, out)
} }
return nil return nil
} }
// RestartIfRunning restarts the systemd unit when it is active (used after
// a forced update staged a new binary). Reports whether a restart
// happened. An inactive or missing unit is not an error. Needs root.
func RestartIfRunning() (bool, error) {
if err := exec.Command("systemctl", "is-active", "--quiet", Name+".service").Run(); err != nil {
return false, nil // inactive or not installed
}
if out, err := exec.Command("systemctl", "restart", Name+".service").CombinedOutput(); err != nil {
return false, fmt.Errorf("systemctl restart (run as root): %w (%s)", err, out)
}
return true, nil
}
+17 -2
View File
@@ -8,16 +8,31 @@ import (
) )
func TestRenderUnit(t *testing.T) { func TestRenderUnit(t *testing.T) {
unit := renderUnit("/usr/local/bin/gpu-turnstile", "/etc/gpu-turnstile.env") unit := renderUnit("/var/lib/gpu-turnstile/gpu-turnstile", "/etc/gpu-turnstile.env", "")
for _, want := range []string{ for _, want := range []string{
"Type=notify", "Type=notify",
"WatchdogSec=30s", "WatchdogSec=30s",
`ExecStart="/usr/local/bin/gpu-turnstile" -config "/etc/gpu-turnstile.env"`, `ExecStart="/var/lib/gpu-turnstile/gpu-turnstile" -config "/etc/gpu-turnstile.env"`,
"Restart=on-failure", "Restart=on-failure",
"WantedBy=multi-user.target", "WantedBy=multi-user.target",
"DynamicUser=yes",
"StateDirectory=gpu-turnstile",
"RuntimeDirectory=gpu-turnstile",
"ProtectSystem=strict",
"NoNewPrivileges=yes",
"RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6",
"SystemCallFilter=@system-service",
} { } {
if !strings.Contains(unit, want) { if !strings.Contains(unit, want) {
t.Fatalf("unit missing %q:\n%s", want, unit) t.Fatalf("unit missing %q:\n%s", want, unit)
} }
} }
if strings.Contains(unit, "BindPaths") {
t.Fatalf("unit without COMFY_DIR must not bind anything:\n%s", unit)
}
unit = renderUnit("/var/lib/gpu-turnstile/gpu-turnstile", "/etc/gpu-turnstile.env", "/home/gpu/ComfyUI")
if !strings.Contains(unit, "BindPaths=/home/gpu/ComfyUI\n") {
t.Fatalf("unit with COMFY_DIR must bind it:\n%s", unit)
}
} }
+13 -3
View File
@@ -28,8 +28,18 @@ func Run(run func(ctx context.Context) error) error {
return run(ctx) return run(ctx)
} }
// Install is unsupported on non-Windows platforms. // Install is unsupported on non-Windows, non-Linux platforms.
func Install(string) error { return errUnsupported } func Install(string, bool, string) error { return errUnsupported }
// Remove is unsupported on non-Windows platforms. // Remove is unsupported on non-Windows, non-Linux platforms.
func Remove() error { return errUnsupported } func Remove() error { return errUnsupported }
// Elevated is always true here: there is no UAC concept, and privilege
// errors surface from the failing operation with a "run as root" hint.
func Elevated() bool { return true }
// RelaunchElevated is unsupported on non-Windows, non-Linux platforms.
func RelaunchElevated([]string) (int, error) { return 0, errUnsupported }
// RestartIfRunning is a no-op on platforms without service integration.
func RestartIfRunning() (bool, error) { return false, nil }
+448 -10
View File
@@ -2,22 +2,54 @@
// Package service integrates gpu-turnstile with the Windows Service // Package service integrates gpu-turnstile with the Windows Service
// Control Manager: running as a service with graceful stop, plus // Control Manager: running as a service with graceful stop, plus
// install/remove helpers. // install/remove helpers. Installed services always run as the virtual
// account NT SERVICE\gpu-turnstile — a per-service low-privilege identity
// managed by the SCM, with no password and no admin rights.
package service package service
import ( import (
"bytes"
"context" "context"
"errors"
"fmt" "fmt"
"io"
"os" "os"
"os/exec"
"path/filepath"
"strings"
"syscall"
"time" "time"
"unsafe"
"golang.org/x/sys/windows"
"golang.org/x/sys/windows/svc" "golang.org/x/sys/windows/svc"
"golang.org/x/sys/windows/svc/mgr" "golang.org/x/sys/windows/svc/mgr"
"gpu-turnstile/internal/supervise"
) )
// Name is the Windows service name. // Name is the Windows service name.
const Name = "gpu-turnstile" const Name = "gpu-turnstile"
// virtualAccount is the per-service identity the service runs as. The SCM
// manages it: no password, automatic "log on as a service" right, gone
// when the service is removed.
const virtualAccount = `NT SERVICE\` + Name
// installDirs returns the canonical install (Program Files) and data
// (ProgramData) directories.
func installDirs() (install, data string) {
pf := os.Getenv("ProgramFiles")
if pf == "" {
pf = `C:\Program Files`
}
pd := os.Getenv("ProgramData")
if pd == "" {
pd = `C:\ProgramData`
}
return filepath.Join(pf, Name), filepath.Join(pd, Name)
}
// IsService reports whether the process is running as a Windows service. // IsService reports whether the process is running as a Windows service.
func IsService() bool { func IsService() bool {
isSvc, err := svc.IsWindowsService() isSvc, err := svc.IsWindowsService()
@@ -64,34 +96,185 @@ func (h *handler) Execute(_ []string, requests <-chan svc.ChangeRequest, status
} }
} }
// Install registers gpu-turnstile as an auto-start Windows service whose // Install registers gpu-turnstile as an auto-start Windows service running
// binPath loads the given config file. Recovery actions restart the // as the NT SERVICE\gpu-turnstile virtual account, whose binPath loads the
// service after 5s on failure — this is also what brings up a staged // given config file. With copyBin it first creates the canonical layout —
// update after the updater exits with a non-zero code. // the binary is copied into %ProgramFiles%\gpu-turnstile and the config
func Install(configPath string) error { // next to it (an existing config there is kept), %ProgramData%\gpu-turnstile
// is created for logs — and registers that copy; with copyBin=false the
// current executable location is registered as-is. Recovery actions restart
// the service after 5s on failure — this is also what brings up a staged
// update after the updater exits with a non-zero code. After registering,
// the virtual account is granted modify access to the install and data
// directories (self-updates rewrite the exe), read access to the
// config file if it lives elsewhere, and — when the config sets COMFY_DIR —
// recursive modify access to the managed ComfyUI's install tree, which may
// live inside a user profile the account otherwise cannot enter. The grants
// must come after CreateService: the virtual account's SID only exists once
// the service is registered.
//
// Re-running install on an existing service converges instead of failing:
// the service is stopped first if running (so the binary can be replaced),
// the installed binary is refreshed only when the content differs, the
// registration is updated only where it drifted, and the service is
// started again only if it was running before.
func Install(configPath string, copyBin bool, version string) error {
exe, err := os.Executable() exe, err := os.Executable()
if err != nil { if err != nil {
return err return err
} }
if abs, absErr := filepath.Abs(exe); absErr == nil {
exe = abs
}
if configPath != "" {
if abs, absErr := filepath.Abs(configPath); absErr == nil {
configPath = abs
}
}
m, err := mgr.Connect() m, err := mgr.Connect()
if err != nil { if err != nil {
return fmt.Errorf("connect to service manager (run as administrator): %w", err) return fmt.Errorf("connect to service manager (run as administrator): %w", err)
} }
defer m.Disconnect() defer m.Disconnect()
// An existing service is converged, not an error. Stop it first so the
// binary copy can be replaced, and remember whether to start it again.
var s *mgr.Service
wasRunning := false
if existing, openErr := m.OpenService(Name); openErr == nil {
s = existing
defer s.Close()
if st, qErr := s.Query(); qErr == nil &&
(st.State == svc.Running || st.State == svc.StartPending) {
wasRunning = true
fmt.Println("stopping the running gpu-turnstile service")
if err := stopAndWait(s); err != nil {
return err
}
}
}
installDir, dataDir := installDirs()
if copyBin {
targetCfg := filepath.Join(installDir, "gpu-turnstile.env")
if !strings.EqualFold(filepath.Dir(exe), installDir) {
if err := os.MkdirAll(installDir, 0o755); err != nil {
return fmt.Errorf("create %s: %w", installDir, err)
}
installedExe := filepath.Join(installDir, "gpu-turnstile.exe")
if same, _ := sameFileContent(exe, installedExe); !same {
fmt.Printf("installing %s\n", installedExe)
if err := copyFile(exe, installedExe); err != nil {
return fmt.Errorf("copy binary to %s: %w", installedExe, err)
}
}
exe = installedExe
if configPath != "" && !strings.EqualFold(configPath, targetCfg) {
if _, statErr := os.Stat(targetCfg); os.IsNotExist(statErr) {
copyFile(configPath, targetCfg) //nolint:errcheck // best effort
}
}
}
// A service has no console: without LOG_FILE the output vanishes, so
// the sample pre-wires it to ProgramData. Missing configs get a
// fresh sample; installer-written ones from older versions get new
// settings appended; anything without CFG_VER is invalid and gets
// replaced (backup kept as .bak).
logPath := filepath.Join(dataDir, "gpu-turnstile.log")
if err := syncEnvFile(targetCfg, version, logPath); err != nil {
return err
}
configPath = targetCfg
}
binPath := fmt.Sprintf(`"%s" -config "%s"`, exe, configPath) binPath := fmt.Sprintf(`"%s" -config "%s"`, exe, configPath)
s, err := m.CreateService(Name, binPath, mgr.Config{
if s == nil {
s, err = m.CreateService(Name, binPath, mgr.Config{
StartType: mgr.StartAutomatic, StartType: mgr.StartAutomatic,
DisplayName: "gpu-turnstile", DisplayName: "gpu-turnstile",
Description: "GPU arbitration proxy for Ollama and ComfyUI", Description: "GPU arbitration proxy for Ollama and ComfyUI",
ServiceStartName: virtualAccount,
}) })
if err != nil { if err != nil {
return fmt.Errorf("create service: %w", err) return fmt.Errorf("create service: %w", err)
} }
defer s.Close() defer s.Close()
restart := mgr.RecoveryAction{Type: mgr.ServiceRestart, Delay: 5 * time.Second} if err := ensureRecovery(s); err != nil {
if err := s.SetRecoveryActions([]mgr.RecoveryAction{restart, restart, restart}, 24*60*60); err != nil { s.Delete() // roll back so a retry starts clean
return err
}
if err := grantAll(exe, configPath); err != nil {
s.Delete() // roll back so a retry starts clean
return err
}
// Best effort: start now instead of waiting for the next boot. A
// missing config (no consumer URLs) fails the start; the service stays
// registered and can be started once the config exists.
fmt.Println("starting the gpu-turnstile service")
s.Start()
return nil
}
// Existing service: update the registration only where it drifted.
cur, err := s.Config()
if err != nil {
return fmt.Errorf("query service config: %w", err)
}
const displayName = "gpu-turnstile"
const description = "GPU arbitration proxy for Ollama and ComfyUI"
if cur.BinaryPathName != binPath ||
cur.StartType != mgr.StartAutomatic ||
cur.ServiceStartName != virtualAccount ||
cur.DisplayName != displayName ||
cur.Description != description {
upd := cur
upd.BinaryPathName = binPath
upd.StartType = mgr.StartAutomatic
upd.ServiceStartName = virtualAccount
upd.DisplayName = displayName
upd.Description = description
if err := s.UpdateConfig(upd); err != nil {
return fmt.Errorf("update service config: %w", err)
}
}
if err := ensureRecovery(s); err != nil {
return err
}
// Idempotent: re-assert the virtual account's ACLs (granting an
// existing ACE is a no-op).
if err := grantAll(exe, configPath); err != nil {
return err
}
if wasRunning {
fmt.Println("starting the gpu-turnstile service")
if err := s.Start(); err != nil {
return fmt.Errorf("start service: %w", err)
}
}
return nil
}
// ensureRecovery makes sure the service restarts 5s after a failure — the
// mechanism that also brings up a staged update. The current actions are
// queried first so an up-to-date service is left untouched.
func ensureRecovery(s *mgr.Service) error {
want := mgr.RecoveryAction{Type: mgr.ServiceRestart, Delay: 5 * time.Second}
actions, err := s.RecoveryActions()
onFailure, flagErr := s.RecoveryActionsOnNonCrashFailures()
if err == nil && flagErr == nil && onFailure && len(actions) == 3 {
ok := true
for _, a := range actions {
if a.Type != want.Type || a.Delay != want.Delay {
ok = false
}
}
if ok {
return nil
}
}
if err := s.SetRecoveryActions([]mgr.RecoveryAction{want, want, want}, 24*60*60); err != nil {
return fmt.Errorf("set recovery actions: %w", err) return fmt.Errorf("set recovery actions: %w", err)
} }
if err := s.SetRecoveryActionsOnNonCrashFailures(true); err != nil { if err := s.SetRecoveryActionsOnNonCrashFailures(true); err != nil {
@@ -100,7 +283,122 @@ func Install(configPath string) error {
return nil return nil
} }
// Remove stops (if running) and unregisters the service. // stopAndWait stops the service and waits up to 30s for the stopped state.
func stopAndWait(s *mgr.Service) error {
if _, err := s.Control(svc.Stop); err != nil {
return fmt.Errorf("stop service: %w", err)
}
deadline := time.Now().Add(30 * time.Second)
for {
st, err := s.Query()
if err != nil {
return fmt.Errorf("query service: %w", err)
}
if st.State == svc.Stopped {
return nil
}
if time.Now().After(deadline) {
return fmt.Errorf("service did not stop within 30s")
}
time.Sleep(300 * time.Millisecond)
}
}
// sameFileContent reports whether two files hold identical bytes. A missing
// destination is simply "different".
func sameFileContent(a, b string) (bool, error) {
ba, err := os.ReadFile(a)
if err != nil {
return false, err
}
bb, err := os.ReadFile(b)
if os.IsNotExist(err) {
return false, nil
}
if err != nil {
return false, err
}
return bytes.Equal(ba, bb), nil
}
// copyFile copies src to dst (0755 on the new file).
func copyFile(src, dst string) error {
in, err := os.Open(src)
if err != nil {
return err
}
defer in.Close()
out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o755)
if err != nil {
return err
}
if _, err := io.Copy(out, in); err != nil {
out.Close()
return err
}
return out.Close()
}
// grantAll gives the virtual account every ACL the service needs: modify
// on the install and ProgramData directories and the LOG_FILE directory
// (if configured elsewhere), read on a config file outside the install
// directory.
func grantAll(exe, configPath string) error {
_, dataDir := installDirs()
if err := os.MkdirAll(dataDir, 0o755); err != nil {
return fmt.Errorf("create %s: %w", dataDir, err)
}
if err := grantAccess(dataDir, "(OI)(CI)(M)"); err != nil {
return err
}
exeDir := filepath.Dir(exe)
if err := grantAccess(exeDir, "(OI)(CI)(M)"); err != nil {
return err
}
if configPath != "" && !strings.HasPrefix(strings.ToLower(configPath), strings.ToLower(exeDir)+`\`) {
if err := grantAccess(configPath, "(R)"); err != nil {
return err
}
}
if logFile := configuredValue(configPath, "LOG_FILE"); logFile != "" {
dir := filepath.Dir(logFile)
if err := os.MkdirAll(dir, 0o755); err == nil {
if err := grantAccess(dir, "(OI)(CI)(M)"); err != nil {
return err
}
}
}
// A managed ComfyUI whose install lives somewhere the virtual account
// may not go (a user profile) needs an explicit grant — recursively,
// since ComfyUI also writes output/temp/user data next to its code. A
// missing directory is skipped: the startup warning covers it.
if comfyDir := configuredValue(configPath, "COMFY_DIR"); comfyDir != "" {
if _, err := os.Stat(comfyDir); err == nil {
if err := grantAccessTree(comfyDir, "(OI)(CI)(M)"); err != nil {
return err
}
// uv venvs (Comfy-Desktop) redirect to a base interpreter that
// can live outside COMFY_DIR; read+execute suffices for it.
if home := comfyVenvHome(comfyDir); home != "" {
if err := grantAccessTree(home, "(OI)(CI)(RX)"); err != nil {
return err
}
}
// Comfy-Desktop keeps models/input/output in a shared dir next
// to the install; the managed instance writes output there.
if shared := supervise.DesktopSharedDir(comfyDir); shared != "" {
if err := grantAccessTree(shared, "(OI)(CI)(M)"); err != nil {
return err
}
}
}
}
return nil
}
// Remove stops (if running) and unregisters the service. The virtual
// account ceases to exist with it; the ACL grants on the install and log
// directories are left in place (harmless without the account).
func Remove() error { func Remove() error {
m, err := mgr.Connect() m, err := mgr.Connect()
if err != nil { if err != nil {
@@ -118,3 +416,143 @@ func Remove() error {
} }
return nil return nil
} }
var procShellExecuteExW = windows.NewLazySystemDLL("shell32.dll").NewProc("ShellExecuteExW")
const seeMaskNoCloseProcess = 0x40
// shellExecuteInfo mirrors SHELLEXECUTEINFOW (64-bit layout).
type shellExecuteInfo struct {
cbSize uint32
fMask uint32
hwnd uintptr
lpVerb *uint16
lpFile *uint16
lpParameters *uint16
lpDirectory *uint16
nShow int32
_ int32
hInstApp uintptr
lpIDList unsafe.Pointer
lpClass *uint16
hkeyClass uintptr
dwHotKey uint32
_ uint32
hIcon uintptr
hProcess windows.Handle
}
// Elevated reports whether the current process token is UAC-elevated.
func Elevated() bool {
var token windows.Token
if err := windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_QUERY, &token); err != nil {
return false
}
defer token.Close()
return token.IsElevated()
}
// RelaunchElevated re-runs the current executable elevated via the UAC
// "runas" verb with the given arguments, waits for the child, and returns
// its exit code. The child gets a fresh console window for its output.
func RelaunchElevated(args []string) (int, error) {
exe, err := os.Executable()
if err != nil {
return 0, err
}
quoted := make([]string, len(args))
for i, a := range args {
quoted[i] = syscall.EscapeArg(a)
}
cwd, _ := os.Getwd()
verb, _ := windows.UTF16PtrFromString("runas")
exeP, _ := windows.UTF16PtrFromString(exe)
params, _ := windows.UTF16PtrFromString(strings.Join(quoted, " "))
dir, _ := windows.UTF16PtrFromString(cwd)
info := shellExecuteInfo{
fMask: seeMaskNoCloseProcess,
lpVerb: verb,
lpFile: exeP,
lpParameters: params,
lpDirectory: dir,
nShow: windows.SW_NORMAL,
}
info.cbSize = uint32(unsafe.Sizeof(info))
r, _, callErr := procShellExecuteExW.Call(uintptr(unsafe.Pointer(&info)))
if r == 0 {
if errors.Is(callErr, syscall.Errno(1223)) { // ERROR_CANCELLED
return 0, ErrUserCancelled
}
return 0, fmt.Errorf("ShellExecuteEx: %w", callErr)
}
defer windows.CloseHandle(windows.Handle(info.hProcess))
windows.WaitForSingleObject(windows.Handle(info.hProcess), windows.INFINITE)
var code uint32
if err := windows.GetExitCodeProcess(windows.Handle(info.hProcess), &code); err != nil {
return 0, err
}
return int(code), nil
}
// grantAccess gives the virtual account the icacls permission set (e.g.
// "(OI)(CI)(M)") on path.
func grantAccess(path, perms string) error {
return runIcacls(path, perms, false)
}
// grantAccessTree is grantAccess with /T: the ACE is applied to the
// existing tree, not just inherited by children created later. Needed when
// the tree already exists, e.g. a ComfyUI install in a user profile. On a
// large tree (a venv has tens of thousands of files) this takes minutes,
// so it says what it is doing instead of looking hung.
func grantAccessTree(path, perms string) error {
return runIcacls(path, perms, true)
}
func runIcacls(path, perms string, recursive bool) error {
args := []string{path, "/grant", virtualAccount + ":" + perms}
if recursive {
fmt.Printf("granting %s %s access to %s (large trees can take minutes)\n", virtualAccount, perms, path)
args = append(args, "/T")
}
start := time.Now()
out, err := exec.Command("icacls", args...).CombinedOutput()
if err != nil {
return fmt.Errorf("grant %s access to %s: %w (%s)", virtualAccount, path, err, strings.TrimSpace(string(out)))
}
if recursive {
fmt.Printf("access granted in %s\n", time.Since(start).Round(time.Second))
}
return nil
}
// RestartIfRunning restarts the service when it is installed and running
// (used after a forced update staged a new binary). Reports whether a
// restart happened. A service that is not installed or not running is not
// an error. Needs elevation.
func RestartIfRunning() (bool, error) {
m, err := mgr.Connect()
if err != nil {
return false, fmt.Errorf("connect to service manager (run as administrator): %w", err)
}
defer m.Disconnect()
s, err := m.OpenService(Name)
if err != nil {
return false, nil // not installed
}
defer s.Close()
st, err := s.Query()
if err != nil {
return false, fmt.Errorf("query service: %w", err)
}
if st.State != svc.Running && st.State != svc.StartPending {
return false, nil
}
if err := stopAndWait(s); err != nil {
return false, err
}
if err := s.Start(); err != nil {
return false, fmt.Errorf("start service: %w", err)
}
return true, nil
}
+38
View File
@@ -0,0 +1,38 @@
package service
import (
"os"
"path/filepath"
"strings"
)
// comfyVenvHome returns the base interpreter directory of the Python venv at
// <comfyDir>/.venv when that directory lives outside comfyDir, "" otherwise.
// uv-created venvs (Comfy-Desktop) ship a redirector python.exe whose real
// interpreter is the pyvenv.cfg "home" tree — typically a sibling of
// COMFY_DIR, which a sandbox/ACL covering COMFY_DIR alone does not reach.
func comfyVenvHome(comfyDir string) string {
data, err := os.ReadFile(filepath.Join(comfyDir, ".venv", "pyvenv.cfg"))
if err != nil {
return ""
}
for _, line := range strings.Split(string(data), "\n") {
k, v, ok := strings.Cut(line, "=")
if !ok || strings.TrimSpace(k) != "home" {
continue
}
home := strings.TrimSpace(v)
if home == "" {
return ""
}
if st, err := os.Stat(home); err != nil || !st.IsDir() {
return ""
}
rel, err := filepath.Rel(comfyDir, home)
if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
return home
}
return "" // inside comfyDir: already covered by the COMFY_DIR grant
}
return ""
}
+45
View File
@@ -0,0 +1,45 @@
package service
import (
"os"
"path/filepath"
"testing"
)
func TestComfyVenvHome(t *testing.T) {
root := t.TempDir()
comfy := filepath.Join(root, "ComfyUI")
outside := filepath.Join(root, "standalone-env")
inside := filepath.Join(comfy, "runtime")
for _, d := range []string{filepath.Join(comfy, ".venv"), outside, inside} {
if err := os.MkdirAll(d, 0o755); err != nil {
t.Fatal(err)
}
}
cfg := filepath.Join(comfy, ".venv", "pyvenv.cfg")
write := func(home string) {
if err := os.WriteFile(cfg, []byte("home = "+home+"\nversion_info = 3.13.0\n"), 0o644); err != nil {
t.Fatal(err)
}
}
write(outside)
if got := comfyVenvHome(comfy); got != outside {
t.Fatalf("outside home: got %q, want %q", got, outside)
}
write(inside)
if got := comfyVenvHome(comfy); got != "" {
t.Fatalf("inside home: got %q, want empty", got)
}
write(filepath.Join(root, "does-not-exist"))
if got := comfyVenvHome(comfy); got != "" {
t.Fatalf("missing home: got %q, want empty", got)
}
if got := comfyVenvHome(filepath.Join(root, "no-venv")); got != "" {
t.Fatalf("no pyvenv.cfg: got %q, want empty", got)
}
}
+372
View File
@@ -0,0 +1,372 @@
// Package supervise runs an upstream server (ComfyUI) as a managed child
// process: started on demand when a request needs it, stopped after an
// idle timeout so the GPU memory it holds is freed, and stopped with the
// parent. Crashes are logged; the next request respawns it.
package supervise
import (
"context"
"fmt"
"io"
"log/slog"
"net/url"
"os"
"os/exec"
"path/filepath"
"regexp"
"runtime"
"strings"
"sync"
"time"
)
// ComfyLayout returns the interpreter and script path of a standard ComfyUI
// venv install rooted at dir for the given GOOS: .venv\Scripts\python.exe
// on Windows, .venv/bin/python elsewhere. The script is main.py — either in
// a ComfyUI subdirectory or directly under dir, whichever exists (the
// subdirectory form wins ties and is the default when neither exists yet,
// so the caller's missing-file warning points at the documented layout).
func ComfyLayout(goos, dir string) (python, script string) {
if goos == "windows" {
python = filepath.Join(dir, ".venv", "Scripts", "python.exe")
} else {
python = filepath.Join(dir, ".venv", "bin", "python")
}
script = filepath.Join(dir, "ComfyUI", "main.py")
if _, err := os.Stat(script); err != nil {
if _, err := os.Stat(filepath.Join(dir, "main.py")); err == nil {
script = filepath.Join(dir, "main.py")
}
}
return python, script
}
// DefaultComfyCommand builds the launch command for the standard venv
// layout (see ComfyLayout): the script is passed relative to dir so dir
// stays the working directory, and --port is taken from comfyURL when the
// URL carries one. On a Comfy-Desktop standalone install the shared data
// directory (models, input, output) is added as --*-directory flags so the
// managed instance sees the desktop app's models.
func DefaultComfyCommand(goos, dir, comfyURL string) string {
python, script := ComfyLayout(goos, dir)
rel, err := filepath.Rel(dir, script)
if err != nil {
rel = script
}
cmd := `"` + python + `" ` + rel
if u, err := url.Parse(comfyURL); err == nil && u.Port() != "" {
cmd += " --port " + u.Port()
}
if shared := DesktopSharedDir(dir); shared != "" {
for _, sub := range []string{"models", "input", "output"} {
p := filepath.Join(shared, sub)
if st, err := os.Stat(p); err == nil && st.IsDir() {
cmd += ` --` + sub + `-directory "` + p + `"`
}
}
}
return cmd
}
// DesktopSharedDir returns the Comfy-Desktop shared data directory
// (<root>/ComfyUI-Shared) when dir looks like a desktop standalone install
// (<root>/ComfyUI-Installs/<name>/ComfyUI) and the shared models directory
// exists; "" otherwise. The desktop app keeps models, input and output
// there rather than inside the ComfyUI tree.
func DesktopSharedDir(dir string) string {
installs := filepath.Dir(filepath.Dir(dir))
if filepath.Base(installs) != "ComfyUI-Installs" {
return ""
}
shared := filepath.Join(filepath.Dir(installs), "ComfyUI-Shared")
if st, err := os.Stat(filepath.Join(shared, "models")); err == nil && st.IsDir() {
return shared
}
return ""
}
// Process is one managed child process.
type Process struct {
name string
argv []string
dir string
probe func(context.Context) error
startTimeout time.Duration
log *slog.Logger
mu sync.Mutex
cmd *exec.Cmd
stopping bool
ready bool
external bool // someone else serves the port; not our process
lastActivity time.Time
}
// New parses cmdLine (double quotes group arguments containing spaces) and
// prepares a managed process. probe reports whether the server answers
// (e.g. the comfy client's Probe); startTimeout bounds WaitReady. dir is
// the child's working directory; empty inherits ours.
func New(name, cmdLine, dir string, probe func(context.Context) error, startTimeout time.Duration, log *slog.Logger) (*Process, error) {
argv, err := splitCommandLine(cmdLine)
if err != nil {
return nil, fmt.Errorf("%s command: %w", name, err)
}
if len(argv) == 0 {
return nil, fmt.Errorf("%s command is empty", name)
}
if log == nil {
log = slog.Default()
}
if startTimeout <= 0 {
startTimeout = 2 * time.Minute
}
return &Process{name: name, argv: argv, dir: dir, probe: probe, startTimeout: startTimeout, log: log}, nil
}
// splitCommandLine splits a command line on whitespace, treating
// double-quoted sections as one argument (quotes removed). Backslashes are
// literal — this matches Windows paths.
func splitCommandLine(s string) ([]string, error) {
var argv []string
var cur strings.Builder
inQuote := false
have := false
flush := func() {
if have {
argv = append(argv, cur.String())
cur.Reset()
have = false
}
}
for _, r := range s {
switch {
case r == '"':
inQuote = !inQuote
have = true
case (r == ' ' || r == '\t') && !inQuote:
flush()
default:
cur.WriteRune(r)
have = true
}
}
if inQuote {
return nil, fmt.Errorf("unterminated quote in %q", s)
}
flush()
return argv, nil
}
// Running reports whether the child process is currently alive.
func (p *Process) Running() bool {
p.mu.Lock()
defer p.mu.Unlock()
return p.cmd != nil
}
// Status describes the child for status displays: "external" when something
// else serves the port, "running" once ready, "starting" while the child
// boots, "stopped" otherwise.
func (p *Process) Status() string {
p.mu.Lock()
defer p.mu.Unlock()
switch {
case p.external:
return "external"
case p.cmd == nil:
return "stopped"
case p.ready:
return "running"
default:
return "starting"
}
}
// Ready reports whether the server has answered a probe since its last
// (re)start. Health checks use it to tell "starting up" from "outage".
func (p *Process) Ready() bool {
p.mu.Lock()
defer p.mu.Unlock()
return p.ready
}
// MarkReady records that the server answered.
func (p *Process) MarkReady() {
p.mu.Lock()
p.ready = true
p.mu.Unlock()
}
// NoteActivity resets the idle clock; called for every request served.
func (p *Process) NoteActivity() {
p.mu.Lock()
p.lastActivity = time.Now()
p.mu.Unlock()
}
// EnsureRunning starts the child if it is not running. It returns as soon
// as the process is spawned; readiness is WaitReady's job (and the proxy's
// retry backoff bridges the gap for plain proxied requests). When the URL
// already answers — e.g. the ComfyUI desktop app grabbed the port — no
// child is spawned: the external server is used as-is, and the idle
// watcher never touches it (it only kills its own child).
func (p *Process) EnsureRunning() error {
p.mu.Lock()
defer p.mu.Unlock()
p.lastActivity = time.Now()
if p.cmd != nil {
return nil
}
pctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
err := p.probe(pctx)
cancel()
if err == nil {
p.ready = true
if !p.external {
p.external = true
p.log.Info(p.name + " is already served externally; not spawning a managed instance")
}
return nil
}
p.external = false
cmd := exec.Command(p.argv[0], p.argv[1:]...)
cmd.Dir = p.dir
// Ask the child not to colorize (ComfyUI ignores this and colors
// anyway, so pipeLog also strips escape sequences).
cmd.Env = append(os.Environ(), "NO_COLOR=1", "TERM=dumb")
stdout, err := cmd.StdoutPipe()
if err != nil {
return err
}
stderr, err := cmd.StderrPipe()
if err != nil {
return err
}
if err := cmd.Start(); err != nil {
return fmt.Errorf("start %s: %w", p.name, err)
}
p.cmd = cmd
p.stopping = false
p.ready = false
go p.pipeLog(stdout)
go p.pipeLog(stderr)
go func() {
err := cmd.Wait()
p.mu.Lock()
p.cmd = nil
p.ready = false
intentional := p.stopping
p.mu.Unlock()
if intentional {
p.log.Info(p.name + " stopped")
} else {
p.log.Warn(p.name+" exited unexpectedly; the next request restarts it", "err", err)
}
}()
p.log.Info(p.name+" starting", "pid", cmd.Process.Pid, "cmd", strings.Join(p.argv, " "))
return nil
}
// WaitReady blocks until the probe succeeds, ctx ends, or the start
// timeout passes.
func (p *Process) WaitReady(ctx context.Context) error {
ctx, cancel := context.WithTimeout(ctx, p.startTimeout)
defer cancel()
for {
pctx, pcancel := context.WithTimeout(ctx, 5*time.Second)
err := p.probe(pctx)
pcancel()
if err == nil {
p.NoteActivity()
p.MarkReady()
return nil
}
select {
case <-ctx.Done():
return fmt.Errorf("%s did not become ready: %w", p.name, ctx.Err())
case <-time.After(500 * time.Millisecond):
}
}
}
// Stop kills the child process (the whole tree on Windows) if running.
func (p *Process) Stop() {
p.mu.Lock()
cmd := p.cmd
if cmd == nil {
p.mu.Unlock()
return
}
p.stopping = true
p.ready = false
p.mu.Unlock()
stopTree(cmd)
}
// WatchIdle stops the child after idleTimeout without activity, but only
// when gpuIdle reports the GPU lock is free (no active or pending work).
// Returns when ctx ends.
func (p *Process) WatchIdle(ctx context.Context, idleTimeout time.Duration, gpuIdle func() bool) {
ticker := time.NewTicker(5 * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
}
p.mu.Lock()
idleFor := time.Since(p.lastActivity)
running := p.cmd != nil
p.mu.Unlock()
if running && idleFor > idleTimeout && gpuIdle() {
p.log.Info(p.name+" idle; stopping to free the GPU", "idle_for", idleFor.Round(time.Second))
p.Stop()
}
}
}
// pipeLog forwards one child output stream to the log at INFO, line by
// line, prefixed with the process name. ANSI escape sequences are
// stripped: ComfyUI colorizes unconditionally, and the escapes only
// render as garbage in a log file.
func (p *Process) pipeLog(r io.Reader) {
buf := make([]byte, 4096)
var line string
for {
n, err := r.Read(buf)
line += string(buf[:n])
for {
i := strings.IndexByte(line, '\n')
if i < 0 {
break
}
p.log.Info(p.name + ": " + stripANSI(strings.TrimRight(line[:i], "\r")))
line = line[i+1:]
}
if err != nil {
if strings.TrimSpace(line) != "" {
p.log.Info(p.name + ": " + stripANSI(line))
}
return
}
}
}
// ansiPattern matches CSI escape sequences (colors, cursor moves, …).
var ansiPattern = regexp.MustCompile("\x1b\\[[0-9;?]*[a-zA-Z]")
func stripANSI(s string) string {
return ansiPattern.ReplaceAllString(s, "")
}
// stopTree kills cmd's process, including its children on Windows (python
// launchers tend to spawn some). The Wait goroutine reaps it.
func stopTree(cmd *exec.Cmd) {
if runtime.GOOS == "windows" {
exec.Command("taskkill", "/T", "/F", "/PID",
fmt.Sprint(cmd.Process.Pid)).Run() //nolint:errcheck // best effort
return
}
cmd.Process.Kill() //nolint:errcheck // best effort
}
+276
View File
@@ -0,0 +1,276 @@
package supervise
import (
"context"
"errors"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
"time"
)
func TestSplitCommandLine(t *testing.T) {
cases := []struct {
in string
want []string
}{
{`python main.py --port 8188`, []string{"python", "main.py", "--port", "8188"}},
{`"C:\Program Files\py\python.exe" main.py`, []string{`C:\Program Files\py\python.exe`, "main.py"}},
{` spaced out `, []string{"spaced", "out"}},
{`a "b c" d`, []string{"a", "b c", "d"}},
{"", nil},
}
for _, c := range cases {
got, err := splitCommandLine(c.in)
if err != nil {
t.Errorf("splitCommandLine(%q): %v", c.in, err)
continue
}
if len(got) != len(c.want) {
t.Errorf("splitCommandLine(%q) = %v, want %v", c.in, got, c.want)
continue
}
for i := range got {
if got[i] != c.want[i] {
t.Errorf("splitCommandLine(%q) = %v, want %v", c.in, got, c.want)
break
}
}
}
if _, err := splitCommandLine(`"unterminated`); err == nil {
t.Error("expected error for unterminated quote")
}
}
// TestHelperProcess is the child executed by the lifecycle tests: it just
// sleeps. The env marker is set only around the spawn, so in the normal
// test run this returns immediately.
func TestHelperProcess(t *testing.T) {
if os.Getenv("GO_HELPER_PROCESS") != "1" {
return
}
time.Sleep(30 * time.Second)
os.Exit(0)
}
func newHelper(t *testing.T, name string) *Process {
t.Helper()
// The probe always fails: nothing external serves the port, so
// EnsureRunning spawns the helper child.
p, err := New(name, `"`+os.Args[0]+`" -test.run=TestHelperProcess`, "",
func(context.Context) error { return errors.New("nothing there") }, 5*time.Second, slog.Default())
if err != nil {
t.Fatal(err)
}
return p
}
// startHelper spawns the child with the marker set; exec.Command inherits
// the environment at spawn time, so it can be unset right after.
func startHelper(t *testing.T, p *Process) {
t.Helper()
os.Setenv("GO_HELPER_PROCESS", "1")
defer os.Unsetenv("GO_HELPER_PROCESS")
if err := p.EnsureRunning(); err != nil {
t.Fatal(err)
}
}
func waitStopped(t *testing.T, p *Process, timeout time.Duration) {
t.Helper()
deadline := time.Now().Add(timeout)
for p.Running() && time.Now().Before(deadline) {
time.Sleep(50 * time.Millisecond)
}
if p.Running() {
t.Fatal("process still running")
}
}
func TestEnsureRunningAndStop(t *testing.T) {
p := newHelper(t, "helper")
if p.Running() {
t.Fatal("Running before start")
}
startHelper(t, p)
if !p.Running() {
t.Fatal("not Running after EnsureRunning")
}
if err := p.EnsureRunning(); err != nil {
t.Fatal("second EnsureRunning must be a no-op")
}
p.Stop()
waitStopped(t, p, 5*time.Second)
}
func TestEnsureRunningPrefersExternalServer(t *testing.T) {
// The port is already served (e.g. the ComfyUI desktop app): no child
// is spawned, the supervisor reports ready, and Stop is a no-op.
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
probe := func(ctx context.Context) error {
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
resp, err := http.DefaultClient.Do(req)
if err != nil {
return err
}
resp.Body.Close()
return nil
}
p, err := New("external", `"`+os.Args[0]+`"`, "", probe, 5*time.Second, slog.Default())
if err != nil {
t.Fatal(err)
}
if err := p.EnsureRunning(); err != nil {
t.Fatal(err)
}
if p.Running() {
t.Fatal("spawned a child even though the port is already served")
}
if !p.Ready() {
t.Fatal("external server should count as ready")
}
p.Stop() // must not touch the external server
}
func TestWaitReady(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
probe := func(ctx context.Context) error {
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
resp, err := http.DefaultClient.Do(req)
if err != nil {
return err
}
resp.Body.Close()
return nil
}
p, err := New("ready", `"`+os.Args[0]+`"`, "", probe, 5*time.Second, slog.Default())
if err != nil {
t.Fatal(err)
}
if err := p.WaitReady(context.Background()); err != nil {
t.Fatal(err)
}
failing, err := New("failing", `"`+os.Args[0]+`"`, "",
func(context.Context) error { return errors.New("no") }, 500*time.Millisecond, slog.Default())
if err != nil {
t.Fatal(err)
}
if err := failing.WaitReady(context.Background()); err == nil {
t.Fatal("expected timeout error from WaitReady")
}
}
func TestWatchIdleStops(t *testing.T) {
p := newHelper(t, "idle")
startHelper(t, p)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go p.WatchIdle(ctx, 200*time.Millisecond, func() bool { return true })
waitStopped(t, p, 10*time.Second)
}
func TestWatchIdleRespectsBusyGPU(t *testing.T) {
p := newHelper(t, "busy")
startHelper(t, p)
defer p.Stop()
ctx, cancel := context.WithCancel(context.Background())
go p.WatchIdle(ctx, 100*time.Millisecond, func() bool { return false })
time.Sleep(600 * time.Millisecond)
cancel()
if !p.Running() {
t.Fatal("process was stopped while the GPU was busy")
}
}
func TestComfyLayoutAndDefaultCommand(t *testing.T) {
// Nested layout (ComfyUI/main.py under dir) wins.
dir := t.TempDir()
nested := filepath.Join(dir, "ComfyUI", "main.py")
if err := os.MkdirAll(filepath.Dir(nested), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(nested, []byte("x"), 0o644); err != nil {
t.Fatal(err)
}
python, script := ComfyLayout("windows", dir)
if want := filepath.Join(dir, ".venv", "Scripts", "python.exe"); python != want {
t.Errorf("python = %s, want %s", python, want)
}
if script != nested {
t.Errorf("script = %s, want %s", script, nested)
}
cmd := DefaultComfyCommand("windows", dir, "http://127.0.0.1:8189")
want := `"` + filepath.Join(dir, ".venv", "Scripts", "python.exe") + `" ` + filepath.Join("ComfyUI", "main.py") + " --port 8189"
if cmd != want {
t.Errorf("cmd = %q, want %q", cmd, want)
}
// Flat layout (main.py directly under dir) is found too.
flat := t.TempDir()
if err := os.WriteFile(filepath.Join(flat, "main.py"), []byte("x"), 0o644); err != nil {
t.Fatal(err)
}
if _, script := ComfyLayout("linux", flat); script != filepath.Join(flat, "main.py") {
t.Errorf("flat script = %s", script)
}
cmd = DefaultComfyCommand("linux", flat, "http://comfy.internal")
if strings.Contains(cmd, "--port") {
t.Errorf("cmd = %q, want no --port for a port-less URL", cmd)
}
if !strings.HasSuffix(cmd, `" main.py`) {
t.Errorf("cmd = %q, want quoted python + relative main.py", cmd)
}
// Neither exists yet: default to the documented nested form so the
// startup warning points there.
empty := t.TempDir()
if _, script := ComfyLayout("windows", empty); script != filepath.Join(empty, "ComfyUI", "main.py") {
t.Errorf("missing-layout script = %s", script)
}
}
func TestStripANSI(t *testing.T) {
cases := map[string]string{
"\x1b[32m[INFO]\x1b[0m Starting server": "[INFO] Starting server",
"\x1b[1m\x1b[31m[ERROR]\x1b[0m boom": "[ERROR] boom",
"plain line": "plain line",
"\x1b[33mWARN\x1b[0m: \x1b[1mbold\x1b[0m": "WARN: bold",
}
for in, want := range cases {
if got := stripANSI(in); got != want {
t.Errorf("stripANSI(%q) = %q, want %q", in, got, want)
}
}
}
func TestDesktopSharedDir(t *testing.T) {
root := t.TempDir()
comfy := filepath.Join(root, "ComfyUI-Installs", "rtx5080", "ComfyUI")
if err := os.MkdirAll(comfy, 0o755); err != nil {
t.Fatal(err)
}
if got := DesktopSharedDir(comfy); got != "" {
t.Fatalf("no shared dir yet: got %q, want empty", got)
}
shared := filepath.Join(root, "ComfyUI-Shared")
if err := os.MkdirAll(filepath.Join(shared, "models"), 0o755); err != nil {
t.Fatal(err)
}
if got := DesktopSharedDir(comfy); got != shared {
t.Fatalf("got %q, want %q", got, shared)
}
if got := DesktopSharedDir(filepath.Join(root, "plain", "ComfyUI")); got != "" {
t.Fatalf("non-desktop layout: got %q, want empty", got)
}
}
+54 -26
View File
@@ -26,11 +26,15 @@ import (
// maxAssetSize bounds release asset downloads. // maxAssetSize bounds release asset downloads.
const maxAssetSize = 512 << 20 const maxAssetSize = 512 << 20
// Updater checks one Gitea repository for newer releases. // Updater checks one Gitea repository for releases.
type Updater struct { type Updater struct {
Repo string // e.g. https://git.rambossek.at/PUBLIC/gpu-turnstile Repo string // e.g. https://git.rambossek.at/PUBLIC/gpu-turnstile
Asset string // e.g. gpu-turnstile.exe Asset string // e.g. gpu-turnstile.exe
Version string // current version, e.g. v0.1.2 ("dev" disables updates) Version string // current binary version, e.g. v0.1.2 ("dev" cannot be compared)
// Desired is the APP_VER policy: "dev" disables updates, "stable" (or
// empty) tracks the latest release, anything else is an exact vX.Y.Z
// release tag to pin — staging it even when that means a downgrade.
Desired string
Log *slog.Logger Log *slog.Logger
Client *http.Client Client *http.Client
} }
@@ -167,41 +171,65 @@ func CleanupOld(exePath string) {
os.Remove(exePath + ".new") os.Remove(exePath + ".new")
} }
// Check performs a single update check. staged is true when a newer, // Check performs a single update check. staged is true when a
// signature-verified binary has been swapped into place at exePath; the // signature-verified binary has been swapped into place at exePath; the
// caller should then restart the process. A nil error with staged=false // caller should then restart the process. to is the release tag the check
// means "no action" (up to date, disabled, or dev build); a non-nil error // resolved (the latest release or the pinned tag), set once the release
// means the check failed and the running binary is untouched. // fetch succeeded — even when staging afterwards fails. A nil error with
func (u *Updater) Check(ctx context.Context, exePath string) (staged bool, err error) { // staged=false means "no action" (up to date, APP_VER=dev, or no embedded
// public key); a non-nil error means the check failed and the running
// binary is untouched.
func (u *Updater) Check(ctx context.Context, exePath string) (staged bool, to string, err error) {
log := u.logger() log := u.logger()
if u.Version == "" || u.Version == "dev" { desired := u.Desired
log.Debug("auto-update: dev build, skipping") if desired == "" {
return false, nil desired = "stable"
}
if desired == "dev" {
log.Debug("auto-update: APP_VER=dev, skipping")
return false, "", nil
} }
if publicKeyPEM == "" { if publicKeyPEM == "" {
log.Debug("auto-update: no public key embedded, skipping") log.Debug("auto-update: no public key embedded, skipping")
return false, nil return false, "", nil
} }
api, err := u.apiURL() api, err := u.apiURL()
if err != nil { if err != nil {
return false, err return false, "", err
} }
body, err := u.get(ctx, api+"/releases/latest") pinned := desired != "stable"
endpoint := api + "/releases/latest"
if pinned {
endpoint = api + "/releases/tags/" + desired
}
body, err := u.get(ctx, endpoint)
if err != nil { if err != nil {
return false, fmt.Errorf("fetch latest release: %w", err) return false, "", fmt.Errorf("fetch release: %w", err)
} }
var rel release var rel release
if err := json.Unmarshal(body, &rel); err != nil { if err := json.Unmarshal(body, &rel); err != nil {
return false, fmt.Errorf("parse release: %w", err) return false, "", fmt.Errorf("parse release: %w", err)
} }
to = rel.TagName
if pinned {
// Pin mode: any difference from the target tag means stage it —
// including downgrades and replacing a dev binary.
if u.Version == rel.TagName {
log.Debug("auto-update: already on pinned version", "version", u.Version)
return false, to, nil
}
} else if u.Version != "" && u.Version != "dev" {
// Stable mode: only strictly newer releases count; a dev binary
// cannot be compared and is always replaced by the latest release.
newer, err := newerVersion(u.Version, rel.TagName) newer, err := newerVersion(u.Version, rel.TagName)
if err != nil { if err != nil {
return false, err return false, to, err
} }
if !newer { if !newer {
log.Debug("auto-update: up to date", "version", u.Version, "latest", rel.TagName) log.Debug("auto-update: up to date", "version", u.Version, "latest", rel.TagName)
return false, nil return false, to, nil
}
} }
urls := make(map[string]string, len(rel.Assets)) urls := make(map[string]string, len(rel.Assets))
@@ -210,42 +238,42 @@ func (u *Updater) Check(ctx context.Context, exePath string) (staged bool, err e
} }
assetURL, ok := urls[u.Asset] assetURL, ok := urls[u.Asset]
if !ok { if !ok {
return false, fmt.Errorf("release %s has no asset %q", rel.TagName, u.Asset) return false, to, fmt.Errorf("release %s has no asset %q", rel.TagName, u.Asset)
} }
sigURL, ok := urls[u.Asset+".sig"] sigURL, ok := urls[u.Asset+".sig"]
if !ok { if !ok {
return false, fmt.Errorf("release %s has no signature asset %q", rel.TagName, u.Asset+".sig") return false, to, fmt.Errorf("release %s has no signature asset %q", rel.TagName, u.Asset+".sig")
} }
data, err := u.get(ctx, assetURL) data, err := u.get(ctx, assetURL)
if err != nil { if err != nil {
return false, fmt.Errorf("download %s: %w", u.Asset, err) return false, to, fmt.Errorf("download %s: %w", u.Asset, err)
} }
sig, err := u.get(ctx, sigURL) sig, err := u.get(ctx, sigURL)
if err != nil { if err != nil {
return false, fmt.Errorf("download signature: %w", err) return false, to, fmt.Errorf("download signature: %w", err)
} }
if sumURL, ok := urls[u.Asset+".sha256"]; ok { if sumURL, ok := urls[u.Asset+".sha256"]; ok {
sumText, err := u.get(ctx, sumURL) sumText, err := u.get(ctx, sumURL)
if err != nil { if err != nil {
return false, fmt.Errorf("download checksum: %w", err) return false, to, fmt.Errorf("download checksum: %w", err)
} }
want := strings.Fields(string(sumText))[0] want := strings.Fields(string(sumText))[0]
got := hex.EncodeToString(sha256Bytes(data)) got := hex.EncodeToString(sha256Bytes(data))
if !strings.EqualFold(want, got) { if !strings.EqualFold(want, got) {
return false, fmt.Errorf("sha256 mismatch: got %s, want %s", got, want) return false, to, fmt.Errorf("sha256 mismatch: got %s, want %s", got, want)
} }
} }
if err := verifySignature(publicKeyPEM, data, sig); err != nil { if err := verifySignature(publicKeyPEM, data, sig); err != nil {
return false, err return false, to, err
} }
if err := stage(exePath, data); err != nil { if err := stage(exePath, data); err != nil {
return false, fmt.Errorf("stage update: %w", err) return false, to, fmt.Errorf("stage update: %w", err)
} }
log.Info("auto-update: new version staged", "from", u.Version, "to", rel.TagName) log.Info("auto-update: new version staged", "from", u.Version, "to", rel.TagName)
return true, nil return true, to, nil
} }
func sha256Bytes(data []byte) []byte { func sha256Bytes(data []byte) []byte {
+72 -9
View File
@@ -43,7 +43,7 @@ func newFakeGitea(t *testing.T, tag string, assetContent []byte) *fakeGitea {
sign := func() []byte { return ed25519.Sign(priv, f.asset) } sign := func() []byte { return ed25519.Sign(priv, f.asset) }
mux := http.NewServeMux() mux := http.NewServeMux()
mux.HandleFunc("/api/v1/repos/o/r/releases/latest", func(w http.ResponseWriter, r *http.Request) { serveRelease := func(w http.ResponseWriter, r *http.Request) {
assets := []map[string]string{ assets := []map[string]string{
{"name": "gpu-turnstile.exe", "browser_download_url": f.srv.URL + "/dl/exe"}, {"name": "gpu-turnstile.exe", "browser_download_url": f.srv.URL + "/dl/exe"},
{"name": "gpu-turnstile.exe.sha256", "browser_download_url": f.srv.URL + "/dl/sha"}, {"name": "gpu-turnstile.exe.sha256", "browser_download_url": f.srv.URL + "/dl/sha"},
@@ -52,7 +52,9 @@ func newFakeGitea(t *testing.T, tag string, assetContent []byte) *fakeGitea {
assets = append(assets, map[string]string{"name": "gpu-turnstile.exe.sig", "browser_download_url": f.srv.URL + "/dl/sig"}) assets = append(assets, map[string]string{"name": "gpu-turnstile.exe.sig", "browser_download_url": f.srv.URL + "/dl/sig"})
} }
json.NewEncoder(w).Encode(map[string]any{"tag_name": f.tag, "assets": assets}) json.NewEncoder(w).Encode(map[string]any{"tag_name": f.tag, "assets": assets})
}) }
mux.HandleFunc("/api/v1/repos/o/r/releases/latest", serveRelease)
mux.HandleFunc("/api/v1/repos/o/r/releases/tags/"+f.tag, serveRelease)
mux.HandleFunc("/dl/exe", func(w http.ResponseWriter, r *http.Request) { w.Write(f.asset) }) mux.HandleFunc("/dl/exe", func(w http.ResponseWriter, r *http.Request) { w.Write(f.asset) })
mux.HandleFunc("/dl/sig", func(w http.ResponseWriter, r *http.Request) { mux.HandleFunc("/dl/sig", func(w http.ResponseWriter, r *http.Request) {
sig := sign() sig := sign()
@@ -73,6 +75,12 @@ func (f *fakeGitea) updater(version string) *Updater {
return &Updater{Repo: f.srv.URL + "/o/r", Asset: "gpu-turnstile.exe", Version: version} return &Updater{Repo: f.srv.URL + "/o/r", Asset: "gpu-turnstile.exe", Version: version}
} }
func (f *fakeGitea) updaterDesired(version, desired string) *Updater {
u := f.updater(version)
u.Desired = desired
return u
}
func fakeExe(t *testing.T) string { func fakeExe(t *testing.T) string {
t.Helper() t.Helper()
exe := filepath.Join(t.TempDir(), "gpu-turnstile.exe") exe := filepath.Join(t.TempDir(), "gpu-turnstile.exe")
@@ -94,13 +102,16 @@ func TestCheckStagesUpdate(t *testing.T) {
withPublicKey(t, f.pubPEM) withPublicKey(t, f.pubPEM)
exe := fakeExe(t) exe := fakeExe(t)
staged, err := f.updater("v0.1.2").Check(context.Background(), exe) staged, to, err := f.updater("v0.1.2").Check(context.Background(), exe)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if !staged { if !staged {
t.Fatal("expected staged update") t.Fatal("expected staged update")
} }
if to != "v9.9.9" {
t.Fatalf("to = %q, want v9.9.9", to)
}
content, _ := os.ReadFile(exe) content, _ := os.ReadFile(exe)
if string(content) != "new-binary" { if string(content) != "new-binary" {
t.Fatalf("exe content = %q", content) t.Fatalf("exe content = %q", content)
@@ -117,7 +128,7 @@ func TestCheckRejectsTamperedSignature(t *testing.T) {
withPublicKey(t, f.pubPEM) withPublicKey(t, f.pubPEM)
exe := fakeExe(t) exe := fakeExe(t)
staged, err := f.updater("v0.1.2").Check(context.Background(), exe) staged, _, err := f.updater("v0.1.2").Check(context.Background(), exe)
if err == nil { if err == nil {
t.Fatal("expected signature error") t.Fatal("expected signature error")
} }
@@ -134,7 +145,7 @@ func TestCheckSkipsOlderOrEqual(t *testing.T) {
for _, tag := range []string{"v0.1.2", "v0.1.1", "v0.0.9"} { for _, tag := range []string{"v0.1.2", "v0.1.1", "v0.0.9"} {
f := newFakeGitea(t, tag, []byte("new-binary")) f := newFakeGitea(t, tag, []byte("new-binary"))
withPublicKey(t, f.pubPEM) withPublicKey(t, f.pubPEM)
staged, err := f.updater("v0.1.2").Check(context.Background(), fakeExe(t)) staged, _, err := f.updater("v0.1.2").Check(context.Background(), fakeExe(t))
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -147,18 +158,70 @@ func TestCheckSkipsOlderOrEqual(t *testing.T) {
func TestCheckSkipsWithoutPublicKey(t *testing.T) { func TestCheckSkipsWithoutPublicKey(t *testing.T) {
f := newFakeGitea(t, "v9.9.9", []byte("new-binary")) f := newFakeGitea(t, "v9.9.9", []byte("new-binary"))
withPublicKey(t, "") withPublicKey(t, "")
staged, err := f.updater("v0.1.2").Check(context.Background(), fakeExe(t)) staged, _, err := f.updater("v0.1.2").Check(context.Background(), fakeExe(t))
if err != nil || staged { if err != nil || staged {
t.Fatalf("staged=%v err=%v, want no action without key", staged, err) t.Fatalf("staged=%v err=%v, want no action without key", staged, err)
} }
} }
func TestCheckSkipsDevBuild(t *testing.T) { func TestCheckDevBuildGetsStable(t *testing.T) {
// A dev binary cannot be compared; APP_VER=stable replaces it with the
// latest release.
f := newFakeGitea(t, "v9.9.9", []byte("new-binary")) f := newFakeGitea(t, "v9.9.9", []byte("new-binary"))
withPublicKey(t, f.pubPEM) withPublicKey(t, f.pubPEM)
staged, err := f.updater("dev").Check(context.Background(), fakeExe(t)) exe := fakeExe(t)
staged, _, err := f.updater("dev").Check(context.Background(), exe)
if err != nil {
t.Fatal(err)
}
if !staged {
t.Fatal("expected dev binary to be replaced by the latest release")
}
content, _ := os.ReadFile(exe)
if string(content) != "new-binary" {
t.Fatalf("exe content = %q", content)
}
}
func TestCheckDesiredDevDisables(t *testing.T) {
f := newFakeGitea(t, "v9.9.9", []byte("new-binary"))
withPublicKey(t, f.pubPEM)
for _, version := range []string{"dev", "v0.1.2"} {
staged, _, err := f.updaterDesired(version, "dev").Check(context.Background(), fakeExe(t))
if err != nil || staged { if err != nil || staged {
t.Fatalf("staged=%v err=%v, want no action for dev build", staged, err) t.Fatalf("version %s: staged=%v err=%v, want no action with APP_VER=dev", version, staged, err)
}
}
}
func TestCheckPinned(t *testing.T) {
// Pin mode stages the exact tag — up or down — and skips when the
// binary already matches.
for _, version := range []string{"v0.1.2", "v9.9.9", "dev"} {
f := newFakeGitea(t, "v0.5.0", []byte("pinned-binary"))
withPublicKey(t, f.pubPEM)
exe := fakeExe(t)
staged, to, err := f.updaterDesired(version, "v0.5.0").Check(context.Background(), exe)
if err != nil {
t.Fatalf("version %s: %v", version, err)
}
if !staged {
t.Fatalf("version %s: expected pinned v0.5.0 to be staged", version)
}
if to != "v0.5.0" {
t.Fatalf("version %s: to = %q, want v0.5.0", version, to)
}
content, _ := os.ReadFile(exe)
if string(content) != "pinned-binary" {
t.Fatalf("version %s: exe content = %q", version, content)
}
}
f := newFakeGitea(t, "v0.5.0", []byte("pinned-binary"))
withPublicKey(t, f.pubPEM)
staged, _, err := f.updaterDesired("v0.5.0", "v0.5.0").Check(context.Background(), fakeExe(t))
if err != nil || staged {
t.Fatalf("staged=%v err=%v, want no action when already on the pinned version", staged, err)
} }
} }