Slurm: a new leader keeps the nodes it has not heard from yet

Tried on the PoC: with a job on a lent node, the leader restarted and its
first slurm.conf had no node in it (none had reported to it yet), so
slurmctld dropped the node and the job ended in NODE_FAIL.

The leader keeps slurm.conf's nodes and the controller's address in the
store's cluster state (slurm-nodes); the next one starts from them and
writes the same files until the nodes report again. With the store
unread, the controller gets no files that leave a node out. A node is
told nothing, rather than that it has no part, until its controller's
address is known: no part stops slurmd, and its jobs with it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

internal/cli/server.go

@@ -450,6 +450,8 @@ func newServerCmd() *cobra.Command {
 								}
 								return slurm.MungeKey(secret)
 							}, log)
+						// What slurm.conf held, for the leader that comes next.
+						slurmSvc.SetMemory(slurmsvc.StoreMemory(st))
 						agentSvc.SetSlurm(slurmSvc)
 						go slurmSvc.Run(ctx)
 						api.SetSlurm(slurmSvc)

internal/control/slurmsvc/failover_test.go

@@ -0,0 +1,117 @@
+package slurmsvc
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"testing"
+
+	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")
+	}
+}

internal/control/slurmsvc/memory.go

@@ -0,0 +1,77 @@
+package slurmsvc
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"slices"
+
+	"gitea.mynodes.xyz/ilker/llmgw/internal/slurm"
+	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
+)
+
+// Memory keeps the nodes slurm.conf last held, and the controller's
+// address, where a leader that takes over finds them. Without it a new
+// leader's first slurm.conf has none of the nodes it has not heard from
+// yet, and the controller, restarted on it, drops them and the jobs
+// running there.
+type Memory interface {
+	Recall(ctx context.Context) (Kept, error)
+	Keep(ctx context.Context, k Kept) error
+}
+
+// Kept is what Memory holds.
+type Kept struct {
+	Controller     string                `json:"controller,omitempty"`
+	ControllerAddr string                `json:"controller_addr,omitempty"`
+	Nodes          map[string]slurm.Node `json:"nodes,omitempty"`
+}
+
+func (k Kept) equal(o Kept) bool {
+	if k.Controller != o.Controller || k.ControllerAddr != o.ControllerAddr || len(k.Nodes) != len(o.Nodes) {
+		return false
+	}
+	for name, n := range k.Nodes {
+		m, ok := o.Nodes[name]
+		if !ok || n.Addr != m.Addr || n.Hardware != m.Hardware || !slices.Equal(n.GPUs, m.GPUs) {
+			return false
+		}
+	}
+	return true
+}
+
+// KV is the store's cluster state.
+type KV interface {
+	PutClusterState(ctx context.Context, key, value string) (bool, error)
+	ClusterState(ctx context.Context, key string) (string, error)
+}
+
+// keptKey is the cluster state key it is kept under.
+const keptKey = "slurm-nodes"
+
+// StoreMemory keeps it in the store's cluster state.
+func StoreMemory(kv KV) Memory { return kvMemory{kv} }
+
+type kvMemory struct{ kv KV }
+
+func (m kvMemory) Recall(ctx context.Context) (Kept, error) {
+	var k Kept
+	v, err := m.kv.ClusterState(ctx, keptKey)
+	if errors.Is(err, store.ErrNotFound) {
+		return k, nil
+	}
+	if err != nil {
+		return k, err
+	}
+	err = json.Unmarshal([]byte(v), &k)
+	return k, err
+}
+
+func (m kvMemory) Keep(ctx context.Context, k Kept) error {
+	b, err := json.Marshal(k)
+	if err != nil {
+		return err
+	}
+	_, err = m.kv.PutClusterState(ctx, keptKey, string(b))
+	return err
+}

internal/control/slurmsvc/service.go

@@ -115,6 +115,14 @@ type Service struct {
 	sent    map[string]*llmgwv1.Slurm
 	mungeKy []byte
 	kick    chan struct{}
+	// mem keeps slurm.conf's nodes and the controller's address for the
+	// next leader; known is what it held when read, and what was kept
+	// since. loaded is false until it was read: until then the controller
+	// is sent no files that leave out a node, which would drop that node
+	// and its jobs.
+	mem    Memory
+	known  Kept
+	loaded bool
 }
 
 // New is a Service; key gives the munge key, made once from a cluster
@@ -129,6 +137,13 @@ func New(settings func() config.SlurmSettings, inv Inventory, agents Agents,
 		kick: make(chan struct{}, 1)}
 }
 
+// SetMemory has slurm.conf's nodes kept in mem, and read from it first.
+func (s *Service) SetMemory(mem Memory) {
+	s.mu.Lock()
+	s.mem = mem
+	s.mu.Unlock()
+}
+
 // current is the settings, slurm.all_nodes resolved to the nodes the
 // inventory has.
 func (s *Service) current() config.SlurmSettings {
@@ -200,6 +215,7 @@ func (s *Service) Tick(ctx context.Context) {
 	set := s.current()
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	s.recall(ctx)
 	phases := s.handOver(set)
 	var key []byte
 	if set.Enabled {
@@ -213,9 +229,18 @@ func (s *Service) Tick(ctx context.Context) {
 		}
 		key = s.mungeKy
 	}
-	files, _ := s.files(set)
+	files, waiting := s.files(set)
+	// Not yet knowing what the last leader knew, a slurm.conf without some
+	// node would take it, and its jobs, from the controller: kept back.
+	hold := s.mem != nil && !s.loaded && len(waiting) > 0
 	for _, node := range s.agents.ConnectedNodes() {
+		if hold && node == set.Controller {
+			continue
+		}
 		m := s.message(set, node, phases[node] == Lent, key, files)
+		if m == nil {
+			continue
+		}
 		if old, ok := s.sent[node]; ok && proto.Equal(old, m) {
 			continue
 		}
@@ -225,6 +250,54 @@ func (s *Service) Tick(ctx context.Context) {
 		}
 		s.sent[node] = m
 	}
+	s.remember(ctx, set)
+}
+
+// recall reads what the last leader kept, once. Under s.mu.
+func (s *Service) recall(ctx context.Context) {
+	if s.mem == nil || s.loaded {
+		return
+	}
+	k, err := s.mem.Recall(ctx)
+	if err != nil {
+		s.log.Warn("could not read Slurm's nodes from the store; the controller's files wait for every node",
+			"err", err)
+		return
+	}
+	s.known, s.loaded = k, true
+}
+
+// remember keeps slurm.conf's nodes and the controller's address when they
+// changed. Under s.mu.
+func (s *Service) remember(ctx context.Context, set config.SlurmSettings) {
+	if s.mem == nil || !s.loaded || !set.Enabled {
+		return
+	}
+	c, _ := s.cluster(set)
+	now := Kept{Controller: set.Controller, ControllerAddr: c.ControllerAddr, Nodes: map[string]slurm.Node{}}
+	for _, n := range c.Nodes {
+		now.Nodes[n.Name] = n
+	}
+	if now.equal(s.known) {
+		return
+	}
+	if err := s.mem.Keep(ctx, now); err != nil {
+		s.log.Warn("could not keep Slurm's nodes in the store", "err", err)
+		return
+	}
+	s.known = now
+}
+
+// controllerAddr is where the nodes reach slurmctld: the controller's
+// management address, else the one kept. Under s.mu.
+func (s *Service) controllerAddr(set config.SlurmSettings) string {
+	if n, ok := s.inv.Get(set.Controller); ok && n.MgmtAddr != "" {
+		return n.MgmtAddr
+	}
+	if s.known.Controller == set.Controller {
+		return s.known.ControllerAddr
+	}
+	return ""
 }
 
 // handOver moves each node of slurm.nodes a step between llmgw and Slurm,
@@ -314,34 +387,52 @@ func (s *Service) ready(node string) {
 // files are the controller's files, and the nodes left out of them for not
 // having reported their hardware. Under s.mu.
 func (s *Service) files(set config.SlurmSettings) (map[string]string, []string) {
-	ctl, ok := s.inv.Get(set.Controller)
-	if !set.Enabled || !ok {
-		return nil, nil
-	}
-	c := slurm.Cluster{Name: set.Cluster(), Controller: set.Controller, ControllerAddr: ctl.MgmtAddr,
-		Partition: set.PartitionName()}
-	var waiting []string
-	for _, node := range set.Nodes {
-		n, ok := s.inv.Get(node)
-		st := s.status[node]
-		if !ok || st == nil || st.Hardware == "" {
-			waiting = append(waiting, node)
-			continue
-		}
-		hw, err := slurm.ParseHardware(st.Hardware)
-		if err != nil {
-			waiting = append(waiting, node)
-			continue
-		}
-		sn := slurm.Node{Name: node, Addr: n.MgmtAddr, Hardware: hw}
-		for _, g := range n.GPUs {
-			sn.GPUs = append(sn.GPUs, slurm.GPU{Minor: g.Index, Name: g.Name})
-		}
-		c.Nodes = append(c.Nodes, sn)
+	c, waiting := s.cluster(set)
+	if c.ControllerAddr == "" {
+		return nil, waiting
 	}
 	return c.Files(), waiting
 }
 
+// cluster is what slurm.conf says: each node as it last reported, else as
+// kept from before, and the nodes left out for having neither. Under s.mu.
+func (s *Service) cluster(set config.SlurmSettings) (slurm.Cluster, []string) {
+	c := slurm.Cluster{Name: set.Cluster(), Controller: set.Controller, Partition: set.PartitionName()}
+	if !set.Enabled {
+		return c, nil
+	}
+	c.ControllerAddr = s.controllerAddr(set)
+	var waiting []string
+	for _, node := range set.Nodes {
+		if sn, ok := s.reported(node); ok {
+			c.Nodes = append(c.Nodes, sn)
+		} else if sn, ok := s.known.Nodes[node]; ok {
+			c.Nodes = append(c.Nodes, sn)
+		} else {
+			waiting = append(waiting, node)
+		}
+	}
+	return c, waiting
+}
+
+// reported is a node as its agent and the inventory say now. Under s.mu.
+func (s *Service) reported(node string) (slurm.Node, bool) {
+	n, ok := s.inv.Get(node)
+	st := s.status[node]
+	if !ok || n.MgmtAddr == "" || st == nil || st.Hardware == "" {
+		return slurm.Node{}, false
+	}
+	hw, err := slurm.ParseHardware(st.Hardware)
+	if err != nil {
+		return slurm.Node{}, false
+	}
+	sn := slurm.Node{Name: node, Addr: n.MgmtAddr, Hardware: hw}
+	for _, g := range n.GPUs {
+		sn.GPUs = append(sn.GPUs, slurm.GPU{Minor: g.Index, Name: g.Name})
+	}
+	return sn, true
+}
+
 // message is what a node is told. Under s.mu.
 func (s *Service) message(set config.SlurmSettings, node string, lent bool, key []byte,
 	files map[string]string) *llmgwv1.Slurm {
@@ -356,21 +447,19 @@ func (s *Service) message(set config.SlurmSettings, node string, lent bool, key
 	}
 	m.Name, m.MungeKey, m.Lent = node, key, lent
 	m.Node = isNode
-	if ctl, ok := s.inv.Get(set.Controller); ok {
-		m.ConfServer = net.JoinHostPort(ctl.MgmtAddr, strconv.Itoa(slurm.CtldPort))
+	if addr := s.controllerAddr(set); addr != "" {
+		m.ConfServer = net.JoinHostPort(addr, strconv.Itoa(slurm.CtldPort))
 	}
 	if isNode && m.ConfServer == "" {
 		// Without its controller a node cannot start; it is told when the
-		// controller registers.
-		m.Node = false
+		// controller registers. Told nothing until then, not "no part": a
+		// node still running Slurm's jobs would stop them.
+		return nil
 	}
 	if isCtl {
 		m.Controller = true
 		m.Files = maps.Clone(files)
 	}
-	if !m.Node && !m.Controller {
-		return &llmgwv1.Slurm{}
-	}
 	return m
 }