CI / test (push) Failing after 7m31s
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>
201 lines
8.1 KiB
Go
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")
|
|
}
|
|
}
|