llmgw/internal/agent/gpucheck_test.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 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")
	}
}