llmgw/internal/control/nvlink/nvlink.go

// Package nvlink works out the cluster's multi-node NVLink domains — an NVL72
// rack, an NVLink Switch system — from what each GPU reports, and what keeps
// one from being usable.
//
// A domain is a scale-up island: its GPUs reach each other over NVLink
// whatever node they are in, so a tensor-parallel group can span its nodes
// at NVLink speed. What it needs to be one is every card joined by the
// fabric manager and IMEX running on every node; either missing, and a
// replica placed across it fails or crawls.
package nvlink

import (
	"fmt"
	"sort"
	"strings"

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

// Level is how much a problem matters.
type Level string

// The levels.
const (
	Fail Level = "fail"
	Warn Level = "warn"
)

// Problem is one thing wrong with a domain.
type Problem struct {
	Level   Level  `json:"level"`
	Domain  string `json:"domain"`
	Message string `json:"message"`
}

// GPU is one card as its domain sees it.
type GPU struct {
	Index     int    `json:"index"`
	Name      string `json:"name"`
	State     string `json:"state"`
	Status    string `json:"status,omitempty"`
	Bandwidth string `json:"bandwidth,omitempty"`
	Joined    bool   `json:"joined"`
}

// Member is one node of a domain.
type Member struct {
	Node string `json:"node"`
	GPUs []GPU  `json:"gpus"`
	// Usable is every card joined and IMEX up: what placement requires.
	Usable       bool   `json:"usable"`
	IMEX         string `json:"imex"`
	IMEXChannels int    `json:"imex_channels"`
}

// Domain is one multi-node NVLink domain.
type Domain struct {
	Key         string   `json:"key"`
	ClusterUUID string   `json:"cluster_uuid"`
	CliqueID    string   `json:"clique_id"`
	Members     []Member `json:"members"`
	GPUs        int      `json:"gpus"`
	Joined      int      `json:"joined"`
	Usable      int      `json:"usable_nodes"`
}

// Report is every domain and what is wrong with them.
type Report struct {
	Domains  []Domain  `json:"domains"`
	Problems []Problem `json:"problems"`
}

// Analyze reads every node's GPUs.
func Analyze(nodes []inventory.Node) Report {
	var r Report
	by := map[string]*Domain{}
	for _, n := range nodes {
		// The domain a node is in is the one any of its cards joined; a card
		// still joining has no domain of its own yet, and is counted against
		// the node's.
		key := ""
		for _, g := range n.GPUs {
			if d := g.Fabric.Domain(); d != "" {
				key = d
				break
			}
		}
		if key == "" {
			continue
		}
		d := by[key]
		if d == nil {
			uuid, clique, _ := strings.Cut(key, "/")
			d = &Domain{Key: key, ClusterUUID: uuid, CliqueID: clique}
			by[key] = d
		}
		m := Member{Node: n.ID, IMEX: n.IMEX.State, IMEXChannels: n.IMEX.Channels}
		joined := 0
		for _, g := range n.GPUs {
			v := GPU{Index: g.Index, Name: g.Name}
			if f := g.Fabric; f != nil {
				v.State, v.Status, v.Bandwidth = f.State, f.Status, f.Bandwidth
				v.Joined = f.Domain() == key
			}
			if v.Joined {
				joined++
			} else {
				why := v.State
				if why == "" {
					why = "no fabric"
				}
				if v.Status != "" && !strings.EqualFold(v.Status, "Success") {
					why += ": " + v.Status
				}
				r.Problems = append(r.Problems, Problem{Level: Fail, Domain: key, Message: fmt.Sprintf(
					"%s GPU %d is not joined to its domain (%s)", n.ID, g.Index, why)})
			}
			if v.Bandwidth != "" && !strings.EqualFold(v.Bandwidth, "Full") {
				r.Problems = append(r.Problems, Problem{Level: Warn, Domain: key, Message: fmt.Sprintf(
					"%s GPU %d has %s NVLink bandwidth", n.ID, g.Index, strings.ToLower(v.Bandwidth))})
			}
			m.GPUs = append(m.GPUs, v)
		}
		imexOK := n.IMEX.State == "active" && n.IMEX.Channels > 0
		switch {
		case n.IMEX.State != "active":
			state := n.IMEX.State
			if state == "" {
				state = "not installed"
			}
			r.Problems = append(r.Problems, Problem{Level: Fail, Domain: key, Message: fmt.Sprintf(
				"%s: nvidia-imex is %s — without it the domain's GPUs cannot share memory across nodes",
				n.ID, state)})
		case n.IMEX.Channels == 0:
			r.Problems = append(r.Problems, Problem{Level: Fail, Domain: key, Message: fmt.Sprintf(
				"%s: nvidia-imex runs but there is no IMEX channel for an engine to open (/dev/nvidia-caps-imex-channels)", n.ID)})
		}
		m.Usable = joined == len(n.GPUs) && joined > 0 && imexOK
		if joined > 0 && joined < len(n.GPUs) {
			r.Problems = append(r.Problems, Problem{Level: Fail, Domain: key, Message: fmt.Sprintf(
				"%s is half joined (%d of %d GPUs), so no replica spans the domain through it", n.ID, joined, len(n.GPUs))})
		}
		d.GPUs += len(n.GPUs)
		d.Joined += joined
		if m.Usable {
			d.Usable++
		}
		d.Members = append(d.Members, m)
	}
	for _, d := range by {
		sort.Slice(d.Members, func(i, j int) bool { return d.Members[i].Node < d.Members[j].Node })
		r.Domains = append(r.Domains, *d)
	}
	sort.Slice(r.Domains, func(i, j int) bool { return r.Domains[i].Key < r.Domains[j].Key })
	sort.SliceStable(r.Problems, func(i, j int) bool {
		if r.Problems[i].Level != r.Problems[j].Level {
			return r.Problems[i].Level == Fail
		}
		return r.Problems[i].Message < r.Problems[j].Message
	})
	return r
}

// Short is a domain's key as a person reads it: the cluster's first group
// and the clique.
func Short(key string) string {
	uuid, clique, _ := strings.Cut(key, "/")
	if i := strings.IndexByte(uuid, '-'); i > 0 {
		uuid = uuid[:i]
	}
	return uuid + "·" + clique
}