From d7566329ae8b7aa1513adecb4c5cfce5e06977ed Mon Sep 17 00:00:00 2001 From: mram Date: Mon, 21 Sep 2026 13:48:07 +0200 Subject: [PATCH] Managed ComfyUI: COMFY_CMD starts it on demand, idle stop frees VRAM (internal/supervise) --- README.md | 23 +++ SPEC.md | 25 +++ cmd/gpu-turnstile/main.go | 69 ++++++- internal/config/config.go | 20 +++ internal/config/config_test.go | 23 +++ internal/config/sample.go | 4 + internal/proxy/proxy.go | 32 ++++ internal/supervise/supervise.go | 257 +++++++++++++++++++++++++++ internal/supervise/supervise_test.go | 158 ++++++++++++++++ 9 files changed, 607 insertions(+), 4 deletions(-) create mode 100644 internal/supervise/supervise.go create mode 100644 internal/supervise/supervise_test.go diff --git a/README.md b/README.md index eb49383..f662b7e 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,10 @@ override file values. Invalid values fail at startup. | `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400–599, e.g. 429) | | `BUSY_RETRY_AFTER` | `30` | Seconds sent as `Retry-After` on busy responses (both modes) | | `WARM_MODEL` | _(empty)_ | Model to reload after an image job (off by default) | +| `COMFY_CMD` | _(empty = unmanaged)_ | Supervise ComfyUI: start on demand, stop when idle to free VRAM. Requires `COMFY_URL` | +| `COMFY_DIR` | _(empty)_ | Working directory for `COMFY_CMD` | +| `COMFY_IDLE_TIMEOUT` | `5m` | Stop the managed ComfyUI after this long idle | +| `COMFY_START_TIMEOUT` | `2m` | Max wait for the managed ComfyUI to come up | | `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` works as an alias | | `LOG_FORMAT` | `text` | `json` for structured JSON logs | | `LOG_FILE` | _(empty)_ | Append logs to this file instead of stderr | @@ -89,6 +93,25 @@ override file values. Invalid values fail at startup. class) in text mode, which renders in `docker compose logs` on Windows Terminal. Set `NO_COLOR` to disable colors. +## Managed ComfyUI (`COMFY_CMD`) + +Don't want ComfyUI running 24/7 (it holds VRAM even when idle — and the +Desktop app kills its server when you close it)? Point `COMFY_CMD` at a +standalone launch command and gpu-turnstile supervises it: the first +request starts it, it stops again after `COMFY_IDLE_TIMEOUT` (default 5m) +without work, freeing the GPU for games or the LLM. Example for a Desktop +install (run it once manually to confirm it works): + +``` +COMFY_URL=http://127.0.0.1:8188 +COMFY_CMD="C:\ComfyUI\.venv\Scripts\python.exe ComfyUI\main.py --port 8188" +COMFY_DIR=C:\ComfyUI +``` + +`POST /prompt` waits for the server to answer before taking the GPU lock +(LLM traffic keeps flowing while torch loads); crashes are logged and the +next request respawns. Shutting gpu-turnstile down stops the child too. + ## Build and run ```sh diff --git a/SPEC.md b/SPEC.md index a379224..a48d71b 100644 --- a/SPEC.md +++ b/SPEC.md @@ -121,6 +121,27 @@ state is `idle`, send `POST /api/generate {"model":WARM_MODEL,"keep_alive":-1}` with empty prompt to reload the chat model so the next chat doesn't pay the load time. Off by default. +### Managed ComfyUI (`COMFY_CMD`) + +When `COMFY_CMD` is set, gpu-turnstile runs ComfyUI as a supervised child +process instead of expecting an always-on server: + +- **Start on demand**: any ComfyUI request spawns it (double quotes in the + command line group arguments with spaces; `COMFY_DIR` sets the working + directory). `POST /prompt` additionally waits for readiness + (`/system_stats`) for up to `COMFY_START_TIMEOUT` *before* taking the GPU + lock, so LLM traffic flows while torch loads. Other requests are bridged + by the normal retry backoff. Spawn failure or a readiness timeout + answers 502. +- **Idle stop**: after `COMFY_IDLE_TIMEOUT` without requests or finished + jobs — and only while the GPU lock is idle — the process tree is killed, + freeing the VRAM ComfyUI holds. The next request restarts it. +- **Crash**: an unexpected exit is logged; the next request respawns. + gpu-turnstile's own shutdown stops the child too. +- Its stdout/stderr is forwarded to the log at INFO. The health check + skips the intentionally-stopped/starting states; a failed probe while + the process is alive and was previously ready is logged as DOWN. + ## Configuration (env) Configuration comes from environment variables and/or an `.env`-style @@ -142,6 +163,10 @@ override file values. A missing file is fine; a malformed one is fatal. | `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400–599, e.g. 429) | | `BUSY_RETRY_AFTER` | `30` | seconds sent as `Retry-After` on busy responses (both modes) | | `WARM_MODEL` | `` | optional model to reload after an image job | +| `COMFY_CMD` | _(empty = unmanaged)_ | spawn and supervise ComfyUI on demand: first request starts it, idle stop after `COMFY_IDLE_TIMEOUT` frees its VRAM. Requires `COMFY_URL` | +| `COMFY_DIR` | `` | working directory for `COMFY_CMD` | +| `COMFY_IDLE_TIMEOUT` | `5m` | stop the managed ComfyUI after this long without requests or jobs | +| `COMFY_START_TIMEOUT` | `2m` | how long a request waits for the managed ComfyUI to come up | | `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` is accepted as an alias | | `LOG_FORMAT` | `text` | `json` for structured JSON logs | | `LOG_FILE` | `` | append logs to this file instead of stderr (useful as a service) | diff --git a/cmd/gpu-turnstile/main.go b/cmd/gpu-turnstile/main.go index d961974..c6f8c67 100644 --- a/cmd/gpu-turnstile/main.go +++ b/cmd/gpu-turnstile/main.go @@ -26,6 +26,7 @@ import ( "gpu-turnstile/internal/ollama" "gpu-turnstile/internal/proxy" "gpu-turnstile/internal/service" + "gpu-turnstile/internal/supervise" "gpu-turnstile/internal/update" ) @@ -436,6 +437,10 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri "backoff_max", cfg.BackoffMax, "prompt_capture_limit", cfg.PromptCaptureLimit, "warm_model", cfg.WarmModel, + "comfy_cmd", cfg.ComfyCmd, + "comfy_dir", cfg.ComfyDir, + "comfy_idle_timeout", cfg.ComfyIdleTimeout, + "comfy_start_timeout", cfg.ComfyStartTimeout, "auto_update", cfg.AutoUpdate, "update_interval", cfg.UpdateInterval, "update_repo", cfg.UpdateRepo, @@ -462,12 +467,31 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri } } + // With COMFY_CMD set, ComfyUI runs as a managed child: started on + // demand by the proxy, stopped after COMFY_IDLE_TIMEOUT idle (and on + // shutdown) so its VRAM is freed. + var comfySup *supervise.Process + if cfg.ComfyCmd != "" { + var err error + comfySup, err = supervise.New("comfy", cfg.ComfyCmd, cfg.ComfyDir, comfyClient.Probe, cfg.ComfyStartTimeout, log) + if err != nil { + return err + } + defer comfySup.Stop() + gpuIdle := func() bool { + state, _, pending := lk.Snapshot() + return state == lock.StateIdle && !pending + } + go comfySup.WatchIdle(ctx, cfg.ComfyIdleTimeout, gpuIdle) + } + srv, err := proxy.New(proxy.Config{ OllamaURL: cfg.OllamaURL, ComfyURL: cfg.ComfyURL, Lock: lk, Ollama: ollamaClient, Comfy: comfyClient, + ComfySup: comfySup, Metrics: metrics.New(), Log: log, LogColor: !cfg.LogJSON && cfg.LogFile == "" && os.Getenv("NO_COLOR") == "", @@ -491,17 +515,38 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri return err } - // Probe the enabled upstreams once; failure is logged, not fatal. + // Probe the enabled upstreams once; failure is logged, not fatal. A + // managed ComfyUI is intentionally down at startup — the first request + // starts it — so neither the startup probe nor the health check treats + // that as an outage. probes := map[string]func(context.Context) error{} if ollamaClient != nil { probes["ollama"] = ollamaClient.Probe } if comfyClient != nil { - probes["comfy"] = comfyClient.Probe + if comfySup == nil { + probes["comfy"] = comfyClient.Probe + } else { + // Managed upstream: an idle-stopped or still-starting server is + // not an outage, so it is skipped until it has answered once + // (Ready resets on every spawn/stop). After that, a failed + // probe while the process lives is a real "DOWN". + probes["comfy"] = func(ctx context.Context) error { + if !comfySup.Ready() { + if comfySup.Running() { + if err := comfyClient.Probe(ctx); err == nil { + comfySup.MarkReady() + } + } + return errManagedDown + } + return comfyClient.Probe(ctx) + } + } } probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout) for name, probe := range probes { - if err := probe(probeCtx); err != nil { + if err := probe(probeCtx); err != nil && !errors.Is(err, errManagedDown) { log.Warn(name+" probe failed", "err", err) } } @@ -566,6 +611,10 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri return nil } +// errManagedDown marks a managed upstream that is intentionally stopped +// (idle); health checks skip it instead of logging an outage. +var errManagedDown = errors.New("managed upstream intentionally stopped") + // healthLoop probes the enabled upstreams every interval and logs status // transitions — "is DOWN" when a previously healthy upstream stops // answering, "recovered" when it comes back. The first round only @@ -574,6 +623,7 @@ func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log * ticker := time.NewTicker(interval) defer ticker.Stop() up := map[string]bool{} + managed := map[string]bool{} for { select { case <-ctx.Done(): @@ -584,8 +634,19 @@ func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log * pctx, cancel := context.WithTimeout(ctx, probeTimeout) err := probe(pctx) cancel() - was, seen := up[name] + if errors.Is(err, errManagedDown) { + managed[name] = true // intentionally stopped; not an outage + continue + } now := err == nil + if managed[name] { + // First real probe after an idle stop only re-baselines — + // an on-demand start is not a "recovery". + managed[name] = false + up[name] = now + continue + } + was, seen := up[name] if seen && now != was { if now { log.Warn(name + " upstream recovered") diff --git a/internal/config/config.go b/internal/config/config.go index 8252561..fb47d42 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -51,6 +51,16 @@ type Config struct { LLMBusyStatus int BusyRetryAfter int + // ComfyCmd spawns and supervises a ComfyUI server on demand (empty = + // unmanaged, the current behavior). ComfyDir is its working directory. + // 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 + WarmModel string LogLevel slog.Level LogJSON bool @@ -89,6 +99,9 @@ func Defaults() Config { LLMBusyStatus: 503, BusyRetryAfter: 30, + ComfyIdleTimeout: 5 * time.Minute, + ComfyStartTimeout: 2 * time.Minute, + LogLevel: slog.LevelWarn, } } @@ -150,6 +163,8 @@ func Load(getenv func(string) string) (Config, error) { {"OLLAMA_URL", &cfg.OllamaURL}, {"COMFY_URL", &cfg.ComfyURL}, {"WARM_MODEL", &cfg.WarmModel}, + {"COMFY_CMD", &cfg.ComfyCmd}, + {"COMFY_DIR", &cfg.ComfyDir}, {"UPDATE_REPO", &cfg.UpdateRepo}, {"UPDATE_ASSET", &cfg.UpdateAsset}, {"LOG_FILE", &cfg.LogFile}, @@ -174,6 +189,8 @@ func Load(getenv func(string) string) (Config, error) { {"SHUTDOWN_TIMEOUT", &cfg.ShutdownTimeout}, {"BACKOFF_INITIAL", &cfg.BackoffInitial}, {"BACKOFF_MAX", &cfg.BackoffMax}, + {"COMFY_IDLE_TIMEOUT", &cfg.ComfyIdleTimeout}, + {"COMFY_START_TIMEOUT", &cfg.ComfyStartTimeout}, {"UPDATE_INTERVAL", &cfg.UpdateInterval}, } { if err := envDuration(getenv, e.name, e.dst); err != nil { @@ -244,6 +261,9 @@ func Load(getenv func(string) string) (Config, error) { default: return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"") } + if cfg.ComfyCmd != "" && cfg.ComfyURL == "" { + return cfg, fmt.Errorf("COMFY_CMD requires COMFY_URL to be set (the proxy needs somewhere to forward)") + } if cfg.OllamaURL == "" && cfg.ComfyURL == "" { return cfg, ErrNoConsumer } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 0491cfd..20fd499 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -74,6 +74,29 @@ func TestAppVersion(t *testing.T) { } } +func TestComfyCmdRequiresURL(t *testing.T) { + _, err := Load(func(k string) string { + if k == "COMFY_CMD" { + return "python main.py" + } + return "" + }) + if err == nil || !strings.Contains(err.Error(), "COMFY_CMD requires COMFY_URL") { + t.Fatalf("err = %v, want COMFY_CMD/COMFY_URL validation error", err) + } + if _, err := Load(func(k string) string { + switch k { + case "COMFY_CMD": + return "python main.py" + case "COMFY_URL": + return "http://127.0.0.1:8188" + } + return "" + }); err != nil { + t.Fatalf("COMFY_CMD with COMFY_URL must load: %v", err) + } +} + func TestParseEnvFile(t *testing.T) { input := `# comment OLLAMA_URL=http://host:11435 diff --git a/internal/config/sample.go b/internal/config/sample.go index 0bee0fb..be1e096 100644 --- a/internal/config/sample.go +++ b/internal/config/sample.go @@ -28,6 +28,10 @@ func sampleEntries(logFile string) []sampleEntry { {"OLLAMA_URL", "http://127.0.0.1:11434", "Ollama upstream URL; setting it enables the Ollama consumer (default: empty = disabled)", false}, {"COMFY_URL", "http://127.0.0.1:8188", "ComfyUI upstream URL; setting it enables the ComfyUI consumer (default: empty = disabled)", false}, {"WARM_MODEL", "", "Optional model to reload after an image job (default: empty = none)", false}, + {"COMFY_CMD", `"C:\ComfyUI\.venv\Scripts\python.exe" ComfyUI\main.py --port 8188`, "Spawn and supervise ComfyUI on demand: the first request starts it, it stops after COMFY_IDLE_TIMEOUT to free VRAM (default: empty = unmanaged)", false}, + {"COMFY_DIR", `C:\ComfyUI`, "Working directory for COMFY_CMD (default: empty = inherit)", false}, + {"COMFY_IDLE_TIMEOUT", "5m", "Stop the managed ComfyUI after this long without requests or jobs (frees VRAM)", false}, + {"COMFY_START_TIMEOUT", "2m", "How long a request waits for the managed ComfyUI to come up", false}, {"UNLOAD_TIMEOUT", "60s", "How long to wait for Ollama to unload a model", false}, {"JOB_TIMEOUT", "15m", "Maximum time to wait for a ComfyUI job", false}, {"LLM_WAIT_TIMEOUT", "10m", "Max time an LLM request waits for the GPU before being answered 503 (wait mode)", false}, diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index d94dcb7..05eb807 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -26,6 +26,7 @@ import ( "gpu-turnstile/internal/lock" "gpu-turnstile/internal/metrics" "gpu-turnstile/internal/ollama" + "gpu-turnstile/internal/supervise" ) // defaultCaptureLimit bounds how much of a /prompt response body is @@ -44,6 +45,11 @@ type Config struct { Metrics *metrics.Metrics Log *slog.Logger + // ComfySup, when non-nil, is the managed ComfyUI process: any ComfyUI + // request starts it on demand, /prompt additionally waits for + // readiness before taking the GPU lock. nil = unmanaged upstream. + ComfySup *supervise.Process + // LogColor enables ANSI colors in per-request log lines. Ignored when // the log level is above INFO (request lines are not emitted at all). LogColor bool @@ -500,6 +506,14 @@ func (s *Server) ComfyHandler() http.Handler { s.handlePrompt(w, r) return } + if s.cfg.ComfySup != nil { + // Any other ComfyUI request also wakes the managed server; the + // retry backoff bridges the time it needs to come up. + if err := s.cfg.ComfySup.EnsureRunning(); err != nil { + http.Error(w, fmt.Sprintf("cannot start ComfyUI: %v", err), http.StatusBadGateway) + return + } + } s.comfyProxy.ServeHTTP(w, r) })) } @@ -539,6 +553,20 @@ func (w *captureWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter } func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) { log := s.log.With("op", "image") + if s.cfg.ComfySup != nil { + // Bring the managed server up *before* taking the GPU lock: torch + // can take a minute to load, and LLM traffic should keep flowing + // in the meantime. + if err := s.cfg.ComfySup.EnsureRunning(); err != nil { + http.Error(w, fmt.Sprintf("cannot start ComfyUI: %v", err), http.StatusBadGateway) + return + } + if err := s.cfg.ComfySup.WaitReady(r.Context()); err != nil { + http.Error(w, fmt.Sprintf("ComfyUI did not become ready: %v", err), http.StatusBadGateway) + return + } + } + start := time.Now() if err := s.cfg.Lock.AcquireImage(r.Context()); err != nil { if errors.Is(err, context.DeadlineExceeded) { @@ -608,6 +636,10 @@ func (s *Server) finishImageJob(promptID string) { s.cfg.Lock.ReleaseImage() log.Info("image lock released") + if s.cfg.ComfySup != nil { + // The idle clock starts when the job ends, not when it began. + s.cfg.ComfySup.NoteActivity() + } if s.cfg.Ollama != nil && s.cfg.WarmModel != "" { if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle { diff --git a/internal/supervise/supervise.go b/internal/supervise/supervise.go new file mode 100644 index 0000000..512fef0 --- /dev/null +++ b/internal/supervise/supervise.go @@ -0,0 +1,257 @@ +// Package supervise runs an upstream server (ComfyUI) as a managed child +// process: started on demand when a request needs it, stopped after an +// idle timeout so the GPU memory it holds is freed, and stopped with the +// parent. Crashes are logged; the next request respawns it. +package supervise + +import ( + "context" + "fmt" + "io" + "log/slog" + "os/exec" + "runtime" + "strings" + "sync" + "time" +) + +// Process is one managed child process. +type Process struct { + name string + argv []string + dir string + probe func(context.Context) error + startTimeout time.Duration + log *slog.Logger + + mu sync.Mutex + cmd *exec.Cmd + stopping bool + ready bool + lastActivity time.Time +} + +// New parses cmdLine (double quotes group arguments containing spaces) and +// prepares a managed process. probe reports whether the server answers +// (e.g. the comfy client's Probe); startTimeout bounds WaitReady. dir is +// the child's working directory; empty inherits ours. +func New(name, cmdLine, dir string, probe func(context.Context) error, startTimeout time.Duration, log *slog.Logger) (*Process, error) { + argv, err := splitCommandLine(cmdLine) + if err != nil { + return nil, fmt.Errorf("%s command: %w", name, err) + } + if len(argv) == 0 { + return nil, fmt.Errorf("%s command is empty", name) + } + if log == nil { + log = slog.Default() + } + if startTimeout <= 0 { + startTimeout = 2 * time.Minute + } + return &Process{name: name, argv: argv, dir: dir, probe: probe, startTimeout: startTimeout, log: log}, nil +} + +// splitCommandLine splits a command line on whitespace, treating +// double-quoted sections as one argument (quotes removed). Backslashes are +// literal — this matches Windows paths. +func splitCommandLine(s string) ([]string, error) { + var argv []string + var cur strings.Builder + inQuote := false + have := false + flush := func() { + if have { + argv = append(argv, cur.String()) + cur.Reset() + have = false + } + } + for _, r := range s { + switch { + case r == '"': + inQuote = !inQuote + have = true + case (r == ' ' || r == '\t') && !inQuote: + flush() + default: + cur.WriteRune(r) + have = true + } + } + if inQuote { + return nil, fmt.Errorf("unterminated quote in %q", s) + } + flush() + return argv, nil +} + +// Running reports whether the child process is currently alive. +func (p *Process) Running() bool { + p.mu.Lock() + defer p.mu.Unlock() + return p.cmd != nil +} + +// Ready reports whether the server has answered a probe since its last +// (re)start. Health checks use it to tell "starting up" from "outage". +func (p *Process) Ready() bool { + p.mu.Lock() + defer p.mu.Unlock() + return p.ready +} + +// MarkReady records that the server answered. +func (p *Process) MarkReady() { + p.mu.Lock() + p.ready = true + p.mu.Unlock() +} + +// NoteActivity resets the idle clock; called for every request served. +func (p *Process) NoteActivity() { + p.mu.Lock() + p.lastActivity = time.Now() + p.mu.Unlock() +} + +// EnsureRunning starts the child if it is not running. It returns as soon +// as the process is spawned; readiness is WaitReady's job (and the proxy's +// retry backoff bridges the gap for plain proxied requests). +func (p *Process) EnsureRunning() error { + p.mu.Lock() + defer p.mu.Unlock() + p.lastActivity = time.Now() + if p.cmd != nil { + return nil + } + cmd := exec.Command(p.argv[0], p.argv[1:]...) + cmd.Dir = p.dir + stdout, err := cmd.StdoutPipe() + if err != nil { + return err + } + stderr, err := cmd.StderrPipe() + if err != nil { + return err + } + if err := cmd.Start(); err != nil { + return fmt.Errorf("start %s: %w", p.name, err) + } + p.cmd = cmd + p.stopping = false + p.ready = false + go p.pipeLog(stdout) + go p.pipeLog(stderr) + go func() { + err := cmd.Wait() + p.mu.Lock() + p.cmd = nil + p.ready = false + intentional := p.stopping + p.mu.Unlock() + if intentional { + p.log.Info(p.name + " stopped") + } else { + p.log.Warn(p.name+" exited unexpectedly; the next request restarts it", "err", err) + } + }() + p.log.Info(p.name+" starting", "pid", cmd.Process.Pid, "cmd", strings.Join(p.argv, " ")) + return nil +} + +// WaitReady blocks until the probe succeeds, ctx ends, or the start +// timeout passes. +func (p *Process) WaitReady(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, p.startTimeout) + defer cancel() + for { + pctx, pcancel := context.WithTimeout(ctx, 5*time.Second) + err := p.probe(pctx) + pcancel() + if err == nil { + p.NoteActivity() + p.MarkReady() + return nil + } + select { + case <-ctx.Done(): + return fmt.Errorf("%s did not become ready: %w", p.name, ctx.Err()) + case <-time.After(500 * time.Millisecond): + } + } +} + +// Stop kills the child process (the whole tree on Windows) if running. +func (p *Process) Stop() { + p.mu.Lock() + cmd := p.cmd + if cmd == nil { + p.mu.Unlock() + return + } + p.stopping = true + p.ready = false + p.mu.Unlock() + stopTree(cmd) +} + +// WatchIdle stops the child after idleTimeout without activity, but only +// when gpuIdle reports the GPU lock is free (no active or pending work). +// Returns when ctx ends. +func (p *Process) WatchIdle(ctx context.Context, idleTimeout time.Duration, gpuIdle func() bool) { + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + p.mu.Lock() + idleFor := time.Since(p.lastActivity) + running := p.cmd != nil + p.mu.Unlock() + if running && idleFor > idleTimeout && gpuIdle() { + p.log.Info(p.name+" idle; stopping to free the GPU", "idle_for", idleFor.Round(time.Second)) + p.Stop() + } + } +} + +// pipeLog forwards one child output stream to the log at INFO, line by +// line, prefixed with the process name. +func (p *Process) pipeLog(r io.Reader) { + buf := make([]byte, 4096) + var line string + for { + n, err := r.Read(buf) + line += string(buf[:n]) + for { + i := strings.IndexByte(line, '\n') + if i < 0 { + break + } + p.log.Info(p.name + ": " + strings.TrimRight(line[:i], "\r")) + line = line[i+1:] + } + if err != nil { + if strings.TrimSpace(line) != "" { + p.log.Info(p.name + ": " + line) + } + return + } + } +} + +// stopTree kills cmd's process, including its children on Windows (python +// launchers tend to spawn some). The Wait goroutine reaps it. +func stopTree(cmd *exec.Cmd) { + if runtime.GOOS == "windows" { + exec.Command("taskkill", "/T", "/F", "/PID", + fmt.Sprint(cmd.Process.Pid)).Run() //nolint:errcheck // best effort + return + } + cmd.Process.Kill() //nolint:errcheck // best effort +} diff --git a/internal/supervise/supervise_test.go b/internal/supervise/supervise_test.go new file mode 100644 index 0000000..bd183d8 --- /dev/null +++ b/internal/supervise/supervise_test.go @@ -0,0 +1,158 @@ +package supervise + +import ( + "context" + "errors" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "testing" + "time" +) + +func TestSplitCommandLine(t *testing.T) { + cases := []struct { + in string + want []string + }{ + {`python main.py --port 8188`, []string{"python", "main.py", "--port", "8188"}}, + {`"C:\Program Files\py\python.exe" main.py`, []string{`C:\Program Files\py\python.exe`, "main.py"}}, + {` spaced out `, []string{"spaced", "out"}}, + {`a "b c" d`, []string{"a", "b c", "d"}}, + {"", nil}, + } + for _, c := range cases { + got, err := splitCommandLine(c.in) + if err != nil { + t.Errorf("splitCommandLine(%q): %v", c.in, err) + continue + } + if len(got) != len(c.want) { + t.Errorf("splitCommandLine(%q) = %v, want %v", c.in, got, c.want) + continue + } + for i := range got { + if got[i] != c.want[i] { + t.Errorf("splitCommandLine(%q) = %v, want %v", c.in, got, c.want) + break + } + } + } + if _, err := splitCommandLine(`"unterminated`); err == nil { + t.Error("expected error for unterminated quote") + } +} + +// TestHelperProcess is the child executed by the lifecycle tests: it just +// sleeps. The env marker is set only around the spawn, so in the normal +// test run this returns immediately. +func TestHelperProcess(t *testing.T) { + if os.Getenv("GO_HELPER_PROCESS") != "1" { + return + } + time.Sleep(30 * time.Second) + os.Exit(0) +} + +func newHelper(t *testing.T, name string) *Process { + t.Helper() + p, err := New(name, `"`+os.Args[0]+`" -test.run=TestHelperProcess`, "", + func(context.Context) error { return nil }, 5*time.Second, slog.Default()) + if err != nil { + t.Fatal(err) + } + return p +} + +// startHelper spawns the child with the marker set; exec.Command inherits +// the environment at spawn time, so it can be unset right after. +func startHelper(t *testing.T, p *Process) { + t.Helper() + os.Setenv("GO_HELPER_PROCESS", "1") + defer os.Unsetenv("GO_HELPER_PROCESS") + if err := p.EnsureRunning(); err != nil { + t.Fatal(err) + } +} + +func waitStopped(t *testing.T, p *Process, timeout time.Duration) { + t.Helper() + deadline := time.Now().Add(timeout) + for p.Running() && time.Now().Before(deadline) { + time.Sleep(50 * time.Millisecond) + } + if p.Running() { + t.Fatal("process still running") + } +} + +func TestEnsureRunningAndStop(t *testing.T) { + p := newHelper(t, "helper") + if p.Running() { + t.Fatal("Running before start") + } + startHelper(t, p) + if !p.Running() { + t.Fatal("not Running after EnsureRunning") + } + if err := p.EnsureRunning(); err != nil { + t.Fatal("second EnsureRunning must be a no-op") + } + p.Stop() + waitStopped(t, p, 5*time.Second) +} + +func TestWaitReady(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + probe := func(ctx context.Context) error { + req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil) + resp, err := http.DefaultClient.Do(req) + if err != nil { + return err + } + resp.Body.Close() + return nil + } + p, err := New("ready", `"`+os.Args[0]+`"`, "", probe, 5*time.Second, slog.Default()) + if err != nil { + t.Fatal(err) + } + if err := p.WaitReady(context.Background()); err != nil { + t.Fatal(err) + } + + failing, err := New("failing", `"`+os.Args[0]+`"`, "", + func(context.Context) error { return errors.New("no") }, 500*time.Millisecond, slog.Default()) + if err != nil { + t.Fatal(err) + } + if err := failing.WaitReady(context.Background()); err == nil { + t.Fatal("expected timeout error from WaitReady") + } +} + +func TestWatchIdleStops(t *testing.T) { + p := newHelper(t, "idle") + startHelper(t, p) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go p.WatchIdle(ctx, 200*time.Millisecond, func() bool { return true }) + waitStopped(t, p, 10*time.Second) +} + +func TestWatchIdleRespectsBusyGPU(t *testing.T) { + p := newHelper(t, "busy") + startHelper(t, p) + defer p.Stop() + ctx, cancel := context.WithCancel(context.Background()) + go p.WatchIdle(ctx, 100*time.Millisecond, func() bool { return false }) + time.Sleep(600 * time.Millisecond) + cancel() + if !p.Running() { + t.Fatal("process was stopped while the GPU was busy") + } +}