gpuinfo.go

// gpuinfo.go — the GPU memory PROVIDER seam (J3). fleet-serve historically
// hard-required nvidia-smi: the startup probe, the 2s health sampler, and the
// vendor/arch advertisement all shelled to it, so a working AMD/Vulkan box
// (the amd-rdna3 tier) could not join the fleet at all — and the refusal text
// blamed the GPU rather than the probe. This file makes the memory source a
// resolved PROVIDER: nvidia-smi where it works, else the Windows-generic
// WDDM source (registry capacity + PDH adapter usage, vram_windows.go), with
// vendor/arch read from the installer's manifest instead of re-derived from
// NVIDIA product names. The serve gate becomes "no working GPU memory
// source", which is the fact the contract actually cares about.
package fleetnode

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/dmmdea/offload-harness/internal/gpuprobe"
)

// MemProbe returns the current GPU memory capacity/usage in GiB. It is the
// unit both the startup gate and the health sampler consume; every provider
// (nvidia-smi, windows-generic) reduces to one of these.
type MemProbe func() (totalGiB, usedGiB float64, err error)

// SmiProbe adapts the nvidia-smi CSV runner (the historical source) to a
// MemProbe via ParseSmiMemory.
func SmiProbe(run func() (string, error)) MemProbe {
	return func() (float64, float64, error) {
		out, err := run()
		if err != nil {
			return 0, 0, fmt.Errorf("nvidia-smi: %w", err)
		}
		return ParseSmiMemory(out)
	}
}

// DeviceProbe returns the full per-device VRAM breakdown (nvidia-smi only
// today — vram_windows.go's GenericWindowsProbe documents why the
// windows-generic WDDM source cannot produce one: its PDH reads sum blindly
// across adapter instances with no per-adapter identity preserved). Unlike
// MemProbe, a DeviceProbe never collapses to a single number itself —
// HeadlineDevice does that deliberately, once, in the sampler, so there is
// exactly one place in the codebase that decides which device is "the"
// device for the health payload's headline pair.
type DeviceProbe func() ([]GPUDevice, error)

// SmiDeviceProbe adapts the nvidia-smi multi-device CSV runner (index, name,
// memory.total, memory.used) to a DeviceProbe via ParseSmiMemoryDevices — the
// devices-aware sibling of SmiProbe/ParseSmiMemory. run must invoke the
// index+name query (a different nvidia-smi invocation than the plain
// memory.total,memory.used one SmiProbe/ParseSmiMemory expect); the two are
// kept separate rather than reusing one query so ParseSmiMemory's existing
// 2-field CSV contract (and its other caller, the pipeline's per-process
// footprint delta sampler) never breaks.
func SmiDeviceProbe(run func() (string, error)) DeviceProbe {
	probe, _ := newSmiProbe(run, gpuprobe.ReadThermal, time.Now, stderrLogf)
	return probe
}

// The thermal cadence. The sampler ticks every 2 s, and a second nvidia-smi per tick would double the spawns on a
// host the harness shares with other work (and on the thermal-care box it would add heat of its own), so the
// thermal query runs at most once per thermalEvery whatever the tick. A failed query is not retried sooner than
// thermalBackoff and says so once. An answer older than thermalMaxAge is dropped, never published: a stale
// "cool" next to a card that has since heat-soaked is the one thing this field must not do. The read is bounded
// by thermalTimeout, which the per-device query (no deadline of its own) does not have. The query runs on a goroutine
// of its own, one at a time, and the probe only reads what the last one stored: the memory reading the placement guards
// stand on never waits for it, and a thermal query that hangs costs the thermal fields and nothing else.
const (
	thermalEvery   = 30 * time.Second
	thermalBackoff = 60 * time.Second
	thermalMaxAge  = 90 * time.Second
	thermalTimeout = 10 * time.Second
)

// thermalReader reads the optional thermal query (gpuprobe.ReadThermal in production, a fake in tests).
type thermalReader func(ctx context.Context) ([]gpuprobe.Thermal, error)

// stderrLogf is where the one line about a failing thermal query goes.
func stderrLogf(format string, args ...any) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
}

// newSmiProbe is SmiDeviceProbe over an injected thermal reader, clock and log, and the cache it merges from (nil when
// there is no thermal reader). A nil reader is no thermal at all: the probe is then exactly the memory probe it was
// before F11.
func newSmiProbe(run func() (string, error), thermal thermalReader, now func() time.Time, logf func(string, ...any)) (DeviceProbe, *thermalCache) {
	var cache *thermalCache
	if thermal != nil {
		if now == nil {
			now = time.Now
		}
		if logf == nil {
			logf = stderrLogf
		}
		cache = &thermalCache{read: thermal, now: now, logf: logf}
	}
	return func() ([]GPUDevice, error) {
		out, err := run()
		if err != nil {
			return nil, fmt.Errorf("nvidia-smi: %w", err)
		}
		devs, err := ParseSmiMemoryDevices(out)
		if err != nil {
			return nil, err
		}
		if cache != nil {
			cache.merge(devs)
		}
		return devs, nil
	}, cache
}

// thermalCache holds the last thermal answer and decides when the next query is due. The probe calls merge on every
// tick; a query that falls due starts on its own goroutine, never more than one at a time, and its answer is what a
// later tick publishes. The lock guards the fields; the query itself runs outside it.
type thermalCache struct {
	mu       sync.Mutex
	read     thermalReader
	now      func() time.Time
	logf     func(string, ...any)
	rows     []gpuprobe.Thermal
	at       time.Time // when the query that gave rows STARTED: an answer is aged from there, never from when it returned
	have     bool
	nextTry  time.Time
	warned   bool           // the current outage has been logged
	inflight bool           // a query is out
	wg       sync.WaitGroup // the query in flight, so a test can wait for it
}

// merge puts the thermal values on devs and starts the next query when one is due and none is out. It never waits for
// a query: a dropped or missing answer leaves the devices exactly as the memory query made them.
func (c *thermalCache) merge(devs []GPUDevice) {
	c.mu.Lock()
	now := c.now()
	if !c.inflight && !now.Before(c.nextTry) {
		c.inflight = true
		c.wg.Add(1)
		go c.refresh(now)
	}
	rows, fresh := c.rows, c.have && now.Sub(c.at) <= thermalMaxAge
	c.mu.Unlock()
	if fresh {
		gpuprobe.MergeThermal(devs, rows)
	}
}

// refresh runs one thermal query started at `started` and stores its outcome. It replaces c.rows wholesale and never
// edits the slice a probe may be reading.
func (c *thermalCache) refresh(started time.Time) {
	defer c.wg.Done()
	ctx, cancel := context.WithTimeout(context.Background(), thermalTimeout)
	rows, err := c.read(ctx)
	cancel()
	if err == nil && len(rows) == 0 {
		err = errors.New("the query answered no card")
	}
	c.mu.Lock()
	c.inflight = false
	announce := false
	if err != nil {
		c.nextTry = started.Add(thermalBackoff)
		announce, c.warned = !c.warned, true
	} else {
		c.rows, c.at, c.have, c.nextTry, c.warned = rows, started, true, started.Add(thermalEvery), false
	}
	c.mu.Unlock()
	if announce {
		c.logf("[fleet-serve] warning: gpu temperature and throttle reasons are not being read (%v); retrying no sooner than every %s, and publishing none while the last answer is older than %s", err, thermalBackoff, thermalMaxAge)
	}
}

// InstalledInfo is the slice of the installer's manifest ($OFFLOAD_HOME/
// installed.json) the fleet node needs: the hardware profile detect.ps1
// classified (the vendor/arch/UMA truth — no re-derivation from product-name
// regexes) and the serving backend.
type InstalledInfo struct {
	Profile string `json:"profile"`
	Backend string `json:"backend"`
	// Accelerators are the additive devices the installer detected (ADR 0024);
	// advertised (minus local-only devices, config.FleetVisibleAccelerators) so a
	// delegator can route NPU-owned work here.
	Accelerators []string `json:"accelerators,omitempty"`
	// AltBackends are the extra serving backends THIS install rendered beside the
	// primary one (today: ["cpu"] when install.sh was given --llama-bin-cpu). The
	// installer writes only what it rendered, so health never advertises a route
	// the box cannot serve.
	AltBackends []string `json:"alt_backends,omitempty"`
}

// Backends is the ordered list health advertises: the primary backend first,
// then the alternates the install rendered. Empty (omitted from health) when
// the manifest carries no backend — every pre-manifest node stays byte-identical.
func (i InstalledInfo) Backends() []string {
	if i.Backend == "" {
		return nil
	}
	out := []string{i.Backend}
	for _, b := range i.AltBackends {
		if b != "" && b != i.Backend {
			out = append(out, b)
		}
	}
	return out
}

// ReadInstalledInfo loads InstalledInfo from an installed.json path. A missing
// or unparsable manifest returns a zero value and the error — callers treat
// that as "no manifest" (pre-installer box), never fatal by itself.
func ReadInstalledInfo(path string) (InstalledInfo, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return InstalledInfo{}, err
	}
	var m InstalledInfo
	if err := json.Unmarshal(b, &m); err != nil {
		return InstalledInfo{}, fmt.Errorf("installed.json: %w", err)
	}
	return m, nil
}

// InstalledJSONPath resolves where the installer wrote its manifest:
// $OFFLOAD_HOME/installed.json, defaulting OFFLOAD_HOME to ~/offload-stack
// exactly like install.ps1 does.
func InstalledJSONPath() string {
	home := os.Getenv("OFFLOAD_HOME")
	if home == "" {
		userHome, err := os.UserHomeDir()
		if err != nil {
			return ""
		}
		home = filepath.Join(userHome, "offload-stack")
	}
	return filepath.Join(home, "installed.json")
}

// VendorArchFromProfile maps an installer profile id to the (vendor, arch)
// pair the health payload advertises. Pure; unit-tested. Unknown/empty
// profiles return ("unknown", "unknown") — honest, never a guessed vendor.
func VendorArchFromProfile(profile string) (vendor, arch string) {
	switch {
	case strings.HasPrefix(profile, "amd-rdna3"): // amd-rdna3 + amd-rdna3-dgpu
		return "amd", "rdna3"
	case profile == "amd-gcn":
		return "amd", "gcn"
	case strings.HasPrefix(profile, "blackwell-"):
		return "nvidia", "blackwell"
	case strings.HasPrefix(profile, "ampere-"):
		return "nvidia", "ampere"
	case profile == "volta-16":
		return "nvidia", "volta"
	case profile == "dual-gpu":
		return "nvidia", "blackwell" // the dual rig's primary; refined when measured
	case profile == "rockchip-rk3588":
		return "rockchip", "rk3588"
	case profile == "cpu":
		return "none", "none"
	default:
		return "unknown", "unknown"
	}
}

// UMAFromProfile reports whether the profile is a unified-memory iGPU class —
// which changes BOTH capacity composition (carve-out + WDDM shared budget)
// and usage composition (Dedicated + Shared). An SoC (rockchip-rk3588) is
// unified through and through: it has no carve-out at all, and its provider
// (MeminfoUMAProbe) composes from RAM alone. Pure; unit-tested. ok is false
// when the profile carries no UMA signal (caller falls back to a capacity
// heuristic).
func UMAFromProfile(profile string) (uma, ok bool) {
	switch profile {
	case "amd-rdna3", "amd-gcn", "rockchip-rk3588":
		return true, true
	case "amd-rdna3-dgpu":
		return false, true
	default:
		if strings.HasPrefix(profile, "blackwell-") || strings.HasPrefix(profile, "ampere-") ||
			profile == "volta-16" || profile == "dual-gpu" {
			return false, true // discrete NVIDIA classes
		}
		return false, false
	}
}

// ResolvedProvider is what fleet-serve runs on after resolution: the probe
// that feeds the gate + sampler, the advertisement identity, the source name
// for the operator log line, and the selection probe's own reading (TotalGiB/
// UsedGiB — so the startup banner never re-executes or discards a probe).
type ResolvedProvider struct {
	Probe    MemProbe
	Vendor   string
	Arch     string
	Source   string // "nvidia-smi" | "windows-generic" | "linux-amdgpu" | "linux-meminfo"
	TotalGiB float64
	UsedGiB  float64
}

// ResolveProvider picks the working GPU memory source. Order: nvidia-smi (the
// proven path — keeps every existing NVIDIA node byte-identical in behavior),
// else the injected generic provider (Windows WDDM; nil off-Windows, on a
// cpu-profile box, or in tests that exclude it). smiIdentity is called ONLY
// when nvidia-smi wins (it typically shells nvidia-smi for the product name —
// pointless on a box where smi just failed). Both probes failing is the new
// gate error: "no working GPU memory source". A note on transience: selection
// happens once at startup — a transiently failing nvidia-smi binds the generic
// provider for the process lifetime (restart re-resolves; documented in ADR
// 0014). Pure of I/O — probes, identity, and manifest info are injected, so
// selection is unit-testable.
func ResolveProvider(smi MemProbe, generic MemProbe, smiIdentity func() (vendor, arch string), info InstalledInfo) (ResolvedProvider, error) {
	var smiErr, genErr error
	if smi != nil {
		if total, used, err := smi(); err == nil {
			vendor, arch := "nvidia", "nvidia"
			if smiIdentity != nil {
				vendor, arch = smiIdentity()
			}
			return ResolvedProvider{Probe: smi, Vendor: vendor, Arch: arch, Source: "nvidia-smi", TotalGiB: total, UsedGiB: used}, nil
		} else {
			smiErr = err
		}
	}
	if generic != nil {
		if total, used, err := generic(); err == nil {
			vendor, arch := VendorArchFromProfile(info.Profile)
			return ResolvedProvider{Probe: generic, Vendor: vendor, Arch: arch, Source: "windows-generic", TotalGiB: total, UsedGiB: used}, nil
		} else {
			genErr = err
		}
	}
	return ResolvedProvider{}, fmt.Errorf("no working GPU memory source: nvidia-smi (%v); windows-generic (%v)", errOrAbsent(smiErr, smi == nil), errOrAbsent(genErr, generic == nil))
}

func errOrAbsent(err error, absent bool) string {
	if absent {
		return "not available on this platform"
	}
	if err == nil {
		return "not tried"
	}
	return err.Error()
}