llmgw/internal/control/placement/placement_test.go

package placement

import (
	"fmt"
	"math/rand/v2"
	"slices"
	"strings"
	"testing"

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

// node builds a node with n GPUs of one class.
func node(id string, class string, n int, fabrics ...string) inventory.Node {
	name, mem := classToHardware(class)

	nd := inventory.Node{ID: id, State: inventory.Ready}
	for i := range n {
		nd.GPUs = append(nd.GPUs, discovery.GPU{
			Index: i, Name: name, MemoryTotalMB: mem, ComputeCapability: "9.0",
			UUID: fmt.Sprintf("GPU-%s-%d", id, i),
		})
	}
	for _, f := range fabrics {
		nd.Fabrics = append(nd.Fabrics, inventory.FabricAttachment{
			Fabric: f, HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200,
		})
	}
	return nd
}

// classToHardware maps a class back to a name and size that derive it.
func classToHardware(class string) (string, int) {
	switch class {
	case "h100-80g":
		return "NVIDIA H100 80GB HBM3", 81559
	case "a100-40g":
		return "Tesla A100-SXM4-40GB", 40960
	default:
		return "NVIDIA L40S", 46068
	}
}

func store(nodes ...inventory.Node) *inventory.Store {
	s := inventory.NewStore()
	for _, n := range nodes {
		s.Upsert(n)
	}
	return s
}

func TestSingleNodePlacement(t *testing.T) {
	s := New(store(node("n1", "h100-80g", 8)))

	res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	if len(res.Placement.Nodes) != 1 {
		t.Fatalf("expected 1 node, got %d", len(res.Placement.Nodes))
	}
	if got := len(res.Placement.Nodes[0].GPUIndices); got != 4 {
		t.Errorf("got %d GPUs, want 4", got)
	}
	if res.Placement.Fabric != "" {
		t.Errorf("a single-node replica was given fabric %q", res.Placement.Fabric)
	}
}

func TestNodeWithoutFabricTakesSingleNodeWork(t *testing.T) {
	// A single-node replica talks to nobody, so it has nothing to be reachable
	// over. Excluding fabric-less nodes here would waste them entirely.
	s := New(store(node("no-ib", "h100-80g", 8)))

	res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
	if res.Placement == nil {
		t.Fatalf("a node without IB was refused single-node work:\n%s", res.Explain())
	}
}

func TestNodeWithoutFabricRefusedMultiNode(t *testing.T) {
	s := New(store(
		node("no-ib-1", "h100-80g", 8),
		node("no-ib-2", "h100-80g", 8),
	))

	res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
	if res.Placement != nil {
		t.Fatal("a multi-node replica was placed on nodes with no fabric")
	}
	if !strings.Contains(res.Explain(), "no node reports a fabric") {
		t.Errorf("the reason is unclear:\n%s", res.Explain())
	}
}

// TestNeverSpansTwoFabrics is the first hard invariant from plan §1.1.
func TestNeverSpansTwoFabrics(t *testing.T) {
	// Two fabrics, one node each: together they would satisfy the count, and a
	// scheduler that only counted nodes would happily combine them.
	s := New(store(
		node("a1", "h100-80g", 8, "ib-a"),
		node("b1", "h100-80g", 8, "ib-b"),
	))

	res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
	if res.Placement != nil {
		var fabrics []string
		for _, n := range res.Placement.Nodes {
			fabrics = append(fabrics, n.NodeID)
		}
		t.Fatalf("a replica was spread across fabrics: %v", fabrics)
	}
	if !strings.Contains(res.Explain(), "only 1 qualify") {
		t.Errorf("explanation should say each fabric was short of nodes:\n%s", res.Explain())
	}
}

func TestMultiNodeWithinOneFabric(t *testing.T) {
	s := New(store(
		node("a1", "h100-80g", 8, "ib-a"),
		node("a2", "h100-80g", 8, "ib-a"),
		node("b1", "h100-80g", 8, "ib-b"),
	))

	res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	if res.Placement.Fabric != "ib-a" {
		t.Errorf("fabric=%q, only ib-a has two nodes", res.Placement.Fabric)
	}
	for _, n := range res.Placement.Nodes {
		if n.NodeID == "b1" {
			t.Error("a node from the other fabric was included")
		}
		if len(n.HCAs) == 0 {
			t.Errorf("node %s got no HCAs; NCCL would pick its own", n.NodeID)
		}
	}
	// Ranks must be assigned, since rank 0 serves HTTP.
	ranks := []int{res.Placement.Nodes[0].Rank, res.Placement.Nodes[1].Rank}
	slices.Sort(ranks)
	if ranks[0] != 0 || ranks[1] != 1 {
		t.Errorf("ranks=%v", ranks)
	}
}

// TestNeverMixesGPUClasses is the second hard invariant from plan §1.1.
func TestNeverMixesGPUClasses(t *testing.T) {
	// Enough GPUs in total, but split across classes.
	n := node("mixed", "h100-80g", 4)
	a100Name, a100Mem := classToHardware("a100-40g")
	for i := 4; i < 8; i++ {
		n.GPUs = append(n.GPUs, discovery.GPU{Index: i, Name: a100Name, MemoryTotalMB: a100Mem})
	}
	s := New(store(n))

	res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
	if res.Placement != nil {
		t.Fatal("a placement used GPUs of more than one class")
	}
	if !strings.Contains(res.Explain(), "has 4 free h100-80g GPUs, needs 8") {
		t.Errorf("explanation should count only the matching class:\n%s", res.Explain())
	}

	// The 4-GPU profile does fit, using only H100s.
	res = s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
	if res.Placement == nil {
		t.Fatalf("the matching-class profile was refused:\n%s", res.Explain())
	}
	for _, idx := range res.Placement.Nodes[0].GPUIndices {
		if idx >= 4 {
			t.Errorf("GPU %d is an A100 and was included in an h100 placement", idx)
		}
	}
}

func TestProfilePriorityOrder(t *testing.T) {
	s := New(store(
		node("h", "h100-80g", 8, "ib-a"),
		node("a", "a100-40g", 8, "ib-a"),
	))

	res := s.Place([]Profile{
		{Name: "a100", GPUClass: "a100-40g", GPUsPerNode: 8, Nodes: 1, Priority: 50},
		{Name: "h100", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1, Priority: 100},
	}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	if res.Placement.Profile.Name != "h100" {
		t.Errorf("chose %q; the higher-priority profile fits", res.Placement.Profile.Name)
	}
}

func TestFallsBackToLowerPriorityProfile(t *testing.T) {
	s := New(store(node("a", "a100-40g", 8)))

	res := s.Place([]Profile{
		{Name: "h100", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1, Priority: 100},
		{Name: "a100", GPUClass: "a100-40g", GPUsPerNode: 8, Nodes: 1, Priority: 50},
	}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	if res.Placement.Profile.Name != "a100" {
		t.Errorf("chose %q", res.Placement.Profile.Name)
	}
	// The rejected first choice must still be explained.
	if !strings.Contains(res.Explain(), "h100") {
		t.Errorf("the higher-priority profile's rejection is missing:\n%s", res.Explain())
	}
}

func TestCordonedNodeExcluded(t *testing.T) {
	st := store(node("n1", "h100-80g", 8))
	st.SetState("n1", inventory.Cordoned)
	s := New(st)

	res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
	if res.Placement != nil {
		t.Fatal("a cordoned node was used")
	}
	if !strings.Contains(res.Explain(), "cordoned") {
		t.Errorf("explanation:\n%s", res.Explain())
	}
}

func TestUnhealthyGPUsExcluded(t *testing.T) {
	n := node("n1", "h100-80g", 8)
	n.GPUs[0].Health.Problems = []string{"12 uncorrected ECC errors"}
	n.GPUs[1].Health.Problems = []string{"40000 MiB already allocated by another process"}
	s := New(store(n))

	res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
	if res.Placement != nil {
		t.Fatal("unhealthy GPUs were counted as available")
	}

	res = s.Place([]Profile{{Name: "tp6", GPUClass: "h100-80g", GPUsPerNode: 6, Nodes: 1}}, nil)
	if res.Placement == nil {
		t.Fatalf("the six healthy GPUs should have been usable:\n%s", res.Explain())
	}
	for _, idx := range res.Placement.Nodes[0].GPUIndices {
		if idx == 0 || idx == 1 {
			t.Errorf("unhealthy GPU %d was placed on", idx)
		}
	}
}

func TestAllocatedGPUsExcluded(t *testing.T) {
	st := store(node("n1", "h100-80g", 8))
	st.Allocate("n1", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
	s := New(st)

	res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
	if res.Placement == nil {
		t.Fatalf("the remaining four GPUs should fit:\n%s", res.Explain())
	}
	for _, idx := range res.Placement.Nodes[0].GPUIndices {
		if idx < 4 {
			t.Errorf("GPU %d is already allocated", idx)
		}
	}

	// Another four will not fit.
	st.Allocate("n1", inventory.Allocation{ShardID: "s2", ReplicaID: "r2", GPUIndices: []int{4, 5, 6, 7}})
	if res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil); res.Placement != nil {
		t.Fatal("a full node was placed on")
	}
}

func TestAvoidSpreadsAcrossNodes(t *testing.T) {
	s := New(store(
		node("n1", "h100-80g", 8),
		node("n2", "h100-80g", 8),
	))
	p := []Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}

	first := s.Place(p, nil)
	if first.Placement == nil {
		t.Fatal("no first placement")
	}
	placed := first.Placement.Nodes[0].NodeID

	second := s.Place(p, []string{placed})
	if second.Placement == nil {
		t.Fatalf("the second replica had nowhere to go:\n%s", second.Explain())
	}
	if second.Placement.Nodes[0].NodeID == placed {
		t.Error("both replicas landed on the same node despite avoid")
	}
}

func TestBinPackingPrefersPartlyUsedNodes(t *testing.T) {
	// Keeping whole nodes free matters: a large model that needs all 8 GPUs
	// cannot be split, so scattering small replicas would strand it.
	st := store(
		node("empty", "h100-80g", 8),
		node("partial", "h100-80g", 8),
	)
	st.Allocate("partial", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
	s := New(st)

	res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	if got := res.Placement.Nodes[0].NodeID; got != "partial" {
		t.Errorf("placed on %q; the partly used node should be filled first", got)
	}
}

func TestNVLinkGroupPreferred(t *testing.T) {
	// GPUs 0-3 and 4-7 form separate NVLink domains. A TP group crossing them
	// falls back to PCIe for its all-reduces and runs several times slower.
	n := node("n1", "h100-80g", 8)
	for i := range n.GPUs {
		var peers []int
		base := (i / 4) * 4
		for j := base; j < base+4; j++ {
			if j != i {
				peers = append(peers, j)
			}
		}
		n.GPUs[i].NVLinkPeers = peers
	}
	st := store(n)
	// Take GPU 0, so a naive "lowest indices" pick would span the two domains.
	st.Allocate("n1", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0}})
	s := New(st)

	res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	got := res.Placement.Nodes[0].GPUIndices
	if !slices.Equal(got, []int{4, 5, 6, 7}) {
		t.Errorf("chose %v; 4-7 is the only intact NVLink domain left", got)
	}
}

func TestExplainIsUsefulWhenNothingFits(t *testing.T) {
	s := New(store(node("n1", "a100-40g", 2)))

	res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
	if res.Placement != nil {
		t.Fatal("unexpected placement")
	}
	exp := res.Explain()
	// It must name the node, the profile, and what was actually wrong.
	for _, want := range []string{"n1", "tp8", "free h100-80g"} {
		if !strings.Contains(exp, want) {
			t.Errorf("explanation is missing %q:\n%s", want, exp)
		}
	}
}

func TestEmptyInventory(t *testing.T) {
	s := New(inventory.NewStore())
	res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
	if res.Placement != nil {
		t.Fatal("something was placed on an empty cluster")
	}
}

// --- property-based invariants (plan §13) ---

// TestInvariantsHold generates random clusters and profiles and checks the two
// hard constraints on every placement produced.
//
// The table-driven tests above cover the cases someone thought of; this covers
// the ones nobody did.
func TestInvariantsHold(t *testing.T) {
	rng := rand.New(rand.NewPCG(1, 2))
	classes := []string{"h100-80g", "a100-40g", "l40s-45g"}
	fabrics := []string{"ib-a", "ib-b", "ib-c"}

	for iter := range 400 {
		st := inventory.NewStore()
		nodeCount := 1 + rng.IntN(6)

		for i := range nodeCount {
			class := classes[rng.IntN(len(classes))]
			gpuCount := 1 + rng.IntN(8)

			var nf []string
			switch rng.IntN(4) {
			case 0: // no fabric at all
			case 1:
				nf = []string{fabrics[rng.IntN(len(fabrics))]}
			default:
				nf = []string{fabrics[rng.IntN(len(fabrics))]}
				if rng.IntN(3) == 0 {
					// A node attached to two fabrics: allowed, and the reason
					// HCAs must be selected per fabric rather than per node.
					nf = append(nf, fabrics[rng.IntN(len(fabrics))])
				}
			}
			n := node(fmt.Sprintf("n%d", i), class, gpuCount, nf...)
			if rng.IntN(8) == 0 {
				n.State = inventory.Cordoned
			}
			st.Upsert(n)
		}

		s := New(st)
		p := Profile{
			Name:        "gen",
			GPUClass:    classes[rng.IntN(len(classes))],
			GPUsPerNode: 1 + rng.IntN(8),
			Nodes:       1 + rng.IntN(3),
		}
		res := s.Place([]Profile{p}, nil)
		if res.Placement == nil {
			continue
		}
		checkInvariants(t, st, p, res.Placement, iter)
	}
}

func checkInvariants(t *testing.T, st *inventory.Store, p Profile, pl *Placement, iter int) {
	t.Helper()

	if len(pl.Nodes) != max(p.Nodes, 1) {
		t.Fatalf("iter %d: placed on %d nodes, profile wants %d", iter, len(pl.Nodes), p.Nodes)
	}

	seenNodes := map[string]struct{}{}
	for _, a := range pl.Nodes {
		if _, dup := seenNodes[a.NodeID]; dup {
			t.Fatalf("iter %d: node %s used twice in one replica", iter, a.NodeID)
		}
		seenNodes[a.NodeID] = struct{}{}

		n, ok := st.Get(a.NodeID)
		if !ok {
			t.Fatalf("iter %d: placed on unknown node %s", iter, a.NodeID)
		}
		if !n.State.Schedulable() {
			t.Fatalf("iter %d: placed on a %s node", iter, n.State)
		}
		if len(a.GPUIndices) != p.GPUsPerNode {
			t.Fatalf("iter %d: node %s got %d GPUs, wants %d", iter, a.NodeID, len(a.GPUIndices), p.GPUsPerNode)
		}

		// Invariant: every GPU is of the profile's class, healthy, and real.
		byIndex := map[int]discovery.GPU{}
		for _, g := range n.GPUs {
			byIndex[g.Index] = g
		}
		for _, idx := range a.GPUIndices {
			g, ok := byIndex[idx]
			if !ok {
				t.Fatalf("iter %d: node %s has no GPU %d", iter, a.NodeID, idx)
			}
			if g.Class() != p.GPUClass {
				t.Fatalf("iter %d: GPU %d on %s is %s, profile wants %s",
					iter, idx, a.NodeID, g.Class(), p.GPUClass)
			}
			if !g.Health.OK() {
				t.Fatalf("iter %d: unhealthy GPU %d was placed on", iter, idx)
			}
		}

		// Invariant: a multi-node replica lives on one fabric, and each shard
		// gets only that fabric's adapters.
		if p.Nodes > 1 {
			if pl.Fabric == "" {
				t.Fatalf("iter %d: multi-node placement with no fabric", iter)
			}
			if !n.HasFabric(pl.Fabric) {
				t.Fatalf("iter %d: node %s is not attached to %s", iter, a.NodeID, pl.Fabric)
			}
			want := n.HCAsFor(pl.Fabric)
			if !slices.Equal(a.HCAs, want) {
				t.Fatalf("iter %d: node %s got HCAs %v, want exactly %v for %s",
					iter, a.NodeID, a.HCAs, want, pl.Fabric)
			}
		}
	}
}