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
}

// 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))
		})
		return &Placement{
			Profile: p,
			Pool:    "none/cpu",
			Nodes:   []NodeAssignment{{NodeID: candidates[0].node.ID, Rank: 0}},
		}, rejections
	}

	best := candidates[0]
	gpus := s.pickGPUs(best.node, best.free, p.GPUsPerNode)
	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)
			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.
// 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
		}
		// 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.
		if p.MinMemoryGB > 0 && n.MemoryTotalMB > 0 {
			if have := n.MemoryTotalMB / 1024; have < p.MinMemoryGB {
				reject(fmt.Sprintf("has %d GB of memory, needs %d", have, p.MinMemoryGB))
				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.
//
// It prefers a set that shares an NVLink domain: within a node, a
// tensor-parallel group that has to fall back to PCIe for its all-reduces runs
// several times slower. NVLink topology is only available from the NVML build,
// so without it this degrades to taking the lowest indices — which is usually
// right, but only by accident of how vendors number slots.
func (s *Scheduler) pickGPUs(n inventory.Node, free []discovery.GPU, want int) []int {
	if len(free) < want {
		return nil
	}
	if want == 1 {
		return []int{free[0].Index}
	}

	freeSet := map[int]struct{}{}
	for _, g := range free {
		freeSet[g.Index] = struct{}{}
	}

	// Try each free GPU as a seed and grow its NVLink-connected set.
	for _, seed := range free {
		group := []int{seed.Index}
		for _, peer := range seed.NVLinkPeers {
			if _, ok := freeSet[peer]; !ok || peer == seed.Index {
				continue
			}
			group = append(group, peer)
			if len(group) == want {
				slices.Sort(group)
				return group
			}
		}
	}

	// No NVLink domain large enough, or no topology at all.
	out := make([]int, 0, want)
	for _, g := range free[:want] {
		out = append(out, g.Index)
	}
	slices.Sort(out)
	return out
}