From 8fccd333aabeeeee41d7afed0572a7bc4f4af3c3 Mon Sep 17 00:00:00 2001 From: mram Date: Tue, 22 Sep 2026 09:24:27 +0200 Subject: [PATCH] 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 --- cmd/gpu-turnstile/main.go | 86 +++++++++++++++++++++-------- cmd/gpu-turnstile/monitor.go | 24 +++++++- cmd/gpu-turnstile/monitor_test.go | 33 ++++++++++- internal/config/config.go | 78 +++++++++++++------------- internal/control/control.go | 49 ++++++++++++++++ internal/control/control_linux.go | 14 ++++- internal/control/control_test.go | 57 +++++++++++++++++++ internal/control/control_windows.go | 13 ++++- 8 files changed, 288 insertions(+), 66 deletions(-) diff --git a/cmd/gpu-turnstile/main.go b/cmd/gpu-turnstile/main.go index 0630090..da7a31f 100644 --- a/cmd/gpu-turnstile/main.go +++ b/cmd/gpu-turnstile/main.go @@ -728,9 +728,10 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri // 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 { + // gw collects the VRAM reading and the last check result for the status + // channel. + gw := &gpuWatch{enabled: len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0} + if gw.enabled { det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUIgnoreProcs, log) go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup, gw) } @@ -836,26 +837,40 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri // (idle); health checks skip it instead of logging an outage. var errManagedDown = errors.New("managed upstream intentionally stopped") -// gpuWatch records the latest VRAM reading from the game detector's poll -// loop, for the status channel. Known stays false when game detection is -// not configured (no nvidia-smi polling happens then). +// gpuWatch records the latest VRAM reading and game-detector finding from +// the game detector's poll loop, for the status channel. Enabled is false +// when game detection is not configured (no polling happens then); Known is +// false until the first successful nvidia-smi reading. type gpuWatch struct { - mu sync.Mutex - usedMB int - total int - known bool + mu sync.Mutex + enabled bool + usedMB int + total int + known bool + foreign string // last Check result: external holders, "" when none + at time.Time } -func (g *gpuWatch) set(used, total int) { +func (g *gpuWatch) setVRAM(used, total int) { g.mu.Lock() g.usedMB, g.total, g.known = used, total, true g.mu.Unlock() } -func (g *gpuWatch) get() (used, total int, known bool) { +func (g *gpuWatch) setCheck(foreign string) { + g.mu.Lock() + g.foreign, g.at = foreign, time.Now() + g.mu.Unlock() +} + +func (g *gpuWatch) get() (used, total int, known bool, foreign string, ageS int64) { g.mu.Lock() defer g.mu.Unlock() - return g.usedMB, g.total, g.known + ageS = -1 + if !g.at.IsZero() { + ageS = int64(time.Since(g.at).Seconds()) + } + return g.usedMB, g.total, g.known, g.foreign, ageS } // gameLoop polls for foreign GPU holders (a game, another ML job). While one @@ -877,8 +892,9 @@ func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *gam if err != nil && ctx.Err() == nil { log.Warn("game detection failed", "err", err) } + gw.setCheck(summarizeHolders(holders)) if used, total, verr := game.QueryVRAMMB(ctx); verr == nil { - gw.set(used, total) + gw.setVRAM(used, total) } switch { case len(holders) > 0 && !held: @@ -1012,12 +1028,21 @@ func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u // restarts. u is nil when AUTO_UPDATE=false. func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string), status func() string, reload func() string) { var mu sync.Mutex - var lastTrigger time.Time + var lastTrigger, lastReload time.Time h := func(cmd string) string { switch cmd { case control.CmdStatus: return "OK " + status() case control.CmdReloadEnv: + // reload re-reads the config file from disk; a short limiter + // keeps a local flood from turning into disk churn. + mu.Lock() + if wait := 2*time.Second - time.Since(lastReload); wait > 0 { + mu.Unlock() + return fmt.Sprintf("ERR rate limited: retry in %ds", int(wait.Seconds())+1) + } + lastReload = time.Now() + mu.Unlock() return reload() case control.CmdUpdateNow: default: @@ -1097,8 +1122,9 @@ type statusSnapshot struct { UptimeS int64 `json:"uptime_s"` Downstreams []statusDownstream `json:"downstreams"` Lock statusLock `json:"lock"` - // GPU carries the latest VRAM reading; Known is false when game - // detection (and with it nvidia-smi polling) is not configured. + // GPU carries the latest VRAM reading and detector finding; Enabled is + // false when game detection (and with it nvidia-smi polling) is not + // configured. GPU statusGPU `json:"gpu"` // MonitorNote is set client-side (never over the wire) when the // monitor's own binary differs from the service's version. @@ -1106,9 +1132,17 @@ type statusSnapshot struct { } type statusGPU struct { + // Enabled reports whether game detection is configured (and with it + // VRAM polling); when false the other fields carry no information. + Enabled bool `json:"enabled"` UsedMB int `json:"used_mb"` TotalMB int `json:"total_mb"` Known bool `json:"known"` + // Foreign is the last detector finding (external GPU holders), empty + // when the last check found none. + Foreign string `json:"foreign,omitempty"` + // AgeS is how long ago the last check ran; -1 before the first check. + AgeS int64 `json:"age_s"` } // reloadHandler re-reads and validates the service's config file for @@ -1130,15 +1164,19 @@ func reloadHandler(current config.Config, configPath string, restartWhenIdle fun } } -// diffConfig lists the names of fields whose values differ between two -// configs. +// diffConfig lists the env names of settings whose values differ between +// two configs (from each field's env tag, so users recognize them). func diffConfig(a, b config.Config) []string { va, vb := reflect.ValueOf(a), reflect.ValueOf(b) t := va.Type() var out []string for i := 0; i < t.NumField(); i++ { if !reflect.DeepEqual(va.Field(i).Interface(), vb.Field(i).Interface()) { - out = append(out, t.Field(i).Name) + name := t.Field(i).Tag.Get("env") + if name == "" { + name = t.Field(i).Name + } + out = append(out, name) } } return out @@ -1151,8 +1189,12 @@ func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Proces Version: version, UptimeS: int64(time.Since(started).Seconds()), } - used, total, known := gw.get() - snap.GPU = statusGPU{UsedMB: used, TotalMB: total, Known: known} + used, total, known, foreign, ageS := gw.get() + snap.GPU = statusGPU{ + Enabled: gw.enabled, + UsedMB: used, TotalMB: total, Known: known, + Foreign: foreign, AgeS: ageS, + } if cfg.OllamaURL != "" { snap.Downstreams = append(snap.Downstreams, statusDownstream{ Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"), diff --git a/cmd/gpu-turnstile/monitor.go b/cmd/gpu-turnstile/monitor.go index 4944503..d8e0066 100644 --- a/cmd/gpu-turnstile/monitor.go +++ b/cmd/gpu-turnstile/monitor.go @@ -196,8 +196,8 @@ func renderMonitor(snap statusSnapshot, width int) string { b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n", cYellow, snap.Lock.ImageQueue, cReset)) } - if snap.GPU.Known { - b.WriteString(" GPU: " + renderVRAM(snap.GPU.UsedMB, snap.GPU.TotalMB) + "\x1b[K\n") + if snap.GPU.Enabled { + b.WriteString(renderGPU(snap.GPU) + "\x1b[K\n") } if snap.MonitorNote != "" { b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n") @@ -207,6 +207,26 @@ func renderMonitor(snap statusSnapshot, width int) string { return b.String() } +// renderGPU renders the GPU line: VRAM usage (when nvidia-smi answered), +// the game detector's last finding, and how long ago it ran. +func renderGPU(g statusGPU) string { + s := " GPU: " + if g.Known { + s += renderVRAM(g.UsedMB, g.TotalMB) + } else { + s += cDim + "VRAM unknown (nvidia-smi not answering)" + cReset + } + if g.Foreign != "" { + s += " · external: " + cRed + g.Foreign + cReset + } else { + s += cDim + " · no external process" + cReset + } + if g.AgeS >= 0 { + s += cDim + fmt.Sprintf(" (checked %s ago)", fmtDur(g.AgeS)) + cReset + } + return s +} + // renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB). func renderVRAM(used, total int) string { format := func(mb int) string { diff --git a/cmd/gpu-turnstile/monitor_test.go b/cmd/gpu-turnstile/monitor_test.go index 556ce4d..b518666 100644 --- a/cmd/gpu-turnstile/monitor_test.go +++ b/cmd/gpu-turnstile/monitor_test.go @@ -2,6 +2,7 @@ package main import ( "gpu-turnstile/internal/config" + "reflect" "strings" "testing" ) @@ -30,6 +31,22 @@ func TestRenderMonitor(t *testing.T) { t.Errorf("frame missing %q:\n%s", want, frame) } } + + snap.GPU = statusGPU{Enabled: true, Known: true, UsedMB: 4300, TotalMB: 16384, Foreign: "cyberpunk2077.exe (pid 1234)", AgeS: 12} + frame = renderMonitor(snap, 80) + for _, want := range []string{"GPU:", "4.2 GiB / 16.0 GiB used", "external:", "checked 12s ago"} { + if !strings.Contains(frame, want) { + t.Errorf("frame missing %q:\n%s", want, frame) + } + } + + snap.GPU = statusGPU{Enabled: true, AgeS: 4} + frame = renderMonitor(snap, 80) + for _, want := range []string{"VRAM unknown", "no external process"} { + if !strings.Contains(frame, want) { + t.Errorf("frame missing %q:\n%s", want, frame) + } + } } func TestFmtDur(t *testing.T) { @@ -50,7 +67,19 @@ func TestDiffConfig(t *testing.T) { b.LogLevel = -4 b.GameProcs = []string{"game.exe"} got := diffConfig(a, b) - if len(got) != 2 || got[0] != "GameProcs" || got[1] != "LogLevel" { - t.Fatalf("got %v, want [GameProcs LogLevel]", got) + if len(got) != 2 || got[0] != "GAME_PROCS" || got[1] != "LOGLEVEL" { + t.Fatalf("got %v, want [GAME_PROCS LOGLEVEL]", got) + } +} + +// Every Config field must carry an env tag so user-facing output (the +// reload diff) can name the setting the user would actually change. +func TestConfigFieldsHaveEnvTags(t *testing.T) { + typ := reflect.TypeOf(config.Config{}) + for i := 0; i < typ.NumField(); i++ { + f := typ.Field(i) + if f.Tag.Get("env") == "" { + t.Errorf("config.Config.%s has no env tag", f.Name) + } } } diff --git a/internal/config/config.go b/internal/config/config.go index a8bb670..2d67044 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -13,43 +13,45 @@ import ( "time" ) -// Config holds every gpu-turnstile setting. +// Config holds every gpu-turnstile setting. Each field's env tag names the +// environment variable / config-file key that sets it; user-facing output +// (e.g. the reload diff) uses those names, never the Go field names. type Config struct { - ListenOllama string - ListenComfy string - OllamaURL string - ComfyURL string - UnloadTimeout time.Duration - JobTimeout time.Duration - LLMWaitTimeout time.Duration + ListenOllama string `env:"LISTEN_OLLAMA"` + ListenComfy string `env:"LISTEN_COMFY"` + OllamaURL string `env:"OLLAMA_URL"` + ComfyURL string `env:"COMFY_URL"` + UnloadTimeout time.Duration `env:"UNLOAD_TIMEOUT"` + JobTimeout time.Duration `env:"JOB_TIMEOUT"` + LLMWaitTimeout time.Duration `env:"LLM_WAIT_TIMEOUT"` - UnloadPollInterval time.Duration - HistoryPollInterval time.Duration - ProbeTimeout time.Duration - HealthInterval time.Duration - FreeTimeout time.Duration - WarmTimeout time.Duration - ShutdownTimeout time.Duration - BackoffInitial time.Duration - BackoffMax time.Duration - PromptCaptureLimit int64 + UnloadPollInterval time.Duration `env:"UNLOAD_POLL_INTERVAL"` + HistoryPollInterval time.Duration `env:"HISTORY_POLL_INTERVAL"` + ProbeTimeout time.Duration `env:"PROBE_TIMEOUT"` + HealthInterval time.Duration `env:"HEALTH_INTERVAL"` + FreeTimeout time.Duration `env:"FREE_TIMEOUT"` + WarmTimeout time.Duration `env:"WARM_TIMEOUT"` + ShutdownTimeout time.Duration `env:"SHUTDOWN_TIMEOUT"` + BackoffInitial time.Duration `env:"BACKOFF_INITIAL"` + BackoffMax time.Duration `env:"BACKOFF_MAX"` + PromptCaptureLimit int64 `env:"PROMPT_CAPTURE_LIMIT"` - AutoUpdate bool - UpdateInterval time.Duration - UpdateRepo string - UpdateAsset string + AutoUpdate bool `env:"AUTO_UPDATE"` + UpdateInterval time.Duration `env:"UPDATE_INTERVAL"` + UpdateRepo string `env:"UPDATE_REPO"` + UpdateAsset string `env:"UPDATE_ASSET"` // AppVersion is the version the user wants to run: "dev" disables // updates, "stable" tracks the latest release, anything else is an // exact vX.Y.Z release to pin. From APP_VER; defaults to "stable". - AppVersion string + AppVersion string `env:"APP_VER"` // LLMBusyMode is "wait" (hold requests until the lock is free or // LLMWaitTimeout expires) or "reject" (immediately answer with // LLMBusyStatus + Retry-After when an image job is active or pending). - LLMBusyMode string - LLMBusyStatus int - BusyRetryAfter int + LLMBusyMode string `env:"LLM_BUSY_MODE"` + LLMBusyStatus int `env:"LLM_BUSY_STATUS"` + BusyRetryAfter int `env:"BUSY_RETRY_AFTER"` // ComfyCmd spawns and supervises a ComfyUI server on demand. When // ComfyCmd is empty but ComfyDir is set, management is enabled with the @@ -59,10 +61,10 @@ type Config struct { // set explicitly. The managed server is stopped after ComfyIdleTimeout // without requests, freeing its VRAM; ComfyStartTimeout bounds how long // a request waits for it to come up. - ComfyCmd string - ComfyDir string - ComfyIdleTimeout time.Duration - ComfyStartTimeout time.Duration + ComfyCmd string `env:"COMFY_CMD"` + ComfyDir string `env:"COMFY_DIR"` + ComfyIdleTimeout time.Duration `env:"COMFY_IDLE_TIMEOUT"` + ComfyStartTimeout time.Duration `env:"COMFY_START_TIMEOUT"` // GameProcs (GAME_PROCS) is a watch list of process names; while any of // them runs, the GPU is treated as held by a foreign process. The @@ -70,15 +72,15 @@ type Config struct { // when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than // that many MiB of VRAM. GamePollInterval (GAME_POLL_INTERVAL) is how // often both checks run. - GameProcs []string - GPUForeignVRAMMB int - GPUIgnoreProcs []string - GamePollInterval time.Duration + GameProcs []string `env:"GAME_PROCS"` + GPUForeignVRAMMB int `env:"GPU_FOREIGN_VRAM_MB"` + GPUIgnoreProcs []string `env:"GPU_IGNORE_PROCS"` + GamePollInterval time.Duration `env:"GAME_POLL_INTERVAL"` - WarmModel string - LogLevel slog.Level - LogJSON bool - LogFile string + WarmModel string `env:"WARM_MODEL"` + LogLevel slog.Level `env:"LOGLEVEL"` + LogJSON bool `env:"LOG_FORMAT"` + LogFile string `env:"LOG_FILE"` } // Defaults returns the configuration used when neither the environment nor diff --git a/internal/control/control.go b/internal/control/control.go index 6dacab8..4c6a2ad 100644 --- a/internal/control/control.go +++ b/internal/control/control.go @@ -7,6 +7,13 @@ // triggers, so the worst a local user can cause is a cheap, throttled // check and a GPU-idle-gated restart onto a signed binary. // +// Abuse hardening: the command read is capped (4 KiB), each connection is +// force-closed after connTimeout so a stalled client cannot pin a goroutine +// (or a Windows pipe instance) forever, and concurrently served connections +// are capped at maxConns — beyond that, connections are closed on arrival. +// On Windows the pipe's ACL additionally denies network logons, so the +// channel cannot be reached from another machine. +// // Protocol: the client writes one command line, the server answers with // one reply line ("OK ..." or "ERR ...") and hangs up. package control @@ -17,6 +24,7 @@ import ( "fmt" "io" "strings" + "time" ) // CmdUpdateNow asks the service to check for, stage and (once the GPU is @@ -37,9 +45,50 @@ var ErrUnavailable = errors.New("control channel unavailable") // line. It must start with "OK " or "ERR ". type Handler func(cmd string) string +// connTimeout bounds one connection's lifetime: a client that stops +// mid-command or never reads the reply would otherwise pin its goroutine +// (and on Windows one of the pipe instances) indefinitely. A var so tests +// can shrink it. +var connTimeout = 10 * time.Second + +// maxConns caps concurrently served connections; beyond it, new +// connections are closed on arrival. Bound on the goroutines a local +// flood can pile up. +const maxConns = 32 + +var connSem = make(chan struct{}, maxConns) + +// serve dispatches connection handling under the concurrency cap. It +// returns false when the cap is reached — the caller must then close the +// connection itself. +func serve(c io.ReadWriteCloser, h Handler) bool { + select { + case connSem <- struct{}{}: + go func() { + defer func() { <-connSem }() + serveConn(c, h) + }() + return true + default: + return false + } +} + +// forceCloser is implemented by connections that can be torn down +// abortively, unblocking pending reads and writes (Windows pipe: +// DisconnectNamedPipe; unix socket: a deadline in the past). The +// connection watchdog uses it; normal closes still flush the reply. +type forceCloser interface { + ForceClose() error +} + // serveConn runs the line protocol on one accepted connection. func serveConn(c io.ReadWriteCloser, h Handler) { defer c.Close() + if fc, ok := c.(forceCloser); ok { + timer := time.AfterFunc(connTimeout, func() { fc.ForceClose() }) + defer timer.Stop() + } line, err := bufio.NewReader(io.LimitReader(c, 4096)).ReadString('\n') cmd := strings.TrimSpace(line) if cmd == "" { diff --git a/internal/control/control_linux.go b/internal/control/control_linux.go index 50b3c79..8d262b4 100644 --- a/internal/control/control_linux.go +++ b/internal/control/control_linux.go @@ -37,12 +37,24 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error { if err != nil { return // shutting down } - go serveConn(c, h) + uc := unixConn{c} + if !serve(uc, h) { + uc.ForceClose() + } } }() return nil } +// unixConn adds an abortive ForceClose to net.Conn: a deadline in the +// past fails pending and future I/O immediately. +type unixConn struct{ net.Conn } + +func (c unixConn) ForceClose() error { + c.SetDeadline(time.Now().Add(-time.Second)) //nolint:errcheck // best effort + return c.Conn.Close() +} + // Ask sends one command to the running service and returns its reply. func Ask(cmd string) (string, error) { c, err := net.DialTimeout("unix", sockPath, 2*time.Second) diff --git a/internal/control/control_test.go b/internal/control/control_test.go index 240f6c6..2c925c3 100644 --- a/internal/control/control_test.go +++ b/internal/control/control_test.go @@ -5,6 +5,7 @@ import ( "net" "strings" "testing" + "time" ) func TestRoundTrip(t *testing.T) { @@ -44,3 +45,59 @@ func TestEmptyReplyIsUnavailable(t *testing.T) { t.Fatalf("err = %v, want ErrUnavailable", err) } } + +func TestServeCap(t *testing.T) { + for i := 0; i < maxConns; i++ { + connSem <- struct{}{} + } + defer func() { + for i := 0; i < maxConns; i++ { + <-connSem + } + }() + server, client := net.Pipe() + defer server.Close() + defer client.Close() + if serve(server, func(string) string { return "OK" }) { + t.Fatal("serve accepted a connection beyond the cap") + } +} + +// forcePipe records ForceClose calls for the watchdog test. +type forcePipe struct { + net.Conn + forced chan struct{} +} + +func (c forcePipe) ForceClose() error { + err := c.Conn.Close() + close(c.forced) + return err +} + +func TestConnWatchdog(t *testing.T) { + old := connTimeout + connTimeout = 50 * time.Millisecond + defer func() { connTimeout = old }() + + server, client := net.Pipe() + defer client.Close() + fc := forcePipe{Conn: server, forced: make(chan struct{})} + done := make(chan struct{}) + go func() { + serveConn(fc, func(string) string { return "OK" }) + close(done) + }() + // The client never sends anything; the watchdog must tear the + // connection down instead of blocking forever. + select { + case <-fc.forced: + case <-time.After(5 * time.Second): + t.Fatal("watchdog did not force-close the stalled connection") + } + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("serveConn still blocked after the force close") + } +} diff --git a/internal/control/control_windows.go b/internal/control/control_windows.go index 7742cce..0ea67ba 100644 --- a/internal/control/control_windows.go +++ b/internal/control/control_windows.go @@ -86,7 +86,10 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error { windows.CloseHandle(pipe) continue } - go serveConn(&pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}, h) + conn := &pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe} + if !serve(conn, h) { + conn.ForceClose() + } } }() return nil @@ -114,6 +117,14 @@ func (c *pipeConn) Close() error { return c.f.Close() } +// ForceClose aborts the connection without flushing: disconnecting +// unblocks pending reads and writes at the cost of possibly discarding an +// unread reply. Used by the connection watchdog; normal closes flush. +func (c *pipeConn) ForceClose() error { + windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort + return c.f.Close() +} + // Ask sends one command to the running service and returns its reply. func Ask(cmd string) (string, error) { name, err := windows.UTF16PtrFromString(pipePath)