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
+40
-38
@@ -13,43 +13,45 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config holds every gpu-turnstile setting.
|
||||
// Config holds every gpu-turnstile setting. Each field's env tag names the
|
||||
// environment variable / config-file key that sets it; user-facing output
|
||||
// (e.g. the reload diff) uses those names, never the Go field names.
|
||||
type Config struct {
|
||||
ListenOllama string
|
||||
ListenComfy string
|
||||
OllamaURL string
|
||||
ComfyURL string
|
||||
UnloadTimeout time.Duration
|
||||
JobTimeout time.Duration
|
||||
LLMWaitTimeout time.Duration
|
||||
ListenOllama string `env:"LISTEN_OLLAMA"`
|
||||
ListenComfy string `env:"LISTEN_COMFY"`
|
||||
OllamaURL string `env:"OLLAMA_URL"`
|
||||
ComfyURL string `env:"COMFY_URL"`
|
||||
UnloadTimeout time.Duration `env:"UNLOAD_TIMEOUT"`
|
||||
JobTimeout time.Duration `env:"JOB_TIMEOUT"`
|
||||
LLMWaitTimeout time.Duration `env:"LLM_WAIT_TIMEOUT"`
|
||||
|
||||
UnloadPollInterval time.Duration
|
||||
HistoryPollInterval time.Duration
|
||||
ProbeTimeout time.Duration
|
||||
HealthInterval time.Duration
|
||||
FreeTimeout time.Duration
|
||||
WarmTimeout time.Duration
|
||||
ShutdownTimeout time.Duration
|
||||
BackoffInitial time.Duration
|
||||
BackoffMax time.Duration
|
||||
PromptCaptureLimit int64
|
||||
UnloadPollInterval time.Duration `env:"UNLOAD_POLL_INTERVAL"`
|
||||
HistoryPollInterval time.Duration `env:"HISTORY_POLL_INTERVAL"`
|
||||
ProbeTimeout time.Duration `env:"PROBE_TIMEOUT"`
|
||||
HealthInterval time.Duration `env:"HEALTH_INTERVAL"`
|
||||
FreeTimeout time.Duration `env:"FREE_TIMEOUT"`
|
||||
WarmTimeout time.Duration `env:"WARM_TIMEOUT"`
|
||||
ShutdownTimeout time.Duration `env:"SHUTDOWN_TIMEOUT"`
|
||||
BackoffInitial time.Duration `env:"BACKOFF_INITIAL"`
|
||||
BackoffMax time.Duration `env:"BACKOFF_MAX"`
|
||||
PromptCaptureLimit int64 `env:"PROMPT_CAPTURE_LIMIT"`
|
||||
|
||||
AutoUpdate bool
|
||||
UpdateInterval time.Duration
|
||||
UpdateRepo string
|
||||
UpdateAsset string
|
||||
AutoUpdate bool `env:"AUTO_UPDATE"`
|
||||
UpdateInterval time.Duration `env:"UPDATE_INTERVAL"`
|
||||
UpdateRepo string `env:"UPDATE_REPO"`
|
||||
UpdateAsset string `env:"UPDATE_ASSET"`
|
||||
|
||||
// AppVersion is the version the user wants to run: "dev" disables
|
||||
// updates, "stable" tracks the latest release, anything else is an
|
||||
// exact vX.Y.Z release to pin. From APP_VER; defaults to "stable".
|
||||
AppVersion string
|
||||
AppVersion string `env:"APP_VER"`
|
||||
|
||||
// LLMBusyMode is "wait" (hold requests until the lock is free or
|
||||
// LLMWaitTimeout expires) or "reject" (immediately answer with
|
||||
// LLMBusyStatus + Retry-After when an image job is active or pending).
|
||||
LLMBusyMode string
|
||||
LLMBusyStatus int
|
||||
BusyRetryAfter int
|
||||
LLMBusyMode string `env:"LLM_BUSY_MODE"`
|
||||
LLMBusyStatus int `env:"LLM_BUSY_STATUS"`
|
||||
BusyRetryAfter int `env:"BUSY_RETRY_AFTER"`
|
||||
|
||||
// ComfyCmd spawns and supervises a ComfyUI server on demand. When
|
||||
// ComfyCmd is empty but ComfyDir is set, management is enabled with the
|
||||
@@ -59,10 +61,10 @@ type Config struct {
|
||||
// set explicitly. The managed server is stopped after ComfyIdleTimeout
|
||||
// without requests, freeing its VRAM; ComfyStartTimeout bounds how long
|
||||
// a request waits for it to come up.
|
||||
ComfyCmd string
|
||||
ComfyDir string
|
||||
ComfyIdleTimeout time.Duration
|
||||
ComfyStartTimeout time.Duration
|
||||
ComfyCmd string `env:"COMFY_CMD"`
|
||||
ComfyDir string `env:"COMFY_DIR"`
|
||||
ComfyIdleTimeout time.Duration `env:"COMFY_IDLE_TIMEOUT"`
|
||||
ComfyStartTimeout time.Duration `env:"COMFY_START_TIMEOUT"`
|
||||
|
||||
// GameProcs (GAME_PROCS) is a watch list of process names; while any of
|
||||
// them runs, the GPU is treated as held by a foreign process. The
|
||||
@@ -70,15 +72,15 @@ type Config struct {
|
||||
// when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than
|
||||
// that many MiB of VRAM. GamePollInterval (GAME_POLL_INTERVAL) is how
|
||||
// often both checks run.
|
||||
GameProcs []string
|
||||
GPUForeignVRAMMB int
|
||||
GPUIgnoreProcs []string
|
||||
GamePollInterval time.Duration
|
||||
GameProcs []string `env:"GAME_PROCS"`
|
||||
GPUForeignVRAMMB int `env:"GPU_FOREIGN_VRAM_MB"`
|
||||
GPUIgnoreProcs []string `env:"GPU_IGNORE_PROCS"`
|
||||
GamePollInterval time.Duration `env:"GAME_POLL_INTERVAL"`
|
||||
|
||||
WarmModel string
|
||||
LogLevel slog.Level
|
||||
LogJSON bool
|
||||
LogFile string
|
||||
WarmModel string `env:"WARM_MODEL"`
|
||||
LogLevel slog.Level `env:"LOGLEVEL"`
|
||||
LogJSON bool `env:"LOG_FORMAT"`
|
||||
LogFile string `env:"LOG_FILE"`
|
||||
}
|
||||
|
||||
// Defaults returns the configuration used when neither the environment nor
|
||||
|
||||
@@ -7,6 +7,13 @@
|
||||
// triggers, so the worst a local user can cause is a cheap, throttled
|
||||
// check and a GPU-idle-gated restart onto a signed binary.
|
||||
//
|
||||
// Abuse hardening: the command read is capped (4 KiB), each connection is
|
||||
// force-closed after connTimeout so a stalled client cannot pin a goroutine
|
||||
// (or a Windows pipe instance) forever, and concurrently served connections
|
||||
// are capped at maxConns — beyond that, connections are closed on arrival.
|
||||
// On Windows the pipe's ACL additionally denies network logons, so the
|
||||
// channel cannot be reached from another machine.
|
||||
//
|
||||
// Protocol: the client writes one command line, the server answers with
|
||||
// one reply line ("OK ..." or "ERR ...") and hangs up.
|
||||
package control
|
||||
@@ -17,6 +24,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
|
||||
@@ -37,9 +45,50 @@ var ErrUnavailable = errors.New("control channel unavailable")
|
||||
// line. It must start with "OK " or "ERR ".
|
||||
type Handler func(cmd string) string
|
||||
|
||||
// connTimeout bounds one connection's lifetime: a client that stops
|
||||
// mid-command or never reads the reply would otherwise pin its goroutine
|
||||
// (and on Windows one of the pipe instances) indefinitely. A var so tests
|
||||
// can shrink it.
|
||||
var connTimeout = 10 * time.Second
|
||||
|
||||
// maxConns caps concurrently served connections; beyond it, new
|
||||
// connections are closed on arrival. Bound on the goroutines a local
|
||||
// flood can pile up.
|
||||
const maxConns = 32
|
||||
|
||||
var connSem = make(chan struct{}, maxConns)
|
||||
|
||||
// serve dispatches connection handling under the concurrency cap. It
|
||||
// returns false when the cap is reached — the caller must then close the
|
||||
// connection itself.
|
||||
func serve(c io.ReadWriteCloser, h Handler) bool {
|
||||
select {
|
||||
case connSem <- struct{}{}:
|
||||
go func() {
|
||||
defer func() { <-connSem }()
|
||||
serveConn(c, h)
|
||||
}()
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// forceCloser is implemented by connections that can be torn down
|
||||
// abortively, unblocking pending reads and writes (Windows pipe:
|
||||
// DisconnectNamedPipe; unix socket: a deadline in the past). The
|
||||
// connection watchdog uses it; normal closes still flush the reply.
|
||||
type forceCloser interface {
|
||||
ForceClose() error
|
||||
}
|
||||
|
||||
// serveConn runs the line protocol on one accepted connection.
|
||||
func serveConn(c io.ReadWriteCloser, h Handler) {
|
||||
defer c.Close()
|
||||
if fc, ok := c.(forceCloser); ok {
|
||||
timer := time.AfterFunc(connTimeout, func() { fc.ForceClose() })
|
||||
defer timer.Stop()
|
||||
}
|
||||
line, err := bufio.NewReader(io.LimitReader(c, 4096)).ReadString('\n')
|
||||
cmd := strings.TrimSpace(line)
|
||||
if cmd == "" {
|
||||
|
||||
@@ -37,12 +37,24 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
||||
if err != nil {
|
||||
return // shutting down
|
||||
}
|
||||
go serveConn(c, h)
|
||||
uc := unixConn{c}
|
||||
if !serve(uc, h) {
|
||||
uc.ForceClose()
|
||||
}
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
}
|
||||
|
||||
// unixConn adds an abortive ForceClose to net.Conn: a deadline in the
|
||||
// past fails pending and future I/O immediately.
|
||||
type unixConn struct{ net.Conn }
|
||||
|
||||
func (c unixConn) ForceClose() error {
|
||||
c.SetDeadline(time.Now().Add(-time.Second)) //nolint:errcheck // best effort
|
||||
return c.Conn.Close()
|
||||
}
|
||||
|
||||
// Ask sends one command to the running service and returns its reply.
|
||||
func Ask(cmd string) (string, error) {
|
||||
c, err := net.DialTimeout("unix", sockPath, 2*time.Second)
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"net"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRoundTrip(t *testing.T) {
|
||||
@@ -44,3 +45,59 @@ func TestEmptyReplyIsUnavailable(t *testing.T) {
|
||||
t.Fatalf("err = %v, want ErrUnavailable", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServeCap(t *testing.T) {
|
||||
for i := 0; i < maxConns; i++ {
|
||||
connSem <- struct{}{}
|
||||
}
|
||||
defer func() {
|
||||
for i := 0; i < maxConns; i++ {
|
||||
<-connSem
|
||||
}
|
||||
}()
|
||||
server, client := net.Pipe()
|
||||
defer server.Close()
|
||||
defer client.Close()
|
||||
if serve(server, func(string) string { return "OK" }) {
|
||||
t.Fatal("serve accepted a connection beyond the cap")
|
||||
}
|
||||
}
|
||||
|
||||
// forcePipe records ForceClose calls for the watchdog test.
|
||||
type forcePipe struct {
|
||||
net.Conn
|
||||
forced chan struct{}
|
||||
}
|
||||
|
||||
func (c forcePipe) ForceClose() error {
|
||||
err := c.Conn.Close()
|
||||
close(c.forced)
|
||||
return err
|
||||
}
|
||||
|
||||
func TestConnWatchdog(t *testing.T) {
|
||||
old := connTimeout
|
||||
connTimeout = 50 * time.Millisecond
|
||||
defer func() { connTimeout = old }()
|
||||
|
||||
server, client := net.Pipe()
|
||||
defer client.Close()
|
||||
fc := forcePipe{Conn: server, forced: make(chan struct{})}
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
serveConn(fc, func(string) string { return "OK" })
|
||||
close(done)
|
||||
}()
|
||||
// The client never sends anything; the watchdog must tear the
|
||||
// connection down instead of blocking forever.
|
||||
select {
|
||||
case <-fc.forced:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("watchdog did not force-close the stalled connection")
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("serveConn still blocked after the force close")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -86,7 +86,10 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
||||
windows.CloseHandle(pipe)
|
||||
continue
|
||||
}
|
||||
go serveConn(&pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}, h)
|
||||
conn := &pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}
|
||||
if !serve(conn, h) {
|
||||
conn.ForceClose()
|
||||
}
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
@@ -114,6 +117,14 @@ func (c *pipeConn) Close() error {
|
||||
return c.f.Close()
|
||||
}
|
||||
|
||||
// ForceClose aborts the connection without flushing: disconnecting
|
||||
// unblocks pending reads and writes at the cost of possibly discarding an
|
||||
// unread reply. Used by the connection watchdog; normal closes flush.
|
||||
func (c *pipeConn) ForceClose() error {
|
||||
windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort
|
||||
return c.f.Close()
|
||||
}
|
||||
|
||||
// Ask sends one command to the running service and returns its reply.
|
||||
func Ask(cmd string) (string, error) {
|
||||
name, err := windows.UTF16PtrFromString(pipePath)
|
||||
|
||||
Reference in New Issue
Block a user