llmgw/internal/control/inventory/inventory.go

// Package inventory holds the control plane's view of the cluster: which nodes
// exist, what hardware they have, and what is already running on them.
package inventory

import (
	"maps"
	"slices"
	"strings"
	"sync"
	"time"

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

// NodeState is where a node stands.
type NodeState int

const (
	// Ready means the node accepts placement.
	Ready NodeState = iota
	// Cordoned means an operator has excluded it from new placement; what runs
	// there keeps running.
	Cordoned
	// Draining means its replicas are being moved off.
	Draining
	// Unreachable means heartbeats stopped but the node may come back.
	Unreachable
	// Lost means it has been unreachable long enough to declare its replicas dead.
	Lost
)

func (s NodeState) String() string {
	switch s {
	case Ready:
		return "ready"
	case Cordoned:
		return "cordoned"
	case Draining:
		return "draining"
	case Unreachable:
		return "unreachable"
	case Lost:
		return "lost"
	default:
		return "unknown"
	}
}

// Schedulable reports whether new replicas may be placed on the node.
func (s NodeState) Schedulable() bool { return s == Ready }

// FabricAttachment is a node's connection to one fabric.
type FabricAttachment struct {
	// Fabric is the operator-assigned name. It is not derived from the GID
	// prefix: OpenSM's default is fe80:: everywhere, so a matching prefix
	// proves nothing (plan §5.2).
	Fabric string `json:"fabric"`
	// HCAs are the adapters attached to this fabric.
	HCAs []string `json:"hcas"`
	// RateGbps is the slowest active port's rate, since a collective runs at
	// the speed of its slowest link.
	RateGbps int `json:"rate_gbps"`
}

// Node is one machine in the cluster.
type Node struct {
	ID       string            `json:"id"`
	Hostname string            `json:"hostname"`
	MgmtAddr string            `json:"mgmt_addr"`
	Labels   map[string]string `json:"labels,omitempty"`

	GPUs []discovery.GPU `json:"gpus"`
	// Fabrics is empty on a node without IB. Such a node is its own
	// single-node island and can never host a multi-node replica.
	Fabrics []FabricAttachment `json:"fabrics,omitempty"`

	State    NodeState `json:"-"`
	LastSeen time.Time `json:"last_seen"`
	// AgentVersion and EngineVersions are reported at registration. The engine
	// version decides which metric names and CLI flags apply (plan §14).
	AgentVersion   string            `json:"agent_version,omitempty"`
	EngineVersions map[string]string `json:"engine_versions,omitempty"`
}

// HasFabric reports whether the node is attached to a fabric.
func (n Node) HasFabric(name string) bool {
	return slices.ContainsFunc(n.Fabrics, func(f FabricAttachment) bool { return f.Fabric == name })
}

// FabricNames lists the node's fabrics, sorted.
func (n Node) FabricNames() []string {
	out := make([]string, 0, len(n.Fabrics))
	for _, f := range n.Fabrics {
		out = append(out, f.Fabric)
	}
	slices.Sort(out)
	return out
}

// HCAsFor returns the adapters attached to one fabric.
//
// A shard is given only these: handing it adapters from a second fabric lets
// NCCL choose one whose peer is unreachable, and the job hangs instead of
// failing (plan §5.3).
func (n Node) HCAsFor(fabric string) []string {
	for _, f := range n.Fabrics {
		if f.Fabric == fabric {
			return slices.Clone(f.HCAs)
		}
	}
	return nil
}

// GPUsByClass groups the node's GPUs by class.
func (n Node) GPUsByClass() map[string][]discovery.GPU {
	out := map[string][]discovery.GPU{}
	for _, g := range n.GPUs {
		out[g.Class()] = append(out[g.Class()], g)
	}
	return out
}

// Allocation records that some of a node's GPUs are in use.
type Allocation struct {
	ShardID    string
	ReplicaID  string
	GPUIndices []int
}

// Store holds the cluster inventory. It is safe for concurrent use.
type Store struct {
	mu    sync.RWMutex
	nodes map[string]*Node
	// allocs maps node id to the shards placed there.
	allocs map[string][]Allocation
}

// NewStore returns an empty inventory.
func NewStore() *Store {
	return &Store{nodes: map[string]*Node{}, allocs: map[string][]Allocation{}}
}

// Upsert records a node's registration, preserving operator-set state such as
// a cordon: an agent reconnecting must not silently undo an operator's decision.
func (s *Store) Upsert(n Node) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if prev, ok := s.nodes[n.ID]; ok {
		if prev.State == Cordoned || prev.State == Draining {
			n.State = prev.State
		}
	}
	n.LastSeen = time.Now()
	s.nodes[n.ID] = &n
}

// SetState changes a node's state.
func (s *Store) SetState(id string, st NodeState) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	n, ok := s.nodes[id]
	if !ok {
		return false
	}
	n.State = st
	return true
}

// Touch records a heartbeat, bringing an unreachable node back.
func (s *Store) Touch(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	n, ok := s.nodes[id]
	if !ok {
		return
	}
	n.LastSeen = time.Now()
	if n.State == Unreachable || n.State == Lost {
		n.State = Ready
	}
}

// Get returns a copy of a node.
func (s *Store) Get(id string) (Node, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	n, ok := s.nodes[id]
	if !ok {
		return Node{}, false
	}
	return *n, true
}

// Nodes returns every node, sorted by id.
func (s *Store) Nodes() []Node {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := make([]Node, 0, len(s.nodes))
	for _, n := range s.nodes {
		out = append(out, *n)
	}
	slices.SortFunc(out, func(a, b Node) int { return strings.Compare(a.ID, b.ID) })
	return out
}

// Remove forgets a node and its allocations.
func (s *Store) Remove(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.nodes, id)
	delete(s.allocs, id)
}

// Allocate records GPUs as taken on a node.
func (s *Store) Allocate(nodeID string, a Allocation) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.allocs[nodeID] = append(s.allocs[nodeID], a)
}

// Release frees the GPUs a shard held.
func (s *Store) Release(shardID string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for node, list := range s.allocs {
		kept := slices.DeleteFunc(slices.Clone(list), func(a Allocation) bool {
			return a.ShardID == shardID
		})
		if len(kept) == 0 {
			delete(s.allocs, node)
		} else {
			s.allocs[node] = kept
		}
	}
}

// ReleaseReplica frees every shard of a replica.
//
// Replicas are freed whole because they fail whole: one dead shard makes the
// others useless, and the control plane rebuilds the replica rather than
// restarting a piece of it (plan §5.4).
func (s *Store) ReleaseReplica(replicaID string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for node, list := range s.allocs {
		kept := slices.DeleteFunc(slices.Clone(list), func(a Allocation) bool {
			return a.ReplicaID == replicaID
		})
		if len(kept) == 0 {
			delete(s.allocs, node)
		} else {
			s.allocs[node] = kept
		}
	}
}

// Allocations returns a node's current allocations.
func (s *Store) Allocations(nodeID string) []Allocation {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return slices.Clone(s.allocs[nodeID])
}

// FreeGPUs returns the indices of a node's unallocated GPUs, sorted.
func (s *Store) FreeGPUs(nodeID string) []int {
	s.mu.RLock()
	defer s.mu.RUnlock()

	n, ok := s.nodes[nodeID]
	if !ok {
		return nil
	}
	taken := map[int]struct{}{}
	for _, a := range s.allocs[nodeID] {
		for _, i := range a.GPUIndices {
			taken[i] = struct{}{}
		}
	}

	var out []int
	for _, g := range n.GPUs {
		if _, used := taken[g.Index]; used {
			continue
		}
		// A GPU with an ECC fault or another process on it is not free in any
		// useful sense.
		if !g.Health.OK() {
			continue
		}
		out = append(out, g.Index)
	}
	slices.Sort(out)
	return out
}

// Fabrics lists every fabric any node is attached to, sorted.
func (s *Store) Fabrics() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	seen := map[string]struct{}{}
	for _, n := range s.nodes {
		for _, f := range n.Fabrics {
			seen[f.Fabric] = struct{}{}
		}
	}
	return slices.Sorted(maps.Keys(seen))
}