llmgw/internal/agent/discovery/host.go

package discovery

import (
	"bufio"
	"os"
	"runtime"
	"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"`
}

// 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.
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 ..."
	fields := strings.Fields(string(b))
	for i, f := range fields {
		if f == "Module" && i+1 < len(fields) {
			return fields[i+1]
		}
	}
	return ""
}