Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
400 lines
14 KiB
Go
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")
|
|
}
|
|
}
|