llmgw/internal/control/adminapi/presets_test.go

package adminapi

import (
	"net/http"
	"path/filepath"
	"slices"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

func v100Node(id string) inventory.Node {
	return inventory.Node{ID: id, State: inventory.Ready, AgentVersion: "v1", GPUs: []discovery.GPU{{
		Index: 0, Name: "Tesla V100-SXM2-32GB", MemoryTotalMB: 32768, ComputeCapability: "7.0"}}}
}

func presetsAPI(t *testing.T, models []catalog.Model, cfg PresetConfig, nodes ...inventory.Node) (*API, *fakeModels, *store.DB) {
	t.Helper()
	a, _ := setup(t, models, nodes...)
	st, err := store.OpenSQLite(t.Context(), filepath.Join(t.TempDir(), "p.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = st.Close() })
	fm := &fakeModels{}
	a.SetModels(fm)
	a.SetPresets(st, func() PresetConfig { return cfg })
	a.SetSettings(&Settings{Store: st, Check: func([]byte) error { return nil }})
	return a, fm, st
}

func TestTheLeaderAddsThePresetsANodeFitsAndNoOtherWhere(t *testing.T) {
	// A customer's portal had no pictures or 3D pages: nobody had written a
	// definition for either, on a cluster with a V100 that runs both.
	a, fm, st := presetsAPI(t, testModels(), PresetConfig{Auto: true, ComfyUI: "0.39.2"}, v100Node("p-g3"))
	added, err := a.PresetsPass(t.Context())
	if err != nil || !slices.Equal(added, []string{"flux-schnell", "triposr"}) {
		t.Fatalf("added %v %v", added, err)
	}
	if !slices.Equal(fm.actors, []string{store.PresetAdder, store.PresetAdder}) {
		t.Errorf("added by %v", fm.actors)
	}
	states, _ := st.PresetStates(t.Context())
	if states["triposr"].AddedBy != store.PresetAdder || states["triposr"].AddedAt.IsZero() {
		t.Errorf("state %+v", states["triposr"])
	}
	events, _ := st.AuditEvents(t.Context(), store.AuditFilter{Limit: 10})
	if len(events) != 2 || events[0].Actor != store.PresetAdder || !strings.Contains(events[0].Detail, "v100-32g") {
		t.Errorf("audit %+v", events)
	}
}

func TestAPresetIsNotAddedByItselfWhereItShouldNotBe(t *testing.T) {
	picture := catalog.Model{Name: "sdxl", Source: "/w/sdxl.safetensors", ComfyUI: &catalog.ComfyUI{Workflow: "flux-schnell"},
		Profiles: []catalog.Profile{{Name: "p", Engine: catalog.EngineComfyUI, GPUClass: "v100-32g", TP: 1}}}
	cases := []struct {
		name   string
		models []catalog.Model
		cfg    PresetConfig
		nodes  []inventory.Node
		want   map[string]string // preset -> part of why it is not added
	}{
		{"auto off", testModels(), PresetConfig{Auto: false, ComfyUI: "0.39.2"}, []inventory.Node{v100Node("g")},
			map[string]string{"flux-schnell": "presets.auto is off", "triposr": "presets.auto is off"}},
		{"no ComfyUI", testModels(), PresetConfig{Auto: true}, []inventory.Node{v100Node("g")},
			map[string]string{"triposr": "engines.comfyui"}},
		{"no card", testModels(), PresetConfig{Auto: true, ComfyUI: "0.39.2"}, nil,
			map[string]string{"triposr": "no node has a card"}},
		{"a picture model already", append(testModels(), picture), PresetConfig{Auto: true, ComfyUI: "0.39.2"},
			[]inventory.Node{v100Node("g")}, map[string]string{"flux-schnell": "sdxl"}},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			a, fm, _ := presetsAPI(t, c.models, c.cfg, c.nodes...)
			v, err := a.Presets(t.Context())
			if err != nil {
				t.Fatal(err)
			}
			if v.ComfyUI != c.cfg.ComfyUI {
				t.Errorf("comfyui %q", v.ComfyUI)
			}
			for _, p := range v.Presets {
				if why, ok := c.want[p.Name]; ok && !strings.Contains(p.Blocked, why) {
					t.Errorf("%s blocked %q, want it to say %q", p.Name, p.Blocked, why)
				}
			}
			_, _ = a.PresetsPass(t.Context())
			for _, m := range fm.added {
				if _, blocked := c.want[m.Name]; blocked {
					t.Errorf("%s was added", m.Name)
				}
			}
		})
	}
}

func TestAPresetRemovedStaysOutUntilAddedAgain(t *testing.T) {
	a, fm, _ := presetsAPI(t, testModels(), PresetConfig{Auto: true, ComfyUI: "0.39.2"}, v100Node("g"))
	if code := call(t, a, http.MethodDelete, "/api/v1/models/triposr", "", nil); code != http.StatusOK {
		t.Fatalf("remove %d", code)
	}
	var v PresetsView
	call(t, a, http.MethodGet, "/api/v1/presets", "", &v)
	i := slices.IndexFunc(v.Presets, func(p PresetView) bool { return p.Name == "triposr" })
	if i < 0 || !v.Presets[i].Declined || !strings.Contains(v.Presets[i].Blocked, "removed") {
		t.Fatalf("after removal %+v", v.Presets)
	}
	added, _ := a.PresetsPass(t.Context())
	if slices.Contains(added, "triposr") {
		t.Error("a removed preset was added back by itself")
	}
	if code := call(t, a, http.MethodPost, "/api/v1/presets/triposr", "", nil); code != http.StatusOK {
		t.Fatalf("add again %d", code)
	}
	if fm.added[len(fm.added)-1].Name != "triposr" {
		t.Error("adding it again did not add it")
	}
	call(t, a, http.MethodGet, "/api/v1/presets", "", &v)
	i = slices.IndexFunc(v.Presets, func(p PresetView) bool { return p.Name == "triposr" })
	if v.Presets[i].Declined {
		t.Error("adding it again left it declined")
	}
	if code := call(t, a, http.MethodPost, "/api/v1/presets/nope", "", nil); code != http.StatusNotFound {
		t.Errorf("unknown preset %d", code)
	}
}

func TestPresetsAutoIsTheClusterSetting(t *testing.T) {
	a, _, st := presetsAPI(t, testModels(), PresetConfig{Auto: true, ComfyUI: "0.39.2"})
	if code := call(t, a, http.MethodPut, "/api/v1/presets/auto", `{"on":false}`, nil); code != http.StatusOK {
		t.Fatalf("put %d", code)
	}
	cur, _ := st.ClusterSettings(t.Context())
	if !strings.Contains(cur.Body, "auto: false") {
		t.Errorf("settings %q", cur.Body)
	}
}