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

Monitor: hotkeys (q quit, u update now) with footer line; show GPU VRAM usage

Browse Source
This commit is contained in:
mram
2026-09-22 09:10:14 +02:00
parent 98d2d714ca
commit b9d3f91403
5 changed files with 206 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/gpu-turnstile/main.go
+42 -4
View File
@@ -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"),
cmd/gpu-turnstile/monitor.go
+103 -5
View File
@@ -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) }
cmd/gpu-turnstile/monitor_unix.go
+18
View File
@@ -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
}
cmd/gpu-turnstile/monitor_windows.go
+20
View File
@@ -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
}
internal/game/game.go
+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) {