feat: a size per card for a model split across cards

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

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.

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:

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
+}

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")
+	}
+}

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])

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/<model>)")
@@ -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
 	}

internal/cli/run_test.go

@@ -120,15 +120,16 @@ 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}
 
 	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)
+	}
+}

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))
 		}