Managed ComfyUI: COMFY_CMD starts it on demand, idle stop frees VRAM (internal/supervise)

This commit is contained in:
mram
2026-09-21 13:48:07 +02:00
parent e43ad02fc4
commit d7566329ae
9 changed files with 607 additions and 4 deletions
+23
View File
@@ -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 (400599, e.g. 429) | | `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400599, e.g. 429) |
| `BUSY_RETRY_AFTER` | `30` | Seconds sent as `Retry-After` on busy responses (both modes) | | `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) | | `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 | | `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_FORMAT` | `text` | `json` for structured JSON logs |
| `LOG_FILE` | _(empty)_ | Append logs to this file instead of stderr | | `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 class) in text mode, which renders in `docker compose logs` on Windows
Terminal. Set `NO_COLOR` to disable colors. 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 ## Build and run
```sh ```sh
+25
View File
@@ -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 with empty prompt to reload the chat model so the next chat doesn't pay the
load time. Off by default. 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 (env)
Configuration comes from environment variables and/or an `.env`-style 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 (400599, e.g. 429) | | `LLM_BUSY_STATUS` | `503` | HTTP status for rejected LLM requests in reject mode (400599, e.g. 429) |
| `BUSY_RETRY_AFTER` | `30` | seconds sent as `Retry-After` on busy responses (both modes) | | `BUSY_RETRY_AFTER` | `30` | seconds sent as `Retry-After` on busy responses (both modes) |
| `WARM_MODEL` | `` | optional model to reload after an image job | | `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 | | `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_FORMAT` | `text` | `json` for structured JSON logs |
| `LOG_FILE` | `` | append logs to this file instead of stderr (useful as a service) | | `LOG_FILE` | `` | append logs to this file instead of stderr (useful as a service) |
+64 -3
View File
@@ -26,6 +26,7 @@ import (
"gpu-turnstile/internal/ollama" "gpu-turnstile/internal/ollama"
"gpu-turnstile/internal/proxy" "gpu-turnstile/internal/proxy"
"gpu-turnstile/internal/service" "gpu-turnstile/internal/service"
"gpu-turnstile/internal/supervise"
"gpu-turnstile/internal/update" "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, "backoff_max", cfg.BackoffMax,
"prompt_capture_limit", cfg.PromptCaptureLimit, "prompt_capture_limit", cfg.PromptCaptureLimit,
"warm_model", cfg.WarmModel, "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, "auto_update", cfg.AutoUpdate,
"update_interval", cfg.UpdateInterval, "update_interval", cfg.UpdateInterval,
"update_repo", cfg.UpdateRepo, "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{ srv, err := proxy.New(proxy.Config{
OllamaURL: cfg.OllamaURL, OllamaURL: cfg.OllamaURL,
ComfyURL: cfg.ComfyURL, ComfyURL: cfg.ComfyURL,
Lock: lk, Lock: lk,
Ollama: ollamaClient, Ollama: ollamaClient,
Comfy: comfyClient, Comfy: comfyClient,
ComfySup: comfySup,
Metrics: metrics.New(), Metrics: metrics.New(),
Log: log, Log: log,
LogColor: !cfg.LogJSON && cfg.LogFile == "" && os.Getenv("NO_COLOR") == "", 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 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{} probes := map[string]func(context.Context) error{}
if ollamaClient != nil { if ollamaClient != nil {
probes["ollama"] = ollamaClient.Probe probes["ollama"] = ollamaClient.Probe
} }
if comfyClient != nil { if comfyClient != nil {
if comfySup == nil {
probes["comfy"] = comfyClient.Probe 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) probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout)
for name, probe := range probes { 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) 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 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 // healthLoop probes the enabled upstreams every interval and logs status
// transitions — "is DOWN" when a previously healthy upstream stops // transitions — "is DOWN" when a previously healthy upstream stops
// answering, "recovered" when it comes back. The first round only // 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) ticker := time.NewTicker(interval)
defer ticker.Stop() defer ticker.Stop()
up := map[string]bool{} up := map[string]bool{}
managed := map[string]bool{}
for { for {
select { select {
case <-ctx.Done(): case <-ctx.Done():
@@ -584,8 +634,19 @@ func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *
pctx, cancel := context.WithTimeout(ctx, probeTimeout) pctx, cancel := context.WithTimeout(ctx, probeTimeout)
err := probe(pctx) err := probe(pctx)
cancel() cancel()
was, seen := up[name] if errors.Is(err, errManagedDown) {
managed[name] = true // intentionally stopped; not an outage
continue
}
now := err == nil 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 seen && now != was {
if now { if now {
log.Warn(name + " upstream recovered") log.Warn(name + " upstream recovered")
+20
View File
@@ -51,6 +51,16 @@ type Config struct {
LLMBusyStatus int LLMBusyStatus int
BusyRetryAfter 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 WarmModel string
LogLevel slog.Level LogLevel slog.Level
LogJSON bool LogJSON bool
@@ -89,6 +99,9 @@ func Defaults() Config {
LLMBusyStatus: 503, LLMBusyStatus: 503,
BusyRetryAfter: 30, BusyRetryAfter: 30,
ComfyIdleTimeout: 5 * time.Minute,
ComfyStartTimeout: 2 * time.Minute,
LogLevel: slog.LevelWarn, LogLevel: slog.LevelWarn,
} }
} }
@@ -150,6 +163,8 @@ func Load(getenv func(string) string) (Config, error) {
{"OLLAMA_URL", &cfg.OllamaURL}, {"OLLAMA_URL", &cfg.OllamaURL},
{"COMFY_URL", &cfg.ComfyURL}, {"COMFY_URL", &cfg.ComfyURL},
{"WARM_MODEL", &cfg.WarmModel}, {"WARM_MODEL", &cfg.WarmModel},
{"COMFY_CMD", &cfg.ComfyCmd},
{"COMFY_DIR", &cfg.ComfyDir},
{"UPDATE_REPO", &cfg.UpdateRepo}, {"UPDATE_REPO", &cfg.UpdateRepo},
{"UPDATE_ASSET", &cfg.UpdateAsset}, {"UPDATE_ASSET", &cfg.UpdateAsset},
{"LOG_FILE", &cfg.LogFile}, {"LOG_FILE", &cfg.LogFile},
@@ -174,6 +189,8 @@ func Load(getenv func(string) string) (Config, error) {
{"SHUTDOWN_TIMEOUT", &cfg.ShutdownTimeout}, {"SHUTDOWN_TIMEOUT", &cfg.ShutdownTimeout},
{"BACKOFF_INITIAL", &cfg.BackoffInitial}, {"BACKOFF_INITIAL", &cfg.BackoffInitial},
{"BACKOFF_MAX", &cfg.BackoffMax}, {"BACKOFF_MAX", &cfg.BackoffMax},
{"COMFY_IDLE_TIMEOUT", &cfg.ComfyIdleTimeout},
{"COMFY_START_TIMEOUT", &cfg.ComfyStartTimeout},
{"UPDATE_INTERVAL", &cfg.UpdateInterval}, {"UPDATE_INTERVAL", &cfg.UpdateInterval},
} { } {
if err := envDuration(getenv, e.name, e.dst); err != nil { if err := envDuration(getenv, e.name, e.dst); err != nil {
@@ -244,6 +261,9 @@ func Load(getenv func(string) string) (Config, error) {
default: default:
return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"") 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 == "" { if cfg.OllamaURL == "" && cfg.ComfyURL == "" {
return cfg, ErrNoConsumer return cfg, ErrNoConsumer
} }
+23
View File
@@ -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) { func TestParseEnvFile(t *testing.T) {
input := `# comment input := `# comment
OLLAMA_URL=http://host:11435 OLLAMA_URL=http://host:11435
+4
View File
@@ -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}, {"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}, {"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}, {"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}, {"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}, {"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}, {"LLM_WAIT_TIMEOUT", "10m", "Max time an LLM request waits for the GPU before being answered 503 (wait mode)", false},
+32
View File
@@ -26,6 +26,7 @@ import (
"gpu-turnstile/internal/lock" "gpu-turnstile/internal/lock"
"gpu-turnstile/internal/metrics" "gpu-turnstile/internal/metrics"
"gpu-turnstile/internal/ollama" "gpu-turnstile/internal/ollama"
"gpu-turnstile/internal/supervise"
) )
// defaultCaptureLimit bounds how much of a /prompt response body is // defaultCaptureLimit bounds how much of a /prompt response body is
@@ -44,6 +45,11 @@ type Config struct {
Metrics *metrics.Metrics Metrics *metrics.Metrics
Log *slog.Logger 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 // LogColor enables ANSI colors in per-request log lines. Ignored when
// the log level is above INFO (request lines are not emitted at all). // the log level is above INFO (request lines are not emitted at all).
LogColor bool LogColor bool
@@ -500,6 +506,14 @@ func (s *Server) ComfyHandler() http.Handler {
s.handlePrompt(w, r) s.handlePrompt(w, r)
return 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) 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) { func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
log := s.log.With("op", "image") 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() start := time.Now()
if err := s.cfg.Lock.AcquireImage(r.Context()); err != nil { if err := s.cfg.Lock.AcquireImage(r.Context()); err != nil {
if errors.Is(err, context.DeadlineExceeded) { if errors.Is(err, context.DeadlineExceeded) {
@@ -608,6 +636,10 @@ func (s *Server) finishImageJob(promptID string) {
s.cfg.Lock.ReleaseImage() s.cfg.Lock.ReleaseImage()
log.Info("image lock released") 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 s.cfg.Ollama != nil && s.cfg.WarmModel != "" {
if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle { if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle {
+257
View File
@@ -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
}
+158
View File
@@ -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")
}
}