auth.ldap checks a username and password against the domain controllers over LDAPS or StartTLS: found with a service account, checked by binding as the person. Groups, nested ones included, give staff roles and places in the tree. Accounts are tied to the objectGUID; the leader reads the directory every 15 minutes, suspends whom it disabled or removed and reinstates them after, and refuses a sync that would suspend most. Entra ID sends no groups for someone in more than 200, only where to fetch them; that was read as none and took an administrator's role away. It is now left as it was, and logged. An account made by SCIM or OIDC can be reinstated again. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1346 lines
47 KiB
Go
1346 lines
47 KiB
Go
// Package router wires the pieces into a working OpenAI-compatible gateway:
|
|
// resolve the model, pick an endpoint, proxy the stream, retry when it is safe.
|
|
//
|
|
// The sub-packages hold the reusable parts (scheduler, endpoint, scrape, proxy,
|
|
// breaker); this package holds the orchestration and the HTTP surface.
|
|
package router
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/mail"
|
|
"log/slog"
|
|
"math"
|
|
"net"
|
|
"net/http"
|
|
"net/netip"
|
|
"slices"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"go.opentelemetry.io/otel/attribute"
|
|
"go.opentelemetry.io/otel/trace"
|
|
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/account"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/accounting"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/comfy"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/admission"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape"
|
|
storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/web"
|
|
)
|
|
|
|
// Router serves the OpenAI API in front of a set of engine replicas.
|
|
type Router struct {
|
|
cfg config.Router
|
|
log *slog.Logger
|
|
|
|
// selfHandler is the client API as Handler built it, key check and all:
|
|
// what the MCP tools make their requests through, as the caller.
|
|
selfHandler atomic.Pointer[http.Handler]
|
|
|
|
reg *endpoint.Registry
|
|
poller *scrape.Poller
|
|
px *proxy.Proxy
|
|
// comfy asks ComfyUI replicas for pictures.
|
|
comfy *comfy.Client
|
|
m *metrics
|
|
// adm is nil when admission control is disabled.
|
|
adm *admission.Controller
|
|
|
|
profiles map[string]*scheduler.Profile
|
|
// models maps every name a client may use — canonical names and aliases —
|
|
// onto its catalog entry.
|
|
//
|
|
// Held atomically because the catalog is no longer fixed at startup: adding
|
|
// a model used to mean restarting the control plane, and this is the
|
|
// router's half of not having to. A request takes the map once and works
|
|
// from it, so a swap mid-request is invisible rather than a race.
|
|
models atomic.Pointer[map[string]catalog.Model]
|
|
// adapters are the LoRA adapters by the name clients ask for them by.
|
|
adapters atomic.Pointer[map[string]adapterRef]
|
|
// canaryTargets maps each model that is another's canary to that one;
|
|
// canaries counts how the two sides served.
|
|
canaryTargets atomic.Pointer[map[string]string]
|
|
canaries canaryStats
|
|
// slo keeps each model's recent token timings, for admission.
|
|
slo sloTracker
|
|
// batchWake has the batch worker look for work at once; jobWake the
|
|
// job worker.
|
|
batchWake chan struct{}
|
|
jobWake chan struct{}
|
|
// reqLogs is each project's keeping of requests, as last read.
|
|
reqLogs requestLogs
|
|
// answers are the answers projects reuse; nil reuses none.
|
|
answers *answerCache
|
|
|
|
// extraMetrics are collectors other components publish on /metrics.
|
|
extraMetrics []prometheus.Collector
|
|
|
|
// activator loads a model that is configured but not currently running.
|
|
activator Activator
|
|
|
|
// metricsHandler is built with the admin registry and reused by any
|
|
// listener that should serve metrics without serving pprof.
|
|
metricsHandler http.Handler
|
|
|
|
// landing configures the public page, which the router serves but does not
|
|
// otherwise care about.
|
|
landing config.Landing
|
|
|
|
// auth, store, acct, recorder and accounts are nil when authentication is
|
|
// disabled.
|
|
auth *auth.Authenticator
|
|
accounts *account.Service
|
|
store *storepkg.DB
|
|
// usageRetention is how long request records are kept; zero is for ever.
|
|
usageRetention time.Duration
|
|
// isLeader says whether this process leads; nil is always.
|
|
isLeader atomic.Pointer[func() bool]
|
|
// federation is which cluster this is, to its peers; nil until said.
|
|
federation atomic.Pointer[federationID]
|
|
acct *accounting.Accountant
|
|
// notifier says what the operator must act on; nil says nothing.
|
|
notifier *notify.Notifier
|
|
// directory is auth.ldap's: sign-in and the accounts kept in step.
|
|
directory *ldapDirectory
|
|
// guard judges requests and answers; nil judges nothing. guardLibrary
|
|
// is the policies it was built from, for the pages that show them.
|
|
guard guard.Guard
|
|
guardLibrary func() guard.Config
|
|
// oidc is sign-in through an identity provider, and scim provisioning
|
|
// by it.
|
|
oidc config.OIDC
|
|
// mail sends what the portal mails: a password reset link.
|
|
mail mail.Config
|
|
// stopped is the operators having stopped the service (llmgw stop).
|
|
stopped atomic.Bool
|
|
scim config.SCIM
|
|
|
|
// threads holds conversations and memory notes. It is store when auth is on,
|
|
// its own file when only the threads block is configured, and nil when
|
|
// nothing is kept. Separate from store because a non-nil store also switches
|
|
// on the accounts, web pages and usage recording, none of which belong to a
|
|
// gateway that only wants to remember things.
|
|
threads *storepkg.DB
|
|
// anonymous means threads has no project to scope to: everything shares the
|
|
// one pool named by anonymousProject.
|
|
anonymous bool
|
|
// personalOff refuses personal use: the chat page then only offers the
|
|
// projects a person is a member of.
|
|
personalOff bool
|
|
// personalDefaults is what personal use is held to when nobody wrote a
|
|
// limit for the person, for the quota page to say so.
|
|
personalDefaults *storepkg.Limits
|
|
// llamaVersion is the llama.cpp build the nodes install, so the mirror page
|
|
// can refuse weights that build cannot load before fetching them.
|
|
llamaVersion string
|
|
// ownsThreads is true when threads is not store and so must be closed here.
|
|
ownsThreads bool
|
|
recorder *accounting.Recorder
|
|
|
|
// prefix scores by prefix: the precise index for models with kv_events,
|
|
// the approximate tree when router.prefix is on. Its tree is nil when it
|
|
// is off.
|
|
prefix *scheduler.PrefixScorer
|
|
// lora knows which replicas hold which adapters that load on demand.
|
|
lora *loraTracker
|
|
// kv is the precise index (plan §7.5 tier 2), fed by the control plane
|
|
// from the agents' KV cache events; kvClient makes its tokenize calls.
|
|
kv *kvindex.Index
|
|
kvClient *http.Client
|
|
|
|
// clusterView is nil until the control plane hands it over, and stays nil
|
|
// in a router-only deployment.
|
|
clusterView web.Source
|
|
clusterUsage web.UsageSource
|
|
|
|
// mirror is nil unless a model mirror is configured, which leaves
|
|
// /admin/mirror unmounted — the right page for a deployment with no mirror
|
|
// is no page, not an empty one pointing at nothing.
|
|
mirror web.MirrorAdmin
|
|
// registry browses a model registry from the mirror page. Nil where the
|
|
// deployment has no way out to the internet.
|
|
registry web.Registry
|
|
// catalogAdmin puts a mirrored model into service. Nil without a store,
|
|
// where a model added at runtime would not survive a restart.
|
|
catalogAdmin web.CatalogAdmin
|
|
// history is the metrics store the dashboard charts read from. Nil where
|
|
// none is configured, which leaves the charts off the page.
|
|
history web.History
|
|
histWindow time.Duration
|
|
|
|
// kbDefCache remembers each caller's default knowledge-base
|
|
// collections for kbDefaultTTL.
|
|
kbDefMu sync.Mutex
|
|
kbDefCache map[string]kbDefaultEntry
|
|
|
|
// settingsCheck and settingsLocal mount the portal's settings page.
|
|
settingsCheck func([]byte) error
|
|
settingsLocal []string
|
|
}
|
|
|
|
// SetClusterView gives the router the data the /admin page renders.
|
|
//
|
|
// Set by the server command once the control plane exists, because the router
|
|
// is built first. Nil leaves /admin unmounted, which is what a router-only
|
|
// deployment should have: no cluster to show.
|
|
func (r *Router) SetClusterView(src web.Source, usage web.UsageSource) {
|
|
r.clusterView, r.clusterUsage = src, usage
|
|
}
|
|
|
|
// SetMirror gives the router a model mirror to administer.
|
|
//
|
|
// Set by the server command when one is configured, for the same reason as the
|
|
// cluster view: the router is built before the things its admin pages show.
|
|
func (r *Router) SetMirror(m web.MirrorAdmin) { r.mirror = m }
|
|
|
|
// SetGuard has every request and answer judged by g.
|
|
func (r *Router) SetGuard(g guard.Guard) { r.guard = g }
|
|
|
|
// SetGuardLibrary is the policies the guard was built from, for the pages.
|
|
func (r *Router) SetGuardLibrary(c func() guard.Config) { r.guardLibrary = c }
|
|
|
|
// SetSettingsCheck enables the portal's settings page: check says whether a
|
|
// document would load as the cluster's settings beside this machine's own
|
|
// server.yaml, and local are the keys that file keeps.
|
|
func (r *Router) SetSettingsCheck(check func([]byte) error, local []string) {
|
|
r.settingsCheck, r.settingsLocal = check, local
|
|
}
|
|
|
|
// SetStopped refuses, or again serves, every request for a model.
|
|
func (r *Router) SetStopped(stopped bool) { r.stopped.Store(stopped) }
|
|
|
|
// SetRegistry gives the mirror page somewhere to look models up.
|
|
func (r *Router) SetRegistry(reg web.Registry) { r.registry = reg }
|
|
|
|
// SetCatalogAdmin lets the mirror page put a model into service.
|
|
func (r *Router) SetCatalogAdmin(c web.CatalogAdmin) { r.catalogAdmin = c }
|
|
|
|
// SetHistory gives the dashboard somewhere to read the cluster's past.
|
|
func (r *Router) SetHistory(h web.History, window time.Duration) {
|
|
r.history, r.histWindow = h, window
|
|
}
|
|
|
|
// Auth exposes the authenticator so ctl and the batch writer can reach it.
|
|
// Nil when authentication is off.
|
|
func (r *Router) Auth() *auth.Authenticator { return r.auth }
|
|
|
|
// Store exposes the key store, or nil when authentication is off.
|
|
func (r *Router) Store() *storepkg.DB { return r.store }
|
|
|
|
// New builds a router from configuration. Start must be called to begin scraping.
|
|
func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
|
if log == nil {
|
|
log = slog.Default()
|
|
}
|
|
|
|
reg := endpoint.NewRegistry(endpoint.Config{
|
|
StaleAfter: cfg.Router.StaleAfter,
|
|
Breaker: breaker.Config{
|
|
FailureThreshold: cfg.Router.FailureThreshold,
|
|
BaseCooldown: cfg.Router.Cooldown,
|
|
},
|
|
})
|
|
reg.Replace(toInfos(cfg.Endpoints))
|
|
|
|
sc := scheduler.ProfileConfig{
|
|
TargetWaiting: cfg.Router.Scoring.TargetWaiting,
|
|
TargetInflight: cfg.Router.Scoring.TargetInflight,
|
|
KVCeiling: cfg.Router.Scoring.KVCeiling,
|
|
SessionLoadFactor: cfg.Router.Scoring.SessionLoadFactor,
|
|
}
|
|
// The metrics come first because the prefix scorer reports into them, and
|
|
// the decision metric is the whole difference between tuning this and
|
|
// guessing at it.
|
|
metrics := newMetrics()
|
|
metrics.tenant.level = cfg.Router.TenantLabel
|
|
metrics.tenant.setTagLabels(cfg.Router.TagLabels)
|
|
// The scorer exists whether or not the tree does: which models use the
|
|
// precise index is the catalog's to say, and the catalog changes while
|
|
// the router runs. With neither in use it scores nothing and records
|
|
// nothing.
|
|
var tree *prefix.Tree
|
|
if cfg.Router.Prefix.Enabled {
|
|
tree = prefix.New(prefix.Config{
|
|
MaxNodes: cfg.Router.Prefix.MaxNodes,
|
|
MinMatch: cfg.Router.Prefix.MinMatch,
|
|
TTL: cfg.Router.Prefix.TTL,
|
|
})
|
|
}
|
|
minBlocks := cfg.Router.KVIndex.MinBlocks
|
|
if minBlocks <= 0 {
|
|
minBlocks = 4
|
|
}
|
|
prefixScorer := scheduler.NewPrefixScorer(tree,
|
|
prefix.Thresholds{
|
|
MatchRatio: cfg.Router.Prefix.MatchRatio,
|
|
BalanceAbs: cfg.Router.Prefix.BalanceAbs,
|
|
BalanceRel: cfg.Router.Prefix.BalanceRel,
|
|
},
|
|
metrics.observePrefix).WithKV(minBlocks, metrics.observeKV)
|
|
sc.Prefix = prefixScorer
|
|
|
|
profiles := map[string]*scheduler.Profile{}
|
|
for _, name := range []string{"chat", "batch", "long-context", "debug"} {
|
|
profiles[name] = scheduler.Builtin(name, sc)
|
|
}
|
|
if _, ok := profiles[cfg.Router.DefaultProfile]; !ok {
|
|
return nil, fmt.Errorf("unknown default_profile %q", cfg.Router.DefaultProfile)
|
|
}
|
|
|
|
// The portal's language for whoever neither chose one nor asked for one
|
|
// it speaks. Checked here so a mistyped code stops the server starting
|
|
// rather than quietly serving English.
|
|
if err := web.SetDefaultLanguage(cfg.Landing.Language); err != nil {
|
|
return nil, fmt.Errorf("landing.language: %w", err)
|
|
}
|
|
|
|
r := &Router{
|
|
batchWake: make(chan struct{}, 1),
|
|
jobWake: make(chan struct{}, 1),
|
|
cfg: cfg.Router,
|
|
landing: cfg.Landing,
|
|
log: log,
|
|
reg: reg,
|
|
profiles: profiles,
|
|
prefix: prefixScorer,
|
|
lora: newLoraTracker(log, metrics.observeAdapter),
|
|
kv: kvindex.New(),
|
|
kvClient: newKVClient(),
|
|
m: metrics,
|
|
comfy: &comfy.Client{HTTP: &http.Client{Transport: &http.Transport{DialContext: (&net.Dialer{Timeout: cfg.Router.Timeouts.Dial}).DialContext}}},
|
|
px: proxy.New(proxy.Config{
|
|
DialTimeout: cfg.Router.Timeouts.Dial,
|
|
ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader,
|
|
GenerationTimeout: cfg.Router.Timeouts.Generation,
|
|
FirstTokenTimeout: cfg.Router.Timeouts.FirstToken,
|
|
IdleTimeout: cfg.Router.Timeouts.Idle,
|
|
TotalTimeout: cfg.Router.Timeouts.Total,
|
|
MaxResponseBytes: cfg.Router.MaxBodyBytes,
|
|
}),
|
|
}
|
|
r.answers = newAnswerCache(cfg.Router.Cache.MaxBytes, cfg.Router.Cache.MaxEntryBytes,
|
|
func(bytes int64, n int) {
|
|
metrics.cacheBytes.Set(float64(bytes))
|
|
metrics.cacheEntries.Set(float64(n))
|
|
})
|
|
r.poller = scrape.NewPoller(reg, scrape.Config{
|
|
Interval: cfg.Router.ScrapeInterval,
|
|
Timeout: cfg.Router.ScrapeTimeout,
|
|
}, log)
|
|
|
|
if !cfg.Auth.Enabled && !loopbackOnly(cfg.Listen) {
|
|
log.Warn("auth is off and the API listens beyond this machine: anyone who can reach "+
|
|
cfg.Listen+" can use every model; set auth.enabled, or listen on 127.0.0.1",
|
|
"listen", cfg.Listen)
|
|
}
|
|
if cfg.Auth.Enabled {
|
|
st, err := storepkg.Open(context.Background(), cfg.Auth.Store)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("auth store: %w", err)
|
|
}
|
|
r.store = st
|
|
r.usageRetention = cfg.Auth.UsageRetention
|
|
// Applied where no scope set one, so a project created without limits
|
|
// is still metered. Zero leaves it unlimited, which was the old
|
|
// behaviour and is now a choice rather than an omission.
|
|
var defRPM, defTPM *int
|
|
if cfg.Auth.DefaultRPM > 0 {
|
|
defRPM = &cfg.Auth.DefaultRPM
|
|
}
|
|
if cfg.Auth.DefaultTPM > 0 {
|
|
defTPM = &cfg.Auth.DefaultTPM
|
|
}
|
|
st.SetDefaultLimits(defRPM, defTPM)
|
|
st.SetPersonalDefaults(personalLimits(cfg.Auth.Personal))
|
|
r.personalOff = cfg.Auth.Personal.Disabled
|
|
r.oidc = cfg.Auth.OIDC
|
|
r.mail = cfg.Mail
|
|
r.scim = cfg.Auth.SCIM
|
|
r.personalDefaults = personalLimits(cfg.Auth.Personal)
|
|
r.llamaVersion = cfg.WantedEngines()["llamacpp"]
|
|
|
|
r.auth = auth.New(st, auth.Config{
|
|
TTL: cfg.Auth.CacheTTL,
|
|
NegativeTTL: cfg.Auth.NegativeTTL,
|
|
})
|
|
|
|
prices := make(map[string]accounting.Price, len(cfg.Auth.Pricing))
|
|
for name, p := range cfg.Auth.Pricing {
|
|
prices[name] = accounting.Price{
|
|
PromptPer1k: p.PromptPer1k, CompletionPer1k: p.CompletionPer1k, PerImage: p.PerImage, PerMesh: p.PerMesh,
|
|
}
|
|
}
|
|
r.recorder = accounting.NewRecorder(st, r.auth, accounting.RecorderConfig{
|
|
FlushInterval: cfg.Auth.FlushInterval,
|
|
BatchSize: cfg.Auth.BatchSize,
|
|
Capacity: cfg.Auth.BufferSize,
|
|
}, log)
|
|
r.accounts = account.New(st, account.Config{
|
|
SignupOpen: cfg.Auth.SignupOpen,
|
|
}, log)
|
|
if cfg.Auth.LDAP.Enabled() {
|
|
d, err := newLDAPDirectory(cfg.Auth.LDAP, st, log)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
r.directory = d
|
|
r.accounts.SetDirectory(d)
|
|
}
|
|
|
|
r.acct = accounting.New(
|
|
accounting.NewPricing(prices, accounting.Price{
|
|
PromptPer1k: cfg.Auth.DefaultPrice.PromptPer1k,
|
|
CompletionPer1k: cfg.Auth.DefaultPrice.CompletionPer1k,
|
|
PerImage: cfg.Auth.DefaultPrice.PerImage,
|
|
PerMesh: cfg.Auth.DefaultPrice.PerMesh,
|
|
}),
|
|
r.recorder, log)
|
|
r.threads = r.store
|
|
} else if cfg.Threads.Store != "" {
|
|
st, err := storepkg.Open(context.Background(), cfg.Threads.Store)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("threads store: %w", err)
|
|
}
|
|
if err := st.EnsureAnonymousProject(context.Background(), anonymousProject, cfg.Threads.RetainHours); err != nil {
|
|
_ = st.Close()
|
|
return nil, fmt.Errorf("threads store: %w", err)
|
|
}
|
|
r.threads, r.anonymous, r.ownsThreads = st, true, true
|
|
}
|
|
|
|
if cfg.Router.Admission.Enabled {
|
|
adm, err := admission.New(admissionConfig(cfg.Router), r.saturation)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("admission: %w", err)
|
|
}
|
|
r.adm = adm
|
|
}
|
|
|
|
// Seeded from configuration, not only pushed. A router running without a
|
|
// control plane — a static endpoint list, which is a supported mode — has
|
|
// nothing to push it a catalog, and would otherwise resolve no model at
|
|
// all. A control plane pushes over the top of this at startup and on every
|
|
// change after.
|
|
if len(cfg.Models) > 0 {
|
|
r.storeModels(cfg.Models)
|
|
}
|
|
return r, nil
|
|
}
|
|
|
|
func admissionConfig(rc config.Router) admission.Config {
|
|
bands := make([]admission.Band, 0, len(rc.Admission.Bands))
|
|
for _, b := range rc.Admission.Bands {
|
|
bands = append(bands, admission.Band{
|
|
Name: b.Name,
|
|
Priority: b.Priority,
|
|
SaturationCeiling: b.SaturationCeiling,
|
|
QueueTTL: b.QueueTTL,
|
|
})
|
|
}
|
|
return admission.Config{
|
|
Bands: bands,
|
|
DefaultBand: rc.Admission.DefaultBand,
|
|
Thresholds: admission.Thresholds{
|
|
// Admission and the queue scorer must agree on what a healthy
|
|
// buffer is, or the two layers fight each other.
|
|
TargetWaitingPerReplica: rc.Scoring.TargetWaiting,
|
|
KVCeiling: rc.Scoring.KVCeiling,
|
|
},
|
|
QueueTTL: rc.Admission.QueueTTL,
|
|
QueueCapacityPerBand: rc.Admission.QueueCapacityPerBand,
|
|
QueueCapacityTotal: rc.Admission.QueueCapacityTotal,
|
|
Interval: rc.Admission.Interval,
|
|
}
|
|
}
|
|
|
|
// saturation reports the pool pressure for a model, for the admission layer.
|
|
//
|
|
// outstanding is what admission has let through. It is compared against the
|
|
// registry's inflight counters rather than the scraped engine stats, because
|
|
// inflight updates the instant a request is dispatched while scraped stats lag
|
|
// by a whole interval. Comparing against scraped data would make every
|
|
// dispatched-but-unscraped request look like queue depth and throttle the
|
|
// steady state down to a few requests per scrape.
|
|
//
|
|
// What is left over — admitted but not yet dispatched — is the thundering herd:
|
|
// a burst that passed the check together before any of it became visible.
|
|
func (r *Router) saturation(model string, outstanding int) (float64, bool) {
|
|
eps := r.reg.Snapshots(model)
|
|
if len(eps) == 0 {
|
|
return 0, false
|
|
}
|
|
dispatched := 0
|
|
for _, e := range eps {
|
|
dispatched += e.Inflight
|
|
}
|
|
sat := admission.Saturation(eps, max(0, outstanding-dispatched), admission.Thresholds{
|
|
TargetWaitingPerReplica: r.cfg.Scoring.TargetWaiting,
|
|
KVCeiling: r.cfg.Scoring.KVCeiling,
|
|
})
|
|
// A model with an SLO is as saturated as it is past it: the queue and
|
|
// the cache can look healthy while answers come too slowly, and that
|
|
// is when the lower bands should wait.
|
|
return max(sat, r.sloPressure(model)), true
|
|
}
|
|
|
|
func toInfos(eps []config.Endpoint) []endpoint.Info {
|
|
out := make([]endpoint.Info, 0, len(eps))
|
|
for _, e := range eps {
|
|
out = append(out, endpoint.Info{
|
|
ID: e.ID,
|
|
URL: strings.TrimRight(e.URL, "/"),
|
|
APIKey: e.APIKey,
|
|
Models: e.Models,
|
|
Engine: e.Engine,
|
|
Pool: e.Pool,
|
|
Weight: e.Weight,
|
|
Priority: e.Priority,
|
|
Workflow: e.Workflow,
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Start begins scraping endpoint metrics and dispatching admissions.
|
|
// It returns immediately.
|
|
func (r *Router) Start(ctx context.Context) {
|
|
r.poller.Sync(ctx)
|
|
if r.lora != nil {
|
|
go r.lora.run(ctx)
|
|
}
|
|
if r.adm != nil {
|
|
r.adm.Start()
|
|
}
|
|
if r.recorder != nil {
|
|
r.recorder.Start(ctx)
|
|
}
|
|
if r.acct != nil && r.store != nil {
|
|
// Budgets are counted in memory; the store is what makes them survive a
|
|
// restart. A failure here is not fatal — refusing to start because last
|
|
// month's spend could not be read would be worse than starting with
|
|
// counters that are low.
|
|
if err := r.acct.RestoreBudgets(ctx, r.store, storepkg.PeriodMonth); err != nil {
|
|
r.log.Warn("could not restore budgets; they start from zero this period", "err", err)
|
|
}
|
|
}
|
|
if r.threads != nil {
|
|
go r.pruneThreads(ctx)
|
|
}
|
|
if r.directory != nil {
|
|
go r.directory.sync(ctx, r.leading)
|
|
}
|
|
if r.store != nil && r.usageRetention > 0 {
|
|
go r.pruneUsage(ctx, r.usageRetention)
|
|
}
|
|
if r.store != nil {
|
|
go r.runBatches(ctx)
|
|
go r.runJobs(ctx)
|
|
go r.pruneRequestBodies(ctx)
|
|
go r.refreshTenants(ctx)
|
|
go r.purgeSessions(ctx)
|
|
}
|
|
}
|
|
|
|
// purgeSessions deletes sessions and key claims that can no longer be used,
|
|
// hourly and once at startup: refused already, but a table that only grows
|
|
// is its own problem, and a session row is a credential's hash kept for
|
|
// nothing.
|
|
func (r *Router) purgeSessions(ctx context.Context) {
|
|
purge := func() {
|
|
now := time.Now()
|
|
if n, err := r.store.PurgeSessions(ctx, now); err != nil {
|
|
r.log.Warn("could not remove expired sessions", "err", err)
|
|
} else if n > 0 {
|
|
r.log.Info("removed expired and idle sessions", "sessions", n)
|
|
}
|
|
if n, err := r.store.PurgeClaims(ctx, now); err != nil {
|
|
r.log.Warn("could not remove used and expired key claims", "err", err)
|
|
} else if n > 0 {
|
|
r.log.Info("removed used and expired key claims", "claims", n)
|
|
}
|
|
}
|
|
purge()
|
|
t := time.NewTicker(time.Hour)
|
|
defer t.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-t.C:
|
|
purge()
|
|
}
|
|
}
|
|
}
|
|
|
|
// refreshTenants keeps the tenant names the request series are labelled with,
|
|
// and the budget gauges, a minute fresh.
|
|
func (r *Router) refreshTenants(ctx context.Context) {
|
|
refresh := func() {
|
|
if err := r.m.tenant.refresh(ctx, r.store, time.Now()); err != nil && ctx.Err() == nil {
|
|
r.log.Warn("could not refresh tenant names and budget gauges", "err", err)
|
|
}
|
|
}
|
|
refresh()
|
|
t := time.NewTicker(time.Minute)
|
|
defer t.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-t.C:
|
|
refresh()
|
|
}
|
|
}
|
|
}
|
|
|
|
// pruneUsage deletes request records older than the retention, hourly and
|
|
// once at startup, as pruneThreads does. The five-minute sums stay, so no
|
|
// total or chart moves; what goes is the trail of each old request.
|
|
func (r *Router) pruneUsage(ctx context.Context, keep time.Duration) {
|
|
prune := func() {
|
|
n, err := r.store.PruneUsage(ctx, time.Now().Add(-keep))
|
|
if err != nil {
|
|
r.log.Warn("could not prune old request records", "err", err)
|
|
return
|
|
}
|
|
if n > 0 {
|
|
r.log.Info("removed request records past their retention", "rows", n, "retention", keep)
|
|
}
|
|
}
|
|
prune()
|
|
t := time.NewTicker(time.Hour)
|
|
defer t.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-t.C:
|
|
prune()
|
|
}
|
|
}
|
|
}
|
|
|
|
// pruneThreads enforces each project's retention window.
|
|
//
|
|
// Without something running this, "we keep conversations for 24 hours" is a
|
|
// comment rather than a promise — the rows would simply accumulate, and the
|
|
// number a project chose would describe nothing.
|
|
//
|
|
// Hourly, and once at startup. A gateway that has been down for a week should
|
|
// not wait another hour before honouring a window that expired during it, and
|
|
// an hour is fine granularity for a limit expressed in hours.
|
|
func (r *Router) pruneThreads(ctx context.Context) {
|
|
prune := func() {
|
|
n, err := r.threads.PruneThreads(ctx, time.Now())
|
|
if err != nil {
|
|
r.log.Warn("could not prune expired conversations", "err", err)
|
|
return
|
|
}
|
|
if n > 0 {
|
|
r.log.Info("removed conversations past their retention", "threads", n)
|
|
}
|
|
m, err := r.threads.PruneMemory(ctx, time.Now())
|
|
if err != nil {
|
|
r.log.Warn("could not prune expired memory notes", "err", err)
|
|
return
|
|
}
|
|
if m > 0 {
|
|
r.log.Info("removed memory notes past their retention", "notes", m)
|
|
}
|
|
}
|
|
prune()
|
|
|
|
t := time.NewTicker(time.Hour)
|
|
defer t.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-t.C:
|
|
prune()
|
|
}
|
|
}
|
|
}
|
|
|
|
// Stop halts scraping and admission.
|
|
func (r *Router) Stop() {
|
|
r.poller.Stop()
|
|
if r.adm != nil {
|
|
r.adm.Close()
|
|
}
|
|
// Drain before closing the database, or the last batch has nowhere to go.
|
|
if r.recorder != nil {
|
|
r.recorder.Stop()
|
|
}
|
|
if r.store != nil {
|
|
_ = r.store.Close()
|
|
}
|
|
if r.ownsThreads {
|
|
_ = r.threads.Close()
|
|
}
|
|
}
|
|
|
|
// SetLeader tells /leader how to know whether this process leads.
|
|
func (r *Router) SetLeader(f func() bool) { r.isLeader.Store(&f) }
|
|
|
|
// leading says whether this process leads, which a process with no rival
|
|
// always does.
|
|
//
|
|
// Jobs and batches are served only where it does. A standby's router has no
|
|
// endpoints — its reconciler placed nothing — and a client's request reaches
|
|
// the leader by being forwarded at the door, which a job replayed inside the
|
|
// process never passes: a 3D job leased by the standby was told "loading"
|
|
// for as long as anyone let it, while the leader served the jobs behind it.
|
|
func (r *Router) leading() bool {
|
|
f := r.isLeader.Load()
|
|
return f == nil || (*f)()
|
|
}
|
|
|
|
// Registry exposes the endpoint registry, for tests and admin endpoints.
|
|
func (r *Router) Registry() *endpoint.Registry { return r.reg }
|
|
|
|
// SetEndpoints replaces the endpoint list and resynchronizes metric scraping.
|
|
//
|
|
// This is how the control plane feeds the router from phase 3 onwards, taking
|
|
// over from the static configuration list (plan §6.4). Endpoints that survive
|
|
// the update keep their live state — inflight counts, breaker state, drain rate
|
|
// — so a replica that was already serving does not start from cold just
|
|
// because another one moved.
|
|
// modelMap returns the name→model map in force.
|
|
//
|
|
// Never nil: a router built before any catalog arrived resolves nothing, which
|
|
// is the same answer it gave for an unknown name before.
|
|
func (r *Router) modelMap() map[string]catalog.Model {
|
|
if m := r.models.Load(); m != nil {
|
|
return *m
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// SetCatalog replaces the models this router knows about.
|
|
//
|
|
// The catalog's counterpart to SetEndpoints, and pushed for the same reason:
|
|
// the control plane owns what should exist, the router needs to know, and
|
|
// neither has to be in the other's process for that to work.
|
|
func (r *Router) SetCatalog(c *catalog.Catalog) { r.storeModels(c.Models()) }
|
|
|
|
// storeModels makes models what names resolve to: each model by its name, its
|
|
// aliases and its adapters. One place for it, whether the models came from
|
|
// configuration at start or from the control plane since.
|
|
func (r *Router) storeModels(models []catalog.Model) {
|
|
m := map[string]catalog.Model{}
|
|
ads := map[string]adapterRef{}
|
|
for _, entry := range models {
|
|
m[entry.Name] = entry
|
|
for _, a := range entry.Aliases {
|
|
m[a] = entry
|
|
}
|
|
// An adapter resolves to the model it runs on, as an alias does, and
|
|
// is applied to the request on the way there.
|
|
for i, a := range entry.Adapters {
|
|
m[a.Name] = entry
|
|
ads[a.Name] = adapterRef{base: entry.Name, owner: a.Owner, index: i, scale: a.EffectiveScale(),
|
|
limit: entry.MaxLoadedAdapters, path: engines.RuntimeLoRAPath(a.Name, a.Source)}
|
|
}
|
|
}
|
|
r.models.Store(&m)
|
|
r.adapters.Store(&ads)
|
|
targets := map[string]string{}
|
|
for _, entry := range models {
|
|
if entry.Canary != nil {
|
|
if t, ok := m[entry.Canary.Model]; ok {
|
|
targets[t.Name] = entry.Name
|
|
}
|
|
}
|
|
}
|
|
r.canaryTargets.Store(&targets)
|
|
if r.m != nil {
|
|
r.m.token.setSLOs(models)
|
|
}
|
|
}
|
|
|
|
// adapterRef is where an adapter runs and how it is applied.
|
|
type adapterRef struct {
|
|
base string
|
|
// owner is who sees it (catalog.Adapter.Owner); empty is everyone.
|
|
owner string
|
|
index int
|
|
scale float64
|
|
// limit is the model's max_loaded_adapters: above zero, the adapter
|
|
// loads on demand, from path.
|
|
limit int
|
|
path string
|
|
}
|
|
|
|
// onDemand reports whether the adapter is loaded when asked for rather than
|
|
// at start.
|
|
func (a adapterRef) onDemand() bool { return a.limit > 0 }
|
|
|
|
// adapter finds one by the name a client asked for.
|
|
func (r *Router) adapter(name string) (adapterRef, bool) {
|
|
if p := r.adapters.Load(); p != nil {
|
|
a, ok := (*p)[name]
|
|
return a, ok
|
|
}
|
|
return adapterRef{}, false
|
|
}
|
|
|
|
// adapterHidden reports whether name is an adapter its owner keeps from this
|
|
// caller: one outside the part of the tree it belongs to, or one with no
|
|
// identity at all. Such a caller is told there is no such model.
|
|
func (r *Router) adapterHidden(p auth.Principal, authed bool, name string) bool {
|
|
ad, ok := r.adapter(name)
|
|
if !ok || ad.owner == "" {
|
|
return false
|
|
}
|
|
if !authed {
|
|
return true
|
|
}
|
|
scope, id, _ := strings.Cut(ad.owner, ":")
|
|
switch scope {
|
|
case "org":
|
|
return p.OrgID != id
|
|
case "team":
|
|
return p.TeamID != id
|
|
case "project":
|
|
return p.ProjectID != id
|
|
case "user":
|
|
return p.UserID != id
|
|
case "key":
|
|
return p.KeyID != id
|
|
}
|
|
return true
|
|
}
|
|
|
|
// LastUsed reports when each endpoint was last sent a request, satisfying
|
|
// control.UsageSource so the control plane can tell a busy replica from one
|
|
// holding its cards out of habit.
|
|
func (r *Router) LastUsed() map[string]time.Time { return r.reg.LastUsed() }
|
|
|
|
// ModelLoad sums what each model's replicas report: how many there are, and
|
|
// requests running and waiting at their engines. Keyed by the catalog name,
|
|
// which the control plane publishes first among each endpoint's names.
|
|
func (r *Router) ModelLoad() map[string]reconciler.Load {
|
|
out := map[string]reconciler.Load{}
|
|
for _, s := range r.reg.Snapshots("") {
|
|
if len(s.Models) == 0 {
|
|
continue
|
|
}
|
|
l := out[s.Models[0]]
|
|
l.Replicas++
|
|
l.Running += s.Stats.Running
|
|
l.Waiting += s.Stats.Waiting
|
|
out[s.Models[0]] = l
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) {
|
|
// A replica that has gone takes its cache with it. Leaving its entries
|
|
// in the tree or the index would steer traffic towards something that is
|
|
// not there, and since nothing reuses a replica id the entries would never
|
|
// be corrected — only expire, minutes later, or never.
|
|
before := map[string]struct{}{}
|
|
for _, s := range r.reg.Snapshots("") {
|
|
before[s.ID] = struct{}{}
|
|
}
|
|
for _, i := range infos {
|
|
delete(before, i.ID)
|
|
}
|
|
for id := range before {
|
|
r.prefix.Forget(id)
|
|
r.kv.Forget(id)
|
|
}
|
|
|
|
if r.lora != nil {
|
|
ids := make(map[string]struct{}, len(infos))
|
|
for _, i := range infos {
|
|
ids[i.ID] = struct{}{}
|
|
}
|
|
r.lora.keep(ids)
|
|
}
|
|
|
|
r.reg.Replace(infos)
|
|
r.poller.Sync(ctx)
|
|
}
|
|
|
|
// --- model resolution ---
|
|
|
|
// resolve maps a client-supplied model name onto the catalog. Unknown names are
|
|
// passed through unchanged: an endpoint may serve a model the catalog does not
|
|
// list, and refusing it would be unhelpful.
|
|
func (r *Router) resolve(name string) (canonical string, m catalog.Model) {
|
|
if e, ok := r.modelMap()[name]; ok {
|
|
return e.Name, e
|
|
}
|
|
return name, catalog.Model{Name: name}
|
|
}
|
|
|
|
func (r *Router) profileFor(m catalog.Model, override string) *scheduler.Profile {
|
|
for _, name := range []string{override, m.RoutingProfile, r.cfg.DefaultProfile} {
|
|
if p, ok := r.profiles[name]; ok && name != "" {
|
|
return p
|
|
}
|
|
}
|
|
return r.profiles["chat"]
|
|
}
|
|
|
|
// --- dispatch ---
|
|
|
|
// attempt describes one upstream try, for logging and metrics.
|
|
type dispatchResult struct {
|
|
endpointID string
|
|
pool string
|
|
servedModel
|
|
outcome proxy.Outcome
|
|
err error
|
|
// guard is what the output filter decided, if there was one, and
|
|
// judged whether there was.
|
|
guard guard.Verdict
|
|
judged bool
|
|
// canaryOf is the model whose canary this request was split by, canary
|
|
// the candidate, and onCanary whether the candidate served it.
|
|
canaryOf, canary string
|
|
onCanary bool
|
|
// cached is an answer from the project's answer cache, which no model
|
|
// was asked for.
|
|
cached bool
|
|
}
|
|
|
|
type servedModel struct {
|
|
model string
|
|
// fallback is true when the served model is not the one that was requested.
|
|
fallback bool
|
|
}
|
|
|
|
// dispatch resolves the model, then tries endpoints until one works, the client
|
|
// has been written to, or the attempt budget runs out.
|
|
//
|
|
// The retry rule (plan §7.7) is strict: retry only while nothing has reached the
|
|
// client. Once bytes are out, a retry would splice two different responses
|
|
// together, which is worse than the error the client would otherwise see.
|
|
func (r *Router) dispatch(ctx context.Context, w http.ResponseWriter, req *parsedRequest) dispatchResult {
|
|
canonical, entry := r.resolve(req.Model)
|
|
|
|
// A canary takes its share before anything else, and falls back to the
|
|
// stable side, not to a fallback, when it has nothing ready.
|
|
// Not an adapter: the candidate's replicas do not run it under the
|
|
// stable model's name.
|
|
_, isAdapter := r.adapter(req.Model)
|
|
if cn := entry.Canary; cn != nil && !isAdapter {
|
|
onCanary := chooseVariant(req, cn)
|
|
target, _ := r.resolve(cn.Model)
|
|
if onCanary && len(r.reg.Snapshots(target)) == 0 {
|
|
onCanary = false
|
|
}
|
|
if _, te := r.resolve(target); onCanary && (req.noStore || req.inCluster) && te.IsExternal() {
|
|
onCanary = false
|
|
}
|
|
setVariantHeader(w.Header(), onCanary)
|
|
if onCanary {
|
|
// The engine is asked for the model by the name it serves.
|
|
creq := *req
|
|
creq.UpstreamBody = withModel(req.UpstreamBody, target)
|
|
res := r.dispatchOne(ctx, w, &creq, target, entry)
|
|
res.model, res.canaryOf, res.canary, res.onCanary = target, canonical, target, true
|
|
if res.err == nil || res.outcome.Wrote || errors.Is(res.err, proxy.ErrClientGone) ||
|
|
!errors.Is(res.err, errNoCapacity) {
|
|
return res
|
|
}
|
|
setVariantHeader(w.Header(), false)
|
|
}
|
|
res := r.dispatchStable(ctx, w, req, canonical, entry)
|
|
res.canaryOf, res.canary = canonical, target
|
|
return res
|
|
}
|
|
return r.dispatchStable(ctx, w, req, canonical, entry)
|
|
}
|
|
|
|
// dispatchStable serves a model as asked, then its fallbacks.
|
|
func (r *Router) dispatchStable(ctx context.Context, w http.ResponseWriter, req *parsedRequest, canonical string, entry catalog.Model) dispatchResult {
|
|
// The requested model first, then its fallbacks. A fallback is a different
|
|
// model, so the client is told about it via x-llmgw-served-model.
|
|
chain := []string{canonical}
|
|
p, authed := auth.FromContext(ctx)
|
|
for _, fb := range entry.Fallbacks {
|
|
// A fallback is another model, and a key allowed only the one it
|
|
// asked for is not served another because that one was full.
|
|
if authed && !r.permits(p, fb) || r.isGuardModel(fb) {
|
|
continue
|
|
}
|
|
// Nor sent outside the cluster under a policy that keeps nothing.
|
|
if _, fe := r.resolve(fb); (req.noStore || req.inCluster) && fe.IsExternal() {
|
|
continue
|
|
}
|
|
chain = append(chain, fb)
|
|
}
|
|
|
|
var last dispatchResult
|
|
for i, model := range chain {
|
|
if i > 0 {
|
|
// Before dispatching, because dispatchOne writes the response and
|
|
// a header set after that is a header nobody receives. It was set
|
|
// in finish until now, which is after the proxy has already called
|
|
// WriteHeader — so this has announced a fallback to nobody since it
|
|
// was written.
|
|
w.Header().Set("x-llmgw-served-model", model)
|
|
}
|
|
res := r.dispatchOne(ctx, w, req, model, entry)
|
|
res.model = model
|
|
res.fallback = i > 0
|
|
|
|
if res.err == nil || res.outcome.Wrote || errors.Is(res.err, proxy.ErrClientGone) {
|
|
return res
|
|
}
|
|
last = res
|
|
// Only a total absence of capacity justifies moving to another model;
|
|
// a bad request would fail identically there.
|
|
if !errors.Is(res.err, errNoCapacity) {
|
|
return res
|
|
}
|
|
}
|
|
return last
|
|
}
|
|
|
|
var errNoCapacity = errors.New("no endpoint available")
|
|
|
|
// errLoading means the model is being started and the wait ran out.
|
|
//
|
|
// Distinct from having no capacity because the right client behaviour differs:
|
|
// this one succeeds on a retry, and the other does not.
|
|
var errLoading = errors.New("model loading")
|
|
|
|
func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *parsedRequest, model string, entry catalog.Model) dispatchResult {
|
|
candidates := r.reg.Snapshots(model)
|
|
if visited := withoutVisited(candidates, req.Header); len(visited) < len(candidates) {
|
|
if len(visited) == 0 {
|
|
return dispatchResult{err: fmt.Errorf("%w: %s is served by a federated cluster this request has already passed through",
|
|
errNoCapacity, model)}
|
|
}
|
|
candidates = visited
|
|
}
|
|
if len(req.avoid) > 0 {
|
|
if rest := slices.DeleteFunc(slices.Clone(candidates), func(e endpoint.Snapshot) bool {
|
|
return slices.Contains(req.avoid, e.ID)
|
|
}); len(rest) > 0 {
|
|
candidates = rest
|
|
}
|
|
}
|
|
if len(candidates) == 0 {
|
|
// Nothing is serving it. That may be because nothing ever will, or
|
|
// because it loads when asked and nobody has asked yet — and the second
|
|
// is not an error, it is the feature working.
|
|
if r.activate(ctx, model) {
|
|
candidates = r.reg.Snapshots(model)
|
|
}
|
|
}
|
|
if len(candidates) == 0 {
|
|
if why, no := r.whyUnservable(model); no && r.activator != nil {
|
|
// Not "loading; retry": it will not load, and a client told to
|
|
// retry keeps retrying something that never arrives.
|
|
return dispatchResult{err: fmt.Errorf("%w: %s", errNoCapacity, why)}
|
|
}
|
|
if r.activator != nil && r.activator.Want(model) {
|
|
// Asked for, still loading. Saying "no endpoint serves this" would
|
|
// be a lie, and a client told to retry will succeed where one told
|
|
// the model does not exist gives up.
|
|
return dispatchResult{err: r.loadingErr(model)}
|
|
}
|
|
return dispatchResult{err: fmt.Errorf("%w: %s", errNoCapacity, noEndpointFor(model, slices.DeleteFunc(r.servedModels(), r.hiddenModel)))}
|
|
}
|
|
|
|
profile := r.profileFor(entry, req.ProfileOverride)
|
|
sctx := &scheduler.Ctx{
|
|
Now: time.Now(),
|
|
Req: scheduler.Request{
|
|
Model: model,
|
|
SessionID: req.SessionID,
|
|
PromptPrefix: req.PromptPrefix,
|
|
PromptTokens: req.PromptTokens,
|
|
MaxTokens: req.MaxTokens,
|
|
},
|
|
}
|
|
// An adapter asked for by name runs on its model's replicas — not on a
|
|
// fallback to another model, which does not have it.
|
|
ad, isAdapter := r.adapter(req.Model)
|
|
isAdapter = isAdapter && ad.base == model
|
|
onDemand := isAdapter && ad.onDemand() && r.lora != nil
|
|
if onDemand {
|
|
sctx.Req.Adapter = req.Model
|
|
r.lora.syncUnknown(ctx, candidates, model)
|
|
}
|
|
// Only where a profile publishes events, and only for a profile with a
|
|
// prefix scorer to read the answer: the batch profile does not care where
|
|
// a prefix is cached, and should not pay a round trip to find out.
|
|
if entry.KVIndex() && profile.Uses("prefix") {
|
|
_, kspan := tracer().Start(ctx, "kv index")
|
|
sctx.Req.KV = r.kvQuery(ctx, req, model, candidates)
|
|
if sctx.Req.KV.Fallback != "" {
|
|
kspan.SetAttributes(attribute.String("llmgw.kv_fallback", string(sctx.Req.KV.Fallback)))
|
|
}
|
|
kspan.End()
|
|
}
|
|
|
|
attempts := max(r.cfg.Retries+1, 1)
|
|
var last dispatchResult
|
|
|
|
for range attempts {
|
|
if len(candidates) == 0 {
|
|
break
|
|
}
|
|
if onDemand {
|
|
sctx.Adapters = r.lora.view(candidates, req.Model, ad.limit)
|
|
}
|
|
_, sspan := tracer().Start(ctx, "schedule", trace.WithAttributes(
|
|
attribute.String("llmgw.model", model), attribute.Int("llmgw.candidates", len(candidates))))
|
|
dec, err := profile.Schedule(sctx, candidates)
|
|
if err == nil {
|
|
sspan.SetAttributes(attribute.String("llmgw.endpoint", dec.Endpoint.ID))
|
|
}
|
|
endSpan(sspan, err)
|
|
if err != nil {
|
|
last = dispatchResult{err: fmt.Errorf("%w: %v", errNoCapacity, err)}
|
|
break
|
|
}
|
|
|
|
release, granted := r.reg.Acquire(dec.Endpoint.ID)
|
|
if !granted {
|
|
// The breaker opened between scoring and dispatch; drop it and
|
|
// score again rather than failing the request.
|
|
candidates = without(candidates, dec.Endpoint.ID)
|
|
continue
|
|
}
|
|
if req.picked != nil {
|
|
req.picked(dec.Endpoint.ID)
|
|
}
|
|
|
|
// The replica loads the adapter first if it does not hold it.
|
|
// Counted against its inflight while it does, so the requests
|
|
// arriving meanwhile see it busy.
|
|
releaseAdapter := func() {}
|
|
if onDemand {
|
|
how, rel, lerr := r.lora.acquire(ctx, dec.Endpoint.Info, model, req.Model, ad.path, ad.limit)
|
|
if lerr != nil {
|
|
release(true)
|
|
if ctx.Err() != nil {
|
|
return dispatchResult{endpointID: dec.Endpoint.ID, err: proxy.ErrClientGone}
|
|
}
|
|
res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool,
|
|
err: fmt.Errorf("adapter %s on %s: %w", req.Model, dec.Endpoint.ID, lerr)}
|
|
if errors.Is(lerr, errAdapterLoad) {
|
|
// Refused for what it is, it would be refused everywhere.
|
|
return res
|
|
}
|
|
r.log.Warn("a replica could not load an adapter, trying another",
|
|
"endpoint", dec.Endpoint.ID, "adapter", req.Model, "err", lerr)
|
|
last = res
|
|
candidates = without(candidates, dec.Endpoint.ID)
|
|
continue
|
|
}
|
|
releaseAdapter = rel
|
|
r.m.adapterAffinity.WithLabelValues(r.metricModel(model), how).Inc()
|
|
}
|
|
|
|
uctx, uspan := tracer().Start(ctx, "upstream", trace.WithSpanKind(trace.SpanKindClient),
|
|
trace.WithAttributes(attribute.String("llmgw.endpoint", dec.Endpoint.ID),
|
|
attribute.String("url.full", dec.Endpoint.URL+req.Path)))
|
|
body := req.UpstreamBody
|
|
// A shared ComfyUI replica is told which of its models' checkpoint;
|
|
// any other endpoint its own upstream name.
|
|
if _, up := dec.Endpoint.ComfyFor(model); up != "" && req.ContentType == "" {
|
|
body = withModel(body, up)
|
|
}
|
|
// llama-server is asked for the model with the adapter applied to
|
|
// this request; vLLM serves each adapter as a model of its own name,
|
|
// which is what the client asked for already.
|
|
if isAdapter && !onDemand && dec.Endpoint.Engine != "vllm" {
|
|
body = withAdapter(body, model, ad.index, ad.scale)
|
|
}
|
|
// Passages share cached blocks only with requests that could have
|
|
// been given the same ones: a prefix found warm would otherwise say
|
|
// what another tenant asked about, by how fast its answer started.
|
|
// vLLM keys its prefix cache by this salt; other engines are not
|
|
// sent a field they may refuse.
|
|
if req.kbSalt != "" && dec.Endpoint.Engine == "vllm" && req.ContentType == "" {
|
|
body = withCacheSalt(body, req.kbSalt)
|
|
}
|
|
// What computes the answer, said before the proxy writes the headers.
|
|
// A model can have replicas on cards and on CPUs at once, and the two
|
|
// differ tenfold in speed: a reply that came slowly is otherwise
|
|
// indistinguishable from a model that is slow everywhere. Set again on
|
|
// a retry, which only happens before anything was written.
|
|
if hw := hardwareOf(dec.Endpoint.Pool); hw != "" {
|
|
w.Header().Set(hardwareHeader, hw)
|
|
} else {
|
|
w.Header().Del(hardwareHeader)
|
|
}
|
|
preq := proxy.Request{
|
|
URL: dec.Endpoint.URL + req.Path,
|
|
APIKey: dec.Endpoint.APIKey,
|
|
Body: body,
|
|
Header: withTraceparent(uctx, req.Header),
|
|
Stream: req.Stream,
|
|
SuppressUsageChunk: req.SuppressUsageChunk,
|
|
CaptureReply: req.captureBytes(),
|
|
ContentType: req.ContentType,
|
|
}
|
|
if onDemand {
|
|
// vLLM's answer for a model it does not have: the adapter was
|
|
// unloaded behind this router's back, by another router or a
|
|
// restart. Not relayed; loaded again and tried again.
|
|
preq.RetryStatus = http.StatusNotFound
|
|
}
|
|
// One of an outside provider's several keys refused: the request
|
|
// is the provider's to answer with another of them, and the key's
|
|
// breaker counts it. With one key the 401 is the client's to see.
|
|
if dec.Endpoint.Engine == endpoint.EngineExternal && strings.Contains(dec.Endpoint.ID, "#") {
|
|
preq.RetryStatus = http.StatusUnauthorized
|
|
}
|
|
if dec.Endpoint.Peer != "" {
|
|
preq.Extra = r.peerHeaders(ctx, dec.Endpoint, req.Header)
|
|
}
|
|
if req.Stream {
|
|
gaps := r.m.token.gap.WithLabelValues(r.metricModel(model))
|
|
preq.OnEvent = func(d time.Duration) { gaps.Observe(d.Seconds()) }
|
|
}
|
|
// A filter per attempt, for the model this attempt is served by.
|
|
filter := r.outputFilter(uctx, req, model, &preq)
|
|
var out proxy.Outcome
|
|
var perr error
|
|
if dec.Endpoint.Engine == endpoint.EngineComfyUI {
|
|
out, perr = r.generateImages(uctx, w, req, dec.Endpoint.Info, body)
|
|
} else {
|
|
out, perr = r.px.Do(uctx, w, preq)
|
|
}
|
|
if dec.Endpoint.Engine == endpoint.EngineExternal && out.Status == http.StatusUnauthorized {
|
|
r.keyRefused(model, dec.Endpoint.ID)
|
|
}
|
|
releaseAdapter()
|
|
uspan.SetAttributes(attribute.Int("http.response.status_code", out.Status))
|
|
endSpan(uspan, perr)
|
|
adapterGone := onDemand && !out.Wrote && out.Status == http.StatusNotFound
|
|
if adapterGone {
|
|
r.lora.forget(dec.Endpoint.ID, req.Model)
|
|
}
|
|
// A retryable failure is the endpoint's fault; a client disconnect is
|
|
// not, and nor is an adapter it was not known to have lost.
|
|
release(perr == nil || errors.Is(perr, proxy.ErrClientGone) || adapterGone)
|
|
|
|
res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool, outcome: out, err: perr}
|
|
if filter != nil {
|
|
res.guard = filter.Verdict()
|
|
res.judged = true
|
|
}
|
|
r.m.observeScores(dec)
|
|
|
|
// Recorded after the decision, not during it: the scorer proposes and
|
|
// the picker disposes, so recording the proposal would fill the tree
|
|
// with replicas that were never sent anything — the one thing that
|
|
// would make the estimate worse than no estimate at all.
|
|
if r.prefix != nil {
|
|
if !req.noStore {
|
|
r.prefix.Observe(req.PromptPrefix, dec.Endpoint.ID)
|
|
}
|
|
}
|
|
|
|
if perr == nil || out.Wrote || !out.Retryable {
|
|
return res
|
|
}
|
|
last = res
|
|
if adapterGone {
|
|
// The replica is fine; the next try may load the adapter on it
|
|
// again.
|
|
r.log.Info("a replica no longer held an adapter; loading it again",
|
|
"endpoint", dec.Endpoint.ID, "adapter", req.Model)
|
|
continue
|
|
}
|
|
r.log.Warn("upstream attempt failed, trying another endpoint",
|
|
"endpoint", dec.Endpoint.ID, "model", model, "err", perr)
|
|
candidates = without(candidates, dec.Endpoint.ID)
|
|
}
|
|
|
|
if last.err == nil {
|
|
last.err = fmt.Errorf("%w: every endpoint for %q failed", errNoCapacity, model)
|
|
} else if !errors.Is(last.err, errNoCapacity) {
|
|
last.err = fmt.Errorf("%w: %v", errNoCapacity, last.err)
|
|
}
|
|
return last
|
|
}
|
|
|
|
// hardwareHeader says what the replica that answered computes on: "cpu", or
|
|
// the GPU class, such as "6000-24g".
|
|
const hardwareHeader = "x-llmgw-hardware"
|
|
|
|
// hardwareOf reads it from a pool, which placement names fabric/class —
|
|
// "none/cpu", "none/6000-24g", "ib-a/h100-80g". Empty for an endpoint
|
|
// configured by hand, which has no pool and says nothing about its hardware.
|
|
func hardwareOf(pool string) string {
|
|
i := strings.LastIndexByte(pool, '/')
|
|
if i < 0 || i == len(pool)-1 {
|
|
return ""
|
|
}
|
|
return pool[i+1:]
|
|
}
|
|
|
|
func without(eps []endpoint.Snapshot, id string) []endpoint.Snapshot {
|
|
return slices.DeleteFunc(slices.Clone(eps), func(e endpoint.Snapshot) bool { return e.ID == id })
|
|
}
|
|
|
|
// personalLimits turns the personal block into limits, nil when it sets none.
|
|
func personalLimits(p config.Personal) *storepkg.Limits {
|
|
var l storepkg.Limits
|
|
if p.RPM > 0 {
|
|
v := p.RPM
|
|
l.RPM = &v
|
|
}
|
|
if p.TPM > 0 {
|
|
v := p.TPM
|
|
l.TPM = &v
|
|
}
|
|
if p.Budget > 0 {
|
|
v := int64(math.Round(p.Budget * 1e6))
|
|
l.Budget = &v
|
|
l.BudgetPeriod = storepkg.PeriodMonth
|
|
if p.Period == "day" {
|
|
l.BudgetPeriod = storepkg.PeriodDay
|
|
}
|
|
}
|
|
if l.RPM == nil && l.TPM == nil && l.Budget == nil {
|
|
return nil
|
|
}
|
|
return &l
|
|
}
|
|
|
|
// loopbackOnly reports whether a listen address is reachable from this
|
|
// machine alone. An empty host (":8080") is every interface.
|
|
func loopbackOnly(listen string) bool {
|
|
host, _, err := net.SplitHostPort(listen)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
if host == "localhost" {
|
|
return true
|
|
}
|
|
a, err := netip.ParseAddr(host)
|
|
return err == nil && a.IsLoopback()
|
|
}
|