Learn per-model context from /props?model=; only status "loaded" is loaded
Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
This commit is contained in:
+12
-12
@@ -11,9 +11,9 @@ import (
|
||||
)
|
||||
|
||||
// movedHeader announces a context-driven move in the response header:
|
||||
// "moved: old><new". The client learns which host served it.
|
||||
// "moved: old>new". The client learns which host served it.
|
||||
func movedHeader(from, to string) string {
|
||||
return "moved:" + from + "><" + to
|
||||
return "moved:" + from + ">" + to
|
||||
}
|
||||
|
||||
// drainer is the drain flag the guard consults when rule 3 prefers a host not
|
||||
@@ -54,7 +54,7 @@ func (p *Handler) guard(w http.ResponseWriter, r *http.Request, hosts []string,
|
||||
// already fits its per-slot context.
|
||||
psc := 0
|
||||
if s, ok := p.health.Get(host); ok {
|
||||
psc = s.PerSlotCtx()
|
||||
psc = s.PerSlotCtxFor(model)
|
||||
}
|
||||
if psc == 0 || estimate <= psc {
|
||||
return host, "", 0, false
|
||||
@@ -84,7 +84,7 @@ func (p *Handler) guard(w http.ResponseWriter, r *http.Request, hosts []string,
|
||||
}
|
||||
// A host with a known per-slot context smaller than the estimate
|
||||
// cannot serve it no matter how it wakes.
|
||||
if psc := s.PerSlotCtx(); psc != 0 && psc < estimate {
|
||||
if psc := s.PerSlotCtxFor(model); psc != 0 && psc < estimate {
|
||||
continue
|
||||
}
|
||||
if p.cfg.Hosts[name].Wake == nil || !p.waker.Wake(r.Context(), name) {
|
||||
@@ -100,7 +100,7 @@ func (p *Handler) guard(w http.ResponseWriter, r *http.Request, hosts []string,
|
||||
|
||||
// Rule 4: no host fits. Answer 400 with the estimate and the largest
|
||||
// available per-slot context, and record the row.
|
||||
p.refuseCtx(w, host, route, model, fp, started, estimate, largestSlotCtx(hosts, p.health))
|
||||
p.refuseCtx(w, host, route, model, fp, started, estimate, largestSlotCtx(hosts, p.health, model))
|
||||
return host, "", estimate, true
|
||||
}
|
||||
|
||||
@@ -120,7 +120,7 @@ func ctxFitHost(hosts []string, model string, estimate int, h Health, cfg *confi
|
||||
if !ok || !s.Healthy || isDraining(drn, name) {
|
||||
continue
|
||||
}
|
||||
if s.PerSlotCtx() >= estimate && contains(s.Loaded, model) {
|
||||
if s.PerSlotCtxFor(model) >= estimate && contains(s.Loaded, model) {
|
||||
return name, true
|
||||
}
|
||||
}
|
||||
@@ -132,23 +132,23 @@ func ctxFitHost(hosts []string, model string, estimate int, h Health, cfg *confi
|
||||
if !ok || !s.Healthy || isDraining(drn, name) {
|
||||
continue
|
||||
}
|
||||
if s.PerSlotCtx() >= estimate && cfg.Serves(name, model) {
|
||||
if s.PerSlotCtxFor(model) >= estimate && cfg.Serves(name, model) {
|
||||
return name, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// largestSlotCtx is the largest PerSlotCtx across the route's healthy hosts, or
|
||||
// 0 when none is healthy or reports a context size.
|
||||
func largestSlotCtx(hosts []string, h Health) int {
|
||||
// largestSlotCtx is the largest per-slot context for model across the route's
|
||||
// healthy hosts that have it loaded, or 0 when none is healthy or reports one.
|
||||
func largestSlotCtx(hosts []string, h Health, model string) int {
|
||||
best := 0
|
||||
for _, name := range hosts {
|
||||
s, ok := h.Get(name)
|
||||
if !ok || !s.Healthy {
|
||||
if !ok || !s.Healthy || !contains(s.Loaded, model) {
|
||||
continue
|
||||
}
|
||||
if psc := s.PerSlotCtx(); psc > best {
|
||||
if psc := s.PerSlotCtxFor(model); psc > best {
|
||||
best = psc
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
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())
|
||||
}
|
||||
var e map[string]any
|
||||
if err := json.Unmarshal([]byte(body), &e); err != nil {
|
||||
t.Fatalf("body %q is not JSON: %v", body, err)
|
||||
}
|
||||
if max, _ := e["max"].(float64); max != 4096 {
|
||||
t.Errorf("max = %v, want 4096: the largest per-slot context among hosts that have shared loaded", e["max"])
|
||||
}
|
||||
}
|
||||
@@ -29,7 +29,7 @@ const (
|
||||
HostHeader = "X-Crossbar-Host"
|
||||
LeaseHeader = "X-Crossbar-Lease" // "new" or "reused"
|
||||
RouteHeader = "X-Crossbar-Route" // client may name the route here instead of the path
|
||||
CtxHeader = "X-Crossbar-Ctx" // "moved:<old>><new>" when the context was relocated
|
||||
CtxHeader = "X-Crossbar-Ctx" // "moved:<old>new" when the context was relocated
|
||||
)
|
||||
|
||||
// errBodyTooLarge is returned when a request body exceeds MaxBody during the model peek.
|
||||
|
||||
Reference in New Issue
Block a user