llmgw/internal/agent/gpucheck.go

package agent

import (
	"context"
	"fmt"
	"os"
	"slices"
	"strconv"
	"strings"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
)

// gpuProcesses and processGroup are variables so tests can stand in for
// nvidia-smi and /proc.
var (
	gpuProcesses = discovery.GPUProcesses
	gpuMemory    = discovery.GPUMemoryByProcess
	processGroup = pgidOf
	parentOf     = ppidOf
	// gpuCheckWait is how long a shard has, once it answers, to show up on
	// its GPUs: the weights are on them by then, but nvidia-smi can lag.
	gpuCheckWait = 15 * time.Second
)

// onItsGPUs checks that a shard placed on GPUs computes on them.
//
// For an engine that falls back to the CPU when its GPU backend does not load,
// "it answers" is not "it is ready": llama.cpp built for CUDA but missing
// cuBLAS comes up on the CPU, answers at a tenth of the speed, and the GPU it
// was given sits idle and reserved. What it computes on is what nvidia-smi
// lists it on, so a shard none of whose processes hold memory on one of its
// GPUs is refused, with what the engine itself can see.
func (a *Agent) onItsGPUs(ctx context.Context, eng engines.Engine, spec engines.ShardSpec, pid int) error {
	fb, ok := eng.(engines.CPUFallback)
	if !ok || !fb.FallsBackToCPU() || len(spec.GPUIndices) == 0 || pid <= 0 {
		return nil
	}
	group := processGroup(pid)
	deadline := time.Now().Add(gpuCheckWait)
	var missing []int
	for {
		procs, err := gpuProcesses(ctx, a.cfg.Discovery.NvidiaSMI)
		if err != nil || procs == nil {
			// No nvidia-smi to ask is no evidence either way, and a node that
			// was given GPUs has one; say so and let the shard serve.
			a.log.Warn("could not check that a GPU shard runs on its GPUs", "shard", spec.ShardID, "err", err)
			return nil
		}
		missing = missing[:0]
		for _, i := range spec.GPUIndices {
			if !holds(procs[i], pid, group) {
				missing = append(missing, i)
			}
		}
		if len(missing) == 0 {
			return nil
		}
		if time.Now().After(deadline) {
			break
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(time.Second):
		}
	}
	seen := "it lists no GPU at all"
	if devs, err := fb.Devices(ctx); err != nil {
		seen = "listing its devices failed: " + err.Error()
	} else if len(devs) > 0 {
		seen = "it lists " + strings.Join(devs, "; ")
	}
	return fmt.Errorf("placed on GPU %s but computing on the CPU: the engine holds no memory there. "+
		"Its GPU backend did not load (%s). Check the build with `ldd` on its libggml-cuda.so, "+
		"or run `llmgw doctor` on this node", joinIndices(missing), seen)
}

// foreignGPUMemory is what each GPU holds for anything other than this
// agent's shards, by index, and each GPU's memory in use. Nil when nvidia-smi
// cannot be asked, which is no reading rather than an empty card.
//
// Everything allocated less what the shards' processes hold, rather than the
// sum of the other processes: a process in another container's PID namespace
// holds memory nvidia-smi lists under no PID this node can see, and it is
// still not memory placement can hand out.
func (a *Agent) foreignGPUMemory(ctx context.Context, shards []int) (foreign, used []int32) {
	mem, err := gpuMemory(ctx, a.cfg.Discovery.NvidiaSMI)
	if err != nil || mem == nil {
		return nil, nil
	}
	groups := make([]int, len(shards))
	for i, pid := range shards {
		groups[i] = processGroup(pid)
	}
	ours := func(p int) bool {
		for i, pid := range shards {
			if holds([]int{p}, pid, groups[i]) {
				return true
			}
		}
		return false
	}
	for _, g := range mem {
		f := g.UsedMB
		for _, p := range g.Procs {
			if ours(p.PID) {
				f -= p.UsedMB
			}
		}
		foreign = append(foreign, int32(max(f, 0)))
		used = append(used, int32(g.UsedMB))
	}
	return foreign, used
}

// holds reports whether any of pids is the shard's process, in its group,
// or a process it started: an engine that computes in a child process is
// still the shard, and Strata's server starts its engine in a group of the
// engine's own.
func holds(pids []int, pid, group int) bool {
	for _, p := range pids {
		if p == pid || (group > 0 && processGroup(p) == group) || descends(p, pid) {
			return true
		}
	}
	return false
}

// descends reports whether process p was started, at any remove, by pid.
func descends(p, pid int) bool {
	// Bounded: a chain this deep is no engine's, and a /proc read racing an
	// exit must not loop.
	for range 16 {
		p = parentOf(p)
		switch {
		case p == pid:
			return true
		case p <= 1:
			return false
		}
	}
	return false
}

// pgidOf reads a process's group from /proc; zero when it cannot.
func pgidOf(pid int) int { return statField(pid, 2) }

// ppidOf reads a process's parent from /proc; zero when it cannot.
func ppidOf(pid int) int { return statField(pid, 1) }

// statField is one of the numbers after the command name in /proc's stat
// line: 1 is the parent, 2 the group.
func statField(pid, field int) int {
	b, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat")
	if err != nil {
		return 0
	}
	// The command name, in parentheses, may hold spaces; the fields after its
	// closing parenthesis are state, ppid, pgrp.
	s := string(b)
	i := strings.LastIndexByte(s, ')')
	if i < 0 {
		return 0
	}
	f := strings.Fields(s[i+1:])
	if len(f) <= field {
		return 0
	}
	n, _ := strconv.Atoi(f[field])
	return n
}

func joinIndices(is []int) string {
	s := make([]string, len(is))
	for k, i := range is {
		s[k] = strconv.Itoa(i)
	}
	return strings.Join(s, ", ")
}

// engineCompute is where an engine said it put its model, from its output
// so far — llama.cpp's, which says — and false for one that said nothing. A
// shard given GPUs that its engine says it computes on the CPU is an error,
// found this way even where nvidia-smi cannot be asked.
func engineCompute(eng engines.Engine, gpus []int, lines []string) (engines.Compute, bool, error) {
	if _, ok := eng.(engines.CPUFallback); !ok {
		return engines.Compute{}, false, nil
	}
	read, ok := engines.ReadLlamaCompute(lines)
	if !ok {
		return engines.Compute{}, false, nil
	}
	if len(gpus) > 0 && read.OnTheCPU() {
		why := "its GPU backend saw no device it could use"
		if read.Total > 0 && !read.NoGPU {
			why = fmt.Sprintf("it offloaded %d of %d layers", read.Offloaded, read.Total)
		}
		return read, true, fmt.Errorf("placed on GPU %v but computing on the CPU: %s, by llama.cpp's own account "+
			"(llmgw ctl engine-check on this node says why)", gpus, why)
	}
	return read, true, nil
}

// computeToWire is a compute as it travels to the control plane.
func computeToWire(c *engines.Compute) *llmgwv1.Compute {
	if c == nil {
		return nil
	}
	return &llmgwv1.Compute{Backend: c.Backend, Devices: c.Devices,
		LayersOffloaded: int32(c.Offloaded), LayersTotal: int32(c.Total)} //nolint:gosec // layer counts
}

// gpuCompute is where a shard computes as nvidia-smi sees it: the CPU for a
// shard given no GPU, else the GPUs whose memory its processes hold, and how
// much. False when nvidia-smi cannot be asked. Engine-agnostic, and needing
// no log line: llama.cpp says where its layers went only at -lv 4.
func (a *Agent) gpuCompute(ctx context.Context, gpus []int, pid int) (engines.Compute, bool) {
	if len(gpus) == 0 {
		return engines.Compute{Backend: "CPU"}, true
	}
	mem, err := gpuMemory(ctx, a.cfg.Discovery.NvidiaSMI)
	if err != nil || mem == nil || pid <= 0 {
		return engines.Compute{}, false
	}
	group := processGroup(pid)
	c := engines.Compute{Backend: "CUDA"}
	for _, g := range mem {
		if !slices.Contains(gpus, g.Index) {
			continue
		}
		used := 0
		for _, p := range g.Procs {
			if holds([]int{p.PID}, pid, group) {
				used += p.UsedMB
			}
		}
		if used > 0 {
			c.Devices = append(c.Devices, fmt.Sprintf("GPU %d: %.1f GB", g.Index, float64(used)/1024))
		}
	}
	if len(c.Devices) == 0 {
		// Placed on GPUs and holding none of their memory.
		c.Backend = "CPU"
	}
	return c, true
}