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

Add --monitor: live status view (downstream health, GPU lock, queue) via the control channel

Browse Source
This commit is contained in:
mram
2026-09-22 08:17:05 +02:00
parent a74cd49fe5
commit 16dc7fe462
9 changed files with 472 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/gpu-turnstile/main.go
+137 -20
View File
@@ -5,6 +5,7 @@ package main
import (
"bufio"
"context"
"encoding/json"
"errors"
"fmt"
"io"
@@ -58,7 +59,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, updateNow, elevatedChild bool, rest []string) {
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, monitor, elevatedChild bool, rest []string) {
rest = args[:0]
for i := 0; i < len(args); i++ {
switch {
@@ -81,13 +82,15 @@ func parseFlags(args []string) (configPath string, install, remove, noCopy, help
forceUpdate = true
case args[i] == "--update-now" || args[i] == "-update-now":
updateNow = true
case args[i] == "--monitor" || args[i] == "-monitor":
monitor = true
case args[i] == "--elevated-child":
elevatedChild = true
default:
rest = append(rest, args[i])
}
}
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, elevatedChild, rest
return configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, monitor, elevatedChild, rest
}
// versionLine is printed at the top of every help and error screen.
@@ -105,6 +108,8 @@ Usage:
(no admin needed when the service runs)
gpu-turnstile --update-now like --force-update, but only
through the running service
gpu-turnstile --monitor live status view (downstreams,
GPU lock, queue); Ctrl+C quits
gpu-turnstile -h | --help this help
Options:
@@ -135,12 +140,12 @@ func fatalUsage(format string, args ...any) {
}
func main() {
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, elevatedChild, args := parseFlags(os.Args[1:])
configPath, install, remove, noCopy, help, showVersion, forceUpdate, updateNow, monitor, elevatedChild, args := parseFlags(os.Args[1:])
if showVersion {
fmt.Println(version)
return
}
bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !elevatedChild && len(args) == 0
bare := configPath == "" && !install && !remove && !forceUpdate && !updateNow && !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
@@ -165,6 +170,8 @@ func main() {
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 monitor && (install || remove || forceUpdate || updateNow):
fatalUsage("error: --monitor cannot be combined with other commands")
case install:
os.Exit(serviceCommand(configPath, true, noCopy, elevatedChild))
case remove:
@@ -173,6 +180,8 @@ func main() {
os.Exit(forceUpdateCommand(configPath, elevatedChild))
case updateNow:
os.Exit(updateNowCommand())
case monitor:
os.Exit(monitorCommand())
}
if len(args) > 0 {
fatalUsage("error: unknown arguments: %s", strings.Join(args, " "))
@@ -567,6 +576,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
@@ -676,13 +687,15 @@ 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
@@ -727,17 +740,20 @@ 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
applyStaged := func(to string) {}
if cfg.AutoUpdate {
exePath, err := os.Executable()
if err != nil {
if p, err := os.Executable(); err != nil {
log.Warn("auto-update disabled: cannot locate executable", "err", err)
} else {
u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
exePath = p
u = &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
// applyStaged is shared by the hourly loop and the control
// channel; the once guard keeps a second trigger from
// double-waiting on the GPU lock.
var once sync.Once
applyStaged := func(to string) {
applyStaged = func(to string) {
if !isService {
log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply", "version", to)
return
@@ -753,12 +769,16 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
})
}
go updateLoop(ctx, cfg.UpdateInterval, log, u, exePath, applyStaged)
if isService {
serveControl(ctx, log, u, exePath, applyStaged)
}
}
}
// The control channel (status for --monitor, update-now trigger) 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))
}
select {
case err := <-errCh:
if err != nil && !errors.Is(err, http.ErrServerClosed) {
@@ -858,7 +878,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{}
@@ -875,9 +896,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".
@@ -920,17 +943,25 @@ func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u
}
// serveControl opens the local control channel (named pipe on Windows,
// unix socket on Linux) so unprivileged local users can trigger an update
// check via --force-update without admin rights. The check payload is
// signature-verified regardless of who asks; triggers are rate-limited to
// one per minute so the channel cannot be used to spam restarts.
func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string)) {
// unix socket on Linux) so unprivileged local users can query status
// (--monitor) and trigger an update check (--force-update/--update-now)
// without admin rights. The update payload is signature-verified regardless
// of who asks; 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) {
var mu sync.Mutex
var lastTrigger time.Time
h := func(cmd string) string {
if cmd != control.CmdUpdateNow {
switch cmd {
case control.CmdStatus:
return "OK " + status()
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()
@@ -955,6 +986,92 @@ func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exeP
}
}
// 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"`
}
// 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 {
return func() string {
snap := statusSnapshot{
Version: version,
UptimeS: int64(time.Since(started).Seconds()),
}
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 {