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

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
 	}

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

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},

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 {

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
+}

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")
+	}
+}