gpu-turnstile/internal/supervise/supervise.go

// 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"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"
)

// ComfyLayout returns the interpreter and script path of a standard ComfyUI
// venv install rooted at dir for the given GOOS: .venv\Scripts\python.exe
// on Windows, .venv/bin/python elsewhere. The script is main.py — either in
// a ComfyUI subdirectory or directly under dir, whichever exists (the
// subdirectory form wins ties and is the default when neither exists yet,
// so the caller's missing-file warning points at the documented layout).
func ComfyLayout(goos, dir string) (python, script string) {
	if goos == "windows" {
		python = filepath.Join(dir, ".venv", "Scripts", "python.exe")
	} else {
		python = filepath.Join(dir, ".venv", "bin", "python")
	}
	script = filepath.Join(dir, "ComfyUI", "main.py")
	if _, err := os.Stat(script); err != nil {
		if _, err := os.Stat(filepath.Join(dir, "main.py")); err == nil {
			script = filepath.Join(dir, "main.py")
		}
	}
	return python, script
}

// DefaultComfyCommand builds the launch command for the standard venv
// layout (see ComfyLayout): the script is passed relative to dir so dir
// stays the working directory, and --port is taken from comfyURL when the
// URL carries one.
func DefaultComfyCommand(goos, dir, comfyURL string) string {
	python, script := ComfyLayout(goos, dir)
	rel, err := filepath.Rel(dir, script)
	if err != nil {
		rel = script
	}
	cmd := `"` + python + `" ` + rel
	if u, err := url.Parse(comfyURL); err == nil && u.Port() != "" {
		cmd += " --port " + u.Port()
	}
	return cmd
}

// 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
	external     bool // someone else serves the port; not our process
	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). When the URL
// already answers — e.g. the ComfyUI desktop app grabbed the port — no
// child is spawned: the external server is used as-is, and the idle
// watcher never touches it (it only kills its own child).
func (p *Process) EnsureRunning() error {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.lastActivity = time.Now()
	if p.cmd != nil {
		return nil
	}
	pctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	err := p.probe(pctx)
	cancel()
	if err == nil {
		p.ready = true
		if !p.external {
			p.external = true
			p.log.Info(p.name + " is already served externally; not spawning a managed instance")
		}
		return nil
	}
	p.external = false
	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
}