diff --git a/CHANGELOG.md b/CHANGELOG.md index ac0808f..5ce1511 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/contrib/INTEGRATION.md b/contrib/INTEGRATION.md index 2e017ee..aa21636 100644 --- a/contrib/INTEGRATION.md +++ b/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 diff --git a/docs/gateway.md b/docs/gateway.md index 8d6ee46..e14b505 100644 --- a/docs/gateway.md +++ b/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) diff --git a/internal/broker/broker.go b/internal/broker/broker.go index e2c946b..d3a9974 100644 --- a/internal/broker/broker.go +++ b/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, diff --git a/internal/broker/park.go b/internal/broker/park.go index 90fbb1d..a263aab 100644 --- a/internal/broker/park.go +++ b/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" diff --git a/internal/broker/park_test.go b/internal/broker/park_test.go index eb17aab..69cc11b 100644 --- a/internal/broker/park_test.go +++ b/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) { diff --git a/internal/config/config.go b/internal/config/config.go index 2ac81ef..c740780 100644 --- a/internal/config/config.go +++ b/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)) diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 02d7b81..e81f3aa 100644 --- a/internal/config/config_test.go +++ b/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) + } +} diff --git a/internal/park/admit.go b/internal/park/admit.go index b98a1e6..73cce74 100644 --- a/internal/park/admit.go +++ b/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 } diff --git a/internal/park/admit_test.go b/internal/park/admit_test.go index 4327ee9..d42c6e3 100644 --- a/internal/park/admit_test.go +++ b/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)}, diff --git a/internal/proto/proto.go b/internal/proto/proto.go index c903c1a..6e6395b 100644 --- a/internal/proto/proto.go +++ b/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 diff --git a/internal/runner/options.go b/internal/runner/options.go index 63b7e2a..f7bcec7 100644 --- a/internal/runner/options.go +++ b/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, diff --git a/internal/runner/options_test.go b/internal/runner/options_test.go index e2261b1..7bb57bb 100644 --- a/internal/runner/options_test.go +++ b/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) diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 01742ce..83d0ab6 100644 --- a/internal/runner/runner.go +++ b/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() diff --git a/internal/server/server.go b/internal/server/server.go index 6e3a815..fff5901 100644 --- a/internal/server/server.go +++ b/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