llmgw/internal/agent/discovery/discover.go

package discovery

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
)

// ErrNoNVML means this binary was built without the NVML backend.
//
// It is declared here rather than in the build-tagged file so callers — and
// tests — can refer to it in either build.
var ErrNoNVML = errors.New("built without the nvml tag; rebuild with -tags nvml for NVLink topology")

// Options configure a full inventory scan.
type Options struct {
	// NvidiaSMI is the path to nvidia-smi; empty means look it up on PATH.
	// Worth configuring on nodes where the driver lives outside the default
	// paths, which containerized driver installs often do.
	NvidiaSMI string `koanf:"nvidia_smi"`
	// IBRoot is the sysfs InfiniBand directory.
	IBRoot string `koanf:"ib_root"`
	// PreferNVML uses the NVML backend when this binary has one. It is the only
	// source of NVLink topology.
	PreferNVML bool `koanf:"prefer_nvml"`
	// SkipGPUs omits GPU discovery. A node without nvidia-smi is already
	// found to have none, so this is for one whose driver is broken.
	SkipGPUs bool `koanf:"skip_gpus"`
	// ModelDirs are scanned for weights. Empty means do not look.
	ModelDirs []string `koanf:"-"`
}

// Discover builds the node's inventory.
//
// When NVML is preferred but fails, it falls back to nvidia-smi and says so in
// Inventory.Source. The fallback is logged rather than silent: on a node that
// is supposed to have working NVML, a quiet downgrade hides a broken driver and
// costs the scheduler its NVLink topology.
func Discover(ctx context.Context, opts Options, log *slog.Logger) (Inventory, error) {
	if log == nil {
		log = slog.Default()
	}
	inv := Inventory{Source: "none"}
	// First, because it has nothing to do with GPUs and every path below can
	// return early. Read after them, a node with no nvidia-smi reported no
	// memory at all: GPU discovery fails, Discover returns, and the fact that
	// has nothing to do with cards goes with it. That is exactly the node whose
	// memory matters most, since a CPU profile is what it will be asked to run.
	inv.MemoryTotalMB = MemoryTotalMB()
	// Also before the GPU paths, and for the same reason: which weights a node
	// holds has nothing to do with its cards, and a node with none is exactly
	// where a CPU profile will look for them.
	inv.Weights = ScanWeights(opts.ModelDirs)
	// And the disk they sit on. A node can hold a model today and be unable to
	// take the next one, which the file list alone never shows.
	inv.Storage = ScanStorage(opts.ModelDirs)

	if !opts.SkipGPUs {
		switch {
		case opts.PreferNVML && nvmlAvailable:
			gpus, err := DiscoverGPUsNVML(ctx)
			if err == nil {
				inv.GPUs, inv.Source = gpus, "nvml"
				break
			}
			log.Warn("NVML discovery failed, falling back to nvidia-smi", "err", err)
			fallthrough
		default:
			gpus, err := DiscoverGPUsSMI(ctx, opts.NvidiaSMI)
			if err != nil {
				return inv, fmt.Errorf("gpu discovery: %w", err)
			}
			inv.GPUs, inv.Source = gpus, "nvidia-smi"
		}
		// NVLink and which adapter sits nearest each card, for either path:
		// NVML knows NVLink but not the adapters, nvidia-smi's query knows
		// neither. A failure costs the topology and nothing else — the cards
		// are still cards.
		if len(inv.GPUs) > 0 {
			topo, err := DiscoverTopology(ctx, opts.NvidiaSMI)
			if err != nil {
				log.Warn("could not read the GPU topology; tensor-parallel groups "+
					"will be chosen without NVLink", "err", err)
			} else {
				topo.Apply(inv.GPUs)
			}
		}
	}

	hcas, err := DiscoverIB(opts.IBRoot)
	if err != nil {
		return inv, fmt.Errorf("infiniband discovery: %w", err)
	}
	inv.HCAs = hcas
	return inv, nil
}

// NVMLSupported reports whether this binary includes the NVML backend.
func NVMLSupported() bool { return nvmlAvailable }