Nodes say where they are with llmgw.io/site, room, row, rack and pdu labels, from agent.yaml or `llmgw ctl label` (kept in the store, over the agent's). Placement spreads a model's replicas across racks and PDUs, keeps a multi-node replica within one domain when its profile says keep_within, places nothing in a domain held with `llmgw ctl facility hold` (optionally draining it), and refuses a shard that would take a domain past its power cap from facility.budgets, counting declared, derived and measured draw. Explain-placement names the domain, and keeps the rejections of a placed replica; `ctl nodes` gains LOCATION and --domain; `ctl facility` lists the domains; the console's Nodes view shows and sets location and holds. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
795 lines
24 KiB
Go
795 lines
24 KiB
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 (
|
|
"fmt"
|
|
"maps"
|
|
"slices"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/facility"
|
|
)
|
|
|
|
// 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 are what placement matches: the agent's own, from agent.yaml,
|
|
// with any an operator set through the control plane over them.
|
|
Labels map[string]string `json:"labels,omitempty"`
|
|
// AgentLabels are the labels the agent reported, kept apart so an
|
|
// operator's label removed brings the agent's value back. Nil on a node
|
|
// built anywhere but a registration; Upsert then takes Labels as the
|
|
// agent's.
|
|
AgentLabels map[string]string `json:"agent_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"`
|
|
// CPUCount is the CPUs the agent may run on. Zero is an agent too old to
|
|
// say, which, like memory, constrains nothing.
|
|
CPUCount int `json:"cpu_count,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
|
|
// What the node's sensors say; zero is not measured.
|
|
CPUTempC float64
|
|
CPUPowerW float64
|
|
InletTempC float64
|
|
GPUs []GPUSensor
|
|
}
|
|
|
|
// GPUSensor is one card's heat and power, and what is throttling it.
|
|
type GPUSensor struct {
|
|
Index int `json:"index"`
|
|
TempC float64 `json:"temp_c,omitempty"`
|
|
MemoryTempC float64 `json:"memory_temp_c,omitempty"`
|
|
PowerW float64 `json:"power_w,omitempty"`
|
|
PowerLimitW float64 `json:"power_limit_w,omitempty"`
|
|
Throttle []string `json:"throttle,omitempty"`
|
|
// How busy the card is, as of this sample; -1 is not reported. Zero
|
|
// MemoryTotalMB is an agent too old to send these, whose figures are
|
|
// the ones it read at registration.
|
|
UtilPercent float64 `json:"util_percent,omitempty"`
|
|
MemUtilPercent float64 `json:"mem_util_percent,omitempty"`
|
|
MemoryUsedMB float64 `json:"memory_used_mb,omitempty"`
|
|
MemoryTotalMB float64 `json:"memory_total_mb,omitempty"`
|
|
SMClockMHz float64 `json:"sm_clock_mhz,omitempty"`
|
|
ECCUncorrected float64 `json:"ecc_uncorrected,omitempty"`
|
|
// Wear is what wears the card out, read every minute; nil from an agent
|
|
// that does not read it.
|
|
Wear *GPUWear `json:"wear,omitempty"`
|
|
}
|
|
|
|
// GPUWear is a card's wear and failures. Counts of -1 are not reported for
|
|
// it: a card before Ampere has no remapped rows, one after no retired pages.
|
|
type GPUWear struct {
|
|
PCIeReplays int64 `json:"pcie_replays"`
|
|
RetiredSBE int64 `json:"retired_pages_sbe"`
|
|
RetiredDBE int64 `json:"retired_pages_dbe"`
|
|
RetirePending bool `json:"retire_pending,omitempty"`
|
|
RemappedCorrectable int64 `json:"remapped_rows_correctable"`
|
|
RemappedUncorrectable int64 `json:"remapped_rows_uncorrectable"`
|
|
RemapPending bool `json:"remap_pending,omitempty"`
|
|
RemapFailure bool `json:"remap_failure,omitempty"`
|
|
NVLinkCRC int64 `json:"nvlink_crc_errors"`
|
|
NVLinkReplay int64 `json:"nvlink_replay_errors"`
|
|
NVLinkRecovery int64 `json:"nvlink_recovery_errors"`
|
|
XIDs int64 `json:"xid_errors"`
|
|
LastXID int `json:"last_xid,omitempty"`
|
|
LastXIDText string `json:"last_xid_text,omitempty"`
|
|
LastXIDAt time.Time `json:"last_xid_at,omitempty"`
|
|
}
|
|
|
|
// Problems says, in a few words each, what is wrong with the card; empty
|
|
// when nothing is.
|
|
func (w *GPUWear) Problems() []string {
|
|
if w == nil {
|
|
return nil
|
|
}
|
|
var out []string
|
|
if w.RemapFailure {
|
|
out = append(out, "out of spare memory rows: replace the card")
|
|
}
|
|
if w.RemapPending {
|
|
out = append(out, "memory rows remapped, waiting for a reset")
|
|
}
|
|
if w.RetirePending {
|
|
out = append(out, "memory pages waiting for a reset to be retired")
|
|
}
|
|
if w.RemappedUncorrectable > 0 {
|
|
out = append(out, fmt.Sprintf("%d rows remapped for uncorrectable errors", w.RemappedUncorrectable))
|
|
}
|
|
if w.RetiredDBE > 0 {
|
|
out = append(out, fmt.Sprintf("%d pages retired for double-bit errors", w.RetiredDBE))
|
|
}
|
|
if w.XIDs > 0 {
|
|
out = append(out, fmt.Sprintf("XID %d (%s)", w.LastXID, w.LastXIDText))
|
|
}
|
|
if w.NVLinkCRC > 0 || w.NVLinkReplay > 0 || w.NVLinkRecovery > 0 {
|
|
out = append(out, fmt.Sprintf("NVLink errors: %d CRC, %d replay, %d recovery", w.NVLinkCRC, w.NVLinkReplay, w.NVLinkRecovery))
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Busy reports whether the sample carries how busy the card is.
|
|
func (s GPUSensor) Busy() bool { return s.MemoryTotalMB > 0 }
|
|
|
|
// 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
|
|
// CPUCores is the cores this shard was given to compute on, for a shard
|
|
// with no GPU; zero for one that shares whatever is free.
|
|
CPUCores int
|
|
// PowerW is the draw this shard's profile declared; zero derives it from
|
|
// the node's cards and the cores, which is what most profiles leave to
|
|
// placement.
|
|
PowerW float64
|
|
}
|
|
|
|
// 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
|
|
// facility is the holds and power budgets placement reads.
|
|
facility facility.Policy
|
|
// operatorLabels are labels set through the control plane, by node,
|
|
// laid over what each agent reports.
|
|
operatorLabels map[string]map[string]string
|
|
}
|
|
|
|
// NewStore returns an empty inventory.
|
|
func NewStore() *Store {
|
|
return &Store{nodes: map[string]*Node{}, allocs: map[string][]Allocation{}}
|
|
}
|
|
|
|
// Clone is a copy to try decisions on without making them: allocations made
|
|
// on it are its own. Nodes are copied by value; the slices inside them are
|
|
// shared, and a clone must not write through them.
|
|
func (s *Store) Clone() *Store {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
c := NewStore()
|
|
for id, n := range s.nodes {
|
|
cp := *n
|
|
c.nodes[id] = &cp
|
|
}
|
|
for id, list := range s.allocs {
|
|
c.allocs[id] = slices.Clone(list)
|
|
}
|
|
c.facility = s.facility.Clone()
|
|
return c
|
|
}
|
|
|
|
// 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
|
|
}
|
|
}
|
|
// A registration carries the agent's labels in Labels; a heartbeat hands
|
|
// back the node it read, whose Labels already have the operator's in
|
|
// them, and its AgentLabels set. Never nil after this, so the second case
|
|
// is always told from the first.
|
|
if n.AgentLabels == nil {
|
|
n.AgentLabels = maps.Clone(n.Labels)
|
|
if n.AgentLabels == nil {
|
|
n.AgentLabels = map[string]string{}
|
|
}
|
|
}
|
|
n.Labels = mergeLabels(n.AgentLabels, s.operatorLabels[n.ID])
|
|
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))
|
|
}
|