Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
  ilkerandClaude Opus 5
|
a3f71aa5ee | ||
|
ilkerandClaude Opus 5
|
f0c1ca87f1 |
No files matched your search
+51
-2
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
)
|
||||
@@ -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=
|
||||
@@ -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
|
||||
}
|
||||
@@ -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{
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
}})
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 })
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+7
-4
@@ -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
|
||||
|
||||
@@ -35,7 +35,7 @@ MILESTONES = [
|
||||
"against. Exit criterion: sim-engine models batching, KV blocks, prefix "
|
||||
"caching and cancellation, and exposes the §7.6 metric contract."
|
||||
),
|
||||
"due_on": "2026-09-16T23:59:59Z",
|
||||
"due_on": "2026-09-16T12:00:00Z",
|
||||
"closed": True,
|
||||
},
|
||||
{
|
||||
@@ -46,7 +46,7 @@ MILESTONES = [
|
||||
"Exit criterion: 1000 concurrent streams behind 3 sim-engines; upstream "
|
||||
"cancellation on client disconnect verified; router overhead p99 < 2 ms."
|
||||
),
|
||||
"due_on": "2026-09-30T23:59:59Z",
|
||||
"due_on": "2026-09-30T12:00:00Z",
|
||||
},
|
||||
{
|
||||
"title": "v0.3.0 — Phase 1.5: admission control",
|
||||
@@ -57,7 +57,7 @@ MILESTONES = [
|
||||
"bronze band, interactive TTFT p99 in the platinum band stays under 1.5x "
|
||||
"its idle value. This test stays in CI permanently."
|
||||
),
|
||||
"due_on": "2026-10-07T23:59:59Z",
|
||||
"due_on": "2026-10-07T12:00:00Z",
|
||||
},
|
||||
{
|
||||
"title": "v0.4.0 — Phase 2: agent",
|
||||
@@ -66,7 +66,7 @@ MILESTONES = [
|
||||
"the control plane stream.\n\n"
|
||||
"Exit criterion: start and stop a model on a real node via ctl."
|
||||
),
|
||||
"due_on": "2026-10-21T23:59:59Z",
|
||||
"due_on": "2026-10-21T12:00:00Z",
|
||||
},
|
||||
{
|
||||
"title": "v0.5.0 — Phase 3: control plane",
|
||||
@@ -76,7 +76,7 @@ MILESTONES = [
|
||||
"Exit criterion: shutting down a node moves its replicas elsewhere while "
|
||||
"the router keeps serving."
|
||||
),
|
||||
"due_on": "2026-11-11T23:59:59Z",
|
||||
"due_on": "2026-11-11T12:00:00Z",
|
||||
},
|
||||
{
|
||||
"title": "v0.6.0 — Phase 4: multi-node and fabrics",
|
||||
@@ -86,7 +86,7 @@ MILESTONES = [
|
||||
"Exit criterion: a 2-node replica on each fabric; killing a node rebuilds "
|
||||
"the replica on the same fabric and never spreads it across two."
|
||||
),
|
||||
"due_on": "2026-11-25T23:59:59Z",
|
||||
"due_on": "2026-11-25T12:00:00Z",
|
||||
},
|
||||
{
|
||||
"title": "v1.0.0 — Phase 5: production features",
|
||||
@@ -95,7 +95,7 @@ MILESTONES = [
|
||||
"prefix tree, fallback models and pool spillover, ctl bench, tracing, "
|
||||
"systemd units and documentation."
|
||||
),
|
||||
"due_on": "2026-12-16T23:59:59Z",
|
||||
"due_on": "2026-12-16T12:00:00Z",
|
||||
},
|
||||
{
|
||||
"title": "Phase 6 — v2 ideas",
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user