diff --git a/CHANGELOG.md b/CHANGELOG.md index 81c4e05..574f414 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,8 +6,8 @@ is only released once its **exit criterion** is met, not when its code lands. | Version | Phase | Scope | State | |---|---|---|---| | v0.1.0 | 0 | Skeleton and sim-engine | released | -| v0.2.0 | 1 | Static router: streaming proxy, filter/score/pick | in progress | -| v0.3.0 | 1.5 | Admission control: priority bands, saturation detector | planned | +| 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.5.0 | 3 | Control plane: catalog, placement, reconciler | planned | | v0.6.0 | 4 | Multi-node replicas and fabric awareness | planned | @@ -16,6 +16,55 @@ 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.2.0 — 2026-09-16 + +The router serves traffic. `llmgw server` is real: it fronts a static set of +engine replicas with an OpenAI-compatible API. + +### Added + +- **Streaming proxy.** Relays SSE event by event with a flush after each, and + propagates cancellation: when a client disconnects, the upstream request is + aborted so the GPU stops working. Separate deadlines for dial, response + headers, first token, inter-token idle and total duration — a wedged engine is + indistinguishable from a slow one without the idle watchdog. +- **Retry and fallback.** Another endpoint is tried only while nothing has + reached the client; once bytes are out, splicing two responses together would + be worse than the short stream. Models can name fallback models, and the + response says which one served it via `x-llmgw-served-model`. +- **OpenAI surface.** `/v1/chat/completions`, `/v1/completions`, + `/v1/embeddings` and `/v1/models`, with aliases resolved from the catalog. + Only routing-relevant fields are parsed; the rest of the body passes through + untouched. +- **Usage accounting that does not depend on the client.** The router can ask + engines for usage the client never requested, then drop that chunk on the way + back, so token counts are recorded either way. +- **Configuration** (`internal/config`): YAML through koanf, overlaid on + defaults, with `LLMGW_`-prefixed environment overrides. Validation rejects + duplicate endpoint ids and aliases that would shadow another model — both + would otherwise misroute silently at request time. +- **Metrics** (plan §8): request counts, duration, TTFT, tokens, plus per-scorer + values and per-filter drop counts. Routing profile weights cannot be tuned + without seeing what the scorers actually produced. Served on a separate admin + listener alongside health and pprof. + +### Phase 1 exit criterion, measured in `test/e2e` + +- 1000 concurrent streams behind 3 sim-engines: all completed, ~6800 req/s. +- Every client disconnecting frees the KV blocks on every engine. +- Router overhead p99 ≈ 0.9 ms against the plan's 2 ms budget. CI asserts a + looser bound by default since a shared machine cannot hold 2 ms reliably; set + `LLMGW_E2E_STRICT=1` on quiet hardware to enforce the real target. + +### Known limitations + +- No admission control, so a batch workload can still starve interactive + traffic. That is v0.3.0, and it is the most valuable thing still missing. +- The endpoint list is static; the control plane arrives in v0.5.0. +- No authentication, rate limiting or budgets yet (v1.0.0). +- Prompt length is estimated as bytes/4. It becomes exact only with a real + tokenizer, which the plan defers to v2 along with tier 2 prefix routing. + ## v0.1.0 — 2026-09-16 First tagged version. It is useful for exactly one thing: developing and testing diff --git a/README.md b/README.md index b233102..a06f11c 100644 --- a/README.md +++ b/README.md @@ -13,15 +13,15 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md). | Version | Phase | Scope | State | |---|---|---|---| | v0.1.0 | 0 | Skeleton and sim-engine | released | -| v0.2.0 | 1 | Static router (filter/score/pick, streaming proxy) | in progress | +| v0.2.0 | 1 | Static router (filter/score/pick, streaming proxy) | released | | v0.3.0 | 1.5 | Admission control (priority bands) | planned | | 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 | -**v0.1.0 does not serve traffic**: `llmgw server` is still a stub. What works -today is `sim-engine` plus the routing libraries it exists to test. +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. ## Building @@ -31,6 +31,31 @@ make test # unit tests make race # with the race detector ``` +## Running the router + +```sh +# Two fake engines and a router in front of them. +./bin/llmgw sim-engine --listen :8001 --model llama & +./bin/llmgw sim-engine --listen :8002 --model llama & +./bin/llmgw server \ + --endpoint e1=http://127.0.0.1:8001=llama \ + --endpoint e2=http://127.0.0.1:8002=llama + +curl -N localhost:8080/v1/chat/completions \ + -H 'content-type: application/json' \ + -H 'x-session-id: demo' \ + -d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}' +``` + +For anything real use a config file; [deploy/examples/server.yaml](deploy/examples/server.yaml) +is commented throughout. Metrics, health and pprof are on the admin port (`:9090` +by default), deliberately separate from the client-facing one. + +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 +`x-llmgw-profile` header. + ## sim-engine Simulates the inference-engine behaviours **the router cares about**, so the diff --git a/deploy/examples/server.yaml b/deploy/examples/server.yaml new file mode 100644 index 0000000..702f7fd --- /dev/null +++ b/deploy/examples/server.yaml @@ -0,0 +1,79 @@ +# llmgw router configuration (phase 1). +# +# The endpoint list is static here. From phase 3 it comes from the control +# plane instead and this section disappears; everything else stays. +# +# Any field can be overridden from the environment with an LLMGW_ prefix and +# __ for nesting: LLMGW_ROUTER__RETRIES=3 + +listen: ":8080" # client-facing OpenAI API +admin_listen: ":9090" # metrics, health, pprof — do not expose this publicly + +router: + scrape_interval: 2s + scrape_timeout: 1s + # How long an endpoint may go unscraped before it is scored down. It is never + # dropped for staleness alone: if the whole pool went stale there would be + # nowhere left to send anything. + stale_after: 10s + + # Additional endpoints to try after a failure, and only while nothing has + # reached the client yet. + retries: 2 + # How long a failing endpoint leaves the candidate pool. A long fallback list + # with no cooldown turns one endpoint's rate limit into a cluster-wide retry + # storm. + cooldown: 15s + # Consecutive failures that open the breaker. Set to 1 for pure cooldown + # behaviour. + failure_threshold: 3 + + default_profile: chat # chat | batch | long-context | debug + session_header: x-session-id # affinity key; falls back to the `user` field + # Ask engines for usage even when the client did not, so token accounting + # works regardless of client behaviour. The extra chunk is dropped on the way + # back to a client that did not ask for it. + inject_usage: true + max_body_bytes: 33554432 + + timeouts: + dial: 5s + response_header: 60s + first_token: 120s + idle: 60s # between streamed tokens; catches a wedged engine + total: 30m + + # 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 + # delay. + scoring: + target_waiting: 4 + target_inflight: 8 + kv_ceiling: 0.9 + session_load_factor: 1.25 + +models: + - name: llama-3.3-70b + aliases: [llama] + fallbacks: [qwen2.5-32b] + routing_profile: chat + + - name: qwen2.5-32b + aliases: [qwen] + routing_profile: batch + +endpoints: + - id: gpu-a-01 + url: http://gpu-a-01:8000 + models: [llama-3.3-70b] + pool: ib-a/h100-80g + weight: 100 # from ctl bench; do not estimate it by hand + priority: 100 # higher pools are tried first + + - id: gpu-b-01 + url: http://gpu-b-01:8000 + models: [llama-3.3-70b, qwen2.5-32b] + pool: ib-b/a100-80g + weight: 45 + priority: 50 diff --git a/go.mod b/go.mod index ed854e7..5104a54 100644 --- a/go.mod +++ b/go.mod @@ -3,15 +3,30 @@ module gitea.mynodes.xyz/ilker/llmgw go 1.26.3 require ( + github.com/cespare/xxhash/v2 v2.3.0 + github.com/knadh/koanf/parsers/yaml v1.1.1 + github.com/knadh/koanf/providers/env/v2 v2.0.1 + github.com/knadh/koanf/providers/file v1.2.1 + github.com/knadh/koanf/v2 v2.3.6 + github.com/prometheus/client_golang v1.24.1 + github.com/prometheus/client_model v0.6.3 github.com/prometheus/common v0.71.0 github.com/spf13/cobra v1.10.2 ) require ( - github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect + github.com/fsnotify/fsnotify v1.9.0 // indirect + github.com/go-viper/mapstructure/v2 v2.4.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/knadh/koanf/maps v0.1.2 // indirect + github.com/mitchellh/copystructure v1.2.0 // indirect + github.com/mitchellh/reflectwalk v1.0.2 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect - github.com/prometheus/client_model v0.6.3 // indirect + github.com/prometheus/procfs v0.21.1 // indirect github.com/spf13/pflag v1.0.9 // indirect + go.yaml.in/yaml/v3 v3.0.5 // indirect + golang.org/x/sync v0.23.0 // indirect + golang.org/x/sys v0.47.0 // indirect google.golang.org/protobuf v1.36.12 // indirect ) diff --git a/go.sum b/go.sum index 82ef0dd..daec55d 100644 --- a/go.sum +++ b/go.sum @@ -1,18 +1,44 @@ +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= +github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k= +github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0= +github.com/go-viper/mapstructure/v2 v2.4.0 h1:EBsztssimR/CONLSZZ04E8qAkxNYq4Qp9LvH92wZUgs= +github.com/go-viper/mapstructure/v2 v2.4.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk= +github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo= +github.com/knadh/koanf/maps v0.1.2/go.mod h1:npD/QZY3V6ghQDdcQzl1W4ICNVTkohC8E73eI2xW4yI= +github.com/knadh/koanf/parsers/yaml v1.1.1 h1:u70vV5IyaM0HvONh8HoqBC97oTgO33KcpZbTLiKVinU= +github.com/knadh/koanf/parsers/yaml v1.1.1/go.mod h1:HHmcHXUrp9cOPcuC+2wrr44GTUB0EC+PyfN3HZD9tFg= +github.com/knadh/koanf/providers/env/v2 v2.0.1 h1:a3KagndPqhcWHQv6Pz4OZmwkI/yMeTjkiZye6ZCkyW0= +github.com/knadh/koanf/providers/env/v2 v2.0.1/go.mod h1:1g01PE+Ve1gBfWNNw2wmULRP0tc8RJrjn5p2N/jNCIc= +github.com/knadh/koanf/providers/file v1.2.1 h1:bEWbtQwYrA+W2DtdBrQWyXqJaJSG3KrP3AESOJYp9wM= +github.com/knadh/koanf/providers/file v1.2.1/go.mod h1:bp1PM5f83Q+TOUu10J/0ApLBd9uIzg+n9UgthfY+nRA= +github.com/knadh/koanf/v2 v2.3.6 h1:JoQPSJmvS4aP0xNc8xMDr5tcrkSEInL23/Il7pITAKo= +github.com/knadh/koanf/v2 v2.3.6/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw= +github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s= +github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ= +github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= -github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= -github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE= +github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU= +github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE= github.com/prometheus/client_model v0.6.3 h1:O0jaTVAYNxTHYInEPFJt5I3+sN8zqBtVMPTB1qyxiEo= github.com/prometheus/client_model v0.6.3/go.mod h1:gpN5P9S7Rr6Yr92PiQ+Ixvhf6JZEkF1dnxsYL2aPBEM= github.com/prometheus/common v0.71.0 h1:9KDAKb7Mj3HEVKyFCK6Dc/HIwlBzZIN2l7/lrHl3KK8= github.com/prometheus/common v0.71.0/go.mod h1:CLJ5H8TEsGX8bl31BdMkfhIZ+QmZ9tBPPotUxUbfcmk= +github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI= +github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= @@ -20,11 +46,17 @@ github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= +golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk= +golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc= google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/cli/server.go b/internal/cli/server.go new file mode 100644 index 0000000..626dff8 --- /dev/null +++ b/internal/cli/server.go @@ -0,0 +1,157 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "os/signal" + "strings" + "syscall" + "time" + + "github.com/spf13/cobra" + "golang.org/x/sync/errgroup" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/router" +) + +func newServerCmd() *cobra.Command { + var ( + configPath string + listen string + adminListen string + endpoints []string + ) + + cmd := &cobra.Command{ + Use: "server", + Short: "Run the router (phase 1; the control plane arrives in phase 3)", + Long: "Serves the OpenAI API in front of a set of engine replicas.\n\n" + + "In this phase the replica list is static: give it in the config file, or\n" + + "pass --endpoint id=URL=model[,model] once per replica for a quick test.", + RunE: func(cmd *cobra.Command, _ []string) error { + log := logger() + + cfg, err := config.Load(configPath) + if err != nil { + return err + } + if listen != "" { + cfg.Listen = listen + } + if adminListen != "" { + cfg.AdminListen = adminListen + } + for _, spec := range endpoints { + ep, err := parseEndpointFlag(spec) + if err != nil { + return err + } + cfg.Endpoints = append(cfg.Endpoints, ep) + } + if err := cfg.Validate(); err != nil { + return err + } + if len(cfg.Endpoints) == 0 { + return errors.New("no endpoints configured; use --endpoint or a config file") + } + + rt, err := router.New(cfg, log) + if err != nil { + return err + } + + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + rt.Start(ctx) + defer rt.Stop() + + adminHandler, err := rt.AdminHandler() + if err != nil { + return err + } + + apiSrv := &http.Server{ + Handler: rt.Handler(), + ReadHeaderTimeout: 10 * time.Second, + // No write timeout: responses stream for minutes and the proxy + // enforces its own per-phase deadlines. + } + adminSrv := &http.Server{ + Handler: adminHandler, + ReadHeaderTimeout: 10 * time.Second, + } + + apiLn, err := net.Listen("tcp", cfg.Listen) + if err != nil { + return fmt.Errorf("listening on %s: %w", cfg.Listen, err) + } + adminLn, err := net.Listen("tcp", cfg.AdminListen) + if err != nil { + _ = apiLn.Close() + return fmt.Errorf("listening on %s: %w", cfg.AdminListen, err) + } + + log.Info("router listening", + "api", apiLn.Addr().String(), + "admin", adminLn.Addr().String(), + "endpoints", len(cfg.Endpoints), + "default_profile", cfg.Router.DefaultProfile, + ) + + g, gctx := errgroup.WithContext(ctx) + g.Go(func() error { return serve(apiSrv, apiLn) }) + g.Go(func() error { return serve(adminSrv, adminLn) }) + g.Go(func() error { + <-gctx.Done() + log.Info("shutting down") + sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + _ = adminSrv.Shutdown(sctx) + return apiSrv.Shutdown(sctx) + }) + return g.Wait() + }, + } + + f := cmd.Flags() + f.StringVar(&configPath, "config", "", "configuration file (optional)") + f.StringVar(&listen, "listen", "", "override the client-facing address") + f.StringVar(&adminListen, "admin-listen", "", "override the metrics/pprof address") + f.StringArrayVar(&endpoints, "endpoint", nil, "static endpoint: id=URL=model[,model] (repeatable)") + + return cmd +} + +func serve(s *http.Server, ln net.Listener) error { + if err := s.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { + return err + } + return nil +} + +// parseEndpointFlag decodes id=URL=model[,model]. +// +// The URL contains '=' only in pathological cases, so splitting into exactly +// three parts keeps the flag readable while staying unambiguous. +func parseEndpointFlag(spec string) (config.Endpoint, error) { + parts := strings.SplitN(spec, "=", 3) + if len(parts) != 3 { + return config.Endpoint{}, fmt.Errorf("--endpoint %q: expected id=URL=model[,model]", spec) + } + models := strings.Split(parts[2], ",") + for i := range models { + models[i] = strings.TrimSpace(models[i]) + } + return config.Endpoint{ + ID: strings.TrimSpace(parts[0]), + URL: strings.TrimRight(strings.TrimSpace(parts[1]), "/"), + Models: models, + Pool: "static", + Weight: 1, + }, nil +} diff --git a/internal/cli/stubs.go b/internal/cli/stubs.go index cfe6caf..54784b3 100644 --- a/internal/cli/stubs.go +++ b/internal/cli/stubs.go @@ -9,19 +9,6 @@ import ( // errNotImplemented marks commands belonging to phases not yet built. var errNotImplemented = errors.New("not implemented yet (see llmgw-plan.md §11)") -func newServerCmd() *cobra.Command { - var configPath string - cmd := &cobra.Command{ - Use: "server", - Short: "Router and control plane (phases 1-3)", - RunE: func(_ *cobra.Command, _ []string) error { - return errNotImplemented - }, - } - cmd.Flags().StringVar(&configPath, "config", "server.yaml", "configuration file") - return cmd -} - func newAgentCmd() *cobra.Command { var configPath string cmd := &cobra.Command{ diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 0000000..600c41f --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,218 @@ +// Package config loads llmgw's YAML configuration. +// +// Phase 1 only needs the router's half of it: the endpoint list is static here +// and will come from the control plane in phase 3 (plan §6.4). The shape is +// already the one the plan describes, so that later change stays additive. +package config + +import ( + "fmt" + "os" + "strings" + "time" + + "github.com/knadh/koanf/parsers/yaml" + env "github.com/knadh/koanf/providers/env/v2" + "github.com/knadh/koanf/providers/file" + "github.com/knadh/koanf/v2" +) + +// Server is the whole server configuration. +type Server struct { + // Listen is the client-facing OpenAI API address. + Listen string `koanf:"listen"` + // AdminListen serves metrics, health and pprof. Kept separate so the + // client-facing port can be exposed without exposing diagnostics. + AdminListen string `koanf:"admin_listen"` + + Router Router `koanf:"router"` + Models []Model `koanf:"models"` + Endpoints []Endpoint `koanf:"endpoints"` +} + +// Router holds routing behaviour. +type Router struct { + ScrapeInterval time.Duration `koanf:"scrape_interval"` + ScrapeTimeout time.Duration `koanf:"scrape_timeout"` + // StaleAfter is how long an endpoint may go unscraped before being scored + // down (plan §7.6). + StaleAfter time.Duration `koanf:"stale_after"` + + // Retries is how many *additional* endpoints may be tried after a failure, + // and only while nothing has been written to the client (plan §7.7). + Retries int `koanf:"retries"` + // Cooldown is how long a failing endpoint leaves the candidate pool. + Cooldown time.Duration `koanf:"cooldown"` + // FailureThreshold is the consecutive failure count that opens a breaker. + // 1 turns the breaker into a pure cooldown. + FailureThreshold int `koanf:"failure_threshold"` + + // DefaultProfile names the routing profile used when a model picks none. + DefaultProfile string `koanf:"default_profile"` + // SessionHeader carries the session id for affinity (plan §7.5 tier 0). + SessionHeader string `koanf:"session_header"` + // InjectUsage asks engines for usage even when the client did not, so token + // accounting works regardless of client behaviour. + InjectUsage bool `koanf:"inject_usage"` + + // MaxBodyBytes bounds the request body the router will read. + MaxBodyBytes int64 `koanf:"max_body_bytes"` + + Timeouts Timeouts `koanf:"timeouts"` + Scoring Scoring `koanf:"scoring"` +} + +// Timeouts bounds each phase of an upstream request. +type Timeouts struct { + Dial time.Duration `koanf:"dial"` + ResponseHeader time.Duration `koanf:"response_header"` + FirstToken time.Duration `koanf:"first_token"` + Idle time.Duration `koanf:"idle"` + Total time.Duration `koanf:"total"` +} + +// Scoring holds the thresholds the scorers normalize against. These want to be +// measured with ctl bench (plan §10), not guessed. +type Scoring struct { + TargetWaiting float64 `koanf:"target_waiting"` + TargetInflight float64 `koanf:"target_inflight"` + KVCeiling float64 `koanf:"kv_ceiling"` + SessionLoadFactor float64 `koanf:"session_load_factor"` +} + +// Model is a catalog entry as clients see it. +type Model struct { + Name string `koanf:"name"` + Aliases []string `koanf:"aliases"` + // Fallbacks are tried, in order, when this model has no usable endpoint. + Fallbacks []string `koanf:"fallbacks"` + // RoutingProfile overrides the router's default profile for this model. + RoutingProfile string `koanf:"routing_profile"` +} + +// Endpoint is one engine replica. Static in phase 1. +type Endpoint struct { + ID string `koanf:"id"` + URL string `koanf:"url"` + Models []string `koanf:"models"` + Pool string `koanf:"pool"` + Weight float64 `koanf:"weight"` + Priority int `koanf:"priority"` +} + +// Default returns the configuration used when nothing is specified. +func Default() Server { + return Server{ + Listen: ":8080", + AdminListen: ":9090", + Router: Router{ + ScrapeInterval: 2 * time.Second, + ScrapeTimeout: time.Second, + StaleAfter: 10 * time.Second, + Retries: 2, + Cooldown: 15 * time.Second, + FailureThreshold: 3, + DefaultProfile: "chat", + SessionHeader: "x-session-id", + InjectUsage: true, + MaxBodyBytes: 32 << 20, + Timeouts: Timeouts{ + Dial: 5 * time.Second, + ResponseHeader: 60 * time.Second, + FirstToken: 120 * time.Second, + Idle: 60 * time.Second, + Total: 30 * time.Minute, + }, + Scoring: Scoring{ + TargetWaiting: 4, + TargetInflight: 8, + KVCeiling: 0.9, + SessionLoadFactor: 1.25, + }, + }, + } +} + +// Load reads configuration from path, overlaying it on the defaults, then +// applies LLMGW_-prefixed environment overrides. +// +// An empty path loads the defaults, which is what makes `llmgw server` with a +// couple of --endpoint flags usable for a smoke test. +func Load(path string) (Server, error) { + cfg := Default() + k := koanf.New(".") + + if path != "" { + if _, err := os.Stat(path); err != nil { + return cfg, fmt.Errorf("config file: %w", err) + } + if err := k.Load(file.Provider(path), yaml.Parser()); err != nil { + return cfg, fmt.Errorf("parsing %s: %w", path, err) + } + } + + // LLMGW_ROUTER__RETRIES=3 maps onto router.retries. + err := k.Load(env.Provider(".", env.Opt{ + Prefix: "LLMGW_", + TransformFunc: func(key, value string) (string, any) { + key = strings.ToLower(strings.TrimPrefix(key, "LLMGW_")) + return strings.ReplaceAll(key, "__", "."), value + }, + }), nil) + if err != nil { + return cfg, fmt.Errorf("environment overrides: %w", err) + } + + if err := k.UnmarshalWithConf("", &cfg, koanf.UnmarshalConf{Tag: "koanf"}); err != nil { + return cfg, fmt.Errorf("decoding configuration: %w", err) + } + return cfg, cfg.Validate() +} + +// Validate rejects configurations that would otherwise fail confusingly at +// request time. +func (s Server) Validate() error { + if s.Listen == "" { + return fmt.Errorf("listen is empty") + } + + seen := make(map[string]struct{}, len(s.Endpoints)) + for i, e := range s.Endpoints { + switch { + case e.ID == "": + return fmt.Errorf("endpoints[%d]: id is required", i) + case e.URL == "": + return fmt.Errorf("endpoint %q: url is required", e.ID) + case !strings.HasPrefix(e.URL, "http://") && !strings.HasPrefix(e.URL, "https://"): + return fmt.Errorf("endpoint %q: url must start with http:// or https://", e.ID) + case len(e.Models) == 0: + return fmt.Errorf("endpoint %q: at least one model is required", e.ID) + } + if _, dup := seen[e.ID]; dup { + return fmt.Errorf("endpoint id %q is used twice", e.ID) + } + seen[e.ID] = struct{}{} + } + + names := make(map[string]struct{}, len(s.Models)) + for i, m := range s.Models { + if m.Name == "" { + return fmt.Errorf("models[%d]: name is required", i) + } + if _, dup := names[m.Name]; dup { + return fmt.Errorf("model %q is declared twice", m.Name) + } + names[m.Name] = struct{}{} + } + // Aliases share a namespace with names, so a collision must be caught here + // rather than silently shadowing a model at request time. + for _, m := range s.Models { + for _, a := range m.Aliases { + if _, clash := names[a]; clash { + return fmt.Errorf("model %q: alias %q collides with another model name", m.Name, a) + } + names[a] = struct{}{} + } + } + return nil +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go new file mode 100644 index 0000000..3a3c2b2 --- /dev/null +++ b/internal/config/config_test.go @@ -0,0 +1,150 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func write(t *testing.T, content string) string { + t.Helper() + p := filepath.Join(t.TempDir(), "server.yaml") + if err := os.WriteFile(p, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + return p +} + +func TestLoadDefaultsWithoutFile(t *testing.T) { + cfg, err := Load("") + if err != nil { + t.Fatal(err) + } + if cfg.Listen != ":8080" || cfg.Router.DefaultProfile != "chat" { + t.Errorf("defaults not applied: %+v", cfg) + } +} + +func TestLoadOverlaysOnDefaults(t *testing.T) { + // Specifying one field must not blank out its neighbours. + p := write(t, ` +listen: ":9999" +endpoints: + - id: a + url: http://127.0.0.1:8000 + models: [m] +`) + cfg, err := Load(p) + if err != nil { + t.Fatal(err) + } + if cfg.Listen != ":9999" { + t.Errorf("listen=%q", cfg.Listen) + } + if cfg.Router.ScrapeInterval != 2*time.Second { + t.Errorf("an unspecified default was lost: %v", cfg.Router.ScrapeInterval) + } + if len(cfg.Endpoints) != 1 || cfg.Endpoints[0].ID != "a" { + t.Errorf("endpoints=%+v", cfg.Endpoints) + } +} + +func TestDurationsParse(t *testing.T) { + p := write(t, ` +router: + scrape_interval: 500ms + timeouts: + idle: 45s +`) + cfg, err := Load(p) + if err != nil { + t.Fatal(err) + } + if cfg.Router.ScrapeInterval != 500*time.Millisecond { + t.Errorf("scrape_interval=%v", cfg.Router.ScrapeInterval) + } + if cfg.Router.Timeouts.Idle != 45*time.Second { + t.Errorf("timeouts.idle=%v", cfg.Router.Timeouts.Idle) + } +} + +func TestEnvOverride(t *testing.T) { + t.Setenv("LLMGW_ROUTER__RETRIES", "7") + cfg, err := Load("") + if err != nil { + t.Fatal(err) + } + if cfg.Router.Retries != 7 { + t.Errorf("retries=%d, expected the environment to win", cfg.Router.Retries) + } +} + +func TestMissingFileIsAnError(t *testing.T) { + if _, err := Load("/nonexistent/server.yaml"); err == nil { + t.Error("expected an error for a missing config file") + } +} + +func TestValidateRejectsBadEndpoints(t *testing.T) { + cases := []struct { + name string + ep Endpoint + want string + }{ + {"no id", Endpoint{URL: "http://x", Models: []string{"m"}}, "id is required"}, + {"no url", Endpoint{ID: "a", Models: []string{"m"}}, "url is required"}, + {"bad scheme", Endpoint{ID: "a", URL: "127.0.0.1:8000", Models: []string{"m"}}, "http://"}, + {"no models", Endpoint{ID: "a", URL: "http://x"}, "at least one model"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + s := Default() + s.Endpoints = []Endpoint{c.ep} + err := s.Validate() + if err == nil { + t.Fatal("expected an error") + } + if !strings.Contains(err.Error(), c.want) { + t.Errorf("error %q does not mention %q", err, c.want) + } + }) + } +} + +func TestValidateRejectsDuplicateEndpointID(t *testing.T) { + s := Default() + s.Endpoints = []Endpoint{ + {ID: "a", URL: "http://x", Models: []string{"m"}}, + {ID: "a", URL: "http://y", Models: []string{"m"}}, + } + if err := s.Validate(); err == nil || !strings.Contains(err.Error(), "twice") { + t.Errorf("duplicate id not rejected: %v", err) + } +} + +func TestValidateRejectsAliasCollision(t *testing.T) { + // An alias shadowing another model would silently misroute every request + // for that model, so it has to fail at load time. + s := Default() + s.Models = []Model{ + {Name: "llama-3.3-70b", Aliases: []string{"qwen"}}, + {Name: "qwen"}, + } + err := s.Validate() + if err == nil || !strings.Contains(err.Error(), "collides") { + t.Errorf("alias collision not rejected: %v", err) + } +} + +func TestValidateAcceptsDistinctAliases(t *testing.T) { + s := Default() + s.Models = []Model{ + {Name: "llama-3.3-70b", Aliases: []string{"llama", "llama3"}}, + {Name: "qwen2.5-32b", Aliases: []string{"qwen"}}, + } + if err := s.Validate(); err != nil { + t.Errorf("valid aliases rejected: %v", err) + } +} diff --git a/internal/router/admin.go b/internal/router/admin.go new file mode 100644 index 0000000..3213f9f --- /dev/null +++ b/internal/router/admin.go @@ -0,0 +1,106 @@ +package router + +import ( + "net/http" + "net/http/pprof" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/collectors" + "github.com/prometheus/client_golang/prometheus/promhttp" +) + +// AdminHandler returns the diagnostics surface: metrics, health and pprof. +// +// It is served on a separate listener so the client-facing port can be exposed +// without also exposing pprof (plan §9.1). +func (r *Router) AdminHandler() (http.Handler, error) { + reg := prometheus.NewRegistry() + reg.MustRegister( + collectors.NewGoCollector(), + collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}), + ) + for _, c := range r.m.collectors() { + if err := reg.Register(c); err != nil { + return nil, err + } + } + if err := reg.Register(newEndpointCollector(r)); err != nil { + return nil, err + } + + mux := http.NewServeMux() + mux.Handle("GET /metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{})) + mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + mux.HandleFunc("GET /debug/pprof/", pprof.Index) + mux.HandleFunc("GET /debug/pprof/cmdline", pprof.Cmdline) + mux.HandleFunc("GET /debug/pprof/profile", pprof.Profile) + mux.HandleFunc("GET /debug/pprof/symbol", pprof.Symbol) + mux.HandleFunc("GET /debug/pprof/trace", pprof.Trace) + return mux, nil +} + +// endpointCollector reports live endpoint state at scrape time rather than +// mirroring it into gauges on every request; the registry already holds it. +type endpointCollector struct { + r *Router + + inflight *prometheus.Desc + waiting *prometheus.Desc + running *prometheus.Desc + kv *prometheus.Desc + stale *prometheus.Desc + drain *prometheus.Desc + breaker *prometheus.Desc +} + +func newEndpointCollector(r *Router) *endpointCollector { + lbl := []string{"replica", "pool"} + return &endpointCollector{ + r: r, + inflight: prometheus.NewDesc("llmgw_endpoint_inflight", "Requests the router has dispatched and not yet finished.", lbl, nil), + waiting: prometheus.NewDesc("llmgw_endpoint_waiting", "Engine queue depth, as last scraped.", lbl, nil), + running: prometheus.NewDesc("llmgw_endpoint_running", "Requests running on the engine, as last scraped.", lbl, nil), + kv: prometheus.NewDesc("llmgw_endpoint_kv_cache_usage", "Engine KV cache occupancy, 0..1.", lbl, nil), + 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), + } +} + +func (c *endpointCollector) Describe(ch chan<- *prometheus.Desc) { + ch <- c.inflight + ch <- c.waiting + ch <- c.running + ch <- c.kv + ch <- c.drain + ch <- c.stale + ch <- c.breaker +} + +func (c *endpointCollector) Collect(ch chan<- prometheus.Metric) { + for _, s := range c.r.reg.Snapshots("") { + pool := s.Pool + if pool == "" { + pool = "none" + } + g := func(d *prometheus.Desc, v float64) { + ch <- prometheus.MustNewConstMetric(d, prometheus.GaugeValue, v, s.ID, pool) + } + g(c.inflight, float64(s.Inflight)) + g(c.waiting, float64(s.Stats.Waiting)) + g(c.running, float64(s.Stats.Running)) + g(c.kv, s.Stats.KVCacheUsage) + g(c.drain, s.DrainRate) + g(c.stale, boolToFloat(s.Stale)) + g(c.breaker, boolToFloat(s.Breaker.String() == "open")) + } +} + +func boolToFloat(b bool) float64 { + if b { + return 1 + } + return 0 +} diff --git a/internal/router/handlers.go b/internal/router/handlers.go new file mode 100644 index 0000000..97bd636 --- /dev/null +++ b/internal/router/handlers.go @@ -0,0 +1,284 @@ +package router + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "net/http" + "slices" + "strings" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" +) + +// 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. +const prefixBytes = 2048 + +// Handler returns the client-facing OpenAI API. +func (r *Router) Handler() http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("POST /v1/chat/completions", r.handleInference("/v1/chat/completions")) + mux.HandleFunc("POST /v1/completions", r.handleInference("/v1/completions")) + mux.HandleFunc("POST /v1/embeddings", r.handleInference("/v1/embeddings")) + mux.HandleFunc("GET /v1/models", r.handleModels) + mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + return mux +} + +// parsedRequest is what the router needs from a request. Everything else in the +// body is forwarded untouched — clients send engine-specific fields we have no +// business interpreting. +type parsedRequest struct { + Path string + Model string + + Stream bool + SuppressUsageChunk bool + UpstreamBody []byte + Header http.Header + + SessionID string + ProfileOverride string + PromptPrefix string + PromptTokens int + MaxTokens int +} + +// inferenceBody covers the fields the router reads across all inference +// endpoints. Unknown fields are ignored here and preserved in UpstreamBody. +type inferenceBody struct { + Model string `json:"model"` + Stream bool `json:"stream"` + User string `json:"user"` + Messages []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } `json:"messages"` + Prompt json.RawMessage `json:"prompt"` + Input json.RawMessage `json:"input"` + MaxTokens int `json:"max_tokens"` + MaxCompletionTokens int `json:"max_completion_tokens"` + StreamOptions *struct { + IncludeUsage bool `json:"include_usage"` + } `json:"stream_options"` +} + +func (r *Router) handleInference(path string) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + start := time.Now() + + parsed, err := r.parse(req, path) + if err != nil { + writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + r.m.observeRequest("", "", http.StatusBadRequest, time.Since(start), nil, 0) + return + } + + res := r.dispatch(req.Context(), w, parsed) + r.finish(w, parsed, res, start) + } +} + +// finish records metrics and, when nothing was written, produces the error body. +func (r *Router) finish(w http.ResponseWriter, req *parsedRequest, res dispatchResult, start time.Time) { + status := res.outcome.Status + switch { + case res.err == nil: + if res.fallback { + // The client asked for one model and got another; say so rather + // than letting them discover it from the response contents. + w.Header().Set("x-llmgw-served-model", res.model) + } + case errors.Is(res.err, proxy.ErrClientGone): + // Nothing to report to a client that is gone; the upstream was cancelled. + status = 499 + r.log.Debug("client disconnected", "model", req.Model, "endpoint", res.endpointID) + case res.outcome.Wrote: + // The response is already partly out; the stream just ends short. + r.log.Warn("stream ended early", "model", req.Model, "endpoint", res.endpointID, "err", res.err) + if status == 0 { + status = http.StatusOK + } + case errors.Is(res.err, errNoCapacity): + status = http.StatusServiceUnavailable + writeError(w, status, "server_error", res.err.Error()) + default: + status = http.StatusBadGateway + writeError(w, status, "server_error", res.err.Error()) + } + + r.m.observeRequest(req.Model, res.pool, status, time.Since(start), res.outcome.Usage, res.outcome.TTFT) +} + +// parse reads the body and extracts the routing-relevant fields. +func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) { + body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes)) + if err != nil { + return nil, err + } + + var b inferenceBody + if err := json.Unmarshal(body, &b); err != nil { + return nil, err + } + if b.Model == "" { + return nil, errors.New("model is required") + } + + p := &parsedRequest{ + Path: path, + Model: b.Model, + Stream: b.Stream, + UpstreamBody: body, + Header: req.Header, + SessionID: sessionID(req, r.cfg.SessionHeader, b.User), + ProfileOverride: req.Header.Get("x-llmgw-profile"), + MaxTokens: max(b.MaxTokens, b.MaxCompletionTokens), + PromptPrefix: promptPrefix(&b), + } + // Bytes over four: a crude stand-in for a tokenizer, good enough for + // scoring. It becomes exact only with plan §7.5 tier 2 (see plan §14). + p.PromptTokens = len(p.PromptPrefix) / 4 + + // Ask the engine for usage the client did not request, so accounting works + // regardless of client behaviour; the extra chunk is dropped on the way back. + if r.cfg.InjectUsage && b.Stream && (b.StreamOptions == nil || !b.StreamOptions.IncludeUsage) { + if injected, ok := injectIncludeUsage(body); ok { + p.UpstreamBody = injected + p.SuppressUsageChunk = true + } + } + return p, nil +} + +// sessionID picks the affinity key: an explicit header first, then OpenAI's +// user field, which most clients already set per end user. +func sessionID(req *http.Request, header, user string) string { + if header != "" { + if v := req.Header.Get(header); v != "" { + return v + } + } + return user +} + +// promptPrefix reconstructs the leading portion of the prompt for prefix-aware +// routing. Only the first prefixBytes matter, so long conversations stop early. +func promptPrefix(b *inferenceBody) string { + var sb strings.Builder + appendRaw := func(raw json.RawMessage) { + if len(raw) == 0 || sb.Len() >= prefixBytes { + return + } + var s string + if err := json.Unmarshal(raw, &s); err != nil { + s = string(raw) + } + sb.WriteString(s) + } + + for _, m := range b.Messages { + if sb.Len() >= prefixBytes { + break + } + sb.WriteString(m.Role) + sb.WriteString(":") + appendRaw(m.Content) + sb.WriteString("\n") + } + appendRaw(b.Prompt) + appendRaw(b.Input) + + s := sb.String() + if len(s) > prefixBytes { + s = s[:prefixBytes] + } + return s +} + +// injectIncludeUsage sets stream_options.include_usage on the outgoing body. +// +// This is the one place the router rewrites a request, so it is done by +// decoding to a generic map and re-encoding: field order changes, but no field +// is lost. If anything looks unexpected it declines rather than guessing. +func injectIncludeUsage(body []byte) ([]byte, bool) { + var m map[string]json.RawMessage + if err := json.Unmarshal(body, &m); err != nil { + return nil, false + } + + opts := map[string]any{} + if raw, ok := m["stream_options"]; ok { + if err := json.Unmarshal(raw, &opts); err != nil { + return nil, false + } + } + opts["include_usage"] = true + + encoded, err := json.Marshal(opts) + if err != nil { + return nil, false + } + m["stream_options"] = encoded + + out, err := json.Marshal(m) + if err != nil { + return nil, false + } + return out, true +} + +// handleModels lists the models that at least one endpoint currently serves. +func (r *Router) handleModels(w http.ResponseWriter, _ *http.Request) { + seen := map[string]struct{}{} + for _, s := range r.reg.Snapshots("") { + for _, m := range s.Models { + seen[m] = struct{}{} + } + } + // Catalog aliases are usable names too, so clients should see them. + for name := range r.models { + if _, ok := seen[r.models[name].Name]; ok { + seen[name] = struct{}{} + } + } + + names := make([]string, 0, len(seen)) + for n := range seen { + names = append(names, n) + } + slices.Sort(names) + + data := make([]map[string]any, 0, len(names)) + for _, n := range names { + data = append(data, map[string]any{ + "id": n, "object": "model", "created": time.Now().Unix(), "owned_by": "llmgw", + }) + } + writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": data}) +} + +func writeJSON(w http.ResponseWriter, code int, v any) { + var buf bytes.Buffer + if err := json.NewEncoder(&buf).Encode(v); err != nil { + http.Error(w, `{"error":{"message":"encoding failed"}}`, http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + _, _ = w.Write(buf.Bytes()) +} + +// writeError emits an OpenAI-shaped error so clients need no special handling +// for gateway failures. +func writeError(w http.ResponseWriter, code int, kind, msg string) { + writeJSON(w, code, map[string]any{"error": map[string]any{ + "message": msg, "type": kind, "code": nil, + }}) +} diff --git a/internal/router/metrics.go b/internal/router/metrics.go new file mode 100644 index 0000000..d07e266 --- /dev/null +++ b/internal/router/metrics.go @@ -0,0 +1,108 @@ +package router + +import ( + "strconv" + "time" + + "github.com/prometheus/client_golang/prometheus" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler" +) + +// metrics holds the router's own Prometheus series (plan §8). +// +// The scorer series matter more than they look: routing profile weights cannot +// be tuned without seeing what each scorer actually produced in production. +type metrics struct { + requests *prometheus.CounterVec + duration *prometheus.HistogramVec + ttft *prometheus.HistogramVec + tokens *prometheus.CounterVec + + scorerValue *prometheus.GaugeVec + filtered *prometheus.CounterVec + considered prometheus.Histogram +} + +func newMetrics() *metrics { + return &metrics{ + requests: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_requests_total", + Help: "Requests handled, by model, pool and response status.", + }, []string{"model", "pool", "status"}), + + duration: prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Name: "llmgw_request_duration_seconds", + Help: "End-to-end request duration.", + Buckets: []float64{0.01, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 120, 300}, + }, []string{"model", "pool"}), + + ttft: prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Name: "llmgw_ttft_seconds", + Help: "Time to first token, as observed by the router.", + Buckets: []float64{0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 30, 60}, + }, []string{"model", "pool"}), + + tokens: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_tokens_total", + Help: "Tokens accounted for, by direction.", + }, []string{"model", "direction"}), + + scorerValue: prometheus.NewGaugeVec(prometheus.GaugeOpts{ + Name: "llmgw_scorer_value", + Help: "Last score each scorer produced for the endpoint that was chosen.", + }, []string{"scorer", "replica"}), + + filtered: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_candidates_filtered_total", + Help: "Candidate endpoints removed, by filter.", + }, []string{"filter"}), + + considered: prometheus.NewHistogram(prometheus.HistogramOpts{ + Name: "llmgw_candidates_considered", + Help: "Candidates remaining after filtering, per scheduling decision.", + Buckets: []float64{1, 2, 3, 5, 8, 16, 32, 64}, + }), + } +} + +// collectors returns everything to register, so registration lives in one place. +func (m *metrics) collectors() []prometheus.Collector { + return []prometheus.Collector{ + m.requests, m.duration, m.ttft, m.tokens, + m.scorerValue, m.filtered, m.considered, + } +} + +func (m *metrics) observeRequest(model, pool string, status int, d time.Duration, u *proxy.Usage, ttft time.Duration) { + if model == "" { + model = "unknown" + } + if pool == "" { + pool = "none" + } + m.requests.WithLabelValues(model, pool, strconv.Itoa(status)).Inc() + m.duration.WithLabelValues(model, pool).Observe(d.Seconds()) + + if ttft > 0 { + m.ttft.WithLabelValues(model, pool).Observe(ttft.Seconds()) + } + if u != nil { + m.tokens.WithLabelValues(model, "prompt").Add(float64(u.PromptTokens)) + m.tokens.WithLabelValues(model, "completion").Add(float64(u.CompletionTokens)) + if u.CachedTokens > 0 { + m.tokens.WithLabelValues(model, "cached").Add(float64(u.CachedTokens)) + } + } +} + +func (m *metrics) observeScores(d scheduler.Decision) { + for name, v := range d.Scores { + m.scorerValue.WithLabelValues(name, d.Endpoint.ID).Set(v) + } + for filter, n := range d.Dropped { + m.filtered.WithLabelValues(filter).Add(float64(n)) + } + m.considered.Observe(float64(d.Considered)) +} diff --git a/internal/router/proxy/proxy.go b/internal/router/proxy/proxy.go new file mode 100644 index 0000000..57a020c --- /dev/null +++ b/internal/router/proxy/proxy.go @@ -0,0 +1,468 @@ +// Package proxy streams a request to one upstream engine and relays the +// response to the client (plan §7.8). +// +// It handles exactly one attempt. Choosing the endpoint, retrying and tripping +// breakers belong to the caller, because those decisions need the scheduler and +// the endpoint registry. What the proxy owes the caller is an honest answer to +// "did anything reach the client yet?" — once the first byte is out, retrying +// elsewhere would corrupt the response. +package proxy + +import ( + "bufio" + "bytes" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "strings" + "sync" + "sync/atomic" + "time" +) + +// Config holds proxy settings. +type Config struct { + // DialTimeout bounds establishing the TCP connection. + DialTimeout time.Duration + // ResponseHeaderTimeout bounds the wait for upstream response headers. + ResponseHeaderTimeout time.Duration + // FirstTokenTimeout bounds the wait for the first streamed event. + FirstTokenTimeout time.Duration + // IdleTimeout bounds the gap between streamed events. A stalled engine + // looks identical to a slow one over a single connection, so without this + // a wedged request holds a client forever. + IdleTimeout time.Duration + // TotalTimeout bounds the whole request. Generous by default: a long + // generation is legitimately slow. + TotalTimeout time.Duration + // MaxIdleConnsPerHost keeps connections warm; engines are few and busy. + MaxIdleConnsPerHost int + // MaxResponseBytes bounds a buffered (non-streaming) response. + MaxResponseBytes int64 +} + +// DefaultConfig returns sensible defaults. +func DefaultConfig() Config { + return Config{ + DialTimeout: 5 * time.Second, + ResponseHeaderTimeout: 60 * time.Second, + FirstTokenTimeout: 120 * time.Second, + IdleTimeout: 60 * time.Second, + TotalTimeout: 30 * time.Minute, + MaxIdleConnsPerHost: 64, + MaxResponseBytes: 64 << 20, + } +} + +func (c *Config) normalize() { + d := DefaultConfig() + if c.DialTimeout <= 0 { + c.DialTimeout = d.DialTimeout + } + if c.ResponseHeaderTimeout <= 0 { + c.ResponseHeaderTimeout = d.ResponseHeaderTimeout + } + if c.FirstTokenTimeout <= 0 { + c.FirstTokenTimeout = d.FirstTokenTimeout + } + if c.IdleTimeout <= 0 { + c.IdleTimeout = d.IdleTimeout + } + if c.TotalTimeout <= 0 { + c.TotalTimeout = d.TotalTimeout + } + if c.MaxIdleConnsPerHost <= 0 { + c.MaxIdleConnsPerHost = d.MaxIdleConnsPerHost + } + if c.MaxResponseBytes <= 0 { + c.MaxResponseBytes = d.MaxResponseBytes + } +} + +// Proxy forwards requests to engine endpoints. +type Proxy struct { + cfg Config + client *http.Client +} + +// New builds a proxy. +func New(cfg Config) *Proxy { + cfg.normalize() + return &Proxy{ + cfg: cfg, + client: &http.Client{ + // No client timeout: streaming responses run for minutes and the + // per-phase watchdogs below bound them more precisely. + Transport: &http.Transport{ + DialContext: (&net.Dialer{Timeout: cfg.DialTimeout}).DialContext, + ResponseHeaderTimeout: cfg.ResponseHeaderTimeout, + MaxIdleConnsPerHost: cfg.MaxIdleConnsPerHost, + MaxIdleConns: cfg.MaxIdleConnsPerHost * 8, + IdleConnTimeout: 90 * time.Second, + // Engines rarely speak h2c, and HTTP/2 over cleartext would + // need explicit setup anyway. + ForceAttemptHTTP2: false, + DisableCompression: true, + }, + }, + } +} + +// Request describes one upstream attempt. +type Request struct { + // URL is the fully qualified upstream URL. + URL string + // Body is the request body to send. + Body []byte + // Header carries headers to forward (authorization is deliberately not). + Header http.Header + // Stream says whether an SSE response is expected. + Stream bool + // SuppressUsageChunk drops the final usage-only chunk from the client while + // still recording it. Set when the router asked for usage the client did + // not request. + SuppressUsageChunk bool +} + +// Usage holds the token counts reported by the engine. +type Usage struct { + PromptTokens int + CompletionTokens int + TotalTokens int + CachedTokens int +} + +// Outcome reports what happened during one attempt. +type Outcome struct { + // Wrote is true once any byte reached the client. Retrying is then unsafe. + Wrote bool + // Status is the upstream status code, 0 if the request never completed. + Status int + // Retryable means the attempt failed in a way another endpoint might survive. + Retryable bool + // TTFT is the time until the first streamed event, zero if none arrived. + TTFT time.Duration + // Usage is the engine-reported token usage, if it was present. + Usage *Usage + // Bytes is how many bytes were relayed to the client. + Bytes int64 +} + +var ( + // ErrClientGone means the client disconnected; the upstream was cancelled. + ErrClientGone = errors.New("client disconnected") + // ErrUpstreamStalled means no event arrived within the idle or first-token timeout. + ErrUpstreamStalled = errors.New("upstream stalled") +) + +// Do performs one attempt against one endpoint. +// +// On a retryable failure it writes nothing to w, leaving the caller free to try +// another endpoint. Once it has written anything, Outcome.Wrote is true and the +// caller must not retry. +func (p *Proxy) Do(ctx context.Context, w http.ResponseWriter, r Request) (Outcome, error) { + ctx, cancel := context.WithTimeout(ctx, p.cfg.TotalTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.URL, bytes.NewReader(r.Body)) + if err != nil { + return Outcome{Retryable: false}, err + } + req.Header = cloneForUpstream(r.Header) + req.Header.Set("Content-Type", "application/json") + req.ContentLength = int64(len(r.Body)) + + resp, err := p.client.Do(req) + if err != nil { + // The client going away is not the endpoint's fault; never blame it. + if ctx.Err() != nil && ctxCauseIsClient(ctx) { + return Outcome{Retryable: false}, ErrClientGone + } + return Outcome{Retryable: true}, fmt.Errorf("upstream request failed: %w", err) + } + defer resp.Body.Close() + + if retryableStatus(resp.StatusCode) { + // Drain a little so the connection can be reused, then let the caller + // pick somewhere else. Nothing has been written to the client. + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10)) + return Outcome{Status: resp.StatusCode, Retryable: true}, + fmt.Errorf("upstream returned %s", resp.Status) + } + + if r.Stream && isEventStream(resp.Header) { + return p.relayStream(ctx, cancel, w, resp, r) + } + return p.relayBuffered(w, resp) +} + +// relayBuffered copies a non-streaming response through, capturing usage. +func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response) (Outcome, error) { + body, err := io.ReadAll(io.LimitReader(resp.Body, p.cfg.MaxResponseBytes)) + if err != nil { + return Outcome{Status: resp.StatusCode, Retryable: true}, err + } + + copyHeaders(w.Header(), resp.Header) + w.WriteHeader(resp.StatusCode) + n, werr := w.Write(body) + + out := Outcome{ + Wrote: n > 0, + Status: resp.StatusCode, + Bytes: int64(n), + Usage: usageFromJSON(body), + } + if werr != nil { + return out, ErrClientGone + } + return out, nil +} + +// relayStream relays an SSE response event by event, flushing as it goes. +// +// abort must cancel the context the upstream request was built with: the idle +// watchdog can only stop a silent engine by aborting that read. +func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w http.ResponseWriter, resp *http.Response, r Request) (Outcome, error) { + flusher, ok := w.(http.Flusher) + if !ok { + return Outcome{Status: resp.StatusCode, Retryable: false}, + errors.New("response writer does not support streaming") + } + + copyHeaders(w.Header(), resp.Header) + w.WriteHeader(resp.StatusCode) + flusher.Flush() + + out := Outcome{Wrote: true, Status: resp.StatusCode} + start := time.Now() + + // One watchdog covers both phases: until the first event it enforces + // FirstTokenTimeout, afterwards IdleTimeout. + wd := newWatchdog(p.cfg.FirstTokenTimeout) + defer wd.run(ctx, abort)() + + br := bufio.NewReaderSize(resp.Body, 16<<10) + var ev eventBuffer + + for { + line, err := br.ReadString('\n') + if len(line) > 0 { + wd.beat() + + ev.add(line) + if ev.complete() { + if out.TTFT == 0 { + out.TTFT = time.Since(start) + wd.setInterval(p.cfg.IdleTimeout) + } + if u := ev.usage(); u != nil { + out.Usage = u + if r.SuppressUsageChunk { + // The client never asked for usage; record it and drop + // the chunk rather than changing what they receive. + ev.reset() + continue + } + } + n, werr := io.WriteString(w, ev.string()) + out.Bytes += int64(n) + ev.reset() + if werr != nil { + return out, ErrClientGone + } + flusher.Flush() + } + } + + if err != nil { + if errors.Is(err, io.EOF) { + // Flush a trailing event that never got its blank line. + if s := ev.string(); s != "" { + n, _ := io.WriteString(w, s) + out.Bytes += int64(n) + flusher.Flush() + } + return out, nil + } + if wd.fired() { + return out, ErrUpstreamStalled + } + if ctx.Err() != nil { + return out, ErrClientGone + } + return out, err + } + } +} + +// --- SSE event assembly --- + +// eventBuffer accumulates the lines of one SSE event. Events are relayed whole +// rather than line by line so a usage-only chunk can be dropped atomically. +type eventBuffer struct { + lines []string + data strings.Builder +} + +func (e *eventBuffer) add(line string) { + e.lines = append(e.lines, line) + if payload, ok := strings.CutPrefix(line, "data: "); ok { + e.data.WriteString(strings.TrimRight(payload, "\r\n")) + } +} + +// complete reports whether the blank line terminating an event has arrived. +func (e *eventBuffer) complete() bool { + if len(e.lines) == 0 { + return false + } + last := e.lines[len(e.lines)-1] + return last == "\n" || last == "\r\n" +} + +func (e *eventBuffer) string() string { + if len(e.lines) == 0 { + return "" + } + return strings.Join(e.lines, "") +} + +func (e *eventBuffer) reset() { + e.lines = e.lines[:0] + e.data.Reset() +} + +// usage returns the token usage carried by this event, if any. +func (e *eventBuffer) usage() *Usage { + s := e.data.String() + if s == "" || s == "[DONE]" || !strings.Contains(s, "\"usage\"") { + return nil + } + return usageFromJSON([]byte(s)) +} + +// --- watchdog --- + +// watchdog aborts the upstream read once an interval passes with no beat, and +// records that it did so, so the resulting read error can be attributed to the +// engine stalling rather than to the client going away. +type watchdog struct { + interval atomic.Int64 + last atomic.Int64 + tripped atomic.Bool +} + +func newWatchdog(interval time.Duration) *watchdog { + w := &watchdog{} + w.interval.Store(int64(interval)) + w.last.Store(time.Now().UnixNano()) + return w +} + +func (w *watchdog) beat() { w.last.Store(time.Now().UnixNano()) } + +func (w *watchdog) setInterval(d time.Duration) { w.interval.Store(int64(d)) } + +func (w *watchdog) fired() bool { return w.tripped.Load() } + +// run starts the watchdog and returns a stop function. abort is called if an +// interval elapses without a beat. +func (w *watchdog) run(ctx context.Context, abort context.CancelFunc) func() { + done := make(chan struct{}) + + go func() { + // Tick well below the shortest interval so the overshoot stays small. + // The interval shrinks once the first token arrives, so derive the tick + // from the smaller of the two rather than from whatever is set now. + tick := max(time.Duration(w.interval.Load())/4, 10*time.Millisecond) + t := time.NewTicker(tick) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-done: + return + case <-t.C: + interval := time.Duration(w.interval.Load()) + if time.Since(time.Unix(0, w.last.Load())) > interval { + w.tripped.Store(true) + abort() + return + } + // Keep the tick proportional as the interval changes. + if want := max(interval/4, 10*time.Millisecond); want != tick { + tick = want + t.Reset(tick) + } + } + } + }() + + var once sync.Once + return func() { once.Do(func() { close(done) }) } +} + +// --- helpers --- + +func retryableStatus(code int) bool { + switch code { + case http.StatusTooManyRequests, + http.StatusBadGateway, + http.StatusServiceUnavailable, + http.StatusGatewayTimeout: + return true + default: + return false + } +} + +func isEventStream(h http.Header) bool { + return strings.HasPrefix(h.Get("Content-Type"), "text/event-stream") +} + +// hopByHop headers must not be forwarded (RFC 9110 §7.6.1). +var hopByHop = map[string]bool{ + "Connection": true, + "Proxy-Connection": true, + "Keep-Alive": true, + "Proxy-Authenticate": true, + "Proxy-Authorization": true, + "Te": true, + "Trailer": true, + "Transfer-Encoding": true, + "Upgrade": true, +} + +func copyHeaders(dst, src http.Header) { + for k, vs := range src { + if hopByHop[http.CanonicalHeaderKey(k)] { + continue + } + for _, v := range vs { + dst.Add(k, v) + } + } +} + +// cloneForUpstream forwards only what the engine needs. The client's +// Authorization header is deliberately dropped: it authenticates the caller to +// llmgw, not llmgw to the engine, and engines are reached over a trusted +// network (plan §14). +func cloneForUpstream(h http.Header) http.Header { + out := make(http.Header, 4) + for _, k := range []string{"Accept", "User-Agent", "Traceparent", "Tracestate"} { + if v := h.Get(k); v != "" { + out.Set(k, v) + } + } + return out +} + +func ctxCauseIsClient(ctx context.Context) bool { + return errors.Is(context.Cause(ctx), context.Canceled) +} diff --git a/internal/router/proxy/proxy_test.go b/internal/router/proxy/proxy_test.go new file mode 100644 index 0000000..9861257 --- /dev/null +++ b/internal/router/proxy/proxy_test.go @@ -0,0 +1,302 @@ +package proxy + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +func testProxy(mut ...func(*Config)) *Proxy { + c := DefaultConfig() + c.FirstTokenTimeout = 2 * time.Second + c.IdleTimeout = 500 * time.Millisecond + c.TotalTimeout = 10 * time.Second + for _, f := range mut { + f(&c) + } + return New(c) +} + +// sseServer serves a scripted SSE stream. +func sseServer(t *testing.T, chunks []string, gap time.Duration) (*httptest.Server, *atomic.Bool) { + t.Helper() + var cancelled atomic.Bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + f := w.(http.Flusher) + f.Flush() + for _, c := range chunks { + select { + case <-r.Context().Done(): + cancelled.Store(true) + return + case <-time.After(gap): + } + if _, err := io.WriteString(w, c); err != nil { + cancelled.Store(true) + return + } + f.Flush() + } + })) + t.Cleanup(srv.Close) + return srv, &cancelled +} + +func dataChunk(s string) string { return "data: " + s + "\n\n" } + +func TestBufferedResponseIsRelayed(t *testing.T) { + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"id":"x","usage":{"prompt_tokens":7,"completion_tokens":3}}`) + })) + defer up.Close() + + rec := httptest.NewRecorder() + out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`)}) + if err != nil { + t.Fatal(err) + } + if out.Status != 200 || !out.Wrote { + t.Errorf("outcome=%+v", out) + } + if out.Usage == nil || out.Usage.PromptTokens != 7 || out.Usage.TotalTokens != 10 { + t.Errorf("usage not captured: %+v", out.Usage) + } + if !strings.Contains(rec.Body.String(), `"id":"x"`) { + t.Errorf("body not relayed: %q", rec.Body.String()) + } +} + +func TestStreamIsRelayedEventByEvent(t *testing.T) { + up, _ := sseServer(t, []string{ + dataChunk(`{"choices":[{"delta":{"content":"a"}}]}`), + dataChunk(`{"choices":[{"delta":{"content":"b"}}]}`), + "data: [DONE]\n\n", + }, time.Millisecond) + + rec := httptest.NewRecorder() + out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`), Stream: true}) + if err != nil { + t.Fatal(err) + } + body := rec.Body.String() + if n := strings.Count(body, "data: "); n != 3 { + t.Errorf("expected 3 data lines, got %d: %q", n, body) + } + if !strings.HasSuffix(body, "[DONE]\n\n") { + t.Errorf("stream did not end with [DONE]: %q", body) + } + if out.TTFT == 0 { + t.Error("TTFT was not recorded") + } +} + +func TestUsageCapturedFromStream(t *testing.T) { + up, _ := sseServer(t, []string{ + dataChunk(`{"choices":[{"delta":{"content":"a"}}]}`), + dataChunk(`{"choices":[],"usage":{"prompt_tokens":11,"completion_tokens":5,"total_tokens":16}}`), + "data: [DONE]\n\n", + }, time.Millisecond) + + rec := httptest.NewRecorder() + out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`), Stream: true}) + if err != nil { + t.Fatal(err) + } + if out.Usage == nil || out.Usage.TotalTokens != 16 { + t.Fatalf("usage not captured from stream: %+v", out.Usage) + } + if !strings.Contains(rec.Body.String(), `"usage"`) { + t.Error("usage chunk should reach a client that asked for it") + } +} + +func TestUsageChunkSuppressed(t *testing.T) { + // The router asked the engine for usage the client never requested: record + // it, but do not change what the client receives. + up, _ := sseServer(t, []string{ + dataChunk(`{"choices":[{"delta":{"content":"a"}}]}`), + dataChunk(`{"choices":[],"usage":{"prompt_tokens":11,"completion_tokens":5,"total_tokens":16}}`), + "data: [DONE]\n\n", + }, time.Millisecond) + + rec := httptest.NewRecorder() + out, err := testProxy().Do(t.Context(), rec, Request{ + URL: up.URL, Body: []byte(`{}`), Stream: true, SuppressUsageChunk: true, + }) + if err != nil { + t.Fatal(err) + } + if out.Usage == nil || out.Usage.TotalTokens != 16 { + t.Fatalf("usage not recorded: %+v", out.Usage) + } + if strings.Contains(rec.Body.String(), `"usage"`) { + t.Errorf("usage chunk leaked to the client: %q", rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "[DONE]") { + t.Error("[DONE] was dropped along with the usage chunk") + } +} + +func TestRetryableStatusWritesNothing(t *testing.T) { + // The caller must be free to try another endpoint, which is only safe if + // the client has seen nothing. + for _, code := range []int{429, 502, 503, 504} { + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "busy", code) + })) + + rec := httptest.NewRecorder() + out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`)}) + up.Close() + + if err == nil { + t.Errorf("%d: expected an error", code) + } + if !out.Retryable { + t.Errorf("%d: not marked retryable", code) + } + if out.Wrote || rec.Body.Len() > 0 { + t.Errorf("%d: wrote %q to the client", code, rec.Body.String()) + } + } +} + +func TestNonRetryableStatusIsRelayed(t *testing.T) { + // A 400 is the client's fault; retrying elsewhere would produce the same + // answer, so relay it. + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + io.WriteString(w, `{"error":{"message":"bad"}}`) + })) + defer up.Close() + + rec := httptest.NewRecorder() + out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`)}) + if err != nil { + t.Fatal(err) + } + if out.Retryable { + t.Error("400 must not be retryable") + } + if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), "bad") { + t.Errorf("error response not relayed: %d %q", rec.Code, rec.Body.String()) + } +} + +func TestConnectionFailureIsRetryable(t *testing.T) { + rec := httptest.NewRecorder() + out, err := testProxy(func(c *Config) { c.DialTimeout = 200 * time.Millisecond }). + Do(t.Context(), rec, Request{URL: "http://127.0.0.1:1/x", Body: []byte(`{}`)}) + if err == nil { + t.Fatal("expected an error") + } + if !out.Retryable || out.Wrote { + t.Errorf("outcome=%+v", out) + } +} + +func TestClientDisconnectCancelsUpstream(t *testing.T) { + // Phase 1's exit criterion at the proxy level: the GPU must stop working + // when the client goes away. + chunks := make([]string, 200) + for i := range chunks { + chunks[i] = dataChunk(fmt.Sprintf(`{"choices":[{"delta":{"content":"%d"}}]}`, i)) + } + up, cancelled := sseServer(t, chunks, 10*time.Millisecond) + + ctx, cancel := context.WithCancel(context.Background()) + rec := httptest.NewRecorder() + + done := make(chan Outcome, 1) + go func() { + out, _ := testProxy().Do(ctx, rec, Request{URL: up.URL, Body: []byte(`{}`), Stream: true}) + done <- out + }() + + time.Sleep(80 * time.Millisecond) // let a few events through + cancel() + + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("proxy did not return after client cancellation") + } + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if cancelled.Load() { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("upstream kept generating after the client disconnected") +} + +func TestStalledUpstreamHitsIdleTimeout(t *testing.T) { + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + f := w.(http.Flusher) + io.WriteString(w, dataChunk(`{"choices":[{"delta":{"content":"a"}}]}`)) + f.Flush() + <-r.Context().Done() // then go silent forever + })) + defer up.Close() + + rec := httptest.NewRecorder() + p := testProxy(func(c *Config) { c.IdleTimeout = 200 * time.Millisecond }) + + start := time.Now() + out, err := p.Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`), Stream: true}) + elapsed := time.Since(start) + + if !errors.Is(err, ErrUpstreamStalled) { + t.Fatalf("expected ErrUpstreamStalled, got %v", err) + } + if !out.Wrote { + t.Error("the events delivered before the stall should be reported as written") + } + if elapsed > 3*time.Second { + t.Errorf("idle timeout took too long: %v", elapsed) + } +} + +func TestHopByHopHeadersNotForwarded(t *testing.T) { + var got http.Header + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.Header.Clone() + io.WriteString(w, `{}`) + })) + defer up.Close() + + h := http.Header{} + h.Set("Authorization", "Bearer sk-llmgw-secret") + h.Set("Accept", "text/event-stream") + h.Set("Traceparent", "00-trace-span-01") + + rec := httptest.NewRecorder() + if _, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`), Header: h}); err != nil { + t.Fatal(err) + } + + if got.Get("Authorization") != "" { + t.Error("the client's Authorization header was forwarded upstream") + } + if got.Get("Accept") != "text/event-stream" { + t.Error("Accept was not forwarded") + } + if got.Get("Traceparent") != "00-trace-span-01" { + t.Error("Traceparent was not forwarded; tracing would break") + } +} diff --git a/internal/router/proxy/usage.go b/internal/router/proxy/usage.go new file mode 100644 index 0000000..fc49ea9 --- /dev/null +++ b/internal/router/proxy/usage.go @@ -0,0 +1,39 @@ +package proxy + +import "encoding/json" + +// usageEnvelope mirrors just the usage block of an OpenAI response. +// Everything else in the payload is passed through untouched, so only these +// fields are ever decoded. +type usageEnvelope struct { + Usage *struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + PromptTokensDetails *struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + +// usageFromJSON extracts token usage from a response or an SSE chunk. +// It returns nil when the payload carries no usage block, which is the normal +// case for every chunk but the last. +func usageFromJSON(b []byte) *Usage { + var env usageEnvelope + if err := json.Unmarshal(b, &env); err != nil || env.Usage == nil { + return nil + } + u := &Usage{ + PromptTokens: env.Usage.PromptTokens, + CompletionTokens: env.Usage.CompletionTokens, + TotalTokens: env.Usage.TotalTokens, + } + if d := env.Usage.PromptTokensDetails; d != nil { + u.CachedTokens = d.CachedTokens + } + if u.TotalTokens == 0 { + u.TotalTokens = u.PromptTokens + u.CompletionTokens + } + return u +} diff --git a/internal/router/router.go b/internal/router/router.go new file mode 100644 index 0000000..25bf7dd --- /dev/null +++ b/internal/router/router.go @@ -0,0 +1,270 @@ +// Package router wires the pieces into a working OpenAI-compatible gateway: +// resolve the model, pick an endpoint, proxy the stream, retry when it is safe. +// +// The sub-packages hold the reusable parts (scheduler, endpoint, scrape, proxy, +// breaker); this package holds the orchestration and the HTTP surface. +package router + +import ( + "context" + "errors" + "fmt" + "log/slog" + "net/http" + "slices" + "strings" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape" +) + +// Router serves the OpenAI API in front of a set of engine replicas. +type Router struct { + cfg config.Router + log *slog.Logger + + reg *endpoint.Registry + poller *scrape.Poller + px *proxy.Proxy + m *metrics + + profiles map[string]*scheduler.Profile + // models maps every name a client may use — canonical names and aliases — + // onto its catalog entry. + models map[string]config.Model +} + +// New builds a router from configuration. Start must be called to begin scraping. +func New(cfg config.Server, log *slog.Logger) (*Router, error) { + if log == nil { + log = slog.Default() + } + + reg := endpoint.NewRegistry(endpoint.Config{ + StaleAfter: cfg.Router.StaleAfter, + Breaker: breaker.Config{ + FailureThreshold: cfg.Router.FailureThreshold, + BaseCooldown: cfg.Router.Cooldown, + }, + }) + reg.Replace(toInfos(cfg.Endpoints)) + + sc := scheduler.ProfileConfig{ + TargetWaiting: cfg.Router.Scoring.TargetWaiting, + TargetInflight: cfg.Router.Scoring.TargetInflight, + KVCeiling: cfg.Router.Scoring.KVCeiling, + SessionLoadFactor: cfg.Router.Scoring.SessionLoadFactor, + } + profiles := map[string]*scheduler.Profile{} + for _, name := range []string{"chat", "batch", "long-context", "debug"} { + profiles[name] = scheduler.Builtin(name, sc) + } + if _, ok := profiles[cfg.Router.DefaultProfile]; !ok { + return nil, fmt.Errorf("unknown default_profile %q", cfg.Router.DefaultProfile) + } + + models := map[string]config.Model{} + for _, m := range cfg.Models { + models[m.Name] = m + for _, a := range m.Aliases { + models[a] = m + } + } + + r := &Router{ + cfg: cfg.Router, + log: log, + reg: reg, + profiles: profiles, + models: models, + m: newMetrics(), + px: proxy.New(proxy.Config{ + DialTimeout: cfg.Router.Timeouts.Dial, + ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader, + FirstTokenTimeout: cfg.Router.Timeouts.FirstToken, + IdleTimeout: cfg.Router.Timeouts.Idle, + TotalTimeout: cfg.Router.Timeouts.Total, + MaxResponseBytes: cfg.Router.MaxBodyBytes, + }), + } + r.poller = scrape.NewPoller(reg, scrape.Config{ + Interval: cfg.Router.ScrapeInterval, + Timeout: cfg.Router.ScrapeTimeout, + }, log) + + return r, nil +} + +func toInfos(eps []config.Endpoint) []endpoint.Info { + out := make([]endpoint.Info, 0, len(eps)) + for _, e := range eps { + out = append(out, endpoint.Info{ + ID: e.ID, + URL: strings.TrimRight(e.URL, "/"), + Models: e.Models, + Pool: e.Pool, + Weight: e.Weight, + Priority: e.Priority, + }) + } + return out +} + +// Start begins scraping endpoint metrics. It returns immediately. +func (r *Router) Start(ctx context.Context) { r.poller.Sync(ctx) } + +// Stop halts scraping. +func (r *Router) Stop() { r.poller.Stop() } + +// Registry exposes the endpoint registry, for tests and admin endpoints. +func (r *Router) Registry() *endpoint.Registry { return r.reg } + +// --- model resolution --- + +// resolve maps a client-supplied model name onto the catalog. Unknown names are +// passed through unchanged: an endpoint may serve a model the catalog does not +// list, and refusing it would be unhelpful. +func (r *Router) resolve(name string) (canonical string, m config.Model) { + if e, ok := r.models[name]; ok { + return e.Name, e + } + return name, config.Model{Name: name} +} + +func (r *Router) profileFor(m config.Model, override string) *scheduler.Profile { + for _, name := range []string{override, m.RoutingProfile, r.cfg.DefaultProfile} { + if p, ok := r.profiles[name]; ok && name != "" { + return p + } + } + return r.profiles["chat"] +} + +// --- dispatch --- + +// attempt describes one upstream try, for logging and metrics. +type dispatchResult struct { + endpointID string + pool string + servedModel + outcome proxy.Outcome + err error +} + +type servedModel struct { + model string + // fallback is true when the served model is not the one that was requested. + fallback bool +} + +// dispatch resolves the model, then tries endpoints until one works, the client +// has been written to, or the attempt budget runs out. +// +// The retry rule (plan §7.7) is strict: retry only while nothing has reached the +// client. Once bytes are out, a retry would splice two different responses +// together, which is worse than the error the client would otherwise see. +func (r *Router) dispatch(ctx context.Context, w http.ResponseWriter, req *parsedRequest) dispatchResult { + canonical, entry := r.resolve(req.Model) + + // The requested model first, then its fallbacks. A fallback is a different + // model, so the client is told about it via x-llmgw-served-model. + chain := append([]string{canonical}, entry.Fallbacks...) + + var last dispatchResult + for i, model := range chain { + res := r.dispatchOne(ctx, w, req, model, entry) + res.model = model + res.fallback = i > 0 + + if res.err == nil || res.outcome.Wrote || errors.Is(res.err, proxy.ErrClientGone) { + return res + } + last = res + // Only a total absence of capacity justifies moving to another model; + // a bad request would fail identically there. + if !errors.Is(res.err, errNoCapacity) { + return res + } + } + return last +} + +var errNoCapacity = errors.New("no endpoint available") + +func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *parsedRequest, model string, entry config.Model) dispatchResult { + candidates := r.reg.Snapshots(model) + if len(candidates) == 0 { + return dispatchResult{err: fmt.Errorf("%w: no endpoint serves %q", errNoCapacity, model)} + } + + profile := r.profileFor(entry, req.ProfileOverride) + sctx := &scheduler.Ctx{ + Now: time.Now(), + Req: scheduler.Request{ + Model: model, + SessionID: req.SessionID, + PromptPrefix: req.PromptPrefix, + PromptTokens: req.PromptTokens, + MaxTokens: req.MaxTokens, + }, + } + + attempts := max(r.cfg.Retries+1, 1) + var last dispatchResult + + for range attempts { + if len(candidates) == 0 { + break + } + dec, err := profile.Schedule(sctx, candidates) + if err != nil { + last = dispatchResult{err: fmt.Errorf("%w: %v", errNoCapacity, err)} + break + } + + release, granted := r.reg.Acquire(dec.Endpoint.ID) + if !granted { + // The breaker opened between scoring and dispatch; drop it and + // score again rather than failing the request. + candidates = without(candidates, dec.Endpoint.ID) + continue + } + + out, perr := r.px.Do(ctx, w, proxy.Request{ + URL: dec.Endpoint.URL + req.Path, + Body: req.UpstreamBody, + Header: req.Header, + Stream: req.Stream, + SuppressUsageChunk: req.SuppressUsageChunk, + }) + // A retryable failure is the endpoint's fault; a client disconnect is not. + release(perr == nil || errors.Is(perr, proxy.ErrClientGone)) + + res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool, outcome: out, err: perr} + r.m.observeScores(dec) + + if perr == nil || out.Wrote || !out.Retryable { + return res + } + last = res + r.log.Warn("upstream attempt failed, trying another endpoint", + "endpoint", dec.Endpoint.ID, "model", model, "err", perr) + candidates = without(candidates, dec.Endpoint.ID) + } + + if last.err == nil { + last.err = fmt.Errorf("%w: every endpoint for %q failed", errNoCapacity, model) + } else if !errors.Is(last.err, errNoCapacity) { + last.err = fmt.Errorf("%w: %v", errNoCapacity, last.err) + } + return last +} + +func without(eps []endpoint.Snapshot, id string) []endpoint.Snapshot { + return slices.DeleteFunc(slices.Clone(eps), func(e endpoint.Snapshot) bool { return e.ID == id }) +} diff --git a/internal/router/router_test.go b/internal/router/router_test.go new file mode 100644 index 0000000..7004b01 --- /dev/null +++ b/internal/router/router_test.go @@ -0,0 +1,315 @@ +package router + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" +) + +// fakeEngine serves the minimum an endpoint needs: the metric contract plus a +// scripted chat response. +type fakeEngine struct { + srv *httptest.Server + calls atomic.Int64 + // status, when non-zero, is returned instead of a completion. + status atomic.Int64 +} + +func newFakeEngine(t *testing.T, model string) *fakeEngine { + t.Helper() + f := &fakeEngine{} + mux := http.NewServeMux() + mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) { + fmt.Fprintf(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n") + fmt.Fprintf(w, "# TYPE vllm:num_requests_waiting gauge\nvllm:num_requests_waiting 0\n") + fmt.Fprintf(w, "# TYPE vllm:kv_cache_usage_perc gauge\nvllm:kv_cache_usage_perc 0.1\n") + }) + mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, _ *http.Request) { + f.calls.Add(1) + if code := int(f.status.Load()); code != 0 { + http.Error(w, "unavailable", code) + return + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"model":%q,"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":2}}`, model) + }) + f.srv = httptest.NewServer(mux) + t.Cleanup(f.srv.Close) + return f +} + +func newTestRouter(t *testing.T, cfg config.Server) *Router { + t.Helper() + rt, err := New(cfg, slog.New(slog.DiscardHandler)) + if err != nil { + t.Fatal(err) + } + rt.Start(t.Context()) + t.Cleanup(rt.Stop) + return rt +} + +func post(t *testing.T, h http.Handler, body any, hdr map[string]string) *httptest.ResponseRecorder { + t.Helper() + b, err := json.Marshal(body) + if err != nil { + t.Fatal(err) + } + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(b)) + req.Header.Set("Content-Type", "application/json") + for k, v := range hdr { + req.Header.Set(k, v) + } + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec +} + +func baseConfig(eps ...config.Endpoint) config.Server { + c := config.Default() + c.Router.ScrapeInterval = 50 * time.Millisecond + c.Router.StaleAfter = 5 * time.Second + c.Endpoints = eps + return c +} + +func TestRoutesToHealthyEndpoint(t *testing.T) { + e := newFakeEngine(t, "m") + rt := newTestRouter(t, baseConfig(config.Endpoint{ + ID: "a", URL: e.srv.URL, Models: []string{"m"}, + })) + + rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil) + if rec.Code != http.StatusOK { + t.Fatalf("status %d: %s", rec.Code, rec.Body.String()) + } + if e.calls.Load() != 1 { + t.Errorf("engine received %d calls", e.calls.Load()) + } +} + +func TestRetriesOnAnotherEndpoint(t *testing.T) { + bad := newFakeEngine(t, "m") + bad.status.Store(http.StatusServiceUnavailable) + good := newFakeEngine(t, "m") + + rt := newTestRouter(t, baseConfig( + config.Endpoint{ID: "bad", URL: bad.srv.URL, Models: []string{"m"}}, + config.Endpoint{ID: "good", URL: good.srv.URL, Models: []string{"m"}}, + )) + + // With two endpoints and two retries, the bad one must not sink the request. + rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil) + if rec.Code != http.StatusOK { + t.Fatalf("status %d: %s", rec.Code, rec.Body.String()) + } + if good.calls.Load() == 0 { + t.Error("the healthy endpoint was never tried") + } +} + +func TestNoRetryAfterBytesReachClient(t *testing.T) { + // An endpoint that starts streaming then dies must not be retried: splicing + // two responses together is worse than a short stream (plan §7.7). + var served atomic.Int64 + mux := http.NewServeMux() + mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) { + io.WriteString(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n"+ + "# TYPE vllm:num_requests_waiting gauge\nvllm:num_requests_waiting 0\n"+ + "# TYPE vllm:kv_cache_usage_perc gauge\nvllm:kv_cache_usage_perc 0\n") + }) + mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, _ *http.Request) { + served.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + f := w.(http.Flusher) + io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n") + f.Flush() + // Then hang up mid-stream. + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + rt := newTestRouter(t, baseConfig( + config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"m"}}, + config.Endpoint{ID: "b", URL: srv.URL, Models: []string{"m"}}, + )) + + rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}, "stream": true}, nil) + if !strings.Contains(rec.Body.String(), `"content":"a"`) { + t.Fatalf("partial stream not delivered: %q", rec.Body.String()) + } + if n := served.Load(); n != 1 { + t.Errorf("request was sent %d times; it must not be retried once streaming began", n) + } +} + +func TestFallbackModel(t *testing.T) { + primary := newFakeEngine(t, "primary") + primary.status.Store(http.StatusServiceUnavailable) + backup := newFakeEngine(t, "backup") + + cfg := baseConfig( + config.Endpoint{ID: "p", URL: primary.srv.URL, Models: []string{"primary"}}, + config.Endpoint{ID: "b", URL: backup.srv.URL, Models: []string{"backup"}}, + ) + cfg.Models = []config.Model{{Name: "primary", Fallbacks: []string{"backup"}}} + rt := newTestRouter(t, cfg) + + rec := post(t, rt.Handler(), map[string]any{"model": "primary", "messages": []any{}}, nil) + if rec.Code != http.StatusOK { + t.Fatalf("status %d: %s", rec.Code, rec.Body.String()) + } + if backup.calls.Load() == 0 { + t.Fatal("the fallback model was never tried") + } + // The client asked for one model and got another; they must be told. + if got := rec.Header().Get("x-llmgw-served-model"); got != "backup" { + t.Errorf("x-llmgw-served-model=%q, expected \"backup\"", got) + } +} + +func TestAliasResolution(t *testing.T) { + e := newFakeEngine(t, "llama-3.3-70b") + cfg := baseConfig(config.Endpoint{ + ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"}, + }) + cfg.Models = []config.Model{{Name: "llama-3.3-70b", Aliases: []string{"llama"}}} + rt := newTestRouter(t, cfg) + + rec := post(t, rt.Handler(), map[string]any{"model": "llama", "messages": []any{}}, nil) + if rec.Code != http.StatusOK { + t.Fatalf("alias was not resolved: %d %s", rec.Code, rec.Body.String()) + } +} + +func TestUnknownModelIs503(t *testing.T) { + e := newFakeEngine(t, "m") + rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})) + + rec := post(t, rt.Handler(), map[string]any{"model": "nope", "messages": []any{}}, nil) + if rec.Code != http.StatusServiceUnavailable { + t.Errorf("status %d, expected 503", rec.Code) + } + var body map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("error body is not JSON: %q", rec.Body.String()) + } + if _, ok := body["error"]; !ok { + t.Errorf("error body is not OpenAI-shaped: %v", body) + } +} + +func TestMissingModelIs400(t *testing.T) { + e := newFakeEngine(t, "m") + rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})) + + rec := post(t, rt.Handler(), map[string]any{"messages": []any{}}, nil) + if rec.Code != http.StatusBadRequest { + t.Errorf("status %d, expected 400", rec.Code) + } +} + +func TestModelsEndpointListsAliases(t *testing.T) { + e := newFakeEngine(t, "llama-3.3-70b") + cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"}}) + cfg.Models = []config.Model{{Name: "llama-3.3-70b", Aliases: []string{"llama"}}} + rt := newTestRouter(t, cfg) + + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + rec := httptest.NewRecorder() + rt.Handler().ServeHTTP(rec, req) + + var out struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatal(err) + } + ids := map[string]bool{} + for _, d := range out.Data { + ids[d.ID] = true + } + if !ids["llama-3.3-70b"] || !ids["llama"] { + t.Errorf("both the model and its alias should be listed: %v", ids) + } +} + +func TestInjectIncludeUsagePreservesFields(t *testing.T) { + // The one place the router rewrites a request body; nothing may be lost. + in := []byte(`{"model":"m","stream":true,"temperature":0.7,"tools":[{"x":1}]}`) + out, ok := injectIncludeUsage(in) + if !ok { + t.Fatal("injection declined a well-formed body") + } + var m map[string]any + if err := json.Unmarshal(out, &m); err != nil { + t.Fatal(err) + } + for _, k := range []string{"model", "stream", "temperature", "tools"} { + if _, present := m[k]; !present { + t.Errorf("field %q was dropped", k) + } + } + opts, _ := m["stream_options"].(map[string]any) + if opts["include_usage"] != true { + t.Errorf("include_usage not set: %v", m["stream_options"]) + } +} + +func TestInjectIncludeUsageKeepsExistingOptions(t *testing.T) { + in := []byte(`{"model":"m","stream":true,"stream_options":{"something":"else"}}`) + out, ok := injectIncludeUsage(in) + if !ok { + t.Fatal("injection declined") + } + var m map[string]any + if err := json.Unmarshal(out, &m); err != nil { + t.Fatal(err) + } + opts, ok := m["stream_options"].(map[string]any) + if !ok { + t.Fatalf("stream_options is not an object: %v", m["stream_options"]) + } + if opts["something"] != "else" { + t.Error("an existing stream_options field was dropped") + } + if opts["include_usage"] != true { + t.Error("include_usage not set") + } +} + +func TestPromptPrefixIsBounded(t *testing.T) { + long := strings.Repeat("x", prefixBytes*3) + b := inferenceBody{Messages: []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + }{{Role: "user", Content: json.RawMessage(`"` + long + `"`)}}} + + if got := len(promptPrefix(&b)); got > prefixBytes { + t.Errorf("prefix is %d bytes, over the %d cap", got, prefixBytes) + } +} + +func TestSessionHeaderFallsBackToUser(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/", nil) + if got := sessionID(req, "x-session-id", "user-42"); got != "user-42" { + t.Errorf("expected the user field to be used, got %q", got) + } + req.Header.Set("x-session-id", "sess-1") + if got := sessionID(req, "x-session-id", "user-42"); got != "sess-1" { + t.Errorf("the header should win, got %q", got) + } +} diff --git a/llmgw-plan.md b/llmgw-plan.md index 1d5251a..0ae8afb 100644 --- a/llmgw-plan.md +++ b/llmgw-plan.md @@ -144,8 +144,7 @@ llmgw/ │ │ ├── scheduler/ # placement algorithm │ │ ├── reconciler/ # desired vs actual loop │ │ └── agentsvc/ # gRPC server -│ ├── router/ -│ │ ├── openai/ # HTTP handlers, request parsing +│ ├── router/ # HTTP handlers, dispatch, metrics │ │ ├── proxy/ # streaming proxy to upstream │ │ ├── admission/ # saturation detector, band queues (§7.2) │ │ ├── scheduler/ # Filter/Scorer/Picker pipeline (§7.3) @@ -166,6 +165,10 @@ llmgw/ └── docs/ ``` +> The OpenAI handlers live in `internal/router` itself rather than an `openai/` +> subpackage: they need the dispatch loop, and splitting them apart would have +> meant an interface whose only implementation is the router. + ### 4.1 Libraries | Need | Choice | @@ -953,8 +956,8 @@ work_dir: /var/lib/llmgw-agent - [x] CI: `go test ./...`, race detector ### Phase 1 — Static router (2 weeks) -- [ ] OpenAI handlers, streaming proxy, cancellation propagation -- [ ] Static endpoint list from configuration +- [x] OpenAI handlers, streaming proxy, cancellation propagation +- [x] Static endpoint list from configuration - [x] **The `Filter`/`Scorer`/`Picker` pipeline (§7.3)** — a plugin pipeline from the start, not one "strategy" - [x] `queue`, `kv`, `inflight` scorers; `max-score` and `p2c` pickers - [x] Session affinity (§7.5 tier 0) — cheap, an early win diff --git a/test/e2e/harness_test.go b/test/e2e/harness_test.go new file mode 100644 index 0000000..e61a7d3 --- /dev/null +++ b/test/e2e/harness_test.go @@ -0,0 +1,198 @@ +// Package e2e runs the router against real sim-engines over real HTTP. +// +// These tests exist to check the phase 1 exit criterion from the plan, so they +// deliberately use the whole stack: no fakes, no in-process shortcuts beyond +// httptest listeners. +package e2e + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/router" + "gitea.mynodes.xyz/ilker/llmgw/internal/sim" +) + +const testModel = "test-model" + +type engine struct { + eng *sim.Engine + srv *httptest.Server +} + +type harness struct { + engines []engine + router *httptest.Server + rt *router.Router +} + +// simConfig is fast enough for tests while still exercising queueing. +func simConfig(name string, seed int64) sim.Config { + c := sim.Default() + c.Model = testModel + c.Seed = seed + c.Tick = time.Millisecond + c.MaxNumSeqs = 512 + c.NumGPUBlocks = 200000 + c.PrefillTokensPerSec = 5e6 + c.DecodeTokensPerSec = 2000 + c.BatchContention = 0 + _ = name + return c +} + +// start brings up n sim-engines behind a router. +func start(t *testing.T, n int, tune ...func(*config.Server)) *harness { + t.Helper() + + h := &harness{} + cfg := config.Default() + cfg.Router.ScrapeInterval = 200 * time.Millisecond + cfg.Router.ScrapeTimeout = 150 * time.Millisecond + cfg.Router.StaleAfter = 2 * time.Second + + for i := range n { + e := sim.New(simConfig(fmt.Sprintf("sim-%d", i), int64(i+1))) + 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{ + ID: fmt.Sprintf("e%d", i), + URL: srv.URL, + Models: []string{testModel}, + Pool: "test", + Weight: 1, + }) + } + for _, f := range tune { + f(&cfg) + } + + rt, err := router.New(cfg, slog.New(slog.DiscardHandler)) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + rt.Start(ctx) + + h.rt = rt + h.router = httptest.NewServer(rt.Handler()) + + t.Cleanup(func() { + h.router.Close() + cancel() + rt.Stop() + for _, e := range h.engines { + e.srv.Close() + e.eng.Close() + } + }) + + // Let the first scrape land so scheduling starts from real data. + h.waitForFreshMetrics(t) + return h +} + +func (h *harness) waitForFreshMetrics(t *testing.T) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + fresh := 0 + for _, s := range h.rt.Registry().Snapshots(testModel) { + if !s.Stale { + fresh++ + } + } + if fresh == len(h.engines) { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatal("endpoint metrics never became fresh") +} + +func (h *harness) totalRunning() int { + n := 0 + for _, e := range h.engines { + s := e.eng.Stats() + n += s.Running + s.Waiting + } + return n +} + +// chatBody builds a streaming chat request. +func chatBody(prompt string, maxTokens int, stream bool) []byte { + b, _ := json.Marshal(map[string]any{ + "model": testModel, + "messages": []map[string]string{{"role": "user", "content": prompt}}, + "max_tokens": maxTokens, + "stream": stream, + }) + return b +} + +// streamRequest sends a streaming request and counts the content chunks. +func streamRequest(ctx context.Context, url string, body []byte, hdr map[string]string) (chunks int, ttfb time.Duration, err error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url+"/v1/chat/completions", bytes.NewReader(body)) + if err != nil { + return 0, 0, err + } + req.Header.Set("Content-Type", "application/json") + for k, v := range hdr { + req.Header.Set(k, v) + } + + start := time.Now() + resp, err := http.DefaultClient.Do(req) + if err != nil { + return 0, 0, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<12)) + return 0, 0, fmt.Errorf("status %d: %s", resp.StatusCode, b) + } + + br := bufio.NewReader(resp.Body) + for { + line, rerr := br.ReadString('\n') + if strings.HasPrefix(line, "data: ") { + if chunks == 0 { + ttfb = time.Since(start) + } + if !strings.Contains(line, "[DONE]") { + chunks++ + } + } + if rerr != nil { + if rerr == io.EOF { + return chunks, ttfb, nil + } + return chunks, ttfb, rerr + } + } +} + +// envInt reads an integer from the environment, falling back to def. +func envInt(key string, def int) int { + v := os.Getenv(key) + if v == "" { + return def + } + var n int + if _, err := fmt.Sscanf(v, "%d", &n); err != nil || n <= 0 { + return def + } + return n +} diff --git a/test/e2e/phase1_test.go b/test/e2e/phase1_test.go new file mode 100644 index 0000000..8ad7d6e --- /dev/null +++ b/test/e2e/phase1_test.go @@ -0,0 +1,197 @@ +package e2e + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "os" + "slices" + "sync" + "sync/atomic" + "testing" + "time" +) + +func bytesReader(b []byte) *bytes.Reader { return bytes.NewReader(b) } + +func drainBody(resp *http.Response) { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() +} + +// TestConcurrentStreams is the first half of the phase 1 exit criterion: +// 1000 concurrent streams behind 3 sim-engines, all of them completing. +func TestConcurrentStreams(t *testing.T) { + concurrency := envInt("LLMGW_E2E_CONCURRENCY", 1000) + if testing.Short() { + concurrency = 50 + } + const tokens = 16 + + h := start(t, 3) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + var ( + ok atomic.Int64 + short atomic.Int64 + failed atomic.Int64 + firstErr atomic.Value + wg sync.WaitGroup + ) + + start := time.Now() + for i := range concurrency { + wg.Add(1) + go func() { + defer wg.Done() + n, _, err := streamRequest(ctx, h.router.URL, + chatBody(fmt.Sprintf("request %d", i), tokens, true), + map[string]string{"x-session-id": fmt.Sprintf("s%d", i%97)}) + switch { + case err != nil: + failed.Add(1) + firstErr.CompareAndSwap(nil, err) + case n < tokens: + short.Add(1) + default: + ok.Add(1) + } + }() + } + wg.Wait() + elapsed := time.Since(start) + + t.Logf("%d streams in %v (%.0f req/s): ok=%d short=%d failed=%d", + concurrency, elapsed.Round(time.Millisecond), + float64(concurrency)/elapsed.Seconds(), ok.Load(), short.Load(), failed.Load()) + + if f := failed.Load(); f > 0 { + t.Fatalf("%d of %d streams failed; first error: %v", f, concurrency, firstErr.Load()) + } + if s := short.Load(); s > 0 { + t.Errorf("%d streams ended before %d tokens", s, tokens) + } + + // Load must actually have been spread; a single engine taking everything + // would pass the checks above while defeating the point of the router. + used := 0 + for _, e := range h.engines { + if st := e.eng.Stats(); st.PrefixQueries > 0 || st.CachedBlocks > 0 { + used++ + } + } + if used < 2 { + t.Errorf("only %d of %d engines saw traffic", used, len(h.engines)) + } +} + +// TestClientDisconnectStopsUpstream is the second half of the exit criterion: +// when a client goes away, the GPU must stop working on its request. +func TestClientDisconnectStopsUpstream(t *testing.T) { + const streams = 40 + + h := start(t, 3) + + ctx, cancel := context.WithCancel(context.Background()) + var wg sync.WaitGroup + for i := range streams { + wg.Add(1) + go func() { + defer wg.Done() + // Long generations, so every request is still running when cancelled. + _, _, _ = streamRequest(ctx, h.router.URL, + chatBody(fmt.Sprintf("long request %d", i), 100000, true), nil) + }() + } + + // Wait until the engines are actually busy before pulling the plug. + deadline := time.Now().Add(10 * time.Second) + for h.totalRunning() < streams && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if busy := h.totalRunning(); busy < streams { + t.Fatalf("only %d of %d requests reached the engines", busy, streams) + } + + cancel() + wg.Wait() + + // The engines must drain: cancellation has to propagate router -> engine. + deadline = time.Now().Add(15 * time.Second) + for time.Now().Before(deadline) { + if h.totalRunning() == 0 { + // And the KV blocks must come back, not just the request slots. + for _, e := range h.engines { + if u := e.eng.Stats().KVCacheUsage; u != 0 { + t.Fatalf("KV cache still held after cancellation: %v", u) + } + } + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatalf("engines still working %d requests after every client disconnected", h.totalRunning()) +} + +// TestRouterOverhead is the third part of the exit criterion: the router should +// add little latency of its own. +// +// It measures the same request against an engine directly and through the +// router, and reports the p99 difference. The assertion is deliberately loose, +// because a shared CI machine cannot hold a 2 ms bound reliably; set +// LLMGW_E2E_STRICT=1 on quiet hardware to enforce the plan's actual target. +func TestRouterOverhead(t *testing.T) { + samples := envInt("LLMGW_E2E_SAMPLES", 300) + if testing.Short() { + samples = 50 + } + + h := start(t, 1) + direct := h.engines[0].srv.URL + body := chatBody("overhead probe", 1, false) + + measure := func(url string) []time.Duration { + out := make([]time.Duration, 0, samples) + for range samples { + start := time.Now() + resp, err := http.Post(url+"/v1/chat/completions", "application/json", bytesReader(body)) //nolint:noctx // test + if err != nil { + t.Fatal(err) + } + drainBody(resp) + out = append(out, time.Since(start)) + } + slices.Sort(out) + return out + } + + // Warm both paths so connection setup does not land in the samples. + measure(direct) + measure(h.router.URL) + + d := measure(direct) + r := measure(h.router.URL) + + p := func(s []time.Duration, q float64) time.Duration { + return s[min(int(float64(len(s))*q), len(s)-1)] + } + overheadP50 := p(r, 0.50) - p(d, 0.50) + overheadP99 := p(r, 0.99) - p(d, 0.99) + + t.Logf("direct p50=%v p99=%v | routed p50=%v p99=%v | overhead p50=%v p99=%v", + p(d, 0.50).Round(time.Microsecond), p(d, 0.99).Round(time.Microsecond), + p(r, 0.50).Round(time.Microsecond), p(r, 0.99).Round(time.Microsecond), + overheadP50.Round(time.Microsecond), overheadP99.Round(time.Microsecond)) + + limit := 25 * time.Millisecond + if os.Getenv("LLMGW_E2E_STRICT") == "1" { + limit = 2 * time.Millisecond + } + if overheadP99 > limit { + t.Errorf("router added %v at p99, over the %v budget", overheadP99, limit) + } +}