llmgw/internal/control/catalog/hardware_test.go

package catalog

import (
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
)

// sf3dByCard is one entry for every site: fp16 on anything with 16 GB, and
// a bf16 workflow with its own weights where the card is Ampere or newer.
func sf3dByCard() Model {
	return Model{Name: "sf3d", Source: "/w/sf3d/model.safetensors",
		ComfyUI: &ComfyUI{Workflow: "/etc/llmgw/workflows/sf3d.json", Output: OutputMeshes},
		Profiles: []Profile{
			{Name: "ampere", Engine: EngineComfyUI, TP: 1, Priority: 10, GPU: &GPUNeed{MinComputeCapability: 8.0, MinMemoryGB: 16},
				Source:  "/w/sf3d-bf16/model-bf16.safetensors",
				ComfyUI: &ProfileComfyUI{Workflow: "/etc/llmgw/workflows/sf3d-bf16.json"}},
			{Name: "any", Engine: EngineComfyUI, TP: 1, GPU: &GPUNeed{MinMemoryGB: 16}},
		}}
}

func TestAProfileThatNamesItsCardByNeedStandsForEveryClassThatMeetsIt(t *testing.T) {
	cat, err := New([]Model{sf3dByCard()})
	if err != nil {
		t.Fatal(err)
	}
	m, _ := cat.Resolve("sf3d")
	turing := GPUKinds([]discovery.GPU{{Name: "Quadro RTX 6000", MemoryTotalMB: 24576, ComputeCapability: "7.5"}})
	if len(turing) != 1 || turing[0].Class != "6000-24g" || turing[0].ComputeCapability != 7.5 || turing[0].MemoryGB != 24 {
		t.Fatalf("kinds %+v", turing)
	}
	ps := m.PlacementProfilesFor(turing)
	if len(ps) != 1 || ps[0].Name != "any@6000-24g" || ps[0].GPUClass != "6000-24g" {
		t.Fatalf("on Turing: %+v", ps)
	}

	both := GPUKinds([]discovery.GPU{
		{Name: "Quadro RTX 6000", MemoryTotalMB: 24576, ComputeCapability: "7.5"},
		{Name: "NVIDIA A100-SXM4-80GB", MemoryTotalMB: 81920, ComputeCapability: "8.0"},
		{Name: "NVIDIA A100-SXM4-80GB", MemoryTotalMB: 81920, ComputeCapability: "8.0"},
		{Name: "Tesla T4", MemoryTotalMB: 15360, ComputeCapability: "7.5"},
	})
	var names []string
	for _, p := range m.PlacementProfilesFor(both) {
		names = append(names, p.Name+"/"+p.GPUClass)
	}
	if got := strings.Join(names, " "); got != "ampere@a100-80g/a100-80g any@6000-24g/6000-24g any@a100-80g/a100-80g" {
		t.Errorf("on a mixed cluster: %s (the T4 has 15 GB, short of 16)", got)
	}
	if len(m.PlacementProfiles()) != 0 {
		t.Error("with no cards known, a profile by need stood for something")
	}

	// A replica placed as ampere@a100-80g is started with that profile's own
	// weights, workflow and class.
	p, ok := m.ProfileByName("ampere@a100-80g")
	if !ok || p.GPUClass != "a100-80g" || p.GPU != nil || m.SourceOf(p) != "/w/sf3d-bf16/model-bf16.safetensors" {
		t.Fatalf("profile %+v %v", p, ok)
	}
	if w, ckpt := m.ComfyWorkflowFor(p); w != "/etc/llmgw/workflows/sf3d-bf16.json" || ckpt != "model-bf16.safetensors" {
		t.Errorf("workflow %q checkpoint %q", w, ckpt)
	}
	if ep, ok := m.EngineProfile("ampere@a100-80g"); !ok || ep.GPUClass != "a100-80g" || ep.ModelPath != "/w/sf3d-bf16/model-bf16.safetensors" {
		t.Errorf("engine profile %+v", ep)
	}
	if w, ckpt := m.ComfyWorkflowFor(m.Profiles[1]); w != "/etc/llmgw/workflows/sf3d.json" || ckpt != "model.safetensors" {
		t.Errorf("the fallback's workflow %q checkpoint %q", w, ckpt)
	}
	if _, ok := m.ProfileByName("nothing@a100-80g"); ok {
		t.Error("a profile that is not there was found")
	}
}

func TestAProfileByNeedIsCheckedLikeAnyOther(t *testing.T) {
	for name, c := range map[string]struct {
		p    Profile
		want string
	}{
		"both":           {Profile{Name: "p", Engine: EngineComfyUI, TP: 1, GPUClass: "a100-80g", GPU: &GPUNeed{MinMemoryGB: 8}}, "one or the other"},
		"an @ in a name": {Profile{Name: "p@x", Engine: EngineComfyUI, TP: 1, GPUClass: "a100-80g"}, "no '@'"},
		"a capability of 80": {Profile{Name: "p", Engine: EngineComfyUI, TP: 1, GPU: &GPUNeed{MinComputeCapability: 80}},
			"compute capability"},
		"comfyui on vllm": {Profile{Name: "p", Engine: "vllm", TP: 1, GPU: &GPUNeed{}, ComfyUI: &ProfileComfyUI{}},
			"engine comfyui"},
	} {
		m := sf3dByCard()
		m.Profiles = []Profile{c.p}
		if c.p.Engine == "vllm" {
			m.ComfyUI = nil
		}
		if _, err := New([]Model{m}); err == nil || !strings.Contains(err.Error(), c.want) {
			t.Errorf("%s: %v", name, err)
		}
	}
	need := GPUNeed{MinComputeCapability: 8.9}
	if need.Fits(GPUKind{Class: "x", MemoryGB: 80}) {
		t.Error("a card that does not say its capability was taken for an Ada")
	}
	if !need.Fits(GPUKind{ComputeCapability: 9.0}) || need.String() != "compute 8.9+" {
		t.Error("a Hopper is not an Ada or newer")
	}
}