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)
	}
}