Compare commits

..
21 Commits
Author SHA1 Message Date
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
23 changed files with 1495 additions and 88 deletions
+394 -48
View File
@@ -5,6 +5,7 @@ package main
import (
"bufio"
"context"
"encoding/json"
"errors"
"fmt"
"io"
@@ -15,13 +16,16 @@ import (
"os"
"os/signal"
"path/filepath"
"reflect"
"runtime"
"strings"
"sync"
"syscall"
"time"
"gpu-turnstile/internal/comfy"
"gpu-turnstile/internal/config"
"gpu-turnstile/internal/control"
"gpu-turnstile/internal/game"
"gpu-turnstile/internal/lock"
"gpu-turnstile/internal/metrics"
@@ -36,9 +40,15 @@ import (
var version = "dev"
// exitCodeUpdate tells the service recovery configuration to restart the
// process: a signed update has been staged and the GPU lock is idle.
// process: a signed update has been staged or a changed config reload was
// requested, and the GPU lock is idle.
const exitCodeUpdate = 3
// exitCodeStaged is returned by an elevated --force-update child when it
// staged a new binary, so the non-elevated parent can tell "updated" from
// "up to date" (it cannot see the child's console).
const exitCodeStaged = 4
// stdoutIsTerminal reports whether stdout is a console (char device), as
// opposed to a pipe or file — which is what Docker containers and services
// see.
@@ -51,7 +61,7 @@ func stdoutIsTerminal() bool {
// --install-service / --remove-service switches, --no-copy, -h/--help,
// -v/--version, --force-update and the hidden --elevated-child marker from
// args.
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, elevatedChild bool, rest []string) {
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild bool, rest []string) {
rest = args[:0]
for i := 0; i < len(args); i++ {
switch {
@@ -60,9 +70,9 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
i++
case strings.HasPrefix(args[i], "-config="):
configPath = strings.TrimPrefix(args[i], "-config=")
case args[i] == "--install-service" || args[i] == "-install-service":
case args[i] == "--install-service" || args[i] == "-install-service" || args[i] == "-i":
install = true
case args[i] == "--remove-service" || args[i] == "-remove-service":
case args[i] == "--remove-service" || args[i] == "-remove-service" || args[i] == "-r":
remove = true
case args[i] == "--no-copy" || args[i] == "-no-copy":
noCopy = true
@@ -72,13 +82,19 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
showVersion = true
case args[i] == "--force-update" || args[i] == "-force-update":
forceUpdate = true
case args[i] == "--update-now" || args[i] == "-update-now":
updateNow = true
case args[i] == "--reload-env" || args[i] == "-reload-env":
reloadEnv = true
case args[i] == "--monitor" || args[i] == "-monitor" || args[i] == "-m":
monitor = true
case args[i] == "--elevated-child":
elevatedChild = true
default:
rest = append(rest, args[i])
}
}
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, elevatedChild, rest
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, rest
}
// versionLine is printed at the top of every help and error screen.
@@ -88,11 +104,20 @@ const usageText = `GPU arbitration proxy for Ollama + ComfyUI
Usage:
gpu-turnstile -config <path> run the proxy
gpu-turnstile --install-service [--no-copy] [-config path] install + start as a service
gpu-turnstile --remove-service stop + uninstall the service
gpu-turnstile -i | --install-service [--no-copy] [-config path]
install + start as a service
gpu-turnstile -r | --remove-service stop + uninstall the service
gpu-turnstile -v | --version print just the version
gpu-turnstile --force-update check for a signed update now,
apply it and restart the service
(no admin needed when the service runs)
gpu-turnstile --update-now like --force-update, but only
through the running service
gpu-turnstile --reload-env make the service re-read and
validate its config file, then
restart onto it if it changed
gpu-turnstile -m | --monitor live status view (downstreams,
GPU lock, queue); Ctrl+C quits
gpu-turnstile -h | --help this help
Options:
@@ -123,12 +148,12 @@ func fatalUsage(format string, args ...any) {
}
func main() {
configPath, install, remove, noCopy, help, showVersion, forceUpdate, elevatedChild, args := parseFlags(os.Args[1:])
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, reloadEnv, monitor, elevatedChild, args := parseFlags(os.Args[1:])
if showVersion {
fmt.Println(version)
return
}
bare := configPath == "" && !install && !remove && !forceUpdate && !elevatedChild && len(args) == 0
bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !reloadEnv && !monitor && !elevatedChild && len(args) == 0
if help || (bare && stdoutIsTerminal()) {
// Bare invocation in a terminal (e.g. double-clicked on Windows)
// shows the help instead of starting a proxy window with no visible
@@ -151,12 +176,24 @@ func main() {
fatalUsage("error: --install-service and --remove-service are mutually exclusive")
case forceUpdate && (install || remove):
fatalUsage("error: --force-update cannot be combined with --install-service/--remove-service")
case updateNow && (install || remove || forceUpdate):
fatalUsage("error: --update-now cannot be combined with other commands")
case reloadEnv && (install || remove || forceUpdate || updateNow):
fatalUsage("error: --reload-env cannot be combined with other commands")
case monitor && (install || remove || forceUpdate || updateNow || reloadEnv):
fatalUsage("error: --monitor cannot be combined with other commands")
case install:
os.Exit(serviceCommand(configPath, true, noCopy, elevatedChild))
case remove:
os.Exit(serviceCommand(configPath, false, noCopy, elevatedChild))
case forceUpdate:
os.Exit(forceUpdateCommand(configPath, elevatedChild))
case updateNow:
os.Exit(updateNowCommand())
case reloadEnv:
os.Exit(reloadEnvCommand())
case monitor:
os.Exit(monitorCommand())
}
if len(args) > 0 {
fatalUsage("error: unknown arguments: %s", strings.Join(args, " "))
@@ -176,7 +213,7 @@ func main() {
syncEnvFile(resolveConfigPath(configPath), cfg.LogFile, log)
if service.IsService() {
if err := service.Run(func(ctx context.Context) error { return run(ctx, cfg, log, logOut, true) }); err != nil {
if err := service.Run(func(ctx context.Context) error { return run(ctx, cfg, log, logOut, true, configPath) }); err != nil {
log.Error("service failed", "err", err)
os.Exit(1)
}
@@ -184,7 +221,7 @@ func main() {
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
if err := run(ctx, cfg, log, logOut, false); err != nil {
if err := run(ctx, cfg, log, logOut, false, configPath); err != nil {
log.Error("listener failed", "err", err)
os.Exit(1)
}
@@ -336,23 +373,35 @@ func forceUpdateCommand(configPath string, elevatedChild bool) int {
fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot locate executable: %v\n", err)
return 1
}
// One-shot CLI: the updater logs to stderr, never to the service's
// LOG_FILE — that file is ACL'd to the service account, and a CLI run
// has nothing worth persisting there.
cfg.LogFile = ""
log, _, logCloser := newLogger(cfg)
defer logCloser.Close()
if cfg.AppVersion == "dev" {
fmt.Printf("%s: APP_VER=dev, updates disabled\n", versionLine())
return 0
}
// A running service can do the privileged work itself (its account owns
// the install dir and it knows when the GPU is idle): ask it over the
// local control channel first, no admin rights needed. Fails fast when
// no service is listening, in which case we do the direct check below.
if reply, err := control.Ask(control.CmdUpdateNow); err == nil {
return printControlReply(reply)
}
u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
// Single-shot: one attempt, fail fast when the server is unreachable
// instead of hanging in a TCP connect for minutes.
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
staged, err := u.Check(ctx, exePath)
staged, to, err := u.Check(ctx, exePath)
if err != nil && isPermission(err) && !service.Elevated() {
code, _ := elevateAndMirror("--force-update")
if code == 0 {
fmt.Println("update applied (elevated)")
reportElevatedUpdate(code, to)
if code == 0 || code == exitCodeStaged {
return 0
}
return code
}
@@ -364,12 +413,13 @@ func forceUpdateCommand(configPath string, elevatedChild bool) int {
fmt.Printf("%s is up to date\n", versionLine())
return 0
}
fmt.Printf("%s: update staged\n", versionLine())
fmt.Printf("gpu-turnstile: updated from %s to %s\n", version, to)
restarted, err := service.RestartIfRunning()
if err != nil && isPermission(err) && !service.Elevated() {
code, _ := elevateAndMirror("--force-update")
if code == 0 {
fmt.Println("update applied (elevated)")
reportElevatedUpdate(code, to)
if code == 0 || code == exitCodeStaged {
return 0
}
return code
}
@@ -378,13 +428,69 @@ func forceUpdateCommand(configPath string, elevatedChild bool) int {
return 1
}
if restarted {
fmt.Println("service restarted on the new version")
fmt.Println("service restarted on " + to)
} else {
fmt.Println("no running service; the new version applies on next start")
fmt.Println("no running service; " + to + " applies on the next start")
}
if elevatedChild {
// Tell the non-elevated parent (which cannot see this console)
// whether anything was staged, so its mirror message is honest.
return exitCodeStaged
}
return 0
}
// printControlReply prints the service's answer to a control-channel
// request: "OK ..." on stdout (exit 0), "ERR ..." on stderr (exit 1).
func printControlReply(reply string) int {
if msg, ok := strings.CutPrefix(reply, "ERR "); ok {
fmt.Fprintf(os.Stderr, "gpu-turnstile: %s\n", msg)
return 1
}
fmt.Println("gpu-turnstile: " + strings.TrimPrefix(reply, "OK "))
return 0
}
// updateNowCommand only goes through the running service's control channel
// (no direct check, no elevation): the unprivileged update trigger.
func updateNowCommand() int {
reply, err := control.Ask(control.CmdUpdateNow)
if err != nil {
fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: no running service to ask — use --force-update for a direct check\n", versionLine())
return 1
}
return printControlReply(reply)
}
// reloadEnvCommand asks the running service, over the control channel, to
// re-read and validate its config file. The client sends no path — the
// service only ever re-reads its own configured file.
func reloadEnvCommand() int {
reply, err := control.Ask(control.CmdReloadEnv)
if err != nil {
fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: no running service to ask\n", versionLine())
return 1
}
return printControlReply(reply)
}
// reportElevatedUpdate prints the parent's summary of an elevated
// --force-update child: exitCodeStaged means the child staged a new binary,
// 0 means it found nothing to do. to is the tag the parent's own check
// resolved before it hit the permission wall ("" when it never got that
// far).
func reportElevatedUpdate(code int, to string) {
switch {
case code == exitCodeStaged && to != "":
fmt.Printf("gpu-turnstile: updated from %s to %s (elevated)\n", version, to)
case code == exitCodeStaged:
fmt.Println("update applied (elevated)")
case code == 0:
fmt.Printf("%s is up to date\n", versionLine())
}
// Non-zero, non-staged codes: elevateAndMirror already printed the failure.
}
// serviceCommand installs (copyBin = register the canonical-layout copy)
// or removes the service and reports the result. On Windows, when the
// shell is not elevated, the command relaunches itself through a UAC
@@ -448,7 +554,7 @@ func managedComfyCommand(cfg config.Config) string {
return ""
}
func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Writer, isService bool) error {
func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Writer, isService bool, configPath string) error {
// ComfyUI can run as a managed child — COMFY_CMD verbatim, or the
// standard venv layout derived from COMFY_DIR alone: started on demand
// by the proxy, stopped after COMFY_IDLE_TIMEOUT idle (and on shutdown)
@@ -498,6 +604,8 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
)
lk := lock.New(log)
health := newHealthTracker()
started := time.Now()
// Each consumer is enabled by setting its URL; a disabled consumer gets
// no client, no listener and no probe.
var ollamaClient *ollama.Client
@@ -607,20 +715,24 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
}
probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout)
for name, probe := range probes {
if err := probe(probeCtx); err != nil && !errors.Is(err, errManagedDown) {
err := probe(probeCtx)
health.set(name, err == nil)
if err != nil && !errors.Is(err, errManagedDown) {
log.Warn(name+" probe failed", "err", err)
}
}
probeCancel()
if cfg.HealthInterval > 0 {
go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes)
go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes, health)
}
// Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS
// and/or GPU_FOREIGN_VRAM_MB — hold the lock externally while they run.
// gw collects the VRAM reading for the status channel.
gw := &gpuWatch{}
if len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0 {
det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUIgnoreProcs, log)
go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup)
go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup, gw)
}
// Bind the listeners up front so a port conflict fails fast and the
@@ -658,8 +770,49 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
service.NotifyReady()
service.StartWatchdog(ctx)
var u *update.Updater
var exePath string
// restartWhenIdle exits with exitCodeUpdate once the GPU lock is idle;
// the service recovery configuration brings the process back. Shared by
// staged updates and config reloads; the once guard makes repeated
// triggers idempotent.
var restartOnce sync.Once
restartWhenIdle := func(reason string) {
log.Warn("restarting once the GPU is idle", "reason", reason)
restartOnce.Do(func() {
go func() {
if waitForIdle(ctx, lk, 24*time.Hour) {
log.Warn("restarting now", "reason", reason)
os.Exit(exitCodeUpdate)
}
}()
})
}
applyStaged := func(to string) {}
if cfg.AutoUpdate {
go updateLoop(ctx, cfg, log, lk, isService)
if p, err := os.Executable(); err != nil {
log.Warn("auto-update disabled: cannot locate executable", "err", err)
} else {
exePath = p
u = &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
applyStaged = func(to string) {
if !isService {
log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply", "version", to)
return
}
restartWhenIdle("update to " + to)
}
go updateLoop(ctx, cfg.UpdateInterval, log, u, exePath, applyStaged)
}
}
// The control channel (status for --monitor, update-now and reload-env
// triggers) is served whenever running as a service, independent of
// AUTO_UPDATE.
if isService {
serveControl(ctx, log, u, exePath, applyStaged,
statusProvider(cfg, lk, comfySup, health, started, gw),
reloadHandler(cfg, configPath, restartWhenIdle))
}
select {
@@ -683,12 +836,34 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
// (idle); health checks skip it instead of logging an outage.
var errManagedDown = errors.New("managed upstream intentionally stopped")
// gpuWatch records the latest VRAM reading from the game detector's poll
// loop, for the status channel. Known stays false when game detection is
// not configured (no nvidia-smi polling happens then).
type gpuWatch struct {
mu sync.Mutex
usedMB int
total int
known bool
}
func (g *gpuWatch) set(used, total int) {
g.mu.Lock()
g.usedMB, g.total, g.known = used, total, true
g.mu.Unlock()
}
func (g *gpuWatch) get() (used, total int, known bool) {
g.mu.Lock()
defer g.mu.Unlock()
return g.usedMB, g.total, g.known
}
// gameLoop polls for foreign GPU holders (a game, another ML job). While one
// is detected it holds the lock externally so new LLM and image requests
// wait (or are rejected per LLM_BUSY_MODE), and — once in-flight work has
// drained — frees VRAM for it: the managed ComfyUI is stopped and Ollama's
// resident models are unloaded.
func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *game.Detector, lk *lock.Lock, ollamaClient *ollama.Client, comfySup *supervise.Process) {
func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *game.Detector, lk *lock.Lock, ollamaClient *ollama.Client, comfySup *supervise.Process, gw *gpuWatch) {
ticker := time.NewTicker(cfg.GamePollInterval)
defer ticker.Stop()
held, freed := false, false
@@ -702,6 +877,9 @@ func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *gam
if err != nil && ctx.Err() == nil {
log.Warn("game detection failed", "err", err)
}
if used, total, verr := game.QueryVRAMMB(ctx); verr == nil {
gw.set(used, total)
}
switch {
case len(holders) > 0 && !held:
held = true
@@ -761,7 +939,8 @@ func freeVRAM(ctx context.Context, unloadTimeout time.Duration, log *slog.Logger
// transitions — "is DOWN" when a previously healthy upstream stops
// answering, "recovered" when it comes back. The first round only
// establishes the baseline; the startup probe already reported that state.
func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *slog.Logger, probes map[string]func(context.Context) error) {
// Every result goes into the tracker for the status channel.
func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *slog.Logger, probes map[string]func(context.Context) error, tracker *healthTracker) {
ticker := time.NewTicker(interval)
defer ticker.Stop()
up := map[string]bool{}
@@ -778,9 +957,11 @@ func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *
cancel()
if errors.Is(err, errManagedDown) {
managed[name] = true // intentionally stopped; not an outage
tracker.set(name, false)
continue
}
now := err == nil
tracker.set(name, now)
if managed[name] {
// First real probe after an idle stop only re-baselines —
// an on-demand start is not a "recovery".
@@ -801,42 +982,207 @@ func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *
}
}
// updateLoop checks for signed updates on startup and every UPDATE_INTERVAL.
// In service mode a staged update is applied by exiting with exitCodeUpdate
// once the GPU lock is idle; the service recovery configuration restarts the
// process with the new binary. Interactively it only logs.
func updateLoop(ctx context.Context, cfg config.Config, log *slog.Logger, lk *lock.Lock, isService bool) {
exePath, err := os.Executable()
if err != nil {
log.Warn("auto-update disabled: cannot locate executable", "err", err)
return
}
u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
// updateLoop checks for signed updates on startup and every interval;
// applyStaged decides what a staged update means (restart when idle as a
// service, log only interactively).
func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string)) {
for {
staged, err := u.Check(ctx, exePath)
staged, to, err := u.Check(ctx, exePath)
if err != nil && ctx.Err() == nil {
log.Warn("auto-update check failed", "err", err)
}
if staged {
if !isService {
log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply")
return
}
log.Warn("auto-update: staged; restarting once the GPU is idle")
if waitForIdle(ctx, lk, 24*time.Hour) {
log.Warn("auto-update: restarting to apply update")
os.Exit(exitCodeUpdate)
}
applyStaged(to)
return
}
select {
case <-ctx.Done():
return
case <-time.After(cfg.UpdateInterval):
case <-time.After(interval):
}
}
}
// serveControl opens the local control channel (named pipe on Windows,
// unix socket on Linux) so unprivileged local users can query status
// (--monitor), trigger an update check (--force-update/--update-now) and
// poke a config reload (--reload-env) without admin rights. The update
// payload is signature-verified regardless of who asks; update triggers
// are rate-limited to one per minute so the channel cannot be used to spam
// restarts. u is nil when AUTO_UPDATE=false.
func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string), status func() string, reload func() string) {
var mu sync.Mutex
var lastTrigger time.Time
h := func(cmd string) string {
switch cmd {
case control.CmdStatus:
return "OK " + status()
case control.CmdReloadEnv:
return reload()
case control.CmdUpdateNow:
default:
return "ERR unknown command: " + cmd
}
if u == nil {
return "ERR auto-update is disabled on this instance"
}
mu.Lock()
if wait := time.Minute - time.Since(lastTrigger); wait > 0 {
mu.Unlock()
return fmt.Sprintf("ERR rate limited: retry in %ds", int(wait.Seconds())+1)
}
lastTrigger = time.Now()
mu.Unlock()
cctx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
staged, to, err := u.Check(cctx, exePath)
if err != nil {
return "ERR update check failed: " + err.Error()
}
if !staged {
return "OK " + version + " is up to date"
}
applyStaged(to)
return "OK updated from " + version + " to " + to + "; the service restarts once the GPU is idle"
}
if err := control.Serve(ctx, h, log); err != nil {
log.Warn("control channel disabled", "err", err)
}
}
// healthTracker records the latest probe result per upstream for the
// status channel.
type healthTracker struct {
mu sync.Mutex
up map[string]bool
}
func newHealthTracker() *healthTracker {
return &healthTracker{up: map[string]bool{}}
}
func (h *healthTracker) set(name string, up bool) {
h.mu.Lock()
h.up[name] = up
h.mu.Unlock()
}
func (h *healthTracker) get(name string) bool {
h.mu.Lock()
defer h.mu.Unlock()
return h.up[name]
}
// statusDownstream/statusLock/statusSnapshot are the JSON the control
// channel serves on CmdStatus; the monitor mode renders them.
type statusDownstream struct {
Name string `json:"name"`
URL string `json:"url"`
Up bool `json:"up"`
Managed string `json:"managed,omitempty"`
}
type statusLock struct {
State string `json:"state"`
Detail string `json:"detail,omitempty"`
LLMInflight int `json:"llm_inflight"`
LLMWaiting int `json:"llm_waiting"`
ImageQueue int `json:"image_queue"`
External string `json:"external,omitempty"`
SinceS int64 `json:"since_s"`
}
type statusSnapshot struct {
Version string `json:"version"`
UptimeS int64 `json:"uptime_s"`
Downstreams []statusDownstream `json:"downstreams"`
Lock statusLock `json:"lock"`
// GPU carries the latest VRAM reading; Known is false when game
// detection (and with it nvidia-smi polling) is not configured.
GPU statusGPU `json:"gpu"`
// MonitorNote is set client-side (never over the wire) when the
// monitor's own binary differs from the service's version.
MonitorNote string `json:"-"`
}
type statusGPU struct {
UsedMB int `json:"used_mb"`
TotalMB int `json:"total_mb"`
Known bool `json:"known"`
}
// reloadHandler re-reads and validates the service's config file for
// CmdReloadEnv. An invalid config is reported and the service keeps running
// untouched; a valid, changed config triggers a GPU-idle-gated restart onto
// it (same mechanism as staged updates); unchanged is a no-op.
func reloadHandler(current config.Config, configPath string, restartWhenIdle func(reason string)) func() string {
return func() string {
ncfg, err := loadMergedConfig(configPath)
if err != nil {
return "ERR config invalid: " + err.Error()
}
changed := diffConfig(current, ncfg)
if len(changed) == 0 {
return "OK config unchanged"
}
restartWhenIdle("config reload (" + strings.Join(changed, ", ") + ")")
return "OK config valid; restarting once the GPU is idle (changed: " + strings.Join(changed, ", ") + ")"
}
}
// diffConfig lists the names of fields whose values differ between two
// configs.
func diffConfig(a, b config.Config) []string {
va, vb := reflect.ValueOf(a), reflect.ValueOf(b)
t := va.Type()
var out []string
for i := 0; i < t.NumField(); i++ {
if !reflect.DeepEqual(va.Field(i).Interface(), vb.Field(i).Interface()) {
out = append(out, t.Field(i).Name)
}
}
return out
}
// statusProvider assembles the one-line JSON snapshot for CmdStatus.
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time, gw *gpuWatch) func() string {
return func() string {
snap := statusSnapshot{
Version: version,
UptimeS: int64(time.Since(started).Seconds()),
}
used, total, known := gw.get()
snap.GPU = statusGPU{UsedMB: used, TotalMB: total, Known: known}
if cfg.OllamaURL != "" {
snap.Downstreams = append(snap.Downstreams, statusDownstream{
Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
})
}
if cfg.ComfyURL != "" {
d := statusDownstream{Name: "comfy", URL: cfg.ComfyURL, Up: health.get("comfy")}
if comfySup != nil {
d.Managed = comfySup.Status()
}
snap.Downstreams = append(snap.Downstreams, d)
}
st := lk.Status()
snap.Lock = statusLock{
State: string(st.State),
Detail: st.Detail,
LLMInflight: st.LLMInflight,
LLMWaiting: st.LLMWaiting,
ImageQueue: st.ImageQueue,
External: st.External,
SinceS: int64(time.Since(st.Since).Seconds()),
}
b, err := json.Marshal(snap)
if err != nil {
return `{"version":"` + version + `"}`
}
return string(b)
}
}
// waitForIdle polls the lock until no LLM or image work is active or
// pending, max at most. Returns false on timeout or cancellation.
func waitForIdle(ctx context.Context, lk *lock.Lock, max time.Duration) bool {
+283
View File
@@ -0,0 +1,283 @@
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" + 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.
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)
updatePending := false
poll := func() string {
frame := renderWaiting()
if reply, err := control.Ask(control.CmdStatus); err == nil {
if msg, ok := strings.CutPrefix(reply, "OK "); ok {
var snap statusSnapshot
if json.Unmarshal([]byte(msg), &snap) == nil {
if snap.Version != "" && snap.Version != version {
if exeChanged(exe, exeStamp) {
fmt.Print("\x1b[2J\x1b[H")
fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version)
restartSelf(exe, "--monitor")
return "" // re-execed; this process exits below
}
snap.MonitorNote = fmt.Sprintf("note: the service runs %s, this monitor is %s", snap.Version, version)
}
if note != "" {
snap.MonitorNote = note
}
frame = renderMonitor(snap, termWidth())
}
}
}
return frame
}
for {
frame := poll()
if frame == "" {
return 0 // restartSelf fired
}
fmt.Print("\x1b[H" + frame + "\x1b[J") // home, frame, clear below
select {
case <-ticker.C:
if note != "" && time.Since(noteAt) > 15*time.Second {
note = ""
}
case k, ok := <-keys:
if !ok {
keys = nil
continue
}
switch k {
case 'q', 'Q', 3: // q or Ctrl+C (raw mode delivers it as a byte)
return 0
case 'u', 'U':
if !updatePending {
updatePending = true
note, noteAt = "checking for updates…", time.Now()
go func() {
reply, err := control.Ask(control.CmdUpdateNow)
if err != nil {
noteCh <- "update: no answer from the service"
return
}
msg := strings.TrimPrefix(reply, "OK ")
msg = strings.TrimPrefix(msg, "ERR ")
noteCh <- "update: " + msg
}()
}
}
case n := <-noteCh:
updatePending = false
note, noteAt = n, time.Now()
}
}
}
// readKeys reads single keypresses from stdin (raw mode was enabled by the
// caller) and delivers them until stdin fails.
func readKeys(keys chan<- byte) {
defer close(keys)
buf := make([]byte, 1)
for {
n, err := os.Stdin.Read(buf)
if n > 0 {
keys <- buf[0]
}
if err != nil {
return
}
}
}
// exeChanged reports whether the executable on disk was replaced since the
// recorded stamp (the updater swaps it via rename, which changes ModTime).
func exeChanged(exe string, stamp time.Time) bool {
if exe == "" || stamp.IsZero() {
return false
}
st, err := os.Stat(exe)
return err == nil && !st.ModTime().Equal(stamp)
}
// restartSelf starts a fresh copy of this executable with the given args on
// the same console; the caller exits right after.
func restartSelf(exe string, args ...string) {
cmd := exec.Command(exe, args...)
cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
cmd.Start() //nolint:errcheck // best effort: on failure we just exit
}
func renderWaiting() string {
return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n\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 {
b.WriteString(renderDownstream(d) + "\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.Known {
b.WriteString(" GPU: " + renderVRAM(snap.GPU.UsedMB, snap.GPU.TotalMB) + "\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()
}
// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
func renderVRAM(used, total int) string {
format := func(mb int) string {
if mb >= 1024 {
return fmt.Sprintf("%.1f GiB", float64(mb)/1024)
}
return fmt.Sprintf("%d MiB", mb)
}
if total > 0 {
return format(used) + " / " + format(total) + " used"
}
return format(used) + " used"
}
// printableLen counts characters without ANSI escapes (ASCII-only content).
func printableLen(s string) int { return len(s) }
func renderDownstream(d statusDownstream) string {
url := cDim + d.URL + 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", cGreen, cReset, d.Name, cGreen, cReset, suffix, 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()
}
+56
View File
@@ -0,0 +1,56 @@
package main
import (
"gpu-turnstile/internal/config"
"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)
}
}
}
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] != "GameProcs" || got[1] != "LogLevel" {
t.Fatalf("got %v, want [GameProcs LogLevel]", got)
}
}
+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).
services:
gpu-turnstile:
image: git.rambossek.at/public/gpu-turnstile:v0.2.0
image: git.rambossek.at/public/gpu-turnstile:v0.2.8
restart: unless-stopped
environment:
# Each consumer is enabled by setting its URL; leave one unset to
+68
View File
@@ -0,0 +1,68 @@
// 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.
//
// 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"
)
// 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
// serveConn runs the line protocol on one accepted connection.
func serveConn(c io.ReadWriteCloser, h Handler) {
defer c.Close()
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
}
+54
View File
@@ -0,0 +1,54 @@
//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
}
go serveConn(c, h)
}
}()
return nil
}
// 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
}
+46
View File
@@ -0,0 +1,46 @@
package control
import (
"errors"
"net"
"strings"
"testing"
)
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)
}
}
+135
View File
@@ -0,0 +1,135 @@
//go:build windows
package control
import (
"context"
"fmt"
"log/slog"
"os"
"syscall"
"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
}
go serveConn(&pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}, h)
}
}()
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
// 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()
}
// Ask sends one command to the running service and returns its reply.
func Ask(cmd string) (string, error) {
name, err := windows.UTF16PtrFromString(pipePath)
if err != nil {
return "", err
}
if err := waitNamedPipe(name, 2000); err != nil {
return "", ErrUnavailable
}
handle, err := windows.CreateFile(name,
windows.GENERIC_READ|windows.GENERIC_WRITE, 0, nil,
windows.OPEN_EXISTING, 0, 0)
if err != nil {
return "", ErrUnavailable
}
f := os.NewFile(uintptr(handle), pipePath)
defer f.Close()
return readReply(f, cmd)
}
+23
View File
@@ -132,6 +132,29 @@ func queryComputeApps(ctx context.Context) ([]computeApp, error) {
return parseComputeApps(string(out))
}
// QueryVRAMMB returns used and total GPU VRAM in MiB via nvidia-smi.
// Unlike the per-process list this works under WDDM too.
func QueryVRAMMB(ctx context.Context) (used, total int, err error) {
out, err := exec.CommandContext(ctx, "nvidia-smi",
"--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits").Output()
if err != nil {
return 0, 0, err
}
usedStr, totalStr, ok := strings.Cut(strings.TrimSpace(string(out)), ",")
if !ok {
return 0, 0, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(string(out)))
}
used, err = strconv.Atoi(strings.TrimSpace(usedStr))
if err != nil {
return 0, 0, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(string(out)))
}
total, err = strconv.Atoi(strings.TrimSpace(totalStr))
if err != nil {
return 0, 0, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(string(out)))
}
return used, total, nil
}
// parseComputeApps parses "pid, used_memory" CSV lines (no header, MiB
// units). Unsupported rows ("N/A" on WDDM) are skipped.
func parseComputeApps(out string) ([]computeApp, error) {
+75 -2
View File
@@ -9,6 +9,7 @@ import (
"context"
"log/slog"
"sync"
"time"
)
// State is the current GPU occupancy state.
@@ -30,10 +31,13 @@ type Lock struct {
change chan struct{} // closed and replaced on every state change
n int // LLM requests in flight
llmWaiting int // LLM requests blocked waiting for the GPU
imageActive bool // an image job holds the GPU
imageQ []imageWaiter
nextID uint64
external string // non-empty: a foreign process (e.g. a game) holds the GPU
external string // non-empty: a foreign process (e.g. a game) holds the GPU
detail string // what the current holder is doing (best effort)
since time.Time // when the current state began
log *slog.Logger
}
@@ -41,7 +45,7 @@ type Lock struct {
// New returns a ready-to-use Lock. log may be nil; if set, every state
// transition is logged at debug level.
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.
@@ -63,6 +67,7 @@ func (l *Lock) logTransition(msg string, args ...any) {
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)
@@ -73,6 +78,7 @@ func (l *Lock) SetExternal(holder string) {
func (l *Lock) ClearExternal() {
l.mu.Lock()
l.external = ""
l.since = time.Now()
l.broadcast()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateIdle)
@@ -90,17 +96,31 @@ func (l *Lock) External() string {
// while waiting; no state is changed in that case.
func (l *Lock) AcquireLLM(ctx context.Context) error {
l.mu.Lock()
waiting := false
for l.imageActive || len(l.imageQ) > 0 || l.external != "" {
if !waiting {
l.llmWaiting++
waiting = true
}
ch := l.change
l.mu.Unlock()
select {
case <-ctx.Done():
l.mu.Lock()
l.llmWaiting--
l.mu.Unlock()
return ctx.Err()
case <-ch:
}
l.mu.Lock()
}
if waiting {
l.llmWaiting--
}
l.n++
if l.n == 1 {
l.since = time.Now()
}
n := l.n
l.mu.Unlock()
l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
@@ -129,6 +149,8 @@ func (l *Lock) ReleaseLLM() {
l.n--
n := l.n
if l.n == 0 {
l.since = time.Now()
l.detail = ""
l.broadcast()
}
l.mu.Unlock()
@@ -156,6 +178,7 @@ func (l *Lock) AcquireImage(ctx context.Context) error {
if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive && l.external == "" {
l.imageQ = l.imageQ[1:]
l.imageActive = true
l.since = time.Now()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateImage)
return nil
@@ -184,6 +207,8 @@ func (l *Lock) AcquireImage(ctx context.Context) error {
func (l *Lock) ReleaseImage() {
l.mu.Lock()
l.imageActive = false
l.since = time.Now()
l.detail = ""
l.broadcast()
l.mu.Unlock()
l.logTransition("lock transition", "state", StateIdle)
@@ -206,3 +231,51 @@ func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) {
}
return state, l.n, l.imageActive || len(l.imageQ) > 0
}
// 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
}
+2
View File
@@ -491,6 +491,7 @@ func (s *Server) OllamaHandler() http.Handler {
return
}
defer s.cfg.Lock.ReleaseLLM()
s.cfg.Lock.SetDetail("ollama: " + r.Method + " " + r.URL.Path)
s.ollamaProxy.ServeHTTP(w, r)
}))
}
@@ -590,6 +591,7 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
}
s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds())
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 {
+15 -2
View File
@@ -15,6 +15,8 @@ import (
"os/signal"
"path/filepath"
"syscall"
"gpu-turnstile/internal/supervise"
)
// Name matches the Windows service name; the systemd unit is Name + ".service".
@@ -64,11 +66,20 @@ func Run(run func(ctx context.Context) error) error {
// 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.
// 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]
Description=gpu-turnstile GPU arbitration proxy for Ollama and ComfyUI
@@ -84,6 +95,8 @@ RestartSec=5s
DynamicUser=yes
StateDirectory=%s
RuntimeDirectory=%s
RuntimeDirectoryMode=0755
%sProtectSystem=strict
ProtectHome=yes
PrivateTmp=yes
@@ -106,7 +119,7 @@ SystemCallErrorNumber=EPERM
[Install]
WantedBy=multi-user.target
`, exePath, configPath, Name, bind)
`, exePath, configPath, Name, Name, bind)
}
// copyFile copies src to dst, creating dst with the given mode.
+1
View File
@@ -17,6 +17,7 @@ func TestRenderUnit(t *testing.T) {
"WantedBy=multi-user.target",
"DynamicUser=yes",
"StateDirectory=gpu-turnstile",
"RuntimeDirectory=gpu-turnstile",
"ProtectSystem=strict",
"NoNewPrivileges=yes",
"RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6",
+17 -1
View File
@@ -24,6 +24,8 @@ import (
"golang.org/x/sys/windows"
"golang.org/x/sys/windows/svc"
"golang.org/x/sys/windows/svc/mgr"
"gpu-turnstile/internal/supervise"
)
// Name is the Windows service name.
@@ -375,6 +377,20 @@ func grantAll(exe, configPath string) error {
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
@@ -496,7 +512,7 @@ func grantAccessTree(path, perms string) error {
func runIcacls(path, perms string, recursive bool) error {
args := []string{path, "/grant", virtualAccount + ":" + perms}
if recursive {
fmt.Printf("granting %s modify access to %s (large trees can take minutes)\n", virtualAccount, path)
fmt.Printf("granting %s %s access to %s (large trees can take minutes)\n", virtualAccount, perms, path)
args = append(args, "/T")
}
start := time.Now()
+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)
}
}
+62 -4
View File
@@ -13,6 +13,7 @@ import (
"os"
"os/exec"
"path/filepath"
"regexp"
"runtime"
"strings"
"sync"
@@ -43,7 +44,9 @@ func ComfyLayout(goos, dir string) (python, script string) {
// 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.
// 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)
@@ -54,9 +57,34 @@ func DefaultComfyCommand(goos, dir, comfyURL string) string {
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
@@ -136,6 +164,24 @@ func (p *Process) Running() bool {
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 {
@@ -185,6 +231,9 @@ func (p *Process) EnsureRunning() error {
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
@@ -278,7 +327,9 @@ func (p *Process) WatchIdle(ctx context.Context, idleTimeout time.Duration, gpuI
}
// pipeLog forwards one child output stream to the log at INFO, line by
// line, prefixed with the process name.
// 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
@@ -290,18 +341,25 @@ func (p *Process) pipeLog(r io.Reader) {
if i < 0 {
break
}
p.log.Info(p.name + ": " + strings.TrimRight(line[:i], "\r"))
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 + ": " + 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) {
+35
View File
@@ -239,3 +239,38 @@ func TestComfyLayoutAndDefaultCommand(t *testing.T) {
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)
}
}
+25 -22
View File
@@ -173,11 +173,13 @@ func CleanupOld(exePath string) {
// Check performs a single update check. staged is true when a
// signature-verified binary has been swapped into place at exePath; the
// caller should then restart the process. A nil error with 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, err error) {
// caller should then restart the process. to is the release tag the check
// resolved (the latest release or the pinned tag), set once the release
// fetch succeeded — even when staging afterwards fails. A nil error with
// 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()
desired := u.Desired
if desired == "" {
@@ -185,15 +187,15 @@ func (u *Updater) Check(ctx context.Context, exePath string) (staged bool, err e
}
if desired == "dev" {
log.Debug("auto-update: APP_VER=dev, skipping")
return false, nil
return false, "", nil
}
if publicKeyPEM == "" {
log.Debug("auto-update: no public key embedded, skipping")
return false, nil
return false, "", nil
}
api, err := u.apiURL()
if err != nil {
return false, err
return false, "", err
}
pinned := desired != "stable"
@@ -203,29 +205,30 @@ func (u *Updater) Check(ctx context.Context, exePath string) (staged bool, err e
}
body, err := u.get(ctx, endpoint)
if err != nil {
return false, fmt.Errorf("fetch release: %w", err)
return false, "", fmt.Errorf("fetch release: %w", err)
}
var rel release
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, nil
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)
if err != nil {
return false, err
return false, to, err
}
if !newer {
log.Debug("auto-update: up to date", "version", u.Version, "latest", rel.TagName)
return false, nil
return false, to, nil
}
}
@@ -235,42 +238,42 @@ func (u *Updater) Check(ctx context.Context, exePath string) (staged bool, err e
}
assetURL, ok := urls[u.Asset]
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"]
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)
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)
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 {
sumText, err := u.get(ctx, sumURL)
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]
got := hex.EncodeToString(sha256Bytes(data))
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 {
return false, err
return false, to, err
}
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)
return true, nil
return true, to, nil
}
func sha256Bytes(data []byte) []byte {
+14 -8
View File
@@ -102,13 +102,16 @@ func TestCheckStagesUpdate(t *testing.T) {
withPublicKey(t, f.pubPEM)
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 {
t.Fatal(err)
}
if !staged {
t.Fatal("expected staged update")
}
if to != "v9.9.9" {
t.Fatalf("to = %q, want v9.9.9", to)
}
content, _ := os.ReadFile(exe)
if string(content) != "new-binary" {
t.Fatalf("exe content = %q", content)
@@ -125,7 +128,7 @@ func TestCheckRejectsTamperedSignature(t *testing.T) {
withPublicKey(t, f.pubPEM)
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 {
t.Fatal("expected signature error")
}
@@ -142,7 +145,7 @@ func TestCheckSkipsOlderOrEqual(t *testing.T) {
for _, tag := range []string{"v0.1.2", "v0.1.1", "v0.0.9"} {
f := newFakeGitea(t, tag, []byte("new-binary"))
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 {
t.Fatal(err)
}
@@ -155,7 +158,7 @@ func TestCheckSkipsOlderOrEqual(t *testing.T) {
func TestCheckSkipsWithoutPublicKey(t *testing.T) {
f := newFakeGitea(t, "v9.9.9", []byte("new-binary"))
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 {
t.Fatalf("staged=%v err=%v, want no action without key", staged, err)
}
@@ -167,7 +170,7 @@ func TestCheckDevBuildGetsStable(t *testing.T) {
f := newFakeGitea(t, "v9.9.9", []byte("new-binary"))
withPublicKey(t, f.pubPEM)
exe := fakeExe(t)
staged, err := f.updater("dev").Check(context.Background(), exe)
staged, _, err := f.updater("dev").Check(context.Background(), exe)
if err != nil {
t.Fatal(err)
}
@@ -184,7 +187,7 @@ 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))
staged, _, err := f.updaterDesired(version, "dev").Check(context.Background(), fakeExe(t))
if err != nil || staged {
t.Fatalf("version %s: staged=%v err=%v, want no action with APP_VER=dev", version, staged, err)
}
@@ -198,13 +201,16 @@ func TestCheckPinned(t *testing.T) {
f := newFakeGitea(t, "v0.5.0", []byte("pinned-binary"))
withPublicKey(t, f.pubPEM)
exe := fakeExe(t)
staged, err := f.updaterDesired(version, "v0.5.0").Check(context.Background(), exe)
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)
@@ -213,7 +219,7 @@ func TestCheckPinned(t *testing.T) {
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))
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)
}