llmgw/internal/control/reconciler/preempt_test.go

package reconciler

import (
	"log/slog"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)

// newFleetHarness is newHarness with a catalog of several models.
func newFleetHarness(t *testing.T, models []catalog.Model, nodes ...inventory.Node) *harness {
	t.Helper()
	cat, err := catalog.New(models)
	if err != nil {
		t.Fatal(err)
	}
	inv := inventory.NewStore()
	c := newClock()
	for _, n := range nodes {
		inv.Upsert(n)
	}
	disp := &fakeDispatcher{}
	cfg := DefaultConfig()
	cfg.UnreachableAfter = 15 * time.Second
	cfg.LostAfter = 60 * time.Second
	r := New(cfg, catalog.NewLive(cat), inv, disp, slog.New(slog.DiscardHandler))
	r.SetClock(c.now)
	c.advance(cfg.SettleFor)
	for _, n := range nodes {
		touchAt(inv, n.ID, c.now())
	}
	return &harness{r: r, inv: inv, disp: disp, clock: c}
}

// gpuModel runs on one whole 8-GPU node.
func gpuModel(name, class string, min int) catalog.Model {
	return catalog.Model{Name: name, Source: "/models/" + name, Class: class,
		Scaling:  catalog.Scaling{Min: min, Max: min},
		Profiles: []catalog.Profile{{Name: "tp8", Engine: "vllm", GPUClass: "h100-80g", TP: 8}}}
}

// settle runs passes until nothing changes, marking what starts ready and
// letting preempted replicas run out their grace.
func (h *harness) settle(t *testing.T, passes int) {
	t.Helper()
	for range passes {
		h.r.Reconcile(t.Context())
		h.markReady()
		h.advance(h.r.cfg.PreemptGrace + time.Second)
	}
}

func replicasOf(h *harness, model string) (live int) {
	for _, rep := range h.r.Replicas() {
		if rep.Model == model && rep.State.Counts() {
			live++
		}
	}
	return live
}

// fill places the standard model, then ages it past MinLifetime so it may be
// taken from, and only then adds the model that wants its node.
func fillThenWant(t *testing.T, want catalog.Model, have catalog.Model, nodes ...inventory.Node) *harness {
	t.Helper()
	h := newFleetHarness(t, []catalog.Model{have, withMin(want, 0)}, nodes...)
	h.settle(t, 3)
	if got := replicasOf(h, have.Name); got != have.Scaling.Min {
		t.Fatalf("%s has %d replicas before anything wanted its room", have.Name, got)
	}
	h.advance(h.r.cfg.MinLifetime + time.Minute)
	cat, err := catalog.New([]catalog.Model{have, want})
	if err != nil {
		t.Fatal(err)
	}
	h.r.cat.Store(cat)
	return h
}

func withMin(m catalog.Model, n int) catalog.Model {
	m.Scaling = catalog.Scaling{Min: n, Max: n}
	return m
}

func TestACriticalModelTakesAStandardModelsNode(t *testing.T) {
	h := fillThenWant(t, gpuModel("vital", catalog.ClassCritical, 1), gpuModel("chat", "", 1), gpuNode("1", 8))

	h.r.Reconcile(t.Context())
	var draining int
	for _, rep := range h.r.Replicas() {
		if rep.Model == "chat" && rep.State == Draining && rep.PreemptedFor == "vital" {
			draining++
		}
	}
	if draining != 1 {
		t.Fatalf("the standard replica was not drained for the critical one: %+v", h.r.Replicas())
	}
	if hs := health(h, "chat"); hs.HeldFor != "vital" || hs.Preempted != 1 {
		t.Errorf("the victim's health does not say it gave way: %+v", hs)
	}

	h.settle(t, 4)
	if replicasOf(h, "vital") != 1 {
		t.Fatalf("the critical model was not placed once room was made: %+v", h.r.Replicas())
	}
	// Held back until the critical model had its replica, then short again
	// because there is no room: it must not have taken the node back — not
	// even for a moment, which would load a model only to preempt it again.
	if replicasOf(h, "chat") != 0 {
		t.Error("the standard model took the node back")
	}
	started := 0
	for _, cmd := range h.disp.startedShards() {
		if strings.HasPrefix(cmd.ReplicaID, "chat-") {
			started++
		}
	}
	if started != 1 {
		t.Errorf("the standard model started %d replicas; it should not have tried again", started)
	}
}

func TestAStandardModelDoesNotTakeAnothersMinimum(t *testing.T) {
	h := fillThenWant(t, gpuModel("second", "", 1), gpuModel("first", "", 1), gpuNode("1", 8))
	h.settle(t, 3)
	if replicasOf(h, "first") != 1 || replicasOf(h, "second") != 0 {
		t.Errorf("a model of the same class took another's minimum: %+v", h.r.Replicas())
	}
	if hs := health(h, "second"); !hs.Unplaceable {
		t.Error("the model left without room should say it cannot be placed")
	}
}

func TestWithinAClassOnlyWhatWasGrownIntoIsTaken(t *testing.T) {
	grown := gpuModel("busy", "", 1)
	grown.Scaling = catalog.Scaling{Min: 1, Max: 2}
	h := newFleetHarness(t, []catalog.Model{grown, withMin(gpuModel("new", "", 1), 0)}, gpuNode("1", 8), gpuNode("2", 8))
	h.settle(t, 2)
	// Load grows it to its max: the second replica is above its minimum.
	h.r.mu.Lock()
	h.r.extra["busy"] = 1
	h.r.mu.Unlock()
	h.settle(t, 3)
	if replicasOf(h, "busy") != 2 {
		t.Fatalf("setup: busy did not grow to 2: %+v", h.r.Replicas())
	}
	h.advance(h.r.cfg.MinLifetime + time.Minute)
	cat, _ := catalog.New([]catalog.Model{grown, gpuModel("new", "", 1)})
	h.r.cat.Store(cat)

	h.settle(t, 5)
	if replicasOf(h, "new") != 1 || replicasOf(h, "busy") != 1 {
		t.Errorf("the grown replica should give way to a minimum: busy %d, new %d",
			replicasOf(h, "busy"), replicasOf(h, "new"))
	}
}

func TestNothingIsStoppedThatWouldNotMakeRoom(t *testing.T) {
	// Two nodes, each half used by an opportunistic model; the critical one
	// needs a whole node. Stopping one frees half a node, which is no use —
	// and both together are on two nodes, which is no use either.
	half := catalog.Model{Name: "half", Source: "/m", Class: catalog.ClassOpportunistic,
		Scaling:  catalog.Scaling{Min: 2, Max: 2},
		Profiles: []catalog.Profile{{Name: "tp4", Engine: "vllm", GPUClass: "h100-80g", TP: 4}}}
	five := func(id string) inventory.Node { return gpuNode(id, 5) }
	h := fillThenWant(t, gpuModel("vital", catalog.ClassCritical, 1), half, five("1"), five("2"))
	h.r.Reconcile(t.Context())
	for _, rep := range h.r.Replicas() {
		if rep.PreemptedFor != "" || rep.State == Draining {
			t.Fatalf("a replica was stopped without making room: %+v", rep)
		}
	}
}

func TestACriticalMinimumIsNeverTaken(t *testing.T) {
	h := fillThenWant(t, gpuModel("also-vital", catalog.ClassCritical, 1), gpuModel("vital", catalog.ClassCritical, 1), gpuNode("1", 8))
	h.settle(t, 3)
	if replicasOf(h, "vital") != 1 {
		t.Error("a critical model's minimum was taken")
	}
}

func TestAReplicaThatJustStartedIsLeftAlone(t *testing.T) {
	h := newFleetHarness(t, []catalog.Model{gpuModel("chat", "", 1), withMin(gpuModel("vital", catalog.ClassCritical, 1), 0)}, gpuNode("1", 8))
	h.settle(t, 2)
	cat, _ := catalog.New([]catalog.Model{gpuModel("chat", "", 1), gpuModel("vital", catalog.ClassCritical, 1)})
	h.r.cat.Store(cat)
	h.r.Reconcile(t.Context())
	for _, rep := range h.r.Replicas() {
		if rep.State == Draining {
			t.Fatal("a replica that had only just loaded was preempted")
		}
	}
}

func health(h *harness, model string) ModelHealth {
	for _, x := range h.r.ModelHealth() {
		if x.Model == model {
			return x
		}
	}
	return ModelHealth{}
}

func TestAHoldLastsUntilTheRoomIsUsedOrItExpires(t *testing.T) {
	h := fillThenWant(t, gpuModel("vital", catalog.ClassCritical, 1), gpuModel("chat", "", 1), gpuNode("1", 8))
	h.r.Reconcile(t.Context())
	if !h.r.isHeld("chat") {
		t.Fatal("the model that gave way is not held")
	}
	if h.r.isHeld("vital") {
		t.Error("the model room was made for is held")
	}
	// Released as soon as the model it was given to has what it asked for.
	h.settle(t, 3)
	if replicasOf(h, "vital") != 1 {
		t.Fatalf("setup: vital was not placed: %+v", h.r.Replicas())
	}
	if h.r.isHeld("chat") {
		t.Error("the hold outlasted its purpose")
	}

	// And, whatever happens, after PreemptHold.
	h.r.mu.Lock()
	h.r.held["other"] = hold{Until: h.r.now().Add(time.Minute), For: "never-placed"}
	h.r.mu.Unlock()
	if !h.r.isHeld("other") {
		t.Fatal("a fresh hold is not in force")
	}
	h.advance(2 * time.Minute)
	if h.r.isHeld("other") {
		t.Error("a hold outlived its time")
	}
}

func TestOnlyTheReplicasThatMakeTheRoomAreStopped(t *testing.T) {
	// Two half-used nodes. "idle" on node 1 is on demand and long unused, so
	// it is first in line; but the critical model needs node 2's other half,
	// which only stopping "chat" frees.
	idle := catalog.Model{Name: "idle", Source: "/m", Scaling: catalog.Scaling{Min: 0, Max: 1},
		Profiles: []catalog.Profile{{Name: "tp4", Engine: "vllm", GPUClass: "h100-80g", TP: 4}}}
	chat := catalog.Model{Name: "chat", Source: "/m", Scaling: catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "tp6", Engine: "vllm", GPUClass: "h100-80g", TP: 6}}}
	vital := catalog.Model{Name: "vital", Source: "/m", Class: catalog.ClassCritical,
		Scaling:  catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "tp8", Engine: "vllm", GPUClass: "h100-80g", TP: 8}}}
	h := newFleetHarness(t, []catalog.Model{idle, chat}, gpuNode("1", 6), gpuNode("2", 8))
	h.r.Want("idle")
	h.settle(t, 4)
	var idleNode, chatNode string
	for _, rep := range h.r.Replicas() {
		switch rep.Model {
		case "idle":
			idleNode = rep.Shards[0].NodeID
		case "chat":
			chatNode = rep.Shards[0].NodeID
		}
	}
	if idleNode == "" || chatNode != "2" {
		t.Fatalf("setup: idle on %q, chat on %q; chat must hold node 2: %+v", idleNode, chatNode, h.r.Replicas())
	}
	h.r.mu.Lock()
	for _, rep := range h.r.replicas {
		rep.LastUsedAt = h.r.now().Add(-time.Hour)
	}
	h.r.mu.Unlock()
	h.advance(h.r.cfg.MinLifetime + time.Minute)
	cat, _ := catalog.New([]catalog.Model{idle, chat, vital})
	h.r.cat.Store(cat)

	h.r.Reconcile(t.Context())
	var stopped []string
	for _, rep := range h.r.Replicas() {
		if rep.PreemptedFor == "vital" {
			stopped = append(stopped, rep.Model)
		}
	}
	if len(stopped) != 1 || stopped[0] != "chat" {
		t.Errorf("stopped %v; only chat's replica makes room for a whole node", stopped)
	}
}

func TestOfTwoThatWouldEachMakeRoomTheSmallerIsStopped(t *testing.T) {
	cpu := inventory.Node{ID: "a", MgmtAddr: "10.0.0.1", State: inventory.Ready, MemoryTotalMB: 31 * 1024, CPUCount: 8}
	big := catalog.Model{Name: "big", Source: "/m", Scaling: catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp", MinMemoryGB: 12, CPUCores: 4}}}
	small := catalog.Model{Name: "small", Source: "/m", Scaling: catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp", MinMemoryGB: 2, CPUCores: 1}}}
	vital := catalog.Model{Name: "vital", Source: "/m", Class: catalog.ClassCritical,
		Scaling:  catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp", MinMemoryGB: 15, CPUCores: 4}}}
	// 8 cores: big 4 + small 1 leaves 3, and vital needs 4. Stopping either
	// is enough; stopping big loses far more.
	for range 5 { // the choice must not depend on map order
		h := newFleetHarness(t, []catalog.Model{big, small}, cpu)
		h.settle(t, 3)
		h.advance(h.r.cfg.MinLifetime + time.Minute)
		cat, _ := catalog.New([]catalog.Model{big, small, vital})
		h.r.cat.Store(cat)
		h.r.Reconcile(t.Context())
		var stopped []string
		for _, rep := range h.r.Replicas() {
			if rep.PreemptedFor == "vital" {
				stopped = append(stopped, rep.Model)
			}
		}
		if len(stopped) != 1 || stopped[0] != "small" {
			t.Fatalf("stopped %v; the small replica alone makes room", stopped)
		}
	}
}

func TestHeadroomCountsWhatAHigherClassCouldTake(t *testing.T) {
	h := fillThenWant(t, gpuModel("vital", catalog.ClassCritical, 1), gpuModel("chat", "", 1), gpuNode("1", 8))
	got, _ := h.r.Headroom("vital", 3)
	if got.More != 0 || got.ByPreempting != 1 || len(got.Preempts) != 1 || got.Preempts[0] != "chat" {
		t.Errorf("a critical model should see one more by stopping chat: %+v", got)
	}
	// Tried on a copy: nothing was stopped.
	if replicasOf(h, "chat") != 1 {
		t.Fatal("asking about room stopped a replica")
	}
	// A model of chat's own class may not take its minimum, so has no such room.
	peer := newFleetHarness(t, []catalog.Model{gpuModel("chat", "", 1), withMin(gpuModel("peer", "", 1), 0)}, gpuNode("1", 8))
	peer.settle(t, 3)
	peer.advance(peer.r.cfg.MinLifetime + time.Minute)
	if got, _ := peer.r.Headroom("peer", 3); got.ByPreempting != 0 {
		t.Errorf("a standard model counted room in another's minimum: %+v", got)
	}
}

func TestARestartBorrowsALowerClassesRoomAndGivesItBack(t *testing.T) {
	// chat fills node 1; the only other node holds an opportunistic replica.
	chat := gpuModel("chat", "", 1)
	extra := gpuModel("extra", catalog.ClassOpportunistic, 1)
	h := newFleetHarness(t, []catalog.Model{chat, extra}, gpuNode("1", 8), gpuNode("2", 8))
	h.settle(t, 3)
	h.advance(h.r.cfg.MinLifetime + time.Minute)
	var old string
	for _, rep := range h.r.Replicas() {
		if rep.Model == "chat" {
			old = rep.ID
		}
	}
	if old == "" || replicasOf(h, "extra") != 1 {
		t.Fatalf("setup: %+v", h.r.Replicas())
	}

	if err := h.r.Restart(old); err != nil {
		t.Fatal(err)
	}
	// Replacement first, beside the old one, in the room extra gives up;
	// the old one keeps serving until then.
	servedThroughout := true
	for range 8 {
		h.r.Reconcile(t.Context())
		h.markReady()
		ready := 0
		for _, rep := range h.r.Replicas() {
			if rep.Model == "chat" && rep.State == Ready {
				ready++
			}
		}
		if ready == 0 {
			servedThroughout = false
		}
		h.advance(h.r.cfg.PreemptGrace + time.Second)
	}
	if !servedThroughout {
		t.Error("chat had no ready replica during its restart")
	}
	var newChat bool
	for _, rep := range h.r.Replicas() {
		if rep.Model == "chat" && rep.ID != old && rep.State == Ready {
			newChat = true
		}
		if rep.ID == old {
			t.Errorf("the old replica is still there: %+v", rep)
		}
	}
	if !newChat {
		t.Fatalf("no new chat replica: %+v", h.r.Replicas())
	}
	// And extra has its replica back, in the room the old chat left.
	h.settle(t, 4)
	if replicasOf(h, "extra") != 1 {
		t.Errorf("the opportunistic model did not get its room back: %+v", h.r.Replicas())
	}
}

func TestHeadroomNamesOnlyTheReplicasThatWouldMakeRoom(t *testing.T) {
	// vital needs a whole node and chat is on the only one; a second
	// critical can never fit. Asking for 3: one by stopping chat, then no
	// candidates left — nothing else may be named.
	h := fillThenWant(t, gpuModel("vital", catalog.ClassCritical, 1), gpuModel("chat", "", 1), gpuNode("1", 8))
	got, _ := h.r.Headroom("vital", 3)
	if got.ByPreempting != 1 || len(got.Preempts) != 1 {
		t.Errorf("%+v", got)
	}
	// Free room for one, and candidates that cannot make room for a second:
	// none of them may be listed.
	small := catalog.Model{Name: "small", Source: "/m", Class: catalog.ClassOpportunistic,
		Scaling:  catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "tp2", Engine: "vllm", GPUClass: "h100-80g", TP: 2}}}
	h2 := fillThenWant(t, gpuModel("big", "", 1), small, gpuNode("1", 8), gpuNode("2", 8))
	got2, _ := h2.r.Headroom("big", 3)
	if got2.More != 1 || got2.ByPreempting != 1 || len(got2.Preempts) != 1 {
		t.Errorf("one free, one by stopping small: %+v", got2)
	}
	h3 := fillThenWant(t, gpuModel("big", "", 1), small, gpuNode("1", 8), gpuNode("2", 6))
	got3, _ := h3.r.Headroom("big", 3)
	if got3.ByPreempting != 0 || len(got3.Preempts) != 0 {
		t.Errorf("stopping small frees a 6-GPU node, too small for big; nobody may be named: %+v", got3)
	}
}