feat(proxy): engine-wide LLM admission control via vLLM gauges (#64)
CI / CI (pull_request) Successful in 34m5s
CI / CI (pull_request) Successful in 34m5s
The per-pod semaphore was not a cap: 3 replicas x 4 slots advertised 12 concurrent requests while the engine runs 4, and nothing coordinated the replicas. Admission is now gated on the engine's own published gauges, which is global by construction because every replica contends for the same engine. - new internal/proxy/engine_state.go: 2s scrape of the backend /metrics, hand-rolled exposition parse (no new dependency), unhealthy after 3 consecutive failed polls, and a startup self-check that raises observer_degraded rather than falling back to fail-open silently - gate: park 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 state fails open on purpose - streaming requests receive SSE keep-alive heartbeats while parked - router: extract only the top-level model field instead of unmarshalling the whole body into a map; body cap from MAX_BODY_BYTES (1 MiB, the model path previously inherited the 100 MB default) - GOMAXPROCS from the cgroup CPU quota: the container is limited to 500m while Go defaulted to the host core count and was CFS-throttled - VLLM_MAX_QUEUE_DEPTH default 20 -> 8 - delete k8s/configmap.yaml: absent from kustomization.yaml, and the live ConfigMap of that name carries an unrelated key - correct the llm-client NetworkPolicy comment; llm-serving-ingress now selects by namespace, not by that label Observers run only for models named in ADMISSION_OBSERVE_MODELS. reasoning is the only vLLM engine here; ornith and qwen-cpu are Ollama and embeddings/rerank are TEI, none of which publish vllm:* series, so observing them would hold observer_degraded at 1 permanently. Companion change (engine flags + dashboard panels) is riotpiao-poimen/homelab#51.
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