diff --git a/internal/cli/server.go b/internal/cli/server.go index 1f4d2d5..432e16b 100644 --- a/internal/cli/server.go +++ b/internal/cli/server.go @@ -95,8 +95,14 @@ func newServerCmd() *cobra.Command { if err != nil { return err } + live := catalog.NewLive(cat) + // The router derives its own view, and it is empty until told. + // Told here rather than at construction so there is exactly one + // path by which it learns the catalog, used at startup and at + // every change after. + rt.SetCatalog(cat) - ctl, svc, err := startControlPlane(ctx, cfg, cat, rt, cfg.AgentTLS.Enabled(), log) + ctl, svc, err := startControlPlane(ctx, cfg, live, rt, cfg.AgentTLS.Enabled(), log) if err != nil { return err } @@ -235,7 +241,7 @@ func newServerCmd() *cobra.Command { // The service is both the AgentService handler and the reconciler's Dispatcher: // it holds the open streams, so it is the only thing that can deliver a command // to a node. -func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Catalog, +func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Live, sink control.EndpointSink, requireIdentity bool, log *slog.Logger, ) (*control.Controller, *agentsvc.Service, error) { ccfg := control.DefaultConfig() diff --git a/internal/control/adminapi/api_test.go b/internal/control/adminapi/api_test.go index c0a5f2c..a5f27d7 100644 --- a/internal/control/adminapi/api_test.go +++ b/internal/control/adminapi/api_test.go @@ -40,7 +40,7 @@ func setup(t *testing.T, models []catalog.Model, nodes ...inventory.Node) (*API, // anything runs, which is the state settling waits for. Leaving it on would // make every test here assert that placement had not happened yet. cfg.Reconciler.SettleFor = -1 - ctl := control.New(cfg, cat, nullDispatcher{}, nullSink{}, slog.New(slog.DiscardHandler)) + ctl := control.New(cfg, catalog.NewLive(cat), nullDispatcher{}, nullSink{}, slog.New(slog.DiscardHandler)) for _, n := range nodes { ctl.Inventory().Upsert(n) } diff --git a/internal/control/catalog/live.go b/internal/control/catalog/live.go new file mode 100644 index 0000000..9841df1 --- /dev/null +++ b/internal/control/catalog/live.go @@ -0,0 +1,36 @@ +package catalog + +import "sync/atomic" + +// Live holds the catalog that is currently in force. +// +// The catalog was built once, from YAML, at process start, which made adding a +// model an operator procedure with downtime: edit a file on the control plane, +// restart it, wait out the settling window, watch replicas churn while adoption +// catches up. That is a reasonable way to manage a fixed set of production +// models and a terrible one for a platform whose whole promise is that people +// can try things. +// +// A Catalog stays immutable — validation happens once, in New, and a holder can +// read one without a lock or a copy. What changes is which immutable catalog is +// current. Readers take the pointer and work from it for the rest of their +// operation, so a swap mid-request is invisible rather than a race: the request +// finishes against the catalog it started with, which is the behaviour anyone +// would want anyway. +type Live struct { + v atomic.Pointer[Catalog] +} + +// NewLive holds an initial catalog. +func NewLive(c *Catalog) *Live { + l := &Live{} + l.v.Store(c) + return l +} + +// Load returns the catalog in force. Never nil once constructed. +func (l *Live) Load() *Catalog { return l.v.Load() } + +// Store replaces it. The caller has already validated the new one, because +// New is the only way to build a Catalog and it refuses an invalid one. +func (l *Live) Store(c *Catalog) { l.v.Store(c) } diff --git a/internal/control/control.go b/internal/control/control.go index 9ca29cb..abf536c 100644 --- a/internal/control/control.go +++ b/internal/control/control.go @@ -25,6 +25,17 @@ type EndpointSink interface { SetEndpoints(ctx context.Context, infos []endpoint.Info) } +// CatalogSink receives catalog changes. The router implements it. +// +// The catalog used to be fixed for the life of the process, so the router could +// derive its view once and keep it. Now that a model can be added while the +// cluster runs, the router has to be told — and it is told rather than asked, +// because a push happens once per change and a pull would happen once per +// request. +type CatalogSink interface { + SetCatalog(c *catalog.Catalog) +} + // UsageSource reports when each endpoint was last sent a request. // // The return path of EndpointSink, and an interface for the same reason: the @@ -62,7 +73,7 @@ type Controller struct { cfg Config log *slog.Logger inv *inventory.Store - cat *catalog.Catalog + cat *catalog.Live rec *reconciler.Reconciler sink EndpointSink usage UsageSource @@ -76,7 +87,7 @@ type Controller struct { } // New builds a controller. -func New(cfg Config, cat *catalog.Catalog, disp reconciler.Dispatcher, sink EndpointSink, log *slog.Logger) *Controller { +func New(cfg Config, cat *catalog.Live, disp reconciler.Dispatcher, sink EndpointSink, log *slog.Logger) *Controller { if log == nil { log = slog.Default() } @@ -118,11 +129,40 @@ func (c *Controller) collectUsage() { c.rec.RecordUsage(c.usage.LastUsed()) } +// ReplaceCatalog puts a new catalog in force, without a restart. +// +// The whole point of the change: adding a model was an operator procedure with +// downtime — edit a file on the control plane, restart, wait out the settling +// window, watch adoption churn. Nothing about that was necessary. The +// reconciler already compares desired against actual every pass, so a new +// catalog is just a different desired, and it is picked up on the next tick +// without anything being torn down. +// +// The catalog is validated by its own constructor, so an invalid one cannot +// reach here; what can reach here is a valid catalog that removes a model +// something is still serving, and that is handled the way it always has been — +// the reconciler stops replicas whose model is gone, and Publish drops +// endpoints for them first, so traffic stops being routed before anything is +// killed. +func (c *Controller) ReplaceCatalog(ctx context.Context, next *catalog.Catalog) { + c.cat.Store(next) + if cs, ok := c.sink.(CatalogSink); ok { + cs.SetCatalog(next) + } + // Endpoints first, then wake the loop: a model that just left the catalog + // should stop receiving traffic before its replicas are stopped, and + // Publish is what takes it out of the router's list. + c.Publish(ctx) + c.rec.Notify() + + c.log.Info("catalog replaced", "models", len(next.Models())) +} + // Reconciler exposes the reconciliation loop, for status reports and ctl. func (c *Controller) Reconciler() *reconciler.Reconciler { return c.rec } // Catalog exposes the model catalog. -func (c *Controller) Catalog() *catalog.Catalog { return c.cat } +func (c *Controller) Catalog() *catalog.Catalog { return c.cat.Load() } // Start runs the reconciler and begins publishing endpoints. func (c *Controller) Start(ctx context.Context) { @@ -187,10 +227,11 @@ func (c *Controller) endpoints() []endpoint.Info { out := make([]endpoint.Info, 0, len(ready)) for _, rep := range ready { - m, ok := c.cat.Resolve(rep.Model) + m, ok := c.cat.Load().Resolve(rep.Model) if !ok { // The model left the catalog while a replica was still running. - // The reconciler will tear it down; until then, do not route to it. + // Withdraw it here first, so traffic stops before the reconciler + // stops the engine on its next pass. continue } // Clients may address a model by any of its names, and the router diff --git a/internal/control/reconciler/reconciler.go b/internal/control/reconciler/reconciler.go index 02f6828..aa28106 100644 --- a/internal/control/reconciler/reconciler.go +++ b/internal/control/reconciler/reconciler.go @@ -137,7 +137,7 @@ func (c *Config) normalize() { // Reconciler compares the catalog against what is running and closes the gap. type Reconciler struct { cfg Config - cat *catalog.Catalog + cat *catalog.Live inv *inventory.Store sched *placement.Scheduler disp Dispatcher @@ -171,7 +171,7 @@ type Reconciler struct { } // New builds a reconciler. -func New(cfg Config, cat *catalog.Catalog, inv *inventory.Store, disp Dispatcher, log *slog.Logger) *Reconciler { +func New(cfg Config, cat *catalog.Live, inv *inventory.Store, disp Dispatcher, log *slog.Logger) *Reconciler { cfg.normalize() if log == nil { log = slog.Default() @@ -252,9 +252,44 @@ func (r *Reconciler) Reconcile(ctx context.Context) { return } - for _, m := range r.cat.Models() { + for _, m := range r.cat.Load().Models() { r.reconcileModel(ctx, m) } + r.stopReplicasOfRemovedModels(ctx) +} + +// stopReplicasOfRemovedModels tears down what the catalog no longer asks for. +// +// reconcileModel only ever visits models that are in the catalog, so a replica +// whose model has been removed was visited by nothing: it kept its GPUs, kept +// counting towards the inventory, and kept running an engine for a name no +// client could ask for. The endpoint publisher already skipped it with the +// comment "the reconciler will tear it down", which was simply not true. +// +// It mattered little while a catalog could only change by restarting the +// process. It matters now that models come and go while the cluster runs, which +// is the whole point of making the catalog live. +// +// Ordered after the model loop so a rebuild elsewhere has already claimed what +// it needs: freeing these GPUs first would let a replacement land on hardware +// that is about to be freed anyway, and churn two nodes instead of one. +func (r *Reconciler) stopReplicasOfRemovedModels(ctx context.Context) { + cat := r.cat.Load() + + r.mu.Lock() + var gone []*Replica + for _, rep := range r.replicas { + if _, ok := cat.Resolve(rep.Model); !ok { + gone = append(gone, rep) + } + } + r.mu.Unlock() + + for _, rep := range gone { + r.log.Info("stopping a replica whose model left the catalog", + "replica", rep.ID, "model", rep.Model) + r.stopReplica(ctx, rep.ID) + } } // fleetHasSettled reports whether placement may run yet. @@ -380,7 +415,7 @@ func (r *Reconciler) expireStuckReplicas() { // does not say — a replica with no deadline at all is the thing this exists to // prevent. func (r *Reconciler) startupTimeout(rep *Replica) time.Duration { - if m, ok := r.cat.Resolve(rep.Model); ok { + if m, ok := r.cat.Load().Resolve(rep.Model); ok { if p, ok := m.ProfileByName(rep.Profile); ok && p.StartupTimeout > 0 { return p.StartupTimeout } @@ -890,7 +925,7 @@ func (r *Reconciler) Adopt(nodeID string, shards []AdoptedShard) (unclaimed []Ad } first := group[0] - m, ok := r.cat.Resolve(first.Model) + m, ok := r.cat.Load().Resolve(first.Model) if !ok { unclaimed = append(unclaimed, group...) r.reservePorts(group) diff --git a/internal/control/reconciler/reconciler_test.go b/internal/control/reconciler/reconciler_test.go index e8acc31..1f5f8d3 100644 --- a/internal/control/reconciler/reconciler_test.go +++ b/internal/control/reconciler/reconciler_test.go @@ -134,7 +134,7 @@ func newHarness(t *testing.T, m catalog.Model, nodes ...inventory.Node) *harness cfg.UnreachableAfter = 15 * time.Second cfg.LostAfter = 60 * time.Second - r := New(cfg, cat, inv, disp, slog.New(slog.DiscardHandler)) + r := New(cfg, catalog.NewLive(cat), inv, disp, slog.New(slog.DiscardHandler)) r.SetClock(c.now) // Nodes were just registered, so their LastSeen is wall-clock now, which is // far from the fake clock's epoch. Re-stamp them onto the fake timeline. @@ -485,7 +485,7 @@ func TestScaleDownDrainsFirst(t *testing.T) { if err != nil { t.Fatal(err) } - h.r.cat = cat + h.r.cat.Store(cat) h.disp.reset() h.r.Reconcile(t.Context()) @@ -1012,7 +1012,7 @@ func settlingHarness(t *testing.T, m catalog.Model) (*Reconciler, *inventory.Sto inv := inventory.NewStore() disp := &fakeDispatcher{} c := newClock() - r := New(DefaultConfig(), cat, inv, disp, slog.New(slog.DiscardHandler)) + r := New(DefaultConfig(), catalog.NewLive(cat), inv, disp, slog.New(slog.DiscardHandler)) r.SetClock(c.now) return r, inv, disp, c } @@ -1229,3 +1229,72 @@ func TestAnEmptyUsageReportLeavesEverythingAlone(t *testing.T) { t.Errorf("idle=%s known=%v; an empty report erased what was known", idle, known) } } + +func TestAModelAddedWhileRunningIsPlaced(t *testing.T) { + // The reason the catalog stopped being fixed at startup. Adding a model was + // an operator procedure with downtime: edit a file on the control plane, + // restart it, wait out the settling window, watch adoption churn. None of + // that was necessary — the loop already compares desired against actual + // every pass, so a new catalog is simply a different desired. + h := newHarness(t, testModel(1), gpuNode("1", 8), gpuNode("2", 8)) + h.r.Reconcile(t.Context()) + if got := len(h.r.Replicas()); got != 1 { + t.Fatalf("expected the original model's replica, got %d", got) + } + + second := testModel(1) + second.Name = "llama-two" + next, err := catalog.New([]catalog.Model{testModel(1), second}) + if err != nil { + t.Fatal(err) + } + h.r.cat.Store(next) + + h.r.Reconcile(t.Context()) + byModel := map[string]int{} + for _, rep := range h.r.Replicas() { + byModel[rep.Model]++ + } + if byModel["llama-two"] == 0 { + t.Errorf("a model added while running was never placed: %v", byModel) + } + if byModel["llama"] == 0 { + t.Errorf("the model that was already running was disturbed: %v", byModel) + } +} + +func TestAModelRemovedWhileRunningIsStopped(t *testing.T) { + // The other direction, and the one with consequences: a replica whose model + // has left the catalog holds GPUs for something nothing can ask for. + first := testModel(1) + second := testModel(1) + second.Name = "llama-two" + both, err := catalog.New([]catalog.Model{first, second}) + if err != nil { + t.Fatal(err) + } + h := newHarness(t, first, gpuNode("1", 8), gpuNode("2", 8)) + h.r.cat.Store(both) + h.r.Reconcile(t.Context()) + h.r.Reconcile(t.Context()) + if len(h.r.Replicas()) < 2 { + t.Fatalf("expected both models placed, got %d", len(h.r.Replicas())) + } + + only, err := catalog.New([]catalog.Model{first}) + if err != nil { + t.Fatal(err) + } + h.r.cat.Store(only) + + // Several passes: stopping is deliberate rather than immediate, and a + // replica is not torn down in the same breath it is noticed. + for range 4 { + h.r.Reconcile(t.Context()) + } + for _, rep := range h.r.Replicas() { + if rep.Model == "llama-two" && rep.State.Counts() { + t.Errorf("a replica of a removed model is still counted: %+v", rep) + } + } +} diff --git a/internal/router/chat.go b/internal/router/chat.go index aebf9bd..159e3de 100644 --- a/internal/router/chat.go +++ b/internal/router/chat.go @@ -148,7 +148,7 @@ func (r *Router) modelsFor(req *http.Request, user store.User) []web.ChatModel { } seen[m] = struct{}{} if r.permits(principal, m) { - out = append(out, web.ChatModel{Name: m, Simulated: r.models[m].Simulated()}) + out = append(out, web.ChatModel{Name: m, Simulated: r.modelMap()[m].Simulated()}) } } } diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 05c7f45..357f702 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -543,7 +543,7 @@ func noEndpointFor(model string, served []string) string { func (r *Router) exampleModels() []string { names := r.servedModels() slices.SortStableFunc(names, func(a, b string) int { - as, bs := r.models[a].Simulated(), r.models[b].Simulated() + as, bs := r.modelMap()[a].Simulated(), r.modelMap()[b].Simulated() switch { case as == bs: return 0 @@ -581,8 +581,8 @@ func (r *Router) handleModels(w http.ResponseWriter, _ *http.Request) { } } // Catalog aliases are usable names too, so clients should see them. - for name := range r.models { - if _, ok := seen[r.models[name].Name]; ok { + for name := range r.modelMap() { + if _, ok := seen[r.modelMap()[name].Name]; ok { seen[name] = struct{}{} } } @@ -600,7 +600,7 @@ func (r *Router) handleModels(w http.ResponseWriter, _ *http.Request) { // "llmgw" here was hiding the one distinction a client needs: this // model returns placeholder text, not inference. owner := "llmgw" - if r.models[n].Simulated() { + if r.modelMap()[n].Simulated() { owner = "llmgw-sim" } entry := map[string]any{ @@ -608,7 +608,7 @@ func (r *Router) handleModels(w http.ResponseWriter, _ *http.Request) { } // Only when known. A context_length of 0 would read as "no context", // which is a worse answer than saying nothing. - m := r.models[n] + m := r.modelMap()[n] if m.ContextLength > 0 { entry["context_length"] = m.ContextLength } diff --git a/internal/router/router.go b/internal/router/router.go index 8d2737e..14fa669 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -13,6 +13,7 @@ import ( "net/http" "slices" "strings" + "sync/atomic" "time" "gitea.mynodes.xyz/ilker/llmgw/internal/account" @@ -46,7 +47,12 @@ type Router struct { profiles map[string]*scheduler.Profile // models maps every name a client may use — canonical names and aliases — // onto its catalog entry. - models map[string]catalog.Model + // + // Held atomically because the catalog is no longer fixed at startup: adding + // a model used to mean restarting the control plane, and this is the + // router's half of not having to. A request takes the map once and works + // from it, so a swap mid-request is invisible rather than a race. + models atomic.Pointer[map[string]catalog.Model] // landing configures the public page, which the router serves but does not // otherwise care about. @@ -135,21 +141,12 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { return nil, fmt.Errorf("unknown default_profile %q", cfg.Router.DefaultProfile) } - models := map[string]catalog.Model{} - for _, m := range cfg.Models { - models[m.Name] = m - for _, a := range m.Aliases { - models[a] = m - } - } - r := &Router{ cfg: cfg.Router, landing: cfg.Landing, log: log, reg: reg, profiles: profiles, - models: models, prefix: prefixScorer, m: metrics, px: proxy.New(proxy.Config{ @@ -220,6 +217,22 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { } r.adm = adm } + + // Seeded from configuration, not only pushed. A router running without a + // control plane — a static endpoint list, which is a supported mode — has + // nothing to push it a catalog, and would otherwise resolve no model at + // all. A control plane pushes over the top of this at startup and on every + // change after. + if len(cfg.Models) > 0 { + seed := map[string]catalog.Model{} + for _, m := range cfg.Models { + seed[m.Name] = m + for _, a := range m.Aliases { + seed[a] = m + } + } + r.models.Store(&seed) + } return r, nil } @@ -335,6 +348,33 @@ func (r *Router) Registry() *endpoint.Registry { return r.reg } // the update keep their live state — inflight counts, breaker state, drain rate // — so a replica that was already serving does not start from cold just // because another one moved. +// modelMap returns the name→model map in force. +// +// Never nil: a router built before any catalog arrived resolves nothing, which +// is the same answer it gave for an unknown name before. +func (r *Router) modelMap() map[string]catalog.Model { + if m := r.models.Load(); m != nil { + return *m + } + return nil +} + +// SetCatalog replaces the models this router knows about. +// +// The catalog's counterpart to SetEndpoints, and pushed for the same reason: +// the control plane owns what should exist, the router needs to know, and +// neither has to be in the other's process for that to work. +func (r *Router) SetCatalog(c *catalog.Catalog) { + m := map[string]catalog.Model{} + for _, entry := range c.Models() { + m[entry.Name] = entry + for _, a := range entry.Aliases { + m[a] = entry + } + } + r.models.Store(&m) +} + // LastUsed reports when each endpoint was last sent a request, satisfying // control.UsageSource so the control plane can tell a busy replica from one // holding its cards out of habit. @@ -368,7 +408,7 @@ func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) { // passed through unchanged: an endpoint may serve a model the catalog does not // list, and refusing it would be unhelpful. func (r *Router) resolve(name string) (canonical string, m catalog.Model) { - if e, ok := r.models[name]; ok { + if e, ok := r.modelMap()[name]; ok { return e.Name, e } return name, catalog.Model{Name: name} diff --git a/test/e2e/phase3_test.go b/test/e2e/phase3_test.go index dfb1833..63d6466 100644 --- a/test/e2e/phase3_test.go +++ b/test/e2e/phase3_test.go @@ -186,7 +186,7 @@ func startControlPlane(t *testing.T, nodes int, minReplicas int) *controlHarness ccfg.Reconciler.SettleFor = 400 * time.Millisecond ccfg.PublishInterval = 100 * time.Millisecond - ctl := control.New(ccfg, cat, disp, rt, slog.New(slog.DiscardHandler)) + ctl := control.New(ccfg, catalog.NewLive(cat), disp, rt, slog.New(slog.DiscardHandler)) disp.onStatus = ctl.Reconciler().OnShardStatus // Every node is this host, so the ports the reconciler allocates are the