PUBLIC/gpu-turnstile
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 19 Wiki Activity

Env names in reload diff; monitor shows last VRAM check result; harden control channel against floods

Browse Source 
- 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
This commit is contained in:
mram
2026-09-22 09:24:27 +02:00
parent 4bd5f34ce7
commit 8fccd333aa
8 changed files with 288 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/gpu-turnstile/main.go
+64 -22
View File
@@ -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"),
cmd/gpu-turnstile/monitor.go
+22 -2
View File
@@ -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 {
cmd/gpu-turnstile/monitor_test.go
+31 -2
View File
@@ -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)
}
}
}