llmgw/internal/control/modelserve/gpus_test.go

package modelserve

import (
	"context"
	"slices"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
)

// gpuSetup is the update setup with qwen on one L40S card of llama.cpp, and
// a pool of one node of four, of which free are free.
func gpuSetup(t *testing.T, free int) (*Server, *updModels, *restartFleet, *[]string) {
	t.Helper()
	s, _, models, _, fleet := updSetup(t)
	m := models.models["qwen"]
	m.ContextLength = 8192
	m.Origin = &catalog.Origin{Repo: "o/r", Files: []catalog.OriginFile{{Path: "q4.gguf", Bytes: 20 << 30}}}
	m.Profiles = []catalog.Profile{
		{Name: "cpu", Engine: "llamacpp", MinMemoryGB: 8},
		{Name: "gpu", Engine: "llamacpp", GPUClass: "l40s-48g", TP: 1, GPUsPerNode: 1,
			Args: []string{"--ctx-size", "8192", "--tensor-split", "1"}},
	}
	models.models["qwen"] = m
	s.deps.Pool = func() []modelplan.PoolNode {
		return []modelplan.PoolNode{{ID: "g1", Class: "l40s-48g", GPUMemGB: 48, GPUs: 4, FreeGPUs: free}}
	}
	var inPlace []string
	s.deps.RestartInPlace = func(ids ...string) error { inPlace = append(inPlace, ids...); return nil }
	return s, models, fleet, &inPlace
}

func TestAModelIsSpreadOverMoreCardsAndItsReplicasStartedAgainBesideTheOld(t *testing.T) {
	s, models, fleet, inPlace := gpuSetup(t, 3)
	got, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2}, "admin")
	if err != nil {
		t.Fatal(err)
	}
	if got.Profile != "gpu" || got.Before != 1 || got.After != 2 || got.Restarted != 2 || got.InPlace {
		t.Fatalf("%+v", got)
	}
	p := models.models["qwen"].Profiles[1]
	// llama.cpp splits layers over a node's cards: gpus_per_node, not tp,
	// and the old card's proportions are gone.
	if p.GPUsPerNode != 2 || p.TP != 1 || p.Nodes != 1 || !slices.Equal(p.Args, []string{"--ctx-size", "8192"}) {
		t.Errorf("profile %+v", p)
	}
	if !strings.Contains(got.PerCard, "a card of 2") {
		t.Errorf("per card %q", got.PerCard)
	}
	if len(*inPlace) != 0 {
		t.Errorf("stopped first though there was room: %v", *inPlace)
	}
	for _, r := range fleet.Replicas() {
		if !strings.Contains(r.ID, "-new") {
			t.Errorf("%s was not started again", r.ID)
		}
	}
	if models.models["qwen"].Profiles[0].GPUsPerNode != 0 {
		t.Error("the CPU profile was changed too")
	}
}

func TestWithNoRoomBesideIt_InPlaceIsAskedFor(t *testing.T) {
	// One card free and the model's replicas on the others: two do not fit
	// beside them.
	s, _, fleet, inPlace := gpuSetup(t, 1)
	fleet.reps[0].Shards = []reconciler.Shard{{NodeID: "g1", GPUIndices: []int{1}}}
	fleet.reps[1].Shards = []reconciler.Shard{{NodeID: "g1", GPUIndices: []int{2}}}
	_, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2}, "admin")
	if err == nil || !strings.Contains(err.Error(), "--in-place") {
		t.Fatalf("not asked for in place: %v", err)
	}
	got, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2, InPlace: true}, "admin")
	if err != nil {
		t.Fatal(err)
	}
	if !got.InPlace || len(*inPlace) != 2 {
		t.Errorf("%+v, stopped first: %v", got, *inPlace)
	}
}

func TestMoreCardsThanANodeHasOrOthersHoldAreRefused(t *testing.T) {
	s, models, _, _ := gpuSetup(t, 4)
	if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 8}, "admin"); err == nil ||
		!strings.Contains(err.Error(), "the most one has is 4") {
		t.Errorf("eight cards on a node of four: %v", err)
	}
	s, _, _, _ = gpuSetup(t, 0)
	if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2, InPlace: true}, "admin"); err == nil ||
		!strings.Contains(err.Error(), "other models hold them") {
		t.Errorf("cards other models hold: %v", err)
	}
	if models.models["qwen"].Profiles[1].GPUsPerNode != 1 {
		t.Error("a refused change was kept")
	}
}

func TestFewerCardsThanTheModelNeedsAreRefusedWithTheFigures(t *testing.T) {
	s, models, _, _ := gpuSetup(t, 3)
	m := models.models["qwen"]
	m.Profiles[1].GPUsPerNode, m.Origin.Files[0].Bytes = 4, 120<<30
	models.models["qwen"] = m
	_, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2}, "admin")
	if err == nil || !strings.Contains(err.Error(), "GB") {
		t.Errorf("120 GB on two 48 GB cards: %v", err)
	}
}

func TestTensorParallelFollowsTheKVHeads(t *testing.T) {
	p := catalog.Profile{Name: "gpu", Engine: "vllm", GPUClass: "h100-80g", TP: 2, GPUsPerNode: 2,
		Args: []string{"--tensor-parallel-size", "2", "--max-model-len", "32768"}}
	if _, err := respread(p, 3, shape{kvHeads: 8}); err == nil || !strings.Contains(err.Error(), "8 KV heads") {
		t.Errorf("three ways for eight heads: %v", err)
	}
	next, err := respread(p, 4, shape{kvHeads: 8})
	if err != nil {
		t.Fatal(err)
	}
	// The agent sets the width from tp; the profile's own spelling went.
	if next.TP != 4 || next.GPUsPerNode != 4 || !slices.Equal(next.Args, []string{"--max-model-len", "32768"}) {
		t.Errorf("%+v", next)
	}
	for _, eng := range []string{"comfyui", "strata"} {
		if _, err := respread(catalog.Profile{Name: "x", Engine: eng}, 2, shape{}); err == nil {
			t.Errorf("%s spread over two cards", eng)
		}
	}
}

func TestAModelOnCPUsOrWithSeveralCardProfilesSaysWhy(t *testing.T) {
	s, models, _, _ := gpuSetup(t, 3)
	m := models.models["qwen"]
	m.Profiles = append(m.Profiles, catalog.Profile{Name: "gpu2", Engine: "vllm", GPUClass: "l40s-48g", TP: 1})
	models.models["qwen"] = m
	if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2}, "admin"); err == nil ||
		!strings.Contains(err.Error(), "name one") {
		t.Errorf("two profiles on cards and none named: %v", err)
	}
	if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2, Profile: "cpu"}, "admin"); err == nil ||
		!strings.Contains(err.Error(), "gpu_class") {
		t.Errorf("a CPU profile: %v", err)
	}
	if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 1, Profile: "gpu"}, "admin"); err == nil ||
		!strings.Contains(err.Error(), "already") {
		t.Errorf("the same count: %v", err)
	}
}