Files
llmgw/internal/control/modelserve/modelserve_test.go
T

400 lines
14 KiB
Go

package modelserve
import (
"bytes"
"context"
"encoding/binary"
"fmt"
"net/http"
"net/http/httptest"
"slices"
"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"
"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
)
// 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
base string
}
func (f fakeRegistry) Base() string { return f.base }
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)
}
}
}
func TestAVisionModelIsServedWithItsProjectorBesideIt(t *testing.T) {
files := []hf.File{
{Path: "gemma-3-4b-it-Q4_K_M.gguf", Bytes: 2 << 30, SHA256: "aa"},
{Path: "mmproj-gemma-3-4b-it-f16.gguf", Bytes: 800 << 20, SHA256: "bb"},
}
s, m, c, _ := setup(t, files, map[string]string{"gemma-3-4b-it-Q4_K_M.gguf": "aa", "mmproj-gemma-3-4b-it-f16.gguf": "bb"})
job, err := s.Start(Request{Repo: "google/gemma-3-4b-it-GGUF", Quant: "Q4_K_M",
Projector: "mmproj-gemma-3-4b-it-f16.gguf", Replicas: 1})
if err != nil {
t.Fatal(err)
}
if j := wait(t, s, job.ID); j.State != Done {
t.Fatalf("%s %+v", j.Error, j.Steps)
}
if len(m.got) != 2 {
t.Fatalf("fetched %+v", m.got)
}
got := c.added[0]
if !strings.HasSuffix(got.Source, "/") || !got.Capabilities.Vision || len(got.Origin.Files) != 2 {
t.Errorf("%s vision=%v origin=%+v", got.Source, got.Capabilities.Vision, got.Origin)
}
// A projector the repository does not have is named with those it has.
job, _ = s.Start(Request{Repo: "google/gemma-3-4b-it-GGUF", Quant: "Q4_K_M", Projector: "mmproj-x.gguf", Name: "g2"})
if j := wait(t, s, job.ID); j.State != Failed || !strings.Contains(j.Error, "mmproj-gemma-3-4b-it-f16.gguf") {
t.Errorf("%+v", j)
}
}
func TestOriginalWeightsAreServedOnVLLMWithTheirConfigAndTokenizer(t *testing.T) {
files := []hf.File{
{Path: "model-00001-of-00002.safetensors", Bytes: 4 << 30, SHA256: "a1"},
{Path: "model-00002-of-00002.safetensors", Bytes: 3 << 30, SHA256: "a2"},
{Path: "model.safetensors.index.json", Bytes: 30000},
{Path: "config.json", Bytes: 700},
{Path: "tokenizer.json", Bytes: 7 << 20},
{Path: "tokenizer_config.json", Bytes: 7000},
{Path: "README.md", Bytes: 9000},
{Path: "original/consolidated.safetensors", Bytes: 7 << 30},
}
s, m, c, _ := setup(t, files, map[string]string{"model-00001-of-00002.safetensors": "a1", "model-00002-of-00002.safetensors": "a2"})
// With no fleet to plan against, the GPUs must be named.
job0, _ := s.Start(Request{Repo: "Qwen/Qwen2.5-7B-Instruct", Format: "safetensors", Name: "x"})
if j := wait(t, s, job0.ID); j.State != Failed || !strings.Contains(j.Error, "GPU class") {
t.Errorf("safetensors with no GPU class: %+v", j)
}
job, err := s.Start(Request{Repo: "Qwen/Qwen2.5-7B-Instruct", Format: "safetensors", GPUClass: "l40s-48g", Replicas: 1})
if err != nil {
t.Fatal(err)
}
if j := wait(t, s, job.ID); j.State != Done {
t.Fatalf("%s %+v", j.Error, j.Steps)
}
var got []string
for _, r := range m.got {
got = append(got, r.File)
}
if len(got) != 6 || slices.Contains(got, "README.md") || slices.Contains(got, "consolidated.safetensors") {
t.Errorf("fetched %v", got)
}
mod := c.added[0]
p := mod.Profiles[0]
if !strings.HasSuffix(mod.Source, "/") || p.Engine != "vllm" || p.GPUClass != "l40s-48g" || p.TP != 1 ||
!slices.Contains(p.Args, "--max-model-len") {
t.Errorf("%s %+v", mod.Source, p)
}
}
func TestTheFileThatSuitsTheCardsIsChosenAndRunsOnThem(t *testing.T) {
files := []hf.File{
{Path: "m-Q4_K_M.gguf", Bytes: 1 << 30, SHA256: "a"},
{Path: "m-Q8_0.gguf", Bytes: 2 << 30, SHA256: "b"},
{Path: "m-F16.gguf", Bytes: 40 << 30, SHA256: "c"},
}
s, m, c, _ := setup(t, files, map[string]string{"m-Q4_K_M.gguf": "a", "m-Q8_0.gguf": "b"})
s.deps.Registry = fakeRegistry{repo: hf.Repo{ID: "o/m-GGUF", Revision: "abcdef123456", Files: files}, base: m.srv.URL}
s.deps.Pool = func() []modelplan.PoolNode {
return []modelplan.PoolNode{
{ID: "g", Class: "6000-24g", GPUMemGB: 24, GPUs: 1, FreeGPUs: 1, HostMemGB: 64, HostFreeGB: 60},
{ID: "c", HostMemGB: 16, HostFreeGB: 14},
}
}
p, err := s.Plan(context.Background(), PlanRequest{Repo: "o/m-GGUF"})
if err != nil {
t.Fatal(err)
}
rec, _ := p.Recommended()
if rec.Artifact != "m-Q8_0.gguf" || rec.CPU || rec.Class != "6000-24g" || p.Artifacts[2].Best == nil || !p.Artifacts[2].Best.CPU {
t.Fatalf("plan: %+v; f16: %+v", rec, p.Artifacts[2].Best)
}
job, err := s.Start(Request{Repo: "o/m-GGUF", Replicas: 1, Fallbacks: true})
if err != nil {
t.Fatal(err)
}
if j := wait(t, s, job.ID); j.State != Done {
t.Fatalf("%s %+v", j.Error, j.Steps)
}
if m.got[0].File != "m-Q8_0.gguf" {
t.Errorf("fetched %s", m.got[0].File)
}
profs := c.added[0].Profiles
if len(profs) != 2 || profs[0].GPUClass != "6000-24g" || profs[0].GPUsPerNode != 1 || profs[0].Engine != "llamacpp" ||
profs[1].Name != "cpu" || profs[1].Priority >= profs[0].Priority || profs[1].MinMemoryGB == 0 {
t.Errorf("profiles: %+v", profs)
}
// A file too big for any card or any node's memory: refused, saying so.
s.deps.Pool = func() []modelplan.PoolNode { return []modelplan.PoolNode{{ID: "c", HostMemGB: 16, HostFreeGB: 14}} }
job, _ = s.Start(Request{Repo: "o/m-GGUF", File: "m-F16.gguf", Name: "big"})
if j := wait(t, s, job.ID); j.State != Failed || !strings.Contains(j.Error, "fits nowhere") {
t.Errorf("%+v", j)
}
}
// busyFleet has the model's cards all in use: nothing placed, and placement
// saying so.
type busyFleet struct{ *fakeFleet }
func (busyFleet) Replicas() []reconciler.Replica { return nil }
func (f busyFleet) ModelHealth() []reconciler.ModelHealth {
return []reconciler.ModelHealth{{Model: "qwen2.5-1.5b-instruct", Unplaceable: true,
Placement: "cuda01: 0 of 1 6000-24g free"}}
}
func TestAModelThatFitsBusyCardsIsAddedToWaitNotRefused(t *testing.T) {
files := []hf.File{{Path: "m-Q8_0.gguf", Bytes: 2 << 30, SHA256: "b"}}
s, m, c, f := setup(t, files, map[string]string{"m-Q8_0.gguf": "b"})
s.deps.Registry = fakeRegistry{repo: hf.Repo{ID: "o/m-GGUF", Revision: "abcdef123456", Files: files}, base: m.srv.URL}
s.deps.Pool = func() []modelplan.PoolNode {
return []modelplan.PoolNode{{ID: "g", Class: "6000-24g", GPUMemGB: 24, GPUs: 1, FreeGPUs: 0, HostMemGB: 8}}
}
s.deps.Fleet = busyFleet{f}
s.deps.UnplacedGrace = 10 * time.Millisecond
job, err := s.Start(Request{Repo: "o/m-GGUF", Name: "qwen2.5-1.5b-instruct", Replicas: 1})
if err != nil {
t.Fatal(err)
}
j := wait(t, s, job.ID)
if j.State != Done || len(c.added) != 1 || c.added[0].Profiles[0].GPUClass != "6000-24g" {
t.Fatalf("%s %+v", j.Error, j.Steps)
}
last := j.Steps[len(j.Steps)-1]
if last.Name != "placement" || last.State != Skipped || !strings.Contains(last.Detail, "starts when there is room") {
t.Errorf("%+v", last)
}
}
func TestOnlyTheLastFinishedJobsAreRemembered(t *testing.T) {
s := New(Deps{})
start := time.Now().Add(-time.Hour)
s.mu.Lock()
running := &Job{ID: "running", State: Running, Started: start}
s.keepJob(running)
for i := range keptJobs + 5 {
s.keepJob(&Job{ID: fmt.Sprintf("done-%d", i), State: "done", Started: start.Add(time.Duration(i+1) * time.Second)})
}
s.mu.Unlock()
if n := len(s.Jobs()); n != keptJobs+1 {
t.Errorf("%d jobs remembered, expected %d finished and the running one", n, keptJobs)
}
if _, ok := s.Get("running"); !ok {
t.Error("the running job, the oldest, was forgotten")
}
if _, ok := s.Get("done-0"); ok {
t.Error("the oldest finished job is still remembered")
}
}