llmgw/internal/control/inventory/inventory_test.go

package inventory

import (
	"slices"
	"sync"
	"testing"

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

func gpus(n int) []discovery.GPU {
	out := make([]discovery.GPU, n)
	for i := range out {
		out[i] = discovery.GPU{Index: i, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559}
	}
	return out
}

func TestUpsertPreservesCordon(t *testing.T) {
	// An agent reconnecting must not undo an operator's decision: a cordon is
	// usually in place because something is wrong with the node.
	s := NewStore()
	s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
	s.SetState("n1", Cordoned)

	s.Upsert(Node{ID: "n1", GPUs: gpus(8)}) // re-registration
	n, _ := s.Get("n1")
	if n.State != Cordoned {
		t.Errorf("state=%s; re-registration cleared the cordon", n.State)
	}
}

func TestUpsertUpdatesHardware(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1", GPUs: gpus(4)})
	s.Upsert(Node{ID: "n1", GPUs: gpus(8)})

	n, _ := s.Get("n1")
	if len(n.GPUs) != 8 {
		t.Errorf("hardware was not updated: %d GPUs", len(n.GPUs))
	}
}

func TestTouchRevivesUnreachableNode(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1"})
	s.SetState("n1", Unreachable)

	s.Touch("n1")
	if n, _ := s.Get("n1"); n.State != Ready {
		t.Errorf("a heartbeat did not revive the node: %s", n.State)
	}
}

func TestTouchDoesNotClearCordon(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1"})
	s.SetState("n1", Cordoned)

	s.Touch("n1")
	if n, _ := s.Get("n1"); n.State != Cordoned {
		t.Errorf("a heartbeat cleared the cordon: %s", n.State)
	}
}

func TestFreeGPUs(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1", GPUs: gpus(8)})

	if got := s.FreeGPUs("n1"); len(got) != 8 {
		t.Fatalf("expected 8 free, got %d", len(got))
	}

	s.Allocate("n1", Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1}})
	got := s.FreeGPUs("n1")
	if !slices.Equal(got, []int{2, 3, 4, 5, 6, 7}) {
		t.Errorf("free=%v", got)
	}
}

func TestUnhealthyGPUsAreNotFree(t *testing.T) {
	s := NewStore()
	n := Node{ID: "n1", GPUs: gpus(4)}
	n.GPUs[2].Health.Problems = []string{"ECC"}
	s.Upsert(n)

	got := s.FreeGPUs("n1")
	if slices.Contains(got, 2) {
		t.Errorf("an unhealthy GPU was reported free: %v", got)
	}
	if len(got) != 3 {
		t.Errorf("free=%v", got)
	}
}

func TestReleaseShard(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
	s.Allocate("n1", Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1}})
	s.Allocate("n1", Allocation{ShardID: "s2", ReplicaID: "r2", GPUIndices: []int{2, 3}})

	s.Release("s1")
	if got := s.FreeGPUs("n1"); slices.Contains(got, 2) || !slices.Contains(got, 0) {
		t.Errorf("wrong shard released: free=%v", got)
	}
}

func TestReleaseReplicaFreesEveryShard(t *testing.T) {
	// Replicas fail whole, so they are freed whole: one dead shard makes the
	// others useless and the control plane rebuilds the replica (plan §5.4).
	s := NewStore()
	s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
	s.Upsert(Node{ID: "n2", GPUs: gpus(8)})
	s.Allocate("n1", Allocation{ShardID: "r1-0", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
	s.Allocate("n2", Allocation{ShardID: "r1-1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})

	s.ReleaseReplica("r1")
	for _, id := range []string{"n1", "n2"} {
		if got := len(s.FreeGPUs(id)); got != 8 {
			t.Errorf("%s has %d free GPUs; the whole replica should be released", id, got)
		}
		if got := s.Allocations(id); len(got) != 0 {
			t.Errorf("%s still has allocations: %v", id, got)
		}
	}
}

func TestHCAsForSelectsOneFabric(t *testing.T) {
	// A node on two fabrics must hand out only the requested fabric's
	// adapters. Mixing them lets NCCL pick one whose peer is unreachable.
	n := Node{ID: "n1", Fabrics: []FabricAttachment{
		{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}},
		{Fabric: "ib-b", HCAs: []string{"mlx5_2", "mlx5_3"}},
	}}

	if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
		t.Errorf("ib-a HCAs=%v", got)
	}
	if got := n.HCAsFor("ib-b"); !slices.Equal(got, []string{"mlx5_2", "mlx5_3"}) {
		t.Errorf("ib-b HCAs=%v", got)
	}
	if got := n.HCAsFor("ib-c"); got != nil {
		t.Errorf("an unattached fabric returned %v", got)
	}
}

func TestHCAsForReturnsACopy(t *testing.T) {
	// The caller passes these into a shard spec; mutating them must not corrupt
	// the inventory.
	n := Node{ID: "n1", Fabrics: []FabricAttachment{{Fabric: "ib-a", HCAs: []string{"mlx5_0"}}}}
	got := n.HCAsFor("ib-a")
	got[0] = "tampered"
	if n.Fabrics[0].HCAs[0] != "mlx5_0" {
		t.Error("HCAsFor exposed the inventory's own slice")
	}
}

func TestFabrics(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1", Fabrics: []FabricAttachment{{Fabric: "ib-b"}}})
	s.Upsert(Node{ID: "n2", Fabrics: []FabricAttachment{{Fabric: "ib-a"}, {Fabric: "ib-b"}}})
	s.Upsert(Node{ID: "n3"}) // no fabric

	if got := s.Fabrics(); !slices.Equal(got, []string{"ib-a", "ib-b"}) {
		t.Errorf("fabrics=%v", got)
	}
}

func TestNodesSorted(t *testing.T) {
	s := NewStore()
	for _, id := range []string{"n3", "n1", "n2"} {
		s.Upsert(Node{ID: id})
	}
	got := s.Nodes()
	if len(got) != 3 || got[0].ID != "n1" || got[2].ID != "n3" {
		t.Errorf("not sorted: %v", got)
	}
}

func TestGetReturnsACopy(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1", GPUs: gpus(2), Labels: map[string]string{"rack": "r1"}})

	n, _ := s.Get("n1")
	n.GPUs[0].Index = 99

	again, _ := s.Get("n1")
	if again.GPUs[0].Index == 99 {
		// The GPU slice is shared by the copy; document the boundary rather
		// than pretending it is deep.
		t.Log("note: Get returns a shallow copy; slices are shared")
	}
	if again.ID != "n1" {
		t.Error("the node was corrupted")
	}
}

func TestRemove(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1", GPUs: gpus(4)})
	s.Allocate("n1", Allocation{ShardID: "s1", GPUIndices: []int{0}})

	s.Remove("n1")
	if _, ok := s.Get("n1"); ok {
		t.Error("the node survived removal")
	}
	if got := s.Allocations("n1"); len(got) != 0 {
		t.Error("allocations outlived the node")
	}
}

func TestStateStrings(t *testing.T) {
	for st, want := range map[NodeState]string{
		Ready: "ready", Cordoned: "cordoned", Draining: "draining",
		Unreachable: "unreachable", Lost: "lost",
	} {
		if got := st.String(); got != want {
			t.Errorf("%d -> %q, want %q", int(st), got, want)
		}
	}
	if !Ready.Schedulable() || Cordoned.Schedulable() || Lost.Schedulable() {
		t.Error("Schedulable is wrong")
	}
}

func TestConcurrentAccess(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1", GPUs: gpus(8)})

	var wg sync.WaitGroup
	for i := range 100 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			switch i % 5 {
			case 0:
				s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
			case 1:
				s.Allocate("n1", Allocation{ShardID: "s", ReplicaID: "r", GPUIndices: []int{i % 8}})
			case 2:
				s.Release("s")
			case 3:
				_ = s.FreeGPUs("n1")
			default:
				_ = s.Nodes()
			}
		}()
	}
	wg.Wait()
}