Files
llmgw/internal/cli/server.go
T

881 lines
32 KiB
Go

package cli
import (
"context"
"crypto/tls"
"errors"
"fmt"
"log/slog"
"maps"
"net"
"net/http"
"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/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/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/mirror"
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
"gitea.mynodes.xyz/ilker/llmgw/internal/promq"
"gitea.mynodes.xyz/ilker/llmgw/internal/router"
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")
}
// Before the router, so its tracer is the configured one.
shutdownTracing, err := obs.SetupTracing(cmd.Context(), obs.TracingOptions{
Endpoint: cfg.Tracing.Endpoint, SampleRatio: cfg.Tracing.SampleRatio,
Service: "llmgw-server", Version: version.Version,
})
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 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()
// 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()
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
)
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
})
}
adminAPI = api
api.SetReleases(cfg.Releases)
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"]))
// 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 /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,
})
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
}
adminSrv := withLimits(&http.Server{
Handler: router.AdminGuard(router.AdminAuth(adminToken, cfg.DebugPprof, auditAPI(rt, adminMux, log))),
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)
// 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()
if agentSrv != nil {
_ = agentSrv.Shutdown(sctx)
}
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())
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 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)
}
// 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))
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)
})
log.Info("nodes install these engines", "engines", cfg.WantedEngines(), "overrides", len(cfg.EngineOverrides))
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
}