Files
llmgw/internal/agent/gpucheck_test.go
T
Ilker ManapandClaude Opus 5.5 1908bfb4b2
CI / test (push) Failing after 7m31s
Strata on cuda01 end to end: its engine is the shard's, its metrics read
Found running it through the cluster:

- Strata's server starts its engine in a process group of its own, so the
  agent saw a GPU replica on the CPU and its 22 GB as somebody else's
  memory. A process the shard started, at any remove, is now the shard's.
- An adopted shard's kept compute was sent at once and never again, so a
  wrong one read again on adoption never reached the control plane: what
  is sent is sent again when it changes, and a GPU shard kept as the CPU
  is read again on adoption.
- /metrics/prometheus is llama.cpp's names without the KV cache; vLLM's,
  with it, are /metrics?format=prometheus. The router now sees running 1,
  waiting 2 under load.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-29 15:48:15 +01:00

201 lines
8.1 KiB
Go

package agent
import (
"context"
"log/slog"
"slices"
"strings"
"testing"
"time"
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
)
// fallbackEngine is an engine that comes up on the CPU when its GPU backend
// does not load, and lists the devices it can see.
type fallbackEngine struct {
engines.Engine
devices []string
}
func (fallbackEngine) FallsBackToCPU() bool { return true }
func (f fallbackEngine) Devices(context.Context) ([]string, error) {
return f.devices, nil
}
func TestAGPUShardComputingOnTheCPUIsNotReady(t *testing.T) {
procs := map[int][]int{}
oldProcs, oldGroup, oldWait := gpuProcesses, processGroup, gpuCheckWait
gpuProcesses = func(context.Context, string) (map[int][]int, error) { return procs, nil }
processGroup = func(pid int) int {
if pid >= 500 && pid < 600 {
return 500 // the shard's group: 500 and its children
}
return pid
}
gpuCheckWait = 0
t.Cleanup(func() { gpuProcesses, processGroup, gpuCheckWait = oldProcs, oldGroup, oldWait })
a := &Agent{log: slog.Default()}
spec := engines.ShardSpec{ShardID: "s", GPUIndices: []int{0, 1}}
eng := fallbackEngine{}
// On no GPU at all: refused, saying what the engine sees.
err := a.onItsGPUs(context.Background(), eng, spec, 500)
if err == nil || !strings.Contains(err.Error(), "computing on the CPU") || !strings.Contains(err.Error(), "lists no GPU") {
t.Fatalf("a shard on no GPU was let through: %v", err)
}
// On one of its two: refused, naming the one it is missing.
procs[0] = []int{500}
if err := a.onItsGPUs(context.Background(), eng, spec, 500); err == nil || !strings.Contains(err.Error(), "GPU 1 ") {
t.Fatalf("a shard on half its GPUs: %v", err)
}
// Another process on GPU 1 is not the shard.
procs[1] = []int{900}
if err := a.onItsGPUs(context.Background(), eng, spec, 500); err == nil {
t.Fatal("another process's memory was taken for the shard's")
}
// A child of the shard holding GPU 1 is the shard.
procs[1] = []int{900, 512}
if err := a.onItsGPUs(context.Background(), eng, spec, 500); err != nil {
t.Fatalf("a shard on both its GPUs: %v", err)
}
// An engine that fails rather than falling back is not checked, nor is a
// shard with no GPUs.
procs = map[int][]int{}
if err := a.onItsGPUs(context.Background(), struct{ engines.Engine }{}, spec, 500); err != nil {
t.Errorf("checked an engine that does not fall back: %v", err)
}
if err := a.onItsGPUs(context.Background(), eng, engines.ShardSpec{}, 500); err != nil {
t.Errorf("checked a CPU shard: %v", err)
}
}
func TestTheCheckWaitsForTheShardToShowUpOnItsGPU(t *testing.T) {
calls := 0
oldProcs, oldWait := gpuProcesses, gpuCheckWait
gpuProcesses = func(context.Context, string) (map[int][]int, error) {
calls++
if calls < 2 {
return map[int][]int{}, nil
}
return map[int][]int{0: {77}}, nil
}
gpuCheckWait = 5 * time.Second
t.Cleanup(func() { gpuProcesses, gpuCheckWait = oldProcs, oldWait })
a := &Agent{log: slog.Default()}
if err := a.onItsGPUs(context.Background(), fallbackEngine{}, engines.ShardSpec{GPUIndices: []int{0}}, 77); err != nil {
t.Fatalf("nvidia-smi lagging by a second failed the shard: %v", err)
}
}
func TestForeignGPUMemoryLeavesOutTheAgentsOwnShards(t *testing.T) {
// cuda01: ComfyUI and one of our engines on the same card. Only ComfyUI's
// share — and whatever no listed process accounts for — is foreign.
oldMem, oldGroup := gpuMemory, processGroup
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) {
return []discovery.GPUMemory{
{UsedMB: 15000, Procs: []discovery.GPUProcess{{PID: 71315, UsedMB: 4932}, {PID: 512, UsedMB: 9500}}},
{UsedMB: 300},
}, nil
}
processGroup = func(pid int) int {
if pid >= 500 && pid < 600 {
return 500
}
return pid
}
t.Cleanup(func() { gpuMemory, processGroup = oldMem, oldGroup })
a := &Agent{log: slog.Default()}
foreign, used := a.foreignGPUMemory(context.Background(), []int{500})
if !slices.Equal(foreign, []int32{5500, 300}) || !slices.Equal(used, []int32{15000, 300}) {
t.Errorf("foreign %v, used %v; want [5500 300] and [15000 300]", foreign, used)
}
// No nvidia-smi to ask is no reading, not an empty card.
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) { return nil, nil }
if foreign, _ := a.foreignGPUMemory(context.Background(), nil); foreign != nil {
t.Errorf("no reading came back as %v", foreign)
}
}
func TestAnEngineStartedInAGroupOfItsOwnIsStillTheShards(t *testing.T) {
// cuda01: Strata's server (259709) starts its engine (261072) in a
// group of the engine's own, and the engine holds the card. It read as
// a shard on the CPU, and as 22 GB of somebody else's memory.
oldMem, oldGroup, oldParent := gpuMemory, processGroup, parentOf
t.Cleanup(func() { gpuMemory, processGroup, parentOf = oldMem, oldGroup, oldParent })
processGroup = func(pid int) int { return pid }
parentOf = func(pid int) int {
return map[int]int{261072: 259709, 259709: 258734, 258734: 1, 999: 1}[pid]
}
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) {
return []discovery.GPUMemory{{Index: 0, UsedMB: 23000,
Procs: []discovery.GPUProcess{{PID: 261072, UsedMB: 22226}, {PID: 999, UsedMB: 700}}}}, nil
}
a := &Agent{log: slog.Default()}
if c, ok := a.gpuCompute(context.Background(), []int{0}, 259709); !ok || c.Backend != "CUDA" ||
len(c.Devices) != 1 || c.Devices[0] != "GPU 0: 21.7 GB" {
t.Errorf("the engine's child on GPU 0: %+v %v", c, ok)
}
if foreign, _ := a.foreignGPUMemory(context.Background(), []int{259709}); !slices.Equal(foreign, []int32{774}) {
t.Errorf("foreign %v; want only what the other process and nobody holds, [774]", foreign)
}
}
func TestLlamaCppsOwnWordOnTheCPUIsEnoughToRefuse(t *testing.T) {
eng := fallbackEngine{}
// Without nvidia-smi to ask, the engine's own output says it.
c, said, err := engineCompute(eng, []int{0}, []string{
"warning: no usable GPU found, --gpu-layers option will be ignored",
"load_tensors: offloaded 0/49 layers to GPU"})
if err == nil || !said || !strings.Contains(err.Error(), "computing on the CPU") || c.Backend != "CPU" {
t.Fatalf("a GPU shard llama.cpp put on the CPU: %v %+v", err, c)
}
c, said, err = engineCompute(eng, []int{0}, []string{
"llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free",
"load_tensors: offloaded 49/49 layers to GPU"})
if err != nil || !said || c.Backend != "CUDA" || c.Offloaded != 49 || computeToWire(&c).GetLayersOffloaded() != 49 {
t.Errorf("on its GPU: %v %+v", err, c)
}
// A CPU shard is not refused for being on the CPU, and says so.
if c, _, err := engineCompute(eng, nil, []string{"load_tensors: offloaded 0/49 layers to GPU"}); err != nil || c.Backend != "CPU" {
t.Errorf("a CPU shard: %v %+v", err, c)
}
// An engine that said nothing says nothing.
if _, said, err := engineCompute(eng, []int{0}, []string{"listening"}); err != nil || said {
t.Errorf("nothing said: %v", err)
}
}
func TestNvidiaSMISaysWhereAShardComputes(t *testing.T) {
oldMem, oldGroup := gpuMemory, processGroup
t.Cleanup(func() { gpuMemory, processGroup = oldMem, oldGroup })
processGroup = func(pid int) int { return pid }
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) {
return []discovery.GPUMemory{
{Index: 0, UsedMB: 20000, Procs: []discovery.GPUProcess{{PID: 700, UsedMB: 2048}, {PID: 900, UsedMB: 17000}}},
{Index: 1, UsedMB: 0},
}, nil
}
a := &Agent{log: slog.Default()}
if c, ok := a.gpuCompute(context.Background(), []int{0}, 700); !ok || c.Backend != "CUDA" ||
len(c.Devices) != 1 || c.Devices[0] != "GPU 0: 2.0 GB" {
t.Errorf("on GPU 0: %+v %v", c, ok)
}
if c, ok := a.gpuCompute(context.Background(), []int{1}, 700); !ok || c.Backend != "CPU" {
t.Errorf("placed on GPU 1, holding none of it: %+v", c)
}
if c, ok := a.gpuCompute(context.Background(), nil, 700); !ok || c.Backend != "CPU" {
t.Errorf("a CPU shard: %+v", c)
}
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) { return nil, nil }
if _, ok := a.gpuCompute(context.Background(), []int{0}, 700); ok {
t.Error("no nvidia-smi read as an answer")
}
}