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

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