From b9d3f914036615b2797308247e7f3c4aa66491fb Mon Sep 17 00:00:00 2001 From: mram Date: Tue, 22 Sep 2026 09:10:14 +0200 Subject: [PATCH] Monitor: hotkeys (q quit, u update now) with footer line; show GPU VRAM usage --- cmd/gpu-turnstile/main.go | 46 +++++++++++- cmd/gpu-turnstile/monitor.go | 108 +++++++++++++++++++++++++-- cmd/gpu-turnstile/monitor_unix.go | 18 +++++ cmd/gpu-turnstile/monitor_windows.go | 20 +++++ internal/game/game.go | 23 ++++++ 5 files changed, 206 insertions(+), 9 deletions(-) diff --git a/cmd/gpu-turnstile/main.go b/cmd/gpu-turnstile/main.go index 9e74b86..0630090 100644 --- a/cmd/gpu-turnstile/main.go +++ b/cmd/gpu-turnstile/main.go @@ -728,9 +728,11 @@ 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 { 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 @@ -809,7 +811,7 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri // AUTO_UPDATE. if isService { serveControl(ctx, log, u, exePath, applyStaged, - statusProvider(cfg, lk, comfySup, health, started), + statusProvider(cfg, lk, comfySup, health, started, gw), reloadHandler(cfg, configPath, restartWhenIdle)) } @@ -834,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 @@ -853,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 @@ -1070,11 +1097,20 @@ 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 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 @@ -1109,12 +1145,14 @@ func diffConfig(a, b config.Config) []string { } // 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) func() string { +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"), diff --git a/cmd/gpu-turnstile/monitor.go b/cmd/gpu-turnstile/monitor.go index 68caf99..4944503 100644 --- a/cmd/gpu-turnstile/monitor.go +++ b/cmd/gpu-turnstile/monitor.go @@ -21,17 +21,26 @@ const ( 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. -// Ctrl+C quits. +// 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() @@ -39,7 +48,18 @@ func monitorCommand() int { if st, err := os.Stat(exe); err == nil { exeStamp = st.ModTime() } - for { + + keys := make(chan byte, 8) + go readKeys(keys) + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + + var note string + var noteAt time.Time + noteCh := make(chan string, 1) + updatePending := false + + poll := func() string { frame := renderWaiting() if reply, err := control.Ask(control.CmdStatus); err == nil { if msg, ok := strings.CutPrefix(reply, "OK "); ok { @@ -50,16 +70,75 @@ func monitorCommand() int { fmt.Print("\x1b[2J\x1b[H") fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version) restartSelf(exe, "--monitor") - return 0 + return "" // re-execed; this process exits below } snap.MonitorNote = fmt.Sprintf("note: the service runs %s, this monitor is %s", snap.Version, version) } + 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 - time.Sleep(time.Second) + select { + case <-ticker.C: + if note != "" && time.Since(noteAt) > 15*time.Second { + note = "" + } + case k, ok := <-keys: + if !ok { + keys = nil + continue + } + switch k { + case 'q', 'Q', 3: // q or Ctrl+C (raw mode delivers it as a byte) + return 0 + case 'u', 'U': + 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 + } } } @@ -82,7 +161,7 @@ func restartSelf(exe string, args ...string) { } func renderWaiting() string { - return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n" + return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n\x1b[K\n" + hotkeysLine } // renderMonitor draws one full frame. Each line ends with \x1b[K (clear to @@ -117,12 +196,31 @@ 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.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) } diff --git a/cmd/gpu-turnstile/monitor_unix.go b/cmd/gpu-turnstile/monitor_unix.go index 4685106..5c617f7 100644 --- a/cmd/gpu-turnstile/monitor_unix.go +++ b/cmd/gpu-turnstile/monitor_unix.go @@ -19,3 +19,21 @@ func termWidth() int { } 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 +} diff --git a/cmd/gpu-turnstile/monitor_windows.go b/cmd/gpu-turnstile/monitor_windows.go index f577d8d..a07996f 100644 --- a/cmd/gpu-turnstile/monitor_windows.go +++ b/cmd/gpu-turnstile/monitor_windows.go @@ -27,3 +27,23 @@ func termWidth() int { } 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 +} diff --git a/internal/game/game.go b/internal/game/game.go index b425cbf..8dd63c2 100644 --- a/internal/game/game.go +++ b/internal/game/game.go @@ -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) {