llmgw/internal/agent/discovery/gpuhealth.go

package discovery

import (
	"context"
	"encoding/xml"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

// GPUWear is what wears a card out or takes it down, beyond temperature:
// read every minute, from nvidia-smi and the kernel's log. The early signs of
// a card failing are here — a growing count of rows its memory had to remap,
// PCIe link retries, XID errors — long before a replica on it fails.
//
// A count of -1 is not reported: a card without row remapping (before
// Ampere) has no such count, which is not the same as none.
type GPUWear struct {
	Index int
	// PCIeReplays are link-level retries on the card's PCIe link: a few are
	// noise, a count that keeps climbing is a riser, slot or cable.
	PCIeReplays int64
	// RetiredSBE and RetiredDBE are memory pages retired for single and
	// double bit errors (cards before Ampere); RetirePending is a page
	// waiting for a reset to be retired.
	RetiredSBE, RetiredDBE int64
	RetirePending          bool
	// RemappedCorrectable and RemappedUncorrectable are memory rows remapped
	// (Ampere and later). RemapPending waits for a reset; RemapFailure is a
	// card out of spare rows — it needs replacing.
	RemappedCorrectable, RemappedUncorrectable int64
	RemapPending, RemapFailure                 bool
	// NVLink errors summed over the card's links, by kind.
	NVLinkCRC, NVLinkReplay, NVLinkRecovery int64
	// XIDs are the driver's error events for this card since the agent
	// started, and the last of them.
	XIDs        int64
	LastXID     int
	LastXIDText string
	LastXIDAt   time.Time
}

// healthEvery is how often health is read: it changes over hours, and the
// kernel log and nvidia-smi's full report are the heavier reads.
const healthEvery = time.Minute

// HealthSampler reads GPUWear, keeping the XID counts between reads.
type HealthSampler struct {
	run func(ctx context.Context, name string, args ...string) ([]byte, error)
	now func() time.Time

	mu      sync.Mutex
	at      time.Time
	last    []GPUWear
	since   time.Time
	xids    map[int]*GPUWear
	started bool
}

// NewHealthSampler reads the machine's own nvidia-smi and kernel log.
func NewHealthSampler() *HealthSampler {
	return &HealthSampler{now: time.Now, xids: map[int]*GPUWear{},
		run: func(ctx context.Context, name string, args ...string) ([]byte, error) {
			return exec.CommandContext(ctx, name, args...).Output() //nolint:gosec // fixed tools
		}}
}

// Sample is the cards' health, read again once healthEvery has passed. Nil
// on a machine without nvidia-smi.
func (h *HealthSampler) Sample(ctx context.Context) []GPUWear {
	h.mu.Lock()
	defer h.mu.Unlock()
	now := h.now()
	if !h.at.IsZero() && now.Sub(h.at) < healthEvery {
		return h.last
	}
	h.at = now
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	report, err := h.run(ctx, "nvidia-smi", "-q", "-x")
	if err != nil {
		h.last = nil
		return nil
	}
	cards, bus := parseHealthReport(report)
	if out, err := h.run(ctx, "nvidia-smi", "nvlink", "-e"); err == nil {
		for i, e := range parseNVLinkErrors(string(out)) {
			if i < len(cards) {
				cards[i].NVLinkCRC, cards[i].NVLinkReplay, cards[i].NVLinkRecovery = e[0], e[1], e[2]
			}
		}
	}
	h.readXIDs(ctx, now, bus)
	for i := range cards {
		if x := h.xids[i]; x != nil {
			cards[i].XIDs, cards[i].LastXID, cards[i].LastXIDText, cards[i].LastXIDAt = x.XIDs, x.LastXID, x.LastXIDText, x.LastXIDAt
		}
	}
	h.last = cards
	return cards
}

// readXIDs adds the XID events logged since the last read. The first read
// starts from now: what happened before the agent started was before, and
// counting the whole log would raise every alert on every restart.
func (h *HealthSampler) readXIDs(ctx context.Context, now time.Time, bus map[string]int) {
	if !h.started {
		h.started, h.since = true, now
		return
	}
	since := h.since
	h.since = now
	out, err := h.run(ctx, "journalctl", "-k", "--no-pager", "-o", "short-unix",
		"--since", "@"+strconv.FormatInt(since.Unix(), 10), "-g", "NVRM: Xid")
	if err != nil {
		// No journal, or nothing matched (journalctl exits 1 for that).
		return
	}
	for _, e := range parseXIDs(string(out)) {
		i, ok := bus[e.bus]
		if !ok {
			continue
		}
		x := h.xids[i]
		if x == nil {
			x = &GPUWear{Index: i}
			h.xids[i] = x
		}
		x.XIDs++
		x.LastXID, x.LastXIDText = e.code, e.text
		x.LastXIDAt = now
		if e.at > 0 {
			x.LastXIDAt = time.Unix(int64(e.at), 0)
		}
	}
}

// healthLog is the part of nvidia-smi's XML report health is read from.
type healthLog struct {
	GPUs []struct {
		ID  string `xml:"id,attr"`
		PCI struct {
			Replay string `xml:"replay_counter"`
		} `xml:"pci"`
		Retired struct {
			SBE     string `xml:"multiple_single_bit_retirement>retired_count"`
			DBE     string `xml:"double_bit_retirement>retired_count"`
			Pending string `xml:"pending_retirement"`
			// Older drivers name it pending_blacklist.
			Blacklist string `xml:"pending_blacklist"`
		} `xml:"retired_pages"`
		Remapped struct {
			Corr    string `xml:"remapped_row_corr"`
			Unc     string `xml:"remapped_row_unc"`
			Pending string `xml:"remapped_row_pending"`
			Failure string `xml:"remapped_row_failure"`
		} `xml:"remapped_rows"`
	} `xml:"gpu"`
}

// parseHealthReport reads each card's health, in index order, and maps each
// card's PCI bus ("3b:00", as the kernel names it) to its index.
func parseHealthReport(report []byte) ([]GPUWear, map[string]int) {
	var log healthLog
	if err := xml.Unmarshal(report, &log); err != nil {
		return nil, nil
	}
	bus := map[string]int{}
	out := make([]GPUWear, 0, len(log.GPUs))
	for i, g := range log.GPUs {
		bus[busOf(g.ID)] = i
		out = append(out, GPUWear{
			Index:                 i,
			PCIeReplays:           count(g.PCI.Replay),
			RetiredSBE:            count(g.Retired.SBE),
			RetiredDBE:            count(g.Retired.DBE),
			RetirePending:         yes(g.Retired.Pending) || yes(g.Retired.Blacklist),
			RemappedCorrectable:   count(g.Remapped.Corr),
			RemappedUncorrectable: count(g.Remapped.Unc),
			RemapPending:          yes(g.Remapped.Pending),
			RemapFailure:          yes(g.Remapped.Failure),
		})
	}
	return out, bus
}

// busOf is a PCI address as the kernel's XID lines carry it: nvidia-smi
// says "00000000:3B:00.0", the kernel "0000:3b:00".
func busOf(id string) string {
	id = strings.ToLower(strings.TrimSpace(id))
	if i := strings.LastIndex(id, "."); i >= 0 {
		id = id[:i]
	}
	parts := strings.Split(id, ":")
	if len(parts) >= 2 {
		return parts[len(parts)-2] + ":" + parts[len(parts)-1]
	}
	return id
}

func count(s string) int64 {
	n, err := strconv.ParseInt(strings.TrimSpace(s), 10, 64)
	if err != nil {
		return -1
	}
	return n
}

func yes(s string) bool { return strings.EqualFold(strings.TrimSpace(s), "yes") }

// nvlinkError matches a line of `nvidia-smi nvlink -e`, whose wording varies
// by driver: "Link 0: Replay Errors: 0", "Link 0: CRC Errors: 0",
// "Link 3: Data CRC Errors: 2".
var (
	nvlinkGPU   = regexp.MustCompile(`^GPU (\d+):`)
	nvlinkError = regexp.MustCompile(`Link \d+:\s*(.*?)\s*Errors?\s*:\s*(\d+)`)
)

// parseNVLinkErrors sums each card's link errors by kind: CRC, replay,
// recovery.
func parseNVLinkErrors(out string) map[int][3]int64 {
	res := map[int][3]int64{}
	gpu := -1
	for _, line := range strings.Split(out, "\n") {
		if m := nvlinkGPU.FindStringSubmatch(strings.TrimSpace(line)); m != nil {
			gpu, _ = strconv.Atoi(m[1])
			continue
		}
		m := nvlinkError.FindStringSubmatch(line)
		if m == nil || gpu < 0 {
			continue
		}
		n, _ := strconv.ParseInt(m[2], 10, 64)
		e := res[gpu]
		kind := strings.ToLower(m[1])
		switch {
		case strings.Contains(kind, "crc"):
			e[0] += n
		case strings.Contains(kind, "replay"):
			e[1] += n
		case strings.Contains(kind, "recovery"):
			e[2] += n
		}
		res[gpu] = e
	}
	return res
}

type xidEvent struct {
	bus  string
	code int
	text string
	at   float64
}

// xidLine matches the driver's report of an XID:
// "NVRM: Xid (PCI:0000:3b:00): 79, pid=..., GPU has fallen off the bus."
var xidLine = regexp.MustCompile(`NVRM: Xid \(PCI:([0-9a-fA-F:.]+)\): (\d+),\s*(.*)$`)

// parseXIDs reads journalctl's short-unix lines: a timestamp, then the
// message.
func parseXIDs(out string) []xidEvent {
	var evs []xidEvent
	for _, line := range strings.Split(out, "\n") {
		m := xidLine.FindStringSubmatch(line)
		if m == nil {
			continue
		}
		code, _ := strconv.Atoi(m[2])
		e := xidEvent{bus: busOf(m[1]), code: code, text: strings.TrimSpace(m[3])}
		if f := strings.Fields(line); len(f) > 0 {
			e.at, _ = strconv.ParseFloat(f[0], 64)
		}
		evs = append(evs, e)
	}
	return evs
}