llmgw/internal/control/placement/topology.go

package placement

import (
	"slices"
	"sync"
	"time"

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

// fabricHints is what a fabric's map says about placing on it: the leaf
// switch each node's first adapter hangs off, and which nodes have a cable
// that trained below the rest.
type fabricHints struct {
	leaf     map[string]string
	degraded map[string]bool
}

// hintCache keeps a parsed map until its fabric is scanned again; parsing
// thousands of lines on every placement would be the scheduler's main cost.
var hintCache struct {
	sync.Mutex
	at    map[string]time.Time
	hints map[string]fabricHints
}

// hints reads every scanned fabric. Nodes on no mapped fabric are absent,
// which leaves them placed as before.
func hints(inv *inventory.Store) fabricHints {
	out := fabricHints{leaf: map[string]string{}, degraded: map[string]bool{}}
	scans := inv.Scans()
	if len(scans) == 0 {
		return out
	}
	var hosts map[string]topology.Host
	hintCache.Lock()
	defer hintCache.Unlock()
	if hintCache.at == nil {
		hintCache.at, hintCache.hints = map[string]time.Time{}, map[string]fabricHints{}
	}
	for key, sc := range scans {
		if sc.Output == "" {
			continue
		}
		h, ok := hintCache.hints[key]
		if !ok || !hintCache.at[key].Equal(sc.At) {
			if hosts == nil {
				hosts = map[string]topology.Host{}
				for _, n := range inv.Nodes() {
					for _, p := range n.IBPorts {
						if p.NodeGUID != "" {
							hosts[p.NodeGUID] = topology.Host{Node: n.ID, HCA: p.HCA}
						}
					}
				}
			}
			f, err := topology.Parse(sc.Output)
			if err != nil {
				continue
			}
			v := topology.Build(f, hosts)
			h = fabricHints{leaf: map[string]string{}, degraded: map[string]bool{}}
			guidNode := map[string]string{}
			for _, hv := range v.Hosts {
				if hv.Node == "" {
					continue
				}
				guidNode[hv.GUID] = hv.Node
				if hv.Rail == 0 {
					h.leaf[hv.Node] = hv.Leaf
				}
			}
			for _, l := range v.Links {
				if l.Host && l.Degraded {
					if n := guidNode[l.To]; n != "" {
						h.degraded[n] = true
					}
				}
			}
			hintCache.hints[key], hintCache.at[key] = h, sc.At
		}
		for n, l := range h.leaf {
			out.leaf[n] = l
		}
		for n := range h.degraded {
			out.degraded[n] = true
		}
	}
	return out
}

// byLeaf orders candidates for a replica spanning nodes: those whose cables
// all trained at their rate first, and among them the nodes of one leaf
// together — the leaf with the most of them first — so a replica's traffic
// stays under one switch when it can, and crosses the spines only when it
// must. Within a leaf the order given (emptiest first) is kept.
func byLeaf(cands []candidate, h fabricHints) []candidate {
	if len(h.leaf) == 0 && len(h.degraded) == 0 {
		return cands
	}
	count := map[string]int{}
	for _, c := range cands {
		if !h.degraded[c.node.ID] {
			count[h.leaf[c.node.ID]]++
		}
	}
	firstSeen := map[string]int{}
	for i, c := range cands {
		if _, ok := firstSeen[h.leaf[c.node.ID]]; !ok {
			firstSeen[h.leaf[c.node.ID]] = i
		}
	}
	out := make([]candidate, len(cands))
	copy(out, cands)
	rank := func(c candidate) (bad int, size int, first int) {
		if h.degraded[c.node.ID] {
			bad = 1
		}
		l := h.leaf[c.node.ID]
		return bad, count[l], firstSeen[l]
	}
	slices.SortStableFunc(out, func(a, b candidate) int {
		ba, sa, fa := rank(a)
		bb, sb, fb := rank(b)
		if ba != bb {
			return ba - bb
		}
		if sa != sb {
			return sb - sa
		}
		return fa - fb
	})
	return out
}