llmgw/internal/control/placement/placement.go

// Package placement decides which nodes and GPUs a replica runs on (plan §6.2).
//
// Two constraints here are not negotiable, and everything else in this package
// exists to keep them true:
//
//   - a multi-node replica never spans two fabrics
//   - a tensor-parallel group never mixes GPU classes
//
// Violating either produces a replica that starts, passes health checks, and
// then hangs or crawls under load — the worst possible failure mode, because it
// looks like success until traffic arrives.
package placement

import (
	"fmt"
	"slices"
	"sort"
	"strings"

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

// Profile is the subset of a catalog profile placement needs.
type Profile struct {
	Name     string
	Engine   string
	GPUClass string
	// GPUsPerNode is how many GPUs each node contributes.
	GPUsPerNode int
	// Nodes is how many nodes the replica spans. Above 1, a fabric is required.
	Nodes int
	// Priority orders profiles; the highest that fits is used.
	Priority int
	// NodeSelector restricts this profile to nodes carrying these labels.
	NodeSelector map[string]string
	// MinMemoryGB is the host RAM this profile needs. Zero is unconstrained.
	MinMemoryGB int
	// MinDiskGB is the free space needed to hold the weights, checked only on
	// nodes that do not already have them. Zero is unconstrained.
	MinDiskGB int
	// Source is the model's weights, as the engine would be given it. Used to
	// prefer nodes that already hold them; empty asks for no preference.
	Source string
}

// CPUOnly reports whether this profile needs no GPUs. The catalog only lets a
// profile omit its GPU class when its engine can serve without one.
func (p Profile) CPUOnly() bool { return p.GPUClass == "" }

// classLabel names the profile's hardware for pools and messages.
func (p Profile) classLabel() string {
	if p.CPUOnly() {
		return "cpu"
	}
	return p.GPUClass
}

// normalize fills in the defaults a catalog entry may omit.
func (p Profile) normalize() Profile {
	if p.Nodes <= 0 {
		p.Nodes = 1
	}
	switch {
	case p.CPUOnly():
		p.GPUsPerNode = 0
	case p.GPUsPerNode <= 0:
		p.GPUsPerNode = 1
	}
	return p
}

// NodeAssignment is one node's share of a replica.
type NodeAssignment struct {
	NodeID string
	Rank   int
	// GPUIndices are the local GPUs the shard uses.
	GPUIndices []int
	// HCAs are the adapters of the chosen fabric only.
	HCAs []string
}

// Placement is a decision: which profile, which fabric, which nodes.
type Placement struct {
	Profile Profile
	// Fabric is empty for a single-node replica.
	Fabric string
	// Pool names the (fabric, gpu class) pair this landed in.
	Pool  string
	Nodes []NodeAssignment
}

// Rejection records why a candidate was not used. Placement failures are
// otherwise impossible to debug: "unschedulable" tells an operator nothing
// about which of six constraints bit.
type Rejection struct {
	Profile string
	Pool    string
	Node    string
	Reason  string
}

func (r Rejection) String() string {
	where := r.Pool
	if r.Node != "" {
		where += "/" + r.Node
	}
	if where == "" {
		where = "-"
	}
	return fmt.Sprintf("%s @ %s: %s", r.Profile, where, r.Reason)
}

// Result is the outcome of a placement attempt.
type Result struct {
	Placement *Placement
	// Rejections explains every candidate that was considered and discarded,
	// in the order they were tried. This is what `ctl explain-placement` prints.
	Rejections []Rejection
}

// Explain renders the rejections as text.
func (r Result) Explain() string {
	if len(r.Rejections) == 0 {
		return "no candidates were considered"
	}
	var b strings.Builder
	for _, rej := range r.Rejections {
		b.WriteString("  ")
		b.WriteString(rej.String())
		b.WriteString("\n")
	}
	return b.String()
}

// Scheduler places replicas onto an inventory.
type Scheduler struct {
	inv *inventory.Store
}

// New builds a scheduler over an inventory.
func New(inv *inventory.Store) *Scheduler { return &Scheduler{inv: inv} }

// Place picks the first profile, in priority order, that fits.
//
// avoid lists nodes already hosting this model, so replicas of one model spread
// across failure domains instead of stacking on one machine.
func (s *Scheduler) Place(profiles []Profile, avoid []string) Result {
	var res Result

	ordered := slices.Clone(profiles)
	sort.SliceStable(ordered, func(i, j int) bool { return ordered[i].Priority > ordered[j].Priority })

	for _, raw := range ordered {
		p := raw.normalize()
		if p.Nodes == 1 {
			if pl, rej := s.placeSingleNode(p, avoid); pl != nil {
				res.Placement = pl
				return res
			} else {
				res.Rejections = append(res.Rejections, rej...)
			}
			continue
		}
		if pl, rej := s.placeMultiNode(p, avoid); pl != nil {
			res.Placement = pl
			return res
		} else {
			res.Rejections = append(res.Rejections, rej...)
		}
	}
	return res
}

// placeSingleNode finds one node with enough free GPUs of the right class.
//
// Nodes without a fabric are eligible here and only here: a single-node replica
// never talks to another node, so it has nothing to be reachable over.
func (s *Scheduler) placeSingleNode(p Profile, avoid []string) (*Placement, []Rejection) {
	var rejections []Rejection

	candidates := s.eligibleNodes(p, "", avoid, &rejections)
	if len(candidates) == 0 {
		return nil, rejections
	}

	// Bin packing: the node with the fewest free GPUs that still fits. Filling
	// partly used nodes first keeps whole nodes free for the large models that
	// cannot be split.
	sort.SliceStable(candidates, func(i, j int) bool {
		return len(candidates[i].free) < len(candidates[j].free)
	})

	if p.CPUOnly() {
		// Nothing to bin pack, so spread instead: two CPU replicas on one
		// machine compete for the same cores, and unlike GPUs there is no
		// allocation to stop them.
		sort.SliceStable(candidates, func(i, j int) bool {
			return len(s.inv.Allocations(candidates[i].node.ID)) <
				len(s.inv.Allocations(candidates[j].node.ID))
		})
		preferNodesHoldingWeights(p.Source, candidates)
		return &Placement{
			Profile: p,
			Pool:    "none/cpu",
			Nodes:   []NodeAssignment{{NodeID: candidates[0].node.ID, Rank: 0}},
		}, rejections
	}

	// Last, so it takes precedence: moving tens of gigabytes costs more than an
	// imperfectly packed node, and with isolated fabrics it may cost a
	// cross-fabric copy over the management network.
	preferNodesHoldingWeights(p.Source, candidates)

	best := candidates[0]
	gpus := s.pickGPUs(best.node, best.free, p.GPUsPerNode, nil)
	if gpus == nil {
		rejections = append(rejections, Rejection{
			Profile: p.Name, Node: best.node.ID,
			Reason: fmt.Sprintf("has %d free %s GPUs but no usable set of %d", len(best.free), p.GPUClass, p.GPUsPerNode),
		})
		return nil, rejections
	}

	return &Placement{
		Profile: p,
		Pool:    "none/" + p.classLabel(),
		Nodes:   []NodeAssignment{{NodeID: best.node.ID, Rank: 0, GPUIndices: gpus}},
	}, rejections
}

// placeMultiNode finds N nodes on a single fabric.
//
// The fabric loop is the whole point: candidates are gathered per fabric and a
// replica is only formed from within one of them. There is deliberately no path
// that combines nodes from different fabrics, however well they would otherwise
// fit.
func (s *Scheduler) placeMultiNode(p Profile, avoid []string) (*Placement, []Rejection) {
	var rejections []Rejection

	fabrics := s.inv.Fabrics()
	if len(fabrics) == 0 {
		return nil, []Rejection{{
			Profile: p.Name,
			Reason:  fmt.Sprintf("needs %d nodes on one fabric, but no node reports a fabric", p.Nodes),
		}}
	}

	for _, fabric := range fabrics {
		pool := fabric + "/" + p.classLabel()
		candidates := s.eligibleNodes(p, fabric, avoid, &rejections)

		if len(candidates) < p.Nodes {
			rejections = append(rejections, Rejection{
				Profile: p.Name, Pool: pool,
				Reason: fmt.Sprintf("needs %d nodes, only %d qualify", p.Nodes, len(candidates)),
			})
			continue
		}

		// Prefer the emptiest nodes for multi-node work: a replica spanning
		// nodes is the least movable thing in the cluster, so give it the
		// cleanest machines rather than squeezing it into the gaps.
		sort.SliceStable(candidates, func(i, j int) bool {
			return len(candidates[i].free) > len(candidates[j].free)
		})

		assignments := make([]NodeAssignment, 0, p.Nodes)
		for rank, c := range candidates[:p.Nodes] {
			gpus := s.pickGPUs(c.node, c.free, p.GPUsPerNode, c.node.HCAsFor(fabric))
			if gpus == nil {
				assignments = nil
				rejections = append(rejections, Rejection{
					Profile: p.Name, Pool: pool, Node: c.node.ID,
					Reason: fmt.Sprintf("no usable set of %d GPUs", p.GPUsPerNode),
				})
				break
			}
			assignments = append(assignments, NodeAssignment{
				NodeID:     c.node.ID,
				Rank:       rank,
				GPUIndices: gpus,
				// Only this fabric's adapters, never the node's others.
				HCAs: c.node.HCAsFor(fabric),
			})
		}
		if assignments == nil {
			continue
		}

		return &Placement{Profile: p, Fabric: fabric, Pool: pool, Nodes: assignments}, rejections
	}
	return nil, rejections
}

// candidate is a node that passed the hard constraints.
type candidate struct {
	node inventory.Node
	free []discovery.GPU
}

// eligibleNodes applies the hard constraints and records why each node failed.
//
// fabric is empty for single-node placement, in which case attachment is not
// checked at all.
// preferNodesHoldingWeights moves nodes that already have the model to the
// front, leaving everything else in the order it was given.
//
// A stable partition rather than a sort, so every other consideration — bin
// packing, spreading, the emptiest machine for a multi-node replica — survives
// intact among nodes that are equal on this one.
func preferNodesHoldingWeights(source string, candidates []candidate) {
	if source == "" {
		return
	}
	sort.SliceStable(candidates, func(i, j int) bool {
		a, aKnown := hasWeights(source, candidates[i].node)
		b, bKnown := hasWeights(source, candidates[j].node)
		// A node that was never asked to look ranks with those that looked and
		// have nothing: unknown is not a reason to prefer or to punish.
		return (a && aKnown) && !(b && bKnown)
	})
}

// hasWeights reports whether a node already holds this profile's model, and
// whether that is known at all.
//
// Never a filter, always a preference. Storage differs at every installation
// and llmgw does not own it: a node that was never asked to look reports
// nothing, and excluding those would break every cluster that has not opted
// in. Even a node that looked and does not have the weights stays eligible,
// because "fetch them" is a valid answer and refusing to place would turn a
// slow start into no start at all.
//
// What it changes is the ordering, and that is worth a lot on its own: with
// isolated fabrics the difference between placing where the weights are and
// placing where they are not is a cross-fabric copy of tens of gigabytes over
// the management network.
func hasWeights(source string, n inventory.Node) (present, known bool) {
	if source == "" || !n.KnowsItsWeights() {
		return false, false
	}
	return n.HasWeights(source), true
}

// unmatchedLabel reports the first selector entry a node does not satisfy.
//
// It returns the label rather than a bare false so the rejection can name it:
// "does not carry ram=32g" is a placement failure someone can act on, and
// "no candidates were considered" is not.
func unmatchedLabel(selector, labels map[string]string) (string, bool) {
	for k, want := range selector {
		if got, ok := labels[k]; !ok || got != want {
			return k + "=" + want, false
		}
	}
	return "", true
}

func (s *Scheduler) eligibleNodes(p Profile, fabric string, avoid []string, rejections *[]Rejection) []candidate {
	pool := "none"
	if fabric != "" {
		pool = fabric
	}
	pool += "/" + p.classLabel()

	var out []candidate
	for _, n := range s.inv.Nodes() {
		reject := func(reason string) {
			*rejections = append(*rejections, Rejection{Profile: p.Name, Pool: pool, Node: n.ID, Reason: reason})
		}

		if !n.State.Schedulable() {
			reject("node is " + n.State.String())
			continue
		}
		if slices.Contains(avoid, n.ID) {
			reject("already hosts this model; spreading across failure domains")
			continue
		}
		if fabric != "" {
			if !n.HasFabric(fabric) {
				// Not worth reporting for every node on every fabric; it is the
				// expected case, not a problem.
				continue
			}
			if len(n.HCAsFor(fabric)) == 0 {
				reject("attached to " + fabric + " but reports no adapters for it")
				continue
			}
		}

		if miss, ok := unmatchedLabel(p.NodeSelector, n.Labels); !ok {
			reject("does not carry " + miss)
			continue
		}
		// The engine has to be installed here. The field has been on this
		// struct since placement was written and was never read, so a node
		// with no vLLM was as good a candidate as one with it: the replica was
		// placed, the node registered healthy, and the shard failed minutes
		// later with "executable file not found" in the engine's own words.
		//
		// A node that reports no engines at all is not excluded. That is what
		// every agent older than this reports, and emptying a fleet on a
		// control plane upgrade is a worse failure than the one being fixed.
		//
		// An agent that says it probed is believed when it reports nothing:
		// that is a node still installing, or one where installing failed,
		// and either way not a place to start an engine. sim is built into
		// the agent and never installed.
		if p.Engine != "" && p.Engine != "sim" && (n.EnginesProbed || len(n.EngineVersions) > 0) {
			if _, ok := n.EngineVersions[p.Engine]; !ok {
				reject("does not have " + p.Engine + " installed" + installing(n, p.Engine))
				continue
			}
		}
		// A node that does not report its memory is not excluded: an agent
		// older than that field would otherwise have its node quietly emptied
		// by a control plane upgrade, which is a worse failure than the one
		// this prevents.
		//
		// What is left, not what the machine has. Comparing against the total
		// meant a node with 31 GB accepted two models that each declared 15,
		// because each was asked about on its own — and the second one is what
		// the kernel kills. GPUs were never handed out twice this way; memory
		// was, because nothing recorded that it had been claimed.
		if p.MinMemoryGB > 0 && n.MemoryTotalMB > 0 {
			total := int(n.MemoryTotalMB / 1024)
			free := total - s.claimedMemory(n.ID)
			if free < p.MinMemoryGB {
				reject(fmt.Sprintf("has %d GB of memory with %d free, needs %d",
					total, free, p.MinMemoryGB))
				continue
			}
		}
		// Room for the weights, and only for a node that would have to fetch
		// them. A node already holding the file needs no space for it.
		//
		// Same two exemptions as memory, for the same reason: zero means the
		// operator declared nothing, and a node reporting no storage is an
		// older agent or one with no model directories, neither of which is a
		// full disk. Without that second exemption a control plane upgrade
		// would empty every node in a fleet that had not upgraded yet.
		if p.MinDiskGB > 0 && p.Source != "" && !n.HasWeights(p.Source) {
			if free, total := n.LocalDisk(); total > 0 {
				if have := int(free / (1 << 30)); have < p.MinDiskGB {
					reject(fmt.Sprintf("has %d GB free for weights, needs %d", have, p.MinDiskGB))
					continue
				}
			}
		}

		if p.CPUOnly() {
			// Every ready node can run it; there is no card to be free.
			out = append(out, candidate{node: n})
			continue
		}

		freeIdx := s.inv.FreeGPUs(n.ID)
		var free []discovery.GPU
		for _, g := range n.GPUs {
			if g.Class() != p.GPUClass || !slices.Contains(freeIdx, g.Index) {
				continue
			}
			free = append(free, g)
		}
		if len(free) < p.GPUsPerNode {
			reject(fmt.Sprintf("has %d free %s GPUs, needs %d", len(free), p.GPUClass, p.GPUsPerNode))
			continue
		}
		out = append(out, candidate{node: n, free: free})
	}
	return out
}

// pickGPUs chooses which of a node's free GPUs to use.
//
// Every set of the right size is scored, which a node's handful of cards
// makes cheap, in this order:
//
//   - NVLink: the most pairs in the set joined by it. A tensor-parallel group
//     that falls back to PCIe for its all-reduces runs several times slower,
//     and "seed and its peers", the rule before, took cards that were each
//     joined to the seed but not to each other.
//   - For a multi-node shard, the fabric: cards near one of its adapters (PIX
//     or PXB rather than across the socket), since every transfer between
//     nodes goes through one.
//   - One NUMA node rather than two.
//   - The lowest indices, so the answer is the same every time.
//
// hcas are the fabric's adapters for a multi-node shard, nil otherwise.
// Without topology — an agent that reported none — every set scores the same
// and it is the lowest indices, as it always was.
func (s *Scheduler) pickGPUs(n inventory.Node, free []discovery.GPU, want int, hcas []string) []int {
	if len(free) < want {
		return nil
	}
	if want == 1 && len(hcas) == 0 {
		return []int{free[0].Index}
	}
	sorted := slices.Clone(free)
	slices.SortFunc(sorted, func(a, b discovery.GPU) int { return a.Index - b.Index })

	// Too many to score every set: take the lowest, as before topology.
	if len(sorted) > 16 {
		out := make([]int, 0, want)
		for _, g := range sorted[:want] {
			out = append(out, g.Index)
		}
		return out
	}

	var best []int
	bestScore := [3]int{-1, 0, 0}
	pick := make([]int, 0, want)
	var walk func(from int)
	walk = func(from int) {
		if len(pick) == want {
			sc := gpuSetScore(sorted, pick, hcas)
			if best == nil || better(sc, bestScore) {
				bestScore = sc
				best = slices.Clone(pick)
			}
			return
		}
		for i := from; i <= len(sorted)-(want-len(pick)); i++ {
			pick = append(pick, i)
			walk(i + 1)
			pick = pick[:len(pick)-1]
		}
	}
	walk(0)

	out := make([]int, len(best))
	for i, k := range best {
		out[i] = sorted[k].Index
	}
	return out
}

// gpuSetScore rates a set, as positions in gpus: NVLink pairs, nearness to
// the fabric (negated, so higher is better) and NUMA nodes spanned (negated).
func gpuSetScore(gpus []discovery.GPU, set []int, hcas []string) [3]int {
	var sc [3]int
	numa := map[int]bool{}
	for a, i := range set {
		g := gpus[i]
		numa[g.NUMANode] = true
		for _, j := range set[a+1:] {
			if slices.Contains(g.NVLinkPeers, gpus[j].Index) {
				sc[0]++
			}
		}
		if len(hcas) > 0 && len(g.NICPaths) > 0 {
			nearest := 5
			for _, h := range hcas {
				if rel, ok := g.NICPaths[h]; ok {
					nearest = min(nearest, discovery.PathRank(rel))
				}
			}
			sc[1] -= nearest
		}
	}
	sc[2] = -len(numa)
	return sc
}

// better reports whether a scores above b, most important first. Equal is
// not better, so the first set found — the lowest indices — keeps a tie.
func better(a, b [3]int) bool {
	for i := range a {
		if a[i] != b[i] {
			return a[i] > b[i]
		}
	}
	return false
}

// claimedMemory is what the shards already on a node asked for.
//
// Declared, not measured. A profile that says it needs 15 GB is treated as
// holding 15 whether the process has touched that much yet or not, which is the
// only figure available when the decision is made — the alternative is placing
// against what a model has grown into so far and being wrong every time one is
// still loading.
func (s *Scheduler) claimedMemory(nodeID string) int {
	total := 0
	for _, a := range s.inv.Allocations(nodeID) {
		total += a.MemoryGB
	}
	return total
}

// installing says why an engine is missing, when the node has said.
func installing(n inventory.Node, engine string) string {
	for _, p := range n.Provision {
		if p.Engine == engine && p.Detail != "" {
			return " (" + p.State + ": " + p.Detail + ")"
		}
	}
	return ""
}