llmgw/internal/control/modelserve/modelserve_test.go

package modelserve

import (
	"bytes"
	"context"
	"encoding/binary"
	"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"
)

// ggufHeader is enough of a GGUF file for its architecture and shape.
func ggufHeader() []byte {
	var b bytes.Buffer
	w := func(v any) { _ = binary.Write(&b, binary.LittleEndian, v) }
	str := func(s string) { w(uint64(len(s))); b.WriteString(s) }
	u32 := func(k string, v uint32) { str(k); w(uint32(4)); w(v) }
	b.WriteString("GGUF")
	w(uint32(3))
	w(uint64(0))
	w(uint64(5))
	str("general.architecture")
	w(uint32(8))
	str("qwen2")
	u32("qwen2.block_count", 28)
	u32("qwen2.attention.head_count", 12)
	u32("qwen2.attention.head_count_kv", 2)
	u32("qwen2.embedding_length", 1536)
	return b.Bytes()
}

type fakeRegistry struct{ repo hf.Repo }

func (f fakeRegistry) Repo(context.Context, string) (hf.Repo, error) { return f.repo, nil }

// fakeMirror serves the header at every file URL and reports each fetched
// file present, with the checksum it was told to.
type fakeMirror struct {
	mu   sync.Mutex
	srv  *httptest.Server
	got  []mirror.FetchRequest
	sums map[string]string
}

func (m *fakeMirror) Fetch(_ context.Context, r mirror.FetchRequest) (mirror.FetchResponse, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.got = append(m.got, r)
	return mirror.FetchResponse{URL: m.srv.URL + "/" + r.Model + "/" + r.Version + "/" + r.File}, nil
}

func (m *fakeMirror) List(context.Context) (mirror.ModelsResponse, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out mirror.ModelsResponse
	for _, r := range m.got {
		out.Models = append(out.Models, mirror.Status{Model: r.Model, Version: r.Version, File: r.File,
			State: mirror.Present, BytesDone: 1 << 30, BytesTotal: 1 << 30, SHA256: m.sums[r.File]})
	}
	return out, nil
}

type fakeCatalog struct {
	mu    sync.Mutex
	added []catalog.Model
}

func (c *fakeCatalog) Add(_ context.Context, m catalog.Model, _ string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.added = append(c.added, m)
	return nil
}

// fakeFleet has a replica ready at engine once the catalog holds the model.
type fakeFleet struct {
	cat    *fakeCatalog
	engine string
	health []reconciler.ModelHealth
}

func (f *fakeFleet) Replicas() []reconciler.Replica {
	f.cat.mu.Lock()
	defer f.cat.mu.Unlock()
	var out []reconciler.Replica
	for _, m := range f.cat.added {
		out = append(out, reconciler.Replica{ID: m.Name + "-r1", Model: m.Name,
			State: reconciler.Ready, Endpoint: f.engine,
			Shards: []reconciler.Shard{{NodeID: "node01"}}})
	}
	return out
}
func (f *fakeFleet) ModelHealth() []reconciler.ModelHealth { return f.health }
func (f *fakeFleet) EngineKey(string) string               { return "k" }

func setup(t *testing.T, files []hf.File, sums map[string]string) (*Server, *fakeMirror, *fakeCatalog, *fakeFleet) {
	t.Helper()
	header := ggufHeader()
	files0 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write(header)
	}))
	t.Cleanup(files0.Close)
	engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer k" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hello!"}}],"usage":{"completion_tokens":3}}`))
	}))
	t.Cleanup(engine.Close)

	m := &fakeMirror{srv: files0, sums: sums}
	c := &fakeCatalog{}
	f := &fakeFleet{cat: c, engine: engine.URL}
	s := New(Deps{
		Registry: fakeRegistry{repo: hf.Repo{ID: "bartowski/Qwen2.5-1.5B-Instruct-GGUF", Revision: "abcdef123456", Files: files}},
		Mirror:   m, Catalog: c, Fleet: f, Poll: 10 * time.Millisecond, ReadyTimeout: 5 * time.Second,
	})
	return s, m, c, f
}

func wait(t *testing.T, s *Server, id string) Job {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if j, _ := s.Get(id); j.State != Running {
			return j
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("the job never finished")
	return Job{}
}

func TestAModelGoesFromTheRegistryToAnAnsweringReplica(t *testing.T) {
	files := []hf.File{
		{Path: "Qwen2.5-1.5B-Instruct-Q8_0.gguf", Bytes: 2 << 30, SHA256: "aa"},
		{Path: "Qwen2.5-1.5B-Instruct-Q4_K_M.gguf", Bytes: 1 << 30, SHA256: "bb"},
		{Path: "mmproj-Q4_K_M.gguf", Bytes: 1 << 20},
	}
	s, m, c, _ := setup(t, files, map[string]string{"Qwen2.5-1.5B-Instruct-Q4_K_M.gguf": "BB"})
	job, err := s.Start(Request{Repo: "bartowski/Qwen2.5-1.5B-Instruct-GGUF", Replicas: 1})
	if err != nil {
		t.Fatal(err)
	}
	j := wait(t, s, job.ID)
	if j.State != Done {
		t.Fatalf("job %s: %s\n%+v", j.State, j.Error, j.Steps)
	}
	if len(m.got) != 1 || m.got[0].File != "Qwen2.5-1.5B-Instruct-Q4_K_M.gguf" || m.got[0].Version != "abcdef12" {
		t.Errorf("fetched %+v; expected the Q4_K_M file at its revision", m.got)
	}
	mod := c.added[0]
	if mod.Name != "qwen2.5-1.5b-instruct" || mod.Profiles[0].MinMemoryGB == 0 || mod.Scaling.Min != 1 {
		t.Errorf("catalog entry %+v", mod)
	}
	if j.Reply != "Hello!" {
		t.Errorf("the smoke test's answer was %q", j.Reply)
	}
	var names []string
	for _, st := range j.Steps {
		names = append(names, st.Name)
	}
	if strings.Join(names, ",") != "registry,mirror,check,catalog,placement,smoke test" {
		t.Errorf("steps %v", names)
	}
}

func TestAQuantTheRepositoryDoesNotHaveIsNamedWithWhatItHas(t *testing.T) {
	s, _, c, _ := setup(t, []hf.File{{Path: "m-Q8_0.gguf", Bytes: 1}}, nil)
	job, _ := s.Start(Request{Repo: "a/b", Replicas: 1})
	j := wait(t, s, job.ID)
	if j.State != Failed || !strings.Contains(j.Error, "m-Q8_0.gguf") {
		t.Errorf("%s: %s", j.State, j.Error)
	}
	if len(c.added) != 0 {
		t.Error("something was added to the catalog")
	}
}

func TestADownloadThatDoesNotMatchTheRegistryIsNotServed(t *testing.T) {
	files := []hf.File{{Path: "m-Q4_K_M.gguf", Bytes: 1 << 30, SHA256: "aaaaaaaaaaaaaaaa"}}
	s, _, c, _ := setup(t, files, map[string]string{"m-Q4_K_M.gguf": "bbbbbbbbbbbbbbbb"})
	job, _ := s.Start(Request{Repo: "a/b", Replicas: 1})
	j := wait(t, s, job.ID)
	if j.State != Failed || !strings.Contains(j.Error, "registry published") {
		t.Errorf("%s: %s", j.State, j.Error)
	}
	if len(c.added) != 0 {
		t.Error("a corrupted download reached the catalog")
	}
}

func TestAParkedModelStopsTheWait(t *testing.T) {
	s, _, _, f := setup(t, []hf.File{{Path: "m-Q4_K_M.gguf", Bytes: 1 << 30}}, nil)
	f.health = []reconciler.ModelHealth{{Model: "qwen2.5-1.5b-instruct", Parked: true, Failures: 5, Reason: "unknown model architecture"}}
	// A fleet where it never becomes ready, only parked.
	s.deps.Fleet = parkedFleet{f}
	job, _ := s.Start(Request{Repo: "a/b", Replicas: 1})
	j := wait(t, s, job.ID)
	if j.State != Failed || !strings.Contains(j.Error, "parked") {
		t.Errorf("%s: %s", j.State, j.Error)
	}
}

type parkedFleet struct{ *fakeFleet }

func (parkedFleet) Replicas() []reconciler.Replica { return nil }

func TestARepositoryThatIsNotOneIsRefused(t *testing.T) {
	s := New(Deps{})
	for _, r := range []string{"", "noslash", "../etc/passwd", "a/b/c"} {
		if _, err := s.Start(Request{Repo: r}); err == nil {
			t.Errorf("%q accepted", r)
		}
	}
}