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.
// The koanf tags are not decoration. Without them the fields bind from their Go
// names, so the documented `name:` key silently produced an empty fabric and
// `rate_gbps` was never read at all — and the agent then refused to start with
// "fabrics[0]: name is required", naming the key the operator had just written.
// Every fabric in this project was configured in a test until the first one was
// configured from a file.
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" koanf:"name"`
	// HCAs are the adapters attached to this fabric.
	HCAs []string `json:"hcas" koanf:"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" koanf:"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"`

	// MemoryTotalMB is the machine's RAM. Zero means the agent did not report
	// it, which is treated as no constraint — the same as every node before
	// this field existed.
	MemoryTotalMB int `json:"memory_total_mb,omitempty"`

	// Weights are the model files this node can see, however they got there.
	// Empty means the node was not asked to look, which is not the same as
	// holding none — so it constrains nothing.
	Weights []discovery.LocalWeights `json:"weights,omitempty"`
	// Storage is the disk those weights sit on, one entry per filesystem.
	//
	// Empty means the node was not asked, not that it has no disk — so an
	// agent too old to report it, or one with no model directories, is left
	// unconstrained rather than treated as full.
	Storage []discovery.Storage `json:"storage,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"`
	// Features are the requests the agent answers beyond the ones every
	// agent does.
	Features       []string          `json:"features,omitempty"`
	EngineVersions map[string]string `json:"engine_versions,omitempty"`
	// EnginesProbed says EngineVersions is an answer, so an empty map means
	// "has none" rather than "an agent too old to say".
	EnginesProbed bool `json:"engines_probed,omitempty"`

	// Host is the operating system and processor, as the agent reads them.
	Host discovery.Host `json:"host"`
	// Provision is where installing each wanted engine stands on this node.
	Provision []EngineInstall `json:"provision,omitempty"`

	// Load is how busy the machine was at the last heartbeat that measured
	// it. Nil before the first one, and from an agent too old to say. Not
	// persisted: it describes a moment, and a restored one would be stale.
	Load *Load `json:"-"`
	// IBPorts are the node's InfiniBand and RoCE ports at the last heartbeat,
	// with their counters. Not persisted, for the same reason as Load.
	IBPorts []IBPort `json:"-"`
	// IBAt is when IBPorts was read, for the rates.
	IBAt time.Time `json:"-"`
	// IMEX is the service a multi-node NVLink domain needs on this node.
	IMEX discovery.IMEX `json:"imex"`
}

// NVLinkDomain is the multi-node NVLink domain this node's GPUs are joined
// to, when every one of them is joined to the same. A node half joined — a
// card still coming up, or one that fell out — is in none: a replica that
// counts on its NVLink would find a card that has none.
func (n Node) NVLinkDomain() string {
	d := ""
	for _, g := range n.GPUs {
		gd := g.Fabric.Domain()
		if gd == "" || (d != "" && gd != d) {
			return ""
		}
		d = gd
	}
	return d
}

// IBPort is one port's state, identity and cumulative counters.
type IBPort struct {
	HCA       string
	Port      int
	State     string
	PhysState string
	RateGbps  int
	LinkLayer string
	Counters  map[string]uint64

	NodeGUID  string
	PortGUID  string
	GIDPrefix string
	// SMGUID and the rest are the master subnet manager this port sees;
	// SMError says why they are empty on an active InfiniBand port.
	SMGUID     string
	SMLID      int
	SMPriority int
	SMState    string
	SMError    string
	// Network is a RoCE port's IP network.
	Network string
	// SMs is every subnet manager on the port's fabric, master and standbys,
	// on the one port per fabric that asked; SMsError says why it could not.
	SMs      []SubnetManager
	SMsError string
	// TxBps and RxBps are bytes a second, from this heartbeat's data
	// counters against the last one's; zero on the first.
	TxBps, RxBps float64
}

// SubnetManager is one subnet manager on a fabric.
type SubnetManager struct {
	GUID     string
	LID      int
	Priority int
	State    string
}

// CanonicalSMs maps every subnet manager GUID the cluster has heard of to
// one name for its fabric: the lowest GUID among the managers listed
// together. A fabric's master changes on a failover and its list does not,
// so ports that see the old master and ports that already see the new one
// stay one fabric through it.
func CanonicalSMs(nodes []Node) map[string]string {
	parent := map[string]string{}
	var find func(string) string
	find = func(x string) string {
		p, ok := parent[x]
		if !ok || p == x {
			parent[x] = x
			return x
		}
		root := find(p)
		parent[x] = root
		return root
	}
	union := func(a, b string) {
		ra, rb := find(a), find(b)
		if ra == rb {
			return
		}
		if ra < rb {
			parent[rb] = ra
		} else {
			parent[ra] = rb
		}
	}
	for _, n := range nodes {
		for _, p := range n.IBPorts {
			if p.SMGUID != "" {
				find(p.SMGUID)
			}
			for _, sm := range p.SMs {
				union(sm.GUID, p.SMGUID)
				if p.SMGUID == "" {
					find(sm.GUID)
				}
			}
		}
	}
	out := make(map[string]string, len(parent))
	for g := range parent {
		out[g] = find(g)
	}
	return out
}

// SMChange is a port seeing a different master subnet manager than before.
type SMChange struct {
	At       time.Time `json:"at"`
	Node     string    `json:"node"`
	Port     string    `json:"port"`
	From, To string
}

// Active reports whether the port can carry traffic.
func (p IBPort) Active() bool { return strings.EqualFold(p.State, "ACTIVE") }

// RoCE reports whether the port is Ethernet carrying RDMA.
func (p IBPort) RoCE() bool { return strings.EqualFold(p.LinkLayer, "Ethernet") }

// FabricKey names the physical fabric the port is on: its master subnet
// manager on InfiniBand, its IP network on RoCE. Empty when not known, and
// for a port that is not up.
func (p IBPort) FabricKey() string { return p.FabricKeyIn(nil) }

// FabricKeyIn is FabricKey with the subnet manager named as canon has it,
// so a fabric keeps one key through a failover.
func (p IBPort) FabricKeyIn(canon map[string]string) string {
	if !p.Active() {
		return ""
	}
	switch {
	case p.RoCE() && p.Network != "":
		return "roce:" + p.Network
	case !p.RoCE() && p.SMGUID != "":
		if c, ok := canon[p.SMGUID]; ok {
			return "sm:" + c
		}
		return "sm:" + p.SMGUID
	}
	return ""
}

// PhysicalFabric is the physical fabric this node's adapters for a declared
// fabric are on. Known is false when none of them says; ok is false when
// they disagree — two adapters declared as one fabric that are cabled to two.
func (n Node) PhysicalFabric(fabric string, canon map[string]string) (key string, known, ok bool) {
	hcas := n.HCAsFor(fabric)
	for _, p := range n.IBPorts {
		if !slices.Contains(hcas, p.HCA) {
			continue
		}
		k := p.FabricKeyIn(canon)
		if k == "" {
			continue
		}
		if key != "" && k != key {
			return "", true, false
		}
		key = k
	}
	return key, key != "", true
}

// Load is a node's use of its CPU and memory, and each engine's share.
type Load struct {
	At            time.Time
	CPUPercent    float64
	CPUCount      int
	Load1         float64
	Load5         float64
	Load15        float64
	MemUsedBytes  int64
	MemAvailBytes int64
	// Shards is each engine's use, by shard ID.
	Shards map[string]ShardLoad
}

// ShardLoad is one engine's use: cores busy, and resident memory.
type ShardLoad struct {
	CPUCores float64
	RSSBytes int64
}

// EngineInstall is where one engine's installation stands on a node.
type EngineInstall struct {
	Engine  string `json:"engine"`
	Version string `json:"version"`
	Variant string `json:"variant,omitempty"`
	// State is "installing", "ready", "failed" or "unsupported".
	State  string `json:"state"`
	Detail string `json:"detail,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
	// MemoryGB is the host RAM this shard's profile asked for.
	//
	// Recorded for the same reason the GPU indices are: it was claimed, and
	// placement has to stop handing it out twice. Without it a node's memory
	// was compared against its total on every decision, so two models each
	// declaring most of a machine were both placed on it — and the second one
	// is what the kernel kills.
	//
	// Zero for a profile that declared nothing, which constrains nothing and
	// is counted as nothing.
	MemoryGB 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
	// smChanges are ports' changes of master subnet manager, newest last.
	smChanges []SMChange
	// scans are the latest ibnetdiscover of each physical fabric, by key.
	scans map[string]FabricScan
}

// 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.
// RecordSMChange keeps a port's change of master subnet manager, the last
// few hundred of them: a failover is the fabric's own event, and nothing
// else would remember it once the ports agree again.
func (s *Store) RecordSMChange(c SMChange) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.smChanges = append(s.smChanges, c)
	if len(s.smChanges) > 500 {
		s.smChanges = slices.Clone(s.smChanges[len(s.smChanges)-500:])
	}
}

// FabricScan is one ibnetdiscover of a physical fabric: its text, and who
// ran it when.
type FabricScan struct {
	Key    string
	At     time.Time
	Node   string
	Port   string
	Output string
	Err    string
	Took   time.Duration
}

// SetScan keeps a fabric's latest scan. A failed one does not replace a
// good one's text: a map from ten minutes ago is worth more than none.
func (s *Store) SetScan(sc FabricScan) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.scans == nil {
		s.scans = map[string]FabricScan{}
	}
	if sc.Err != "" {
		if old, ok := s.scans[sc.Key]; ok && old.Output != "" {
			old.Err = sc.Err
			s.scans[sc.Key] = old
			return
		}
	}
	s.scans[sc.Key] = sc
}

// Scans returns the latest scan of every fabric.
func (s *Store) Scans() map[string]FabricScan {
	s.mu.Lock()
	defer s.mu.Unlock()
	return maps.Clone(s.scans)
}

// SMChanges returns the recorded changes, oldest first.
func (s *Store) SMChanges() []SMChange {
	s.mu.Lock()
	defer s.mu.Unlock()
	return slices.Clone(s.smChanges)
}

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
}

// LocalDisk sums the space on filesystems this node fills itself.
//
// Shared filesystems are left out on purpose. Space on a filer belongs to
// everyone mounting it, and counting it per node would let twenty nodes each
// believe they have the same free terabyte — then place twenty models into it.
// A node whose weights are all on a shared mount therefore reports no local
// disk, which is correct: it has none to spend.
//
// Both zero means nothing is known — an agent too old to report storage, or a
// node with no model directories — and callers must read that as unknown
// rather than as a full disk.
func (n Node) LocalDisk() (free, total int64) {
	for _, s := range n.Storage {
		if s.Shared {
			continue
		}
		free += s.FreeBytes
		total += s.TotalBytes
	}
	return free, total
}

// HasWeights reports whether this node can see the given source.
//
// An exact path match, deliberately. A source that is not a path — a model id
// an engine will resolve itself — cannot be checked from here, and guessing
// would be worse than the honest answer that nothing is known.
func (n Node) HasWeights(source string) bool {
	for _, w := range n.Weights {
		if w.Path == source {
			return true
		}
	}
	return false
}

// KnowsItsWeights reports whether this node was asked to look at all.
//
// The distinction that keeps the feature optional: a node reporting no weights
// because nobody configured a directory must not be excluded from placement,
// while a node that looked and genuinely has none should be.
func (n Node) KnowsItsWeights() bool { return len(n.Weights) > 0 }

// 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))
}