Files
llmgw/internal/cli/server.go
T
Ilker ManapandClaude Opus 5.5 964b659849 A stopping control plane gives up the lease before it lets its agents go
With the agent server closed rather than drained, the process could exit
before Run's release had reached the store, and the standby waited out the
lease. The lease is now given up first, synchronously, and nothing takes it
again after.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-05 02:31:13 +02:00

1052 lines
39 KiB
Go

package cli
import (
"context"
"crypto/rand"
"crypto/tls"
"encoding/hex"
"errors"
"fmt"
"log/slog"
"maps"
"net"
"net/http"
"os"
"os/signal"
"slices"
"strings"
"sync"
"syscall"
"time"
"gitea.mynodes.xyz/ilker/llmgw/internal/clientip"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
"gitea.mynodes.xyz/ilker/llmgw/internal/sourcecheck"
"github.com/prometheus/client_golang/prometheus"
"github.com/spf13/cobra"
"golang.org/x/net/http2"
"golang.org/x/net/http2/h2c"
"golang.org/x/net/netutil"
"golang.org/x/sync/errgroup"
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
"gitea.mynodes.xyz/ilker/llmgw/internal/audit"
"gitea.mynodes.xyz/ilker/llmgw/internal/bench/eval"
"gitea.mynodes.xyz/ilker/llmgw/internal/bundle"
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fabrictest"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/federation"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fleetmetrics"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
"gitea.mynodes.xyz/ilker/llmgw/internal/promq"
"gitea.mynodes.xyz/ilker/llmgw/internal/router"
"gitea.mynodes.xyz/ilker/llmgw/internal/secrets"
storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
"gitea.mynodes.xyz/ilker/llmgw/internal/web"
)
func newServerCmd() *cobra.Command {
var (
configPath string
listen string
adminListen string
metricsListen string
endpoints []string
)
cmd := &cobra.Command{
Use: "server",
Short: "Run the router (phase 1; the control plane arrives in phase 3)",
Long: "Serves the OpenAI API in front of a set of engine replicas.\n\n" +
"In this phase the replica list is static: give it in the config file, or\n" +
"pass --endpoint id=URL=model[,model] once per replica for a quick test.",
RunE: func(cmd *cobra.Command, _ []string) error {
log := logger()
cfg, err := config.Load(configPath)
if err != nil {
return err
}
if listen != "" {
cfg.Listen = listen
}
if adminListen != "" {
cfg.AdminListen = adminListen
}
if metricsListen != "" {
cfg.MetricsListen = metricsListen
}
for _, spec := range endpoints {
ep, err := parseEndpointFlag(spec)
if err != nil {
return err
}
cfg.Endpoints = append(cfg.Endpoints, ep)
}
if err := cfg.Validate(); err != nil {
return err
}
controlMode := cfg.ControlPlaneMode()
if !controlMode && len(cfg.Endpoints) == 0 {
return errors.New("nothing to serve: declare models with profiles, " +
"or give --endpoint for a router-only run")
}
// Offline, nothing is asked of the internet: releases come from
// the mirror's copy, and no registry is asked about newer weights.
if cfg.Offline.Enabled() {
cfg.Releases = cfg.ReleasesBase()
cfg.Mirror.UpdateCheck = 0
}
// Before the router, so its tracer is the configured one. The
// host tells two control planes' traces apart.
host, _ := os.Hostname()
shutdownTracing, err := obs.SetupTracing(cmd.Context(), obs.TracingOptions{
Endpoint: cfg.Tracing.Endpoint, SampleRatio: cfg.Tracing.SampleRatio,
Service: "llmgw-server", Version: version.Version, Instance: host,
})
if err != nil {
return err
}
defer func() {
// Bounded: a collector that is gone must not hold up a shutdown.
sctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_ = shutdownTracing(sctx)
}()
rt, err := router.New(cfg, log)
if err != nil {
return err
}
// The store's own secrets sealed under a key kept outside it,
// before anything reads them.
if st := rt.Store(); st != nil && cfg.Secrets.MasterKey != "" {
sealer, err := secrets.NewSealer(cmd.Context(), cfg.Secrets.MasterKey)
if err != nil {
return err
}
st.SetSealer(sealer)
}
// The longest any key may live, where every way of making one
// reads it: the store.
if st := rt.Store(); st != nil {
if err := st.SetMaxKeyTTL(cmd.Context(), storepkg.ScopeGateway, "", cfg.Auth.MaxKeyTTL); err != nil {
log.Warn("could not record auth.max_key_ttl", "err", err)
}
}
// Guardrails, when server.yaml defines any. Attachments to
// parts of the tenancy tree live in the store; without one, only
// guard.default applies.
var gate *guard.Gate
if len(cfg.Guard.Policies) > 0 {
var att guard.Attachments
if st := rt.Store(); st != nil {
att = st
}
if gate, err = guard.NewGate(cfg.Guard, att, log); err != nil {
return fmt.Errorf("guard: %w", err)
}
gate.SetJudge(rt.Judge, rt.GuardFormat)
rt.SetGuard(gate)
rt.SetGuardLibrary(cfg.Guard)
if err := rt.RegisterMetrics(gate); err != nil {
return fmt.Errorf("guard metrics: %w", err)
}
log.Info("guardrails on", "policies", len(cfg.Guard.Policies), "default", cfg.Guard.Default)
}
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
rt.Start(ctx)
defer rt.Stop()
// Which cluster this is, to federated peers: the id in the hop
// header that stops a request going round two clusters that
// federate each other, and the name in the tags a peer may be
// sent. The id is the store's, so every control plane of one
// cluster says the same.
clusterID, err := federationID(ctx, rt.Store())
if err != nil {
return err
}
clusterName := cfg.Federation.Name
if clusterName == "" {
clusterName = host
}
rt.SetFederationIdentity(clusterID, clusterName)
// Webhooks for what needs a person, when server.yaml names any.
notifier, err := notify.New(ctx, cfg.Notify, log)
if err != nil {
return err
}
rt.SetNotifier(notifier)
if gate != nil && notifier != nil {
gate.SetJudgeAlarm(func(model string, err error) {
notifier.Notify(notify.JudgeUnavailable, model,
fmt.Sprintf("the guardrails judge %s gives no verdict (%v): policies with fail: closed "+
"refuse what it should read, fail: open serve it unjudged", model, err),
map[string]any{"model": model, "error": err.Error()})
})
}
if notifier != nil {
log.Info("notifications on", "webhooks", len(cfg.Notify.Webhooks))
}
// The control plane replaces the static endpoint list: agents
// register, the reconciler places replicas, and Ready ones are
// published to the router (plan §6.4).
var agentHandler http.Handler
var agentPath string
var admin *control.Controller
updateMetrics := &modelserve.UpdateCollector{}
// The store's cluster, when it is one; read by the metrics and
// by `ctl status`.
var storeHA *pgha.Collector
if controlMode {
cat, err := cfg.Catalog()
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, live, rt, cfg.AgentTLS.Enabled(), log)
if err != nil {
return err
}
defer ctl.Close()
ctl.Reconciler().SetNotifier(notifier)
admin = ctl
agentPath, agentHandler = svc.Handler()
// Before AdminHandler, which is where the registry is built —
// registering after it silently publishes nothing, which is
// exactly what the first version of this did.
//
// On the same /metrics as the traffic figures: correlating
// "the cluster was full" with "requests were queueing" should
// not need two systems.
// When the newest backup was written, so its absence alerts
// rather than being found on the day it is needed.
backupDir := cfg.BackupDir
if backupDir == "" {
backupDir = DefaultBackupDir
}
if err := rt.RegisterMetrics(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
Name: "llmgw_backup_last_success_timestamp_seconds",
Help: "When the newest backup archive was written, as a Unix time; 0 when there is none.",
}, func() float64 {
if t := NewestBackup(backupDir); !t.IsZero() {
return float64(t.Unix())
}
return 0
})); err != nil {
return err
}
if err := rt.RegisterMetrics(tlsconf.ExpiryCollector(map[string]string{
"agent-server": cfg.AgentTLS.Cert, "agent-ca": cfg.AgentTLS.CA,
})); err != nil {
return err
}
if err := rt.RegisterMetrics(fleetmetrics.New(ctl)); err != nil {
return err
}
// Whether models' repositories have new weights; given its
// source once the mirror is set up, below.
if err := rt.RegisterMetrics(updateMetrics); err != nil {
return err
}
// A store of several hosts is a replicated cluster: what state
// it is in is published beside everything else, read from its
// Patroni. One host is a single PostgreSQL, which has none.
if addr, err := storepkg.ResolveAddress(cfg.Auth.Store); err == nil && storepkg.IsPostgres(addr) {
if hosts, _, _, err := storepkg.SplitPostgres(strings.TrimSpace(addr)); err == nil && len(hosts) > 1 {
if cfg.AgentTLS.CA != "" {
// Members' APIs, once on TLS, verified against it.
if err := pgha.SetRESTCA(cfg.AgentTLS.CA); err != nil {
log.Warn("the store's API cannot be verified", "err", err)
}
}
storeHA = pgha.NewCollector(hosts)
if err := rt.RegisterMetrics(storeHA); err != nil {
return err
}
}
}
// So a request for a model that is configured but not running
// starts it, rather than being told nothing serves it.
rt.SetActivator(ctl.Reconciler())
log.Info("control plane running",
"models", len(cfg.Models), "agent_listen", cfg.AgentListen,
"agent_mtls", cfg.AgentTLS.Enabled())
}
adminMux, err := rt.AdminHandler()
if err != nil {
return err
}
// Kept past this block: serving a model from the registry needs
// both the admin API and the mirror, and the mirror is set up
// below.
var (
adminAPI *adminapi.API
catalogSvc *modelsvc.Service
evals *eval.Service
)
if admin != nil {
api := adminapi.New(admin)
if storeHA != nil {
api.SetStoreStatus(func() ([]adminapi.StoreMember, error) {
c, err := storeHA.Read()
var out []adminapi.StoreMember
for _, m := range c.Members {
out = append(out, adminapi.StoreMember{Name: m.Name, Host: m.Host, Role: m.Role, State: m.State, Lag: m.Lag})
}
return out, err
})
}
if st := rt.Store(); st != nil {
api.SetControlPlanes(newControlPlaneProbe(st, cfg.Listen).Status)
}
adminAPI = api
api.SetReleases(cfg.Releases)
api.SetNotifier(notifier)
api.SetReplayer(rt.Replay)
// Quality suites, asked of a model the way a replay is.
evals = newEvalService(cfg, rt, admin, log)
api.SetEval(evals)
if st := rt.Store(); st != nil && gate != nil {
api.SetGuard(&adminapi.GuardAdmin{Library: cfg.Guard, Store: st, Invalidate: gate.Invalidate})
}
if st := rt.Store(); st != nil {
// Removing a node revokes its certificate.
api.SetRevoker(st.RevokeNode)
api.SetAudit(st)
}
// Catalog changes need somewhere to persist, so this is wired
// only when there is a store. Without one, the endpoints say so
// rather than accepting a model that vanishes on restart.
if st := rt.Store(); st != nil {
svc := modelsvc.New(cfg.Models, st, admin, log)
svc.SetCheck(servableCheck(cfg.WantedEngines()["llamacpp"]))
// Other llmgw clusters' models, offered as PEER/MODEL:
// read before the catalog is built, so a model that falls
// back to one is whole at startup.
fed := federation.New(federation.Config{Self: clusterID, Interval: cfg.Federation.Interval},
st, svc, admin.Catalog, log)
if err := fed.Load(ctx); err != nil {
log.Warn("could not read the federation peers; continuing without them", "err", err)
}
svc.SetFederated(fed.Models)
admin.SetPeers(fed)
api.SetFederation(fed)
if err := rt.RegisterMetrics(fed); err != nil {
return err
}
go fed.Run(ctx)
// Rebuild the catalog from configuration plus what was
// stored, and put that in force — otherwise models added
// before the last restart would be silently absent until
// someone added them again.
if built, err := svc.Load(ctx); err != nil {
log.Error("could not load stored models; continuing with configuration only", "err", err)
} else {
admin.ReplaceCatalog(ctx, built)
}
api.SetModels(svc)
catalogSvc = svc
// And the mirror page, so a model that has just been
// fetched can be served without writing YAML by hand.
rt.SetCatalogAdmin(svc)
// Benchmark results live here too. A measurement costs real
// GPU time; without somewhere to put it the number lives in
// whoever ran it's terminal and is gone by morning.
api.SetCapacity(st)
// And what each group of the tenancy tree spent, which is
// the other half of a quota being reviewable.
api.SetUsage(st)
// Capacity quotas, set and read over this API.
api.SetQuotas(st)
}
api.Register(adminMux)
// The dashboard lives here and only here. This listener is
// bound to loopback by default and also serves pprof, which can
// dump the process heap — where API keys sit in plaintext.
// Reaching the page from elsewhere is an SSH tunnel's job:
// ssh -L 9090:127.0.0.1:9090 <host>
var usage web.UsageSource
if st := rt.Store(); st != nil {
usage = st
}
// The dashboard's operational charts. llmgw publishes these
// metrics; what it cannot do is remember them, so the history
// is read back from wherever they were scraped to.
//
// Built before either copy of the page is registered, and given
// to both. The first version built it afterwards and handed it
// only to the one on the client port, so the copy an operator
// reaches by SSH tunnel — the one on this very listener — had
// no charts at all and said nothing about why.
var history web.History
if h := promq.New(cfg.Monitoring.Prometheus, cfg.Monitoring.Timeout); h != nil {
history = h
// And the console's charts, over the admin API.
if adminAPI != nil {
adminAPI.SetHistory(h)
}
log.Info("dashboard history configured",
"prometheus", cfg.Monitoring.Prometheus)
}
rt.SetHistory(history, cfg.Monitoring.Window)
adminMux.HandleFunc("GET /", web.Dashboard(
api, usage, version.Version, 5*time.Second, true,
web.NavExtras{History: history, Window: cfg.Monitoring.Window}, log))
adminMux.Handle("GET /_static/", web.Static())
adminMux.HandleFunc("GET /_lang/{code}", web.SwitchLanguage)
adminMux.HandleFunc("GET /nodes/{id}/logs", web.NodeLogsPage(api, true, version.Version))
adminMux.HandleFunc("GET /nodes/{id}", web.NodePage(api, history, cfg.Monitoring.Window, true, version.Version, log))
adminMux.HandleFunc("GET /nvlink", web.NVLinkPage(api, true, version.Version))
adminMux.HandleFunc("GET /fabrics", web.FabricsPage(api, true, version.Version))
adminMux.HandleFunc("GET /fabrics/{id}", web.FabricsPage(api, true, version.Version))
adminMux.HandleFunc("GET /fabrics/{id}/topology.json", web.FabricTopologyJSON(api))
// The same page on the client-facing port at /admin, behind a
// session and an admin flag. Handing the router the data rather
// than the port keeps pprof where it is.
rt.SetClusterView(api, usage)
}
// The model mirror, when one is configured. Built here and not at
// first use so a mistyped address stops the control plane starting,
// rather than turning into a broken admin page hours later when
// somebody finally opens it.
if cfg.Mirror.API != "" {
mc, err := mirror.NewClient(cfg.Mirror.API, cfg.Mirror.TLS, cfg.Mirror.Timeout)
if err != nil {
return fmt.Errorf("the model mirror: %w", err)
}
rt.SetMirror(mc)
// Browsing the registry is what turns "paste a resolve URL"
// into "find the model". It needs a way out to the internet
// from this machine, which the mirror's own fetching does not.
registry := hf.New(cfg.Mirror.Registry, cfg.Mirror.Timeout)
rt.SetRegistry(registry)
// One command from a registry repository to a replica that
// answers: `llmgw ctl model serve`.
if adminAPI != nil && catalogSvc != nil {
serve := modelserve.New(modelserve.Deps{
Registry: registry, Mirror: mc, Catalog: catalogSvc, Fleet: admin,
LlamaVersion: cfg.WantedEngines()["llamacpp"], Log: log,
// Updates: the catalog in force, the service that
// changes it, and what replaces replicas.
Current: func() []catalog.Model { return admin.Catalog().Models() },
Models: catalogSvc, Restart: admin.Reconciler().Restart,
Pool: adminAPI.PoolNodes, Notify: notifier,
// Canaries: how the router saw each side serve.
CanaryStats: rt.CanaryStats, ResetCanary: rt.ResetCanary,
// Updates and promotions held to a quality suite.
Eval: evals,
})
go serve.WatchCanaries(ctx, 30*time.Second)
adminAPI.SetServe(serve)
updateMetrics.Set(serve)
go serve.WatchUpdates(ctx, cfg.Mirror.UpdateCheck)
}
log.Info("model mirror configured",
"api", cfg.Mirror.API, "mtls", cfg.Mirror.TLS.Enabled())
}
apiSrv := withLimits(&http.Server{
Handler: rt.Handler(),
// Reading the request: long enough for a body up to
// max_body_bytes on a slow link, short enough that a client
// trickling one cannot hold a connection and its buffer.
ReadTimeout: 5 * time.Minute,
// No write timeout: responses stream for minutes and the proxy
// enforces its own per-phase deadlines.
})
if err := checkAdminListen(cfg); err != nil {
return err
}
if bad := clientip.SetTrusted(cfg.Router.TrustedProxies); len(bad) > 0 {
return fmt.Errorf("router.trusted_proxies: %q is not an address or a CIDR prefix", bad[0])
}
adminToken, err := ensureAdminToken(adminTokenPath(configPath, cfg))
if err != nil {
return err
}
adminHandler := auditAPI(rt, adminMux, log)
// With a rival, a standby's admin API is the leader's: what the
// console and ctl show lives in the leader's memory.
var forward *leaderForward
if admin != nil && admin.Leadership() != nil && rt.Store() != nil && cfg.AgentTLS.Enabled() {
self, _ := os.Hostname()
forward = &leaderForward{st: rt.Store(), lead: admin.Leadership(), self: self, tls: cfg.AgentTLS, log: log}
}
served := adminHandler
if forward != nil {
served = forward.standby(adminHandler)
}
adminSrv := withLimits(&http.Server{
Handler: router.AdminGuard(router.AdminAuth(adminToken, cfg.DebugPprof, served)),
ReadTimeout: time.Minute,
})
apiLn, err := net.Listen("tcp", cfg.Listen)
if err != nil {
return fmt.Errorf("listening on %s: %w", cfg.Listen, err)
}
if cfg.MaxConnections > 0 {
// Past this, new connections wait in the kernel's backlog
// rather than each taking a goroutine and its buffers.
apiLn = netutil.LimitListener(apiLn, cfg.MaxConnections)
}
adminLn, err := net.Listen("tcp", cfg.AdminListen)
if err != nil {
_ = apiLn.Close()
return fmt.Errorf("listening on %s: %w", cfg.AdminListen, err)
}
// A metrics-only listener, when one is configured. Started after
// AdminHandler, which is where the registry it serves is built.
var metricsSrv *http.Server
var metricsLn net.Listener
if cfg.MetricsListen != "" {
h := rt.MetricsHandler()
if h == nil {
_ = apiLn.Close()
_ = adminLn.Close()
return fmt.Errorf("metrics handler is not ready; this is a wiring fault, not a configuration one")
}
metricsLn, err = net.Listen("tcp", cfg.MetricsListen)
if err != nil {
_ = apiLn.Close()
_ = adminLn.Close()
return fmt.Errorf("listening on %s: %w", cfg.MetricsListen, err)
}
// Beside the metrics, the machines to scrape for their node
// exporters, for Prometheus' http_sd_configs: on this port
// because Prometheus already reaches it.
if adminAPI != nil {
mux := http.NewServeMux()
mux.HandleFunc("GET /prometheus/targets", adminAPI.PromSD)
mux.Handle("/", h)
h = mux
}
metricsSrv = withLimits(&http.Server{Handler: h, ReadTimeout: 30 * time.Second})
go func() {
if err := metricsSrv.Serve(metricsLn); err != nil && !errors.Is(err, http.ErrServerClosed) {
log.Error("metrics listener stopped", "err", err)
}
}()
}
log.Info("router listening",
"api", apiLn.Addr().String(),
"admin", adminLn.Addr().String(),
"metrics", metricsAddr(metricsLn),
"endpoints", len(cfg.Endpoints),
"default_profile", cfg.Router.DefaultProfile,
)
var agentSrv *http.Server
var agentLn net.Listener
if agentHandler != nil {
mux := http.NewServeMux()
mux.Handle(agentPath, agentHandler)
if forward != nil {
mux.Handle(forwardPrefix+"/", router.AdminGuard(forward.leaderSide(adminHandler)))
}
// No read timeout: an agent's stream is one request that
// lasts as long as the agent does.
agentSrv = withLimits(&http.Server{Handler: mux})
if !cfg.AgentTLS.Enabled() {
// Without TLS, agents speak plaintext HTTP/2, which Go will
// not do unless asked. Nothing then stops a node from
// claiming to be another one, so this belongs on a trusted
// network only.
agentSrv.Handler = h2c.NewHandler(mux, &http2.Server{})
}
agentAddr, aerr := agentListenAddr(cfg)
if aerr != nil {
_ = apiLn.Close()
_ = adminLn.Close()
return aerr
}
if !cfg.AgentTLS.Enabled() {
log.Warn("the agent port has no TLS (agent_insecure): agents are not authenticated and "+
"traffic is plaintext; use only on one machine or a trusted network", "listen", agentAddr)
}
if agentLn, err = net.Listen("tcp", agentAddr); err != nil {
_ = apiLn.Close()
_ = adminLn.Close()
return fmt.Errorf("listening on %s: %w", agentAddr, err)
}
if cfg.AgentTLS.Enabled() {
tc, terr := cfg.AgentTLS.ServerConfig()
if terr != nil {
_ = apiLn.Close()
_ = adminLn.Close()
_ = agentLn.Close()
return terr
}
// Serve on an already-wrapped listener rather than
// ServeTLS: agentSrv.TLSConfig stays nil, which is exactly
// the case where net/http sets up HTTP/2 for ALPN itself.
agentLn = tls.NewListener(agentLn, tc)
}
}
g, gctx := errgroup.WithContext(ctx)
g.Go(func() error { return serve(apiSrv, apiLn) })
g.Go(func() error { return serve(adminSrv, adminLn) })
if agentSrv != nil {
g.Go(func() error { return serve(agentSrv, agentLn) })
}
g.Go(func() error {
<-gctx.Done()
log.Info("shutting down")
sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
// The lease first: a standby polls for it, and every agent
// about to be let go will be refused by a standby that has
// not taken it yet.
if admin != nil && admin.Leadership() != nil {
admin.Leadership().StepDown()
}
if agentSrv != nil {
// Closed, not shut down gracefully: an agent's stream
// never goes idle, so Shutdown would wait out its whole
// timeout while every node stayed on a control plane
// that is leaving, and a standby could not be found
// until it gave up. Closed, they move at once.
_ = agentSrv.Close()
}
if metricsSrv != nil {
_ = metricsSrv.Shutdown(sctx)
}
_ = adminSrv.Shutdown(sctx)
return apiSrv.Shutdown(sctx)
})
return g.Wait()
},
}
f := cmd.Flags()
f.StringVar(&configPath, "config", "", "configuration file (optional)")
f.StringVar(&listen, "listen", "", "override the client-facing address")
f.StringVar(&adminListen, "admin-listen", "", "override the dashboard/pprof address")
f.StringVar(&metricsListen, "metrics-listen", "", "serve metrics only on this address, safe to expose")
f.StringArrayVar(&endpoints, "endpoint", nil, "static endpoint: id=URL=model[,model] (repeatable)")
cmd.AddCommand(newServerInitCmd(), newServerAddCmd())
cmd.AddCommand(newBackupCmd(), newRestoreCmd())
return cmd
}
// startControlPlane builds the control plane and its agent service.
//
// 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.Live,
sink control.EndpointSink, requireIdentity bool, log *slog.Logger,
) (*control.Controller, *agentsvc.Service, error) {
ccfg := control.DefaultConfig()
ccfg.Reconciler = reconciler.Config{
Interval: cfg.Control.Interval,
BaseBackoff: cfg.Control.BaseBackoff,
MaxBackoff: cfg.Control.MaxBackoff,
UnreachableAfter: cfg.Control.UnreachableAfter,
LostAfter: cfg.Control.LostAfter,
MgmtInterface: cfg.Control.MgmtInterface,
}
if cfg.Control.Timezone != "" {
loc, err := time.LoadLocation(cfg.Control.Timezone)
if err != nil {
return nil, nil, fmt.Errorf("control.timezone: %w", err)
}
ccfg.Reconciler.Location = loc
}
if len(cfg.Control.PortRange) == 2 {
ccfg.Reconciler.PortRange = [2]int{cfg.Control.PortRange[0], cfg.Control.PortRange[1]}
}
// Each replica's engine demands a key only the router has. The secret it
// is derived from lives in the store, so every control plane against it —
// and this one after a restart — computes the same key for a replica.
if h, ok := sink.(interface{ Store() *storepkg.DB }); ok && h.Store() != nil {
secret, err := h.Store().ClusterSecret(ctx, "engine-key")
if err != nil {
return nil, nil, fmt.Errorf("reading the engine key secret: %w", err)
}
ccfg.Reconciler.EngineKey = control.EngineKeys([]byte(secret))
} else {
log.Warn("engines start without a key: there is no store to keep the secret " +
"they are derived from, so anyone who can reach a node's ports reaches its models")
}
// The service needs the inventory and the status sink, both of which the
// controller owns; the controller needs the service as its dispatcher. The
// cycle is broken by giving the controller a dispatcher that forwards to
// the service once it exists.
fwd := &dispatcherRef{}
ctl := control.New(ccfg, cat, fwd, sink, log)
// The router is both where endpoints go and where activity comes back
// from. It only reports usage when it is the thing serving requests, which
// is why this is a type assertion rather than a parameter.
if u, ok := sink.(control.UsageSource); ok {
ctl.SetUsageSource(u)
}
// Power caps and estimates; Validate has already refused a bad one.
if pol, err := cfg.Facility.Policy(); err == nil {
ctl.Facility().SetPolicy(pol)
}
// With a shared store, two control planes can see each other, so this one
// competes rather than assuming it is alone. Without one there is nobody
// to compete with and it places unconditionally, exactly as before.
// The same shape as the usage assertion above, and for the same reason:
// the sink is the router when the router is serving, and nothing when a
// control plane runs on its own.
var lead *control.Leadership
if h, ok := sink.(interface{ Store() *storepkg.DB }); ok {
if st := h.Store(); st != nil {
lead = control.NewLeadership(st, cfg.Control.LeaseTTL, log)
if r, ok := sink.(interface{ SetLeader(func() bool) }); ok {
r.SetLeader(lead.IsLeader)
}
ctl.SetLeadership(lead)
// The tenancy tree's capacity quotas: what each part is
// guaranteed and may hold.
ctl.Reconciler().SetQuotas(control.QuotaPolicy(st, log))
// And what each model was measured to serve, for scaling and for
// choosing between profiles.
ctl.Reconciler().SetMeasurements(control.Measurements(st, log))
// The engine builds chosen per node on the cluster page.
ctl.SetEngineChoices(control.NewEngineChoices(ctx, st, log))
// Domains held for maintenance and labels set on nodes.
ctl.Facility().UseStore(ctx, st)
log.Info("competing to be the control plane that places work",
"owner", lead.Owner())
}
}
svc := agentsvc.New(ctl.Inventory(), ctl.Reconciler(), requireIdentity, log)
svc.SetEnginesFor(func(n inventory.Node) map[string]string {
return cfg.EnginesFor(n.ID, slices.Sorted(maps.Keys(n.GPUsByClass())), n.Labels)
})
// The KV cache events agents forward go to the router's precise index,
// when there is a router in this process to have one.
if k, ok := sink.(interface{ ApplyKV(kvevents.Update) }); ok {
svc.SetKVSink(k.ApplyKV)
}
log.Info("nodes install these engines", "engines", cfg.WantedEngines(), "overrides", len(cfg.EngineOverrides))
if cfg.Offline.Enabled() {
// From the mirror, each file held to the checksum the bundle gave,
// which this control plane keeps and the mirror has no say in.
trusted := bundle.NewTrusted(cfg.Offline.ChecksumsFile())
svc.SetOffline(func(engine, version string) (string, map[string]string) {
return bundle.EngineBase(cfg.Offline.Mirror, engine, version), trusted.For(engine, version)
})
if v, ok := cfg.WantedEngines()["llamacpp"]; ok && len(trusted.For("llamacpp", v)) == 0 {
log.Warn("offline, and no bundle imported here carries llama.cpp at this version: nodes cannot "+
"install it until one does (llmgw bundle create --engine llamacpp="+v+", then bundle import or trust)",
"version", v, "checksums", cfg.Offline.ChecksumsFile(), "err", trusted.Err())
}
log.Info("offline: nodes install engines from the mirror", "mirror", cfg.Offline.Mirror,
"releases", cfg.Releases)
}
if ch := ctl.EngineChoices(); ch != nil {
svc.SetEngineChoices(ch.For)
ch.OnChange(svc.Reprovision)
}
// A standby must refuse agents as well as refusing to place. Commands
// travel down the stream the agent opened, so a standby holding one is a
// node the leader cannot command — and the leader's placements would fail
// with "no agent connected" for a node that is plainly connected.
if lead != nil {
svc.SetLeadership(lead)
}
// The control planes agents may dial, from the store every control plane
// shares: told to each agent as it registers, and to all of them when
// `server add` records another.
if h, ok := sink.(interface{ Store() *storepkg.DB }); ok && h.Store() != nil {
st := h.Store()
svc.SetControlPlanes(func(ctx context.Context) ([]string, error) {
cps, err := st.ControlPlanes(ctx)
urls := make([]string, len(cps))
for i, c := range cps {
urls[i] = c.URL
}
return urls, err
})
go svc.WatchControlPlanes(ctx, 30*time.Second)
// A removed node's certificate is refused.
svc.SetRevocations(st.NodeRevokedAt)
}
fwd.set(svc)
// Every InfiniBand fabric mapped every five minutes, through one node.
go svc.ScanFabrics(ctx, 5*time.Minute)
ctl.Start(ctx)
return ctl, svc, nil
}
// auditAPI records every change made through the admin API, when there is a
// store to record it in.
func auditAPI(rt *router.Router, next http.Handler, log *slog.Logger) http.Handler {
st := rt.Store()
if st == nil {
return next
}
return audit.API(st, adminapi.Actor, log, next)
}
// metricsAddr renders the metrics listener for the startup line, saying so
// plainly when there is none rather than printing an empty string.
func metricsAddr(ln net.Listener) string {
if ln == nil {
return "off (admin listener only)"
}
return ln.Addr().String()
}
// dispatcherRef lets the controller be built before the service exists.
type dispatcherRef struct {
mu sync.RWMutex
d reconciler.Dispatcher
}
func (r *dispatcherRef) set(d reconciler.Dispatcher) {
r.mu.Lock()
r.d = d
r.mu.Unlock()
}
func (r *dispatcherRef) get() (reconciler.Dispatcher, error) {
r.mu.RLock()
defer r.mu.RUnlock()
if r.d == nil {
return nil, errors.New("the agent service is not ready yet")
}
return r.d, nil
}
// RemoveWeights forwards a capability beyond scheduling.
//
// It has to be spelled out here rather than inherited, because this type
// satisfies reconciler.Dispatcher by forwarding each method by hand — so
// anything the controller reaches by type assertion stops at this boundary
// unless it is forwarded too. A capability added to the service and not to
// this type is simply absent at runtime, with nothing at compile time to say
// so.
func (r *dispatcherRef) RemoveWeights(ctx context.Context, nodeID string, paths []string) error {
d, err := r.get()
if err != nil {
return err
}
rm, ok := d.(control.WeightRemover)
if !ok {
return errors.New("the agent service cannot remove weights")
}
return rm.RemoveWeights(ctx, nodeID, paths)
}
// Logs forwards a node log request, the same way RemoveWeights does: the
// controller finds it by type assertion, so a forwarder without it would make
// the capability silently absent at runtime.
func (r *dispatcherRef) Logs(ctx context.Context, nodeID, source string, lines int) ([]string, []string, error) {
d, err := r.get()
if err != nil {
return nil, nil, err
}
lr, ok := d.(control.LogReader)
if !ok {
return nil, nil, errors.New("this dispatcher cannot read node logs")
}
return lr.Logs(ctx, nodeID, source, lines)
}
func (r *dispatcherRef) FabricTest(ctx context.Context, nodeID string, req *llmgwv1.FabricTestRequest) (*llmgwv1.FabricTestReply, error) {
d, err := r.get()
if err != nil {
return nil, err
}
ft, ok := d.(fabrictest.Asker)
if !ok {
return nil, errors.New("this dispatcher cannot run a bandwidth test")
}
return ft.FabricTest(ctx, nodeID, req)
}
// Upgrade forwards a node upgrade, as Logs does.
func (r *dispatcherRef) Upgrade(ctx context.Context, nodeID, version, releases string) error {
d, err := r.get()
if err != nil {
return err
}
u, ok := d.(control.Upgrader)
if !ok {
return errors.New("this dispatcher cannot upgrade nodes")
}
return u.Upgrade(ctx, nodeID, version, releases)
}
func (r *dispatcherRef) StartShard(ctx context.Context, cmd reconciler.ShardCommand) error {
d, err := r.get()
if err != nil {
return err
}
return d.StartShard(ctx, cmd)
}
func (r *dispatcherRef) StopShard(ctx context.Context, nodeID, shardID string) error {
d, err := r.get()
if err != nil {
return err
}
return d.StopShard(ctx, nodeID, shardID)
}
func (r *dispatcherRef) DrainShard(ctx context.Context, nodeID, shardID string) error {
d, err := r.get()
if err != nil {
return err
}
return d.DrainShard(ctx, nodeID, shardID)
}
func serve(s *http.Server, ln net.Listener) error {
if err := s.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
return err
}
return nil
}
// parseEndpointFlag decodes id=URL=model[,model].
//
// The URL contains '=' only in pathological cases, so splitting into exactly
// three parts keeps the flag readable while staying unambiguous.
func parseEndpointFlag(spec string) (config.Endpoint, error) {
parts := strings.SplitN(spec, "=", 3)
if len(parts) != 3 {
return config.Endpoint{}, fmt.Errorf("--endpoint %q: expected id=URL=model[,model]", spec)
}
models := strings.Split(parts[2], ",")
for i := range models {
models[i] = strings.TrimSpace(models[i])
}
return config.Endpoint{
ID: strings.TrimSpace(parts[0]),
URL: strings.TrimRight(strings.TrimSpace(parts[1]), "/"),
Models: models,
Pool: "static",
Weight: 1,
}, nil
}
// servableCheck refuses a model whose llama.cpp weights the nodes' build
// cannot load, read from the weights' own header before anything is placed.
func servableCheck(llamaVersion string) func(context.Context, catalog.Model) error {
return func(ctx context.Context, m catalog.Model) error {
if llamaVersion == "" {
return nil
}
usesLlama := false
for _, p := range m.Profiles {
if p.Engine == "llamacpp" {
usesLlama = true
}
}
url, ok := sourcecheck.HTTPSource(m.Source)
if !usesLlama || !ok || !gguf.IsGGUF(url) {
return nil
}
if v := gguf.Check(ctx, url, llamaVersion); !v.Servable {
return fmt.Errorf("%s cannot be served: %s", m.Name, v.Reason)
}
return nil
}
}
// errAgentInsecure is a control plane asked to open its agent port with no
// TLS and not told that is meant.
var errAgentInsecure = errors.New("agent_tls is not set, so the agent port would be plaintext and anyone " +
"reaching it could register as a node: set agent_tls (llmgw server init writes it), or " +
"agent_insecure: true for one machine or a test")
// agentListenAddr is where the agent port listens: agent_listen with TLS;
// without it, refused unless agent_insecure says so, and then on loopback
// when agent_listen names only a port.
func agentListenAddr(cfg config.Server) (string, error) {
if cfg.AgentTLS.Enabled() {
return cfg.AgentListen, nil
}
// agent_tls.insecure says the same, in the older place.
if !cfg.AgentInsecure && !cfg.AgentTLS.Insecure {
return "", errAgentInsecure
}
if host, port, err := net.SplitHostPort(cfg.AgentListen); err == nil && host == "" {
return net.JoinHostPort("127.0.0.1", port), nil
}
return cfg.AgentListen, nil
}
// withLimits gives a server what every listener here wants: a bound on how
// long the headers may take, how long a connection may sit idle between
// requests, and how large the headers may be.
func withLimits(s *http.Server) *http.Server {
s.ReadHeaderTimeout = 10 * time.Second
s.IdleTimeout = 2 * time.Minute
s.MaxHeaderBytes = 1 << 20
return s
}
// federationID is this cluster's id among federated peers: kept in the store,
// so every control plane of the cluster and every restart says the same. A
// gateway with no store makes one for its lifetime, which is all a router
// with no peers of its own needs to refuse a request that comes back to it.
func federationID(ctx context.Context, st *storepkg.DB) (string, error) {
if st == nil {
b := make([]byte, 8)
if _, err := rand.Read(b); err != nil {
return "", err
}
return hex.EncodeToString(b), nil
}
secret, err := st.ClusterSecret(ctx, "cluster-id")
if err != nil {
return "", fmt.Errorf("this cluster's id for federation: %w", err)
}
// A prefix is plenty to tell clusters apart, and the id is sent to
// peers: nothing is derived from it, but there is no reason to send
// more of a secret than needed.
return secret[:16], nil
}