Managed ComfyUI: COMFY_CMD starts it on demand, idle stop frees VRAM (internal/supervise)
This commit is contained in:
mram
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user