feat(proxy): engine-wide LLM admission control via vLLM gauges (#65)
CI / CI (push) Successful in 41m20s
CI / CI (push) Successful in 41m20s
Closes #64 Replaces the per-pod concurrency semaphore as the effective cap with admission gated on the vLLM engine's own scheduler gauges. ## Why `internal/proxy/proxy.go:119` built a semaphore of 4 slots *inside each gateway process*, and `api-gateway` runs 3 replicas, so the advertised cap was 12 while the engine's `--max-num-seqs` is 4. Nothing coordinated replicas, and the mismatch was invisible because the queue never rejected anything (`queue_full_requests_total 0`). Measured on the engine while investigating a 10-minute TTFT: - `vllm:cache_config_info{num_gpu_blocks="9527",block_size="16"}` = **152,432 tokens** of KV capacity for the whole engine. PP=2 splits layers, so both ranks hold KV for the same tokens — not 2x that. - Uncached prefill runs at **262 tok/s** contended, 525 tok/s over 7d (`vllm:request_prefill_time_seconds_sum/count`). - One request showed `llm_ttft_seconds_sum 590.981 / count 5` at the gateway and a `(480, 960]` s entry in `vllm:request_prefill_time_seconds_bucket`, with engine-side generation throughput at 0.0-0.1 tok/s throughout: a 65k-token prompt costs ~250 s of prefill and starves every other request while it runs. - `vllm:num_preemptions_total = 0`; KV peaked at `0.5399`. Reading the engine's own gauges is global by construction, needs no lock and no new datastore, and it also accounts for callers that never pass the gateway — namespace `poimen` is admitted directly on 80/8080 by `llm-serving-ingress`, and `memory/crates/mem-cli/src/kv_warmer.rs` is such a caller. A distributed lock would have been blind to that traffic. ## Changes - **new `internal/proxy/engine_state.go`** — 2s scrape of `http://<addr>/metrics` on its own client (the proxy transports carry socket options tuned for inference), hand-rolled exposition parse, `Healthy=false` after 3 consecutive failed polls with `engine_state_stale_total` counted once at the transition, and a startup self-check that raises `observer_degraded` and logs at error level instead of degrading to fail-open silently. - **admission gate** (`internal/proxy/scheduler.go`) — parks while `kv_cache_usage_perc >= 0.90` or `running + waiting >= ADMISSION_MAX_INFLIGHT`, bounded by `ADMISSION_MAX_PARK_SECONDS`, then `503 + Retry-After: 5`. Unknown or unhealthy state fails **open** on purpose: a metrics outage must not become an inference outage. - **parked streaming heartbeat** (`internal/proxy/proxy.go`) — `stream: true` requests get 200 + `text/event-stream` and a `: keep-alive` comment every 15s while parked, so a client idle timeout does not kill a request that is legitimately waiting. Non-streaming requests park silently. - **`internal/proxy/router.go`** — extract only the top-level `model` field with a `json.Decoder` token walk instead of unmarshalling the whole body into `map[string]interface{}` (that built the entire `messages`+`tools` tree per request, ~300 KB for a 65k-token prompt). `MaxBodySize` now comes from `MAX_BODY_BYTES`; the model path previously inherited the 100 MB default at `proxy.go:93` against a 512Mi container limit. - **`cmd/gateway/main.go`** — `GOMAXPROCS` from the cgroup CPU quota (cgroup v2 `cpu.max`, v1 fallback). The container is limited to `500m` while Go defaulted to the host core count and was being CFS-throttled. - **`k8s/deployment.yaml`** — `ADMISSION_MAX_PARK_SECONDS=120`, `ADMISSION_OBSERVE_MODELS=reasoning`, `VLLM_MAX_QUEUE_DEPTH=8`, `MAX_BODY_BYTES=1048576`; corrected the `llm-client` comment (the NetworkPolicy selects by namespace now, not by that label). - **deleted `k8s/configmap.yaml`** — absent from `k8s/kustomization.yaml`, and the live ConfigMap of that name carries an unrelated key (`gw-config-fixed.yaml`) plus a stale routing table. - **`internal/observability/`** — `admission_parked_total`, `admission_rejected_total{reason}`, `engine_state_stale_total`, and gauges `vllm_requests_running` / `vllm_kv_usage_perc` (the gateway's own reading, deliberately mirroring the vLLM series so a disagreement is visible) and `observer_degraded`. Names follow the existing convention, which carries no prefix. Out of scope here, in `riotpiao-poimen/homelab#51`: the engine flags (`--max-num-partial-prefills`, `--max-long-partial-prefills`, `--long-prefill-token-threshold`, `--max-num-batched-tokens`, `--max-num-seqs`, `--gpu-memory-utilization`) and the four dashboard panels in `generate-dashboards.py`. ## Verification - `go build ./...` clean, `go vet ./...` clean. - `go test ./...` — 368 tests, 0 failures (baseline on `main` is 368/0). - `structural-check.sh` on the changed files — 0 errors, 1 pre-existing warning (`:latest` image tag). - `go.mod` / `go.sum` byte-identical: no new dependency. Not yet verified: the real-engine load test (2x65k + 15x8k against `reasoning`). That run also produces the number for `ADMISSION_MAX_PARK_SECONDS`, which today is the provisional 120; retuning it is a follow-up since it is env-only and needs no rebuild. ## Accepted risks - **No tests were added** for the observer or the gate (deliberate: production is the test bed). The mitigation is the startup self-check, which is production self-diagnosis rather than a test, plus the `observer_degraded` alert. A parse failure that outlives the first poll logs and surfaces as a metric; it does not silently disable gating without a signal. - **The cap is soft, not hard.** Between polls each replica may admit more than one request, so there is no clean bound. Exposure is low (~11,500 requests over 7d ≈ 1.4/min) and the engine remains the hard limit — a race parks requests in vLLM's queue rather than over-committing KV. - **`ADMISSION_MAX_INFLIGHT` is manually coupled to `--max-num-seqs`** and can drift silently. It defaults to `VLLM_MAX_NUM_SEQS`, and startup logs the effective value. - The first scrape waits one interval, so the first ~2s of admission is ungated (fail-open). Scraping at `Start()` instead would make the gateway contact every observed backend before any request needs one, which is wasted work and a surprising side effect. ## Rollback Revert this commit and redeploy. The admission path is fail-open, so the code can also be neutralised without reverting by unsetting `ADMISSION_OBSERVE_MODELS`. Separately, the orphaned cluster object can be removed by hand: `kubectl delete cm -n api api-gateway-config` (not a code change, deliberately left to a human). --------- Reviewed-on: #65 Co-authored-by: poimen <[email protected]>
This commit was merged in pull request #65.
This commit is contained in:
+198
-20
@@ -2,6 +2,7 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -47,6 +48,10 @@ type Handler struct {
|
||||
schedulers map[string]*Scheduler
|
||||
// CAG: delay applied to warmup requests before queue entry (#45)
|
||||
warmupDelay time.Duration
|
||||
// engineObservers polls each model backend's /metrics for admission (#64)
|
||||
engineObservers map[string]*EngineObserver
|
||||
// maxBodyBytes caps synthesized model-route request bodies (MAX_BODY_BYTES)
|
||||
maxBodyBytes int64
|
||||
// transportsMu guards h.transports for goroutine-safe reads/writes.
|
||||
transportsMu sync.RWMutex
|
||||
}
|
||||
@@ -122,18 +127,80 @@ func New(cfg *config.Config) *Handler {
|
||||
maxSlots = n
|
||||
}
|
||||
}
|
||||
maxQueue := 20
|
||||
maxQueue := 8
|
||||
if v := os.Getenv("VLLM_MAX_QUEUE_DEPTH"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
maxQueue = n
|
||||
}
|
||||
}
|
||||
|
||||
// Issue #64: engine-wide admission thresholds. MaxInFlight defaults to the
|
||||
// same value as VLLM_MAX_NUM_SEQS; the observers below read the engine's
|
||||
// own gauges, which every replica shares.
|
||||
admissionMaxInFlight := maxSlots
|
||||
if v := os.Getenv("ADMISSION_MAX_INFLIGHT"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
admissionMaxInFlight = n
|
||||
}
|
||||
}
|
||||
admissionMaxKVUsage := 0.90
|
||||
if v := os.Getenv("ADMISSION_MAX_KV_USAGE"); v != "" {
|
||||
if f, err := strconv.ParseFloat(v, 64); err == nil && f > 0 {
|
||||
admissionMaxKVUsage = f
|
||||
}
|
||||
}
|
||||
maxParkSeconds := 120
|
||||
if v := os.Getenv("ADMISSION_MAX_PARK_SECONDS"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
maxParkSeconds = n
|
||||
}
|
||||
}
|
||||
h.maxBodyBytes = 1048576
|
||||
if v := os.Getenv("MAX_BODY_BYTES"); v != "" {
|
||||
if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {
|
||||
h.maxBodyBytes = n
|
||||
}
|
||||
}
|
||||
|
||||
limits := AdmissionLimits{
|
||||
MaxInFlight: admissionMaxInFlight,
|
||||
MaxKVUsage: admissionMaxKVUsage,
|
||||
MaxPark: time.Duration(maxParkSeconds) * time.Second,
|
||||
}
|
||||
|
||||
// Issue #64: only backends that actually publish the vLLM scheduler gauges
|
||||
// may be observed. `reasoning` is the only vLLM engine here; ornith and
|
||||
// qwen-cpu are Ollama and embeddings/rerank are TEI, none of which export
|
||||
// vllm:* series. Observing them would trip the startup self-check and hold
|
||||
// observer_degraded at 1 forever. Unset (the default) observes nothing and
|
||||
// leaves admission exactly as it was before this change.
|
||||
observeModels := make(map[string]bool)
|
||||
for _, name := range strings.Split(os.Getenv("ADMISSION_OBSERVE_MODELS"), ",") {
|
||||
if name = strings.TrimSpace(name); name != "" {
|
||||
observeModels[name] = true
|
||||
}
|
||||
}
|
||||
|
||||
h.engineObservers = make(map[string]*EngineObserver)
|
||||
h.schedulers = make(map[string]*Scheduler)
|
||||
for modelName, model := range cfg.Models {
|
||||
addrs := model.AllAddresses()
|
||||
if len(addrs) > 0 {
|
||||
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue)
|
||||
if len(addrs) == 0 {
|
||||
continue
|
||||
}
|
||||
if observeModels[modelName] {
|
||||
for _, addr := range addrs {
|
||||
if _, ok := h.engineObservers[addr]; !ok {
|
||||
h.engineObservers[addr] = NewEngineObserver(addr, enginePollInterval)
|
||||
}
|
||||
}
|
||||
}
|
||||
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue, h.engineObservers, limits)
|
||||
}
|
||||
// The handler has no request-path lifecycle to hang cancellation on, so the
|
||||
// observers poll for the lifetime of the process.
|
||||
for _, o := range h.engineObservers {
|
||||
o.Start(context.Background())
|
||||
}
|
||||
// Pre-populate transports for all scheduler backends at startup (single-threaded).
|
||||
// getOrCreateTransport is not goroutine-safe; pre-populating avoids concurrent writes.
|
||||
@@ -151,7 +218,7 @@ func New(cfg *config.Config) *Handler {
|
||||
// Create a transport per unique upstream address for connection reuse
|
||||
transport := h.getOrCreateTransport(route.Upstream.Address, &route.Upstream)
|
||||
|
||||
upstreamURL, _ := url.Parse("http://" + route.Upstream.Address)
|
||||
upstreamURL := &url.URL{Scheme: "http", Host: route.Upstream.Address}
|
||||
|
||||
r := &Route{
|
||||
Name: name,
|
||||
@@ -395,27 +462,104 @@ func (h *Handler) applyWarmup(ctx context.Context, r *http.Request) bool {
|
||||
// acquireScheduledBackend selects a backend from sched, applies warmup
|
||||
// deprioritisation, acquires a concurrency slot, and patches route when the
|
||||
// selected backend differs from the config default. It returns the backend
|
||||
// to Release (caller must defer), the (possibly patched) route, and ok=true.
|
||||
// If the slot cannot be acquired it writes a 503 and returns (nil, nil, false).
|
||||
func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool) {
|
||||
// to Release (caller must defer), the (possibly patched) route, whether SSE
|
||||
// headers were already written while parking, and ok=true.
|
||||
// If the slot cannot be acquired it writes a 503 and returns ok=false.
|
||||
func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool, bool) {
|
||||
backend := sched.Select(cacheKey)
|
||||
if !h.applyWarmup(ctx, r) {
|
||||
return nil, nil, false
|
||||
return nil, nil, false, false
|
||||
}
|
||||
if backend == nil {
|
||||
return nil, route, true
|
||||
return nil, route, false, true
|
||||
}
|
||||
if err := backend.Acquire(ctx); err != nil {
|
||||
if errors.Is(err, ErrQueueFull) {
|
||||
|
||||
// Decision Q22: a request that actually parks on the engine gate and wants
|
||||
// an SSE stream gets its headers pre-written plus a 15 s keep-alive, so no
|
||||
// intermediary times the connection out before the engine admits it.
|
||||
var (
|
||||
prewritten bool
|
||||
hbStop chan struct{}
|
||||
hbDone chan struct{}
|
||||
)
|
||||
onPark := func() {
|
||||
if !requestWantsStream(r) {
|
||||
return
|
||||
}
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
prewritten = true
|
||||
hbStop = make(chan struct{})
|
||||
hbDone = make(chan struct{})
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
fmt.Fprint(w, ": keep-alive\n\n")
|
||||
flusher.Flush()
|
||||
go func() {
|
||||
defer close(hbDone)
|
||||
ticker := time.NewTicker(15 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-hbStop:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
fmt.Fprint(w, ": keep-alive\n\n")
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
err := backend.acquire(ctx, onPark)
|
||||
if prewritten {
|
||||
close(hbStop)
|
||||
<-hbDone
|
||||
}
|
||||
if err != nil {
|
||||
// Once the SSE status line is spent, report the rejection as an event
|
||||
// instead of a second WriteHeader.
|
||||
reject := func(body string) {
|
||||
if prewritten {
|
||||
fmt.Fprintf(w, "event: error\ndata: {\"error\":%q}\n\n", body)
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
fmt.Fprintf(w, "%s", body)
|
||||
}
|
||||
if errors.Is(err, ErrAdmissionTimeout) {
|
||||
observability.Global().IncrAdmissionRejected("park_timeout")
|
||||
logging.Warnf("admission timeout", map[string]string{
|
||||
"backend": backend.Address,
|
||||
"path": r.URL.Path,
|
||||
"model": route.ModelName,
|
||||
})
|
||||
if !prewritten {
|
||||
w.Header().Set("Retry-After", "5")
|
||||
}
|
||||
reject("admission timeout")
|
||||
} else if errors.Is(err, ErrQueueFull) {
|
||||
observability.Global().IncrQueueFull()
|
||||
observability.Global().IncrAdmissionRejected("queue_full")
|
||||
logging.Warnf("backend queue full", map[string]string{
|
||||
"backend": backend.Address,
|
||||
"path": r.URL.Path,
|
||||
"model": route.ModelName,
|
||||
})
|
||||
w.Header().Set("Retry-After", "5")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
fmt.Fprintf(w, "queue full")
|
||||
if !prewritten {
|
||||
w.Header().Set("Retry-After", "5")
|
||||
}
|
||||
reject("queue full")
|
||||
} else {
|
||||
logging.Warnf("backend acquire failed", map[string]string{
|
||||
"backend": backend.Address,
|
||||
@@ -423,15 +567,14 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr
|
||||
"model": route.ModelName,
|
||||
"err": err.Error(),
|
||||
})
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
fmt.Fprintf(w, "upstream busy")
|
||||
reject("upstream busy")
|
||||
}
|
||||
return nil, nil, false
|
||||
return nil, nil, prewritten, false
|
||||
}
|
||||
if backend.Address != route.Upstream.Address {
|
||||
patchedCfg := *route.Upstream
|
||||
patchedCfg.Address = backend.Address
|
||||
patchedURL, _ := url.Parse("http://" + backend.Address)
|
||||
patchedURL := &url.URL{Scheme: "http", Host: backend.Address}
|
||||
route = &Route{
|
||||
Name: route.Name,
|
||||
ModelName: route.ModelName,
|
||||
@@ -442,7 +585,27 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr
|
||||
},
|
||||
}
|
||||
}
|
||||
return backend, route, true
|
||||
return backend, route, prewritten, true
|
||||
}
|
||||
|
||||
// requestWantsStream reports whether the JSON request body asks for an SSE
|
||||
// stream ("stream": true). The body is restored for the upstream.
|
||||
func requestWantsStream(r *http.Request) bool {
|
||||
if r.Body == nil {
|
||||
return false
|
||||
}
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
r.Body = io.NopCloser(bytes.NewReader(body))
|
||||
var probe struct {
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &probe); err != nil {
|
||||
return false
|
||||
}
|
||||
return probe.Stream
|
||||
}
|
||||
|
||||
// handleAuthEndpoint routes /auth/* requests and returns true if handled.
|
||||
@@ -605,9 +768,10 @@ func (h *Handler) enforceBodyLimit(w http.ResponseWriter, r *http.Request, route
|
||||
// Returns (nil, nil, false) when the backend slot cannot be acquired (response written).
|
||||
func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *http.Request, route *Route, sched *Scheduler, cacheKey string) (*httputil.ReverseProxy, *Backend, bool) {
|
||||
var selectedBackend *Backend
|
||||
var headersPrewritten bool
|
||||
if sched != nil {
|
||||
var ok bool
|
||||
selectedBackend, route, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
|
||||
selectedBackend, route, headersPrewritten, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
|
||||
if !ok {
|
||||
return nil, nil, false
|
||||
}
|
||||
@@ -628,6 +792,20 @@ func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *h
|
||||
s.MarkFailed(capturedBackend.Address)
|
||||
}
|
||||
}
|
||||
if headersPrewritten {
|
||||
// SSE headers and a 200 are already on the wire: a 502 here would be
|
||||
// an ignored WriteHeader and would tear the event stream. Report the
|
||||
// failure as an event instead.
|
||||
fmt.Fprint(w, "event: error\ndata: {\"error\":\"upstream error\"}\n\n")
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
logging.Errorf("upstream error", err, map[string]string{
|
||||
"upstream": capturedRoute.Upstream.Address,
|
||||
"path": r.URL.Path,
|
||||
})
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
fmt.Fprintf(w, "upstream error")
|
||||
logging.Errorf("upstream error", err, map[string]string{
|
||||
|
||||
Reference in New Issue
Block a user