ctl federation add/list/remove/sync keeps peers in the store (migration 0029); their /v1/models are offered here as external models named PEER/MODEL and kept in step, a model still named by a fallback kept withdrawn. Peers are checked at /ready, timed, and taken out of routing while down. Requests to a peer carry x-llmgw-hop and are refused (508) when they come back to a cluster they passed through; residency and no_store keep requests off peers; X-Request-Id and, with --send-tags, attribution tags go along. ctl model fallbacks sets overflow to a peer. Console view, metrics, admin API routes and docs. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
595 lines
21 KiB
Go
595 lines
21 KiB
Go
package router
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"slices"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/federation"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
|
)
|
|
|
|
// fedCluster is one llmgw cluster, whole but for its nodes: a router serving
|
|
// on a loopback port, the control plane that publishes to it, the model
|
|
// service and federation. Its one model is served by a fake engine standing
|
|
// in for the replica a node would run.
|
|
type fedCluster struct {
|
|
name, id string
|
|
model string
|
|
rt *Router
|
|
srv *httptest.Server
|
|
ctl *control.Controller
|
|
models *modelsvc.Service
|
|
fed *federation.Service
|
|
engine *fakeEngine
|
|
// key is a client's key here, in project.
|
|
key, project string
|
|
keys int
|
|
}
|
|
|
|
type fedNoDispatch struct{}
|
|
|
|
func (fedNoDispatch) StartShard(context.Context, reconciler.ShardCommand) error { return nil }
|
|
func (fedNoDispatch) StopShard(context.Context, string, string) error { return nil }
|
|
func (fedNoDispatch) DrainShard(context.Context, string, string) error { return nil }
|
|
|
|
// withLocal is the sink between a control plane with no nodes and its
|
|
// router: what the control plane publishes, and the cluster's own engine.
|
|
type withLocal struct {
|
|
rt *Router
|
|
local endpoint.Info
|
|
}
|
|
|
|
func (s withLocal) SetEndpoints(ctx context.Context, infos []endpoint.Info) {
|
|
s.rt.SetEndpoints(ctx, append(slices.Clone(infos), s.local))
|
|
}
|
|
|
|
func (s withLocal) SetCatalog(c *catalog.Catalog) { s.rt.SetCatalog(c) }
|
|
|
|
func newFedCluster(t *testing.T, name, model string) *fedCluster {
|
|
t.Helper()
|
|
c := &fedCluster{name: name, id: name + "-cluster-id", model: model}
|
|
c.engine = newFakeEngine(t, model)
|
|
|
|
cfg := baseConfig()
|
|
cfg.Auth.Enabled = true
|
|
cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db")
|
|
cfg.Auth.FlushInterval = 20 * time.Millisecond
|
|
cfg.Auth.BatchSize = 1
|
|
cfg.Router.Retries = 0
|
|
c.rt = newTestRouter(t, cfg)
|
|
c.rt.SetFederationIdentity(c.id, name)
|
|
|
|
empty, err := catalog.New(nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
log := slog.New(slog.DiscardHandler)
|
|
sink := withLocal{rt: c.rt, local: endpoint.Info{ID: name + "-replica", URL: c.engine.srv.URL, Models: []string{model}}}
|
|
c.ctl = control.New(control.DefaultConfig(), catalog.NewLive(empty), fedNoDispatch{}, sink, log)
|
|
c.models = modelsvc.New(nil, c.rt.store, c.ctl, log)
|
|
c.fed = federation.New(federation.Config{Self: c.id, Interval: time.Hour}, c.rt.store, c.models, c.ctl.Catalog, log)
|
|
if err := c.fed.Load(t.Context()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
c.models.SetFederated(c.fed.Models)
|
|
c.ctl.SetPeers(c.fed)
|
|
c.srv = httptest.NewServer(c.rt.Handler())
|
|
t.Cleanup(c.srv.Close)
|
|
|
|
if err := c.models.Add(t.Context(), catModel(model, nil), "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// The control plane has published nothing yet — an empty list is no
|
|
// change — so the engine is put in by hand the first time.
|
|
sink.SetEndpoints(t.Context(), nil)
|
|
|
|
org, err := c.rt.store.CreateOrg(t.Context(), "acme-"+name)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
team, err := c.rt.store.CreateTeam(t.Context(), org, "platform")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if c.project, err = c.rt.store.CreateProject(t.Context(), team, "gateway"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
c.key = c.issueKey(t)
|
|
return c
|
|
}
|
|
|
|
// issueKey makes a key in the cluster's project, as `llmgw keys create`
|
|
// would, and returns its secret.
|
|
func (c *fedCluster) issueKey(t *testing.T) string {
|
|
t.Helper()
|
|
c.keys++
|
|
secret := store.KeyPrefix + strings.Repeat(string(rune('a'+c.keys)), 43)
|
|
sum := sha256.Sum256([]byte(secret))
|
|
if err := c.rt.store.CreateKey(t.Context(), store.Key{
|
|
ID: fmt.Sprintf("%s-key-%d", c.name, c.keys), Hash: sum[:], Prefix: "pfx" + c.name, Name: "test",
|
|
ProjectID: c.project, CreatedAt: time.Now(),
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return secret
|
|
}
|
|
|
|
// federate makes other a peer of c, with a key other issued to c.
|
|
func (c *fedCluster) federate(t *testing.T, other *fedCluster, sendTags bool) federation.Status {
|
|
t.Helper()
|
|
keyFile := filepath.Join(t.TempDir(), other.name+".key")
|
|
if err := os.WriteFile(keyFile, []byte(other.issueKey(t)+"\n"), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
st, err := c.fed.Add(t.Context(), federation.AddRequest{Name: other.name, URL: other.srv.URL,
|
|
KeyFile: keyFile, SendTags: sendTags}, "test")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
c.ctl.Publish(t.Context())
|
|
return st
|
|
}
|
|
|
|
// ask sends a chat completion to the cluster as its client.
|
|
func (c *fedCluster) ask(t *testing.T, model string, hdr map[string]string) *httptest.ResponseRecorder {
|
|
t.Helper()
|
|
h := map[string]string{"Authorization": "Bearer " + c.key}
|
|
for k, v := range hdr {
|
|
h[k] = v
|
|
}
|
|
return post(t, c.rt.Handler(), map[string]any{"model": model,
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}}}, h)
|
|
}
|
|
|
|
// usageOf waits for the cluster's record of a request.
|
|
func (c *fedCluster) usageOf(t *testing.T, requestID string) store.Usage {
|
|
t.Helper()
|
|
deadline := time.Now().Add(5 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
u, ok, err := c.rt.store.UsageByRequest(t.Context(), requestID)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if ok {
|
|
return u
|
|
}
|
|
time.Sleep(20 * time.Millisecond)
|
|
}
|
|
t.Fatalf("%s has no record of request %s", c.name, requestID)
|
|
return store.Usage{}
|
|
}
|
|
|
|
func listedModels(t *testing.T, c *fedCluster) (map[string]map[string]any, http.Header) {
|
|
t.Helper()
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
req.Header.Set("Authorization", "Bearer "+c.key)
|
|
rec := httptest.NewRecorder()
|
|
c.rt.Handler().ServeHTTP(rec, req)
|
|
var list struct {
|
|
Data []map[string]any `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &list); err != nil {
|
|
t.Fatalf("%d %s", rec.Code, rec.Body.String())
|
|
}
|
|
out := map[string]map[string]any{}
|
|
for _, m := range list.Data {
|
|
out[fmt.Sprint(m["id"])] = m
|
|
}
|
|
return out, rec.Header()
|
|
}
|
|
|
|
func TestAPeersModelsAreOfferedHereUnderItsNameAndFollowItsList(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
|
|
st := west.federate(t, east, false)
|
|
if !st.Up || st.Cluster != east.id || st.Latency <= 0 {
|
|
t.Errorf("east after adding: %+v", st)
|
|
}
|
|
listed, hdr := listedModels(t, west)
|
|
if listed["east/qwen-east"] == nil || listed["east/qwen-east"]["peer"] != "east" {
|
|
t.Fatalf("west offers %v; east's model should be there as east/qwen-east", slices.Collect(mapsKeys(listed)))
|
|
}
|
|
if hdr.Get(clusterHeader) != west.id {
|
|
t.Errorf("west's /v1/models says cluster %q", hdr.Get(clusterHeader))
|
|
}
|
|
if m, ok := west.ctl.Catalog().Resolve("east/qwen-east"); !ok || !m.IsPeer() || m.External.Model != "qwen-east" {
|
|
t.Errorf("west's catalog has %+v", m)
|
|
}
|
|
|
|
// Asked for by its name here, answered by east's engine under east's name.
|
|
rec := west.ask(t, "east/qwen-east", nil)
|
|
if rec.Code != http.StatusOK || east.engine.calls.Load() != 1 {
|
|
t.Fatalf("%d %s; east's engine was called %d times", rec.Code, rec.Body.String(), east.engine.calls.Load())
|
|
}
|
|
var sent map[string]any
|
|
_ = json.Unmarshal(east.engine.last.Load().([]byte), &sent)
|
|
if sent["model"] != "qwen-east" {
|
|
t.Errorf("east's engine was asked for %v", sent["model"])
|
|
}
|
|
|
|
// East adds a model, then drops its first: west follows at its next check.
|
|
if err := east.models.Add(t.Context(), catModel("qwen-east-2", nil), "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := east.models.Remove(t.Context(), "qwen-east"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
east.rt.SetEndpoints(t.Context(), []endpoint.Info{{ID: "e2", URL: east.engine.srv.URL, Models: []string{"qwen-east-2"}}})
|
|
if _, err := west.fed.Sync(t.Context(), "east"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
if _, ok := west.ctl.Catalog().Resolve("east/qwen-east-2"); !ok {
|
|
t.Error("a model east added is not offered here")
|
|
}
|
|
if _, ok := west.ctl.Catalog().Resolve("east/qwen-east"); ok {
|
|
t.Error("a model east dropped, which nothing here names, is still offered")
|
|
}
|
|
}
|
|
|
|
func mapsKeys(m map[string]map[string]any) func(func(string) bool) {
|
|
return func(yield func(string) bool) {
|
|
for k := range m {
|
|
if !yield(k) {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestOverflowGoesToThePeerWhichSeesOneKeyAndTheRequestID(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
west.federate(t, east, false)
|
|
if _, err := west.models.SetFallbacks(t.Context(), "qwen-west", []string{"east/qwen-east"}, "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
|
|
// West has capacity: east is not asked.
|
|
if rec := west.ask(t, "qwen-west", nil); rec.Code != http.StatusOK || east.engine.calls.Load() != 0 {
|
|
t.Fatalf("%d; east called %d times", rec.Code, east.engine.calls.Load())
|
|
}
|
|
|
|
// West's replica is full: the overflow goes east.
|
|
west.engine.status.Store(http.StatusServiceUnavailable)
|
|
rec := west.ask(t, "qwen-west", map[string]string{"X-Request-Id": "overflow-1"})
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("%d %s", rec.Code, rec.Body.String())
|
|
}
|
|
if got := rec.Header().Values("x-llmgw-served-model"); !slices.Equal(got, []string{"east/qwen-east"}) {
|
|
t.Errorf("served by %q", got)
|
|
}
|
|
// East's own answer headers do not reach the client beside west's.
|
|
if got := rec.Header().Values("X-Request-Id"); len(got) != 1 {
|
|
t.Errorf("the client was given request ids %q", got)
|
|
}
|
|
if east.engine.calls.Load() != 1 {
|
|
t.Errorf("east's engine was called %d times", east.engine.calls.Load())
|
|
}
|
|
|
|
// West records which peer served; east, the one key west uses there,
|
|
// the request's id, and no tags, which are off unless asked for.
|
|
here := west.usageOf(t, "overflow-1")
|
|
if here.ServedModel != "east/qwen-east" || here.Replica != "peer:east/qwen-east" {
|
|
t.Errorf("west's record: served %q by %q", here.ServedModel, here.Replica)
|
|
}
|
|
there := east.usageOf(t, "overflow-1")
|
|
if there.KeyID != "east-key-2" || there.Model != "qwen-east" {
|
|
t.Errorf("east's record: key %q, model %q", there.KeyID, there.Model)
|
|
}
|
|
if len(there.Tags) != 0 {
|
|
t.Errorf("east was told %v without send-tags", there.Tags)
|
|
}
|
|
|
|
// With tags on, east can say which cluster and project asked.
|
|
west.federate(t, east, true)
|
|
rec = west.ask(t, "qwen-west", map[string]string{"X-Request-Id": "overflow-2"})
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("%d %s", rec.Code, rec.Body.String())
|
|
}
|
|
there = east.usageOf(t, "overflow-2")
|
|
if there.Tags["federated_from"] != "west" || there.Tags["project"] != west.project {
|
|
t.Errorf("east's record carries tags %v", there.Tags)
|
|
}
|
|
}
|
|
|
|
func TestTwoClustersFederatingEachOtherDoNotSendARequestRoundAgain(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
west.federate(t, east, false)
|
|
east.federate(t, west, false)
|
|
|
|
// Neither offers the other's peer models back: no east/west/... names.
|
|
if _, ok := west.ctl.Catalog().Resolve("east/west/qwen-west"); ok {
|
|
t.Error("west took back its own model through east")
|
|
}
|
|
if _, ok := east.ctl.Catalog().Resolve("west/east/qwen-east"); ok {
|
|
t.Error("east took back its own model through west")
|
|
}
|
|
|
|
// Each falls back to the other, and both are full.
|
|
if _, err := west.models.SetFallbacks(t.Context(), "qwen-west", []string{"east/qwen-east"}, "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := east.models.SetFallbacks(t.Context(), "qwen-east", []string{"west/qwen-west"}, "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
east.ctl.Publish(t.Context())
|
|
west.engine.status.Store(http.StatusServiceUnavailable)
|
|
east.engine.status.Store(http.StatusServiceUnavailable)
|
|
|
|
rec := west.ask(t, "qwen-west", nil)
|
|
if rec.Code == http.StatusOK {
|
|
t.Fatalf("served: %s", rec.Body.String())
|
|
}
|
|
// West's replica once, east's once; east did not send it back west.
|
|
if w, e := west.engine.calls.Load(), east.engine.calls.Load(); w != 1 || e != 1 {
|
|
t.Errorf("west's engine was called %d times, east's %d; each should be once", w, e)
|
|
}
|
|
|
|
// And a request that says it has passed through a cluster already is
|
|
// refused there, before anything else.
|
|
rec = west.ask(t, "qwen-west", map[string]string{hopHeader: east.id + ", " + west.id})
|
|
if rec.Code != http.StatusLoopDetected {
|
|
t.Errorf("a request that came back: %d %s", rec.Code, rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestAProjectThatStaysInTheClusterIsNeverSentToAPeer(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
west.federate(t, east, false)
|
|
if _, err := west.models.SetFallbacks(t.Context(), "qwen-west", []string{"east/qwen-east"}, "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
west.engine.status.Store(http.StatusServiceUnavailable)
|
|
|
|
if err := west.rt.store.SetInCluster(t.Context(), west.project, true); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if rec := west.ask(t, "east/qwen-east", nil); rec.Code != http.StatusForbidden {
|
|
t.Errorf("a peer's model, asked for by a project that stays in the cluster: %d %s", rec.Code, rec.Body.String())
|
|
}
|
|
if rec := west.ask(t, "qwen-west", nil); rec.Code == http.StatusOK {
|
|
t.Errorf("served by %q", rec.Header().Get("x-llmgw-served-model"))
|
|
}
|
|
if east.engine.calls.Load() != 0 {
|
|
t.Errorf("east was sent %d requests", east.engine.calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestAPolicyThatKeepsNothingIsNeverSentToAPeer(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
west.federate(t, east, false)
|
|
if _, err := west.models.SetFallbacks(t.Context(), "qwen-west", []string{"east/qwen-east"}, "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
west.engine.status.Store(http.StatusServiceUnavailable)
|
|
g := &switchGuard{}
|
|
g.noStore.Store(true)
|
|
west.rt.SetGuard(g)
|
|
|
|
if rec := west.ask(t, "east/qwen-east", nil); rec.Code != http.StatusBadRequest {
|
|
t.Errorf("a peer's model under no_store: %d %s", rec.Code, rec.Body.String())
|
|
}
|
|
if rec := west.ask(t, "qwen-west", nil); rec.Code == http.StatusOK {
|
|
t.Errorf("served by %q", rec.Header().Get("x-llmgw-served-model"))
|
|
}
|
|
if east.engine.calls.Load() != 0 {
|
|
t.Errorf("east was sent %d requests", east.engine.calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestAPeerThatIsDownIsTakenOutOfRoutingAndBackWhenItAnswers(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
west.federate(t, east, false)
|
|
if len(west.rt.reg.Snapshots("east/qwen-east")) != 1 {
|
|
t.Fatal("east's model is not routed to while east is up")
|
|
}
|
|
|
|
// East has nothing of its own to serve with: its /ready says so.
|
|
east.rt.SetEndpoints(t.Context(), nil)
|
|
if st, _ := west.fed.Sync(t.Context(), "east"); st.Up || st.Error == "" {
|
|
t.Errorf("east with nothing to serve: %+v", st)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
if n := len(west.rt.reg.Snapshots("east/qwen-east")); n != 0 {
|
|
t.Errorf("east is down and still routed to: %d endpoints", n)
|
|
}
|
|
if _, ok := west.ctl.Catalog().Resolve("east/qwen-east"); !ok {
|
|
t.Error("east's model left the catalog when east went down; down is not gone")
|
|
}
|
|
if rec := west.ask(t, "east/qwen-east", nil); rec.Code == http.StatusOK {
|
|
t.Error("served by a peer that is down")
|
|
}
|
|
|
|
// Back.
|
|
east.ctl.Publish(t.Context())
|
|
east.rt.SetEndpoints(t.Context(), []endpoint.Info{{ID: "e", URL: east.engine.srv.URL, Models: []string{"qwen-east"}}})
|
|
if st, _ := west.fed.Sync(t.Context(), "east"); !st.Up {
|
|
t.Errorf("east is back: %+v", st)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
if rec := west.ask(t, "east/qwen-east", nil); rec.Code != http.StatusOK {
|
|
t.Errorf("east is back: %d %s", rec.Code, rec.Body.String())
|
|
}
|
|
|
|
// Gone altogether.
|
|
east.srv.Close()
|
|
if st, _ := west.fed.Sync(t.Context(), "east"); st.Up {
|
|
t.Errorf("east is gone: %+v", st)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
if n := len(west.rt.reg.Snapshots("east/qwen-east")); n != 0 {
|
|
t.Errorf("east is gone and still routed to: %d endpoints", n)
|
|
}
|
|
}
|
|
|
|
func TestAPeerModelAFallbackNamesIsKeptWithdrawnAndThePeerCannotBeRemoved(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
west.federate(t, east, false)
|
|
if _, err := west.models.SetFallbacks(t.Context(), "qwen-west", []string{"east/qwen-east"}, "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if err := west.fed.Remove(t.Context(), "east"); err == nil || !strings.Contains(err.Error(), "qwen-west") {
|
|
t.Errorf("removing a peer a fallback names: %v", err)
|
|
}
|
|
|
|
// East stops offering it: kept, withdrawn, and not routed to.
|
|
if err := east.models.Remove(t.Context(), "qwen-east"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// Its replica goes too; east still serves something, so it is up.
|
|
east.rt.SetEndpoints(t.Context(), []endpoint.Info{{ID: "other", URL: east.engine.srv.URL, Models: []string{"other"}}})
|
|
if _, err := west.fed.Sync(t.Context(), "east"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
west.ctl.Publish(t.Context())
|
|
if _, ok := west.ctl.Catalog().Resolve("qwen-west"); !ok {
|
|
t.Fatal("the model that falls back to east's lost its place in the catalog")
|
|
}
|
|
st, _ := west.fed.Status("east")
|
|
if !slices.ContainsFunc(st.Models, func(m federation.ModelStatus) bool { return m.Upstream == "qwen-east" && m.Withdrawn }) {
|
|
t.Errorf("east's models: %+v", st.Models)
|
|
}
|
|
if n := len(west.rt.reg.Snapshots("east/qwen-east")); n != 0 {
|
|
t.Errorf("a withdrawn model is routed to: %d endpoints", n)
|
|
}
|
|
|
|
// Once nothing names it, it goes, and so may the peer.
|
|
if _, err := west.models.SetFallbacks(t.Context(), "qwen-west", nil, "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := west.fed.Sync(t.Context(), "east"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, ok := west.ctl.Catalog().Resolve("east/qwen-east"); ok {
|
|
t.Error("a withdrawn model nothing names is still in the catalog")
|
|
}
|
|
if err := west.fed.Remove(t.Context(), "east"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(west.fed.List()) != 0 {
|
|
t.Error("east is still a peer")
|
|
}
|
|
}
|
|
|
|
func TestPeersAndTheirModelsSurviveARestart(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
west.federate(t, east, false)
|
|
if _, err := west.models.SetFallbacks(t.Context(), "qwen-west", []string{"east/qwen-east"}, "test"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// A new control plane against the same store, before east is asked
|
|
// anything: the model falling back to east's is whole.
|
|
log := slog.New(slog.DiscardHandler)
|
|
empty, _ := catalog.New(nil)
|
|
ctl := control.New(control.DefaultConfig(), catalog.NewLive(empty), fedNoDispatch{}, withLocal{rt: west.rt}, log)
|
|
models := modelsvc.New(nil, west.rt.store, ctl, log)
|
|
fed := federation.New(federation.Config{Self: west.id}, west.rt.store, models, ctl.Catalog, log)
|
|
if err := fed.Load(t.Context()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
models.SetFederated(fed.Models)
|
|
built, err := models.Load(t.Context())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
m, ok := built.Resolve("qwen-west")
|
|
if !ok || !slices.Equal(m.Fallbacks, []string{"east/qwen-east"}) {
|
|
t.Errorf("after a restart: %+v", m)
|
|
}
|
|
if _, ok := built.Resolve("east/qwen-east"); !ok {
|
|
t.Error("east's model is not in the catalog after a restart")
|
|
}
|
|
// Down until checked.
|
|
if _, _, ok := fed.PeerRoute("east", "qwen-east"); ok {
|
|
t.Error("a peer not yet checked is routed to")
|
|
}
|
|
}
|
|
|
|
func TestAClusterCannotBeItsOwnPeer(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
keyFile := filepath.Join(t.TempDir(), "self.key")
|
|
if err := os.WriteFile(keyFile, []byte(west.issueKey(t)), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err := west.fed.Add(t.Context(), federation.AddRequest{Name: "self", URL: west.srv.URL, KeyFile: keyFile}, "test")
|
|
if err == nil || !strings.Contains(err.Error(), "this cluster") {
|
|
t.Errorf("adding itself: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestAPeerIsCheckedBeforeItIsKept(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
keyFile := filepath.Join(t.TempDir(), "wrong.key")
|
|
if err := os.WriteFile(keyFile, []byte(store.KeyPrefix+strings.Repeat("z", 43)), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for name, req := range map[string]federation.AddRequest{
|
|
"a key the peer refuses": {Name: "east", URL: east.srv.URL, KeyFile: keyFile},
|
|
"nothing at the address": {Name: "east", URL: "http://127.0.0.1:1", KeyFile: keyFile},
|
|
"a name with a slash": {Name: "ea/st", URL: east.srv.URL, KeyFile: keyFile},
|
|
"no key file": {Name: "east", URL: east.srv.URL},
|
|
} {
|
|
if _, err := west.fed.Add(t.Context(), req, "test"); err == nil {
|
|
t.Errorf("%s: accepted", name)
|
|
}
|
|
}
|
|
if len(west.fed.List()) != 0 {
|
|
t.Errorf("kept: %+v", west.fed.List())
|
|
}
|
|
}
|
|
|
|
// The federation's checks run beside requests and catalog rebuilds.
|
|
func TestFederationChecksAreSafeBesideTraffic(t *testing.T) {
|
|
west := newFedCluster(t, "west", "qwen-west")
|
|
east := newFedCluster(t, "east", "qwen-east")
|
|
west.federate(t, east, false)
|
|
var wg sync.WaitGroup
|
|
for range 4 {
|
|
wg.Go(func() {
|
|
for range 5 {
|
|
west.fed.SyncAll(t.Context())
|
|
west.ctl.Publish(t.Context())
|
|
}
|
|
})
|
|
wg.Go(func() {
|
|
for range 5 {
|
|
west.ask(t, "east/qwen-east", nil)
|
|
_ = west.fed.List()
|
|
}
|
|
})
|
|
}
|
|
wg.Wait()
|
|
}
|