llmgw/internal/agent/discovery/sensors.go

package discovery

import (
	"context"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Sensors is what a node's own hardware says about heat and power: the CPUs,
// the board where it has an inlet sensor, and each GPU. A zero is not
// measured, never a reading of zero — a VM has none of these, and must not
// look like a machine running at 0 °C on no power.
type Sensors struct {
	// CPUTempC is the hottest CPU sensor: the package, or the hottest core
	// where no package reading exists.
	CPUTempC float64
	// CPUPowerW is the CPU packages' draw, from RAPL: the change in their
	// energy counters over the time between two samples.
	CPUPowerW float64
	// InletTempC is the air coming in, where the board has a sensor it calls
	// that. Most do not, and a CPU's temperature is not a substitute: it
	// follows the load, not the room.
	InletTempC float64
	GPUs       []GPUSensor
}

// GPUSensor is one card's heat, power and whether it is slowing itself down.
type GPUSensor struct {
	Index       int
	TempC       float64
	MemoryTempC float64
	PowerW      float64
	PowerLimitW float64
	// Throttle names what is holding the card's clocks down now:
	// sw_power_cap, hw_slowdown, sw_thermal, hw_thermal, power_brake. Empty
	// is running free. Idle and application clock settings are left out;
	// they are not the card being held back.
	Throttle []string
	// How busy it is, read with the rest every sample. Discovery reads the
	// same once, at start, which made a busy card look as idle as it was
	// when the agent came up. -1 is not reported by the driver; zero
	// MemoryTotalMB is a sample without these at all.
	UtilPercent, MemUtilPercent float64
	MemoryUsedMB, MemoryTotalMB float64
	SMClockMHz                  float64
	ECCUncorrected              float64
	// Profile is how much of the chip the work uses, where something on the
	// node can say: nil on most cards. UtilPercent is not that.
	Profile *GPUProfile
}

// sensorEvery bounds how often nvidia-smi is run. Heartbeats come more often
// than temperatures change, and each run costs a few hundred milliseconds.
const sensorEvery = 10 * time.Second

// SensorSampler reads Sensors. The first sample has no CPU power, which needs
// two readings of the energy counters.
type SensorSampler struct {
	sys string // /sys, or a test's copy of it
	smi func(ctx context.Context, fields string) ([]byte, error)
	now func() time.Time
	// prof reads the cards' profiling counters; nil reads none.
	prof *ProfSampler

	mu         sync.Mutex
	last       Sensors
	lastAt     time.Time
	prevEnergy map[string]uint64
	prevAt     time.Time
}

// NewSensorSampler reads this machine.
func NewSensorSampler() *SensorSampler {
	return &SensorSampler{sys: "/sys", smi: runSMI, now: time.Now}
}

// WithProfile has each sample carry the cards' profiles as well.
func (s *SensorSampler) WithProfile(p *ProfSampler) *SensorSampler {
	s.prof = p
	return s
}

// SampleGPUs reads every card's heat, power and how busy it is now, without
// a sampler's ten seconds of memory: for a screen that asks every few seconds
// and wants what it asked for. Nil on a machine without nvidia-smi or a card.
func SampleGPUs(ctx context.Context) []GPUSensor {
	return (&SensorSampler{smi: runSMI}).gpus(ctx)
}

func runSMI(ctx context.Context, fields string) ([]byte, error) {
	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	//nolint:gosec // fixed binary, fields from this file
	return exec.CommandContext(ctx, "nvidia-smi", "--query-gpu="+fields,
		"--format=csv,noheader,nounits").Output()
}

// Sample reads the sensors, or returns the last reading when that is recent.
func (s *SensorSampler) Sample(ctx context.Context) Sensors {
	s.mu.Lock()
	defer s.mu.Unlock()
	now := s.now()
	if !s.lastAt.IsZero() && now.Sub(s.lastAt) < sensorEvery {
		return s.last
	}
	out := Sensors{}
	out.CPUTempC, out.InletTempC = s.hwmon()
	out.CPUPowerW = s.rapl(now)
	out.GPUs = s.gpus(ctx)
	if s.prof != nil {
		profs := s.prof.Sample(ctx)
		for i := range out.GPUs {
			if p, ok := profs[out.GPUs[i].Index]; ok {
				out.GPUs[i].Profile = &p
			}
		}
	}
	s.last, s.lastAt = out, now
	return out
}

// hwmon reads the hottest CPU sensor and an inlet sensor, from every hwmon
// device the kernel exposes — what lm_sensors reads, without it installed.
func (s *SensorSampler) hwmon() (cpu, inlet float64) {
	devs, _ := filepath.Glob(filepath.Join(s.sys, "class/hwmon/hwmon*"))
	var pkg, core float64
	for _, d := range devs {
		name := readTrim(filepath.Join(d, "name"))
		inputs, _ := filepath.Glob(filepath.Join(d, "temp*_input"))
		for _, in := range inputs {
			milli, err := strconv.ParseFloat(readTrim(in), 64)
			if err != nil || milli <= 0 {
				continue
			}
			c := milli / 1000
			label := strings.ToLower(readTrim(strings.TrimSuffix(in, "_input") + "_label"))
			switch {
			case isCPUSensor(name) && (strings.HasPrefix(label, "package") || label == "tctl" || label == "tdie"):
				pkg = max(pkg, c)
			case isCPUSensor(name):
				core = max(core, c)
			case label == "cpu":
				// A board's own reading of the CPU, as Dell's SMM gives.
				core = max(core, c)
			case strings.Contains(label, "inlet") || strings.Contains(label, "ambient") ||
				strings.Contains(label, "intake") || label == "systin":
				inlet = max(inlet, c)
			}
		}
	}
	if pkg > 0 {
		return pkg, inlet
	}
	return core, inlet
}

func isCPUSensor(name string) bool {
	switch name {
	case "coretemp", "k10temp", "zenpower", "cpu_thermal":
		return true
	}
	return false
}

// rapl is the CPU packages' power from their energy counters. Package
// domains only: the dram and core subdomains are parts of a package, and
// adding them would count the same energy twice.
func (s *SensorSampler) rapl(now time.Time) float64 {
	doms, _ := filepath.Glob(filepath.Join(s.sys, "class/powercap/intel-rapl:*"))
	energy := map[string]uint64{}
	wrap := map[string]uint64{}
	for _, d := range doms {
		if strings.Count(filepath.Base(d), ":") != 1 {
			continue // intel-rapl:0:0 is a subdomain of intel-rapl:0
		}
		e, err := strconv.ParseUint(readTrim(filepath.Join(d, "energy_uj")), 10, 64)
		if err != nil {
			continue // unreadable without root on newer kernels
		}
		energy[d] = e
		wrap[d], _ = strconv.ParseUint(readTrim(filepath.Join(d, "max_energy_range_uj")), 10, 64)
	}
	prev, prevAt := s.prevEnergy, s.prevAt
	s.prevEnergy, s.prevAt = energy, now
	dt := now.Sub(prevAt).Seconds()
	if prev == nil || dt <= 0 || len(energy) == 0 {
		return 0
	}
	var joules float64
	for d, e := range energy {
		p, ok := prev[d]
		if !ok {
			return 0 // a domain appeared: no fair difference this time
		}
		if e < p { // the counter wrapped
			e += wrap[d]
		}
		joules += float64(e-p) / 1e6
	}
	return joules / dt
}

// The throttle reasons nvidia-smi reports, as bits. Only those that mean the
// card is being held back are named.
var throttleBits = []struct {
	bit  uint64
	name string
}{
	{0x4, "sw_power_cap"},
	{0x8, "hw_slowdown"},
	{0x20, "sw_thermal"},
	{0x40, "hw_thermal"},
	{0x80, "power_brake"},
}

// gpus asks nvidia-smi for each card's heat and power. The memory temperature
// is not known to every driver, and a field nvidia-smi does not know fails the
// whole query, so it is asked for first and dropped on failure.
// gpuFields are what each sample asks nvidia-smi for, in order. The busy
// figures follow the throttle reasons; the memory temperature is last, being
// the one field a driver may reject outright.
const (
	gpuBase  = "index,temperature.gpu,power.draw,enforced.power.limit,clocks_throttle_reasons.active"
	gpuBusy  = ",utilization.gpu,utilization.memory,memory.used,memory.total,clocks.sm,ecc.errors.uncorrected.volatile.total"
	gpuHBMTC = ",temperature.memory"
)

func (s *SensorSampler) gpus(ctx context.Context) []GPUSensor {
	// Most fields first, fewer on a driver that refuses one: a query naming
	// a field the driver does not know fails as a whole.
	var out []byte
	var busy bool
	for _, q := range []struct {
		fields string
		busy   bool
	}{{gpuBase + gpuBusy + gpuHBMTC, true}, {gpuBase + gpuBusy, true}, {gpuBase + gpuHBMTC, false}, {gpuBase, false}} {
		var err error
		if out, err = s.smi(ctx, q.fields); err == nil {
			busy = q.busy
			break
		}
		out = nil
	}
	if out == nil {
		return nil
	}
	var gs []GPUSensor
	for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
		f := strings.Split(line, ",")
		if len(f) < 5 {
			continue
		}
		for i := range f {
			f[i] = strings.TrimSpace(f[i])
		}
		idx, err := strconv.Atoi(f[0])
		if err != nil {
			continue
		}
		g := GPUSensor{Index: idx, TempC: num(f[1]), PowerW: num(f[2]), PowerLimitW: num(f[3])}
		if bits, err := strconv.ParseUint(strings.TrimPrefix(f[4], "0x"), 16, 64); err == nil {
			for _, b := range throttleBits {
				if bits&b.bit != 0 {
					g.Throttle = append(g.Throttle, b.name)
				}
			}
		}
		rest := f[5:]
		if busy && len(rest) >= 6 {
			g.UtilPercent, g.MemUtilPercent = numOrUnknown(rest[0]), numOrUnknown(rest[1])
			g.MemoryUsedMB, g.MemoryTotalMB = num(rest[2]), num(rest[3])
			g.SMClockMHz, g.ECCUncorrected = num(rest[4]), numOrUnknown(rest[5])
			rest = rest[6:]
		}
		if len(rest) > 0 {
			g.MemoryTempC = num(rest[0])
		}
		gs = append(gs, g)
	}
	return gs
}

// numOrUnknown is a figure, or -1 where nvidia-smi says it has none: a card
// without ECC has no error count, which is not the same as none.
func numOrUnknown(s string) float64 {
	v, err := strconv.ParseFloat(s, 64)
	if err != nil || v < 0 {
		return -1
	}
	return v
}

// num reads a figure nvidia-smi may give as "[N/A]" or "N/A".
func num(s string) float64 {
	v, err := strconv.ParseFloat(s, 64)
	if err != nil || v < 0 {
		return 0
	}
	return v
}

func readTrim(path string) string {
	b, err := os.ReadFile(path) //nolint:gosec // sysfs
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}