From dc38dd609639693075b3e9118937fd12339b7b3e Mon Sep 17 00:00:00 2001 From: ilker Date: Wed, 16 Sep 2026 01:53:38 +0200 Subject: [PATCH] Phase 1.5: admission control MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Inserted between quota checks and scheduling: admission answers when a request may proceed, the scheduler still answers where it goes. Keeping them apart is the point — an engine queue is FIFO and knows nothing about tenants, so once a batch job is in it, no endpoint selection rescues the request behind it. - saturation detector: one number per pool, the worse of queue depth and KV occupancy, since either alone misleads. - priority bands with per-band saturation ceilings, so the cluster degrades from the bottom up. - tenant round-robin within a band: one caller sending a thousand requests takes one slot in the rotation, not a thousand. - bounded queues with capacity (503) and TTL (504); cancelled requests leave the queue instead of holding capacity. - metrics for queue depth, wait time, rejections and pool saturation. Exit criterion, measured in test/e2e: with 500 concurrent bronze requests running continuously, interactive p99 in the platinum band stays at 1.04-1.24x idle against a 1.5x budget. With admission off the same load pushes it to 55-60x. Both directions are asserted, so the test fails if admission regresses and also if the load generator stops being a real load. Two lags had to be closed, and they needed different answers. Engine metrics are scraped and lag an interval; router inflight does not, but it counts running requests, and treating those as queue makes a healthy engine look overloaded every interval. What is actually queued is inflight minus batch capacity, and capacity is not in the metric contract, so it is learned as the high-water mark of observed running. Separately, between admission granting a request and the router dispatching it the request is invisible everywhere, so a simultaneous burst passed the check en masse; the controller now counts its own outstanding leases against dispatched inflight. One plan correction: a saturation ceiling of 1.0 was documented as "never held back". It is not — it means "flow up to the healthy limit". A band that must never wait needs a ceiling above 1. Plan §7.2.2 updated. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 63 ++- README.md | 14 +- deploy/examples/server.yaml | 26 ++ internal/config/config.go | 44 +- internal/router/admin.go | 21 + internal/router/admission/admission_test.go | 436 ++++++++++++++++++++ internal/router/admission/controller.go | 412 ++++++++++++++++++ internal/router/admission/queue.go | 168 ++++++++ internal/router/admission/saturation.go | 98 +++++ internal/router/endpoint/endpoint.go | 33 +- internal/router/handlers.go | 77 ++++ internal/router/metrics.go | 29 ++ internal/router/router.go | 79 +++- llmgw-plan.md | 19 +- test/e2e/harness_test.go | 46 ++- test/e2e/phase15_test.go | 203 +++++++++ 16 files changed, 1744 insertions(+), 24 deletions(-) create mode 100644 internal/router/admission/admission_test.go create mode 100644 internal/router/admission/controller.go create mode 100644 internal/router/admission/queue.go create mode 100644 internal/router/admission/saturation.go create mode 100644 test/e2e/phase15_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 574f414..eeb28d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,8 +7,8 @@ is only released once its **exit criterion** is met, not when its code lands. |---|---|---|---| | v0.1.0 | 0 | Skeleton and sim-engine | released | | v0.2.0 | 1 | Static router: streaming proxy, filter/score/pick | released | -| v0.3.0 | 1.5 | Admission control: priority bands, saturation detector | in progress | -| v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | planned | +| v0.3.0 | 1.5 | Admission control: priority bands, saturation detector | released | +| v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | in progress | | v0.5.0 | 3 | Control plane: catalog, placement, reconciler | planned | | v0.6.0 | 4 | Multi-node replicas and fabric awareness | planned | | v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned | @@ -16,6 +16,65 @@ is only released once its **exit criterion** is met, not when its code lands. Until v1.0.0 the minor version is the breaking one: configuration formats and Go APIs change between minors without a deprecation period. +## v0.3.0 — 2026-09-16 + +Flow control. A batch workload no longer starves interactive traffic. + +### Added + +- **Admission control** (`internal/router/admission`), inserted between quota + checks and scheduling. It answers *when* a request may proceed; the scheduler + still answers *where* it goes. Keeping them apart is the point: an engine + queue is FIFO and knows nothing about tenants, so once a batch job is in it, + no amount of endpoint selection rescues the request behind it. +- **Saturation detector.** One number per pool, the worse of queue depth and KV + occupancy. Either signal alone misleads — a pool can have an empty queue and a + nearly full cache, or the reverse. +- **Priority bands.** Each band has a saturation ceiling; above it, that band + waits. The cluster then degrades from the bottom up. A ceiling of 1.0 means + "flow up to the healthy limit"; above 1 means never wait. +- **Tenant fairness.** Within a band, dispatch turns rotate between tenants, so + one caller sending a thousand requests occupies one slot in the rotation + rather than a thousand. +- **Bounded queues.** Capacity per band (immediate 503 when full) and a TTL + (504), because sending the GPU work the client has already abandoned helps + nobody. Cancelled requests leave the queue instead of holding capacity. +- **Metrics:** `llmgw_admission_wait_seconds`, `llmgw_admission_rejected_total`, + `llmgw_admission_queue_depth`, `llmgw_pool_saturation`. Responses carry + `x-llmgw-queue-ms` when a request waited. + +### Phase 1.5 exit criterion, measured in `test/e2e` + +While 500 concurrent bronze requests run continuously, interactive p99 in the +platinum band stays at **1.04–1.24x** its idle value, against a 1.5x budget. +With admission switched off the same load pushes it to **55–60x**. Both +directions are asserted, so the test fails if admission regresses *and* if the +load generator ever stops being a real load. + +### Two lags, two different fixes + +Getting this right took longer than the queueing itself: + +- Engine metrics are scraped, so they lag by an interval. The router's inflight + count does not — but it includes *running* requests, and treating those as + queue depth makes a healthy engine look overloaded every interval. What is + genuinely queued is `inflight − batch capacity`, and batch capacity is not in + the metric contract, so it is learned as the high-water mark of observed + `running`. +- Between admission granting a request and the router dispatching it, the + request is invisible to every signal. A simultaneous burst therefore passed + the check en masse — exactly the batch arrival pattern this layer exists to + survive. The controller now counts its own outstanding leases against + dispatched inflight. + +### Known limitations + +- The band comes from a trusted header until authentication arrives in v1.0.0. + Anyone who can reach the port can promote themselves to platinum. +- The tenant key is the session id, for the same reason. +- Displacement (dropping queued low-band work to make room) is not implemented; + the plan defers it, and bounded queues plus TTL cover most of the need. + ## v0.2.0 — 2026-09-16 The router serves traffic. `llmgw server` is real: it fronts a static set of diff --git a/README.md b/README.md index a06f11c..7b39862 100644 --- a/README.md +++ b/README.md @@ -14,14 +14,16 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md). |---|---|---|---| | v0.1.0 | 0 | Skeleton and sim-engine | released | | v0.2.0 | 1 | Static router (filter/score/pick, streaming proxy) | released | -| v0.3.0 | 1.5 | Admission control (priority bands) | planned | +| v0.3.0 | 1.5 | Admission control (priority bands) | released | | v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | planned | | v0.5.0 | 3 | Control plane: catalog, scheduler, reconciler | planned | | v0.6.0 | 4 | Multi-node placement and fabric awareness | planned | | v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned | -As of v0.2.0 the router serves traffic. There is still no admission control, so -a batch workload can starve interactive traffic — that is what v0.3.0 is for. +As of v0.3.0 the router serves traffic and protects it: a batch workload no +longer starves interactive requests. Still missing before v1.0.0: authentication, +rate limits and budgets, and the control plane that replaces the static endpoint +list. ## Building @@ -51,6 +53,12 @@ For anything real use a config file; [deploy/examples/server.yaml](deploy/exampl is commented throughout. Metrics, health and pprof are on the admin port (`:9090` by default), deliberately separate from the client-facing one. +Requests carry a priority band (`x-llmgw-band`). Under load, low bands wait at +the gateway while high bands keep flowing, so the cluster degrades from the +bottom up instead of everywhere at once. Measured effect: interactive p99 stays +within 1.24x of idle while 500 concurrent batch requests run, against 58x with +admission off. + Routing profiles select how endpoints are scored: `chat` favours session affinity, `batch` favours the shortest queue, `long-context` favours free KV cache. Pick one per model in the config, or per request with the diff --git a/deploy/examples/server.yaml b/deploy/examples/server.yaml index 702f7fd..5367abb 100644 --- a/deploy/examples/server.yaml +++ b/deploy/examples/server.yaml @@ -43,6 +43,32 @@ router: idle: 60s # between streamed tokens; catches a wedged engine total: 30m + # Flow control: decides *when* a request may proceed, before scheduling + # decides where it goes. Without it a batch workload starves interactive + # traffic, because an engine queue is FIFO and knows nothing about tenants. + admission: + enabled: true + # Until authentication lands, the band comes from a header and is trusted. + # Do not expose this port to callers who should not be able to promote + # themselves. + band_header: x-llmgw-band + default_band: gold + # A band waits once pool saturation passes its ceiling, so the cluster + # degrades from the bottom up. A ceiling of 1.0 means "flow up to the + # healthy limit"; use a value above 1 for a band that must never wait. + bands: + - { name: platinum, priority: 400, saturation_ceiling: 1.00 } + - { name: gold, priority: 300, saturation_ceiling: 0.85 } + - { name: silver, priority: 200, saturation_ceiling: 0.60 } + - { name: bronze, priority: 100, saturation_ceiling: 0.30 } + # Drop a request that has waited this long: the client has usually given up, + # and sending the GPU work nobody is waiting for helps no one. + queue_ttl: 30s + # Reject immediately once a band's queue is this deep, rather than letting + # it accumulate silently. + queue_capacity_per_band: 2000 + interval: 10ms + # What the scorers normalize against. Measure these with `ctl bench` rather # than guessing: target_waiting in particular governs how the whole system # behaves under load. Too low starves the GPU, too high brings back queueing diff --git a/internal/config/config.go b/internal/config/config.go index 600c41f..07960f3 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -58,8 +58,34 @@ type Router struct { // MaxBodyBytes bounds the request body the router will read. MaxBodyBytes int64 `koanf:"max_body_bytes"` - Timeouts Timeouts `koanf:"timeouts"` - Scoring Scoring `koanf:"scoring"` + Timeouts Timeouts `koanf:"timeouts"` + Scoring Scoring `koanf:"scoring"` + Admission Admission `koanf:"admission"` +} + +// Admission configures the flow control layer (plan §7.2). +type Admission struct { + // Enabled turns the layer on. Off, every request goes straight through, + // which is the v0.2.0 behaviour. + Enabled bool `koanf:"enabled"` + + Bands []Band `koanf:"bands"` + DefaultBand string `koanf:"default_band"` + // BandHeader lets a client declare its band. Until authentication lands in + // v1.0.0 this is trusted, so the port must not be exposed to callers who + // should not be able to promote themselves. + BandHeader string `koanf:"band_header"` + + QueueTTL time.Duration `koanf:"queue_ttl"` + QueueCapacityPerBand int `koanf:"queue_capacity_per_band"` + Interval time.Duration `koanf:"interval"` +} + +// Band is one priority class. +type Band struct { + Name string `koanf:"name"` + Priority int `koanf:"priority"` + SaturationCeiling float64 `koanf:"saturation_ceiling"` } // Timeouts bounds each phase of an upstream request. @@ -129,6 +155,20 @@ func Default() Server { KVCeiling: 0.9, SessionLoadFactor: 1.25, }, + Admission: Admission{ + Enabled: true, + BandHeader: "x-llmgw-band", + DefaultBand: "gold", + Bands: []Band{ + {Name: "platinum", Priority: 400, SaturationCeiling: 1.00}, + {Name: "gold", Priority: 300, SaturationCeiling: 0.85}, + {Name: "silver", Priority: 200, SaturationCeiling: 0.60}, + {Name: "bronze", Priority: 100, SaturationCeiling: 0.30}, + }, + QueueTTL: 30 * time.Second, + QueueCapacityPerBand: 2000, + Interval: 10 * time.Millisecond, + }, }, } } diff --git a/internal/router/admin.go b/internal/router/admin.go index 3213f9f..88793d7 100644 --- a/internal/router/admin.go +++ b/internal/router/admin.go @@ -53,6 +53,9 @@ type endpointCollector struct { stale *prometheus.Desc drain *prometheus.Desc breaker *prometheus.Desc + + queueDepth *prometheus.Desc + saturation *prometheus.Desc } func newEndpointCollector(r *Router) *endpointCollector { @@ -66,6 +69,9 @@ func newEndpointCollector(r *Router) *endpointCollector { drain: prometheus.NewDesc("llmgw_endpoint_drain_rate", "Observed completions per second.", lbl, nil), stale: prometheus.NewDesc("llmgw_endpoint_metrics_stale", "1 when the endpoint's metrics are stale.", lbl, nil), breaker: prometheus.NewDesc("llmgw_endpoint_breaker_open", "1 when the endpoint's circuit breaker is open.", lbl, nil), + + queueDepth: prometheus.NewDesc("llmgw_admission_queue_depth", "Requests waiting in each admission band.", []string{"band"}, nil), + saturation: prometheus.NewDesc("llmgw_pool_saturation", "Pool saturation per model; 1 is the healthy limit.", []string{"model"}, nil), } } @@ -77,6 +83,8 @@ func (c *endpointCollector) Describe(ch chan<- *prometheus.Desc) { ch <- c.drain ch <- c.stale ch <- c.breaker + ch <- c.queueDepth + ch <- c.saturation } func (c *endpointCollector) Collect(ch chan<- prometheus.Metric) { @@ -96,6 +104,19 @@ func (c *endpointCollector) Collect(ch chan<- prometheus.Metric) { g(c.stale, boolToFloat(s.Stale)) g(c.breaker, boolToFloat(s.Breaker.String() == "open")) } + + if c.r.adm != nil { + for _, d := range c.r.adm.Depths() { + ch <- prometheus.MustNewConstMetric(c.queueDepth, prometheus.GaugeValue, float64(d.Depth), d.Band) + } + } + // Saturation is reported per model, which is the granularity admission + // decisions are made at. + for _, model := range c.r.servedModels() { + if sat, ok := c.r.saturation(model, 0); ok { + ch <- prometheus.MustNewConstMetric(c.saturation, prometheus.GaugeValue, sat, model) + } + } } func boolToFloat(b bool) float64 { diff --git a/internal/router/admission/admission_test.go b/internal/router/admission/admission_test.go new file mode 100644 index 0000000..073d580 --- /dev/null +++ b/internal/router/admission/admission_test.go @@ -0,0 +1,436 @@ +package admission + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +func testConfig() Config { + c := DefaultConfig() + c.Interval = time.Millisecond + c.QueueTTL = 2 * time.Second + return c +} + +// fixedSat lets a test drive saturation directly. +type fixedSat struct{ v atomic.Uint64 } + +func (f *fixedSat) set(v float64) { f.v.Store(uint64(v * 1000)) } +func (f *fixedSat) fn(string, int) (float64, bool) { + return float64(f.v.Load()) / 1000, true +} + +func newController(t *testing.T, cfg Config, sat SaturationFunc) *Controller { + t.Helper() + c, err := New(cfg, sat) + if err != nil { + t.Fatal(err) + } + c.Start() + t.Cleanup(c.Close) + return c +} + +func TestAdmitsImmediatelyWhenIdle(t *testing.T) { + s := &fixedSat{} + c := newController(t, testConfig(), s.fn) + + start := time.Now() + lease, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze"}) + if err != nil { + t.Fatal(err) + } + if d := time.Since(start); d > 50*time.Millisecond { + t.Errorf("idle admission took %v; the fast path is not being used", d) + } + lease.Release() +} + +func TestLowBandHeldWhileHighBandFlows(t *testing.T) { + // The core of plan §7.2.2: the cluster degrades from the bottom up. + s := &fixedSat{} + s.set(0.5) // above bronze (0.30), below gold (0.85) + c := newController(t, testConfig(), s.fn) + + ctx, cancel := context.WithTimeout(t.Context(), 300*time.Millisecond) + defer cancel() + if _, err := c.Acquire(ctx, Request{Model: "m", Band: "bronze"}); err == nil { + t.Error("bronze was admitted above its ceiling") + } + + lease, err := c.Acquire(t.Context(), Request{Model: "m", Band: "gold"}) + if err != nil { + t.Fatalf("gold was held below its ceiling: %v", err) + } + lease.Release() +} + +func TestTopBandFlowsUpToTheHealthyLimit(t *testing.T) { + // A ceiling of 1.0 means "hold only once the pool is past its healthy + // limit", not "never hold". At exactly the limit the top band still flows. + s := &fixedSat{} + s.set(1.0) + c := newController(t, testConfig(), s.fn) + + lease, err := c.Acquire(t.Context(), Request{Model: "m", Band: "platinum"}) + if err != nil { + t.Fatalf("top band held at exactly its ceiling: %v", err) + } + lease.Release() +} + +func TestTopBandIsHeldWhenOverloaded(t *testing.T) { + // Holding even the top band is deliberate: with nobody above it to protect, + // the benefit is a bounded engine queue. Dumping a flood straight into a + // FIFO engine queue makes every request in that band slow, whereas queueing + // at the gateway keeps ordering, TTL and tenant fairness. + s := &fixedSat{} + s.set(5) + cfg := testConfig() + cfg.QueueTTL = 200 * time.Millisecond + c := newController(t, cfg, s.fn) + + if _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "platinum"}); !errors.Is(err, ErrQueueTimeout) { + t.Fatalf("expected the top band to queue while overloaded, got %v", err) + } +} + +func TestCeilingAboveOneNeverHolds(t *testing.T) { + // The escape hatch for a band that must never wait, whatever the load. + s := &fixedSat{} + s.set(50) + cfg := testConfig() + cfg.Bands = []Band{{Name: "critical", Priority: 900, SaturationCeiling: 1000}} + cfg.DefaultBand = "critical" + c := newController(t, cfg, s.fn) + + lease, err := c.Acquire(t.Context(), Request{Model: "m", Band: "critical"}) + if err != nil { + t.Fatalf("a band with an effectively unbounded ceiling was held: %v", err) + } + lease.Release() +} + +func TestQueuedRequestAdmittedWhenSaturationDrops(t *testing.T) { + s := &fixedSat{} + s.set(0.9) + c := newController(t, testConfig(), s.fn) + + done := make(chan error, 1) + go func() { + _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze"}) + done <- err + }() + + // Give it time to actually queue, then open the gate. + time.Sleep(50 * time.Millisecond) + if d := c.Depths(); d[3].Depth != 1 { + t.Fatalf("expected 1 queued bronze request, got %+v", d) + } + s.set(0.1) + + select { + case err := <-done: + if err != nil { + t.Fatalf("queued request failed: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("queued request was not admitted after saturation dropped") + } +} + +func TestQueueTTL(t *testing.T) { + s := &fixedSat{} + s.set(1.0) + cfg := testConfig() + cfg.QueueTTL = 150 * time.Millisecond + c := newController(t, cfg, s.fn) + + start := time.Now() + _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze"}) + if !errors.Is(err, ErrQueueTimeout) { + t.Fatalf("expected ErrQueueTimeout, got %v", err) + } + if d := time.Since(start); d > time.Second { + t.Errorf("TTL fired after %v", d) + } +} + +func TestQueueCapacity(t *testing.T) { + s := &fixedSat{} + s.set(1.0) + cfg := testConfig() + cfg.QueueCapacityPerBand = 3 + cfg.QueueTTL = 5 * time.Second + c := newController(t, cfg, s.fn) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + for range 3 { + go c.Acquire(ctx, Request{Model: "m", Band: "bronze"}) //nolint:errcheck // queued on purpose + } + // Wait for the queue to fill. + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) && c.Depths()[3].Depth < 3 { + time.Sleep(2 * time.Millisecond) + } + + _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze"}) + if !errors.Is(err, ErrQueueFull) { + t.Errorf("expected ErrQueueFull, got %v", err) + } +} + +func TestClientCancellationLeavesQueue(t *testing.T) { + s := &fixedSat{} + s.set(1.0) + c := newController(t, testConfig(), s.fn) + + ctx, cancel := context.WithCancel(context.Background()) + go c.Acquire(ctx, Request{Model: "m", Band: "bronze"}) //nolint:errcheck // cancelled on purpose + + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) && c.Depths()[3].Depth == 0 { + time.Sleep(2 * time.Millisecond) + } + cancel() + + deadline = time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if c.Depths()[3].Depth == 0 { + return + } + time.Sleep(2 * time.Millisecond) + } + t.Fatal("a cancelled request stayed in the queue, occupying capacity") +} + +func TestTenantFairnessWithinBand(t *testing.T) { + // One tenant flooding the band must not take every dispatch turn. + s := &fixedSat{} + s.set(1.0) // hold everything initially + cfg := testConfig() + cfg.QueueTTL = 5 * time.Second + c := newController(t, cfg, s.fn) + + const flood = 40 + var wg sync.WaitGroup + order := make(chan string, flood+4) + + for i := range flood { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze", Tenant: "noisy"}); err == nil { + order <- "noisy" + } + _ = i + }() + } + // Let the flood queue up first, so the quiet tenant is strictly behind it. + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) && c.Depths()[3].Depth < flood { + time.Sleep(2 * time.Millisecond) + } + + for range 4 { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze", Tenant: "quiet"}); err == nil { + order <- "quiet" + } + }() + } + time.Sleep(30 * time.Millisecond) + + s.set(0.0) // open the gate + wg.Wait() + close(order) + + var seen []string + for v := range order { + seen = append(seen, v) + } + if len(seen) != flood+4 { + t.Fatalf("expected %d admissions, got %d", flood+4, len(seen)) + } + + // With FIFO the quiet tenant would be admitted 41st at best. Round-robin + // should interleave it into the first handful. + firstQuiet := -1 + for i, v := range seen { + if v == "quiet" { + firstQuiet = i + break + } + } + if firstQuiet < 0 { + t.Fatal("the quiet tenant was never admitted") + } + if firstQuiet > 8 { + t.Errorf("quiet tenant admitted at position %d; a flood is monopolising the band", firstQuiet) + } +} + +func TestUnknownBandUsesDefault(t *testing.T) { + s := &fixedSat{} + s.set(0.5) + cfg := testConfig() + cfg.DefaultBand = "bronze" + c := newController(t, cfg, s.fn) + + ctx, cancel := context.WithTimeout(t.Context(), 200*time.Millisecond) + defer cancel() + // bronze's ceiling is 0.30, so an unknown band must be held here. + if _, err := c.Acquire(ctx, Request{Model: "m", Band: "no-such-band"}); err == nil { + t.Error("unknown band did not fall back to the default") + } +} + +func TestUnknownModelIsLetThrough(t *testing.T) { + // Better a precise 503 from dispatch than a 30-second queue timeout. + c := newController(t, testConfig(), func(string, int) (float64, bool) { return 0, false }) + lease, err := c.Acquire(t.Context(), Request{Model: "nope", Band: "bronze"}) + if err != nil { + t.Fatalf("unknown model was queued instead of passed through: %v", err) + } + lease.Release() +} + +func TestConfigValidation(t *testing.T) { + if _, err := New(Config{Bands: []Band{{Name: "a"}, {Name: "a"}}}, nil); err == nil { + t.Error("duplicate band names accepted") + } + if _, err := New(Config{Bands: []Band{{Name: "a"}}, DefaultBand: "b"}, nil); err == nil { + t.Error("default_band outside the band list accepted") + } +} + +func TestBandsSortedByPriority(t *testing.T) { + c, err := New(Config{Bands: []Band{ + {Name: "low", Priority: 1, SaturationCeiling: 0.2}, + {Name: "high", Priority: 9, SaturationCeiling: 1}, + }, DefaultBand: "low"}, func(string, int) (float64, bool) { return 0, true }) + if err != nil { + t.Fatal(err) + } + if got := c.Bands()[0].Name; got != "high" { + t.Errorf("bands are not ordered by priority: first is %q", got) + } +} + +// --- saturation --- + +func snap(id string, waiting, running, inflight int, kv float64, open bool) endpoint.Snapshot { + s := endpoint.Snapshot{ + Info: endpoint.Info{ID: id}, + Inflight: inflight, + Breaker: breaker.Closed, + } + s.Stats.Waiting = waiting + s.Stats.Running = running + s.Stats.KVCacheUsage = kv + if open { + s.Breaker = breaker.Open + } + return s +} + +func TestSaturationCombinesQueueAndKV(t *testing.T) { + th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9} + + // Empty queue but a nearly full cache still counts as saturated. + if got := Saturation([]endpoint.Snapshot{snap("a", 0, 1, 1, 0.89, false)}, 0, th); got < 0.9 { + t.Errorf("KV pressure ignored: %v", got) + } + // Full queue with an empty cache likewise. + if got := Saturation([]endpoint.Snapshot{snap("a", 4, 1, 5, 0, false)}, 0, th); got < 1 { + t.Errorf("queue depth ignored: %v", got) + } +} + +func TestSaturationCountsUnreportedInflight(t *testing.T) { + // A burst the engine has not reported yet is real load; missing it would + // leave admission blind exactly when it needs to bite. + th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9} + cold := Saturation([]endpoint.Snapshot{snap("a", 0, 0, 0, 0, false)}, 0, th) + burst := Saturation([]endpoint.Snapshot{snap("a", 0, 0, 8, 0, false)}, 0, th) + if !(burst > cold) { + t.Errorf("unreported inflight not counted: cold=%v burst=%v", cold, burst) + } +} + +func TestSaturationSkipsOpenBreakers(t *testing.T) { + th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9} + eps := []endpoint.Snapshot{ + snap("busy", 8, 4, 12, 0.5, false), + snap("dead", 0, 0, 0, 0, true), // would halve the mean if counted + } + withDead := Saturation(eps, 0, th) + alone := Saturation(eps[:1], 0, th) + if fmt.Sprintf("%.4f", withDead) != fmt.Sprintf("%.4f", alone) { + t.Errorf("an endpoint with an open breaker was counted as capacity: %v vs %v", withDead, alone) + } +} + +func TestSaturationCountsUnreportedAdmissions(t *testing.T) { + // The thundering herd guard: a simultaneous burst must register as load + // before any of it reaches an engine, or the whole burst walks through the + // fast path while every signal still reads zero. + th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9} + idle := []endpoint.Snapshot{snap("a", 0, 0, 0, 0, false), snap("b", 0, 0, 0, 0, false)} + + if got := Saturation(idle, 0, th); got != 0 { + t.Fatalf("idle pool reported %v", got) + } + if got := Saturation(idle, 16, th); got <= 1 { + t.Errorf("16 unreported admissions across 2 replicas should saturate, got %v", got) + } +} + +func TestBurstDoesNotAllPassTheFastPath(t *testing.T) { + // End to end for the same guard: with saturation derived from outstanding + // leases, a burst of 200 must not be admitted wholesale. + var admitted atomic.Int64 + c := newController(t, testConfig(), func(_ string, outstanding int) (float64, bool) { + // Two replicas, healthy buffer of 4 each: 8 in flight is the limit. + return float64(outstanding) / 8, true + }) + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + var wg sync.WaitGroup + for range 200 { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := c.Acquire(ctx, Request{Model: "m", Band: "bronze"}); err == nil { + admitted.Add(1) + } + }() + } + wg.Wait() + + // bronze's ceiling is 0.30, so roughly 8*0.3 = 2-3 should get in. Allow + // generous slack for scheduling, but nothing near 200. + if n := admitted.Load(); n > 20 { + t.Errorf("%d of 200 requests were admitted at once; the burst guard is not working", n) + } +} + +func TestSaturationWithNoUsableEndpoints(t *testing.T) { + th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9} + if got := Saturation(nil, 0, th); got != 1 { + t.Errorf("an empty pool should report full saturation, got %v", got) + } +} diff --git a/internal/router/admission/controller.go b/internal/router/admission/controller.go new file mode 100644 index 0000000..77e67aa --- /dev/null +++ b/internal/router/admission/controller.go @@ -0,0 +1,412 @@ +package admission + +import ( + "context" + "errors" + "fmt" + "slices" + "sync" + "time" +) + +// Band is a priority class. Requests in a band are held once pool saturation +// exceeds the band's ceiling, so the cluster degrades from the bottom up: +// batch traffic slows first and interactive traffic keeps its latency. +type Band struct { + Name string `koanf:"name"` + // Priority orders bands; higher is served first. + Priority int `koanf:"priority"` + // SaturationCeiling is the pool saturation above which this band waits. + // 1.0 means "never hold this band back". + SaturationCeiling float64 `koanf:"saturation_ceiling"` +} + +// Config configures the controller. +type Config struct { + Bands []Band `koanf:"bands"` + DefaultBand string `koanf:"default_band"` + + Thresholds Thresholds `koanf:"-"` + + // QueueTTL drops a request that has waited this long. Sending the GPU work + // the client has already given up on helps nobody. + QueueTTL time.Duration `koanf:"queue_ttl"` + // QueueCapacityPerBand bounds each band's queue. Rejecting immediately + // beats accumulating silently. + QueueCapacityPerBand int `koanf:"queue_capacity_per_band"` + // Interval is how often the dispatcher re-checks saturation. Releases also + // wake it, so this is only the backstop for saturation falling on its own. + Interval time.Duration `koanf:"interval"` +} + +// DefaultBands returns the four-band scheme from plan §7.2.2. +func DefaultBands() []Band { + return []Band{ + {Name: "platinum", Priority: 400, SaturationCeiling: 1.00}, + {Name: "gold", Priority: 300, SaturationCeiling: 0.85}, + {Name: "silver", Priority: 200, SaturationCeiling: 0.60}, + {Name: "bronze", Priority: 100, SaturationCeiling: 0.30}, + } +} + +// DefaultConfig returns a usable configuration. +func DefaultConfig() Config { + return Config{ + Bands: DefaultBands(), + DefaultBand: "gold", + Thresholds: Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9}, + QueueTTL: 30 * time.Second, + QueueCapacityPerBand: 2000, + Interval: 10 * time.Millisecond, + } +} + +func (c *Config) normalize() error { + if len(c.Bands) == 0 { + c.Bands = DefaultBands() + } + slices.SortFunc(c.Bands, func(a, b Band) int { return b.Priority - a.Priority }) + + seen := map[string]struct{}{} + for i := range c.Bands { + b := &c.Bands[i] + if b.Name == "" { + return fmt.Errorf("bands[%d]: name is required", i) + } + if _, dup := seen[b.Name]; dup { + return fmt.Errorf("band %q is declared twice", b.Name) + } + seen[b.Name] = struct{}{} + if b.SaturationCeiling <= 0 { + b.SaturationCeiling = 1 + } + } + if c.DefaultBand == "" { + c.DefaultBand = c.Bands[len(c.Bands)-1].Name + } + if _, ok := seen[c.DefaultBand]; !ok { + return fmt.Errorf("default_band %q is not one of the configured bands", c.DefaultBand) + } + c.Thresholds.normalize() + if c.QueueTTL <= 0 { + c.QueueTTL = 30 * time.Second + } + if c.QueueCapacityPerBand <= 0 { + c.QueueCapacityPerBand = 2000 + } + if c.Interval <= 0 { + c.Interval = 10 * time.Millisecond + } + return nil +} + +// SaturationFunc reports a model's pool saturation and whether it is known. +// +// outstanding is how many requests this controller has admitted and not yet +// seen released, so the caller can fold in load the engines have not reported. +type SaturationFunc func(model string, outstanding int) (float64, bool) + +// Request identifies who is asking for what. +type Request struct { + Model string + // Band selects the priority class; unknown names fall back to the default. + Band string + // Tenant is the fairness key within a band. Two requests sharing a tenant + // compete with each other rather than with everyone else. + Tenant string +} + +var ( + // ErrQueueFull means the band's queue is at capacity. + ErrQueueFull = errors.New("admission queue is full") + // ErrQueueTimeout means the request waited longer than the TTL. + ErrQueueTimeout = errors.New("admission queue timeout") + // ErrClosed means the controller is shutting down. + ErrClosed = errors.New("admission controller closed") +) + +// Stats is a snapshot of one band's queue, for metrics. +type Stats struct { + Band string + Depth int +} + +// Controller admits requests according to band priority and tenant fairness. +type Controller struct { + cfg Config + satFn SaturationFunc + bandOf map[string]int + + mu sync.Mutex + queues []*bandQueue + closed bool + + wake chan struct{} + done chan struct{} + stop chan struct{} + once sync.Once + + // outstanding counts admitted-but-unreleased requests per model. It is the + // only load signal available with zero lag, which is what stops a + // simultaneous burst from passing the check en masse. + outstanding map[string]int + + // Observed counters, read through Metrics. + admitted map[string]int64 + rejected map[string]int64 + waitTotal map[string]time.Duration +} + +// New builds a controller. Start must be called before it will admit anything +// that has to wait. +func New(cfg Config, satFn SaturationFunc) (*Controller, error) { + if err := cfg.normalize(); err != nil { + return nil, err + } + c := &Controller{ + cfg: cfg, + satFn: satFn, + bandOf: make(map[string]int, len(cfg.Bands)), + wake: make(chan struct{}, 1), + done: make(chan struct{}), + stop: make(chan struct{}), + outstanding: map[string]int{}, + admitted: map[string]int64{}, + rejected: map[string]int64{}, + waitTotal: map[string]time.Duration{}, + } + for i, b := range cfg.Bands { + c.bandOf[b.Name] = i + c.queues = append(c.queues, newBandQueue(b)) + } + return c, nil +} + +// Bands returns the configured bands, highest priority first. +func (c *Controller) Bands() []Band { return slices.Clone(c.cfg.Bands) } + +// Start launches the dispatcher. +func (c *Controller) Start() { + go c.dispatchLoop() +} + +// Close stops the dispatcher and fails everything still queued. +func (c *Controller) Close() { + c.once.Do(func() { + close(c.stop) + <-c.done + }) +} + +// Lease is held for the lifetime of an admitted request. +type Lease struct { + c *Controller + band string + model string + once sync.Once +} + +// Release returns the slot and wakes the dispatcher, since finishing a request +// is the most likely reason saturation just dropped. +func (l *Lease) Release() { + if l == nil || l.c == nil { + return + } + l.once.Do(func() { + l.c.mu.Lock() + if n := l.c.outstanding[l.model]; n > 1 { + l.c.outstanding[l.model] = n - 1 + } else { + delete(l.c.outstanding, l.model) + } + l.c.mu.Unlock() + l.c.signal() + }) +} + +// Band reports which band admitted the request. +func (l *Lease) Band() string { + if l == nil { + return "" + } + return l.band +} + +// Acquire blocks until the request may proceed, the context ends, or the TTL +// expires. +func (c *Controller) Acquire(ctx context.Context, r Request) (*Lease, error) { + idx, ok := c.bandOf[r.Band] + if !ok { + idx = c.bandOf[c.cfg.DefaultBand] + } + band := c.cfg.Bands[idx] + + // Fast path: under the ceiling with nobody already waiting in this band. + // Skipping the queue here is what keeps the common case free of latency. + c.mu.Lock() + if c.closed { + c.mu.Unlock() + return nil, ErrClosed + } + if c.queues[idx].len() == 0 && c.admits(band, r.Model) { + c.grant(band.Name, r.Model) + c.mu.Unlock() + return &Lease{c: c, band: band.Name, model: r.Model}, nil + } + + if c.queues[idx].len() >= c.cfg.QueueCapacityPerBand { + c.rejected[band.Name]++ + c.mu.Unlock() + return nil, fmt.Errorf("%w: band %q", ErrQueueFull, band.Name) + } + + w := &waiter{ + tenant: tenantKey(r), + model: r.Model, + ready: make(chan struct{}), + enqueued: time.Now(), + deadline: time.Now().Add(c.cfg.QueueTTL), + } + c.queues[idx].push(w) + c.mu.Unlock() + + c.signal() + + select { + case <-w.ready: + c.mu.Lock() + if c.closed { + c.mu.Unlock() + return nil, ErrClosed + } + c.waitTotal[band.Name] += time.Since(w.enqueued) + c.mu.Unlock() + return &Lease{c: c, band: band.Name, model: r.Model}, nil + + case <-ctx.Done(): + c.remove(idx, w, "client_gone") + return nil, ctx.Err() + + case <-time.After(time.Until(w.deadline)): + c.remove(idx, w, "ttl") + return nil, fmt.Errorf("%w after %v", ErrQueueTimeout, c.cfg.QueueTTL) + } +} + +// remove takes a waiter out of its queue unless it was admitted first. +func (c *Controller) remove(idx int, w *waiter, reason string) { + c.mu.Lock() + defer c.mu.Unlock() + if w.admitted { + // It won the race; the caller will see ready closed next time round. + return + } + c.queues[idx].remove(w) + c.rejected[c.cfg.Bands[idx].Name+"/"+reason]++ +} + +// grant records an admission. The caller holds the lock. +func (c *Controller) grant(band, model string) { + c.admitted[band]++ + c.outstanding[model]++ +} + +// admits reports whether a band may proceed right now. The caller holds the lock. +func (c *Controller) admits(b Band, model string) bool { + sat, known := c.satFn(model, c.outstanding[model]) + if !known { + // Nothing is known about this model's pool; let it through so dispatch + // can return a precise error instead of a queue timeout. + return true + } + return sat <= b.SaturationCeiling +} + +func (c *Controller) signal() { + select { + case c.wake <- struct{}{}: + default: + } +} + +func (c *Controller) dispatchLoop() { + defer close(c.done) + t := time.NewTicker(c.cfg.Interval) + defer t.Stop() + + for { + select { + case <-c.stop: + c.drain() + return + case <-c.wake: + c.dispatch() + case <-t.C: + c.dispatch() + } + } +} + +// dispatch walks the bands from highest priority down, admitting whatever the +// current saturation allows. +func (c *Controller) dispatch() { + now := time.Now() + + c.mu.Lock() + defer c.mu.Unlock() + + for i, q := range c.queues { + q.dropExpired(now, func() { c.rejected[c.cfg.Bands[i].Name+"/ttl"]++ }) + + band := c.cfg.Bands[i] + // Each pass admits at most the queue length, so a band cannot spin + // forever while saturation is recomputed from stale metrics. + for range q.len() { + w := q.peek() + if w == nil || !c.admits(band, w.model) { + break + } + q.pop() + w.admitted = true + c.grant(band.Name, w.model) + close(w.ready) + } + } +} + +func (c *Controller) drain() { + c.mu.Lock() + defer c.mu.Unlock() + c.closed = true + for _, q := range c.queues { + for { + w := q.pop() + if w == nil { + break + } + // Closing without marking admitted lets Acquire report ErrClosed. + close(w.ready) + w.admitted = true + } + } +} + +// Depths returns each band's current queue depth. +func (c *Controller) Depths() []Stats { + c.mu.Lock() + defer c.mu.Unlock() + out := make([]Stats, 0, len(c.queues)) + for i, q := range c.queues { + out = append(out, Stats{Band: c.cfg.Bands[i].Name, Depth: q.len()}) + } + return out +} + +// tenantKey picks the fairness key, falling back to the band so that unknown +// tenants at least compete as one group rather than each looking unique. +func tenantKey(r Request) string { + if r.Tenant != "" { + return r.Tenant + } + return "anonymous" +} diff --git a/internal/router/admission/queue.go b/internal/router/admission/queue.go new file mode 100644 index 0000000..8aa1e33 --- /dev/null +++ b/internal/router/admission/queue.go @@ -0,0 +1,168 @@ +package admission + +import "time" + +// waiter is one queued request. +type waiter struct { + tenant string + model string + ready chan struct{} + enqueued time.Time + deadline time.Time + // admitted guards the race between the dispatcher admitting a waiter and + // the waiter giving up on its own. + admitted bool +} + +// bandQueue holds one band's waiters, grouped by tenant. +// +// Grouping is what makes fairness possible: a single tenant sending a thousand +// requests occupies one slot in the rotation, not a thousand, so a traffic +// spike from one caller cannot consume the band. +type bandQueue struct { + band Band + + byTenant map[string][]*waiter + // rotation is the round-robin order of tenants with pending work. + rotation []string + next int + count int +} + +func newBandQueue(b Band) *bandQueue { + return &bandQueue{band: b, byTenant: map[string][]*waiter{}} +} + +func (q *bandQueue) len() int { return q.count } + +func (q *bandQueue) push(w *waiter) { + if _, ok := q.byTenant[w.tenant]; !ok { + q.byTenant[w.tenant] = nil + q.rotation = append(q.rotation, w.tenant) + } + q.byTenant[w.tenant] = append(q.byTenant[w.tenant], w) + q.count++ +} + +// peek returns the waiter the next pop would return, without removing it. +func (q *bandQueue) peek() *waiter { + if q.count == 0 { + return nil + } + idx, tenant := q.nextTenant() + if idx < 0 { + return nil + } + _ = tenant + return q.byTenant[q.rotation[idx]][0] +} + +// pop removes and returns the next waiter in tenant rotation. +func (q *bandQueue) pop() *waiter { + if q.count == 0 { + return nil + } + idx, _ := q.nextTenant() + if idx < 0 { + return nil + } + + tenant := q.rotation[idx] + ws := q.byTenant[tenant] + w := ws[0] + q.byTenant[tenant] = ws[1:] + q.count-- + + if len(q.byTenant[tenant]) == 0 { + delete(q.byTenant, tenant) + q.rotation = append(q.rotation[:idx], q.rotation[idx+1:]...) + if idx < q.next { + q.next-- + } + } else { + q.next = idx + 1 + } + if len(q.rotation) > 0 { + q.next %= len(q.rotation) + } else { + q.next = 0 + } + return w +} + +// nextTenant finds the next rotation entry that still has waiters. +func (q *bandQueue) nextTenant() (int, string) { + n := len(q.rotation) + for i := range n { + idx := (q.next + i) % n + if len(q.byTenant[q.rotation[idx]]) > 0 { + return idx, q.rotation[idx] + } + } + return -1, "" +} + +// remove takes a specific waiter out, for cancellation and TTL expiry. +func (q *bandQueue) remove(target *waiter) bool { + ws, ok := q.byTenant[target.tenant] + if !ok { + return false + } + for i, w := range ws { + if w != target { + continue + } + q.byTenant[target.tenant] = append(ws[:i:i], ws[i+1:]...) + q.count-- + if len(q.byTenant[target.tenant]) == 0 { + delete(q.byTenant, target.tenant) + q.dropTenant(target.tenant) + } + return true + } + return false +} + +func (q *bandQueue) dropTenant(tenant string) { + for i, t := range q.rotation { + if t != tenant { + continue + } + q.rotation = append(q.rotation[:i], q.rotation[i+1:]...) + if i < q.next { + q.next-- + } + break + } + if len(q.rotation) > 0 { + q.next %= len(q.rotation) + } else { + q.next = 0 + } +} + +// dropExpired removes waiters past their deadline, calling onDrop for each. +// Their Acquire call is already unblocking on its own timer; this keeps the +// queue from carrying corpses that would occupy capacity. +func (q *bandQueue) dropExpired(now time.Time, onDrop func()) { + if q.count == 0 { + return + } + for tenant, ws := range q.byTenant { + kept := ws[:0] + for _, w := range ws { + if now.After(w.deadline) && !w.admitted { + q.count-- + onDrop() + continue + } + kept = append(kept, w) + } + if len(kept) == 0 { + delete(q.byTenant, tenant) + q.dropTenant(tenant) + continue + } + q.byTenant[tenant] = kept + } +} diff --git a/internal/router/admission/saturation.go b/internal/router/admission/saturation.go new file mode 100644 index 0000000..41e4335 --- /dev/null +++ b/internal/router/admission/saturation.go @@ -0,0 +1,98 @@ +// Package admission answers "may this request proceed now?" (plan §7.2). +// +// It is deliberately separate from scheduling, which answers "where should it +// go?". Kept apart, the router can hold requests in its own queue and apply +// priority and tenant fairness *before* they enter an engine queue. An engine +// queue is FIFO and knows nothing about tenants: once a batch job has put 500 +// requests into it, no amount of clever endpoint selection will rescue the +// interactive request that arrives next. +package admission + +import ( + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// Thresholds define what "saturated" means for a pool. +type Thresholds struct { + // TargetWaitingPerReplica is the healthy queue depth: enough for continuous + // batching to stay efficient, not enough for queueing delay to show up. + // Measure it with ctl bench (plan §10); this single number governs how the + // whole system behaves under load. + TargetWaitingPerReplica float64 + // KVCeiling is the KV occupancy treated as full. + KVCeiling float64 +} + +func (t *Thresholds) normalize() { + if t.TargetWaitingPerReplica <= 0 { + t.TargetWaitingPerReplica = 4 + } + if t.KVCeiling <= 0 || t.KVCeiling > 1 { + t.KVCeiling = 0.9 + } +} + +// Saturation reduces a pool's state to a single number, where 1 means "at the +// healthy limit" and above 1 means overloaded. +// +// Two signals are combined because either alone misleads: a pool can have an +// empty queue while its KV cache is nearly full (long contexts), or plenty of +// KV headroom while requests pile up (slow decode). The worse of the two wins. +// +// unreported is load that has been admitted but that no endpoint has accounted +// for yet. Without it a simultaneous burst all passes the check before any of +// it becomes visible — which is precisely when admission needs to bite. +func Saturation(eps []endpoint.Snapshot, unreported int, t Thresholds) float64 { + t.normalize() + + var sum float64 + var n int + for _, e := range eps { + // A replica the breaker has taken out is not available capacity, so + // counting it would make the pool look healthier than it is. + if e.Breaker.String() == "open" { + continue + } + + // The engine's own queue is the truth, but it is one scrape interval + // old. Fill that gap with what the router knows without any lag: how + // many requests it has dispatched here, minus what this engine can run + // at once. Anything beyond that capacity is queued by definition. + // + // Subtracting the capacity is the whole point. Counting raw inflight + // would make a fully busy but perfectly healthy engine look like it had + // a queue as long as its batch, every single scrape interval. + waiting := float64(e.Stats.Waiting) + if e.RunningCapacity > 0 { + if queued := e.Inflight - e.RunningCapacity; float64(queued) > waiting { + waiting = float64(queued) + } + } else if e.Inflight > 0 { + // Nothing observed yet, so capacity is unknown. Treat dispatched + // work as queued: conservative for the first scrape, and it relaxes + // as soon as the engine reports anything. + waiting = max(waiting, float64(e.Inflight)) + } + + q := waiting / t.TargetWaitingPerReplica + kv := e.Stats.KVCacheUsage / t.KVCeiling + sum += max(q, kv) + n++ + } + _ = n + + if n == 0 { + // Every replica is out. Report full saturation so low bands stop + // queueing work that has nowhere to go; the dispatch layer will return + // a clearer error than a queue timeout. + return 1 + } + // Spread unaccounted admissions evenly; they will land somewhere. + if unreported > 0 { + sum += float64(unreported) / t.TargetWaitingPerReplica + } + + // The mean, not the minimum: one idle replica does not make a busy pool + // idle, and admission is about the pool's capacity as a whole. + return sum / float64(n) +} diff --git a/internal/router/endpoint/endpoint.go b/internal/router/endpoint/endpoint.go index 092d46f..012d229 100644 --- a/internal/router/endpoint/endpoint.go +++ b/internal/router/endpoint/endpoint.go @@ -71,6 +71,14 @@ type Snapshot struct { // unavoidable within a 2s scrape interval. Inflight int + // RunningCapacity is the highest concurrent request count ever observed on + // this endpoint — in practice the engine's batch capacity, learned rather + // than configured. Without it, requests the router has dispatched but the + // engine has not yet reported are indistinguishable from queue depth, and + // during a scrape interval that makes a perfectly healthy engine look + // overloaded. + RunningCapacity int + // Stale means no successful scrape arrived within StaleAfter. Stale // endpoints are not dropped, only scored down. Stale bool @@ -101,10 +109,11 @@ func DefaultConfig() Config { type live struct { info Info - stats Stats - scrapedAt time.Time - lastSucc float64 - drainRate float64 + stats Stats + scrapedAt time.Time + lastSucc float64 + drainRate float64 + maxRunning int inflight int brk *breaker.Breaker @@ -199,6 +208,7 @@ func (r *Registry) Observe(id string, s Stats) { l.stats = s l.lastSucc = s.SuccessTotal l.scrapedAt = now + l.maxRunning = max(l.maxRunning, s.Running) } // Snapshots returns a view of every endpoint serving the given model, or of all @@ -220,13 +230,14 @@ func (r *Registry) Snapshots(model string) []Snapshot { 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(), + Info: l.info, + Stats: l.stats, + DrainRate: l.drainRate, + Inflight: l.inflight, + RunningCapacity: l.maxRunning, + Stale: stale, + Age: age, + Breaker: l.brk.State(), }) } slices.SortFunc(out, func(a, b Snapshot) int { return strcmp(a.ID, b.ID) }) diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 97bd636..9be63c1 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -2,17 +2,63 @@ package router import ( "bytes" + "context" "encoding/json" "errors" "io" "net/http" "slices" + "strconv" "strings" "time" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/admission" "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" ) +// admit waits for the flow control layer to let the request through. With +// admission disabled it returns immediately, which is the v0.2.0 behaviour. +func (r *Router) admit(ctx context.Context, req *parsedRequest) (*admission.Lease, error) { + if r.adm == nil { + return nil, nil + } + return r.adm.Acquire(ctx, admission.Request{ + Model: req.Model, + Band: req.Band, + // Until authentication arrives the session is the best tenant key we + // have: it groups one caller's traffic without pretending to identify + // them (plan §7.9). + Tenant: req.SessionID, + }) +} + +// rejectAdmission turns a flow control refusal into a status the client can act +// on: 503 when the queue is full (try later), 504 when it timed out waiting. +func (r *Router) rejectAdmission(w http.ResponseWriter, req *parsedRequest, err error, start time.Time) { + status := http.StatusServiceUnavailable + switch { + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + // The client gave up while queued; nobody is listening. + r.m.observeAdmissionRejected(req.Band, "client_gone") + r.m.observeRequest(req.Model, "", 499, time.Since(start), nil, 0) + return + case errors.Is(err, admission.ErrQueueTimeout): + status = http.StatusGatewayTimeout + } + reason := "unavailable" + switch { + case errors.Is(err, admission.ErrQueueFull): + reason = "queue_full" + case errors.Is(err, admission.ErrQueueTimeout): + reason = "ttl" + case errors.Is(err, admission.ErrClosed): + reason = "shutdown" + } + r.m.observeAdmissionRejected(req.Band, reason) + writeError(w, status, "server_error", err.Error()) + r.m.observeRequest(req.Model, "", status, time.Since(start), nil, 0) +} + // prefixBytes is how much of the prompt is kept for prefix-aware routing. // Large enough to cover a system prompt plus the first message, small enough // that hashing it stays cheap. @@ -44,6 +90,7 @@ type parsedRequest struct { Header http.Header SessionID string + Band string ProfileOverride string PromptPrefix string PromptTokens int @@ -80,6 +127,19 @@ func (r *Router) handleInference(path string) http.HandlerFunc { return } + // Step 5 of the lifecycle: "now?" before "where?" (plan §7.1). + lease, err := r.admit(req.Context(), parsed) + if err != nil { + r.rejectAdmission(w, parsed, err, start) + return + } + defer lease.Release() + queued := time.Since(start) + r.m.observeAdmission(lease.Band(), parsed.Model, queued) + if queued > time.Millisecond { + w.Header().Set("x-llmgw-queue-ms", strconv.FormatInt(queued.Milliseconds(), 10)) + } + res := r.dispatch(req.Context(), w, parsed) r.finish(w, parsed, res, start) } @@ -138,6 +198,7 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) { UpstreamBody: body, Header: req.Header, SessionID: sessionID(req, r.cfg.SessionHeader, b.User), + Band: req.Header.Get(r.cfg.Admission.BandHeader), ProfileOverride: req.Header.Get("x-llmgw-profile"), MaxTokens: max(b.MaxTokens, b.MaxCompletionTokens), PromptPrefix: promptPrefix(&b), @@ -234,6 +295,22 @@ func injectIncludeUsage(body []byte) ([]byte, bool) { return out, true } +// servedModels lists the canonical models at least one endpoint serves. +func (r *Router) servedModels() []string { + seen := map[string]struct{}{} + for _, s := range r.reg.Snapshots("") { + for _, m := range s.Models { + seen[m] = struct{}{} + } + } + out := make([]string, 0, len(seen)) + for m := range seen { + out = append(out, m) + } + slices.Sort(out) + return out +} + // handleModels lists the models that at least one endpoint currently serves. func (r *Router) handleModels(w http.ResponseWriter, _ *http.Request) { seen := map[string]struct{}{} diff --git a/internal/router/metrics.go b/internal/router/metrics.go index d07e266..8558514 100644 --- a/internal/router/metrics.go +++ b/internal/router/metrics.go @@ -23,6 +23,9 @@ type metrics struct { scorerValue *prometheus.GaugeVec filtered *prometheus.CounterVec considered prometheus.Histogram + + admissionWait *prometheus.HistogramVec + admissionRejected *prometheus.CounterVec } func newMetrics() *metrics { @@ -64,6 +67,17 @@ func newMetrics() *metrics { Help: "Candidates remaining after filtering, per scheduling decision.", Buckets: []float64{1, 2, 3, 5, 8, 16, 32, 64}, }), + + admissionWait: prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Name: "llmgw_admission_wait_seconds", + Help: "Time spent waiting in the admission queue.", + Buckets: []float64{0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 2.5, 5, 10, 30}, + }, []string{"band", "model"}), + + admissionRejected: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_admission_rejected_total", + Help: "Requests refused by flow control, by reason.", + }, []string{"band", "reason"}), } } @@ -72,9 +86,24 @@ func (m *metrics) collectors() []prometheus.Collector { return []prometheus.Collector{ m.requests, m.duration, m.ttft, m.tokens, m.scorerValue, m.filtered, m.considered, + m.admissionWait, m.admissionRejected, } } +func (m *metrics) observeAdmission(band, model string, wait time.Duration) { + if band == "" { + band = "none" + } + m.admissionWait.WithLabelValues(band, model).Observe(wait.Seconds()) +} + +func (m *metrics) observeAdmissionRejected(band, reason string) { + if band == "" { + band = "unknown" + } + m.admissionRejected.WithLabelValues(band, reason).Inc() +} + func (m *metrics) observeRequest(model, pool string, status int, d time.Duration, u *proxy.Usage, ttft time.Duration) { if model == "" { model = "unknown" diff --git a/internal/router/router.go b/internal/router/router.go index 25bf7dd..8484297 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -16,6 +16,7 @@ import ( "time" "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "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/proxy" @@ -32,6 +33,8 @@ type Router struct { poller *scrape.Poller px *proxy.Proxy 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 — @@ -97,9 +100,66 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { Timeout: cfg.Router.ScrapeTimeout, }, log) + 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 + } 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, + }) + } + 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, + 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 + } + return admission.Saturation(eps, max(0, outstanding-dispatched), admission.Thresholds{ + TargetWaitingPerReplica: r.cfg.Scoring.TargetWaiting, + KVCeiling: r.cfg.Scoring.KVCeiling, + }), true +} + func toInfos(eps []config.Endpoint) []endpoint.Info { out := make([]endpoint.Info, 0, len(eps)) for _, e := range eps { @@ -115,11 +175,22 @@ func toInfos(eps []config.Endpoint) []endpoint.Info { return out } -// Start begins scraping endpoint metrics. It returns immediately. -func (r *Router) Start(ctx context.Context) { r.poller.Sync(ctx) } +// Start begins scraping endpoint metrics and dispatching admissions. +// It returns immediately. +func (r *Router) Start(ctx context.Context) { + r.poller.Sync(ctx) + if r.adm != nil { + r.adm.Start() + } +} -// Stop halts scraping. -func (r *Router) Stop() { r.poller.Stop() } +// Stop halts scraping and admission. +func (r *Router) Stop() { + r.poller.Stop() + if r.adm != nil { + r.adm.Close() + } +} // Registry exposes the endpoint registry, for tests and admin endpoints. func (r *Router) Registry() *endpoint.Registry { return r.reg } diff --git a/llmgw-plan.md b/llmgw-plan.md index 0ae8afb..a53ce2e 100644 --- a/llmgw-plan.md +++ b/llmgw-plan.md @@ -524,18 +524,35 @@ batching to stay efficient, but not enough for the queue to balloon. Set `TargetWaitingPerReplica` too low and the GPU starves; too high and queueing delay returns. Both thresholds must be tunable per model/profile, and `ctl bench` (§10) should propose them. +**Two lags have to be closed, and they need different answers.** The engine's queue is scraped, +so it is up to one interval old; the router's own inflight count is not. But inflight includes +*running* requests, and treating those as queue makes a perfectly healthy engine look overloaded +every interval. What is actually queued is `inflight − batch capacity`, and batch capacity is not +in the metric contract — so llmgw learns it as the high-water mark of observed `running`. The +second lag is shorter and subtler: between admission granting a request and the router +dispatching it, the request is invisible everywhere. A simultaneous burst can therefore pass the +check en masse, which is exactly the batch-job arrival pattern this layer exists to survive. The +controller closes it by counting its own outstanding leases against dispatched inflight. + #### 7.2.2 Priority bands Every API key maps to a band, and the band sets the saturation ceiling: ```yaml bands: - - { name: platinum, priority: 400, saturation_ceiling: 1.00 } # never held back + - { name: platinum, priority: 400, saturation_ceiling: 1.00 } # flows to the healthy limit - { name: gold, priority: 300, saturation_ceiling: 0.85 } - { name: silver, priority: 200, saturation_ceiling: 0.60 } - { name: bronze, priority: 100, saturation_ceiling: 0.30 } # batch jobs ``` +A ceiling of 1.00 means "hold only once the pool is past its healthy limit", +not "never hold". Holding even the top band is deliberate: there is nobody above +it to protect, so the benefit is a bounded engine queue. A flood dumped straight +into a FIFO engine queue makes every request in that band slow, while queueing it +at the gateway preserves ordering, TTL and tenant fairness. A band that must +never wait — a health probe, say — gets a ceiling above 1 (`1000`). + The dispatch loop scans bands from highest to lowest. Once pool saturation exceeds a band's ceiling, that band's requests wait in the queue. The result: as the cluster fills, batch traffic slows down first and interactive traffic is protected. (In llm-d's measurements platinum, gold diff --git a/test/e2e/harness_test.go b/test/e2e/harness_test.go index e61a7d3..09a6d21 100644 --- a/test/e2e/harness_test.go +++ b/test/e2e/harness_test.go @@ -53,10 +53,39 @@ func simConfig(name string, seed int64) sim.Config { return c } +// engineOpt tunes the sim-engines a harness starts. +type engineOpt func(*sim.Config) + +// engineCapacity is the batch capacity used by withLoadedEngines, chosen to sit +// in the range a real vLLM deployment runs at. Making it much smaller would +// shift the bottleneck to slot starvation inside the engine, which admission +// cannot do anything about — the test would then be measuring the wrong thing. +const engineCapacity = 64 + +// withLoadedEngines sizes the engines so that a few hundred concurrent requests +// genuinely overflow them, which is what the priority isolation test needs. +func withLoadedEngines() func(*config.Server) { + return nil // marker; see start() +} + // start brings up n sim-engines behind a router. +// +// A nil tune function is the "small engines" marker, which keeps the call sites +// readable without a second variadic parameter. func start(t *testing.T, n int, tune ...func(*config.Server)) *harness { t.Helper() + small := false + kept := tune[:0] + for _, f := range tune { + if f == nil { + small = true + continue + } + kept = append(kept, f) + } + tune = kept + h := &harness{} cfg := config.Default() cfg.Router.ScrapeInterval = 200 * time.Millisecond @@ -64,7 +93,13 @@ func start(t *testing.T, n int, tune ...func(*config.Server)) *harness { cfg.Router.StaleAfter = 2 * time.Second for i := range n { - e := sim.New(simConfig(fmt.Sprintf("sim-%d", i), int64(i+1))) + sc := simConfig(fmt.Sprintf("sim-%d", i), int64(i+1)) + if small { + sc.MaxNumSeqs = engineCapacity + sc.DecodeTokensPerSec = 800 + sc.BatchContention = 0.005 + } + e := sim.New(sc) srv := httptest.NewServer(sim.NewServer(e, slog.New(slog.DiscardHandler))) h.engines = append(h.engines, engine{eng: e, srv: srv}) cfg.Endpoints = append(cfg.Endpoints, config.Endpoint{ @@ -122,6 +157,15 @@ func (h *harness) waitForFreshMetrics(t *testing.T) { t.Fatal("endpoint metrics never became fresh") } +// running counts requests actively generating across all engines. +func (h *harness) running() int { + n := 0 + for _, e := range h.engines { + n += e.eng.Stats().Running + } + return n +} + func (h *harness) totalRunning() int { n := 0 for _, e := range h.engines { diff --git a/test/e2e/phase15_test.go b/test/e2e/phase15_test.go new file mode 100644 index 0000000..c271e12 --- /dev/null +++ b/test/e2e/phase15_test.go @@ -0,0 +1,203 @@ +package e2e + +import ( + "context" + "fmt" + "slices" + "sync" + "sync/atomic" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" +) + +// TestPriorityIsolation is the phase 1.5 exit criterion. +// +// While a sustained flood of batch requests runs in the bronze band, +// interactive requests in the platinum band must keep their latency. This is +// the whole reason admission exists: an engine queue is FIFO, so once a batch +// job is in it no amount of clever endpoint selection can protect anyone. +// +// It is meant to stay in CI permanently. +func TestPriorityIsolation(t *testing.T) { + workers := envInt("LLMGW_E2E_BATCH", 500) + probes := envInt("LLMGW_E2E_PROBES", 40) + if testing.Short() { + workers, probes = 80, 15 + } + + h := start(t, 3, func(c *config.Server) { + c.Router.Admission.Enabled = true + }, withLoadedEngines()) + + idle := measureProbes(t, h, probes, "platinum") + t.Logf("idle platinum: p50=%v p99=%v", idle.p50.Round(time.Millisecond), idle.p99.Round(time.Millisecond)) + + load := startBronzeLoad(t, h, workers, "bronze") + defer load.stop() + + loaded := measureProbes(t, h, probes, "platinum") + t.Logf("loaded platinum: p50=%v p99=%v (%s)", + loaded.p50.Round(time.Millisecond), loaded.p99.Round(time.Millisecond), load.summary()) + + // The criterion: interactive p99 under load stays within 1.5x of idle. + // The floor keeps a few-millisecond baseline from making the bound absurd. + budget := max(time.Duration(float64(idle.p99)*1.5), idle.p99+50*time.Millisecond) + t.Logf("platinum p99 %v against a %v budget (%.2fx idle)", + loaded.p99.Round(time.Millisecond), budget.Round(time.Millisecond), + float64(loaded.p99)/float64(max(idle.p99, time.Microsecond))) + + if loaded.p99 > budget { + t.Errorf("interactive p99 rose to %v under bronze load, over the %v budget; "+ + "admission control is not protecting the high band", + loaded.p99.Round(time.Millisecond), budget.Round(time.Millisecond)) + } +} + +// TestAdmissionDisabledShowsTheProblem documents why the layer exists: the same +// load does wreck interactive latency once admission is switched off. If this +// ever stops showing degradation, the load generator has stopped being a load +// and the test above proves less than it appears to. +func TestAdmissionDisabledShowsTheProblem(t *testing.T) { + if testing.Short() { + t.Skip("load-dependent") + } + workers := envInt("LLMGW_E2E_BATCH", 500) + + h := start(t, 3, func(c *config.Server) { + c.Router.Admission.Enabled = false + }, withLoadedEngines()) + + idle := measureProbes(t, h, 20, "") + + load := startBronzeLoad(t, h, workers, "") + defer load.stop() + + loaded := measureProbes(t, h, 20, "") + ratio := float64(loaded.p99) / float64(max(idle.p99, time.Microsecond)) + t.Logf("without admission: idle p99=%v loaded p99=%v (%.1fx)", + idle.p99.Round(time.Millisecond), loaded.p99.Round(time.Millisecond), ratio) + + if ratio < 2 { + t.Skipf("the load did not degrade latency (%.1fx); this machine is too fast "+ + "for %d workers, so the isolation test proves less than it should", ratio, workers) + } +} + +// --- load generation --- + +// bronzeLoad is a sustained batch workload. +// +// Workers loop rather than firing once: a one-shot burst drains while the +// probes are still running, and the measurement then catches a recovering +// cluster instead of a loaded one. +type bronzeLoad struct { + cancel context.CancelFunc + wg sync.WaitGroup + done atomic.Int64 + failed atomic.Int64 + firstEr atomic.Value +} + +func startBronzeLoad(t *testing.T, h *harness, workers int, band string) *bronzeLoad { + t.Helper() + + ctx, cancel := context.WithCancel(context.Background()) + l := &bronzeLoad{cancel: cancel} + + for i := range workers { + l.wg.Add(1) + go func() { + defer l.wg.Done() + hdr := map[string]string{"x-session-id": fmt.Sprintf("batch-%d", i%8)} + if band != "" { + hdr["x-llmgw-band"] = band + } + for ctx.Err() == nil { + _, _, err := streamRequest(ctx, h.router.URL, + chatBody(fmt.Sprintf("batch job %d", i), 200, true), hdr) + switch { + case err == nil: + l.done.Add(1) + case ctx.Err() != nil: // shutting down, expected + default: + l.failed.Add(1) + l.firstEr.CompareAndSwap(nil, err) + } + } + }() + } + + l.waitUntilBusy(t, h) + return l +} + +// waitUntilBusy blocks until the engines are demonstrably working, so probes +// are not taken while the load is still ramping up. +// +// It waits on engines being busy rather than full: when admission is doing its +// job the pool is deliberately held below capacity, so a fullness check would +// time out exactly when the system is behaving correctly. +func (l *bronzeLoad) waitUntilBusy(t *testing.T, h *harness) { + t.Helper() + const busyFor = 500 * time.Millisecond + + deadline := time.Now().Add(20 * time.Second) + var since time.Time + for time.Now().Before(deadline) { + if h.running() > 0 { + if since.IsZero() { + since = time.Now() + } else if time.Since(since) >= busyFor { + return + } + } else { + since = time.Time{} + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("engines never became busy; %s", l.summary()) +} + +func (l *bronzeLoad) summary() string { + return fmt.Sprintf("bronze: completed=%d failed=%d firstErr=%v", + l.done.Load(), l.failed.Load(), l.firstEr.Load()) +} + +func (l *bronzeLoad) stop() { + l.cancel() + l.wg.Wait() +} + +// --- probes --- + +type latency struct{ p50, p99 time.Duration } + +// measureProbes sends sequential short requests and reports their latency. +// Sequential on purpose: these stand in for a single interactive user, and +// concurrency would measure something else. +func measureProbes(t *testing.T, h *harness, n int, band string) latency { + t.Helper() + hdr := map[string]string{"x-session-id": "probe"} + if band != "" { + hdr["x-llmgw-band"] = band + } + + samples := make([]time.Duration, 0, n) + for range n { + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + start := time.Now() + _, _, err := streamRequest(ctx, h.router.URL, chatBody("ping", 4, true), hdr) + cancel() + if err != nil { + t.Fatalf("probe failed: %v", err) + } + samples = append(samples, time.Since(start)) + } + slices.Sort(samples) + return latency{ + p50: samples[len(samples)/2], + p99: samples[min(int(float64(len(samples))*0.99), len(samples)-1)], + } +}