llmgw/internal/agent/discovery/mig.go

package discovery

// MIG — Multi-Instance GPU — carves one A100, H100 or later card into up to
// seven instances, each with its own memory, its own share of the
// multiprocessors and its own UUID. To CUDA an instance is a device of its
// own, and to placement it is one too: a card in MIG mode is reported as its
// instances instead of as the whole card, each a GPU of a class named from
// the card's and the instance's profile ("a100-80g-mig-3g.40gb").
//
// llmgw does not create instances. Which split a card gets is an operator's
// decision, made with nvidia-smi mig; discovery reads what is there.
//
// Two facts about MIG shape everything below. A process is given one MIG
// device: CUDA enumerates a single instance however many are made visible.
// And instances have no NVLink or peer-to-peer access between them. A replica
// on MIG therefore runs on exactly one instance, never a tensor-parallel group
// of them, and the catalog refuses a profile that asks for more.

import (
	"context"
	"encoding/xml"
	"fmt"
	"math"
	"os/exec"
	"regexp"
	"slices"
	"strconv"
	"strings"
)

// MIGDevice is what a MIG instance carries beyond any GPU's fields: which
// card it is carved from, and how.
type MIGDevice struct {
	// ParentIndex is the card's nvidia-smi index, which is what its sensors,
	// topology and wear are reported by.
	ParentIndex int    `json:"parent_index"`
	ParentUUID  string `json:"parent_uuid,omitempty"`
	// ParentMemoryTotalMB is the whole card's memory: the class is named from
	// the card, so an instance of an 80 GB card is not mistaken for one of a
	// 40 GB card that happens to have the same profile name.
	ParentMemoryTotalMB int `json:"parent_memory_total_mb"`
	// Device is the instance's number on its card, as nvidia-smi -L lists it.
	Device int `json:"device"`
	// Profile is NVIDIA's name for the slice, like "3g.40gb": three of the
	// card's seven compute slices and forty gigabytes of its memory.
	Profile string `json:"profile"`
	// GPUInstance and ComputeInstance are the instance's ids in nvidia-smi
	// mig -lgi and -lci, which is how nvidia-smi attributes processes to it.
	GPUInstance     int `json:"gpu_instance"`
	ComputeInstance int `json:"compute_instance"`
}

// migIndexBase starts the numbers MIG instances are given.
//
// Placement, the agent protocol and every running shard name a node's devices
// by an integer, which for a whole card is its nvidia-smi index. An instance
// has none of its own, so it is given 1000 + 100 × card + device: card 0's
// instance 2 is 1002, card 3's instance 0 is 1300. Readable, stable while the
// instances are, and never one a whole card could have.
const migIndexBase = 1000

// MIGIndex is the device number placement knows a MIG instance by.
func MIGIndex(card, device int) int { return migIndexBase + 100*card + device }

// IsMIGIndex reports whether a device number names a MIG instance.
func IsMIGIndex(i int) bool { return i >= migIndexBase }

// DeviceLabel is how a device number is shown to people: "3" for a whole
// card, "0/mig2" for card 0's MIG instance 2.
func DeviceLabel(i int) string {
	if !IsMIGIndex(i) {
		return strconv.Itoa(i)
	}
	i -= migIndexBase
	return fmt.Sprintf("%d/mig%d", i/100, i%100)
}

// DeviceLabels is DeviceLabel for each of a list.
func DeviceLabels(is []int) []string {
	out := make([]string, len(is))
	for k, i := range is {
		out[k] = DeviceLabel(i)
	}
	return out
}

// IsMIG reports whether this GPU is a MIG instance rather than a whole card.
func (g GPU) IsMIG() bool { return g.MIG != nil }

// CardIndex is the nvidia-smi index of the physical card this GPU is on: its
// own index for a whole card, its parent's for a MIG instance.
func (g GPU) CardIndex() int {
	if g.MIG != nil {
		return g.MIG.ParentIndex
	}
	return g.Index
}

// migClassMarker joins a card's class to an instance's profile in a MIG
// class: "a100-80g" + "-mig-" + "3g.40gb".
const migClassMarker = "-mig-"

// IsMIGClass reports whether a GPU class names MIG instances.
func IsMIGClass(class string) bool { return strings.Contains(class, migClassMarker) }

// migProfileClean keeps a profile name to what a class may hold: "1g.10gb+me"
// becomes "1g.10gb-me".
var migProfileClean = regexp.MustCompile(`[^a-z0-9.]+`)

// migClass is the class of a MIG instance: the card's class and the profile.
func migClass(parentClass, profile string) string {
	p := migProfileClean.ReplaceAllString(strings.ToLower(strings.TrimSpace(profile)), "-")
	return parentClass + migClassMarker + strings.Trim(p, "-")
}

// MIGProfileMemoryGB reads the memory out of a MIG class or profile:
// 40 from "a100-80g-mig-3g.40gb" or "3g.40gb". Zero when there is none.
func MIGProfileMemoryGB(s string) int {
	if i := strings.LastIndex(s, migClassMarker); i >= 0 {
		s = s[i+len(migClassMarker):]
	}
	m := migProfileMemory.FindStringSubmatch(strings.ToLower(s))
	if m == nil {
		return 0
	}
	gb, _ := strconv.Atoi(m[1])
	return gb
}

var migProfileMemory = regexp.MustCompile(`^\d+g\.(\d+)gb`)

// --- nvidia-smi -L ---

// listedMIG is one MIG instance as nvidia-smi -L lists it.
type listedMIG struct {
	Device  int
	Profile string
	UUID    string
}

var (
	// "GPU 0: NVIDIA A100-SXM4-80GB (UUID: GPU-5d5ba0d6-...)"
	smiListGPU = regexp.MustCompile(`^GPU (\d+):`)
	// "  MIG 3g.40gb     Device  0: (UUID: MIG-c6d4f1ef-...)". Drivers before
	// R470 name the instance MIG-GPU-<card uuid>/<gi>/<ci>, which CUDA of that
	// era takes in the same place, so the UUID is read as whatever it is.
	smiListMIG = regexp.MustCompile(`^\s+MIG\s+(\S+)\s+Device\s+(\d+):\s*\(UUID:\s*([^)\s]+)\)`)
)

// parseSMIList reads nvidia-smi -L into each card's MIG instances, by the
// card's index, in device order. A card with none is absent.
func parseSMIList(out string) map[int][]listedMIG {
	res := map[int][]listedMIG{}
	card := -1
	for _, line := range strings.Split(out, "\n") {
		if m := smiListGPU.FindStringSubmatch(line); m != nil {
			card, _ = strconv.Atoi(m[1])
			continue
		}
		m := smiListMIG.FindStringSubmatch(line)
		if m == nil || card < 0 {
			continue
		}
		dev, _ := strconv.Atoi(m[2])
		res[card] = append(res[card], listedMIG{Device: dev, Profile: m[1], UUID: m[3]})
	}
	for c := range res {
		slices.SortFunc(res[c], func(a, b listedMIG) int { return a.Device - b.Device })
	}
	return res
}

// --- nvidia-smi -q -x ---

// migLog is the part of nvidia-smi's XML report MIG is read from: each card's
// mode, and each instance's ids and memory.
type migLog struct {
	GPUs []migLogGPU `xml:"gpu"`
}

type migLogGPU struct {
	ID      string         `xml:"id,attr"`
	Current string         `xml:"mig_mode>current_mig"`
	Pending string         `xml:"mig_mode>pending_mig"`
	Devices []migLogDevice `xml:"mig_devices>mig_device"`
	// Processes carry the GPU and compute instance they run in on a card
	// in MIG mode, and "N/A" for both otherwise.
	Processes []struct {
		GPUInstance     string `xml:"gpu_instance_id"`
		ComputeInstance string `xml:"compute_instance_id"`
		PID             string `xml:"pid"`
		UsedMB          string `xml:"used_memory"`
	} `xml:"processes>process_info"`
	Used string `xml:"fb_memory_usage>used"`
}

type migLogDevice struct {
	Index           string `xml:"index"`
	GPUInstance     string `xml:"gpu_instance_id"`
	ComputeInstance string `xml:"compute_instance_id"`
	Total           string `xml:"fb_memory_usage>total"`
	Used            string `xml:"fb_memory_usage>used"`
}

// migState is one card's MIG mode and instances from the report.
type migState struct {
	BusID          string
	Enabled        bool
	PendingEnabled bool
	Devices        []migReported
}

// migReported is one instance as the report describes it.
type migReported struct {
	Device, GPUInstance, ComputeInstance int
	TotalMB, UsedMB                      int
}

// enabled reads nvidia-smi's mode words: "Enabled", "Disabled", "N/A".
func enabled(s string) bool { return strings.EqualFold(strings.TrimSpace(s), "Enabled") }

func atoiOr(s string, def int) int {
	n, err := strconv.Atoi(strings.TrimSpace(s))
	if err != nil {
		return def
	}
	return n
}

// parseMIGReport reads each card's MIG state from the XML report, in the
// report's order, which is nvidia-smi's index order.
func parseMIGReport(report []byte) ([]migState, error) {
	var log migLog
	if err := xml.Unmarshal(report, &log); err != nil {
		return nil, fmt.Errorf("reading nvidia-smi's report: %w", err)
	}
	out := make([]migState, len(log.GPUs))
	for i, g := range log.GPUs {
		out[i] = migState{BusID: g.ID, Enabled: enabled(g.Current), PendingEnabled: enabled(g.Pending)}
		for _, d := range g.Devices {
			out[i].Devices = append(out[i].Devices, migReported{
				Device: atoiOr(d.Index, -1), GPUInstance: atoiOr(d.GPUInstance, -1),
				ComputeInstance: atoiOr(d.ComputeInstance, -1),
				TotalMB:         mib(d.Total), UsedMB: mib(d.Used),
			})
		}
		slices.SortFunc(out[i].Devices, func(a, b migReported) int { return a.Device - b.Device })
	}
	return out, nil
}

// --- putting it together ---

// migProblemNoInstances keeps a card in MIG mode with nothing carved from it
// out of placement: CUDA cannot use the whole card while MIG is on, and
// there is no instance to use instead.
const migProblemNoInstances = "MIG is enabled with no instances created: nothing can run on it " +
	"until instances are created (nvidia-smi mig -cgi … -C) or MIG is turned off"

// migProblemUnnamed is a card whose instances exist but could not be named:
// an engine is given an instance by its UUID, and without one it cannot be.
const migProblemUnnamed = "MIG is enabled and has instances, but nvidia-smi -L did not name them, " +
	"so none can be given to an engine"

// expandMIG replaces each card in MIG mode with its instances.
//
// listed is nvidia-smi -L's instances by card, which names them; report is
// the XML report's state per card, which gives their memory and ids. Either
// may be missing: -L alone still says which instances exist, and the report
// alone still says a card is in MIG mode, so a card it says has instances
// that -L did not name is kept as the whole card with the reason, rather
// than placed as a card CUDA will not open.
func expandMIG(cards []GPU, listed map[int][]listedMIG, report []migState) []GPU {
	if len(listed) == 0 && !slices.ContainsFunc(report, func(s migState) bool { return s.Enabled }) {
		return cards
	}
	byBus := map[string]migState{}
	for _, s := range report {
		byBus[normalizeBusID(s.BusID)] = s
	}
	out := make([]GPU, 0, len(cards))
	for pos, c := range cards {
		st, ok := byBus[normalizeBusID(c.PCIBusID)]
		if !ok && pos < len(report) {
			st = report[pos]
		}
		devs := listed[c.Index]
		if !st.Enabled && len(devs) == 0 {
			out = append(out, c)
			continue
		}
		if len(devs) == 0 {
			p := migProblemNoInstances
			if len(st.Devices) > 0 {
				p = migProblemUnnamed
			}
			c.Health.Problems = append(c.Health.Problems, p)
			out = append(out, c)
			continue
		}
		for _, d := range devs {
			var rep migReported
			found := false
			for _, r := range st.Devices {
				if r.Device == d.Device {
					rep, found = r, true
				}
			}
			mig := migInstance(c, d.Device, d.Profile, d.UUID)
			if found {
				mig.MIG.GPUInstance, mig.MIG.ComputeInstance = rep.GPUInstance, rep.ComputeInstance
				mig.MemoryTotalMB, mig.MemoryUsedMB = rep.TotalMB, rep.UsedMB
			}
			if mig.MemoryTotalMB == 0 {
				// The profile says it to the gigabyte, which is close enough
				// to place by when the report could not be read.
				mig.MemoryTotalMB = MIGProfileMemoryGB(d.Profile) * 1024
			}
			assessHealth(&mig)
			out = append(out, mig)
		}
	}
	return out
}

// migInstance builds a MIG instance of card: everything about the card that
// is shared — driver, compute capability, NUMA node, bus, temperature and
// ECC errors — and nothing that is not. Not NVLink, which an instance cannot
// use; not utilisation, which no driver reports per instance; and not the
// card's memory or problems, which the instance's own replace.
func migInstance(card GPU, device int, profile, uuid string) GPU {
	g := GPU{
		Index: MIGIndex(card.Index, device), UUID: uuid,
		Name:              strings.TrimSpace(card.Name + " MIG " + profile),
		ComputeCapability: card.ComputeCapability, NUMANode: card.NUMANode,
		PCIBusID: card.PCIBusID, DriverVersion: card.DriverVersion, CUDAVersion: card.CUDAVersion,
		NICPaths: card.NICPaths, Fabric: card.Fabric,
		Health: GPUHealth{ECCErrors: card.Health.ECCErrors, TemperatureC: card.Health.TemperatureC},
		MIG: &MIGDevice{
			ParentIndex: card.Index, ParentUUID: card.UUID, ParentMemoryTotalMB: card.MemoryTotalMB,
			Device: device, Profile: profile, GPUInstance: -1, ComputeInstance: -1,
		},
	}
	return g
}

// readMIG asks nvidia-smi for what expandMIG needs. A failure of either
// command leaves its half empty, which expandMIG is written to take.
func readMIG(ctx context.Context, binary string) (map[int][]listedMIG, []migState) {
	var listed map[int][]listedMIG
	if out, err := exec.CommandContext(ctx, binary, "-L").Output(); err == nil {
		listed = parseSMIList(string(out))
	}
	var report []migState
	if out, err := exec.CommandContext(ctx, binary, "-q", "-x").Output(); err == nil {
		report, _ = parseMIGReport(out)
	}
	return listed, report
}

// --- NVML ---

// nvmlMIGCard is the part of an NVML device MIG discovery reads, behind an
// interface so the logic is tested without a driver: gpu_nvml.go adapts
// nvml.Device to it.
type nvmlMIGCard interface {
	// MIGMode is nvmlDeviceGetMigMode; ok is false on a card without MIG.
	MIGMode() (current, pending, ok bool)
	// MaxMIGDevices is nvmlDeviceGetMaxMigDeviceCount.
	MaxMIGDevices() int
	// MIGDevice is nvmlDeviceGetMigDeviceHandleByIndex; ok is false for an
	// index with no instance behind it, which is how NVML says so.
	MIGDevice(i int) (nvmlMIGInstance, bool)
}

// nvmlMIGInstance is one MIG device handle.
type nvmlMIGInstance interface {
	UUID() string
	// Name is "NVIDIA A100-SXM4-80GB MIG 3g.40gb" on the drivers that name
	// MIG devices, which is where the profile is read from.
	Name() string
	Memory() (totalMB, usedMB int, ok bool)
	Instances() (gpuInstance, computeInstance int)
	// Slices and MemoryMB are the instance's attributes, which name the
	// profile on a driver whose device name does not.
	Attributes() (slices, memoryMB int, ok bool)
}

var nvmlMIGName = regexp.MustCompile(`MIG\s+(\S+)\s*$`)

// migFromNVML is what one card is to placement under NVML: its instances if
// it is in MIG mode, itself if it is not. isMIG says which.
func migFromNVML(card GPU, dev nvmlMIGCard) (out []GPU, isMIG bool) {
	current, _, ok := dev.MIGMode()
	if !ok || !current {
		return []GPU{card}, false
	}
	n := dev.MaxMIGDevices()
	for i := range n {
		h, ok := dev.MIGDevice(i)
		if !ok {
			continue
		}
		profile := ""
		if m := nvmlMIGName.FindStringSubmatch(h.Name()); m != nil {
			profile = m[1]
		} else if n, mem, ok := h.Attributes(); ok {
			// The convention NVIDIA's names follow: compute slices, then the
			// memory in whole gigabytes.
			profile = fmt.Sprintf("%dg.%dgb", n, int(math.Round(float64(mem)/1024)))
		}
		g := migInstance(card, i, profile, h.UUID())
		g.MIG.GPUInstance, g.MIG.ComputeInstance = h.Instances()
		if total, used, ok := h.Memory(); ok {
			g.MemoryTotalMB, g.MemoryUsedMB = total, used
		} else {
			g.MemoryTotalMB = MIGProfileMemoryGB(profile) * 1024
		}
		assessHealth(&g)
		out = append(out, g)
	}
	if len(out) == 0 {
		card.Health.Problems = append(card.Health.Problems, migProblemNoInstances)
		return []GPU{card}, true
	}
	return out, true
}

// --- what the doctor reports ---

// MIGCard is one card's MIG state, as `llmgw doctor` reports it.
type MIGCard struct {
	Index int    `json:"index"`
	Name  string `json:"name,omitempty"`
	// Enabled is the mode in force; PendingEnabled the mode after the next
	// reset of the card, which differs from it until the card is reset.
	Enabled        bool `json:"enabled"`
	PendingEnabled bool `json:"pending_enabled"`
	// Instances are the MIG devices an engine can be given, as nvidia-smi -L
	// lists them.
	Instances []MIGListed `json:"instances,omitempty"`
	// GPUInstances and ComputeInstances are nvidia-smi mig -lgi and -lci: a
	// GPU instance with no compute instance in it is not a device yet.
	GPUInstances     []MIGInstance `json:"gpu_instances,omitempty"`
	ComputeInstances []MIGInstance `json:"compute_instances,omitempty"`
}

// MIGListed is one instance an engine can be given.
type MIGListed struct {
	Device  int    `json:"device"`
	Profile string `json:"profile"`
	UUID    string `json:"uuid"`
}

// MIGInstance is one row of nvidia-smi mig -lgi or -lci.
type MIGInstance struct {
	GPU     int    `json:"gpu"`
	Profile string `json:"profile"`
	// ProfileID is NVIDIA's number for the profile.
	ProfileID int `json:"profile_id"`
	// GPUInstance is the GPU instance's id; ComputeInstance the compute
	// instance's within it, -1 in a -lgi row.
	GPUInstance     int    `json:"gpu_instance"`
	ComputeInstance int    `json:"compute_instance"`
	Placement       string `json:"placement,omitempty"`
}

var (
	// -lgi: "|   0  MIG 3g.40gb          9        2          4:4     |"
	migLGIRow = regexp.MustCompile(`^\|\s*(\d+)\s+MIG\s+(\S+)\s+(\d+)\s+(\d+)\s+(\d+:\d+)\s*\|`)
	// -lci: "|   0      1       MIG 3g.40gb          2         0          0:3     |"
	migLCIRow = regexp.MustCompile(`^\|\s*(\d+)\s+(\d+)\s+MIG\s+(\S+)\s+(\d+)\s+(\d+)\s+(\d+:\d+)\s*\|`)
)

// ParseMIGInstances reads nvidia-smi mig -lgi or -lci, whichever it is
// given: the two tables differ by the GPU instance column -lci adds. "No GPU
// instances found" and the like are no rows, not an error.
func ParseMIGInstances(out string) []MIGInstance {
	var res []MIGInstance
	for _, line := range strings.Split(out, "\n") {
		line = strings.TrimSpace(line)
		if m := migLCIRow.FindStringSubmatch(line); m != nil {
			res = append(res, MIGInstance{
				GPU: atoiOr(m[1], -1), GPUInstance: atoiOr(m[2], -1), Profile: m[3],
				ProfileID: atoiOr(m[4], -1), ComputeInstance: atoiOr(m[5], -1), Placement: m[6],
			})
			continue
		}
		if m := migLGIRow.FindStringSubmatch(line); m != nil {
			res = append(res, MIGInstance{
				GPU: atoiOr(m[1], -1), Profile: m[2], ProfileID: atoiOr(m[3], -1),
				GPUInstance: atoiOr(m[4], -1), ComputeInstance: -1, Placement: m[5],
			})
		}
	}
	return res
}

// DiscoverMIG reads every card's MIG state for the doctor: the mode from the
// XML report, the devices from -L, the instances from nvidia-smi mig. Only
// cards that support MIG are listed, and an error is returned only when the
// report itself cannot be read.
func DiscoverMIG(ctx context.Context, binary string) ([]MIGCard, error) {
	out, err := exec.CommandContext(ctx, binary, "-q", "-x").Output()
	if err != nil {
		return nil, fmt.Errorf("running %s -q -x: %w", binary, err)
	}
	var list, lgi, lci string
	if b, err := exec.CommandContext(ctx, binary, "-L").Output(); err == nil {
		list = string(b)
	}
	// nvidia-smi mig exits non-zero when there is nothing to list, with the
	// table or the reason on its output either way.
	b, _ := exec.CommandContext(ctx, binary, "mig", "-lgi").CombinedOutput()
	lgi = string(b)
	b, _ = exec.CommandContext(ctx, binary, "mig", "-lci").CombinedOutput()
	lci = string(b)
	return buildMIGCards(out, list, lgi, lci)
}

// buildMIGCards is DiscoverMIG from the commands' output.
func buildMIGCards(report []byte, list, lgi, lci string) ([]MIGCard, error) {
	var log struct {
		GPUs []struct {
			Name    string `xml:"product_name"`
			Current string `xml:"mig_mode>current_mig"`
			Pending string `xml:"mig_mode>pending_mig"`
		} `xml:"gpu"`
	}
	if err := xml.Unmarshal(report, &log); err != nil {
		return nil, fmt.Errorf("reading nvidia-smi's report: %w", err)
	}
	listed := parseSMIList(list)
	gis, cis := ParseMIGInstances(lgi), ParseMIGInstances(lci)
	var out []MIGCard
	for i, g := range log.GPUs {
		cur, pend := strings.TrimSpace(g.Current), strings.TrimSpace(g.Pending)
		// A card without MIG says N/A, or nothing on a driver that predates
		// it; neither is worth a line.
		if !enabled(cur) && !enabled(pend) && !strings.EqualFold(cur, "Disabled") {
			continue
		}
		c := MIGCard{Index: i, Name: strings.TrimSpace(g.Name), Enabled: enabled(cur), PendingEnabled: enabled(pend)}
		for _, l := range listed[i] {
			c.Instances = append(c.Instances, MIGListed(l))
		}
		for _, gi := range gis {
			if gi.GPU == i && gi.ComputeInstance < 0 {
				c.GPUInstances = append(c.GPUInstances, gi)
			}
		}
		for _, ci := range cis {
			if ci.GPU == i && ci.ComputeInstance >= 0 {
				c.ComputeInstances = append(c.ComputeInstances, ci)
			}
		}
		out = append(out, c)
	}
	return out, nil
}