llmgw/internal/control/slurmsvc/failover_test.go

package slurmsvc

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

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

// fakeKV is the store's cluster state.
type fakeKV struct {
	m   map[string]string
	err error
}

func (f *fakeKV) PutClusterState(_ context.Context, k, v string) (bool, error) {
	if f.err != nil {
		return false, f.err
	}
	f.m[k] = v
	return true, nil
}

func (f *fakeKV) ClusterState(_ context.Context, k string) (string, error) {
	if f.err != nil {
		return "", f.err
	}
	v, ok := f.m[k]
	if !ok {
		return "", store.ErrNotFound
	}
	return v, nil
}

// takeOver is the leader that comes next: the same settings and store, an
// inventory with only the nodes connected to it so far.
func takeOver(r *rig, kv *fakeKV, connected ...string) *rig {
	n := &rig{set: r.set, inv: &fakeInv{nodes: map[string]*inventory.Node{}, allocs: map[string]int{}},
		agents: &fakeAgents{conn: connected, got: map[string][]*llmgwv1.Slurm{}}}
	for _, c := range connected {
		node := *r.inv.nodes[c]
		node.State = inventory.Ready
		n.inv.nodes[c] = &node
	}
	n.svc = New(func() config.SlurmSettings { return n.set }, n.inv, n.agents,
		func(context.Context) ([]byte, error) { return []byte(strings.Repeat("k", 1024)), nil }, nil)
	n.svc.SetMemory(StoreMemory(kv))
	return n
}

func TestANewLeaderKeepsTheNodesItHasNotHeardFromInTheControllersFiles(t *testing.T) {
	kv := &fakeKV{m: map[string]string{}}
	r := newRig()
	r.svc.SetMemory(StoreMemory(kv))
	r.set.Lent = []string{"cuda01"}
	r.inv.allocs = map[string]int{}
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "MIXED", Jobs: 1})
	r.tick()
	want := r.agents.last("node06").GetFiles()["slurm.conf"]
	if !strings.Contains(want, "NodeName=cuda01") || !strings.Contains(want, "Gres=gpu:rtx3090:1") {
		t.Fatalf("slurm.conf:\n%s", want)
	}

	// The leader goes; the next hears from the controller first.
	n := takeOver(r, kv, "node06")
	n.tick()
	if got := n.agents.last("node06").GetFiles()["slurm.conf"]; got != want {
		t.Fatalf("the new leader's slurm.conf differs; the controller would drop the node and its job:\n%s\nwas:\n%s",
			got, want)
	}
}

func TestANodeIsToldNothingUntilItsControllerIsKnown(t *testing.T) {
	r := newRig()
	r.svc.SetMemory(StoreMemory(&fakeKV{m: map[string]string{}}))
	delete(r.inv.nodes, "node06")
	r.agents.conn = []string{"cuda01"}
	r.tick()
	if m := r.agents.last("cuda01"); m != nil {
		t.Fatalf("a node was told %v before its controller was known; an empty part stops its slurmd and its jobs", m)
	}
}

func TestANewLeaderTellsANodeItsControllerFromTheStore(t *testing.T) {
	kv := &fakeKV{m: map[string]string{}}
	r := newRig()
	r.svc.SetMemory(StoreMemory(kv))
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "IDLE+DRAIN", Reason: "llmgw: not lent yet"})
	r.tick()

	n := takeOver(r, kv, "cuda01")
	n.tick()
	if m := n.agents.last("cuda01"); !m.GetNode() || m.GetConfServer() != "10.0.0.146:6817" {
		t.Fatalf("the node was told %v", m)
	}
}

func TestWithTheStoreUnreadTheControllerIsSentNothingThatLeavesANodeOut(t *testing.T) {
	kv := &fakeKV{m: map[string]string{}, err: errors.New("store down")}
	r := newRig()
	r.svc.SetMemory(StoreMemory(kv))
	r.tick()
	if m := r.agents.last("node06"); m != nil {
		t.Fatalf("the controller was sent %v without cuda01, with the store unread", m)
	}
	// Once every node has reported, nothing is left out: sent.
	r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw, State: "IDLE+DRAIN", Reason: "llmgw: not lent yet"})
	r.tick()
	if !strings.Contains(r.agents.last("node06").GetFiles()["slurm.conf"], "NodeName=cuda01") {
		t.Fatal("with every node reported, the controller still got no files")
	}
}

func TestWithNothingKeptANewLeaderWaitsForTheNodesAWhile(t *testing.T) {
	r := newRig()
	r.svc.SetMemory(StoreMemory(&fakeKV{m: map[string]string{}}))
	r.tick()
	if m := r.agents.last("node06"); m != nil {
		t.Fatalf("with nothing kept, the controller was sent %v before cuda01 reported", m)
	}
	defer func(d time.Duration) { HoldFor = d }(HoldFor)
	HoldFor = 0
	r.tick()
	if m := r.agents.last("node06"); !m.GetController() || strings.Contains(m.GetFiles()["slurm.conf"], "NodeName=cuda01") {
		t.Fatalf("after the wait, the controller was told %v", m)
	}
}