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
This commit is contained in:
mram
+64
-22
@@ -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"),
|
||||
|
||||
Reference in New Issue
Block a user