llmgw/internal/control/modelserve/updates_test.go

package modelserve

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

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

const (
	oldRev = "abc123456789aaaa"
	newRev = "def987654321bbbb"
)

type movingRegistry struct {
	mu   sync.Mutex
	repo hf.Repo
}

func (r *movingRegistry) Repo(context.Context, string) (hf.Repo, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.repo, nil
}

// updMirror is a mirror that records where it fetched each file from.
type updMirror struct {
	mu   sync.Mutex
	base string
	held []mirror.Status
	sums map[string]string
}

func (m *updMirror) Fetch(_ context.Context, r mirror.FetchRequest) (mirror.FetchResponse, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.held = append(m.held, mirror.Status{Model: r.Model, Version: r.Version, File: r.File, Source: r.Source,
		State: mirror.Present, BytesDone: 1 << 30, SHA256: m.sums[r.Version+"/"+r.File]})
	return mirror.FetchResponse{URL: mirror.Entry{Model: r.Model, Version: r.Version, File: r.File}.URL(m.base)}, nil
}

func (m *updMirror) List(context.Context) (mirror.ModelsResponse, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return mirror.ModelsResponse{Base: m.base, Models: append([]mirror.Status{}, m.held...)}, nil
}

type updModels struct {
	mu       sync.Mutex
	models   map[string]catalog.Model
	declared map[string]bool
}

func (c *updModels) Declared(n string) bool { return c.declared[n] }
func (c *updModels) SetSource(_ context.Context, n, src string, o *catalog.Origin, _ string) (catalog.Model, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	m := c.models[n]
	m.Source, m.Origin = src, o
	c.models[n] = m
	return m, nil
}
func (c *updModels) SetOrigin(_ context.Context, n string, o *catalog.Origin) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	m := c.models[n]
	m.Origin = o
	c.models[n] = m
	return nil
}
func (c *updModels) all() []catalog.Model {
	c.mu.Lock()
	defer c.mu.Unlock()
	var out []catalog.Model
	for _, m := range c.models {
		out = append(out, m)
	}
	return out
}

// restartFleet replaces a restarted replica with a new one, ready.
type restartFleet struct {
	mu     sync.Mutex
	reps   []reconciler.Replica
	engine string
	seq    int
}

func (f *restartFleet) Replicas() []reconciler.Replica {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]reconciler.Replica{}, f.reps...)
}
func (f *restartFleet) ModelHealth() []reconciler.ModelHealth { return nil }
func (f *restartFleet) EngineKey(string) string               { return "" }
func (f *restartFleet) Restart(ids ...string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, id := range ids {
		for i, r := range f.reps {
			if r.ID == id {
				f.seq++
				f.reps[i] = reconciler.Replica{ID: r.Model + "-new" + string(rune('0'+f.seq)), Model: r.Model,
					State: reconciler.Ready, Endpoint: f.engine}
			}
		}
	}
	return nil
}

func updSetup(t *testing.T) (*Server, *updMirror, *updModels, *movingRegistry, *restartFleet) {
	t.Helper()
	header := ggufHeader()
	files := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write(header) }))
	t.Cleanup(files.Close)
	engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hi"}}],"usage":{"completion_tokens":1}}`))
	}))
	t.Cleanup(engine.Close)

	mir := &updMirror{base: files.URL, sums: map[string]string{
		"abc12345/q4.gguf": "aaaa", "def98765/q4.gguf": "bbbb"}}
	// What was fetched before origins were kept: the mirror alone says where from.
	_, _ = mir.Fetch(context.Background(), mirror.FetchRequest{Model: "qwen", Version: "abc12345", File: "q4.gguf",
		Source: "https://huggingface.co/o/r/resolve/" + oldRev + "/q4.gguf"})
	models := &updModels{models: map[string]catalog.Model{"qwen": {Name: "qwen",
		Source:   files.URL + "/qwen/abc12345/q4.gguf",
		Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp", MinMemoryGB: 8}}}}, declared: map[string]bool{}}
	reg := &movingRegistry{repo: hf.Repo{ID: "o/r", Revision: newRev, Files: []hf.File{
		{Path: "q4.gguf", SHA256: "bbbb", Bytes: 2 << 30}, {Path: "README.md", Bytes: 10}}}}
	fleet := &restartFleet{engine: engine.URL, reps: []reconciler.Replica{
		{ID: "qwen-1", Model: "qwen", State: reconciler.Ready, Endpoint: engine.URL},
		{ID: "qwen-2", Model: "qwen", State: reconciler.Ready, Endpoint: engine.URL}}}
	s := New(Deps{Registry: reg, Mirror: mir, Catalog: &fakeCatalog{}, Fleet: fleet,
		Current: models.all, Models: models, Restart: fleet.Restart,
		Poll: 10 * time.Millisecond, ReadyTimeout: 5 * time.Second})
	return s, mir, models, reg, fleet
}

func TestNewWeightsInTheRepositoryAreFoundForAModelServedBeforeOriginsWereKept(t *testing.T) {
	s, _, models, _, _ := updSetup(t)
	s.CheckUpdates(context.Background())
	u := s.Updates()["qwen"]
	if !u.Available || u.Repo != "o/r" || u.Current != oldRev || u.Latest != newRev ||
		len(u.Changed) != 1 || u.Changed[0] != "q4.gguf" || u.Bytes != 2<<30 {
		t.Fatalf("%+v", u)
	}
	if o := models.models["qwen"].Origin; o == nil || o.Repo != "o/r" || o.Files[0].SHA256 != "aaaa" {
		t.Errorf("the origin worked out was not recorded: %+v", o)
	}
}

func TestANewCommitThatLeavesTheWeightsAloneIsNotAnUpdate(t *testing.T) {
	s, mir, _, _, _ := updSetup(t)
	mir.held[0].SHA256 = "bbbb"
	u, err := s.CheckModel(context.Background(), "qwen")
	if err != nil || u.Available || u.Latest != newRev || updateSays(u) != "same" {
		t.Fatalf("%+v %v", u, err)
	}
}

func updateSays(u UpdateInfo) string {
	if !u.Available && u.Latest != u.Current && len(u.Missing) == 0 {
		return "same"
	}
	return ""
}

func TestAnUpdateFetchesBesideTheOldReplacesTheReplicasAndCanGoBack(t *testing.T) {
	s, mir, models, _, fleet := updSetup(t)
	oldSource := models.models["qwen"].Source
	job, err := s.StartUpdate("qwen", "tester")
	if err != nil {
		t.Fatal(err)
	}
	j := wait(t, s, job.ID)
	if j.State != Done {
		t.Fatalf("%s: %+v", j.Error, j.Steps)
	}
	m := models.models["qwen"]
	if !strings.HasSuffix(m.Source, "/qwen/def98765/q4.gguf") || m.Origin.Revision != newRev ||
		m.Origin.Previous == nil || m.Origin.Previous.Source != oldSource {
		t.Fatalf("catalog: %s %+v", m.Source, m.Origin)
	}
	if len(mir.held) != 2 {
		t.Errorf("the old files did not stay on the mirror beside the new: %+v", mir.held)
	}
	for _, r := range fleet.Replicas() {
		if !strings.Contains(r.ID, "new") {
			t.Errorf("%s was not replaced", r.ID)
		}
	}
	if u := s.Updates()["qwen"]; u.Available || u.Current != newRev {
		t.Errorf("still offered: %+v", u)
	}

	// And back.
	if _, err := s.Rollback(context.Background(), "qwen", "tester"); err != nil {
		t.Fatal(err)
	}
	m = models.models["qwen"]
	if m.Source != oldSource || m.Origin.Revision != oldRev || m.Origin.Previous.Revision != newRev {
		t.Fatalf("rollback: %s %+v", m.Source, m.Origin)
	}
	if u := s.Updates()["qwen"]; !u.Available {
		t.Errorf("the update it went back from is not offered again: %+v", u)
	}

	// And forward again, from files already on the mirror.
	job, err = s.StartUpdate("qwen", "tester")
	if err != nil {
		t.Fatal(err)
	}
	if j := wait(t, s, job.ID); j.State != Done || models.models["qwen"].Origin.Revision != newRev {
		t.Fatalf("%s: %+v", j.Error, j.Steps)
	}
	// Then there is nothing newer: refused, saying why.
	job, _ = s.StartUpdate("qwen", "tester")
	if j := wait(t, s, job.ID); j.State != Failed || !strings.Contains(j.Error, "already runs") {
		t.Fatalf("%+v", j)
	}
}

func TestAModelFromConfigurationIsNotUpdatedBehindItsBack(t *testing.T) {
	s, _, models, _, _ := updSetup(t)
	models.declared["qwen"] = true
	if _, err := s.StartUpdate("qwen", "t"); err == nil || !strings.Contains(err.Error(), "server.yaml") {
		t.Fatalf("%v", err)
	}
}

func TestAnUpdateToFilesTheRepositoryDroppedIsRefused(t *testing.T) {
	s, _, _, reg, _ := updSetup(t)
	reg.repo.Files = []hf.File{{Path: "q4-v2.gguf", SHA256: "cccc"}}
	job, _ := s.StartUpdate("qwen", "t")
	if j := wait(t, s, job.ID); j.State != Failed || !strings.Contains(j.Error, "no longer has q4.gguf") {
		t.Fatalf("%+v", j)
	}
}

func TestRegistryDownloadURLsAreRead(t *testing.T) {
	for u, want := range map[string][3]string{
		"https://huggingface.co/bartowski/Qwen-GGUF/resolve/abc/sub/f.gguf": {"bartowski/Qwen-GGUF", "abc", "sub/f.gguf"},
		"https://hf.internal/mirror/x/resolve/main/f.gguf":                  {"mirror/x", "main", "f.gguf"},
	} {
		repo, rev, path, ok := parseResolve(u)
		if !ok || [3]string{repo, rev, path} != want {
			t.Errorf("%s: %s %s %s %v", u, repo, rev, path, ok)
		}
	}
	for _, u := range []string{"https://example.com/f.gguf", "https://huggingface.co/a/b/c/resolve/x/f", ":"} {
		if _, _, _, ok := parseResolve(u); ok {
			t.Errorf("%s was read", u)
		}
	}
}

// findingRegistry is a registry with history: searchable, and readable at a
// past commit.
type findingRegistry struct {
	movingRegistry
	found []hf.Model
	at    map[string]hf.Repo // id@rev
}

func (r *findingRegistry) Explore(context.Context, hf.Query) ([]hf.Model, error) { return r.found, nil }
func (r *findingRegistry) RepoAt(_ context.Context, id, rev string) (hf.Repo, error) {
	if repo, ok := r.at[id+"@"+rev]; ok {
		return repo, nil
	}
	return hf.Repo{}, errors.New("404")
}

func sibling(names ...string) []struct {
	RFilename string `json:"rfilename"`
} {
	var out []struct {
		RFilename string `json:"rfilename"`
	}
	for _, n := range names {
		out = append(out, struct {
			RFilename string `json:"rfilename"`
		}{n})
	}
	return out
}

func TestWeightsFetchedBeforeOriginsWereKeptAreTracedToTheirRepository(t *testing.T) {
	s, mir, models, reg, _ := updSetup(t)
	// The mirror kept neither source nor checksum.
	mir.held[0].Source, mir.held[0].SHA256 = "", ""
	fr := &findingRegistry{movingRegistry: movingRegistry{repo: reg.repo},
		found: []hf.Model{
			{ID: "copycat/r", Siblings: sibling("q4.gguf")},   // the right name, not the commit
			{ID: "other/thing", Siblings: sibling("q8.gguf")}, // not the file
			{ID: "o/r", Siblings: sibling("README.md", "q4.gguf")},
		},
		at: map[string]hf.Repo{
			"copycat/r@abc12345": {ID: "copycat/r", Revision: "ffff", Files: []hf.File{{Path: "q4.gguf", Bytes: 1 << 30}}},
			"o/r@abc12345": {ID: "o/r", Revision: oldRev, Files: []hf.File{
				{Path: "q4.gguf", SHA256: "aaaa", Bytes: 1 << 30}}},
		}}
	s.deps.Registry = fr
	s.CheckUpdates(context.Background())
	o := models.models["qwen"].Origin
	if o == nil || o.Repo != "o/r" || o.Revision != oldRev || o.Files[0].SHA256 != "aaaa" {
		t.Fatalf("traced to %+v", o)
	}
	if u := s.Updates()["qwen"]; !u.Available {
		t.Errorf("%+v", u)
	}
}

func TestAnOriginGivenByHandIsChecked(t *testing.T) {
	s, mir, models, reg, _ := updSetup(t)
	mir.held[0].Source = ""
	fr := &findingRegistry{movingRegistry: movingRegistry{repo: reg.repo}, at: map[string]hf.Repo{
		"o/r@abc12345":   {ID: "o/r", Revision: oldRev, Files: []hf.File{{Path: "q4.gguf", SHA256: "aaaa", Bytes: 1 << 30}}},
		"bad/r@abc12345": {ID: "bad/r", Revision: oldRev, Files: []hf.File{{Path: "q4.gguf", Bytes: 5}}},
	}}
	s.deps.Registry = fr
	if _, err := s.SetOriginOf(context.Background(), "qwen", "bad/r"); err == nil {
		t.Error("a repository whose file is another size was taken")
	}
	if _, err := s.SetOriginOf(context.Background(), "qwen", "o/r"); err != nil {
		t.Fatal(err)
	}
	if o := models.models["qwen"].Origin; o == nil || o.Repo != "o/r" {
		t.Fatalf("%+v", o)
	}
}