diff --git a/CHANGELOG.md b/CHANGELOG.md index d66904b..cc4d2d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,16 @@ has a matching GitHub release whose body comes from its section below. ### Added +- Gateway fallback routing: an `[[engine]]` may name a `fallback`, another engine in the file that + needs no card (a CPU llama-server, an Ollama on another box), and `fallback_after` (10 s by + default). A request that has waited that long for its own engine is proxied to the fallback when + the fallback answers its probe at that moment, with the body's model renamed when `fallback_model` + is set. A non-streaming answer carries `X-Berth-Fallback: `; a streaming one, whose SSE + preamble was already sent, gets the fallback's stream after one `: berth: fallback ` + comment. The gateway never starts, leases or drains the fallback, and the primary's start is never + cancelled, so later requests land on it. The request reaches the fallback without its + `Authorization` header, so the node's API key never leaves the gateway. One hop: `berth config check` + refuses chains. - Cgroup attribution closes the one known stowaway window (chaos scenario 3: a child whose parent exits within one 2 s tick). On Linux with a systemd user manager, `berth run` execs the engine in its own transient cgroup scope (`systemd-run --user diff --git a/README.md b/README.md index 56aadc0..47a7421 100644 --- a/README.md +++ b/README.md @@ -630,6 +630,39 @@ than `--ready` (10 minutes); if the holder has not left by then, the request fai dropped, so it cannot take the card hours later. Asking llama-swap for one of its models later drains the gateway's engine the same way, so either side can take the card back. +**A slow engine can name a fallback.** An interactive client should not sit through a 90 s cold start +when an engine that needs no card can answer now. Give an engine a `fallback` (another engine in the +same file, such as a CPU llama-server) and a request that has waited `fallback_after` (10 s by +default) goes there instead, when the fallback answers its probe at that moment. The fallback is +never started, leased or drained by the gateway: one that is not answering means no fallback, and the +request keeps waiting for its own engine. The primary's start is never cancelled either way, so +requests that come later land on it. `fallback_model` names the model the fallback lists the request +under when the two disagree, a non-streaming answer carries `X-Berth-Fallback: `, and a +streaming answer, whose headers went out with the keep-alive preamble, carries one +`: berth: fallback ` comment before its data. The fallback gets the request without its +`Authorization` header: the node's API key stays with the gateway. One hop, per engine, opt-in: `berth config check` +refuses chains and a fallback that is not an engine in the file. + +```toml +[[engine]] +name = "strata" +type = "llamacpp" +url = "http://127.0.0.1:8080" +start = ["/path/to/strata/.venv/bin/python", "serve/server.py", "--engine", "strata", "--port", "8080"] +serves = ["qwen3.8-flash-next-iq2_xs"] +priority = "interactive" +fallback = "strata-cpu" +fallback_after = "10s" +fallback_model = "qwen3.8-cpu" # what the CPU server lists the model as + +[[engine]] +name = "strata-cpu" +type = "llamacpp" +url = "http://127.0.0.1:8010" +start = ["llama-server", "-m", "/models/qwen3.8-cpu.gguf", "--port", "8010"] # no GPU +serves = ["qwen3.8-cpu"] +``` + **A start reads the engine's files into the page cache while the holder drains.** Give an engine a `prefetch` list of absolute files and directories, and the gateway asks the kernel to read them the moment it decides to start that engine, before its `berth run` waits for the lease. The kernel reads diff --git a/internal/config/config.go b/internal/config/config.go index 67e00a8..66e88ad 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -101,6 +101,19 @@ type Engine struct { // gateway's wake page (/gui/) starts the engine on a visit and // follows it there once it answers. GUI string `toml:"gui,omitempty" json:"gui,omitempty" yaml:"gui,omitempty"` + // Fallback is the name of another engine in this file: while a request + // waits for this engine to come up, past FallbackAfter, the gateway + // sends it to the fallback instead, if the fallback answers its probe + // at that moment. The gateway never starts the fallback, and the start + // of this engine is never cancelled: later requests land on it. No + // chains: a fallback has no fallback of its own. + Fallback string `toml:"fallback,omitempty" json:"fallback,omitempty" yaml:"fallback,omitempty"` + // FallbackAfter is how long a request waits for this engine before it + // may go to the fallback. 0 = 10 s when Fallback is set. + FallbackAfter Duration `toml:"fallback_after,omitzero" json:"fallback_after,omitempty" yaml:"fallback_after,omitempty"` + // FallbackModel is the model name sent to the fallback; unset keeps + // the request's model name. + FallbackModel string `toml:"fallback_model,omitempty" json:"fallback_model,omitempty" yaml:"fallback_model,omitempty"` } // Types are the engine readers that exist. @@ -613,6 +626,32 @@ func (c *Config) Validate() error { } } } + // A fallback is resolved by name after every engine is known, so it is + // checked in a second pass. An error names both engines. + byName := map[string]int{} + for i, e := range c.Engine { + if e.Name != "" { + byName[e.Name] = i + } + } + for i, e := range c.Engine { + if e.Fallback == "" { + continue + } + where := fmt.Sprintf("engine[%d] %q", i, e.Name) + if e.FallbackAfter.Duration < 0 { + errs = append(errs, fmt.Errorf("%s: fallback_after %s must be >= 0", where, e.FallbackAfter)) + } + j, ok := byName[e.Fallback] + switch { + case !ok: + errs = append(errs, fmt.Errorf("%s: fallback %q is not an engine in this file", where, e.Fallback)) + case j == i: + errs = append(errs, fmt.Errorf("%s: fallback names the engine itself", where)) + case c.Engine[j].Fallback != "": + errs = append(errs, fmt.Errorf("%s: fallback %q has a fallback of its own; no chains", where, e.Fallback)) + } + } for _, p := range c.Engines.Ports { if p <= 0 || p > 65535 { errs = append(errs, fmt.Errorf("engines.discover_ports: %d is not a port", p)) diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 7ad645f..ace08ed 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -4,6 +4,7 @@ package config import ( + "fmt" "os" "path/filepath" "strconv" @@ -718,3 +719,63 @@ func TestLibraryFloorDefaults(t *testing.T) { t.Errorf("explicit 0 = %g, want the guard off", c.LibraryFloor()) } } + +func TestEngineFallbackParsedAndValidated(t *testing.T) { + const engines = ` +[[engine]] +name = "strata" +type = "llamacpp" +url = "http://127.0.0.1:8080" +serves = ["qwen3.8"] +%s +[[engine]] +name = "strata-cpu" +type = "llamacpp" +url = "http://127.0.0.1:8081" +serves = ["qwen3.8-cpu"] +` + c, err := Load(write(t, fmt.Sprintf(engines, ` +fallback = "strata-cpu" +fallback_after = "30s" +fallback_model = "qwen3.8-cpu" +`))) + if err != nil { + t.Fatalf("a valid fallback was refused: %v", err) + } + if c.Engine[0].Fallback != "strata-cpu" || c.Engine[0].FallbackAfter.Seconds() != 30 || c.Engine[0].FallbackModel != "qwen3.8-cpu" { + t.Fatalf("fallback keys not loaded: %+v", c.Engine[0]) + } + // Without fallback_after the gateway applies its own default. + c, err = Load(write(t, fmt.Sprintf(engines, `fallback = "strata-cpu"`))) + if err != nil || c.Engine[0].FallbackAfter.Duration != 0 { + t.Fatalf("no fallback_after: %+v err=%v", c.Engine[0], err) + } + + for _, tc := range []struct { + name string + keys string + wantErr string + }{ + {"unknown engine", `fallback = "nope"`, `engine[0] "strata": fallback "nope" is not an engine in this file`}, + {"itself", `fallback = "strata"`, `engine[0] "strata": fallback names the engine itself`}, + {"negative after", `fallback = "strata-cpu" +fallback_after = "-5s"`, `engine[0] "strata": fallback_after -5s must be >= 0`}, + } { + _, err := Load(write(t, fmt.Sprintf(engines, tc.keys))) + if err == nil || !strings.Contains(err.Error(), tc.wantErr) { + t.Errorf("%s: error lacks %q: %v", tc.name, tc.wantErr, err) + } + } + // A fallback that has a fallback of its own is refused, and the error + // names both engines. + _, err = Load(write(t, fmt.Sprintf(engines, `fallback = "strata-cpu"`)+` +[[engine]] +name = "third" +type = "llamacpp" +url = "http://127.0.0.1:8082" +fallback = "strata" +`)) + if err == nil || !strings.Contains(err.Error(), `engine[2] "third": fallback "strata" has a fallback of its own; no chains`) { + t.Errorf("a chain was accepted or names neither engine: %v", err) + } +} diff --git a/internal/gateway/fallback.go b/internal/gateway/fallback.go new file mode 100644 index 0000000..9fd0bfd --- /dev/null +++ b/internal/gateway/fallback.go @@ -0,0 +1,121 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +// Fallback routing: an engine may name another engine to take one request +// while it is still coming up, so an interactive client does not sit +// through a 90 s cold start when something that needs no card can answer +// now. Opt-in per engine, one hop, and the primary's start is never +// cancelled by it. +package gateway + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "time" + + "github.com/T-Crypt/berth/internal/config" +) + +// defaultFallbackAfter is how long a request waits for its engine before +// it may go to the fallback, when the engine names one but sets no time. +const defaultFallbackAfter = 10 * time.Second + +func (g *Gateway) fallbackAfter(e *config.Engine) time.Duration { + if e.FallbackAfter.Duration > 0 { + return e.FallbackAfter.Duration + } + return defaultFallbackAfter +} + +// waitPrimary runs ensure for the engine a request resolved to, and waits +// it out. When the engine names a fallback, the wait watches the clock: +// the engine still not answering after fallback_after, and the fallback +// answering its probe right then, this request goes to the fallback +// instead. The probe is all the gateway ever sends a fallback it did not +// see up: one that is not answering means no fallback, and the request +// keeps waiting for the primary as today. ensure is never cancelled +// either way, so the primary's start goes on and later requests land on +// it. +// +// The return says where the request goes: fb nil means the primary (err +// is ensure's outcome); fb set means proxy it there, waited long. A +// streaming response (w non-nil) gets its loading comments every +// keepAliveEvery while it waits, as before. +func (g *Gateway) waitPrimary(ctx context.Context, e *config.Engine, w http.ResponseWriter) (fb *config.Engine, waited time.Duration, err error) { + start := time.Now() + // A request that switches to the fallback returns without its engine's + // wait being over, and the handler's context goes with it: the wait for + // the primary rides on one that survives, so the start is never + // cancelled and answering late still counts for the next request. + ensureCtx := ctx + if e.Fallback != "" { + ensureCtx = context.WithoutCancel(ctx) + } + ch := make(chan error, 1) + go func() { ch <- g.ensure(ensureCtx, e) }() + var fbAt <-chan time.Time + if e.Fallback != "" { + at := time.NewTimer(g.fallbackAfter(e)) + defer at.Stop() + fbAt = at.C // a Timer fires once: the fallback gets one look + } + var tick <-chan time.Time + if w != nil { + every := time.NewTicker(keepAliveEvery) + defer every.Stop() + tick = every.C + } + for { + select { + case err := <-ch: + return nil, time.Since(start), err + case <-ctx.Done(): + return nil, time.Since(start), ctx.Err() + case <-tick: + fmt.Fprintf(w, ": berth: loading %s, %d s\n\n", e.Name, int64(time.Since(start).Seconds())) + _ = http.NewResponseController(w).Flush() + case <-fbAt: + f := g.engineByName(e.Fallback) + if f != nil && g.answers(ctx, f.URL) { + return f, time.Since(start), nil + } + } + } +} + +// tookFallback hands one request to the fallback: it journals the switch, +// rewrites the body's model when the fallback lists it under another +// name, and marks the response, a header while the headers are still open +// and an SSE comment once the preamble is already on the wire. It returns +// the base URL to proxy to and the body to send. +func (g *Gateway) tookFallback(w http.ResponseWriter, preamble bool, e, fb *config.Engine, waited time.Duration, body []byte) (string, []byte) { + if e.FallbackModel != "" { + body = rewriteModel(body, e.FallbackModel) + } + g.log().Info("gateway request fell back", "engine", e.Name, "fallback", fb.Name, "waited", waited.Round(time.Millisecond)) + if preamble { + fmt.Fprintf(w, ": berth: fallback %s\n\n", fb.Name) + _ = http.NewResponseController(w).Flush() + } else { + w.Header().Set("X-Berth-Fallback", fb.Name) + } + return fb.URL, body +} + +// rewriteModel replaces the request body's model name. Everything else in +// the body goes through byte for byte. +func rewriteModel(body []byte, model string) []byte { + var req map[string]json.RawMessage + if json.Unmarshal(body, &req) != nil { + return body + } + name, _ := json.Marshal(model) + req["model"] = name + out, err := json.Marshal(req) + if err != nil { + return body + } + return out +} diff --git a/internal/gateway/fallback_test.go b/internal/gateway/fallback_test.go new file mode 100644 index 0000000..6b984d3 --- /dev/null +++ b/internal/gateway/fallback_test.go @@ -0,0 +1,294 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +//go:build unix + +package gateway + +import ( + "bufio" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync" + "testing" + "time" + + "github.com/T-Crypt/berth/internal/broker" + "github.com/T-Crypt/berth/internal/config" +) + +// fallbackSrv is a fallback that is up at once: it answers the probe, records +// the model the proxied request carried and streams its own reply, so a test +// can tell whose answer arrived and under which name. +func fallbackSrv(t *testing.T) (*httptest.Server, func() string) { + t.Helper() + var mu sync.Mutex + model := "" + mux := http.NewServeMux() + mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"object":"list","data":[{"id":"cpu"}]}`) + }) + mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + var req struct { + Model string `json:"model"` + } + _ = json.Unmarshal(b, &req) + mu.Lock() + model = req.Model + mu.Unlock() + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"from the cpu fallback\"}}]}\n\n") + fmt.Fprint(w, "data: [DONE]\n\n") + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv, func() string { mu.Lock(); defer mu.Unlock(); return model } +} + +// fallbackGateway pairs the primary (the fake child engine, which answers +// only after load) with a fallback at fbURL the request may take after wait. +// It returns the gateway's test server, the primary's base URL and the +// daemon, whose leases say who was really started. +func fallbackGateway(t *testing.T, load, wait time.Duration, fbModel, fbURL string, log io.Writer) (*httptest.Server, string, *broker.Broker) { + t.Helper() + sock, b := daemon(t) + primary := freeAddr(t) + t.Setenv("BERTH_GW_ENGINE", primary+",big,"+load.String()) + g := &Gateway{Socket: sock, Ready: 10 * time.Second, Key: testKey, + Log: slog.New(slog.NewTextHandler(log, nil)), + Engines: []config.Engine{ + {Name: "strata", Type: "llamacpp", URL: "http://" + primary, Start: []string{os.Args[0]}, + Serves: []string{"big"}, Priority: "interactive", + Fallback: "strata-cpu", FallbackAfter: config.Duration{Duration: wait}, FallbackModel: fbModel}, + {Name: "strata-cpu", Type: "llamacpp", URL: fbURL, Serves: []string{"cpu"}}, + }} + gw := httptest.NewServer(g.Handler()) + t.Cleanup(gw.Close) + return gw, "http://" + primary, b +} + +// A request whose engine is slow to come up goes to a fallback that answers +// now: the answer carries X-Berth-Fallback, the model arrives under the name +// the fallback lists, the switch is journalled once, and the primary's own +// start is not cancelled by it. +func TestFallbackAnswersWhenThePrimaryIsSlow(t *testing.T) { + fb, saw := fallbackSrv(t) + log := &syncBuf{} + gw, primary, b := fallbackGateway(t, 3*time.Second, 300*time.Millisecond, "cpu", fb.URL, log) + + start := time.Now() + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"big"}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status %d: %s", resp.StatusCode, body) + } + if got := resp.Header.Get("X-Berth-Fallback"); got != "strata-cpu" { + t.Fatalf("X-Berth-Fallback %q, want strata-cpu", got) + } + if !strings.Contains(string(body), "from the cpu fallback") { + t.Fatalf("the fallback did not answer: %s", body) + } + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("the request waited %s; it should have gone to the fallback", elapsed.Round(time.Millisecond)) + } + if got := saw(); got != "cpu" { + t.Fatalf("the fallback saw model %q, want the rewritten \"cpu\"", got) + } + s := log.String() + if !strings.Contains(s, "gateway request fell back") || !strings.Contains(s, "engine=strata") || + !strings.Contains(s, "fallback=strata-cpu") || !strings.Contains(s, "waited=") { + t.Fatalf("the switch is not journalled: %s", s) + } + if n := strings.Count(s, "gateway request fell back"); n != 1 { + t.Fatalf("%d fallback lines for one request: %s", n, s) + } + + // The gateway never leases the fallback, and the primary's start went on + // answering for the requests that come after this one. + for _, l := range b.Status().Leases { + if l.Engine == "strata-cpu" { + t.Fatal("the gateway leased the fallback") + } + } + deadline := time.Now().Add(6 * time.Second) + for { + up, err := http.Get(primary + "/v1/models") + if err == nil { + ok := up.StatusCode == http.StatusOK + up.Body.Close() + if ok { + break + } + } + if time.Now().After(deadline) { + t.Fatal("the primary's start was cancelled by the fallback: it never answered") + } + time.Sleep(100 * time.Millisecond) + } +} + +// Without fallback_model the fallback sees the request's own model name. +func TestFallbackWithoutAModelKeepsTheRequest(t *testing.T) { + fb, saw := fallbackSrv(t) + gw, _, _ := fallbackGateway(t, 3*time.Second, 300*time.Millisecond, "", fb.URL, &syncBuf{}) + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"big"}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || resp.Header.Get("X-Berth-Fallback") != "strata-cpu" || + !strings.Contains(string(body), "from the cpu fallback") { + t.Fatalf("%d %s", resp.StatusCode, body) + } + if got := saw(); got != "big" { + t.Fatalf("the fallback saw model %q, want the request's own \"big\"", got) + } +} + +// The primary answering before fallback_after keeps the request: no header +// marks it, and the fallback is never asked. +func TestNoFallbackWhenThePrimaryAnswersFirst(t *testing.T) { + fb, saw := fallbackSrv(t) + start := time.Now() + gw, _, _ := fallbackGateway(t, 300*time.Millisecond, 5*time.Second, "", fb.URL, &syncBuf{}) + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"big"}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), "never collide") { + t.Fatalf("%d %s", resp.StatusCode, body) + } + if resp.Header.Get("X-Berth-Fallback") != "" { + t.Fatal("the primary's answer is marked as a fallback") + } + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("the request waited %s, longer than the engine's load", elapsed.Round(time.Millisecond)) + } + if got := saw(); got != "" { + t.Fatalf("the fallback was asked: model %q", got) + } +} + +// A fallback that is not answering means no fallback: the request keeps +// waiting for its engine and completes on it. +func TestNoFallbackWhenTheFallbackIsDown(t *testing.T) { + log := &syncBuf{} + gw, _, _ := fallbackGateway(t, time.Second, 200*time.Millisecond, "", "http://"+freeAddr(t), log) + start := time.Now() + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"big"}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), "never collide") { + t.Fatalf("%d %s", resp.StatusCode, body) + } + if resp.Header.Get("X-Berth-Fallback") != "" { + t.Fatal("the primary's answer is marked as a fallback") + } + if elapsed := time.Since(start); elapsed < 900*time.Millisecond { + t.Fatalf("the request answered after %s, before the engine was up", elapsed.Round(time.Millisecond)) + } + if strings.Contains(log.String(), "gateway request fell back") { + t.Fatalf("a down fallback was journalled as a switch: %s", log.String()) + } +} + +// A streaming request whose keep-alive preamble is already on the wire gets +// the fallback's stream through the same body-only path the primary would +// use: the loading comments stop, one comment marks the switch, and the +// fallback's own data follows. +func TestFallbackStreamAfterThePreamble(t *testing.T) { + shortKeepAlive(t, 100*time.Millisecond) + fb, saw := fallbackSrv(t) + gw, _, _ := fallbackGateway(t, 3*time.Second, 400*time.Millisecond, "cpu", fb.URL, &syncBuf{}) + resp := call(t, "POST", gw.URL+"/v1/chat/completions", + `{"model":"big","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + defer resp.Body.Close() + if ct := resp.Header.Get("Content-Type"); ct != "text/event-stream" { + t.Fatalf("Content-Type %q, want text/event-stream", ct) + } + if resp.Header.Get("X-Berth-Fallback") != "" { + t.Fatal("the headers went out with the preamble; a fallback header cannot follow") + } + var loading, data int + switched := false + sc := bufio.NewScanner(resp.Body) + for sc.Scan() { + line := sc.Text() + switch { + case strings.HasPrefix(line, ": berth: loading"): + loading++ + case strings.HasPrefix(line, ": berth: fallback"): + if line != ": berth: fallback strata-cpu" { + t.Fatalf("fallback comment = %q", line) + } + switched = true + case strings.HasPrefix(line, "data: "): + data++ + } + } + if loading == 0 { + t.Error("no loading comment arrived before the switch") + } + if !switched { + t.Error("no fallback comment arrived") + } + if data != 2 { + t.Errorf("data lines = %d, want the fallback's reply and DONE", data) + } + if got := saw(); got != "cpu" { + t.Errorf("the fallback saw model %q, want the rewritten \"cpu\"", got) + } +} + +// An engine that names a fallback but sets no time waits 10 s. +func TestFallbackAfterDefaults(t *testing.T) { + g := &Gateway{} + for _, c := range []struct { + e config.Engine + want time.Duration + }{ + {config.Engine{Fallback: "cpu"}, 10 * time.Second}, + {config.Engine{Fallback: "cpu", FallbackAfter: config.Duration{Duration: 30 * time.Second}}, 30 * time.Second}, + } { + if got := g.fallbackAfter(&c.e); got != c.want { + t.Errorf("fallbackAfter(%+v) = %s, want %s", c.e, got, c.want) + } + } +} + +// The node's API key authenticates the client to the gateway and goes no +// further: a fallback may run on another box, so the request it gets +// carries no Authorization header. +func TestFallbackDropsTheNodeKey(t *testing.T) { + var mu sync.Mutex + auth := "unset" + mux := http.NewServeMux() + mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"object":"list","data":[{"id":"cpu"}]}`) + }) + mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + auth = r.Header.Get("Authorization") + mu.Unlock() + fmt.Fprint(w, `{"choices":[]}`) + }) + fb := httptest.NewServer(mux) + t.Cleanup(fb.Close) + gw, _, _ := fallbackGateway(t, 3*time.Second, 300*time.Millisecond, "cpu", fb.URL, io.Discard) + + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"big"}`) + resp.Body.Close() + if resp.Header.Get("X-Berth-Fallback") != "strata-cpu" { + t.Fatalf("status %d: the request did not go to the fallback", resp.StatusCode) + } + mu.Lock() + defer mu.Unlock() + if auth != "" { + t.Fatalf("the fallback got Authorization %q, want none", auth) + } +} diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index ecc75d4..ea7f001 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -511,20 +511,35 @@ func (g *Gateway) forward(w http.ResponseWriter, r *http.Request) { return } preamble := false + var fb *config.Engine + var waited time.Duration if !t.live { if req.Stream { preamble = true ssePreamble(w) - if err := g.keepAlive(w, r, t.engine); err != nil { + var err error + fb, waited, err = g.waitPrimary(r.Context(), t.engine, w) + if err != nil { g.sseFailure(w, t.engine.Name, req.Model, err) return } - } else if err := g.ensure(r.Context(), t.engine); err != nil { - g.startFailed(w, t.engine.Name, req.Model, err) - return + } else { + var err error + fb, waited, err = g.waitPrimary(r.Context(), t.engine, nil) + if err != nil { + g.startFailed(w, t.engine.Name, req.Model, err) + return + } } } - u, err := url.Parse(t.url) + base := t.url + if fb != nil { + base, body = g.tookFallback(w, preamble, t.engine, fb, waited, body) + // The node key is the client's credential for this gateway; a + // fallback may run on another box and never gets it. + r.Header.Del("Authorization") + } + u, err := url.Parse(base) if err != nil { apiError(w, http.StatusInternalServerError, "engine url: "+err.Error()) return @@ -554,29 +569,6 @@ func ssePreamble(w http.ResponseWriter) { _ = http.NewResponseController(w).Flush() } -// keepAlive runs ensure for a streaming request, writing an SSE comment -// every keepAliveEvery with the seconds gone, so proxies and clients hold -// the connection and a reader sees the wait moving. -func (g *Gateway) keepAlive(w http.ResponseWriter, r *http.Request, e *config.Engine) error { - type outcome struct{ err error } - ch := make(chan outcome, 1) - go func() { ch <- outcome{g.ensure(r.Context(), e)} }() - start := time.Now() - tick := time.NewTicker(keepAliveEvery) - defer tick.Stop() - for { - select { - case o := <-ch: - return o.err - case <-r.Context().Done(): - return r.Context().Err() - case <-tick.C: - fmt.Fprintf(w, ": berth: loading %s, %d s\n\n", e.Name, int64(time.Since(start).Seconds())) - _ = http.NewResponseController(w).Flush() - } - } -} - // sseFailure answers a streaming request whose engine would not come up: // the same OpenAI-shaped error the 503 path carries, as one SSE event. func (g *Gateway) sseFailure(w http.ResponseWriter, engine, what string, err error) {