llmgw/internal/agent/discovery/host.go

package discovery

import (
	"bufio"
	"os"
	"regexp"
	"runtime"
	"strconv"
	"strings"
)

// Host is what a node is, as far as installing software on it is concerned.
type Host struct {
	// OSID, OSVersion and OSCodename come from /etc/os-release: "ubuntu",
	// "24.04", "noble".
	OSID       string `json:"os_id,omitempty"`
	OSVersion  string `json:"os_version,omitempty"`
	OSCodename string `json:"os_codename,omitempty"`
	// OSLike is ID_LIKE: the families this system says it follows, "debian"
	// for Ubuntu, "rhel centos fedora" for Rocky. It is what decides the
	// package manager, which the name alone does not for every derivative.
	OSLike []string `json:"os_like,omitempty"`
	// Arch is Debian's spelling, "amd64" or "arm64", because that is what
	// package names and release archives are keyed by.
	Arch string `json:"arch,omitempty"`
	// CPUFlags are the vector extensions inference cares about, of those this
	// processor has.
	CPUFlags []string `json:"cpu_flags,omitempty"`
	// NVIDIADriver is the loaded driver's version, empty when there is none.
	NVIDIADriver string `json:"nvidia_driver,omitempty"`
	// MinComputeCapability is the oldest GPU's, 0 when there are none or it
	// is not known. Not read here but from the GPU inventory, by the agent.
	MinComputeCapability float64 `json:"min_compute_capability,omitempty"`
	// MemBandwidthGBs is how fast the memory reads, in GB/s, 0 when not
	// measured: what decides a model's speed on the CPU. Not read here but
	// measured by the agent when it starts (MeasureMemoryBandwidth).
	MemBandwidthGBs float64 `json:"mem_bandwidth_gbs,omitempty"`
	// MemBandwidthCoreGBs is what one core reads alone, in GB/s, 0 when not
	// measured: a replica given n threads reads at about n times this, up to
	// MemBandwidthGBs.
	MemBandwidthCoreGBs float64 `json:"mem_bandwidth_core_gbs,omitempty"`
}

// Family is the packaging family: "deb", "rpm", or "" when neither is known.
func (h Host) Family() string {
	for _, id := range append([]string{h.OSID}, h.OSLike...) {
		switch id {
		case "debian", "ubuntu":
			return "deb"
		case "rhel", "fedora", "centos", "rocky", "almalinux", "suse", "opensuse":
			return "rpm"
		}
	}
	return ""
}

// HasCPUFlag reports whether the processor has an extension.
func (h Host) HasCPUFlag(f string) bool {
	for _, have := range h.CPUFlags {
		if have == f {
			return true
		}
	}
	return false
}

// interestingFlags are the extensions that decide which build of an engine
// runs well. A VM on QEMU's default CPU model has none of them — and still
// runs, five times slower, with nothing saying why.
var interestingFlags = []string{"avx", "avx2", "fma", "f16c", "avx512f", "avx512_bf16", "amx_tile", "asimd", "sve"}

// DiscoverHost reads what the machine is.
func DiscoverHost() Host {
	h := hostFromOSRelease("/etc/os-release")
	h.Arch = debianArch(runtime.GOARCH)
	h.CPUFlags = cpuFlagsFrom("/proc/cpuinfo")
	h.NVIDIADriver = nvidiaDriverFrom("/proc/driver/nvidia/version")
	return h
}

func hostFromOSRelease(path string) Host {
	var h Host
	f, err := os.Open(path) //nolint:gosec // a fixed system path
	if err != nil {
		return h
	}
	defer f.Close()
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		k, v, ok := strings.Cut(sc.Text(), "=")
		if !ok {
			continue
		}
		v = strings.Trim(v, `"'`)
		switch k {
		case "ID":
			h.OSID = v
		case "VERSION_ID":
			h.OSVersion = v
		case "VERSION_CODENAME":
			h.OSCodename = v
		case "ID_LIKE":
			h.OSLike = strings.Fields(v)
		}
	}
	return h
}

func debianArch(goarch string) string {
	switch goarch {
	case "amd64", "arm64":
		return goarch
	case "386":
		return "i386"
	case "arm":
		return "armhf"
	}
	return goarch
}

func cpuFlagsFrom(path string) []string {
	b, err := os.ReadFile(path) //nolint:gosec // a fixed system path
	if err != nil {
		return nil
	}
	have := map[string]bool{}
	for _, line := range strings.Split(string(b), "\n") {
		k, v, ok := strings.Cut(line, ":")
		if !ok {
			continue
		}
		// "flags" on x86, "Features" on ARM. The first processor is enough:
		// a machine does not mix instruction sets across its cores.
		if k = strings.TrimSpace(k); k != "flags" && k != "Features" {
			continue
		}
		for _, f := range strings.Fields(v) {
			have[f] = true
		}
		break
	}
	var out []string
	for _, f := range interestingFlags {
		if have[f] {
			out = append(out, f)
		}
	}
	return out
}

// nvidiaDriverFrom reads the loaded kernel module's version, which exists
// exactly when a driver is loaded — unlike nvidia-smi, which can be installed
// on a machine whose module failed to load.
var driverVersion = regexp.MustCompile(`^[0-9]+\.[0-9]+(\.[0-9]+)?$`)

func nvidiaDriverFrom(path string) string {
	b, err := os.ReadFile(path) //nolint:gosec // a fixed system path
	if err != nil {
		return ""
	}
	// "NVRM version: NVIDIA UNIX x86_64 Kernel Module  570.86.15  Thu ...", or
	// for the open module, the default since 560, "NVIDIA UNIX Open Kernel
	// Module for x86_64  580.82.07  Release Build ...": the first version
	// after "Module", not the word after it. Read as "for", the driver was
	// unknown, and every build was tried whatever driver it needs.
	fields := strings.Fields(string(b))
	for i, f := range fields {
		if f != "Module" {
			continue
		}
		for _, v := range fields[i+1:] {
			if driverVersion.MatchString(v) {
				return v
			}
		}
		return ""
	}
	return ""
}

// MinComputeCapability is the oldest GPU's compute capability, 0 when none
// is known: an engine built for newer GPUs is refused by the oldest.
func MinComputeCapability(gpus []GPU) float64 {
	low := 0.0
	for _, g := range gpus {
		cc, err := strconv.ParseFloat(strings.TrimSpace(g.ComputeCapability), 64)
		if err != nil || cc <= 0 {
			continue
		}
		if low == 0 || cc < low {
			low = cc
		}
	}
	return low
}