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:
+112
-12
@@ -69,9 +69,10 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
// Restore the body so it can be read again by the upstream
|
||||
r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
||||
|
||||
// Parse the JSON to find the model field
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal(bodyBytes, &payload); err != nil {
|
||||
// Extract only the top-level "model" value with a token walk: the common
|
||||
// path must not materialise the whole body (messages and all) as a map.
|
||||
modelRaw, hasModel, err := extractTopLevelModel(bodyBytes)
|
||||
if err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "invalid_json",
|
||||
Message: "request body is not valid JSON",
|
||||
@@ -79,7 +80,6 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
}
|
||||
|
||||
// Extract the model name
|
||||
modelVal, hasModel := payload["model"]
|
||||
if !hasModel {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "missing_model",
|
||||
@@ -88,7 +88,7 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
}
|
||||
|
||||
// Handle null model
|
||||
if modelVal == nil {
|
||||
if bytes.Equal(bytes.TrimSpace(modelRaw), []byte("null")) {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "null_model",
|
||||
Message: "'model' field is null",
|
||||
@@ -96,8 +96,8 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
}
|
||||
|
||||
// Extract as string
|
||||
modelName, ok := modelVal.(string)
|
||||
if !ok {
|
||||
var modelName string
|
||||
if err := json.Unmarshal(modelRaw, &modelName); err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "missing_model",
|
||||
Message: "'model' field must be a string",
|
||||
@@ -122,10 +122,26 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// If upstream expects a different model name, rewrite the body
|
||||
// If upstream expects a different model name, rewrite the body. This path is
|
||||
// uncommon, so it may decode the body as raw fields; the common path above
|
||||
// does not build a full map.
|
||||
if modelUpstream.UpstreamModel != "" && modelUpstream.UpstreamModel != modelName {
|
||||
payload["model"] = modelUpstream.UpstreamModel
|
||||
newBody, err := json.Marshal(payload)
|
||||
var fields map[string]json.RawMessage
|
||||
if err := json.Unmarshal(bodyBytes, &fields); err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "invalid_request",
|
||||
Message: fmt.Sprintf("failed to rewrite model name: %v", err),
|
||||
}
|
||||
}
|
||||
newModel, err := json.Marshal(modelUpstream.UpstreamModel)
|
||||
if err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "invalid_request",
|
||||
Message: fmt.Sprintf("failed to rewrite model name: %v", err),
|
||||
}
|
||||
}
|
||||
fields["model"] = newModel
|
||||
newBody, err := json.Marshal(fields)
|
||||
if err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "invalid_request",
|
||||
@@ -151,11 +167,11 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
ConnectTimeout: h.defaultConnectTimeout,
|
||||
ReadTimeout: h.defaultReadTimeout,
|
||||
WriteTimeout: h.defaultWriteTimeout,
|
||||
MaxBodySize: h.defaultMaxBodySize,
|
||||
MaxBodySize: h.maxBodyBytes,
|
||||
AuthRequired: false,
|
||||
}
|
||||
|
||||
targetURL, _ := url.Parse("http://" + modelUpstream.Address)
|
||||
targetURL := &url.URL{Scheme: "http", Host: modelUpstream.Address}
|
||||
|
||||
route := &Route{
|
||||
Name: "v1-chat-" + modelName,
|
||||
@@ -169,3 +185,87 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
|
||||
return route, nil
|
||||
}
|
||||
|
||||
// extractTopLevelModel walks a JSON object body with a streaming decoder and
|
||||
// returns the raw value of the top-level "model" field. It never builds the
|
||||
// full document as a map, so a large "messages" array is skipped by token.
|
||||
// hasModel is false when the key is absent or the document is a bare null.
|
||||
func extractTopLevelModel(body []byte) (model json.RawMessage, hasModel bool, err error) {
|
||||
dec := json.NewDecoder(bytes.NewReader(body))
|
||||
tok, err := dec.Token()
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if tok == nil {
|
||||
// A bare `null` body decoded to a nil map before, i.e. no model field.
|
||||
if _, err := dec.Token(); err != io.EOF {
|
||||
if err == nil {
|
||||
return nil, false, fmt.Errorf("unexpected data after JSON null")
|
||||
}
|
||||
return nil, false, err
|
||||
}
|
||||
return nil, false, nil
|
||||
}
|
||||
if delim, ok := tok.(json.Delim); !ok || delim != '{' {
|
||||
return nil, false, fmt.Errorf("request body is not a JSON object")
|
||||
}
|
||||
for dec.More() {
|
||||
keyTok, err := dec.Token()
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
key, _ := keyTok.(string)
|
||||
if key == "model" {
|
||||
if err := dec.Decode(&model); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
hasModel = true
|
||||
continue
|
||||
}
|
||||
if err := skipJSONValue(dec); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
}
|
||||
// Consume the closing brace and reject trailing data, preserving the
|
||||
// invalid_json answer the old whole-body decode produced.
|
||||
if _, err := dec.Token(); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if _, err := dec.Token(); err != io.EOF {
|
||||
if err == nil {
|
||||
return nil, false, fmt.Errorf("unexpected data after JSON object")
|
||||
}
|
||||
return nil, false, err
|
||||
}
|
||||
return model, hasModel, nil
|
||||
}
|
||||
|
||||
// skipJSONValue advances dec past the next JSON value without materialising it.
|
||||
func skipJSONValue(dec *json.Decoder) error {
|
||||
tok, err := dec.Token()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
delim, ok := tok.(json.Delim)
|
||||
if !ok {
|
||||
return nil // scalar
|
||||
}
|
||||
if delim != '{' && delim != '[' {
|
||||
return fmt.Errorf("unexpected delimiter %v", delim)
|
||||
}
|
||||
depth := 1
|
||||
for depth > 0 {
|
||||
tok, err := dec.Token()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if d, ok := tok.(json.Delim); ok {
|
||||
if d == '{' || d == '[' {
|
||||
depth++
|
||||
} else {
|
||||
depth--
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user