llmgw/internal/agent/discovery/prof.go

package discovery

import (
	"context"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"
)

// GPUProfile is how much of a card's chip the work on it uses. nvidia-smi's
// utilization.gpu cannot say: it is the share of time any kernel ran, and a
// kernel keeping one SM of a hundred busy reads 100%. Measured on a Quadro
// RTX 6000 under decode, it read 92% while the SMs had work 69% of the time
// and issued an instruction on 18% of their cycles.
//
// These come from the card's profiling counters, which only two things read
// without stopping the work: NVML's GPM, on Hopper and later, and DCGM, which
// a site may run as dcgm-exporter. Percentages; -1 is not reported.
type GPUProfile struct {
	// SMActive is the share of time the SMs had work, averaged over all of
	// them; SMOccupancy how full those SMs were with warps.
	SMActive, SMOccupancy float64
	// TensorActive is the tensor cores. Decode at batch one hardly uses them
	// even when nothing is wrong: a quantised matrix-vector product runs on
	// the integer units. Prefill and batches are where they show.
	TensorActive float64
	// DRAMActive is the memory interface, which decode is bound by.
	DRAMActive float64
	// Source is "gpm" or "dcgm".
	Source string
}

// DefaultDCGMExporter is where dcgm-exporter listens unless told otherwise.
const DefaultDCGMExporter = "http://127.0.0.1:9400/metrics"

// dcgmRetry is how long a dcgm-exporter that did not answer is left alone.
// Most nodes run none, and asking every ten seconds would be asking nobody.
const dcgmRetry = 5 * time.Minute

// gpmSource reads GPM, where this build and the card have it.
type gpmSource interface {
	read() map[int]GPUProfile
}

// ProfSampler reads each card's GPUProfile: from GPM where the card has it,
// otherwise from a dcgm-exporter on the node.
type ProfSampler struct {
	gpm  gpmSource
	dcgm string
	get  func(ctx context.Context, url string) ([]byte, error)
	now  func() time.Time

	dcgmAfter time.Time
}

// NewProfSampler reads this machine. dcgmURL is the dcgm-exporter to ask:
// empty is DefaultDCGMExporter, "off" asks none.
func NewProfSampler(dcgmURL string) *ProfSampler {
	switch dcgmURL {
	case "":
		dcgmURL = DefaultDCGMExporter
	case "off":
		dcgmURL = ""
	}
	return &ProfSampler{gpm: newGPM(), dcgm: dcgmURL, get: httpGet, now: time.Now}
}

// Sample is each card's profile by index; a card missing has none. GPM needs
// two samples, so its cards appear from the second call.
func (p *ProfSampler) Sample(ctx context.Context) map[int]GPUProfile {
	out := map[int]GPUProfile{}
	if p.gpm != nil {
		for i, v := range p.gpm.read() {
			out[i] = v
		}
	}
	if p.dcgm == "" || p.now().Before(p.dcgmAfter) {
		return out
	}
	body, err := p.get(ctx, p.dcgm)
	if err != nil {
		p.dcgmAfter = p.now().Add(dcgmRetry)
		return out
	}
	for i, v := range parseDCGM(body) {
		if _, ok := out[i]; !ok {
			out[i] = v
		}
	}
	return out
}

func httpGet(ctx context.Context, url string) ([]byte, error) {
	ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		return nil, io.ErrUnexpectedEOF
	}
	return io.ReadAll(io.LimitReader(resp.Body, 4<<20))
}

// The dcgm-exporter fields read, as ratios of one. Its default field list
// has them on every card DCGM can profile.
var dcgmFields = map[string]func(*GPUProfile, float64){
	"DCGM_FI_PROF_SM_ACTIVE":          func(p *GPUProfile, v float64) { p.SMActive = v },
	"DCGM_FI_PROF_SM_OCCUPANCY":       func(p *GPUProfile, v float64) { p.SMOccupancy = v },
	"DCGM_FI_PROF_PIPE_TENSOR_ACTIVE": func(p *GPUProfile, v float64) { p.TensorActive = v },
	"DCGM_FI_PROF_DRAM_ACTIVE":        func(p *GPUProfile, v float64) { p.DRAMActive = v },
}

// parseDCGM reads dcgm-exporter's text exposition: a line per field and
// card, the card in the gpu label. A card with none of the fields — DCGM
// cannot profile it — is left out. A MIG instance's lines carry a
// GPU_I_ID and are left out too: these are whole cards.
func parseDCGM(body []byte) map[int]GPUProfile {
	out := map[int]GPUProfile{}
	for _, line := range strings.Split(string(body), "\n") {
		if line == "" || line[0] == '#' {
			continue
		}
		open := strings.IndexByte(line, '{')
		end := strings.LastIndexByte(line, '}')
		if open < 0 || end < open {
			continue
		}
		set, ok := dcgmFields[line[:open]]
		if !ok {
			continue
		}
		labels := line[open+1 : end]
		if strings.Contains(labels, `GPU_I_ID="`) && !strings.Contains(labels, `GPU_I_ID=""`) {
			continue
		}
		idx, err := strconv.Atoi(labelValue(labels, "gpu"))
		if err != nil {
			continue
		}
		val := strings.Fields(line[end+1:])
		if len(val) == 0 {
			continue
		}
		v, err := strconv.ParseFloat(val[0], 64)
		if err != nil || v < 0 {
			continue
		}
		p, seen := out[idx]
		if !seen {
			p = GPUProfile{SMActive: -1, SMOccupancy: -1, TensorActive: -1, DRAMActive: -1, Source: "dcgm"}
		}
		set(&p, 100*v)
		out[idx] = p
	}
	return out
}

// labelValue is one label's value from a Prometheus label set.
func labelValue(labels, name string) string {
	for _, kv := range strings.Split(labels, ",") {
		k, v, ok := strings.Cut(strings.TrimSpace(kv), "=")
		if ok && k == name {
			return strings.Trim(v, `"`)
		}
	}
	return ""
}