diff --git a/CHANGELOG.md b/CHANGELOG.md index 9479b56..43083c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,14 @@ has a matching GitHub release whose body comes from its section below. ### Added +- Measured switch costs on the gateway: it times every start it runs, from the lease + grant to the engine answering, and keeps the last 20 times per engine and tier in + `~/.local/state/berth/swapcost.json`. `GET /v1/berth/cost?engine=NAME` answers one + estimate per tier from those times (an estimate with `known:false` has no samples + behind it), and `POST /v1/berth/plan` prices a declared chain: the per-step estimate, a + lower-bound total with its unmeasured count, the engine to prefetch at each step, and + the engines the chain comes back to. A failed start records nothing, and every route + takes the node's key like the rest of the gateway. - `prefetch` on an `[[engine]]` entry: absolute file and directory paths the gateway asks the kernel to read into the page cache the moment it decides to start that engine, so the disk reads overlap the drain of the engine it replaces (when the machine has the RAM to keep them). A directory means its own files, diff --git a/README.md b/README.md index abe7428..0e7e3fb 100644 --- a/README.md +++ b/README.md @@ -691,6 +691,62 @@ needs the key. The gateway listens on loopback only, refuses a Host that is not a loopback name (DNS rebinding) and refuses any request that carries a browser `Origin`, so a web page cannot start an engine. ComfyUI already hands off at a job boundary under a lease (`--drain-comfyui`); a drain first unloads its models via `POST /free`, and `--keep-comfyui-up` leaves the server on its port for OpenWebUI instead of stopping it. +#### What a switch costs + +The gateway runs a start and sees both ends of it: the grant its run printed, which is +when the card showed free, and the engine's first answer. It times that pair for every +start it runs and keeps the last 20 times per engine and tier in a book +(`~/.local/state/berth/swapcost.json`), so it can price the next switch without guessing. +The wait for a free card belongs to the queue, not to the engine, so it is not part of the +figure. A failed start records nothing, and so does an engine the gateway does not start itself +(llama-swap's models load under llama-swap, so the gateway has nothing to time for them). + +`GET /v1/berth/cost?engine=NAME` answers one estimate per tier, hot first. On the reference machine, +after five gateway starts of Strata from disk: + +```bash +curl -s "http://127.0.0.1:9799/v1/berth/cost?engine=strata" -H "Authorization: Bearer $(berth key)" +``` + +```json +{"engine":"strata","estimates":[ + {"tier":"hot","took_ms":0,"samples":0,"known":true}, + {"tier":"warm","took_ms":0,"samples":0,"known":false}, + {"tier":"cool","took_ms":0,"samples":0,"known":false}, + {"tier":"cold","took_ms":72400,"samples":5,"known":true}]} +``` + +Hot, the engine already holding the card, is the free switch. Warm is a parked engine, +its weights in host RAM. Cool and cold split at half of the engine's files still held by +the kernel page cache. An estimate with `known:false` has no samples behind it, and its +`took_ms: 0` is not a prediction: a figure the gateway gives is one it measured. + +`POST /v1/berth/plan` prices a chain an orchestrator declares ahead of time, one step per +switch, each naming the engine and the tier to find it in: + +```bash +curl -s http://127.0.0.1:9799/v1/berth/plan -H "Authorization: Bearer $(berth key)" \ + -d '{"chain":[{"engine":"strata","tier":"cold"},{"engine":"comfyui","tier":"cold"},{"engine":"strata","tier":"cold"}]}' +``` + +```json +{"steps":[ + {"engine":"strata","tier":"cold","took_ms":72400,"samples":5,"known":true}, + {"engine":"comfyui","tier":"cold","took_ms":0,"samples":0,"known":false}, + {"engine":"strata","tier":"cold","took_ms":72400,"samples":5,"known":true}], + "total_ms":144800, + "unmeasured":1, + "prefetch":["comfyui","strata",""], + "keep_parked":["strata"]} +``` + +`total_ms` sums the known steps only: with `unmeasured` above 0 it is a lower bound, read +as "at least 145 s, 1 step unmeasured", never a promise. `prefetch` is per step, the engine +to warm while that step holds the card, and `keep_parked` lists the engines the chain comes +back to, which are worth parking in RAM instead of stopping and reloading. An unknown tier +or an engine `berth.toml` does not name is a 400 naming the step; a missing or wrong key +and any browser `Origin` are refused, as on every other route. + #### Open a stopped engine's web UI A bookmark to ComfyUI or Strata's chat page is a dead link while that engine is stopped. Give the engine a @@ -1104,7 +1160,7 @@ In the next v0.3.x releases. Each item has its building blocks merged or a spec whose parent exits within one 2 s tick). - KV snapshots wired into the drain and start paths, so a sequence of agents skips the prefill after each switch. -- Measured switch costs on the live daemon and the gateway (`/v1/berth/cost`, `/v1/berth/plan`). +- Measured switch costs on the live daemon. - Gateway fallback routing, and models started on demand for a standalone llama-server. - Parked resume numbers in [Swap numbers](#swap-numbers), which take the vLLM and SGLang park adapters out of experimental, and a measurement of NVIDIA's `cuda-checkpoint` for unmodified diff --git a/internal/cli/gateway.go b/internal/cli/gateway.go index c3cef64..f302451 100644 --- a/internal/cli/gateway.go +++ b/internal/cli/gateway.go @@ -77,7 +77,8 @@ func runGateway(c *Context) int { return c.fail(Systemf("api key: %v", err)) } g := &gateway.Gateway{Socket: *sock, Engines: cfg.Engine, Ready: *ready, Key: key, Log: log, HeadroomFloor: *headroomMiB << 20, - AllowHosts: allow, StateFile: defaultStateFile("gateway-engines.json"), ConfigPath: *cfgPath} + AllowHosts: allow, StateFile: defaultStateFile("gateway-engines.json"), + CostFile: defaultStateFile("swapcost.json"), ConfigPath: *cfgPath} ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() srv := &http.Server{Addr: addr, Handler: g.Handler(), ReadHeaderTimeout: 10 * time.Second} diff --git a/internal/gateway/cost.go b/internal/gateway/cost.go new file mode 100644 index 0000000..4082c2f --- /dev/null +++ b/internal/gateway/cost.go @@ -0,0 +1,248 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +// The gateway runs a start, so it is the one witness that sees both ends +// of it: the grant its run printed, which is when the card showed free, +// and the engine's first answer. That is the pair swapcost measures, and +// it is a pair, not a request's wait: the time behind the grant belongs +// to whoever held the card, not to the engine loading. + +package gateway + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "time" + + "github.com/T-Crypt/berth/internal/config" + "github.com/T-Crypt/berth/internal/pagecache" + "github.com/T-Crypt/berth/internal/swapcost" +) + +// costBook returns the gateway's book of measured switch times, read from +// CostFile once per gateway; an empty CostFile keeps it in memory only. +// A file that is not a book is one WARN and an empty book: a figure the +// gateway gives is one it measured after that, not a guess carried over. +func (g *Gateway) costBook() *swapcost.Book { + g.costOnce.Do(func() { + b, err := swapcost.NewBook(g.CostFile) + if err != nil { + g.log().Warn("gateway could not read its switch-cost book; starting empty", "path", g.CostFile, "err", err) + } + g.costs = b + }) + return g.costs +} + +// position is where an engine sat before a start: whether its weights +// were parked in host RAM, and what share of its files the kernel page +// cache held. The tier a sample is recorded under is read from it. +type position struct { + parked bool + cached float64 +} + +// readPosition takes the engine's position now, at the moment its start +// is decided and before its prefetch reads begin: the node's lease table +// says whether the engine is parked, and the page cache says how much of +// its files it holds. No prefetch list, or a file whose pages cannot be +// checked, leaves the share 0: a part of the answer would be a wrong +// fraction, and cold is the honest reading of not knowing. +func (g *Gateway) readPosition(e *config.Engine) *position { + pos := &position{parked: parkedIn(g.nodeStatus(), e.Name)} + files, _ := prefetchFiles(e.Prefetch) + if len(files) == 0 { + return pos + } + var cached, size int64 + for _, f := range files { + c, sz, err := pagecache.Resident(f) + if err != nil { + return pos + } + cached += c + size += sz + } + if size > 0 { + pos.cached = float64(cached) / float64(size) + } + return pos +} + +// stampGrant notes the moment the run printed its grant, the point the +// card was free. Engines this gateway adopted carry no position: their +// starts were run by a previous gateway, and grants the follower sees on +// them complete the engine record but are never measured. +func (g *Gateway) stampGrant(s *started) { + g.mu.Lock() + defer g.mu.Unlock() + if s.pos != nil && s.grantAt.IsZero() { + s.grantAt = time.Now() + } +} + +// markAnswered closes a start whose engine just answered: the failure +// cooldown clears, and when this gateway ran the start and saw its grant, +// the switch from that grant to this answer goes into the book once, +// whichever waiting request noticed the answer first. +func (g *Gateway) markAnswered(e *config.Engine, s *started) { + var sample *swapcost.Sample + g.mu.Lock() + s.answered = true + if !s.recorded && s.pos != nil && !s.grantAt.IsZero() { + s.recorded = true + now := time.Now() + sample = &swapcost.Sample{ + Engine: e.Name, + Tier: swapcost.TierFor(s.pos.parked, false, s.pos.cached), + Took: now.Sub(s.grantAt), + At: now, + } + } + delete(g.failed, e.Name) + g.mu.Unlock() + if sample != nil { + g.saveSwitch(*sample) + } +} + +// saveSwitch adds one measured switch to the book and writes the book to +// its file. A failed write is one WARN line: the gateway keeps serving, +// and the next recorded switch carries this one too. +func (g *Gateway) saveSwitch(s swapcost.Sample) { + book := g.costBook() + g.costMu.Lock() + defer g.costMu.Unlock() + book.Record(s) + if err := book.Save(); err != nil { + g.log().Warn("gateway could not save its switch-cost book", "path", g.CostFile, "err", err) + } +} + +// tierOrder is the order /v1/berth/cost answers tiers in: the book's +// four positions, nearest the card first. +var tierOrder = []swapcost.Tier{swapcost.Hot, swapcost.Warm, swapcost.Cool, swapcost.Cold} + +// costReply is one engine's estimates, one per tier, times in ms. An +// estimate with known false has no samples: its took_ms is 0 and is not +// a prediction. +type costReply struct { + Engine string `json:"engine"` + Estimates []estimateReply `json:"estimates"` +} + +type estimateReply struct { + Tier string `json:"tier"` + TookMs int64 `json:"took_ms"` + Samples int `json:"samples"` + Known bool `json:"known"` +} + +// cost answers what the gateway measured about starting one engine from +// each tier. An engine berth.toml does not name is a 404: the gateway +// starts what it is configured to start and has measured nothing else. +func (g *Gateway) cost(w http.ResponseWriter, r *http.Request) { + name := r.URL.Query().Get("engine") + if name == "" { + apiError(w, http.StatusBadRequest, "name the engine to price: /v1/berth/cost?engine=NAME") + return + } + if g.engineByName(name) == nil { + apiError(w, http.StatusNotFound, fmt.Sprintf("no engine %q in the config", name)) + return + } + book := g.costBook() + out := costReply{Engine: name, Estimates: make([]estimateReply, 0, len(tierOrder))} + g.costMu.Lock() + for _, tier := range tierOrder { + es := book.Estimate(name, tier) + out.Estimates = append(out.Estimates, estimateReply{ + Tier: string(es.Tier), TookMs: es.Took.Milliseconds(), Samples: es.Samples, Known: es.Known}) + } + g.costMu.Unlock() + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(out) +} + +// planReply is swapcost.Plan as the gateway serves it: one step per +// chain entry, and total_ms the sum of the known steps only. With +// unmeasured above 0 the total is a lower bound, "at least this, N +// steps unmeasured", never a promise. prefetch is per step: the engine +// to warm while that step holds the card, empty where the next step is +// the same engine or the chain ends; keep_parked lists the engines the +// chain comes back to, so the caller parks them instead of paying a +// reload. +type planReply struct { + Steps []stepReply `json:"steps"` + TotalMs int64 `json:"total_ms"` + Unmeasured int `json:"unmeasured"` + Prefetch []string `json:"prefetch"` + KeepParked []string `json:"keep_parked"` +} + +type stepReply struct { + Engine string `json:"engine"` + Tier string `json:"tier"` + TookMs int64 `json:"took_ms"` + Samples int `json:"samples"` + Known bool `json:"known"` +} + +// plan works a declared chain out against the book: what each switch is +// likely to take, what the total is a lower bound on, which engine to +// prefetch at each step, and which engines to keep parked. +func (g *Gateway) plan(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(io.LimitReader(r.Body, 1<<20)) + if err != nil { + apiError(w, http.StatusBadRequest, "read body: "+err.Error()) + return + } + var req struct { + Chain []struct { + Engine string `json:"engine"` + Tier string `json:"tier"` + } `json:"chain"` + } + if json.Unmarshal(body, &req) != nil { + apiError(w, http.StatusBadRequest, `the body must be JSON like {"chain":[{"engine":"NAME","tier":"cold"}]}`) + return + } + if len(req.Chain) == 0 { + apiError(w, http.StatusBadRequest, "the chain is empty: name at least one step") + return + } + steps := make([]swapcost.Step, 0, len(req.Chain)) + for i, s := range req.Chain { + switch swapcost.Tier(s.Tier) { + case swapcost.Hot, swapcost.Warm, swapcost.Cool, swapcost.Cold: + default: + apiError(w, http.StatusBadRequest, fmt.Sprintf("chain step %d names tier %q: it must be hot, warm, cool or cold", i+1, s.Tier)) + return + } + if g.engineByName(s.Engine) == nil { + apiError(w, http.StatusBadRequest, fmt.Sprintf("chain step %d names engine %q, which is not in the config", i+1, s.Engine)) + return + } + steps = append(steps, swapcost.Step{Engine: s.Engine, Tier: swapcost.Tier(s.Tier)}) + } + book := g.costBook() + g.costMu.Lock() + p := book.Plan(steps, time.Now()) + g.costMu.Unlock() + out := planReply{ + Steps: make([]stepReply, 0, len(p.Steps)), + TotalMs: p.Total.Milliseconds(), + Unmeasured: p.Unmeasured, + Prefetch: p.Prefetch, + KeepParked: p.KeepParked, + } + for _, sc := range p.Steps { + out.Steps = append(out.Steps, stepReply{ + Engine: sc.Step.Engine, Tier: string(sc.Step.Tier), + TookMs: sc.Estimate.Took.Milliseconds(), Samples: sc.Estimate.Samples, Known: sc.Estimate.Known}) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(out) +} diff --git a/internal/gateway/cost_test.go b/internal/gateway/cost_test.go new file mode 100644 index 0000000..44a3401 --- /dev/null +++ b/internal/gateway/cost_test.go @@ -0,0 +1,317 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +//go:build unix + +package gateway + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/T-Crypt/berth/internal/config" + "github.com/T-Crypt/berth/internal/runner" + "github.com/T-Crypt/berth/internal/swapcost" +) + +// costTestEngines is the config the route tests run against; the book is +// seeded per test, not here. +func costTestEngines() []config.Engine { + return []config.Engine{ + {Name: "strata", Type: "llamacpp", URL: "http://127.0.0.1:8080", Serves: []string{"router-strata"}}, + {Name: "comfyui", Type: "comfyui", URL: "http://127.0.0.1:8188"}, + } +} + +func costTestGateway(t *testing.T) (*Gateway, *httptest.Server) { + t.Helper() + g := &Gateway{Key: testKey, Log: discard(), Engines: costTestEngines()} + gw := httptest.NewServer(g.Handler()) + t.Cleanup(gw.Close) + return g, gw +} + +type estimateRow struct { + Tier string `json:"tier"` + TookMs int64 `json:"took_ms"` + Samples int `json:"samples"` + Known bool `json:"known"` +} + +type costBody struct { + Engine string `json:"engine"` + Estimates []estimateRow `json:"estimates"` +} + +// The cost route answers one estimate per tier: a pair the book has +// measured comes back with its median, one it has not comes back +// known:false rather than a guess, and hot is always the free switch. +func TestCostRoute(t *testing.T) { + g, gw := costTestGateway(t) + g.costBook().Record(swapcost.Sample{Engine: "strata", Tier: swapcost.Warm, Took: 8 * time.Second, At: time.Now()}) + g.costBook().Record(swapcost.Sample{Engine: "strata", Tier: swapcost.Cool, Took: 45 * time.Second, At: time.Now()}) + + cases := []struct { + name string + url string + want int + }{ + {"measured engine", "/v1/berth/cost?engine=strata", http.StatusOK}, + {"unmeasured engine", "/v1/berth/cost?engine=comfyui", http.StatusOK}, + {"no engine", "/v1/berth/cost", http.StatusBadRequest}, + {"engine not in config", "/v1/berth/cost?engine=nope", http.StatusNotFound}, + } + for _, c := range cases { + resp := call(t, http.MethodGet, gw.URL+c.url, "") + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != c.want { + t.Fatalf("%s: status %d: %s", c.name, resp.StatusCode, body) + } + } + + resp := call(t, http.MethodGet, gw.URL+"/v1/berth/cost?engine=strata", "") + var got costBody + if err := json.NewDecoder(resp.Body).Decode(&got); err != nil { + t.Fatal(err) + } + resp.Body.Close() + want := costBody{Engine: "strata", Estimates: []estimateRow{ + {Tier: "hot", TookMs: 0, Samples: 0, Known: true}, + {Tier: "warm", TookMs: 8000, Samples: 1, Known: true}, + {Tier: "cool", TookMs: 45000, Samples: 1, Known: true}, + {Tier: "cold", TookMs: 0, Samples: 0, Known: false}, + }} + if len(got.Estimates) != len(want.Estimates) || got.Engine != want.Engine { + t.Fatalf("cost = %+v, want %+v", got, want) + } + for i := range want.Estimates { + if got.Estimates[i] != want.Estimates[i] { + t.Fatalf("estimate %d = %+v, want %+v", i, got.Estimates[i], want.Estimates[i]) + } + } + + resp = call(t, http.MethodGet, gw.URL+"/v1/berth/cost?engine=comfyui", "") + if err := json.NewDecoder(resp.Body).Decode(&got); err != nil { + t.Fatal(err) + } + resp.Body.Close() + for _, e := range got.Estimates[1:] { + if e.Known || e.Samples != 0 { + t.Fatalf("unmeasured %s = %+v, want known:false", e.Tier, e) + } + } +} + +type planStep struct { + Engine string `json:"engine"` + Tier string `json:"tier"` + TookMs int64 `json:"took_ms"` + Samples int `json:"samples"` + Known bool `json:"known"` +} + +type planBody struct { + Steps []planStep `json:"steps"` + TotalMs int64 `json:"total_ms"` + Unmeasured int `json:"unmeasured"` + Prefetch []string `json:"prefetch"` + KeepParked []string `json:"keep_parked"` +} + +// The plan route prices a declared chain from the book: known steps sum +// into the total, unknown ones are counted, and the steps the chain +// comes back to are listed for keeping parked. +func TestPlanRoute(t *testing.T) { + g, gw := costTestGateway(t) + g.costBook().Record(swapcost.Sample{Engine: "strata", Tier: swapcost.Cool, Took: 45 * time.Second, At: time.Now()}) + + resp := call(t, http.MethodPost, gw.URL+"/v1/berth/plan", + `{"chain":[{"engine":"comfyui","tier":"cold"},{"engine":"strata","tier":"cool"},{"engine":"comfyui","tier":"cold"}]}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status %d: %s", resp.StatusCode, body) + } + var got planBody + if err := json.Unmarshal(body, &got); err != nil { + t.Fatal(err) + } + want := planBody{ + Steps: []planStep{ + {Engine: "comfyui", Tier: "cold", Known: false}, + {Engine: "strata", Tier: "cool", TookMs: 45000, Samples: 1, Known: true}, + {Engine: "comfyui", Tier: "cold", Known: false}, + }, + TotalMs: 45000, + Unmeasured: 2, + Prefetch: []string{"strata", "comfyui", ""}, + KeepParked: []string{"comfyui"}, + } + if got.TotalMs != want.TotalMs || got.Unmeasured != want.Unmeasured || + strings.Join(got.Prefetch, ",") != strings.Join(want.Prefetch, ",") || + strings.Join(got.KeepParked, ",") != strings.Join(want.KeepParked, ",") || + len(got.Steps) != len(want.Steps) { + t.Fatalf("plan = %+v, want %+v", got, want) + } + for i := range want.Steps { + if got.Steps[i] != want.Steps[i] { + t.Fatalf("step %d = %+v, want %+v", i, got.Steps[i], want.Steps[i]) + } + } + + // Every bad body names what is wrong. + for _, c := range []struct{ body, says string }{ + {`{`, "JSON"}, + {`{}`, "chain is empty"}, + {`{"chain":[]}`, "chain is empty"}, + {`{"chain":[{"engine":"strata","tier":"frozen"}]}`, `tier \"frozen\"`}, + {`{"chain":[{"engine":"strata","tier":"cold"},{"engine":"nope","tier":"hot"}]}`, `engine \"nope\"`}, + } { + resp := call(t, http.MethodPost, gw.URL+"/v1/berth/plan", c.body) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest || !strings.Contains(string(body), c.says) { + t.Errorf("body %s: status %d: %s", c.body, resp.StatusCode, body) + } + } +} + +// The cost and plan routes sit behind the same key and Origin checks as +// the rest of /v1. +func TestCostPlanRoutesGuarded(t *testing.T) { + _, gw := costTestGateway(t) + routes := []struct{ method, path, body string }{ + {http.MethodGet, "/v1/berth/cost?engine=strata", ""}, + {http.MethodPost, "/v1/berth/plan", `{"chain":[{"engine":"strata","tier":"cold"}]}`}, + } + for _, r := range routes { + for _, c := range []struct { + header map[string]string + want int + }{ + {nil, http.StatusUnauthorized}, + {map[string]string{"Authorization": "Bearer nope"}, http.StatusUnauthorized}, + {map[string]string{"Authorization": "Bearer " + testKey, "Origin": "https://evil.example"}, http.StatusForbidden}, + } { + req, _ := http.NewRequest(r.method, gw.URL+r.path, strings.NewReader(r.body)) + for k, v := range c.header { + req.Header.Set(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != c.want { + t.Errorf("%s %s %v: status %d, want %d", r.method, r.path, c.header, resp.StatusCode, c.want) + } + } + } +} + +// waitFor fails unless out holds want before the deadline. +func waitFor(t *testing.T, out *syncBuf, want, fail string) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for !strings.Contains(out.String(), want) { + if time.Now().After(deadline) { + t.Fatalf("%s: %s", fail, out.String()) + } + time.Sleep(10 * time.Millisecond) + } +} + +// A start the gateway ran goes into the book exactly once, from the +// grant to the engine answering. The queue the start waits behind the +// foreground holder is not the engine's switch cost: the sample and the +// whole request differ by that queue, so elapsed minus took must show +// more than a second the gateway waited before the card was free. +func TestStartRecordsOneSampleFromGrantToAnswer(t *testing.T) { + sock, _ := daemon(t) + holderOut := &syncBuf{} + holder := make(chan int, 1) + go func() { // foreground beats the gateway's start: the start waits behind, no drain + holder <- runner.Run(runner.Options{Socket: sock, Engine: "ninfer", Priority: "foreground", + Argv: []string{"sleep", "2"}, Stderr: holderOut}) + }() + waitFor(t, holderOut, "granted", "holder never granted") + + addr := freeAddr(t) + t.Setenv("BERTH_GW_ENGINE", addr+",cost-model,3s") + costPath := filepath.Join(t.TempDir(), "swapcost.json") + g := &Gateway{Socket: sock, Ready: 15 * time.Second, Key: testKey, Log: discard(), CostFile: costPath, + Engines: []config.Engine{{Name: "cost-engine", Type: "llamacpp", URL: "http://" + addr, + Start: []string{os.Args[0]}, Serves: []string{"cost-model"}}}} + gw := httptest.NewServer(g.Handler()) + defer gw.Close() + + req, _ := http.NewRequest(http.MethodPost, gw.URL+"/v1/chat/completions", strings.NewReader(`{"model":"cost-model"}`)) + req.Header.Set("Authorization", "Bearer "+testKey) + start := time.Now() + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + elapsed := time.Since(start) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status %d: %s", resp.StatusCode, body) + } + <-holder + + // The book on disk: exactly one sample, on the cold tier (no + // prefetch list means the page cache places nothing), and a time + // that is the switch alone. + raw, err := os.ReadFile(costPath) + if err != nil { + t.Fatal(err) + } + var disk struct { + Samples map[string]map[string][]struct { + Engine string + Tier string + Took time.Duration + } + } + if err := json.Unmarshal(raw, &disk); err != nil { + t.Fatalf("book = %s", raw) + } + tiers := disk.Samples["cost-engine"] + total := 0 + for _, ring := range tiers { + total += len(ring) + } + if total != 1 || len(tiers["cold"]) != 1 { + t.Fatalf("samples = %s, want exactly one on the cold tier", raw) + } + took := tiers["cold"][0].Took + if took <= 0 { + t.Fatalf("took = %s, want a positive measured switch", took) + } + if gap := elapsed - took; gap < time.Second { + t.Fatalf("took = %s of the whole request (%s): the wait before the grant leaked in", took, elapsed) + } + + // A second request is proxied: no new start, no new sample. + resp = call(t, http.MethodPost, gw.URL+"/v1/chat/completions", `{"model":"cost-model"}`) + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("second call: status %d", resp.StatusCode) + } + book, err := swapcost.NewBook(costPath) + if err != nil { + t.Fatal(err) + } + if es := book.Estimate("cost-engine", swapcost.Cold); es.Samples != 1 || !es.Known { + t.Fatalf("estimate after the second call = %+v, want the one sample", es) + } +} diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index fa73128..ecc75d4 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -40,6 +40,7 @@ import ( "github.com/T-Crypt/berth/internal/proto" "github.com/T-Crypt/berth/internal/runner" "github.com/T-Crypt/berth/internal/server" + "github.com/T-Crypt/berth/internal/swapcost" ) type Gateway struct { @@ -67,6 +68,10 @@ type Gateway struct { // processes, so the next gateway adopts the ones still running // instead of starting them again. Empty disables the record. StateFile string + // CostFile is where the gateway's book of measured switch times + // lives, so the times survive a restart of the gateway. Empty keeps + // the book in memory only. + CostFile string // ConfigPath is the berth.toml the spawned runs read their entry's // park settings from (the park hook is not a flag). ConfigPath string @@ -86,6 +91,10 @@ type Gateway struct { logOnce sync.Once logDir string logErr error // one engineLogs() decision per gateway + + costOnce sync.Once + costs *swapcost.Book + costMu sync.Mutex // one book access at a time: Record and Save together, Estimate undivided } type started struct { @@ -95,6 +104,15 @@ type started struct { output *tailBuf // the tail of the engine's own stdout and stderr answered bool // the engine answered at least once (under g.mu) why *startError + + // The switch cost of a start this gateway ran: the engine's position + // when the start was decided, when its run printed its grant, and + // whether the sample went into the book. All three under g.mu; pos + // is set before the run spawns, so a grant the follower sees always + // has somewhere to land. + pos *position + grantAt time.Time + recorded bool } // failure is an engine whose start died before it answered. A start that @@ -187,6 +205,8 @@ func (g *Gateway) Handler() http.Handler { mux := http.NewServeMux() mux.HandleFunc("GET /v1/models", g.models) mux.HandleFunc("GET /v1/berth/wait", g.wait) + mux.HandleFunc("GET /v1/berth/cost", g.cost) + mux.HandleFunc("POST /v1/berth/plan", g.plan) for _, p := range []string{"/v1/chat/completions", "/v1/completions", "/v1/embeddings"} { mux.HandleFunc("POST "+p, g.forward) } @@ -728,6 +748,7 @@ func (g *Gateway) ensure(ctx context.Context, e *config.Engine) error { g.running = map[string]*started{} } s, ok := g.running[e.Name] + starting := false if !ok { if len(e.Start) == 0 { g.mu.Unlock() @@ -748,6 +769,19 @@ func (g *Gateway) ensure(ctx context.Context, e *config.Engine) error { } s = &started{exited: make(chan struct{}), stderr: &lockedBuf{}, output: &tailBuf{}} g.running[e.Name] = s + starting = true + } + g.mu.Unlock() + if starting { + // This gateway runs the start, so its switch gets measured. The + // engine's position is read now, before its reads begin: the + // prefetch below warms the page cache as soon as it can, and a + // share taken after that would call an engine the start warmed + // one that started warm. + pos := g.readPosition(e) + g.mu.Lock() + s.pos = pos + g.mu.Unlock() if len(e.Prefetch) > 0 { // The run takes the lease and the drain starts; the disk is // free while that happens, so the reads go first. @@ -755,15 +789,11 @@ func (g *Gateway) ensure(ctx context.Context, e *config.Engine) error { } go g.run(e, s, ready) } - g.mu.Unlock() deadline := time.Now().Add(ready) for { if probe(ctx, e.URL) { - g.mu.Lock() - s.answered = true - delete(g.failed, e.Name) - g.mu.Unlock() + g.markAnswered(e, s) return nil } select { diff --git a/internal/gateway/log.go b/internal/gateway/log.go index 62c5f4b..47f56a6 100644 --- a/internal/gateway/log.go +++ b/internal/gateway/log.go @@ -82,6 +82,7 @@ func (g *Gateway) followEngine(name, path string, from int64, stop <-chan struct see: func(line string) { if lease, ok := grantedLease(line); ok { g.setLease(name, lease) + g.stampGrant(s) } }} buf := make([]byte, 16<<10)