park_ttl: an engine's own park ttl, so ComfyUI nodes that call other engines finish

CHANGELOG.md

@@ -40,6 +40,11 @@ has a matching GitHub release whose body comes from its section below.
   restore at 30 s, so whatever goes wrong along the way costs one line and the switch goes on as
   before. A snapshot holds the conversation's tokens: the ledger is owner-only, and the README asks
   for a private directory.
+- `park_ttl` on an `[[engine]]` entry sets how long that engine stays parked before Berth stops it,
+  in place of `[park] ttl`. With `keep_comfyui_up = true` and a short `park_ttl`, a ComfyUI node that
+  calls another engine (a chat panel calling llama-swap) gets its answer. The drain leaves ComfyUI
+  up, and Berth stops the server once the ttl passes
+  ([#121](https://github.com/T-Crypt/berth/issues/121)).
 
 ### Changed
 

contrib/INTEGRATION.md

@@ -98,6 +98,12 @@ berth run --as comfyui --priority batch --wait 87600h --drain-comfyui http://127
   python main.py --listen 127.0.0.1 --port 8188
 ```
 
+If a ComfyUI node calls another engine (a chat panel, a prompt expander), that call can drain
+ComfyUI while the node waits for its answer, and stopping the server kills the request. Set
+`keep_comfyui_up = true` and a short `park_ttl` on the ComfyUI engine entry, as
+[the gateway docs](https://github.com/T-Crypt/berth/blob/main/docs/gateway.md#comfyui-nodes-that-call-other-engines)
+show.
+
 Reading: `type = "comfyui"` (found on :8188 without config) shows jobs running and queued.
 
 ## vLLM and SGLang

docs/gateway.md

@@ -321,6 +321,31 @@ Completed-session release ends the process; daemon-requested parking and resume
 keep their native protocol and lease routing. Without these options, behavior
 is unchanged.
 
+## ComfyUI nodes that call other engines
+
+Berth hands ComfyUI's card over once its queue is empty. Some nodes do their work outside the
+queue. A chat panel or a prompt expander inside ComfyUI can call a model on llama-swap, and that
+call can drain ComfyUI's own lease. Without `keep_comfyui_up`, the drain stops the ComfyUI server,
+and the request that asked for the card dies with it.
+
+Keep ComfyUI up on a drain and give it a short park ttl:
+
+```toml
+[[engine]]
+name            = "comfyui"
+type            = "comfyui"
+url             = "http://127.0.0.1:8188"
+drain_comfyui   = "http://127.0.0.1:8188"
+keep_comfyui_up = true
+park_ttl        = "10m"
+```
+
+A drain then unloads ComfyUI's models through `POST /free` and leaves the server running, so the
+node's request finishes. The lease stays parked. After `park_ttl` Berth stops the server, which
+frees the CUDA context it still holds (504 MiB on the reference card). That stop is a normal drain,
+so a job ComfyUI is running at that moment finishes first. `park_ttl` works on any parked engine
+and replaces `[park] ttl` (30 minutes by default) for that engine.
+
 ## Park in RAM (warm switch)
 
 A drain today ends with the engine's weights on disk, and the next request
@@ -344,6 +369,7 @@ url        = "http://127.0.0.1:8000"
 park       = "hook"
 park_url   = "http://127.0.0.1:8000/sleep?level=1"   # optional: each beats the built-in
 resume_url = "http://127.0.0.1:8000/wake_up"
+park_ttl   = "1h"                                     # optional: this engine's own ttl
 
 [park]
 floor_gib = 4.0        # RAM a park must leave the machine (default 4)

internal/broker/broker.go

@@ -184,6 +184,9 @@ type Request struct {
 	// the turn ends the normal rules apply. A re-attach starts a fresh
 	// turn. 0 = off.
 	MinTurn time.Duration
+	// ParkTTL is how long the lease may stay parked before the broker
+	// stops its engine, in place of the node's ParkTTL. 0 = the node's.
+	ParkTTL time.Duration
 	// Container is the engine's container: Acquire takes the name the
 	// operator gave (--container), the attach resolves the full ID. A GPU
 	// process inside it belongs to the lease.
@@ -206,6 +209,7 @@ type lease struct {
 	reattach bool // re-attached after a daemon restart: granted is the re-attach, not a fresh switch
 	pinned   bool // never drained (Request.Pinned)
 	minTurn  time.Duration
+	parkTTL  time.Duration // 0 = the node's
 	prio     Priority
 	needs    uint64
 	state    state
@@ -290,6 +294,7 @@ func newLease(id string, r Request, now time.Time) *lease {
 		needs:    r.Needs,
 		pinned:   r.Pinned,
 		minTurn:  r.MinTurn,
+		parkTTL:  r.ParkTTL,
 		gpu:      map[ident]gpu.Proc{},
 		desc:     map[ident]struct{}{},
 		arrived:  now,

internal/broker/park.go

@@ -12,6 +12,7 @@ package broker
 // an eviction is a drain the runner stops, exactly as today.
 
 import (
+	"cmp"
 	"fmt"
 	"sort"
 	"strings"
@@ -161,7 +162,7 @@ func (b *Broker) evictParked(now time.Time) {
 	var entries []park.Parked
 	for _, l := range b.leases {
 		if l.state == parked {
-			entries = append(entries, park.Parked{Name: l.id, Bytes: b.footprintOf(l), Since: l.parkedAt})
+			entries = append(entries, park.Parked{Name: l.id, Bytes: b.footprintOf(l), Since: l.parkedAt, TTL: l.parkTTL})
 		}
 	}
 	if len(entries) == 0 {
@@ -181,7 +182,7 @@ func (b *Broker) evictParked(now time.Time) {
 		}
 		why := "the memory pressure is high"
 		switch {
-		case now.Sub(l.parkedAt) > b.cfg.ParkTTL:
+		case now.Sub(l.parkedAt) > cmp.Or(l.parkTTL, b.cfg.ParkTTL):
 			why = fmt.Sprintf("parked %s, past the ttl", now.Sub(l.parkedAt).Round(time.Second))
 		case avail < b.cfg.ParkFloor:
 			why = "available RAM is under the floor"

internal/broker/park_test.go

@@ -292,6 +292,30 @@ func TestParkedEnginesAreEvicted(t *testing.T) {
 	}
 }
 
+// An engine's park_ttl stops it sooner than the node's ttl: ComfyUI kept up
+// on a drain gives its context back minutes later, not half an hour.
+func TestParkedEngineOwnTTL(t *testing.T) {
+	tree := parkTree(100, 101)
+	b, c, _ := parkSetup(t, tree, 16*GiB, 0)
+	l := warmLease(t, b, c, "comfyui", 8*GiB, 100, 101)
+	b.cfg.MemAvailable = func() uint64 { return 16 * GiB }
+	b.cfg.MemPressure = func() float64 { return 0 }
+	b.mu.Lock()
+	l.l.parkTTL = 5 * time.Minute
+	l.l.parkedAt = time.Now().Add(-6 * time.Minute)
+	b.mu.Unlock()
+
+	b.Tick(context.Background())
+	select {
+	case why := <-l.Drains():
+		if !strings.Contains(why, "past the ttl") {
+			t.Fatalf("eviction drain reason %q, want past the ttl", why)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("the parked engine outlived its own ttl")
+	}
+}
+
 // A resume puts the parked lease back in the queue at its own priority,
 // and the grant closes the resume signal for the runner's connection.
 func TestResumeGrantsTheParkedLease(t *testing.T) {

internal/config/config.go

@@ -91,6 +91,12 @@ type Engine struct {
 	// time, whatever its priority, only a hold takes the card out. When
 	// the turn ends the normal rules apply. 0 = off, the default.
 	MinTurn Duration `toml:"min_turn,omitzero" json:"min_turn,omitempty" yaml:"min_turn,omitempty"`
+	// ParkTTL is how long this engine stays parked before Berth stops it,
+	// in place of [park] ttl. A short one suits keep_comfyui_up: a drain
+	// leaves ComfyUI up, so work its own nodes run outside the queue (a
+	// chat call to another engine) finishes, and its CUDA context goes
+	// soon after. 0 = [park] ttl.
+	ParkTTL Duration `toml:"park_ttl,omitzero" json:"park_ttl,omitempty" yaml:"park_ttl,omitempty"`
 	// Container is the engine's container, as berth run --container: the
 	// name of `docker run --name` (or of an existing container for
 	// `docker start -a`). The drain stops it through the runtime CLI
@@ -601,6 +607,9 @@ func (c *Config) Validate() error {
 		if e.MinTurn.Duration < 0 {
 			errs = append(errs, fmt.Errorf("%s: min_turn %s must be >= 0", where, e.MinTurn))
 		}
+		if e.ParkTTL.Duration < 0 {
+			errs = append(errs, fmt.Errorf("%s: park_ttl %s must be >= 0", where, e.ParkTTL))
+		}
 		if e.GUI != "" {
 			if u, err := url.Parse(e.GUI); err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
 				errs = append(errs, fmt.Errorf("%s: gui %q must be an http(s) URL", where, e.GUI))

internal/config/config_test.go

@@ -820,3 +820,17 @@ fallback = "strata"
 		t.Errorf("a chain was accepted or names neither engine: %v", err)
 	}
 }
+
+func TestEngineParkTTL(t *testing.T) {
+	const entry = "[[engine]]\nname = \"comfyui\"\ntype = \"comfyui\"\nurl = \"http://127.0.0.1:8188\"\n"
+	c, err := Load(write(t, entry+"park_ttl = \"10m\"\n"))
+	if err != nil {
+		t.Fatalf("a valid park_ttl was refused: %v", err)
+	}
+	if got := c.Engine[0].ParkTTL.Duration; got != 10*time.Minute {
+		t.Errorf("park_ttl = %s, want 10m", got)
+	}
+	if _, err := Load(write(t, entry+"park_ttl = \"-1m\"\n")); err == nil || !strings.Contains(err.Error(), `engine[0] "comfyui": park_ttl -1m0s must be >= 0`) {
+		t.Errorf("a negative park_ttl: %v", err)
+	}
+}

internal/park/admit.go

@@ -4,6 +4,7 @@
 package park
 
 import (
+	"cmp"
 	"fmt"
 	"slices"
 	"strconv"
@@ -25,6 +26,7 @@ type Parked struct {
 	Name  string
 	Bytes uint64
 	Since time.Time
+	TTL   time.Duration // its own ttl; 0 = Evict's
 }
 
 const gib = uint64(1) << 30
@@ -69,7 +71,7 @@ func Evict(parked []Parked, b Budget, pressure, threshold float64, now time.Time
 	}
 	avail := b.Available
 	for _, p := range byAge {
-		if now.Sub(p.Since) > ttl {
+		if now.Sub(p.Since) > cmp.Or(p.TTL, ttl) {
 			stop(p)
 			avail += p.Bytes
 		}

internal/park/admit_test.go

@@ -105,6 +105,16 @@ func TestEvict(t *testing.T) {
 			pressure: 0,
 			want:     []string{"staler", "stale"},
 		},
+		{
+			name: "an engine's own ttl beats the node's, both ways",
+			parked: []Parked{
+				{Name: "short", Bytes: g, Since: now.Add(-6 * time.Minute), TTL: 5 * time.Minute},
+				{Name: "long", Bytes: g, Since: now.Add(-40 * time.Minute), TTL: time.Hour},
+				m("node", g, 6*time.Minute),
+			},
+			b:    Budget{Available: 10 * g, Floor: 4 * g},
+			want: []string{"short"},
+		},
 		{
 			name:     "below the floor: least recently parked until back at it",
 			parked:   []Parked{m("new", 4*g, time.Minute), m("mid", 4*g, 5*time.Minute), m("old", 4*g, 10*time.Minute)},

internal/proto/proto.go

@@ -75,6 +75,9 @@ type Msg struct {
 	// younger than it after the grant, no waiter drains it but a hold.
 	// OpAcquire and OpReattach; 0 = off.
 	MinTurn float64 `json:"min_turn,omitempty"`
+	// ParkTTL is how long the lease may stay parked, in seconds, in place
+	// of the node's [park] ttl. OpAcquire and OpReattach; 0 = the node's.
+	ParkTTL float64 `json:"park_ttl,omitempty"`
 	// Container is the container the engine runs in: OpAcquire and
 	// OpReattach carry the name the operator gave (--container), OpAttach
 	// the full ID berth run resolved it to. The broker matches GPU

internal/runner/options.go

@@ -54,7 +54,7 @@ func OptionsFor(e *config.Engine, o Over) Options {
 		Socket: o.Socket, Engine: e.Name, Group: o.Group, Priority: prio,
 		Needs: e.NeedsMiB << 20, Pinned: o.Pinned, Wait: o.Wait,
 		DrainPost: e.DrainPost, DrainSlots: e.DrainSlots, DrainComfyUI: e.DrainComfyUI,
-		DrainMax: e.DrainMax.Duration, MinTurn: e.MinTurn.Duration, Container: e.Container, APIKeyEnv: e.APIKeyEnv, APIKeyFile: e.APIKeyFile,
+		DrainMax: e.DrainMax.Duration, MinTurn: e.MinTurn.Duration, ParkTTL: e.ParkTTL.Duration, Container: e.Container, APIKeyEnv: e.APIKeyEnv, APIKeyFile: e.APIKeyFile,
 		ExitComfyUIIdle: e.ExitComfyUIIdle, ExitComfyUIIdleAfter: e.ExitComfyUIIdleAfter.Duration,
 		Grace: grace, FailClosed: o.FailClosed,
 		Argv: argv, Stderr: o.Stderr, Output: o.Output, Log: o.Log,

internal/runner/options_test.go

@@ -18,13 +18,14 @@ func TestOptionsFor(t *testing.T) {
 	e := &config.Engine{Name: "strata", Type: "strata", Priority: "foreground", NeedsMiB: 8192,
 		Start: []string{"python", "serve.py"}, DrainPost: "http://127.0.0.1:1/unload",
 		DrainSlots: "http://127.0.0.1:2", DrainComfyUI: "http://127.0.0.1:3",
-		DrainMax: config.Duration{Duration: time.Minute}, MinTurn: config.Duration{Duration: 5 * time.Minute}}
+		DrainMax: config.Duration{Duration: time.Minute}, MinTurn: config.Duration{Duration: 5 * time.Minute},
+		ParkTTL: config.Duration{Duration: 10 * time.Minute}}
 
 	o := OptionsFor(e, Over{Socket: "s", Group: "strata", Wait: 10 * time.Second, Argv: []string{"true"}})
 	want := Options{
 		Socket: "s", Engine: "strata", Group: "strata", Priority: "foreground", Needs: 8192 << 20,
 		Wait: 10 * time.Second, Grace: 30 * time.Second, DrainPost: e.DrainPost, DrainSlots: e.DrainSlots,
-		DrainComfyUI: e.DrainComfyUI, DrainMax: e.DrainMax.Duration, MinTurn: 5 * time.Minute, Argv: []string{"true"},
+		DrainComfyUI: e.DrainComfyUI, DrainMax: e.DrainMax.Duration, MinTurn: 5 * time.Minute, ParkTTL: 10 * time.Minute, Argv: []string{"true"},
 	}
 	if !reflect.DeepEqual(o, want) {
 		t.Errorf("entry settings = %+v, want %+v", o, want)

internal/runner/runner.go

@@ -46,7 +46,10 @@ type Options struct {
 	// neither does a higher one; only a hold (a game, a call, berth hold)
 	// takes the card out. When the turn ends the normal rules apply; a
 	// re-attach starts a fresh turn. 0 = off.
-	MinTurn   time.Duration
+	MinTurn time.Duration
+	// ParkTTL is how long the lease may stay parked before the broker stops
+	// the engine, in place of the node's [park] ttl. 0 = the node's.
+	ParkTTL   time.Duration
 	Wait      time.Duration
 	DrainPost string // asked first on drain, e.g. llama-swap's unload URL
 	// DrainSlots is an engine base URL with a llama.cpp-style /slots: on
@@ -167,7 +170,7 @@ func Run(o Options) int {
 	}
 	defer raw.Close()
 
-	if err := conn.Send(proto.Msg{Op: proto.OpAcquire, Engine: o.Engine, Group: o.Group, Priority: o.Priority, Needs: o.Needs, Wait: o.Wait.Seconds(), Pinned: o.Pinned, MinTurn: o.MinTurn.Seconds()}); err != nil {
+	if err := conn.Send(proto.Msg{Op: proto.OpAcquire, Engine: o.Engine, Group: o.Group, Priority: o.Priority, Needs: o.Needs, Wait: o.Wait.Seconds(), Pinned: o.Pinned, MinTurn: o.MinTurn.Seconds(), ParkTTL: o.ParkTTL.Seconds()}); err != nil {
 		say("acquire: %v", err)
 		return ExitRefused
 	}
@@ -340,7 +343,7 @@ func reattach(o Options, cmd *exec.Cmd, done <-chan struct{}, cid *containerRef)
 		if err != nil {
 			continue
 		}
-		err = conn.Send(proto.Msg{Op: proto.OpReattach, Engine: o.Engine, Group: o.Group, Priority: o.Priority, Needs: o.Needs, Pinned: o.Pinned, MinTurn: o.MinTurn.Seconds(), PID: pid, PIDStart: start, Container: cid.id()})
+		err = conn.Send(proto.Msg{Op: proto.OpReattach, Engine: o.Engine, Group: o.Group, Priority: o.Priority, Needs: o.Needs, Pinned: o.Pinned, MinTurn: o.MinTurn.Seconds(), ParkTTL: o.ParkTTL.Seconds(), PID: pid, PIDStart: start, Container: cid.id()})
 		var m proto.Msg
 		if err == nil {
 			m, err = conn.Recv()

internal/server/server.go

@@ -226,6 +226,7 @@ func hold(ctx context.Context, conn *proto.Conn, req proto.Msg, b *broker.Broker
 		Wait:      time.Duration(req.Wait * float64(time.Second)),
 		Pinned:    req.Pinned,
 		MinTurn:   time.Duration(req.MinTurn * float64(time.Second)),
+		ParkTTL:   time.Duration(req.ParkTTL * float64(time.Second)),
 		Container: req.Container,
 	}
 	var lease *broker.Lease