llmgw/internal/control/modelsvc/modelsvc_test.go

package modelsvc

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"strings"
	"testing"
	"time"

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

type memStore struct {
	rows map[string]store.StoredModel
	err  error
}

func newMemStore() *memStore { return &memStore{rows: map[string]store.StoredModel{}} }

func (m *memStore) PutModel(_ context.Context, s store.StoredModel) error {
	if m.err != nil {
		return m.err
	}
	m.rows[s.Name] = s
	return nil
}

func (m *memStore) DeleteModel(_ context.Context, name string) error {
	delete(m.rows, name)
	return nil
}

func (m *memStore) Models(context.Context) ([]store.StoredModel, error) {
	out := make([]store.StoredModel, 0, len(m.rows))
	for _, r := range m.rows {
		out = append(out, r)
	}
	return out, nil
}

type recorder struct{ last *catalog.Catalog }

func (r *recorder) ReplaceCatalog(_ context.Context, c *catalog.Catalog) { r.last = c }

func model(name string) catalog.Model {
	return catalog.Model{
		Name: name, Source: "/models/" + name,
		Scaling:  catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1, MinMemoryGB: 2}},
	}
}

func newService(t *testing.T, declared ...catalog.Model) (*Service, *memStore, *recorder) {
	t.Helper()
	st, rec := newMemStore(), &recorder{}
	return New(declared, st, rec, slog.New(slog.DiscardHandler)), st, rec
}

func TestAddedModelReachesTheCatalogAndSurvivesARestart(t *testing.T) {
	// The whole point: trying a model should not mean a config change, a
	// restart and an outage — and it should still be there tomorrow.
	s, st, rec := newService(t, model("declared"))

	if err := s.Add(t.Context(), model("added"), "ilker"); err != nil {
		t.Fatal(err)
	}
	if _, ok := rec.last.Resolve("added"); !ok {
		t.Error("the new model is not in the catalog that was applied")
	}
	if _, ok := rec.last.Resolve("declared"); !ok {
		t.Error("the declared model was lost")
	}
	if _, ok := st.rows["added"]; !ok {
		t.Error("the model was not persisted")
	}

	// A fresh service over the same store, as after a restart.
	s2, _, _ := newService(t, model("declared"))
	s2.store = st
	built, err := s2.Load(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if _, ok := built.Resolve("added"); !ok {
		t.Error("the model did not survive a restart")
	}
}

func TestConfigurationWinsOverTheDatabase(t *testing.T) {
	// The dangerous direction: a request for a reviewed production model must
	// never be answered by something added from a laptop.
	s, _, _ := newService(t, model("shared"))
	err := s.Add(t.Context(), model("shared"), "ilker")
	if err == nil {
		t.Fatal("a stored model was allowed to shadow a declared one")
	}
	if !strings.Contains(err.Error(), "declared in the cluster's settings") {
		t.Errorf("the error does not say why: %v", err)
	}
}

func TestADeclaredModelCannotBeRemovedAtRuntime(t *testing.T) {
	// It would come back on the next restart, so the cluster would disagree
	// with its own config file until then.
	s, _, _ := newService(t, model("declared"))
	err := s.Remove(t.Context(), "declared")
	if err == nil {
		t.Fatal("a declared model was removed at runtime")
	}
	if !strings.Contains(err.Error(), "llmgw ctl config edit") {
		t.Errorf("the error does not explain the consequence: %v", err)
	}
}

func TestNothingIsPersistedWhenTheCatalogWouldBeInvalid(t *testing.T) {
	// A model that is written and then skipped on every start is a model that
	// looks added and never appears.
	s, st, rec := newService(t)
	bad := model("bad")
	bad.Fallbacks = []string{"a-model-that-does-not-exist"}

	if err := s.Add(t.Context(), bad, "ilker"); err == nil {
		t.Fatal("a model naming a missing fallback was accepted")
	}
	if len(st.rows) != 0 {
		t.Errorf("it was persisted anyway: %v", st.rows)
	}
	if rec.last != nil {
		t.Error("a catalog was applied for a model that was rejected")
	}
}

func TestRemovingAModelOthersFallBackToIsRefused(t *testing.T) {
	// Leaving a fallback pointing at nothing is discovered during the outage it
	// was meant to cover.
	s, _, _ := newService(t)
	if err := s.Add(t.Context(), model("small"), "ilker"); err != nil {
		t.Fatal(err)
	}
	big := model("big")
	big.Fallbacks = []string{"small"}
	if err := s.Add(t.Context(), big, "ilker"); err != nil {
		t.Fatal(err)
	}

	if err := s.Remove(t.Context(), "small"); err == nil {
		t.Fatal("removed a model another falls back to")
	}
}

func TestOneUnreadableStoredModelDoesNotStopTheCluster(t *testing.T) {
	// Refusing to start over one bad row would let a model somebody added from
	// a laptop take the cluster down, and the declared models are precisely the
	// ones that must survive that.
	s, st, _ := newService(t, model("declared"))
	good, _ := json.Marshal(model("good"))
	st.rows["good"] = store.StoredModel{Name: "good", Definition: good, CreatedAt: time.Now()}
	st.rows["broken"] = store.StoredModel{
		Name: "broken", Definition: []byte(`{"name":`), CreatedAt: time.Now(),
	}

	built, err := s.Load(t.Context())
	if err != nil {
		t.Fatalf("one bad row stopped the catalog from loading: %v", err)
	}
	for _, name := range []string{"declared", "good"} {
		if _, ok := built.Resolve(name); !ok {
			t.Errorf("%s was lost alongside the bad row", name)
		}
	}
	if _, ok := built.Resolve("broken"); ok {
		t.Error("the unreadable model was loaded anyway")
	}
}

func TestAStoredModelThatNoLongerValidatesIsSkippedAlone(t *testing.T) {
	// Checked one at a time and against what is already accepted, so one bad
	// model cannot take a good one with it.
	s, st, _ := newService(t)
	bad := model("bad")
	bad.Profiles = nil // a model with no profile cannot be placed
	badJSON, _ := json.Marshal(bad)
	goodJSON, _ := json.Marshal(model("good"))
	st.rows["bad"] = store.StoredModel{Name: "bad", Definition: badJSON, CreatedAt: time.Now()}
	st.rows["good"] = store.StoredModel{Name: "good", Definition: goodJSON, CreatedAt: time.Now().Add(time.Second)}

	built, err := s.Load(t.Context())
	if err != nil {
		t.Fatalf("load failed: %v", err)
	}
	if _, ok := built.Resolve("good"); !ok {
		t.Error("a valid stored model was dropped because another was invalid")
	}
}

func TestScalingChangesTheCountAndSurvivesARestart(t *testing.T) {
	s, st, rec := newService(t)
	if err := s.Add(t.Context(), model("added"), "ilker"); err != nil {
		t.Fatal(err)
	}

	got, err := s.Scale(t.Context(), "added", catalog.Scaling{Min: 3}, "ilker")
	if err != nil {
		t.Fatal(err)
	}
	// Max follows min when it is left out, and the answer reports what the
	// catalog settled on rather than what was asked for — a caller echoing its
	// own request would print a max of zero.
	if got.Scaling.Min != 3 || got.Scaling.Max != 3 {
		t.Errorf("expected 3/3 in force, got %d/%d", got.Scaling.Min, got.Scaling.Max)
	}
	if m, _ := rec.last.Resolve("added"); m.Scaling.Min != 3 {
		t.Errorf("the applied catalog still says %d", m.Scaling.Min)
	}

	s2, _, _ := newService(t)
	s2.store = st
	built, err := s2.Load(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if m, _ := built.Resolve("added"); m.Scaling.Min != 3 {
		t.Errorf("the new count did not survive a restart; got %d", m.Scaling.Min)
	}
}

func TestScalingChangesNothingButTheCount(t *testing.T) {
	s, _, rec := newService(t)
	m := model("added")
	m.Aliases = []string{"a"}
	m.Profiles[0].Args = []string{"--ctx-size", "16384"}
	if err := s.Add(t.Context(), m, "ilker"); err != nil {
		t.Fatal(err)
	}

	if _, err := s.Scale(t.Context(), "added", catalog.Scaling{Min: 2}, "ilker"); err != nil {
		t.Fatal(err)
	}

	// Changing a number must not quietly become a redefinition: the source and
	// the engine arguments somebody tuned are read back and written out again.
	got, _ := rec.last.Resolve("added")
	if got.Source != "/models/added" {
		t.Errorf("the source changed to %q", got.Source)
	}
	if len(got.Profiles) != 1 || strings.Join(got.Profiles[0].Args, " ") != "--ctx-size 16384" {
		t.Errorf("the profile arguments were lost: %+v", got.Profiles)
	}
	if len(got.Aliases) != 1 || got.Aliases[0] != "a" {
		t.Errorf("the aliases were lost: %v", got.Aliases)
	}
}

func TestADeclaredModelCannotBeScaledHere(t *testing.T) {
	s, _, _ := newService(t, model("declared"))

	_, err := s.Scale(t.Context(), "declared", catalog.Scaling{Min: 5}, "ilker")

	// The YAML is read on every start, so the number would hold until the next
	// restart and then silently revert. Refusing and saying where to edit is
	// the smaller surprise — the same call Remove makes.
	if err == nil {
		t.Fatal("expected scaling a declared model to be refused")
	}
	if !strings.Contains(err.Error(), "cluster's settings") {
		t.Errorf("the error should say where to change it; got %q", err)
	}
}

func TestScalingAModelThatWasNeverAddedIsAnError(t *testing.T) {
	s, _, _ := newService(t)
	if _, err := s.Scale(t.Context(), "ghost", catalog.Scaling{Min: 1}, "ilker"); err == nil {
		t.Error("expected an error for a model that is not there")
	}
}

func TestAMaxBelowAMinIsRefusedBeforeAnythingIsWritten(t *testing.T) {
	s, st, _ := newService(t)
	if err := s.Add(t.Context(), model("added"), "ilker"); err != nil {
		t.Fatal(err)
	}

	if _, err := s.Scale(t.Context(), "added", catalog.Scaling{Min: 3, Max: 1}, "ilker"); err == nil {
		t.Fatal("expected a max below a min to be refused")
	}

	// Refused before the write, so the stored definition is untouched: a
	// rejected change that had already been persisted would come back on the
	// next restart as a catalog that does not build.
	var stored catalog.Model
	if err := json.Unmarshal(st.rows["added"].Definition, &stored); err != nil {
		t.Fatal(err)
	}
	if stored.Scaling.Min != 1 {
		t.Errorf("the stored definition was changed to %d despite the refusal", stored.Scaling.Min)
	}
}

func TestScalingToZeroKeepsTheModelInTheCatalog(t *testing.T) {
	s, _, rec := newService(t)
	if err := s.Add(t.Context(), model("added"), "ilker"); err != nil {
		t.Fatal(err)
	}

	got, err := s.Scale(t.Context(), "added", catalog.Scaling{Min: 0, Max: 0}, "ilker")
	if err != nil {
		t.Fatal(err)
	}
	// Configured and deliberately not running is a different and equally
	// useful thing from removed: the entry stays so it can be brought back
	// with one number.
	if got.Scaling.OnDemand() {
		t.Error("min 0 max 0 is not on-demand; it is off")
	}
	if _, ok := rec.last.Resolve("added"); !ok {
		t.Error("scaling to zero must not remove the model")
	}
}

func TestMinZeroWithAMaxIsOnDemand(t *testing.T) {
	s, _, _ := newService(t)
	if err := s.Add(t.Context(), model("added"), "ilker"); err != nil {
		t.Fatal(err)
	}

	got, err := s.Scale(t.Context(), "added",
		catalog.Scaling{Min: 0, Max: 2, IdleTimeout: 10 * time.Minute}, "ilker")
	if err != nil {
		t.Fatal(err)
	}
	if !got.Scaling.OnDemand() {
		t.Error("min 0 with a max above it should be on-demand")
	}
	if got.Scaling.IdleTimeout != 10*time.Minute {
		t.Errorf("the idle timeout was lost; got %v", got.Scaling.IdleTimeout)
	}
}

func TestACPUModelAddedNowMustSayHowMuchMemoryItNeeds(t *testing.T) {
	s, st, _ := newService(t)
	m := model("bare")
	m.Profiles[0].MinMemoryGB = 0
	err := s.Add(t.Context(), m, "tester")
	if err == nil || !strings.Contains(err.Error(), "declares no memory") {
		t.Fatalf("a CPU model with no memory was added: %v", err)
	}
	if len(st.rows) != 0 {
		t.Error("it was stored anyway")
	}
}

func TestResourcesAreSetOnAStoredModelAndSurviveARestart(t *testing.T) {
	s, _, _ := newService(t)
	if err := s.Add(t.Context(), model("m"), "tester"); err != nil {
		t.Fatal(err)
	}
	got, err := s.Resources(t.Context(), "m", "", 6, 3, "tester")
	if err != nil {
		t.Fatal(err)
	}
	if p := got.Profiles[0]; p.MinMemoryGB != 6 || p.CPUCores != 3 {
		t.Errorf("profile %+v", p)
	}
	if _, err := s.Resources(t.Context(), "m", "nope", 6, 3, "tester"); err == nil {
		t.Error("a profile the model does not have was accepted")
	}
	if _, err := s.Resources(t.Context(), "m", "", -1, 0, "tester"); err == nil {
		t.Error("negative memory was accepted")
	}
	cat, err := s.Load(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if m, _ := cat.Resolve("m"); m.Profiles[0].CPUCores != 3 {
		t.Error("the cores were lost on reload")
	}
}

// The leader serving a judge gives way to the one init added while it
// fetched: AddNew leaves a model already there as it is.
func TestAddNewLeavesAModelAddedMeanwhile(t *testing.T) {
	s, _, rec := newService(t)
	first := model("judge")
	first.Description = "init's"
	if err := s.Add(t.Context(), first, "init"); err != nil {
		t.Fatal(err)
	}
	second := model("judge")
	second.Description = "the leader's"
	if err := s.AddNew(t.Context(), second, "llmgw"); !errors.Is(err, catalog.ErrExists) {
		t.Fatalf("err %v", err)
	}
	if m, _ := rec.last.Resolve("judge"); m.Description != "init's" {
		t.Errorf("replaced: %q", m.Description)
	}
	if err := s.AddNew(t.Context(), model("fresh"), "llmgw"); err != nil {
		t.Errorf("a new model: %v", err)
	}
}

func TestAStoppedModelRunsNoneAndStartsAsItWas(t *testing.T) {
	s, _, rec := newService(t)
	if err := s.Add(t.Context(), model("added"), "ilker"); err != nil {
		t.Fatal(err)
	}
	if _, err := s.Scale(t.Context(), "added", catalog.Scaling{Min: 2, Max: 4, IdleTimeout: 20 * time.Minute}, "ilker"); err != nil {
		t.Fatal(err)
	}
	got, err := s.Stop(t.Context(), "added", "ilker")
	if err != nil {
		t.Fatal(err)
	}
	if got.Scaling.Min != 0 || got.Scaling.Max != 0 || got.Stopped == nil || got.Stopped.Min != 2 {
		t.Fatalf("stopped: %+v %+v", got.Scaling, got.Stopped)
	}
	if m, _ := rec.last.Resolve("added"); m.Scaling.Max != 0 || m.Scaling.OnDemand() {
		t.Errorf("the catalog in force runs it: %+v", m.Scaling)
	}
	if _, err := s.Stop(t.Context(), "added", "ilker"); err == nil {
		t.Error("stopped twice")
	}
	got, err = s.Start(t.Context(), "added", "ilker")
	if err != nil {
		t.Fatal(err)
	}
	if got.Scaling.Min != 2 || got.Scaling.Max != 4 || got.Scaling.IdleTimeout != 20*time.Minute || got.Stopped != nil {
		t.Errorf("started: %+v %+v", got.Scaling, got.Stopped)
	}
	if _, err := s.Start(t.Context(), "added", "ilker"); err == nil {
		t.Error("started what was not stopped")
	}
	// A scale while stopped starts it at that number.
	_, _ = s.Stop(t.Context(), "added", "ilker")
	if got, _ := s.Scale(t.Context(), "added", catalog.Scaling{Min: 1}, "ilker"); got.Stopped != nil {
		t.Error("still stopped after a scale")
	}
}