package cli import ( "context" "crypto/rand" "crypto/tls" "encoding/hex" "encoding/json" "errors" "fmt" "log/slog" "maps" "math" "net" "net/http" "os" "os/signal" "slices" "strings" "sync" "sync/atomic" "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/accounting" "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/held" "gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc" "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler" "gitea.mynodes.xyz/ilker/llmgw/internal/control/slurmsvc" "gitea.mynodes.xyz/ilker/llmgw/internal/control/train" "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" "gitea.mynodes.xyz/ilker/llmgw/internal/slurm" 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", Args: cobra.NoArgs, 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() // server.yaml for this machine's own keys, the store's settings // for the cluster's: the same on every control plane. boot, err := loadClusterSettings(cmd.Context(), configPath, log) if err != nil { return err } cfg := boot.cfg 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) }() // The ports before anything else: a second server on this machine, // whose ports are taken, must stop here, not go on to compete for // the lease and lead a cluster no request can reach. 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) } defer apiLn.Close() defer adminLn.Close() 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(gate.Library) 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() // Cancelled to start again on changed settings. ctx, cancelRun := context.WithCancel(ctx) defer cancelRun() var restartReason atomic.Pointer[string] rt.Start(ctx) defer rt.Stop() rt.SetStopped(cfg.Serving.Stopped) // 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 } if st := rt.Store(); st != nil { // Each control plane raises the same events; one sends each. notifier.SetClaimer(st) } 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). // The settings in force, for what reads them at each use. var current atomic.Pointer[config.Server] current.Store(&cfg) rt.SetSlurm(func() config.SlurmSettings { return current.Load().Slurm }) var agentSvc *agentsvc.Service var agentHandler http.Handler var agentPath string var admin *control.Controller // What the settings say of the firewall, for the mirror and the // store's members; nil without a control plane. var fwPub *firewallPublisher // The Slurm cluster beside llmgw, when there is a store (BACKLOG P). var slurmSvc *slurmsvc.Service 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() agentSvc = svc // Stopped by the operators (llmgw stop): nothing placed. ctl.Reconciler().SetStopped(cfg.Serving.Stopped) ctl.Reconciler().SetEvictionTimes(cfg.OnDemand.EvictAfterIdle, cfg.OnDemand.MinLifetime) // Whether nodes keep llmgw's firewall, as the settings say now. svc.SetFirewall(func() *llmgwv1.Firewall { c := current.Load() from, to := c.EnginePorts() return &llmgwv1.Firewall{Enabled: c.Firewall.Enabled, Networks: c.FirewallNetworks, EnginePortsFrom: int32(from), EnginePortsTo: int32(to)} //nolint:gosec // ports }) // What each node's agent runs on, as the settings say now. svc.SetAgentSettings(func(node string) string { b, err := config.Marshal(current.Load().Agents.For(node)) if err != nil { log.Warn("could not write a node's agent settings", "node", node, "err", err) return "" } return string(b) }) ctl.Reconciler().SetNotifier(notifier) admin = ctl agentPath, agentHandler = svc.Handler() fwPub = newFirewallPublisher(func() config.Server { return *current.Load() }, rt.Store(), func() []string { var out []string for _, n := range ctl.Inventory().Nodes() { out = append(out, n.MgmtAddr) } return out }, func() bool { return ctl.Leadership() == nil || ctl.Leadership().IsLeader() }, log) go fwPub.run(ctx, 5*time.Minute) if l := ctl.Leadership(); l != nil { l.OnLead(fwPub.Kick) } // 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.SetWebhooks(func() int { return len(current.Load().Notify.Webhooks) }) 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: gate.Library, Store: st, Invalidate: gate.Invalidate, JudgeErrors: gate.JudgeErrors, Ask: gate.AskJudge}) } if st := rt.Store(); st != nil { // Removing a node revokes its certificate. api.SetRevoker(st.RevokeNode) if agentSvc != nil { api.SetAgentRestarter(agentSvc.RestartAgents) } api.SetAudit(st) // Prices from what the cluster costs: what each model // holds, kept by the hour on the leader. api.SetPrices(st, func() adminapi.PriceConfig { au := current.Load().Auth cur := make(map[string]accounting.Price, len(au.Pricing)) for m, p := range au.Pricing { cur[m] = accountingPrice(p) } return adminapi.PriceConfig{Basis: accounting.Basis{GPUHour: au.GPUHourCost, CoreHour: au.CoreHourCost}, Current: cur, Default: accountingPrice(au.DefaultPrice)} }) go (&held.Recorder{Inv: admin.Inventory(), Store: st, Log: log, Models: func() map[string]string { m := map[string]string{} for _, rep := range admin.Reconciler().Replicas() { m[rep.ID] = rep.Model } return m }, Leader: func() bool { return admin.Leadership() == nil || admin.Leadership().IsLeader() }, }).Run(ctx) // Knowledge bases, with the cluster's own models to embed a // probe (a collection's dimension) and an operator's search. api.SetKB(&adminapi.KBAdmin{Store: st, Embed: embedder(rt.Replay)}) // LoRA adapters trained on the nodes. if agentSvc != nil { trainer := train.New(st, agentSvc, log) agentSvc.SetTrain(trainer) go trainer.Run(ctx) api.SetTrain(&adminapi.TrainAdmin{Store: st, Service: trainer}) // Slurm on the nodes the settings name, its munge key // made from a cluster secret. slurmSvc = slurmsvc.New(func() config.SlurmSettings { return current.Load().Slurm }, admin.Inventory(), agentSvc, func(ctx context.Context) ([]byte, error) { secret, err := st.ClusterSecret(ctx, "slurm-munge") if err != nil { return nil, err } return slurm.MungeKey(secret) }, log) // What slurm.conf held, for the leader that comes next. slurmSvc.SetMemory(slurmsvc.StoreMemory(st)) // slurm.users "llmgw": a Linux user for each active account. slurmSvc.SetUsers(func(ctx context.Context) ([]slurmsvc.User, error) { us, err := st.SlurmUsers(ctx) out := make([]slurmsvc.User, 0, len(us)) for _, u := range us { out = append(out, adminapi.SlurmUserView(u)) } return out, err }) api.SetSlurmUsers(st) // Slurm's accounts from the projects, and the jobs // that ended priced as the cluster's hours are. slurmSvc.SetAccounts(func(ctx context.Context) ([]slurm.Account, error) { return slurmAccounts(ctx, st) }) slurmSvc.SetJobs(func(ctx context.Context, jobs []slurmsvc.JobRecord) error { au := current.Load().Auth return recordSlurmJobs(ctx, st, current.Load().Slurm.Cluster(), au.GPUHourCost, au.CoreHourCost, jobs, log) }) agentSvc.SetSlurm(slurmSvc) go slurmSvc.Run(ctx) api.SetSlurm(slurmSvc) } // The cluster's settings, each change checked against this // machine's own keys before it is kept. api.SetMCP(st) rt.SetAdminMCP(http.HandlerFunc(api.MCP)) api.SetAuth(&adminapi.AuthAdmin{Status: rt.AuthStatus, TestLDAP: rt.TestDirectory, Secrets: st, Log: log}) api.SetSettings(&adminapi.Settings{Store: st, LocalKeys: config.LocalKeys(), Check: func(body []byte) error { return checkClusterSettings(configPath, body) }}) // And the same document on the portal, for administrators. rt.SetSettingsCheck(func(body []byte) error { return checkClusterSettings(configPath, body) }, config.LocalKeys()) } // 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"])) // What is added from the registry is pinned to the // commit it is at, with each file's sha256 for the nodes. svc.SetRegistry(hf.New(cfg.Mirror.Registry, 30*time.Second)) // 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) } // And again on becoming the leader: models added or // removed through the leader before are in the store. if l := admin.Leadership(); l != nil { l.OnLead(func() { if built, err := svc.Load(ctx); err != nil { log.Error("could not read the stored models on becoming the leader", "err", err) } else { admin.ReplaceCatalog(ctx, built) } }) } api.SetModels(svc) catalogSvc = svc // The models that come with llmgw: added by the leader // when a node fits one and the catalog makes nothing of // its kind, unless somebody removed it or turned it off. api.SetPresets(st, func() adminapi.PresetConfig { c := current.Load() return adminapi.PresetConfig{Auto: c.Presets.Auto, ComfyUI: c.Engines[catalog.EngineComfyUI]} }) go runPresetsPass(ctx, api, admin.Leadership(), log) // 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.SetTenancy(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 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 != "" { mirrors := append([]string{cfg.Mirror.API}, cfg.Mirror.Peers...) mc, err := mirror.NewFailoverClient(mirrors, cfg.Mirror.TLS, cfg.Mirror.Timeout) if err != nil { return fmt.Errorf("the model mirror: %w", err) } // Mirrors kept in step: a node is given a model's file on one // that answers while the one the catalog names is down. if len(cfg.Mirror.Peers) > 0 { src, err := mirror.NewSources(mirrors, cfg.Mirror.TLS, log) if err != nil { return fmt.Errorf("the model mirrors: %w", err) } go src.Run(ctx) agentsvc.SetSourceRewrite(src.Rewrite) } rt.SetMirror(mc) if adminAPI != nil { // A model removed with its files is removed from it too. adminAPI.SetPurgeMirror(mc) } if fwPub != nil { fwPub.setMirror(mc) } // Adapters trained on the nodes are kept there too; sending one // takes longer than the mirror's usual answers. if adminAPI != nil { if up, err := mirror.NewFailoverClient(mirrors, cfg.Mirror.TLS, 10*time.Minute); err == nil { adminAPI.SetTrainMirror(up) } } // 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, RestartInPlace: admin.Reconciler().RestartInPlace, 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) if gate != nil { // A judge a policy names and the catalog lacks is // fetched and served, on the leader. keeper := &judgeKeeper{ grace: judgeGrace, library: gate.Library, models: func() []catalog.Model { return admin.Catalog().Models() }, serve: serve, leader: func() bool { return admin.Leadership() == nil || admin.Leadership().IsLeader() }, gpus: func() []adminapi.ClassUse { return adminAPI.Summary().GPUs }, log: log, // An earlier install's judge, moved to today's figures. profiles: catalogSvc, restart: admin.Reconciler().Restart, replicas: func(model string) []string { var ids []string for _, r := range admin.Reconciler().ReadyReplicas() { if r.Model == model { ids = append(ids, r.ID) } } return ids }, nodeCPUs: func() []int { var out []int for _, n := range admin.Nodes() { out = append(out, n.CPUCount) } return out }, } go keeper.run(ctx, 30*time.Second) } 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} // And with router.active_active, a standby answers the client // API itself, to the leader's replicas. if cfg.Router.ActiveActive { var refresh func(context.Context) error if catalogSvc != nil { refresh = catalogSvc.Refresh } routes := newStandbyRoutes(forward, rt, refresh, admin.Publish, log) forward.routes = routes admin.SetReplicaFeed(routes) rt.SetActivator(standbyActivator{Reconciler: admin.Reconciler(), routes: routes}) go routes.run(ctx) go rt.ShareBudgets(ctx) log.Info("standbys answer the client API themselves (router.active_active)") } } served := adminHandler if forward != nil { served = forward.standby(adminHandler) _, port, _ := net.SplitHostPort(cfg.Listen) apiSrv.Handler = forward.public(apiSrv.Handler, port) go forward.trustControlPlanes(ctx, cfg.Router.TrustedProxies) } adminSrv := withLimits(&http.Server{ Handler: router.AdminGuard(router.AdminAuth(adminToken, cfg.DebugPprof, served)), ReadTimeout: time.Minute, }) // 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) } } // Every change to the settings in force: what can change live // does, the rest starts this process again, which systemd brings // back on them. // And this machine's own firewall, through its sync unit. if err := writeFirewallState(cfg); err != nil { log.Warn("could not record the firewall's settings for this machine", "err", err) } watch := &settingsWatch{path: configPath, st: rt.Store(), cur: boot, log: log, live: func(next config.Server, changed []string) bool { if !onlyLive(changed) { return false } current.Store(&next) if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "serving") }) { rt.SetStopped(next.Serving.Stopped) if admin != nil { admin.Reconciler().SetStopped(next.Serving.Stopped) } log.Warn("the service is " + map[bool]string{true: "stopped", false: "started"}[next.Serving.Stopped] + " by the cluster's settings") } if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "firewall") }) { if err := writeFirewallState(next); err != nil { log.Warn("could not record the firewall's settings for this machine", "err", err) } if agentSvc != nil { agentSvc.BroadcastFirewall() } if fwPub != nil { fwPub.Kick() } } if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "slurm") }) && slurmSvc != nil { slurmSvc.Kick() } if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "on_demand") }) && admin != nil { admin.Reconciler().SetEvictionTimes(next.OnDemand.EvictAfterIdle, next.OnDemand.MinLifetime) } if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "agents") }) && agentSvc != nil { // A moment apart per node, so off this goroutine. go agentSvc.BroadcastAgentSettings() } if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "guard") }) { if gate == nil { // Built at start, and only when there are policies. return len(next.Guard.Policies) == 0 } if err := gate.Replace(next.Guard); err != nil { log.Error("the changed guard policies are not valid; the old ones stay in force", "err", err) } } return true }, restart: func(reason string) { if os.Getenv("INVOCATION_ID") == "" { log.Warn("the settings changed in a way this process applies only when it starts; "+ "it is not run by systemd, so start it again by hand", "why", reason) return } // The leader last, so its standbys are already on the new // settings when it hands over to one of them. delay := time.Duration(0) if admin != nil && admin.Leadership() != nil && admin.Leadership().IsLeader() && rt.Store() != nil { if cps, err := rt.Store().ControlPlanes(ctx); err == nil && len(cps) > 1 { delay = 30 * time.Second } } log.Warn("starting again to apply the settings", "why", reason, "in", delay) restartReason.Store(&reason) time.AfterFunc(delay, cancelRun) }, } go watch.run(ctx) 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) }) err = g.Wait() if r := restartReason.Load(); r != nil { // Not a failure, but said as one, so systemd starts it again. return fmt.Errorf("stopped to start again on changed settings (%s)", *r) } return err }, } 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) // A standby read these when it started; the leader before it may // have changed them since. lead.OnLead(func() { if ch := ctl.EngineChoices(); ch != nil { ch.Reload(ctx) } ctl.Facility().Reload(ctx) }) 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 } // Each source a llama.cpp profile loads: its own, or the model's. for _, p := range m.Profiles { if p.Engine != "llamacpp" { continue } url, ok := sourcecheck.HTTPSource(m.SourceOf(p)) if !ok || !gguf.IsGGUF(url) { continue } 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 } // embedder has a model embed one text through the gateway itself, the way a // kept request is replayed: for the administrator's knowledge-base // endpoints, not for anybody's traffic. func embedder(replay func(ctx context.Context, path, model string, body []byte) (int, []byte, time.Duration, error)) func(context.Context, string, string) ([]float32, error) { return func(ctx context.Context, model, text string) ([]float32, error) { body, _ := json.Marshal(map[string]any{"model": model, "input": []string{text}}) status, out, _, err := replay(ctx, "/v1/embeddings", model, body) if err != nil { return nil, err } if status != http.StatusOK { msg := strings.TrimSpace(string(out)) var e struct { Error struct { Message string `json:"message"` } `json:"error"` } if json.Unmarshal(out, &e) == nil && e.Error.Message != "" { msg = e.Error.Message } return nil, fmt.Errorf("%d: %s", status, msg) } var res struct { Data []struct { Embedding []float32 `json:"embedding"` } `json:"data"` } if err := json.Unmarshal(out, &res); err != nil { return nil, err } if len(res.Data) != 1 || len(res.Data[0].Embedding) == 0 { return nil, fmt.Errorf("%s answered no embedding", model) } return res.Data[0].Embedding, nil } } // presetsEvery is how often the leader looks for presets to add; the first // look waits as long, for the nodes to have registered their cards. const presetsEvery = time.Minute // runPresetsPass has the leader add the presets a node fits and the catalog // lacks, now and then. func runPresetsPass(ctx context.Context, api *adminapi.API, lead *control.Leadership, log *slog.Logger) { t := time.NewTicker(presetsEvery) defer t.Stop() for { select { case <-ctx.Done(): return case <-t.C: } if lead != nil && !lead.IsLeader() { continue } added, err := api.PresetsPass(ctx) for _, name := range added { log.Info("added a preset by itself: a node fits it and the catalog made nothing of its kind", "model", name, "off_switch", "presets.auto") } if err != nil { log.Warn("could not add a preset", "err", err) } } } // accountingPrice is a configured price as accounting has it. func accountingPrice(p config.Price) accounting.Price { return accounting.Price{PromptPer1k: p.PromptPer1k, CompletionPer1k: p.CompletionPer1k, CachedPer1k: p.CachedPer1k, PerImage: p.PerImage, PerMesh: p.PerMesh} } // slurmAccounts are Slurm's accounts from the store's projects, members and // Slurm users. func slurmAccounts(ctx context.Context, st *storepkg.DB) ([]slurm.Account, error) { ps, err := st.SlurmProjects(ctx) if err != nil { return nil, err } ms, err := st.AllMemberships(ctx) if err != nil { return nil, err } us, err := st.SlurmUsers(ctx) if err != nil { return nil, err } logins := make(map[string]string, len(us)) for _, u := range us { logins[u.UserID] = u.Login } tps := make([]slurm.TenancyProject, 0, len(ps)) for _, p := range ps { tps = append(tps, slurm.TenancyProject{ID: p.ID, Name: p.Name, TeamID: p.TeamID, TeamName: p.TeamName, OrgID: p.OrgID, OrgName: p.OrgName}) } tms := make([]slurm.TenancyMember, 0, len(ms)) for _, m := range ms { tms = append(tms, slurm.TenancyMember{UserID: m.UserID, Scope: string(m.Scope), ScopeID: m.ScopeID}) } return slurm.Accounts(tps, tms, logins), nil } // recordSlurmJobs records the jobs that ended, each with its project (from // its account) and person (from its login), priced as the hours the // cluster's models hold are: a GPU job its cards' hours, any other its // cores'. func recordSlurmJobs(ctx context.Context, st *storepkg.DB, cluster string, gpuHour, coreHour float64, jobs []slurmsvc.JobRecord, log *slog.Logger) error { as, err := slurmAccounts(ctx, st) if err != nil { return err } project := map[string]string{} for _, a := range as { if a.Project != "" { project[a.Name] = a.Project } } us, err := st.SlurmUsers(ctx) if err != nil { return err } person := map[string]string{} for _, u := range us { person[u.Login] = u.UserID } rows := make([]storepkg.SlurmJob, 0, len(jobs)) for _, j := range jobs { secs := j.Elapsed.Seconds() cost := float64(j.CPUs) * secs / 3600 * coreHour if j.GPUs > 0 { cost = float64(j.GPUs) * secs / 3600 * gpuHour } rows = append(rows, storepkg.SlurmJob{Key: cluster + "/" + j.ID, UserID: person[j.Login], ProjectID: project[j.Account], Login: j.Login, Account: j.Account, Partition: j.Partition, State: j.State, Start: j.Start, End: j.End, ElapsedSeconds: int64(secs), CPUs: int64(j.CPUs), GPUs: int64(j.GPUs), CostMicros: int64(math.Ceil(cost * 1e6))}) } n, err := st.RecordSlurmJobs(ctx, rows) if n > 0 { log.Info("Slurm jobs that ended recorded, with their cost", "jobs", n) } return err }