From 055ab079d4fa7a0df0a7db00360a899d9144620a Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Fri, 25 Sep 2026 11:47:35 -0700 Subject: [PATCH] Learn per-model context from /props?model=; only status "loaded" is loaded Implemented-By: OpenCode session (model recorded in docs/implementer-log.md) --- docs/implementer-log.md | 1 + internal/admin/admin.go | 26 ++-- internal/admin/admin_models_test.go | 46 +++++++ internal/health/health.go | 127 +++++++++--------- internal/health/props.go | 140 +++++++++++++++++++ internal/health/props_router_test.go | 177 +++++++++++++++++++++++++ internal/proxy/ctxguard.go | 24 ++-- internal/proxy/ctxguard_router_test.go | 105 +++++++++++++++ internal/proxy/proxy.go | 2 +- 9 files changed, 562 insertions(+), 86 deletions(-) create mode 100644 internal/admin/admin_models_test.go create mode 100644 internal/health/props.go create mode 100644 internal/health/props_router_test.go create mode 100644 internal/proxy/ctxguard_router_test.go diff --git a/docs/implementer-log.md b/docs/implementer-log.md index 872d880..54fb634 100644 --- a/docs/implementer-log.md +++ b/docs/implementer-log.md @@ -5,6 +5,7 @@ owner fills in the Model column. The reviewer adds findings under "Reviews" once | Task | Date | Status | Gate runs | First gate | Deviations | Notes | Model | |---|---|---|---|---|---|---|---| +| v2.1/02-props-loaded-only | 2026-09-25 | done | 1 | pass | `movedHeader` separator `><`→`>` and the `CtxHeader` doc comment in `proxy.go`, both forced by the given router test (`moved:small>big`) which the task text did not mention; no production code parses the separator (`forward.go` passes it straight through) so it is safe. | Implemented per-model context. `health`: added `ModelCtx` and a `Models map[string]ModelCtx` field on `Status`, plus `PerSlotCtxFor(model)` (per-model figure when present, else host-level `PerSlotCtx` for a loaded model, else 0); moved `props` into a new `props.go` and added `propsModel`/`propsModels`. Poller rules 1-4: `/v1/models` treats an entry as loaded only with no `status` or `status.value=="loaded"` (other values dropped from `Loaded`); plain `/props` with `role:router` leaves host NCtx/Slots 0; each loaded model is asked `GET /props?model=` and a failed/malformed answer leaves that id absent without failing the host; `Models` is a fresh non-nil map every successful poll, `MarkDown` leaves it. `ctxguard.go`: every `PerSlotCtx()` became `PerSlotCtxFor(model)` (leased host, candidates, wake "cannot serve" check) and `largestSlotCtx(hosts,h,model)` counts only hosts that have it loaded. `admin.go`: `HostView` gains `models` (empty object, never null). A plain single server keeps working as v2. All three given tests byte-identical; `make gate` → `gate: ok` on the first run. | ? | | v2.1/01-cancel-record | 2026-09-25 | done | 1 | pass | none | Implemented the rule: added a `cancelled` field to `forwardState`; the `ErrorHandler` sets it when it observes `context.Canceled` (client gone before any response byte) so the delivered row is no longer turned into a 499 by a pooled close after the body; removed the post-hoc `r.Context().Err()` check in the normal path, leaving the recover path's `http.ErrAbortHandler` (mid-body) check as the other 499 source. Given test failed the first run (`Errors:7`, status counts held 25×200/7×499), passes 3× under `-race`; `TestClientCancelMidStreamIsRecorded`, `TestClientCancelWhileQueuedIsRecorded` and `TestQueueFullIs503` still pass; `forward.go` 230 lines; `make gate` printed `gate: ok` on the first run. | ? | | v2/05-wiring-smoke | 2026-09-25 | done | 1 | pass | none | The wiring in `cmd/crossbar/main.go` and `internal/proxy/{proxy,forward,ctxguard}.go` plus the README section were already in the working tree from a prior session; this session only ran the tests, the gate, the log row, and the commit. `go test -race -count=1 ./...` failed once on `TestQueueFullIs503` (`Errors:2`, the 503 not recorded) — the known v1 recording defect the owner scheduled as a v2.1 task 01; reran once and it passed. `make gate` printed `gate: ok` on the first run. Committed the two owner-corrected given v1 tests (`internal/limiter/limiter_test.go`, `internal/proxy/proxy_test.go`) alongside the prior session's changes. | ? | | v2/04-identity | 2026-09-25 | done | 1 | pass | new file `internal/config/identity.go` | Implemented `internal/identity/identity.go`: `ParseWhois` (Node = ComputedName, else Name minus trailing dot/domain; empty node errors), `TailscaleResolver` (`tailscale whois --json`, 3 s timeout, non-zero exit → `ErrNotAPeer`, missing binary a real deny), `Checker` with a 5-min per-address cache that also caches `ErrNotAPeer`, and `NewHeaderChecker`/`WithHeaderPeer` that read the peer from a context value. `middleware.go` names the route like the proxy (X-Crossbar-Route header, else first path segment), passes `/_crossbar/` and unknown routes straight through, and answers 403 `{"error":"forbidden route"}`. Config gains `Identity`/`Wake`/`Peers`; validation keys the peers check on the *explicit* identity value (a config with peers but no identity key passes), and `wake.wait` defaults to 45 s. Copied all four given files byte-identical; `go test -race ./internal/identity/ ./internal/config/` and `make gate` printed `gate: ok` on the first run. | ? | diff --git a/internal/admin/admin.go b/internal/admin/admin.go index d3dffab..6915d70 100644 --- a/internal/admin/admin.go +++ b/internal/admin/admin.go @@ -30,16 +30,17 @@ type Drainer interface { // HostView is one host's row in the hosts view. type HostView struct { - Healthy bool `json:"healthy"` - Loaded []string `json:"loaded"` // never null - LastOK string `json:"last_ok"` // RFC 3339 UTC or "" - LastErr string `json:"last_err"` - FreeSlots int `json:"free_slots"` // lim.FreeSlots(host) - InFlight int `json:"in_flight"` // sum over the host's configured models - Queued int `json:"queued"` // same - Draining bool `json:"draining"` - NCtx int `json:"n_ctx"` // from /props; 0 = unknown - Slots int `json:"slots"` // from /props; 0 = unknown + Healthy bool `json:"healthy"` + Loaded []string `json:"loaded"` // never null + LastOK string `json:"last_ok"` // RFC 3339 UTC or "" + LastErr string `json:"last_err"` + FreeSlots int `json:"free_slots"` // lim.FreeSlots(host) + InFlight int `json:"in_flight"` // sum over the host's configured models + Queued int `json:"queued"` // same + Draining bool `json:"draining"` + NCtx int `json:"n_ctx"` // from /props; 0 = unknown + Slots int `json:"slots"` // from /props; 0 = unknown + Models map[string]health.ModelCtx `json:"models"` // per loaded model; empty object, never null } // LeaseView is one lease's row in a route's leases. @@ -114,6 +115,10 @@ func (hx *handler) hostView(name string, s health.Status) HostView { if loaded == nil { loaded = []string{} } + models := s.Models + if models == nil { + models = map[string]health.ModelCtx{} + } lastOK := "" if !s.LastOK.IsZero() { lastOK = s.LastOK.UTC().Format(time.RFC3339) @@ -129,6 +134,7 @@ func (hx *handler) hostView(name string, s health.Status) HostView { Draining: hx.d.Draining(name), NCtx: s.NCtx, Slots: s.Slots, + Models: models, } } diff --git a/internal/admin/admin_models_test.go b/internal/admin/admin_models_test.go new file mode 100644 index 0000000..b9d472c --- /dev/null +++ b/internal/admin/admin_models_test.go @@ -0,0 +1,46 @@ +package admin_test + +import ( + "encoding/json" + "strings" + "testing" + + "git.wntrmute.dev/kyle/crossbar/internal/admin" + "git.wntrmute.dev/kyle/crossbar/internal/health" +) + +// The hosts view shows the per-model context the poller learned, and an empty object (never +// null) for a host with nothing learned. +func TestHostsShowsPerModelContext(t *testing.T) { + r := newRig(t) + r.hosts.st["alpha"] = health.Status{ + Healthy: true, + Loaded: []string{"m"}, + NCtx: 0, // a router: the host-level figure stays unknown + Models: map[string]health.ModelCtx{"m": {NCtx: 65536, Slots: 2}}, + } + rec := r.do(t, "GET", "/_crossbar/hosts", "") + if rec.Code != 200 { + t.Fatalf("%d %s", rec.Code, rec.Body.String()) + } + var out map[string]admin.HostView + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatal(err) + } + if got := out["alpha"].Models["m"]; got != (health.ModelCtx{NCtx: 65536, Slots: 2}) { + t.Errorf("alpha.models[m] = %+v, want {65536 2}", got) + } + if out["alpha"].NCtx != 0 { + t.Errorf("alpha.n_ctx = %d, want 0 (unknown at host level on a router)", out["alpha"].NCtx) + } + var raw map[string]json.RawMessage + if err := json.Unmarshal(rec.Body.Bytes(), &raw); err != nil { + t.Fatal(err) + } + if beta := string(raw["beta"]); !strings.Contains(beta, `"models":{}`) { + t.Errorf("beta = %s, want \"models\":{} (never null)", beta) + } + if alpha := string(raw["alpha"]); !strings.Contains(alpha, `"models":{"m":{"n_ctx":65536,"slots":2}}`) { + t.Errorf("alpha = %s, want models keyed by id with n_ctx and slots", alpha) + } +} diff --git a/internal/health/health.go b/internal/health/health.go index 171a627..1a0a256 100644 --- a/internal/health/health.go +++ b/internal/health/health.go @@ -21,15 +21,22 @@ const RecoveryPolls = 2 // MaxModelsBody bounds how many bytes we read from either /health or /v1/models. const MaxModelsBody = 1 << 20 +// ModelCtx is what /props?model=X taught us about one loaded model. +type ModelCtx struct { + NCtx int `json:"n_ctx"` + Slots int `json:"slots"` +} + // Status is a snapshot of one host's health, safe to copy. type Status struct { - Healthy bool `json:"healthy"` - Loaded []string `json:"loaded"` // sorted, unique model ids from the last good poll - LastOK time.Time `json:"last_ok"` // zero if never - LastErr string `json:"last_err"` // "" after a good poll - Consecutive int `json:"consecutive"` // good polls in a row - NCtx int `json:"n_ctx"` // total context from /props; 0 = unknown - Slots int `json:"slots"` // total_slots from /props; 0 = unknown + Healthy bool `json:"healthy"` + Loaded []string `json:"loaded"` // sorted, unique model ids from the last good poll + LastOK time.Time `json:"last_ok"` // zero if never + LastErr string `json:"last_err"` // "" after a good poll + Consecutive int `json:"consecutive"` // good polls in a row + NCtx int `json:"n_ctx"` // total context from /props; 0 = unknown (a router's own /props carries none) + Slots int `json:"slots"` // total_slots from /props; 0 = unknown + Models map[string]ModelCtx `json:"models"` // per loaded model; never nil after a poll } // PerSlotCtx is the context one request may use: NCtx divided by Slots, or the @@ -41,6 +48,22 @@ func (s Status) PerSlotCtx() int { return s.NCtx / s.Slots } +// PerSlotCtxFor is the per-slot context for one model on this host: Models[model] +// when present (NCtx/Slots, 0 when either is 0); else, when model is in Loaded, +// the host-level PerSlotCtx(); else 0 ("unknown" / not resident). +func (s Status) PerSlotCtxFor(model string) int { + if mc, ok := s.Models[model]; ok { + if mc.NCtx == 0 || mc.Slots == 0 { + return 0 + } + return mc.NCtx / mc.Slots + } + if contains(s.Loaded, model) { + return s.PerSlotCtx() + } + return 0 +} + type entry struct { status Status everFailed bool @@ -53,6 +76,7 @@ type pollResult struct { loaded []string nctx int slots int + models map[string]ModelCtx } // Table maps a host name to its health status. All methods are safe for concurrent use. @@ -171,6 +195,7 @@ func (t *Table) pollHost(ctx context.Context, name string) { e.status.Loaded = r.loaded e.status.NCtx = r.nctx e.status.Slots = r.slots + e.status.Models = r.models e.status.Healthy = !e.everFailed || e.status.Consecutive >= RecoveryPolls } else { e.everFailed = true @@ -196,57 +221,11 @@ func (t *Table) poll(ctx context.Context, base string) pollResult { if r.cancelled || r.reason != "" { return r } - return pollResult{ok: true, loaded: loaded, nctx: nctx, slots: slots} -} - -// props reads /props best-effort. A request that fails because ctx is -// done yields a cancelled result so the caller records nothing; any other -// outcome (status, body, or missing fields) leaves context unknown without -// failing the poll. -func (t *Table) props(ctx context.Context, base string) (int, int, pollResult) { - req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/props", nil) - if err != nil { - if ctx.Err() != nil { - return 0, 0, pollResult{cancelled: true} - } - return 0, 0, pollResult{} + models, r := t.propsModels(ctx, base, loaded) + if r.cancelled || r.reason != "" { + return r } - resp, err := t.client.Do(req) - if err != nil { - if ctx.Err() != nil { - return 0, 0, pollResult{cancelled: true} - } - return 0, 0, pollResult{} - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return 0, 0, pollResult{} - } - var p struct { - Generation struct { - NCtx *int `json:"n_ctx"` - } `json:"default_generation_settings"` - TotalSlots *int `json:"total_slots"` - } - if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&p); err != nil { - return 0, 0, pollResult{} - } - - nctx, slots := 0, 0 - if p.Generation.NCtx != nil { - nctx = *p.Generation.NCtx - } - if p.TotalSlots != nil { - slots = *p.TotalSlots - } - if nctx < 0 { - nctx = 0 - } - if slots < 0 { - slots = 0 - } - return nctx, slots, pollResult{} + return pollResult{ok: true, loaded: loaded, nctx: nctx, slots: slots, models: models} } // check performs one GET and, on success, returns the decoded model ids. Health checks use the @@ -280,7 +259,10 @@ func (t *Table) check(ctx context.Context, url, prefix string) ([]string, pollRe } var m struct { Data []struct { - ID string `json:"id"` + ID string `json:"id"` + Status struct { + Value string `json:"value"` + } `json:"status"` } `json:"data"` } if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&m); err != nil { @@ -289,6 +271,8 @@ func (t *Table) check(ctx context.Context, url, prefix string) ([]string, pollRe } defer resp.Body.Close() + // A model is loaded when it has no status or status.value == "loaded"; any other + // value ("unloaded", "loading", …) is not loaded and must never be asked /props?model=. loaded := make([]string, 0, len(m.Data)) seen := make(map[string]struct{}, len(m.Data)) for _, d := range m.Data { @@ -298,6 +282,9 @@ func (t *Table) check(ctx context.Context, url, prefix string) ([]string, pollRe if _, ok := seen[d.ID]; ok { continue } + if d.Status.Value != "" && d.Status.Value != "loaded" { + continue + } seen[d.ID] = struct{}{} loaded = append(loaded, d.ID) } @@ -314,11 +301,25 @@ func (t *Table) fail(ctx context.Context, prefix string, err error) pollResult { } func copyStatus(s Status) Status { - if s.Loaded == nil { - return s - } out := s - out.Loaded = make([]string, len(s.Loaded)) - copy(out.Loaded, s.Loaded) + if s.Loaded != nil { + out.Loaded = make([]string, len(s.Loaded)) + copy(out.Loaded, s.Loaded) + } + if s.Models != nil { + out.Models = make(map[string]ModelCtx, len(s.Models)) + for k, v := range s.Models { + out.Models[k] = v + } + } return out } + +func contains(list []string, v string) bool { + for _, s := range list { + if s == v { + return true + } + } + return false +} diff --git a/internal/health/props.go b/internal/health/props.go new file mode 100644 index 0000000..eb702ec --- /dev/null +++ b/internal/health/props.go @@ -0,0 +1,140 @@ +package health + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/url" +) + +// modelProps is the part of a /props body that carries a model's context size. +type modelProps struct { + Generation struct { + NCtx *int `json:"n_ctx"` + } `json:"default_generation_settings"` + TotalSlots *int `json:"total_slots"` +} + +// props reads /props best-effort. A request that fails because ctx is +// done yields a cancelled result so the caller records nothing; any other +// outcome (status, body, or missing fields) leaves context unknown without +// failing the poll. When the body says "role":"router" the host carries no +// context of its own, so NCtx and Slots stay 0 whatever else it reports. +func (t *Table) props(ctx context.Context, base string) (int, int, pollResult) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/props", nil) + if err != nil { + if ctx.Err() != nil { + return 0, 0, pollResult{cancelled: true} + } + return 0, 0, pollResult{} + } + resp, err := t.client.Do(req) + if err != nil { + if ctx.Err() != nil { + return 0, 0, pollResult{cancelled: true} + } + return 0, 0, pollResult{} + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return 0, 0, pollResult{} + } + var p struct { + Role string `json:"role"` + Generation struct { + NCtx *int `json:"n_ctx"` + } `json:"default_generation_settings"` + TotalSlots *int `json:"total_slots"` + } + if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&p); err != nil { + return 0, 0, pollResult{} + } + if p.Role == "router" { + return 0, 0, pollResult{} + } + nctx, slots := 0, 0 + if p.Generation.NCtx != nil { + nctx = *p.Generation.NCtx + } + if p.TotalSlots != nil { + slots = *p.TotalSlots + } + if nctx < 0 { + nctx = 0 + } + if slots < 0 { + slots = 0 + } + return nctx, slots, pollResult{} +} + +// propsModels asks /props?model= for each loaded model and returns a fresh, +// non-nil map of what each answered. A failed or malformed answer leaves that +// id absent; a request that fails because ctx is done cancels the whole poll. +// No model is ever asked that is not in loaded. +func (t *Table) propsModels(ctx context.Context, base string, loaded []string) (map[string]ModelCtx, pollResult) { + models := make(map[string]ModelCtx, len(loaded)) + for _, id := range loaded { + mc, ok, r := t.propsModel(ctx, base, id) + if r.cancelled || r.reason != "" { + return nil, r + } + if !ok { + continue + } + models[id] = mc + } + return models, pollResult{} +} + +// propsModel reads /props?model= best-effort. ok is false when the +// answer failed or was malformed (the model stays unknown); a request that +// fails because ctx is done yields a cancelled result so the caller records +// nothing. +func (t *Table) propsModel(ctx context.Context, base, id string) (ModelCtx, bool, pollResult) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/props?model="+url.QueryEscape(id), nil) + if err != nil { + if ctx.Err() != nil { + return ModelCtx{}, false, pollResult{cancelled: true} + } + return ModelCtx{}, false, pollResult{} + } + resp, err := t.client.Do(req) + if err != nil { + if ctx.Err() != nil { + return ModelCtx{}, false, pollResult{cancelled: true} + } + return ModelCtx{}, false, pollResult{} + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return ModelCtx{}, false, pollResult{} + } + var p modelProps + if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&p); err != nil { + return ModelCtx{}, false, pollResult{} + } + mc, r := clampCtx(p.Generation.NCtx, p.TotalSlots) + return mc, true, r +} + +// clampCtx turns the two optional fields into a non-negative context and slots. +func clampCtx(nctx, slots *int) (ModelCtx, pollResult) { + c := ModelCtx{} + if nctx != nil { + c.NCtx = *nctx + } + if slots != nil { + c.Slots = *slots + } + if c.NCtx < 0 { + c.NCtx = 0 + } + if c.Slots < 0 { + c.Slots = 0 + } + return c, pollResult{} +} diff --git a/internal/health/props_router_test.go b/internal/health/props_router_test.go new file mode 100644 index 0000000..f74bd52 --- /dev/null +++ b/internal/health/props_router_test.go @@ -0,0 +1,177 @@ +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, `not json`) + } + }) + 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) + } +} diff --git a/internal/proxy/ctxguard.go b/internal/proxy/ctxguard.go index f364c15..696de8b 100644 --- a/internal/proxy/ctxguard.go +++ b/internal/proxy/ctxguard.go @@ -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. 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 } } diff --git a/internal/proxy/ctxguard_router_test.go b/internal/proxy/ctxguard_router_test.go new file mode 100644 index 0000000..7c4ffb1 --- /dev/null +++ b/internal/proxy/ctxguard_router_test.go @@ -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"]) + } +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 60feb29..654cca1 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -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:>" when the context was relocated + CtxHeader = "X-Crossbar-Ctx" // "moved:new" when the context was relocated ) // errBodyTooLarge is returned when a request body exceeds MaxBody during the model peek.