Files
homelab-frontend/cmd/gateway/main.go
T
poimenandrock 369d51a5c2
CI / CI (push) Successful in 41m20s
feat(proxy): engine-wide LLM admission control via vLLM gauges (#65)
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]>
2026-09-30 10:56:10 +00:00

238 lines
8.3 KiB
Go

package main
import (
"context"
"fmt"
"log"
"math"
"net/http"
"os"
"os/signal"
"runtime"
"strconv"
"strings"
"syscall"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/notification"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/proxy"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/server"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/serviceadapter"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/temporal"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/tracing"
)
func main() {
// Size the Go scheduler to the container's CPU quota before any server
// goroutine starts.
configureGOMAXPROCS()
ctx := context.Background()
// Initialize OpenTelemetry tracing
tracingCfg := tracing.DefaultConfig()
shutdownTracer, err := tracing.Init(ctx, tracingCfg)
if err != nil {
log.Printf("warning: failed to initialize tracing: %v", err)
} else {
log.Printf("tracing initialized: service=%s endpoint=%s", tracingCfg.ServiceName, tracingCfg.OTLPEndpoint)
defer func() {
if err := shutdownTracer(ctx); err != nil {
log.Printf("error shutting down tracer: %v", err)
}
}()
}
// Load configuration
cfg, err := config.Load()
if err != nil {
fmt.Fprintf(os.Stderr, "failed to load config: %v\n", err)
os.Exit(1)
}
// Determine if auth is enabled by checking if any route requires it
authEnabled := false
for _, route := range cfg.Routes {
if route.Upstream.AuthRequired {
authEnabled = true
break
}
}
// Create the reverse proxy handler that routes requests based on configuration
upstreamHandler := proxy.New(cfg)
// Create the Temporal workflow handler
// Temporal server address can be configured via environment variable
temporalHostPort := os.Getenv("TEMPORAL_HOST_PORT")
if temporalHostPort == "" {
temporalHostPort = "localhost:7233"
}
log.Printf("Temporal server: %s", temporalHostPort)
temporalHandler := temporal.NewHandler(temporalHostPort)
// Create server with health checker
srv := server.New(cfg.ListenAddr, cfg.ShutdownTimeout, nil)
// Initialize health checker with config validity and auth status
healthChecker := server.NewHealthChecker(true, authEnabled)
srv.SetHealthChecker(healthChecker)
// Create ServiceAdapter registry and dispatcher (phase 8)
registry := serviceadapter.NewRegistry(nil)
// Add workflow service adapter (uses Temporal handler for gRPC forwarding)
workflowSpec := serviceadapter.GetWorkflowSpec()
workflowAdapterHandler := serviceadapter.NewWorkflowAdapter(temporalHandler)
workflowAdapter := &serviceadapter.ServiceAdapter{
Namespace: "temporal",
ServiceName: "workflow",
Handler: workflowAdapterHandler,
Spec: *workflowSpec,
}
_ = registry.Add(workflowAdapter)
// Add notification service adapter (internal handler, no upstream proxy)
notifHandler := notification.NewHandler()
notifAdapter := &serviceadapter.ServiceAdapter{
Namespace: "notification",
ServiceName: "notification",
Handler: notifHandler,
Spec: serviceadapter.Spec{
ServiceName: "notification",
Auth: serviceadapter.Auth{Required: true},
Resources: []serviceadapter.Resource{
{Name: "send-email", Methods: []serviceadapter.Method{{Verb: "POST", UpstreamPath: "/send-email"}}},
{Name: "send-message", Methods: []serviceadapter.Method{{Verb: "POST", UpstreamPath: "/send-message"}}},
{Name: "list-messages", Methods: []serviceadapter.Method{{Verb: "GET", UpstreamPath: "/list-messages"}}},
{Name: "delete-message", Methods: []serviceadapter.Method{{Verb: "DELETE", UpstreamPath: "/delete-message"}}},
{Name: "delete-all-messages", Methods: []serviceadapter.Method{{Verb: "DELETE", UpstreamPath: "/delete-all-messages"}}},
{Name: "list-applications", Methods: []serviceadapter.Method{{Verb: "GET", UpstreamPath: "/list-applications"}}},
{Name: "create-application", Methods: []serviceadapter.Method{{Verb: "POST", UpstreamPath: "/create-application"}}},
{Name: "delete-application", Methods: []serviceadapter.Method{{Verb: "DELETE", UpstreamPath: "/delete-application"}}},
},
},
}
_ = registry.Add(notifAdapter)
// Add Poimen orchestrator adapter (typed REST for AI orchestrator frontend)
poimenAdapter := serviceadapter.NewPoimenAdapter(temporalHandler)
_ = registry.Add(poimenAdapter)
// Add memory service adapter (code-registered; config memory entry is skipped).
// Includes standing-context resource for CAG support (issue #44).
_ = registry.Add(serviceadapter.NewMemoryAdapter())
// Add sessions service adapter — REST CRUD bridged to poimen-memory ingest
// and poimen-graphql reads. Issue #50.
_ = registry.Add(serviceadapter.NewSessionAdapter())
// Add other adapters from config (skip if already registered in code)
for _, a := range cfg.Adapters {
if existing := registry.Get(a.ServiceName); existing != nil {
log.Printf("skip config adapter '%s': already registered with internal handler", a.ServiceName)
continue
}
_ = registry.Add(a)
}
log.Printf("%d service adapters loaded", registry.Count())
// Create shared JWT validator for X-Service auth enforcement
var jwtValidator *auth.Validator
if cfg.Auth.Enabled && cfg.Auth.JWKSURL != "" {
jwtValidator = auth.NewValidator(cfg.Auth.Issuer, cfg.Auth.Audience, cfg.Auth.JWKSURL)
}
dispatcher := serviceadapter.NewDispatcher(registry, jwtValidator)
// Wire workflow adapter to temporal handler for proper request forwarding
// workflowAdapterHandler (above) handles JSON-to-gRPC translation for workflow service
// Create router that handles health endpoints, X-Service (ServiceAdapter) routing,
// temporal endpoints, and passes others to upstream handler
router := server.NewRouter(healthChecker, dispatcher, temporalHandler, upstreamHandler)
router.SetMetricsHandler(observability.NewMetricsHandler(observability.Global()))
// Wrap router with tracing middleware
tracedRouter := tracing.Middleware(router)
srv.SetHandler(tracedRouter)
// Set up signal handling
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
// Start server in a goroutine
var serverErr error
go func() {
log.Printf("gateway listening on %s", srv.Addr())
serverErr = srv.ListenAndServe()
if serverErr != nil && serverErr != http.ErrServerClosed {
log.Printf("server error: %v", serverErr)
}
}()
// Wait for shutdown signal
sig := <-sigChan
log.Printf("received signal: %v", sig)
// Gracefully shutdown the server
if err := srv.Shutdown(context.Background()); err != nil {
fmt.Fprintf(os.Stderr, "shutdown error: %v\n", err)
os.Exit(1)
}
log.Printf("gateway shutdown complete")
os.Exit(0)
}
// configureGOMAXPROCS sets GOMAXPROCS from the cgroup CPU quota so the runtime
// sizes its scheduler to the container limit rather than the host CPU count.
// The default is left alone when no quota is configured.
func configureGOMAXPROCS() {
quota, period := readCPUQuota()
if quota <= 0 || period <= 0 {
return
}
n := int(math.Ceil(float64(quota) / float64(period)))
if n < 1 {
n = 1
}
runtime.GOMAXPROCS(n)
log.Printf("GOMAXPROCS set to %d (cgroup cpu quota %d/%d)", n, quota, period)
}
// readCPUQuota returns the CPU quota and period from cgroup v2, falling back to
// cgroup v1. A quota of 0 means no limit.
func readCPUQuota() (int64, int64) {
// cgroup v2: "<quota> <period>" or "max <period>".
if data, err := os.ReadFile("/sys/fs/cgroup/cpu.max"); err == nil {
fields := strings.Fields(string(data))
if len(fields) >= 2 {
period, perr := strconv.ParseInt(fields[1], 10, 64)
if perr == nil && period > 0 {
if fields[0] == "max" {
return 0, period
}
if quota, qerr := strconv.ParseInt(fields[0], 10, 64); qerr == nil {
return quota, period
}
}
}
}
// cgroup v1: cpu.cfs_quota_us is -1 when unlimited.
period := int64(100000)
if data, err := os.ReadFile("/sys/fs/cgroup/cpu/cpu.cfs_period_us"); err == nil {
if v, perr := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64); perr == nil && v > 0 {
period = v
}
}
if data, err := os.ReadFile("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"); err == nil {
if v, qerr := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64); qerr == nil {
return v, period
}
}
return 0, 0
}