llmgw/internal/agent/discovery/gpm_nvml.go

//go:build nvml

package discovery

import (
	"sync"

	"github.com/NVIDIA/go-nvml/pkg/nvml"
)

// nvmlGPM reads GPM, NVML's own profiling counters, on the cards that have
// them: Hopper and later. Each read takes a sample and works out the metrics
// against the one before, so NVML stays initialised between reads.
type nvmlGPM struct {
	mu     sync.Mutex
	opened bool
	failed bool
	cards  map[int]*gpmCard
}

type gpmCard struct {
	dev          nvml.Device
	prev, cur    nvml.GpmSample
	havePrevious bool
}

func newGPM() gpmSource { return &nvmlGPM{} }

// The GPM metrics read, in the order they are asked for.
var gpmMetrics = []struct {
	id  nvml.GpmMetricId
	set func(*GPUProfile, float64)
}{
	{nvml.GPM_METRIC_SM_UTIL, func(p *GPUProfile, v float64) { p.SMActive = v }},
	{nvml.GPM_METRIC_SM_OCCUPANCY, func(p *GPUProfile, v float64) { p.SMOccupancy = v }},
	{nvml.GPM_METRIC_ANY_TENSOR_UTIL, func(p *GPUProfile, v float64) { p.TensorActive = v }},
	{nvml.GPM_METRIC_DRAM_BW_UTIL, func(p *GPUProfile, v float64) { p.DRAMActive = v }},
}

func (g *nvmlGPM) read() map[int]GPUProfile {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.failed {
		return nil
	}
	if !g.opened && !g.open() {
		g.failed = true
		return nil
	}
	out := map[int]GPUProfile{}
	for idx, c := range g.cards {
		if c.cur.Get(c.dev) != nvml.SUCCESS {
			// A card in MIG mode, or one that has gone: no sample, and the
			// next must not be measured against an old one.
			c.havePrevious = false
			continue
		}
		if c.havePrevious {
			if p, ok := gpmBetween(c.prev, c.cur); ok {
				out[idx] = p
			}
		}
		c.prev, c.cur = c.cur, c.prev
		c.havePrevious = true
	}
	return out
}

// open finds the cards with GPM. False is NVML not starting, or no card
// having it, and the reader then never tries again.
func (g *nvmlGPM) open() bool {
	if nvml.Init() != nvml.SUCCESS {
		return false
	}
	g.cards = map[int]*gpmCard{}
	n, ret := nvml.DeviceGetCount()
	if ret != nvml.SUCCESS {
		n = 0
	}
	for i := range n {
		dev, ret := nvml.DeviceGetHandleByIndex(i)
		if ret != nvml.SUCCESS {
			continue
		}
		sup, ret := dev.GpmQueryDeviceSupport()
		if ret != nvml.SUCCESS || sup.IsSupportedDevice == 0 {
			continue
		}
		a, ret := nvml.GpmSampleAlloc()
		if ret != nvml.SUCCESS {
			continue
		}
		b, ret := nvml.GpmSampleAlloc()
		if ret != nvml.SUCCESS {
			_ = a.Free()
			continue
		}
		g.cards[i] = &gpmCard{dev: dev, prev: a, cur: b}
	}
	if len(g.cards) == 0 {
		_ = nvml.Shutdown()
		return false
	}
	g.opened = true
	return true
}

// gpmBetween is the profile over the time between two samples.
func gpmBetween(earlier, later nvml.GpmSample) (GPUProfile, bool) {
	q := nvml.GpmMetricsGetType{NumMetrics: uint32(len(gpmMetrics)), Sample1: earlier, Sample2: later}
	for i, m := range gpmMetrics {
		q.Metrics[i].MetricId = uint32(m.id)
	}
	if nvml.GpmMetricsGet(&q) != nvml.SUCCESS {
		return GPUProfile{}, false
	}
	p := GPUProfile{SMActive: -1, SMOccupancy: -1, TensorActive: -1, DRAMActive: -1, Source: "gpm"}
	for i, m := range gpmMetrics {
		if r := q.Metrics[i]; r.NvmlReturn == uint32(nvml.SUCCESS) && r.Value >= 0 {
			m.set(&p, r.Value)
		}
	}
	return p, true
}