llmgw/internal/agent/discovery/gpuprocs.go

package discovery

import (
	"context"
	"encoding/xml"
	"fmt"
	"os/exec"
	"slices"
	"strconv"
	"strings"
)

// GPUProcesses reports which processes hold memory on each GPU, by the GPU's
// nvidia-smi index: what a process actually computes on, as opposed to what it
// was told to use. A llama.cpp started with --n-gpu-layers on a GPU whose
// backend did not load runs on the CPU without a word, and this is where that
// shows — the process is on no GPU's list.
//
// From the full query rather than --query-compute-apps, which lists compute
// processes only: llama.cpp's Vulkan build is listed as a graphics one, and
// would have read as not on its GPU at all.
//
// A node with no nvidia-smi returns nil and no error, as discovery does.
func GPUProcesses(ctx context.Context, binary string) (map[int][]int, error) {
	if binary == "" {
		if _, err := exec.LookPath("nvidia-smi"); err != nil {
			return nil, nil
		}
		binary = "nvidia-smi"
	}
	out, err := exec.CommandContext(ctx, binary, "-q", "-x").Output()
	if err != nil {
		return nil, fmt.Errorf("running %s: %w", binary, err)
	}
	return parseGPUProcesses(out)
}

// GPUMemory is one device's memory as nvidia-smi reports it at a moment: what
// is allocated on it, and which processes hold how much of that.
//
// A device is what discovery reports: a whole card, or each MIG instance of a
// card in MIG mode, in the same order, so the readings line up with the
// inventory one for one.
type GPUMemory struct {
	// Index is the device's number as placement knows it: the card's
	// nvidia-smi index, or MIGIndex's for an instance.
	Index  int
	UsedMB int
	Procs  []GPUProcess
}

// GPUProcess is one process holding memory on a card.
type GPUProcess struct {
	PID    int
	UsedMB int
}

// GPUMemoryByProcess reports each GPU's memory, by index, with the processes
// holding it — from the same report as GPUProcesses, so the two never
// disagree about who is on a card.
//
// A node with no nvidia-smi returns nil and no error, as discovery does.
func GPUMemoryByProcess(ctx context.Context, binary string) ([]GPUMemory, error) {
	if binary == "" {
		if _, err := exec.LookPath("nvidia-smi"); err != nil {
			return nil, nil
		}
		binary = "nvidia-smi"
	}
	out, err := exec.CommandContext(ctx, binary, "-q", "-x").Output()
	if err != nil {
		return nil, fmt.Errorf("running %s: %w", binary, err)
	}
	return parseGPUMemory(out)
}

// parseGPUProcesses reads the report, by device as placement numbers them:
// nvidia-smi's index for a whole card, MIGIndex's for a MIG instance.
func parseGPUProcesses(report []byte) (map[int][]int, error) {
	mem, err := parseGPUMemory(report)
	if err != nil {
		return nil, err
	}
	out := map[int][]int{}
	for _, g := range mem {
		for _, p := range g.Procs {
			out[g.Index] = append(out[g.Index], p.PID)
		}
	}
	return out, nil
}

// parseGPUMemory reads the report into each device's memory and its holders.
//
// Cards are listed in index order, which is nvidia-smi's own numbering. A
// card in MIG mode is read as its instances, each with its own memory and
// the processes nvidia-smi attributes to its GPU and compute instance: a
// process on one instance is not on its neighbours, and the card's total
// would have charged every instance with all of them.
func parseGPUMemory(report []byte) ([]GPUMemory, error) {
	var log migLog
	if err := xml.Unmarshal(report, &log); err != nil {
		return nil, fmt.Errorf("reading nvidia-smi's report: %w", err)
	}
	out := make([]GPUMemory, 0, len(log.GPUs))
	for i, g := range log.GPUs {
		procs := func(gi, ci int) []GPUProcess {
			var ps []GPUProcess
			for _, p := range g.Processes {
				pid, err := strconv.Atoi(strings.TrimSpace(p.PID))
				if err != nil {
					continue
				}
				if gi >= 0 && (atoiOr(p.GPUInstance, -1) != gi || atoiOr(p.ComputeInstance, -1) != ci) {
					continue
				}
				ps = append(ps, GPUProcess{PID: pid, UsedMB: mib(p.UsedMB)})
			}
			return ps
		}
		if !enabled(g.Current) || len(g.Devices) == 0 {
			out = append(out, GPUMemory{Index: i, UsedMB: mib(g.Used), Procs: procs(-1, -1)})
			continue
		}
		devs := slices.Clone(g.Devices)
		slices.SortFunc(devs, func(a, b migLogDevice) int { return atoiOr(a.Index, -1) - atoiOr(b.Index, -1) })
		for _, d := range devs {
			out = append(out, GPUMemory{
				Index:  MIGIndex(i, atoiOr(d.Index, 0)),
				UsedMB: mib(d.Used),
				Procs:  procs(atoiOr(d.GPUInstance, -1), atoiOr(d.ComputeInstance, -1)),
			})
		}
	}
	return out, nil
}

// mib reads "4932 MiB"; anything else, "N/A" included, is zero.
func mib(s string) int {
	n, _ := strconv.Atoi(strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(s), "MiB")))
	return n
}