178 lines
6.8 KiB
Go
178 lines
6.8 KiB
Go
package health_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.wntrmute.dev/kyle/crossbar/internal/health"
|
|
)
|
|
|
|
// routerFake is shaped like llama-server's router mode: /v1/models lists every configured model
|
|
// with a status, a plain /props answers as the router itself (no context), and /props?model=X
|
|
// answers for one loaded child server. It counts the per-model /props queries it receives.
|
|
type routerFake struct {
|
|
srv *httptest.Server
|
|
mu sync.Mutex
|
|
queries map[string]int
|
|
}
|
|
|
|
func newRouterFake(t *testing.T) *routerFake {
|
|
f := &routerFake{queries: map[string]int{}}
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, `{"status":"ok"}`) })
|
|
mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
|
|
fmt.Fprint(w, `{"object":"list","data":[
|
|
{"id":"big","object":"model","status":{"value":"loaded","args":["--ctx-size","262144"]}},
|
|
{"id":"small","object":"model","status":{"value":"loaded"}},
|
|
{"id":"cold","object":"model","status":{"value":"unloaded"}},
|
|
{"id":"warming","object":"model","status":{"value":"loading"}}]}`)
|
|
})
|
|
mux.HandleFunc("/props", func(w http.ResponseWriter, r *http.Request) {
|
|
model := r.URL.Query().Get("model")
|
|
if model == "" {
|
|
fmt.Fprint(w, `{"role":"router","model_alias":"llama-server","model_path":"none","default_generation_settings":{"params":null,"n_ctx":0}}`)
|
|
return
|
|
}
|
|
f.mu.Lock()
|
|
f.queries[model]++
|
|
f.mu.Unlock()
|
|
switch model {
|
|
case "big":
|
|
fmt.Fprint(w, `{"default_generation_settings":{"n_ctx":262144,"params":{}},"total_slots":4,"model_alias":"big"}`)
|
|
case "small":
|
|
fmt.Fprint(w, `{"default_generation_settings":{"n_ctx":32768,"params":{}},"total_slots":1,"model_alias":"small"}`)
|
|
default:
|
|
// Asking a router for an unloaded model would make it load the model. The fake
|
|
// answers 500 so a wrong query is visible in the counts and cannot look like success.
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
fmt.Fprint(w, `{"error":"the poller must not ask for a model that is not loaded"}`)
|
|
}
|
|
})
|
|
f.srv = httptest.NewServer(mux)
|
|
t.Cleanup(f.srv.Close)
|
|
return f
|
|
}
|
|
|
|
func (f *routerFake) count(model string) int {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
return f.queries[model]
|
|
}
|
|
|
|
func TestRouterLoadedMeansStatusLoaded(t *testing.T) {
|
|
f := newRouterFake(t)
|
|
tbl := health.New(map[string]string{"r": f.srv.URL}, time.Hour, nil)
|
|
tbl.PollOnce(context.Background())
|
|
s, ok := tbl.Get("r")
|
|
if !ok || !s.Healthy {
|
|
t.Fatalf("status = %+v, want a healthy host", s)
|
|
}
|
|
if len(s.Loaded) != 2 || s.Loaded[0] != "big" || s.Loaded[1] != "small" {
|
|
t.Errorf("Loaded = %v, want [big small]: unloaded and loading models are not loaded", s.Loaded)
|
|
}
|
|
}
|
|
|
|
func TestRouterContextIsLearnedPerModel(t *testing.T) {
|
|
f := newRouterFake(t)
|
|
tbl := health.New(map[string]string{"r": f.srv.URL}, time.Hour, nil)
|
|
tbl.PollOnce(context.Background())
|
|
s, _ := tbl.Get("r")
|
|
if s.NCtx != 0 || s.Slots != 0 || s.PerSlotCtx() != 0 {
|
|
t.Errorf("a router's own /props carries no context; host-level must stay unknown: %+v", s)
|
|
}
|
|
if got := s.Models["big"]; got != (health.ModelCtx{NCtx: 262144, Slots: 4}) {
|
|
t.Errorf("Models[big] = %+v, want {262144 4}", got)
|
|
}
|
|
if got := s.Models["small"]; got != (health.ModelCtx{NCtx: 32768, Slots: 1}) {
|
|
t.Errorf("Models[small] = %+v, want {32768 1}", got)
|
|
}
|
|
if got := s.PerSlotCtxFor("big"); got != 65536 {
|
|
t.Errorf("PerSlotCtxFor(big) = %d, want 262144/4", got)
|
|
}
|
|
if got := s.PerSlotCtxFor("small"); got != 32768 {
|
|
t.Errorf("PerSlotCtxFor(small) = %d, want 32768/1", got)
|
|
}
|
|
if got := s.PerSlotCtxFor("cold"); got != 0 {
|
|
t.Errorf("PerSlotCtxFor(cold) = %d, want 0: nothing is known about an unloaded model", got)
|
|
}
|
|
if _, present := s.Models["cold"]; present {
|
|
t.Errorf("Models must not carry an entry for an unloaded model: %+v", s.Models)
|
|
}
|
|
}
|
|
|
|
func TestRouterUnloadedModelsAreNeverQueried(t *testing.T) {
|
|
f := newRouterFake(t)
|
|
tbl := health.New(map[string]string{"r": f.srv.URL}, time.Hour, nil)
|
|
for i := 0; i < 3; i++ {
|
|
tbl.PollOnce(context.Background())
|
|
}
|
|
if f.count("cold") != 0 || f.count("warming") != 0 {
|
|
t.Fatalf("/props?model= was asked for a model that is not loaded (cold %d, warming %d): on a real router that loads the model", f.count("cold"), f.count("warming"))
|
|
}
|
|
if f.count("big") == 0 || f.count("small") == 0 {
|
|
t.Errorf("loaded models must be asked: big %d, small %d", f.count("big"), f.count("small"))
|
|
}
|
|
}
|
|
|
|
func TestPlainServerStillReadsHostLevelContext(t *testing.T) {
|
|
// A single llama-server (no status field, no router role) behaves as in v2: every listed model
|
|
// is loaded, the host-level context comes from the plain /props, and the per-model view falls
|
|
// back to it for any loaded model.
|
|
srv := propsFake(t, `{"default_generation_settings":{"n_ctx":131072,"params":{}},"total_slots":4,"model_path":"/x/m.gguf"}`, 200)
|
|
tbl := health.New(map[string]string{"a": srv.URL}, time.Hour, nil)
|
|
tbl.PollOnce(context.Background())
|
|
s, _ := tbl.Get("a")
|
|
if !s.Healthy || len(s.Loaded) != 1 || s.Loaded[0] != "m" || s.NCtx != 131072 || s.Slots != 4 {
|
|
t.Fatalf("status = %+v, want healthy, Loaded [m], NCtx 131072, Slots 4", s)
|
|
}
|
|
if got := s.PerSlotCtxFor("m"); got != 32768 {
|
|
t.Errorf("PerSlotCtxFor(m) = %d, want the host-level 131072/4", got)
|
|
}
|
|
if got := s.PerSlotCtxFor("other"); got != 0 {
|
|
t.Errorf("PerSlotCtxFor(other) = %d, want 0 for a model the host does not list", got)
|
|
}
|
|
if s.Models == nil {
|
|
t.Errorf("Models must be an empty map after a poll, never nil")
|
|
}
|
|
}
|
|
|
|
func TestPerModelPropsFailureLeavesTheModelUnknown(t *testing.T) {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, `{"status":"ok"}`) })
|
|
mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
|
|
fmt.Fprint(w, `{"data":[{"id":"ok","status":{"value":"loaded"}},{"id":"broken","status":{"value":"loaded"}}]}`)
|
|
})
|
|
mux.HandleFunc("/props", func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Query().Get("model") {
|
|
case "":
|
|
fmt.Fprint(w, `{"role":"router","default_generation_settings":{"n_ctx":0}}`)
|
|
case "ok":
|
|
fmt.Fprint(w, `{"default_generation_settings":{"n_ctx":8192},"total_slots":2}`)
|
|
default:
|
|
fmt.Fprint(w, `<html>not json</html>`)
|
|
}
|
|
})
|
|
srv := httptest.NewServer(mux)
|
|
t.Cleanup(srv.Close)
|
|
tbl := health.New(map[string]string{"r": srv.URL}, time.Hour, nil)
|
|
tbl.PollOnce(context.Background())
|
|
s, _ := tbl.Get("r")
|
|
if !s.Healthy {
|
|
t.Fatalf("a broken per-model /props must not make the host unhealthy: %+v", s)
|
|
}
|
|
if len(s.Loaded) != 2 {
|
|
t.Errorf("Loaded = %v, want both models: /props is advisory", s.Loaded)
|
|
}
|
|
if got := s.PerSlotCtxFor("ok"); got != 4096 {
|
|
t.Errorf("PerSlotCtxFor(ok) = %d, want 8192/2", got)
|
|
}
|
|
if _, present := s.Models["broken"]; present || s.PerSlotCtxFor("broken") != 0 {
|
|
t.Errorf("a model whose /props failed stays unknown: %+v", s.Models)
|
|
}
|
|
}
|