Ilker ManapandClaude Opus 5.5
92733317a9
llmgw-kb sends many passages in one /v1/embeddings request. Their sizes were added together and compared with the context, so a batch of passages that each fit was refused as one over-long prompt. The longest input is what has to fit; the total still counts for usage and limits. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
811 lines
29 KiB
Go
811 lines
29 KiB
Go
package router
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"slices"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
|
)
|
|
|
|
// catModel builds a minimal catalog entry for router tests, which care only
|
|
// about names, aliases and fallbacks.
|
|
func catModel(name string, aliases []string, fallbacks ...string) catalog.Model {
|
|
return catalog.Model{
|
|
Name: name, Source: "/models/" + name, Aliases: aliases, Fallbacks: fallbacks,
|
|
Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}},
|
|
}
|
|
}
|
|
|
|
// fakeEngine serves the minimum an endpoint needs: the metric contract plus a
|
|
// scripted chat response.
|
|
type fakeEngine struct {
|
|
srv *httptest.Server
|
|
calls atomic.Int64
|
|
// status, when non-zero, is returned instead of a completion.
|
|
status atomic.Int64
|
|
// last is the body of the last completion asked for.
|
|
last atomic.Value
|
|
// traceparent is the one the last completion came with.
|
|
traceparent atomic.Value
|
|
}
|
|
|
|
func newFakeEngine(t *testing.T, model string) *fakeEngine {
|
|
t.Helper()
|
|
f := &fakeEngine{}
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
|
|
fmt.Fprintf(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n")
|
|
fmt.Fprintf(w, "# TYPE vllm:num_requests_waiting gauge\nvllm:num_requests_waiting 0\n")
|
|
fmt.Fprintf(w, "# TYPE vllm:kv_cache_usage_perc gauge\nvllm:kv_cache_usage_perc 0.1\n")
|
|
})
|
|
mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
|
|
f.calls.Add(1)
|
|
f.traceparent.Store(r.Header.Get("traceparent"))
|
|
if b, err := io.ReadAll(r.Body); err == nil {
|
|
f.last.Store(b)
|
|
}
|
|
if code := int(f.status.Load()); code != 0 {
|
|
http.Error(w, "unavailable", code)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprintf(w, `{"model":%q,"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":2}}`, model)
|
|
})
|
|
// Three dimensions: the input's length in bytes, its number of words,
|
|
// and 1.
|
|
mux.HandleFunc("POST /v1/embeddings", func(w http.ResponseWriter, r *http.Request) {
|
|
var in struct {
|
|
Input []string `json:"input"`
|
|
}
|
|
_ = json.NewDecoder(r.Body).Decode(&in)
|
|
var rows []string
|
|
for i, s := range in.Input {
|
|
rows = append(rows, fmt.Sprintf(`{"index":%d,"embedding":[%d,%d,1]}`, i, len(s), len(strings.Fields(s))))
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprintf(w, `{"model":%q,"data":[%s],"usage":{"prompt_tokens":%d}}`, model, strings.Join(rows, ","), len(in.Input))
|
|
})
|
|
f.srv = httptest.NewServer(mux)
|
|
t.Cleanup(f.srv.Close)
|
|
return f
|
|
}
|
|
|
|
func newTestRouter(t *testing.T, cfg config.Server) *Router {
|
|
t.Helper()
|
|
rt, err := New(cfg, slog.New(slog.DiscardHandler))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rt.Start(t.Context())
|
|
t.Cleanup(rt.Stop)
|
|
return rt
|
|
}
|
|
|
|
func post(t *testing.T, h http.Handler, body any, hdr map[string]string) *httptest.ResponseRecorder {
|
|
t.Helper()
|
|
b, err := json.Marshal(body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(b))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
for k, v := range hdr {
|
|
req.Header.Set(k, v)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
h.ServeHTTP(rec, req)
|
|
return rec
|
|
}
|
|
|
|
func baseConfig(eps ...config.Endpoint) config.Server {
|
|
c := config.Default()
|
|
c.Router.ScrapeInterval = 50 * time.Millisecond
|
|
c.Router.StaleAfter = 5 * time.Second
|
|
c.Endpoints = eps
|
|
return c
|
|
}
|
|
|
|
func TestRoutesToHealthyEndpoint(t *testing.T) {
|
|
e := newFakeEngine(t, "m")
|
|
rt := newTestRouter(t, baseConfig(config.Endpoint{
|
|
ID: "a", URL: e.srv.URL, Models: []string{"m"},
|
|
}))
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
|
}
|
|
if e.calls.Load() != 1 {
|
|
t.Errorf("engine received %d calls", e.calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestRetriesOnAnotherEndpoint(t *testing.T) {
|
|
bad := newFakeEngine(t, "m")
|
|
bad.status.Store(http.StatusServiceUnavailable)
|
|
good := newFakeEngine(t, "m")
|
|
|
|
rt := newTestRouter(t, baseConfig(
|
|
config.Endpoint{ID: "bad", URL: bad.srv.URL, Models: []string{"m"}},
|
|
config.Endpoint{ID: "good", URL: good.srv.URL, Models: []string{"m"}},
|
|
))
|
|
|
|
// With two endpoints and two retries, the bad one must not sink the request.
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
|
}
|
|
if good.calls.Load() == 0 {
|
|
t.Error("the healthy endpoint was never tried")
|
|
}
|
|
}
|
|
|
|
func TestNoRetryAfterBytesReachClient(t *testing.T) {
|
|
// An endpoint that starts streaming then dies must not be retried: splicing
|
|
// two responses together is worse than a short stream (plan §7.7).
|
|
var served atomic.Int64
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
|
|
io.WriteString(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n"+
|
|
"# TYPE vllm:num_requests_waiting gauge\nvllm:num_requests_waiting 0\n"+
|
|
"# TYPE vllm:kv_cache_usage_perc gauge\nvllm:kv_cache_usage_perc 0\n")
|
|
})
|
|
mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, _ *http.Request) {
|
|
served.Add(1)
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f := w.(http.Flusher)
|
|
io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
|
|
f.Flush()
|
|
// Then hang up mid-stream.
|
|
})
|
|
srv := httptest.NewServer(mux)
|
|
defer srv.Close()
|
|
|
|
rt := newTestRouter(t, baseConfig(
|
|
config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"m"}},
|
|
config.Endpoint{ID: "b", URL: srv.URL, Models: []string{"m"}},
|
|
))
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}, "stream": true}, nil)
|
|
if !strings.Contains(rec.Body.String(), `"content":"a"`) {
|
|
t.Fatalf("partial stream not delivered: %q", rec.Body.String())
|
|
}
|
|
if n := served.Load(); n != 1 {
|
|
t.Errorf("request was sent %d times; it must not be retried once streaming began", n)
|
|
}
|
|
}
|
|
|
|
func TestFallbackModel(t *testing.T) {
|
|
primary := newFakeEngine(t, "primary")
|
|
primary.status.Store(http.StatusServiceUnavailable)
|
|
backup := newFakeEngine(t, "backup")
|
|
|
|
cfg := baseConfig(
|
|
config.Endpoint{ID: "p", URL: primary.srv.URL, Models: []string{"primary"}},
|
|
config.Endpoint{ID: "b", URL: backup.srv.URL, Models: []string{"backup"}},
|
|
)
|
|
cfg.Models = []catalog.Model{catModel("primary", nil, "backup")}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "primary", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
|
}
|
|
if backup.calls.Load() == 0 {
|
|
t.Fatal("the fallback model was never tried")
|
|
}
|
|
// The client asked for one model and got another; they must be told.
|
|
if got := rec.Header().Get("x-llmgw-served-model"); got != "backup" {
|
|
t.Errorf("x-llmgw-served-model=%q, expected \"backup\"", got)
|
|
}
|
|
}
|
|
|
|
func TestAliasResolution(t *testing.T) {
|
|
e := newFakeEngine(t, "llama-3.3-70b")
|
|
cfg := baseConfig(config.Endpoint{
|
|
ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"},
|
|
})
|
|
cfg.Models = []catalog.Model{catModel("llama-3.3-70b", []string{"llama"})}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "llama", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("alias was not resolved: %d %s", rec.Code, rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestUnknownModelIs503(t *testing.T) {
|
|
e := newFakeEngine(t, "m")
|
|
rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "nope", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusServiceUnavailable {
|
|
t.Errorf("status %d, expected 503", rec.Code)
|
|
}
|
|
var body map[string]any
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
|
t.Fatalf("error body is not JSON: %q", rec.Body.String())
|
|
}
|
|
if _, ok := body["error"]; !ok {
|
|
t.Errorf("error body is not OpenAI-shaped: %v", body)
|
|
}
|
|
}
|
|
|
|
func TestMissingModelIs400(t *testing.T) {
|
|
e := newFakeEngine(t, "m")
|
|
rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"messages": []any{}}, nil)
|
|
if rec.Code != http.StatusBadRequest {
|
|
t.Errorf("status %d, expected 400", rec.Code)
|
|
}
|
|
}
|
|
|
|
func TestModelsEndpointListsAliases(t *testing.T) {
|
|
e := newFakeEngine(t, "llama-3.3-70b")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"}})
|
|
cfg.Models = []catalog.Model{catModel("llama-3.3-70b", []string{"llama"})}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
rec := httptest.NewRecorder()
|
|
rt.Handler().ServeHTTP(rec, req)
|
|
|
|
var out struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ids := map[string]bool{}
|
|
for _, d := range out.Data {
|
|
ids[d.ID] = true
|
|
}
|
|
if !ids["llama-3.3-70b"] || !ids["llama"] {
|
|
t.Errorf("both the model and its alias should be listed: %v", ids)
|
|
}
|
|
}
|
|
|
|
func TestModelsEndpointMarksSimulatedModels(t *testing.T) {
|
|
// A client picking a model from this list has nothing else to go on. Both
|
|
// entries are OpenAI-shaped and only owned_by distinguishes the one that
|
|
// returns placeholder text, so it has to be right for aliases too.
|
|
e := newFakeEngine(t, "real")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"real", "fake"}})
|
|
sim := catModel("fake", []string{"fake-alias"})
|
|
sim.Profiles = []catalog.Profile{{Name: "p", Engine: "sim", GPUClass: "h100-80g", TP: 1}}
|
|
cfg.Models = []catalog.Model{catModel("real", []string{"real-alias"}), sim}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
rec := httptest.NewRecorder()
|
|
rt.Handler().ServeHTTP(rec, req)
|
|
|
|
var out struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
OwnedBy string `json:"owned_by"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
owner := map[string]string{}
|
|
for _, d := range out.Data {
|
|
owner[d.ID] = d.OwnedBy
|
|
}
|
|
for _, name := range []string{"fake", "fake-alias"} {
|
|
if owner[name] != "llmgw-sim" {
|
|
t.Errorf("%s: owned_by = %q, want llmgw-sim", name, owner[name])
|
|
}
|
|
}
|
|
for _, name := range []string{"real", "real-alias"} {
|
|
if owner[name] != "llmgw" {
|
|
t.Errorf("%s: owned_by = %q, want llmgw", name, owner[name])
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestInjectIncludeUsagePreservesFields(t *testing.T) {
|
|
// The one place the router rewrites a request body; nothing may be lost.
|
|
in := []byte(`{"model":"m","stream":true,"temperature":0.7,"tools":[{"x":1}]}`)
|
|
out, ok := injectIncludeUsage(in)
|
|
if !ok {
|
|
t.Fatal("injection declined a well-formed body")
|
|
}
|
|
var m map[string]any
|
|
if err := json.Unmarshal(out, &m); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, k := range []string{"model", "stream", "temperature", "tools"} {
|
|
if _, present := m[k]; !present {
|
|
t.Errorf("field %q was dropped", k)
|
|
}
|
|
}
|
|
opts, _ := m["stream_options"].(map[string]any)
|
|
if opts["include_usage"] != true {
|
|
t.Errorf("include_usage not set: %v", m["stream_options"])
|
|
}
|
|
}
|
|
|
|
func TestInjectIncludeUsageKeepsExistingOptions(t *testing.T) {
|
|
in := []byte(`{"model":"m","stream":true,"stream_options":{"something":"else"}}`)
|
|
out, ok := injectIncludeUsage(in)
|
|
if !ok {
|
|
t.Fatal("injection declined")
|
|
}
|
|
var m map[string]any
|
|
if err := json.Unmarshal(out, &m); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
opts, ok := m["stream_options"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("stream_options is not an object: %v", m["stream_options"])
|
|
}
|
|
if opts["something"] != "else" {
|
|
t.Error("an existing stream_options field was dropped")
|
|
}
|
|
if opts["include_usage"] != true {
|
|
t.Error("include_usage not set")
|
|
}
|
|
}
|
|
|
|
func TestPromptPrefixIsBounded(t *testing.T) {
|
|
long := strings.Repeat("x", prefixBytes*3)
|
|
b := inferenceBody{Messages: []struct {
|
|
Role string `json:"role"`
|
|
Content json.RawMessage `json:"content"`
|
|
}{{Role: "user", Content: json.RawMessage(`"` + long + `"`)}}}
|
|
|
|
if got := len(promptPrefix(&b)); got > prefixBytes {
|
|
t.Errorf("prefix is %d bytes, over the %d cap", got, prefixBytes)
|
|
}
|
|
}
|
|
|
|
func TestSessionHeaderFallsBackToUser(t *testing.T) {
|
|
req := httptest.NewRequest(http.MethodPost, "/", nil)
|
|
if got := sessionID(req, "x-session-id", "user-42"); got != "user-42" {
|
|
t.Errorf("expected the user field to be used, got %q", got)
|
|
}
|
|
req.Header.Set("x-session-id", "sess-1")
|
|
if got := sessionID(req, "x-session-id", "user-42"); got != "sess-1" {
|
|
t.Errorf("the header should win, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestUnknownModelErrorNamesWhatIsAvailable(t *testing.T) {
|
|
// Reported by a user who copied the curl from the account page and sent
|
|
// its placeholder unchanged. "no endpoint serves \"MODEL\"" is true and
|
|
// useless: it reads the same for a typo, a placeholder and a real outage.
|
|
e := newFakeEngine(t, "m")
|
|
rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "MODEL", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusServiceUnavailable {
|
|
t.Fatalf("status %d, expected 503", rec.Code)
|
|
}
|
|
var body struct {
|
|
Error struct {
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
|
t.Fatalf("error body is not JSON: %q", rec.Body.String())
|
|
}
|
|
if !strings.Contains(body.Error.Message, `"MODEL"`) {
|
|
t.Errorf("the message does not repeat what was asked for: %q", body.Error.Message)
|
|
}
|
|
if !strings.Contains(body.Error.Message, "m") ||
|
|
!strings.Contains(body.Error.Message, "available") {
|
|
t.Errorf("the message does not say what is available: %q", body.Error.Message)
|
|
}
|
|
}
|
|
|
|
func TestTheErrorDoesNotBecomeAPage(t *testing.T) {
|
|
// A large catalog must not turn one error into a wall of names.
|
|
var many []string
|
|
for i := range 40 {
|
|
many = append(many, fmt.Sprintf("model-%02d", i))
|
|
}
|
|
msg := noEndpointFor("nope", many)
|
|
if !strings.Contains(msg, "and 30 more") {
|
|
t.Errorf("the list was not capped: %q", msg)
|
|
}
|
|
if strings.Contains(msg, "model-11") {
|
|
t.Errorf("more than the cap was listed: %q", msg)
|
|
}
|
|
}
|
|
|
|
func TestTheErrorIsHonestWhenNothingIsServed(t *testing.T) {
|
|
// Listing nothing as "available models are " would read as a truncated
|
|
// message rather than an empty cluster.
|
|
msg := noEndpointFor("m", nil)
|
|
if !strings.Contains(msg, "no model is currently being served") {
|
|
t.Errorf("an empty fleet is not explained: %q", msg)
|
|
}
|
|
}
|
|
|
|
func TestAnOverLongRequestIsRefusedBeforeItReachesAnEngine(t *testing.T) {
|
|
// The whole point of checking here: an over-long prompt used to be
|
|
// admitted, scheduled, given a replica and a slot, and refused by the
|
|
// engine — paying the cost of the request to learn something knowable from
|
|
// the request itself.
|
|
e := newFakeEngine(t, "m")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
|
|
m := catModel("m", nil)
|
|
m.ContextLength = 2048
|
|
cfg.Models = []catalog.Model{m}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
before := e.calls.Load()
|
|
rec := post(t, rt.Handler(), map[string]any{
|
|
"model": "m",
|
|
"messages": []any{map[string]any{"role": "user", "content": strings.Repeat("x", 40_000)}},
|
|
}, nil)
|
|
|
|
if rec.Code != http.StatusBadRequest {
|
|
t.Fatalf("status %d, expected 400; body %s", rec.Code, rec.Body.String())
|
|
}
|
|
if got := e.calls.Load(); got != before {
|
|
t.Errorf("the engine was called %d times for a request that could not fit", got-before)
|
|
}
|
|
if !strings.Contains(rec.Body.String(), "2048") {
|
|
t.Errorf("the error does not name the limit: %s", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestModelsEndpointPublishesContextAndCapabilities(t *testing.T) {
|
|
// A client sizing its prompts had nothing to read, and a client deciding
|
|
// whether to send tools had nothing either.
|
|
e := newFakeEngine(t, "m")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
|
|
m := catModel("m", nil)
|
|
m.ContextLength, m.MaxOutput = 32768, 4096
|
|
m.Capabilities = catalog.Capabilities{Tools: true}
|
|
cfg.Models = []catalog.Model{m}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
rec := httptest.NewRecorder()
|
|
rt.Handler().ServeHTTP(rec, req)
|
|
|
|
var out struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
Context int `json:"context_length"`
|
|
MaxOutput int `json:"max_output_tokens"`
|
|
Capabilities map[string]bool `json:"capabilities"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var found bool
|
|
for _, d := range out.Data {
|
|
if d.ID != "m" {
|
|
continue
|
|
}
|
|
found = true
|
|
if d.Context != 32768 || d.MaxOutput != 4096 {
|
|
t.Errorf("context %d, max output %d", d.Context, d.MaxOutput)
|
|
}
|
|
if !d.Capabilities["tools"] || d.Capabilities["vision"] {
|
|
t.Errorf("capabilities wrong: %v", d.Capabilities)
|
|
}
|
|
}
|
|
if !found {
|
|
t.Fatal("the model was not listed")
|
|
}
|
|
}
|
|
|
|
func TestAModelWithNoContextPublishesNoContextField(t *testing.T) {
|
|
// Zero would read as "no context", which is worse than saying nothing.
|
|
e := newFakeEngine(t, "m")
|
|
rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
rec := httptest.NewRecorder()
|
|
rt.Handler().ServeHTTP(rec, req)
|
|
if strings.Contains(rec.Body.String(), "context_length") {
|
|
t.Errorf("an unknown context was published as a number: %s", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
type fakeActivator struct {
|
|
onDemand map[string]bool
|
|
asked []string
|
|
}
|
|
|
|
func (f *fakeActivator) Want(model string) bool {
|
|
f.asked = append(f.asked, model)
|
|
return f.onDemand[model]
|
|
}
|
|
|
|
func TestALoadingModelIsToldToRetryRatherThanRefused(t *testing.T) {
|
|
// "no endpoint serves this" would be a lie while a model is starting, and
|
|
// the difference matters: a client told to retry succeeds, one told the
|
|
// model does not exist gives up.
|
|
e := newFakeEngine(t, "m")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
|
|
// A very short wait, because the point here is what happens when it runs
|
|
// out rather than how long it is.
|
|
cfg.Router.ActivationWait = 50 * time.Millisecond
|
|
rt := newTestRouter(t, cfg)
|
|
rt.SetActivator(&fakeActivator{onDemand: map[string]bool{"sleepy": true}})
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "sleepy", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusServiceUnavailable {
|
|
t.Fatalf("status %d, expected 503", rec.Code)
|
|
}
|
|
if got := rec.Header().Get("Retry-After"); got == "" {
|
|
t.Error("no Retry-After; a well-behaved client backs off as it would from a permanent failure")
|
|
}
|
|
if !strings.Contains(rec.Body.String(), "loading") {
|
|
t.Errorf("the client is not told what is happening: %s", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestAModelThatIsSimplyUnknownStillSaysSo(t *testing.T) {
|
|
// The other half. On-demand must not turn every typo into "loading",
|
|
// because then nothing is ever wrong and the useful error disappears.
|
|
e := newFakeEngine(t, "m")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
|
|
rt := newTestRouter(t, cfg)
|
|
rt.SetActivator(&fakeActivator{onDemand: map[string]bool{}})
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "typo", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusServiceUnavailable {
|
|
t.Fatalf("status %d", rec.Code)
|
|
}
|
|
if rec.Header().Get("Retry-After") != "" {
|
|
t.Error("an unknown model was presented as retryable")
|
|
}
|
|
if !strings.Contains(rec.Body.String(), "available models") {
|
|
t.Errorf("the useful error was lost: %s", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestAServedModelNeverWaitsForActivation(t *testing.T) {
|
|
// A model that is up must not pay a millisecond for the feature.
|
|
e := newFakeEngine(t, "m")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
|
|
cfg.Router.ActivationWait = 10 * time.Second
|
|
rt := newTestRouter(t, cfg)
|
|
act := &fakeActivator{onDemand: map[string]bool{"m": true}}
|
|
rt.SetActivator(act)
|
|
|
|
start := time.Now()
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
|
}
|
|
if elapsed := time.Since(start); elapsed > time.Second {
|
|
t.Errorf("a served request took %s; activation is on the hot path", elapsed)
|
|
}
|
|
if len(act.asked) != 0 {
|
|
t.Errorf("the activator was consulted for a model that is already serving: %v", act.asked)
|
|
}
|
|
}
|
|
|
|
func TestWithoutAnActivatorNothingChanges(t *testing.T) {
|
|
// Every deployment before this existed, and a router running without a
|
|
// control plane.
|
|
e := newFakeEngine(t, "m")
|
|
rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "nope", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusServiceUnavailable {
|
|
t.Errorf("status %d", rec.Code)
|
|
}
|
|
if !strings.Contains(rec.Body.String(), "available models") {
|
|
t.Errorf("body = %s", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestTheServedModelHeaderArrivesBeforeTheBody(t *testing.T) {
|
|
// It used to be set in finish, which runs after the proxy has written the
|
|
// response — so it announced a fallback to nobody. Nothing failed, because
|
|
// the only thing asserting it was a comment.
|
|
// Nothing serves the primary at all, which is the only condition that
|
|
// justifies moving to another model — an upstream returning 500 is a
|
|
// failure the fallback would share, and dispatch rightly declines to try.
|
|
live := newFakeEngine(t, "backup")
|
|
cfg := baseConfig(config.Endpoint{ID: "b", URL: live.srv.URL, Models: []string{"backup"}})
|
|
m := catModel("primary", nil, "backup")
|
|
cfg.Models = []catalog.Model{m, catModel("backup", nil)}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "primary", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
|
}
|
|
if got := rec.Header().Get("x-llmgw-served-model"); got != "backup" {
|
|
t.Errorf("served-model header = %q, want backup", got)
|
|
}
|
|
}
|
|
|
|
type unservableActivator struct {
|
|
fakeActivator
|
|
why map[string]string
|
|
}
|
|
|
|
func (u *unservableActivator) Unservable(model string) (string, bool) {
|
|
w, ok := u.why[model]
|
|
return w, ok
|
|
}
|
|
|
|
func TestAModelThatCannotComeUpIsSaidSoAtOnce(t *testing.T) {
|
|
// A model that fits on no node, or failed every time it started, was
|
|
// answered "loading; retry shortly" — and a client that retried did so
|
|
// forever.
|
|
e := newFakeEngine(t, "m")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
|
|
cfg.Router.ActivationWait = 10 * time.Second
|
|
rt := newTestRouter(t, cfg)
|
|
rt.SetActivator(&unservableActivator{
|
|
fakeActivator: fakeActivator{onDemand: map[string]bool{"phi-4": true}},
|
|
why: map[string]string{"phi-4": "phi-4 cannot be placed on any node: needs 15 GB"},
|
|
})
|
|
|
|
start := time.Now()
|
|
rec := post(t, rt.Handler(), map[string]any{"model": "phi-4", "messages": []any{}}, nil)
|
|
if took := time.Since(start); took > 2*time.Second {
|
|
t.Errorf("the request waited %s for a model that will not load", took)
|
|
}
|
|
if rec.Code != http.StatusServiceUnavailable || rec.Header().Get("Retry-After") != "" {
|
|
t.Errorf("status %d, Retry-After %q", rec.Code, rec.Header().Get("Retry-After"))
|
|
}
|
|
if !strings.Contains(rec.Body.String(), "needs 15 GB") {
|
|
t.Errorf("the reason was not given: %s", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestAMadeUpModelNameDoesNotBecomeAMetricLabel(t *testing.T) {
|
|
// Counted as sent, every request for a nonsense model minted a new series.
|
|
e := newFakeEngine(t, "m")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
|
|
cfg.Models = []catalog.Model{{Name: "m", Aliases: []string{"mm"}}}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
if got := rt.metricModel("mm"); got != "m" {
|
|
t.Errorf("an alias is counted as %q, not the model it names", got)
|
|
}
|
|
for i := 0; i < 20; i++ {
|
|
post(t, rt.Handler(), map[string]any{"model": fmt.Sprintf("junk-%d", i), "messages": []any{}}, nil)
|
|
}
|
|
// Twenty names, one series: all of them "unknown".
|
|
reg := prometheus.NewRegistry()
|
|
reg.MustRegister(rt.m.requests)
|
|
mfs, err := reg.Gather()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
series := 0
|
|
for _, mf := range mfs {
|
|
series += len(mf.GetMetric())
|
|
}
|
|
if series != 1 {
|
|
t.Errorf("%d series after twenty made-up names", series)
|
|
}
|
|
}
|
|
|
|
func TestAFallbackTheKeyMayNotUseIsNotServed(t *testing.T) {
|
|
primary := newFakeEngine(t, "primary")
|
|
primary.status.Store(http.StatusServiceUnavailable)
|
|
backup := newFakeEngine(t, "backup")
|
|
cfg := baseConfig(
|
|
config.Endpoint{ID: "p", URL: primary.srv.URL, Models: []string{"primary"}},
|
|
config.Endpoint{ID: "b", URL: backup.srv.URL, Models: []string{"backup"}},
|
|
)
|
|
cfg.Models = []catalog.Model{catModel("primary", nil, "backup"), catModel("backup", nil)}
|
|
rt := newTestRouter(t, cfg)
|
|
asKey := func(allowed ...string) http.Handler {
|
|
p := auth.Principal{KeyID: "k1", Limits: store.Effective{AllowedModels: allowed}}
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
rt.Handler().ServeHTTP(w, r.WithContext(auth.WithPrincipal(r.Context(), p)))
|
|
})
|
|
}
|
|
|
|
rec := post(t, asKey("primary"), map[string]any{"model": "primary", "messages": []any{}}, nil)
|
|
if rec.Code == http.StatusOK || backup.calls.Load() != 0 {
|
|
t.Fatalf("a key allowed only primary was served by backup: %d", rec.Code)
|
|
}
|
|
rec = post(t, asKey("primary", "backup"), map[string]any{"model": "primary", "messages": []any{}}, nil)
|
|
if rec.Code != http.StatusOK || rec.Header().Get("x-llmgw-served-model") != "backup" {
|
|
t.Fatalf("a key allowed both was not given the fallback: %d", rec.Code)
|
|
}
|
|
}
|
|
|
|
func TestModelsEndpointListsOnlyWhatTheKeyMayUse(t *testing.T) {
|
|
e := newFakeEngine(t, "real")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"real", "other"}})
|
|
cfg.Models = []catalog.Model{catModel("real", []string{"real-alias"}), catModel("other", nil)}
|
|
rt := newTestRouter(t, cfg)
|
|
|
|
p := auth.Principal{Limits: store.Effective{AllowedModels: []string{"real"}}}
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
req = req.WithContext(auth.WithPrincipal(req.Context(), p))
|
|
rec := httptest.NewRecorder()
|
|
rt.handleModels(rec, req)
|
|
|
|
var out struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var ids []string
|
|
for _, d := range out.Data {
|
|
ids = append(ids, d.ID)
|
|
}
|
|
if !slices.Equal(ids, []string{"real", "real-alias"}) {
|
|
t.Errorf("listed %v; the key may use real and its alias only", ids)
|
|
}
|
|
}
|
|
|
|
func TestAStoppedServiceAnswers503AndStillListsModels(t *testing.T) {
|
|
e := newFakeEngine(t, "m")
|
|
rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
|
|
h := rt.Handler()
|
|
rt.SetStopped(true)
|
|
rec := post(t, h, map[string]any{"model": "m", "messages": []map[string]string{{"role": "user", "content": "hi"}}}, nil)
|
|
if rec.Code != http.StatusServiceUnavailable || !strings.Contains(rec.Body.String(), "system_stopped") ||
|
|
rec.Header().Get("Retry-After") == "" {
|
|
t.Fatalf("stopped: %d %s", rec.Code, rec.Body.String())
|
|
}
|
|
list := httptest.NewRecorder()
|
|
h.ServeHTTP(list, httptest.NewRequest(http.MethodGet, "/v1/models", nil))
|
|
if list.Code == http.StatusServiceUnavailable {
|
|
t.Error("listing models was refused too")
|
|
}
|
|
rt.SetStopped(false)
|
|
if rec := post(t, h, map[string]any{"model": "m", "messages": []map[string]string{{"role": "user", "content": "hi"}}}, nil); rec.Code == http.StatusServiceUnavailable {
|
|
t.Errorf("started again, still refused: %s", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestABatchOfEmbeddingsIsCheckedTextByText(t *testing.T) {
|
|
// llmgw-kb sends many passages in one request; each is embedded on its
|
|
// own, so their sum over the context is no reason to refuse them. One
|
|
// passage over the context still is.
|
|
e := newFakeEngine(t, "emb")
|
|
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"emb"}})
|
|
m := catModel("emb", nil)
|
|
m.ContextLength = 2048
|
|
cfg.Models = []catalog.Model{m}
|
|
rt := newTestRouter(t, cfg)
|
|
send := func(input []string) *httptest.ResponseRecorder {
|
|
b, _ := json.Marshal(map[string]any{"model": "emb", "input": input})
|
|
rec := httptest.NewRecorder()
|
|
rt.Handler().ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/embeddings", bytes.NewReader(b)))
|
|
return rec
|
|
}
|
|
passage := strings.Repeat("x", 6000) // about 1500 tokens
|
|
if rec := send([]string{passage, passage, passage, passage}); rec.Code == http.StatusBadRequest {
|
|
t.Errorf("four passages that each fit were refused: %s", rec.Body.String())
|
|
}
|
|
if rec := send([]string{"short", strings.Repeat("x", 40_000)}); rec.Code != http.StatusBadRequest {
|
|
t.Errorf("a passage over the context was let through: %d", rec.Code)
|
|
}
|
|
}
|