diff --git a/go.mod b/go.mod index 0f50395..ed854e7 100644 --- a/go.mod +++ b/go.mod @@ -8,9 +8,10 @@ require ( ) require ( + github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect - github.com/prometheus/client_model v0.6.2 // indirect + github.com/prometheus/client_model v0.6.3 // indirect github.com/spf13/pflag v1.0.9 // indirect google.golang.org/protobuf v1.36.12 // indirect ) diff --git a/go.sum b/go.sum index 41032db..82ef0dd 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= @@ -7,6 +9,8 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE= +github.com/prometheus/client_model v0.6.3 h1:O0jaTVAYNxTHYInEPFJt5I3+sN8zqBtVMPTB1qyxiEo= +github.com/prometheus/client_model v0.6.3/go.mod h1:gpN5P9S7Rr6Yr92PiQ+Ixvhf6JZEkF1dnxsYL2aPBEM= github.com/prometheus/common v0.71.0 h1:9KDAKb7Mj3HEVKyFCK6Dc/HIwlBzZIN2l7/lrHl3KK8= github.com/prometheus/common v0.71.0/go.mod h1:CLJ5H8TEsGX8bl31BdMkfhIZ+QmZ9tBPPotUxUbfcmk= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= diff --git a/internal/router/breaker/breaker.go b/internal/router/breaker/breaker.go new file mode 100644 index 0000000..a02c2ad --- /dev/null +++ b/internal/router/breaker/breaker.go @@ -0,0 +1,198 @@ +// Package breaker provides a per-endpoint circuit breaker. +// +// It folds two concepts from plan §7.7 into a single component: +// +// - Cooldown: an endpoint that errors leaves the candidate pool for a while. +// (Set Threshold to 1 to get exactly this behaviour.) +// - Circuit breaker: opens once consecutive failures cross a threshold, then +// moves to half-open probes with exponential backoff. +// +// Keeping them apart would mean two timers and two separate answers to +// "is this endpoint usable right now?"; merging them produces fewer surprises. +package breaker + +import ( + "sync" + "time" +) + +// State is the breaker's current state. +type State int + +const ( + // Closed is normal operation: requests pass through. + Closed State = iota + // Open means the endpoint is out of service; no requests are sent. + Open + // HalfOpen allows a limited number of probe requests. + HalfOpen +) + +func (s State) String() string { + switch s { + case Closed: + return "closed" + case Open: + return "open" + case HalfOpen: + return "half-open" + default: + return "unknown" + } +} + +// Config holds the breaker's settings. +type Config struct { + // FailureThreshold is the number of consecutive failures that opens the + // circuit. 1 means "one error puts the endpoint in cooldown". + FailureThreshold int + // BaseCooldown is how long the circuit stays open the first time it trips. + BaseCooldown time.Duration + // MaxCooldown caps the exponential backoff. + MaxCooldown time.Duration + // HalfOpenProbes is how many concurrent probes are allowed while half-open. + HalfOpenProbes int +} + +// DefaultConfig returns sensible defaults. +func DefaultConfig() Config { + return Config{ + FailureThreshold: 3, + BaseCooldown: 15 * time.Second, + MaxCooldown: 2 * time.Minute, + HalfOpenProbes: 1, + } +} + +func (c *Config) normalize() { + if c.FailureThreshold <= 0 { + c.FailureThreshold = 3 + } + if c.BaseCooldown <= 0 { + c.BaseCooldown = 15 * time.Second + } + if c.MaxCooldown < c.BaseCooldown { + c.MaxCooldown = max(c.BaseCooldown, 2*time.Minute) + } + if c.HalfOpenProbes <= 0 { + c.HalfOpenProbes = 1 + } +} + +// Breaker is one endpoint's circuit breaker. It is safe for concurrent use. +type Breaker struct { + cfg Config + now func() time.Time + + mu sync.Mutex + state State + consecFails int + openUntil time.Time + openRounds int // how many times it has opened, for backoff + probesInUse int + totalTrips int64 +} + +// New creates a breaker with the given settings. +func New(cfg Config) *Breaker { + cfg.normalize() + return &Breaker{cfg: cfg, now: time.Now, state: Closed} +} + +// SetClock replaces the time source, for tests. +func (b *Breaker) SetClock(now func() time.Time) { + b.mu.Lock() + b.now = now + b.mu.Unlock() +} + +// Allow reports whether a request may be sent to this endpoint. +// If it returns true, the outcome must be reported via Success or Failure. +func (b *Breaker) Allow() bool { + b.mu.Lock() + defer b.mu.Unlock() + + switch b.state { + case Closed: + return true + case Open: + if b.now().Before(b.openUntil) { + return false + } + // Cooldown elapsed: move to half-open and let this probe through. + b.state = HalfOpen + b.probesInUse = 1 + return true + case HalfOpen: + if b.probesInUse >= b.cfg.HalfOpenProbes { + return false + } + b.probesInUse++ + return true + default: + return true + } +} + +// Success reports a successful request. +func (b *Breaker) Success() { + b.mu.Lock() + defer b.mu.Unlock() + + b.consecFails = 0 + if b.state == HalfOpen { + b.probesInUse = 0 + b.openRounds = 0 + } + b.state = Closed +} + +// Failure reports a failed request. +func (b *Breaker) Failure() { + b.mu.Lock() + defer b.mu.Unlock() + + if b.state == HalfOpen { + // The probe failed too: reopen and lengthen the backoff. + b.probesInUse = 0 + b.trip() + return + } + + b.consecFails++ + if b.consecFails >= b.cfg.FailureThreshold { + b.trip() + } +} + +// trip opens the circuit. The caller must hold the lock. +func (b *Breaker) trip() { + cooldown := b.cfg.BaseCooldown << min(b.openRounds, 16) + if cooldown > b.cfg.MaxCooldown || cooldown <= 0 { + cooldown = b.cfg.MaxCooldown + } + b.state = Open + b.openUntil = b.now().Add(cooldown) + b.openRounds++ + b.consecFails = 0 + b.totalTrips++ +} + +// State returns the current state. An open circuit whose cooldown has elapsed +// is reported as half-open without changing state: scoring must not have side +// effects. +func (b *Breaker) State() State { + b.mu.Lock() + defer b.mu.Unlock() + if b.state == Open && !b.now().Before(b.openUntil) { + return HalfOpen + } + return b.state +} + +// Trips returns how many times the circuit has opened. +func (b *Breaker) Trips() int64 { + b.mu.Lock() + defer b.mu.Unlock() + return b.totalTrips +} diff --git a/internal/router/breaker/breaker_test.go b/internal/router/breaker/breaker_test.go new file mode 100644 index 0000000..04f8a44 --- /dev/null +++ b/internal/router/breaker/breaker_test.go @@ -0,0 +1,178 @@ +package breaker + +import ( + "sync" + "testing" + "time" +) + +// fakeClock lets tests advance time by hand. +type fakeClock struct { + mu sync.Mutex + t time.Time +} + +func newClock() *fakeClock { + return &fakeClock{t: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)} +} + +func (c *fakeClock) now() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.t +} + +func (c *fakeClock) advance(d time.Duration) { + c.mu.Lock() + c.t = c.t.Add(d) + c.mu.Unlock() +} + +func newTestBreaker(cfg Config) (*Breaker, *fakeClock) { + b := New(cfg) + c := newClock() + b.SetClock(c.now) + return b, c +} + +func TestOpensOnceThresholdIsCrossed(t *testing.T) { + b, _ := newTestBreaker(Config{FailureThreshold: 3, BaseCooldown: time.Minute}) + + for i := range 2 { + b.Failure() + if !b.Allow() { + t.Fatalf("circuit opened after %d failures; threshold is 3", i+1) + } + } + b.Failure() + + if b.Allow() { + t.Error("threshold crossed but the circuit is still closed") + } + if got := b.State(); got != Open { + t.Errorf("state %v, expected open", got) + } +} + +func TestSuccessResetsFailureCount(t *testing.T) { + b, _ := newTestBreaker(Config{FailureThreshold: 3, BaseCooldown: time.Minute}) + + b.Failure() + b.Failure() + b.Success() // must reset the counter + b.Failure() + b.Failure() + + if !b.Allow() { + t.Error("consecutive counter survived an intervening success") + } +} + +func TestSingleFailureCooldown(t *testing.T) { + // LiteLLM-style cooldown: one error takes the endpoint out for a while. + b, clock := newTestBreaker(Config{FailureThreshold: 1, BaseCooldown: 15 * time.Second}) + + b.Failure() + if b.Allow() { + t.Fatal("a single failure did not trigger cooldown") + } + + clock.advance(14 * time.Second) + if b.Allow() { + t.Error("allowed a request before the cooldown elapsed") + } + + clock.advance(2 * time.Second) + if !b.Allow() { + t.Error("cooldown elapsed but requests are still blocked") + } +} + +func TestHalfOpenClosesOnSuccess(t *testing.T) { + b, clock := newTestBreaker(Config{FailureThreshold: 1, BaseCooldown: time.Second, HalfOpenProbes: 1}) + + b.Failure() + clock.advance(2 * time.Second) + + if !b.Allow() { + t.Fatal("half-open probe was not allowed") + } + // A second concurrent probe must be blocked. + if b.Allow() { + t.Error("second probe allowed despite HalfOpenProbes=1") + } + + b.Success() + if got := b.State(); got != Closed { + t.Errorf("state is %v after a successful probe", got) + } + if !b.Allow() { + t.Error("requests still blocked after the circuit closed") + } +} + +func TestHalfOpenFailureBacksOffExponentially(t *testing.T) { + b, clock := newTestBreaker(Config{ + FailureThreshold: 1, + BaseCooldown: time.Second, + MaxCooldown: time.Minute, + }) + + b.Failure() // first trip: 1s + clock.advance(1100 * time.Millisecond) + + if !b.Allow() { + t.Fatal("no probe allowed after the first cooldown") + } + b.Failure() // probe failed: the second trip should last 2s + + clock.advance(1100 * time.Millisecond) + if b.Allow() { + t.Error("backoff did not grow; allowed again after 1s") + } + clock.advance(time.Second) + if !b.Allow() { + t.Error("2s backoff elapsed but requests are still blocked") + } +} + +func TestBackoffIsCapped(t *testing.T) { + b, clock := newTestBreaker(Config{ + FailureThreshold: 1, + BaseCooldown: time.Second, + MaxCooldown: 4 * time.Second, + }) + + for range 10 { + b.Failure() + clock.advance(5 * time.Second) + b.Allow() + } + b.Failure() + + clock.advance(5 * time.Second) + if !b.Allow() { + t.Error("cooldown exceeded the MaxCooldown cap") + } +} + +func TestConcurrentUse(t *testing.T) { + b := New(DefaultConfig()) + var wg sync.WaitGroup + for i := range 50 { + wg.Add(1) + go func() { + defer wg.Done() + if b.Allow() { + if i%2 == 0 { + b.Success() + } else { + b.Failure() + } + } + _ = b.State() + _ = b.Trips() + }() + } + wg.Wait() +} diff --git a/internal/router/endpoint/endpoint.go b/internal/router/endpoint/endpoint.go new file mode 100644 index 0000000..092d46f --- /dev/null +++ b/internal/router/endpoint/endpoint.go @@ -0,0 +1,297 @@ +// Package endpoint tracks the upstream replicas the router can send to, along +// with their live state. +// +// In phase 1 the endpoint list comes from configuration. In phase 3 the control +// plane's EndpointSource (plan §6.4) will feed the same Registry; the scheduler +// and proxy must not notice the difference, which is why the API is shaped +// around wholesale list replacement from the start. +package endpoint + +import ( + "maps" + "slices" + "sync" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker" +) + +// Info is a replica's static description. +type Info struct { + // ID uniquely identifies the replica and labels its metrics. + ID string + // URL is the engine's root address (e.g. http://gpu-a-01:8000), no trailing slash. + URL string + // Models lists the model names this replica serves. + Models []string + // Pool names the (fabric, gpu_class) pool. Used for observability and pool + // preference only; phase 1 may well have a single pool. + Pool string + // Weight is the capacity weight measured by ctl bench (§6.2.1). 0 means 1. + Weight float64 + // Priority expresses pool preference; higher pools are tried first. + Priority int +} + +// Serves reports whether the replica serves the given model. +func (i Info) Serves(model string) bool { + return slices.Contains(i.Models, model) +} + +// Stats is the router-side view of the engine metric contract (plan §7.6). +type Stats struct { + // The three mandatory fields. + Running int + Waiting int + KVCacheUsage float64 + + // Optional. + BlockSize int + NumGPUBlocks int + + // Raw counters, kept so derived rates can be computed. + PromptTokens float64 + GenerationTokens float64 + SuccessTotal float64 +} + +// Snapshot is a complete view of one endpoint at scheduling time. +// It is a value type so the scheduler can work without holding a lock. +type Snapshot struct { + Info + Stats Stats + + // DrainRate is the observed completion rate in requests/sec, exponentially + // smoothed. This is AIBrix's load-balance signal: Running alone cannot tell + // that a replica is slow. + DrainRate float64 + + // Inflight counts requests the router has dispatched here that have not + // finished. It compensates for scrape lag (§7.4); without it, herding is + // unavoidable within a 2s scrape interval. + Inflight int + + // Stale means no successful scrape arrived within StaleAfter. Stale + // endpoints are not dropped, only scored down. + Stale bool + // Age is the time since the last successful scrape. + Age time.Duration + // Breaker is the circuit breaker's current state, read without side effects. + Breaker breaker.State +} + +// Config holds Registry settings. +type Config struct { + // StaleAfter is how long an endpoint may go without a successful scrape + // before being considered stale. + StaleAfter time.Duration + // Breaker holds per-endpoint circuit breaker settings. + Breaker breaker.Config +} + +// DefaultConfig returns sensible defaults. +func DefaultConfig() Config { + return Config{ + StaleAfter: 10 * time.Second, + Breaker: breaker.DefaultConfig(), + } +} + +// live is one endpoint's mutable state. +type live struct { + info Info + + stats Stats + scrapedAt time.Time + lastSucc float64 + drainRate float64 + + inflight int + brk *breaker.Breaker +} + +// Registry holds the live state of all endpoints. It is safe for concurrent use. +type Registry struct { + cfg Config + now func() time.Time + + mu sync.RWMutex + m map[string]*live +} + +// NewRegistry creates an empty registry. +func NewRegistry(cfg Config) *Registry { + if cfg.StaleAfter <= 0 { + cfg.StaleAfter = 10 * time.Second + } + return &Registry{cfg: cfg, now: time.Now, m: make(map[string]*live)} +} + +// SetClock replaces the time source, for tests. +func (r *Registry) SetClock(now func() time.Time) { + r.mu.Lock() + r.now = now + r.mu.Unlock() +} + +// Replace reconciles the endpoint list with the given set. Endpoints that +// remain keep their live state, new ones get a breaker, removed ones are dropped. +func (r *Registry) Replace(infos []Info) { + r.mu.Lock() + defer r.mu.Unlock() + + seen := make(map[string]struct{}, len(infos)) + for _, in := range infos { + seen[in.ID] = struct{}{} + if l, ok := r.m[in.ID]; ok { + l.info = in + continue + } + r.m[in.ID] = &live{info: in, brk: breaker.New(r.cfg.Breaker)} + } + for id := range r.m { + if _, ok := seen[id]; !ok { + delete(r.m, id) + } + } +} + +// IDs returns the registered endpoint IDs in sorted order. +func (r *Registry) IDs() []string { + r.mu.RLock() + defer r.mu.RUnlock() + return slices.Sorted(maps.Keys(r.m)) +} + +// Infos returns the registered endpoint descriptions, sorted by ID. +func (r *Registry) Infos() []Info { + r.mu.RLock() + defer r.mu.RUnlock() + out := make([]Info, 0, len(r.m)) + for _, l := range r.m { + out = append(out, l.info) + } + slices.SortFunc(out, func(a, b Info) int { return strcmp(a.ID, b.ID) }) + return out +} + +// Observe records a successful scrape and updates the drain rate. +func (r *Registry) Observe(id string, s Stats) { + r.mu.Lock() + defer r.mu.Unlock() + + l, ok := r.m[id] + if !ok { + return + } + now := r.now() + if !l.scrapedAt.IsZero() { + dt := now.Sub(l.scrapedAt).Seconds() + // A counter going backwards means the engine restarted; skip the derivative. + if dt > 0 && s.SuccessTotal >= l.lastSucc { + rate := (s.SuccessTotal - l.lastSucc) / dt + const alpha = 0.3 // exponential smoothing factor + l.drainRate = alpha*rate + (1-alpha)*l.drainRate + } else if s.SuccessTotal < l.lastSucc { + l.drainRate = 0 + } + } + l.stats = s + l.lastSucc = s.SuccessTotal + l.scrapedAt = now +} + +// Snapshots returns a view of every endpoint serving the given model, or of all +// endpoints when model is empty. Breaker state is read, never consumed. +func (r *Registry) Snapshots(model string) []Snapshot { + r.mu.RLock() + defer r.mu.RUnlock() + + now := r.now() + out := make([]Snapshot, 0, len(r.m)) + for _, l := range r.m { + if model != "" && !l.info.Serves(model) { + continue + } + var age time.Duration + stale := true + if !l.scrapedAt.IsZero() { + age = now.Sub(l.scrapedAt) + stale = age > r.cfg.StaleAfter + } + out = append(out, Snapshot{ + Info: l.info, + Stats: l.stats, + DrainRate: l.drainRate, + Inflight: l.inflight, + Stale: stale, + Age: age, + Breaker: l.brk.State(), + }) + } + slices.SortFunc(out, func(a, b Snapshot) int { return strcmp(a.ID, b.ID) }) + return out +} + +// Acquire asks permission to send a request to an endpoint. +// +// It returns (nil, false) if the circuit breaker refuses. On success the +// inflight counter is incremented and the returned function must be called to +// report the outcome: release(err == nil). +func (r *Registry) Acquire(id string) (release func(ok bool), granted bool) { + r.mu.Lock() + l, exists := r.m[id] + r.mu.Unlock() + if !exists { + return nil, false + } + if !l.brk.Allow() { + return nil, false + } + + r.mu.Lock() + l.inflight++ + r.mu.Unlock() + + var once sync.Once + return func(ok bool) { + once.Do(func() { + r.mu.Lock() + l.inflight-- + r.mu.Unlock() + if ok { + l.brk.Success() + } else { + l.brk.Failure() + } + }) + }, true +} + +// ScrapeFailed reports a failed metrics scrape. +// +// A scrape failure is not a request failure, so it does not touch the breaker: +// the endpoint may be serving fine while only its metrics endpoint is down. +// It will go stale on its own and score lower (§7.6). +func (r *Registry) ScrapeFailed(string) {} + +// Inflight returns an endpoint's current inflight count, for tests and metrics. +func (r *Registry) Inflight(id string) int { + r.mu.RLock() + defer r.mu.RUnlock() + if l, ok := r.m[id]; ok { + return l.inflight + } + return 0 +} + +func strcmp(a, b string) int { + switch { + case a < b: + return -1 + case a > b: + return 1 + default: + return 0 + } +} diff --git a/internal/router/endpoint/endpoint_test.go b/internal/router/endpoint/endpoint_test.go new file mode 100644 index 0000000..29e4c11 --- /dev/null +++ b/internal/router/endpoint/endpoint_test.go @@ -0,0 +1,245 @@ +package endpoint + +import ( + "sync" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker" +) + +type clock struct { + mu sync.Mutex + t time.Time +} + +func newClock() *clock { + return &clock{t: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)} +} +func (c *clock) now() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.t +} +func (c *clock) advance(d time.Duration) { + c.mu.Lock() + c.t = c.t.Add(d) + c.mu.Unlock() +} + +func newReg(t *testing.T) (*Registry, *clock) { + t.Helper() + r := NewRegistry(Config{StaleAfter: 10 * time.Second, Breaker: breaker.DefaultConfig()}) + c := newClock() + r.SetClock(c.now) + return r, c +} + +func TestReplaceKeepsLiveState(t *testing.T) { + r, _ := newReg(t) + r.Replace([]Info{{ID: "a", URL: "http://a", Models: []string{"m"}}}) + r.Observe("a", Stats{Running: 5}) + + // Re-supplying the same endpoint must not lose its stats. + r.Replace([]Info{ + {ID: "a", URL: "http://a", Models: []string{"m"}, Weight: 2}, + {ID: "b", URL: "http://b", Models: []string{"m"}}, + }) + + snaps := r.Snapshots("m") + if len(snaps) != 2 { + t.Fatalf("expected 2 endpoints, got %d", len(snaps)) + } + if snaps[0].Stats.Running != 5 { + t.Error("Replace reset the live state") + } + if snaps[0].Weight != 2 { + t.Error("Replace did not update the static info") + } +} + +func TestReplaceRemovesMissingEndpoints(t *testing.T) { + r, _ := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}, {ID: "b", Models: []string{"m"}}}) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + + if got := r.IDs(); len(got) != 1 || got[0] != "a" { + t.Errorf("endpoint dropped from the list was not removed: %v", got) + } +} + +func TestFilterByModel(t *testing.T) { + r, _ := newReg(t) + r.Replace([]Info{ + {ID: "a", Models: []string{"llama"}}, + {ID: "b", Models: []string{"qwen"}}, + {ID: "c", Models: []string{"llama", "qwen"}}, + }) + + if got := len(r.Snapshots("llama")); got != 2 { + t.Errorf("got %d endpoints for llama", got) + } + if got := len(r.Snapshots("")); got != 3 { + t.Errorf("empty model should return all, got %d", got) + } +} + +func TestStalenessDetection(t *testing.T) { + r, c := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + + // With no scrape at all the endpoint counts as stale. + if !r.Snapshots("m")[0].Stale { + t.Error("not marked stale despite having no measurement") + } + + r.Observe("a", Stats{Running: 1}) + if r.Snapshots("m")[0].Stale { + t.Error("a fresh measurement was marked stale") + } + + c.advance(11 * time.Second) + s := r.Snapshots("m")[0] + if !s.Stale { + t.Error("not marked stale after StaleAfter elapsed") + } + if s.Age < 11*time.Second { + t.Errorf("Age is wrong: %v", s.Age) + } +} + +func TestDrainRateCalculation(t *testing.T) { + r, c := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + + r.Observe("a", Stats{SuccessTotal: 100}) + c.advance(time.Second) + r.Observe("a", Stats{SuccessTotal: 110}) // 10 requests/sec + + if got := r.Snapshots("m")[0].DrainRate; got <= 0 { + t.Fatalf("drain rate was not computed: %v", got) + } + + // Exponential smoothing: one sample approaches the raw rate but stays below it. + if got := r.Snapshots("m")[0].DrainRate; got > 10 { + t.Errorf("drain rate exceeded the raw value: %v", got) + } +} + +func TestDrainRateResetsWhenCounterGoesBackwards(t *testing.T) { + // Counters reset when an engine restarts; that must not yield a negative rate. + r, c := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + + r.Observe("a", Stats{SuccessTotal: 1000}) + c.advance(time.Second) + r.Observe("a", Stats{SuccessTotal: 1100}) + c.advance(time.Second) + r.Observe("a", Stats{SuccessTotal: 5}) // restarted + + if got := r.Snapshots("m")[0].DrainRate; got != 0 { + t.Errorf("drain rate was not reset after the counter went backwards: %v", got) + } +} + +func TestAcquireCountsInflight(t *testing.T) { + r, _ := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + + rel1, ok := r.Acquire("a") + if !ok { + t.Fatal("Acquire was refused") + } + rel2, ok := r.Acquire("a") + if !ok { + t.Fatal("second Acquire was refused") + } + if got := r.Inflight("a"); got != 2 { + t.Errorf("inflight=%d, expected 2", got) + } + + rel1(true) + rel2(true) + if got := r.Inflight("a"); got != 0 { + t.Errorf("inflight=%d after release", got) + } +} + +func TestReleaseIsIdempotent(t *testing.T) { + r, _ := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + + rel, _ := r.Acquire("a") + rel(true) + rel(true) // must be idempotent; the counter must not go negative + + if got := r.Inflight("a"); got != 0 { + t.Errorf("double release corrupted the counter: %d", got) + } +} + +func TestFailuresOpenTheCircuit(t *testing.T) { + r, _ := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + + for range 3 { // DefaultConfig threshold is 3 + rel, ok := r.Acquire("a") + if !ok { + t.Fatal("circuit opened too early") + } + rel(false) + } + + if _, ok := r.Acquire("a"); ok { + t.Error("circuit did not open after consecutive failures") + } + if r.Snapshots("m")[0].Breaker != breaker.Open { + t.Error("snapshot does not reflect the breaker state") + } +} + +func TestScrapeFailureDoesNotTripBreaker(t *testing.T) { + // A scrape failure is not a request failure: the endpoint may be serving + // fine while only its metrics endpoint is unreachable. + r, _ := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + + for range 10 { + r.ScrapeFailed("a") + } + if _, ok := r.Acquire("a"); !ok { + t.Error("scrape failures opened the circuit") + } +} + +func TestUnknownEndpoint(t *testing.T) { + r, _ := newReg(t) + if _, ok := r.Acquire("missing"); ok { + t.Error("an unregistered endpoint was allowed") + } + r.Observe("missing", Stats{}) // must not panic +} + +func TestConcurrentAccess(t *testing.T) { + r, _ := newReg(t) + r.Replace([]Info{{ID: "a", Models: []string{"m"}}, {ID: "b", Models: []string{"m"}}}) + + var wg sync.WaitGroup + for i := range 100 { + wg.Add(1) + go func() { + defer wg.Done() + id := "a" + if i%2 == 0 { + id = "b" + } + if rel, ok := r.Acquire(id); ok { + rel(i%3 != 0) + } + r.Observe(id, Stats{Running: i}) + _ = r.Snapshots("m") + _ = r.Inflight(id) + }() + } + wg.Wait() +} diff --git a/internal/router/scheduler/filters.go b/internal/router/scheduler/filters.go new file mode 100644 index 0000000..617b8a6 --- /dev/null +++ b/internal/router/scheduler/filters.go @@ -0,0 +1,84 @@ +package scheduler + +import ( + "gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// filterFunc is a helper for expressing simple filters. +type filterFunc struct { + name string + keep func(c *Ctx, e endpoint.Snapshot) bool +} + +func (f filterFunc) Name() string { return f.name } + +func (f filterFunc) Filter(c *Ctx, eps []endpoint.Snapshot) []endpoint.Snapshot { + out := make([]endpoint.Snapshot, 0, len(eps)) + for _, e := range eps { + if f.keep(c, e) { + out = append(out, e) + } + } + return out +} + +// BreakerOpen drops endpoints whose circuit is open. +// +// Half-open ones survive: the probe request has to get through, otherwise the +// circuit would never close again. +func BreakerOpen() Filter { + return filterFunc{ + name: "breaker-open", + keep: func(_ *Ctx, e endpoint.Snapshot) bool { + return e.Breaker != breaker.Open + }, + } +} + +// KVSaturated drops endpoints whose KV cache usage is at or above a ceiling. +// +// Note this is a hard constraint: if the whole pool is above the ceiling, no +// candidates remain. The built-in profiles rely on the kv scorer's penalty +// instead; use this filter only when "never send to a full replica" is meant +// literally. +func KVSaturated(ceiling float64) Filter { + return filterFunc{ + name: "kv-saturated", + keep: func(_ *Ctx, e endpoint.Snapshot) bool { + return e.Stats.KVCacheUsage < ceiling + }, + } +} + +// PoolPriority keeps only the endpoints of the highest-priority pool +// (plan §7.3, the two-stage decision: pool first, then replica). +// +// Spillover arrives in phase 5, at which point this filter will fall through +// to the next pool once the preferred one is saturated. +func PoolPriority() Filter { + return poolPriorityFilter{} +} + +type poolPriorityFilter struct{} + +func (poolPriorityFilter) Name() string { return "pool-priority" } + +func (poolPriorityFilter) Filter(_ *Ctx, eps []endpoint.Snapshot) []endpoint.Snapshot { + if len(eps) == 0 { + return eps + } + best := eps[0].Priority + for _, e := range eps[1:] { + if e.Priority > best { + best = e.Priority + } + } + out := make([]endpoint.Snapshot, 0, len(eps)) + for _, e := range eps { + if e.Priority == best { + out = append(out, e) + } + } + return out +} diff --git a/internal/router/scheduler/hashring.go b/internal/router/scheduler/hashring.go new file mode 100644 index 0000000..bbeea0e --- /dev/null +++ b/internal/router/scheduler/hashring.go @@ -0,0 +1,108 @@ +package scheduler + +import ( + "slices" + + "github.com/cespare/xxhash/v2" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// vnodesPerEndpoint is how many virtual nodes each endpoint places on the ring. +// More nodes improve the distribution; the only cost is the sort at pick time. +const vnodesPerEndpoint = 64 + +type vnode struct { + hash uint64 + idx int // index into the eps slice +} + +// pickBoundedLoad maps a key onto an endpoint by consistent hashing, walking +// the ring past any endpoint that is overloaded. +// +// "Overloaded" is measured by inflight count: an endpoint carrying more than +// loadFactor times the average is skipped. If every endpoint is overloaded it +// returns to the first target — insisting on the target beats losing affinity +// entirely, because at that point the whole pool is equally loaded anyway. +// +// It returns -1 only when there are no candidates. +func pickBoundedLoad(key string, eps []endpoint.Snapshot, loadFactor float64) int { + if len(eps) == 0 { + return -1 + } + if len(eps) == 1 { + return 0 + } + + ring := buildRing(eps) + h := xxhash.Sum64String(key) + + // Average load and the resulting ceiling. + totalInflight := 0 + for _, e := range eps { + totalInflight += e.Inflight + } + avg := float64(totalInflight) / float64(len(eps)) + ceiling := avg * loadFactor + if ceiling < 1 { + // Do not let the ceiling fall below 1 at very low load, or affinity + // would vanish entirely on an idle pool. + ceiling = 1 + } + + start, _ := slices.BinarySearchFunc(ring, h, func(v vnode, target uint64) int { + switch { + case v.hash < target: + return -1 + case v.hash > target: + return 1 + default: + return 0 + } + }) + + first := -1 + seen := make(map[int]struct{}, len(eps)) + for i := range ring { + v := ring[(start+i)%len(ring)] + if _, dup := seen[v.idx]; dup { + continue + } + seen[v.idx] = struct{}{} + if first < 0 { + first = v.idx + } + if float64(eps[v.idx].Inflight) <= ceiling { + return v.idx + } + if len(seen) == len(eps) { + break + } + } + return first +} + +func buildRing(eps []endpoint.Snapshot) []vnode { + ring := make([]vnode, 0, len(eps)*vnodesPerEndpoint) + var buf []byte + for i, e := range eps { + for v := range vnodesPerEndpoint { + buf = buf[:0] + buf = append(buf, e.ID...) + buf = append(buf, '#') + buf = append(buf, byte('0'+v/10), byte('0'+v%10)) + ring = append(ring, vnode{hash: xxhash.Sum64(buf), idx: i}) + } + } + slices.SortFunc(ring, func(a, b vnode) int { + switch { + case a.hash < b.hash: + return -1 + case a.hash > b.hash: + return 1 + default: + return 0 + } + }) + return ring +} diff --git a/internal/router/scheduler/pickers.go b/internal/router/scheduler/pickers.go new file mode 100644 index 0000000..7a05151 --- /dev/null +++ b/internal/router/scheduler/pickers.go @@ -0,0 +1,147 @@ +package scheduler + +import ( + "math" + "math/rand/v2" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// RandSource is the pickers' randomness source; tests can pin it. +type RandSource interface { + Float64() float64 + IntN(n int) int +} + +type globalRand struct{} + +func (globalRand) Float64() float64 { return rand.Float64() } +func (globalRand) IntN(n int) int { return rand.IntN(n) } + +// MaxScore picks the highest-scoring candidate. +// +// Deterministic and easy to explain, but prone to herding: with equally scored +// replicas, consecutive requests all choose the same one. Prefer P2C once there +// are many replicas. +func MaxScore() Picker { return maxScorePicker{} } + +type maxScorePicker struct{} + +func (maxScorePicker) Name() string { return "max-score" } + +func (maxScorePicker) Pick(_ *Ctx, _ []endpoint.Snapshot, scores []float64) int { + best, bestIdx := math.Inf(-1), 0 + for i, s := range scores { + if s > best { + best, bestIdx = s, i + } + } + return bestIdx +} + +// P2C samples two candidates at random and takes the better one +// (power of two choices). +// +// It breaks herding: even when two independent requests see the same "best" +// replica, they may draw different pairs and land elsewhere. This is the +// default picker once there are many replicas. +func P2C() Picker { return &p2cPicker{rng: globalRand{}} } + +// P2CWith returns a P2C picker with a specific randomness source, for tests. +func P2CWith(r RandSource) Picker { return &p2cPicker{rng: r} } + +type p2cPicker struct{ rng RandSource } + +func (*p2cPicker) Name() string { return "p2c" } + +func (p *p2cPicker) Pick(_ *Ctx, _ []endpoint.Snapshot, scores []float64) int { + n := len(scores) + if n == 0 { + return -1 + } + if n == 1 { + return 0 + } + a := p.rng.IntN(n) + b := p.rng.IntN(n - 1) + if b >= a { + b++ // ensure it differs from a + } + if scores[a] >= scores[b] { + return a + } + return b +} + +// Softmax treats scores as logits and samples from them (Dynamo's +// router_temperature). +// +// A temperature of 0 degenerates to argmax. A small temperature breaks the +// oscillation that synchronized decisions cause among equally scored replicas; +// a large one approaches uniform random. +func Softmax(temperature float64) Picker { + return &softmaxPicker{temp: temperature, rng: globalRand{}} +} + +// SoftmaxWith returns a Softmax picker with a specific randomness source, for tests. +func SoftmaxWith(temperature float64, r RandSource) Picker { + return &softmaxPicker{temp: temperature, rng: r} +} + +type softmaxPicker struct { + temp float64 + rng RandSource +} + +func (*softmaxPicker) Name() string { return "softmax" } + +func (p *softmaxPicker) Pick(c *Ctx, eps []endpoint.Snapshot, scores []float64) int { + if p.temp <= 0 { + return maxScorePicker{}.Pick(c, eps, scores) + } + n := len(scores) + if n == 0 { + return -1 + } + + // Subtract the maximum to avoid overflow. + maxS := math.Inf(-1) + for _, s := range scores { + maxS = max(maxS, s) + } + weights := make([]float64, n) + var sum float64 + for i, s := range scores { + w := math.Exp((s - maxS) / p.temp) + weights[i] = w + sum += w + } + if sum <= 0 || math.IsInf(sum, 0) || math.IsNaN(sum) { + return maxScorePicker{}.Pick(c, eps, scores) + } + + r := p.rng.Float64() * sum + for i, w := range weights { + r -= w + if r <= 0 { + return i + } + } + return n - 1 +} + +// RoundRobin cycles through candidates. For debugging and simple setups. +func RoundRobin() Picker { return &rrPicker{} } + +type rrPicker struct{ n int } + +func (*rrPicker) Name() string { return "round-robin" } + +func (p *rrPicker) Pick(_ *Ctx, _ []endpoint.Snapshot, scores []float64) int { + if len(scores) == 0 { + return -1 + } + i := p.n % len(scores) + p.n++ + return i +} diff --git a/internal/router/scheduler/profiles.go b/internal/router/scheduler/profiles.go new file mode 100644 index 0000000..eabfcfd --- /dev/null +++ b/internal/router/scheduler/profiles.go @@ -0,0 +1,118 @@ +package scheduler + +// ProfileConfig holds thresholds shared by the built-in profiles. +type ProfileConfig struct { + // TargetWaiting is the queue scorer's "healthy buffer" target. + TargetWaiting float64 + // TargetInflight normalizes the inflight scorer. + TargetInflight float64 + // KVCeiling is the KV occupancy ceiling. + KVCeiling float64 + // SessionLoadFactor is the load ceiling multiplier for session affinity. + SessionLoadFactor float64 +} + +// DefaultProfileConfig returns sensible starting values. +// These numbers should be measured with ctl bench (§10); these are only seeds. +func DefaultProfileConfig() ProfileConfig { + return ProfileConfig{ + TargetWaiting: 4, + TargetInflight: 8, + KVCeiling: 0.9, + SessionLoadFactor: 1.25, + } +} + +func (c *ProfileConfig) normalize() { + d := DefaultProfileConfig() + if c.TargetWaiting <= 0 { + c.TargetWaiting = d.TargetWaiting + } + if c.TargetInflight <= 0 { + c.TargetInflight = d.TargetInflight + } + if c.KVCeiling <= 0 || c.KVCeiling > 1 { + c.KVCeiling = d.KVCeiling + } + if c.SessionLoadFactor < 1 { + c.SessionLoadFactor = d.SessionLoadFactor + } +} + +// Chat is the profile for multi-turn conversations and agent workloads. +// Session affinity dominates, to preserve prefix cache reuse. +func Chat(c ProfileConfig) *Profile { + c.normalize() + return &Profile{ + Name: "chat", + Filters: []Filter{BreakerOpen(), PoolPriority()}, + Scorers: []WeightedScorer{ + {Session(c.SessionLoadFactor), 3.0}, + {Queue(c.TargetWaiting), 1.0}, + {KV(c.KVCeiling), 1.0}, + {Inflight(c.TargetInflight), 1.0}, + {Weight(), 0.5}, + }, + Picker: P2C(), + } +} + +// Batch is the profile for bulk work: throughput matters, TTFT does not, so +// finding the shortest queue outranks stickiness. +func Batch(c ProfileConfig) *Profile { + c.normalize() + return &Profile{ + Name: "batch", + Filters: []Filter{BreakerOpen(), PoolPriority()}, + Scorers: []WeightedScorer{ + {Queue(c.TargetWaiting), 3.0}, + {KV(c.KVCeiling), 1.0}, + {Load(30), 1.0}, + {Inflight(c.TargetInflight), 1.0}, + {Weight(), 1.0}, + }, + Picker: MaxScore(), + } +} + +// LongContext is the profile for long-prompt workloads. KV pressure is the +// bottleneck there, so the kv scorer dominates. +func LongContext(c ProfileConfig) *Profile { + c.normalize() + return &Profile{ + Name: "long-context", + Filters: []Filter{BreakerOpen(), PoolPriority()}, + Scorers: []WeightedScorer{ + {KV(c.KVCeiling), 3.0}, + {Session(c.SessionLoadFactor), 2.0}, + {Queue(c.TargetWaiting), 1.0}, + {Inflight(c.TargetInflight), 1.0}, + }, + Picker: Softmax(0.2), + } +} + +// Debug is a round-robin profile for troubleshooting. +func Debug() *Profile { + return &Profile{ + Name: "debug", + Filters: []Filter{BreakerOpen()}, + Picker: RoundRobin(), + } +} + +// Builtin returns a built-in profile by name, or nil if the name is unknown. +func Builtin(name string, c ProfileConfig) *Profile { + switch name { + case "chat": + return Chat(c) + case "batch": + return Batch(c) + case "long-context": + return LongContext(c) + case "debug": + return Debug() + default: + return nil + } +} diff --git a/internal/router/scheduler/scheduler.go b/internal/router/scheduler/scheduler.go new file mode 100644 index 0000000..4ab93fe --- /dev/null +++ b/internal/router/scheduler/scheduler.go @@ -0,0 +1,170 @@ +// Package scheduler answers "where should this request go?" (plan §7.3). +// +// Instead of one monolithic "strategy" function it is a plugin pipeline: +// +// filter -> drops candidates (hard constraints) +// score -> rates each candidate 0..1, combined by weight +// pick -> chooses one from the combined scores +// +// A "strategy" then reduces to a set of scorer weights: the difference between +// the chat and batch profiles is weights, not code. +package scheduler + +import ( + "errors" + "fmt" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// Request holds the request fields a routing decision depends on. +type Request struct { + // Model is the resolved model name. + Model string + // SessionID drives session affinity (x-session-id header or the user field). + SessionID string + // PromptPrefix is the leading portion of the prompt, for prefix awareness. + PromptPrefix string + // PromptTokens is the estimated prompt length in tokens. + PromptTokens int + // MaxTokens is the requested generation length. + MaxTokens int +} + +// Ctx is the context of a single scheduling decision. +type Ctx struct { + Req Request + Now time.Time +} + +// Filter drops candidates. +type Filter interface { + Name() string + Filter(c *Ctx, eps []endpoint.Snapshot) []endpoint.Snapshot +} + +// Scorer rates each candidate between 0 and 1. The returned slice must match +// eps in both order and length. +type Scorer interface { + Name() string + Score(c *Ctx, eps []endpoint.Snapshot) []float64 +} + +// Picker chooses one candidate from the combined scores and returns its index. +type Picker interface { + Name() string + Pick(c *Ctx, eps []endpoint.Snapshot, scores []float64) int +} + +// WeightedScorer pairs a scorer with its weight in a profile. +type WeightedScorer struct { + Scorer Scorer + Weight float64 +} + +// Profile is a scorer set with weights for one class of request (plan §7.3). +type Profile struct { + Name string + Filters []Filter + Scorers []WeightedScorer + Picker Picker + + // StalePenalty multiplies the total score of an endpoint with stale + // metrics. Stale endpoints are never dropped — if they were, a fully stale + // pool would leave no candidates — but they rank below fresh ones. + // Zero means 0.5. + StalePenalty float64 +} + +// Decision is a routing choice together with its rationale. +// +// Recording the rationale is not decoration: tuning the prefix and queue +// weights is impossible without seeing why decisions came out the way they +// did (§7.5). +type Decision struct { + Endpoint endpoint.Snapshot + Score float64 + // Scores maps scorer name to raw score for the chosen endpoint. + Scores map[string]float64 + // Dropped maps filter name to how many candidates it removed. + Dropped map[string]int + // Considered is how many candidates survived the filters. + Considered int +} + +var ( + // ErrNoEndpoints means no endpoint serves the requested model. + ErrNoEndpoints = errors.New("no endpoint serves this model") + // ErrAllFiltered means every candidate was removed by the filters. + ErrAllFiltered = errors.New("all candidates were filtered out") +) + +// Schedule chooses one endpoint from the candidates. +func (p *Profile) Schedule(c *Ctx, eps []endpoint.Snapshot) (Decision, error) { + if len(eps) == 0 { + return Decision{}, ErrNoEndpoints + } + + dropped := make(map[string]int) + cand := eps + for _, f := range p.Filters { + before := len(cand) + cand = f.Filter(c, cand) + if n := before - len(cand); n > 0 { + dropped[f.Name()] = n + } + if len(cand) == 0 { + return Decision{Dropped: dropped}, fmt.Errorf("%w: %v", ErrAllFiltered, dropped) + } + } + + total := make([]float64, len(cand)) + raw := make([]map[string]float64, len(cand)) + for i := range raw { + raw[i] = make(map[string]float64, len(p.Scorers)) + } + + for _, ws := range p.Scorers { + if ws.Weight == 0 || ws.Scorer == nil { + continue + } + s := ws.Scorer.Score(c, cand) + if len(s) != len(cand) { + // A misbehaving scorer must not break the whole decision; skip it. + continue + } + name := ws.Scorer.Name() + for i, v := range s { + raw[i][name] = v + total[i] += ws.Weight * v + } + } + + penalty := p.StalePenalty + if penalty <= 0 { + penalty = 0.5 + } + for i := range cand { + if cand[i].Stale { + total[i] *= penalty + } + } + + picker := p.Picker + if picker == nil { + picker = MaxScore() + } + idx := picker.Pick(c, cand, total) + if idx < 0 || idx >= len(cand) { + idx = 0 + } + + return Decision{ + Endpoint: cand[idx], + Score: total[idx], + Scores: raw[idx], + Dropped: dropped, + Considered: len(cand), + }, nil +} diff --git a/internal/router/scheduler/scheduler_test.go b/internal/router/scheduler/scheduler_test.go new file mode 100644 index 0000000..1509b9c --- /dev/null +++ b/internal/router/scheduler/scheduler_test.go @@ -0,0 +1,385 @@ +package scheduler + +import ( + "fmt" + "math" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +func ep(id string, mut ...func(*endpoint.Snapshot)) endpoint.Snapshot { + s := endpoint.Snapshot{ + Info: endpoint.Info{ID: id, URL: "http://" + id, Models: []string{"m"}, Weight: 1}, + Breaker: breaker.Closed, + } + for _, f := range mut { + f(&s) + } + return s +} + +func waiting(n int) func(*endpoint.Snapshot) { + return func(s *endpoint.Snapshot) { s.Stats.Waiting = n } +} +func kv(v float64) func(*endpoint.Snapshot) { + return func(s *endpoint.Snapshot) { s.Stats.KVCacheUsage = v } +} +func inflight(n int) func(*endpoint.Snapshot) { + return func(s *endpoint.Snapshot) { s.Inflight = n } +} + +func newCtx(req Request) *Ctx { return &Ctx{Req: req, Now: time.Now()} } + +func TestEmptyCandidateList(t *testing.T) { + p := Chat(ProfileConfig{}) + if _, err := p.Schedule(newCtx(Request{Model: "m"}), nil); err != ErrNoEndpoints { + t.Errorf("expected ErrNoEndpoints, got %v", err) + } +} + +func TestPicksShortestQueue(t *testing.T) { + p := &Profile{ + Name: "t", + Scorers: []WeightedScorer{{Queue(4), 1}}, + Picker: MaxScore(), + } + eps := []endpoint.Snapshot{ + ep("a", waiting(3)), + ep("b", waiting(0)), + ep("c", waiting(2)), + } + d, err := p.Schedule(newCtx(Request{Model: "m"}), eps) + if err != nil { + t.Fatal(err) + } + if d.Endpoint.ID != "b" { + t.Errorf("did not pick the emptiest queue: %s", d.Endpoint.ID) + } + if d.Scores["queue"] != 1 { + t.Errorf("queue score %v, expected 1", d.Scores["queue"]) + } +} + +func TestKVCeilingAppliesHardPenalty(t *testing.T) { + // A replica just above the ceiling must rank far below one just under it. + s := KV(0.9) + got := s.Score(newCtx(Request{}), []endpoint.Snapshot{ + ep("under", kv(0.89)), + ep("over", kv(0.91)), + }) + if got[0] <= got[1]*2 { + t.Errorf("ceiling penalty too weak: under=%v over=%v", got[0], got[1]) + } +} + +func TestInflightCompensatesScrapeLag(t *testing.T) { + // Both replicas report identical (stale) engine metrics, but the router has + // already dispatched to one of them. The inflight scorer must notice. + p := &Profile{ + Name: "t", + Scorers: []WeightedScorer{{Queue(4), 1}, {Inflight(8), 1}}, + Picker: MaxScore(), + } + eps := []endpoint.Snapshot{ + ep("busy", waiting(0), inflight(5)), + ep("idle", waiting(0), inflight(0)), + } + d, err := p.Schedule(newCtx(Request{Model: "m"}), eps) + if err != nil { + t.Fatal(err) + } + if d.Endpoint.ID != "idle" { + t.Errorf("inflight was ignored, picked %s", d.Endpoint.ID) + } +} + +func TestOpenBreakerIsFiltered(t *testing.T) { + p := &Profile{ + Name: "t", + Filters: []Filter{BreakerOpen()}, + Scorers: []WeightedScorer{{Queue(4), 1}}, + Picker: MaxScore(), + } + eps := []endpoint.Snapshot{ + ep("open", waiting(0), func(s *endpoint.Snapshot) { s.Breaker = breaker.Open }), + ep("closed", waiting(3)), + } + d, err := p.Schedule(newCtx(Request{Model: "m"}), eps) + if err != nil { + t.Fatal(err) + } + if d.Endpoint.ID != "closed" { + t.Errorf("picked an endpoint with an open circuit: %s", d.Endpoint.ID) + } + if d.Dropped["breaker-open"] != 1 { + t.Errorf("the drop was not recorded: %v", d.Dropped) + } +} + +func TestHalfOpenBreakerSurvivesFilter(t *testing.T) { + // The probe request must get through, or the circuit never closes. + p := &Profile{Name: "t", Filters: []Filter{BreakerOpen()}, Picker: MaxScore()} + eps := []endpoint.Snapshot{ + ep("half", func(s *endpoint.Snapshot) { s.Breaker = breaker.HalfOpen }), + } + if _, err := p.Schedule(newCtx(Request{Model: "m"}), eps); err != nil { + t.Errorf("half-open endpoint was filtered out: %v", err) + } +} + +func TestAllCandidatesFiltered(t *testing.T) { + p := &Profile{Name: "t", Filters: []Filter{BreakerOpen()}, Picker: MaxScore()} + eps := []endpoint.Snapshot{ + ep("a", func(s *endpoint.Snapshot) { s.Breaker = breaker.Open }), + } + _, err := p.Schedule(newCtx(Request{Model: "m"}), eps) + if err == nil { + t.Fatal("expected an error") + } + if got := fmt.Sprint(err); got == "" { + t.Error("error message is empty") + } +} + +func TestStaleEndpointIsPenalizedNotDropped(t *testing.T) { + p := &Profile{ + Name: "t", + Scorers: []WeightedScorer{{Queue(4), 1}}, + Picker: MaxScore(), + StalePenalty: 0.5, + } + eps := []endpoint.Snapshot{ + ep("stale", waiting(0), func(s *endpoint.Snapshot) { s.Stale = true }), + ep("fresh", waiting(1)), + } + d, err := p.Schedule(newCtx(Request{Model: "m"}), eps) + if err != nil { + t.Fatal(err) + } + // stale: 1.0*0.5 = 0.5 ; fresh: 0.75 + if d.Endpoint.ID != "fresh" { + t.Errorf("stale endpoint was preferred: %s", d.Endpoint.ID) + } + + // When everything is stale a choice must still be made. + allStale := []endpoint.Snapshot{ + ep("a", func(s *endpoint.Snapshot) { s.Stale = true }), + ep("b", func(s *endpoint.Snapshot) { s.Stale = true }), + } + if _, err := p.Schedule(newCtx(Request{Model: "m"}), allStale); err != nil { + t.Errorf("no candidates left when the whole pool is stale: %v", err) + } +} + +func TestPoolPriority(t *testing.T) { + p := &Profile{ + Name: "t", + Filters: []Filter{PoolPriority()}, + Scorers: []WeightedScorer{{Queue(4), 1}}, + Picker: MaxScore(), + } + eps := []endpoint.Snapshot{ + ep("h100", waiting(3), func(s *endpoint.Snapshot) { s.Priority = 100; s.Pool = "ib-a/h100" }), + ep("a100", waiting(0), func(s *endpoint.Snapshot) { s.Priority = 50; s.Pool = "ib-b/a100" }), + } + d, err := p.Schedule(newCtx(Request{Model: "m"}), eps) + if err != nil { + t.Fatal(err) + } + // The lower-priority pool must not win despite being emptier (spillover is phase 5). + if d.Endpoint.ID != "h100" { + t.Errorf("pool priority was ignored: %s", d.Endpoint.ID) + } +} + +func TestSessionStickiness(t *testing.T) { + p := &Profile{ + Name: "t", + Scorers: []WeightedScorer{{Session(1.25), 3}, {Queue(4), 1}}, + Picker: MaxScore(), + } + eps := []endpoint.Snapshot{ep("a"), ep("b"), ep("c")} + + first, err := p.Schedule(newCtx(Request{Model: "m", SessionID: "session-1"}), eps) + if err != nil { + t.Fatal(err) + } + for range 20 { + d, err := p.Schedule(newCtx(Request{Model: "m", SessionID: "session-1"}), eps) + if err != nil { + t.Fatal(err) + } + if d.Endpoint.ID != first.Endpoint.ID { + t.Fatalf("same session landed on a different replica: %s != %s", d.Endpoint.ID, first.Endpoint.ID) + } + } +} + +func TestSessionsSpreadAcrossReplicas(t *testing.T) { + p := &Profile{Name: "t", Scorers: []WeightedScorer{{Session(1.25), 1}}, Picker: MaxScore()} + eps := []endpoint.Snapshot{ep("a"), ep("b"), ep("c")} + + counts := map[string]int{} + for i := range 300 { + d, err := p.Schedule(newCtx(Request{Model: "m", SessionID: fmt.Sprintf("s-%d", i)}), eps) + if err != nil { + t.Fatal(err) + } + counts[d.Endpoint.ID]++ + } + if len(counts) != 3 { + t.Fatalf("sessions did not spread across all replicas: %v", counts) + } + for id, n := range counts { + if n < 50 { + t.Errorf("%s received only %d sessions; distribution too skewed: %v", id, n, counts) + } + } +} + +func TestStickinessYieldsUnderOverload(t *testing.T) { + // Bounded load: when the target replica is far above average, affinity must + // move on. Otherwise prefix affinity drowns a single replica (plan §7.5). + eps := []endpoint.Snapshot{ + ep("a", inflight(100)), + ep("b", inflight(0)), + ep("c", inflight(0)), + } + s := Session(1.25) + var chosen string + for i := range 50 { + got := s.Score(newCtx(Request{SessionID: fmt.Sprintf("k%d", i)}), eps) + for j, v := range got { + if v == 1 { + chosen = eps[j].ID + } + } + if chosen == "a" { + t.Fatalf("overloaded replica (inflight=100) was chosen anyway") + } + } +} + +func TestNoSessionIsNeutral(t *testing.T) { + s := Session(1.25) + got := s.Score(newCtx(Request{}), []endpoint.Snapshot{ep("a"), ep("b")}) + for _, v := range got { + if v != 0 { + t.Errorf("scored without a session id: %v", got) + } + } +} + +// --- picker tests --- + +type fixedRand struct { + floats []float64 + ints []int + fi, ii int +} + +func (r *fixedRand) Float64() float64 { + v := r.floats[r.fi%len(r.floats)] + r.fi++ + return v +} + +func (r *fixedRand) IntN(n int) int { + v := r.ints[r.ii%len(r.ints)] + r.ii++ + return v % n +} + +func TestP2CPicksTheBetterOfTwo(t *testing.T) { + p := P2CWith(&fixedRand{ints: []int{0, 0}}) // a ve (0>=0 -> b) + eps := []endpoint.Snapshot{ep("a"), ep("b"), ep("c")} + idx := p.Pick(newCtx(Request{}), eps, []float64{0.1, 0.9, 0.5}) + if idx != 1 { + t.Errorf("P2C did not pick the higher score: idx=%d", idx) + } +} + +func TestP2CNeverDrawsTheSameCandidateTwice(t *testing.T) { + // When IntN(n-1) lands on a it must be bumped by one. + for a := range 5 { + p := &p2cPicker{rng: &fixedRand{ints: []int{a, a}}} + eps := make([]endpoint.Snapshot, 5) + for i := range eps { + eps[i] = ep(fmt.Sprintf("e%d", i)) + } + scores := make([]float64, 5) + scores[a] = -1 // make a worst: picking it means the same candidate was drawn twice + if idx := p.Pick(newCtx(Request{}), eps, scores); idx == a { + t.Errorf("same candidate drawn twice for a=%d", a) + } + } +} + +func TestSoftmaxAtZeroTemperatureIsArgmax(t *testing.T) { + p := SoftmaxWith(0, &fixedRand{floats: []float64{0.99}}) + idx := p.Pick(newCtx(Request{}), nil, []float64{0.2, 0.8, 0.5}) + if idx != 1 { + t.Errorf("temperature=0 did not produce argmax: %d", idx) + } +} + +func TestSoftmaxSpreadsAtHighTemperature(t *testing.T) { + // At a high temperature a low-scoring candidate must be reachable. + p := SoftmaxWith(10, &fixedRand{floats: []float64{0.99}}) + idx := p.Pick(newCtx(Request{}), nil, []float64{1.0, 0.0}) + if idx != 1 { + t.Errorf("no spread at high temperature: %d", idx) + } +} + +func TestSoftmaxDoesNotOverflow(t *testing.T) { + p := SoftmaxWith(0.001, &fixedRand{floats: []float64{0.5}}) + idx := p.Pick(newCtx(Request{}), nil, []float64{1000, 999}) + if idx < 0 || math.IsNaN(float64(idx)) { + t.Errorf("softmax broke on large scores: %d", idx) + } +} + +func TestRoundRobinCycles(t *testing.T) { + p := RoundRobin() + scores := []float64{0, 0, 0} + got := []int{} + for range 5 { + got = append(got, p.Pick(newCtx(Request{}), nil, scores)) + } + want := []int{0, 1, 2, 0, 1} + for i := range want { + if got[i] != want[i] { + t.Fatalf("round-robin order is wrong: %v", got) + } + } +} + +func TestBuiltinProfiles(t *testing.T) { + for _, name := range []string{"chat", "batch", "long-context", "debug"} { + if p := Builtin(name, ProfileConfig{}); p == nil { + t.Errorf("built-in profile not found: %s", name) + } + } + if p := Builtin("no-such-profile", ProfileConfig{}); p != nil { + t.Error("expected nil for an unknown profile") + } +} + +func TestBuiltinProfilesSchedule(t *testing.T) { + eps := []endpoint.Snapshot{ep("a", waiting(1)), ep("b", waiting(0), kv(0.5))} + for _, name := range []string{"chat", "batch", "long-context", "debug"} { + p := Builtin(name, ProfileConfig{}) + d, err := p.Schedule(newCtx(Request{Model: "m", SessionID: "s1"}), eps) + if err != nil { + t.Errorf("%s: %v", name, err) + continue + } + if d.Endpoint.ID == "" { + t.Errorf("%s: picked an empty endpoint", name) + } + } +} diff --git a/internal/router/scheduler/scorers.go b/internal/router/scheduler/scorers.go new file mode 100644 index 0000000..ab4471d --- /dev/null +++ b/internal/router/scheduler/scorers.go @@ -0,0 +1,180 @@ +package scheduler + +import ( + "math" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// scorerFunc helps define scorers that rate each endpoint independently. +type scorerFunc struct { + name string + f func(c *Ctx, e endpoint.Snapshot) float64 +} + +func (s scorerFunc) Name() string { return s.name } + +func (s scorerFunc) Score(c *Ctx, eps []endpoint.Snapshot) []float64 { + out := make([]float64, len(eps)) + for i, e := range eps { + out[i] = clamp01(s.f(c, e)) + } + return out +} + +func clamp01(v float64) float64 { + if math.IsNaN(v) { + return 0 + } + return min(max(v, 0), 1) +} + +// Queue scores by queue depth: 1 - min(1, waiting/target). +// +// target is the "healthy buffer" size (plan §7.2.1): enough queued requests +// for continuous batching to stay efficient, but not so many that queueing +// delay creeps back in. Measure it with ctl bench rather than guessing. +func Queue(target float64) Scorer { + if target <= 0 { + target = 4 + } + return scorerFunc{ + name: "queue", + f: func(_ *Ctx, e endpoint.Snapshot) float64 { + return 1 - float64(e.Stats.Waiting)/target + }, + } +} + +// KV scores by KV cache occupancy: 1 - usage. +// +// Above the ceiling a hard penalty applies: going from 90% to 95% is not the +// same as going from 10% to 15% — a replica near the ceiling starts preempting +// and dropping requests. +func KV(ceiling float64) Scorer { + if ceiling <= 0 || ceiling > 1 { + ceiling = 0.9 + } + return scorerFunc{ + name: "kv", + f: func(_ *Ctx, e endpoint.Snapshot) float64 { + u := e.Stats.KVCacheUsage + s := 1 - u + if u >= ceiling { + s *= 0.1 + } + return s + }, + } +} + +// Inflight scores by the unfinished requests the router itself is tracking. +// +// With a 2s scrape interval the engine metrics are stale at decision time, so +// back-to-back requests all see the same "idle" replica and all pile onto it. +// This scorer closes that gap, which is why every profile keeps it on (§7.6). +func Inflight(target float64) Scorer { + if target <= 0 { + target = 8 + } + return scorerFunc{ + name: "inflight", + f: func(_ *Ctx, e endpoint.Snapshot) float64 { + w := e.Weight + if w <= 0 { + w = 1 + } + return 1 - float64(e.Inflight)/(target*w) + }, + } +} + +// Load is AIBrix's load-balance score: running / drain_rate. +// +// Running alone cannot spot a slow replica — one processing the same number of +// requests at half the speed is twice as loaded. While drain_rate is still +// unmeasured (the first seconds) it returns a neutral 1 and leaves the decision +// to the other scorers. +func Load(targetSeconds float64) Scorer { + if targetSeconds <= 0 { + targetSeconds = 30 + } + return scorerFunc{ + name: "load", + f: func(_ *Ctx, e endpoint.Snapshot) float64 { + if e.DrainRate <= 0 { + return 1 + } + // Estimated seconds to drain, normalized against the target. + drainTime := float64(e.Stats.Running) / e.DrainRate + return 1 - drainTime/targetSeconds + }, + } +} + +// Weight scores by static capacity weight. +// +// In heterogeneous pools (A100s alongside H100s) raw queue length misleads: +// the same queue drains faster on the faster card. The weight comes from +// ctl bench (§6.2.1) and should never be guessed by hand. +func Weight() Scorer { + return weightScorer{} +} + +type weightScorer struct{} + +func (weightScorer) Name() string { return "weight" } + +func (weightScorer) Score(_ *Ctx, eps []endpoint.Snapshot) []float64 { + maxW := 0.0 + for _, e := range eps { + w := e.Weight + if w <= 0 { + w = 1 + } + maxW = max(maxW, w) + } + out := make([]float64, len(eps)) + if maxW <= 0 { + return out + } + for i, e := range eps { + w := e.Weight + if w <= 0 { + w = 1 + } + out[i] = w / maxW + } + return out +} + +// Session scores session affinity (plan §7.5, tier 0). +// +// Bounded-load consistent hashing: a session maps to a replica, but if that +// replica carries more than loadFactor times the average it moves on around the +// ring. Unbounded stickiness drowns a single replica: cache hits go up while +// p99 collapses. +// +// With no session id every candidate scores 0, leaving the decision to the +// other scorers. +func Session(loadFactor float64) Scorer { + if loadFactor < 1 { + loadFactor = 1.25 + } + return sessionScorer{loadFactor: loadFactor} +} + +type sessionScorer struct{ loadFactor float64 } + +func (sessionScorer) Name() string { return "session" } + +func (s sessionScorer) Score(c *Ctx, eps []endpoint.Snapshot) []float64 { + out := make([]float64, len(eps)) + if c.Req.SessionID == "" || len(eps) == 0 { + return out + } + if idx := pickBoundedLoad(c.Req.SessionID, eps, s.loadFactor); idx >= 0 { + out[idx] = 1 + } + return out +} diff --git a/internal/router/scrape/parse.go b/internal/router/scrape/parse.go new file mode 100644 index 0000000..a697129 --- /dev/null +++ b/internal/router/scrape/parse.go @@ -0,0 +1,173 @@ +// Package scrape reads engines' /metrics output and turns it into something +// the router can make decisions from (plan §7.6). +package scrape + +import ( + "fmt" + "io" + "strconv" + + dto "github.com/prometheus/client_model/go" + "github.com/prometheus/common/expfmt" + "github.com/prometheus/common/model" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// Mapping says which metric names the contract fields are read from. +// +// Being configurable is essential: engine versions rename their metrics +// (plan §14). Once the agent reports an engine version the control plane can +// pick the right mapping; in phase 1 it comes from configuration. +type Mapping struct { + Running string + Waiting string + KVCacheUsage string + + CacheConfig string + PromptTokens string + GenerationTokens string + SuccessTotal string +} + +// VLLMMapping returns vLLM's current metric names. +func VLLMMapping() Mapping { + return Mapping{ + Running: "vllm:num_requests_running", + Waiting: "vllm:num_requests_waiting", + KVCacheUsage: "vllm:kv_cache_usage_perc", + CacheConfig: "vllm:cache_config_info", + PromptTokens: "vllm:prompt_tokens_total", + GenerationTokens: "vllm:generation_tokens_total", + SuccessTotal: "vllm:request_success_total", + } +} + +func (m *Mapping) normalize() { + d := VLLMMapping() + if m.Running == "" { + m.Running = d.Running + } + if m.Waiting == "" { + m.Waiting = d.Waiting + } + if m.KVCacheUsage == "" { + m.KVCacheUsage = d.KVCacheUsage + } + if m.CacheConfig == "" { + m.CacheConfig = d.CacheConfig + } + if m.PromptTokens == "" { + m.PromptTokens = d.PromptTokens + } + if m.GenerationTokens == "" { + m.GenerationTokens = d.GenerationTokens + } + if m.SuccessTotal == "" { + m.SuccessTotal = d.SuccessTotal + } +} + +// MissingError reports that mandatory contract metrics are absent. +type MissingError struct { + Missing []string +} + +func (e *MissingError) Error() string { + return fmt.Sprintf("engine metric contract incomplete: %v", e.Missing) +} + +// Parse converts Prometheus text-format metrics into Stats. +// +// If any of the mandatory three (running, waiting, kv usage) is absent it +// returns *MissingError: routing to that endpoint would be blind. +func Parse(r io.Reader, m Mapping) (endpoint.Stats, error) { + m.normalize() + + // TextParser's zero value is invalid; the scheme must be explicit. + // The legacy scheme accepts vLLM's names containing ':'. + p := expfmt.NewTextParser(model.LegacyValidation) + fams, err := p.TextToMetricFamilies(r) + if err != nil { + return endpoint.Stats{}, fmt.Errorf("failed to parse metrics: %w", err) + } + + var s endpoint.Stats + var missing []string + + if v, ok := firstValue(fams, m.Running); ok { + s.Running = int(v) + } else { + missing = append(missing, m.Running) + } + if v, ok := firstValue(fams, m.Waiting); ok { + s.Waiting = int(v) + } else { + missing = append(missing, m.Waiting) + } + if v, ok := firstValue(fams, m.KVCacheUsage); ok { + s.KVCacheUsage = v + } else { + missing = append(missing, m.KVCacheUsage) + } + if len(missing) > 0 { + return endpoint.Stats{}, &MissingError{Missing: missing} + } + + // Optional fields: silently stay zero when absent. + s.PromptTokens, _ = sumValues(fams, m.PromptTokens) + s.GenerationTokens, _ = sumValues(fams, m.GenerationTokens) + // request_success_total is split by finished_reason; sum across labels. + s.SuccessTotal, _ = sumValues(fams, m.SuccessTotal) + + if f, ok := fams[m.CacheConfig]; ok && len(f.Metric) > 0 { + for _, lp := range f.Metric[0].Label { + switch lp.GetName() { + case "block_size": + s.BlockSize, _ = strconv.Atoi(lp.GetValue()) + case "num_gpu_blocks": + s.NumGPUBlocks, _ = strconv.Atoi(lp.GetValue()) + } + } + } + + return s, nil +} + +// metricValue extracts a numeric value regardless of the metric type. +func metricValue(mt *dto.Metric) (float64, bool) { + switch { + case mt.Gauge != nil: + return mt.Gauge.GetValue(), true + case mt.Counter != nil: + return mt.Counter.GetValue(), true + case mt.Untyped != nil: + return mt.Untyped.GetValue(), true + default: + return 0, false + } +} + +func firstValue(fams map[string]*dto.MetricFamily, name string) (float64, bool) { + f, ok := fams[name] + if !ok || len(f.Metric) == 0 { + return 0, false + } + return metricValue(f.Metric[0]) +} + +func sumValues(fams map[string]*dto.MetricFamily, name string) (float64, bool) { + f, ok := fams[name] + if !ok { + return 0, false + } + var sum float64 + var any bool + for _, mt := range f.Metric { + if v, ok := metricValue(mt); ok { + sum += v + any = true + } + } + return sum, any +} diff --git a/internal/router/scrape/parse_test.go b/internal/router/scrape/parse_test.go new file mode 100644 index 0000000..19ccf9d --- /dev/null +++ b/internal/router/scrape/parse_test.go @@ -0,0 +1,115 @@ +package scrape + +import ( + "errors" + "strings" + "testing" +) + +const vllmSample = ` +# HELP vllm:num_requests_running Number of requests currently running on GPU. +# TYPE vllm:num_requests_running gauge +vllm:num_requests_running{model_name="llama"} 7 +# TYPE vllm:num_requests_waiting gauge +vllm:num_requests_waiting{model_name="llama"} 3 +# TYPE vllm:kv_cache_usage_perc gauge +vllm:kv_cache_usage_perc{model_name="llama"} 0.42 +# TYPE vllm:cache_config_info gauge +vllm:cache_config_info{block_size="16",num_gpu_blocks="4096"} 1 +# TYPE vllm:generation_tokens_total counter +vllm:generation_tokens_total{model_name="llama"} 12345 +# TYPE vllm:request_success_total counter +vllm:request_success_total{model_name="llama",finished_reason="length"} 100 +vllm:request_success_total{model_name="llama",finished_reason="stop"} 50 +` + +func TestParseContractFields(t *testing.T) { + s, err := Parse(strings.NewReader(vllmSample), VLLMMapping()) + if err != nil { + t.Fatal(err) + } + if s.Running != 7 { + t.Errorf("Running=%d", s.Running) + } + if s.Waiting != 3 { + t.Errorf("Waiting=%d", s.Waiting) + } + if s.KVCacheUsage != 0.42 { + t.Errorf("KVCacheUsage=%v", s.KVCacheUsage) + } + if s.BlockSize != 16 || s.NumGPUBlocks != 4096 { + t.Errorf("cache_config_info not read: block=%d blocks=%d", s.BlockSize, s.NumGPUBlocks) + } + // A counter split across labels must be summed. + if s.SuccessTotal != 150 { + t.Errorf("SuccessTotal=%v, expected 150 (labels must be summed)", s.SuccessTotal) + } + if s.GenerationTokens != 12345 { + t.Errorf("GenerationTokens=%v", s.GenerationTokens) + } +} + +func TestParseRejectsMissingMandatoryMetric(t *testing.T) { + // kv_cache_usage is absent: routing would be blind, so this must error. + input := ` +# TYPE vllm:num_requests_running gauge +vllm:num_requests_running 1 +# TYPE vllm:num_requests_waiting gauge +vllm:num_requests_waiting 0 +` + _, err := Parse(strings.NewReader(input), VLLMMapping()) + var me *MissingError + if !errors.As(err, &me) { + t.Fatalf("expected MissingError, got %v", err) + } + if len(me.Missing) != 1 || me.Missing[0] != "vllm:kv_cache_usage_perc" { + t.Errorf("missing list is wrong: %v", me.Missing) + } +} + +func TestParseToleratesMissingOptionals(t *testing.T) { + input := ` +# TYPE vllm:num_requests_running gauge +vllm:num_requests_running 2 +# TYPE vllm:num_requests_waiting gauge +vllm:num_requests_waiting 1 +# TYPE vllm:kv_cache_usage_perc gauge +vllm:kv_cache_usage_perc 0.1 +` + s, err := Parse(strings.NewReader(input), VLLMMapping()) + if err != nil { + t.Fatalf("errored although only optional metrics were missing: %v", err) + } + if s.BlockSize != 0 || s.SuccessTotal != 0 { + t.Error("missing optional fields should have stayed zero") + } +} + +func TestParseWithCustomMapping(t *testing.T) { + // A different engine with different names (plan §14: names drift across versions). + input := ` +# TYPE sglang:running gauge +sglang:running 4 +# TYPE sglang:queued gauge +sglang:queued 2 +# TYPE sglang:kv_usage gauge +sglang:kv_usage 0.8 +` + s, err := Parse(strings.NewReader(input), Mapping{ + Running: "sglang:running", + Waiting: "sglang:queued", + KVCacheUsage: "sglang:kv_usage", + }) + if err != nil { + t.Fatal(err) + } + if s.Running != 4 || s.Waiting != 2 || s.KVCacheUsage != 0.8 { + t.Errorf("custom mapping did not work: %+v", s) + } +} + +func TestParseRejectsGarbage(t *testing.T) { + if _, err := Parse(strings.NewReader("this is not metrics {{{"), VLLMMapping()); err == nil { + t.Error("expected an error for malformed input") + } +} diff --git a/internal/router/scrape/poller.go b/internal/router/scrape/poller.go new file mode 100644 index 0000000..d40dee2 --- /dev/null +++ b/internal/router/scrape/poller.go @@ -0,0 +1,185 @@ +package scrape + +import ( + "context" + "fmt" + "log/slog" + "math/rand/v2" + "net/http" + "sync" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// Config holds poller settings. +type Config struct { + // Interval is the scrape period. Plan §7.3 suggests 1-2s. + Interval time.Duration + // Timeout bounds a single scrape; it must be shorter than Interval. + Timeout time.Duration + // Path is the metrics endpoint. + Path string + // Mapping selects the metric names. + Mapping Mapping +} + +// DefaultConfig returns sensible defaults. +func DefaultConfig() Config { + return Config{ + Interval: 2 * time.Second, + Timeout: time.Second, + Path: "/metrics", + Mapping: VLLMMapping(), + } +} + +func (c *Config) normalize() { + if c.Interval <= 0 { + c.Interval = 2 * time.Second + } + if c.Timeout <= 0 || c.Timeout >= c.Interval { + c.Timeout = c.Interval / 2 + } + if c.Path == "" { + c.Path = "/metrics" + } + c.Mapping.normalize() +} + +// Poller scrapes each registered endpoint periodically, one goroutine each. +// +// Per-endpoint goroutines plus jitter: scraping every endpoint at the same +// instant loads both the engines and the router in bursts and synchronizes the +// measurements, which makes the scores oscillate in lockstep. +type Poller struct { + cfg Config + reg *endpoint.Registry + http *http.Client + log *slog.Logger + + mu sync.Mutex + running map[string]context.CancelFunc + wg sync.WaitGroup +} + +// NewPoller creates a poller. Call Sync to start it. +func NewPoller(reg *endpoint.Registry, cfg Config, log *slog.Logger) *Poller { + cfg.normalize() + if log == nil { + log = slog.Default() + } + return &Poller{ + cfg: cfg, + reg: reg, + log: log, + http: &http.Client{ + Timeout: cfg.Timeout, + Transport: &http.Transport{ + MaxIdleConnsPerHost: 2, + DisableCompression: true, + }, + }, + running: make(map[string]context.CancelFunc), + } +} + +// Sync starts and stops scrape goroutines to match the registry's endpoint +// list. Call it again whenever that list changes. +func (p *Poller) Sync(ctx context.Context) { + infos := p.reg.Infos() + + p.mu.Lock() + defer p.mu.Unlock() + + want := make(map[string]endpoint.Info, len(infos)) + for _, in := range infos { + want[in.ID] = in + } + + for id, cancel := range p.running { + if _, ok := want[id]; !ok { + cancel() + delete(p.running, id) + } + } + for id, in := range want { + if _, ok := p.running[id]; ok { + continue + } + cctx, cancel := context.WithCancel(ctx) + p.running[id] = cancel + p.wg.Add(1) + go func() { + defer p.wg.Done() + p.loop(cctx, in) + }() + } +} + +// Stop halts every scrape goroutine and waits for them to finish. +func (p *Poller) Stop() { + p.mu.Lock() + for id, cancel := range p.running { + cancel() + delete(p.running, id) + } + p.mu.Unlock() + p.wg.Wait() +} + +func (p *Poller) loop(ctx context.Context, in endpoint.Info) { + // Jitter keeps goroutines that start together from staying in lockstep. + jitter := time.Duration(rand.Int64N(int64(p.cfg.Interval))) //nolint:gosec // jitter + select { + case <-ctx.Done(): + return + case <-time.After(jitter): + } + + t := time.NewTicker(p.cfg.Interval) + defer t.Stop() + + p.once(ctx, in) + for { + select { + case <-ctx.Done(): + return + case <-t.C: + p.once(ctx, in) + } + } +} + +func (p *Poller) once(ctx context.Context, in endpoint.Info) { + stats, err := p.fetch(ctx, in.URL) + if err != nil { + if ctx.Err() != nil { + return + } + p.reg.ScrapeFailed(in.ID) + p.log.Debug("scrape failed", "endpoint", in.ID, "err", err) + return + } + p.reg.Observe(in.ID, stats) +} + +func (p *Poller) fetch(ctx context.Context, base string) (endpoint.Stats, error) { + ctx, cancel := context.WithTimeout(ctx, p.cfg.Timeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+p.cfg.Path, nil) + if err != nil { + return endpoint.Stats{}, err + } + resp, err := p.http.Do(req) + if err != nil { + return endpoint.Stats{}, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return endpoint.Stats{}, fmt.Errorf("unexpected status: %s", resp.Status) + } + return Parse(resp.Body, p.cfg.Mapping) +}