llmgw/internal/control/slurmsvc/service_test.go

package slurmsvc

import (
	"context"
	"slices"
	"strings"
	"testing"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)

type fakeInv struct {
	nodes  map[string]*inventory.Node
	allocs map[string]int
}

func (f *fakeInv) Get(id string) (inventory.Node, bool) {
	n, ok := f.nodes[id]
	if !ok {
		return inventory.Node{}, false
	}
	return *n, true
}

func (f *fakeInv) SetState(id string, st inventory.NodeState) bool {
	n, ok := f.nodes[id]
	if ok {
		n.State = st
	}
	return ok
}

func (f *fakeInv) Allocations(id string) []inventory.Allocation {
	return make([]inventory.Allocation, f.allocs[id])
}

func (f *fakeInv) Nodes() []inventory.Node {
	var out []inventory.Node
	for _, n := range f.nodes {
		out = append(out, *n)
	}
	slices.SortFunc(out, func(a, b inventory.Node) int { return strings.Compare(a.ID, b.ID) })
	return out
}

type fakeAgents struct {
	conn []string
	got  map[string][]*llmgwv1.Slurm
}

func (f *fakeAgents) Send(node string, m *llmgwv1.ControlMessage) error {
	f.got[node] = append(f.got[node], m.GetSlurm())
	return nil
}

func (f *fakeAgents) ConnectedNodes() []string { return f.conn }

func (f *fakeAgents) last(node string) *llmgwv1.Slurm {
	ms := f.got[node]
	if len(ms) == 0 {
		return nil
	}
	return ms[len(ms)-1]
}

type rig struct {
	set    config.SlurmSettings
	inv    *fakeInv
	agents *fakeAgents
	svc    *Service
}

func newRig() *rig {
	r := &rig{
		set: config.SlurmSettings{Enabled: true, Controller: "node06", Nodes: []string{"cuda01"}},
		inv: &fakeInv{nodes: map[string]*inventory.Node{
			"node06": {ID: "node06", MgmtAddr: "10.0.0.146"},
			"cuda01": {ID: "cuda01", MgmtAddr: "10.0.0.147", GPUs: []discovery.GPU{{Index: 0, Name: "NVIDIA GeForce RTX 3090"}}},
		}, allocs: map[string]int{"cuda01": 2}},
		agents: &fakeAgents{conn: []string{"node06", "cuda01"}, got: map[string][]*llmgwv1.Slurm{}},
	}
	r.svc = New(func() config.SlurmSettings { return r.set }, r.inv, r.agents,
		func(context.Context) ([]byte, error) { return []byte(strings.Repeat("k", 1024)), nil }, nil)
	return r
}

func (r *rig) report(node string, st *llmgwv1.SlurmStatus) { r.svc.OnSlurmStatus(node, st) }

func (r *rig) tick() { r.svc.Tick(context.Background()) }

const hw = "NodeName=cuda01 CPUs=16 RealMemory=64000"

func TestMessagesAndFiles(t *testing.T) {
	r := newRig()
	r.tick()
	ctl, node := r.agents.last("node06"), r.agents.last("cuda01")
	if !ctl.GetController() || ctl.GetNode() || len(ctl.GetMungeKey()) != 1024 {
		t.Fatalf("controller told %v", ctl)
	}
	if !node.GetNode() || node.GetController() || node.GetConfServer() != "10.0.0.146:6817" || node.GetLent() ||
		node.GetName() != "cuda01" || len(node.GetFiles()) != 0 {
		t.Fatalf("node told %v", node)
	}
	if strings.Contains(ctl.GetFiles()["slurm.conf"], "NodeName=cuda01") {
		t.Fatal("a node with no hardware reported is in slurm.conf")
	}
	if w := r.svc.Status().Waiting; len(w) != 1 || w[0] != "cuda01" {
		t.Fatalf("waiting %v", w)
	}
	// Same again: nobody is told twice.
	r.tick()
	if len(r.agents.got["node06"]) != 1 || len(r.agents.got["cuda01"]) != 1 {
		t.Fatal("told again with nothing changed")
	}
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "IDLE+DRAIN", Reason: "llmgw: not lent yet"})
	r.tick()
	conf := r.agents.last("node06").GetFiles()["slurm.conf"]
	if !strings.Contains(conf, "NodeName=cuda01 NodeAddr=10.0.0.147 CPUs=16 RealMemory=62080 Gres=gpu:rtx3090:1") {
		t.Fatalf("slurm.conf:\n%s", conf)
	}
	if !strings.Contains(r.agents.last("node06").GetFiles()["gres.conf"], "File=/dev/nvidia0") {
		t.Fatal("gres.conf has no card")
	}
	if len(r.agents.got["cuda01"]) != 1 {
		t.Fatal("a node was told again when only the controller's files changed")
	}
	// A new stream is told again.
	r.svc.OnRegister("cuda01")
	r.tick()
	if len(r.agents.got["cuda01"]) != 2 {
		t.Fatal("a new stream was not told")
	}
	if r.inv.nodes["cuda01"].State != inventory.Ready {
		t.Fatal("a node drained in Slurm with no job was taken from llmgw")
	}
}

func TestLendAndReclaim(t *testing.T) {
	r := newRig()
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "IDLE+DRAIN", Reason: "llmgw: not lent yet"})
	r.set.Lent = []string{"cuda01"}
	r.tick()
	if r.inv.nodes["cuda01"].State != inventory.Draining {
		t.Fatal("a lent node was not drained of its replicas")
	}
	if r.agents.last("cuda01").GetLent() {
		t.Fatal("lent to Slurm with replicas still on it")
	}
	if p := r.svc.Status().Nodes[0].Phase; p != Lending {
		t.Fatalf("phase %s", p)
	}
	r.inv.allocs["cuda01"] = 0
	r.tick()
	if !r.agents.last("cuda01").GetLent() || r.svc.Status().Nodes[0].Phase != Lent {
		t.Fatal("not lent once its replicas were gone")
	}
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "ALLOCATED", Jobs: 3})

	// Taken back: Slurm drains it; llmgw keeps off until its jobs end.
	r.set.Lent = nil
	r.tick()
	if r.agents.last("cuda01").GetLent() {
		t.Fatal("still lent")
	}
	if s := r.inv.nodes["cuda01"].State; s != inventory.Draining && s != inventory.Cordoned {
		t.Fatalf("llmgw may place on a node with Slurm jobs: %v", s)
	}
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "ALLOCATED+DRAIN", Jobs: 1})
	r.tick()
	if r.svc.Status().Nodes[0].Phase != Reclaiming || r.inv.nodes["cuda01"].State == inventory.Ready {
		t.Fatal("given back with a job left")
	}
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "IDLE+DRAIN", Reason: "llmgw: serving models"})
	r.tick()
	if r.inv.nodes["cuda01"].State != inventory.Ready || r.svc.Status().Nodes[0].Phase != Llmgw {
		t.Fatal("not llmgw's when its last job ended")
	}
}

func TestNewLeaderKeepsOffUntilHeard(t *testing.T) {
	r := newRig()
	r.inv.allocs["cuda01"] = 0
	r.tick()
	// Nothing heard yet: cordoned, nothing moved.
	if r.inv.nodes["cuda01"].State != inventory.Cordoned {
		t.Fatalf("state %v before the agent said anything", r.inv.nodes["cuda01"].State)
	}
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "MIXED", Jobs: 2})
	r.tick()
	if r.inv.nodes["cuda01"].State == inventory.Ready {
		t.Fatal("ready with Slurm jobs on it")
	}
}

func TestBrokenSlurmLeavesTheNodeToLlmgw(t *testing.T) {
	r := newRig()
	r.report("cuda01", &llmgwv1.SlurmStatus{Error: "installing Slurm: no package", Jobs: 0})
	r.tick()
	if r.inv.nodes["cuda01"].State != inventory.Ready {
		t.Fatal("a Slurm that does not work took a node from llmgw")
	}
}

func TestOperatorsCordonIsLeft(t *testing.T) {
	r := newRig()
	r.inv.nodes["cuda01"].State = inventory.Cordoned
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "IDLE+DRAIN"})
	r.tick()
	if r.inv.nodes["cuda01"].State != inventory.Cordoned {
		t.Fatal("llmgw uncordoned a node an operator cordoned")
	}
}

func TestOffAndOut(t *testing.T) {
	r := newRig()
	r.inv.allocs["cuda01"] = 0
	r.set.Lent = []string{"cuda01"}
	r.tick()
	r.set.Enabled = false
	r.tick()
	if m := r.agents.last("cuda01"); m.GetNode() || m.GetController() {
		t.Fatalf("told %v with Slurm off", m)
	}
	if r.inv.nodes["cuda01"].State != inventory.Ready {
		t.Fatal("not llmgw's with Slurm off")
	}

	r = newRig()
	r.inv.allocs["cuda01"] = 0
	r.set.Lent = []string{"cuda01"}
	r.tick()
	r.set.Nodes, r.set.Lent = nil, nil
	r.tick()
	if r.inv.nodes["cuda01"].State != inventory.Ready || r.agents.last("cuda01").GetNode() {
		t.Fatal("a node taken out of slurm.nodes is not llmgw's")
	}
	if !r.agents.last("node06").GetController() {
		t.Fatal("the controller stopped with no nodes")
	}
}

func TestNoControllerYet(t *testing.T) {
	r := newRig()
	delete(r.inv.nodes, "node06")
	r.agents.conn = []string{"cuda01"}
	r.tick()
	if m := r.agents.last("cuda01"); m.GetNode() {
		t.Fatalf("a node was started with no controller to fetch its files from: %v", m)
	}
}

func TestAllNodes(t *testing.T) {
	r := newRig()
	r.set.Nodes, r.set.AllNodes = nil, true
	r.inv.nodes["cuda02"] = &inventory.Node{ID: "cuda02", MgmtAddr: "10.0.0.150"}
	r.agents.conn = append(r.agents.conn, "cuda02")
	r.tick()
	for _, n := range []string{"cuda01", "cuda02", "node06"} {
		if !r.agents.last(n).GetNode() {
			t.Fatalf("%s, a node that joined, does not run slurmd", n)
		}
	}
	if got := len(r.svc.Status().Nodes); got != 3 {
		t.Fatalf("%d nodes in the status", got)
	}
	r.set.Lent = []string{"cuda02"}
	r.tick()
	if p := r.svc.Status().Nodes[1]; p.Node != "cuda02" || p.Phase != Lent {
		t.Fatalf("cuda02 not lent: %+v", p)
	}
}

func TestLendingSaysWhenAReplicaWithNowhereToGoStops(t *testing.T) {
	// Seen on the PoC: a GPU model with no other card waited out the drain's
	// grace, and five minutes of "1 replica(s) still to move off" read as a hang.
	at := time.Date(2026, 11, 14, 19, 38, 34, 0, time.Local)
	n := NodeStatus{Node: "cuda01", Phase: Lending, Replicas: 1, StopsBy: at}
	if d := n.Describe(); !strings.Contains(d, "nowhere else to go is stopped at 19:38:34") {
		t.Errorf("%q", d)
	}
	n.StopsBy = time.Time{}
	if d := n.Describe(); d != "being lent: 1 replica(s) still to move off" {
		t.Errorf("%q", d)
	}
}

func TestPyxisIsLoadedOnlyOnceEveryMachineHasIt(t *testing.T) {
	// "required" in plugstack.conf fails every job on a machine without the
	// plugin; sent early, it would stop Slurm's jobs everywhere it is missing.
	r := newRig()
	r.set.Containers = true
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw})
	r.tick()
	if !r.agents.last("cuda01").GetContainers() || !r.agents.last("node06").GetContainers() {
		t.Fatal("the machines were not told to set containers up")
	}
	if _, ok := r.agents.last("node06").GetFiles()["plugstack.conf"]; ok {
		t.Fatal("pyxis loaded before any machine has it")
	}
	if s := r.svc.Status().Containers; s != "installing" {
		t.Errorf("status %q", s)
	}
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, Containers: "pyxis 0.24.0, enroot 4.2.1"})
	r.tick()
	if _, ok := r.agents.last("node06").GetFiles()["plugstack.conf"]; ok {
		t.Fatal("pyxis loaded with the controller still without it")
	}
	r.report("node06", &llmgwv1.SlurmStatus{Containers: "pyxis 0.24.0, enroot 4.2.1"})
	r.tick()
	if p := r.agents.last("node06").GetFiles()["plugstack.conf"]; !strings.Contains(p, "required /usr/local/lib/slurm/spank_pyxis.so") {
		t.Fatalf("plugstack.conf %q", p)
	}
	if s := r.svc.Status().Containers; s != "on" {
		t.Errorf("status %q", s)
	}
	r.set.Containers = false
	r.tick()
	if _, ok := r.agents.last("node06").GetFiles()["plugstack.conf"]; ok || r.svc.Status().Containers != "" {
		t.Fatal("pyxis still loaded with containers off")
	}
}