111 lines
4.3 KiB
Go
111 lines
4.3 KiB
Go
package proxy_test
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"git.wntrmute.dev/kyle/crossbar/internal/proxy"
|
|
)
|
|
|
|
// routerUpstream is a fake shaped like llama-server's router mode: the plain /props carries no
|
|
// context (role router, n_ctx 0), /v1/models lists models with a status, and /props?model=X
|
|
// answers for one loaded model. The guard must work from the per-model figures.
|
|
func routerUpstream(t *testing.T, name string, models map[string][2]int, unloaded ...string) *upstream {
|
|
u := &upstream{name: name}
|
|
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":[`)
|
|
first := true
|
|
for id := range models {
|
|
if !first {
|
|
fmt.Fprint(w, ",")
|
|
}
|
|
first = false
|
|
fmt.Fprintf(w, `{"id":%q,"status":{"value":"loaded"}}`, id)
|
|
}
|
|
for _, id := range unloaded {
|
|
fmt.Fprintf(w, `,{"id":%q,"status":{"value":"unloaded"}}`, id)
|
|
}
|
|
fmt.Fprint(w, `]}`)
|
|
})
|
|
mux.HandleFunc("/props", func(w http.ResponseWriter, r *http.Request) {
|
|
model := r.URL.Query().Get("model")
|
|
if model == "" {
|
|
fmt.Fprint(w, `{"role":"router","default_generation_settings":{"n_ctx":0}}`)
|
|
return
|
|
}
|
|
m, ok := models[model]
|
|
if !ok {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
fmt.Fprint(w, `{"error":"asked for a model that is not loaded"}`)
|
|
return
|
|
}
|
|
fmt.Fprintf(w, `{"default_generation_settings":{"n_ctx":%d},"total_slots":%d}`, m[0], m[1])
|
|
})
|
|
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
|
|
u.hits.Add(1)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
|
})
|
|
u.srv = httptest.NewServer(mux)
|
|
t.Cleanup(u.srv.Close)
|
|
return u
|
|
}
|
|
|
|
// On routers the guard reads the per-model per-slot context: `small` serves "shared" from a
|
|
// 8192-context child with two slots (4096 per slot), `big` from a 131072-context child with one.
|
|
// The plain /props of both says nothing, so a v2 guard that only knew host-level figures would
|
|
// stay inert and let the oversized prompt overflow `small`.
|
|
func TestRouterGuardUsesPerModelContext(t *testing.T) {
|
|
small := routerUpstream(t, "small", map[string][2]int{"shared": {8192, 2}})
|
|
big := routerUpstream(t, "big", map[string][2]int{"shared": {131072, 1}})
|
|
r := newRig(t, ctxHosts, small, big)
|
|
|
|
resp := r.post("/r/v1/chat/completions", bodyOfTokens(100))
|
|
drain(resp)
|
|
if got := resp.Header.Get(proxy.HostHeader); got != "small" {
|
|
t.Fatalf("small prompt went to %q, want small (weight 10)", got)
|
|
}
|
|
resp = r.post("/r/v1/chat/completions", bodyOfTokens(6000))
|
|
drain(resp)
|
|
if resp.StatusCode != 200 || resp.Header.Get(proxy.HostHeader) != "big" {
|
|
t.Fatalf("6000-token prompt: status %d host %q, want 200 on big", resp.StatusCode, resp.Header.Get(proxy.HostHeader))
|
|
}
|
|
if got := resp.Header.Get("X-Crossbar-Ctx"); got != "moved:small>big" {
|
|
t.Errorf("X-Crossbar-Ctx = %q, want moved:small>big", got)
|
|
}
|
|
}
|
|
|
|
// A model the router lists as unloaded is not resident there: the guard must not move a prompt
|
|
// to that host, and the refusal names the largest per-slot context among hosts that do serve it.
|
|
func TestRouterUnloadedModelIsNotACandidate(t *testing.T) {
|
|
small := routerUpstream(t, "small", map[string][2]int{"shared": {8192, 2}})
|
|
big := routerUpstream(t, "big", map[string][2]int{"other": {131072, 1}}, "shared") // shared unloaded on big
|
|
r := newRig(t, ctxHosts, small, big)
|
|
|
|
resp := r.post("/r/v1/chat/completions", bodyOfTokens(6000))
|
|
body := drain(resp)
|
|
if resp.StatusCode != 400 {
|
|
t.Fatalf("status %d body %s, want 400: shared is loaded only on small, where it does not fit", resp.StatusCode, body)
|
|
}
|
|
if big.hits.Load() != 0 {
|
|
t.Errorf("big served %d requests for a model it does not have loaded", big.hits.Load())
|
|
}
|
|
// v2.3: the refusal is llama-server's exceed_context_size_error shape; n_ctx is what "max" was.
|
|
var e struct {
|
|
Error struct {
|
|
NCtx float64 `json:"n_ctx"`
|
|
} `json:"error"`
|
|
}
|
|
if err := json.Unmarshal([]byte(body), &e); err != nil {
|
|
t.Fatalf("body %q is not JSON: %v", body, err)
|
|
}
|
|
if e.Error.NCtx != 4096 {
|
|
t.Errorf("error.n_ctx = %v, want 4096: the largest per-slot context among hosts that have shared loaded", e.Error.NCtx)
|
|
}
|
|
}
|