From cb865ae70d2637b619ee1de0cdac604c9ae0295e Mon Sep 17 00:00:00 2001 From: T-Crypt <132922589+T-Crypt@users.noreply.github.com> Date: Mon, 2 Nov 2026 08:30:49 -0700 Subject: [PATCH 1/2] feat: a size per card for a lease on several cards A model split unevenly across cards (llama.cpp --tensor-split 3,1) needs a different amount on each. needs_mib and berth run --needs-mib took one size for every card, so the lighter card reserved as much as the heavier one and nothing else could use the room left on it. needs_mib now reads a number or a list (needs_mib = [18000, 6000]), and --needs-mib reads "18000,6000", one size per card in the order gpu names them. The run sends the list as needs_each and the largest size as needs, so a berthd that predates needs_each reserves too much on the smaller card, never too little. berthd refuses a list that does not match the cards, and each card's share of the lease reserves its own size. One value sizes every card, as before, and writes back as a plain number. --- internal/broker/gang.go | 33 ++++-- internal/broker/gang_test.go | 37 ++++++- internal/cli/config.go | 4 +- internal/cli/run.go | 15 +-- internal/cli/run_test.go | 27 +++-- internal/config/config.go | 9 +- internal/config/mibs.go | 162 ++++++++++++++++++++++++++++++ internal/config/mibs_test.go | 97 ++++++++++++++++++ internal/gateway/gateway_test.go | 5 +- internal/gateway/resume_test.go | 2 +- internal/gateway/session_test.go | 15 +++ internal/gateway/spawn.go | 10 +- internal/node/node_test.go | 2 +- internal/proto/proto.go | 5 + internal/runner/auth_test.go | 2 +- internal/runner/card_test.go | 56 +++++++++++ internal/runner/front_test.go | 6 +- internal/runner/lend_test.go | 2 +- internal/runner/options.go | 22 +++- internal/runner/options_test.go | 4 +- internal/runner/park_auth_test.go | 2 +- internal/runner/park_test.go | 8 +- internal/runner/runner.go | 8 +- internal/server/server.go | 11 +- 24 files changed, 490 insertions(+), 54 deletions(-) create mode 100644 internal/config/mibs.go create mode 100644 internal/config/mibs_test.go diff --git a/internal/broker/gang.go b/internal/broker/gang.go index 76d0ad2..f04438e 100644 --- a/internal/broker/gang.go +++ b/internal/broker/gang.go @@ -82,11 +82,13 @@ func (g *gang) commit(now time.Time) bool { // AcquireAll asks for one lease on every card in cards, granted on all of // them at once or on none, and returns its share on each card in the order -// given. One card is a plain Acquire. Each share needs r.Needs on its card. -// The cards must share the node's IDs and Seq. -func AcquireAll(ctx context.Context, cards []*Broker, r Request, onWait func(string)) ([]*Lease, error) { +// given. One card is a plain Acquire. needs is the size on each card, in +// the order of cards; nil gives every share r.Needs. The cards must share +// the node's IDs and Seq. +func AcquireAll(ctx context.Context, cards []*Broker, needs []uint64, r Request, onWait func(string)) ([]*Lease, error) { + per := shareRequests(cards, needs, r) if len(cards) == 1 { - l, err := cards[0].Acquire(ctx, r, onWait) + l, err := cards[0].Acquire(ctx, per[cards[0]], onWait) if err != nil { return nil, err } @@ -96,7 +98,7 @@ func AcquireAll(ctx context.Context, cards []*Broker, r Request, onWait func(str slices.SortFunc(sorted, func(a, b *Broker) int { return strings.Compare(a.cfg.Card, b.cfg.Card) }) for _, b := range sorted { b.mu.Lock() - err := b.refuse(r) + err := b.refuse(per[b]) b.mu.Unlock() if err != nil { return nil, err @@ -110,7 +112,7 @@ func AcquireAll(ctx context.Context, cards []*Broker, r Request, onWait func(str defer close(done) for _, b := range sorted { b.mu.Lock() - l := newLease(id, r, b.now()) + l := newLease(id, per[b], b.now()) l.arrivedN, l.gang = arrived, g b.queue = append(b.queue, l) b.log.Info("acquire", "lease", l.id, "engine", l.engine, "priority", l.prio, "needs_mib", l.needs>>20, "cards", len(sorted)) @@ -189,11 +191,12 @@ func AcquireAll(ctx context.Context, cards []*Broker, r Request, onWait func(str // ReattachAll leases every card in cards again, under one lease ID, to an // engine that held a lease on all of them and outlived its daemon. It is // all or none: a card that refuses hands back the shares already taken. -func ReattachAll(cards []*Broker, r Request, pid int, start uint64) ([]*Lease, error) { +func ReattachAll(cards []*Broker, needs []uint64, r Request, pid int, start uint64) ([]*Lease, error) { + per := shareRequests(cards, needs, r) id := fmt.Sprintf("%s-%d", r.Engine, cards[0].cfg.IDs.Add(1)) out := make([]*Lease, 0, len(cards)) for _, b := range cards { - l, err := b.reattach(id, r, pid, start) + l, err := b.reattach(id, per[b], pid, start) if err != nil { for _, got := range out { got.Release() @@ -204,3 +207,17 @@ func ReattachAll(cards []*Broker, r Request, pid int, start uint64) ([]*Lease, e } return out, nil } + +// shareRequests is the request each card's share is made from: r, with that +// card's size from needs (in the order of cards) when needs is given. +func shareRequests(cards []*Broker, needs []uint64, r Request) map[*Broker]Request { + per := make(map[*Broker]Request, len(cards)) + for i, b := range cards { + rr := r + if needs != nil { + rr.Needs = needs[i] + } + per[b] = rr + } + return per +} diff --git a/internal/broker/gang_test.go b/internal/broker/gang_test.go index 8abfe4d..a13a169 100644 --- a/internal/broker/gang_test.go +++ b/internal/broker/gang_test.go @@ -38,7 +38,7 @@ type gangResult struct { func acquireAll(cards []*Broker, r Request) chan gangResult { ch := make(chan gangResult, 1) go func() { - ls, err := AcquireAll(context.Background(), cards, r, nil) + ls, err := AcquireAll(context.Background(), cards, nil, r, nil) ch <- gangResult{ls, err} }() return ch @@ -148,7 +148,7 @@ func TestTwoGangsInOppositeCardOrderBothFinish(t *testing.T) { // lease ID after a berthd restart, or on neither. func TestReattachAllIsOneLeaseOnEveryCard(t *testing.T) { cards := node(t, "gpu0", "gpu1") - ls, err := ReattachAll(cards, Request{Engine: "maya"}, 100, 7) + ls, err := ReattachAll(cards, nil, Request{Engine: "maya"}, 100, 7) if err != nil { t.Fatal(err) } @@ -165,7 +165,7 @@ func TestReattachAllIsOneLeaseOnEveryCard(t *testing.T) { if _, err := cards[1].Reattach(Request{Engine: "voice"}, 100, 7); err != nil { t.Fatal(err) } - if _, err := ReattachAll(cards, Request{Engine: "maya"}, 100, 7); err == nil { + if _, err := ReattachAll(cards, nil, Request{Engine: "maya"}, 100, 7); err == nil { t.Fatal("pid 100 already holds a lease on gpu1, and the re-attach was granted") } tickAll(cards) @@ -227,3 +227,34 @@ func TestPinnedLeaseOnSeveralCardsIsNeverDrained(t *testing.T) { t.Fatal("gpu1 stayed shut after the pinned run released") } } + +// A model split unevenly across two cards (llama.cpp --tensor-split 3,1) +// leases its own size on each, so the room left on the lighter card still +// serves another engine without a drain. +func TestSeveralCardsEachTakeTheirOwnSize(t *testing.T) { + bs := node(t, "gpu0", "gpu1") + ch := make(chan gangResult, 1) + go func() { + ls, err := AcquireAll(context.Background(), bs, []uint64{18 * GiB, 6 * GiB}, Request{Engine: "llama-70b", Priority: Foreground}, nil) + ch <- gangResult{ls, err} + }() + select { + case r := <-ch: + if r.err != nil || len(r.ls) != 2 { + t.Fatalf("split lease = %+v", r) + } + case <-time.After(2 * time.Second): + t.Fatal("the split lease was not granted on two free cards") + } + for i, want := range []uint64{18 * GiB, 6 * GiB} { + for _, l := range bs[i].Status().Leases { + if l.Engine == "llama-70b" && l.Needs != want { + t.Errorf("%s share needs %d MiB, want %d MiB", bs[i].cfg.Card, l.Needs>>20, want>>20) + } + } + } + must(t, acquire(bs[1], Request{Engine: "comfyui", Priority: Batch, Needs: 16 * GiB})) + if !holds(bs[1], "llama-70b") { + t.Fatal("comfyui's grant on gpu1 drained the split model") + } +} diff --git a/internal/cli/config.go b/internal/cli/config.go index 5cba9ac..86e6e86 100644 --- a/internal/cli/config.go +++ b/internal/cli/config.go @@ -143,8 +143,8 @@ func engineExtra(e config.Engine) string { if len(e.Serves) > 0 { parts = append(parts, "serves: "+strings.Join(e.Serves, ", ")) } - if e.NeedsMiB > 0 { - parts = append(parts, fmt.Sprintf("needs_mib: %d", e.NeedsMiB)) + if len(e.NeedsMiB.Bytes()) > 0 { + parts = append(parts, "needs_mib: "+e.NeedsMiB.String()) } if len(e.Start) > 0 { parts = append(parts, "start: "+e.Start[0]) diff --git a/internal/cli/run.go b/internal/cli/run.go index 8c61f84..29fda41 100644 --- a/internal/cli/run.go +++ b/internal/cli/run.go @@ -8,7 +8,6 @@ import ( "fmt" "path/filepath" "slices" - "strconv" "strings" "time" @@ -29,7 +28,8 @@ func runRun(c *Context) int { prio := fs.String("priority", "interactive", "batch, interactive, foreground or hold") pinned := fs.Bool("pinned", false, "never drained: the card stays CMD's until it exits (an app like Blender, or a training run; not an engine a request should swap out)") switched := fs.Bool("switched", false, "a person's switch (the gateway sets it for the wake page's Launch and /v1/berth/switch): wins a juggle") - needs := fs.Uint64("needs-mib", 0, "memory this engine needs; 0 = the whole card") + var needs config.MiBs + fs.Var(&needs, "needs-mib", "memory this engine needs in MiB; 0 = the whole card; with --gpu gpu0,gpu1, one size for every card or one per card in the same order (18000,6000)") wait := fs.Duration("wait", 0, "how long to wait for a berth; 0 = berthd's default; negative = no deadline") drainPost := fs.String("drain-post", "", "URL to POST when asked to drain, before signalling (e.g. llama-swap's unload)") drainSlots := fs.String("drain-slots", "", "engine base URL whose /slots must show idle before a drain unloads or stops it (Strata, or a llama-swap model's /upstream/)") @@ -85,7 +85,7 @@ func runRun(c *Context) int { return c.fail(required) } applyEngineEntry(c, &e, engineFlags{ - priority: prio, needsMiB: needs, drainPost: drainPost, drainSlots: drainSlots, + priority: prio, needsMiB: &needs, drainPost: drainPost, drainSlots: drainSlots, drainComfyUI: drainComfy, keepComfyUI: keepComfy, drainMax: drainMax, minTurn: minTurn, window: &win, container: container, compose: composeFile, }) @@ -148,6 +148,9 @@ func runRun(c *Context) int { }); err != nil { return c.fail(Systemf("%v", err)) } + if err := config.CheckNeeds(e.NeedsMiB, e.GPU); err != nil { + return c.fail(Systemf("--needs-mib: %v", err)) + } // Nothing here goes to stdout: the engine owns it. berth's own notices are // diagnostics on stderr, and --quiet never hides the unprotected warning. return runner.Run(runner.OptionsFor(&e, runner.Over{ @@ -203,7 +206,7 @@ func engineEntry(cfg config.Config, as, group string) *config.Engine { // engineFlags are the `berth run` flags that fold onto an engine's entry. type engineFlags struct { priority *string - needsMiB *uint64 + needsMiB *config.MiBs drainPost *string drainSlots *string drainComfyUI *string @@ -235,8 +238,8 @@ func applyEngineEntry(c *Context, e *config.Engine, f engineFlags) { e.Priority = *prio } if c.Changed("needs-mib") { - if e.NeedsMiB > 0 && e.NeedsMiB != *needs { - notice("needs-mib", "needs_mib", strconv.FormatUint(e.NeedsMiB, 10)) + if len(e.NeedsMiB.Bytes()) > 0 && !slices.Equal(e.NeedsMiB, *needs) { + notice("needs-mib", "needs_mib", e.NeedsMiB.String()) } e.NeedsMiB = *needs } diff --git a/internal/cli/run_test.go b/internal/cli/run_test.go index 8f5e144..71225ff 100644 --- a/internal/cli/run_test.go +++ b/internal/cli/run_test.go @@ -120,7 +120,7 @@ func TestRunFailOpenStartsWithAWarningEvenQuiet(t *testing.T) { func TestApplyEngineEntry(t *testing.T) { night := proto.Window{Start: 23 * 60, End: 7 * 60} noon := proto.Window{Start: 12 * 60, End: 13 * 60} - e := &config.Engine{Name: "comfyui", Priority: "batch", NeedsMiB: 8192, + e := &config.Engine{Name: "comfyui", Priority: "batch", NeedsMiB: config.MiBs{8192}, DrainPost: "http://127.0.0.1:1/free", 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}, Window: night} @@ -128,7 +128,8 @@ func TestApplyEngineEntry(t *testing.T) { var errBuf bytes.Buffer c := &Context{setFlags: map[string]bool{}, Err: &errBuf} prio, keep := "interactive", false - needs, dp, ds, dc := uint64(0), "", "", "" + var needs config.MiBs + dp, ds, dc := "", "", "" var dm, mt time.Duration win, cont, comp := noon, "", "" // The flags ride as pointers, so the local changes below reach the same @@ -137,10 +138,10 @@ func TestApplyEngineEntry(t *testing.T) { drainComfyUI: &dc, keepComfyUI: &keep, drainMax: &dm, minTurn: &mt, window: &win, container: &cont, compose: &comp} applyEngineEntry(c, e, f) - if e.Priority != "batch" || e.NeedsMiB != 8192 || e.DrainPost != "http://127.0.0.1:1/free" || + if e.Priority != "batch" || e.NeedsMiB.String() != "8192" || e.DrainPost != "http://127.0.0.1:1/free" || e.DrainSlots != "http://127.0.0.1:2" || e.DrainComfyUI != "http://127.0.0.1:3" || e.DrainMax.Duration != time.Minute || e.MinTurn.Duration != 5*time.Minute || e.Window != night { - t.Errorf("defaults from the entry = %q %d %q %q %q %s %s %s", e.Priority, e.NeedsMiB, e.DrainPost, e.DrainSlots, e.DrainComfyUI, e.DrainMax, e.MinTurn, e.Window) + t.Errorf("defaults from the entry = %q %s %q %q %q %s %s %s", e.Priority, e.NeedsMiB, e.DrainPost, e.DrainSlots, e.DrainComfyUI, e.DrainMax, e.MinTurn, e.Window) } if errBuf.String() != "" { t.Errorf("no flag set, but a notice was printed: %q", errBuf.String()) @@ -149,13 +150,13 @@ func TestApplyEngineEntry(t *testing.T) { errBuf.Reset() c = &Context{setFlags: map[string]bool{"priority": true, "needs-mib": true, "drain-post": true, "drain-slots": true, "drain-comfyui": true, "drain-max": true, "min-turn": true, "window": true}, Err: &errBuf} - prio, needs, dp, ds, dc = "foreground", 4096, "http://127.0.0.1:9", "http://127.0.0.1:9", "http://127.0.0.1:9" + prio, needs, dp, ds, dc = "foreground", config.MiBs{4096}, "http://127.0.0.1:9", "http://127.0.0.1:9", "http://127.0.0.1:9" dm = 2 * time.Minute mt = time.Minute applyEngineEntry(c, e, f) - if e.Priority != "foreground" || e.NeedsMiB != 4096 || e.DrainPost != "http://127.0.0.1:9" || + if e.Priority != "foreground" || e.NeedsMiB.String() != "4096" || e.DrainPost != "http://127.0.0.1:9" || e.DrainMax.Duration != 2*time.Minute || e.MinTurn.Duration != time.Minute || e.Window != noon { - t.Errorf("explicit flags lost to the entry: %q %d %q %s %s %s", e.Priority, e.NeedsMiB, e.DrainPost, e.DrainMax, e.MinTurn, e.Window) + t.Errorf("explicit flags lost to the entry: %q %s %q %s %s %s", e.Priority, e.NeedsMiB, e.DrainPost, e.DrainMax, e.MinTurn, e.Window) } for _, want := range []string{ "berth: --priority overrides berth.toml priority = batch for comfyui", @@ -409,3 +410,15 @@ func TestRunRejectsABadCard(t *testing.T) { t.Errorf("stderr = %q", errOut) } } + +// Sizes per card must match the cards --gpu names, before berthd is dialed. +func TestRunRejectsSizesThatDoNotMatchTheCards(t *testing.T) { + isolate(t) + code, _, errOut := runCLI(t, "run", "--as", "llama-70b", "--gpu", "gpu0", "--needs-mib", "18000,6000", "--", "true") + if code != ExitSystem { + t.Fatalf("exit = %d, want %d", code, ExitSystem) + } + if want := `needs_mib 18000,6000 has 2 sizes for 1 card(s) in gpu "gpu0"`; !strings.Contains(errOut, want) { + t.Errorf("stderr = %q, want it to hold %q", errOut, want) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 7303740..a971fee 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -56,9 +56,9 @@ type Engine struct { // model it serves and it is not running. Start is the engine's command; // it runs under a lease exactly as `berth run` would. Start []string `toml:"start,omitempty" json:"start,omitempty" yaml:"start,omitempty"` - Serves []string `toml:"serves,omitempty" json:"serves,omitempty" yaml:"serves,omitempty"` // model names this engine answers for - Priority string `toml:"priority,omitempty" json:"priority,omitempty" yaml:"priority,omitempty"` // batch, interactive (default) or foreground - NeedsMiB uint64 `toml:"needs_mib,omitzero" json:"needs_mib,omitempty" yaml:"needs_mib,omitempty"` // 0 = the whole card + Serves []string `toml:"serves,omitempty" json:"serves,omitempty" yaml:"serves,omitempty"` // model names this engine answers for + Priority string `toml:"priority,omitempty" json:"priority,omitempty" yaml:"priority,omitempty"` // batch, interactive (default) or foreground + NeedsMiB MiBs `toml:"needs_mib,omitempty" json:"needs_mib,omitempty" yaml:"needs_mib,omitempty"` // per card, in gpu order; none = the whole card // Background: the gateway starts this engine at batch priority whenever // the card has been free (no lease, no waiter, the card serving this // machine) for IdleAfter. Any other request drains it as any batch @@ -703,6 +703,9 @@ func (c *Config) Validate() error { if e.DrainMax.Duration < 0 { errs = append(errs, fmt.Errorf("%s: drain_max %s must be >= 0", where, e.DrainMax)) } + if err := CheckNeeds(e.NeedsMiB, e.GPU); err != nil { + errs = append(errs, fmt.Errorf("%s: %w", where, err)) + } if e.MinTurn.Duration < 0 { errs = append(errs, fmt.Errorf("%s: min_turn %s must be >= 0", where, e.MinTurn)) } diff --git a/internal/config/mibs.go b/internal/config/mibs.go new file mode 100644 index 0000000..bf7219d --- /dev/null +++ b/internal/config/mibs.go @@ -0,0 +1,162 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +package config + +import ( + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" + + "gopkg.in/yaml.v3" +) + +// MiBs is an engine's card memory in MiB, one value per card it runs on, in +// the order its gpu key names them: needs_mib = [18000, 6000] for an engine +// split unevenly across gpu0 and gpu1. One value is the size on every card, +// and none (or 0) is the whole card. It reads as a number or a list in TOML, +// JSON and YAML, and as "18000" or "18000,6000" on the command line; one +// value writes back as a plain number. +type MiBs []uint64 + +var errMiBs = errors.New("needs_mib is a number of MiB, or a list of them with one per card") + +// Bytes is the size on each card in bytes, nil for the whole card. +func (m MiBs) Bytes() []uint64 { + if m.whole() { + return nil + } + out := make([]uint64, len(m)) + for i, v := range m { + out[i] = v << 20 + } + return out +} + +// whole is true when no card is sized: no value, or the single value 0. +func (m MiBs) whole() bool { return len(m) == 0 || len(m) == 1 && m[0] == 0 } + +// String is the command-line spelling: "18000" or "18000,6000". +func (m MiBs) String() string { + parts := make([]string, len(m)) + for i, v := range m { + parts[i] = strconv.FormatUint(v, 10) + } + return strings.Join(parts, ",") +} + +// Set parses the command-line spelling, for flag.Var. +func (m *MiBs) Set(s string) error { + var out MiBs + for _, p := range strings.Split(s, ",") { + v, err := strconv.ParseUint(strings.TrimSpace(p), 10, 64) + if err != nil { + return fmt.Errorf("%q: %w", s, errMiBs) + } + out = append(out, v) + } + *m = out + return nil +} + +// UnmarshalTOML reads a TOML integer or array of integers. +func (m *MiBs) UnmarshalTOML(v any) error { + switch v := v.(type) { + case int64: + return m.set([]int64{v}) + case []any: + ns := make([]int64, len(v)) + for i, e := range v { + n, ok := e.(int64) + if !ok { + return errMiBs + } + ns[i] = n + } + return m.set(ns) + } + return errMiBs +} + +func (m *MiBs) set(ns []int64) error { + if len(ns) == 0 { + return errMiBs + } + out := make(MiBs, len(ns)) + for i, n := range ns { + if n < 0 { + return errMiBs + } + out[i] = uint64(n) + } + *m = out + return nil +} + +func (m MiBs) MarshalTOML() ([]byte, error) { + if len(m) == 1 { + return []byte(strconv.FormatUint(m[0], 10)), nil + } + return []byte("[" + strings.ReplaceAll(m.String(), ",", ", ") + "]"), nil +} + +func (m MiBs) MarshalJSON() ([]byte, error) { + if len(m) == 1 { + return json.Marshal(m[0]) + } + return json.Marshal([]uint64(m)) +} + +func (m *MiBs) UnmarshalJSON(b []byte) error { + var one uint64 + if json.Unmarshal(b, &one) == nil { + *m = MiBs{one} + return nil + } + var each []uint64 + if err := json.Unmarshal(b, &each); err != nil || len(each) == 0 { + return errMiBs + } + *m = each + return nil +} + +func (m MiBs) MarshalYAML() (any, error) { + if len(m) == 1 { + return m[0], nil + } + return []uint64(m), nil +} + +func (m *MiBs) UnmarshalYAML(node *yaml.Node) error { + var one uint64 + if node.Kind == yaml.ScalarNode && node.Decode(&one) == nil { + *m = MiBs{one} + return nil + } + var each []uint64 + if node.Kind != yaml.SequenceNode || node.Decode(&each) != nil || len(each) == 0 { + return errMiBs + } + *m = each + return nil +} + +// CheckNeeds is the rule for needs_mib and berth run --needs-mib: a list of +// sizes has one size per card that gpu names. +func CheckNeeds(m MiBs, gpu string) error { + if len(m) < 2 { + return nil + } + cards := 0 + if gpu != "" { + cards = len(strings.Split(gpu, ",")) + } + if cards != len(m) { + return fmt.Errorf("needs_mib %s has %d sizes for %d card(s) in gpu %q: give one size, or one per card in the same order", + m, len(m), cards, gpu) + } + return nil +} diff --git a/internal/config/mibs_test.go b/internal/config/mibs_test.go new file mode 100644 index 0000000..2dfeb48 --- /dev/null +++ b/internal/config/mibs_test.go @@ -0,0 +1,97 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +package config + +import ( + "bytes" + "encoding/json" + "slices" + "strings" + "testing" + + "github.com/BurntSushi/toml" + "gopkg.in/yaml.v3" +) + +// needs_mib on an entry is one size for every card, or a list with one size +// per card in gpu's order, for a model split across cards. +func TestEngineNeedsPerCard(t *testing.T) { + const entry = "[[engine]]\nname = \"llama-70b\"\ntype = \"custom\"\nurl = \"http://127.0.0.1:8080\"\nhealth = \"/health\"\n" + for _, tc := range []struct { + body string + want []uint64 + }{ + {"needs_mib = 8192\n", []uint64{8192 << 20}}, + {"gpu = \"gpu0,gpu1\"\nneeds_mib = 8192\n", []uint64{8192 << 20}}, + {"gpu = \"gpu0,gpu1\"\nneeds_mib = [18000, 6000]\n", []uint64{18000 << 20, 6000 << 20}}, + {"", nil}, + {"needs_mib = 0\n", nil}, + } { + c, err := Load(write(t, entry+tc.body)) + if err != nil { + t.Fatalf("%q: %v", tc.body, err) + } + if got := c.Engine[0].NeedsMiB.Bytes(); !slices.Equal(got, tc.want) { + t.Errorf("%q: bytes = %v, want %v", tc.body, got, tc.want) + } + } + for _, bad := range []struct{ body, wantErr string }{ + {"gpu = \"gpu0,gpu1\"\nneeds_mib = [18000, 6000, 2000]\n", `needs_mib 18000,6000,2000 has 3 sizes for 2 card(s) in gpu "gpu0,gpu1"`}, + {"needs_mib = [18000, 6000]\n", `needs_mib 18000,6000 has 2 sizes for 0 card(s) in gpu ""`}, + {"needs_mib = -1\n", "needs_mib is a number of MiB, or a list of them with one per card"}, + {"needs_mib = \"8G\"\n", "needs_mib is a number of MiB, or a list of them with one per card"}, + } { + if _, err := Load(write(t, entry+bad.body)); err == nil || !strings.Contains(err.Error(), bad.wantErr) { + t.Errorf("%q: error %v, want it to say %q", bad.body, err, bad.wantErr) + } + } +} + +// One size writes back as the plain number it always was, and a list as a +// list, in every encoding the config has. +func TestNeedsKeepTheirSpelling(t *testing.T) { + for _, tc := range []struct { + m MiBs + toml, json, yaml string + }{ + {MiBs{8192}, "needs_mib = 8192", `"needs_mib":8192`, "needs_mib: 8192"}, + {MiBs{18000, 6000}, "needs_mib = [18000, 6000]", `"needs_mib":[18000,6000]`, "needs_mib:\n - 18000\n - 6000"}, + } { + e := Engine{Name: "llama-70b", NeedsMiB: tc.m} + var tb bytes.Buffer + if err := toml.NewEncoder(&tb).Encode(e); err != nil { + t.Fatal(err) + } + jb, _ := json.Marshal(e) + yb, _ := yaml.Marshal(e) + for _, got := range []struct{ enc, out, want string }{{"toml", tb.String(), tc.toml}, {"json", string(jb), tc.json}, {"yaml", string(yb), tc.yaml}} { + if !strings.Contains(got.out, got.want) { + t.Errorf("%s of %v = %s, want it to hold %s", got.enc, tc.m, got.out, got.want) + } + } + var back Engine + if err := json.Unmarshal(jb, &back); err != nil || !slices.Equal(back.NeedsMiB, tc.m) { + t.Errorf("json back = %v (%v), want %v", back.NeedsMiB, err, tc.m) + } + if err := yaml.Unmarshal(yb, &back); err != nil || !slices.Equal(back.NeedsMiB, tc.m) { + t.Errorf("yaml back = %v (%v), want %v", back.NeedsMiB, err, tc.m) + } + } +} + +// berth run --needs-mib reads the same spellings. +func TestNeedsFlag(t *testing.T) { + for _, tc := range []struct { + in string + want MiBs + wantErr bool + }{{"8192", MiBs{8192}, false}, {"18000,6000", MiBs{18000, 6000}, false}, {"18000, 6000", MiBs{18000, 6000}, false}, + {"8G", nil, true}, {"18000,", nil, true}} { + var m MiBs + err := m.Set(tc.in) + if (err != nil) != tc.wantErr || !tc.wantErr && !slices.Equal(m, tc.want) { + t.Errorf("Set(%q) = %v, %v; want %v, error %v", tc.in, m, err, tc.want, tc.wantErr) + } + } +} diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index 1f8f58a..00f4a06 100644 --- a/internal/gateway/gateway_test.go +++ b/internal/gateway/gateway_test.go @@ -92,7 +92,8 @@ func testRun(argv []string) { as := fs.String("as", "", "") group := fs.String("group", "", "") prio := fs.String("priority", "", "") - needs := fs.Uint64("needs-mib", 0, "") + var needs config.MiBs + fs.Var(&needs, "needs-mib", "") wait := fs.Duration("wait", 0, "") grace := fs.Duration("grace", 0, "") card := fs.String("gpu", "", "") @@ -113,7 +114,7 @@ func testRun(argv []string) { os.Exit(2) } os.Exit(runner.Run(runner.Options{ - Socket: *socket, Engine: *as, Group: *group, Priority: *prio, Needs: *needs << 20, Card: *card, + Socket: *socket, Engine: *as, Group: *group, Priority: *prio, Needs: needs.Bytes(), Card: *card, Switched: *switched, Wait: *wait, Grace: *grace, Argv: fs.Args(), Stderr: os.Stderr, })) } diff --git a/internal/gateway/resume_test.go b/internal/gateway/resume_test.go index dc29fbe..8ce37e0 100644 --- a/internal/gateway/resume_test.go +++ b/internal/gateway/resume_test.go @@ -155,7 +155,7 @@ func TestGatewayResumesAParkedEngine(t *testing.T) { t.Fatal("the vllm adapter is missing") } engineLog := &syncBuf{} - go runner.Run(runner.Options{Socket: sock, Engine: "vllm", Priority: "interactive", Needs: 8 << 30, + go runner.Run(runner.Options{Socket: sock, Engine: "vllm", Priority: "interactive", Needs: []uint64{8 << 30}, Park: &adapter, ParkBase: srv.URL, ParkProbe: "/v1/models", Grace: 200 * time.Millisecond, Argv: []string{"sleep", "30"}, Stderr: engineLog}) diff --git a/internal/gateway/session_test.go b/internal/gateway/session_test.go index 3a829d7..55e5201 100644 --- a/internal/gateway/session_test.go +++ b/internal/gateway/session_test.go @@ -52,3 +52,18 @@ func TestRunArgvCarriesSlice(t *testing.T) { t.Fatalf("an entry with no slice gained one: %q", bare) } } + +// A gateway start of an entry split across cards hands berth run its size +// on each card, in gpu's order. +func TestRunArgvCarriesEachCardsSize(t *testing.T) { + for _, tc := range []struct { + needs config.MiBs + want string + }{{config.MiBs{18000, 6000}, "18000,6000"}, {config.MiBs{8192}, "8192"}} { + args := runArgv(runner.OptionsFor(&config.Engine{Name: "llama-70b", GPU: "gpu0,gpu1", NeedsMiB: tc.needs}, runner.Over{})) + i := slices.Index(args, "--needs-mib") + if i < 0 || i+1 >= len(args) || args[i+1] != tc.want { + t.Errorf("needs %v: args %q, want --needs-mib %s", tc.needs, args, tc.want) + } + } +} diff --git a/internal/gateway/spawn.go b/internal/gateway/spawn.go index ad86307..d788397 100644 --- a/internal/gateway/spawn.go +++ b/internal/gateway/spawn.go @@ -14,10 +14,10 @@ import ( "io" "os" "os/exec" - "strconv" "strings" "time" + "github.com/T-Crypt/berth/internal/config" "github.com/T-Crypt/berth/internal/procs" "github.com/T-Crypt/berth/internal/runner" "github.com/T-Crypt/berth/internal/sysd" @@ -96,8 +96,12 @@ func runArgv(o runner.Options) []string { if o.Priority != "" { args = append(args, "--priority", o.Priority) } - if o.Needs != 0 { - args = append(args, "--needs-mib", strconv.FormatUint(o.Needs>>20, 10)) + if len(o.Needs) > 0 { + mibs := make(config.MiBs, len(o.Needs)) + for i, b := range o.Needs { + mibs[i] = b >> 20 + } + args = append(args, "--needs-mib", mibs.String()) } if o.Wait != 0 { args = append(args, "--wait", o.Wait.String()) diff --git a/internal/node/node_test.go b/internal/node/node_test.go index 7c4d26e..d3b0684 100644 --- a/internal/node/node_test.go +++ b/internal/node/node_test.go @@ -229,7 +229,7 @@ func TestStopNamesALeaseOnSeveralCardsOnce(t *testing.T) { if err != nil { t.Fatal(err) } - ls, err := broker.AcquireAll(ctx, []*broker.Broker{cs[0].B, cs[1].B}, broker.Request{Engine: "maya", Priority: broker.Foreground}, nil) + ls, err := broker.AcquireAll(ctx, []*broker.Broker{cs[0].B, cs[1].B}, nil, broker.Request{Engine: "maya", Priority: broker.Foreground}, nil) if err != nil { t.Fatal(err) } diff --git a/internal/proto/proto.go b/internal/proto/proto.go index a05bce2..fa91943 100644 --- a/internal/proto/proto.go +++ b/internal/proto/proto.go @@ -110,6 +110,11 @@ type Msg struct { Group string `json:"group,omitempty"` // the engine family the lease counts for in juggles and usage Priority string `json:"priority,omitempty"` Needs uint64 `json:"needs,omitempty"` // bytes; 0 = the whole card + // NeedsEach is the size in bytes on each card Card names, in that + // order, when the cards differ. Needs is then the largest, so a berthd + // that predates NeedsEach reserves too much on a small card, never too + // little. + NeedsEach []uint64 `json:"needs_each,omitempty"` // Card is the card an OpAcquire or OpReattach is for: its name on the // node (gpu1) or its PCI address, or a comma list of them for one lease // on several cards. "" is the node's first card. diff --git a/internal/runner/auth_test.go b/internal/runner/auth_test.go index e5ef028..093bc3d 100644 --- a/internal/runner/auth_test.go +++ b/internal/runner/auth_test.go @@ -53,7 +53,7 @@ func TestRunAuthenticatesDrainRequests(t *testing.T) { defer srv.Close() card := &parkCard{total: 24 << 30} sock, b := parkDaemon(t, card, 16<<30) - e := config.Engine{Name: "fixture", Priority: "interactive", NeedsMiB: 8192, APIKeyEnv: "FIXTURE_DRAIN_KEY"} + e := config.Engine{Name: "fixture", Priority: "interactive", NeedsMiB: config.MiBs{8192}, APIKeyEnv: "FIXTURE_DRAIN_KEY"} switch kind { case "file": e.APIKeyEnv = "" diff --git a/internal/runner/card_test.go b/internal/runner/card_test.go index d5ec914..205afae 100644 --- a/internal/runner/card_test.go +++ b/internal/runner/card_test.go @@ -6,11 +6,14 @@ package runner import ( + "bytes" "context" "io" "log/slog" "os" "path/filepath" + "slices" + "strings" "testing" "time" @@ -180,3 +183,56 @@ func TestReattachNamesTheCard(t *testing.T) { t.Fatalf("reattach msg = op %s card %q window %q, want reattach on gpu1 with its window", m.Op, m.Card, m.Window) } } + +// A run split unevenly across two cards leases each card's own size, in the +// order --gpu names them, through the socket and berthd; sizes that do not +// match the cards are refused. +func TestRunSizesEachCard(t *testing.T) { + sock, n := twoCardDaemon(t) + done := make(chan int, 1) + go func() { + done <- Run(Options{Socket: sock, Engine: "llama-70b", Priority: "interactive", Card: "gpu1,gpu0", + Needs: []uint64{6 << 30, 18 << 30}, Grace: time.Second, Argv: []string{"sleep", "2"}, Stderr: io.Discard}) + }() + want := map[string]uint64{"gpu0": 18 << 30, "gpu1": 6 << 30} + got := map[string]uint64{} + for deadline := time.Now().Add(3 * time.Second); len(got) < 2 && time.Now().Before(deadline); time.Sleep(20 * time.Millisecond) { + for name := range want { + c, err := n.Card(name) + if err != nil { + t.Fatal(err) + } + for _, l := range c.B.Status().Leases { + if l.Engine == "llama-70b" { + got[name] = l.Needs + } + } + } + } + if got["gpu0"] != want["gpu0"] || got["gpu1"] != want["gpu1"] { + t.Errorf("shares = %v, want %v", got, want) + } + if code := <-done; code != 0 { + t.Fatalf("exit %d", code) + } + + var errBuf bytes.Buffer + code := Run(Options{Socket: sock, Engine: "llama-70b", Priority: "interactive", Card: "gpu0", + Needs: []uint64{6 << 30, 18 << 30}, Grace: time.Second, Argv: []string{"true"}, Stderr: &errBuf}) + if code != ExitRefused || !strings.Contains(errBuf.String(), `2 sizes for 1 card(s) in "gpu0"`) { + t.Errorf("exit %d, output %q; want a refusal naming 2 sizes for 1 card", code, errBuf.String()) + } +} + +// A re-attach carries the per-card sizes too, with the largest in Needs for +// a berthd that reads only that. +func TestReattachCarriesEachCardsSize(t *testing.T) { + m := reattachMsg(Options{Engine: "llama-70b", Card: "gpu0,gpu1", Needs: []uint64{18 << 30, 6 << 30}}, 100, 7, &containerRef{}, false, newLine(nil)) + if m.Needs != 18<<30 || !slices.Equal(m.NeedsEach, []uint64{18 << 30, 6 << 30}) { + t.Errorf("needs %d, each %v; want 18 GiB and [18 GiB 6 GiB]", m.Needs, m.NeedsEach) + } + m = reattachMsg(Options{Engine: "llama-70b", Card: "gpu0,gpu1", Needs: []uint64{8 << 30}}, 100, 7, &containerRef{}, false, newLine(nil)) + if m.Needs != 8<<30 || m.NeedsEach != nil { + t.Errorf("one size: needs %d, each %v; want 8 GiB and none", m.Needs, m.NeedsEach) + } +} diff --git a/internal/runner/front_test.go b/internal/runner/front_test.go index 2bb7de0..a20a649 100644 --- a/internal/runner/front_test.go +++ b/internal/runner/front_test.go @@ -91,7 +91,7 @@ func TestCheckpointParkHoldsRequestsUntilTheRestore(t *testing.T) { log := &syncBuf{} exited := make(chan int, 1) go func() { - exited <- Run(Options{Socket: sock, Engine: "ninfer", Priority: "interactive", Needs: 8 << 30, + exited <- Run(Options{Socket: sock, Engine: "ninfer", Priority: "interactive", Needs: []uint64{8 << 30}, DrainPost: swap.URL + "/api/models/unload/ninfer", Front: addr, Checkpoint: true, ParkProbe: "/v1/models", Grace: 2 * time.Second, Argv: []string{os.Args[0], "--port", PortPlaceholder}, Stderr: log}) @@ -184,7 +184,7 @@ func TestCheckpointWithoutTheToolDrainsAndSaysWhy(t *testing.T) { exited := make(chan int, 1) log := &syncBuf{} go func() { - exited <- Run(Options{Socket: sock, Engine: "ninfer", Priority: "interactive", Needs: 8 << 30, + exited <- Run(Options{Socket: sock, Engine: "ninfer", Priority: "interactive", Needs: []uint64{8 << 30}, Front: freeFront(t), Checkpoint: true, ParkProbe: "/v1/models", Grace: time.Second, Argv: []string{os.Args[0], "--port", PortPlaceholder}, Stderr: log}) }() @@ -288,7 +288,7 @@ func TestStopDuringAParkStopsTheEngine(t *testing.T) { log := &syncBuf{} exited := make(chan int, 1) go func() { - exited <- Run(Options{Socket: sock, Engine: "ninfer", Priority: "interactive", Needs: 8 << 30, + exited <- Run(Options{Socket: sock, Engine: "ninfer", Priority: "interactive", Needs: []uint64{8 << 30}, Front: freeFront(t), Checkpoint: true, ParkProbe: "/v1/models", Grace: 2 * time.Second, Argv: []string{os.Args[0], "--port", PortPlaceholder}, Stderr: log}) }() diff --git a/internal/runner/lend_test.go b/internal/runner/lend_test.go index 1baffaa..4f115c3 100644 --- a/internal/runner/lend_test.go +++ b/internal/runner/lend_test.go @@ -314,7 +314,7 @@ func TestLendAskOnARunWithoutLendCallsIsRefused(t *testing.T) { // has lent nothing. func TestReattachNamesTheOpenLend(t *testing.T) { o := lendEntry("http://127.0.0.1:1") - o.Needs = 16 << 30 + o.Needs = []uint64{16 << 30} line := newLine(nil) line.record(lend.Ticket{ID: "17", GPU: "0000:09:00.0", MiB: 4096}) diff --git a/internal/runner/options.go b/internal/runner/options.go index ad96b99..1070c64 100644 --- a/internal/runner/options.go +++ b/internal/runner/options.go @@ -6,6 +6,7 @@ package runner import ( "io" "log/slog" + "slices" "time" "github.com/T-Crypt/berth/internal/config" @@ -60,7 +61,7 @@ func OptionsFor(e *config.Engine, o Over) Options { } opts := Options{ Socket: o.Socket, Engine: e.Name, Card: e.GPU, Group: o.Group, Priority: prio, - Needs: e.NeedsMiB << 20, Pinned: o.Pinned, Switched: o.Switched, Wait: o.Wait, + Needs: e.NeedsMiB.Bytes(), Pinned: o.Pinned, Switched: o.Switched, Wait: o.Wait, DrainPost: e.DrainPost, DrainSlots: e.DrainSlots, DrainComfyUI: e.DrainComfyUI, DrainMax: e.DrainMax.Duration, MinTurn: e.MinTurn.Duration, ParkTTL: e.ParkTTL.Duration, Window: e.Window, Container: e.Container, Compose: e.Compose, APIKeyEnv: e.APIKeyEnv, APIKeyFile: e.APIKeyFile, @@ -122,3 +123,22 @@ func OptionsFor(e *config.Engine, o Over) Options { } return opts } + +// largest is the size a berthd that reads only Needs gets: the biggest card +// share, so an older daemon over-reserves the smaller cards rather than +// under-reserving the biggest. +func largest(needs []uint64) uint64 { + if len(needs) == 0 { + return 0 + } + return slices.Max(needs) +} + +// each is NeedsEach on the wire: the per-card sizes when there is more than +// one, nil when a single size covers every card. +func each(needs []uint64) []uint64 { + if len(needs) < 2 { + return nil + } + return needs +} diff --git a/internal/runner/options_test.go b/internal/runner/options_test.go index b6bb46b..1233ab3 100644 --- a/internal/runner/options_test.go +++ b/internal/runner/options_test.go @@ -16,7 +16,7 @@ import ( // OptionsFor carries the engine entry's settings, fills the defaults it // leaves open, and defers every zero override to the entry. func TestOptionsFor(t *testing.T) { - e := &config.Engine{Name: "strata", Type: "strata", Priority: "foreground", NeedsMiB: 8192, + e := &config.Engine{Name: "strata", Type: "strata", Priority: "foreground", NeedsMiB: config.MiBs{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}, @@ -25,7 +25,7 @@ func TestOptionsFor(t *testing.T) { 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, + Socket: "s", Engine: "strata", Group: "strata", Priority: "foreground", Needs: []uint64{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, ParkTTL: 10 * time.Minute, Argv: []string{"true"}, WholeMachine: true, Window: e.Window, diff --git a/internal/runner/park_auth_test.go b/internal/runner/park_auth_test.go index 7fd8fbc..e08803c 100644 --- a/internal/runner/park_auth_test.go +++ b/internal/runner/park_auth_test.go @@ -48,7 +48,7 @@ func TestRunAuthenticatesParkAndResume(t *testing.T) { defer srv.Close() card := &parkCard{total: 24 << 30} sock, b := parkDaemon(t, card, 24<<30) - e := config.Engine{Name: "fixture", Type: kind, URL: srv.URL, Park: "hook", NeedsMiB: 8192, APIKeyEnv: "FIXTURE_PARK_KEY"} + e := config.Engine{Name: "fixture", Type: kind, URL: srv.URL, Park: "hook", NeedsMiB: config.MiBs{8192}, APIKeyEnv: "FIXTURE_PARK_KEY"} var logs syncBuf exited := make(chan int, 1) o := OptionsFor(&e, Over{Socket: sock, Grace: 20 * time.Millisecond, Argv: []string{"sleep", "30"}, Stderr: &logs}) diff --git a/internal/runner/park_test.go b/internal/runner/park_test.go index 1dea18d..a891bb1 100644 --- a/internal/runner/park_test.go +++ b/internal/runner/park_test.go @@ -222,7 +222,7 @@ func runPark(t *testing.T, sock, engine, stubURL string, hook bool) (chan int, * }) exited := make(chan int, 1) go func() { - exited <- Run(Options{Socket: sock, Engine: engine, Priority: "interactive", Needs: 8 << 30, + exited <- Run(Options{Socket: sock, Engine: engine, Priority: "interactive", Needs: []uint64{8 << 30}, Park: &adapter, ParkBase: stubURL, ParkProbe: "/v1/models", Hook: hook, Grace: 200 * time.Millisecond, Argv: []string{"sleep", "30"}, Stderr: log}) }() @@ -426,7 +426,7 @@ func TestRunParksThroughBerthProtocol(t *testing.T) { }) exited := make(chan int, 1) go func() { - exited <- Run(Options{Socket: sock, Engine: "custom", Priority: "interactive", Needs: 8 << 30, + exited <- Run(Options{Socket: sock, Engine: "custom", Priority: "interactive", Needs: []uint64{8 << 30}, Hook: true, ParkBase: srv.URL, ParkProbe: "/v1/models", Grace: 200 * time.Millisecond, Argv: []string{"sleep", "30"}, Stderr: log}) }() @@ -473,7 +473,7 @@ func TestRunParksComfyUIKeepUp(t *testing.T) { }) exited := make(chan int, 1) go func() { - exited <- Run(Options{Socket: sock, Engine: "comfyui", Priority: "interactive", Needs: 8 << 30, + exited <- Run(Options{Socket: sock, Engine: "comfyui", Priority: "interactive", Needs: []uint64{8 << 30}, DrainComfyUI: srv.URL, Park: &adapter, ParkBase: srv.URL, ParkProbe: "/system_stats", Grace: 200 * time.Millisecond, Argv: []string{"sleep", "30"}, Stderr: log}) }() @@ -549,7 +549,7 @@ func TestRunHookProbeNegative(t *testing.T) { }) exited := make(chan int, 1) go func() { - exited <- Run(Options{Socket: sock, Engine: "custom", Priority: "interactive", Needs: 8 << 30, + exited <- Run(Options{Socket: sock, Engine: "custom", Priority: "interactive", Needs: []uint64{8 << 30}, Hook: true, ParkBase: srv.URL, ParkProbe: "/v1/models", Grace: 200 * time.Millisecond, Argv: []string{"sleep", "30"}, Stderr: log}) }() diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 38e1d42..888f9b7 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -48,7 +48,9 @@ type Options struct { Card string Group string // the engine family for juggles and usage (llama-swap's models: "llama-swap") Priority string - Needs uint64 + // Needs is the card memory in bytes, one value per card in Card's + // order; one value sizes every card, none takes the whole card. + Needs []uint64 Pinned bool // never drained: an app that must not be stopped (Blender, a capture) Switched bool // a person's switch (the wake page's Launch, /v1/berth/switch): wins a juggle // MinTurn is a minimum turn: while the lease is younger than it after @@ -224,7 +226,7 @@ func Run(o Options) int { } defer raw.Close() - if err := conn.Send(proto.Msg{Op: proto.OpAcquire, Card: o.Card, Engine: o.Engine, Group: o.Group, Priority: o.Priority, Needs: o.Needs, Wait: o.Wait.Seconds(), Pinned: o.Pinned, Switched: o.Switched, Lend: o.Lend != nil, WholeMachine: o.WholeMachine, MinTurn: o.MinTurn.Seconds(), ParkTTL: o.ParkTTL.Seconds(), Window: o.Window.String()}); err != nil { + if err := conn.Send(proto.Msg{Op: proto.OpAcquire, Card: o.Card, Engine: o.Engine, Group: o.Group, Priority: o.Priority, Needs: largest(o.Needs), NeedsEach: each(o.Needs), Wait: o.Wait.Seconds(), Pinned: o.Pinned, Switched: o.Switched, Lend: o.Lend != nil, WholeMachine: o.WholeMachine, MinTurn: o.MinTurn.Seconds(), ParkTTL: o.ParkTTL.Seconds(), Window: o.Window.String()}); err != nil { say("acquire: %v", err) return ExitRefused } @@ -466,7 +468,7 @@ func reattach(o Options, cmd *exec.Cmd, done <-chan struct{}, cid *containerRef, // waiters behind memory that is already gone. func reattachMsg(o Options, pid int, start uint64, cid *containerRef, parked bool, line *leaseLine) proto.Msg { m := proto.Msg{Op: proto.OpReattach, Card: o.Card, Engine: o.Engine, Group: o.Group, Priority: o.Priority, - Needs: o.Needs, Pinned: o.Pinned, Lend: o.Lend != nil, WholeMachine: o.WholeMachine, + Needs: largest(o.Needs), NeedsEach: each(o.Needs), Pinned: o.Pinned, Lend: o.Lend != nil, WholeMachine: o.WholeMachine, MinTurn: o.MinTurn.Seconds(), ParkTTL: o.ParkTTL.Seconds(), Window: o.Window.String(), PID: pid, PIDStart: start, Containers: cid.list(), Parked: parked} if t, ok := line.openLend(); ok { diff --git a/internal/server/server.go b/internal/server/server.go index dc2ccb1..cb8625e 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -268,6 +268,13 @@ func hold(ctx context.Context, conn *proto.Conn, req proto.Msg, n *node.Node, fi Window: win, Containers: req.ContainerIDs(), } + if len(req.NeedsEach) > 0 && len(req.NeedsEach) != len(cards) { + reply(conn, log, proto.Msg{Op: proto.OpRefused, Reason: fmt.Sprintf("%d sizes for %d card(s) in %q: give one size, or one per card", len(req.NeedsEach), len(cards), req.Card)}) + return + } + if len(req.NeedsEach) == 1 { + r.Needs = req.NeedsEach[0] + } if len(cards) > 1 { holdAll(ctx, u, req, r, cards, incoming, cancel, log) return @@ -386,9 +393,9 @@ func holdAll(ctx context.Context, u *consumer, req proto.Msg, r broker.Request, var leases []*broker.Lease var err error if req.Op == proto.OpReattach { - leases, err = broker.ReattachAll(bs, r, req.PID, req.PIDStart) + leases, err = broker.ReattachAll(bs, req.NeedsEach, r, req.PID, req.PIDStart) } else { - leases, err = broker.AcquireAll(ctx, bs, r, func(reason string) { + leases, err = broker.AcquireAll(ctx, bs, req.NeedsEach, r, func(reason string) { if !u.send(proto.Msg{Op: proto.OpWaiting, Reason: reason}) { cancel() } From 085964cf0c612c69e9c1e8dd0272d930af205285 Mon Sep 17 00:00:00 2001 From: T-Crypt <132922589+T-Crypt@users.noreply.github.com> Date: Mon, 2 Nov 2026 08:30:49 -0700 Subject: [PATCH 2/2] docs: split one model across cards with llama.cpp --tensor-split --- CHANGELOG.md | 6 ++++++ README.md | 2 +- docs/leasing.md | 30 ++++++++++++++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 71bb67c..e9bd886 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ has a matching GitHub release whose body comes from its section below. ### Added +- `needs_mib = [18000, 6000]` on an `[[engine]]` entry with `gpu = "gpu0,gpu1"`, and `berth run --needs-mib 18000,6000`: + a size per card, in the order the cards are named, for a model split unevenly across cards (llama.cpp + `--tensor-split`). Each card's share of the lease reserves its own size, so the room left on the lighter card can + serve another engine. One value still sizes every card. A berthd from before this release reads only the largest + size and reserves it on every card. docs/leasing.md has a new section on splitting one model across cards. + - `window = "23:00-07:00"` on an `[[engine]]` entry, and `berth run --window 23:00-07:00`: a lease window, for work that should only run in cheap hours. berthd grants the lease inside them and drains its holder when they close, at the same boundary as any other drain. A run that asks outside its window waits for it with no deadline unless `--wait` sets one, holds up no waiter behind it, and `berth diff --git a/README.md b/README.md index 426dbfb..f25416c 100644 --- a/README.md +++ b/README.md @@ -216,7 +216,7 @@ berth run --as maya --gpu gpu1 -- llama-server -m maya.gguf --port 8091 ``` An engine that spans several cards names them all: `berth run --gpu gpu0,gpu1`, or `gpu = "gpu0,gpu1"`. It holds -one lease across those cards, with `--needs-mib` on each. The lease is granted on every card at once, and only when each one fits, so no card +one lease across those cards, with `--needs-mib` on each (one size for all, or one per card: `--needs-mib 18000,6000`). The lease is granted on every card at once, and only when each one fits, so no card sits reserved while another is awaited. `CUDA_VISIBLE_DEVICES` lists each card's UUID in the order named. A drain on any of the cards stops the engine, and its exit frees every card. A lease on several cards cannot park or lend yet: a drain stops it the usual way. diff --git a/docs/leasing.md b/docs/leasing.md index 189c3ce..c83bd9e 100644 --- a/docs/leasing.md +++ b/docs/leasing.md @@ -176,6 +176,36 @@ A timer can start a one-off batch job at the window's open instead: a systemd us Task Scheduler task triggered daily at 23:00. The window is what keeps the job out of the expensive hours when it overruns or the machine wakes late. +## Split one model across cards + +llama.cpp splits a model that is too big for one card across several with `--tensor-split` (`-ts`). Run it under one +lease on those cards: + +```sh +berth run --as llama-70b --priority foreground --gpu gpu0,gpu1 \ + -- llama-server -m llama-70b-q4.gguf -ngl 99 --tensor-split 3,1 --port 8080 +``` + +- berthd grants both cards at once, or neither, so the model never loads onto one card while it waits for the other. +- The engine starts with `CUDA_VISIBLE_DEVICES` set to the cards' UUIDs in the order `--gpu` names them. The first + `--tensor-split` value goes to the first card named, the second to the second. +- A request that drains either card stops the engine, and its exit frees both. A lease on several cards cannot park + yet, so the next start loads the model from disk again. +- Without `--needs-mib` the lease takes the whole of each card. That is the safe choice when nothing else runs there. +- To leave room beside an uneven split, give one size per card, in the same order: + `--needs-mib 18000,6000`, or on the engine entry: + + ```toml + [[engine]] + name = "llama-70b" + gpu = "gpu0,gpu1" + needs_mib = [18000, 6000] # MiB on gpu0, then gpu1 + ``` + + A size is what the engine holds on that card at its peak, KV cache and compute buffers included, not the share of + the weights. Start it once with the whole cards, read each card's HOLDS in `berth status`, and add some headroom. + One value, as before, is the size on every card. + ## Training runs A training run that loses its cards loses everything since its last checkpoint. Run it under one pinned lease: