llmgw/internal/agent/gpucheck.go

package agent

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

	"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
	// 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 or in its group:
// an engine that computes in a child process is still the shard.
func holds(pids []int, pid, group int) bool {
	for _, p := range pids {
		if p == pid || (group > 0 && processGroup(p) == group) {
			return true
		}
	}
	return false
}

// pgidOf reads a process's group from /proc; zero when it cannot.
func pgidOf(pid 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) < 3 {
		return 0
	}
	n, _ := strconv.Atoi(f[2])
	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, ", ")
}