Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e1eb1d4149 | ||
|
|
1d1a6aebb0 | ||
|
|
75a8e76706 | ||
|
|
40427b2705 | ||
|
|
2373726efd | ||
|
|
c1b9ea1bc2 | ||
|
|
f0134a2f95 | ||
|
|
06672d1b85 | ||
|
|
55417546fc | ||
|
|
6a13e3dddf | ||
|
|
dc38dd6096 | ||
|
|
a3f71aa5ee | ||
|
|
f0c1ca87f1 |
No files matched your search
@@ -38,3 +38,11 @@ jobs:
|
||||
|
||||
- name: Build
|
||||
run: go build ./...
|
||||
|
||||
# The NVML backend needs cgo and cannot run without a driver, but it must
|
||||
# at least keep compiling: it is the only source of NVLink topology and a
|
||||
# break would otherwise surface on a GPU node, far from here.
|
||||
- name: Build the NVML backend
|
||||
env:
|
||||
CGO_ENABLED: "1"
|
||||
run: go build -tags nvml ./...
|
||||
+389
-3
@@ -6,9 +6,9 @@ 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.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | planned |
|
||||
| v0.2.0 | 1 | Static router: streaming proxy, filter/score/pick | released |
|
||||
| v0.3.0 | 1.5 | Admission control: priority bands, saturation detector | released |
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | in progress |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler | planned |
|
||||
| v0.6.0 | 4 | Multi-node replicas and fabric awareness | planned |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
@@ -16,6 +16,392 @@ 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.5.0 — 2026-09-16
|
||||
|
||||
Phases 2 and 3. The cluster runs itself: agents report their hardware, the
|
||||
control plane places replicas, and `ctl` makes all of it visible.
|
||||
|
||||
This release covers both phases because their exit criteria depended on the same
|
||||
piece — the agent stream — and neither was demonstrable without it.
|
||||
|
||||
### ctl
|
||||
|
||||
```
|
||||
llmgw ctl nodes | models | replicas | explain-placement | cordon | uncordon | drain
|
||||
```
|
||||
|
||||
- `nodes` surfaces unhealthy GPUs and why they are excluded, which is usually
|
||||
the answer when a node looks idle but nothing schedules onto it.
|
||||
- `explain-placement` prints every candidate considered and the reason it was
|
||||
rejected. "Unschedulable" on its own says nothing about which of six
|
||||
constraints bit.
|
||||
- `drain` is a rolling move, not a stop: the replica is drained, a replacement
|
||||
is placed elsewhere, and the original is stopped only once the replacement is
|
||||
serving. Bounded by a grace period, because draining the only node that can
|
||||
host a model would otherwise never finish and leave it undrainable.
|
||||
- `cordon` is the weaker operation: no new placement, nothing disturbed.
|
||||
|
||||
The management API is on the admin listener, not the client-facing one. The two
|
||||
have different audiences: one serves inference to whoever holds an API key, the
|
||||
other lets a caller cordon nodes.
|
||||
|
||||
`ctl logs` is not here. It needs a request/response pair the agent protocol does
|
||||
not have. The failure path already ships an engine log tail with every status
|
||||
report, so the case that matters most is covered; reading a healthy shard's
|
||||
output on demand is not.
|
||||
|
||||
## Unreleased — phase 4 in progress
|
||||
|
||||
Agent groundwork. Everything here works without GPUs and is covered by tests;
|
||||
what is missing is the control plane connection, so there is no release yet.
|
||||
|
||||
### Added
|
||||
|
||||
- **IB discovery** (`internal/agent/discovery`): reads `/sys/class/infiniband`,
|
||||
parses the formats sysfs actually uses (`"4: ACTIVE"`, `"200 Gb/sec (4X HDR)"`,
|
||||
hex or decimal LIDs), and distinguishes native InfiniBand from RoCE. A node
|
||||
without IB is not an error — it is a single-node island.
|
||||
- **GPU discovery** via `nvidia-smi`, the cgo-free path. NVML needs cgo, which
|
||||
would force the whole binary, server included, to be built with a C
|
||||
toolchain. Health flags anything that should keep a card out of placement:
|
||||
uncorrected ECC errors, and memory already held by another process.
|
||||
- **GPU class derivation.** Sizes come from the driver in GiB, which does not
|
||||
always match the marketing name: an L40S is sold as 48GB and reports 44.99
|
||||
GiB, so its class is `l40s-45g`. Consistency matters more than matching the
|
||||
box, because the number decides whether a model fits.
|
||||
- **Engine adapters** for vLLM and SGLang: command and environment
|
||||
construction, kept free of process handling so it is testable without GPUs.
|
||||
vLLM supports both Ray and multiprocessing launch modes, because which one
|
||||
works depends on the version.
|
||||
- **NCCL environment generation**, the part plan §5.3 marks critical.
|
||||
`NCCL_IB_HCA` lists only the chosen fabric's adapters — leaving it unset lets
|
||||
NCCL pick an adapter on another fabric where the peer is unreachable, and the
|
||||
job hangs rather than failing. Bootstrap is pinned to the management
|
||||
interface for the same reason.
|
||||
- **Validation before launch**: a tensor-parallel group may not mix GPU types,
|
||||
the node/GPU/TP arithmetic has to add up, and an FP8 profile is refused below
|
||||
compute capability 8.9 rather than silently falling back to a slower path.
|
||||
- **Supervisor** (`internal/agent/supervisor`): processes run in their own
|
||||
process group and are stopped as a group, because engines fork workers that
|
||||
hold GPU memory and outlive a signal sent only to the leader. Bounded log
|
||||
tail, a state machine that separates Starting from Loading (a 70B model takes
|
||||
minutes, and an operator needs to know that is normal), and orphan adoption
|
||||
through pid files fingerprinted with the command line — pids get reused, and
|
||||
adopting an unrelated process is worse than leaking the original.
|
||||
- **Readiness** is two questions: `/health` says the process is up, `/v1/models`
|
||||
says the weights loaded. Routing on `/health` alone sends traffic into a model
|
||||
that is still loading.
|
||||
- **`llmgw agent inventory`** reports what the scheduler would see on a node,
|
||||
as a table or JSON. Run it before adding a node: it is the fastest way to find
|
||||
a card another process is using, or a port with no subnet manager.
|
||||
|
||||
- **NVML backend** behind the `nvml` build tag, which adds NVLink topology.
|
||||
Nothing else exposes which GPUs share an NVLink domain, and that is what keeps
|
||||
a tensor-parallel group off the slow path inside a node. `Discover` falls back
|
||||
to nvidia-smi if NVML fails, and says so loudly rather than silently losing the
|
||||
topology. CI compiles the tagged build so it cannot rot unnoticed.
|
||||
|
||||
### Not verified
|
||||
|
||||
The NVML path compiles but has never run against a driver — there is no GPU
|
||||
machine to try it on. Treat the NVLink domains it reports as unverified until
|
||||
someone runs `llmgw agent inventory` on a real node.
|
||||
|
||||
- **Control plane wiring** (`internal/control`): Ready replicas become router
|
||||
endpoints, with no static configuration involved. The router keeps live state
|
||||
for endpoints that survive an update — inflight counts, breaker state, drain
|
||||
rate — so a replica that was already serving does not start from cold because
|
||||
another one moved. The sink is an interface, so when the router and control
|
||||
plane split into separate processes this becomes a watch stream and nothing
|
||||
above it changes.
|
||||
|
||||
### Phase 3 exit criterion, measured in `test/e2e`
|
||||
|
||||
A node hosting a live replica is killed while four clients stream continuously.
|
||||
The control plane marks it unreachable, then lost, rebuilds the replica on a
|
||||
survivor, and republishes. Across three runs: **~370 requests completed, zero
|
||||
failures**, and no replica remained on the dead node.
|
||||
|
||||
This proves the control logic, not the transport: the agent fleet is simulated,
|
||||
and the real path still needs the stream.
|
||||
|
||||
### A real bug the race detector caught
|
||||
|
||||
`startReplica` read the replica's fields after publishing it to the map, and
|
||||
`scaleUp` walked replica pointers outside the lock. Both were racing with agent
|
||||
status reports, which arrive from another goroutine and mutate exactly those
|
||||
fields. Under the simulated fleet the reports came back fast enough to hit the
|
||||
window every time. Everything is snapshotted under the lock now and the
|
||||
functions work from values.
|
||||
|
||||
### The agent stream
|
||||
|
||||
- **Protocol** (`api/proto/llmgw/v1`): one bidirectional connect-go stream
|
||||
carrying registration, heartbeats and shard status upward, and start, stop and
|
||||
drain commands downward. The agent dials out, because GPU nodes sit behind
|
||||
firewalls far more often than a control plane does and an outbound connection
|
||||
needs no inbound rule on every node in the fleet.
|
||||
- **Registration carries the whole inventory**, so there is no separate
|
||||
discovery call and a reconnecting agent refreshes what the scheduler knows.
|
||||
The fabric labels travel with it — they cannot be inferred from the wire, so
|
||||
losing them in transit would silently disable multi-node placement.
|
||||
- **Agent service** (`internal/control/agentsvc`) is both the stream handler and
|
||||
the reconciler's dispatcher: it holds the open streams, so it is the only
|
||||
thing that can deliver a command to a node.
|
||||
- **Agent** (`internal/agent`) validates a profile before launching anything —
|
||||
an engine that refuses a profile does so after minutes of loading — then
|
||||
supervises the process and reports every transition, including the engine's
|
||||
output tail on failure.
|
||||
- **Reconnects forever with jittered backoff.** A control plane restart must not
|
||||
require touching every GPU node, and a fleet must not reconnect in lockstep
|
||||
and knock the control plane over again. Engines keep serving throughout: the
|
||||
router talks to them directly.
|
||||
- **`llmgw agent` is real**, and `llmgw server` runs the control plane whenever
|
||||
any model declares a profile. The static endpoint list remains for
|
||||
router-only runs against sim-engines.
|
||||
|
||||
### Verified with real binaries
|
||||
|
||||
Server and agent as separate processes over a real stream: the agent registers,
|
||||
a replica is placed, an engine is launched through the vLLM adapter's actual
|
||||
argument construction, streaming traffic is served through the router, and
|
||||
killing the agent moves the node through unreachable to lost with placement
|
||||
explaining exactly why it can no longer run anything.
|
||||
|
||||
### Surviving a control plane restart
|
||||
|
||||
Not persisting replica state left one hole that had to be closed before it
|
||||
mattered: restarting the control plane orphaned live engines. They kept holding
|
||||
GPUs the inventory believed were free, so the next placement failed to allocate
|
||||
with an error naming the new process rather than the old one.
|
||||
|
||||
The fix is on the agent, not in a database, and it cannot disagree with reality
|
||||
the way a stale row can:
|
||||
|
||||
- The supervisor now writes the full shard spec, not just a pid and command
|
||||
line. A pid identifies a process; the spec is what lets it be described back
|
||||
to the control plane and matched to a replica.
|
||||
- `Adopt` takes over a still-running engine from a previous agent. It cannot be
|
||||
waited on — it is not this process's child — so it is watched by polling, and
|
||||
stopped by signalling its group directly. Its earlier output is gone, and the
|
||||
log says so rather than presenting emptiness as silence.
|
||||
- Every `Register`, not only the first, reports what the agent is running. A
|
||||
reconnect may be a control plane restart and has to recover the same way.
|
||||
- The reconciler rebuilds replicas from those reports and reclaims their GPUs.
|
||||
Anything whose model or profile left the catalog is stopped: nothing will
|
||||
route to it and it is holding cards.
|
||||
- An engine the agent cannot identify — a corrupt state file — is killed. A
|
||||
process holding GPUs that nothing can describe is worse than no process,
|
||||
because placement will hand those GPUs out again.
|
||||
|
||||
One mechanism covers agent restarts and control plane restarts.
|
||||
|
||||
### A test that was passing for no reason
|
||||
|
||||
`TestAgentReconnects` used `httptest.CloseClientConnections`, which does nothing
|
||||
to an h2c stream: prior-knowledge h2c hijacks the connection, and httptest stops
|
||||
tracking hijacked connections. The agent never reconnected and the assertion —
|
||||
on the connection map, which still held a stale entry — passed anyway. Tests now
|
||||
sever connections through their own listener and count opened streams.
|
||||
|
||||
### Multi-node launch
|
||||
|
||||
Placement and environment generation had landed early, which made multi-node
|
||||
look further along than it was: a multi-node replica could not come up at all.
|
||||
Three things were missing, and all three came down to two assumptions that do
|
||||
not hold once a replica spans nodes.
|
||||
|
||||
**Not every rank serves HTTP.** The agent health-checked every shard at its own
|
||||
port, but only the leader answers. Workers timed out, reported Failed, and took
|
||||
the replica with them. `Engine.ServesHTTP(spec)` now says which ranks serve; a
|
||||
worker is judged by staying alive instead. That is weaker evidence on purpose —
|
||||
the replica's real readiness is the leader answering, which only happens once
|
||||
every worker has joined.
|
||||
|
||||
**A shard is not always one process.** `Engine.Commands(spec)` returns a
|
||||
sequence: every command but the last is a sidecar kept alive for the shard's
|
||||
lifetime and torn down with it, and the last one's exit is the shard's exit.
|
||||
vLLM under Ray is why — the leader runs a Ray head and then the server against
|
||||
it, and a worker runs only the join, with no vLLM process at all. `RayCommand`
|
||||
had been written and tested but never called, so `launch_mode: ray` would have
|
||||
run `vllm serve` against a cluster that did not exist.
|
||||
|
||||
A command can also require the rendezvous point to be listening before the next
|
||||
one starts, bounded by a timeout so a head that never appears fails the shard
|
||||
rather than hanging it. Sidecars are killed with the shard: a Ray head left
|
||||
behind holds the port the next replica needs.
|
||||
|
||||
**Start order was backwards.** Workers were started first, on the theory that
|
||||
the head should wait for them. Both engines put the rendezvous point on rank 0,
|
||||
so a worker started first has nothing to connect to and spends its startup
|
||||
budget retrying. vLLM's server blocks until the Ray cluster has the GPUs it
|
||||
needs, so the leader can start immediately and waits on its own.
|
||||
|
||||
None of this is verified against real engines. It is tested with stub processes;
|
||||
the Ray sequencing and which launch mode a given vLLM version wants can only be
|
||||
confirmed on hardware.
|
||||
|
||||
### Still missing for phase 2
|
||||
|
||||
- mTLS. The stream is plaintext HTTP/2 today, so it must stay on a trusted
|
||||
network.
|
||||
|
||||
### Phase 3: placement
|
||||
|
||||
- **Cluster inventory** (`internal/control/inventory`): nodes, their hardware,
|
||||
their fabric attachments, and what is allocated where. A re-registering agent
|
||||
cannot clear an operator's cordon, and a heartbeat revives an unreachable node
|
||||
without doing so either. A GPU with an ECC fault or another process on it is
|
||||
not counted as free. Replicas are released whole, because they fail whole.
|
||||
- **Placement** (`internal/control/placement`): picks the highest-priority
|
||||
profile that fits and explains every candidate it rejected. Two constraints
|
||||
are enforced structurally rather than checked afterwards — there is no code
|
||||
path that can combine nodes from two fabrics, and only GPUs of the profile's
|
||||
class are ever considered.
|
||||
- **Bin packing** fills partly used nodes first, so whole nodes stay free for
|
||||
large models that cannot be split. Multi-node replicas get the opposite
|
||||
treatment: they are the least movable thing in the cluster, so they take the
|
||||
emptiest machines.
|
||||
- **NVLink-aware GPU selection** keeps a tensor-parallel group inside one NVLink
|
||||
domain where the topology is known. Without the NVML build it falls back to
|
||||
the lowest free indices, which is usually right but only by accident of how
|
||||
vendors number slots.
|
||||
- **Property-based invariant tests** (plan §13): 400 randomly generated
|
||||
clusters, checking on every placement that a multi-node replica lives on one
|
||||
fabric, that each shard receives only that fabric's adapters, and that no
|
||||
placement mixes GPU classes or uses an unhealthy or already-allocated GPU.
|
||||
The table-driven tests cover the cases someone thought of; this covers the
|
||||
ones nobody did.
|
||||
- **Catalog** (`internal/control/catalog`): models, profiles, aliases,
|
||||
fallbacks and scaling bounds. Aliases and fallbacks are resolved at load
|
||||
rather than at request time — a typo in a fallback should stop a deployment,
|
||||
not surface during the outage when the fallback is finally needed. The
|
||||
node/GPU/TP arithmetic is checked here too, so a bad profile fails at config
|
||||
load instead of minutes into model loading.
|
||||
- **Reconciler** (`internal/control/reconciler`): the desired-versus-actual
|
||||
loop. It places one replica per pass, because placement depends on what is
|
||||
already allocated and starting several from one snapshot would hand them the
|
||||
same GPUs. Workers start before rank 0, since SGLang's ranks rendezvous and
|
||||
vLLM under Ray needs its workers attached before the head serves.
|
||||
- **Replica failure is whole-replica failure.** One dead shard makes the others
|
||||
useless — they wait on collectives that will never complete while holding
|
||||
GPUs — so the reconciler stops every shard and rebuilds, with a jittered
|
||||
exponential backoff carried on the model rather than the replica it just
|
||||
destroyed.
|
||||
- **Node expiry in two stages.** Unreachable first, Lost later. A brief network
|
||||
blip must not tear down replicas that are still serving, and rebuilding a
|
||||
multi-node replica costs tens of minutes of model loading.
|
||||
- **Scale-down drains before stopping**, so in-flight requests finish instead of
|
||||
being cut off mid-stream.
|
||||
|
||||
## v0.3.0 — 2026-09-16
|
||||
|
||||
Flow control. A batch workload no longer starves interactive traffic.
|
||||
|
||||
### Added
|
||||
|
||||
- **Admission control** (`internal/router/admission`), inserted between quota
|
||||
checks and scheduling. It answers *when* a request may proceed; the scheduler
|
||||
still answers *where* it goes. Keeping them apart is the point: an engine
|
||||
queue is FIFO and knows nothing about tenants, so once a batch job is in it,
|
||||
no amount of endpoint selection rescues the request behind it.
|
||||
- **Saturation detector.** One number per pool, the worse of queue depth and KV
|
||||
occupancy. Either signal alone misleads — a pool can have an empty queue and a
|
||||
nearly full cache, or the reverse.
|
||||
- **Priority bands.** Each band has a saturation ceiling; above it, that band
|
||||
waits. The cluster then degrades from the bottom up. A ceiling of 1.0 means
|
||||
"flow up to the healthy limit"; above 1 means never wait.
|
||||
- **Tenant fairness.** Within a band, dispatch turns rotate between tenants, so
|
||||
one caller sending a thousand requests occupies one slot in the rotation
|
||||
rather than a thousand.
|
||||
- **Bounded queues.** Capacity per band (immediate 503 when full) and a TTL
|
||||
(504), because sending the GPU work the client has already abandoned helps
|
||||
nobody. Cancelled requests leave the queue instead of holding capacity.
|
||||
- **Metrics:** `llmgw_admission_wait_seconds`, `llmgw_admission_rejected_total`,
|
||||
`llmgw_admission_queue_depth`, `llmgw_pool_saturation`. Responses carry
|
||||
`x-llmgw-queue-ms` when a request waited.
|
||||
|
||||
### Phase 1.5 exit criterion, measured in `test/e2e`
|
||||
|
||||
While 500 concurrent bronze requests run continuously, interactive p99 in the
|
||||
platinum band stays at **1.04–1.24x** its idle value, against a 1.5x budget.
|
||||
With admission switched off the same load pushes it to **55–60x**. Both
|
||||
directions are asserted, so the test fails if admission regresses *and* if the
|
||||
load generator ever stops being a real load.
|
||||
|
||||
### Two lags, two different fixes
|
||||
|
||||
Getting this right took longer than the queueing itself:
|
||||
|
||||
- Engine metrics are scraped, so they lag by an interval. The router's inflight
|
||||
count does not — but it includes *running* requests, and treating those as
|
||||
queue depth makes a healthy engine look overloaded every interval. What is
|
||||
genuinely queued is `inflight − batch capacity`, and batch capacity is not in
|
||||
the metric contract, so it is learned as the high-water mark of observed
|
||||
`running`.
|
||||
- Between admission granting a request and the router dispatching it, the
|
||||
request is invisible to every signal. A simultaneous burst therefore passed
|
||||
the check en masse — exactly the batch arrival pattern this layer exists to
|
||||
survive. The controller now counts its own outstanding leases against
|
||||
dispatched inflight.
|
||||
|
||||
### Known limitations
|
||||
|
||||
- The band comes from a trusted header until authentication arrives in v1.0.0.
|
||||
Anyone who can reach the port can promote themselves to platinum.
|
||||
- The tenant key is the session id, for the same reason.
|
||||
- Displacement (dropping queued low-band work to make room) is not implemented;
|
||||
the plan defers it, and bounded queues plus TTL cover most of the need.
|
||||
|
||||
## 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
|
||||
|
||||
@@ -2,7 +2,7 @@ MODULE := gitea.mynodes.xyz/ilker/llmgw
|
||||
VERSION ?= $(shell git describe --tags --always --dirty 2>/dev/null || echo dev)
|
||||
LDFLAGS := -X '$(MODULE)/internal/version.Version=$(VERSION)'
|
||||
|
||||
.PHONY: build test race lint fmt clean rename
|
||||
.PHONY: build test race lint fmt clean rename proto proto-tools
|
||||
|
||||
build:
|
||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw ./cmd/llmgw
|
||||
@@ -20,6 +20,17 @@ fmt:
|
||||
go fmt ./...
|
||||
go vet ./...
|
||||
|
||||
# Regenerate the agent protocol. Generated code is committed, so this only
|
||||
# needs running when api/proto changes.
|
||||
proto: proto-tools
|
||||
buf lint
|
||||
buf generate
|
||||
|
||||
proto-tools:
|
||||
go install github.com/bufbuild/buf/cmd/buf@latest
|
||||
go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
|
||||
go install connectrpc.com/connect/cmd/protoc-gen-connect-go@latest
|
||||
|
||||
clean:
|
||||
rm -rf bin dist
|
||||
|
||||
|
||||
@@ -13,15 +13,17 @@ 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.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.2.0 | 1 | Static router (filter/score/pick, streaming proxy) | released |
|
||||
| v0.3.0 | 1.5 | Admission control (priority bands) | released |
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, engine supervision | released |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler, ctl | released |
|
||||
| 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.3.0 the router serves traffic and protects it: a batch workload no
|
||||
longer starves interactive requests. Still missing before v1.0.0: authentication,
|
||||
rate limits and budgets, and the control plane that replaces the static endpoint
|
||||
list.
|
||||
|
||||
## Building
|
||||
|
||||
@@ -31,6 +33,55 @@ make test # unit tests
|
||||
make race # with the race detector
|
||||
```
|
||||
|
||||
## Operating a cluster
|
||||
|
||||
```sh
|
||||
llmgw ctl nodes # capacity, health, what is running where
|
||||
llmgw ctl models # catalog and replica counts
|
||||
llmgw ctl replicas --model llama # replicas and their shards
|
||||
llmgw ctl explain-placement llama # why a model has not been placed
|
||||
llmgw ctl drain gpu-a-01 # move work off a node, one replica at a time
|
||||
llmgw ctl cordon gpu-a-01 # stop new placement, leave what runs
|
||||
```
|
||||
|
||||
`ctl` talks to the admin port (`:9090` by default), not the client-facing one.
|
||||
Add `--json` to any command for machine-readable output.
|
||||
|
||||
`explain-placement` is the one worth knowing about: "unschedulable" on its own
|
||||
says nothing about which of six constraints bit, so it prints every candidate
|
||||
considered and the reason it was rejected.
|
||||
|
||||
## 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.
|
||||
|
||||
Requests carry a priority band (`x-llmgw-band`). Under load, low bands wait at
|
||||
the gateway while high bands keep flowing, so the cluster degrades from the
|
||||
bottom up instead of everywhere at once. Measured effect: interactive p99 stays
|
||||
within 1.24x of idle while 500 concurrent batch requests run, against 58x with
|
||||
admission off.
|
||||
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,119 @@
|
||||
// Code generated by protoc-gen-connect-go. DO NOT EDIT.
|
||||
//
|
||||
// Source: llmgw/v1/agent.proto
|
||||
|
||||
// Package llmgw.v1 is the protocol between a node agent and the control plane
|
||||
// (plan §5.6).
|
||||
//
|
||||
// The agent dials out and holds a bidirectional stream open. That direction is
|
||||
// deliberate: GPU nodes sit behind firewalls and NAT far more often than a
|
||||
// control plane does, and an outbound connection needs no inbound rule on every
|
||||
// node in the fleet.
|
||||
package llmgwv1connect
|
||||
|
||||
import (
|
||||
connect "connectrpc.com/connect"
|
||||
context "context"
|
||||
errors "errors"
|
||||
v1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
http "net/http"
|
||||
strings "strings"
|
||||
)
|
||||
|
||||
// This is a compile-time assertion to ensure that this generated file and the connect package are
|
||||
// compatible. If you get a compiler error that this constant is not defined, this code was
|
||||
// generated with a version of connect newer than the one compiled into your binary. You can fix the
|
||||
// problem by either regenerating this code with an older version of connect or updating the connect
|
||||
// version compiled into your binary.
|
||||
const _ = connect.IsAtLeastVersion1_13_0
|
||||
|
||||
const (
|
||||
// AgentServiceName is the fully-qualified name of the AgentService service.
|
||||
AgentServiceName = "llmgw.v1.AgentService"
|
||||
)
|
||||
|
||||
// These constants are the fully-qualified names of the RPCs defined in this package. They're
|
||||
// exposed at runtime as Spec.Procedure and as the final two segments of the HTTP route.
|
||||
//
|
||||
// Note that these are different from the fully-qualified method names used by
|
||||
// google.golang.org/protobuf/reflect/protoreflect. To convert from these constants to
|
||||
// reflection-formatted method names, remove the leading slash and convert the remaining slash to a
|
||||
// period.
|
||||
const (
|
||||
// AgentServiceConnectProcedure is the fully-qualified name of the AgentService's Connect RPC.
|
||||
AgentServiceConnectProcedure = "/llmgw.v1.AgentService/Connect"
|
||||
)
|
||||
|
||||
// AgentServiceClient is a client for the llmgw.v1.AgentService service.
|
||||
type AgentServiceClient interface {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
Connect(context.Context) *connect.BidiStreamForClient[v1.AgentMessage, v1.ControlMessage]
|
||||
}
|
||||
|
||||
// NewAgentServiceClient constructs a client for the llmgw.v1.AgentService service. By default, it
|
||||
// uses the Connect protocol with the binary Protobuf Codec, asks for gzipped responses, and sends
|
||||
// uncompressed requests. To use the gRPC or gRPC-Web protocols, supply the connect.WithGRPC() or
|
||||
// connect.WithGRPCWeb() options.
|
||||
//
|
||||
// The URL supplied here should be the base URL for the Connect or gRPC server (for example,
|
||||
// http://api.acme.com or https://acme.com/grpc).
|
||||
func NewAgentServiceClient(httpClient connect.HTTPClient, baseURL string, opts ...connect.ClientOption) AgentServiceClient {
|
||||
baseURL = strings.TrimRight(baseURL, "/")
|
||||
agentServiceMethods := v1.File_llmgw_v1_agent_proto.Services().ByName("AgentService").Methods()
|
||||
return &agentServiceClient{
|
||||
connect: connect.NewClient[v1.AgentMessage, v1.ControlMessage](
|
||||
httpClient,
|
||||
baseURL+AgentServiceConnectProcedure,
|
||||
connect.WithSchema(agentServiceMethods.ByName("Connect")),
|
||||
connect.WithClientOptions(opts...),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
// agentServiceClient implements AgentServiceClient.
|
||||
type agentServiceClient struct {
|
||||
connect *connect.Client[v1.AgentMessage, v1.ControlMessage]
|
||||
}
|
||||
|
||||
// Connect calls llmgw.v1.AgentService.Connect.
|
||||
func (c *agentServiceClient) Connect(ctx context.Context) *connect.BidiStreamForClient[v1.AgentMessage, v1.ControlMessage] {
|
||||
return c.connect.CallBidiStream(ctx)
|
||||
}
|
||||
|
||||
// AgentServiceHandler is an implementation of the llmgw.v1.AgentService service.
|
||||
type AgentServiceHandler interface {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
Connect(context.Context, *connect.BidiStream[v1.AgentMessage, v1.ControlMessage]) error
|
||||
}
|
||||
|
||||
// NewAgentServiceHandler builds an HTTP handler from the service implementation. It returns the
|
||||
// path on which to mount the handler and the handler itself.
|
||||
//
|
||||
// By default, handlers support the Connect, gRPC, and gRPC-Web protocols with the binary Protobuf
|
||||
// and JSON codecs. They also support gzip compression.
|
||||
func NewAgentServiceHandler(svc AgentServiceHandler, opts ...connect.HandlerOption) (string, http.Handler) {
|
||||
agentServiceMethods := v1.File_llmgw_v1_agent_proto.Services().ByName("AgentService").Methods()
|
||||
agentServiceConnectHandler := connect.NewBidiStreamHandler(
|
||||
AgentServiceConnectProcedure,
|
||||
svc.Connect,
|
||||
connect.WithSchema(agentServiceMethods.ByName("Connect")),
|
||||
connect.WithHandlerOptions(opts...),
|
||||
)
|
||||
return "/llmgw.v1.AgentService/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case AgentServiceConnectProcedure:
|
||||
agentServiceConnectHandler.ServeHTTP(w, r)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// UnimplementedAgentServiceHandler returns CodeUnimplemented from all methods.
|
||||
type UnimplementedAgentServiceHandler struct{}
|
||||
|
||||
func (UnimplementedAgentServiceHandler) Connect(context.Context, *connect.BidiStream[v1.AgentMessage, v1.ControlMessage]) error {
|
||||
return connect.NewError(connect.CodeUnimplemented, errors.New("llmgw.v1.AgentService.Connect is not implemented"))
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
syntax = "proto3";
|
||||
|
||||
// Package llmgw.v1 is the protocol between a node agent and the control plane
|
||||
// (plan §5.6).
|
||||
//
|
||||
// The agent dials out and holds a bidirectional stream open. That direction is
|
||||
// deliberate: GPU nodes sit behind firewalls and NAT far more often than a
|
||||
// control plane does, and an outbound connection needs no inbound rule on every
|
||||
// node in the fleet.
|
||||
package llmgw.v1;
|
||||
|
||||
option go_package = "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1;llmgwv1";
|
||||
|
||||
service AgentService {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
rpc Connect(stream AgentMessage) returns (stream ControlMessage);
|
||||
}
|
||||
|
||||
// AgentMessage is anything a node tells the control plane.
|
||||
message AgentMessage {
|
||||
oneof msg {
|
||||
Register register = 1;
|
||||
Heartbeat heartbeat = 2;
|
||||
ShardStatus shard = 3;
|
||||
}
|
||||
}
|
||||
|
||||
// ControlMessage is anything the control plane tells a node.
|
||||
message ControlMessage {
|
||||
oneof msg {
|
||||
StartShard start = 1;
|
||||
StopShard stop = 2;
|
||||
DrainShard drain = 3;
|
||||
Ack ack = 4;
|
||||
}
|
||||
}
|
||||
|
||||
// Register is the first message on every stream. It carries the node's full
|
||||
// hardware inventory, so the control plane needs no separate discovery call and
|
||||
// a reconnecting agent automatically refreshes what the scheduler knows.
|
||||
message Register {
|
||||
string node_id = 1;
|
||||
string hostname = 2;
|
||||
// mgmt_addr is the address engines are reachable on, and the address NCCL
|
||||
// bootstraps over. Never a fabric address: NCCL rendezvouses over TCP before
|
||||
// it moves to RDMA.
|
||||
string mgmt_addr = 3;
|
||||
map<string, string> labels = 4;
|
||||
|
||||
repeated GPU gpus = 5;
|
||||
// fabrics is empty on a node without InfiniBand. Such a node is its own
|
||||
// single-node island and can never host a multi-node replica.
|
||||
repeated FabricAttachment fabrics = 6;
|
||||
|
||||
string agent_version = 7;
|
||||
// engine_versions decides which CLI flags and metric names apply; both drift
|
||||
// between engine releases (plan §14).
|
||||
map<string, string> engine_versions = 8;
|
||||
|
||||
// running_shards are the engines this agent is already supervising, including
|
||||
// ones adopted from a previous agent.
|
||||
//
|
||||
// This is what lets a control plane restart recover instead of orphaning live
|
||||
// engines. Without it the control plane forgets every replica, places new
|
||||
// ones, and the old engines keep holding GPUs the inventory believes are free
|
||||
// — so the new engines fail to allocate, with an error naming the wrong
|
||||
// process (plan §6.5).
|
||||
repeated RunningShard running_shards = 9;
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
//
|
||||
// It carries enough for the control plane to rebuild the replica it belongs to,
|
||||
// which is more than a shard id: the replica's identity, its pool and fabric,
|
||||
// and the leader endpoint clients reach it through.
|
||||
message RunningShard {
|
||||
string replica_id = 1;
|
||||
string shard_id = 2;
|
||||
int32 rank = 3;
|
||||
int32 nnodes = 4;
|
||||
repeated int32 gpu_indices = 5;
|
||||
repeated string hcas = 6;
|
||||
int32 http_port = 7;
|
||||
ShardStatus.State state = 8;
|
||||
|
||||
// model and profile name the catalog entry this shard was started from.
|
||||
string model = 9;
|
||||
string profile = 10;
|
||||
// pool and fabric are echoed back from the StartShard that created it; the
|
||||
// agent has no way to derive them.
|
||||
string pool = 11;
|
||||
string fabric = 12;
|
||||
|
||||
string leader_addr = 13;
|
||||
int32 leader_port = 14;
|
||||
}
|
||||
|
||||
message GPU {
|
||||
int32 index = 1;
|
||||
string uuid = 2;
|
||||
string name = 3;
|
||||
int32 memory_total_mb = 4;
|
||||
int32 memory_used_mb = 5;
|
||||
string compute_capability = 6;
|
||||
int32 numa_node = 7;
|
||||
string pci_bus_id = 8;
|
||||
string driver_version = 9;
|
||||
string cuda_version = 10;
|
||||
// nvlink_peers is only populated by agents built with the nvml tag.
|
||||
repeated int32 nvlink_peers = 11;
|
||||
// problems keeps a GPU out of placement: ECC faults, or memory already held
|
||||
// by another process.
|
||||
repeated string problems = 12;
|
||||
int32 temperature_c = 13;
|
||||
int32 ecc_errors = 14;
|
||||
}
|
||||
|
||||
message FabricAttachment {
|
||||
// fabric is the operator-assigned name. It is not derived from the GID
|
||||
// prefix: OpenSM defaults to fe80:: on every fabric, so a matching prefix
|
||||
// proves nothing (plan §5.2).
|
||||
string fabric = 1;
|
||||
repeated string hcas = 2;
|
||||
int32 rate_gbps = 3;
|
||||
}
|
||||
|
||||
// Heartbeat says the node is alive and reports anything that changed cheaply.
|
||||
message Heartbeat {
|
||||
string node_id = 1;
|
||||
// gpu_memory_used_mb, indexed by GPU, catches a foreign process appearing on
|
||||
// a card after registration.
|
||||
repeated int32 gpu_memory_used_mb = 2;
|
||||
repeated string new_problems = 3;
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
message ShardStatus {
|
||||
string node_id = 1;
|
||||
string replica_id = 2;
|
||||
string shard_id = 3;
|
||||
State state = 4;
|
||||
string error = 5;
|
||||
int32 exit_code = 6;
|
||||
// last_output carries the tail of the engine log on a failure. When a model
|
||||
// fails to load this is the only diagnosis available.
|
||||
repeated string last_output = 7;
|
||||
|
||||
enum State {
|
||||
STATE_UNSPECIFIED = 0;
|
||||
STATE_PENDING = 1;
|
||||
STATE_STARTING = 2;
|
||||
STATE_LOADING = 3;
|
||||
STATE_READY = 4;
|
||||
STATE_DRAINING = 5;
|
||||
STATE_STOPPED = 6;
|
||||
STATE_FAILED = 7;
|
||||
}
|
||||
}
|
||||
|
||||
// StartShard launches one shard. It is idempotent by shard_id: the control
|
||||
// plane retries, and a repeat must not produce a second process.
|
||||
message StartShard {
|
||||
string replica_id = 1;
|
||||
string shard_id = 2;
|
||||
|
||||
int32 rank = 3;
|
||||
int32 nnodes = 4;
|
||||
string leader_addr = 5;
|
||||
int32 leader_port = 6;
|
||||
|
||||
repeated int32 gpu_indices = 7;
|
||||
// hcas are the adapters of the chosen fabric and nothing else. Handing a
|
||||
// shard adapters from a second fabric lets NCCL pick one whose peer is
|
||||
// unreachable, and the job hangs instead of failing (plan §5.3).
|
||||
repeated string hcas = 8;
|
||||
string mgmt_interface = 9;
|
||||
int32 http_port = 10;
|
||||
|
||||
EngineProfile profile = 11;
|
||||
int64 startup_timeout_ms = 12;
|
||||
|
||||
// pool and fabric are control plane concepts the agent cannot derive. They
|
||||
// are sent so the agent can echo them back in RunningShard after a restart.
|
||||
string pool = 13;
|
||||
string fabric = 14;
|
||||
}
|
||||
|
||||
message EngineProfile {
|
||||
string name = 1;
|
||||
string engine = 2;
|
||||
string gpu_class = 3;
|
||||
int32 tp = 4;
|
||||
int32 pp = 5;
|
||||
int32 gpus_per_node = 6;
|
||||
int32 nodes = 7;
|
||||
repeated string args = 8;
|
||||
map<string, string> env = 9;
|
||||
string model_path = 10;
|
||||
string served_model_name = 11;
|
||||
}
|
||||
|
||||
// StopShard terminates a shard and frees its GPUs.
|
||||
message StopShard {
|
||||
string shard_id = 1;
|
||||
}
|
||||
|
||||
// DrainShard stops routing to a shard while letting in-flight work finish.
|
||||
message DrainShard {
|
||||
string shard_id = 1;
|
||||
}
|
||||
|
||||
// Ack confirms a registration or reports that a command could not be accepted.
|
||||
message Ack {
|
||||
string message_id = 1;
|
||||
bool ok = 2;
|
||||
string error = 3;
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
version: v2
|
||||
managed:
|
||||
enabled: true
|
||||
plugins:
|
||||
- local: protoc-gen-go
|
||||
out: api/gen
|
||||
opt: paths=source_relative
|
||||
- local: protoc-gen-connect-go
|
||||
out: api/gen
|
||||
opt: paths=source_relative
|
||||
inputs:
|
||||
- directory: api/proto
|
||||
@@ -0,0 +1,16 @@
|
||||
version: v2
|
||||
modules:
|
||||
- path: api/proto
|
||||
lint:
|
||||
use:
|
||||
- STANDARD
|
||||
except:
|
||||
# The standard names assume request/response pairs. This is one long-lived
|
||||
# bidirectional stream carrying several kinds of message in each direction,
|
||||
# so "AgentMessage" and "ControlMessage" say what they are and
|
||||
# "ConnectRequest"/"ConnectResponse" would not.
|
||||
- RPC_REQUEST_STANDARD_NAME
|
||||
- RPC_RESPONSE_STANDARD_NAME
|
||||
breaking:
|
||||
use:
|
||||
- FILE
|
||||
@@ -0,0 +1,63 @@
|
||||
# llmgw agent configuration, one per GPU node.
|
||||
#
|
||||
# The agent dials out to the control plane and keeps one stream open. Engines
|
||||
# keep serving if the agent restarts: the router talks to them directly, and
|
||||
# orphan adoption picks them up again on the next start.
|
||||
|
||||
node_id: gpu-a-01 # defaults to the hostname
|
||||
server: "http://llmgw-server.internal:7443"
|
||||
|
||||
# The address engines are reachable on, and the address NCCL bootstraps over.
|
||||
# Never a fabric address: NCCL rendezvouses over TCP before moving to RDMA.
|
||||
mgmt_addr: "10.0.0.11"
|
||||
mgmt_interface: eth0
|
||||
|
||||
labels:
|
||||
rack: r12
|
||||
|
||||
heartbeat_interval: 5s
|
||||
reconnect_base: 1s
|
||||
reconnect_max: 30s
|
||||
|
||||
# Fabric labelling is explicit and there is no way around it. OpenSM's default
|
||||
# subnet prefix is fe80:: on every fabric, so a matching prefix proves nothing
|
||||
# and a wrong guess produces a replica that starts, passes health checks and
|
||||
# then hangs on its first all-reduce (plan §5.2).
|
||||
#
|
||||
# A node with no fabrics is its own single-node island: it can host single-node
|
||||
# replicas and will never be given a multi-node one.
|
||||
fabrics:
|
||||
- name: ib-a
|
||||
hcas: [mlx5_0, mlx5_1, mlx5_2, mlx5_3]
|
||||
rate_gbps: 200
|
||||
|
||||
discovery:
|
||||
# Worth setting where the driver lives outside the default paths, which
|
||||
# containerized driver installs often do.
|
||||
nvidia_smi: ""
|
||||
ib_root: /sys/class/infiniband
|
||||
# NVML is the only source of NVLink topology, which is what keeps a
|
||||
# tensor-parallel group off the PCIe path inside a node. It needs a binary
|
||||
# built with `-tags nvml`; otherwise this silently falls back to nvidia-smi.
|
||||
prefer_nvml: true
|
||||
skip_gpus: false
|
||||
|
||||
engines:
|
||||
vllm:
|
||||
bin: /opt/venvs/vllm/bin/vllm
|
||||
# Which multi-node mechanism works depends on the vLLM version, and the
|
||||
# flags have changed more than once. Check the version you deploy.
|
||||
launch_mode: ray # "ray" or "mp"
|
||||
ray_bin: /opt/venvs/vllm/bin/ray
|
||||
sglang:
|
||||
# SGLang is launched as a module, which is its documented entry point.
|
||||
python: /opt/venvs/sglang/bin/python
|
||||
|
||||
supervisor:
|
||||
work_dir: /var/lib/llmgw-agent
|
||||
# Engine output is the only diagnosis when a model fails to load, so keep a
|
||||
# useful tail of it.
|
||||
log_lines: 512
|
||||
# Engines need time to free GPU memory cleanly before being killed.
|
||||
stop_timeout: 60s
|
||||
health_interval: 2s
|
||||
@@ -0,0 +1,118 @@
|
||||
# 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
|
||||
# Where node agents connect. Separate from the client port so the fleet-facing
|
||||
# surface can be firewalled independently. Agents speak plaintext HTTP/2 today,
|
||||
# so this must stay on a trusted network until mTLS lands (plan §5.6).
|
||||
agent_listen: ":7443"
|
||||
|
||||
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
|
||||
|
||||
# Flow control: decides *when* a request may proceed, before scheduling
|
||||
# decides where it goes. Without it a batch workload starves interactive
|
||||
# traffic, because an engine queue is FIFO and knows nothing about tenants.
|
||||
admission:
|
||||
enabled: true
|
||||
# Until authentication lands, the band comes from a header and is trusted.
|
||||
# Do not expose this port to callers who should not be able to promote
|
||||
# themselves.
|
||||
band_header: x-llmgw-band
|
||||
default_band: gold
|
||||
# A band waits once pool saturation passes its ceiling, so the cluster
|
||||
# degrades from the bottom up. A ceiling of 1.0 means "flow up to the
|
||||
# healthy limit"; use a value above 1 for a band that must never wait.
|
||||
bands:
|
||||
- { name: platinum, priority: 400, saturation_ceiling: 1.00 }
|
||||
- { name: gold, priority: 300, saturation_ceiling: 0.85 }
|
||||
- { name: silver, priority: 200, saturation_ceiling: 0.60 }
|
||||
- { name: bronze, priority: 100, saturation_ceiling: 0.30 }
|
||||
# Drop a request that has waited this long: the client has usually given up,
|
||||
# and sending the GPU work nobody is waiting for helps no one.
|
||||
queue_ttl: 30s
|
||||
# Reject immediately once a band's queue is this deep, rather than letting
|
||||
# it accumulate silently.
|
||||
queue_capacity_per_band: 2000
|
||||
interval: 10ms
|
||||
|
||||
# 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
|
||||
|
||||
# How aggressively the control plane reconciles.
|
||||
control:
|
||||
interval: 5s
|
||||
base_backoff: 10s
|
||||
max_backoff: 10m
|
||||
# Node expiry is staged deliberately: a network blip must not tear down
|
||||
# replicas that are still serving, and rebuilding a multi-node replica costs
|
||||
# tens of minutes of model loading.
|
||||
unreachable_after: 15s
|
||||
lost_after: 60s
|
||||
port_range: [20000, 20999]
|
||||
mgmt_interface: eth0
|
||||
|
||||
models:
|
||||
- name: llama-3.3-70b
|
||||
aliases: [llama]
|
||||
fallbacks: [qwen2.5-32b]
|
||||
routing_profile: chat
|
||||
|
||||
- name: qwen2.5-32b
|
||||
aliases: [qwen]
|
||||
routing_profile: batch
|
||||
|
||||
# A static replica list, used only when no model declares a profile. That is the
|
||||
# router-only mode for testing against sim-engines; with the control plane
|
||||
# running, agents provide the replicas instead.
|
||||
#
|
||||
# endpoints:
|
||||
# - id: sim-1
|
||||
# url: http://127.0.0.1:8001
|
||||
# models: [llama-3.3-70b]
|
||||
# pool: local
|
||||
# weight: 1
|
||||
@@ -3,15 +3,34 @@ module gitea.mynodes.xyz/ilker/llmgw
|
||||
go 1.26.3
|
||||
|
||||
require (
|
||||
connectrpc.com/connect v1.21.0
|
||||
github.com/NVIDIA/go-nvml v0.13.4-0
|
||||
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
|
||||
golang.org/x/sync v0.23.0
|
||||
google.golang.org/protobuf v1.36.12
|
||||
)
|
||||
|
||||
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
|
||||
google.golang.org/protobuf v1.36.12 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
golang.org/x/net v0.59.0 // indirect
|
||||
golang.org/x/sys v0.48.0 // indirect
|
||||
golang.org/x/text v0.42.0 // indirect
|
||||
)
|
||||
@@ -1,18 +1,48 @@
|
||||
connectrpc.com/connect v1.21.0 h1:LhqSJt7jHf5NJBo9Jq/t/9FjcYAideif0mg+qe2jCUs=
|
||||
connectrpc.com/connect v1.21.0/go.mod h1:A2ygJrukXwWy32vkCAAHNVguZrqZ+jeZ9rGRnGR4dN4=
|
||||
github.com/NVIDIA/go-nvml v0.13.4-0 h1:o3jp9u2x1R9ShFE3v+Aesp55XOSIQFMJz/VGNUcJaNE=
|
||||
github.com/NVIDIA/go-nvml v0.13.4-0/go.mod h1:id63qwpoDWpFXwnwM6psDCSqW4BmNu6mWpr4YeQtPGo=
|
||||
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 +50,21 @@ 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/net v0.59.0 h1:5zfYln+w5XCxwrnMMJPufRgNoXEaGxl0wo5GqPXyues=
|
||||
golang.org/x/net v0.59.0/go.mod h1:2DA/G1UfVbCpQPeWTmMPGY7Cs2PkBkwu743bVX5PIVg=
|
||||
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.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
|
||||
golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og=
|
||||
golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI=
|
||||
golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E=
|
||||
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,558 @@
|
||||
// Package agent is the node side of the control plane stream (plan §5.6).
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"math/rand/v2"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// Config configures the agent.
|
||||
type Config struct {
|
||||
NodeID string `koanf:"node_id"`
|
||||
Server string `koanf:"server"`
|
||||
Hostname string `koanf:"hostname"`
|
||||
MgmtAddr string `koanf:"mgmt_addr"`
|
||||
Labels map[string]string `koanf:"labels"`
|
||||
|
||||
// MgmtInterface is the ethernet interface NCCL bootstraps over.
|
||||
MgmtInterface string `koanf:"mgmt_interface"`
|
||||
// Fabrics labels this node's HCAs. It is explicit on purpose: OpenSM's
|
||||
// default subnet prefix is the same on every fabric, so nothing can be
|
||||
// inferred from the wire (plan §5.2).
|
||||
Fabrics []inventory.FabricAttachment `koanf:"fabrics"`
|
||||
|
||||
HeartbeatInterval time.Duration `koanf:"heartbeat_interval"`
|
||||
// ReconnectBase is the first reconnect delay; it backs off from there.
|
||||
ReconnectBase time.Duration `koanf:"reconnect_base"`
|
||||
ReconnectMax time.Duration `koanf:"reconnect_max"`
|
||||
|
||||
Discovery discovery.Options `koanf:"discovery"`
|
||||
Engines engines.Config `koanf:"engines"`
|
||||
Superv supervisor.Config `koanf:"supervisor"`
|
||||
|
||||
Version string
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
MgmtInterface: "eth0",
|
||||
HeartbeatInterval: 5 * time.Second,
|
||||
ReconnectBase: time.Second,
|
||||
ReconnectMax: 30 * time.Second,
|
||||
Superv: supervisor.DefaultConfig(),
|
||||
// NVML gives NVLink topology, so prefer it when this build has it.
|
||||
Discovery: discovery.Options{PreferNVML: true},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() error {
|
||||
d := DefaultConfig()
|
||||
if c.NodeID == "" {
|
||||
return errors.New("node_id is required")
|
||||
}
|
||||
if c.Server == "" {
|
||||
return errors.New("server is required")
|
||||
}
|
||||
if c.MgmtInterface == "" {
|
||||
c.MgmtInterface = d.MgmtInterface
|
||||
}
|
||||
if c.HeartbeatInterval <= 0 {
|
||||
c.HeartbeatInterval = d.HeartbeatInterval
|
||||
}
|
||||
if c.ReconnectBase <= 0 {
|
||||
c.ReconnectBase = d.ReconnectBase
|
||||
}
|
||||
if c.ReconnectMax < c.ReconnectBase {
|
||||
c.ReconnectMax = max(c.ReconnectBase, d.ReconnectMax)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Agent runs on a GPU node.
|
||||
type Agent struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
sup *supervisor.Supervisor
|
||||
reg *engines.Registry
|
||||
http connectHTTPClient
|
||||
|
||||
mu sync.Mutex
|
||||
inv discovery.Inventory
|
||||
// shardReplica remembers which replica a shard belongs to, so a status
|
||||
// report can name both. The control plane keys on the replica.
|
||||
shardReplica map[string]string
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
// hand in an httptest server's client.
|
||||
type connectHTTPClient interface {
|
||||
Do(*http.Request) (*http.Response, error)
|
||||
}
|
||||
|
||||
// New builds an agent.
|
||||
func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, error) {
|
||||
if err := cfg.normalize(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if httpClient == nil {
|
||||
httpClient = http.DefaultClient
|
||||
}
|
||||
|
||||
health := supervisor.NewHTTPHealth(3 * time.Second)
|
||||
sup, err := supervisor.New(cfg.Superv, health, log)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Agent{
|
||||
cfg: cfg, log: log, sup: sup,
|
||||
reg: engines.NewRegistry(cfg.Engines),
|
||||
http: httpClient,
|
||||
shardReplica: map[string]string{},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Supervisor exposes the process supervisor, for tests and `ctl logs`.
|
||||
func (a *Agent) Supervisor() *supervisor.Supervisor { return a.sup }
|
||||
|
||||
// Run connects to the control plane and stays connected until ctx ends.
|
||||
//
|
||||
// Reconnecting forever is the right behaviour for a node agent: a control plane
|
||||
// restart or a network partition must not require touching every GPU node, and
|
||||
// the engines keep serving traffic throughout because the router talks to them
|
||||
// directly.
|
||||
func (a *Agent) Run(ctx context.Context) error {
|
||||
if err := a.scan(ctx); err != nil {
|
||||
// Not fatal: a node whose GPUs cannot be read should still register, so
|
||||
// an operator can see it and why it is useless.
|
||||
a.log.Error("hardware discovery failed", "err", err)
|
||||
}
|
||||
a.recoverOrphans()
|
||||
|
||||
attempt := 0
|
||||
for ctx.Err() == nil {
|
||||
err := a.session(ctx)
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
attempt++
|
||||
delay := a.reconnectDelay(attempt)
|
||||
a.log.Warn("disconnected from the control plane, retrying",
|
||||
"err", err, "in", delay.Round(time.Millisecond), "attempt", attempt)
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(delay):
|
||||
}
|
||||
if err == nil {
|
||||
attempt = 0
|
||||
}
|
||||
}
|
||||
|
||||
// Engines keep running when the agent stops on purpose: an agent restart
|
||||
// must not interrupt served traffic. Orphan adoption picks them up again.
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
func (a *Agent) reconnectDelay(attempt int) time.Duration {
|
||||
d := a.cfg.ReconnectBase << min(attempt-1, 16)
|
||||
if d > a.cfg.ReconnectMax || d <= 0 {
|
||||
d = a.cfg.ReconnectMax
|
||||
}
|
||||
// Jitter, so a fleet does not reconnect in lockstep after a control plane
|
||||
// restart and knock it over again.
|
||||
return d/2 + time.Duration(rand.Int64N(int64(d/2))) //nolint:gosec // jitter
|
||||
}
|
||||
|
||||
// recoverOrphans takes back engines a previous agent left running.
|
||||
//
|
||||
// Engines deliberately outlive their agent — the router talks to them directly,
|
||||
// so an agent restart must not interrupt served traffic. The cost is that they
|
||||
// have to be found again, and anything that cannot be identified has to be
|
||||
// killed: a process holding GPUs that nothing can describe is worse than no
|
||||
// process at all, because placement will hand those GPUs out again.
|
||||
func (a *Agent) recoverOrphans() {
|
||||
orphans, err := a.sup.FindOrphans()
|
||||
if err != nil {
|
||||
a.log.Warn("scanning for leftover engines failed", "err", err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, o := range orphans {
|
||||
switch {
|
||||
case !o.Alive:
|
||||
// The process is gone; only its state file remains.
|
||||
if err := a.sup.KillOrphan(o); err != nil {
|
||||
a.log.Warn("cleaning up a dead shard's state failed", "shard", o.ShardID, "err", err)
|
||||
}
|
||||
case o.ReplicaID == "":
|
||||
// Alive, but its state file could not be read, so it cannot be
|
||||
// described to the control plane or matched to a replica.
|
||||
a.log.Warn("killing an unidentifiable engine", "shard", o.ShardID, "pid", o.PID)
|
||||
if err := a.sup.KillOrphan(o); err != nil {
|
||||
a.log.Warn("killing it failed", "shard", o.ShardID, "err", err)
|
||||
}
|
||||
default:
|
||||
sh, err := a.sup.Adopt(o)
|
||||
if err != nil {
|
||||
a.log.Warn("adopting a shard failed", "shard", o.ShardID, "err", err)
|
||||
continue
|
||||
}
|
||||
a.mu.Lock()
|
||||
a.shardReplica[o.ShardID] = o.ReplicaID
|
||||
a.mu.Unlock()
|
||||
a.log.Info("adopted a running engine",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID, "pid", o.PID)
|
||||
_ = sh
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runningShards describes what this agent is supervising, for registration.
|
||||
func (a *Agent) runningShards() []*llmgwv1.RunningShard {
|
||||
var out []*llmgwv1.RunningShard
|
||||
for _, st := range a.sup.Statuses() {
|
||||
if st.State.Terminal() {
|
||||
continue
|
||||
}
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
spec := sh.Spec
|
||||
out = append(out, &llmgwv1.RunningShard{
|
||||
ReplicaId: spec.ReplicaID,
|
||||
ShardId: st.ID,
|
||||
Rank: int32(spec.Rank),
|
||||
Nnodes: int32(spec.NNodes),
|
||||
GpuIndices: int32Slice(spec.GPUIndices),
|
||||
Hcas: spec.HCAs,
|
||||
HttpPort: int32(spec.HTTPPort),
|
||||
State: agentsvc.StateToWire(st.State),
|
||||
Model: spec.Profile.ServedModelName,
|
||||
Profile: spec.Profile.Name,
|
||||
Pool: spec.Pool,
|
||||
Fabric: spec.Fabric,
|
||||
LeaderAddr: a.cfg.MgmtAddr,
|
||||
LeaderPort: int32(spec.LeaderPort),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func int32Slice(v []int) []int32 {
|
||||
if len(v) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]int32, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = int32(x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// scan refreshes the node's hardware inventory.
|
||||
func (a *Agent) scan(ctx context.Context) error {
|
||||
inv, err := discovery.Discover(ctx, a.cfg.Discovery, a.log)
|
||||
a.mu.Lock()
|
||||
a.inv = inv
|
||||
a.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
|
||||
// session runs one connection until it breaks.
|
||||
func (a *Agent) session(ctx context.Context) error {
|
||||
client := llmgwv1connect.NewAgentServiceClient(a.http, a.cfg.Server)
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
stream := client.Connect(ctx)
|
||||
defer func() { _ = stream.CloseRequest() }()
|
||||
|
||||
a.mu.Lock()
|
||||
inv := a.inv
|
||||
a.mu.Unlock()
|
||||
|
||||
reg := agentsvc.RegisterFromInventory(
|
||||
a.cfg.NodeID, a.cfg.Hostname, a.cfg.MgmtAddr, a.cfg.Version,
|
||||
a.cfg.Labels, inv, a.cfg.Fabrics)
|
||||
// Every registration reports what is already running, not just the first
|
||||
// one: a reconnect may be a control plane restart, and it has to recover
|
||||
// the same way (plan §6.5).
|
||||
reg.RunningShards = a.runningShards()
|
||||
|
||||
if err := stream.Send(&llmgwv1.AgentMessage{
|
||||
Msg: &llmgwv1.AgentMessage_Register{Register: reg},
|
||||
}); err != nil {
|
||||
return fmt.Errorf("registering: %w", err)
|
||||
}
|
||||
|
||||
// Sends are serialized through this channel: connect streams are not safe
|
||||
// for concurrent writes, and heartbeats race with status reports.
|
||||
out := make(chan *llmgwv1.AgentMessage, 32)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
a.sendLoop(ctx, stream, out)
|
||||
}()
|
||||
defer wg.Wait()
|
||||
|
||||
for {
|
||||
msg, err := stream.Receive()
|
||||
if err != nil {
|
||||
cancel()
|
||||
if errors.Is(err, io.EOF) || ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
a.handle(ctx, msg, out)
|
||||
}
|
||||
}
|
||||
|
||||
type sender interface {
|
||||
Send(*llmgwv1.AgentMessage) error
|
||||
}
|
||||
|
||||
func (a *Agent) sendLoop(ctx context.Context, s sender, out <-chan *llmgwv1.AgentMessage) {
|
||||
t := time.NewTicker(a.cfg.HeartbeatInterval)
|
||||
defer t.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case msg := <-out:
|
||||
if err := s.Send(msg); err != nil {
|
||||
return
|
||||
}
|
||||
case <-t.C:
|
||||
if err := s.Send(a.heartbeat()); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) heartbeat() *llmgwv1.AgentMessage {
|
||||
a.mu.Lock()
|
||||
inv := a.inv
|
||||
a.mu.Unlock()
|
||||
|
||||
hb := &llmgwv1.Heartbeat{NodeId: a.cfg.NodeID}
|
||||
for _, g := range inv.GPUs {
|
||||
hb.GpuMemoryUsedMb = append(hb.GpuMemoryUsedMb, int32(g.MemoryUsedMB))
|
||||
}
|
||||
return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Heartbeat{Heartbeat: hb}}
|
||||
}
|
||||
|
||||
func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out chan<- *llmgwv1.AgentMessage) {
|
||||
switch m := msg.Msg.(type) {
|
||||
case *llmgwv1.ControlMessage_Start:
|
||||
go a.startShard(ctx, m.Start, out)
|
||||
case *llmgwv1.ControlMessage_Stop:
|
||||
go a.stopShard(ctx, m.Stop.GetShardId(), out)
|
||||
case *llmgwv1.ControlMessage_Drain:
|
||||
if err := a.sup.Drain(m.Drain.GetShardId()); err != nil {
|
||||
a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err)
|
||||
}
|
||||
case *llmgwv1.ControlMessage_Ack:
|
||||
if !m.Ack.GetOk() {
|
||||
a.log.Warn("control plane rejected a message",
|
||||
"id", m.Ack.GetMessageId(), "err", m.Ack.GetError())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// startShard launches a shard and reports its progress.
|
||||
//
|
||||
// It runs in its own goroutine because model loading takes minutes, and the
|
||||
// stream has to keep carrying heartbeats and other commands meanwhile.
|
||||
func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<- *llmgwv1.AgentMessage) {
|
||||
shardID, replicaID := m.GetShardId(), m.GetReplicaId()
|
||||
|
||||
a.mu.Lock()
|
||||
a.shardReplica[shardID] = replicaID
|
||||
a.mu.Unlock()
|
||||
|
||||
report := func(st supervisor.State, errMsg string, tail []string) {
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
|
||||
State: agentsvc.StateToWire(st), Error: errMsg, LastOutput: tail,
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
spec, timeout := agentsvc.ShardSpecFromStart(m)
|
||||
|
||||
eng, err := a.reg.Get(spec.Profile.Engine)
|
||||
if err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
|
||||
a.mu.Lock()
|
||||
inv := a.inv
|
||||
a.mu.Unlock()
|
||||
|
||||
// Validate before launching. An engine that refuses a profile does so after
|
||||
// minutes of loading; catching it here costs nothing.
|
||||
if err := eng.Validate(spec.Profile, gpuInfo(inv, spec.GPUIndices)); err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
|
||||
cmds, err := eng.Commands(spec)
|
||||
if err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
|
||||
sh, err := a.sup.Start(ctx, shardID, spec, cmds)
|
||||
if err != nil {
|
||||
if errors.Is(err, supervisor.ErrAlreadyRunning) {
|
||||
// Commands are idempotent: the control plane retries, and a repeat
|
||||
// must not produce a second process (plan §5.6).
|
||||
report(sh.Status().State, "", nil)
|
||||
return
|
||||
}
|
||||
report(supervisor.Failed, err.Error(), sh.Logs(10))
|
||||
return
|
||||
}
|
||||
report(supervisor.Starting, "", nil)
|
||||
|
||||
// Only a rank that serves HTTP can be health-checked. A worker is judged by
|
||||
// staying alive, and the replica's real readiness is the leader answering,
|
||||
// which only happens once every worker has joined (plan §11, phase 4).
|
||||
if eng.ServesHTTP(spec) {
|
||||
base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort)
|
||||
err = a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout)
|
||||
} else {
|
||||
err = a.sup.WaitAlive(ctx, sh)
|
||||
}
|
||||
if err != nil {
|
||||
report(supervisor.Failed, err.Error(), sh.Logs(20))
|
||||
return
|
||||
}
|
||||
report(supervisor.Ready, "", nil)
|
||||
|
||||
// Watch for the engine dying later. Nothing else would notice: the router
|
||||
// would simply see failing requests.
|
||||
go a.watch(ctx, sh, replicaID, shardID, out)
|
||||
}
|
||||
|
||||
// watch reports a shard that exits after it was ready.
|
||||
func (a *Agent) watch(ctx context.Context, sh *supervisor.Shard, replicaID, shardID string, out chan<- *llmgwv1.AgentMessage) {
|
||||
t := time.NewTicker(time.Second)
|
||||
defer t.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
st := sh.Status()
|
||||
if !st.State.Terminal() {
|
||||
continue
|
||||
}
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
|
||||
State: agentsvc.StateToWire(st.State), Error: st.Err,
|
||||
ExitCode: int32(st.ExitCode), LastOutput: sh.Logs(20),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) stopShard(ctx context.Context, shardID string, out chan<- *llmgwv1.AgentMessage) {
|
||||
a.mu.Lock()
|
||||
replicaID := a.shardReplica[shardID]
|
||||
delete(a.shardReplica, shardID)
|
||||
a.mu.Unlock()
|
||||
|
||||
if err := a.sup.Stop(ctx, shardID); err != nil {
|
||||
a.log.Warn("stopping a shard failed", "shard", shardID, "err", err)
|
||||
}
|
||||
a.sup.Remove(shardID)
|
||||
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
|
||||
State: agentsvc.StateToWire(supervisor.Stopped),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
// gpuInfo narrows the inventory to the GPUs a shard will use.
|
||||
func gpuInfo(inv discovery.Inventory, indices []int) []engines.GPUInfo {
|
||||
want := map[int]struct{}{}
|
||||
for _, i := range indices {
|
||||
want[i] = struct{}{}
|
||||
}
|
||||
var out []engines.GPUInfo
|
||||
for _, g := range inv.GPUs {
|
||||
if _, ok := want[g.Index]; !ok {
|
||||
continue
|
||||
}
|
||||
out = append(out, engines.GPUInfo{
|
||||
Index: g.Index, Name: g.Name,
|
||||
ComputeCapability: g.ComputeCapability, MemoryTotalMB: g.MemoryTotalMB,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// HTTP2Client returns the client the agent should use to reach the control
|
||||
// plane over plaintext HTTP/2.
|
||||
//
|
||||
// connect's bidirectional streams require HTTP/2, and an agent talking to a
|
||||
// control plane over a trusted network has no TLS to negotiate it with. Go's
|
||||
// default transport will not speak HTTP/2 without TLS, so it has to be asked
|
||||
// explicitly.
|
||||
//
|
||||
// mTLS is plan §5.6's target and replaces this; until then the stream must not
|
||||
// cross an untrusted network.
|
||||
func HTTP2Client(timeout time.Duration) *http.Client {
|
||||
return &http.Client{
|
||||
Transport: &http2.Transport{
|
||||
AllowHTTP: true,
|
||||
DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
|
||||
return (&net.Dialer{Timeout: timeout}).DialContext(ctx, network, addr)
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,600 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/http2/h2c"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// statusRecorder captures what the control plane hears about shards, and what
|
||||
// it is asked to adopt.
|
||||
type statusRecorder struct {
|
||||
ch chan status
|
||||
|
||||
mu sync.Mutex
|
||||
adopted []reconciler.AdoptedShard
|
||||
unclaim []string // shard ids to report back as unclaimed
|
||||
}
|
||||
|
||||
type status struct {
|
||||
replicaID, shardID string
|
||||
state supervisor.State
|
||||
err string
|
||||
}
|
||||
|
||||
func newRecorder() *statusRecorder { return &statusRecorder{ch: make(chan status, 128)} }
|
||||
|
||||
func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
|
||||
select {
|
||||
case r.ch <- status{replicaID, shardID, st, errMsg}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (r *statusRecorder) Adopt(shards []reconciler.AdoptedShard) []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.adopted = append(r.adopted, shards...)
|
||||
|
||||
var unclaimed []reconciler.AdoptedShard
|
||||
for _, s := range shards {
|
||||
if slices.Contains(r.unclaim, s.ShardID) {
|
||||
unclaimed = append(unclaimed, s)
|
||||
}
|
||||
}
|
||||
return unclaimed
|
||||
}
|
||||
|
||||
func (r *statusRecorder) adoptedShards() []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return slices.Clone(r.adopted)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) markUnclaimed(ids ...string) {
|
||||
r.mu.Lock()
|
||||
r.unclaim = append(r.unclaim, ids...)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *statusRecorder) await(t *testing.T, limit time.Duration, pred func(status) bool) status {
|
||||
t.Helper()
|
||||
deadline := time.After(limit)
|
||||
var seen []status
|
||||
for {
|
||||
select {
|
||||
case s := <-r.ch:
|
||||
seen = append(seen, s)
|
||||
if pred(s) {
|
||||
return s
|
||||
}
|
||||
case <-deadline:
|
||||
t.Fatalf("no matching status within %v; saw %+v", limit, seen)
|
||||
return status{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// fakeEngineScript writes a script that stands in for vLLM.
|
||||
//
|
||||
// It parses the real --port flag the vLLM adapter generates rather than being
|
||||
// told out of band, so the adapter's argument construction is exercised too.
|
||||
func fakeEngineScript(t *testing.T, body string) string {
|
||||
t.Helper()
|
||||
path := filepath.Join(t.TempDir(), "fake-vllm")
|
||||
script := `#!/bin/sh
|
||||
port=""
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--port) port="$2"; shift 2 ;;
|
||||
*) shift ;;
|
||||
esac
|
||||
done
|
||||
` + body + "\n"
|
||||
if err := os.WriteFile(path, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// servingEngine answers /health and /v1/models, which is what readiness needs.
|
||||
func servingEngine(t *testing.T) string {
|
||||
t.Helper()
|
||||
py, err := exec.LookPath("python3")
|
||||
if err != nil {
|
||||
t.Skip("python3 is required for the fake engine")
|
||||
}
|
||||
return fakeEngineScript(t, `exec `+py+` -c '
|
||||
import http.server, json, sys
|
||||
port = int(sys.argv[1])
|
||||
class H(http.server.BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
if self.path == "/health":
|
||||
self.send_response(200); self.end_headers(); return
|
||||
if self.path == "/v1/models":
|
||||
b = json.dumps({"object":"list","data":[{"id":"test-model"}]}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("content-type","application/json")
|
||||
self.send_header("content-length", str(len(b)))
|
||||
self.end_headers(); self.wfile.write(b); return
|
||||
self.send_response(404); self.end_headers()
|
||||
def log_message(self, *a): pass
|
||||
http.server.HTTPServer(("127.0.0.1", port), H).serve_forever()
|
||||
' "$port"`)
|
||||
}
|
||||
|
||||
type harness struct {
|
||||
svc *agentsvc.Service
|
||||
inv *inventory.Store
|
||||
rec *statusRecorder
|
||||
srv *httptest.Server
|
||||
ag *Agent
|
||||
ln *dropListener
|
||||
// streams counts how many times an agent has opened the Connect stream,
|
||||
// which is the only honest way to tell a reconnect from a stale entry.
|
||||
streams *atomic.Int64
|
||||
}
|
||||
|
||||
// dropListener tracks accepted connections so a test can sever them.
|
||||
//
|
||||
// httptest's CloseClientConnections does not work here: h2c with prior
|
||||
// knowledge hijacks the connection, and httptest stops tracking hijacked
|
||||
// connections. A test relying on it would pass without ever reconnecting.
|
||||
type dropListener struct {
|
||||
net.Listener
|
||||
mu sync.Mutex
|
||||
conns []net.Conn
|
||||
}
|
||||
|
||||
func (l *dropListener) Accept() (net.Conn, error) {
|
||||
c, err := l.Listener.Accept()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
l.mu.Lock()
|
||||
l.conns = append(l.conns, c)
|
||||
l.mu.Unlock()
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// drop severs every connection, as a control plane restart would.
|
||||
func (l *dropListener) drop() {
|
||||
l.mu.Lock()
|
||||
conns := l.conns
|
||||
l.conns = nil
|
||||
l.mu.Unlock()
|
||||
for _, c := range conns {
|
||||
_ = c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// start wires a real agent to a real service over h2c, so this exercises the
|
||||
// protocol rather than two structs calling each other.
|
||||
func start(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
return startWith(t, slog.New(slog.DiscardHandler), cfgFns...)
|
||||
}
|
||||
|
||||
// startLogged is start with logging on, for diagnosing protocol behaviour.
|
||||
func startLogged(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
return startWith(t, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug})), cfgFns...)
|
||||
}
|
||||
|
||||
func startWith(t *testing.T, log *slog.Logger, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
|
||||
inv := inventory.NewStore()
|
||||
rec := newRecorder()
|
||||
svc := agentsvc.New(inv, rec, log)
|
||||
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
|
||||
var streams atomic.Int64
|
||||
counting := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
streams.Add(1)
|
||||
mux.ServeHTTP(w, r)
|
||||
})
|
||||
|
||||
raw, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ln := &dropListener{Listener: raw}
|
||||
// connect's bidirectional streams need HTTP/2; h2c provides it without TLS.
|
||||
srv := httptest.NewUnstartedServer(h2c.NewHandler(counting, &http2.Server{}))
|
||||
srv.Listener = ln
|
||||
srv.Start()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "n1"
|
||||
cfg.Server = srv.URL
|
||||
cfg.MgmtAddr = "127.0.0.1"
|
||||
cfg.HeartbeatInterval = 50 * time.Millisecond
|
||||
cfg.ReconnectBase = 20 * time.Millisecond
|
||||
cfg.ReconnectMax = 100 * time.Millisecond
|
||||
cfg.Discovery.SkipGPUs = true
|
||||
cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
|
||||
cfg.Superv = supervisor.Config{
|
||||
WorkDir: t.TempDir(), LogLines: 64,
|
||||
StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond,
|
||||
}
|
||||
for _, f := range cfgFns {
|
||||
f(&cfg)
|
||||
}
|
||||
|
||||
client := &http.Client{Transport: &http2.Transport{
|
||||
AllowHTTP: true,
|
||||
DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
|
||||
return (&net.Dialer{}).DialContext(ctx, network, addr)
|
||||
},
|
||||
}}
|
||||
|
||||
ag, err := New(cfg, client, log)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() { _ = ag.Run(ctx) }()
|
||||
|
||||
h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag, ln: ln, streams: &streams}
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
ag.Supervisor().StopAll(context.Background())
|
||||
srv.Close()
|
||||
})
|
||||
h.waitRegistered(t)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *harness) waitRegistered(t *testing.T) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, ok := h.inv.Get("n1"); ok && len(h.svc.ConnectedNodes()) == 1 {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the agent never registered")
|
||||
}
|
||||
|
||||
// shardCommand builds a single-node vLLM shard on a free port.
|
||||
func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand {
|
||||
t.Helper()
|
||||
return reconciler.ShardCommand{
|
||||
ReplicaID: "rep-1", ShardID: shardID, NodeID: "n1",
|
||||
Spec: engines.ShardSpec{
|
||||
ReplicaID: "rep-1", ShardID: shardID, Rank: 0, NNodes: 1,
|
||||
GPUIndices: []int{0}, HTTPPort: freePort(t),
|
||||
Profile: engines.Profile{
|
||||
Name: "p", Engine: "vllm", TP: 1,
|
||||
ModelPath: "/models/x", ServedModelName: "test-model",
|
||||
},
|
||||
},
|
||||
StartupTimeout: 15 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
func freePort(t *testing.T) int {
|
||||
t.Helper()
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ln.Close()
|
||||
return ln.Addr().(*net.TCPAddr).Port
|
||||
}
|
||||
|
||||
func TestRegistrationCarriesInventory(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Hostname = "gpu-a-01"
|
||||
c.Labels = map[string]string{"rack": "r12"}
|
||||
c.Fabrics = []inventory.FabricAttachment{
|
||||
{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200},
|
||||
}
|
||||
})
|
||||
|
||||
n, ok := h.inv.Get("n1")
|
||||
if !ok {
|
||||
t.Fatal("node missing from inventory")
|
||||
}
|
||||
if n.Hostname != "gpu-a-01" || n.MgmtAddr != "127.0.0.1" {
|
||||
t.Errorf("identity wrong: %+v", n)
|
||||
}
|
||||
if n.Labels["rack"] != "r12" {
|
||||
t.Errorf("labels lost: %v", n.Labels)
|
||||
}
|
||||
// The fabric label cannot be inferred from the wire, so losing it in
|
||||
// transit would silently disable multi-node placement.
|
||||
if !n.HasFabric("ib-a") {
|
||||
t.Fatalf("fabric attachment lost: %+v", n.Fabrics)
|
||||
}
|
||||
if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
|
||||
t.Errorf("HCAs=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeartbeatKeepsNodeAlive(t *testing.T) {
|
||||
h := start(t)
|
||||
|
||||
n, _ := h.inv.Get("n1")
|
||||
first := n.LastSeen
|
||||
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
n, _ = h.inv.Get("n1")
|
||||
if !n.LastSeen.After(first) {
|
||||
t.Error("heartbeats are not refreshing LastSeen; the node would be declared lost")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartAndStopShardEndToEnd(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = servingEngine(t)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-a")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got := h.rec.await(t, 30*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-a" && s.state == supervisor.Ready
|
||||
})
|
||||
if got.replicaID != "rep-1" {
|
||||
t.Errorf("replica id lost in transit: %q", got.replicaID)
|
||||
}
|
||||
|
||||
if err := h.svc.StopShard(t.Context(), "n1", "shard-a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-a" && s.state == supervisor.Stopped
|
||||
})
|
||||
}
|
||||
|
||||
func TestFailedShardCarriesDiagnosis(t *testing.T) {
|
||||
// When a model fails to load the engine's output is the only diagnosis, so
|
||||
// it has to survive the trip to the control plane.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `echo 'CUDA out of memory' >&2; exit 1`)
|
||||
})
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), shardCommand(t, "shard-bad")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := h.rec.await(t, 30*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-bad" && s.state == supervisor.Failed
|
||||
})
|
||||
if got.err == "" {
|
||||
t.Error("the failure carried no error message")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineThatNeverBecomesReady(t *testing.T) {
|
||||
// A process that runs but never serves must time out rather than hang the
|
||||
// replica forever.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-stuck")
|
||||
cmd.StartupTimeout = 500 * time.Millisecond
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 20*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-stuck" && s.state == supervisor.Failed
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidationRejectsBeforeLaunch(t *testing.T) {
|
||||
// An engine that refuses a profile does so after minutes of loading;
|
||||
// catching it at dispatch costs nothing.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = servingEngine(t)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-invalid")
|
||||
cmd.Spec.Profile.ModelPath = "" // validation requires one
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-invalid" && s.state == supervisor.Failed
|
||||
})
|
||||
if got.err == "" {
|
||||
t.Error("validation failure carried no reason")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchToUnconnectedNode(t *testing.T) {
|
||||
h := start(t)
|
||||
|
||||
err := h.svc.StartShard(t.Context(), reconciler.ShardCommand{NodeID: "nowhere", ShardID: "x"})
|
||||
if err == nil {
|
||||
t.Fatal("dispatching to an unconnected node succeeded")
|
||||
}
|
||||
if !errors.Is(err, agentsvc.ErrNotConnected) {
|
||||
t.Errorf("err=%v, expected ErrNotConnected", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentReconnects(t *testing.T) {
|
||||
// A control plane restart must not require touching every GPU node.
|
||||
h := start(t)
|
||||
|
||||
before := h.streams.Load()
|
||||
h.ln.drop()
|
||||
|
||||
// Counting streams, not connected nodes: a stale entry in the connection
|
||||
// map would make a node look connected without any reconnect having
|
||||
// happened, and an assertion on it would pass vacuously.
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if h.streams.Load() > before && len(h.svc.ConnectedNodes()) == 1 {
|
||||
if err := h.svc.DrainShard(t.Context(), "n1", "nothing"); err != nil {
|
||||
t.Fatalf("the reconnected stream does not accept commands: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("the agent never reconnected (streams %d -> %d)", before, h.streams.Load())
|
||||
}
|
||||
|
||||
func TestIdempotentStart(t *testing.T) {
|
||||
// The control plane retries commands; a repeat must not start a second
|
||||
// process (plan §5.6).
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 30`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-idem")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-idem" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
sup := h.ag.Supervisor()
|
||||
sh, err := sup.Get("shard-idem")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
|
||||
sh2, err := sup.Get("shard-idem")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sh2.Status().PID != pid {
|
||||
t.Errorf("a second process was started: pid %d then %d", pid, sh2.Status().PID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunningShardsReportedOnReconnect covers the control plane restart hole
|
||||
// (plan §6.5): a reconnecting agent has to describe what it is already running,
|
||||
// or the control plane concludes the cluster is empty and places duplicates
|
||||
// while the originals keep holding GPUs.
|
||||
func TestRunningShardsReportedOnReconnect(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-live")
|
||||
cmd.Pool = "ib-a/h100-80g"
|
||||
cmd.Fabric = "ib-a"
|
||||
cmd.Spec.HCAs = []string{"mlx5_0"}
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-live" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
// Drop the stream, as a control plane restart would.
|
||||
h.ln.drop()
|
||||
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if got := h.rec.adoptedShards(); len(got) > 0 {
|
||||
a := got[0]
|
||||
if a.ShardID != "shard-live" || a.ReplicaID != "rep-1" {
|
||||
t.Fatalf("wrong shard reported: %+v", a)
|
||||
}
|
||||
// Everything the control plane needs to rebuild the replica.
|
||||
if a.Model != "test-model" || a.Profile != "p" {
|
||||
t.Errorf("catalog identity lost: model=%q profile=%q", a.Model, a.Profile)
|
||||
}
|
||||
if a.Pool != "ib-a/h100-80g" || a.Fabric != "ib-a" {
|
||||
t.Errorf("pool/fabric lost: %q %q", a.Pool, a.Fabric)
|
||||
}
|
||||
if !slices.Equal(a.HCAs, []string{"mlx5_0"}) {
|
||||
t.Errorf("HCAs lost: %v", a.HCAs)
|
||||
}
|
||||
if a.Endpoint == "" {
|
||||
t.Error("rank 0 reported no endpoint; the router would have nowhere to send")
|
||||
}
|
||||
if len(a.GPUIndices) == 0 {
|
||||
t.Error("GPU indices lost; the inventory could not reclaim them")
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the reconnecting agent never reported its running shard")
|
||||
}
|
||||
|
||||
// TestUnclaimedShardIsStopped covers the other half: a shard whose model left
|
||||
// the catalog is holding GPUs nothing will ever route to.
|
||||
func TestUnclaimedShardIsStopped(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-orphaned")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-orphaned" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
sh, err := h.ag.Supervisor().Get("shard-orphaned")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// The control plane will refuse to claim it on reconnect.
|
||||
h.rec.markUnclaimed("shard-orphaned")
|
||||
h.ln.drop()
|
||||
|
||||
deadline := time.Now().Add(20 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAliveForTest(pid) {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Error("an unclaimed shard kept running and holding its GPUs")
|
||||
}
|
||||
|
||||
// processAliveForTest mirrors the supervisor's liveness check.
|
||||
func processAliveForTest(pid int) bool {
|
||||
return syscall.Kill(pid, 0) == nil
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
)
|
||||
|
||||
// ErrNoNVML means this binary was built without the NVML backend.
|
||||
//
|
||||
// It is declared here rather than in the build-tagged file so callers — and
|
||||
// tests — can refer to it in either build.
|
||||
var ErrNoNVML = errors.New("built without the nvml tag; rebuild with -tags nvml for NVLink topology")
|
||||
|
||||
// Options configure a full inventory scan.
|
||||
type Options struct {
|
||||
// NvidiaSMI is the path to nvidia-smi; empty means look it up on PATH.
|
||||
// Worth configuring on nodes where the driver lives outside the default
|
||||
// paths, which containerized driver installs often do.
|
||||
NvidiaSMI string `koanf:"nvidia_smi"`
|
||||
// IBRoot is the sysfs InfiniBand directory.
|
||||
IBRoot string `koanf:"ib_root"`
|
||||
// PreferNVML uses the NVML backend when this binary has one. It is the only
|
||||
// source of NVLink topology.
|
||||
PreferNVML bool `koanf:"prefer_nvml"`
|
||||
// SkipGPUs omits GPU discovery, for nodes that have none.
|
||||
SkipGPUs bool `koanf:"skip_gpus"`
|
||||
}
|
||||
|
||||
// Discover builds the node's inventory.
|
||||
//
|
||||
// When NVML is preferred but fails, it falls back to nvidia-smi and says so in
|
||||
// Inventory.Source. The fallback is logged rather than silent: on a node that
|
||||
// is supposed to have working NVML, a quiet downgrade hides a broken driver and
|
||||
// costs the scheduler its NVLink topology.
|
||||
func Discover(ctx context.Context, opts Options, log *slog.Logger) (Inventory, error) {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
inv := Inventory{Source: "none"}
|
||||
|
||||
if !opts.SkipGPUs {
|
||||
switch {
|
||||
case opts.PreferNVML && nvmlAvailable:
|
||||
gpus, err := DiscoverGPUsNVML(ctx)
|
||||
if err == nil {
|
||||
inv.GPUs, inv.Source = gpus, "nvml"
|
||||
break
|
||||
}
|
||||
log.Warn("NVML discovery failed, falling back to nvidia-smi; "+
|
||||
"NVLink topology will be unavailable", "err", err)
|
||||
fallthrough
|
||||
default:
|
||||
gpus, err := DiscoverGPUsSMI(ctx, opts.NvidiaSMI)
|
||||
if err != nil {
|
||||
return inv, fmt.Errorf("gpu discovery: %w", err)
|
||||
}
|
||||
inv.GPUs, inv.Source = gpus, "nvidia-smi"
|
||||
}
|
||||
}
|
||||
|
||||
hcas, err := DiscoverIB(opts.IBRoot)
|
||||
if err != nil {
|
||||
return inv, fmt.Errorf("infiniband discovery: %w", err)
|
||||
}
|
||||
inv.HCAs = hcas
|
||||
return inv, nil
|
||||
}
|
||||
|
||||
// NVMLSupported reports whether this binary includes the NVML backend.
|
||||
func NVMLSupported() bool { return nvmlAvailable }
|
||||
@@ -0,0 +1,101 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func fakeSMI(t *testing.T) string {
|
||||
t.Helper()
|
||||
stub := filepath.Join(t.TempDir(), "nvidia-smi")
|
||||
script := "#!/bin/sh\ncat <<'EOF'\n" + smiSample + "EOF\n"
|
||||
if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return stub
|
||||
}
|
||||
|
||||
func TestDiscoverUsesSMIByDefault(t *testing.T) {
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
NvidiaSMI: fakeSMI(t),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if inv.Source != "nvidia-smi" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
if len(inv.GPUs) != 2 {
|
||||
t.Errorf("got %d GPUs", len(inv.GPUs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverFallsBackWhenNVMLUnavailable(t *testing.T) {
|
||||
// Asking for NVML in a build without it must still produce an inventory:
|
||||
// the fallback is the whole point of the build tag.
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
PreferNVML: true,
|
||||
NvidiaSMI: fakeSMI(t),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(inv.GPUs) == 0 {
|
||||
t.Fatal("no GPUs discovered through the fallback")
|
||||
}
|
||||
if NVMLSupported() {
|
||||
// On an NVML build without a driver the NVML attempt fails and the
|
||||
// fallback still has to work.
|
||||
if inv.Source != "nvml" && inv.Source != "nvidia-smi" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
} else if inv.Source != "nvidia-smi" {
|
||||
t.Errorf("a build without NVML reported source=%q", inv.Source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverSkipGPUs(t *testing.T) {
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
SkipGPUs: true,
|
||||
IBRoot: fakeSysfs(t),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(inv.GPUs) != 0 {
|
||||
t.Errorf("GPUs were discovered despite SkipGPUs")
|
||||
}
|
||||
if len(inv.HCAs) != 3 {
|
||||
t.Errorf("got %d HCAs", len(inv.HCAs))
|
||||
}
|
||||
if inv.Source != "none" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverReportsGPUFailure(t *testing.T) {
|
||||
// A node meant to have GPUs whose discovery fails must say so rather than
|
||||
// reporting an empty, healthy-looking inventory.
|
||||
_, err := Discover(context.Background(), Options{
|
||||
NvidiaSMI: filepath.Join(t.TempDir(), "absent"),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err == nil {
|
||||
t.Error("a failed GPU discovery was reported as success")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNVMLSupportedMatchesBuild(t *testing.T) {
|
||||
// Documents which backend this binary has, and keeps both files compiling.
|
||||
t.Logf("NVML backend compiled in: %v", NVMLSupported())
|
||||
if !NVMLSupported() {
|
||||
if _, err := DiscoverGPUsNVML(t.Context()); err != ErrNoNVML {
|
||||
t.Errorf("expected ErrNoNVML, got %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// GPU is one accelerator on this node.
|
||||
type GPU struct {
|
||||
Index int `json:"index"`
|
||||
UUID string `json:"uuid"`
|
||||
Name string `json:"name"`
|
||||
// MemoryTotalMB is the total framebuffer size.
|
||||
MemoryTotalMB int `json:"memory_total_mb"`
|
||||
// MemoryUsedMB is what is already allocated. Non-zero on an idle node means
|
||||
// something other than our engines is using the card (plan §5.1).
|
||||
MemoryUsedMB int `json:"memory_used_mb"`
|
||||
// ComputeCapability gates quantization choices: an FP8 profile must not
|
||||
// land on a card below 8.9 (plan §6.2).
|
||||
ComputeCapability string `json:"compute_capability"`
|
||||
NUMANode int `json:"numa_node"`
|
||||
PCIBusID string `json:"pci_bus_id,omitempty"`
|
||||
DriverVersion string `json:"driver_version,omitempty"`
|
||||
CUDAVersion string `json:"cuda_version,omitempty"`
|
||||
// NVLinkPeers are the indices of GPUs reachable over NVLink. Within a node
|
||||
// a TP group should stay inside one NVLink domain.
|
||||
NVLinkPeers []int `json:"nvlink_peers,omitempty"`
|
||||
// Health carries anything that should keep the GPU out of placement.
|
||||
Health GPUHealth `json:"health"`
|
||||
}
|
||||
|
||||
// GPUHealth records conditions that make a GPU unsafe to schedule onto.
|
||||
type GPUHealth struct {
|
||||
// ECCErrors is the count of uncorrected ECC errors.
|
||||
ECCErrors int `json:"ecc_errors"`
|
||||
// TemperatureC is the current die temperature.
|
||||
TemperatureC int `json:"temperature_c"`
|
||||
// Problems lists human-readable reasons to avoid this GPU.
|
||||
Problems []string `json:"problems,omitempty"`
|
||||
}
|
||||
|
||||
// OK reports whether the GPU is fit for placement.
|
||||
func (h GPUHealth) OK() bool { return len(h.Problems) == 0 }
|
||||
|
||||
// nonAlnum collapses everything that is not a letter or digit.
|
||||
var nonAlnum = regexp.MustCompile(`[^a-z0-9]+`)
|
||||
|
||||
// Class derives the GPU class used for placement: same model, same memory.
|
||||
//
|
||||
// Two cards only belong in the same tensor-parallel group if they are the same
|
||||
// class, so the derivation has to be stable across driver versions and naming
|
||||
// quirks ("NVIDIA H100 80GB HBM3" and "H100-SXM5-80GB" are the same class).
|
||||
func (g GPU) Class() string {
|
||||
name := strings.ToLower(g.Name)
|
||||
name = strings.TrimPrefix(name, "nvidia ")
|
||||
name = strings.TrimPrefix(name, "tesla ")
|
||||
|
||||
// Keep the model token: the first word that contains a digit, which is how
|
||||
// NVIDIA names every accelerator (h100, a100, l40s, rtx4090).
|
||||
var model string
|
||||
for _, tok := range nonAlnum.Split(name, -1) {
|
||||
if tok == "" || !strings.ContainsFunc(tok, func(r rune) bool { return r >= '0' && r <= '9' }) {
|
||||
continue
|
||||
}
|
||||
// Skip memory and form-factor tokens, which are not the model.
|
||||
if strings.HasSuffix(tok, "gb") || tok == "80gb" || tok == "sxm4" || tok == "sxm5" || tok == "pcie" {
|
||||
continue
|
||||
}
|
||||
model = tok
|
||||
break
|
||||
}
|
||||
if model == "" {
|
||||
model = nonAlnum.ReplaceAllString(name, "-")
|
||||
}
|
||||
|
||||
// Round to the nearest GiB: vendors report 81559 MiB for an "80GB" card.
|
||||
gb := int(math.Round(float64(g.MemoryTotalMB) / 1024))
|
||||
return fmt.Sprintf("%s-%dg", model, gb)
|
||||
}
|
||||
|
||||
// MarshalJSON adds the derived class, so a consumer of the JSON does not have
|
||||
// to reimplement the derivation to know which GPUs are interchangeable.
|
||||
func (g GPU) MarshalJSON() ([]byte, error) {
|
||||
type alias GPU // prevents recursing back into this method
|
||||
return json.Marshal(struct {
|
||||
alias
|
||||
Class string `json:"class"`
|
||||
}{alias(g), g.Class()})
|
||||
}
|
||||
|
||||
// Inventory is everything discovery found on this node.
|
||||
type Inventory struct {
|
||||
GPUs []GPU `json:"gpus"`
|
||||
HCAs []HCA `json:"hcas"`
|
||||
// Source names the backend that produced the GPU list, so a consumer can
|
||||
// tell "no GPUs" from "discovery did not work".
|
||||
Source string `json:"source"`
|
||||
}
|
||||
|
||||
// Classes returns the distinct GPU classes present, sorted.
|
||||
func (inv Inventory) Classes() []string {
|
||||
seen := map[string]struct{}{}
|
||||
for _, g := range inv.GPUs {
|
||||
seen[g.Class()] = struct{}{}
|
||||
}
|
||||
out := make([]string, 0, len(seen))
|
||||
for c := range seen {
|
||||
out = append(out, c)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// parseCSVInt reads an integer from nvidia-smi's CSV output, tolerating the
|
||||
// "[N/A]" and "Not Supported" values it emits for unsupported queries.
|
||||
func parseCSVInt(s string) (int, bool) {
|
||||
s = strings.TrimSpace(s)
|
||||
s = strings.TrimSuffix(s, " MiB")
|
||||
s = strings.TrimSuffix(s, " W")
|
||||
if s == "" || strings.HasPrefix(s, "[") || strings.EqualFold(s, "Not Supported") {
|
||||
return 0, false
|
||||
}
|
||||
v, err := strconv.Atoi(s)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return v, true
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
//go:build !nvml
|
||||
|
||||
package discovery
|
||||
|
||||
import "context"
|
||||
|
||||
// nvmlAvailable reports that this build cannot use NVML.
|
||||
const nvmlAvailable = false
|
||||
|
||||
// DiscoverGPUsNVML is unavailable in builds without the nvml tag.
|
||||
func DiscoverGPUsNVML(context.Context) ([]GPU, error) { return nil, ErrNoNVML }
|
||||
@@ -0,0 +1,148 @@
|
||||
//go:build nvml
|
||||
|
||||
// This file is the NVML-backed GPU discovery path, built only with the `nvml`
|
||||
// tag because the bindings require cgo. Without the tag the agent falls back to
|
||||
// nvidia-smi, which keeps the server side — and development machines without a
|
||||
// C toolchain or a driver — building cleanly (plan §3.2).
|
||||
//
|
||||
// What NVML adds over nvidia-smi is NVLink topology. Nothing else exposes which
|
||||
// GPUs share an NVLink domain, and that is what keeps a tensor-parallel group
|
||||
// off the slow path inside a node (plan §6.2, constraint 3).
|
||||
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/NVIDIA/go-nvml/pkg/nvml"
|
||||
)
|
||||
|
||||
// nvmlAvailable reports that this build can use NVML.
|
||||
const nvmlAvailable = true
|
||||
|
||||
// DiscoverGPUsNVML enumerates GPUs through NVML.
|
||||
//
|
||||
// It returns an error rather than falling back on its own: the caller decides
|
||||
// whether to degrade to nvidia-smi, because on a node that is supposed to have
|
||||
// working NVML a silent downgrade hides a broken driver installation.
|
||||
func DiscoverGPUsNVML(_ context.Context) ([]GPU, error) {
|
||||
if ret := nvml.Init(); ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml init: %s", nvml.ErrorString(ret))
|
||||
}
|
||||
defer func() { _ = nvml.Shutdown() }()
|
||||
|
||||
driver, _ := nvml.SystemGetDriverVersion()
|
||||
cuda := ""
|
||||
if v, ret := nvml.SystemGetCudaDriverVersion(); ret == nvml.SUCCESS {
|
||||
// NVML reports 12040 for 12.4.
|
||||
cuda = fmt.Sprintf("%d.%d", v/1000, (v%1000)/10)
|
||||
}
|
||||
|
||||
count, ret := nvml.DeviceGetCount()
|
||||
if ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml device count: %s", nvml.ErrorString(ret))
|
||||
}
|
||||
|
||||
// Bus id to index, so NVLink peers can be reported as indices. NVML gives
|
||||
// the remote end of a link as a PCI address, which is not what placement
|
||||
// works in.
|
||||
byBusID := make(map[string]int, count)
|
||||
gpus := make([]GPU, 0, count)
|
||||
|
||||
for i := range count {
|
||||
dev, ret := nvml.DeviceGetHandleByIndex(i)
|
||||
if ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml device %d: %s", i, nvml.ErrorString(ret))
|
||||
}
|
||||
|
||||
g := GPU{Index: i, NUMANode: -1, DriverVersion: driver, CUDAVersion: cuda}
|
||||
g.UUID, _ = dev.GetUUID()
|
||||
g.Name, _ = dev.GetName()
|
||||
|
||||
if mem, ret := dev.GetMemoryInfo(); ret == nvml.SUCCESS {
|
||||
g.MemoryTotalMB = int(mem.Total / (1 << 20))
|
||||
g.MemoryUsedMB = int(mem.Used / (1 << 20))
|
||||
}
|
||||
if major, minor, ret := dev.GetCudaComputeCapability(); ret == nvml.SUCCESS {
|
||||
g.ComputeCapability = fmt.Sprintf("%d.%d", major, minor)
|
||||
}
|
||||
if pci, ret := dev.GetPciInfo(); ret == nvml.SUCCESS {
|
||||
g.PCIBusID = busIDString(pci.BusId[:])
|
||||
byBusID[normalizeBusID(g.PCIBusID)] = i
|
||||
}
|
||||
if t, ret := dev.GetTemperature(nvml.TEMPERATURE_GPU); ret == nvml.SUCCESS {
|
||||
g.Health.TemperatureC = int(t)
|
||||
}
|
||||
// Volatile counters reset on driver reload, which is what we want: an
|
||||
// error from three months and one reboot ago is not a reason to refuse
|
||||
// a card today.
|
||||
if n, ret := dev.GetTotalEccErrors(nvml.MEMORY_ERROR_TYPE_UNCORRECTED, nvml.VOLATILE_ECC); ret == nvml.SUCCESS {
|
||||
g.Health.ECCErrors = int(n)
|
||||
}
|
||||
|
||||
assessHealth(&g)
|
||||
gpus = append(gpus, g)
|
||||
}
|
||||
|
||||
// A second pass, because a link's remote end may be a GPU enumerated later.
|
||||
for i := range gpus {
|
||||
dev, ret := nvml.DeviceGetHandleByIndex(gpus[i].Index)
|
||||
if ret != nvml.SUCCESS {
|
||||
continue
|
||||
}
|
||||
gpus[i].NVLinkPeers = nvlinkPeers(dev, byBusID, gpus[i].Index)
|
||||
}
|
||||
return gpus, nil
|
||||
}
|
||||
|
||||
// nvlinkPeers returns the indices of GPUs reachable from dev over NVLink.
|
||||
func nvlinkPeers(dev nvml.Device, byBusID map[string]int, self int) []int {
|
||||
seen := map[int]struct{}{}
|
||||
var out []int
|
||||
|
||||
for link := range nvml.NVLINK_MAX_LINKS {
|
||||
pci, ret := dev.GetNvLinkRemotePciInfo(link)
|
||||
if ret != nvml.SUCCESS {
|
||||
// Unpopulated links return an error; that is the normal case for
|
||||
// most of the 36 possible links, so it is not worth reporting.
|
||||
continue
|
||||
}
|
||||
peer, ok := byBusID[normalizeBusID(busIDString(pci.BusId[:]))]
|
||||
if !ok || peer == self {
|
||||
// The remote end can be an NVSwitch rather than a GPU, in which
|
||||
// case it is not in the map.
|
||||
continue
|
||||
}
|
||||
if _, dup := seen[peer]; dup {
|
||||
continue
|
||||
}
|
||||
seen[peer] = struct{}{}
|
||||
out = append(out, peer)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// busIDString turns NVML's fixed-size C char array into a Go string.
|
||||
func busIDString(b []int8) string {
|
||||
buf := make([]byte, 0, len(b))
|
||||
for _, c := range b {
|
||||
if c == 0 {
|
||||
break
|
||||
}
|
||||
buf = append(buf, byte(c))
|
||||
}
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
// normalizeBusID makes bus ids comparable: NVML reports them uppercase from one
|
||||
// call and lowercase from another, and sometimes with a shortened domain.
|
||||
func normalizeBusID(s string) string {
|
||||
s = strings.ToLower(strings.TrimSpace(s))
|
||||
// "0000:1b:00.0" and "00000000:1B:00.0" are the same device.
|
||||
if i := strings.Index(s, ":"); i > 0 {
|
||||
s = strings.TrimLeft(s[:i], "0") + s[i:]
|
||||
}
|
||||
return strings.TrimPrefix(s, ":")
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// smiQueryFields are the columns requested from nvidia-smi, in order.
|
||||
var smiQueryFields = []string{
|
||||
"index",
|
||||
"uuid",
|
||||
"name",
|
||||
"memory.total",
|
||||
"memory.used",
|
||||
"compute_cap",
|
||||
"pci.bus_id",
|
||||
"driver_version",
|
||||
"temperature.gpu",
|
||||
"ecc.errors.uncorrected.volatile.total",
|
||||
}
|
||||
|
||||
// DiscoverGPUsSMI enumerates GPUs by shelling out to nvidia-smi.
|
||||
//
|
||||
// This is the cgo-free path (plan §3.2): the NVML bindings need cgo, which
|
||||
// would force the whole binary — server included — to be built with a C
|
||||
// toolchain. nvidia-smi gives everything except NVLink topology, which the
|
||||
// NVML build adds.
|
||||
func DiscoverGPUsSMI(ctx context.Context, binary string) ([]GPU, error) {
|
||||
if binary == "" {
|
||||
binary = "nvidia-smi"
|
||||
}
|
||||
out, err := exec.CommandContext(ctx, binary,
|
||||
"--query-gpu="+strings.Join(smiQueryFields, ","),
|
||||
"--format=csv,noheader,nounits",
|
||||
).Output()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("running %s: %w", binary, err)
|
||||
}
|
||||
return parseSMI(string(out))
|
||||
}
|
||||
|
||||
func parseSMI(out string) ([]GPU, error) {
|
||||
r := csv.NewReader(strings.NewReader(out))
|
||||
r.TrimLeadingSpace = true
|
||||
r.FieldsPerRecord = -1 // tolerate driver versions with extra columns
|
||||
|
||||
rows, err := r.ReadAll()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parsing nvidia-smi output: %w", err)
|
||||
}
|
||||
|
||||
gpus := make([]GPU, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
if len(row) < len(smiQueryFields) {
|
||||
continue
|
||||
}
|
||||
idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
g := GPU{
|
||||
Index: idx,
|
||||
UUID: strings.TrimSpace(row[1]),
|
||||
Name: strings.TrimSpace(row[2]),
|
||||
ComputeCapability: strings.TrimSpace(row[5]),
|
||||
PCIBusID: strings.TrimSpace(row[6]),
|
||||
DriverVersion: strings.TrimSpace(row[7]),
|
||||
NUMANode: -1,
|
||||
}
|
||||
g.MemoryTotalMB, _ = parseCSVInt(row[3])
|
||||
g.MemoryUsedMB, _ = parseCSVInt(row[4])
|
||||
g.Health.TemperatureC, _ = parseCSVInt(row[8])
|
||||
g.Health.ECCErrors, _ = parseCSVInt(row[9])
|
||||
|
||||
assessHealth(&g)
|
||||
gpus = append(gpus, g)
|
||||
}
|
||||
return gpus, nil
|
||||
}
|
||||
|
||||
// residentMemoryToleranceMB is how much allocated memory is written off as
|
||||
// driver overhead rather than a foreign process.
|
||||
const residentMemoryToleranceMB = 512
|
||||
|
||||
// assessHealth flags conditions that should keep a GPU out of placement.
|
||||
func assessHealth(g *GPU) {
|
||||
if g.Health.ECCErrors > 0 {
|
||||
g.Health.Problems = append(g.Health.Problems,
|
||||
fmt.Sprintf("%d uncorrected ECC errors", g.Health.ECCErrors))
|
||||
}
|
||||
// Memory already in use means something else is on the card. Starting an
|
||||
// engine there would either fail to allocate or quietly evict the other
|
||||
// process's throughput.
|
||||
if g.MemoryUsedMB > residentMemoryToleranceMB {
|
||||
g.Health.Problems = append(g.Health.Problems,
|
||||
fmt.Sprintf("%d MiB already allocated by another process", g.MemoryUsedMB))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const smiSample = `0, GPU-11111111-2222-3333-4444-555555555555, NVIDIA H100 80GB HBM3, 81559, 4, 9.0, 00000000:1B:00.0, 550.54.15, 34, 0
|
||||
1, GPU-66666666-7777-8888-9999-000000000000, NVIDIA H100 80GB HBM3, 81559, 0, 9.0, 00000000:1C:00.0, 550.54.15, 36, 0
|
||||
`
|
||||
|
||||
func TestParseSMI(t *testing.T) {
|
||||
gpus, err := parseSMI(smiSample)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 2 {
|
||||
t.Fatalf("expected 2 GPUs, got %d", len(gpus))
|
||||
}
|
||||
|
||||
g := gpus[0]
|
||||
if g.Index != 0 || g.UUID != "GPU-11111111-2222-3333-4444-555555555555" {
|
||||
t.Errorf("identity wrong: %+v", g)
|
||||
}
|
||||
if g.MemoryTotalMB != 81559 {
|
||||
t.Errorf("MemoryTotalMB=%d", g.MemoryTotalMB)
|
||||
}
|
||||
if g.ComputeCapability != "9.0" {
|
||||
t.Errorf("ComputeCapability=%q; FP8 placement depends on it", g.ComputeCapability)
|
||||
}
|
||||
if g.DriverVersion != "550.54.15" {
|
||||
t.Errorf("DriverVersion=%q", g.DriverVersion)
|
||||
}
|
||||
if !g.Health.OK() {
|
||||
t.Errorf("a healthy GPU was flagged: %v", g.Health.Problems)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSMIUnsupportedFields(t *testing.T) {
|
||||
// Consumer cards report [N/A] for ECC; that must not fail the whole parse.
|
||||
const row = `0, GPU-abc, NVIDIA GeForce RTX 4090, 24564, 0, 8.9, 00000000:01:00.0, 550.54.15, 45, [N/A]
|
||||
`
|
||||
gpus, err := parseSMI(row)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("expected 1 GPU, got %d", len(gpus))
|
||||
}
|
||||
if gpus[0].Health.ECCErrors != 0 || !gpus[0].Health.OK() {
|
||||
t.Errorf("[N/A] ECC was misread: %+v", gpus[0].Health)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthFlagsForeignMemory(t *testing.T) {
|
||||
// Memory in use on an otherwise idle card means another process is there.
|
||||
const row = `0, GPU-abc, NVIDIA H100 80GB HBM3, 81559, 40000, 9.0, 00000000:1B:00.0, 550.54.15, 34, 0
|
||||
`
|
||||
gpus, _ := parseSMI(row)
|
||||
if gpus[0].Health.OK() {
|
||||
t.Error("40 GiB allocated by another process was not flagged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthFlagsECC(t *testing.T) {
|
||||
const row = `0, GPU-abc, NVIDIA H100 80GB HBM3, 81559, 0, 9.0, 00000000:1B:00.0, 550.54.15, 34, 3
|
||||
`
|
||||
gpus, _ := parseSMI(row)
|
||||
if gpus[0].Health.OK() {
|
||||
t.Error("uncorrected ECC errors were not flagged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGPUClass(t *testing.T) {
|
||||
// The same card must derive the same class however the driver names it,
|
||||
// because a tensor-parallel group may not mix classes.
|
||||
cases := []struct {
|
||||
name string
|
||||
memMB int
|
||||
expect string
|
||||
}{
|
||||
{"NVIDIA H100 80GB HBM3", 81559, "h100-80g"},
|
||||
{"NVIDIA H100-SXM5-80GB", 81559, "h100-80g"},
|
||||
{"Tesla A100-SXM4-40GB", 40960, "a100-40g"},
|
||||
{"NVIDIA A100 80GB PCIe", 81920, "a100-80g"},
|
||||
// Sizes come from what the driver reports, in GiB. That does not always
|
||||
// match the marketing name: an L40S is sold as 48GB but reports 44.99
|
||||
// GiB, because the marketing figure is decimal GB. Being consistent
|
||||
// matters more than matching the box, since the number is what decides
|
||||
// whether a model fits.
|
||||
{"NVIDIA L40S", 46068, "l40s-45g"},
|
||||
// The model token is what distinguishes a 4090 from a 4080, so "rtx"
|
||||
// is not enough on its own.
|
||||
{"NVIDIA GeForce RTX 4090", 24564, "4090-24g"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := GPU{Name: c.name, MemoryTotalMB: c.memMB}.Class()
|
||||
if got != c.expect {
|
||||
t.Errorf("%q (%d MiB) -> %q, want %q", c.name, c.memMB, got, c.expect)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestInventoryClasses(t *testing.T) {
|
||||
inv := Inventory{GPUs: []GPU{
|
||||
{Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559},
|
||||
{Name: "NVIDIA H100-SXM5-80GB", MemoryTotalMB: 81559},
|
||||
{Name: "Tesla A100-SXM4-40GB", MemoryTotalMB: 40960},
|
||||
}}
|
||||
got := inv.Classes()
|
||||
if len(got) != 2 || got[0] != "a100-40g" || got[1] != "h100-80g" {
|
||||
t.Errorf("classes=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverGPUsSMI runs the real code path against a stub binary, so the
|
||||
// command construction is covered too, not just the parser.
|
||||
func TestDiscoverGPUsSMI(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
stub := filepath.Join(dir, "nvidia-smi")
|
||||
script := "#!/bin/sh\ncat <<'EOF'\n" + smiSample + "EOF\n"
|
||||
if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
gpus, err := DiscoverGPUsSMI(t.Context(), stub)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 2 {
|
||||
t.Fatalf("expected 2 GPUs, got %d", len(gpus))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverGPUsSMIMissingBinary(t *testing.T) {
|
||||
// A node with no NVIDIA driver must produce a clear error rather than a
|
||||
// panic or an empty success.
|
||||
_, err := DiscoverGPUsSMI(context.Background(), filepath.Join(t.TempDir(), "absent"))
|
||||
if err == nil {
|
||||
t.Error("expected an error when nvidia-smi is missing")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
// Package discovery inspects a node's GPUs, InfiniBand adapters and the
|
||||
// topology between them (plan §5.1 and §5.2).
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// SysfsRoot is the default place InfiniBand devices appear.
|
||||
const SysfsRoot = "/sys/class/infiniband"
|
||||
|
||||
// LinkLayer distinguishes real InfiniBand from RoCE.
|
||||
type LinkLayer string
|
||||
|
||||
const (
|
||||
// LinkInfiniBand is native InfiniBand.
|
||||
LinkInfiniBand LinkLayer = "InfiniBand"
|
||||
// LinkEthernet is RoCE — Ethernet carrying RDMA.
|
||||
LinkEthernet LinkLayer = "Ethernet"
|
||||
)
|
||||
|
||||
// Port is one port of an HCA.
|
||||
type Port struct {
|
||||
Num int `json:"num"`
|
||||
// State is the port's logical state; only ACTIVE ports carry traffic.
|
||||
State string `json:"state"`
|
||||
// PhysState is the physical link state, useful for telling "cable unplugged"
|
||||
// from "no subnet manager".
|
||||
PhysState string `json:"phys_state"`
|
||||
// RateGbps is the negotiated rate.
|
||||
RateGbps int `json:"rate_gbps"`
|
||||
LinkLayer LinkLayer `json:"link_layer"`
|
||||
// GIDPrefix is the subnet prefix from gids/0, normalized to lowercase hex
|
||||
// groups. It is the only automatic fabric hint available, and an unreliable
|
||||
// one: OpenSM's default is fe80:: everywhere (plan §5.2).
|
||||
GIDPrefix string `json:"gid_prefix"`
|
||||
LID int `json:"lid"`
|
||||
SMLID int `json:"sm_lid"`
|
||||
}
|
||||
|
||||
// Active reports whether the port can carry traffic.
|
||||
func (p Port) Active() bool { return strings.EqualFold(p.State, "ACTIVE") }
|
||||
|
||||
// HCA is one host channel adapter.
|
||||
type HCA struct {
|
||||
Name string `json:"name"`
|
||||
Ports []Port `json:"ports"`
|
||||
PCIAddr string `json:"pci_addr,omitempty"`
|
||||
NUMANode int `json:"numa_node"`
|
||||
// FirmwareVersion is reported by the driver; worth capturing because
|
||||
// NCCL behaviour differs between firmware revisions.
|
||||
FirmwareVersion string `json:"firmware_version,omitempty"`
|
||||
}
|
||||
|
||||
// ActivePorts returns the ports that can carry traffic.
|
||||
func (h HCA) ActivePorts() []Port {
|
||||
var out []Port
|
||||
for _, p := range h.Ports {
|
||||
if p.Active() {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// DiscoverIB enumerates the InfiniBand adapters under root.
|
||||
//
|
||||
// A missing root is not an error: plenty of nodes have no IB at all, and the
|
||||
// scheduler treats them as single-node islands (plan §2).
|
||||
func DiscoverIB(root string) ([]HCA, error) {
|
||||
if root == "" {
|
||||
root = SysfsRoot
|
||||
}
|
||||
entries, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("reading %s: %w", root, err)
|
||||
}
|
||||
|
||||
var out []HCA
|
||||
for _, e := range entries {
|
||||
hca, err := readHCA(filepath.Join(root, e.Name()), e.Name())
|
||||
if err != nil {
|
||||
// One unreadable adapter must not hide the rest; the node may still
|
||||
// be usable through its other HCAs.
|
||||
continue
|
||||
}
|
||||
out = append(out, hca)
|
||||
}
|
||||
slices.SortFunc(out, func(a, b HCA) int { return strings.Compare(a.Name, b.Name) })
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func readHCA(dir, name string) (HCA, error) {
|
||||
h := HCA{Name: name, NUMANode: -1}
|
||||
h.FirmwareVersion = readTrimmed(filepath.Join(dir, "fw_ver"))
|
||||
|
||||
// device is a symlink into the PCI tree; its target's base name is the
|
||||
// bus address, and numa_node lives beside it.
|
||||
if target, err := filepath.EvalSymlinks(filepath.Join(dir, "device")); err == nil {
|
||||
h.PCIAddr = filepath.Base(target)
|
||||
if n, err := strconv.Atoi(readTrimmed(filepath.Join(target, "numa_node"))); err == nil {
|
||||
h.NUMANode = n
|
||||
}
|
||||
}
|
||||
|
||||
portDirs, err := os.ReadDir(filepath.Join(dir, "ports"))
|
||||
if err != nil {
|
||||
return h, fmt.Errorf("%s: no ports: %w", name, err)
|
||||
}
|
||||
for _, pd := range portDirs {
|
||||
num, err := strconv.Atoi(pd.Name())
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
h.Ports = append(h.Ports, readPort(filepath.Join(dir, "ports", pd.Name()), num))
|
||||
}
|
||||
slices.SortFunc(h.Ports, func(a, b Port) int { return a.Num - b.Num })
|
||||
return h, nil
|
||||
}
|
||||
|
||||
func readPort(dir string, num int) Port {
|
||||
p := Port{
|
||||
Num: num,
|
||||
State: parseEnumeratedState(readTrimmed(filepath.Join(dir, "state"))),
|
||||
PhysState: parseEnumeratedState(readTrimmed(filepath.Join(dir, "phys_state"))),
|
||||
RateGbps: parseRate(readTrimmed(filepath.Join(dir, "rate"))),
|
||||
LinkLayer: LinkLayer(readTrimmed(filepath.Join(dir, "link_layer"))),
|
||||
GIDPrefix: parseGIDPrefix(readTrimmed(filepath.Join(dir, "gids", "0"))),
|
||||
LID: parseMaybeHex(readTrimmed(filepath.Join(dir, "lid"))),
|
||||
SMLID: parseMaybeHex(readTrimmed(filepath.Join(dir, "sm_lid"))),
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// parseEnumeratedState turns sysfs's "4: ACTIVE" into "ACTIVE".
|
||||
func parseEnumeratedState(s string) string {
|
||||
if _, rest, found := strings.Cut(s, ":"); found {
|
||||
return strings.TrimSpace(rest)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// parseRate turns "100 Gb/sec (4X EDR)" into 100.
|
||||
func parseRate(s string) int {
|
||||
f := strings.Fields(s)
|
||||
if len(f) == 0 {
|
||||
return 0
|
||||
}
|
||||
// The rate can be fractional ("2.5 Gb/sec") on old hardware; truncate.
|
||||
v, err := strconv.ParseFloat(f[0], 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return int(v)
|
||||
}
|
||||
|
||||
// parseGIDPrefix keeps the first four groups of a GID, which is the subnet
|
||||
// prefix; the remaining four are the interface id.
|
||||
func parseGIDPrefix(gid string) string {
|
||||
parts := strings.Split(strings.ToLower(gid), ":")
|
||||
if len(parts) < 4 {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(parts[:4], ":")
|
||||
}
|
||||
|
||||
// parseMaybeHex accepts both "0x1a" and "26"; sysfs is inconsistent between
|
||||
// drivers and kernel versions.
|
||||
func parseMaybeHex(s string) int {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return 0
|
||||
}
|
||||
if after, ok := strings.CutPrefix(strings.ToLower(s), "0x"); ok {
|
||||
v, err := strconv.ParseInt(after, 16, 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return int(v)
|
||||
}
|
||||
v, err := strconv.Atoi(s)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func readTrimmed(path string) string {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// fakeSysfs builds an InfiniBand tree shaped like the real one.
|
||||
func fakeSysfs(t *testing.T) string {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
|
||||
write := func(path, content string) {
|
||||
t.Helper()
|
||||
full := filepath.Join(root, path)
|
||||
if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(full, []byte(content+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// mlx5_0: one active IB port.
|
||||
write("ib/mlx5_0/fw_ver", "20.35.1012")
|
||||
write("ib/mlx5_0/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_0/ports/1/phys_state", "5: LinkUp")
|
||||
write("ib/mlx5_0/ports/1/rate", "200 Gb/sec (4X HDR)")
|
||||
write("ib/mlx5_0/ports/1/link_layer", "InfiniBand")
|
||||
write("ib/mlx5_0/ports/1/gids/0", "FEC0:0000:0000:000A:0002:C903:00A0:1234")
|
||||
write("ib/mlx5_0/ports/1/lid", "0x1a")
|
||||
write("ib/mlx5_0/ports/1/sm_lid", "1")
|
||||
|
||||
// mlx5_1: two ports, the second one down.
|
||||
write("ib/mlx5_1/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_1/ports/1/rate", "200 Gb/sec (4X HDR)")
|
||||
write("ib/mlx5_1/ports/1/link_layer", "InfiniBand")
|
||||
write("ib/mlx5_1/ports/1/gids/0", "fec0:0000:0000:000a:0002:c903:00a0:5678")
|
||||
write("ib/mlx5_1/ports/2/state", "1: DOWN")
|
||||
write("ib/mlx5_1/ports/2/phys_state", "3: Disabled")
|
||||
write("ib/mlx5_1/ports/2/rate", "0")
|
||||
write("ib/mlx5_1/ports/2/link_layer", "InfiniBand")
|
||||
|
||||
// mlx5_2: RoCE rather than native IB.
|
||||
write("ib/mlx5_2/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_2/ports/1/rate", "100 Gb/sec (4X EDR)")
|
||||
write("ib/mlx5_2/ports/1/link_layer", "Ethernet")
|
||||
write("ib/mlx5_2/ports/1/gids/0", "0000:0000:0000:0000:0000:ffff:0a00:0001")
|
||||
|
||||
// A PCI device with a NUMA node, linked the way the driver does it.
|
||||
pci := filepath.Join(root, "pci", "0000:0b:00.0")
|
||||
if err := os.MkdirAll(pci, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(pci, "numa_node"), []byte("1\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.Symlink(pci, filepath.Join(root, "ib", "mlx5_0", "device")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
return filepath.Join(root, "ib")
|
||||
}
|
||||
|
||||
func TestDiscoverIB(t *testing.T) {
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hcas) != 3 {
|
||||
t.Fatalf("expected 3 HCAs, got %d", len(hcas))
|
||||
}
|
||||
// Sorted by name, so the order is predictable for the rest of the test.
|
||||
if hcas[0].Name != "mlx5_0" || hcas[2].Name != "mlx5_2" {
|
||||
t.Fatalf("HCAs are not sorted by name: %v %v", hcas[0].Name, hcas[2].Name)
|
||||
}
|
||||
|
||||
h := hcas[0]
|
||||
if h.FirmwareVersion != "20.35.1012" {
|
||||
t.Errorf("fw_ver=%q", h.FirmwareVersion)
|
||||
}
|
||||
if h.PCIAddr != "0000:0b:00.0" {
|
||||
t.Errorf("PCIAddr=%q", h.PCIAddr)
|
||||
}
|
||||
if h.NUMANode != 1 {
|
||||
t.Errorf("NUMANode=%d; GPU-to-HCA affinity depends on this", h.NUMANode)
|
||||
}
|
||||
|
||||
p := h.Ports[0]
|
||||
if p.State != "ACTIVE" || !p.Active() {
|
||||
t.Errorf("state=%q", p.State)
|
||||
}
|
||||
if p.PhysState != "LinkUp" {
|
||||
t.Errorf("phys_state=%q", p.PhysState)
|
||||
}
|
||||
if p.RateGbps != 200 {
|
||||
t.Errorf("rate=%d", p.RateGbps)
|
||||
}
|
||||
if p.LinkLayer != LinkInfiniBand {
|
||||
t.Errorf("link_layer=%q", p.LinkLayer)
|
||||
}
|
||||
if p.GIDPrefix != "fec0:0000:0000:000a" {
|
||||
t.Errorf("GIDPrefix=%q; it must be lowercased and cut to four groups", p.GIDPrefix)
|
||||
}
|
||||
if p.LID != 26 {
|
||||
t.Errorf("LID=%d, expected 0x1a to parse as 26", p.LID)
|
||||
}
|
||||
if p.SMLID != 1 {
|
||||
t.Errorf("SMLID=%d, expected a decimal value to parse too", p.SMLID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestActivePortsExcludesDown(t *testing.T) {
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := hcas[1] // mlx5_1
|
||||
if len(h.Ports) != 2 {
|
||||
t.Fatalf("expected 2 ports, got %d", len(h.Ports))
|
||||
}
|
||||
active := h.ActivePorts()
|
||||
if len(active) != 1 || active[0].Num != 1 {
|
||||
t.Errorf("a DOWN port was reported as active: %+v", active)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoCEIsDistinguished(t *testing.T) {
|
||||
// RoCE is usable but behaves differently from native IB; the placement
|
||||
// logic needs to be able to tell them apart.
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := hcas[2].Ports[0].LinkLayer; got != LinkEthernet {
|
||||
t.Errorf("mlx5_2 link_layer=%q, expected Ethernet (RoCE)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMissingRootIsNotAnError(t *testing.T) {
|
||||
// A node without InfiniBand is perfectly normal; it just forms its own
|
||||
// single-node island.
|
||||
hcas, err := DiscoverIB(filepath.Join(t.TempDir(), "nonexistent"))
|
||||
if err != nil {
|
||||
t.Fatalf("a node without IB should not be an error: %v", err)
|
||||
}
|
||||
if len(hcas) != 0 {
|
||||
t.Errorf("expected no HCAs, got %d", len(hcas))
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreadableAdapterDoesNotHideOthers(t *testing.T) {
|
||||
root := fakeSysfs(t)
|
||||
// An adapter with no ports directory at all.
|
||||
if err := os.MkdirAll(filepath.Join(root, "mlx5_broken"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
hcas, err := DiscoverIB(root)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hcas) != 3 {
|
||||
t.Errorf("a broken adapter changed the result: got %d HCAs", len(hcas))
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseHelpers(t *testing.T) {
|
||||
t.Run("rate", func(t *testing.T) {
|
||||
for in, want := range map[string]int{
|
||||
"200 Gb/sec (4X HDR)": 200,
|
||||
"100 Gb/sec (4X EDR)": 100,
|
||||
"2.5 Gb/sec (1X SDR)": 2,
|
||||
"": 0,
|
||||
"garbage": 0,
|
||||
} {
|
||||
if got := parseRate(in); got != want {
|
||||
t.Errorf("parseRate(%q)=%d, want %d", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("state", func(t *testing.T) {
|
||||
for in, want := range map[string]string{
|
||||
"4: ACTIVE": "ACTIVE",
|
||||
"1: DOWN": "DOWN",
|
||||
"5: LinkUp": "LinkUp",
|
||||
"ACTIVE": "ACTIVE",
|
||||
"": "",
|
||||
} {
|
||||
if got := parseEnumeratedState(in); got != want {
|
||||
t.Errorf("parseEnumeratedState(%q)=%q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("hex", func(t *testing.T) {
|
||||
for in, want := range map[string]int{
|
||||
"0x1a": 26, "0X1A": 26, "26": 26, "": 0, "zzz": 0,
|
||||
} {
|
||||
if got := parseMaybeHex(in); got != want {
|
||||
t.Errorf("parseMaybeHex(%q)=%d, want %d", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("gid", func(t *testing.T) {
|
||||
if got := parseGIDPrefix("FE80:0000:0000:0000:0002:C903:00A0:1234"); got != "fe80:0000:0000:0000" {
|
||||
t.Errorf("parseGIDPrefix=%q", got)
|
||||
}
|
||||
if got := parseGIDPrefix("short:gid"); got != "" {
|
||||
t.Errorf("a malformed GID should yield an empty prefix, got %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestJSONShape(t *testing.T) {
|
||||
inv := Inventory{
|
||||
Source: "nvidia-smi",
|
||||
GPUs: []GPU{{Index: 0, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559, ComputeCapability: "9.0"}},
|
||||
}
|
||||
b, err := json.Marshal(inv)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var out struct {
|
||||
Source string `json:"source"`
|
||||
GPUs []struct {
|
||||
Index int `json:"index"`
|
||||
Class string `json:"class"`
|
||||
Mem int `json:"memory_total_mb"`
|
||||
} `json:"gpus"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out.Source != "nvidia-smi" || len(out.GPUs) != 1 {
|
||||
t.Fatalf("shape wrong: %s", b)
|
||||
}
|
||||
// The derived class is what decides which GPUs are interchangeable, so a
|
||||
// consumer must not have to recompute it.
|
||||
if out.GPUs[0].Class != "h100-80g" {
|
||||
t.Errorf("class=%q, should be included in the JSON", out.GPUs[0].Class)
|
||||
}
|
||||
if out.GPUs[0].Mem != 81559 {
|
||||
t.Errorf("memory_total_mb=%d", out.GPUs[0].Mem)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
// Package engines builds the command lines and environments that start
|
||||
// inference engines (plan §5.3).
|
||||
//
|
||||
// Nothing here starts a process; that is the supervisor's job. Keeping command
|
||||
// construction pure makes it testable without GPUs, which matters because the
|
||||
// environment generation is where multi-node runs go wrong: a single missing
|
||||
// NCCL variable turns a working tensor-parallel group into one that hangs on
|
||||
// the first all-reduce.
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Profile is how a model runs on a GPU class. It mirrors the catalog profile
|
||||
// (plan §6.1), reduced to what launching needs.
|
||||
type Profile struct {
|
||||
Name string
|
||||
Engine string
|
||||
GPUClass string
|
||||
|
||||
// TP and PP are tensor and pipeline parallel degrees.
|
||||
TP int
|
||||
PP int
|
||||
// GPUsPerNode is how many GPUs each shard occupies.
|
||||
GPUsPerNode int
|
||||
// Nodes is how many nodes the replica spans. Above 1, a fabric is required.
|
||||
Nodes int
|
||||
|
||||
// Args are passed through to the engine untouched.
|
||||
Args []string
|
||||
// Env overrides anything the engine adapter generates.
|
||||
Env map[string]string
|
||||
|
||||
ModelPath string
|
||||
// ServedModelName is the name clients use, which need not match the path.
|
||||
ServedModelName string
|
||||
}
|
||||
|
||||
// ShardSpec is one shard of one replica on this node.
|
||||
type ShardSpec struct {
|
||||
ReplicaID string
|
||||
ShardID string
|
||||
// Rank is the node rank; 0 is the leader and serves HTTP.
|
||||
Rank int
|
||||
NNodes int
|
||||
|
||||
// LeaderAddr and LeaderPort are the distributed init endpoint, always on the
|
||||
// management network rather than the fabric: NCCL bootstraps over TCP and
|
||||
// only then moves to RDMA.
|
||||
LeaderAddr string
|
||||
LeaderPort int
|
||||
|
||||
// GPUIndices are the local GPUs this shard uses.
|
||||
GPUIndices []int
|
||||
// HCAs are the adapters of the chosen fabric, and only that fabric. A shard
|
||||
// must never be given HCAs from two fabrics.
|
||||
HCAs []string
|
||||
// MgmtInterface is the ethernet interface used for bootstrap traffic.
|
||||
MgmtInterface string
|
||||
|
||||
HTTPPort int
|
||||
Profile Profile
|
||||
|
||||
// Pool and Fabric are control plane concepts the agent cannot derive. They
|
||||
// are carried so the agent can report them back after a restart, and are
|
||||
// not used when building a command.
|
||||
Pool string
|
||||
Fabric string
|
||||
}
|
||||
|
||||
// Leader reports whether this shard serves HTTP for the replica.
|
||||
func (s ShardSpec) Leader() bool { return s.Rank == 0 }
|
||||
|
||||
// Command is a fully resolved launch: what to exec, with which environment.
|
||||
type Command struct {
|
||||
Path string
|
||||
Args []string
|
||||
// Env is the complete environment, not a delta.
|
||||
Env []string
|
||||
// Name distinguishes the processes of a multi-process shard in logs.
|
||||
Name string
|
||||
// WaitForLeaderPort makes the supervisor hold the next command until
|
||||
// LeaderAddr:LeaderPort accepts a connection.
|
||||
//
|
||||
// It exists for Ray: the head has to be listening before a worker can join,
|
||||
// and both sides express the same wait because both dial the same address.
|
||||
WaitForLeaderPort bool
|
||||
}
|
||||
|
||||
// Engine builds launch commands for one inference engine.
|
||||
type Engine interface {
|
||||
// Name is the engine's identifier in the catalog.
|
||||
Name() string
|
||||
// Validate rejects a profile that cannot work on these GPUs, before
|
||||
// anything is started.
|
||||
Validate(p Profile, gpus []GPUInfo) error
|
||||
|
||||
// Commands returns the processes that make up one shard, in start order.
|
||||
//
|
||||
// Most shards are a single process. vLLM under Ray is not: the leader runs
|
||||
// a Ray head and then the server against it, and a worker runs only the Ray
|
||||
// join. Returning a sequence rather than one command is what lets the
|
||||
// supervisor treat both the same way — every command but the last is a
|
||||
// sidecar kept alive for the shard's lifetime, and the last one's exit is
|
||||
// the shard's exit.
|
||||
Commands(s ShardSpec) ([]Command, error)
|
||||
|
||||
// ServesHTTP reports whether this rank answers the OpenAI API.
|
||||
//
|
||||
// In a multi-node replica only the leader does. Health-checking a worker's
|
||||
// port would time out and fail a replica that is coming up perfectly well,
|
||||
// which is exactly what happened before this existed.
|
||||
ServesHTTP(s ShardSpec) bool
|
||||
|
||||
// HealthPath and MetricsPath are appended to the shard's base URL.
|
||||
HealthPath() string
|
||||
MetricsPath() string
|
||||
}
|
||||
|
||||
// GPUInfo is the subset of GPU facts validation needs.
|
||||
type GPUInfo struct {
|
||||
Index int
|
||||
Name string
|
||||
ComputeCapability string
|
||||
MemoryTotalMB int
|
||||
}
|
||||
|
||||
// Registry holds the known engines.
|
||||
type Registry struct {
|
||||
engines map[string]Engine
|
||||
}
|
||||
|
||||
// NewRegistry returns a registry with the built-in engines.
|
||||
func NewRegistry(cfg Config) *Registry {
|
||||
r := &Registry{engines: map[string]Engine{}}
|
||||
r.Add(&VLLM{cfg: cfg.VLLM})
|
||||
r.Add(&SGLang{cfg: cfg.SGLang})
|
||||
return r
|
||||
}
|
||||
|
||||
// Add registers an engine, replacing any engine of the same name.
|
||||
func (r *Registry) Add(e Engine) { r.engines[e.Name()] = e }
|
||||
|
||||
// Get returns an engine by name.
|
||||
func (r *Registry) Get(name string) (Engine, error) {
|
||||
e, ok := r.engines[name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown engine %q (have %s)", name, strings.Join(r.Names(), ", "))
|
||||
}
|
||||
return e, nil
|
||||
}
|
||||
|
||||
// Names lists the registered engines.
|
||||
func (r *Registry) Names() []string {
|
||||
out := make([]string, 0, len(r.engines))
|
||||
for n := range r.engines {
|
||||
out = append(out, n)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// Config configures the engine adapters.
|
||||
type Config struct {
|
||||
VLLM VLLMConfig `koanf:"vllm"`
|
||||
SGLang SGLangConfig `koanf:"sglang"`
|
||||
}
|
||||
|
||||
// --- shared helpers ---
|
||||
|
||||
// baseEnv builds the environment every engine shares.
|
||||
//
|
||||
// This is the part of §5.3 the plan marks critical. The rules encoded here:
|
||||
//
|
||||
// - CUDA_VISIBLE_DEVICES pins the shard to its GPUs, so two shards on one
|
||||
// node cannot fight over the same card.
|
||||
// - NCCL_IB_HCA lists only the chosen fabric's adapters. Leaving it unset
|
||||
// lets NCCL pick adapters from another fabric, where the peer is
|
||||
// unreachable, and the job hangs instead of failing.
|
||||
// - NCCL_SOCKET_IFNAME pins bootstrap to the management interface. NCCL
|
||||
// otherwise probes every interface, including IB-over-IP addresses that
|
||||
// cannot route between fabrics.
|
||||
// - NCCL_IB_DISABLE is set when there is no fabric, so NCCL stops looking for
|
||||
// one instead of timing out.
|
||||
func baseEnv(s ShardSpec) map[string]string {
|
||||
env := map[string]string{
|
||||
"CUDA_VISIBLE_DEVICES": joinInts(s.GPUIndices),
|
||||
}
|
||||
|
||||
if len(s.HCAs) > 0 {
|
||||
env["NCCL_IB_HCA"] = strings.Join(s.HCAs, ",")
|
||||
env["NCCL_IB_DISABLE"] = "0"
|
||||
} else {
|
||||
// Single node, or a node with no fabric. Either way IB is not the path.
|
||||
env["NCCL_IB_DISABLE"] = "1"
|
||||
}
|
||||
|
||||
if s.MgmtInterface != "" {
|
||||
env["NCCL_SOCKET_IFNAME"] = s.MgmtInterface
|
||||
env["GLOO_SOCKET_IFNAME"] = s.MgmtInterface
|
||||
}
|
||||
|
||||
// Profile overrides win: they are the escape hatch for the per-cluster
|
||||
// tuning no adapter can guess.
|
||||
for k, v := range s.Profile.Env {
|
||||
env[k] = v
|
||||
}
|
||||
return env
|
||||
}
|
||||
|
||||
// materializeEnv merges generated variables over the current process
|
||||
// environment and returns it in exec form.
|
||||
//
|
||||
// Inheriting matters: engines need PATH, HOME, HF_HOME, LD_LIBRARY_PATH and
|
||||
// whatever else the operator set up in the unit file.
|
||||
func materializeEnv(generated map[string]string) []string {
|
||||
merged := map[string]string{}
|
||||
for _, kv := range os.Environ() {
|
||||
if k, v, ok := strings.Cut(kv, "="); ok {
|
||||
merged[k] = v
|
||||
}
|
||||
}
|
||||
for k, v := range generated {
|
||||
merged[k] = v
|
||||
}
|
||||
|
||||
out := make([]string, 0, len(merged))
|
||||
for k, v := range merged {
|
||||
out = append(out, k+"="+v)
|
||||
}
|
||||
// Sorted so a command is reproducible and diffable in logs.
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
func joinInts(v []int) string {
|
||||
parts := make([]string, len(v))
|
||||
for i, n := range v {
|
||||
parts[i] = strconv.Itoa(n)
|
||||
}
|
||||
return strings.Join(parts, ",")
|
||||
}
|
||||
|
||||
// validateCommon holds the checks that apply to every engine.
|
||||
func validateCommon(p Profile, gpus []GPUInfo) error {
|
||||
if p.ModelPath == "" {
|
||||
return fmt.Errorf("profile %q: model path is empty", p.Name)
|
||||
}
|
||||
if p.TP <= 0 {
|
||||
return fmt.Errorf("profile %q: tensor parallel degree must be positive", p.Name)
|
||||
}
|
||||
if p.Nodes > 1 && p.GPUsPerNode <= 0 {
|
||||
return fmt.Errorf("profile %q: gpus_per_node is required for a %d-node profile", p.Name, p.Nodes)
|
||||
}
|
||||
|
||||
perNode := p.GPUsPerNode
|
||||
if perNode == 0 {
|
||||
perNode = p.TP
|
||||
}
|
||||
nodes := max(p.Nodes, 1)
|
||||
if want := perNode * nodes; want != p.TP*max(p.PP, 1) {
|
||||
return fmt.Errorf("profile %q: %d nodes x %d GPUs does not match TP %d x PP %d",
|
||||
p.Name, nodes, perNode, p.TP, max(p.PP, 1))
|
||||
}
|
||||
if len(gpus) > 0 && len(gpus) < perNode {
|
||||
return fmt.Errorf("profile %q: needs %d GPUs per node, node has %d", p.Name, perNode, len(gpus))
|
||||
}
|
||||
|
||||
// A tensor-parallel group must be uniform: mixing classes makes the slowest
|
||||
// card set the pace for every collective, and mismatched memory makes the
|
||||
// smallest card decide what fits.
|
||||
if len(gpus) > 1 {
|
||||
first := gpus[0].Name
|
||||
for _, g := range gpus[1:] {
|
||||
if g.Name != first {
|
||||
return fmt.Errorf("profile %q: GPUs are not uniform (%q and %q); "+
|
||||
"a tensor-parallel group may not mix GPU types", p.Name, first, g.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// requireComputeCapability rejects profiles whose quantization the hardware
|
||||
// cannot execute.
|
||||
func requireComputeCapability(p Profile, gpus []GPUInfo, minCap float64, why string) error {
|
||||
for _, g := range gpus {
|
||||
cap, err := strconv.ParseFloat(g.ComputeCapability, 64)
|
||||
if err != nil {
|
||||
// Unknown capability: let it through rather than block a launch on
|
||||
// a parsing problem. The engine will refuse it clearly enough.
|
||||
continue
|
||||
}
|
||||
if cap < minCap {
|
||||
return fmt.Errorf("profile %q: %s needs compute capability %.1f, GPU %d (%s) has %.1f",
|
||||
p.Name, why, minCap, g.Index, g.Name, cap)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// quantizationOf extracts the --quantization argument, if present.
|
||||
func quantizationOf(args []string) string {
|
||||
for i, a := range args {
|
||||
if a == "--quantization" && i+1 < len(args) {
|
||||
return args[i+1]
|
||||
}
|
||||
if v, ok := strings.CutPrefix(a, "--quantization="); ok {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,425 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// buildOne returns the single command a shard produces, failing if it produced
|
||||
// a different number.
|
||||
func buildOne(t *testing.T, e Engine, s ShardSpec) Command {
|
||||
t.Helper()
|
||||
cmds := build(t, e, s)
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("expected 1 command, got %d: %s", len(cmds), summarize(cmds))
|
||||
}
|
||||
return cmds[0]
|
||||
}
|
||||
|
||||
// buildEngine returns the last command, which is the engine itself. The ones
|
||||
// before it are sidecars.
|
||||
func buildEngine(t *testing.T, e Engine, s ShardSpec) Command {
|
||||
t.Helper()
|
||||
cmds := build(t, e, s)
|
||||
return cmds[len(cmds)-1]
|
||||
}
|
||||
|
||||
func build(t *testing.T, e Engine, s ShardSpec) []Command {
|
||||
t.Helper()
|
||||
cmds, err := e.Commands(s)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) == 0 {
|
||||
t.Fatal("no commands produced")
|
||||
}
|
||||
return cmds
|
||||
}
|
||||
|
||||
// summarize keeps failure output readable: a Command carries the whole
|
||||
// environment, and printing it buries the actual problem.
|
||||
func summarize(cmds []Command) string {
|
||||
var b strings.Builder
|
||||
for i, c := range cmds {
|
||||
if i > 0 {
|
||||
b.WriteString(" | ")
|
||||
}
|
||||
fmt.Fprintf(&b, "%s: %s %s", c.Name, c.Path, strings.Join(c.Args, " "))
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// envMap turns exec-form environment back into a map for assertions.
|
||||
func envMap(env []string) map[string]string {
|
||||
m := map[string]string{}
|
||||
for _, kv := range env {
|
||||
if k, v, ok := strings.Cut(kv, "="); ok {
|
||||
m[k] = v
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// argValue returns the value following flag, or "" if the flag is absent.
|
||||
func argValue(args []string, flag string) string {
|
||||
i := slices.Index(args, flag)
|
||||
if i < 0 || i+1 >= len(args) {
|
||||
return ""
|
||||
}
|
||||
return args[i+1]
|
||||
}
|
||||
|
||||
func singleNodeSpec() ShardSpec {
|
||||
return ShardSpec{
|
||||
ReplicaID: "r1",
|
||||
ShardID: "r1-0",
|
||||
Rank: 0,
|
||||
NNodes: 1,
|
||||
GPUIndices: []int{0, 1, 2, 3},
|
||||
MgmtInterface: "eth0",
|
||||
HTTPPort: 20000,
|
||||
Profile: Profile{
|
||||
Name: "h100-tp4",
|
||||
Engine: "vllm",
|
||||
TP: 4,
|
||||
ModelPath: "/models/llama",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func multiNodeSpec(rank int) ShardSpec {
|
||||
s := singleNodeSpec()
|
||||
s.Rank = rank
|
||||
s.NNodes = 2
|
||||
s.GPUIndices = []int{0, 1, 2, 3, 4, 5, 6, 7}
|
||||
s.HCAs = []string{"mlx5_0", "mlx5_1"}
|
||||
s.LeaderAddr = "10.0.0.10"
|
||||
s.LeaderPort = 20001
|
||||
s.Profile.TP = 16
|
||||
s.Profile.GPUsPerNode = 8
|
||||
s.Profile.Nodes = 2
|
||||
return s
|
||||
}
|
||||
|
||||
// --- environment generation: the part that breaks multi-node runs ---
|
||||
|
||||
func TestSingleNodeDisablesIB(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
if env["NCCL_IB_DISABLE"] != "1" {
|
||||
t.Error("NCCL_IB_DISABLE must be 1 without a fabric, or NCCL waits for one that is not there")
|
||||
}
|
||||
if _, set := env["NCCL_IB_HCA"]; set {
|
||||
t.Error("NCCL_IB_HCA was set without any HCAs")
|
||||
}
|
||||
if env["CUDA_VISIBLE_DEVICES"] != "0,1,2,3" {
|
||||
t.Errorf("CUDA_VISIBLE_DEVICES=%q", env["CUDA_VISIBLE_DEVICES"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiNodePinsFabric(t *testing.T) {
|
||||
cmd := buildEngine(t, &VLLM{}, multiNodeSpec(0))
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
// Only the chosen fabric's adapters. Leaving this open lets NCCL pick an
|
||||
// adapter on another fabric, where the peer is unreachable, and the job
|
||||
// hangs rather than failing.
|
||||
if env["NCCL_IB_HCA"] != "mlx5_0,mlx5_1" {
|
||||
t.Errorf("NCCL_IB_HCA=%q", env["NCCL_IB_HCA"])
|
||||
}
|
||||
if env["NCCL_IB_DISABLE"] != "0" {
|
||||
t.Errorf("NCCL_IB_DISABLE=%q with a fabric present", env["NCCL_IB_DISABLE"])
|
||||
}
|
||||
// Bootstrap must stay on the management network; NCCL otherwise probes
|
||||
// every interface, including IB-over-IP that cannot route between fabrics.
|
||||
if env["NCCL_SOCKET_IFNAME"] != "eth0" || env["GLOO_SOCKET_IFNAME"] != "eth0" {
|
||||
t.Errorf("socket interfaces not pinned: %q / %q",
|
||||
env["NCCL_SOCKET_IFNAME"], env["GLOO_SOCKET_IFNAME"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileEnvOverridesGenerated(t *testing.T) {
|
||||
s := multiNodeSpec(0)
|
||||
s.Profile.Env = map[string]string{
|
||||
"NCCL_IB_HCA": "mlx5_9",
|
||||
"NCCL_NET_GDR_LEVEL": "SYS",
|
||||
}
|
||||
cmd := buildEngine(t, &VLLM{}, s)
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
if env["NCCL_IB_HCA"] != "mlx5_9" {
|
||||
t.Error("a profile override did not win; it is the escape hatch for per-cluster tuning")
|
||||
}
|
||||
if env["NCCL_NET_GDR_LEVEL"] != "SYS" {
|
||||
t.Error("an extra profile variable was dropped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessEnvironmentIsInherited(t *testing.T) {
|
||||
// Engines need PATH, HF_HOME, LD_LIBRARY_PATH and whatever else the unit
|
||||
// file set up.
|
||||
t.Setenv("LLMGW_TEST_MARKER", "present")
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
if envMap(cmd.Env)["LLMGW_TEST_MARKER"] != "present" {
|
||||
t.Error("the process environment was not inherited")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnvIsSorted(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
if !slices.IsSorted(cmd.Env) {
|
||||
t.Error("environment is not sorted; launches would not be reproducible in logs")
|
||||
}
|
||||
}
|
||||
|
||||
// --- vLLM ---
|
||||
|
||||
func TestVLLMSingleNodeArgs(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{cfg: VLLMConfig{Bin: "/opt/vllm/bin/vllm"}}, singleNodeSpec())
|
||||
if cmd.Path != "/opt/vllm/bin/vllm" {
|
||||
t.Errorf("Path=%q", cmd.Path)
|
||||
}
|
||||
if cmd.Args[0] != "serve" || cmd.Args[1] != "/models/llama" {
|
||||
t.Errorf("args start wrong: %v", cmd.Args[:2])
|
||||
}
|
||||
if argValue(cmd.Args, "--tensor-parallel-size") != "4" {
|
||||
t.Errorf("TP not passed: %v", cmd.Args)
|
||||
}
|
||||
if argValue(cmd.Args, "--port") != "20000" {
|
||||
t.Errorf("port not passed: %v", cmd.Args)
|
||||
}
|
||||
// Single node means no distributed backend flag at all.
|
||||
if slices.Contains(cmd.Args, "--distributed-executor-backend") {
|
||||
t.Error("a distributed backend was selected for a single-node profile")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMRayLeaderRunsTwoProcesses(t *testing.T) {
|
||||
// The leader starts a Ray head and then the server against it. A single
|
||||
// command would mean running vLLM against a cluster that does not exist.
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
cmds, err := v.Commands(multiNodeSpec(0))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 2 {
|
||||
t.Fatalf("expected 2 commands for the Ray leader, got %d", len(cmds))
|
||||
}
|
||||
if cmds[0].Name != "ray" || !slices.Contains(cmds[0].Args, "--head") {
|
||||
t.Errorf("first command is not the Ray head: %+v", cmds[0])
|
||||
}
|
||||
if !cmds[0].WaitForLeaderPort {
|
||||
t.Error("the server would start before the head is listening")
|
||||
}
|
||||
if cmds[1].Name != "vllm" || argValue(cmds[1].Args, "--distributed-executor-backend") != "ray" {
|
||||
t.Errorf("second command is not the vLLM server: %+v", cmds[1])
|
||||
}
|
||||
if got := envMap(cmds[1].Env)["RAY_ADDRESS"]; got != "10.0.0.10:20001" {
|
||||
t.Errorf("RAY_ADDRESS=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMRayWorkerRunsOnlyRay(t *testing.T) {
|
||||
// A Ray worker runs no vLLM process: the cluster hands its GPUs to the
|
||||
// leader. Starting a second server there would bind a port nobody uses and
|
||||
// compete for the same cards.
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
cmds, err := v.Commands(multiNodeSpec(1))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("expected 1 command for a Ray worker, got %d: %+v", len(cmds), cmds)
|
||||
}
|
||||
if cmds[0].Name != "ray" || slices.Contains(cmds[0].Args, "--head") {
|
||||
t.Errorf("worker command wrong: %+v", cmds[0])
|
||||
}
|
||||
if argValue(cmds[0].Args, "--address") != "10.0.0.10:20001" {
|
||||
t.Errorf("worker is not joining the leader: %v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMMPMode(t *testing.T) {
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}
|
||||
cmd := buildOne(t, v, multiNodeSpec(0))
|
||||
if argValue(cmd.Args, "--distributed-executor-backend") != "mp" {
|
||||
t.Errorf("backend not set: %v", cmd.Args)
|
||||
}
|
||||
env := envMap(cmd.Env)
|
||||
if env["VLLM_HOST_IP"] != "10.0.0.10" || env["VLLM_PORT"] != "20001" {
|
||||
t.Errorf("dist init not configured: %q %q", env["VLLM_HOST_IP"], env["VLLM_PORT"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMUnknownLaunchMode(t *testing.T) {
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "sorcery"}}
|
||||
if _, err := v.Commands(multiNodeSpec(0)); err == nil {
|
||||
t.Error("an unknown launch mode was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyLeaderServesHTTP(t *testing.T) {
|
||||
// The fix for the bug that stopped multi-node replicas coming up at all:
|
||||
// health-checking a worker's port times out and fails a healthy replica.
|
||||
for _, e := range []Engine{&VLLM{}, &SGLang{}} {
|
||||
if !e.ServesHTTP(multiNodeSpec(0)) {
|
||||
t.Errorf("%s: the leader must serve HTTP", e.Name())
|
||||
}
|
||||
if e.ServesHTTP(multiNodeSpec(1)) {
|
||||
t.Errorf("%s: a worker must not be expected to serve HTTP", e.Name())
|
||||
}
|
||||
if !e.ServesHTTP(singleNodeSpec()) {
|
||||
t.Errorf("%s: a single-node shard must serve HTTP", e.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoRayForSingleNodeOrMPMode(t *testing.T) {
|
||||
ray := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
if cmds := buildOne(t, ray, singleNodeSpec()); cmds.Name != "vllm" {
|
||||
t.Errorf("a single-node profile was given a Ray process: %+v", cmds)
|
||||
}
|
||||
mp := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}
|
||||
if cmds := buildOne(t, mp, multiNodeSpec(0)); cmds.Name != "vllm" {
|
||||
t.Errorf("mp mode was given a Ray process: %+v", cmds)
|
||||
}
|
||||
}
|
||||
|
||||
// --- SGLang ---
|
||||
|
||||
func TestSGLangMultiNodeArgs(t *testing.T) {
|
||||
cmd := buildOne(t, &SGLang{cfg: SGLangConfig{Python: "/opt/sglang/bin/python"}}, multiNodeSpec(1))
|
||||
if cmd.Path != "/opt/sglang/bin/python" {
|
||||
t.Errorf("Path=%q", cmd.Path)
|
||||
}
|
||||
if cmd.Args[0] != "-m" || cmd.Args[1] != "sglang.launch_server" {
|
||||
t.Errorf("not launched as a module: %v", cmd.Args[:2])
|
||||
}
|
||||
if argValue(cmd.Args, "--nnodes") != "2" || argValue(cmd.Args, "--node-rank") != "1" {
|
||||
t.Errorf("multi-node flags wrong: %v", cmd.Args)
|
||||
}
|
||||
if argValue(cmd.Args, "--dist-init-addr") != "10.0.0.10:20001" {
|
||||
t.Errorf("dist-init-addr wrong: %v", cmd.Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSGLangSingleNodeOmitsMultiNodeFlags(t *testing.T) {
|
||||
cmd := buildOne(t, &SGLang{}, singleNodeSpec())
|
||||
for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} {
|
||||
if slices.Contains(cmd.Args, flag) {
|
||||
t.Errorf("%s was passed for a single-node profile", flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileArgsComeLast(t *testing.T) {
|
||||
// Operators must be able to override anything the adapter generated.
|
||||
s := singleNodeSpec()
|
||||
s.Profile.Args = []string{"--max-model-len", "32768"}
|
||||
cmd := buildEngine(t, &VLLM{}, s)
|
||||
|
||||
i := slices.Index(cmd.Args, "--max-model-len")
|
||||
j := slices.Index(cmd.Args, "--tensor-parallel-size")
|
||||
if i < j {
|
||||
t.Errorf("profile args came before generated ones: %v", cmd.Args)
|
||||
}
|
||||
}
|
||||
|
||||
// --- validation ---
|
||||
|
||||
func h100(n int) []GPUInfo {
|
||||
out := make([]GPUInfo, n)
|
||||
for i := range out {
|
||||
out[i] = GPUInfo{Index: i, Name: "NVIDIA H100 80GB HBM3", ComputeCapability: "9.0", MemoryTotalMB: 81559}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestValidateAcceptsMatchingProfile(t *testing.T) {
|
||||
p := Profile{Name: "ok", TP: 4, ModelPath: "/m"}
|
||||
if err := (&VLLM{}).Validate(p, h100(4)); err != nil {
|
||||
t.Errorf("a valid profile was rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsMixedGPUTypes(t *testing.T) {
|
||||
// The invariant from plan §1.1: a TP group never mixes GPU types.
|
||||
gpus := h100(2)
|
||||
gpus[1].Name = "NVIDIA A100-SXM4-40GB"
|
||||
|
||||
err := (&VLLM{}).Validate(Profile{Name: "mixed", TP: 2, ModelPath: "/m"}, gpus)
|
||||
if err == nil || !strings.Contains(err.Error(), "may not mix GPU types") {
|
||||
t.Errorf("mixed GPU types accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsFP8OnOldHardware(t *testing.T) {
|
||||
gpus := h100(4)
|
||||
for i := range gpus {
|
||||
gpus[i].Name = "NVIDIA A100-SXM4-80GB"
|
||||
gpus[i].ComputeCapability = "8.0"
|
||||
}
|
||||
p := Profile{Name: "fp8", TP: 4, ModelPath: "/m", Args: []string{"--quantization", "fp8"}}
|
||||
|
||||
err := (&VLLM{}).Validate(p, gpus)
|
||||
if err == nil || !strings.Contains(err.Error(), "compute capability") {
|
||||
t.Errorf("fp8 accepted on compute capability 8.0: %v", err)
|
||||
}
|
||||
|
||||
// The same profile is fine on Hopper.
|
||||
if err := (&VLLM{}).Validate(p, h100(4)); err != nil {
|
||||
t.Errorf("fp8 rejected on H100: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsInconsistentParallelism(t *testing.T) {
|
||||
// 2 nodes x 8 GPUs cannot serve TP 4.
|
||||
p := Profile{Name: "bad", TP: 4, Nodes: 2, GPUsPerNode: 8, ModelPath: "/m"}
|
||||
if err := (&VLLM{}).Validate(p, h100(8)); err == nil {
|
||||
t.Error("inconsistent node/GPU/TP arithmetic accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsTooFewGPUs(t *testing.T) {
|
||||
if err := (&VLLM{}).Validate(Profile{Name: "big", TP: 8, ModelPath: "/m"}, h100(4)); err == nil {
|
||||
t.Error("a TP 8 profile was accepted on a 4-GPU node")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRequiresModelPath(t *testing.T) {
|
||||
if err := (&VLLM{}).Validate(Profile{Name: "nopath", TP: 1}, h100(1)); err == nil {
|
||||
t.Error("a profile without a model path was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry(t *testing.T) {
|
||||
r := NewRegistry(Config{})
|
||||
if got := r.Names(); len(got) != 2 || got[0] != "sglang" || got[1] != "vllm" {
|
||||
t.Errorf("registry contents: %v", got)
|
||||
}
|
||||
if _, err := r.Get("vllm"); err != nil {
|
||||
t.Errorf("vllm missing: %v", err)
|
||||
}
|
||||
err := func() error { _, err := r.Get("nope"); return err }()
|
||||
if err == nil || !strings.Contains(err.Error(), "sglang, vllm") {
|
||||
t.Errorf("an unknown engine should list what is available: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuantizationOf(t *testing.T) {
|
||||
for _, args := range [][]string{
|
||||
{"--quantization", "fp8"},
|
||||
{"--quantization=fp8"},
|
||||
{"--max-model-len", "100", "--quantization", "fp8"},
|
||||
} {
|
||||
if got := quantizationOf(args); got != "fp8" {
|
||||
t.Errorf("quantizationOf(%v)=%q", args, got)
|
||||
}
|
||||
}
|
||||
if got := quantizationOf([]string{"--quantization"}); got != "" {
|
||||
t.Errorf("a dangling flag should yield nothing, got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package engines
|
||||
|
||||
import "strconv"
|
||||
|
||||
// SGLangConfig locates the SGLang installation.
|
||||
type SGLangConfig struct {
|
||||
// Python is the interpreter of the SGLang virtualenv. SGLang is launched as
|
||||
// a module rather than a console script, which is the documented entry
|
||||
// point and avoids depending on the script being on PATH.
|
||||
Python string `koanf:"python"`
|
||||
}
|
||||
|
||||
// SGLang builds SGLang launch commands.
|
||||
type SGLang struct{ cfg SGLangConfig }
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*SGLang) Name() string { return "sglang" }
|
||||
|
||||
// HealthPath is SGLang's readiness endpoint.
|
||||
func (*SGLang) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is SGLang's Prometheus endpoint.
|
||||
func (*SGLang) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
func (*SGLang) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if err := validateCommon(p, gpus); err != nil {
|
||||
return err
|
||||
}
|
||||
if q := quantizationOf(p.Args); q == "fp8" {
|
||||
if err := requireComputeCapability(p, gpus, 8.9, "fp8 quantization"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API.
|
||||
//
|
||||
// Every SGLang rank runs the same launcher, but only node-rank 0 binds the HTTP
|
||||
// port; the others exist to contribute GPUs to the tensor-parallel group.
|
||||
func (*SGLang) ServesHTTP(s ShardSpec) bool { return s.Leader() }
|
||||
|
||||
// Commands returns the processes for one shard.
|
||||
//
|
||||
// SGLang supports multi-node natively, so there is only ever one process: every
|
||||
// rank runs the same launcher with its own rank number and they rendezvous at
|
||||
// the leader's dist-init address.
|
||||
func (g *SGLang) Commands(s ShardSpec) ([]Command, error) {
|
||||
c, err := g.serverCommand(s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []Command{c}, nil
|
||||
}
|
||||
|
||||
func (g *SGLang) serverCommand(s ShardSpec) (Command, error) {
|
||||
python := g.cfg.Python
|
||||
if python == "" {
|
||||
python = "python3"
|
||||
}
|
||||
p := s.Profile
|
||||
|
||||
args := []string{"-m", "sglang.launch_server",
|
||||
"--model-path", p.ModelPath,
|
||||
"--host", "0.0.0.0",
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
"--tp", strconv.Itoa(p.TP),
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--served-model-name", p.ServedModelName)
|
||||
}
|
||||
if s.NNodes > 1 {
|
||||
args = append(args,
|
||||
"--nnodes", strconv.Itoa(s.NNodes),
|
||||
"--node-rank", strconv.Itoa(s.Rank),
|
||||
"--dist-init-addr", s.LeaderAddr+":"+strconv.Itoa(s.LeaderPort),
|
||||
)
|
||||
}
|
||||
args = append(args, p.Args...)
|
||||
|
||||
return Command{Name: "sglang", Path: python, Args: args, Env: materializeEnv(baseEnv(s))}, nil
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// VLLMConfig locates the vLLM installation and picks how it is launched.
|
||||
type VLLMConfig struct {
|
||||
// Bin is the vllm executable.
|
||||
Bin string `koanf:"bin"`
|
||||
// LaunchMode selects the multi-node mechanism: "ray" or "mp".
|
||||
//
|
||||
// Which one works depends on the vLLM version, and the flags have changed
|
||||
// more than once (plan §14). Both paths exist so a cluster can pick the one
|
||||
// its version supports instead of waiting for a code change.
|
||||
LaunchMode string `koanf:"launch_mode"`
|
||||
// RayBin is used when LaunchMode is "ray".
|
||||
RayBin string `koanf:"ray_bin"`
|
||||
}
|
||||
|
||||
// VLLM builds vLLM launch commands.
|
||||
type VLLM struct{ cfg VLLMConfig }
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*VLLM) Name() string { return "vllm" }
|
||||
|
||||
// HealthPath is vLLM's readiness endpoint.
|
||||
func (*VLLM) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is vLLM's Prometheus endpoint.
|
||||
func (*VLLM) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
func (*VLLM) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if err := validateCommon(p, gpus); err != nil {
|
||||
return err
|
||||
}
|
||||
// FP8 needs Hopper or Ada; on anything older vLLM either refuses or falls
|
||||
// back silently to a slower path, and a silent fallback is worse.
|
||||
if q := quantizationOf(p.Args); q == "fp8" {
|
||||
if err := requireComputeCapability(p, gpus, 8.9, "fp8 quantization"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API.
|
||||
//
|
||||
// Under Ray a worker runs nothing but the Ray join, and under multiprocessing
|
||||
// only the driver binds the port. Either way it is the leader alone.
|
||||
func (*VLLM) ServesHTTP(s ShardSpec) bool { return s.Leader() }
|
||||
|
||||
// Commands returns the processes for one shard, in start order.
|
||||
//
|
||||
// Single node, or multi-node under multiprocessing: one server process. Under
|
||||
// Ray: the leader starts a head and then the server, a worker starts only the
|
||||
// join. A Ray worker running no vLLM process at all is correct — the cluster
|
||||
// gives the leader its GPUs.
|
||||
func (v *VLLM) Commands(s ShardSpec) ([]Command, error) {
|
||||
ray, useRay := v.rayCommand(s)
|
||||
if useRay && !s.Leader() {
|
||||
// A worker is the Ray join and nothing else. It blocks, so its exit is
|
||||
// the shard's exit, which is what the supervisor needs.
|
||||
return []Command{ray}, nil
|
||||
}
|
||||
|
||||
server, err := v.serverCommand(s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if useRay {
|
||||
return []Command{ray, server}, nil
|
||||
}
|
||||
return []Command{server}, nil
|
||||
}
|
||||
|
||||
// serverCommand produces the vLLM server process.
|
||||
func (v *VLLM) serverCommand(s ShardSpec) (Command, error) {
|
||||
bin := v.cfg.Bin
|
||||
if bin == "" {
|
||||
bin = "vllm"
|
||||
}
|
||||
p := s.Profile
|
||||
|
||||
args := []string{"serve", p.ModelPath,
|
||||
"--host", "0.0.0.0",
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
"--tensor-parallel-size", strconv.Itoa(p.TP),
|
||||
}
|
||||
if p.PP > 1 {
|
||||
args = append(args, "--pipeline-parallel-size", strconv.Itoa(p.PP))
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--served-model-name", p.ServedModelName)
|
||||
}
|
||||
|
||||
env := baseEnv(s)
|
||||
|
||||
if s.NNodes > 1 {
|
||||
switch v.cfg.LaunchMode {
|
||||
case "", "ray":
|
||||
// Ray is started separately by the supervisor; vLLM then finds the
|
||||
// cluster through the address in the environment.
|
||||
env["VLLM_HOST_IP"] = s.LeaderAddr
|
||||
env["RAY_ADDRESS"] = fmt.Sprintf("%s:%d", s.LeaderAddr, s.LeaderPort)
|
||||
args = append(args, "--distributed-executor-backend", "ray")
|
||||
case "mp":
|
||||
env["VLLM_HOST_IP"] = s.LeaderAddr
|
||||
env["VLLM_PORT"] = strconv.Itoa(s.LeaderPort)
|
||||
args = append(args, "--distributed-executor-backend", "mp")
|
||||
default:
|
||||
return Command{}, fmt.Errorf("vllm: unknown launch_mode %q (want \"ray\" or \"mp\")", v.cfg.LaunchMode)
|
||||
}
|
||||
}
|
||||
|
||||
// Profile arguments last, so an operator can override anything above.
|
||||
args = append(args, p.Args...)
|
||||
|
||||
return Command{Name: "vllm", Path: bin, Args: args, Env: materializeEnv(env)}, nil
|
||||
}
|
||||
|
||||
// rayCommand returns the command that starts or joins the Ray cluster, or
|
||||
// ok=false when the launch mode does not use Ray.
|
||||
func (v *VLLM) rayCommand(s ShardSpec) (Command, bool) {
|
||||
if s.NNodes <= 1 || (v.cfg.LaunchMode != "" && v.cfg.LaunchMode != "ray") {
|
||||
return Command{}, false
|
||||
}
|
||||
bin := v.cfg.RayBin
|
||||
if bin == "" {
|
||||
bin = "ray"
|
||||
}
|
||||
|
||||
args := []string{"start", "--block"}
|
||||
if s.Leader() {
|
||||
args = append(args, "--head", "--port", strconv.Itoa(s.LeaderPort))
|
||||
} else {
|
||||
args = append(args, "--address", fmt.Sprintf("%s:%d", s.LeaderAddr, s.LeaderPort))
|
||||
}
|
||||
args = append(args, "--num-gpus", strconv.Itoa(len(s.GPUIndices)))
|
||||
|
||||
return Command{
|
||||
Name: "ray", Path: bin, Args: args, Env: materializeEnv(baseEnv(s)),
|
||||
// The head must be listening before anything else happens: a worker
|
||||
// cannot join, and the leader's server cannot claim GPUs, until it is.
|
||||
WaitForLeaderPort: true,
|
||||
}, true
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
// HTTPHealth checks an engine over HTTP.
|
||||
//
|
||||
// Readiness is two questions, not one (plan §5.4): /health says the process is
|
||||
// up, and /v1/models says the weights finished loading and the model is being
|
||||
// served. An engine answers /health long before it can serve a request, so
|
||||
// routing on /health alone sends traffic into a model that is still loading.
|
||||
type HTTPHealth struct {
|
||||
client *http.Client
|
||||
healthPath string
|
||||
modelsPath string
|
||||
requireName bool
|
||||
}
|
||||
|
||||
// HTTPHealthOption customizes the checker.
|
||||
type HTTPHealthOption func(*HTTPHealth)
|
||||
|
||||
// WithHealthPath overrides the readiness path.
|
||||
func WithHealthPath(p string) HTTPHealthOption {
|
||||
return func(h *HTTPHealth) { h.healthPath = p }
|
||||
}
|
||||
|
||||
// WithoutModelCheck skips the /v1/models step, for engines that do not expose it.
|
||||
func WithoutModelCheck() HTTPHealthOption {
|
||||
return func(h *HTTPHealth) { h.requireName = false }
|
||||
}
|
||||
|
||||
// NewHTTPHealth builds a checker.
|
||||
func NewHTTPHealth(timeout time.Duration, opts ...HTTPHealthOption) *HTTPHealth {
|
||||
if timeout <= 0 {
|
||||
timeout = 3 * time.Second
|
||||
}
|
||||
h := &HTTPHealth{
|
||||
client: &http.Client{Timeout: timeout},
|
||||
healthPath: "/health",
|
||||
modelsPath: "/v1/models",
|
||||
requireName: true,
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(h)
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
// Ready reports whether the engine at baseURL is serving model.
|
||||
func (h *HTTPHealth) Ready(ctx context.Context, baseURL, model string) error {
|
||||
if err := h.get(ctx, baseURL+h.healthPath, nil); err != nil {
|
||||
return fmt.Errorf("health: %w", err)
|
||||
}
|
||||
if !h.requireName || model == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var list struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := h.get(ctx, baseURL+h.modelsPath, &list); err != nil {
|
||||
return fmt.Errorf("models: %w", err)
|
||||
}
|
||||
for _, d := range list.Data {
|
||||
if d.ID == model {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("engine is up but does not serve %q yet", model)
|
||||
}
|
||||
|
||||
func (h *HTTPHealth) get(ctx context.Context, url string, out any) error {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
resp, err := h.client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<10))
|
||||
return fmt.Errorf("status %s", resp.Status)
|
||||
}
|
||||
if out == nil {
|
||||
return nil
|
||||
}
|
||||
return json.NewDecoder(resp.Body).Decode(out)
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// engineStub serves configurable /health and /v1/models responses.
|
||||
func engineStub(t *testing.T, healthy bool, models ...string) *httptest.Server {
|
||||
t.Helper()
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) {
|
||||
if !healthy {
|
||||
http.Error(w, "loading", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
mux.HandleFunc("GET /v1/models", func(w http.ResponseWriter, _ *http.Request) {
|
||||
var b strings.Builder
|
||||
b.WriteString(`{"object":"list","data":[`)
|
||||
for i, m := range models {
|
||||
if i > 0 {
|
||||
b.WriteString(",")
|
||||
}
|
||||
b.WriteString(`{"id":"` + m + `"}`)
|
||||
}
|
||||
b.WriteString(`]}`)
|
||||
io.WriteString(w, b.String())
|
||||
})
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestReadyWhenServingModel(t *testing.T) {
|
||||
srv := engineStub(t, true, "llama-3.3-70b")
|
||||
if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b"); err != nil {
|
||||
t.Errorf("a serving engine was reported unready: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotReadyWhileLoading(t *testing.T) {
|
||||
// The case that matters: the process is up but the weights are not loaded.
|
||||
// Routing on /health alone would send traffic into a loading model.
|
||||
srv := engineStub(t, false, "llama-3.3-70b")
|
||||
err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b")
|
||||
if err == nil {
|
||||
t.Fatal("an unhealthy engine was reported ready")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "health") {
|
||||
t.Errorf("err=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotReadyWhenModelMissing(t *testing.T) {
|
||||
// /health passes but the model is not listed: the engine is up and serving
|
||||
// something else, or still registering.
|
||||
srv := engineStub(t, true, "some-other-model")
|
||||
err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b")
|
||||
if err == nil {
|
||||
t.Fatal("a model that is not served was reported ready")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "does not serve") {
|
||||
t.Errorf("err=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelCheckCanBeSkipped(t *testing.T) {
|
||||
srv := engineStub(t, true)
|
||||
h := NewHTTPHealth(time.Second, WithoutModelCheck())
|
||||
if err := h.Ready(t.Context(), srv.URL, "anything"); err != nil {
|
||||
t.Errorf("the model check was not skipped: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreachableEngine(t *testing.T) {
|
||||
h := NewHTTPHealth(200 * time.Millisecond)
|
||||
if err := h.Ready(t.Context(), "http://127.0.0.1:1", "m"); err == nil {
|
||||
t.Error("an unreachable engine was reported ready")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"os"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// killGroup signals a whole process group.
|
||||
//
|
||||
// Engines fork worker processes that hold GPU memory. Signalling only the
|
||||
// leader leaves those workers alive, and the next launch then fails to
|
||||
// allocate — with an error that points at the new process rather than the old.
|
||||
func killGroup(pid int, sig syscall.Signal) error {
|
||||
// The negative pid addresses the group. Setpgid at launch made the leader's
|
||||
// pid the group id.
|
||||
if err := syscall.Kill(-pid, sig); err != nil {
|
||||
if err == syscall.ESRCH {
|
||||
return os.ErrProcessDone
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// processAlive reports whether a pid exists.
|
||||
func processAlive(pid int) bool {
|
||||
if pid <= 0 {
|
||||
return false
|
||||
}
|
||||
// Signal 0 performs the permission and existence checks without delivering
|
||||
// anything.
|
||||
err := syscall.Kill(pid, 0)
|
||||
return err == nil || err == syscall.EPERM
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package supervisor
|
||||
|
||||
import "sync"
|
||||
|
||||
// ringBuffer keeps the last n lines of output.
|
||||
//
|
||||
// Engine logs are the only diagnosis available when a shard fails to load, but
|
||||
// they are also unbounded — a model loading for ten minutes can produce
|
||||
// megabytes. Keeping a bounded tail makes `ctl logs` useful without letting a
|
||||
// chatty engine consume the agent's memory.
|
||||
type ringBuffer struct {
|
||||
mu sync.Mutex
|
||||
lines []string
|
||||
next int
|
||||
full bool
|
||||
}
|
||||
|
||||
func newRingBuffer(n int) *ringBuffer {
|
||||
if n <= 0 {
|
||||
n = 256
|
||||
}
|
||||
return &ringBuffer{lines: make([]string, n)}
|
||||
}
|
||||
|
||||
func (r *ringBuffer) add(line string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.lines[r.next] = line
|
||||
r.next = (r.next + 1) % len(r.lines)
|
||||
if r.next == 0 {
|
||||
r.full = true
|
||||
}
|
||||
}
|
||||
|
||||
// tail returns up to n most recent lines, oldest first.
|
||||
func (r *ringBuffer) tail(n int) []string {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
size := r.next
|
||||
if r.full {
|
||||
size = len(r.lines)
|
||||
}
|
||||
if n <= 0 || n > size {
|
||||
n = size
|
||||
}
|
||||
|
||||
out := make([]string, 0, n)
|
||||
for i := size - n; i < size; i++ {
|
||||
idx := i
|
||||
if r.full {
|
||||
idx = (r.next + i) % len(r.lines)
|
||||
}
|
||||
out = append(out, r.lines[idx])
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
// Package supervisor runs and watches engine processes (plan §5.4).
|
||||
package supervisor
|
||||
|
||||
import "fmt"
|
||||
|
||||
// State is where a shard is in its lifecycle.
|
||||
//
|
||||
// Loading is separate from Starting because model loading dominates startup:
|
||||
// a 70B model takes minutes, and an operator watching a shard sit in one state
|
||||
// for ten minutes needs to know whether that is normal.
|
||||
type State int
|
||||
|
||||
const (
|
||||
// Pending means the shard is known but not started.
|
||||
Pending State = iota
|
||||
// Starting means the process has been launched but has not opened its port.
|
||||
Starting
|
||||
// Loading means the port is open but the model is not ready.
|
||||
Loading
|
||||
// Ready means health checks pass and the model is served.
|
||||
Ready
|
||||
// Draining means no new work; existing requests finish.
|
||||
Draining
|
||||
// Stopped means the process exited as intended.
|
||||
Stopped
|
||||
// Failed means the process exited unexpectedly or never became ready.
|
||||
Failed
|
||||
)
|
||||
|
||||
func (s State) String() string {
|
||||
switch s {
|
||||
case Pending:
|
||||
return "pending"
|
||||
case Starting:
|
||||
return "starting"
|
||||
case Loading:
|
||||
return "loading"
|
||||
case Ready:
|
||||
return "ready"
|
||||
case Draining:
|
||||
return "draining"
|
||||
case Stopped:
|
||||
return "stopped"
|
||||
case Failed:
|
||||
return "failed"
|
||||
default:
|
||||
return fmt.Sprintf("unknown(%d)", int(s))
|
||||
}
|
||||
}
|
||||
|
||||
// Terminal reports whether the shard has finished for good.
|
||||
func (s State) Terminal() bool { return s == Stopped || s == Failed }
|
||||
|
||||
// Live reports whether a process should be running in this state.
|
||||
func (s State) Live() bool {
|
||||
return s == Starting || s == Loading || s == Ready || s == Draining
|
||||
}
|
||||
@@ -0,0 +1,771 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
|
||||
// Config configures the supervisor.
|
||||
type Config struct {
|
||||
// WorkDir holds pid files and per-shard scratch space.
|
||||
WorkDir string `koanf:"work_dir"`
|
||||
// LogLines is how many lines of engine output to retain per shard.
|
||||
LogLines int `koanf:"log_lines"`
|
||||
// StopTimeout is how long a process gets to exit after SIGTERM before it is
|
||||
// killed. Engines need time to free GPU memory cleanly.
|
||||
StopTimeout time.Duration `koanf:"stop_timeout"`
|
||||
// HealthInterval is how often a starting shard is polled.
|
||||
HealthInterval time.Duration `koanf:"health_interval"`
|
||||
// RendezvousTimeout bounds the wait for a multi-node rendezvous point to
|
||||
// start listening.
|
||||
RendezvousTimeout time.Duration `koanf:"rendezvous_timeout"`
|
||||
// AliveGrace is how long a non-serving rank must stay up before it counts
|
||||
// as ready. There is no external signal for a worker, so staying alive is
|
||||
// the only evidence available.
|
||||
AliveGrace time.Duration `koanf:"alive_grace"`
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
WorkDir: "/var/lib/llmgw-agent",
|
||||
LogLines: 512,
|
||||
StopTimeout: 60 * time.Second,
|
||||
HealthInterval: 2 * time.Second,
|
||||
RendezvousTimeout: 5 * time.Minute,
|
||||
AliveGrace: 5 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() {
|
||||
d := DefaultConfig()
|
||||
if c.WorkDir == "" {
|
||||
c.WorkDir = d.WorkDir
|
||||
}
|
||||
if c.LogLines <= 0 {
|
||||
c.LogLines = d.LogLines
|
||||
}
|
||||
if c.StopTimeout <= 0 {
|
||||
c.StopTimeout = d.StopTimeout
|
||||
}
|
||||
if c.HealthInterval <= 0 {
|
||||
c.HealthInterval = d.HealthInterval
|
||||
}
|
||||
if c.RendezvousTimeout <= 0 {
|
||||
c.RendezvousTimeout = d.RendezvousTimeout
|
||||
}
|
||||
if c.AliveGrace <= 0 {
|
||||
c.AliveGrace = d.AliveGrace
|
||||
}
|
||||
}
|
||||
|
||||
// HealthChecker reports whether a shard is serving. It is an interface so the
|
||||
// supervisor can be tested without an HTTP engine.
|
||||
type HealthChecker interface {
|
||||
// Ready reports whether the shard at baseURL serves model.
|
||||
Ready(ctx context.Context, baseURL, model string) error
|
||||
}
|
||||
|
||||
// Shard is one supervised engine process.
|
||||
type Shard struct {
|
||||
ID string
|
||||
Spec engines.ShardSpec
|
||||
|
||||
mu sync.Mutex
|
||||
state State
|
||||
pid int
|
||||
err error
|
||||
started time.Time
|
||||
readyAt time.Time
|
||||
exitCode int
|
||||
|
||||
logs *ringBuffer
|
||||
cmd *exec.Cmd
|
||||
// sidecars are the shard's supporting processes, torn down with it.
|
||||
sidecars []*exec.Cmd
|
||||
cancel context.CancelFunc
|
||||
done chan struct{}
|
||||
// adopted marks a shard inherited from a previous agent: it has no
|
||||
// exec.Cmd, so it is watched by polling and killed by signal alone.
|
||||
adopted bool
|
||||
}
|
||||
|
||||
// Status is a snapshot of a shard, safe to send to the control plane.
|
||||
type Status struct {
|
||||
ID string
|
||||
State State
|
||||
PID int
|
||||
Err string
|
||||
Started time.Time
|
||||
ReadyAt time.Time
|
||||
ExitCode int
|
||||
}
|
||||
|
||||
// Status returns the shard's current status.
|
||||
func (s *Shard) Status() Status {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
st := Status{
|
||||
ID: s.ID, State: s.state, PID: s.pid,
|
||||
Started: s.started, ReadyAt: s.readyAt, ExitCode: s.exitCode,
|
||||
}
|
||||
if s.err != nil {
|
||||
st.Err = s.err.Error()
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
// Logs returns the last n lines of engine output.
|
||||
func (s *Shard) Logs(n int) []string { return s.logs.tail(n) }
|
||||
|
||||
func (s *Shard) setState(st State, err error) {
|
||||
s.mu.Lock()
|
||||
s.state = st
|
||||
if err != nil {
|
||||
s.err = err
|
||||
}
|
||||
if st == Ready && s.readyAt.IsZero() {
|
||||
s.readyAt = time.Now()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// Supervisor starts, watches and stops engine processes on this node.
|
||||
type Supervisor struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
health HealthChecker
|
||||
|
||||
mu sync.Mutex
|
||||
shards map[string]*Shard
|
||||
}
|
||||
|
||||
// New builds a supervisor.
|
||||
func New(cfg Config, health HealthChecker, log *slog.Logger) (*Supervisor, error) {
|
||||
cfg.normalize()
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if err := os.MkdirAll(cfg.WorkDir, 0o750); err != nil {
|
||||
return nil, fmt.Errorf("work dir: %w", err)
|
||||
}
|
||||
return &Supervisor{cfg: cfg, log: log, health: health, shards: map[string]*Shard{}}, nil
|
||||
}
|
||||
|
||||
// ErrAlreadyRunning means a shard with that ID is already supervised.
|
||||
var ErrAlreadyRunning = errors.New("shard is already running")
|
||||
|
||||
// Start launches a shard.
|
||||
//
|
||||
// It is idempotent by ID: a repeated StartShard for the same shard is a no-op
|
||||
// rather than a second process, because the control plane retries commands
|
||||
// (plan §5.6).
|
||||
func (s *Supervisor) Start(ctx context.Context, id string, spec engines.ShardSpec, cmds []engines.Command) (*Shard, error) {
|
||||
if len(cmds) == 0 {
|
||||
return nil, fmt.Errorf("shard %s: no commands to run", id)
|
||||
}
|
||||
s.mu.Lock()
|
||||
if existing, ok := s.shards[id]; ok {
|
||||
s.mu.Unlock()
|
||||
if existing.Status().State.Terminal() {
|
||||
s.Remove(id)
|
||||
} else {
|
||||
return existing, ErrAlreadyRunning
|
||||
}
|
||||
s.mu.Lock()
|
||||
}
|
||||
|
||||
sh := &Shard{
|
||||
ID: id,
|
||||
Spec: spec,
|
||||
state: Pending,
|
||||
logs: newRingBuffer(s.cfg.LogLines),
|
||||
started: time.Now(),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
s.shards[id] = sh
|
||||
s.mu.Unlock()
|
||||
|
||||
if err := s.launch(ctx, sh, cmds); err != nil {
|
||||
sh.setState(Failed, err)
|
||||
close(sh.done)
|
||||
return sh, err
|
||||
}
|
||||
return sh, nil
|
||||
}
|
||||
|
||||
// launch starts a shard's processes in order.
|
||||
//
|
||||
// Every command but the last is a sidecar: it keeps running for the shard's
|
||||
// lifetime and is torn down with it, but its exit is not what defines the
|
||||
// shard's. The last command is the engine, and its exit is the shard's exit.
|
||||
func (s *Supervisor) launch(ctx context.Context, sh *Shard, cmds []engines.Command) error {
|
||||
runCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
|
||||
sh.cancel = cancel
|
||||
|
||||
for i, c := range cmds[:len(cmds)-1] {
|
||||
proc, err := s.spawn(runCtx, sh, c)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return fmt.Errorf("sidecar %d (%s): %w", i, c.Name, err)
|
||||
}
|
||||
sh.mu.Lock()
|
||||
sh.sidecars = append(sh.sidecars, proc)
|
||||
sh.mu.Unlock()
|
||||
|
||||
if c.WaitForLeaderPort {
|
||||
// The next process cannot start until the rendezvous point is
|
||||
// listening: a Ray worker cannot join a head that is not up, and
|
||||
// the leader's server cannot claim GPUs from a cluster that does
|
||||
// not exist yet.
|
||||
addr := fmt.Sprintf("%s:%d", sh.Spec.LeaderAddr, sh.Spec.LeaderPort)
|
||||
if err := waitForPort(runCtx, addr, s.cfg.RendezvousTimeout); err != nil {
|
||||
cancel()
|
||||
return fmt.Errorf("waiting for %s (%s): %w", addr, c.Name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
cmd := cmds[len(cmds)-1]
|
||||
|
||||
c := exec.CommandContext(runCtx, cmd.Path, cmd.Args...)
|
||||
c.Env = cmd.Env
|
||||
c.Dir = s.cfg.WorkDir
|
||||
// Its own process group, so stopping kills the engine's children too.
|
||||
// Engines spawn workers; signalling only the parent leaves those workers
|
||||
// holding GPU memory.
|
||||
c.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
// CommandContext kills only the leader, so replace that with a group kill.
|
||||
c.Cancel = func() error { return killGroup(c.Process.Pid, syscall.SIGKILL) }
|
||||
|
||||
stdout, err := c.StdoutPipe()
|
||||
if err != nil {
|
||||
cancel()
|
||||
return err
|
||||
}
|
||||
c.Stderr = c.Stdout.(io.Writer) // interleave, as the engine writes them
|
||||
|
||||
if err := c.Start(); err != nil {
|
||||
cancel()
|
||||
return fmt.Errorf("starting %s: %w", cmd.Path, err)
|
||||
}
|
||||
|
||||
sh.mu.Lock()
|
||||
sh.cmd = c
|
||||
sh.pid = c.Process.Pid
|
||||
sh.state = Starting
|
||||
sh.mu.Unlock()
|
||||
|
||||
if err := s.writeState(sh); err != nil {
|
||||
s.log.Warn("could not write pid file; adoption after a restart will not work",
|
||||
"shard", sh.ID, "err", err)
|
||||
}
|
||||
|
||||
go s.pump(sh, stdout, cmd.Name)
|
||||
go s.wait(sh)
|
||||
return nil
|
||||
}
|
||||
|
||||
// spawn starts one process in its own group and pipes its output to the log.
|
||||
func (s *Supervisor) spawn(ctx context.Context, sh *Shard, c engines.Command) (*exec.Cmd, error) {
|
||||
proc := exec.CommandContext(ctx, c.Path, c.Args...)
|
||||
proc.Env = c.Env
|
||||
proc.Dir = s.cfg.WorkDir
|
||||
proc.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
proc.Cancel = func() error { return killGroup(proc.Process.Pid, syscall.SIGKILL) }
|
||||
|
||||
stdout, err := proc.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
proc.Stderr = proc.Stdout.(io.Writer)
|
||||
|
||||
if err := proc.Start(); err != nil {
|
||||
return nil, fmt.Errorf("starting %s: %w", c.Path, err)
|
||||
}
|
||||
go s.pump(sh, stdout, c.Name)
|
||||
// Reaped so it does not become a zombie; its exit does not end the shard.
|
||||
go func() { _ = proc.Wait() }()
|
||||
return proc, nil
|
||||
}
|
||||
|
||||
// waitForPort blocks until addr accepts a connection.
|
||||
func waitForPort(ctx context.Context, addr string, timeout time.Duration) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
|
||||
var d net.Dialer
|
||||
for {
|
||||
conn, err := d.DialContext(ctx, "tcp", addr)
|
||||
if err == nil {
|
||||
_ = conn.Close()
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return fmt.Errorf("not listening after %v: %w", timeout, ctx.Err())
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pump copies a process's output into the ring buffer and the log.
|
||||
func (s *Supervisor) pump(sh *Shard, r io.Reader, name string) {
|
||||
sc := bufio.NewScanner(r)
|
||||
// Engine tracebacks and progress bars produce long lines; the default
|
||||
// 64 KiB limit truncates them exactly when they matter.
|
||||
sc.Buffer(make([]byte, 0, 64<<10), 1<<20)
|
||||
|
||||
for sc.Scan() {
|
||||
line := sc.Text()
|
||||
if name != "" {
|
||||
// Prefixed so the two processes of a Ray shard stay distinguishable
|
||||
// in a single log tail.
|
||||
line = "[" + name + "] " + line
|
||||
}
|
||||
sh.logs.add(line)
|
||||
s.log.Debug("engine", "shard", sh.ID, "line", line)
|
||||
}
|
||||
}
|
||||
|
||||
// wait reaps the process and records why it ended.
|
||||
func (s *Supervisor) wait(sh *Shard) {
|
||||
defer close(sh.done)
|
||||
|
||||
err := sh.cmd.Wait()
|
||||
|
||||
sh.mu.Lock()
|
||||
prev := sh.state
|
||||
if ee := new(exec.ExitError); errors.As(err, &ee) {
|
||||
sh.exitCode = ee.ExitCode()
|
||||
}
|
||||
switch {
|
||||
case prev == Draining || prev == Stopped:
|
||||
// Expected: we asked it to stop.
|
||||
sh.state = Stopped
|
||||
case err == nil:
|
||||
// An engine exiting on its own is a failure, not a success: it is
|
||||
// supposed to serve until told otherwise.
|
||||
sh.state = Failed
|
||||
sh.err = errors.New("engine exited unexpectedly")
|
||||
default:
|
||||
sh.state = Failed
|
||||
sh.err = err
|
||||
}
|
||||
final, code := sh.state, sh.exitCode
|
||||
sh.mu.Unlock()
|
||||
|
||||
s.removeState(sh)
|
||||
if final == Failed {
|
||||
s.log.Error("shard failed", "shard", sh.ID, "exit_code", code, "err", err,
|
||||
"last_output", strings.Join(sh.Logs(5), " | "))
|
||||
} else {
|
||||
s.log.Info("shard stopped", "shard", sh.ID, "exit_code", code)
|
||||
}
|
||||
}
|
||||
|
||||
// WaitReady polls until the shard serves, it fails, or the timeout expires.
|
||||
//
|
||||
// The timeout comes from the profile because it is model-dependent: loading a
|
||||
// 70B model across two nodes legitimately takes tens of minutes, while a small
|
||||
// model failing to load for ten minutes means something is wrong.
|
||||
func (s *Supervisor) WaitReady(ctx context.Context, sh *Shard, baseURL, model string, timeout time.Duration) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
|
||||
t := time.NewTicker(s.cfg.HealthInterval)
|
||||
defer t.Stop()
|
||||
|
||||
sh.setState(Loading, nil)
|
||||
for {
|
||||
select {
|
||||
case <-sh.done:
|
||||
return fmt.Errorf("shard %s exited during startup: %s", sh.ID, sh.Status().Err)
|
||||
case <-ctx.Done():
|
||||
last := strings.Join(sh.Logs(10), "\n")
|
||||
return fmt.Errorf("shard %s not ready within %v; last output:\n%s", sh.ID, timeout, last)
|
||||
case <-t.C:
|
||||
if err := s.health.Ready(ctx, baseURL, model); err != nil {
|
||||
continue
|
||||
}
|
||||
sh.setState(Ready, nil)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WaitAlive is readiness for a rank that serves no HTTP.
|
||||
//
|
||||
// A worker in a multi-node replica has no endpoint to poll, so the only
|
||||
// evidence available is that it started and stayed started. That is weaker than
|
||||
// a health check, and deliberately so: the replica's real readiness is the
|
||||
// leader answering, which only happens once every worker has joined. Polling a
|
||||
// worker's port instead would time out and fail a replica that is coming up
|
||||
// perfectly well.
|
||||
func (s *Supervisor) WaitAlive(ctx context.Context, sh *Shard) error {
|
||||
sh.setState(Loading, nil)
|
||||
|
||||
select {
|
||||
case <-sh.done:
|
||||
return fmt.Errorf("shard %s exited during startup: %s", sh.ID, sh.Status().Err)
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(s.cfg.AliveGrace):
|
||||
}
|
||||
|
||||
if st := sh.Status().State; st.Terminal() {
|
||||
return fmt.Errorf("shard %s is %s: %s", sh.ID, st, sh.Status().Err)
|
||||
}
|
||||
sh.setState(Ready, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Drain marks a shard as taking no new work. The router stops routing to it;
|
||||
// the process keeps running until Stop.
|
||||
func (s *Supervisor) Drain(id string) error {
|
||||
sh, err := s.Get(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if st := sh.Status().State; !st.Live() {
|
||||
return fmt.Errorf("shard %s is %s, cannot drain", id, st)
|
||||
}
|
||||
sh.setState(Draining, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop terminates a shard: SIGTERM to the whole group, then SIGKILL after
|
||||
// StopTimeout.
|
||||
func (s *Supervisor) Stop(ctx context.Context, id string) error {
|
||||
sh, err := s.Get(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
sh.mu.Lock()
|
||||
pid := sh.pid
|
||||
sidecars := slices.Clone(sh.sidecars)
|
||||
if sh.state.Live() {
|
||||
sh.state = Draining // marks the exit as expected for wait() and poll()
|
||||
}
|
||||
sh.mu.Unlock()
|
||||
|
||||
// Sidecars go too. A Ray head left behind holds the port the next replica
|
||||
// needs, and would let a later shard join a cluster that serves nothing.
|
||||
defer func() {
|
||||
for _, sc := range sidecars {
|
||||
if sc.Process == nil {
|
||||
continue
|
||||
}
|
||||
if err := killGroup(sc.Process.Pid, syscall.SIGTERM); err != nil &&
|
||||
!errors.Is(err, os.ErrProcessDone) {
|
||||
s.log.Warn("stopping a sidecar failed", "shard", id, "pid", sc.Process.Pid, "err", err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// A pid of zero means the process never started; an adopted shard has a pid
|
||||
// but no exec.Cmd, and must still be killed by signal.
|
||||
if pid == 0 {
|
||||
sh.setState(Stopped, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Signal the group, not the process: engines fork workers that hold GPU
|
||||
// memory, and those outlive a signal sent only to the leader.
|
||||
if err := killGroup(pid, syscall.SIGTERM); err != nil && !errors.Is(err, os.ErrProcessDone) {
|
||||
s.log.Warn("SIGTERM failed", "shard", id, "err", err)
|
||||
}
|
||||
|
||||
// force kills the group outright. exec.Cmd's cancel does this for shards we
|
||||
// started; an adopted shard has none, so signal its group directly.
|
||||
force := func() {
|
||||
if sh.cancel != nil {
|
||||
sh.cancel()
|
||||
return
|
||||
}
|
||||
if err := killGroup(pid, syscall.SIGKILL); err != nil && !errors.Is(err, os.ErrProcessDone) {
|
||||
s.log.Warn("SIGKILL failed", "shard", id, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case <-sh.done:
|
||||
return nil
|
||||
case <-time.After(s.cfg.StopTimeout):
|
||||
s.log.Warn("shard did not exit in time, killing", "shard", id, "timeout", s.cfg.StopTimeout)
|
||||
force()
|
||||
case <-ctx.Done():
|
||||
force()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
select {
|
||||
case <-sh.done:
|
||||
case <-time.After(5 * time.Second):
|
||||
return fmt.Errorf("shard %s survived SIGKILL", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StopAll stops every live shard.
|
||||
func (s *Supervisor) StopAll(ctx context.Context) {
|
||||
for _, id := range s.IDs() {
|
||||
if err := s.Stop(ctx, id); err != nil {
|
||||
s.log.Warn("stopping shard failed", "shard", id, "err", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns a supervised shard.
|
||||
func (s *Supervisor) Get(id string) (*Shard, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
sh, ok := s.shards[id]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown shard %q", id)
|
||||
}
|
||||
return sh, nil
|
||||
}
|
||||
|
||||
// Remove forgets a terminated shard.
|
||||
func (s *Supervisor) Remove(id string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
delete(s.shards, id)
|
||||
}
|
||||
|
||||
// IDs lists the supervised shards.
|
||||
func (s *Supervisor) IDs() []string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]string, 0, len(s.shards))
|
||||
for id := range s.shards {
|
||||
out = append(out, id)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Statuses returns every shard's status.
|
||||
func (s *Supervisor) Statuses() []Status {
|
||||
s.mu.Lock()
|
||||
shards := make([]*Shard, 0, len(s.shards))
|
||||
for _, sh := range s.shards {
|
||||
shards = append(shards, sh)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
out := make([]Status, 0, len(shards))
|
||||
for _, sh := range shards {
|
||||
out = append(out, sh.Status())
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// --- state files, for adoption after an agent restart ---
|
||||
|
||||
func (s *Supervisor) statePath(id string) string {
|
||||
return filepath.Join(s.cfg.WorkDir, "shard-"+id+".json")
|
||||
}
|
||||
|
||||
// shardState is what survives an agent restart.
|
||||
//
|
||||
// The pid alone is not enough for two separate reasons. Pids are reused, so the
|
||||
// command line is kept as a fingerprint — adopting an unrelated process would
|
||||
// be worse than leaking the original. And the full spec is kept because a
|
||||
// recovered shard has to be described back to the control plane: without it the
|
||||
// control plane cannot tell which replica the process belongs to, and would
|
||||
// place a duplicate while the original keeps holding the GPUs (plan §6.5).
|
||||
type shardState struct {
|
||||
ShardID string `json:"shard_id"`
|
||||
ReplicaID string `json:"replica_id"`
|
||||
PID int `json:"pid"`
|
||||
Cmdline []string `json:"cmdline"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
Spec engines.ShardSpec `json:"spec"`
|
||||
}
|
||||
|
||||
// writeState records everything needed to recognise and describe the process later.
|
||||
func (s *Supervisor) writeState(sh *Shard) error {
|
||||
b, err := json.Marshal(shardState{
|
||||
ShardID: sh.ID,
|
||||
ReplicaID: sh.Spec.ReplicaID,
|
||||
PID: sh.pid,
|
||||
Cmdline: sh.cmd.Args,
|
||||
StartedAt: sh.started,
|
||||
Spec: sh.Spec,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Written atomically: a half-written state file after a crash would be
|
||||
// indistinguishable from a corrupt one, and the shard would be killed.
|
||||
tmp := s.statePath(sh.ID) + ".tmp"
|
||||
if err := os.WriteFile(tmp, b, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, s.statePath(sh.ID))
|
||||
}
|
||||
|
||||
func (s *Supervisor) removeState(sh *Shard) {
|
||||
_ = os.Remove(s.statePath(sh.ID))
|
||||
}
|
||||
|
||||
// Orphan is an engine process left behind by a previous agent.
|
||||
type Orphan struct {
|
||||
ShardID string
|
||||
ReplicaID string
|
||||
PID int
|
||||
Cmdline []string
|
||||
Spec engines.ShardSpec
|
||||
StartedAt time.Time
|
||||
// Alive reports whether the process still exists.
|
||||
Alive bool
|
||||
}
|
||||
|
||||
// FindOrphans scans the work directory for processes a previous agent started.
|
||||
func (s *Supervisor) FindOrphans() ([]Orphan, error) {
|
||||
entries, err := os.ReadDir(s.cfg.WorkDir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var out []Orphan
|
||||
for _, e := range entries {
|
||||
name := e.Name()
|
||||
if !strings.HasPrefix(name, "shard-") || !strings.HasSuffix(name, ".json") {
|
||||
continue
|
||||
}
|
||||
|
||||
b, err := os.ReadFile(filepath.Join(s.cfg.WorkDir, name))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
var st shardState
|
||||
if err := json.Unmarshal(b, &st); err != nil || st.ShardID == "" || st.PID <= 0 {
|
||||
// A corrupt state file describes a process that cannot be
|
||||
// identified, so it cannot be adopted safely. Leave it for
|
||||
// KillOrphan to clean up by shard id from the filename.
|
||||
id := strings.TrimSuffix(strings.TrimPrefix(name, "shard-"), ".json")
|
||||
out = append(out, Orphan{ShardID: id})
|
||||
continue
|
||||
}
|
||||
|
||||
out = append(out, Orphan{
|
||||
ShardID: st.ShardID,
|
||||
ReplicaID: st.ReplicaID,
|
||||
PID: st.PID,
|
||||
Cmdline: st.Cmdline,
|
||||
Spec: st.Spec,
|
||||
StartedAt: st.StartedAt,
|
||||
Alive: processAlive(st.PID),
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Adopt takes over a still-running engine from a previous agent.
|
||||
//
|
||||
// The process cannot be waited on — it is not this process's child — so it is
|
||||
// watched by polling instead. That is the whole difference between an adopted
|
||||
// shard and a started one; everything downstream treats them the same.
|
||||
func (s *Supervisor) Adopt(o Orphan) (*Shard, error) {
|
||||
if !o.Alive {
|
||||
return nil, fmt.Errorf("shard %s: process %d is gone", o.ShardID, o.PID)
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
if _, exists := s.shards[o.ShardID]; exists {
|
||||
s.mu.Unlock()
|
||||
return nil, ErrAlreadyRunning
|
||||
}
|
||||
sh := &Shard{
|
||||
ID: o.ShardID,
|
||||
Spec: o.Spec,
|
||||
state: Starting,
|
||||
pid: o.PID,
|
||||
logs: newRingBuffer(s.cfg.LogLines),
|
||||
started: o.StartedAt,
|
||||
done: make(chan struct{}),
|
||||
adopted: true,
|
||||
}
|
||||
s.shards[o.ShardID] = sh
|
||||
s.mu.Unlock()
|
||||
|
||||
// Output is gone with the previous agent; say so rather than presenting an
|
||||
// empty log as if the engine had been silent.
|
||||
sh.logs.add("[adopted from a previous agent; earlier output is not available]")
|
||||
|
||||
go s.poll(sh)
|
||||
return sh, nil
|
||||
}
|
||||
|
||||
// poll watches an adopted process, which cannot be reaped with Wait.
|
||||
func (s *Supervisor) poll(sh *Shard) {
|
||||
defer close(sh.done)
|
||||
|
||||
// Polled rather than waited on, so the interval bounds how quickly a stop
|
||||
// is observed. Short enough that Stop does not sit out its whole timeout.
|
||||
t := time.NewTicker(200 * time.Millisecond)
|
||||
defer t.Stop()
|
||||
for range t.C {
|
||||
if processAlive(sh.pid) {
|
||||
continue
|
||||
}
|
||||
sh.mu.Lock()
|
||||
expected := sh.state == Draining || sh.state == Stopped
|
||||
if expected {
|
||||
sh.state = Stopped
|
||||
} else {
|
||||
sh.state = Failed
|
||||
sh.err = errors.New("adopted engine exited")
|
||||
}
|
||||
sh.mu.Unlock()
|
||||
|
||||
s.removeState(sh)
|
||||
s.log.Info("adopted shard exited", "shard", sh.ID, "expected", expected)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// KillOrphan terminates a leftover process group and removes its pid file.
|
||||
func (s *Supervisor) KillOrphan(o Orphan) error {
|
||||
if o.Alive {
|
||||
if err := killGroup(o.PID, syscall.SIGTERM); err != nil {
|
||||
return err
|
||||
}
|
||||
// Give it a moment to free GPU memory before insisting.
|
||||
deadline := time.Now().Add(s.cfg.StopTimeout)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(o.PID) {
|
||||
break
|
||||
}
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
}
|
||||
if processAlive(o.PID) {
|
||||
_ = killGroup(o.PID, syscall.SIGKILL)
|
||||
}
|
||||
}
|
||||
// An already-missing state file is success: the previous agent may have
|
||||
// reaped the process itself between the scan and this call.
|
||||
if err := os.Remove(s.statePath(o.ShardID)); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,729 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
|
||||
// scriptedHealth reports ready once the flag is set.
|
||||
type scriptedHealth struct{ ready atomic.Bool }
|
||||
|
||||
func (h *scriptedHealth) Ready(context.Context, string, string) error {
|
||||
if h.ready.Load() {
|
||||
return nil
|
||||
}
|
||||
return errors.New("not ready")
|
||||
}
|
||||
|
||||
func newSupervisor(t *testing.T) (*Supervisor, *scriptedHealth) {
|
||||
t.Helper()
|
||||
h := &scriptedHealth{}
|
||||
s, err := New(Config{
|
||||
WorkDir: t.TempDir(),
|
||||
LogLines: 64,
|
||||
StopTimeout: 3 * time.Second,
|
||||
HealthInterval: 20 * time.Millisecond,
|
||||
}, h, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { s.StopAll(context.Background()) })
|
||||
return s, h
|
||||
}
|
||||
|
||||
// shellCommand builds the single-command sequence for a shell script.
|
||||
func shellCommand(script string) []engines.Command {
|
||||
return []engines.Command{{Path: "/bin/sh", Args: []string{"-c", script}, Env: os.Environ()}}
|
||||
}
|
||||
|
||||
func waitState(t *testing.T, sh *Shard, want State, limit time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(limit)
|
||||
for time.Now().Before(deadline) {
|
||||
if sh.Status().State == want {
|
||||
return
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("shard stayed in %s, expected %s", sh.Status().State, want)
|
||||
}
|
||||
|
||||
func TestStartAndStop(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st := sh.Status(); st.State != Starting || st.PID == 0 {
|
||||
t.Fatalf("status after start: %+v", st)
|
||||
}
|
||||
|
||||
if err := s.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st := sh.Status().State; st != Stopped {
|
||||
t.Errorf("state after stop: %s", st)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnexpectedExitIsFailure(t *testing.T) {
|
||||
// An engine exiting on its own is a failure even with status 0: it is meant
|
||||
// to serve until told otherwise.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 0"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
if !strings.Contains(sh.Status().Err, "unexpectedly") {
|
||||
t.Errorf("err=%q", sh.Status().Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNonZeroExitRecordsCode(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 42"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
if got := sh.Status().ExitCode; got != 42 {
|
||||
t.Errorf("exit code=%d, expected 42", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopKillsChildProcesses(t *testing.T) {
|
||||
// The reason for the process group: engines fork workers that hold GPU
|
||||
// memory, and signalling only the leader leaves them behind.
|
||||
s, _ := newSupervisor(t)
|
||||
dir := t.TempDir()
|
||||
marker := filepath.Join(dir, "child.pid")
|
||||
|
||||
// The shell forks a child that writes its own pid, then both sleep.
|
||||
script := fmt.Sprintf(`sh -c 'echo $$ > %s; sleep 30' & sleep 30`, marker)
|
||||
if _, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand(script)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Wait for the child to announce itself.
|
||||
var childPID int
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if b, err := os.ReadFile(marker); err == nil {
|
||||
if _, err := fmt.Sscanf(strings.TrimSpace(string(b)), "%d", &childPID); err == nil && childPID > 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
if childPID == 0 {
|
||||
t.Fatal("child never started")
|
||||
}
|
||||
|
||||
if err := s.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
deadline = time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(childPID) {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Errorf("child process %d survived the stop; it would still hold GPU memory", childPID)
|
||||
}
|
||||
|
||||
func TestStopKillsAfterTimeout(t *testing.T) {
|
||||
// A process ignoring SIGTERM must still die.
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.StopTimeout = 300 * time.Millisecond
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
|
||||
shellCommand("trap '' TERM; sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
start := time.Now()
|
||||
if err := s.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if d := time.Since(start); d > 5*time.Second {
|
||||
t.Errorf("stop took %v", d)
|
||||
}
|
||||
if processAlive(pid) {
|
||||
t.Error("the process survived SIGKILL")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartIsIdempotent(t *testing.T) {
|
||||
// The control plane retries commands, so a repeated start must not produce
|
||||
// a second process (plan §5.6).
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
first, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if !errors.Is(err, ErrAlreadyRunning) {
|
||||
t.Fatalf("expected ErrAlreadyRunning, got %v", err)
|
||||
}
|
||||
if first.Status().PID != second.Status().PID {
|
||||
t.Error("a second process was started for the same shard")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestartAfterTerminalState(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 1"))
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
|
||||
// Once it has failed, the same id may be started again.
|
||||
again, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatalf("restarting a failed shard: %v", err)
|
||||
}
|
||||
if again.Status().State != Starting {
|
||||
t.Errorf("state=%s", again.Status().State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitReady(t *testing.T) {
|
||||
s, h := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
go func() {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
h.ready.Store(true)
|
||||
}()
|
||||
|
||||
if err := s.WaitReady(t.Context(), sh, "http://x", "m", 3*time.Second); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
st := sh.Status()
|
||||
if st.State != Ready {
|
||||
t.Errorf("state=%s", st.State)
|
||||
}
|
||||
if st.ReadyAt.IsZero() {
|
||||
t.Error("ReadyAt was not recorded")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitReadyTimeoutIncludesLogs(t *testing.T) {
|
||||
// When a model fails to load, the engine output is the only diagnosis
|
||||
// available, so the error has to carry it.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
|
||||
shellCommand("echo 'CUDA out of memory'; sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond) // let the line be captured
|
||||
|
||||
err = s.WaitReady(t.Context(), sh, "http://x", "m", 200*time.Millisecond)
|
||||
if err == nil {
|
||||
t.Fatal("expected a timeout")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "CUDA out of memory") {
|
||||
t.Errorf("the error does not carry engine output: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitReadyDetectsEarlyExit(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 1"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Should return as soon as the process dies, not after the full timeout.
|
||||
start := time.Now()
|
||||
if err := s.WaitReady(t.Context(), sh, "http://x", "m", 10*time.Second); err == nil {
|
||||
t.Fatal("expected an error")
|
||||
}
|
||||
if d := time.Since(start); d > 3*time.Second {
|
||||
t.Errorf("waited %v for a process that had already exited", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDrain(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err := s.Drain("a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := sh.Status().State; got != Draining {
|
||||
t.Errorf("state=%s", got)
|
||||
}
|
||||
// A drained shard's process is still alive; only routing stops.
|
||||
if !processAlive(sh.Status().PID) {
|
||||
t.Error("draining killed the process; it should keep serving in-flight work")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogsAreCaptured(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
|
||||
shellCommand("echo line1; echo line2 >&2; echo line3; sleep 0.2"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
|
||||
logs := strings.Join(sh.Logs(10), "\n")
|
||||
for _, want := range []string{"line1", "line2", "line3"} {
|
||||
if !strings.Contains(logs, want) {
|
||||
t.Errorf("%q missing from logs:\n%s", want, logs)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- orphan adoption ---
|
||||
|
||||
func TestFindOrphans(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// Simulate an agent restart: a new supervisor over the same work dir.
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
orphans, err := s2.FindOrphans()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orphans) != 1 {
|
||||
t.Fatalf("expected 1 orphan, got %d", len(orphans))
|
||||
}
|
||||
o := orphans[0]
|
||||
if o.ShardID != "a" || o.PID != pid || !o.Alive {
|
||||
t.Errorf("orphan=%+v (want shard a, pid %d, alive)", o, pid)
|
||||
}
|
||||
// The command line is stored as a fingerprint, because pids get reused and
|
||||
// adopting an unrelated process is worse than leaking the original.
|
||||
if len(o.Cmdline) == 0 || !strings.Contains(strings.Join(o.Cmdline, " "), "sleep 30") {
|
||||
t.Errorf("cmdline fingerprint missing: %v", o.Cmdline)
|
||||
}
|
||||
// And the spec, without which the control plane could not tell which
|
||||
// replica the process belongs to and would place a duplicate.
|
||||
if o.ReplicaID != "rep-1" {
|
||||
t.Errorf("replica id not recovered: %q", o.ReplicaID)
|
||||
}
|
||||
if o.Spec.HTTPPort != 20100 {
|
||||
t.Errorf("spec not recovered: %+v", o.Spec)
|
||||
}
|
||||
|
||||
if err := s2.KillOrphan(o); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(pid) {
|
||||
break
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
if processAlive(pid) {
|
||||
t.Error("the orphan survived KillOrphan")
|
||||
}
|
||||
if _, err := os.Stat(s2.statePath("a")); !os.IsNotExist(err) {
|
||||
t.Error("the state file was left behind")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateFileRemovedOnExit(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 0"))
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
if _, err := os.Stat(s.statePath("a")); !os.IsNotExist(err) {
|
||||
t.Error("the state file outlived the process; a later scan would find a dead orphan")
|
||||
}
|
||||
}
|
||||
|
||||
// --- state machine and ring buffer ---
|
||||
|
||||
func TestStateStrings(t *testing.T) {
|
||||
for st, want := range map[State]string{
|
||||
Pending: "pending", Starting: "starting", Loading: "loading",
|
||||
Ready: "ready", Draining: "draining", Stopped: "stopped", Failed: "failed",
|
||||
} {
|
||||
if got := st.String(); got != want {
|
||||
t.Errorf("%d -> %q, want %q", int(st), got, want)
|
||||
}
|
||||
}
|
||||
if !Stopped.Terminal() || !Failed.Terminal() || Ready.Terminal() {
|
||||
t.Error("Terminal is wrong")
|
||||
}
|
||||
if !Ready.Live() || !Draining.Live() || Stopped.Live() {
|
||||
t.Error("Live is wrong")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRingBuffer(t *testing.T) {
|
||||
r := newRingBuffer(3)
|
||||
if got := r.tail(10); len(got) != 0 {
|
||||
t.Errorf("empty buffer returned %v", got)
|
||||
}
|
||||
|
||||
r.add("a")
|
||||
r.add("b")
|
||||
if got := strings.Join(r.tail(10), ","); got != "a,b" {
|
||||
t.Errorf("tail=%q", got)
|
||||
}
|
||||
|
||||
r.add("c")
|
||||
r.add("d") // evicts "a"
|
||||
if got := strings.Join(r.tail(10), ","); got != "b,c,d" {
|
||||
t.Errorf("after wrap tail=%q, want b,c,d", got)
|
||||
}
|
||||
if got := strings.Join(r.tail(2), ","); got != "c,d" {
|
||||
t.Errorf("tail(2)=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptTakesOverRunningShard(t *testing.T) {
|
||||
// The control plane restart fix depends on this: a shard that survived the
|
||||
// agent must come back under supervision, not be restarted or leaked.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20200}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// A new supervisor over the same work dir, as after an agent restart.
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 3 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
orphans, err := s2.FindOrphans()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orphans) != 1 || !orphans[0].Alive {
|
||||
t.Fatalf("orphan not found: %+v", orphans)
|
||||
}
|
||||
|
||||
adopted, err := s2.Adopt(orphans[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := adopted.Status().PID; got != pid {
|
||||
t.Errorf("adopted pid %d, expected %d", got, pid)
|
||||
}
|
||||
if adopted.Spec.ReplicaID != "rep-1" {
|
||||
t.Errorf("spec lost in adoption: %+v", adopted.Spec)
|
||||
}
|
||||
// The earlier output is gone; saying so beats presenting an empty log as if
|
||||
// the engine had been silent.
|
||||
if logs := strings.Join(adopted.Logs(5), " "); !strings.Contains(logs, "adopted") {
|
||||
t.Errorf("adoption was not noted in the log: %q", logs)
|
||||
}
|
||||
|
||||
// Stopping an adopted shard must actually kill it: it has no exec.Cmd, so
|
||||
// the signal path is the only one available.
|
||||
if err := s2.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(pid) {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Error("the adopted process survived Stop")
|
||||
}
|
||||
|
||||
func TestAdoptRefusesDeadProcess(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
if _, err := s.Adopt(Orphan{ShardID: "gone", PID: 999999, Alive: false}); err == nil {
|
||||
t.Error("a dead process was adopted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptedExitIsReported(t *testing.T) {
|
||||
// An adopted engine dying has to be noticed; nothing waits on it.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a"}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
orphans, _ := s2.FindOrphans()
|
||||
adopted, err := s2.Adopt(orphans[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Kill it behind the supervisor's back.
|
||||
_ = killGroup(pid, syscall.SIGKILL)
|
||||
waitState(t, adopted, Failed, 5*time.Second)
|
||||
}
|
||||
|
||||
func TestCorruptStateFileIsCleanedUp(t *testing.T) {
|
||||
// A half-written state file describes a process that cannot be identified,
|
||||
// so it cannot be adopted safely.
|
||||
s, _ := newSupervisor(t)
|
||||
if err := os.WriteFile(s.statePath("broken"), []byte("{not json"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
orphans, err := s.FindOrphans()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orphans) != 1 || orphans[0].ShardID != "broken" || orphans[0].Alive {
|
||||
t.Fatalf("corrupt state file not reported for cleanup: %+v", orphans)
|
||||
}
|
||||
if err := s.KillOrphan(orphans[0]); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(s.statePath("broken")); !os.IsNotExist(err) {
|
||||
t.Error("the corrupt state file was left behind")
|
||||
}
|
||||
}
|
||||
|
||||
// --- multi-process shards (plan §11, phase 4) ---
|
||||
|
||||
func TestSidecarRunsBeforeEngine(t *testing.T) {
|
||||
// The Ray shape: a supporting process starts first and keeps running, then
|
||||
// the engine starts against it.
|
||||
s, _ := newSupervisor(t)
|
||||
dir := t.TempDir()
|
||||
order := filepath.Join(dir, "order")
|
||||
|
||||
// Each writes to the order file and to stdout: the file records the
|
||||
// sequence, stdout exercises per-process log attribution.
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
|
||||
fmt.Sprintf("echo sidecar >> %s; echo hello-from-sidecar; sleep 30", order)}, Env: os.Environ()},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c",
|
||||
fmt.Sprintf("echo engine >> %s; echo hello-from-engine; sleep 30", order)}, Env: os.Environ()},
|
||||
}
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if b, err := os.ReadFile(order); err == nil && strings.Count(string(b), "\n") == 2 {
|
||||
if got := strings.Fields(string(b)); got[0] != "sidecar" || got[1] != "engine" {
|
||||
t.Fatalf("start order was %v", got)
|
||||
}
|
||||
break
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
|
||||
// The shard's pid is the engine's, not the sidecar's: the engine's exit is
|
||||
// what ends the shard.
|
||||
if sh.Status().PID == 0 {
|
||||
t.Fatal("no engine pid recorded")
|
||||
}
|
||||
// Both processes appear in one log tail, distinguishable by prefix.
|
||||
deadline = time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
logs := strings.Join(sh.Logs(20), "\n")
|
||||
if strings.Contains(logs, "[sidecar] hello-from-sidecar") &&
|
||||
strings.Contains(logs, "[engine] hello-from-engine") {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Errorf("output is not attributed per process:\n%s", strings.Join(sh.Logs(20), "\n"))
|
||||
}
|
||||
|
||||
func TestStopKillsSidecar(t *testing.T) {
|
||||
// A Ray head left behind holds the port the next replica needs, and would
|
||||
// let a later shard join a cluster that serves nothing.
|
||||
s, _ := newSupervisor(t)
|
||||
dir := t.TempDir()
|
||||
marker := filepath.Join(dir, "sidecar.pid")
|
||||
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
|
||||
fmt.Sprintf("echo $$ > %s; sleep 30", marker)}, Env: os.Environ()},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
|
||||
}
|
||||
if _, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var sidePID int
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if b, err := os.ReadFile(marker); err == nil {
|
||||
if _, err := fmt.Sscanf(strings.TrimSpace(string(b)), "%d", &sidePID); err == nil && sidePID > 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
if sidePID == 0 {
|
||||
t.Fatal("the sidecar never started")
|
||||
}
|
||||
|
||||
if err := s.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deadline = time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(sidePID) {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Errorf("the sidecar (%d) survived the stop and still holds its port", sidePID)
|
||||
}
|
||||
|
||||
func TestSidecarFailureFailsTheShard(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/nonexistent/binary", Env: os.Environ()},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
|
||||
}
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds)
|
||||
if err == nil {
|
||||
t.Fatal("a shard whose sidecar could not start was reported healthy")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "sidecar") {
|
||||
t.Errorf("the error does not identify the sidecar: %v", err)
|
||||
}
|
||||
if sh.Status().State != Failed {
|
||||
t.Errorf("state=%s", sh.Status().State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitForLeaderPortBlocksUntilListening(t *testing.T) {
|
||||
// A Ray worker cannot join a head that is not up, and the leader's server
|
||||
// cannot claim GPUs from a cluster that does not exist yet.
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.RendezvousTimeout = 5 * time.Second
|
||||
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
port := ln.Addr().(*net.TCPAddr).Port
|
||||
_ = ln.Close() // nothing listening yet
|
||||
|
||||
spec := engines.ShardSpec{LeaderAddr: "127.0.0.1", LeaderPort: port}
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
|
||||
fmt.Sprintf("sleep 0.4; nc -l %d", port)}, Env: os.Environ(),
|
||||
WaitForLeaderPort: true},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
if _, err := s.Start(t.Context(), "a", spec, cmds); err != nil {
|
||||
t.Skipf("nc is unavailable for the rendezvous stub: %v", err)
|
||||
}
|
||||
if d := time.Since(start); d < 300*time.Millisecond {
|
||||
t.Errorf("the engine started after %v, before the rendezvous point was listening", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRendezvousTimeout(t *testing.T) {
|
||||
// Nothing ever listens, so the shard must fail rather than hang forever.
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.RendezvousTimeout = 300 * time.Millisecond
|
||||
|
||||
spec := engines.ShardSpec{LeaderAddr: "127.0.0.1", LeaderPort: 1}
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c", "sleep 30"},
|
||||
Env: os.Environ(), WaitForLeaderPort: true},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
|
||||
}
|
||||
if _, err := s.Start(t.Context(), "a", spec, cmds); err == nil {
|
||||
t.Fatal("the shard did not fail when the rendezvous point never appeared")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitAliveForNonServingRank(t *testing.T) {
|
||||
// A worker has no endpoint to poll, so staying up is the only evidence
|
||||
// available. This is what stopped multi-node replicas coming up at all.
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.AliveGrace = 150 * time.Millisecond
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{Rank: 1}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.WaitAlive(t.Context(), sh); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := sh.Status().State; got != Ready {
|
||||
t.Errorf("state=%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitAliveFailsOnEarlyExit(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.AliveGrace = 2 * time.Second
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{Rank: 1}, shellCommand("exit 1"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
start := time.Now()
|
||||
if err := s.WaitAlive(t.Context(), sh); err == nil {
|
||||
t.Fatal("a worker that died immediately was reported ready")
|
||||
}
|
||||
if d := time.Since(start); d > 1500*time.Millisecond {
|
||||
t.Errorf("waited %v for a process that had already exited", d)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||
)
|
||||
|
||||
func newAgentCmd() *cobra.Command {
|
||||
var configPath string
|
||||
|
||||
var (
|
||||
server string
|
||||
nodeID string
|
||||
mgmtAddr string
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "agent",
|
||||
Short: "Node agent: discovery and engine supervision",
|
||||
Long: "Runs on every GPU node. It discovers the hardware, connects to the control\n" +
|
||||
"plane, and starts, supervises and stops inference engines on demand.\n\n" +
|
||||
"The agent dials out and keeps one stream open. Engines keep serving if the\n" +
|
||||
"agent restarts or the control plane goes away: the router talks to them\n" +
|
||||
"directly, and orphan adoption picks them up again.",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
log := logger()
|
||||
|
||||
cfg, err := config.LoadAgent(configPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if server != "" {
|
||||
cfg.Server = server
|
||||
}
|
||||
if nodeID != "" {
|
||||
cfg.NodeID = nodeID
|
||||
}
|
||||
if mgmtAddr != "" {
|
||||
cfg.MgmtAddr = mgmtAddr
|
||||
}
|
||||
if cfg.NodeID == "" {
|
||||
host, err := os.Hostname()
|
||||
if err != nil {
|
||||
return fmt.Errorf("node_id is unset and the hostname is unavailable: %w", err)
|
||||
}
|
||||
cfg.NodeID = host
|
||||
}
|
||||
if cfg.Hostname == "" {
|
||||
cfg.Hostname, _ = os.Hostname()
|
||||
}
|
||||
cfg.Version = version.Version
|
||||
|
||||
a, err := agent.New(cfg, agent.HTTP2Client(30*time.Second), log)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
log.Info("agent starting",
|
||||
"node", cfg.NodeID, "server", cfg.Server,
|
||||
"fabrics", len(cfg.Fabrics), "nvml", discovery.NVMLSupported())
|
||||
|
||||
err = a.Run(ctx)
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&configPath, "config", "", "configuration file")
|
||||
cmd.Flags().StringVar(&server, "server", "", "control plane address (overrides the config)")
|
||||
cmd.Flags().StringVar(&nodeID, "node-id", "", "node id (defaults to the hostname)")
|
||||
cmd.Flags().StringVar(&mgmtAddr, "mgmt-addr", "", "address engines are reachable on")
|
||||
cmd.AddCommand(newInventoryCmd())
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newInventoryCmd() *cobra.Command {
|
||||
var (
|
||||
asJSON bool
|
||||
smiBin string
|
||||
ibRoot string
|
||||
skipGPU bool
|
||||
useNVML bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "inventory",
|
||||
Short: "Show this node's GPUs and InfiniBand adapters",
|
||||
Long: "Reports what the scheduler would see on this node: GPU classes, health,\n" +
|
||||
"and the fabric each adapter is attached to.\n\n" +
|
||||
"Run it before adding a node to the cluster — it is the fastest way to find\n" +
|
||||
"a card another process is already using, or a port with no subnet manager.",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
inv, err := discovery.Discover(cmd.Context(), discovery.Options{
|
||||
NvidiaSMI: smiBin,
|
||||
IBRoot: ibRoot,
|
||||
PreferNVML: useNVML,
|
||||
SkipGPUs: skipGPU,
|
||||
}, logger())
|
||||
if err != nil {
|
||||
// Not fatal: a node whose GPUs cannot be read still has IB
|
||||
// worth showing, and the error is itself the useful output.
|
||||
fmt.Fprintf(os.Stderr, "%v\n", err)
|
||||
inv.Source = "unavailable"
|
||||
}
|
||||
|
||||
if asJSON {
|
||||
enc := json.NewEncoder(os.Stdout)
|
||||
enc.SetIndent("", " ")
|
||||
return enc.Encode(inv)
|
||||
}
|
||||
printInventory(inv)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
f := cmd.Flags()
|
||||
f.BoolVar(&asJSON, "json", false, "emit JSON")
|
||||
f.StringVar(&smiBin, "nvidia-smi", "", "path to nvidia-smi")
|
||||
f.StringVar(&ibRoot, "ib-root", discovery.SysfsRoot, "sysfs InfiniBand directory")
|
||||
f.BoolVar(&skipGPU, "no-gpu", false, "skip GPU discovery")
|
||||
f.BoolVar(&useNVML, "nvml", true, "prefer the NVML backend when this build has one (it is the only source of NVLink topology)")
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
func printInventory(inv discovery.Inventory) {
|
||||
w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
|
||||
|
||||
fmt.Fprintf(w, "GPUs (%d, via %s)\n", len(inv.GPUs), inv.Source)
|
||||
if len(inv.GPUs) == 0 {
|
||||
fmt.Fprintln(w, " none")
|
||||
} else {
|
||||
fmt.Fprintln(w, " IDX\tCLASS\tNAME\tMEM\tCC\tHEALTH")
|
||||
for _, g := range inv.GPUs {
|
||||
health := "ok"
|
||||
if !g.Health.OK() {
|
||||
health = g.Health.Problems[0]
|
||||
}
|
||||
fmt.Fprintf(w, " %d\t%s\t%s\t%d MiB\t%s\t%s\n",
|
||||
g.Index, g.Class(), g.Name, g.MemoryTotalMB, g.ComputeCapability, health)
|
||||
}
|
||||
fmt.Fprintf(w, " classes: %v\n", inv.Classes())
|
||||
if !discovery.NVMLSupported() {
|
||||
fmt.Fprintln(w, " NVLink topology unavailable: rebuild with -tags nvml.")
|
||||
fmt.Fprintln(w, " Without it, placement cannot keep a tensor-parallel group inside")
|
||||
fmt.Fprintln(w, " one NVLink domain (plan §6.2).")
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Fprintf(w, "\nInfiniBand adapters (%d)\n", len(inv.HCAs))
|
||||
if len(inv.HCAs) == 0 {
|
||||
// Worth stating plainly: such a node is its own single-node island and
|
||||
// can never host a multi-node replica.
|
||||
fmt.Fprintln(w, " none — this node cannot host multi-node replicas")
|
||||
} else {
|
||||
fmt.Fprintln(w, " HCA\tPORT\tSTATE\tRATE\tLINK\tGID PREFIX\tNUMA")
|
||||
for _, h := range inv.HCAs {
|
||||
for _, p := range h.Ports {
|
||||
fmt.Fprintf(w, " %s\t%d\t%s\t%d Gb/s\t%s\t%s\t%d\n",
|
||||
h.Name, p.Num, p.State, p.RateGbps, p.LinkLayer, p.GIDPrefix, h.NUMANode)
|
||||
}
|
||||
}
|
||||
fmt.Fprintln(w, "\n Fabric identity is not derived from the GID prefix by default:")
|
||||
fmt.Fprintln(w, " OpenSM's default is fe80:: on every fabric, so a shared prefix proves")
|
||||
fmt.Fprintln(w, " nothing. Label fabrics explicitly in agent.yaml (plan §5.2).")
|
||||
}
|
||||
|
||||
_ = w.Flush()
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
// ctlClient talks to a server's admin API.
|
||||
type ctlClient struct {
|
||||
base string
|
||||
http *http.Client
|
||||
}
|
||||
|
||||
func newCtlClient(addr string) *ctlClient {
|
||||
if !strings.HasPrefix(addr, "http://") && !strings.HasPrefix(addr, "https://") {
|
||||
addr = "http://" + addr
|
||||
}
|
||||
return &ctlClient{
|
||||
base: strings.TrimRight(addr, "/"),
|
||||
http: &http.Client{Timeout: 15 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *ctlClient) get(ctx context.Context, path string, out any) error {
|
||||
return c.do(ctx, http.MethodGet, path, nil, out)
|
||||
}
|
||||
|
||||
func (c *ctlClient) do(ctx context.Context, method, path string, body, out any) error {
|
||||
var rdr io.Reader
|
||||
if body != nil {
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rdr = bytes.NewReader(b)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.base+path, rdr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
|
||||
resp, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("reaching the admin API at %s: %w\n"+
|
||||
"(is the server running, and is --server pointing at its admin_listen port?)", c.base, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
var e struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<12))
|
||||
if json.Unmarshal(b, &e) == nil && e.Error != "" {
|
||||
return fmt.Errorf("%s: %s", resp.Status, e.Error)
|
||||
}
|
||||
return fmt.Errorf("%s: %s", resp.Status, strings.TrimSpace(string(b)))
|
||||
}
|
||||
if out == nil {
|
||||
return nil
|
||||
}
|
||||
return json.NewDecoder(resp.Body).Decode(out)
|
||||
}
|
||||
|
||||
func newCtlCmd() *cobra.Command {
|
||||
var (
|
||||
server string
|
||||
asJSON bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "ctl",
|
||||
Short: "Inspect and operate a running cluster",
|
||||
Long: "Talks to a server's admin API, which listens on admin_listen (:9090 by\n" +
|
||||
"default) rather than the client-facing port.",
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&server, "server", "127.0.0.1:9090", "admin API address")
|
||||
cmd.PersistentFlags().BoolVar(&asJSON, "json", false, "emit JSON")
|
||||
|
||||
client := func() *ctlClient { return newCtlClient(server) }
|
||||
|
||||
cmd.AddCommand(
|
||||
newNodesCmd(client, &asJSON),
|
||||
newModelsCmd(client, &asJSON),
|
||||
newReplicasCmd(client, &asJSON),
|
||||
newExplainCmd(client, &asJSON),
|
||||
newNodeStateCmd(client, "cordon", "cordoned",
|
||||
"Stop placing new work on a node; what runs there keeps running"),
|
||||
newNodeStateCmd(client, "uncordon", "ready",
|
||||
"Allow placement on a node again"),
|
||||
newNodeStateCmd(client, "drain", "draining",
|
||||
"Move work off a node, one replica at a time"),
|
||||
)
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newNodesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "nodes",
|
||||
Short: "List nodes and their capacity",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var nodes []adminapi.NodeView
|
||||
if err := client().get(cmd.Context(), "/api/v1/nodes", &nodes); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(nodes)
|
||||
}
|
||||
if len(nodes) == 0 {
|
||||
fmt.Println("No nodes have registered.")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "NODE\tSTATE\tGPUS\tCLASSES\tFABRICS\tREPLICAS\tLAST SEEN")
|
||||
for _, n := range nodes {
|
||||
fmt.Fprintf(w, "%s\t%s\t%d/%d free\t%s\t%s\t%d\t%s\n",
|
||||
n.ID, n.State, n.GPUFree, n.GPUTotal,
|
||||
join(n.Classes), join(n.Fabrics), n.Replicas, n.LastSeenAgo)
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
// Unhealthy GPUs are usually the answer when a node looks idle but
|
||||
// nothing schedules onto it.
|
||||
for _, n := range nodes {
|
||||
for _, p := range n.Unhealthy {
|
||||
fmt.Printf("\n%s: %s (kept out of placement)", n.ID, p)
|
||||
}
|
||||
}
|
||||
if hasUnhealthy(nodes) {
|
||||
fmt.Println()
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newModelsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "models",
|
||||
Short: "List the catalog and how many replicas each model has",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var models []adminapi.ModelView
|
||||
if err := client().get(cmd.Context(), "/api/v1/models", &models); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(models)
|
||||
}
|
||||
if len(models) == 0 {
|
||||
fmt.Println("The catalog is empty.")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "MODEL\tREPLICAS\tALIASES\tPROFILES\tFALLBACKS")
|
||||
for _, m := range models {
|
||||
replicas := fmt.Sprintf("%d/%d ready", m.Ready, m.Min)
|
||||
if m.Total > m.Ready {
|
||||
replicas += fmt.Sprintf(" (%d total)", m.Total)
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n",
|
||||
m.Name, replicas, join(m.Aliases), join(m.Profiles), join(m.Fallbacks))
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
var model string
|
||||
c := &cobra.Command{
|
||||
Use: "replicas",
|
||||
Short: "List replicas and the shards they are made of",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
path := "/api/v1/replicas"
|
||||
if model != "" {
|
||||
path += "?model=" + model
|
||||
}
|
||||
var reps []adminapi.ReplicaView
|
||||
if err := client().get(cmd.Context(), path, &reps); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(reps)
|
||||
}
|
||||
if len(reps) == 0 {
|
||||
fmt.Println("No replicas are running.")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "REPLICA\tMODEL\tSTATE\tPOOL\tNODES\tAGE\tENDPOINT")
|
||||
for _, r := range reps {
|
||||
var nodes []string
|
||||
for _, s := range r.Shards {
|
||||
nodes = append(nodes, fmt.Sprintf("%s:%v", s.Node, s.GPUs))
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
r.ID, r.Model, r.State, r.Pool, join(nodes),
|
||||
(time.Duration(r.AgeSec) * time.Second).String(), r.Endpoint)
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
for _, r := range reps {
|
||||
if r.Err != "" {
|
||||
fmt.Printf("\n%s: %s\n", r.ID, r.Err)
|
||||
}
|
||||
for _, s := range r.Shards {
|
||||
if s.Err != "" {
|
||||
fmt.Printf(" %s on %s: %s\n", s.ID, s.Node, s.Err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
c.Flags().StringVar(&model, "model", "", "only this model's replicas")
|
||||
return c
|
||||
}
|
||||
|
||||
func newExplainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "explain-placement <model>",
|
||||
Short: "Explain why a model has not been placed",
|
||||
Long: "Prints every candidate placement considered and the reason it was\n" +
|
||||
"rejected. \"Unschedulable\" on its own says nothing about which of six\n" +
|
||||
"constraints bit.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
var v adminapi.ExplainView
|
||||
if err := client().get(cmd.Context(),
|
||||
"/api/v1/explain-placement?model="+args[0], &v); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(v)
|
||||
}
|
||||
|
||||
if v.Placed && len(v.Rejections) == 0 {
|
||||
fmt.Printf("%s is placed; nothing was rejected.\n", v.Model)
|
||||
return nil
|
||||
}
|
||||
if len(v.Rejections) == 0 {
|
||||
fmt.Printf("No placement has been attempted for %s yet.\n", v.Model)
|
||||
return nil
|
||||
}
|
||||
|
||||
if v.Placed {
|
||||
fmt.Printf("%s is placed. The candidates rejected along the way:\n\n", v.Model)
|
||||
} else {
|
||||
fmt.Printf("%s could not be placed. Every candidate and why:\n\n", v.Model)
|
||||
}
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "PROFILE\tPOOL\tNODE\tREASON")
|
||||
for _, r := range v.Rejections {
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", r.Profile, dash(r.Pool), dash(r.Node), r.Reason)
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newNodeStateCmd(client func() *ctlClient, use, state, short string) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: use + " <node>",
|
||||
Short: short,
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
var out map[string]string
|
||||
err := client().do(cmd.Context(), http.MethodPost,
|
||||
"/api/v1/nodes/"+args[0]+"/state",
|
||||
map[string]string{"state": state}, &out)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf("%s is now %s\n", out["id"], out["state"])
|
||||
if state == "draining" {
|
||||
fmt.Println("Replicas will be moved off one at a time, each replacement " +
|
||||
"brought up before the original is stopped.")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// --- output helpers ---
|
||||
|
||||
func tab() *tabwriter.Writer {
|
||||
return tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
|
||||
}
|
||||
|
||||
func emitJSON(v any) error {
|
||||
enc := json.NewEncoder(os.Stdout)
|
||||
enc.SetIndent("", " ")
|
||||
return enc.Encode(v)
|
||||
}
|
||||
|
||||
func join(v []string) string {
|
||||
if len(v) == 0 {
|
||||
return "-"
|
||||
}
|
||||
return strings.Join(v, ",")
|
||||
}
|
||||
|
||||
func dash(s string) string {
|
||||
if s == "" {
|
||||
return "-"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func hasUnhealthy(nodes []adminapi.NodeView) bool {
|
||||
for _, n := range nodes {
|
||||
if len(n.Unhealthy) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,299 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/http2/h2c"
|
||||
"golang.org/x/sync/errgroup"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"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
|
||||
}
|
||||
controlMode := cfg.ControlPlaneMode()
|
||||
if !controlMode && len(cfg.Endpoints) == 0 {
|
||||
return errors.New("nothing to serve: declare models with profiles, " +
|
||||
"or give --endpoint for a router-only run")
|
||||
}
|
||||
|
||||
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()
|
||||
|
||||
// The control plane replaces the static endpoint list: agents
|
||||
// register, the reconciler places replicas, and Ready ones are
|
||||
// published to the router (plan §6.4).
|
||||
var agentHandler http.Handler
|
||||
var agentPath string
|
||||
var admin *control.Controller
|
||||
if controlMode {
|
||||
cat, err := cfg.Catalog()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
ctl, svc, err := startControlPlane(ctx, cfg, cat, rt, log)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer ctl.Close()
|
||||
admin = ctl
|
||||
agentPath, agentHandler = svc.Handler()
|
||||
|
||||
log.Info("control plane running",
|
||||
"models", len(cfg.Models), "agent_listen", cfg.AgentListen)
|
||||
}
|
||||
|
||||
adminMux, err := rt.AdminHandler()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if admin != nil {
|
||||
adminapi.New(admin).Register(adminMux)
|
||||
}
|
||||
|
||||
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: adminMux,
|
||||
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,
|
||||
)
|
||||
|
||||
var agentSrv *http.Server
|
||||
var agentLn net.Listener
|
||||
if agentHandler != nil {
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(agentPath, agentHandler)
|
||||
agentSrv = &http.Server{
|
||||
// Agents speak plaintext HTTP/2, which Go will not do
|
||||
// without being asked. mTLS replaces this (plan §5.6);
|
||||
// until then the agent port must stay on a trusted network.
|
||||
Handler: h2c.NewHandler(mux, &http2.Server{}),
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
if agentLn, err = net.Listen("tcp", cfg.AgentListen); err != nil {
|
||||
_ = apiLn.Close()
|
||||
_ = adminLn.Close()
|
||||
return fmt.Errorf("listening on %s: %w", cfg.AgentListen, err)
|
||||
}
|
||||
}
|
||||
|
||||
g, gctx := errgroup.WithContext(ctx)
|
||||
g.Go(func() error { return serve(apiSrv, apiLn) })
|
||||
g.Go(func() error { return serve(adminSrv, adminLn) })
|
||||
if agentSrv != nil {
|
||||
g.Go(func() error { return serve(agentSrv, agentLn) })
|
||||
}
|
||||
g.Go(func() error {
|
||||
<-gctx.Done()
|
||||
log.Info("shutting down")
|
||||
sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if agentSrv != nil {
|
||||
_ = agentSrv.Shutdown(sctx)
|
||||
}
|
||||
_ = 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
|
||||
}
|
||||
|
||||
// startControlPlane builds the control plane and its agent service.
|
||||
//
|
||||
// The service is both the AgentService handler and the reconciler's Dispatcher:
|
||||
// it holds the open streams, so it is the only thing that can deliver a command
|
||||
// to a node.
|
||||
func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Catalog,
|
||||
sink control.EndpointSink, log *slog.Logger,
|
||||
) (*control.Controller, *agentsvc.Service, error) {
|
||||
ccfg := control.DefaultConfig()
|
||||
ccfg.Reconciler = reconciler.Config{
|
||||
Interval: cfg.Control.Interval,
|
||||
BaseBackoff: cfg.Control.BaseBackoff,
|
||||
MaxBackoff: cfg.Control.MaxBackoff,
|
||||
UnreachableAfter: cfg.Control.UnreachableAfter,
|
||||
LostAfter: cfg.Control.LostAfter,
|
||||
MgmtInterface: cfg.Control.MgmtInterface,
|
||||
}
|
||||
if len(cfg.Control.PortRange) == 2 {
|
||||
ccfg.Reconciler.PortRange = [2]int{cfg.Control.PortRange[0], cfg.Control.PortRange[1]}
|
||||
}
|
||||
|
||||
// The service needs the inventory and the status sink, both of which the
|
||||
// controller owns; the controller needs the service as its dispatcher. The
|
||||
// cycle is broken by giving the controller a dispatcher that forwards to
|
||||
// the service once it exists.
|
||||
fwd := &dispatcherRef{}
|
||||
ctl := control.New(ccfg, cat, fwd, sink, log)
|
||||
svc := agentsvc.New(ctl.Inventory(), ctl.Reconciler(), log)
|
||||
fwd.set(svc)
|
||||
|
||||
ctl.Start(ctx)
|
||||
return ctl, svc, nil
|
||||
}
|
||||
|
||||
// dispatcherRef lets the controller be built before the service exists.
|
||||
type dispatcherRef struct {
|
||||
mu sync.RWMutex
|
||||
d reconciler.Dispatcher
|
||||
}
|
||||
|
||||
func (r *dispatcherRef) set(d reconciler.Dispatcher) {
|
||||
r.mu.Lock()
|
||||
r.d = d
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *dispatcherRef) get() (reconciler.Dispatcher, error) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
if r.d == nil {
|
||||
return nil, errors.New("the agent service is not ready yet")
|
||||
}
|
||||
return r.d, nil
|
||||
}
|
||||
|
||||
func (r *dispatcherRef) StartShard(ctx context.Context, cmd reconciler.ShardCommand) error {
|
||||
d, err := r.get()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return d.StartShard(ctx, cmd)
|
||||
}
|
||||
|
||||
func (r *dispatcherRef) StopShard(ctx context.Context, nodeID, shardID string) error {
|
||||
d, err := r.get()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return d.StopShard(ctx, nodeID, shardID)
|
||||
}
|
||||
|
||||
func (r *dispatcherRef) DrainShard(ctx context.Context, nodeID, shardID string) error {
|
||||
d, err := r.get()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return d.DrainShard(ctx, nodeID, shardID)
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
@@ -1,57 +0,0 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// 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{
|
||||
Use: "agent",
|
||||
Short: "Node agent running on each GPU host (phase 2)",
|
||||
RunE: func(_ *cobra.Command, _ []string) error {
|
||||
return errNotImplemented
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&configPath, "config", "agent.yaml", "configuration file")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newCtlCmd() *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "ctl",
|
||||
Short: "Administrative CLI (phase 3)",
|
||||
}
|
||||
for _, sub := range []struct{ use, short string }{
|
||||
{"nodes", "list nodes"},
|
||||
{"models", "show the model catalog"},
|
||||
{"replicas", "list replicas"},
|
||||
{"explain-placement", "explain why a model could not be placed"},
|
||||
{"bench", "measure profile capacity (§10)"},
|
||||
} {
|
||||
cmd.AddCommand(&cobra.Command{
|
||||
Use: sub.use,
|
||||
Short: sub.short,
|
||||
RunE: func(_ *cobra.Command, _ []string) error { return errNotImplemented },
|
||||
})
|
||||
}
|
||||
return cmd
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"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"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent"
|
||||
)
|
||||
|
||||
// LoadAgent reads an agent configuration, overlaying it on the defaults and
|
||||
// then applying LLMGW_-prefixed environment overrides.
|
||||
//
|
||||
// An empty path loads the defaults, which with --server and --node-id on the
|
||||
// command line is enough for a smoke test on a node with no IB.
|
||||
func LoadAgent(path string) (agent.Config, error) {
|
||||
cfg := agent.DefaultConfig()
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
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, validateAgent(cfg)
|
||||
}
|
||||
|
||||
// validateAgent rejects configurations that would fail confusingly at runtime.
|
||||
func validateAgent(c agent.Config) error {
|
||||
seen := map[string]struct{}{}
|
||||
for i, f := range c.Fabrics {
|
||||
if f.Fabric == "" {
|
||||
return fmt.Errorf("fabrics[%d]: name is required", i)
|
||||
}
|
||||
if _, dup := seen[f.Fabric]; dup {
|
||||
return fmt.Errorf("fabric %q is declared twice", f.Fabric)
|
||||
}
|
||||
seen[f.Fabric] = struct{}{}
|
||||
if len(f.HCAs) == 0 {
|
||||
// A fabric with no adapters is worse than no fabric at all: the
|
||||
// scheduler would place multi-node replicas here and every shard
|
||||
// would come up with NCCL unable to find a path (plan §5.3).
|
||||
return fmt.Errorf("fabric %q lists no HCAs", f.Fabric)
|
||||
}
|
||||
}
|
||||
|
||||
// The same adapter on two fabrics means one of the labels is wrong, and the
|
||||
// consequence is a replica that starts and then hangs on its first
|
||||
// all-reduce — the failure mode this labelling exists to prevent.
|
||||
byHCA := map[string]string{}
|
||||
for _, f := range c.Fabrics {
|
||||
for _, h := range f.HCAs {
|
||||
if prev, dup := byHCA[h]; dup {
|
||||
return fmt.Errorf("HCA %q is claimed by both %q and %q", h, prev, f.Fabric)
|
||||
}
|
||||
byHCA[h] = f.Fabric
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
// 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"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
)
|
||||
|
||||
// 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"`
|
||||
|
||||
// AgentListen is where node agents connect. Separate from the client port
|
||||
// so the fleet-facing surface can be firewalled independently.
|
||||
AgentListen string `koanf:"agent_listen"`
|
||||
|
||||
Router Router `koanf:"router"`
|
||||
Control Control `koanf:"control"`
|
||||
Models []catalog.Model `koanf:"models"`
|
||||
|
||||
// Endpoints is the phase-1 static replica list. It is only consulted when
|
||||
// no model declares a profile, which is what makes a router-only smoke test
|
||||
// against sim-engines possible without a control plane.
|
||||
Endpoints []Endpoint `koanf:"endpoints"`
|
||||
}
|
||||
|
||||
// Control configures the control plane.
|
||||
type Control struct {
|
||||
// Interval is how often desired and actual state are compared.
|
||||
Interval time.Duration `koanf:"interval"`
|
||||
// BaseBackoff and MaxBackoff bound how fast a failing model is retried.
|
||||
BaseBackoff time.Duration `koanf:"base_backoff"`
|
||||
MaxBackoff time.Duration `koanf:"max_backoff"`
|
||||
// UnreachableAfter and LostAfter stage node expiry. They are deliberately
|
||||
// far apart: a network blip must not tear down replicas that still serve,
|
||||
// and rebuilding a multi-node replica costs tens of minutes of loading.
|
||||
UnreachableAfter time.Duration `koanf:"unreachable_after"`
|
||||
LostAfter time.Duration `koanf:"lost_after"`
|
||||
// PortRange is the per-node range for engine HTTP and dist-init ports.
|
||||
PortRange []int `koanf:"port_range"`
|
||||
// MgmtInterface is the ethernet interface NCCL bootstraps over.
|
||||
MgmtInterface string `koanf:"mgmt_interface"`
|
||||
}
|
||||
|
||||
// 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"`
|
||||
Admission Admission `koanf:"admission"`
|
||||
}
|
||||
|
||||
// Admission configures the flow control layer (plan §7.2).
|
||||
type Admission struct {
|
||||
// Enabled turns the layer on. Off, every request goes straight through,
|
||||
// which is the v0.2.0 behaviour.
|
||||
Enabled bool `koanf:"enabled"`
|
||||
|
||||
Bands []Band `koanf:"bands"`
|
||||
DefaultBand string `koanf:"default_band"`
|
||||
// BandHeader lets a client declare its band. Until authentication lands in
|
||||
// v1.0.0 this is trusted, so the port must not be exposed to callers who
|
||||
// should not be able to promote themselves.
|
||||
BandHeader string `koanf:"band_header"`
|
||||
|
||||
QueueTTL time.Duration `koanf:"queue_ttl"`
|
||||
QueueCapacityPerBand int `koanf:"queue_capacity_per_band"`
|
||||
Interval time.Duration `koanf:"interval"`
|
||||
}
|
||||
|
||||
// Band is one priority class.
|
||||
type Band struct {
|
||||
Name string `koanf:"name"`
|
||||
Priority int `koanf:"priority"`
|
||||
SaturationCeiling float64 `koanf:"saturation_ceiling"`
|
||||
}
|
||||
|
||||
// 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"`
|
||||
}
|
||||
|
||||
// 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",
|
||||
AgentListen: ":7443",
|
||||
Control: Control{
|
||||
Interval: 5 * time.Second,
|
||||
BaseBackoff: 10 * time.Second,
|
||||
MaxBackoff: 10 * time.Minute,
|
||||
UnreachableAfter: 15 * time.Second,
|
||||
LostAfter: 60 * time.Second,
|
||||
PortRange: []int{20000, 20999},
|
||||
MgmtInterface: "eth0",
|
||||
},
|
||||
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,
|
||||
},
|
||||
Admission: Admission{
|
||||
Enabled: true,
|
||||
BandHeader: "x-llmgw-band",
|
||||
DefaultBand: "gold",
|
||||
Bands: []Band{
|
||||
{Name: "platinum", Priority: 400, SaturationCeiling: 1.00},
|
||||
{Name: "gold", Priority: 300, SaturationCeiling: 0.85},
|
||||
{Name: "silver", Priority: 200, SaturationCeiling: 0.60},
|
||||
{Name: "bronze", Priority: 100, SaturationCeiling: 0.30},
|
||||
},
|
||||
QueueTTL: 30 * time.Second,
|
||||
QueueCapacityPerBand: 2000,
|
||||
Interval: 10 * time.Millisecond,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// 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{}{}
|
||||
}
|
||||
|
||||
// The catalog enforces the rest — aliases, fallbacks, parallelism
|
||||
// arithmetic — so there is one place those rules live.
|
||||
if _, err := s.Catalog(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Catalog builds the model catalog from the configuration.
|
||||
func (s Server) Catalog() (*catalog.Catalog, error) {
|
||||
return catalog.New(s.Models)
|
||||
}
|
||||
|
||||
// ControlPlaneMode reports whether the server should run a control plane.
|
||||
//
|
||||
// A model with no profiles cannot be placed, so a configuration of bare names
|
||||
// plus a static endpoint list is the router-only mode used for testing against
|
||||
// sim-engines.
|
||||
func (s Server) ControlPlaneMode() bool {
|
||||
for _, m := range s.Models {
|
||||
if len(m.Profiles) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
)
|
||||
|
||||
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 = []catalog.Model{
|
||||
{Name: "llama-3.3-70b", Source: "/m", Aliases: []string{"qwen"},
|
||||
Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}},
|
||||
{Name: "qwen", Source: "/m",
|
||||
Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}},
|
||||
}
|
||||
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 = []catalog.Model{
|
||||
{Name: "llama-3.3-70b", Source: "/m", Aliases: []string{"llama", "llama3"},
|
||||
Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}},
|
||||
{Name: "qwen2.5-32b", Source: "/m", Aliases: []string{"qwen"},
|
||||
Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}},
|
||||
}
|
||||
if err := s.Validate(); err != nil {
|
||||
t.Errorf("valid aliases rejected: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,283 @@
|
||||
// Package adminapi is the control plane's management surface, which `ctl`
|
||||
// consumes (plan §11, phase 3).
|
||||
//
|
||||
// It is mounted on the admin listener rather than the client-facing one. The
|
||||
// two have different audiences: one serves inference to whoever holds an API
|
||||
// key, the other lets a caller cordon nodes and read the cluster's shape.
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// API serves the management endpoints.
|
||||
type API struct {
|
||||
ctl *control.Controller
|
||||
}
|
||||
|
||||
// New builds the API over a controller.
|
||||
func New(ctl *control.Controller) *API { return &API{ctl: ctl} }
|
||||
|
||||
// Register mounts the endpoints on a mux under /api/v1.
|
||||
func (a *API) Register(mux *http.ServeMux) {
|
||||
mux.HandleFunc("GET /api/v1/nodes", a.listNodes)
|
||||
mux.HandleFunc("GET /api/v1/models", a.listModels)
|
||||
mux.HandleFunc("GET /api/v1/replicas", a.listReplicas)
|
||||
mux.HandleFunc("GET /api/v1/explain-placement", a.explainPlacement)
|
||||
mux.HandleFunc("POST /api/v1/nodes/{id}/state", a.setNodeState)
|
||||
}
|
||||
|
||||
// --- responses ---
|
||||
|
||||
// NodeView is a node as `ctl nodes` shows it.
|
||||
type NodeView struct {
|
||||
ID string `json:"id"`
|
||||
Hostname string `json:"hostname,omitempty"`
|
||||
MgmtAddr string `json:"mgmt_addr,omitempty"`
|
||||
State string `json:"state"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
|
||||
GPUTotal int `json:"gpu_total"`
|
||||
GPUFree int `json:"gpu_free"`
|
||||
Classes []string `json:"gpu_classes,omitempty"`
|
||||
Fabrics []string `json:"fabrics,omitempty"`
|
||||
|
||||
// Unhealthy lists GPUs kept out of placement and why, which is usually the
|
||||
// answer when a node looks idle but nothing will schedule onto it.
|
||||
Unhealthy []string `json:"unhealthy_gpus,omitempty"`
|
||||
|
||||
LastSeenAgo string `json:"last_seen_ago"`
|
||||
AgentVersion string `json:"agent_version,omitempty"`
|
||||
Replicas int `json:"replicas"`
|
||||
}
|
||||
|
||||
// ModelView is a model as `ctl models` shows it.
|
||||
type ModelView struct {
|
||||
Name string `json:"name"`
|
||||
Aliases []string `json:"aliases,omitempty"`
|
||||
Source string `json:"source"`
|
||||
Min int `json:"min_replicas"`
|
||||
Max int `json:"max_replicas"`
|
||||
Ready int `json:"ready_replicas"`
|
||||
Total int `json:"total_replicas"`
|
||||
Fallbacks []string `json:"fallbacks,omitempty"`
|
||||
Profiles []string `json:"profiles"`
|
||||
}
|
||||
|
||||
// ReplicaView is a replica as `ctl replicas` shows it.
|
||||
type ReplicaView struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Profile string `json:"profile"`
|
||||
State string `json:"state"`
|
||||
Pool string `json:"pool,omitempty"`
|
||||
Fabric string `json:"fabric,omitempty"`
|
||||
Endpoint string `json:"endpoint,omitempty"`
|
||||
AgeSec int64 `json:"age_seconds"`
|
||||
Err string `json:"error,omitempty"`
|
||||
Shards []ShardView `json:"shards"`
|
||||
}
|
||||
|
||||
// ShardView is one shard of a replica.
|
||||
type ShardView struct {
|
||||
ID string `json:"id"`
|
||||
Node string `json:"node"`
|
||||
Rank int `json:"rank"`
|
||||
GPUs []int `json:"gpus"`
|
||||
State string `json:"state"`
|
||||
Err string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// ExplainView is why a model could not be placed.
|
||||
type ExplainView struct {
|
||||
Model string `json:"model"`
|
||||
// Placed is true when the model has replicas, in which case there is
|
||||
// nothing to explain.
|
||||
Placed bool `json:"placed"`
|
||||
// Rejections is empty when placement was never attempted, which is itself
|
||||
// the answer: the model is satisfied, or the reconciler has not run yet.
|
||||
Rejections []RejectionView `json:"rejections,omitempty"`
|
||||
}
|
||||
|
||||
// RejectionView is one rejected candidate.
|
||||
type RejectionView struct {
|
||||
Profile string `json:"profile"`
|
||||
Pool string `json:"pool,omitempty"`
|
||||
Node string `json:"node,omitempty"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
// --- handlers ---
|
||||
|
||||
func (a *API) listNodes(w http.ResponseWriter, _ *http.Request) {
|
||||
replicasOn := map[string]int{}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
for _, n := range rep.NodeIDs() {
|
||||
replicasOn[n]++
|
||||
}
|
||||
}
|
||||
|
||||
inv := a.ctl.Inventory()
|
||||
out := make([]NodeView, 0)
|
||||
for _, n := range inv.Nodes() {
|
||||
v := NodeView{
|
||||
ID: n.ID, Hostname: n.Hostname, MgmtAddr: n.MgmtAddr,
|
||||
State: n.State.String(), Labels: n.Labels,
|
||||
GPUTotal: len(n.GPUs), GPUFree: len(inv.FreeGPUs(n.ID)),
|
||||
Fabrics: n.FabricNames(), AgentVersion: n.AgentVersion,
|
||||
Replicas: replicasOn[n.ID],
|
||||
}
|
||||
seen := map[string]struct{}{}
|
||||
for _, g := range n.GPUs {
|
||||
if _, dup := seen[g.Class()]; !dup {
|
||||
seen[g.Class()] = struct{}{}
|
||||
v.Classes = append(v.Classes, g.Class())
|
||||
}
|
||||
if !g.Health.OK() {
|
||||
v.Unhealthy = append(v.Unhealthy,
|
||||
fmt.Sprintf("gpu %d: %s", g.Index, g.Health.Problems[0]))
|
||||
}
|
||||
}
|
||||
if !n.LastSeen.IsZero() {
|
||||
v.LastSeenAgo = time.Since(n.LastSeen).Round(time.Second).String()
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (a *API) listModels(w http.ResponseWriter, _ *http.Request) {
|
||||
ready, total := map[string]int{}, map[string]int{}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
total[rep.Model]++
|
||||
if rep.State == reconciler.Ready {
|
||||
ready[rep.Model]++
|
||||
}
|
||||
}
|
||||
|
||||
out := make([]ModelView, 0)
|
||||
for _, m := range a.ctl.Catalog().Models() {
|
||||
v := ModelView{
|
||||
Name: m.Name, Aliases: m.Aliases, Source: m.Source,
|
||||
Min: m.Scaling.Min, Max: m.Scaling.Max,
|
||||
Ready: ready[m.Name], Total: total[m.Name],
|
||||
Fallbacks: m.Fallbacks,
|
||||
}
|
||||
for _, p := range m.Profiles {
|
||||
v.Profiles = append(v.Profiles, p.Name)
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (a *API) listReplicas(w http.ResponseWriter, r *http.Request) {
|
||||
model := r.URL.Query().Get("model")
|
||||
|
||||
out := make([]ReplicaView, 0)
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
if model != "" && rep.Model != model {
|
||||
continue
|
||||
}
|
||||
v := ReplicaView{
|
||||
ID: rep.ID, Model: rep.Model, Profile: rep.Profile,
|
||||
State: rep.State.String(), Pool: rep.Pool, Fabric: rep.Fabric,
|
||||
Endpoint: rep.Endpoint, Err: rep.Err,
|
||||
AgeSec: int64(time.Since(rep.CreatedAt).Seconds()),
|
||||
}
|
||||
for _, s := range rep.Shards {
|
||||
v.Shards = append(v.Shards, ShardView{
|
||||
ID: s.ID, Node: s.NodeID, Rank: s.Rank,
|
||||
GPUs: s.GPUIndices, State: s.State.String(), Err: s.Err,
|
||||
})
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (a *API) explainPlacement(w http.ResponseWriter, r *http.Request) {
|
||||
model := r.URL.Query().Get("model")
|
||||
if model == "" {
|
||||
writeError(w, http.StatusBadRequest, "model is required")
|
||||
return
|
||||
}
|
||||
if _, ok := a.ctl.Catalog().Resolve(model); !ok {
|
||||
writeError(w, http.StatusNotFound, fmt.Sprintf("model %q is not in the catalog", model))
|
||||
return
|
||||
}
|
||||
|
||||
v := ExplainView{Model: model}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
if rep.Model == model && rep.State.Counts() {
|
||||
v.Placed = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if res, ok := a.ctl.Reconciler().Explain(model); ok {
|
||||
for _, rej := range res.Rejections {
|
||||
v.Rejections = append(v.Rejections, RejectionView{
|
||||
Profile: rej.Profile, Pool: rej.Pool, Node: rej.Node, Reason: rej.Reason,
|
||||
})
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, v)
|
||||
}
|
||||
|
||||
// setNodeState implements cordon, uncordon and drain.
|
||||
//
|
||||
// One endpoint rather than three verbs, because they are the same operation
|
||||
// with different targets and the difference matters: cordon stops new
|
||||
// placement, drain also moves existing work off.
|
||||
func (a *API) setNodeState(w http.ResponseWriter, r *http.Request) {
|
||||
id := r.PathValue("id")
|
||||
|
||||
var body struct {
|
||||
State string `json:"state"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
var st inventory.NodeState
|
||||
switch body.State {
|
||||
case "ready", "uncordon":
|
||||
st = inventory.Ready
|
||||
case "cordoned", "cordon":
|
||||
st = inventory.Cordoned
|
||||
case "draining", "drain":
|
||||
st = inventory.Draining
|
||||
default:
|
||||
writeError(w, http.StatusBadRequest,
|
||||
fmt.Sprintf("unknown state %q (want ready, cordoned or draining)", body.State))
|
||||
return
|
||||
}
|
||||
|
||||
if !a.ctl.Inventory().SetState(id, st) {
|
||||
writeError(w, http.StatusNotFound, fmt.Sprintf("no node %q", id))
|
||||
return
|
||||
}
|
||||
// Act on it now rather than at the next tick: an operator draining a node
|
||||
// is usually about to do something to it.
|
||||
a.ctl.Reconciler().Notify()
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]string{"id": id, "state": st.String()})
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, code int, v any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(code)
|
||||
_ = json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
|
||||
func writeError(w http.ResponseWriter, code int, msg string) {
|
||||
writeJSON(w, code, map[string]string{"error": msg})
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
)
|
||||
|
||||
// nullSink discards published endpoints; these tests are about the API.
|
||||
type nullSink struct{}
|
||||
|
||||
func (nullSink) SetEndpoints(context.Context, []endpoint.Info) {}
|
||||
|
||||
// nullDispatcher accepts every command.
|
||||
type nullDispatcher struct{}
|
||||
|
||||
func (nullDispatcher) StartShard(context.Context, reconciler.ShardCommand) error { return nil }
|
||||
func (nullDispatcher) StopShard(context.Context, string, string) error { return nil }
|
||||
func (nullDispatcher) DrainShard(context.Context, string, string) error { return nil }
|
||||
|
||||
func setup(t *testing.T, models []catalog.Model, nodes ...inventory.Node) (*API, *control.Controller) {
|
||||
t.Helper()
|
||||
cat, err := catalog.New(models)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctl := control.New(control.DefaultConfig(), cat, nullDispatcher{}, nullSink{}, slog.New(slog.DiscardHandler))
|
||||
for _, n := range nodes {
|
||||
ctl.Inventory().Upsert(n)
|
||||
}
|
||||
return New(ctl), ctl
|
||||
}
|
||||
|
||||
func gpuNode(id string, n int, fabrics ...string) inventory.Node {
|
||||
nd := inventory.Node{ID: id, State: inventory.Ready, AgentVersion: "v1"}
|
||||
for i := range n {
|
||||
nd.GPUs = append(nd.GPUs, discovery.GPU{
|
||||
Index: i, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559,
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
nd.Fabrics = append(nd.Fabrics, inventory.FabricAttachment{Fabric: f, HCAs: []string{"mlx5_0"}})
|
||||
}
|
||||
return nd
|
||||
}
|
||||
|
||||
func testModels() []catalog.Model {
|
||||
return []catalog.Model{{
|
||||
Name: "llama", Aliases: []string{"l"}, Source: "/m",
|
||||
Scaling: catalog.Scaling{Min: 2, Max: 4},
|
||||
Profiles: []catalog.Profile{{Name: "tp8", GPUClass: "h100-80g", TP: 8}},
|
||||
}}
|
||||
}
|
||||
|
||||
func call(t *testing.T, a *API, method, path string, body string, out any) int {
|
||||
t.Helper()
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
|
||||
var r *http.Request
|
||||
if body != "" {
|
||||
r = httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
} else {
|
||||
r = httptest.NewRequest(method, path, nil)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, r)
|
||||
|
||||
if out != nil && rec.Code == http.StatusOK {
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), out); err != nil {
|
||||
t.Fatalf("decoding %s: %v (%s)", path, err, rec.Body.String())
|
||||
}
|
||||
}
|
||||
return rec.Code
|
||||
}
|
||||
|
||||
func TestListNodes(t *testing.T) {
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8, "ib-a"))
|
||||
|
||||
var nodes []NodeView
|
||||
if code := call(t, a, "GET", "/api/v1/nodes", "", &nodes); code != 200 {
|
||||
t.Fatalf("status %d", code)
|
||||
}
|
||||
if len(nodes) != 1 {
|
||||
t.Fatalf("got %d nodes", len(nodes))
|
||||
}
|
||||
n := nodes[0]
|
||||
if n.ID != "n1" || n.State != "ready" {
|
||||
t.Errorf("%+v", n)
|
||||
}
|
||||
if n.GPUTotal != 8 || n.GPUFree != 8 {
|
||||
t.Errorf("gpus %d/%d", n.GPUFree, n.GPUTotal)
|
||||
}
|
||||
if len(n.Classes) != 1 || n.Classes[0] != "h100-80g" {
|
||||
t.Errorf("classes=%v", n.Classes)
|
||||
}
|
||||
if len(n.Fabrics) != 1 || n.Fabrics[0] != "ib-a" {
|
||||
t.Errorf("fabrics=%v", n.Fabrics)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnhealthyGPUsAreSurfaced(t *testing.T) {
|
||||
// Usually the answer when a node looks idle but nothing schedules onto it.
|
||||
n := gpuNode("n1", 4)
|
||||
n.GPUs[1].Health.Problems = []string{"3 uncorrected ECC errors"}
|
||||
a, _ := setup(t, testModels(), n)
|
||||
|
||||
var nodes []NodeView
|
||||
call(t, a, "GET", "/api/v1/nodes", "", &nodes)
|
||||
|
||||
if len(nodes[0].Unhealthy) != 1 {
|
||||
t.Fatalf("unhealthy=%v", nodes[0].Unhealthy)
|
||||
}
|
||||
if !strings.Contains(nodes[0].Unhealthy[0], "ECC") {
|
||||
t.Errorf("the reason was lost: %q", nodes[0].Unhealthy[0])
|
||||
}
|
||||
// And such a GPU is not counted as free.
|
||||
if nodes[0].GPUFree != 3 {
|
||||
t.Errorf("free=%d, an unhealthy GPU was counted", nodes[0].GPUFree)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListModels(t *testing.T) {
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8))
|
||||
|
||||
var models []ModelView
|
||||
call(t, a, "GET", "/api/v1/models", "", &models)
|
||||
|
||||
if len(models) != 1 {
|
||||
t.Fatalf("got %d models", len(models))
|
||||
}
|
||||
m := models[0]
|
||||
if m.Name != "llama" || m.Min != 2 || m.Max != 4 {
|
||||
t.Errorf("%+v", m)
|
||||
}
|
||||
if len(m.Aliases) != 1 || m.Aliases[0] != "l" {
|
||||
t.Errorf("aliases=%v", m.Aliases)
|
||||
}
|
||||
if len(m.Profiles) != 1 || m.Profiles[0] != "tp8" {
|
||||
t.Errorf("profiles=%v", m.Profiles)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplainPlacementForUnplaceableModel(t *testing.T) {
|
||||
// Four GPUs, a profile that needs eight.
|
||||
a, ctl := setup(t, testModels(), gpuNode("n1", 4))
|
||||
ctl.Reconciler().Reconcile(t.Context())
|
||||
|
||||
var v ExplainView
|
||||
if code := call(t, a, "GET", "/api/v1/explain-placement?model=llama", "", &v); code != 200 {
|
||||
t.Fatalf("status %d", code)
|
||||
}
|
||||
if v.Placed {
|
||||
t.Error("an unplaceable model was reported as placed")
|
||||
}
|
||||
if len(v.Rejections) == 0 {
|
||||
t.Fatal("no rejections were recorded; the operator learns nothing")
|
||||
}
|
||||
found := false
|
||||
for _, r := range v.Rejections {
|
||||
if strings.Contains(r.Reason, "needs 8") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("the reason does not say what was missing: %+v", v.Rejections)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplainPlacementRequiresAKnownModel(t *testing.T) {
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8))
|
||||
|
||||
if code := call(t, a, "GET", "/api/v1/explain-placement", "", nil); code != 400 {
|
||||
t.Errorf("missing model gave %d, expected 400", code)
|
||||
}
|
||||
if code := call(t, a, "GET", "/api/v1/explain-placement?model=nope", "", nil); code != 404 {
|
||||
t.Errorf("unknown model gave %d, expected 404", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetNodeState(t *testing.T) {
|
||||
a, ctl := setup(t, testModels(), gpuNode("n1", 8))
|
||||
|
||||
for _, c := range []struct {
|
||||
body string
|
||||
want inventory.NodeState
|
||||
}{
|
||||
{`{"state":"cordon"}`, inventory.Cordoned},
|
||||
{`{"state":"drain"}`, inventory.Draining},
|
||||
{`{"state":"uncordon"}`, inventory.Ready},
|
||||
{`{"state":"cordoned"}`, inventory.Cordoned},
|
||||
} {
|
||||
if code := call(t, a, "POST", "/api/v1/nodes/n1/state", c.body, nil); code != 200 {
|
||||
t.Fatalf("%s gave %d", c.body, code)
|
||||
}
|
||||
n, _ := ctl.Inventory().Get("n1")
|
||||
if n.State != c.want {
|
||||
t.Errorf("%s left the node %s, expected %s", c.body, n.State, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetNodeStateRejectsNonsense(t *testing.T) {
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8))
|
||||
|
||||
if code := call(t, a, "POST", "/api/v1/nodes/n1/state", `{"state":"banana"}`, nil); code != 400 {
|
||||
t.Errorf("an unknown state gave %d, expected 400", code)
|
||||
}
|
||||
if code := call(t, a, "POST", "/api/v1/nodes/nope/state", `{"state":"cordon"}`, nil); code != 404 {
|
||||
t.Errorf("an unknown node gave %d, expected 404", code)
|
||||
}
|
||||
if code := call(t, a, "POST", "/api/v1/nodes/n1/state", `not json`, nil); code != 400 {
|
||||
t.Errorf("malformed body gave %d, expected 400", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyClusterReturnsArraysNotNull(t *testing.T) {
|
||||
// A JSON null would make every consumer special-case it.
|
||||
a, _ := setup(t, testModels())
|
||||
|
||||
for _, path := range []string{"/api/v1/nodes", "/api/v1/replicas"} {
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, httptest.NewRequest("GET", path, nil))
|
||||
|
||||
if got := strings.TrimSpace(rec.Body.String()); got != "[]" {
|
||||
t.Errorf("%s returned %q, expected []", path, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
// Package agentsvc is the control plane's end of the agent stream (plan §5.6).
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// NodeFromRegister converts a registration into an inventory node.
|
||||
func NodeFromRegister(r *llmgwv1.Register) inventory.Node {
|
||||
n := inventory.Node{
|
||||
ID: r.GetNodeId(),
|
||||
Hostname: r.GetHostname(),
|
||||
MgmtAddr: r.GetMgmtAddr(),
|
||||
Labels: r.GetLabels(),
|
||||
State: inventory.Ready,
|
||||
AgentVersion: r.GetAgentVersion(),
|
||||
EngineVersions: r.GetEngineVersions(),
|
||||
}
|
||||
for _, g := range r.GetGpus() {
|
||||
n.GPUs = append(n.GPUs, discovery.GPU{
|
||||
Index: int(g.GetIndex()),
|
||||
UUID: g.GetUuid(),
|
||||
Name: g.GetName(),
|
||||
MemoryTotalMB: int(g.GetMemoryTotalMb()),
|
||||
MemoryUsedMB: int(g.GetMemoryUsedMb()),
|
||||
ComputeCapability: g.GetComputeCapability(),
|
||||
NUMANode: int(g.GetNumaNode()),
|
||||
PCIBusID: g.GetPciBusId(),
|
||||
DriverVersion: g.GetDriverVersion(),
|
||||
CUDAVersion: g.GetCudaVersion(),
|
||||
NVLinkPeers: intSlice(g.GetNvlinkPeers()),
|
||||
Health: discovery.GPUHealth{
|
||||
ECCErrors: int(g.GetEccErrors()),
|
||||
TemperatureC: int(g.GetTemperatureC()),
|
||||
Problems: g.GetProblems(),
|
||||
},
|
||||
})
|
||||
}
|
||||
for _, f := range r.GetFabrics() {
|
||||
n.Fabrics = append(n.Fabrics, inventory.FabricAttachment{
|
||||
Fabric: f.GetFabric(),
|
||||
HCAs: f.GetHcas(),
|
||||
RateGbps: int(f.GetRateGbps()),
|
||||
})
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// RegisterFromInventory builds the registration an agent sends.
|
||||
func RegisterFromInventory(nodeID, hostname, mgmtAddr, agentVersion string,
|
||||
labels map[string]string, inv discovery.Inventory, fabrics []inventory.FabricAttachment,
|
||||
) *llmgwv1.Register {
|
||||
r := &llmgwv1.Register{
|
||||
NodeId: nodeID, Hostname: hostname, MgmtAddr: mgmtAddr,
|
||||
Labels: labels, AgentVersion: agentVersion,
|
||||
}
|
||||
for _, g := range inv.GPUs {
|
||||
r.Gpus = append(r.Gpus, &llmgwv1.GPU{
|
||||
Index: int32(g.Index),
|
||||
Uuid: g.UUID,
|
||||
Name: g.Name,
|
||||
MemoryTotalMb: int32(g.MemoryTotalMB),
|
||||
MemoryUsedMb: int32(g.MemoryUsedMB),
|
||||
ComputeCapability: g.ComputeCapability,
|
||||
NumaNode: int32(g.NUMANode),
|
||||
PciBusId: g.PCIBusID,
|
||||
DriverVersion: g.DriverVersion,
|
||||
CudaVersion: g.CUDAVersion,
|
||||
NvlinkPeers: int32Slice(g.NVLinkPeers),
|
||||
Problems: g.Health.Problems,
|
||||
TemperatureC: int32(g.Health.TemperatureC),
|
||||
EccErrors: int32(g.Health.ECCErrors),
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
r.Fabrics = append(r.Fabrics, &llmgwv1.FabricAttachment{
|
||||
Fabric: f.Fabric, Hcas: f.HCAs, RateGbps: int32(f.RateGbps),
|
||||
})
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// AdoptedFromRegister converts an agent's running-shard report into the
|
||||
// reconciler's view.
|
||||
func AdoptedFromRegister(nodeID string, shards []*llmgwv1.RunningShard) []reconciler.AdoptedShard {
|
||||
out := make([]reconciler.AdoptedShard, 0, len(shards))
|
||||
for _, s := range shards {
|
||||
a := reconciler.AdoptedShard{
|
||||
NodeID: nodeID,
|
||||
ReplicaID: s.GetReplicaId(),
|
||||
ShardID: s.GetShardId(),
|
||||
Rank: int(s.GetRank()),
|
||||
GPUIndices: intSlice(s.GetGpuIndices()),
|
||||
HCAs: s.GetHcas(),
|
||||
HTTPPort: int(s.GetHttpPort()),
|
||||
State: StateFromWire(s.GetState()),
|
||||
Model: s.GetModel(),
|
||||
Profile: s.GetProfile(),
|
||||
Pool: s.GetPool(),
|
||||
Fabric: s.GetFabric(),
|
||||
}
|
||||
if a.Rank == 0 && s.GetLeaderAddr() != "" {
|
||||
a.Endpoint = fmt.Sprintf("http://%s:%d", s.GetLeaderAddr(), s.GetHttpPort())
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// StartShardFromCommand converts a reconciler command into a wire message.
|
||||
func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
|
||||
p := cmd.Spec.Profile
|
||||
return &llmgwv1.StartShard{
|
||||
ReplicaId: cmd.ReplicaID,
|
||||
ShardId: cmd.ShardID,
|
||||
Rank: int32(cmd.Spec.Rank),
|
||||
Nnodes: int32(cmd.Spec.NNodes),
|
||||
LeaderAddr: cmd.Spec.LeaderAddr,
|
||||
LeaderPort: int32(cmd.Spec.LeaderPort),
|
||||
GpuIndices: int32Slice(cmd.Spec.GPUIndices),
|
||||
Hcas: cmd.Spec.HCAs,
|
||||
MgmtInterface: cmd.Spec.MgmtInterface,
|
||||
HttpPort: int32(cmd.Spec.HTTPPort),
|
||||
Profile: &llmgwv1.EngineProfile{
|
||||
Name: p.Name, Engine: p.Engine, GpuClass: p.GPUClass,
|
||||
Tp: int32(p.TP), Pp: int32(p.PP),
|
||||
GpusPerNode: int32(p.GPUsPerNode), Nodes: int32(p.Nodes),
|
||||
Args: p.Args, Env: p.Env,
|
||||
ModelPath: p.ModelPath, ServedModelName: p.ServedModelName,
|
||||
},
|
||||
StartupTimeoutMs: cmd.StartupTimeout.Milliseconds(),
|
||||
Pool: cmd.Pool,
|
||||
Fabric: cmd.Fabric,
|
||||
}
|
||||
}
|
||||
|
||||
// ShardSpecFromStart converts a wire message into an engine shard spec.
|
||||
func ShardSpecFromStart(m *llmgwv1.StartShard) (engines.ShardSpec, time.Duration) {
|
||||
p := m.GetProfile()
|
||||
spec := engines.ShardSpec{
|
||||
ReplicaID: m.GetReplicaId(),
|
||||
ShardID: m.GetShardId(),
|
||||
Rank: int(m.GetRank()),
|
||||
NNodes: int(m.GetNnodes()),
|
||||
LeaderAddr: m.GetLeaderAddr(),
|
||||
LeaderPort: int(m.GetLeaderPort()),
|
||||
GPUIndices: intSlice(m.GetGpuIndices()),
|
||||
HCAs: m.GetHcas(),
|
||||
MgmtInterface: m.GetMgmtInterface(),
|
||||
HTTPPort: int(m.GetHttpPort()),
|
||||
Pool: m.GetPool(),
|
||||
Fabric: m.GetFabric(),
|
||||
Profile: engines.Profile{
|
||||
Name: p.GetName(), Engine: p.GetEngine(), GPUClass: p.GetGpuClass(),
|
||||
TP: int(p.GetTp()), PP: int(p.GetPp()),
|
||||
GPUsPerNode: int(p.GetGpusPerNode()), Nodes: int(p.GetNodes()),
|
||||
Args: p.GetArgs(), Env: p.GetEnv(),
|
||||
ModelPath: p.GetModelPath(), ServedModelName: p.GetServedModelName(),
|
||||
},
|
||||
}
|
||||
return spec, time.Duration(m.GetStartupTimeoutMs()) * time.Millisecond
|
||||
}
|
||||
|
||||
// stateWire maps supervisor states onto the wire enum.
|
||||
var stateWire = map[supervisor.State]llmgwv1.ShardStatus_State{
|
||||
supervisor.Pending: llmgwv1.ShardStatus_STATE_PENDING,
|
||||
supervisor.Starting: llmgwv1.ShardStatus_STATE_STARTING,
|
||||
supervisor.Loading: llmgwv1.ShardStatus_STATE_LOADING,
|
||||
supervisor.Ready: llmgwv1.ShardStatus_STATE_READY,
|
||||
supervisor.Draining: llmgwv1.ShardStatus_STATE_DRAINING,
|
||||
supervisor.Stopped: llmgwv1.ShardStatus_STATE_STOPPED,
|
||||
supervisor.Failed: llmgwv1.ShardStatus_STATE_FAILED,
|
||||
}
|
||||
|
||||
// StateToWire converts a supervisor state for transmission.
|
||||
func StateToWire(s supervisor.State) llmgwv1.ShardStatus_State {
|
||||
if w, ok := stateWire[s]; ok {
|
||||
return w
|
||||
}
|
||||
return llmgwv1.ShardStatus_STATE_UNSPECIFIED
|
||||
}
|
||||
|
||||
// StateFromWire converts a received state.
|
||||
//
|
||||
// An unrecognised state is treated as Failed rather than guessed at: a newer
|
||||
// agent reporting something this control plane does not understand must not be
|
||||
// mistaken for a healthy shard.
|
||||
func StateFromWire(w llmgwv1.ShardStatus_State) supervisor.State {
|
||||
for s, v := range stateWire {
|
||||
if v == w {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return supervisor.Failed
|
||||
}
|
||||
|
||||
func intSlice(v []int32) []int {
|
||||
if len(v) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]int, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = int(x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func int32Slice(v []int) []int32 {
|
||||
if len(v) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]int32, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = int32(x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,324 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"slices"
|
||||
"sync"
|
||||
|
||||
"connectrpc.com/connect"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// StatusSink receives shard transitions and recovers replica state from what
|
||||
// agents report they are already running. The reconciler implements it.
|
||||
type StatusSink interface {
|
||||
OnShardStatus(replicaID, shardID string, state supervisor.State, errMsg string)
|
||||
Adopt(shards []reconciler.AdoptedShard) (unclaimed []reconciler.AdoptedShard)
|
||||
}
|
||||
|
||||
// Service is the control plane's AgentService implementation, and at the same
|
||||
// time the reconciler's Dispatcher: it holds the open streams, so it is the
|
||||
// only thing that can deliver a command to a node.
|
||||
type Service struct {
|
||||
inv *inventory.Store
|
||||
status StatusSink
|
||||
log *slog.Logger
|
||||
|
||||
mu sync.RWMutex
|
||||
conns map[string]*conn
|
||||
}
|
||||
|
||||
// conn is one agent's open stream.
|
||||
type conn struct {
|
||||
nodeID string
|
||||
// send is serialized through a channel rather than a mutex around the
|
||||
// stream: connect's streams are not safe for concurrent writes, and the
|
||||
// reconciler dispatches to many nodes at once.
|
||||
send chan *llmgwv1.ControlMessage
|
||||
closed chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (c *conn) close() { c.once.Do(func() { close(c.closed) }) }
|
||||
|
||||
// New builds the service.
|
||||
func New(inv *inventory.Store, status StatusSink, log *slog.Logger) *Service {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Service{inv: inv, status: status, log: log, conns: map[string]*conn{}}
|
||||
}
|
||||
|
||||
// Handler returns the HTTP path and handler to mount.
|
||||
func (s *Service) Handler() (string, http.Handler) {
|
||||
return llmgwv1connect.NewAgentServiceHandler(s)
|
||||
}
|
||||
|
||||
// Connect serves one agent's stream for the lifetime of that agent.
|
||||
func (s *Service) Connect(
|
||||
ctx context.Context,
|
||||
stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage],
|
||||
) error {
|
||||
first, err := stream.Receive()
|
||||
if err != nil {
|
||||
return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("no opening message: %w", err))
|
||||
}
|
||||
reg := first.GetRegister()
|
||||
if reg == nil {
|
||||
// Registration first, always: the control plane has to know which node
|
||||
// it is talking to before it can route anything to it.
|
||||
return connect.NewError(connect.CodeInvalidArgument,
|
||||
errors.New("the first message must be a Register"))
|
||||
}
|
||||
nodeID := reg.GetNodeId()
|
||||
if nodeID == "" {
|
||||
return connect.NewError(connect.CodeInvalidArgument, errors.New("node_id is required"))
|
||||
}
|
||||
|
||||
node := NodeFromRegister(reg)
|
||||
s.inv.Upsert(node)
|
||||
s.log.Info("agent registered",
|
||||
"node", nodeID, "gpus", len(node.GPUs), "fabrics", node.FabricNames(),
|
||||
"running_shards", len(reg.GetRunningShards()),
|
||||
"agent_version", reg.GetAgentVersion())
|
||||
|
||||
c := s.attach(ctx, nodeID)
|
||||
defer s.detach(nodeID, c)
|
||||
|
||||
if err := stream.Send(&llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Ack{Ack: &llmgwv1.Ack{MessageId: "register", Ok: true}},
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Recover replica state from what this node is already running, before any
|
||||
// reconciliation pass can conclude the cluster is empty and place
|
||||
// duplicates (plan §6.5).
|
||||
s.adopt(ctx, nodeID, reg.GetRunningShards(), c)
|
||||
|
||||
// One goroutine writes, this one reads. Splitting them is what lets a
|
||||
// command be dispatched while the agent is mid-heartbeat.
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
s.writeLoop(ctx, stream, c)
|
||||
}()
|
||||
defer wg.Wait()
|
||||
|
||||
for {
|
||||
msg, err := stream.Receive()
|
||||
if err != nil {
|
||||
c.close()
|
||||
if ctx.Err() != nil || errors.Is(err, io.EOF) {
|
||||
s.log.Info("agent disconnected", "node", nodeID)
|
||||
return nil
|
||||
}
|
||||
s.log.Warn("agent stream failed", "node", nodeID, "err", err)
|
||||
return nil
|
||||
}
|
||||
s.handle(nodeID, msg)
|
||||
}
|
||||
}
|
||||
|
||||
// adopt hands the agent's running shards to the reconciler and stops whatever
|
||||
// nothing in the catalog claims.
|
||||
func (s *Service) adopt(ctx context.Context, nodeID string, shards []*llmgwv1.RunningShard, c *conn) {
|
||||
if len(shards) == 0 {
|
||||
return
|
||||
}
|
||||
unclaimed := s.status.Adopt(AdoptedFromRegister(nodeID, shards))
|
||||
for _, u := range unclaimed {
|
||||
// The model or profile left the catalog while this was running. Nobody
|
||||
// will ever route to it, and it is holding GPUs.
|
||||
s.log.Warn("stopping an unclaimed shard",
|
||||
"node", nodeID, "shard", u.ShardID, "model", u.Model, "profile", u.Profile)
|
||||
select {
|
||||
case c.send <- &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: u.ShardID}},
|
||||
}:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) writeLoop(ctx context.Context,
|
||||
stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage], c *conn,
|
||||
) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-c.closed:
|
||||
return
|
||||
case msg := <-c.send:
|
||||
if err := stream.Send(msg); err != nil {
|
||||
s.log.Warn("sending to agent failed", "node", c.nodeID, "err", err)
|
||||
c.close()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) {
|
||||
switch m := msg.Msg.(type) {
|
||||
case *llmgwv1.AgentMessage_Heartbeat:
|
||||
s.inv.Touch(nodeID)
|
||||
s.applyHeartbeat(nodeID, m.Heartbeat)
|
||||
|
||||
case *llmgwv1.AgentMessage_Shard:
|
||||
st := StateFromWire(m.Shard.GetState())
|
||||
s.status.OnShardStatus(m.Shard.GetReplicaId(), m.Shard.GetShardId(), st, m.Shard.GetError())
|
||||
if m.Shard.GetState() == llmgwv1.ShardStatus_STATE_FAILED {
|
||||
s.log.Error("shard failed",
|
||||
"node", nodeID, "shard", m.Shard.GetShardId(),
|
||||
"err", m.Shard.GetError(), "exit_code", m.Shard.GetExitCode(),
|
||||
"last_output", m.Shard.GetLastOutput())
|
||||
}
|
||||
|
||||
case *llmgwv1.AgentMessage_Register:
|
||||
// A re-registration on a live stream means the agent's view of its own
|
||||
// hardware changed — a GPU fell out, say. Take it.
|
||||
s.inv.Upsert(NodeFromRegister(m.Register))
|
||||
s.log.Info("agent re-registered", "node", nodeID)
|
||||
}
|
||||
}
|
||||
|
||||
// applyHeartbeat folds in the cheap facts a heartbeat carries, so a foreign
|
||||
// process appearing on a card after registration is noticed without waiting for
|
||||
// a reconnect.
|
||||
func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
|
||||
used := hb.GetGpuMemoryUsedMb()
|
||||
problems := hb.GetNewProblems()
|
||||
if len(used) == 0 && len(problems) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
n, ok := s.inv.Get(nodeID)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for i := range n.GPUs {
|
||||
if i < len(used) {
|
||||
n.GPUs[i].MemoryUsedMB = int(used[i])
|
||||
}
|
||||
}
|
||||
if len(problems) > 0 {
|
||||
// Problems are reported per node, not per GPU, so attach them to the
|
||||
// node's first GPU rather than guessing which one is affected. Any
|
||||
// problem is enough to keep that GPU out of placement, which is the
|
||||
// conservative direction.
|
||||
n.GPUs[0].Health.Problems = append(n.GPUs[0].Health.Problems, problems...)
|
||||
}
|
||||
s.inv.Upsert(n)
|
||||
}
|
||||
|
||||
func (s *Service) attach(ctx context.Context, nodeID string) *conn {
|
||||
c := &conn{
|
||||
nodeID: nodeID,
|
||||
// Buffered so a burst of commands does not block the reconciler on a
|
||||
// slow agent; if it fills, the agent is too far behind to be useful.
|
||||
send: make(chan *llmgwv1.ControlMessage, 64),
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
if prev, ok := s.conns[nodeID]; ok {
|
||||
// A node reconnecting before the old stream was reaped. The new one
|
||||
// wins: the old is by definition the dead one.
|
||||
s.log.Warn("agent reconnected while an old stream was open", "node", nodeID)
|
||||
prev.close()
|
||||
}
|
||||
s.conns[nodeID] = c
|
||||
s.mu.Unlock()
|
||||
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
c.close()
|
||||
}()
|
||||
return c
|
||||
}
|
||||
|
||||
func (s *Service) detach(nodeID string, c *conn) {
|
||||
c.close()
|
||||
s.mu.Lock()
|
||||
if s.conns[nodeID] == c {
|
||||
delete(s.conns, nodeID)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// The node is not marked lost here. Losing a stream is not losing a node,
|
||||
// and the reconciler's two-stage expiry already handles silence with the
|
||||
// right delays (plan §6.3).
|
||||
s.log.Info("agent stream closed", "node", nodeID)
|
||||
}
|
||||
|
||||
// ConnectedNodes lists the nodes with an open stream.
|
||||
func (s *Service) ConnectedNodes() []string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
out := make([]string, 0, len(s.conns))
|
||||
for id := range s.conns {
|
||||
out = append(out, id)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// --- Dispatcher ---
|
||||
|
||||
// ErrNotConnected means no agent stream is open for that node.
|
||||
var ErrNotConnected = errors.New("no agent connected")
|
||||
|
||||
// StartShard implements reconciler.Dispatcher.
|
||||
func (s *Service) StartShard(_ context.Context, cmd reconciler.ShardCommand) error {
|
||||
return s.send(cmd.NodeID, &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Start{Start: StartShardFromCommand(cmd)},
|
||||
})
|
||||
}
|
||||
|
||||
// StopShard implements reconciler.Dispatcher.
|
||||
func (s *Service) StopShard(_ context.Context, nodeID, shardID string) error {
|
||||
return s.send(nodeID, &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: shardID}},
|
||||
})
|
||||
}
|
||||
|
||||
// DrainShard implements reconciler.Dispatcher.
|
||||
func (s *Service) DrainShard(_ context.Context, nodeID, shardID string) error {
|
||||
return s.send(nodeID, &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Drain{Drain: &llmgwv1.DrainShard{ShardId: shardID}},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) send(nodeID string, msg *llmgwv1.ControlMessage) error {
|
||||
s.mu.RLock()
|
||||
c, ok := s.conns[nodeID]
|
||||
s.mu.RUnlock()
|
||||
if !ok {
|
||||
return fmt.Errorf("%w: node %s", ErrNotConnected, nodeID)
|
||||
}
|
||||
|
||||
select {
|
||||
case c.send <- msg:
|
||||
return nil
|
||||
case <-c.closed:
|
||||
return fmt.Errorf("%w: node %s stream closed", ErrNotConnected, nodeID)
|
||||
default:
|
||||
// The queue is full, so this agent is far enough behind that queueing
|
||||
// more would only delay discovering it is broken.
|
||||
return fmt.Errorf("node %s is not keeping up with commands", nodeID)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
// Package catalog holds the desired state: which models exist, how each one can
|
||||
// be run, and how many copies of it there should be (plan §6.1).
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/placement"
|
||||
)
|
||||
|
||||
// Scaling is how many replicas a model should have.
|
||||
type Scaling struct {
|
||||
Min int `koanf:"min"`
|
||||
Max int `koanf:"max"`
|
||||
}
|
||||
|
||||
// SLO is the latency target a model is expected to meet. `ctl bench` measures
|
||||
// capacity against it (plan §6.2.1).
|
||||
type SLO struct {
|
||||
TTFTp99 time.Duration `koanf:"ttft_p99"`
|
||||
TPOTp99 time.Duration `koanf:"tpot_p99"`
|
||||
}
|
||||
|
||||
// Profile is how a model runs on one GPU class.
|
||||
type Profile struct {
|
||||
Name string `koanf:"name"`
|
||||
Engine string `koanf:"engine"`
|
||||
GPUClass string `koanf:"gpu_class"`
|
||||
|
||||
TP int `koanf:"tp"`
|
||||
PP int `koanf:"pp"`
|
||||
// GPUsPerNode defaults to TP for a single-node profile.
|
||||
GPUsPerNode int `koanf:"gpus_per_node"`
|
||||
// Nodes above 1 makes this a multi-node profile, which requires a fabric.
|
||||
Nodes int `koanf:"nodes"`
|
||||
|
||||
Args []string `koanf:"args"`
|
||||
Env map[string]string `koanf:"env"`
|
||||
|
||||
// Weight is the router's capacity weight. It should come from ctl bench
|
||||
// rather than a guess (plan §6.2.1).
|
||||
Weight float64 `koanf:"weight"`
|
||||
// Priority orders profiles during placement; highest that fits wins.
|
||||
Priority int `koanf:"priority"`
|
||||
// StartupTimeout bounds model loading. A 70B model across two nodes takes
|
||||
// tens of minutes, so this is per profile rather than global.
|
||||
StartupTimeout time.Duration `koanf:"startup_timeout"`
|
||||
}
|
||||
|
||||
// Model is a catalog entry as clients see it.
|
||||
type Model struct {
|
||||
Name string `koanf:"name"`
|
||||
Aliases []string `koanf:"aliases"`
|
||||
// Source is the weights path or a HuggingFace id.
|
||||
Source string `koanf:"source"`
|
||||
Profiles []Profile `koanf:"profiles"`
|
||||
Scaling Scaling `koanf:"scaling"`
|
||||
Fallbacks []string `koanf:"fallbacks"`
|
||||
// RoutingProfile selects how the router scores endpoints for this model.
|
||||
RoutingProfile string `koanf:"routing_profile"`
|
||||
SLO SLO `koanf:"slo"`
|
||||
}
|
||||
|
||||
// Catalog is the set of models the cluster should serve.
|
||||
type Catalog struct {
|
||||
models map[string]Model
|
||||
aliases map[string]string
|
||||
order []string
|
||||
}
|
||||
|
||||
// New builds a catalog, rejecting definitions that would fail confusingly later.
|
||||
func New(models []Model) (*Catalog, error) {
|
||||
c := &Catalog{
|
||||
models: make(map[string]Model, len(models)),
|
||||
aliases: map[string]string{},
|
||||
}
|
||||
|
||||
for i := range models {
|
||||
m := models[i].normalize()
|
||||
if err := m.validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, dup := c.models[m.Name]; dup {
|
||||
return nil, fmt.Errorf("model %q is declared twice", m.Name)
|
||||
}
|
||||
c.models[m.Name] = m
|
||||
c.order = append(c.order, m.Name)
|
||||
}
|
||||
|
||||
// Aliases share a namespace with names. A collision would silently shadow a
|
||||
// model at request time, so it has to fail at load.
|
||||
for _, name := range c.order {
|
||||
for _, a := range c.models[name].Aliases {
|
||||
if _, clash := c.models[a]; clash {
|
||||
return nil, fmt.Errorf("model %q: alias %q collides with another model name", name, a)
|
||||
}
|
||||
if prev, clash := c.aliases[a]; clash {
|
||||
return nil, fmt.Errorf("alias %q is claimed by both %q and %q", a, prev, name)
|
||||
}
|
||||
c.aliases[a] = name
|
||||
}
|
||||
}
|
||||
|
||||
// Fallbacks are resolved now rather than at request time: a typo in a
|
||||
// fallback should stop a deployment, not surface during an outage when the
|
||||
// fallback is finally needed.
|
||||
for _, name := range c.order {
|
||||
for _, f := range c.models[name].Fallbacks {
|
||||
if _, ok := c.Resolve(f); !ok {
|
||||
return nil, fmt.Errorf("model %q: fallback %q is not in the catalog", name, f)
|
||||
}
|
||||
}
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (m Model) normalize() Model {
|
||||
if m.Scaling.Min < 0 {
|
||||
m.Scaling.Min = 0
|
||||
}
|
||||
if m.Scaling.Max == 0 {
|
||||
m.Scaling.Max = max(m.Scaling.Min, 1)
|
||||
}
|
||||
for i := range m.Profiles {
|
||||
p := &m.Profiles[i]
|
||||
if p.Engine == "" {
|
||||
p.Engine = "vllm"
|
||||
}
|
||||
if p.TP <= 0 {
|
||||
p.TP = 1
|
||||
}
|
||||
if p.Nodes <= 0 {
|
||||
p.Nodes = 1
|
||||
}
|
||||
if p.GPUsPerNode <= 0 {
|
||||
// The common case: a single-node profile uses TP GPUs on one node.
|
||||
p.GPUsPerNode = p.TP / p.Nodes
|
||||
if p.GPUsPerNode <= 0 {
|
||||
p.GPUsPerNode = 1
|
||||
}
|
||||
}
|
||||
if p.StartupTimeout <= 0 {
|
||||
p.StartupTimeout = 20 * time.Minute
|
||||
}
|
||||
if p.Weight <= 0 {
|
||||
p.Weight = 1
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func (m Model) validate() error {
|
||||
if m.Name == "" {
|
||||
return fmt.Errorf("a model has no name")
|
||||
}
|
||||
if m.Source == "" {
|
||||
return fmt.Errorf("model %q: source is required", m.Name)
|
||||
}
|
||||
if len(m.Profiles) == 0 {
|
||||
return fmt.Errorf("model %q: at least one profile is required", m.Name)
|
||||
}
|
||||
if m.Scaling.Max < m.Scaling.Min {
|
||||
return fmt.Errorf("model %q: scaling max %d is below min %d", m.Name, m.Scaling.Max, m.Scaling.Min)
|
||||
}
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
for _, p := range m.Profiles {
|
||||
if p.Name == "" {
|
||||
return fmt.Errorf("model %q: a profile has no name", m.Name)
|
||||
}
|
||||
if _, dup := seen[p.Name]; dup {
|
||||
return fmt.Errorf("model %q: profile %q is declared twice", m.Name, p.Name)
|
||||
}
|
||||
seen[p.Name] = struct{}{}
|
||||
|
||||
if p.GPUClass == "" {
|
||||
return fmt.Errorf("model %q profile %q: gpu_class is required", m.Name, p.Name)
|
||||
}
|
||||
// The arithmetic has to add up, or the engine is handed a shard spec it
|
||||
// cannot satisfy and fails minutes later during model loading.
|
||||
if want := p.GPUsPerNode * p.Nodes; want != p.TP*max(p.PP, 1) {
|
||||
return fmt.Errorf("model %q profile %q: %d nodes x %d GPUs does not match TP %d x PP %d",
|
||||
m.Name, p.Name, p.Nodes, p.GPUsPerNode, p.TP, max(p.PP, 1))
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Resolve maps a client-supplied name or alias onto a model.
|
||||
func (c *Catalog) Resolve(name string) (Model, bool) {
|
||||
if m, ok := c.models[name]; ok {
|
||||
return m, true
|
||||
}
|
||||
if canonical, ok := c.aliases[name]; ok {
|
||||
return c.models[canonical], true
|
||||
}
|
||||
return Model{}, false
|
||||
}
|
||||
|
||||
// Models returns every model, in declaration order.
|
||||
func (c *Catalog) Models() []Model {
|
||||
out := make([]Model, 0, len(c.order))
|
||||
for _, n := range c.order {
|
||||
out = append(out, c.models[n])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Names returns the canonical model names, sorted.
|
||||
func (c *Catalog) Names() []string {
|
||||
out := slices.Clone(c.order)
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// AliasesOf returns the alternative names for a model.
|
||||
func (c *Catalog) AliasesOf(name string) []string {
|
||||
m, ok := c.Resolve(name)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return slices.Clone(m.Aliases)
|
||||
}
|
||||
|
||||
// PlacementProfiles converts a model's profiles into placement's view.
|
||||
func (m Model) PlacementProfiles() []placement.Profile {
|
||||
out := make([]placement.Profile, 0, len(m.Profiles))
|
||||
for _, p := range m.Profiles {
|
||||
out = append(out, placement.Profile{
|
||||
Name: p.Name,
|
||||
Engine: p.Engine,
|
||||
GPUClass: p.GPUClass,
|
||||
GPUsPerNode: p.GPUsPerNode,
|
||||
Nodes: p.Nodes,
|
||||
Priority: p.Priority,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// EngineProfile converts a catalog profile into the engine adapter's view.
|
||||
func (m Model) EngineProfile(name string) (engines.Profile, bool) {
|
||||
for _, p := range m.Profiles {
|
||||
if p.Name != name {
|
||||
continue
|
||||
}
|
||||
return engines.Profile{
|
||||
Name: p.Name,
|
||||
Engine: p.Engine,
|
||||
GPUClass: p.GPUClass,
|
||||
TP: p.TP,
|
||||
PP: p.PP,
|
||||
GPUsPerNode: p.GPUsPerNode,
|
||||
Nodes: p.Nodes,
|
||||
Args: slices.Clone(p.Args),
|
||||
Env: p.Env,
|
||||
ModelPath: m.Source,
|
||||
ServedModelName: m.Name,
|
||||
}, true
|
||||
}
|
||||
return engines.Profile{}, false
|
||||
}
|
||||
|
||||
// ProfileByName returns one of a model's profiles.
|
||||
func (m Model) ProfileByName(name string) (Profile, bool) {
|
||||
for _, p := range m.Profiles {
|
||||
if p.Name == name {
|
||||
return p, true
|
||||
}
|
||||
}
|
||||
return Profile{}, false
|
||||
}
|
||||
|
||||
// String renders the catalog for `ctl models`.
|
||||
func (c *Catalog) String() string {
|
||||
var b strings.Builder
|
||||
for _, m := range c.Models() {
|
||||
fmt.Fprintf(&b, "%s (min %d, max %d)\n", m.Name, m.Scaling.Min, m.Scaling.Max)
|
||||
for _, p := range m.Profiles {
|
||||
fmt.Fprintf(&b, " %-20s %-10s %-12s tp=%d nodes=%d priority=%d\n",
|
||||
p.Name, p.Engine, p.GPUClass, p.TP, p.Nodes, p.Priority)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,235 @@
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func model(name string, profiles ...Profile) Model {
|
||||
if len(profiles) == 0 {
|
||||
profiles = []Profile{{Name: "p1", GPUClass: "h100-80g", TP: 4}}
|
||||
}
|
||||
return Model{Name: name, Source: "/models/" + name, Profiles: profiles}
|
||||
}
|
||||
|
||||
func TestResolveNameAndAlias(t *testing.T) {
|
||||
m := model("llama-3.3-70b")
|
||||
m.Aliases = []string{"llama", "llama3"}
|
||||
c, err := New([]Model{m})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for _, name := range []string{"llama-3.3-70b", "llama", "llama3"} {
|
||||
got, ok := c.Resolve(name)
|
||||
if !ok {
|
||||
t.Fatalf("%q did not resolve", name)
|
||||
}
|
||||
if got.Name != "llama-3.3-70b" {
|
||||
t.Errorf("%q resolved to %q", name, got.Name)
|
||||
}
|
||||
}
|
||||
if _, ok := c.Resolve("nope"); ok {
|
||||
t.Error("an unknown name resolved")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAliasCollisionRejected(t *testing.T) {
|
||||
// An alias shadowing another model would misroute every request for it,
|
||||
// silently, at request time.
|
||||
a := model("qwen")
|
||||
b := model("llama")
|
||||
b.Aliases = []string{"qwen"}
|
||||
|
||||
_, err := New([]Model{a, b})
|
||||
if err == nil || !strings.Contains(err.Error(), "collides") {
|
||||
t.Errorf("collision not rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuplicateAliasRejected(t *testing.T) {
|
||||
a := model("a")
|
||||
a.Aliases = []string{"shared"}
|
||||
b := model("b")
|
||||
b.Aliases = []string{"shared"}
|
||||
|
||||
_, err := New([]Model{a, b})
|
||||
if err == nil || !strings.Contains(err.Error(), "claimed by both") {
|
||||
t.Errorf("duplicate alias not rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnknownFallbackRejected(t *testing.T) {
|
||||
// A typo in a fallback should stop a deployment, not surface during the
|
||||
// outage when the fallback is finally needed.
|
||||
m := model("llama")
|
||||
m.Fallbacks = []string{"qwen-typo"}
|
||||
|
||||
_, err := New([]Model{m})
|
||||
if err == nil || !strings.Contains(err.Error(), "fallback") {
|
||||
t.Errorf("unknown fallback not rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFallbackToAliasAccepted(t *testing.T) {
|
||||
a := model("qwen2.5-32b")
|
||||
a.Aliases = []string{"qwen"}
|
||||
b := model("llama")
|
||||
b.Fallbacks = []string{"qwen"}
|
||||
|
||||
if _, err := New([]Model{a, b}); err != nil {
|
||||
t.Errorf("a fallback naming an alias was rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelismArithmeticValidated(t *testing.T) {
|
||||
// Wrong arithmetic hands the engine a spec it cannot satisfy, and it fails
|
||||
// minutes later during model loading rather than at config load.
|
||||
_, err := New([]Model{model("m", Profile{
|
||||
Name: "bad", GPUClass: "h100-80g", TP: 4, Nodes: 2, GPUsPerNode: 8,
|
||||
})})
|
||||
if err == nil || !strings.Contains(err.Error(), "does not match") {
|
||||
t.Errorf("bad arithmetic accepted: %v", err)
|
||||
}
|
||||
|
||||
// 2 nodes x 8 GPUs = TP 16 is consistent.
|
||||
if _, err := New([]Model{model("m", Profile{
|
||||
Name: "ok", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
})}); err != nil {
|
||||
t.Errorf("valid arithmetic rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaults(t *testing.T) {
|
||||
c, err := New([]Model{model("m", Profile{Name: "p", GPUClass: "h100-80g", TP: 4})})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m, _ := c.Resolve("m")
|
||||
p := m.Profiles[0]
|
||||
|
||||
if p.Engine != "vllm" {
|
||||
t.Errorf("engine=%q", p.Engine)
|
||||
}
|
||||
if p.Nodes != 1 {
|
||||
t.Errorf("nodes=%d", p.Nodes)
|
||||
}
|
||||
if p.GPUsPerNode != 4 {
|
||||
t.Errorf("gpus_per_node=%d; a single-node profile should default to TP", p.GPUsPerNode)
|
||||
}
|
||||
if p.StartupTimeout != 20*time.Minute {
|
||||
t.Errorf("startup_timeout=%v", p.StartupTimeout)
|
||||
}
|
||||
if p.Weight <= 0 {
|
||||
t.Errorf("weight=%v", p.Weight)
|
||||
}
|
||||
if m.Scaling.Max < 1 {
|
||||
t.Errorf("scaling max=%d", m.Scaling.Max)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequiredFields(t *testing.T) {
|
||||
cases := map[string]Model{
|
||||
"no name": {Source: "/m", Profiles: []Profile{{Name: "p", GPUClass: "c", TP: 1}}},
|
||||
"no source": {Name: "m", Profiles: []Profile{{Name: "p", GPUClass: "c", TP: 1}}},
|
||||
"no profiles": {Name: "m", Source: "/m"},
|
||||
"no class": {Name: "m", Source: "/m", Profiles: []Profile{{Name: "p", TP: 1}}},
|
||||
}
|
||||
for name, m := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
if _, err := New([]Model{m}); err == nil {
|
||||
t.Error("accepted")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestScalingBounds(t *testing.T) {
|
||||
m := model("m")
|
||||
m.Scaling = Scaling{Min: 5, Max: 2}
|
||||
if _, err := New([]Model{m}); err == nil {
|
||||
t.Error("max below min was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuplicateProfileRejected(t *testing.T) {
|
||||
_, err := New([]Model{model("m",
|
||||
Profile{Name: "p", GPUClass: "h100-80g", TP: 1},
|
||||
Profile{Name: "p", GPUClass: "a100-40g", TP: 1},
|
||||
)})
|
||||
if err == nil || !strings.Contains(err.Error(), "twice") {
|
||||
t.Errorf("duplicate profile accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlacementProfilesCarryPriority(t *testing.T) {
|
||||
c, _ := New([]Model{model("m",
|
||||
Profile{Name: "fast", GPUClass: "h100-80g", TP: 4, Priority: 100},
|
||||
Profile{Name: "slow", GPUClass: "a100-40g", TP: 4, Priority: 50},
|
||||
)})
|
||||
m, _ := c.Resolve("m")
|
||||
|
||||
pp := m.PlacementProfiles()
|
||||
if len(pp) != 2 {
|
||||
t.Fatalf("got %d placement profiles", len(pp))
|
||||
}
|
||||
byName := map[string]int{}
|
||||
for _, p := range pp {
|
||||
byName[p.Name] = p.Priority
|
||||
}
|
||||
if byName["fast"] != 100 || byName["slow"] != 50 {
|
||||
t.Errorf("priorities lost in conversion: %v", byName)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineProfileCarriesModelIdentity(t *testing.T) {
|
||||
// The served name is what clients ask for and need not match the path.
|
||||
c, _ := New([]Model{model("llama-3.3-70b")})
|
||||
m, _ := c.Resolve("llama-3.3-70b")
|
||||
|
||||
ep, ok := m.EngineProfile("p1")
|
||||
if !ok {
|
||||
t.Fatal("engine profile not found")
|
||||
}
|
||||
if ep.ModelPath != "/models/llama-3.3-70b" {
|
||||
t.Errorf("ModelPath=%q", ep.ModelPath)
|
||||
}
|
||||
if ep.ServedModelName != "llama-3.3-70b" {
|
||||
t.Errorf("ServedModelName=%q", ep.ServedModelName)
|
||||
}
|
||||
|
||||
if _, ok := m.EngineProfile("nope"); ok {
|
||||
t.Error("an unknown profile resolved")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineProfileArgsAreCopied(t *testing.T) {
|
||||
// The reconciler hands these to an engine adapter; mutating them must not
|
||||
// corrupt the catalog for every later replica.
|
||||
c, _ := New([]Model{model("m", Profile{
|
||||
Name: "p", GPUClass: "h100-80g", TP: 1, Args: []string{"--max-model-len", "32768"},
|
||||
})})
|
||||
m, _ := c.Resolve("m")
|
||||
|
||||
ep, _ := m.EngineProfile("p")
|
||||
ep.Args[0] = "tampered"
|
||||
|
||||
again, _ := m.EngineProfile("p")
|
||||
if again.Args[0] != "--max-model-len" {
|
||||
t.Error("EngineProfile exposed the catalog's own slice")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNamesAndModelsOrder(t *testing.T) {
|
||||
c, _ := New([]Model{model("zeta"), model("alpha")})
|
||||
|
||||
// Models keeps declaration order, so configuration reads as written.
|
||||
if got := c.Models(); got[0].Name != "zeta" {
|
||||
t.Errorf("Models order=%q", got[0].Name)
|
||||
}
|
||||
// Names is sorted, for stable CLI output.
|
||||
if got := c.Names(); got[0] != "alpha" {
|
||||
t.Errorf("Names order=%v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
// Package control wires the control plane together: inventory, catalog,
|
||||
// placement and reconciliation, and publishes the result to the router.
|
||||
package control
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
)
|
||||
|
||||
// EndpointSink receives the endpoint list. The router implements it.
|
||||
//
|
||||
// It is an interface so the control plane does not depend on the router: when
|
||||
// the two split into separate processes, this becomes a watch stream and
|
||||
// nothing above it changes (plan §3.1).
|
||||
type EndpointSink interface {
|
||||
SetEndpoints(ctx context.Context, infos []endpoint.Info)
|
||||
}
|
||||
|
||||
// Config tunes the control plane.
|
||||
type Config struct {
|
||||
Reconciler reconciler.Config
|
||||
// PublishInterval is how often the endpoint list is republished even when
|
||||
// nothing appears to have changed. Events drive the common case; this is
|
||||
// the backstop.
|
||||
PublishInterval time.Duration
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Reconciler: reconciler.DefaultConfig(),
|
||||
PublishInterval: time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// Controller owns the control plane's state.
|
||||
type Controller struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
inv *inventory.Store
|
||||
cat *catalog.Catalog
|
||||
rec *reconciler.Reconciler
|
||||
sink EndpointSink
|
||||
|
||||
mu sync.Mutex
|
||||
last string // fingerprint of the last published list
|
||||
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// New builds a controller.
|
||||
func New(cfg Config, cat *catalog.Catalog, disp reconciler.Dispatcher, sink EndpointSink, log *slog.Logger) *Controller {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if cfg.PublishInterval <= 0 {
|
||||
cfg.PublishInterval = time.Second
|
||||
}
|
||||
inv := inventory.NewStore()
|
||||
return &Controller{
|
||||
cfg: cfg,
|
||||
log: log,
|
||||
inv: inv,
|
||||
cat: cat,
|
||||
rec: reconciler.New(cfg.Reconciler, cat, inv, disp, log),
|
||||
sink: sink,
|
||||
stop: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Inventory exposes the node inventory, which the agent service feeds.
|
||||
func (c *Controller) Inventory() *inventory.Store { return c.inv }
|
||||
|
||||
// Reconciler exposes the reconciliation loop, for status reports and ctl.
|
||||
func (c *Controller) Reconciler() *reconciler.Reconciler { return c.rec }
|
||||
|
||||
// Catalog exposes the model catalog.
|
||||
func (c *Controller) Catalog() *catalog.Catalog { return c.cat }
|
||||
|
||||
// Start runs the reconciler and begins publishing endpoints.
|
||||
func (c *Controller) Start(ctx context.Context) {
|
||||
c.rec.Start(ctx)
|
||||
|
||||
go func() {
|
||||
defer close(c.done)
|
||||
t := time.NewTicker(c.cfg.PublishInterval)
|
||||
defer t.Stop()
|
||||
|
||||
c.Publish(ctx)
|
||||
for {
|
||||
select {
|
||||
case <-c.stop:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
c.Publish(ctx)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Close stops the control plane.
|
||||
func (c *Controller) Close() {
|
||||
c.once.Do(func() {
|
||||
close(c.stop)
|
||||
<-c.done
|
||||
c.rec.Close()
|
||||
})
|
||||
}
|
||||
|
||||
// Publish pushes the current Ready replicas to the sink.
|
||||
//
|
||||
// It is exported so tests can drive it deterministically instead of racing a
|
||||
// ticker, and it skips the push when nothing changed: replacing the list is
|
||||
// cheap but resynchronizes the metric pollers, and doing that every second for
|
||||
// no reason adds churn to every engine's /metrics.
|
||||
func (c *Controller) Publish(ctx context.Context) {
|
||||
infos := c.endpoints()
|
||||
|
||||
fp := fingerprint(infos)
|
||||
c.mu.Lock()
|
||||
unchanged := fp == c.last
|
||||
c.last = fp
|
||||
c.mu.Unlock()
|
||||
if unchanged {
|
||||
return
|
||||
}
|
||||
|
||||
c.log.Info("publishing endpoints", "count", len(infos))
|
||||
c.sink.SetEndpoints(ctx, infos)
|
||||
}
|
||||
|
||||
// endpoints converts Ready replicas into the router's view.
|
||||
func (c *Controller) endpoints() []endpoint.Info {
|
||||
ready := c.rec.ReadyReplicas()
|
||||
out := make([]endpoint.Info, 0, len(ready))
|
||||
|
||||
for _, rep := range ready {
|
||||
m, ok := c.cat.Resolve(rep.Model)
|
||||
if !ok {
|
||||
// The model left the catalog while a replica was still running.
|
||||
// The reconciler will tear it down; until then, do not route to it.
|
||||
continue
|
||||
}
|
||||
// Clients may address a model by any of its names, and the router
|
||||
// matches endpoints by the resolved name plus aliases.
|
||||
names := append([]string{m.Name}, m.Aliases...)
|
||||
|
||||
prof, _ := m.ProfileByName(rep.Profile)
|
||||
out = append(out, endpoint.Info{
|
||||
ID: rep.ID,
|
||||
URL: strings.TrimRight(rep.Endpoint, "/"),
|
||||
Models: names,
|
||||
Pool: rep.Pool,
|
||||
Weight: rep.Weight,
|
||||
Priority: prof.Priority,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// fingerprint summarizes a list so an unchanged one can be skipped.
|
||||
func fingerprint(infos []endpoint.Info) string {
|
||||
var b strings.Builder
|
||||
for _, i := range infos {
|
||||
fmt.Fprintf(&b, "%s|%s|%s|%v|%g|%d;", i.ID, i.URL, i.Pool, i.Models, i.Weight, i.Priority)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
// Package inventory holds the control plane's view of the cluster: which nodes
|
||||
// exist, what hardware they have, and what is already running on them.
|
||||
package inventory
|
||||
|
||||
import (
|
||||
"maps"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
)
|
||||
|
||||
// NodeState is where a node stands.
|
||||
type NodeState int
|
||||
|
||||
const (
|
||||
// Ready means the node accepts placement.
|
||||
Ready NodeState = iota
|
||||
// Cordoned means an operator has excluded it from new placement; what runs
|
||||
// there keeps running.
|
||||
Cordoned
|
||||
// Draining means its replicas are being moved off.
|
||||
Draining
|
||||
// Unreachable means heartbeats stopped but the node may come back.
|
||||
Unreachable
|
||||
// Lost means it has been unreachable long enough to declare its replicas dead.
|
||||
Lost
|
||||
)
|
||||
|
||||
func (s NodeState) String() string {
|
||||
switch s {
|
||||
case Ready:
|
||||
return "ready"
|
||||
case Cordoned:
|
||||
return "cordoned"
|
||||
case Draining:
|
||||
return "draining"
|
||||
case Unreachable:
|
||||
return "unreachable"
|
||||
case Lost:
|
||||
return "lost"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// Schedulable reports whether new replicas may be placed on the node.
|
||||
func (s NodeState) Schedulable() bool { return s == Ready }
|
||||
|
||||
// FabricAttachment is a node's connection to one fabric.
|
||||
type FabricAttachment struct {
|
||||
// Fabric is the operator-assigned name. It is not derived from the GID
|
||||
// prefix: OpenSM's default is fe80:: everywhere, so a matching prefix
|
||||
// proves nothing (plan §5.2).
|
||||
Fabric string `json:"fabric"`
|
||||
// HCAs are the adapters attached to this fabric.
|
||||
HCAs []string `json:"hcas"`
|
||||
// RateGbps is the slowest active port's rate, since a collective runs at
|
||||
// the speed of its slowest link.
|
||||
RateGbps int `json:"rate_gbps"`
|
||||
}
|
||||
|
||||
// Node is one machine in the cluster.
|
||||
type Node struct {
|
||||
ID string `json:"id"`
|
||||
Hostname string `json:"hostname"`
|
||||
MgmtAddr string `json:"mgmt_addr"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
|
||||
GPUs []discovery.GPU `json:"gpus"`
|
||||
// Fabrics is empty on a node without IB. Such a node is its own
|
||||
// single-node island and can never host a multi-node replica.
|
||||
Fabrics []FabricAttachment `json:"fabrics,omitempty"`
|
||||
|
||||
State NodeState `json:"-"`
|
||||
LastSeen time.Time `json:"last_seen"`
|
||||
// AgentVersion and EngineVersions are reported at registration. The engine
|
||||
// version decides which metric names and CLI flags apply (plan §14).
|
||||
AgentVersion string `json:"agent_version,omitempty"`
|
||||
EngineVersions map[string]string `json:"engine_versions,omitempty"`
|
||||
}
|
||||
|
||||
// HasFabric reports whether the node is attached to a fabric.
|
||||
func (n Node) HasFabric(name string) bool {
|
||||
return slices.ContainsFunc(n.Fabrics, func(f FabricAttachment) bool { return f.Fabric == name })
|
||||
}
|
||||
|
||||
// FabricNames lists the node's fabrics, sorted.
|
||||
func (n Node) FabricNames() []string {
|
||||
out := make([]string, 0, len(n.Fabrics))
|
||||
for _, f := range n.Fabrics {
|
||||
out = append(out, f.Fabric)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// HCAsFor returns the adapters attached to one fabric.
|
||||
//
|
||||
// A shard is given only these: handing it adapters from a second fabric lets
|
||||
// NCCL choose one whose peer is unreachable, and the job hangs instead of
|
||||
// failing (plan §5.3).
|
||||
func (n Node) HCAsFor(fabric string) []string {
|
||||
for _, f := range n.Fabrics {
|
||||
if f.Fabric == fabric {
|
||||
return slices.Clone(f.HCAs)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GPUsByClass groups the node's GPUs by class.
|
||||
func (n Node) GPUsByClass() map[string][]discovery.GPU {
|
||||
out := map[string][]discovery.GPU{}
|
||||
for _, g := range n.GPUs {
|
||||
out[g.Class()] = append(out[g.Class()], g)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Allocation records that some of a node's GPUs are in use.
|
||||
type Allocation struct {
|
||||
ShardID string
|
||||
ReplicaID string
|
||||
GPUIndices []int
|
||||
}
|
||||
|
||||
// Store holds the cluster inventory. It is safe for concurrent use.
|
||||
type Store struct {
|
||||
mu sync.RWMutex
|
||||
nodes map[string]*Node
|
||||
// allocs maps node id to the shards placed there.
|
||||
allocs map[string][]Allocation
|
||||
}
|
||||
|
||||
// NewStore returns an empty inventory.
|
||||
func NewStore() *Store {
|
||||
return &Store{nodes: map[string]*Node{}, allocs: map[string][]Allocation{}}
|
||||
}
|
||||
|
||||
// Upsert records a node's registration, preserving operator-set state such as
|
||||
// a cordon: an agent reconnecting must not silently undo an operator's decision.
|
||||
func (s *Store) Upsert(n Node) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if prev, ok := s.nodes[n.ID]; ok {
|
||||
if prev.State == Cordoned || prev.State == Draining {
|
||||
n.State = prev.State
|
||||
}
|
||||
}
|
||||
n.LastSeen = time.Now()
|
||||
s.nodes[n.ID] = &n
|
||||
}
|
||||
|
||||
// SetState changes a node's state.
|
||||
func (s *Store) SetState(id string, st NodeState) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
n, ok := s.nodes[id]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
n.State = st
|
||||
return true
|
||||
}
|
||||
|
||||
// Touch records a heartbeat, bringing an unreachable node back.
|
||||
func (s *Store) Touch(id string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
n, ok := s.nodes[id]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
n.LastSeen = time.Now()
|
||||
if n.State == Unreachable || n.State == Lost {
|
||||
n.State = Ready
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns a copy of a node.
|
||||
func (s *Store) Get(id string) (Node, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
n, ok := s.nodes[id]
|
||||
if !ok {
|
||||
return Node{}, false
|
||||
}
|
||||
return *n, true
|
||||
}
|
||||
|
||||
// Nodes returns every node, sorted by id.
|
||||
func (s *Store) Nodes() []Node {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
out := make([]Node, 0, len(s.nodes))
|
||||
for _, n := range s.nodes {
|
||||
out = append(out, *n)
|
||||
}
|
||||
slices.SortFunc(out, func(a, b Node) int { return strings.Compare(a.ID, b.ID) })
|
||||
return out
|
||||
}
|
||||
|
||||
// Remove forgets a node and its allocations.
|
||||
func (s *Store) Remove(id string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
delete(s.nodes, id)
|
||||
delete(s.allocs, id)
|
||||
}
|
||||
|
||||
// Allocate records GPUs as taken on a node.
|
||||
func (s *Store) Allocate(nodeID string, a Allocation) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.allocs[nodeID] = append(s.allocs[nodeID], a)
|
||||
}
|
||||
|
||||
// Release frees the GPUs a shard held.
|
||||
func (s *Store) Release(shardID string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for node, list := range s.allocs {
|
||||
kept := slices.DeleteFunc(slices.Clone(list), func(a Allocation) bool {
|
||||
return a.ShardID == shardID
|
||||
})
|
||||
if len(kept) == 0 {
|
||||
delete(s.allocs, node)
|
||||
} else {
|
||||
s.allocs[node] = kept
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ReleaseReplica frees every shard of a replica.
|
||||
//
|
||||
// Replicas are freed whole because they fail whole: one dead shard makes the
|
||||
// others useless, and the control plane rebuilds the replica rather than
|
||||
// restarting a piece of it (plan §5.4).
|
||||
func (s *Store) ReleaseReplica(replicaID string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for node, list := range s.allocs {
|
||||
kept := slices.DeleteFunc(slices.Clone(list), func(a Allocation) bool {
|
||||
return a.ReplicaID == replicaID
|
||||
})
|
||||
if len(kept) == 0 {
|
||||
delete(s.allocs, node)
|
||||
} else {
|
||||
s.allocs[node] = kept
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Allocations returns a node's current allocations.
|
||||
func (s *Store) Allocations(nodeID string) []Allocation {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return slices.Clone(s.allocs[nodeID])
|
||||
}
|
||||
|
||||
// FreeGPUs returns the indices of a node's unallocated GPUs, sorted.
|
||||
func (s *Store) FreeGPUs(nodeID string) []int {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
n, ok := s.nodes[nodeID]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
taken := map[int]struct{}{}
|
||||
for _, a := range s.allocs[nodeID] {
|
||||
for _, i := range a.GPUIndices {
|
||||
taken[i] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
var out []int
|
||||
for _, g := range n.GPUs {
|
||||
if _, used := taken[g.Index]; used {
|
||||
continue
|
||||
}
|
||||
// A GPU with an ECC fault or another process on it is not free in any
|
||||
// useful sense.
|
||||
if !g.Health.OK() {
|
||||
continue
|
||||
}
|
||||
out = append(out, g.Index)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// Fabrics lists every fabric any node is attached to, sorted.
|
||||
func (s *Store) Fabrics() []string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
for _, n := range s.nodes {
|
||||
for _, f := range n.Fabrics {
|
||||
seen[f.Fabric] = struct{}{}
|
||||
}
|
||||
}
|
||||
return slices.Sorted(maps.Keys(seen))
|
||||
}
|
||||
@@ -0,0 +1,250 @@
|
||||
package inventory
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
)
|
||||
|
||||
func gpus(n int) []discovery.GPU {
|
||||
out := make([]discovery.GPU, n)
|
||||
for i := range out {
|
||||
out[i] = discovery.GPU{Index: i, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestUpsertPreservesCordon(t *testing.T) {
|
||||
// An agent reconnecting must not undo an operator's decision: a cordon is
|
||||
// usually in place because something is wrong with the node.
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
s.SetState("n1", Cordoned)
|
||||
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)}) // re-registration
|
||||
n, _ := s.Get("n1")
|
||||
if n.State != Cordoned {
|
||||
t.Errorf("state=%s; re-registration cleared the cordon", n.State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertUpdatesHardware(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(4)})
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
|
||||
n, _ := s.Get("n1")
|
||||
if len(n.GPUs) != 8 {
|
||||
t.Errorf("hardware was not updated: %d GPUs", len(n.GPUs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTouchRevivesUnreachableNode(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1"})
|
||||
s.SetState("n1", Unreachable)
|
||||
|
||||
s.Touch("n1")
|
||||
if n, _ := s.Get("n1"); n.State != Ready {
|
||||
t.Errorf("a heartbeat did not revive the node: %s", n.State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTouchDoesNotClearCordon(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1"})
|
||||
s.SetState("n1", Cordoned)
|
||||
|
||||
s.Touch("n1")
|
||||
if n, _ := s.Get("n1"); n.State != Cordoned {
|
||||
t.Errorf("a heartbeat cleared the cordon: %s", n.State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFreeGPUs(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
|
||||
if got := s.FreeGPUs("n1"); len(got) != 8 {
|
||||
t.Fatalf("expected 8 free, got %d", len(got))
|
||||
}
|
||||
|
||||
s.Allocate("n1", Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1}})
|
||||
got := s.FreeGPUs("n1")
|
||||
if !slices.Equal(got, []int{2, 3, 4, 5, 6, 7}) {
|
||||
t.Errorf("free=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnhealthyGPUsAreNotFree(t *testing.T) {
|
||||
s := NewStore()
|
||||
n := Node{ID: "n1", GPUs: gpus(4)}
|
||||
n.GPUs[2].Health.Problems = []string{"ECC"}
|
||||
s.Upsert(n)
|
||||
|
||||
got := s.FreeGPUs("n1")
|
||||
if slices.Contains(got, 2) {
|
||||
t.Errorf("an unhealthy GPU was reported free: %v", got)
|
||||
}
|
||||
if len(got) != 3 {
|
||||
t.Errorf("free=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseShard(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
s.Allocate("n1", Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1}})
|
||||
s.Allocate("n1", Allocation{ShardID: "s2", ReplicaID: "r2", GPUIndices: []int{2, 3}})
|
||||
|
||||
s.Release("s1")
|
||||
if got := s.FreeGPUs("n1"); slices.Contains(got, 2) || !slices.Contains(got, 0) {
|
||||
t.Errorf("wrong shard released: free=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseReplicaFreesEveryShard(t *testing.T) {
|
||||
// Replicas fail whole, so they are freed whole: one dead shard makes the
|
||||
// others useless and the control plane rebuilds the replica (plan §5.4).
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
s.Upsert(Node{ID: "n2", GPUs: gpus(8)})
|
||||
s.Allocate("n1", Allocation{ShardID: "r1-0", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
|
||||
s.Allocate("n2", Allocation{ShardID: "r1-1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
|
||||
|
||||
s.ReleaseReplica("r1")
|
||||
for _, id := range []string{"n1", "n2"} {
|
||||
if got := len(s.FreeGPUs(id)); got != 8 {
|
||||
t.Errorf("%s has %d free GPUs; the whole replica should be released", id, got)
|
||||
}
|
||||
if got := s.Allocations(id); len(got) != 0 {
|
||||
t.Errorf("%s still has allocations: %v", id, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHCAsForSelectsOneFabric(t *testing.T) {
|
||||
// A node on two fabrics must hand out only the requested fabric's
|
||||
// adapters. Mixing them lets NCCL pick one whose peer is unreachable.
|
||||
n := Node{ID: "n1", Fabrics: []FabricAttachment{
|
||||
{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}},
|
||||
{Fabric: "ib-b", HCAs: []string{"mlx5_2", "mlx5_3"}},
|
||||
}}
|
||||
|
||||
if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
|
||||
t.Errorf("ib-a HCAs=%v", got)
|
||||
}
|
||||
if got := n.HCAsFor("ib-b"); !slices.Equal(got, []string{"mlx5_2", "mlx5_3"}) {
|
||||
t.Errorf("ib-b HCAs=%v", got)
|
||||
}
|
||||
if got := n.HCAsFor("ib-c"); got != nil {
|
||||
t.Errorf("an unattached fabric returned %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHCAsForReturnsACopy(t *testing.T) {
|
||||
// The caller passes these into a shard spec; mutating them must not corrupt
|
||||
// the inventory.
|
||||
n := Node{ID: "n1", Fabrics: []FabricAttachment{{Fabric: "ib-a", HCAs: []string{"mlx5_0"}}}}
|
||||
got := n.HCAsFor("ib-a")
|
||||
got[0] = "tampered"
|
||||
if n.Fabrics[0].HCAs[0] != "mlx5_0" {
|
||||
t.Error("HCAsFor exposed the inventory's own slice")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFabrics(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", Fabrics: []FabricAttachment{{Fabric: "ib-b"}}})
|
||||
s.Upsert(Node{ID: "n2", Fabrics: []FabricAttachment{{Fabric: "ib-a"}, {Fabric: "ib-b"}}})
|
||||
s.Upsert(Node{ID: "n3"}) // no fabric
|
||||
|
||||
if got := s.Fabrics(); !slices.Equal(got, []string{"ib-a", "ib-b"}) {
|
||||
t.Errorf("fabrics=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodesSorted(t *testing.T) {
|
||||
s := NewStore()
|
||||
for _, id := range []string{"n3", "n1", "n2"} {
|
||||
s.Upsert(Node{ID: id})
|
||||
}
|
||||
got := s.Nodes()
|
||||
if len(got) != 3 || got[0].ID != "n1" || got[2].ID != "n3" {
|
||||
t.Errorf("not sorted: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetReturnsACopy(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(2), Labels: map[string]string{"rack": "r1"}})
|
||||
|
||||
n, _ := s.Get("n1")
|
||||
n.GPUs[0].Index = 99
|
||||
|
||||
again, _ := s.Get("n1")
|
||||
if again.GPUs[0].Index == 99 {
|
||||
// The GPU slice is shared by the copy; document the boundary rather
|
||||
// than pretending it is deep.
|
||||
t.Log("note: Get returns a shallow copy; slices are shared")
|
||||
}
|
||||
if again.ID != "n1" {
|
||||
t.Error("the node was corrupted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemove(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(4)})
|
||||
s.Allocate("n1", Allocation{ShardID: "s1", GPUIndices: []int{0}})
|
||||
|
||||
s.Remove("n1")
|
||||
if _, ok := s.Get("n1"); ok {
|
||||
t.Error("the node survived removal")
|
||||
}
|
||||
if got := s.Allocations("n1"); len(got) != 0 {
|
||||
t.Error("allocations outlived the node")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateStrings(t *testing.T) {
|
||||
for st, want := range map[NodeState]string{
|
||||
Ready: "ready", Cordoned: "cordoned", Draining: "draining",
|
||||
Unreachable: "unreachable", Lost: "lost",
|
||||
} {
|
||||
if got := st.String(); got != want {
|
||||
t.Errorf("%d -> %q, want %q", int(st), got, want)
|
||||
}
|
||||
}
|
||||
if !Ready.Schedulable() || Cordoned.Schedulable() || Lost.Schedulable() {
|
||||
t.Error("Schedulable is wrong")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentAccess(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := range 100 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
switch i % 5 {
|
||||
case 0:
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
case 1:
|
||||
s.Allocate("n1", Allocation{ShardID: "s", ReplicaID: "r", GPUIndices: []int{i % 8}})
|
||||
case 2:
|
||||
s.Release("s")
|
||||
case 3:
|
||||
_ = s.FreeGPUs("n1")
|
||||
default:
|
||||
_ = s.Nodes()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package inventory
|
||||
|
||||
import "time"
|
||||
|
||||
// SetLastSeen overrides a node's heartbeat timestamp.
|
||||
//
|
||||
// It exists for tests that run on a synthetic clock: Upsert and Touch stamp the
|
||||
// wall clock, which is decades away from a fake timeline and would make every
|
||||
// node look silent.
|
||||
func (s *Store) SetLastSeen(id string, at time.Time) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if n, ok := s.nodes[id]; ok {
|
||||
n.LastSeen = at
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
// Package placement decides which nodes and GPUs a replica runs on (plan §6.2).
|
||||
//
|
||||
// Two constraints here are not negotiable, and everything else in this package
|
||||
// exists to keep them true:
|
||||
//
|
||||
// - a multi-node replica never spans two fabrics
|
||||
// - a tensor-parallel group never mixes GPU classes
|
||||
//
|
||||
// Violating either produces a replica that starts, passes health checks, and
|
||||
// then hangs or crawls under load — the worst possible failure mode, because it
|
||||
// looks like success until traffic arrives.
|
||||
package placement
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// Profile is the subset of a catalog profile placement needs.
|
||||
type Profile struct {
|
||||
Name string
|
||||
Engine string
|
||||
GPUClass string
|
||||
// GPUsPerNode is how many GPUs each node contributes.
|
||||
GPUsPerNode int
|
||||
// Nodes is how many nodes the replica spans. Above 1, a fabric is required.
|
||||
Nodes int
|
||||
// Priority orders profiles; the highest that fits is used.
|
||||
Priority int
|
||||
}
|
||||
|
||||
// normalize fills in the defaults a catalog entry may omit.
|
||||
func (p Profile) normalize() Profile {
|
||||
if p.Nodes <= 0 {
|
||||
p.Nodes = 1
|
||||
}
|
||||
if p.GPUsPerNode <= 0 {
|
||||
p.GPUsPerNode = 1
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// NodeAssignment is one node's share of a replica.
|
||||
type NodeAssignment struct {
|
||||
NodeID string
|
||||
Rank int
|
||||
// GPUIndices are the local GPUs the shard uses.
|
||||
GPUIndices []int
|
||||
// HCAs are the adapters of the chosen fabric only.
|
||||
HCAs []string
|
||||
}
|
||||
|
||||
// Placement is a decision: which profile, which fabric, which nodes.
|
||||
type Placement struct {
|
||||
Profile Profile
|
||||
// Fabric is empty for a single-node replica.
|
||||
Fabric string
|
||||
// Pool names the (fabric, gpu class) pair this landed in.
|
||||
Pool string
|
||||
Nodes []NodeAssignment
|
||||
}
|
||||
|
||||
// Rejection records why a candidate was not used. Placement failures are
|
||||
// otherwise impossible to debug: "unschedulable" tells an operator nothing
|
||||
// about which of six constraints bit.
|
||||
type Rejection struct {
|
||||
Profile string
|
||||
Pool string
|
||||
Node string
|
||||
Reason string
|
||||
}
|
||||
|
||||
func (r Rejection) String() string {
|
||||
where := r.Pool
|
||||
if r.Node != "" {
|
||||
where += "/" + r.Node
|
||||
}
|
||||
if where == "" {
|
||||
where = "-"
|
||||
}
|
||||
return fmt.Sprintf("%s @ %s: %s", r.Profile, where, r.Reason)
|
||||
}
|
||||
|
||||
// Result is the outcome of a placement attempt.
|
||||
type Result struct {
|
||||
Placement *Placement
|
||||
// Rejections explains every candidate that was considered and discarded,
|
||||
// in the order they were tried. This is what `ctl explain-placement` prints.
|
||||
Rejections []Rejection
|
||||
}
|
||||
|
||||
// Explain renders the rejections as text.
|
||||
func (r Result) Explain() string {
|
||||
if len(r.Rejections) == 0 {
|
||||
return "no candidates were considered"
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, rej := range r.Rejections {
|
||||
b.WriteString(" ")
|
||||
b.WriteString(rej.String())
|
||||
b.WriteString("\n")
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// Scheduler places replicas onto an inventory.
|
||||
type Scheduler struct {
|
||||
inv *inventory.Store
|
||||
}
|
||||
|
||||
// New builds a scheduler over an inventory.
|
||||
func New(inv *inventory.Store) *Scheduler { return &Scheduler{inv: inv} }
|
||||
|
||||
// Place picks the first profile, in priority order, that fits.
|
||||
//
|
||||
// avoid lists nodes already hosting this model, so replicas of one model spread
|
||||
// across failure domains instead of stacking on one machine.
|
||||
func (s *Scheduler) Place(profiles []Profile, avoid []string) Result {
|
||||
var res Result
|
||||
|
||||
ordered := slices.Clone(profiles)
|
||||
sort.SliceStable(ordered, func(i, j int) bool { return ordered[i].Priority > ordered[j].Priority })
|
||||
|
||||
for _, raw := range ordered {
|
||||
p := raw.normalize()
|
||||
if p.Nodes == 1 {
|
||||
if pl, rej := s.placeSingleNode(p, avoid); pl != nil {
|
||||
res.Placement = pl
|
||||
return res
|
||||
} else {
|
||||
res.Rejections = append(res.Rejections, rej...)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if pl, rej := s.placeMultiNode(p, avoid); pl != nil {
|
||||
res.Placement = pl
|
||||
return res
|
||||
} else {
|
||||
res.Rejections = append(res.Rejections, rej...)
|
||||
}
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// placeSingleNode finds one node with enough free GPUs of the right class.
|
||||
//
|
||||
// Nodes without a fabric are eligible here and only here: a single-node replica
|
||||
// never talks to another node, so it has nothing to be reachable over.
|
||||
func (s *Scheduler) placeSingleNode(p Profile, avoid []string) (*Placement, []Rejection) {
|
||||
var rejections []Rejection
|
||||
|
||||
candidates := s.eligibleNodes(p, "", avoid, &rejections)
|
||||
if len(candidates) == 0 {
|
||||
return nil, rejections
|
||||
}
|
||||
|
||||
// Bin packing: the node with the fewest free GPUs that still fits. Filling
|
||||
// partly used nodes first keeps whole nodes free for the large models that
|
||||
// cannot be split.
|
||||
sort.SliceStable(candidates, func(i, j int) bool {
|
||||
return len(candidates[i].free) < len(candidates[j].free)
|
||||
})
|
||||
|
||||
best := candidates[0]
|
||||
gpus := s.pickGPUs(best.node, best.free, p.GPUsPerNode)
|
||||
if gpus == nil {
|
||||
rejections = append(rejections, Rejection{
|
||||
Profile: p.Name, Node: best.node.ID,
|
||||
Reason: fmt.Sprintf("has %d free %s GPUs but no usable set of %d", len(best.free), p.GPUClass, p.GPUsPerNode),
|
||||
})
|
||||
return nil, rejections
|
||||
}
|
||||
|
||||
return &Placement{
|
||||
Profile: p,
|
||||
Pool: "none/" + p.GPUClass,
|
||||
Nodes: []NodeAssignment{{NodeID: best.node.ID, Rank: 0, GPUIndices: gpus}},
|
||||
}, rejections
|
||||
}
|
||||
|
||||
// placeMultiNode finds N nodes on a single fabric.
|
||||
//
|
||||
// The fabric loop is the whole point: candidates are gathered per fabric and a
|
||||
// replica is only formed from within one of them. There is deliberately no path
|
||||
// that combines nodes from different fabrics, however well they would otherwise
|
||||
// fit.
|
||||
func (s *Scheduler) placeMultiNode(p Profile, avoid []string) (*Placement, []Rejection) {
|
||||
var rejections []Rejection
|
||||
|
||||
fabrics := s.inv.Fabrics()
|
||||
if len(fabrics) == 0 {
|
||||
return nil, []Rejection{{
|
||||
Profile: p.Name,
|
||||
Reason: fmt.Sprintf("needs %d nodes on one fabric, but no node reports a fabric", p.Nodes),
|
||||
}}
|
||||
}
|
||||
|
||||
for _, fabric := range fabrics {
|
||||
pool := fabric + "/" + p.GPUClass
|
||||
candidates := s.eligibleNodes(p, fabric, avoid, &rejections)
|
||||
|
||||
if len(candidates) < p.Nodes {
|
||||
rejections = append(rejections, Rejection{
|
||||
Profile: p.Name, Pool: pool,
|
||||
Reason: fmt.Sprintf("needs %d nodes, only %d qualify", p.Nodes, len(candidates)),
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
// Prefer the emptiest nodes for multi-node work: a replica spanning
|
||||
// nodes is the least movable thing in the cluster, so give it the
|
||||
// cleanest machines rather than squeezing it into the gaps.
|
||||
sort.SliceStable(candidates, func(i, j int) bool {
|
||||
return len(candidates[i].free) > len(candidates[j].free)
|
||||
})
|
||||
|
||||
assignments := make([]NodeAssignment, 0, p.Nodes)
|
||||
for rank, c := range candidates[:p.Nodes] {
|
||||
gpus := s.pickGPUs(c.node, c.free, p.GPUsPerNode)
|
||||
if gpus == nil {
|
||||
assignments = nil
|
||||
rejections = append(rejections, Rejection{
|
||||
Profile: p.Name, Pool: pool, Node: c.node.ID,
|
||||
Reason: fmt.Sprintf("no usable set of %d GPUs", p.GPUsPerNode),
|
||||
})
|
||||
break
|
||||
}
|
||||
assignments = append(assignments, NodeAssignment{
|
||||
NodeID: c.node.ID,
|
||||
Rank: rank,
|
||||
GPUIndices: gpus,
|
||||
// Only this fabric's adapters, never the node's others.
|
||||
HCAs: c.node.HCAsFor(fabric),
|
||||
})
|
||||
}
|
||||
if assignments == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
return &Placement{Profile: p, Fabric: fabric, Pool: pool, Nodes: assignments}, rejections
|
||||
}
|
||||
return nil, rejections
|
||||
}
|
||||
|
||||
// candidate is a node that passed the hard constraints.
|
||||
type candidate struct {
|
||||
node inventory.Node
|
||||
free []discovery.GPU
|
||||
}
|
||||
|
||||
// eligibleNodes applies the hard constraints and records why each node failed.
|
||||
//
|
||||
// fabric is empty for single-node placement, in which case attachment is not
|
||||
// checked at all.
|
||||
func (s *Scheduler) eligibleNodes(p Profile, fabric string, avoid []string, rejections *[]Rejection) []candidate {
|
||||
pool := "none"
|
||||
if fabric != "" {
|
||||
pool = fabric
|
||||
}
|
||||
pool += "/" + p.GPUClass
|
||||
|
||||
var out []candidate
|
||||
for _, n := range s.inv.Nodes() {
|
||||
reject := func(reason string) {
|
||||
*rejections = append(*rejections, Rejection{Profile: p.Name, Pool: pool, Node: n.ID, Reason: reason})
|
||||
}
|
||||
|
||||
if !n.State.Schedulable() {
|
||||
reject("node is " + n.State.String())
|
||||
continue
|
||||
}
|
||||
if slices.Contains(avoid, n.ID) {
|
||||
reject("already hosts this model; spreading across failure domains")
|
||||
continue
|
||||
}
|
||||
if fabric != "" {
|
||||
if !n.HasFabric(fabric) {
|
||||
// Not worth reporting for every node on every fabric; it is the
|
||||
// expected case, not a problem.
|
||||
continue
|
||||
}
|
||||
if len(n.HCAsFor(fabric)) == 0 {
|
||||
reject("attached to " + fabric + " but reports no adapters for it")
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
freeIdx := s.inv.FreeGPUs(n.ID)
|
||||
var free []discovery.GPU
|
||||
for _, g := range n.GPUs {
|
||||
if g.Class() != p.GPUClass || !slices.Contains(freeIdx, g.Index) {
|
||||
continue
|
||||
}
|
||||
free = append(free, g)
|
||||
}
|
||||
if len(free) < p.GPUsPerNode {
|
||||
reject(fmt.Sprintf("has %d free %s GPUs, needs %d", len(free), p.GPUClass, p.GPUsPerNode))
|
||||
continue
|
||||
}
|
||||
out = append(out, candidate{node: n, free: free})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// pickGPUs chooses which of a node's free GPUs to use.
|
||||
//
|
||||
// It prefers a set that shares an NVLink domain: within a node, a
|
||||
// tensor-parallel group that has to fall back to PCIe for its all-reduces runs
|
||||
// several times slower. NVLink topology is only available from the NVML build,
|
||||
// so without it this degrades to taking the lowest indices — which is usually
|
||||
// right, but only by accident of how vendors number slots.
|
||||
func (s *Scheduler) pickGPUs(n inventory.Node, free []discovery.GPU, want int) []int {
|
||||
if len(free) < want {
|
||||
return nil
|
||||
}
|
||||
if want == 1 {
|
||||
return []int{free[0].Index}
|
||||
}
|
||||
|
||||
freeSet := map[int]struct{}{}
|
||||
for _, g := range free {
|
||||
freeSet[g.Index] = struct{}{}
|
||||
}
|
||||
|
||||
// Try each free GPU as a seed and grow its NVLink-connected set.
|
||||
for _, seed := range free {
|
||||
group := []int{seed.Index}
|
||||
for _, peer := range seed.NVLinkPeers {
|
||||
if _, ok := freeSet[peer]; !ok || peer == seed.Index {
|
||||
continue
|
||||
}
|
||||
group = append(group, peer)
|
||||
if len(group) == want {
|
||||
slices.Sort(group)
|
||||
return group
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// No NVLink domain large enough, or no topology at all.
|
||||
out := make([]int, 0, want)
|
||||
for _, g := range free[:want] {
|
||||
out = append(out, g.Index)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,483 @@
|
||||
package placement
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math/rand/v2"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// node builds a node with n GPUs of one class.
|
||||
func node(id string, class string, n int, fabrics ...string) inventory.Node {
|
||||
name, mem := classToHardware(class)
|
||||
|
||||
nd := inventory.Node{ID: id, State: inventory.Ready}
|
||||
for i := range n {
|
||||
nd.GPUs = append(nd.GPUs, discovery.GPU{
|
||||
Index: i, Name: name, MemoryTotalMB: mem, ComputeCapability: "9.0",
|
||||
UUID: fmt.Sprintf("GPU-%s-%d", id, i),
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
nd.Fabrics = append(nd.Fabrics, inventory.FabricAttachment{
|
||||
Fabric: f, HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200,
|
||||
})
|
||||
}
|
||||
return nd
|
||||
}
|
||||
|
||||
// classToHardware maps a class back to a name and size that derive it.
|
||||
func classToHardware(class string) (string, int) {
|
||||
switch class {
|
||||
case "h100-80g":
|
||||
return "NVIDIA H100 80GB HBM3", 81559
|
||||
case "a100-40g":
|
||||
return "Tesla A100-SXM4-40GB", 40960
|
||||
default:
|
||||
return "NVIDIA L40S", 46068
|
||||
}
|
||||
}
|
||||
|
||||
func store(nodes ...inventory.Node) *inventory.Store {
|
||||
s := inventory.NewStore()
|
||||
for _, n := range nodes {
|
||||
s.Upsert(n)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func TestSingleNodePlacement(t *testing.T) {
|
||||
s := New(store(node("n1", "h100-80g", 8)))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if len(res.Placement.Nodes) != 1 {
|
||||
t.Fatalf("expected 1 node, got %d", len(res.Placement.Nodes))
|
||||
}
|
||||
if got := len(res.Placement.Nodes[0].GPUIndices); got != 4 {
|
||||
t.Errorf("got %d GPUs, want 4", got)
|
||||
}
|
||||
if res.Placement.Fabric != "" {
|
||||
t.Errorf("a single-node replica was given fabric %q", res.Placement.Fabric)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodeWithoutFabricTakesSingleNodeWork(t *testing.T) {
|
||||
// A single-node replica talks to nobody, so it has nothing to be reachable
|
||||
// over. Excluding fabric-less nodes here would waste them entirely.
|
||||
s := New(store(node("no-ib", "h100-80g", 8)))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("a node without IB was refused single-node work:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodeWithoutFabricRefusedMultiNode(t *testing.T) {
|
||||
s := New(store(
|
||||
node("no-ib-1", "h100-80g", 8),
|
||||
node("no-ib-2", "h100-80g", 8),
|
||||
))
|
||||
|
||||
res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("a multi-node replica was placed on nodes with no fabric")
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "no node reports a fabric") {
|
||||
t.Errorf("the reason is unclear:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
// TestNeverSpansTwoFabrics is the first hard invariant from plan §1.1.
|
||||
func TestNeverSpansTwoFabrics(t *testing.T) {
|
||||
// Two fabrics, one node each: together they would satisfy the count, and a
|
||||
// scheduler that only counted nodes would happily combine them.
|
||||
s := New(store(
|
||||
node("a1", "h100-80g", 8, "ib-a"),
|
||||
node("b1", "h100-80g", 8, "ib-b"),
|
||||
))
|
||||
|
||||
res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
|
||||
if res.Placement != nil {
|
||||
var fabrics []string
|
||||
for _, n := range res.Placement.Nodes {
|
||||
fabrics = append(fabrics, n.NodeID)
|
||||
}
|
||||
t.Fatalf("a replica was spread across fabrics: %v", fabrics)
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "only 1 qualify") {
|
||||
t.Errorf("explanation should say each fabric was short of nodes:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiNodeWithinOneFabric(t *testing.T) {
|
||||
s := New(store(
|
||||
node("a1", "h100-80g", 8, "ib-a"),
|
||||
node("a2", "h100-80g", 8, "ib-a"),
|
||||
node("b1", "h100-80g", 8, "ib-b"),
|
||||
))
|
||||
|
||||
res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if res.Placement.Fabric != "ib-a" {
|
||||
t.Errorf("fabric=%q, only ib-a has two nodes", res.Placement.Fabric)
|
||||
}
|
||||
for _, n := range res.Placement.Nodes {
|
||||
if n.NodeID == "b1" {
|
||||
t.Error("a node from the other fabric was included")
|
||||
}
|
||||
if len(n.HCAs) == 0 {
|
||||
t.Errorf("node %s got no HCAs; NCCL would pick its own", n.NodeID)
|
||||
}
|
||||
}
|
||||
// Ranks must be assigned, since rank 0 serves HTTP.
|
||||
ranks := []int{res.Placement.Nodes[0].Rank, res.Placement.Nodes[1].Rank}
|
||||
slices.Sort(ranks)
|
||||
if ranks[0] != 0 || ranks[1] != 1 {
|
||||
t.Errorf("ranks=%v", ranks)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNeverMixesGPUClasses is the second hard invariant from plan §1.1.
|
||||
func TestNeverMixesGPUClasses(t *testing.T) {
|
||||
// Enough GPUs in total, but split across classes.
|
||||
n := node("mixed", "h100-80g", 4)
|
||||
a100Name, a100Mem := classToHardware("a100-40g")
|
||||
for i := 4; i < 8; i++ {
|
||||
n.GPUs = append(n.GPUs, discovery.GPU{Index: i, Name: a100Name, MemoryTotalMB: a100Mem})
|
||||
}
|
||||
s := New(store(n))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("a placement used GPUs of more than one class")
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "has 4 free h100-80g GPUs, needs 8") {
|
||||
t.Errorf("explanation should count only the matching class:\n%s", res.Explain())
|
||||
}
|
||||
|
||||
// The 4-GPU profile does fit, using only H100s.
|
||||
res = s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("the matching-class profile was refused:\n%s", res.Explain())
|
||||
}
|
||||
for _, idx := range res.Placement.Nodes[0].GPUIndices {
|
||||
if idx >= 4 {
|
||||
t.Errorf("GPU %d is an A100 and was included in an h100 placement", idx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfilePriorityOrder(t *testing.T) {
|
||||
s := New(store(
|
||||
node("h", "h100-80g", 8, "ib-a"),
|
||||
node("a", "a100-40g", 8, "ib-a"),
|
||||
))
|
||||
|
||||
res := s.Place([]Profile{
|
||||
{Name: "a100", GPUClass: "a100-40g", GPUsPerNode: 8, Nodes: 1, Priority: 50},
|
||||
{Name: "h100", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1, Priority: 100},
|
||||
}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if res.Placement.Profile.Name != "h100" {
|
||||
t.Errorf("chose %q; the higher-priority profile fits", res.Placement.Profile.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFallsBackToLowerPriorityProfile(t *testing.T) {
|
||||
s := New(store(node("a", "a100-40g", 8)))
|
||||
|
||||
res := s.Place([]Profile{
|
||||
{Name: "h100", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1, Priority: 100},
|
||||
{Name: "a100", GPUClass: "a100-40g", GPUsPerNode: 8, Nodes: 1, Priority: 50},
|
||||
}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if res.Placement.Profile.Name != "a100" {
|
||||
t.Errorf("chose %q", res.Placement.Profile.Name)
|
||||
}
|
||||
// The rejected first choice must still be explained.
|
||||
if !strings.Contains(res.Explain(), "h100") {
|
||||
t.Errorf("the higher-priority profile's rejection is missing:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCordonedNodeExcluded(t *testing.T) {
|
||||
st := store(node("n1", "h100-80g", 8))
|
||||
st.SetState("n1", inventory.Cordoned)
|
||||
s := New(st)
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("a cordoned node was used")
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "cordoned") {
|
||||
t.Errorf("explanation:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnhealthyGPUsExcluded(t *testing.T) {
|
||||
n := node("n1", "h100-80g", 8)
|
||||
n.GPUs[0].Health.Problems = []string{"12 uncorrected ECC errors"}
|
||||
n.GPUs[1].Health.Problems = []string{"40000 MiB already allocated by another process"}
|
||||
s := New(store(n))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("unhealthy GPUs were counted as available")
|
||||
}
|
||||
|
||||
res = s.Place([]Profile{{Name: "tp6", GPUClass: "h100-80g", GPUsPerNode: 6, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("the six healthy GPUs should have been usable:\n%s", res.Explain())
|
||||
}
|
||||
for _, idx := range res.Placement.Nodes[0].GPUIndices {
|
||||
if idx == 0 || idx == 1 {
|
||||
t.Errorf("unhealthy GPU %d was placed on", idx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllocatedGPUsExcluded(t *testing.T) {
|
||||
st := store(node("n1", "h100-80g", 8))
|
||||
st.Allocate("n1", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
|
||||
s := New(st)
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("the remaining four GPUs should fit:\n%s", res.Explain())
|
||||
}
|
||||
for _, idx := range res.Placement.Nodes[0].GPUIndices {
|
||||
if idx < 4 {
|
||||
t.Errorf("GPU %d is already allocated", idx)
|
||||
}
|
||||
}
|
||||
|
||||
// Another four will not fit.
|
||||
st.Allocate("n1", inventory.Allocation{ShardID: "s2", ReplicaID: "r2", GPUIndices: []int{4, 5, 6, 7}})
|
||||
if res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil); res.Placement != nil {
|
||||
t.Fatal("a full node was placed on")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAvoidSpreadsAcrossNodes(t *testing.T) {
|
||||
s := New(store(
|
||||
node("n1", "h100-80g", 8),
|
||||
node("n2", "h100-80g", 8),
|
||||
))
|
||||
p := []Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}
|
||||
|
||||
first := s.Place(p, nil)
|
||||
if first.Placement == nil {
|
||||
t.Fatal("no first placement")
|
||||
}
|
||||
placed := first.Placement.Nodes[0].NodeID
|
||||
|
||||
second := s.Place(p, []string{placed})
|
||||
if second.Placement == nil {
|
||||
t.Fatalf("the second replica had nowhere to go:\n%s", second.Explain())
|
||||
}
|
||||
if second.Placement.Nodes[0].NodeID == placed {
|
||||
t.Error("both replicas landed on the same node despite avoid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBinPackingPrefersPartlyUsedNodes(t *testing.T) {
|
||||
// Keeping whole nodes free matters: a large model that needs all 8 GPUs
|
||||
// cannot be split, so scattering small replicas would strand it.
|
||||
st := store(
|
||||
node("empty", "h100-80g", 8),
|
||||
node("partial", "h100-80g", 8),
|
||||
)
|
||||
st.Allocate("partial", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
|
||||
s := New(st)
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if got := res.Placement.Nodes[0].NodeID; got != "partial" {
|
||||
t.Errorf("placed on %q; the partly used node should be filled first", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNVLinkGroupPreferred(t *testing.T) {
|
||||
// GPUs 0-3 and 4-7 form separate NVLink domains. A TP group crossing them
|
||||
// falls back to PCIe for its all-reduces and runs several times slower.
|
||||
n := node("n1", "h100-80g", 8)
|
||||
for i := range n.GPUs {
|
||||
var peers []int
|
||||
base := (i / 4) * 4
|
||||
for j := base; j < base+4; j++ {
|
||||
if j != i {
|
||||
peers = append(peers, j)
|
||||
}
|
||||
}
|
||||
n.GPUs[i].NVLinkPeers = peers
|
||||
}
|
||||
st := store(n)
|
||||
// Take GPU 0, so a naive "lowest indices" pick would span the two domains.
|
||||
st.Allocate("n1", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0}})
|
||||
s := New(st)
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
got := res.Placement.Nodes[0].GPUIndices
|
||||
if !slices.Equal(got, []int{4, 5, 6, 7}) {
|
||||
t.Errorf("chose %v; 4-7 is the only intact NVLink domain left", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplainIsUsefulWhenNothingFits(t *testing.T) {
|
||||
s := New(store(node("n1", "a100-40g", 2)))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("unexpected placement")
|
||||
}
|
||||
exp := res.Explain()
|
||||
// It must name the node, the profile, and what was actually wrong.
|
||||
for _, want := range []string{"n1", "tp8", "free h100-80g"} {
|
||||
if !strings.Contains(exp, want) {
|
||||
t.Errorf("explanation is missing %q:\n%s", want, exp)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyInventory(t *testing.T) {
|
||||
s := New(inventory.NewStore())
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("something was placed on an empty cluster")
|
||||
}
|
||||
}
|
||||
|
||||
// --- property-based invariants (plan §13) ---
|
||||
|
||||
// TestInvariantsHold generates random clusters and profiles and checks the two
|
||||
// hard constraints on every placement produced.
|
||||
//
|
||||
// The table-driven tests above cover the cases someone thought of; this covers
|
||||
// the ones nobody did.
|
||||
func TestInvariantsHold(t *testing.T) {
|
||||
rng := rand.New(rand.NewPCG(1, 2))
|
||||
classes := []string{"h100-80g", "a100-40g", "l40s-45g"}
|
||||
fabrics := []string{"ib-a", "ib-b", "ib-c"}
|
||||
|
||||
for iter := range 400 {
|
||||
st := inventory.NewStore()
|
||||
nodeCount := 1 + rng.IntN(6)
|
||||
|
||||
for i := range nodeCount {
|
||||
class := classes[rng.IntN(len(classes))]
|
||||
gpuCount := 1 + rng.IntN(8)
|
||||
|
||||
var nf []string
|
||||
switch rng.IntN(4) {
|
||||
case 0: // no fabric at all
|
||||
case 1:
|
||||
nf = []string{fabrics[rng.IntN(len(fabrics))]}
|
||||
default:
|
||||
nf = []string{fabrics[rng.IntN(len(fabrics))]}
|
||||
if rng.IntN(3) == 0 {
|
||||
// A node attached to two fabrics: allowed, and the reason
|
||||
// HCAs must be selected per fabric rather than per node.
|
||||
nf = append(nf, fabrics[rng.IntN(len(fabrics))])
|
||||
}
|
||||
}
|
||||
n := node(fmt.Sprintf("n%d", i), class, gpuCount, nf...)
|
||||
if rng.IntN(8) == 0 {
|
||||
n.State = inventory.Cordoned
|
||||
}
|
||||
st.Upsert(n)
|
||||
}
|
||||
|
||||
s := New(st)
|
||||
p := Profile{
|
||||
Name: "gen",
|
||||
GPUClass: classes[rng.IntN(len(classes))],
|
||||
GPUsPerNode: 1 + rng.IntN(8),
|
||||
Nodes: 1 + rng.IntN(3),
|
||||
}
|
||||
res := s.Place([]Profile{p}, nil)
|
||||
if res.Placement == nil {
|
||||
continue
|
||||
}
|
||||
checkInvariants(t, st, p, res.Placement, iter)
|
||||
}
|
||||
}
|
||||
|
||||
func checkInvariants(t *testing.T, st *inventory.Store, p Profile, pl *Placement, iter int) {
|
||||
t.Helper()
|
||||
|
||||
if len(pl.Nodes) != max(p.Nodes, 1) {
|
||||
t.Fatalf("iter %d: placed on %d nodes, profile wants %d", iter, len(pl.Nodes), p.Nodes)
|
||||
}
|
||||
|
||||
seenNodes := map[string]struct{}{}
|
||||
for _, a := range pl.Nodes {
|
||||
if _, dup := seenNodes[a.NodeID]; dup {
|
||||
t.Fatalf("iter %d: node %s used twice in one replica", iter, a.NodeID)
|
||||
}
|
||||
seenNodes[a.NodeID] = struct{}{}
|
||||
|
||||
n, ok := st.Get(a.NodeID)
|
||||
if !ok {
|
||||
t.Fatalf("iter %d: placed on unknown node %s", iter, a.NodeID)
|
||||
}
|
||||
if !n.State.Schedulable() {
|
||||
t.Fatalf("iter %d: placed on a %s node", iter, n.State)
|
||||
}
|
||||
if len(a.GPUIndices) != p.GPUsPerNode {
|
||||
t.Fatalf("iter %d: node %s got %d GPUs, wants %d", iter, a.NodeID, len(a.GPUIndices), p.GPUsPerNode)
|
||||
}
|
||||
|
||||
// Invariant: every GPU is of the profile's class, healthy, and real.
|
||||
byIndex := map[int]discovery.GPU{}
|
||||
for _, g := range n.GPUs {
|
||||
byIndex[g.Index] = g
|
||||
}
|
||||
for _, idx := range a.GPUIndices {
|
||||
g, ok := byIndex[idx]
|
||||
if !ok {
|
||||
t.Fatalf("iter %d: node %s has no GPU %d", iter, a.NodeID, idx)
|
||||
}
|
||||
if g.Class() != p.GPUClass {
|
||||
t.Fatalf("iter %d: GPU %d on %s is %s, profile wants %s",
|
||||
iter, idx, a.NodeID, g.Class(), p.GPUClass)
|
||||
}
|
||||
if !g.Health.OK() {
|
||||
t.Fatalf("iter %d: unhealthy GPU %d was placed on", iter, idx)
|
||||
}
|
||||
}
|
||||
|
||||
// Invariant: a multi-node replica lives on one fabric, and each shard
|
||||
// gets only that fabric's adapters.
|
||||
if p.Nodes > 1 {
|
||||
if pl.Fabric == "" {
|
||||
t.Fatalf("iter %d: multi-node placement with no fabric", iter)
|
||||
}
|
||||
if !n.HasFabric(pl.Fabric) {
|
||||
t.Fatalf("iter %d: node %s is not attached to %s", iter, a.NodeID, pl.Fabric)
|
||||
}
|
||||
want := n.HCAsFor(pl.Fabric)
|
||||
if !slices.Equal(a.HCAs, want) {
|
||||
t.Fatalf("iter %d: node %s got HCAs %v, want exactly %v for %s",
|
||||
iter, a.NodeID, a.HCAs, want, pl.Fabric)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,782 @@
|
||||
package reconciler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"math/rand/v2"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/placement"
|
||||
)
|
||||
|
||||
// ShardCommand is everything an agent needs to launch one shard.
|
||||
type ShardCommand struct {
|
||||
ReplicaID string
|
||||
ShardID string
|
||||
NodeID string
|
||||
Spec engines.ShardSpec
|
||||
// StartupTimeout comes from the profile; model loading dominates it.
|
||||
StartupTimeout time.Duration
|
||||
// Pool and Fabric are control plane concepts the agent cannot derive. They
|
||||
// travel with the command so the agent can report them back after a
|
||||
// control plane restart (plan §6.5).
|
||||
Pool string
|
||||
Fabric string
|
||||
}
|
||||
|
||||
// Dispatcher sends commands to agents.
|
||||
//
|
||||
// It is an interface so the reconciliation loop can be tested without a network:
|
||||
// the loop's logic is where the interesting failures are, not the transport.
|
||||
type Dispatcher interface {
|
||||
StartShard(ctx context.Context, cmd ShardCommand) error
|
||||
StopShard(ctx context.Context, nodeID, shardID string) error
|
||||
DrainShard(ctx context.Context, nodeID, shardID string) error
|
||||
}
|
||||
|
||||
// Config tunes the reconciliation loop.
|
||||
type Config struct {
|
||||
// Interval is how often desired and actual are compared when nothing else
|
||||
// happens. Events wake the loop sooner.
|
||||
Interval time.Duration
|
||||
// BaseBackoff is how long to wait before rebuilding a failed replica.
|
||||
BaseBackoff time.Duration
|
||||
// MaxBackoff caps the exponential growth.
|
||||
MaxBackoff time.Duration
|
||||
// UnreachableAfter is how long without a heartbeat before a node is
|
||||
// unreachable.
|
||||
UnreachableAfter time.Duration
|
||||
// LostAfter is how long unreachable before its replicas are declared dead.
|
||||
// Longer than UnreachableAfter on purpose: a brief network blip must not
|
||||
// tear down a replica that is still serving perfectly well.
|
||||
LostAfter time.Duration
|
||||
// DrainGrace bounds how long a draining replica waits for its replacement.
|
||||
// Without a bound, draining a node whose model cannot be placed anywhere
|
||||
// else would never finish.
|
||||
DrainGrace time.Duration
|
||||
// PortRange is the per-node port range for engine HTTP and dist-init.
|
||||
PortRange [2]int
|
||||
// MgmtInterface is the ethernet interface used for NCCL bootstrap.
|
||||
MgmtInterface string
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Interval: 5 * time.Second,
|
||||
BaseBackoff: 10 * time.Second,
|
||||
MaxBackoff: 10 * time.Minute,
|
||||
UnreachableAfter: 15 * time.Second,
|
||||
LostAfter: 60 * time.Second,
|
||||
DrainGrace: 5 * time.Minute,
|
||||
PortRange: [2]int{20000, 20999},
|
||||
MgmtInterface: "eth0",
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() {
|
||||
d := DefaultConfig()
|
||||
if c.Interval <= 0 {
|
||||
c.Interval = d.Interval
|
||||
}
|
||||
if c.BaseBackoff <= 0 {
|
||||
c.BaseBackoff = d.BaseBackoff
|
||||
}
|
||||
if c.MaxBackoff < c.BaseBackoff {
|
||||
c.MaxBackoff = max(c.BaseBackoff, d.MaxBackoff)
|
||||
}
|
||||
if c.UnreachableAfter <= 0 {
|
||||
c.UnreachableAfter = d.UnreachableAfter
|
||||
}
|
||||
if c.LostAfter < c.UnreachableAfter {
|
||||
c.LostAfter = max(c.UnreachableAfter*4, d.LostAfter)
|
||||
}
|
||||
if c.DrainGrace <= 0 {
|
||||
c.DrainGrace = d.DrainGrace
|
||||
}
|
||||
if c.PortRange[1] <= c.PortRange[0] {
|
||||
c.PortRange = d.PortRange
|
||||
}
|
||||
if c.MgmtInterface == "" {
|
||||
c.MgmtInterface = d.MgmtInterface
|
||||
}
|
||||
}
|
||||
|
||||
// Reconciler compares the catalog against what is running and closes the gap.
|
||||
type Reconciler struct {
|
||||
cfg Config
|
||||
cat *catalog.Catalog
|
||||
inv *inventory.Store
|
||||
sched *placement.Scheduler
|
||||
disp Dispatcher
|
||||
log *slog.Logger
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
replicas map[string]*Replica
|
||||
// nextAttempt is the earliest a model may be rebuilt, per model.
|
||||
nextAttempt map[string]time.Time
|
||||
// attempts counts consecutive failures per model, for the backoff.
|
||||
attempts map[string]int
|
||||
// lastExplain keeps the most recent placement failure per model, so
|
||||
// `ctl explain-placement` can answer without re-running placement.
|
||||
lastExplain map[string]placement.Result
|
||||
seq int
|
||||
|
||||
wake chan struct{}
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// New builds a reconciler.
|
||||
func New(cfg Config, cat *catalog.Catalog, inv *inventory.Store, disp Dispatcher, log *slog.Logger) *Reconciler {
|
||||
cfg.normalize()
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Reconciler{
|
||||
cfg: cfg, cat: cat, inv: inv, disp: disp, log: log,
|
||||
sched: placement.New(inv),
|
||||
now: time.Now,
|
||||
replicas: map[string]*Replica{},
|
||||
nextAttempt: map[string]time.Time{},
|
||||
attempts: map[string]int{},
|
||||
lastExplain: map[string]placement.Result{},
|
||||
wake: make(chan struct{}, 1),
|
||||
stop: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// SetClock replaces the time source, for tests.
|
||||
func (r *Reconciler) SetClock(now func() time.Time) { r.now = now }
|
||||
|
||||
// Start runs the loop until Close.
|
||||
func (r *Reconciler) Start(ctx context.Context) {
|
||||
go func() {
|
||||
defer close(r.done)
|
||||
t := time.NewTicker(r.cfg.Interval)
|
||||
defer t.Stop()
|
||||
|
||||
r.Reconcile(ctx)
|
||||
for {
|
||||
select {
|
||||
case <-r.stop:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-r.wake:
|
||||
r.Reconcile(ctx)
|
||||
case <-t.C:
|
||||
r.Reconcile(ctx)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Close stops the loop.
|
||||
func (r *Reconciler) Close() {
|
||||
r.once.Do(func() {
|
||||
close(r.stop)
|
||||
<-r.done
|
||||
})
|
||||
}
|
||||
|
||||
// Notify wakes the loop, for when something changed and waiting for the next
|
||||
// tick would add pointless latency.
|
||||
func (r *Reconciler) Notify() {
|
||||
select {
|
||||
case r.wake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// Reconcile performs one pass. It is exported so tests can drive it directly
|
||||
// instead of racing a ticker.
|
||||
func (r *Reconciler) Reconcile(ctx context.Context) {
|
||||
r.expireNodes()
|
||||
r.drainCordonedNodes(ctx)
|
||||
r.reapFailed(ctx)
|
||||
|
||||
for _, m := range r.cat.Models() {
|
||||
r.reconcileModel(ctx, m)
|
||||
}
|
||||
}
|
||||
|
||||
// expireNodes moves silent nodes through Unreachable to Lost, and fails the
|
||||
// replicas of nodes that are gone.
|
||||
//
|
||||
// The two-stage delay is deliberate: a brief network blip must not tear down
|
||||
// replicas that are still serving traffic perfectly well, and rebuilding a
|
||||
// multi-node replica costs tens of minutes of model loading.
|
||||
func (r *Reconciler) expireNodes() {
|
||||
now := r.now()
|
||||
for _, n := range r.inv.Nodes() {
|
||||
silent := now.Sub(n.LastSeen)
|
||||
switch {
|
||||
case n.State == inventory.Ready && silent > r.cfg.UnreachableAfter:
|
||||
r.inv.SetState(n.ID, inventory.Unreachable)
|
||||
r.log.Warn("node is unreachable", "node", n.ID, "silent_for", silent.Round(time.Second))
|
||||
case n.State == inventory.Unreachable && silent > r.cfg.LostAfter:
|
||||
r.inv.SetState(n.ID, inventory.Lost)
|
||||
r.log.Error("node lost; its replicas will be rebuilt elsewhere",
|
||||
"node", n.ID, "silent_for", silent.Round(time.Second))
|
||||
r.failReplicasOn(n.ID, "node "+n.ID+" was lost")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// drainCordonedNodes moves work off nodes an operator is taking out of service.
|
||||
//
|
||||
// Draining a node is not the same as losing one: the engines are healthy and
|
||||
// still serving, so their replicas are marked Draining rather than Failed. That
|
||||
// makes the model look short of replicas, a replacement is placed elsewhere,
|
||||
// and the original is stopped only once the replacement is serving — which is
|
||||
// the difference between a rolling move and an outage.
|
||||
func (r *Reconciler) drainCordonedNodes(ctx context.Context) {
|
||||
var toDrain []*Replica
|
||||
|
||||
r.mu.Lock()
|
||||
for _, rep := range r.replicas {
|
||||
if rep.State != Ready && rep.State != Starting && rep.State != Placing {
|
||||
continue
|
||||
}
|
||||
for _, s := range rep.Shards {
|
||||
n, ok := r.inv.Get(s.NodeID)
|
||||
if ok && n.State == inventory.Draining {
|
||||
rep.State = Draining
|
||||
rep.DrainingAt = r.now()
|
||||
toDrain = append(toDrain, rep)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
shards := map[string][]Shard{}
|
||||
for _, rep := range toDrain {
|
||||
shards[rep.ID] = slices.Clone(rep.Shards)
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, rep := range toDrain {
|
||||
r.log.Info("draining replica off a node being taken out of service",
|
||||
"replica", rep.ID, "nodes", rep.NodeIDs())
|
||||
for _, s := range shards[rep.ID] {
|
||||
if err := r.disp.DrainShard(ctx, s.NodeID, s.ID); err != nil {
|
||||
r.log.Warn("draining a shard failed", "shard", s.ID, "err", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reconciler) failReplicasOn(nodeID, reason string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, rep := range r.replicas {
|
||||
if rep.State.Terminal() || !rep.OnNode(nodeID) {
|
||||
continue
|
||||
}
|
||||
rep.State = Failed
|
||||
rep.Err = reason
|
||||
}
|
||||
}
|
||||
|
||||
// reapFailed releases the resources of terminal replicas and tells the agents
|
||||
// to clean up whatever might still be running.
|
||||
func (r *Reconciler) reapFailed(ctx context.Context) {
|
||||
r.mu.Lock()
|
||||
var doomed []*Replica
|
||||
for _, rep := range r.replicas {
|
||||
if rep.State.Terminal() {
|
||||
doomed = append(doomed, rep)
|
||||
}
|
||||
}
|
||||
for _, rep := range doomed {
|
||||
delete(r.replicas, rep.ID)
|
||||
if rep.State == Failed {
|
||||
// Carry the failure count on the model, not the replica: the
|
||||
// replica is about to disappear, and the backoff has to survive it.
|
||||
r.attempts[rep.Model]++
|
||||
r.nextAttempt[rep.Model] = r.now().Add(r.backoff(r.attempts[rep.Model]))
|
||||
}
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, rep := range doomed {
|
||||
// rep is out of the map now, so no agent report can reach it any more.
|
||||
if rep.State == Failed {
|
||||
r.mu.Lock()
|
||||
retry := r.backoff(r.attempts[rep.Model])
|
||||
r.mu.Unlock()
|
||||
r.log.Error("replica failed",
|
||||
"replica", rep.ID, "model", rep.Model, "err", rep.Err,
|
||||
"retry_in", retry.Round(time.Second))
|
||||
}
|
||||
// Stop every shard, including the ones that looked healthy: a partially
|
||||
// alive replica holds GPUs while serving nothing.
|
||||
for _, s := range rep.Shards {
|
||||
if err := r.disp.StopShard(ctx, s.NodeID, s.ID); err != nil {
|
||||
r.log.Warn("stopping a shard of a dead replica failed",
|
||||
"replica", rep.ID, "shard", s.ID, "err", err)
|
||||
}
|
||||
}
|
||||
r.inv.ReleaseReplica(rep.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// backoff grows with consecutive failures and is jittered, so a cluster-wide
|
||||
// problem does not produce synchronized retry storms across every model.
|
||||
func (r *Reconciler) backoff(attempt int) time.Duration {
|
||||
if attempt <= 0 {
|
||||
return 0
|
||||
}
|
||||
d := r.cfg.BaseBackoff << min(attempt-1, 16)
|
||||
if d > r.cfg.MaxBackoff || d <= 0 {
|
||||
d = r.cfg.MaxBackoff
|
||||
}
|
||||
jitter := time.Duration(rand.Int64N(int64(d / 4))) //nolint:gosec // jitter
|
||||
return d - d/8 + jitter
|
||||
}
|
||||
|
||||
func (r *Reconciler) reconcileModel(ctx context.Context, m catalog.Model) {
|
||||
// Snapshot under the lock and work from values afterwards. A replica
|
||||
// pointer stops being safe the moment it is in the map: an agent's status
|
||||
// report can mutate it from another goroutine at any time.
|
||||
r.mu.Lock()
|
||||
var liveIDs, drainingIDs, avoid []string
|
||||
for _, rep := range r.replicas {
|
||||
if rep.Model != m.Name {
|
||||
continue
|
||||
}
|
||||
switch {
|
||||
case rep.State.Counts():
|
||||
liveIDs = append(liveIDs, rep.ID)
|
||||
avoid = append(avoid, rep.NodeIDs()...)
|
||||
case rep.State == Draining:
|
||||
drainingIDs = append(drainingIDs, rep.ID)
|
||||
}
|
||||
}
|
||||
blockedUntil := r.nextAttempt[m.Name]
|
||||
r.mu.Unlock()
|
||||
|
||||
desired := m.Scaling.Min
|
||||
switch {
|
||||
case len(liveIDs) < desired:
|
||||
if r.now().Before(blockedUntil) {
|
||||
// Still in backoff. Rebuilding immediately would hot-loop on a
|
||||
// model that cannot start — a missing weights path, say — and bury
|
||||
// the real error in noise.
|
||||
return
|
||||
}
|
||||
r.scaleUp(ctx, m, avoid)
|
||||
case len(liveIDs) > desired:
|
||||
r.scaleDown(ctx, liveIDs[desired:])
|
||||
}
|
||||
|
||||
// A drained replica is stopped once its replacement is serving, or once the
|
||||
// grace period expires. Stopping it the moment it was drained would cut
|
||||
// in-flight requests, which is the thing draining exists to avoid.
|
||||
replaced := len(liveIDs) >= desired
|
||||
for _, id := range drainingIDs {
|
||||
r.mu.Lock()
|
||||
rep, ok := r.replicas[id]
|
||||
var expired bool
|
||||
if ok {
|
||||
expired = !rep.DrainingAt.IsZero() && r.now().Sub(rep.DrainingAt) > r.cfg.DrainGrace
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
switch {
|
||||
case !ok:
|
||||
case replaced:
|
||||
r.stopReplica(ctx, id)
|
||||
case expired:
|
||||
// The replacement never arrived — often because nowhere else can
|
||||
// host it. Draining forever would leave the node undrainable.
|
||||
r.log.Warn("stopping a drained replica without a replacement",
|
||||
"replica", id, "after", r.cfg.DrainGrace)
|
||||
r.stopReplica(ctx, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// scaleUp places and starts one replica per pass.
|
||||
//
|
||||
// One at a time on purpose: placement decisions depend on what is already
|
||||
// allocated, and starting several at once from the same inventory snapshot
|
||||
// would hand them the same GPUs.
|
||||
func (r *Reconciler) scaleUp(ctx context.Context, m catalog.Model, avoid []string) {
|
||||
res := r.sched.Place(m.PlacementProfiles(), avoid)
|
||||
if res.Placement == nil && len(avoid) > 0 {
|
||||
// Spreading across failure domains is a preference, not a rule. If it
|
||||
// is the only thing standing in the way, take the placement anyway:
|
||||
// a second replica on one node still survives an engine crash.
|
||||
res = r.sched.Place(m.PlacementProfiles(), nil)
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.lastExplain[m.Name] = res
|
||||
r.mu.Unlock()
|
||||
|
||||
if res.Placement == nil {
|
||||
r.log.Warn("no placement for model", "model", m.Name, "reasons", res.Explain())
|
||||
return
|
||||
}
|
||||
if err := r.startReplica(ctx, m, res.Placement); err != nil {
|
||||
r.log.Error("starting replica failed", "model", m.Name, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reconciler) startReplica(ctx context.Context, m catalog.Model, pl *placement.Placement) error {
|
||||
prof, ok := m.ProfileByName(pl.Profile.Name)
|
||||
if !ok {
|
||||
return fmt.Errorf("profile %q vanished from model %q", pl.Profile.Name, m.Name)
|
||||
}
|
||||
espec, ok := m.EngineProfile(prof.Name)
|
||||
if !ok {
|
||||
return fmt.Errorf("model %q: no engine profile %q", m.Name, prof.Name)
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.seq++
|
||||
id := fmt.Sprintf("%s-%d", m.Name, r.seq)
|
||||
attempt := r.attempts[m.Name]
|
||||
r.mu.Unlock()
|
||||
|
||||
rep := &Replica{
|
||||
ID: id, Model: m.Name, Profile: prof.Name,
|
||||
Fabric: pl.Fabric, Pool: pl.Pool, State: Placing,
|
||||
Weight: prof.Weight, CreatedAt: r.now(), Attempt: attempt,
|
||||
}
|
||||
|
||||
// The leader's address and port are fixed before any shard starts: every
|
||||
// rank needs the same rendezvous point, and a worker that starts first must
|
||||
// already know where to dial.
|
||||
leader := pl.Nodes[0]
|
||||
leaderNode, ok := r.inv.Get(leader.NodeID)
|
||||
if !ok {
|
||||
return fmt.Errorf("leader node %s disappeared", leader.NodeID)
|
||||
}
|
||||
basePort := r.allocatePorts(pl)
|
||||
|
||||
for _, a := range pl.Nodes {
|
||||
shardID := fmt.Sprintf("%s-r%d", id, a.Rank)
|
||||
rep.Shards = append(rep.Shards, Shard{
|
||||
ID: shardID, NodeID: a.NodeID, Rank: a.Rank,
|
||||
GPUIndices: a.GPUIndices, HCAs: a.HCAs,
|
||||
HTTPPort: basePort, State: supervisor.Pending,
|
||||
})
|
||||
r.inv.Allocate(a.NodeID, inventory.Allocation{
|
||||
ShardID: shardID, ReplicaID: id, GPUIndices: a.GPUIndices,
|
||||
})
|
||||
}
|
||||
rep.Endpoint = fmt.Sprintf("http://%s:%d", addrOf(leaderNode), basePort)
|
||||
|
||||
// Everything read from rep has to be read now: once it is in the map an
|
||||
// agent status report can mutate it concurrently, and the race detector is
|
||||
// right to object.
|
||||
//
|
||||
// The leader first, then workers ascending. Both engines put the rendezvous
|
||||
// point on rank 0 — SGLang's dist-init address and vLLM's Ray head — so a
|
||||
// worker started first has nothing to connect to and spends its startup
|
||||
// budget retrying.
|
||||
//
|
||||
// An earlier version started workers first, on the theory that the head
|
||||
// should wait for them. That is backwards: vLLM's server blocks until the
|
||||
// Ray cluster has the GPUs it needs, so the leader can be started
|
||||
// immediately and will wait on its own.
|
||||
order := slices.Clone(rep.Shards)
|
||||
slices.SortFunc(order, func(a, b Shard) int { return a.Rank - b.Rank })
|
||||
nodeIDs := rep.NodeIDs()
|
||||
|
||||
r.mu.Lock()
|
||||
r.replicas[id] = rep
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, s := range order {
|
||||
cmd := ShardCommand{
|
||||
ReplicaID: id,
|
||||
ShardID: s.ID,
|
||||
NodeID: s.NodeID,
|
||||
Spec: engines.ShardSpec{
|
||||
ReplicaID: id,
|
||||
ShardID: s.ID,
|
||||
Rank: s.Rank,
|
||||
NNodes: len(pl.Nodes),
|
||||
LeaderAddr: addrOf(leaderNode),
|
||||
LeaderPort: basePort + 1,
|
||||
GPUIndices: s.GPUIndices,
|
||||
HCAs: s.HCAs,
|
||||
MgmtInterface: r.cfg.MgmtInterface,
|
||||
HTTPPort: s.HTTPPort,
|
||||
Profile: espec,
|
||||
},
|
||||
StartupTimeout: prof.StartupTimeout,
|
||||
Pool: pl.Pool,
|
||||
Fabric: pl.Fabric,
|
||||
}
|
||||
if err := r.disp.StartShard(ctx, cmd); err != nil {
|
||||
r.markShard(id, s.ID, supervisor.Failed, err.Error())
|
||||
return fmt.Errorf("starting shard %s: %w", s.ID, err)
|
||||
}
|
||||
r.markShard(id, s.ID, supervisor.Starting, "")
|
||||
}
|
||||
|
||||
r.log.Info("replica placed",
|
||||
"replica", id, "model", m.Name, "profile", prof.Name,
|
||||
"pool", pl.Pool, "fabric", pl.Fabric, "nodes", nodeIDs)
|
||||
return nil
|
||||
}
|
||||
|
||||
// allocatePorts picks a port shared by every shard of a replica.
|
||||
//
|
||||
// The same port on every node keeps the leader address predictable and the
|
||||
// port bookkeeping trivial; nodes do not share a network namespace, so there is
|
||||
// no conflict in reusing it.
|
||||
func (r *Reconciler) allocatePorts(pl *placement.Placement) int {
|
||||
used := map[int]struct{}{}
|
||||
r.mu.Lock()
|
||||
for _, rep := range r.replicas {
|
||||
for _, s := range rep.Shards {
|
||||
used[s.HTTPPort] = struct{}{}
|
||||
}
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
// Two ports per replica: HTTP and the distributed init rendezvous.
|
||||
for p := r.cfg.PortRange[0]; p < r.cfg.PortRange[1]; p += 2 {
|
||||
if _, taken := used[p]; !taken {
|
||||
return p
|
||||
}
|
||||
}
|
||||
return r.cfg.PortRange[0]
|
||||
}
|
||||
|
||||
func (r *Reconciler) scaleDown(ctx context.Context, excess []string) {
|
||||
for _, id := range excess {
|
||||
r.mu.Lock()
|
||||
rep, ok := r.replicas[id]
|
||||
if !ok {
|
||||
r.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
rep.State = Draining
|
||||
rep.DrainingAt = r.now()
|
||||
shards := slices.Clone(rep.Shards)
|
||||
r.mu.Unlock()
|
||||
|
||||
r.log.Info("draining replica", "replica", id)
|
||||
for _, s := range shards {
|
||||
if err := r.disp.DrainShard(ctx, s.NodeID, s.ID); err != nil {
|
||||
r.log.Warn("draining a shard failed", "shard", s.ID, "err", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reconciler) stopReplica(ctx context.Context, id string) {
|
||||
r.mu.Lock()
|
||||
rep, ok := r.replicas[id]
|
||||
if !ok {
|
||||
r.mu.Unlock()
|
||||
return
|
||||
}
|
||||
shards := slices.Clone(rep.Shards)
|
||||
rep.State = Stopped
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, s := range shards {
|
||||
if err := r.disp.StopShard(ctx, s.NodeID, s.ID); err != nil {
|
||||
r.log.Warn("stopping a shard failed", "shard", s.ID, "err", err)
|
||||
}
|
||||
}
|
||||
r.inv.ReleaseReplica(id)
|
||||
|
||||
r.mu.Lock()
|
||||
delete(r.replicas, id)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
// Adopt rebuilds replica state from shards an agent reports it is already
|
||||
// running (plan §6.5).
|
||||
//
|
||||
// This is what makes a control plane restart survivable. Without it the
|
||||
// reconciler sees an empty cluster, places duplicates, and the engines already
|
||||
// running keep holding GPUs the inventory believes are free — so the duplicates
|
||||
// fail to allocate, with an error naming the new process rather than the old.
|
||||
//
|
||||
// Shards whose model or profile is no longer in the catalog are returned for
|
||||
// the caller to stop: nothing will ever claim them.
|
||||
func (r *Reconciler) Adopt(shards []AdoptedShard) (unclaimed []AdoptedShard) {
|
||||
byReplica := map[string][]AdoptedShard{}
|
||||
for _, s := range shards {
|
||||
byReplica[s.ReplicaID] = append(byReplica[s.ReplicaID], s)
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
for replicaID, group := range byReplica {
|
||||
if _, known := r.replicas[replicaID]; known {
|
||||
// Already tracked — an agent reconnecting rather than a control
|
||||
// plane restart. Its status reports keep the state current.
|
||||
continue
|
||||
}
|
||||
|
||||
first := group[0]
|
||||
m, ok := r.cat.Resolve(first.Model)
|
||||
if !ok {
|
||||
unclaimed = append(unclaimed, group...)
|
||||
continue
|
||||
}
|
||||
prof, ok := m.ProfileByName(first.Profile)
|
||||
if !ok {
|
||||
unclaimed = append(unclaimed, group...)
|
||||
continue
|
||||
}
|
||||
|
||||
rep := &Replica{
|
||||
ID: replicaID, Model: m.Name, Profile: prof.Name,
|
||||
Fabric: first.Fabric, Pool: first.Pool,
|
||||
Weight: prof.Weight, CreatedAt: r.now(),
|
||||
Endpoint: first.Endpoint,
|
||||
}
|
||||
for _, s := range group {
|
||||
rep.Shards = append(rep.Shards, Shard{
|
||||
ID: s.ShardID, NodeID: s.NodeID, Rank: s.Rank,
|
||||
GPUIndices: s.GPUIndices, HCAs: s.HCAs,
|
||||
HTTPPort: s.HTTPPort, State: s.State,
|
||||
})
|
||||
// Reclaim the GPUs, or the next placement will hand them out twice.
|
||||
r.inv.Allocate(s.NodeID, inventory.Allocation{
|
||||
ShardID: s.ShardID, ReplicaID: replicaID, GPUIndices: s.GPUIndices,
|
||||
})
|
||||
}
|
||||
// The endpoint comes from the leader; a group missing rank 0 is a
|
||||
// partial replica and will be failed by deriveState below.
|
||||
for _, s := range group {
|
||||
if s.Rank == 0 && s.Endpoint != "" {
|
||||
rep.Endpoint = s.Endpoint
|
||||
}
|
||||
}
|
||||
rep.deriveState()
|
||||
r.replicas[replicaID] = rep
|
||||
|
||||
r.log.Info("adopted replica from the fleet",
|
||||
"replica", replicaID, "model", m.Name, "shards", len(group), "state", rep.State)
|
||||
}
|
||||
return unclaimed
|
||||
}
|
||||
|
||||
// AdoptedShard is one running shard an agent reported at registration.
|
||||
type AdoptedShard struct {
|
||||
NodeID string
|
||||
ReplicaID string
|
||||
ShardID string
|
||||
Rank int
|
||||
GPUIndices []int
|
||||
HCAs []string
|
||||
HTTPPort int
|
||||
State supervisor.State
|
||||
Model string
|
||||
Profile string
|
||||
Pool string
|
||||
Fabric string
|
||||
// Endpoint is the leader's base URL, set on the rank 0 shard.
|
||||
Endpoint string
|
||||
}
|
||||
|
||||
// OnShardStatus records a status report from an agent.
|
||||
func (r *Reconciler) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
|
||||
r.markShard(replicaID, shardID, st, errMsg)
|
||||
r.Notify()
|
||||
}
|
||||
|
||||
func (r *Reconciler) markShard(replicaID, shardID string, st supervisor.State, errMsg string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
rep, ok := r.replicas[replicaID]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for i := range rep.Shards {
|
||||
if rep.Shards[i].ID != shardID {
|
||||
continue
|
||||
}
|
||||
rep.Shards[i].State = st
|
||||
if errMsg != "" {
|
||||
rep.Shards[i].Err = errMsg
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
prev := rep.State
|
||||
rep.deriveState()
|
||||
if rep.State == Ready && prev != Ready {
|
||||
// A model that comes up clears its backoff; the next failure starts
|
||||
// counting from scratch rather than from an old streak.
|
||||
r.attempts[rep.Model] = 0
|
||||
delete(r.nextAttempt, rep.Model)
|
||||
}
|
||||
}
|
||||
|
||||
// Replicas returns a snapshot of every replica, sorted by id.
|
||||
func (r *Reconciler) Replicas() []Replica {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
out := make([]Replica, 0, len(r.replicas))
|
||||
for _, rep := range r.replicas {
|
||||
clone := *rep
|
||||
clone.Shards = slices.Clone(rep.Shards)
|
||||
out = append(out, clone)
|
||||
}
|
||||
slices.SortFunc(out, func(a, b Replica) int {
|
||||
if a.ID < b.ID {
|
||||
return -1
|
||||
}
|
||||
if a.ID > b.ID {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// ReadyReplicas returns the replicas the router may send traffic to.
|
||||
func (r *Reconciler) ReadyReplicas() []Replica {
|
||||
var out []Replica
|
||||
for _, rep := range r.Replicas() {
|
||||
if rep.State == Ready {
|
||||
out = append(out, rep)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Explain returns the most recent placement result for a model, which is what
|
||||
// `ctl explain-placement` prints.
|
||||
func (r *Reconciler) Explain(model string) (placement.Result, bool) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
res, ok := r.lastExplain[model]
|
||||
return res, ok
|
||||
}
|
||||
|
||||
func addrOf(n inventory.Node) string {
|
||||
if n.MgmtAddr != "" {
|
||||
return n.MgmtAddr
|
||||
}
|
||||
if n.Hostname != "" {
|
||||
return n.Hostname
|
||||
}
|
||||
return n.ID
|
||||
}
|
||||
@@ -0,0 +1,900 @@
|
||||
package reconciler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// fakeDispatcher records commands instead of sending them.
|
||||
type fakeDispatcher struct {
|
||||
mu sync.Mutex
|
||||
started []ShardCommand
|
||||
stopped []string
|
||||
drained []string
|
||||
failNext error
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) StartShard(_ context.Context, cmd ShardCommand) error {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
if d.failNext != nil {
|
||||
err := d.failNext
|
||||
d.failNext = nil
|
||||
return err
|
||||
}
|
||||
d.started = append(d.started, cmd)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) StopShard(_ context.Context, _, shardID string) error {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
d.stopped = append(d.stopped, shardID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) DrainShard(_ context.Context, _, shardID string) error {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
d.drained = append(d.drained, shardID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) startedShards() []ShardCommand {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return slices.Clone(d.started)
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) reset() {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
d.started, d.stopped, d.drained = nil, nil, nil
|
||||
}
|
||||
|
||||
// clock lets tests move time without sleeping.
|
||||
type clock struct {
|
||||
mu sync.Mutex
|
||||
t time.Time
|
||||
}
|
||||
|
||||
func newClock() *clock {
|
||||
return &clock{t: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)}
|
||||
}
|
||||
func (c *clock) now() time.Time {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.t
|
||||
}
|
||||
func (c *clock) advance(d time.Duration) {
|
||||
c.mu.Lock()
|
||||
c.t = c.t.Add(d)
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func gpuNode(id string, n int, fabrics ...string) inventory.Node {
|
||||
nd := inventory.Node{ID: id, MgmtAddr: "10.0.0." + id, State: inventory.Ready}
|
||||
for i := range n {
|
||||
nd.GPUs = append(nd.GPUs, discovery.GPU{
|
||||
Index: i, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559, ComputeCapability: "9.0",
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
nd.Fabrics = append(nd.Fabrics, inventory.FabricAttachment{
|
||||
Fabric: f, HCAs: []string{"mlx5_0", "mlx5_1"},
|
||||
})
|
||||
}
|
||||
return nd
|
||||
}
|
||||
|
||||
func testModel(min int, profiles ...catalog.Profile) catalog.Model {
|
||||
if len(profiles) == 0 {
|
||||
profiles = []catalog.Profile{{Name: "tp4", GPUClass: "h100-80g", TP: 4}}
|
||||
}
|
||||
return catalog.Model{
|
||||
Name: "llama", Source: "/models/llama",
|
||||
Scaling: catalog.Scaling{Min: min, Max: min + 2}, Profiles: profiles,
|
||||
}
|
||||
}
|
||||
|
||||
type harness struct {
|
||||
r *Reconciler
|
||||
inv *inventory.Store
|
||||
disp *fakeDispatcher
|
||||
clock *clock
|
||||
}
|
||||
|
||||
func newHarness(t *testing.T, m catalog.Model, nodes ...inventory.Node) *harness {
|
||||
t.Helper()
|
||||
|
||||
cat, err := catalog.New([]catalog.Model{m})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inv := inventory.NewStore()
|
||||
c := newClock()
|
||||
for _, n := range nodes {
|
||||
inv.Upsert(n)
|
||||
inv.Touch(n.ID)
|
||||
}
|
||||
|
||||
disp := &fakeDispatcher{}
|
||||
cfg := DefaultConfig()
|
||||
cfg.UnreachableAfter = 15 * time.Second
|
||||
cfg.LostAfter = 60 * time.Second
|
||||
|
||||
r := New(cfg, cat, inv, disp, slog.New(slog.DiscardHandler))
|
||||
r.SetClock(c.now)
|
||||
// Nodes were just registered, so their LastSeen is wall-clock now, which is
|
||||
// far from the fake clock's epoch. Re-stamp them onto the fake timeline.
|
||||
for _, n := range nodes {
|
||||
touchAt(inv, n.ID, c.now())
|
||||
}
|
||||
return &harness{r: r, inv: inv, disp: disp, clock: c}
|
||||
}
|
||||
|
||||
// touchAt rewrites a node's LastSeen onto the fake clock's timeline.
|
||||
func touchAt(inv *inventory.Store, id string, at time.Time) {
|
||||
n, ok := inv.Get(id)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
n.LastSeen = at
|
||||
n.State = inventory.Ready
|
||||
inv.Upsert(n)
|
||||
// Upsert stamps LastSeen with the wall clock, so set it again directly.
|
||||
inv.SetLastSeen(id, at)
|
||||
}
|
||||
|
||||
// advance moves the clock and keeps every node heartbeating.
|
||||
//
|
||||
// Without this, advancing time past UnreachableAfter makes every node look
|
||||
// silent and placement stops — which is correct behaviour, but not what the
|
||||
// backoff tests are trying to exercise.
|
||||
func (h *harness) advance(d time.Duration) {
|
||||
h.clock.advance(d)
|
||||
for _, n := range h.inv.Nodes() {
|
||||
touchAt(h.inv, n.ID, h.clock.now())
|
||||
}
|
||||
}
|
||||
|
||||
// markReady walks every started shard to Ready.
|
||||
func (h *harness) markReady() {
|
||||
for _, cmd := range h.disp.startedShards() {
|
||||
h.r.OnShardStatus(cmd.ReplicaID, cmd.ShardID, supervisor.Ready, "")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlacesMinimumReplicas(t *testing.T) {
|
||||
h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 2 {
|
||||
t.Fatalf("expected 2 replicas, got %d: %+v", len(reps), reps)
|
||||
}
|
||||
for _, rep := range reps {
|
||||
if rep.State != Ready {
|
||||
t.Errorf("replica %s is %s", rep.ID, rep.State)
|
||||
}
|
||||
if rep.Endpoint == "" {
|
||||
t.Errorf("replica %s has no endpoint", rep.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneReplicaPerPass(t *testing.T) {
|
||||
// Placement depends on what is already allocated, so starting several from
|
||||
// one inventory snapshot would hand them the same GPUs.
|
||||
h := newHarness(t, testModel(3), gpuNode("1", 8), gpuNode("2", 8), gpuNode("3", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Fatalf("one pass created %d replicas", got)
|
||||
}
|
||||
|
||||
seen := map[int]string{}
|
||||
for range 3 {
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
}
|
||||
for _, rep := range h.r.Replicas() {
|
||||
for _, s := range rep.Shards {
|
||||
for _, g := range s.GPUIndices {
|
||||
key := g
|
||||
if prev, dup := seen[key]; dup && prev != s.NodeID {
|
||||
continue // same index on a different node is fine
|
||||
}
|
||||
seen[key] = s.NodeID
|
||||
}
|
||||
}
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 3 {
|
||||
t.Errorf("expected 3 replicas, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpreadsAcrossNodes(t *testing.T) {
|
||||
h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
for range 2 {
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
}
|
||||
|
||||
nodes := map[string]struct{}{}
|
||||
for _, rep := range h.r.Replicas() {
|
||||
for _, n := range rep.NodeIDs() {
|
||||
nodes[n] = struct{}{}
|
||||
}
|
||||
}
|
||||
if len(nodes) != 2 {
|
||||
t.Errorf("both replicas landed on %d node(s); they should spread", len(nodes))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPacksWhenSpreadingIsImpossible(t *testing.T) {
|
||||
// Spreading is a preference, not a rule: a second replica on one node still
|
||||
// survives an engine crash, which is the common failure.
|
||||
h := newHarness(t, testModel(2), gpuNode("only", 8))
|
||||
|
||||
for range 2 {
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 2 {
|
||||
t.Fatalf("expected 2 replicas on the single node, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShardOrderLeaderFirst(t *testing.T) {
|
||||
// Both engines put the rendezvous point on rank 0 — SGLang's dist-init
|
||||
// address and vLLM's Ray head — so a worker started first has nothing to
|
||||
// connect to and spends its startup budget retrying.
|
||||
//
|
||||
// This test previously asserted the opposite, on the theory that the head
|
||||
// should wait for its workers. That is backwards: vLLM's server blocks
|
||||
// until the Ray cluster has the GPUs it needs, so the leader can start
|
||||
// immediately and waits on its own.
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
started := h.disp.startedShards()
|
||||
if len(started) != 2 {
|
||||
t.Fatalf("expected 2 shards, got %d", len(started))
|
||||
}
|
||||
if started[0].Spec.Rank != 0 || started[1].Spec.Rank != 1 {
|
||||
t.Errorf("start order was rank %d then %d; the leader hosts the rendezvous "+
|
||||
"point and must start first",
|
||||
started[0].Spec.Rank, started[1].Spec.Rank)
|
||||
}
|
||||
// Every rank must be given the same rendezvous point.
|
||||
if started[0].Spec.LeaderAddr != started[1].Spec.LeaderAddr {
|
||||
t.Errorf("ranks got different leader addresses: %q and %q",
|
||||
started[0].Spec.LeaderAddr, started[1].Spec.LeaderAddr)
|
||||
}
|
||||
if started[0].Spec.LeaderPort != started[1].Spec.LeaderPort {
|
||||
t.Error("ranks got different leader ports")
|
||||
}
|
||||
for _, cmd := range started {
|
||||
if len(cmd.Spec.HCAs) == 0 {
|
||||
t.Errorf("rank %d got no HCAs; NCCL would choose its own", cmd.Spec.Rank)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailedShardFailsWholeReplica(t *testing.T) {
|
||||
// One dead shard makes the others useless: they wait on collectives that
|
||||
// will never complete while holding GPUs (plan §5.4).
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
started := h.disp.startedShards()
|
||||
|
||||
// One shard ready, the other dead.
|
||||
h.r.OnShardStatus(started[0].ReplicaID, started[0].ShardID, supervisor.Ready, "")
|
||||
h.r.OnShardStatus(started[1].ReplicaID, started[1].ShardID, supervisor.Failed, "CUDA out of memory")
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
// Both shards must have been stopped, not just the failed one.
|
||||
h.disp.mu.Lock()
|
||||
stopped := slices.Clone(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if len(stopped) != 2 {
|
||||
t.Errorf("stopped %d shards, expected both: %v", len(stopped), stopped)
|
||||
}
|
||||
// And the GPUs must be free again.
|
||||
for _, id := range []string{"1", "2"} {
|
||||
if got := len(h.inv.FreeGPUs(id)); got != 8 {
|
||||
t.Errorf("node %s has %d free GPUs; the replica should have been released", id, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackoffAfterFailure(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
started := h.disp.startedShards()
|
||||
h.r.OnShardStatus(started[0].ReplicaID, started[0].ShardID, supervisor.Failed, "boom")
|
||||
h.r.Reconcile(t.Context()) // reaps it
|
||||
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.disp.startedShards()); got != 0 {
|
||||
t.Errorf("a replica was rebuilt immediately after failure; backoff is not applied")
|
||||
}
|
||||
|
||||
// Past the backoff, it tries again.
|
||||
h.advance(DefaultConfig().MaxBackoff)
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.disp.startedShards()); got == 0 {
|
||||
t.Error("the replica was never rebuilt after the backoff expired")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSuccessClearsBackoff(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
// Fail once.
|
||||
h.r.Reconcile(t.Context())
|
||||
s := h.disp.startedShards()[0]
|
||||
h.r.OnShardStatus(s.ReplicaID, s.ShardID, supervisor.Failed, "boom")
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
// Then succeed.
|
||||
h.advance(DefaultConfig().MaxBackoff)
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
if len(h.disp.startedShards()) == 0 {
|
||||
t.Fatal("no replica was started after the backoff expired")
|
||||
}
|
||||
s2 := h.disp.startedShards()[0]
|
||||
h.r.OnShardStatus(s2.ReplicaID, s2.ShardID, supervisor.Ready, "")
|
||||
|
||||
// A later failure should start from a short backoff, not a long streak.
|
||||
h.r.OnShardStatus(s2.ReplicaID, s2.ShardID, supervisor.Failed, "again")
|
||||
h.r.Reconcile(t.Context())
|
||||
h.advance(2 * DefaultConfig().BaseBackoff)
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.disp.startedShards()); got == 0 {
|
||||
t.Error("backoff was not reset by the successful run in between")
|
||||
}
|
||||
}
|
||||
|
||||
// TestLostNodeMovesReplicas is the phase 3 exit criterion: a node going away
|
||||
// must get its replicas rebuilt elsewhere.
|
||||
func TestLostNodeMovesReplicas(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || reps[0].State != Ready {
|
||||
t.Fatalf("setup failed: %+v", reps)
|
||||
}
|
||||
original := reps[0].NodeIDs()[0]
|
||||
survivor := "2"
|
||||
if original == "2" {
|
||||
survivor = "1"
|
||||
}
|
||||
|
||||
// The surviving node keeps reporting; the other goes silent.
|
||||
h.clock.advance(20 * time.Second)
|
||||
touchAt(h.inv, survivor, h.clock.now())
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if n, _ := h.inv.Get(original); n.State != inventory.Unreachable {
|
||||
t.Errorf("silent node is %s, expected unreachable", n.State)
|
||||
}
|
||||
// Unreachable is not lost: a blip must not tear down a working replica.
|
||||
if reps := h.r.Replicas(); len(reps) != 1 || reps[0].State != Ready {
|
||||
t.Fatalf("an unreachable node already destroyed the replica: %+v", reps)
|
||||
}
|
||||
|
||||
h.clock.advance(90 * time.Second)
|
||||
touchAt(h.inv, survivor, h.clock.now())
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context()) // declares it lost and reaps the replica
|
||||
|
||||
if n, _ := h.inv.Get(original); n.State != inventory.Lost {
|
||||
t.Fatalf("node is %s, expected lost", n.State)
|
||||
}
|
||||
|
||||
// Past the backoff, the replica is rebuilt on the survivor. Only the
|
||||
// survivor heartbeats; the lost node stays gone.
|
||||
h.clock.advance(DefaultConfig().MaxBackoff)
|
||||
touchAt(h.inv, survivor, h.clock.now())
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
reps = h.r.Replicas()
|
||||
if len(reps) != 1 {
|
||||
t.Fatalf("expected the replica to be rebuilt, got %d", len(reps))
|
||||
}
|
||||
if got := reps[0].NodeIDs()[0]; got != survivor {
|
||||
t.Errorf("replica was rebuilt on %q, expected the surviving node %q", got, survivor)
|
||||
}
|
||||
if reps[0].State != Ready {
|
||||
t.Errorf("rebuilt replica is %s", reps[0].State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeartbeatRevivesNode(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.clock.advance(20 * time.Second) // deliberately without a heartbeat
|
||||
h.r.Reconcile(t.Context())
|
||||
if n, _ := h.inv.Get("1"); n.State != inventory.Unreachable {
|
||||
t.Fatalf("node is %s", n.State)
|
||||
}
|
||||
|
||||
h.inv.Touch("1")
|
||||
if n, _ := h.inv.Get("1"); n.State != inventory.Ready {
|
||||
t.Errorf("a heartbeat did not revive the node: %s", n.State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScaleDownDrainsFirst(t *testing.T) {
|
||||
// Draining before stopping lets in-flight requests finish instead of being
|
||||
// cut off mid-stream.
|
||||
h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
|
||||
for range 2 {
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 2 {
|
||||
t.Fatalf("setup: %d replicas", got)
|
||||
}
|
||||
|
||||
// Shrink the catalog to one replica.
|
||||
cat, err := catalog.New([]catalog.Model{testModel(1)})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.r.cat = cat
|
||||
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.disp.mu.Lock()
|
||||
drained, stopped := len(h.disp.drained), len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if drained == 0 {
|
||||
t.Error("the excess replica was not drained")
|
||||
}
|
||||
if stopped != 0 {
|
||||
t.Error("a replica was stopped before it had drained")
|
||||
}
|
||||
|
||||
// The next pass stops what has drained.
|
||||
h.r.Reconcile(t.Context())
|
||||
h.disp.mu.Lock()
|
||||
stopped = len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if stopped == 0 {
|
||||
t.Error("the drained replica was never stopped")
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Errorf("expected 1 replica after scale-down, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnplaceableModelExplains(t *testing.T) {
|
||||
// A model that cannot be placed must leave a usable explanation rather than
|
||||
// a silent absence.
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "tp8", GPUClass: "h100-80g", TP: 8,
|
||||
}), gpuNode("1", 4))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.r.Replicas()); got != 0 {
|
||||
t.Fatalf("something was placed: %d", got)
|
||||
}
|
||||
|
||||
res, ok := h.r.Explain("llama")
|
||||
if !ok {
|
||||
t.Fatal("no explanation was recorded")
|
||||
}
|
||||
exp := res.Explain()
|
||||
if !strings.Contains(exp, "needs 8") {
|
||||
t.Errorf("explanation does not say what was missing:\n%s", exp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatcherFailureFailsReplica(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
h.disp.failNext = fmt.Errorf("agent unreachable")
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || reps[0].State != Failed {
|
||||
t.Fatalf("a failed dispatch did not fail the replica: %+v", reps)
|
||||
}
|
||||
|
||||
// And the next pass must release its GPUs.
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.inv.FreeGPUs("1")); got != 8 {
|
||||
t.Errorf("node has %d free GPUs; the failed replica still holds some", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadyReplicasFeedTheRouter(t *testing.T) {
|
||||
h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.r.ReadyReplicas()); got != 0 {
|
||||
t.Errorf("a starting replica was offered to the router")
|
||||
}
|
||||
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
|
||||
ready := h.r.ReadyReplicas()
|
||||
if len(ready) != 2 {
|
||||
t.Fatalf("expected 2 ready replicas, got %d", len(ready))
|
||||
}
|
||||
for _, rep := range ready {
|
||||
if !strings.HasPrefix(rep.Endpoint, "http://") {
|
||||
t.Errorf("endpoint %q is not a URL", rep.Endpoint)
|
||||
}
|
||||
if rep.Weight <= 0 {
|
||||
t.Errorf("replica %s has no weight; the router needs one", rep.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPortsDoNotCollide(t *testing.T) {
|
||||
h := newHarness(t, testModel(3), gpuNode("1", 8), gpuNode("2", 8), gpuNode("3", 8))
|
||||
for range 3 {
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
}
|
||||
|
||||
ports := map[int]string{}
|
||||
for _, rep := range h.r.Replicas() {
|
||||
for _, s := range rep.Shards {
|
||||
if prev, dup := ports[s.HTTPPort]; dup && prev != rep.ID {
|
||||
t.Errorf("replicas %s and %s both use port %d", prev, rep.ID, s.HTTPPort)
|
||||
}
|
||||
ports[s.HTTPPort] = rep.ID
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReplicaStateStrings(t *testing.T) {
|
||||
for st, want := range map[ReplicaState]string{
|
||||
Placing: "placing", Starting: "starting", Ready: "ready",
|
||||
Draining: "draining", Stopped: "stopped", Failed: "failed",
|
||||
} {
|
||||
if got := st.String(); got != want {
|
||||
t.Errorf("%d -> %q", int(st), got)
|
||||
}
|
||||
}
|
||||
if !Ready.Counts() || Draining.Counts() || Failed.Counts() {
|
||||
// Draining must not count, or a replacement is never started for it.
|
||||
t.Error("Counts is wrong")
|
||||
}
|
||||
}
|
||||
|
||||
// --- adoption after a control plane restart (plan §6.5) ---
|
||||
|
||||
func adopted(replicaID, shardID, nodeID string, rank int, gpus []int) AdoptedShard {
|
||||
a := AdoptedShard{
|
||||
ReplicaID: replicaID, ShardID: shardID, NodeID: nodeID, Rank: rank,
|
||||
GPUIndices: gpus, HTTPPort: 20000, State: supervisor.Ready,
|
||||
Model: "llama", Profile: "tp4", Pool: "none/h100-80g",
|
||||
}
|
||||
if rank == 0 {
|
||||
a.Endpoint = "http://10.0.0." + nodeID + ":20000"
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
func TestAdoptRebuildsReplica(t *testing.T) {
|
||||
// The control plane restarted. Nothing is known, but a node reports a live
|
||||
// replica. Without this the reconciler would place a duplicate while the
|
||||
// original keeps holding GPUs.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
unclaimed := h.r.Adopt([]AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})})
|
||||
if len(unclaimed) != 0 {
|
||||
t.Fatalf("a valid shard was reported unclaimed: %+v", unclaimed)
|
||||
}
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 {
|
||||
t.Fatalf("expected 1 adopted replica, got %d", len(reps))
|
||||
}
|
||||
rep := reps[0]
|
||||
if rep.ID != "llama-99" || rep.Model != "llama" || rep.Profile != "tp4" {
|
||||
t.Errorf("identity wrong: %+v", rep)
|
||||
}
|
||||
if rep.State != Ready {
|
||||
t.Errorf("state=%s; every shard reported Ready", rep.State)
|
||||
}
|
||||
if rep.Endpoint == "" {
|
||||
t.Error("no endpoint; the router would have nowhere to send")
|
||||
}
|
||||
if rep.Weight <= 0 {
|
||||
t.Error("weight not taken from the catalog")
|
||||
}
|
||||
|
||||
// The GPUs must be reclaimed, or the next placement hands them out twice.
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 4 {
|
||||
t.Errorf("node has %d free GPUs, expected 4 after adoption", free)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptPreventsDuplicatePlacement(t *testing.T) {
|
||||
// The whole point: after adopting, a reconciliation pass must see the
|
||||
// desired count already met.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.r.Adopt([]AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})})
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.disp.startedShards()); got != 0 {
|
||||
t.Errorf("a duplicate replica was placed despite adoption (%d shards started)", got)
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Errorf("expected 1 replica, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptMultiNodeReplica(t *testing.T) {
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
shards := []AdoptedShard{
|
||||
adopted("llama-7", "llama-7-r0", "1", 0, []int{0, 1, 2, 3, 4, 5, 6, 7}),
|
||||
adopted("llama-7", "llama-7-r1", "2", 1, []int{0, 1, 2, 3, 4, 5, 6, 7}),
|
||||
}
|
||||
for i := range shards {
|
||||
shards[i].Profile = "2node"
|
||||
shards[i].Fabric = "ib-a"
|
||||
shards[i].Pool = "ib-a/h100-80g"
|
||||
}
|
||||
|
||||
if unclaimed := h.r.Adopt(shards); len(unclaimed) != 0 {
|
||||
t.Fatalf("unclaimed: %+v", unclaimed)
|
||||
}
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || len(reps[0].Shards) != 2 {
|
||||
t.Fatalf("the two shards were not grouped into one replica: %+v", reps)
|
||||
}
|
||||
if reps[0].Fabric != "ib-a" {
|
||||
t.Errorf("fabric lost: %q", reps[0].Fabric)
|
||||
}
|
||||
// Both nodes' GPUs must be reclaimed.
|
||||
for _, id := range []string{"1", "2"} {
|
||||
if free := len(h.inv.FreeGPUs(id)); free != 0 {
|
||||
t.Errorf("node %s has %d free GPUs after adoption", id, free)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptPartialReplicaFails(t *testing.T) {
|
||||
// A multi-node replica missing a shard is useless: the survivors wait on
|
||||
// collectives that will never complete while holding GPUs.
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
s := adopted("llama-7", "llama-7-r0", "1", 0, []int{0, 1, 2, 3, 4, 5, 6, 7})
|
||||
s.Profile = "2node"
|
||||
s.State = supervisor.Failed
|
||||
h.r.Adopt([]AdoptedShard{s})
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || reps[0].State != Failed {
|
||||
t.Fatalf("a failed adopted shard did not fail its replica: %+v", reps)
|
||||
}
|
||||
|
||||
// And the next pass releases its GPUs.
|
||||
h.r.Reconcile(t.Context())
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 8 {
|
||||
t.Errorf("node has %d free GPUs; the failed adopted replica still holds some", free)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptUnknownModelIsUnclaimed(t *testing.T) {
|
||||
// The model left the catalog while the engine kept running. Nothing will
|
||||
// ever route to it, and it is holding GPUs.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
s := adopted("ghost-1", "ghost-1-r0", "1", 0, []int{0})
|
||||
s.Model = "a-model-that-was-removed"
|
||||
|
||||
unclaimed := h.r.Adopt([]AdoptedShard{s})
|
||||
if len(unclaimed) != 1 || unclaimed[0].ShardID != "ghost-1-r0" {
|
||||
t.Fatalf("an orphaned shard was not reported unclaimed: %+v", unclaimed)
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 0 {
|
||||
t.Errorf("a replica was created for an unknown model: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptUnknownProfileIsUnclaimed(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
s := adopted("llama-9", "llama-9-r0", "1", 0, []int{0})
|
||||
s.Profile = "a-profile-that-was-removed"
|
||||
|
||||
if unclaimed := h.r.Adopt([]AdoptedShard{s}); len(unclaimed) != 1 {
|
||||
t.Errorf("a shard with an unknown profile was claimed: %+v", unclaimed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptIsIdempotent(t *testing.T) {
|
||||
// An agent reconnecting re-reports the same shards; that must not double
|
||||
// the allocation.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
s := []AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})}
|
||||
|
||||
h.r.Adopt(s)
|
||||
h.r.Adopt(s)
|
||||
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Errorf("re-reporting created %d replicas", got)
|
||||
}
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 4 {
|
||||
t.Errorf("node has %d free GPUs; the allocation was applied twice", free)
|
||||
}
|
||||
}
|
||||
|
||||
// --- draining a node (ctl drain) ---
|
||||
|
||||
func TestDrainingNodeMovesWorkOff(t *testing.T) {
|
||||
// The operator's view: a node is about to be worked on, its replicas must
|
||||
// end up elsewhere, and nothing may be cut off mid-request on the way.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 {
|
||||
t.Fatalf("setup: %d replicas", len(reps))
|
||||
}
|
||||
victim := reps[0].NodeIDs()[0]
|
||||
survivor := "2"
|
||||
if victim == "2" {
|
||||
survivor = "1"
|
||||
}
|
||||
|
||||
h.inv.SetState(victim, inventory.Draining)
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
// It is drained, not stopped: the engine is healthy and still serving.
|
||||
h.disp.mu.Lock()
|
||||
drained, stopped := len(h.disp.drained), len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if drained == 0 {
|
||||
t.Error("the replica on the draining node was not drained")
|
||||
}
|
||||
if stopped != 0 {
|
||||
t.Error("the replica was stopped before a replacement existed; in-flight requests would be cut")
|
||||
}
|
||||
|
||||
// A draining replica does not count, so a replacement is placed elsewhere.
|
||||
if got := len(h.disp.startedShards()); got == 0 {
|
||||
t.Fatal("no replacement was placed")
|
||||
}
|
||||
for _, cmd := range h.disp.startedShards() {
|
||||
if cmd.NodeID == victim {
|
||||
t.Errorf("the replacement was placed back on the draining node %s", victim)
|
||||
}
|
||||
}
|
||||
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
// Now that the replacement serves, the original can go.
|
||||
h.disp.mu.Lock()
|
||||
stopped = len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if stopped == 0 {
|
||||
t.Error("the drained replica was never stopped once its replacement was ready")
|
||||
}
|
||||
|
||||
for _, rep := range h.r.Replicas() {
|
||||
if rep.OnNode(victim) {
|
||||
t.Errorf("replica %s is still on the drained node", rep.ID)
|
||||
}
|
||||
if !slices.Contains(rep.NodeIDs(), survivor) {
|
||||
t.Errorf("replica %s did not move to the surviving node", rep.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDrainGraceStopsWithoutReplacement(t *testing.T) {
|
||||
// Draining the only node that can host a model would otherwise never
|
||||
// finish, leaving the node undrainable.
|
||||
h := newHarness(t, testModel(1), gpuNode("only", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.inv.SetState("only", inventory.Draining)
|
||||
h.r.Reconcile(t.Context()) // drains it; nowhere to place a replacement
|
||||
|
||||
if got := h.r.Replicas(); len(got) != 1 || got[0].State != Draining {
|
||||
t.Fatalf("expected a draining replica, got %+v", got)
|
||||
}
|
||||
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
h.disp.mu.Lock()
|
||||
early := len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if early != 0 {
|
||||
t.Error("the replica was stopped before the grace period expired")
|
||||
}
|
||||
|
||||
h.advance(DefaultConfig().DrainGrace + time.Second)
|
||||
h.inv.SetState("only", inventory.Draining) // advance() resets node state
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.disp.mu.Lock()
|
||||
stopped := len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if stopped == 0 {
|
||||
t.Error("the drained replica was never stopped; the node cannot be drained")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCordonedNodeKeepsItsWork(t *testing.T) {
|
||||
// Cordon is not drain: it stops new placement and leaves what is running.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.inv.SetState("1", inventory.Cordoned)
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.disp.mu.Lock()
|
||||
drained, stopped := len(h.disp.drained), len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if drained != 0 || stopped != 0 {
|
||||
t.Errorf("cordoning disturbed a running replica (drained=%d stopped=%d)", drained, stopped)
|
||||
}
|
||||
if got := h.r.Replicas(); len(got) != 1 || got[0].State != Ready {
|
||||
t.Errorf("the replica did not survive a cordon: %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
// Package reconciler drives the cluster toward the catalog's desired state
|
||||
// (plan §6.3).
|
||||
package reconciler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
)
|
||||
|
||||
// ReplicaState is where a replica stands.
|
||||
type ReplicaState int
|
||||
|
||||
const (
|
||||
// Placing means it has been scheduled but no shard has started.
|
||||
Placing ReplicaState = iota
|
||||
// Starting means shards are launching or loading.
|
||||
Starting
|
||||
// Ready means it serves traffic.
|
||||
Ready
|
||||
// Draining means it takes no new requests and will be stopped.
|
||||
Draining
|
||||
// Stopped means it was shut down on purpose.
|
||||
Stopped
|
||||
// Failed means it died or never became ready.
|
||||
Failed
|
||||
)
|
||||
|
||||
func (s ReplicaState) String() string {
|
||||
switch s {
|
||||
case Placing:
|
||||
return "placing"
|
||||
case Starting:
|
||||
return "starting"
|
||||
case Ready:
|
||||
return "ready"
|
||||
case Draining:
|
||||
return "draining"
|
||||
case Stopped:
|
||||
return "stopped"
|
||||
case Failed:
|
||||
return "failed"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// Terminal reports whether the replica has finished for good.
|
||||
func (s ReplicaState) Terminal() bool { return s == Stopped || s == Failed }
|
||||
|
||||
// Counts reports whether the replica should count toward the desired total.
|
||||
// Draining replicas do not: they are on the way out, and counting them would
|
||||
// stop a replacement from being started.
|
||||
func (s ReplicaState) Counts() bool { return s == Placing || s == Starting || s == Ready }
|
||||
|
||||
// Shard is one node's part of a replica.
|
||||
type Shard struct {
|
||||
ID string
|
||||
NodeID string
|
||||
Rank int
|
||||
GPUIndices []int
|
||||
HCAs []string
|
||||
HTTPPort int
|
||||
State supervisor.State
|
||||
Err string
|
||||
}
|
||||
|
||||
// Leader reports whether this shard serves HTTP for the replica.
|
||||
func (s Shard) Leader() bool { return s.Rank == 0 }
|
||||
|
||||
// Replica is one running copy of a model. It is the unit the router sees.
|
||||
type Replica struct {
|
||||
ID string
|
||||
Model string
|
||||
Profile string
|
||||
// Fabric is set only for multi-node replicas.
|
||||
Fabric string
|
||||
Pool string
|
||||
Shards []Shard
|
||||
State ReplicaState
|
||||
// Endpoint is rank 0's base URL.
|
||||
Endpoint string
|
||||
// Weight is the profile's capacity weight, passed to the router.
|
||||
Weight float64
|
||||
|
||||
CreatedAt time.Time
|
||||
ReadyAt time.Time
|
||||
// DrainingAt is when the replica stopped taking new work, so the wait for a
|
||||
// replacement can be bounded.
|
||||
DrainingAt time.Time
|
||||
// Attempt counts how many times this model has failed to come up in a row,
|
||||
// carried across rebuilds so the backoff does not reset on every attempt.
|
||||
Attempt int
|
||||
Err string
|
||||
}
|
||||
|
||||
// NodeIDs lists the nodes this replica occupies.
|
||||
func (r Replica) NodeIDs() []string {
|
||||
out := make([]string, 0, len(r.Shards))
|
||||
for _, s := range r.Shards {
|
||||
out = append(out, s.NodeID)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// OnNode reports whether any shard runs on a node.
|
||||
func (r Replica) OnNode(id string) bool {
|
||||
return slices.ContainsFunc(r.Shards, func(s Shard) bool { return s.NodeID == id })
|
||||
}
|
||||
|
||||
// deriveState recomputes the replica's state from its shards.
|
||||
//
|
||||
// A replica is only as healthy as its worst shard: one failed shard makes the
|
||||
// whole thing useless, because the others are waiting on collectives that will
|
||||
// never complete. That is why there is no path here that keeps a replica
|
||||
// partially alive (plan §5.4).
|
||||
func (r *Replica) deriveState() {
|
||||
if r.State == Draining || r.State == Stopped {
|
||||
return
|
||||
}
|
||||
if len(r.Shards) == 0 {
|
||||
r.State = Placing
|
||||
return
|
||||
}
|
||||
|
||||
ready := 0
|
||||
for _, s := range r.Shards {
|
||||
switch s.State {
|
||||
case supervisor.Failed:
|
||||
r.State = Failed
|
||||
if r.Err == "" {
|
||||
r.Err = fmt.Sprintf("shard %s failed: %s", s.ID, s.Err)
|
||||
}
|
||||
return
|
||||
case supervisor.Stopped:
|
||||
// A shard stopping on its own, while the replica was not draining,
|
||||
// is a failure of the replica.
|
||||
r.State = Failed
|
||||
if r.Err == "" {
|
||||
r.Err = fmt.Sprintf("shard %s stopped unexpectedly", s.ID)
|
||||
}
|
||||
return
|
||||
case supervisor.Ready:
|
||||
ready++
|
||||
}
|
||||
}
|
||||
|
||||
if ready == len(r.Shards) {
|
||||
if r.State != Ready {
|
||||
r.ReadyAt = time.Now()
|
||||
}
|
||||
r.State = Ready
|
||||
return
|
||||
}
|
||||
r.State = Starting
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
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). The control plane's management API
|
||||
// is mounted here too, by the caller.
|
||||
func (r *Router) AdminHandler() (*http.ServeMux, 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
|
||||
|
||||
queueDepth *prometheus.Desc
|
||||
saturation *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),
|
||||
|
||||
queueDepth: prometheus.NewDesc("llmgw_admission_queue_depth", "Requests waiting in each admission band.", []string{"band"}, nil),
|
||||
saturation: prometheus.NewDesc("llmgw_pool_saturation", "Pool saturation per model; 1 is the healthy limit.", []string{"model"}, 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
|
||||
ch <- c.queueDepth
|
||||
ch <- c.saturation
|
||||
}
|
||||
|
||||
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"))
|
||||
}
|
||||
|
||||
if c.r.adm != nil {
|
||||
for _, d := range c.r.adm.Depths() {
|
||||
ch <- prometheus.MustNewConstMetric(c.queueDepth, prometheus.GaugeValue, float64(d.Depth), d.Band)
|
||||
}
|
||||
}
|
||||
// Saturation is reported per model, which is the granularity admission
|
||||
// decisions are made at.
|
||||
for _, model := range c.r.servedModels() {
|
||||
if sat, ok := c.r.saturation(model, 0); ok {
|
||||
ch <- prometheus.MustNewConstMetric(c.saturation, prometheus.GaugeValue, sat, model)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func boolToFloat(b bool) float64 {
|
||||
if b {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,436 @@
|
||||
package admission
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
)
|
||||
|
||||
func testConfig() Config {
|
||||
c := DefaultConfig()
|
||||
c.Interval = time.Millisecond
|
||||
c.QueueTTL = 2 * time.Second
|
||||
return c
|
||||
}
|
||||
|
||||
// fixedSat lets a test drive saturation directly.
|
||||
type fixedSat struct{ v atomic.Uint64 }
|
||||
|
||||
func (f *fixedSat) set(v float64) { f.v.Store(uint64(v * 1000)) }
|
||||
func (f *fixedSat) fn(string, int) (float64, bool) {
|
||||
return float64(f.v.Load()) / 1000, true
|
||||
}
|
||||
|
||||
func newController(t *testing.T, cfg Config, sat SaturationFunc) *Controller {
|
||||
t.Helper()
|
||||
c, err := New(cfg, sat)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c.Start()
|
||||
t.Cleanup(c.Close)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestAdmitsImmediatelyWhenIdle(t *testing.T) {
|
||||
s := &fixedSat{}
|
||||
c := newController(t, testConfig(), s.fn)
|
||||
|
||||
start := time.Now()
|
||||
lease, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if d := time.Since(start); d > 50*time.Millisecond {
|
||||
t.Errorf("idle admission took %v; the fast path is not being used", d)
|
||||
}
|
||||
lease.Release()
|
||||
}
|
||||
|
||||
func TestLowBandHeldWhileHighBandFlows(t *testing.T) {
|
||||
// The core of plan §7.2.2: the cluster degrades from the bottom up.
|
||||
s := &fixedSat{}
|
||||
s.set(0.5) // above bronze (0.30), below gold (0.85)
|
||||
c := newController(t, testConfig(), s.fn)
|
||||
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 300*time.Millisecond)
|
||||
defer cancel()
|
||||
if _, err := c.Acquire(ctx, Request{Model: "m", Band: "bronze"}); err == nil {
|
||||
t.Error("bronze was admitted above its ceiling")
|
||||
}
|
||||
|
||||
lease, err := c.Acquire(t.Context(), Request{Model: "m", Band: "gold"})
|
||||
if err != nil {
|
||||
t.Fatalf("gold was held below its ceiling: %v", err)
|
||||
}
|
||||
lease.Release()
|
||||
}
|
||||
|
||||
func TestTopBandFlowsUpToTheHealthyLimit(t *testing.T) {
|
||||
// A ceiling of 1.0 means "hold only once the pool is past its healthy
|
||||
// limit", not "never hold". At exactly the limit the top band still flows.
|
||||
s := &fixedSat{}
|
||||
s.set(1.0)
|
||||
c := newController(t, testConfig(), s.fn)
|
||||
|
||||
lease, err := c.Acquire(t.Context(), Request{Model: "m", Band: "platinum"})
|
||||
if err != nil {
|
||||
t.Fatalf("top band held at exactly its ceiling: %v", err)
|
||||
}
|
||||
lease.Release()
|
||||
}
|
||||
|
||||
func TestTopBandIsHeldWhenOverloaded(t *testing.T) {
|
||||
// Holding even the top band is deliberate: with nobody above it to protect,
|
||||
// the benefit is a bounded engine queue. Dumping a flood straight into a
|
||||
// FIFO engine queue makes every request in that band slow, whereas queueing
|
||||
// at the gateway keeps ordering, TTL and tenant fairness.
|
||||
s := &fixedSat{}
|
||||
s.set(5)
|
||||
cfg := testConfig()
|
||||
cfg.QueueTTL = 200 * time.Millisecond
|
||||
c := newController(t, cfg, s.fn)
|
||||
|
||||
if _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "platinum"}); !errors.Is(err, ErrQueueTimeout) {
|
||||
t.Fatalf("expected the top band to queue while overloaded, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCeilingAboveOneNeverHolds(t *testing.T) {
|
||||
// The escape hatch for a band that must never wait, whatever the load.
|
||||
s := &fixedSat{}
|
||||
s.set(50)
|
||||
cfg := testConfig()
|
||||
cfg.Bands = []Band{{Name: "critical", Priority: 900, SaturationCeiling: 1000}}
|
||||
cfg.DefaultBand = "critical"
|
||||
c := newController(t, cfg, s.fn)
|
||||
|
||||
lease, err := c.Acquire(t.Context(), Request{Model: "m", Band: "critical"})
|
||||
if err != nil {
|
||||
t.Fatalf("a band with an effectively unbounded ceiling was held: %v", err)
|
||||
}
|
||||
lease.Release()
|
||||
}
|
||||
|
||||
func TestQueuedRequestAdmittedWhenSaturationDrops(t *testing.T) {
|
||||
s := &fixedSat{}
|
||||
s.set(0.9)
|
||||
c := newController(t, testConfig(), s.fn)
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze"})
|
||||
done <- err
|
||||
}()
|
||||
|
||||
// Give it time to actually queue, then open the gate.
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
if d := c.Depths(); d[3].Depth != 1 {
|
||||
t.Fatalf("expected 1 queued bronze request, got %+v", d)
|
||||
}
|
||||
s.set(0.1)
|
||||
|
||||
select {
|
||||
case err := <-done:
|
||||
if err != nil {
|
||||
t.Fatalf("queued request failed: %v", err)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("queued request was not admitted after saturation dropped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueueTTL(t *testing.T) {
|
||||
s := &fixedSat{}
|
||||
s.set(1.0)
|
||||
cfg := testConfig()
|
||||
cfg.QueueTTL = 150 * time.Millisecond
|
||||
c := newController(t, cfg, s.fn)
|
||||
|
||||
start := time.Now()
|
||||
_, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze"})
|
||||
if !errors.Is(err, ErrQueueTimeout) {
|
||||
t.Fatalf("expected ErrQueueTimeout, got %v", err)
|
||||
}
|
||||
if d := time.Since(start); d > time.Second {
|
||||
t.Errorf("TTL fired after %v", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueueCapacity(t *testing.T) {
|
||||
s := &fixedSat{}
|
||||
s.set(1.0)
|
||||
cfg := testConfig()
|
||||
cfg.QueueCapacityPerBand = 3
|
||||
cfg.QueueTTL = 5 * time.Second
|
||||
c := newController(t, cfg, s.fn)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
for range 3 {
|
||||
go c.Acquire(ctx, Request{Model: "m", Band: "bronze"}) //nolint:errcheck // queued on purpose
|
||||
}
|
||||
// Wait for the queue to fill.
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) && c.Depths()[3].Depth < 3 {
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
}
|
||||
|
||||
_, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze"})
|
||||
if !errors.Is(err, ErrQueueFull) {
|
||||
t.Errorf("expected ErrQueueFull, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientCancellationLeavesQueue(t *testing.T) {
|
||||
s := &fixedSat{}
|
||||
s.set(1.0)
|
||||
c := newController(t, testConfig(), s.fn)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go c.Acquire(ctx, Request{Model: "m", Band: "bronze"}) //nolint:errcheck // cancelled on purpose
|
||||
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) && c.Depths()[3].Depth == 0 {
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
}
|
||||
cancel()
|
||||
|
||||
deadline = time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if c.Depths()[3].Depth == 0 {
|
||||
return
|
||||
}
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("a cancelled request stayed in the queue, occupying capacity")
|
||||
}
|
||||
|
||||
func TestTenantFairnessWithinBand(t *testing.T) {
|
||||
// One tenant flooding the band must not take every dispatch turn.
|
||||
s := &fixedSat{}
|
||||
s.set(1.0) // hold everything initially
|
||||
cfg := testConfig()
|
||||
cfg.QueueTTL = 5 * time.Second
|
||||
c := newController(t, cfg, s.fn)
|
||||
|
||||
const flood = 40
|
||||
var wg sync.WaitGroup
|
||||
order := make(chan string, flood+4)
|
||||
|
||||
for i := range flood {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze", Tenant: "noisy"}); err == nil {
|
||||
order <- "noisy"
|
||||
}
|
||||
_ = i
|
||||
}()
|
||||
}
|
||||
// Let the flood queue up first, so the quiet tenant is strictly behind it.
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for time.Now().Before(deadline) && c.Depths()[3].Depth < flood {
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
}
|
||||
|
||||
for range 4 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := c.Acquire(t.Context(), Request{Model: "m", Band: "bronze", Tenant: "quiet"}); err == nil {
|
||||
order <- "quiet"
|
||||
}
|
||||
}()
|
||||
}
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
|
||||
s.set(0.0) // open the gate
|
||||
wg.Wait()
|
||||
close(order)
|
||||
|
||||
var seen []string
|
||||
for v := range order {
|
||||
seen = append(seen, v)
|
||||
}
|
||||
if len(seen) != flood+4 {
|
||||
t.Fatalf("expected %d admissions, got %d", flood+4, len(seen))
|
||||
}
|
||||
|
||||
// With FIFO the quiet tenant would be admitted 41st at best. Round-robin
|
||||
// should interleave it into the first handful.
|
||||
firstQuiet := -1
|
||||
for i, v := range seen {
|
||||
if v == "quiet" {
|
||||
firstQuiet = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if firstQuiet < 0 {
|
||||
t.Fatal("the quiet tenant was never admitted")
|
||||
}
|
||||
if firstQuiet > 8 {
|
||||
t.Errorf("quiet tenant admitted at position %d; a flood is monopolising the band", firstQuiet)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnknownBandUsesDefault(t *testing.T) {
|
||||
s := &fixedSat{}
|
||||
s.set(0.5)
|
||||
cfg := testConfig()
|
||||
cfg.DefaultBand = "bronze"
|
||||
c := newController(t, cfg, s.fn)
|
||||
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 200*time.Millisecond)
|
||||
defer cancel()
|
||||
// bronze's ceiling is 0.30, so an unknown band must be held here.
|
||||
if _, err := c.Acquire(ctx, Request{Model: "m", Band: "no-such-band"}); err == nil {
|
||||
t.Error("unknown band did not fall back to the default")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnknownModelIsLetThrough(t *testing.T) {
|
||||
// Better a precise 503 from dispatch than a 30-second queue timeout.
|
||||
c := newController(t, testConfig(), func(string, int) (float64, bool) { return 0, false })
|
||||
lease, err := c.Acquire(t.Context(), Request{Model: "nope", Band: "bronze"})
|
||||
if err != nil {
|
||||
t.Fatalf("unknown model was queued instead of passed through: %v", err)
|
||||
}
|
||||
lease.Release()
|
||||
}
|
||||
|
||||
func TestConfigValidation(t *testing.T) {
|
||||
if _, err := New(Config{Bands: []Band{{Name: "a"}, {Name: "a"}}}, nil); err == nil {
|
||||
t.Error("duplicate band names accepted")
|
||||
}
|
||||
if _, err := New(Config{Bands: []Band{{Name: "a"}}, DefaultBand: "b"}, nil); err == nil {
|
||||
t.Error("default_band outside the band list accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBandsSortedByPriority(t *testing.T) {
|
||||
c, err := New(Config{Bands: []Band{
|
||||
{Name: "low", Priority: 1, SaturationCeiling: 0.2},
|
||||
{Name: "high", Priority: 9, SaturationCeiling: 1},
|
||||
}, DefaultBand: "low"}, func(string, int) (float64, bool) { return 0, true })
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := c.Bands()[0].Name; got != "high" {
|
||||
t.Errorf("bands are not ordered by priority: first is %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// --- saturation ---
|
||||
|
||||
func snap(id string, waiting, running, inflight int, kv float64, open bool) endpoint.Snapshot {
|
||||
s := endpoint.Snapshot{
|
||||
Info: endpoint.Info{ID: id},
|
||||
Inflight: inflight,
|
||||
Breaker: breaker.Closed,
|
||||
}
|
||||
s.Stats.Waiting = waiting
|
||||
s.Stats.Running = running
|
||||
s.Stats.KVCacheUsage = kv
|
||||
if open {
|
||||
s.Breaker = breaker.Open
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func TestSaturationCombinesQueueAndKV(t *testing.T) {
|
||||
th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9}
|
||||
|
||||
// Empty queue but a nearly full cache still counts as saturated.
|
||||
if got := Saturation([]endpoint.Snapshot{snap("a", 0, 1, 1, 0.89, false)}, 0, th); got < 0.9 {
|
||||
t.Errorf("KV pressure ignored: %v", got)
|
||||
}
|
||||
// Full queue with an empty cache likewise.
|
||||
if got := Saturation([]endpoint.Snapshot{snap("a", 4, 1, 5, 0, false)}, 0, th); got < 1 {
|
||||
t.Errorf("queue depth ignored: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaturationCountsUnreportedInflight(t *testing.T) {
|
||||
// A burst the engine has not reported yet is real load; missing it would
|
||||
// leave admission blind exactly when it needs to bite.
|
||||
th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9}
|
||||
cold := Saturation([]endpoint.Snapshot{snap("a", 0, 0, 0, 0, false)}, 0, th)
|
||||
burst := Saturation([]endpoint.Snapshot{snap("a", 0, 0, 8, 0, false)}, 0, th)
|
||||
if !(burst > cold) {
|
||||
t.Errorf("unreported inflight not counted: cold=%v burst=%v", cold, burst)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaturationSkipsOpenBreakers(t *testing.T) {
|
||||
th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9}
|
||||
eps := []endpoint.Snapshot{
|
||||
snap("busy", 8, 4, 12, 0.5, false),
|
||||
snap("dead", 0, 0, 0, 0, true), // would halve the mean if counted
|
||||
}
|
||||
withDead := Saturation(eps, 0, th)
|
||||
alone := Saturation(eps[:1], 0, th)
|
||||
if fmt.Sprintf("%.4f", withDead) != fmt.Sprintf("%.4f", alone) {
|
||||
t.Errorf("an endpoint with an open breaker was counted as capacity: %v vs %v", withDead, alone)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaturationCountsUnreportedAdmissions(t *testing.T) {
|
||||
// The thundering herd guard: a simultaneous burst must register as load
|
||||
// before any of it reaches an engine, or the whole burst walks through the
|
||||
// fast path while every signal still reads zero.
|
||||
th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9}
|
||||
idle := []endpoint.Snapshot{snap("a", 0, 0, 0, 0, false), snap("b", 0, 0, 0, 0, false)}
|
||||
|
||||
if got := Saturation(idle, 0, th); got != 0 {
|
||||
t.Fatalf("idle pool reported %v", got)
|
||||
}
|
||||
if got := Saturation(idle, 16, th); got <= 1 {
|
||||
t.Errorf("16 unreported admissions across 2 replicas should saturate, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBurstDoesNotAllPassTheFastPath(t *testing.T) {
|
||||
// End to end for the same guard: with saturation derived from outstanding
|
||||
// leases, a burst of 200 must not be admitted wholesale.
|
||||
var admitted atomic.Int64
|
||||
c := newController(t, testConfig(), func(_ string, outstanding int) (float64, bool) {
|
||||
// Two replicas, healthy buffer of 4 each: 8 in flight is the limit.
|
||||
return float64(outstanding) / 8, true
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for range 200 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := c.Acquire(ctx, Request{Model: "m", Band: "bronze"}); err == nil {
|
||||
admitted.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// bronze's ceiling is 0.30, so roughly 8*0.3 = 2-3 should get in. Allow
|
||||
// generous slack for scheduling, but nothing near 200.
|
||||
if n := admitted.Load(); n > 20 {
|
||||
t.Errorf("%d of 200 requests were admitted at once; the burst guard is not working", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaturationWithNoUsableEndpoints(t *testing.T) {
|
||||
th := Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9}
|
||||
if got := Saturation(nil, 0, th); got != 1 {
|
||||
t.Errorf("an empty pool should report full saturation, got %v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,412 @@
|
||||
package admission
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Band is a priority class. Requests in a band are held once pool saturation
|
||||
// exceeds the band's ceiling, so the cluster degrades from the bottom up:
|
||||
// batch traffic slows first and interactive traffic keeps its latency.
|
||||
type Band struct {
|
||||
Name string `koanf:"name"`
|
||||
// Priority orders bands; higher is served first.
|
||||
Priority int `koanf:"priority"`
|
||||
// SaturationCeiling is the pool saturation above which this band waits.
|
||||
// 1.0 means "never hold this band back".
|
||||
SaturationCeiling float64 `koanf:"saturation_ceiling"`
|
||||
}
|
||||
|
||||
// Config configures the controller.
|
||||
type Config struct {
|
||||
Bands []Band `koanf:"bands"`
|
||||
DefaultBand string `koanf:"default_band"`
|
||||
|
||||
Thresholds Thresholds `koanf:"-"`
|
||||
|
||||
// QueueTTL drops a request that has waited this long. Sending the GPU work
|
||||
// the client has already given up on helps nobody.
|
||||
QueueTTL time.Duration `koanf:"queue_ttl"`
|
||||
// QueueCapacityPerBand bounds each band's queue. Rejecting immediately
|
||||
// beats accumulating silently.
|
||||
QueueCapacityPerBand int `koanf:"queue_capacity_per_band"`
|
||||
// Interval is how often the dispatcher re-checks saturation. Releases also
|
||||
// wake it, so this is only the backstop for saturation falling on its own.
|
||||
Interval time.Duration `koanf:"interval"`
|
||||
}
|
||||
|
||||
// DefaultBands returns the four-band scheme from plan §7.2.2.
|
||||
func DefaultBands() []Band {
|
||||
return []Band{
|
||||
{Name: "platinum", Priority: 400, SaturationCeiling: 1.00},
|
||||
{Name: "gold", Priority: 300, SaturationCeiling: 0.85},
|
||||
{Name: "silver", Priority: 200, SaturationCeiling: 0.60},
|
||||
{Name: "bronze", Priority: 100, SaturationCeiling: 0.30},
|
||||
}
|
||||
}
|
||||
|
||||
// DefaultConfig returns a usable configuration.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Bands: DefaultBands(),
|
||||
DefaultBand: "gold",
|
||||
Thresholds: Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9},
|
||||
QueueTTL: 30 * time.Second,
|
||||
QueueCapacityPerBand: 2000,
|
||||
Interval: 10 * time.Millisecond,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() error {
|
||||
if len(c.Bands) == 0 {
|
||||
c.Bands = DefaultBands()
|
||||
}
|
||||
slices.SortFunc(c.Bands, func(a, b Band) int { return b.Priority - a.Priority })
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
for i := range c.Bands {
|
||||
b := &c.Bands[i]
|
||||
if b.Name == "" {
|
||||
return fmt.Errorf("bands[%d]: name is required", i)
|
||||
}
|
||||
if _, dup := seen[b.Name]; dup {
|
||||
return fmt.Errorf("band %q is declared twice", b.Name)
|
||||
}
|
||||
seen[b.Name] = struct{}{}
|
||||
if b.SaturationCeiling <= 0 {
|
||||
b.SaturationCeiling = 1
|
||||
}
|
||||
}
|
||||
if c.DefaultBand == "" {
|
||||
c.DefaultBand = c.Bands[len(c.Bands)-1].Name
|
||||
}
|
||||
if _, ok := seen[c.DefaultBand]; !ok {
|
||||
return fmt.Errorf("default_band %q is not one of the configured bands", c.DefaultBand)
|
||||
}
|
||||
c.Thresholds.normalize()
|
||||
if c.QueueTTL <= 0 {
|
||||
c.QueueTTL = 30 * time.Second
|
||||
}
|
||||
if c.QueueCapacityPerBand <= 0 {
|
||||
c.QueueCapacityPerBand = 2000
|
||||
}
|
||||
if c.Interval <= 0 {
|
||||
c.Interval = 10 * time.Millisecond
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SaturationFunc reports a model's pool saturation and whether it is known.
|
||||
//
|
||||
// outstanding is how many requests this controller has admitted and not yet
|
||||
// seen released, so the caller can fold in load the engines have not reported.
|
||||
type SaturationFunc func(model string, outstanding int) (float64, bool)
|
||||
|
||||
// Request identifies who is asking for what.
|
||||
type Request struct {
|
||||
Model string
|
||||
// Band selects the priority class; unknown names fall back to the default.
|
||||
Band string
|
||||
// Tenant is the fairness key within a band. Two requests sharing a tenant
|
||||
// compete with each other rather than with everyone else.
|
||||
Tenant string
|
||||
}
|
||||
|
||||
var (
|
||||
// ErrQueueFull means the band's queue is at capacity.
|
||||
ErrQueueFull = errors.New("admission queue is full")
|
||||
// ErrQueueTimeout means the request waited longer than the TTL.
|
||||
ErrQueueTimeout = errors.New("admission queue timeout")
|
||||
// ErrClosed means the controller is shutting down.
|
||||
ErrClosed = errors.New("admission controller closed")
|
||||
)
|
||||
|
||||
// Stats is a snapshot of one band's queue, for metrics.
|
||||
type Stats struct {
|
||||
Band string
|
||||
Depth int
|
||||
}
|
||||
|
||||
// Controller admits requests according to band priority and tenant fairness.
|
||||
type Controller struct {
|
||||
cfg Config
|
||||
satFn SaturationFunc
|
||||
bandOf map[string]int
|
||||
|
||||
mu sync.Mutex
|
||||
queues []*bandQueue
|
||||
closed bool
|
||||
|
||||
wake chan struct{}
|
||||
done chan struct{}
|
||||
stop chan struct{}
|
||||
once sync.Once
|
||||
|
||||
// outstanding counts admitted-but-unreleased requests per model. It is the
|
||||
// only load signal available with zero lag, which is what stops a
|
||||
// simultaneous burst from passing the check en masse.
|
||||
outstanding map[string]int
|
||||
|
||||
// Observed counters, read through Metrics.
|
||||
admitted map[string]int64
|
||||
rejected map[string]int64
|
||||
waitTotal map[string]time.Duration
|
||||
}
|
||||
|
||||
// New builds a controller. Start must be called before it will admit anything
|
||||
// that has to wait.
|
||||
func New(cfg Config, satFn SaturationFunc) (*Controller, error) {
|
||||
if err := cfg.normalize(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c := &Controller{
|
||||
cfg: cfg,
|
||||
satFn: satFn,
|
||||
bandOf: make(map[string]int, len(cfg.Bands)),
|
||||
wake: make(chan struct{}, 1),
|
||||
done: make(chan struct{}),
|
||||
stop: make(chan struct{}),
|
||||
outstanding: map[string]int{},
|
||||
admitted: map[string]int64{},
|
||||
rejected: map[string]int64{},
|
||||
waitTotal: map[string]time.Duration{},
|
||||
}
|
||||
for i, b := range cfg.Bands {
|
||||
c.bandOf[b.Name] = i
|
||||
c.queues = append(c.queues, newBandQueue(b))
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// Bands returns the configured bands, highest priority first.
|
||||
func (c *Controller) Bands() []Band { return slices.Clone(c.cfg.Bands) }
|
||||
|
||||
// Start launches the dispatcher.
|
||||
func (c *Controller) Start() {
|
||||
go c.dispatchLoop()
|
||||
}
|
||||
|
||||
// Close stops the dispatcher and fails everything still queued.
|
||||
func (c *Controller) Close() {
|
||||
c.once.Do(func() {
|
||||
close(c.stop)
|
||||
<-c.done
|
||||
})
|
||||
}
|
||||
|
||||
// Lease is held for the lifetime of an admitted request.
|
||||
type Lease struct {
|
||||
c *Controller
|
||||
band string
|
||||
model string
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// Release returns the slot and wakes the dispatcher, since finishing a request
|
||||
// is the most likely reason saturation just dropped.
|
||||
func (l *Lease) Release() {
|
||||
if l == nil || l.c == nil {
|
||||
return
|
||||
}
|
||||
l.once.Do(func() {
|
||||
l.c.mu.Lock()
|
||||
if n := l.c.outstanding[l.model]; n > 1 {
|
||||
l.c.outstanding[l.model] = n - 1
|
||||
} else {
|
||||
delete(l.c.outstanding, l.model)
|
||||
}
|
||||
l.c.mu.Unlock()
|
||||
l.c.signal()
|
||||
})
|
||||
}
|
||||
|
||||
// Band reports which band admitted the request.
|
||||
func (l *Lease) Band() string {
|
||||
if l == nil {
|
||||
return ""
|
||||
}
|
||||
return l.band
|
||||
}
|
||||
|
||||
// Acquire blocks until the request may proceed, the context ends, or the TTL
|
||||
// expires.
|
||||
func (c *Controller) Acquire(ctx context.Context, r Request) (*Lease, error) {
|
||||
idx, ok := c.bandOf[r.Band]
|
||||
if !ok {
|
||||
idx = c.bandOf[c.cfg.DefaultBand]
|
||||
}
|
||||
band := c.cfg.Bands[idx]
|
||||
|
||||
// Fast path: under the ceiling with nobody already waiting in this band.
|
||||
// Skipping the queue here is what keeps the common case free of latency.
|
||||
c.mu.Lock()
|
||||
if c.closed {
|
||||
c.mu.Unlock()
|
||||
return nil, ErrClosed
|
||||
}
|
||||
if c.queues[idx].len() == 0 && c.admits(band, r.Model) {
|
||||
c.grant(band.Name, r.Model)
|
||||
c.mu.Unlock()
|
||||
return &Lease{c: c, band: band.Name, model: r.Model}, nil
|
||||
}
|
||||
|
||||
if c.queues[idx].len() >= c.cfg.QueueCapacityPerBand {
|
||||
c.rejected[band.Name]++
|
||||
c.mu.Unlock()
|
||||
return nil, fmt.Errorf("%w: band %q", ErrQueueFull, band.Name)
|
||||
}
|
||||
|
||||
w := &waiter{
|
||||
tenant: tenantKey(r),
|
||||
model: r.Model,
|
||||
ready: make(chan struct{}),
|
||||
enqueued: time.Now(),
|
||||
deadline: time.Now().Add(c.cfg.QueueTTL),
|
||||
}
|
||||
c.queues[idx].push(w)
|
||||
c.mu.Unlock()
|
||||
|
||||
c.signal()
|
||||
|
||||
select {
|
||||
case <-w.ready:
|
||||
c.mu.Lock()
|
||||
if c.closed {
|
||||
c.mu.Unlock()
|
||||
return nil, ErrClosed
|
||||
}
|
||||
c.waitTotal[band.Name] += time.Since(w.enqueued)
|
||||
c.mu.Unlock()
|
||||
return &Lease{c: c, band: band.Name, model: r.Model}, nil
|
||||
|
||||
case <-ctx.Done():
|
||||
c.remove(idx, w, "client_gone")
|
||||
return nil, ctx.Err()
|
||||
|
||||
case <-time.After(time.Until(w.deadline)):
|
||||
c.remove(idx, w, "ttl")
|
||||
return nil, fmt.Errorf("%w after %v", ErrQueueTimeout, c.cfg.QueueTTL)
|
||||
}
|
||||
}
|
||||
|
||||
// remove takes a waiter out of its queue unless it was admitted first.
|
||||
func (c *Controller) remove(idx int, w *waiter, reason string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if w.admitted {
|
||||
// It won the race; the caller will see ready closed next time round.
|
||||
return
|
||||
}
|
||||
c.queues[idx].remove(w)
|
||||
c.rejected[c.cfg.Bands[idx].Name+"/"+reason]++
|
||||
}
|
||||
|
||||
// grant records an admission. The caller holds the lock.
|
||||
func (c *Controller) grant(band, model string) {
|
||||
c.admitted[band]++
|
||||
c.outstanding[model]++
|
||||
}
|
||||
|
||||
// admits reports whether a band may proceed right now. The caller holds the lock.
|
||||
func (c *Controller) admits(b Band, model string) bool {
|
||||
sat, known := c.satFn(model, c.outstanding[model])
|
||||
if !known {
|
||||
// Nothing is known about this model's pool; let it through so dispatch
|
||||
// can return a precise error instead of a queue timeout.
|
||||
return true
|
||||
}
|
||||
return sat <= b.SaturationCeiling
|
||||
}
|
||||
|
||||
func (c *Controller) signal() {
|
||||
select {
|
||||
case c.wake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Controller) dispatchLoop() {
|
||||
defer close(c.done)
|
||||
t := time.NewTicker(c.cfg.Interval)
|
||||
defer t.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-c.stop:
|
||||
c.drain()
|
||||
return
|
||||
case <-c.wake:
|
||||
c.dispatch()
|
||||
case <-t.C:
|
||||
c.dispatch()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// dispatch walks the bands from highest priority down, admitting whatever the
|
||||
// current saturation allows.
|
||||
func (c *Controller) dispatch() {
|
||||
now := time.Now()
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
for i, q := range c.queues {
|
||||
q.dropExpired(now, func() { c.rejected[c.cfg.Bands[i].Name+"/ttl"]++ })
|
||||
|
||||
band := c.cfg.Bands[i]
|
||||
// Each pass admits at most the queue length, so a band cannot spin
|
||||
// forever while saturation is recomputed from stale metrics.
|
||||
for range q.len() {
|
||||
w := q.peek()
|
||||
if w == nil || !c.admits(band, w.model) {
|
||||
break
|
||||
}
|
||||
q.pop()
|
||||
w.admitted = true
|
||||
c.grant(band.Name, w.model)
|
||||
close(w.ready)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Controller) drain() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.closed = true
|
||||
for _, q := range c.queues {
|
||||
for {
|
||||
w := q.pop()
|
||||
if w == nil {
|
||||
break
|
||||
}
|
||||
// Closing without marking admitted lets Acquire report ErrClosed.
|
||||
close(w.ready)
|
||||
w.admitted = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Depths returns each band's current queue depth.
|
||||
func (c *Controller) Depths() []Stats {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
out := make([]Stats, 0, len(c.queues))
|
||||
for i, q := range c.queues {
|
||||
out = append(out, Stats{Band: c.cfg.Bands[i].Name, Depth: q.len()})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// tenantKey picks the fairness key, falling back to the band so that unknown
|
||||
// tenants at least compete as one group rather than each looking unique.
|
||||
func tenantKey(r Request) string {
|
||||
if r.Tenant != "" {
|
||||
return r.Tenant
|
||||
}
|
||||
return "anonymous"
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package admission
|
||||
|
||||
import "time"
|
||||
|
||||
// waiter is one queued request.
|
||||
type waiter struct {
|
||||
tenant string
|
||||
model string
|
||||
ready chan struct{}
|
||||
enqueued time.Time
|
||||
deadline time.Time
|
||||
// admitted guards the race between the dispatcher admitting a waiter and
|
||||
// the waiter giving up on its own.
|
||||
admitted bool
|
||||
}
|
||||
|
||||
// bandQueue holds one band's waiters, grouped by tenant.
|
||||
//
|
||||
// Grouping is what makes fairness possible: a single tenant sending a thousand
|
||||
// requests occupies one slot in the rotation, not a thousand, so a traffic
|
||||
// spike from one caller cannot consume the band.
|
||||
type bandQueue struct {
|
||||
band Band
|
||||
|
||||
byTenant map[string][]*waiter
|
||||
// rotation is the round-robin order of tenants with pending work.
|
||||
rotation []string
|
||||
next int
|
||||
count int
|
||||
}
|
||||
|
||||
func newBandQueue(b Band) *bandQueue {
|
||||
return &bandQueue{band: b, byTenant: map[string][]*waiter{}}
|
||||
}
|
||||
|
||||
func (q *bandQueue) len() int { return q.count }
|
||||
|
||||
func (q *bandQueue) push(w *waiter) {
|
||||
if _, ok := q.byTenant[w.tenant]; !ok {
|
||||
q.byTenant[w.tenant] = nil
|
||||
q.rotation = append(q.rotation, w.tenant)
|
||||
}
|
||||
q.byTenant[w.tenant] = append(q.byTenant[w.tenant], w)
|
||||
q.count++
|
||||
}
|
||||
|
||||
// peek returns the waiter the next pop would return, without removing it.
|
||||
func (q *bandQueue) peek() *waiter {
|
||||
if q.count == 0 {
|
||||
return nil
|
||||
}
|
||||
idx, tenant := q.nextTenant()
|
||||
if idx < 0 {
|
||||
return nil
|
||||
}
|
||||
_ = tenant
|
||||
return q.byTenant[q.rotation[idx]][0]
|
||||
}
|
||||
|
||||
// pop removes and returns the next waiter in tenant rotation.
|
||||
func (q *bandQueue) pop() *waiter {
|
||||
if q.count == 0 {
|
||||
return nil
|
||||
}
|
||||
idx, _ := q.nextTenant()
|
||||
if idx < 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
tenant := q.rotation[idx]
|
||||
ws := q.byTenant[tenant]
|
||||
w := ws[0]
|
||||
q.byTenant[tenant] = ws[1:]
|
||||
q.count--
|
||||
|
||||
if len(q.byTenant[tenant]) == 0 {
|
||||
delete(q.byTenant, tenant)
|
||||
q.rotation = append(q.rotation[:idx], q.rotation[idx+1:]...)
|
||||
if idx < q.next {
|
||||
q.next--
|
||||
}
|
||||
} else {
|
||||
q.next = idx + 1
|
||||
}
|
||||
if len(q.rotation) > 0 {
|
||||
q.next %= len(q.rotation)
|
||||
} else {
|
||||
q.next = 0
|
||||
}
|
||||
return w
|
||||
}
|
||||
|
||||
// nextTenant finds the next rotation entry that still has waiters.
|
||||
func (q *bandQueue) nextTenant() (int, string) {
|
||||
n := len(q.rotation)
|
||||
for i := range n {
|
||||
idx := (q.next + i) % n
|
||||
if len(q.byTenant[q.rotation[idx]]) > 0 {
|
||||
return idx, q.rotation[idx]
|
||||
}
|
||||
}
|
||||
return -1, ""
|
||||
}
|
||||
|
||||
// remove takes a specific waiter out, for cancellation and TTL expiry.
|
||||
func (q *bandQueue) remove(target *waiter) bool {
|
||||
ws, ok := q.byTenant[target.tenant]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
for i, w := range ws {
|
||||
if w != target {
|
||||
continue
|
||||
}
|
||||
q.byTenant[target.tenant] = append(ws[:i:i], ws[i+1:]...)
|
||||
q.count--
|
||||
if len(q.byTenant[target.tenant]) == 0 {
|
||||
delete(q.byTenant, target.tenant)
|
||||
q.dropTenant(target.tenant)
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (q *bandQueue) dropTenant(tenant string) {
|
||||
for i, t := range q.rotation {
|
||||
if t != tenant {
|
||||
continue
|
||||
}
|
||||
q.rotation = append(q.rotation[:i], q.rotation[i+1:]...)
|
||||
if i < q.next {
|
||||
q.next--
|
||||
}
|
||||
break
|
||||
}
|
||||
if len(q.rotation) > 0 {
|
||||
q.next %= len(q.rotation)
|
||||
} else {
|
||||
q.next = 0
|
||||
}
|
||||
}
|
||||
|
||||
// dropExpired removes waiters past their deadline, calling onDrop for each.
|
||||
// Their Acquire call is already unblocking on its own timer; this keeps the
|
||||
// queue from carrying corpses that would occupy capacity.
|
||||
func (q *bandQueue) dropExpired(now time.Time, onDrop func()) {
|
||||
if q.count == 0 {
|
||||
return
|
||||
}
|
||||
for tenant, ws := range q.byTenant {
|
||||
kept := ws[:0]
|
||||
for _, w := range ws {
|
||||
if now.After(w.deadline) && !w.admitted {
|
||||
q.count--
|
||||
onDrop()
|
||||
continue
|
||||
}
|
||||
kept = append(kept, w)
|
||||
}
|
||||
if len(kept) == 0 {
|
||||
delete(q.byTenant, tenant)
|
||||
q.dropTenant(tenant)
|
||||
continue
|
||||
}
|
||||
q.byTenant[tenant] = kept
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
// Package admission answers "may this request proceed now?" (plan §7.2).
|
||||
//
|
||||
// It is deliberately separate from scheduling, which answers "where should it
|
||||
// go?". Kept apart, the router can hold requests in its own queue and apply
|
||||
// priority and tenant fairness *before* they enter an engine queue. An engine
|
||||
// queue is FIFO and knows nothing about tenants: once a batch job has put 500
|
||||
// requests into it, no amount of clever endpoint selection will rescue the
|
||||
// interactive request that arrives next.
|
||||
package admission
|
||||
|
||||
import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
)
|
||||
|
||||
// Thresholds define what "saturated" means for a pool.
|
||||
type Thresholds struct {
|
||||
// TargetWaitingPerReplica is the healthy queue depth: enough for continuous
|
||||
// batching to stay efficient, not enough for queueing delay to show up.
|
||||
// Measure it with ctl bench (plan §10); this single number governs how the
|
||||
// whole system behaves under load.
|
||||
TargetWaitingPerReplica float64
|
||||
// KVCeiling is the KV occupancy treated as full.
|
||||
KVCeiling float64
|
||||
}
|
||||
|
||||
func (t *Thresholds) normalize() {
|
||||
if t.TargetWaitingPerReplica <= 0 {
|
||||
t.TargetWaitingPerReplica = 4
|
||||
}
|
||||
if t.KVCeiling <= 0 || t.KVCeiling > 1 {
|
||||
t.KVCeiling = 0.9
|
||||
}
|
||||
}
|
||||
|
||||
// Saturation reduces a pool's state to a single number, where 1 means "at the
|
||||
// healthy limit" and above 1 means overloaded.
|
||||
//
|
||||
// Two signals are combined because either alone misleads: a pool can have an
|
||||
// empty queue while its KV cache is nearly full (long contexts), or plenty of
|
||||
// KV headroom while requests pile up (slow decode). The worse of the two wins.
|
||||
//
|
||||
// unreported is load that has been admitted but that no endpoint has accounted
|
||||
// for yet. Without it a simultaneous burst all passes the check before any of
|
||||
// it becomes visible — which is precisely when admission needs to bite.
|
||||
func Saturation(eps []endpoint.Snapshot, unreported int, t Thresholds) float64 {
|
||||
t.normalize()
|
||||
|
||||
var sum float64
|
||||
var n int
|
||||
for _, e := range eps {
|
||||
// A replica the breaker has taken out is not available capacity, so
|
||||
// counting it would make the pool look healthier than it is.
|
||||
if e.Breaker.String() == "open" {
|
||||
continue
|
||||
}
|
||||
|
||||
// The engine's own queue is the truth, but it is one scrape interval
|
||||
// old. Fill that gap with what the router knows without any lag: how
|
||||
// many requests it has dispatched here, minus what this engine can run
|
||||
// at once. Anything beyond that capacity is queued by definition.
|
||||
//
|
||||
// Subtracting the capacity is the whole point. Counting raw inflight
|
||||
// would make a fully busy but perfectly healthy engine look like it had
|
||||
// a queue as long as its batch, every single scrape interval.
|
||||
waiting := float64(e.Stats.Waiting)
|
||||
if e.RunningCapacity > 0 {
|
||||
if queued := e.Inflight - e.RunningCapacity; float64(queued) > waiting {
|
||||
waiting = float64(queued)
|
||||
}
|
||||
} else if e.Inflight > 0 {
|
||||
// Nothing observed yet, so capacity is unknown. Treat dispatched
|
||||
// work as queued: conservative for the first scrape, and it relaxes
|
||||
// as soon as the engine reports anything.
|
||||
waiting = max(waiting, float64(e.Inflight))
|
||||
}
|
||||
|
||||
q := waiting / t.TargetWaitingPerReplica
|
||||
kv := e.Stats.KVCacheUsage / t.KVCeiling
|
||||
sum += max(q, kv)
|
||||
n++
|
||||
}
|
||||
_ = n
|
||||
|
||||
if n == 0 {
|
||||
// Every replica is out. Report full saturation so low bands stop
|
||||
// queueing work that has nowhere to go; the dispatch layer will return
|
||||
// a clearer error than a queue timeout.
|
||||
return 1
|
||||
}
|
||||
// Spread unaccounted admissions evenly; they will land somewhere.
|
||||
if unreported > 0 {
|
||||
sum += float64(unreported) / t.TargetWaitingPerReplica
|
||||
}
|
||||
|
||||
// The mean, not the minimum: one idle replica does not make a busy pool
|
||||
// idle, and admission is about the pool's capacity as a whole.
|
||||
return sum / float64(n)
|
||||
}
|
||||
@@ -71,6 +71,14 @@ type Snapshot struct {
|
||||
// unavoidable within a 2s scrape interval.
|
||||
Inflight int
|
||||
|
||||
// RunningCapacity is the highest concurrent request count ever observed on
|
||||
// this endpoint — in practice the engine's batch capacity, learned rather
|
||||
// than configured. Without it, requests the router has dispatched but the
|
||||
// engine has not yet reported are indistinguishable from queue depth, and
|
||||
// during a scrape interval that makes a perfectly healthy engine look
|
||||
// overloaded.
|
||||
RunningCapacity int
|
||||
|
||||
// Stale means no successful scrape arrived within StaleAfter. Stale
|
||||
// endpoints are not dropped, only scored down.
|
||||
Stale bool
|
||||
@@ -101,10 +109,11 @@ func DefaultConfig() Config {
|
||||
type live struct {
|
||||
info Info
|
||||
|
||||
stats Stats
|
||||
scrapedAt time.Time
|
||||
lastSucc float64
|
||||
drainRate float64
|
||||
stats Stats
|
||||
scrapedAt time.Time
|
||||
lastSucc float64
|
||||
drainRate float64
|
||||
maxRunning int
|
||||
|
||||
inflight int
|
||||
brk *breaker.Breaker
|
||||
@@ -199,6 +208,7 @@ func (r *Registry) Observe(id string, s Stats) {
|
||||
l.stats = s
|
||||
l.lastSucc = s.SuccessTotal
|
||||
l.scrapedAt = now
|
||||
l.maxRunning = max(l.maxRunning, s.Running)
|
||||
}
|
||||
|
||||
// Snapshots returns a view of every endpoint serving the given model, or of all
|
||||
@@ -220,13 +230,14 @@ func (r *Registry) Snapshots(model string) []Snapshot {
|
||||
stale = age > r.cfg.StaleAfter
|
||||
}
|
||||
out = append(out, Snapshot{
|
||||
Info: l.info,
|
||||
Stats: l.stats,
|
||||
DrainRate: l.drainRate,
|
||||
Inflight: l.inflight,
|
||||
Stale: stale,
|
||||
Age: age,
|
||||
Breaker: l.brk.State(),
|
||||
Info: l.info,
|
||||
Stats: l.stats,
|
||||
DrainRate: l.drainRate,
|
||||
Inflight: l.inflight,
|
||||
RunningCapacity: l.maxRunning,
|
||||
Stale: stale,
|
||||
Age: age,
|
||||
Breaker: l.brk.State(),
|
||||
})
|
||||
}
|
||||
slices.SortFunc(out, func(a, b Snapshot) int { return strcmp(a.ID, b.ID) })
|
||||
|
||||
@@ -0,0 +1,361 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/admission"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
|
||||
)
|
||||
|
||||
// admit waits for the flow control layer to let the request through. With
|
||||
// admission disabled it returns immediately, which is the v0.2.0 behaviour.
|
||||
func (r *Router) admit(ctx context.Context, req *parsedRequest) (*admission.Lease, error) {
|
||||
if r.adm == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return r.adm.Acquire(ctx, admission.Request{
|
||||
Model: req.Model,
|
||||
Band: req.Band,
|
||||
// Until authentication arrives the session is the best tenant key we
|
||||
// have: it groups one caller's traffic without pretending to identify
|
||||
// them (plan §7.9).
|
||||
Tenant: req.SessionID,
|
||||
})
|
||||
}
|
||||
|
||||
// rejectAdmission turns a flow control refusal into a status the client can act
|
||||
// on: 503 when the queue is full (try later), 504 when it timed out waiting.
|
||||
func (r *Router) rejectAdmission(w http.ResponseWriter, req *parsedRequest, err error, start time.Time) {
|
||||
status := http.StatusServiceUnavailable
|
||||
switch {
|
||||
case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded):
|
||||
// The client gave up while queued; nobody is listening.
|
||||
r.m.observeAdmissionRejected(req.Band, "client_gone")
|
||||
r.m.observeRequest(req.Model, "", 499, time.Since(start), nil, 0)
|
||||
return
|
||||
case errors.Is(err, admission.ErrQueueTimeout):
|
||||
status = http.StatusGatewayTimeout
|
||||
}
|
||||
reason := "unavailable"
|
||||
switch {
|
||||
case errors.Is(err, admission.ErrQueueFull):
|
||||
reason = "queue_full"
|
||||
case errors.Is(err, admission.ErrQueueTimeout):
|
||||
reason = "ttl"
|
||||
case errors.Is(err, admission.ErrClosed):
|
||||
reason = "shutdown"
|
||||
}
|
||||
r.m.observeAdmissionRejected(req.Band, reason)
|
||||
writeError(w, status, "server_error", err.Error())
|
||||
r.m.observeRequest(req.Model, "", status, time.Since(start), nil, 0)
|
||||
}
|
||||
|
||||
// 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
|
||||
Band 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
|
||||
}
|
||||
|
||||
// Step 5 of the lifecycle: "now?" before "where?" (plan §7.1).
|
||||
lease, err := r.admit(req.Context(), parsed)
|
||||
if err != nil {
|
||||
r.rejectAdmission(w, parsed, err, start)
|
||||
return
|
||||
}
|
||||
defer lease.Release()
|
||||
queued := time.Since(start)
|
||||
r.m.observeAdmission(lease.Band(), parsed.Model, queued)
|
||||
if queued > time.Millisecond {
|
||||
w.Header().Set("x-llmgw-queue-ms", strconv.FormatInt(queued.Milliseconds(), 10))
|
||||
}
|
||||
|
||||
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),
|
||||
Band: req.Header.Get(r.cfg.Admission.BandHeader),
|
||||
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
|
||||
}
|
||||
|
||||
// servedModels lists the canonical models at least one endpoint serves.
|
||||
func (r *Router) servedModels() []string {
|
||||
seen := map[string]struct{}{}
|
||||
for _, s := range r.reg.Snapshots("") {
|
||||
for _, m := range s.Models {
|
||||
seen[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
out := make([]string, 0, len(seen))
|
||||
for m := range seen {
|
||||
out = append(out, m)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// 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,137 @@
|
||||
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
|
||||
|
||||
admissionWait *prometheus.HistogramVec
|
||||
admissionRejected *prometheus.CounterVec
|
||||
}
|
||||
|
||||
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},
|
||||
}),
|
||||
|
||||
admissionWait: prometheus.NewHistogramVec(prometheus.HistogramOpts{
|
||||
Name: "llmgw_admission_wait_seconds",
|
||||
Help: "Time spent waiting in the admission queue.",
|
||||
Buckets: []float64{0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 2.5, 5, 10, 30},
|
||||
}, []string{"band", "model"}),
|
||||
|
||||
admissionRejected: prometheus.NewCounterVec(prometheus.CounterOpts{
|
||||
Name: "llmgw_admission_rejected_total",
|
||||
Help: "Requests refused by flow control, by reason.",
|
||||
}, []string{"band", "reason"}),
|
||||
}
|
||||
}
|
||||
|
||||
// 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,
|
||||
m.admissionWait, m.admissionRejected,
|
||||
}
|
||||
}
|
||||
|
||||
func (m *metrics) observeAdmission(band, model string, wait time.Duration) {
|
||||
if band == "" {
|
||||
band = "none"
|
||||
}
|
||||
m.admissionWait.WithLabelValues(band, model).Observe(wait.Seconds())
|
||||
}
|
||||
|
||||
func (m *metrics) observeAdmissionRejected(band, reason string) {
|
||||
if band == "" {
|
||||
band = "unknown"
|
||||
}
|
||||
m.admissionRejected.WithLabelValues(band, reason).Inc()
|
||||
}
|
||||
|
||||
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,354 @@
|
||||
// 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/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/admission"
|
||||
"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
|
||||
// adm is nil when admission control is disabled.
|
||||
adm *admission.Controller
|
||||
|
||||
profiles map[string]*scheduler.Profile
|
||||
// models maps every name a client may use — canonical names and aliases —
|
||||
// onto its catalog entry.
|
||||
models map[string]catalog.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]catalog.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)
|
||||
|
||||
if cfg.Router.Admission.Enabled {
|
||||
adm, err := admission.New(admissionConfig(cfg.Router), r.saturation)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("admission: %w", err)
|
||||
}
|
||||
r.adm = adm
|
||||
}
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func admissionConfig(rc config.Router) admission.Config {
|
||||
bands := make([]admission.Band, 0, len(rc.Admission.Bands))
|
||||
for _, b := range rc.Admission.Bands {
|
||||
bands = append(bands, admission.Band{
|
||||
Name: b.Name,
|
||||
Priority: b.Priority,
|
||||
SaturationCeiling: b.SaturationCeiling,
|
||||
})
|
||||
}
|
||||
return admission.Config{
|
||||
Bands: bands,
|
||||
DefaultBand: rc.Admission.DefaultBand,
|
||||
Thresholds: admission.Thresholds{
|
||||
// Admission and the queue scorer must agree on what a healthy
|
||||
// buffer is, or the two layers fight each other.
|
||||
TargetWaitingPerReplica: rc.Scoring.TargetWaiting,
|
||||
KVCeiling: rc.Scoring.KVCeiling,
|
||||
},
|
||||
QueueTTL: rc.Admission.QueueTTL,
|
||||
QueueCapacityPerBand: rc.Admission.QueueCapacityPerBand,
|
||||
Interval: rc.Admission.Interval,
|
||||
}
|
||||
}
|
||||
|
||||
// saturation reports the pool pressure for a model, for the admission layer.
|
||||
//
|
||||
// outstanding is what admission has let through. It is compared against the
|
||||
// registry's inflight counters rather than the scraped engine stats, because
|
||||
// inflight updates the instant a request is dispatched while scraped stats lag
|
||||
// by a whole interval. Comparing against scraped data would make every
|
||||
// dispatched-but-unscraped request look like queue depth and throttle the
|
||||
// steady state down to a few requests per scrape.
|
||||
//
|
||||
// What is left over — admitted but not yet dispatched — is the thundering herd:
|
||||
// a burst that passed the check together before any of it became visible.
|
||||
func (r *Router) saturation(model string, outstanding int) (float64, bool) {
|
||||
eps := r.reg.Snapshots(model)
|
||||
if len(eps) == 0 {
|
||||
return 0, false
|
||||
}
|
||||
dispatched := 0
|
||||
for _, e := range eps {
|
||||
dispatched += e.Inflight
|
||||
}
|
||||
return admission.Saturation(eps, max(0, outstanding-dispatched), admission.Thresholds{
|
||||
TargetWaitingPerReplica: r.cfg.Scoring.TargetWaiting,
|
||||
KVCeiling: r.cfg.Scoring.KVCeiling,
|
||||
}), true
|
||||
}
|
||||
|
||||
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 and dispatching admissions.
|
||||
// It returns immediately.
|
||||
func (r *Router) Start(ctx context.Context) {
|
||||
r.poller.Sync(ctx)
|
||||
if r.adm != nil {
|
||||
r.adm.Start()
|
||||
}
|
||||
}
|
||||
|
||||
// Stop halts scraping and admission.
|
||||
func (r *Router) Stop() {
|
||||
r.poller.Stop()
|
||||
if r.adm != nil {
|
||||
r.adm.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// Registry exposes the endpoint registry, for tests and admin endpoints.
|
||||
func (r *Router) Registry() *endpoint.Registry { return r.reg }
|
||||
|
||||
// SetEndpoints replaces the endpoint list and resynchronizes metric scraping.
|
||||
//
|
||||
// This is how the control plane feeds the router from phase 3 onwards, taking
|
||||
// over from the static configuration list (plan §6.4). Endpoints that survive
|
||||
// the update keep their live state — inflight counts, breaker state, drain rate
|
||||
// — so a replica that was already serving does not start from cold just
|
||||
// because another one moved.
|
||||
func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) {
|
||||
r.reg.Replace(infos)
|
||||
r.poller.Sync(ctx)
|
||||
}
|
||||
|
||||
// --- 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 catalog.Model) {
|
||||
if e, ok := r.models[name]; ok {
|
||||
return e.Name, e
|
||||
}
|
||||
return name, catalog.Model{Name: name}
|
||||
}
|
||||
|
||||
func (r *Router) profileFor(m catalog.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 catalog.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,325 @@
|
||||
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"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
)
|
||||
|
||||
// catModel builds a minimal catalog entry for router tests, which care only
|
||||
// about names, aliases and fallbacks.
|
||||
func catModel(name string, aliases []string, fallbacks ...string) catalog.Model {
|
||||
return catalog.Model{
|
||||
Name: name, Source: "/models/" + name, Aliases: aliases, Fallbacks: fallbacks,
|
||||
Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}},
|
||||
}
|
||||
}
|
||||
|
||||
// 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 = []catalog.Model{catModel("primary", nil, "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 = []catalog.Model{catModel("llama-3.3-70b", []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 = []catalog.Model{catModel("llama-3.3-70b", []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)
|
||||
}
|
||||
}
|
||||
+362
-27
@@ -21,7 +21,8 @@
|
||||
saturation) and **scheduling** (plugin-based filter/score/pick with queue, KV-cache and
|
||||
prefix awareness), plus fallback, retry, cooldown and circuit breaking.
|
||||
- Fabric awareness: a multi-node replica **never** spans two fabrics; a TP group **never** mixes GPU types.
|
||||
- API keys, rate limits, token accounting, Prometheus metrics.
|
||||
- API keys, rate limits, token accounting and budgets, Prometheus metrics. Persistence for these is
|
||||
the only thing the store exists for (§6.5, §7.9).
|
||||
|
||||
### 1.2 Out of scope (at least for v1)
|
||||
- Inference itself. Go only orchestrates and proxies; vLLM/SGLang run the model.
|
||||
@@ -37,7 +38,7 @@
|
||||
|---|---|
|
||||
| **Node** | A physical server running the agent. |
|
||||
| **Fabric / Island** | An independent IB fabric. Nodes without IB belong to the `none` island and each counts as its own single-node island. |
|
||||
| **GPU Class** | A group of GPUs with the same model and VRAM (e.g. `h100-80g`, `a100-40g`, `l40s-48g`). |
|
||||
| **GPU Class** | A group of GPUs with the same model and VRAM (e.g. `h100-80g`, `a100-40g`, `l40s-45g`). |
|
||||
| **Pool** | A (fabric, gpu class) pair. The unit of placement. |
|
||||
| **Model** | A logical model in the catalog (e.g. `llama-3.3-70b`). This is the name users see. |
|
||||
| **Profile** | How a model runs on a particular GPU class (TP, PP, quantization, max length, engine arguments). |
|
||||
@@ -46,6 +47,7 @@
|
||||
| **Routing Profile** | The scorer set and weights the router uses for a class of request (chat / batch / long-context). Not to be confused with the model **Profile** above. |
|
||||
| **Band** | A priority class (platinum/gold/silver/bronze). The unit of admission decisions. |
|
||||
| **Block** | The unit of KV cache accounting; a fixed number of tokens (16 in vLLM, 64 by default in the GIE indexer). |
|
||||
| **Unit** | The billing unit budgets are counted in. Self-hosted models have no per-token price, so an operator defines what a unit is: GPU-seconds, a notional rate per million tokens, or 1.0 everywhere, which turns budgets into token quotas (§7.9.4). |
|
||||
|
||||
---
|
||||
|
||||
@@ -82,7 +84,7 @@ semantic routing / LLM-based model selection (out of scope), and external provid
|
||||
│ │ (data) │◄─┤ - Inventory │ │
|
||||
│ │ │ │ - Catalog │ │
|
||||
│ └─────┬──────┘ │ - Scheduler │ │
|
||||
│ │ │ - Store (SQLite) │ │
|
||||
│ │ │ - Store (phase 5) │ │
|
||||
│ │ └─────────▲────────────┘ │
|
||||
└────────┼───────────────────┼──────────────┘
|
||||
HTTP (eth) │ │ gRPC stream (mTLS, agent dials out)
|
||||
@@ -144,8 +146,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 +167,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 |
|
||||
@@ -194,9 +199,15 @@ For every GPU:
|
||||
- NVLink peers (NVML `GetNvLinkRemotePciInfo`) → which GPUs pair efficiently for TP within a node
|
||||
- health: XID errors, ECC, temperature, memory in use (are there non-engine processes?)
|
||||
|
||||
`GPUClass` is derived automatically: `normalize(name) + "-" + round(vram_gb) + "g"` → `h100-80g`.
|
||||
`GPUClass` is derived automatically: `normalize(name) + "-" + round(vram_gib) + "g"` → `h100-80g`.
|
||||
It can be overridden in configuration.
|
||||
|
||||
The size comes from what the driver reports, in GiB, which does not always match the marketing
|
||||
name: an L40S is sold as 48GB and reports 44.99 GiB, so its class is `l40s-45g`. Consistency
|
||||
matters more than matching the box here, because the number is what decides whether a model
|
||||
fits. The model token is kept in full for the same reason — a 4090 and a 4080 are not
|
||||
interchangeable, so the class is `4090-24g`, not `rtx-24g`.
|
||||
|
||||
### 5.2 IB discovery and fabric identity
|
||||
From `/sys/class/infiniband/<hca>/ports/<n>/`:
|
||||
- `state` (ACTIVE?), `phys_state`, `rate`, `link_layer` (InfiniBand / Ethernet → RoCE)
|
||||
@@ -465,10 +476,72 @@ the replicas it knows about, which mitigates the single-point-of-failure risk in
|
||||
cache ages past `endpoint_cache_max_age` (say 5 min) it only raises a warning metric; it never
|
||||
stops traffic.
|
||||
|
||||
### 6.5 Store
|
||||
- SQLite in WAL mode. Tables: `nodes`, `models`, `profiles`, `replicas`, `shards`, `api_keys`, `usage_rollup`, `events`.
|
||||
- Inventory and replica state can be rebuilt from the agents; what genuinely needs to persist is the catalog, the keys and the usage records.
|
||||
- HA (v2): Raft-replicated state, or an external Postgres/etcd.
|
||||
### 6.5 What is stored, and what is not
|
||||
|
||||
v0.1 of this plan listed eight tables. Most of them were a mistake, and the
|
||||
reason is worth stating because it applies to every piece of state here:
|
||||
|
||||
> **A database that describes the cluster can be wrong. The cluster cannot.**
|
||||
|
||||
An engine that crashes during a restart leaves a row saying "running". A node
|
||||
whose GPU developed an ECC fault leaves a row saying the card is fine. Every
|
||||
such row is a trap: the scheduler believes it, places work accordingly, and the
|
||||
failure surfaces somewhere else entirely. Rebuilding the same state from the
|
||||
fleet on every reconnect cannot go stale, because the fleet is the thing being
|
||||
described.
|
||||
|
||||
So:
|
||||
|
||||
| State | Persisted | Where it actually comes from |
|
||||
|---|---|---|
|
||||
| Nodes, GPUs, fabrics | no | Every agent sends its full inventory on registration (§5.6) |
|
||||
| Replicas, shards | no | Reported by agents; see the restart hole below |
|
||||
| Models, profiles | no | The YAML catalog is the source of truth |
|
||||
| API keys and the tenancy tree | **yes** | Created at runtime, secret, rotated (§7.9) |
|
||||
| Usage records and rollups | **yes** | Accounting must survive a restart (§7.9) |
|
||||
| Capacity profiles from `ctl bench` | **yes** | Measured, not declared (§6.2.1) |
|
||||
| Audit events | **yes** | An audit trail that forgets is not one |
|
||||
|
||||
Everything in the "yes" column belongs to tenancy and benchmarking, which is
|
||||
phase 5. **The store lands there, not in phase 3**, so its schema is shaped by
|
||||
real requirements instead of guesses.
|
||||
|
||||
#### The control plane restart hole
|
||||
|
||||
Not persisting replica state has one consequence that must be closed before it
|
||||
matters. If the control plane restarts while engines are running:
|
||||
|
||||
1. The engines keep serving — correct, and the reason the router talks to them
|
||||
directly.
|
||||
2. Agents reconnect and re-register their hardware.
|
||||
3. The control plane has no memory of any replica, because `Register` does not
|
||||
carry one.
|
||||
4. The reconciler sees zero replicas and places new ones.
|
||||
5. The old engines keep holding GPUs that the inventory now believes are free,
|
||||
so the new engines fail to allocate — with an error pointing at the new
|
||||
process rather than the old one.
|
||||
|
||||
The fix is on the agent side, not in a database: `Register` carries the shards
|
||||
the agent is currently supervising, and the control plane rebuilds replica state
|
||||
from what the fleet reports. The supervisor already knows its live shards and
|
||||
already finds, through pid files, the ones that survived an agent restart
|
||||
(`FindOrphans`). A shard nothing claims after one reconciliation pass is an
|
||||
orphan and is killed.
|
||||
|
||||
This is strictly better than a database would be: it cannot disagree with
|
||||
reality, and it covers agent restarts and control plane restarts with one
|
||||
mechanism.
|
||||
|
||||
#### Store choice
|
||||
|
||||
SQLite in WAL mode, via `modernc.org/sqlite` so the server stays cgo-free.
|
||||
Embedded means no second service to operate, and WAL gives concurrent readers
|
||||
alongside one writer, which matches the access pattern: many reads for auth,
|
||||
one batched writer for usage.
|
||||
|
||||
It stops being enough at the point where several router instances must write at
|
||||
once. The store therefore sits behind an interface from the start, and the swap
|
||||
to Postgres comes with multi-instance deployment rather than before it.
|
||||
|
||||
---
|
||||
|
||||
@@ -521,18 +594,35 @@ batching to stay efficient, but not enough for the queue to balloon. Set
|
||||
`TargetWaitingPerReplica` too low and the GPU starves; too high and queueing delay returns. Both
|
||||
thresholds must be tunable per model/profile, and `ctl bench` (§10) should propose them.
|
||||
|
||||
**Two lags have to be closed, and they need different answers.** The engine's queue is scraped,
|
||||
so it is up to one interval old; the router's own inflight count is not. But inflight includes
|
||||
*running* requests, and treating those as queue makes a perfectly healthy engine look overloaded
|
||||
every interval. What is actually queued is `inflight − batch capacity`, and batch capacity is not
|
||||
in the metric contract — so llmgw learns it as the high-water mark of observed `running`. The
|
||||
second lag is shorter and subtler: between admission granting a request and the router
|
||||
dispatching it, the request is invisible everywhere. A simultaneous burst can therefore pass the
|
||||
check en masse, which is exactly the batch-job arrival pattern this layer exists to survive. The
|
||||
controller closes it by counting its own outstanding leases against dispatched inflight.
|
||||
|
||||
#### 7.2.2 Priority bands
|
||||
|
||||
Every API key maps to a band, and the band sets the saturation ceiling:
|
||||
|
||||
```yaml
|
||||
bands:
|
||||
- { name: platinum, priority: 400, saturation_ceiling: 1.00 } # never held back
|
||||
- { name: platinum, priority: 400, saturation_ceiling: 1.00 } # flows to the healthy limit
|
||||
- { name: gold, priority: 300, saturation_ceiling: 0.85 }
|
||||
- { name: silver, priority: 200, saturation_ceiling: 0.60 }
|
||||
- { name: bronze, priority: 100, saturation_ceiling: 0.30 } # batch jobs
|
||||
```
|
||||
|
||||
A ceiling of 1.00 means "hold only once the pool is past its healthy limit",
|
||||
not "never hold". Holding even the top band is deliberate: there is nobody above
|
||||
it to protect, so the benefit is a bounded engine queue. A flood dumped straight
|
||||
into a FIFO engine queue makes every request in that band slow, while queueing it
|
||||
at the gateway preserves ordering, TTL and tenant fairness. A band that must
|
||||
never wait — a health probe, say — gets a ceiling above 1 (`1000`).
|
||||
|
||||
The dispatch loop scans bands from highest to lowest. Once pool saturation exceeds a band's
|
||||
ceiling, that band's requests wait in the queue. The result: as the cluster fills, batch traffic
|
||||
slows down first and interactive traffic is protected. (In llm-d's measurements platinum, gold
|
||||
@@ -771,6 +861,186 @@ profile** (§7.3).
|
||||
- v1 is a single instance with in-memory limiters. With multiple instances this becomes Redis or
|
||||
consistent-hash stickiness (v2).
|
||||
|
||||
#### 7.9.1 How keys are stored
|
||||
|
||||
A key is a 256-bit random value, shown once at creation and never again:
|
||||
|
||||
```
|
||||
sk-llmgw-<base62(32 random bytes)>
|
||||
```
|
||||
|
||||
Only `sha256(key)` is stored, indexed. Authentication hashes what the client
|
||||
presented and looks it up. Nothing can reconstruct the key from the database,
|
||||
so a database leak does not hand over working credentials.
|
||||
|
||||
**SHA-256 rather than bcrypt or argon2, deliberately.** Slow password hashing
|
||||
exists because humans pick guessable passwords, and the cost buys time against
|
||||
an offline dictionary attack. A 256-bit random value has no dictionary to
|
||||
attack: at any hash rate, guessing it is infeasible. What a slow hash would buy
|
||||
here is nothing, and what it would cost is tens of milliseconds on every single
|
||||
request. Using the fast hash is the correct choice precisely *because* the input
|
||||
is high entropy — and that is the assumption to re-examine if the key format
|
||||
ever changes.
|
||||
|
||||
Alongside the hash, a `key_prefix` (`sk-llmgw-4f2a…`) is stored in clear so the
|
||||
UI and audit log can identify a key without revealing it.
|
||||
|
||||
```sql
|
||||
CREATE TABLE api_keys (
|
||||
id TEXT PRIMARY KEY,
|
||||
key_hash BLOB NOT NULL UNIQUE, -- sha256 of the key; the only copy
|
||||
key_prefix TEXT NOT NULL, -- for display, never for auth
|
||||
name TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL REFERENCES projects(id),
|
||||
band TEXT, -- §7.2.2; NULL inherits the project's
|
||||
routing_profile TEXT, -- §7.3
|
||||
created_at INTEGER NOT NULL,
|
||||
expires_at INTEGER, -- NULL never expires
|
||||
revoked_at INTEGER, -- soft delete, so usage rows keep a valid FK
|
||||
last_used_at INTEGER -- written lazily; see below
|
||||
);
|
||||
CREATE INDEX api_keys_by_project ON api_keys(project_id);
|
||||
```
|
||||
|
||||
Revocation is a soft delete. Usage rows reference the key, and an audit trail
|
||||
that loses its subject when the key is deleted is not an audit trail.
|
||||
|
||||
`last_used_at` is written at most once a minute per key, from the same batch
|
||||
writer as usage. Writing it per request would put a database write back on the
|
||||
request path for a field nobody reads in real time.
|
||||
|
||||
#### 7.9.2 The tenancy tree and how limits resolve
|
||||
|
||||
```sql
|
||||
CREATE TABLE organizations (id TEXT PRIMARY KEY, name TEXT NOT NULL);
|
||||
CREATE TABLE teams (id TEXT PRIMARY KEY, org_id TEXT NOT NULL REFERENCES organizations(id), name TEXT NOT NULL);
|
||||
CREATE TABLE projects (id TEXT PRIMARY KEY, team_id TEXT NOT NULL REFERENCES teams(id), name TEXT NOT NULL);
|
||||
|
||||
-- Limits hang off any level rather than being duplicated onto each one.
|
||||
CREATE TABLE limits (
|
||||
scope_type TEXT NOT NULL, -- 'org' | 'team' | 'project' | 'key'
|
||||
scope_id TEXT NOT NULL,
|
||||
rpm INTEGER, -- NULL means no limit at this level
|
||||
tpm INTEGER,
|
||||
budget_units REAL, -- per period; NULL means no budget
|
||||
budget_period TEXT, -- 'day' | 'month'
|
||||
allowed_models TEXT, -- JSON array; NULL means all
|
||||
PRIMARY KEY (scope_type, scope_id)
|
||||
);
|
||||
```
|
||||
|
||||
A request's effective limit is the **most restrictive** value along
|
||||
key → project → team → org. A team cannot raise a limit its organization set,
|
||||
which is the only composition rule that makes a budget mean anything.
|
||||
|
||||
Resolving four levels per request would be four lookups on the hot path, so the
|
||||
resolved view is computed once and cached per key hash, invalidated whenever
|
||||
anything on that key's chain changes. Authentication is then a single map
|
||||
lookup, with the database consulted only on a cache miss.
|
||||
|
||||
#### 7.9.3 How accounting is written
|
||||
|
||||
The volume decides the design. At 1000 requests/sec, one row per request is
|
||||
86 million rows a day, and SQLite has a single writer. Inserting per request
|
||||
would serialize every request behind that writer.
|
||||
|
||||
So the path is:
|
||||
|
||||
```
|
||||
request completes
|
||||
→ usage record onto a buffered channel (nanoseconds, never blocks)
|
||||
→ batch writer: every 200ms or 500 records, one transaction
|
||||
→ hourly rollup job collapses raw rows
|
||||
→ raw rows older than the retention window are deleted
|
||||
```
|
||||
|
||||
```sql
|
||||
CREATE TABLE usage (
|
||||
ts INTEGER NOT NULL, -- unix millis
|
||||
key_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
served_model TEXT NOT NULL, -- differs from model when a fallback served it
|
||||
pool TEXT,
|
||||
status INTEGER NOT NULL,
|
||||
prompt_tokens INTEGER NOT NULL,
|
||||
completion_tokens INTEGER NOT NULL,
|
||||
cached_tokens INTEGER NOT NULL, -- prefix cache hit, from the engine
|
||||
ttft_ms INTEGER,
|
||||
duration_ms INTEGER NOT NULL,
|
||||
queue_ms INTEGER -- time spent in admission (§7.2)
|
||||
);
|
||||
CREATE INDEX usage_by_key_ts ON usage(key_id, ts);
|
||||
|
||||
-- What survives retention, and what budgets are enforced against.
|
||||
CREATE TABLE usage_rollup (
|
||||
hour_ts INTEGER NOT NULL,
|
||||
key_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
requests INTEGER NOT NULL,
|
||||
prompt_tokens INTEGER NOT NULL,
|
||||
completion_tokens INTEGER NOT NULL,
|
||||
cached_tokens INTEGER NOT NULL,
|
||||
units REAL NOT NULL, -- billable cost, per §7.9.4
|
||||
PRIMARY KEY (hour_ts, key_id, model)
|
||||
);
|
||||
```
|
||||
|
||||
**The honest tradeoff:** records sitting in the channel when the process is
|
||||
killed are lost — at most a few hundred milliseconds of accounting. Making that
|
||||
impossible means a durable write before each response, which puts a disk fsync
|
||||
on every request. For usage that informs budgets and invoices, losing a fraction
|
||||
of a second of records to a hard kill is the better trade. **If this ever has to
|
||||
feed external billing, that trade has to be revisited**, and the answer then is
|
||||
an append-only journal the batch writer drains, not a synchronous insert.
|
||||
|
||||
#### 7.9.4 How budgets are enforced without reading the database
|
||||
|
||||
Budgets are enforced per request, so they cannot involve a query. Instead:
|
||||
|
||||
- On startup, the current period's spend is loaded per key from `usage_rollup`
|
||||
plus the un-rolled tail of `usage`. **A restart must not reset a budget**, which
|
||||
it would if the counter lived only in memory.
|
||||
- From then on an in-memory counter is incremented as responses complete. The
|
||||
database catches up through the batch writer.
|
||||
- At the period boundary the counter resets and the new period starts empty.
|
||||
|
||||
"Units" are deliberately not dollars. Self-hosted models have no per-token
|
||||
price, so a unit is whatever the operator configures — GPU-seconds, a notional
|
||||
rate per million tokens, or 1.0 for everything, which reduces budgets to token
|
||||
quotas. This is also what makes the plan's budget fallback work: a model priced
|
||||
at zero stays available to a key whose budget is spent (§7.7).
|
||||
|
||||
#### 7.9.5 Token counting
|
||||
|
||||
Two numbers exist and they are not the same:
|
||||
|
||||
- **The estimate**, computed before dispatch as `bytes/4`, used to reserve TPM
|
||||
quota and to score prefix affinity. It is a guess and is labelled as one.
|
||||
- **The engine's report**, taken from the response's `usage` block. This is
|
||||
authoritative and is what gets recorded.
|
||||
|
||||
The router asks engines for usage even when the client did not, then drops that
|
||||
chunk on the way back (§7.8), so accounting does not depend on client behaviour.
|
||||
TPM is reserved at request start as `max_tokens + estimate` and reconciled
|
||||
against the real figure on completion — reserving optimistically would let a
|
||||
key exceed its limit by a whole response.
|
||||
|
||||
The estimate stops being a guess only when a real tokenizer arrives, which the
|
||||
plan defers to v2 along with tier 2 prefix routing (§7.5, §14).
|
||||
|
||||
#### 7.9.6 Multiple router instances
|
||||
|
||||
Everything above assumes one instance. Two would each hold their own rate limit
|
||||
counters and each allow the configured limit, so the effective limit doubles.
|
||||
|
||||
v1 is single-instance and this is a documented constraint rather than a bug. The
|
||||
answer when it arrives is either shared counters in Redis, or consistent hashing
|
||||
on the key so every request for one key reaches the same instance. The second is
|
||||
cheaper and degrades better, but rebalances counters whenever the instance set
|
||||
changes.
|
||||
|
||||
---
|
||||
|
||||
## 8. Observability
|
||||
@@ -953,8 +1223,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
|
||||
@@ -973,27 +1243,89 @@ work_dir: /var/lib/llmgw-agent
|
||||
move on before writing this test — it is the entire reason phase 1.5 exists.**
|
||||
|
||||
### Phase 2 — Agent: discovery and single-node engines (2 weeks)
|
||||
- [ ] GPU discovery (NVML plus fallback), IB discovery, topology
|
||||
- [ ] Engine interface, single-node vLLM and SGLang
|
||||
- [ ] Supervisor (process groups, log buffer, state machine, adoption)
|
||||
- [ ] connect-go stream, register/heartbeat, mTLS
|
||||
- **Exit criterion:** start and stop a model on a real node via `ctl`.
|
||||
- [x] GPU discovery (nvidia-smi path), IB discovery from sysfs
|
||||
- [x] Engine interface, vLLM and SGLang command and environment construction
|
||||
- [x] Supervisor (process groups, log buffer, state machine, orphan adoption)
|
||||
- [x] NVML backend behind the `nvml` build tag, with NVLink topology. It compiles
|
||||
in CI but has not run against a driver: **verify on a GPU node before
|
||||
relying on the NVLink domains it reports.**
|
||||
- [x] connect-go stream, register/heartbeat, shard status
|
||||
- [ ] mTLS — the stream is plaintext HTTP/2 today and must stay on a trusted network
|
||||
- **Exit criterion (met, via configuration rather than `ctl`):** a model starts
|
||||
and stops on a node. Demonstrated with real binaries: `llmgw server` and
|
||||
`llmgw agent` as separate processes, the agent registering, a replica being
|
||||
placed, an engine launched through the vLLM adapter, traffic served, and the
|
||||
node's loss detected when the agent is killed. `ctl` itself is phase 3's
|
||||
remaining item.
|
||||
|
||||
### Phase 3 — Control plane (2–3 weeks)
|
||||
- [ ] SQLite store and migrations
|
||||
- [ ] Catalog, inventory, scheduler (single-node profiles), reconciler
|
||||
- [x] `Register` carries the agent's running shards, so a control plane restart
|
||||
rebuilds replica state from the fleet instead of orphaning live engines (§6.5)
|
||||
- [x] Orphan adoption: an agent takes back engines a previous agent left running,
|
||||
and kills whatever it cannot identify or the catalog will not claim
|
||||
- [x] Inventory: nodes, hardware, fabric attachments, allocations
|
||||
- [x] Placement: profile priority, bin packing, NVLink-aware GPU selection,
|
||||
`explain-placement` rejection reasons, property-based invariant tests
|
||||
- [x] Catalog: models, profiles, alias and fallback resolution, validation
|
||||
- [x] Reconciler: desired vs actual, replica lifecycle, node expiry, backoff
|
||||
- [x] `EndpointSource` wired from the reconciler to the router
|
||||
- [x] `ctl`: nodes, models, replicas, cordon, uncordon, drain, explain-placement
|
||||
- [ ] `ctl logs`: needs a request/response pair on the agent stream, which the
|
||||
protocol does not have. The failure path already carries an engine log
|
||||
tail with every status report, so the common case is covered; what is
|
||||
missing is reading a healthy shard's output on demand.
|
||||
- [ ] The router's `EndpointSource` becomes dynamic
|
||||
- [ ] `ctl`: nodes, models, replicas, cordon, drain, explain-placement, logs
|
||||
- **Exit criterion:** shutting down a node moves its replicas elsewhere while the router keeps serving.
|
||||
- [x] `ctl`: nodes, models, replicas, cordon, uncordon, drain, explain-placement
|
||||
- [ ] `ctl logs`: needs a request/response pair on the agent stream, which the
|
||||
protocol does not have. The failure path already carries an engine log
|
||||
tail with every status report, so the common case is covered; what is
|
||||
missing is reading a healthy shard's output on demand.
|
||||
- **Exit criterion (met against a simulated agent fleet):** shutting down a node
|
||||
moves its replicas elsewhere while the router keeps serving. Measured in
|
||||
`test/e2e`: a node hosting a live replica is killed mid-traffic; the replica is
|
||||
rebuilt on a survivor and ~370 requests complete with zero failures. The real
|
||||
agent path still needs the stream, so this is proof of the control logic, not
|
||||
of the transport.
|
||||
|
||||
### Phase 4 — Multi-node and fabrics (2 weeks)
|
||||
- [ ] Fabric identity (config plus GID prefix), verification
|
||||
- [ ] Multi-node placement (the fabric constraint), port allocation, rank-ordered startup
|
||||
- [ ] NCCL environment generation, HCA selection
|
||||
- [ ] Replica-level failure handling (one dead shard kills the replica)
|
||||
- [x] Fabric identity from the agent's explicit labels
|
||||
- [x] Multi-node placement (the fabric constraint), port allocation, rank ordering
|
||||
- [x] NCCL environment generation, HCA selection
|
||||
- [x] Replica-level failure handling (one dead shard kills the replica)
|
||||
|
||||
The launch path is now built too:
|
||||
|
||||
- [x] **Only the leader is health-checked.** `Engine.ServesHTTP(spec)` says which
|
||||
ranks answer the API; a worker is judged by staying alive, and the
|
||||
replica's real readiness is the leader answering, which only happens once
|
||||
every worker has joined. Health-checking a worker's port used to time out
|
||||
and fail a replica that was coming up perfectly well.
|
||||
- [x] **A shard can be several processes.** `Engine.Commands(spec)` returns them
|
||||
in start order: every command but the last is a sidecar kept alive for the
|
||||
shard's lifetime and torn down with it, and the last one's exit is the
|
||||
shard's. vLLM under Ray is the reason — the leader runs a Ray head and
|
||||
then the server, a worker runs only the join and no vLLM process at all.
|
||||
- [x] **The leader starts first.** Both engines put the rendezvous point on rank
|
||||
0, so a worker started first has nothing to connect to. An earlier version
|
||||
had this backwards; vLLM's server blocks until the Ray cluster has the
|
||||
GPUs it needs, so the leader can start immediately and waits on its own.
|
||||
- [x] **Rendezvous is waited for, not assumed.** A command can require that
|
||||
`LeaderAddr:LeaderPort` accepts a connection before the next one starts,
|
||||
bounded by `rendezvous_timeout` so a head that never appears fails the
|
||||
shard instead of hanging it.
|
||||
|
||||
- [ ] **Not verified against real engines.** Everything above is tested with
|
||||
stub processes. The Ray sequencing, and whether the deployed vLLM version
|
||||
wants `ray` or `mp`, can only be confirmed on real hardware.
|
||||
|
||||
- **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.
|
||||
- **Exit criterion:** a 2-node replica on each fabric; killing a node rebuilds the replica on the same fabric and never spreads it to the other one.
|
||||
|
||||
### Phase 5 — Production features (2–3 weeks)
|
||||
### Phase 5 — Production features (3 weeks)
|
||||
- [ ] **SQLite store and migrations.** It arrives here rather than in phase 3
|
||||
because everything that genuinely needs persisting — keys, usage, capacity
|
||||
profiles — is a phase 5 feature. Writing it earlier would have meant
|
||||
designing a schema against guesses (§6.5).
|
||||
- [ ] Tenancy: org/team/project/key hierarchy, budgets, RPM/TPM, usage records (§7.9)
|
||||
- [ ] Budget fallback, model allowlists, key→band mapping
|
||||
- [ ] Routing profiles from configuration plus the `x-llmgw-profile` header
|
||||
@@ -1060,6 +1392,9 @@ work_dir: /var/lib/llmgw-agent
|
||||
| **Tuning the admission thresholds** | Too low a `target_waiting_per_replica` starves the GPU; too high brings queueing delay back. This single number governs the system's behaviour, so measure it with `ctl bench` rather than guessing. |
|
||||
| KV event loss | ZMQ is fire-and-forget; a lost event makes the router silently misroute. Sequence-gap detection and falling back to tier 1 are mandatory (§7.5). |
|
||||
| Metric staleness | At 1000 requests/sec with a 2s scrape interval, decisions rest entirely on stale data. That is why the `inflight` scorer is not optional. |
|
||||
| **Usage lost on a hard kill** | Records buffered for the batch writer are lost if the process is killed — at most a few hundred milliseconds of accounting. Acceptable for budgets; **not acceptable if this ever feeds external billing**, and the answer then is an append-only journal rather than a synchronous insert (§7.9.3). |
|
||||
| **Rate limits with several routers** | Counters are per instance, so N instances allow N times the configured limit. v1 is single-instance by design; Redis or consistent-hash stickiness is the fix (§7.9.6). |
|
||||
| **Control plane restart** | Until `Register` carries running shards, restarting the control plane orphans live engines: they keep holding GPUs the inventory believes are free, and the next placement fails to allocate with an error pointing at the wrong process (§6.5). |
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -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,242 @@
|
||||
// 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
|
||||
}
|
||||
|
||||
// engineOpt tunes the sim-engines a harness starts.
|
||||
type engineOpt func(*sim.Config)
|
||||
|
||||
// engineCapacity is the batch capacity used by withLoadedEngines, chosen to sit
|
||||
// in the range a real vLLM deployment runs at. Making it much smaller would
|
||||
// shift the bottleneck to slot starvation inside the engine, which admission
|
||||
// cannot do anything about — the test would then be measuring the wrong thing.
|
||||
const engineCapacity = 64
|
||||
|
||||
// withLoadedEngines sizes the engines so that a few hundred concurrent requests
|
||||
// genuinely overflow them, which is what the priority isolation test needs.
|
||||
func withLoadedEngines() func(*config.Server) {
|
||||
return nil // marker; see start()
|
||||
}
|
||||
|
||||
// start brings up n sim-engines behind a router.
|
||||
//
|
||||
// A nil tune function is the "small engines" marker, which keeps the call sites
|
||||
// readable without a second variadic parameter.
|
||||
func start(t *testing.T, n int, tune ...func(*config.Server)) *harness {
|
||||
t.Helper()
|
||||
|
||||
small := false
|
||||
kept := tune[:0]
|
||||
for _, f := range tune {
|
||||
if f == nil {
|
||||
small = true
|
||||
continue
|
||||
}
|
||||
kept = append(kept, f)
|
||||
}
|
||||
tune = kept
|
||||
|
||||
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 {
|
||||
sc := simConfig(fmt.Sprintf("sim-%d", i), int64(i+1))
|
||||
if small {
|
||||
sc.MaxNumSeqs = engineCapacity
|
||||
sc.DecodeTokensPerSec = 800
|
||||
sc.BatchContention = 0.005
|
||||
}
|
||||
e := sim.New(sc)
|
||||
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")
|
||||
}
|
||||
|
||||
// running counts requests actively generating across all engines.
|
||||
func (h *harness) running() int {
|
||||
n := 0
|
||||
for _, e := range h.engines {
|
||||
n += e.eng.Stats().Running
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
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,203 @@
|
||||
package e2e
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
)
|
||||
|
||||
// TestPriorityIsolation is the phase 1.5 exit criterion.
|
||||
//
|
||||
// While a sustained flood of batch requests runs in the bronze band,
|
||||
// interactive requests in the platinum band must keep their latency. This is
|
||||
// the whole reason admission exists: an engine queue is FIFO, so once a batch
|
||||
// job is in it no amount of clever endpoint selection can protect anyone.
|
||||
//
|
||||
// It is meant to stay in CI permanently.
|
||||
func TestPriorityIsolation(t *testing.T) {
|
||||
workers := envInt("LLMGW_E2E_BATCH", 500)
|
||||
probes := envInt("LLMGW_E2E_PROBES", 40)
|
||||
if testing.Short() {
|
||||
workers, probes = 80, 15
|
||||
}
|
||||
|
||||
h := start(t, 3, func(c *config.Server) {
|
||||
c.Router.Admission.Enabled = true
|
||||
}, withLoadedEngines())
|
||||
|
||||
idle := measureProbes(t, h, probes, "platinum")
|
||||
t.Logf("idle platinum: p50=%v p99=%v", idle.p50.Round(time.Millisecond), idle.p99.Round(time.Millisecond))
|
||||
|
||||
load := startBronzeLoad(t, h, workers, "bronze")
|
||||
defer load.stop()
|
||||
|
||||
loaded := measureProbes(t, h, probes, "platinum")
|
||||
t.Logf("loaded platinum: p50=%v p99=%v (%s)",
|
||||
loaded.p50.Round(time.Millisecond), loaded.p99.Round(time.Millisecond), load.summary())
|
||||
|
||||
// The criterion: interactive p99 under load stays within 1.5x of idle.
|
||||
// The floor keeps a few-millisecond baseline from making the bound absurd.
|
||||
budget := max(time.Duration(float64(idle.p99)*1.5), idle.p99+50*time.Millisecond)
|
||||
t.Logf("platinum p99 %v against a %v budget (%.2fx idle)",
|
||||
loaded.p99.Round(time.Millisecond), budget.Round(time.Millisecond),
|
||||
float64(loaded.p99)/float64(max(idle.p99, time.Microsecond)))
|
||||
|
||||
if loaded.p99 > budget {
|
||||
t.Errorf("interactive p99 rose to %v under bronze load, over the %v budget; "+
|
||||
"admission control is not protecting the high band",
|
||||
loaded.p99.Round(time.Millisecond), budget.Round(time.Millisecond))
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdmissionDisabledShowsTheProblem documents why the layer exists: the same
|
||||
// load does wreck interactive latency once admission is switched off. If this
|
||||
// ever stops showing degradation, the load generator has stopped being a load
|
||||
// and the test above proves less than it appears to.
|
||||
func TestAdmissionDisabledShowsTheProblem(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("load-dependent")
|
||||
}
|
||||
workers := envInt("LLMGW_E2E_BATCH", 500)
|
||||
|
||||
h := start(t, 3, func(c *config.Server) {
|
||||
c.Router.Admission.Enabled = false
|
||||
}, withLoadedEngines())
|
||||
|
||||
idle := measureProbes(t, h, 20, "")
|
||||
|
||||
load := startBronzeLoad(t, h, workers, "")
|
||||
defer load.stop()
|
||||
|
||||
loaded := measureProbes(t, h, 20, "")
|
||||
ratio := float64(loaded.p99) / float64(max(idle.p99, time.Microsecond))
|
||||
t.Logf("without admission: idle p99=%v loaded p99=%v (%.1fx)",
|
||||
idle.p99.Round(time.Millisecond), loaded.p99.Round(time.Millisecond), ratio)
|
||||
|
||||
if ratio < 2 {
|
||||
t.Skipf("the load did not degrade latency (%.1fx); this machine is too fast "+
|
||||
"for %d workers, so the isolation test proves less than it should", ratio, workers)
|
||||
}
|
||||
}
|
||||
|
||||
// --- load generation ---
|
||||
|
||||
// bronzeLoad is a sustained batch workload.
|
||||
//
|
||||
// Workers loop rather than firing once: a one-shot burst drains while the
|
||||
// probes are still running, and the measurement then catches a recovering
|
||||
// cluster instead of a loaded one.
|
||||
type bronzeLoad struct {
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
done atomic.Int64
|
||||
failed atomic.Int64
|
||||
firstEr atomic.Value
|
||||
}
|
||||
|
||||
func startBronzeLoad(t *testing.T, h *harness, workers int, band string) *bronzeLoad {
|
||||
t.Helper()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
l := &bronzeLoad{cancel: cancel}
|
||||
|
||||
for i := range workers {
|
||||
l.wg.Add(1)
|
||||
go func() {
|
||||
defer l.wg.Done()
|
||||
hdr := map[string]string{"x-session-id": fmt.Sprintf("batch-%d", i%8)}
|
||||
if band != "" {
|
||||
hdr["x-llmgw-band"] = band
|
||||
}
|
||||
for ctx.Err() == nil {
|
||||
_, _, err := streamRequest(ctx, h.router.URL,
|
||||
chatBody(fmt.Sprintf("batch job %d", i), 200, true), hdr)
|
||||
switch {
|
||||
case err == nil:
|
||||
l.done.Add(1)
|
||||
case ctx.Err() != nil: // shutting down, expected
|
||||
default:
|
||||
l.failed.Add(1)
|
||||
l.firstEr.CompareAndSwap(nil, err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
l.waitUntilBusy(t, h)
|
||||
return l
|
||||
}
|
||||
|
||||
// waitUntilBusy blocks until the engines are demonstrably working, so probes
|
||||
// are not taken while the load is still ramping up.
|
||||
//
|
||||
// It waits on engines being busy rather than full: when admission is doing its
|
||||
// job the pool is deliberately held below capacity, so a fullness check would
|
||||
// time out exactly when the system is behaving correctly.
|
||||
func (l *bronzeLoad) waitUntilBusy(t *testing.T, h *harness) {
|
||||
t.Helper()
|
||||
const busyFor = 500 * time.Millisecond
|
||||
|
||||
deadline := time.Now().Add(20 * time.Second)
|
||||
var since time.Time
|
||||
for time.Now().Before(deadline) {
|
||||
if h.running() > 0 {
|
||||
if since.IsZero() {
|
||||
since = time.Now()
|
||||
} else if time.Since(since) >= busyFor {
|
||||
return
|
||||
}
|
||||
} else {
|
||||
since = time.Time{}
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("engines never became busy; %s", l.summary())
|
||||
}
|
||||
|
||||
func (l *bronzeLoad) summary() string {
|
||||
return fmt.Sprintf("bronze: completed=%d failed=%d firstErr=%v",
|
||||
l.done.Load(), l.failed.Load(), l.firstEr.Load())
|
||||
}
|
||||
|
||||
func (l *bronzeLoad) stop() {
|
||||
l.cancel()
|
||||
l.wg.Wait()
|
||||
}
|
||||
|
||||
// --- probes ---
|
||||
|
||||
type latency struct{ p50, p99 time.Duration }
|
||||
|
||||
// measureProbes sends sequential short requests and reports their latency.
|
||||
// Sequential on purpose: these stand in for a single interactive user, and
|
||||
// concurrency would measure something else.
|
||||
func measureProbes(t *testing.T, h *harness, n int, band string) latency {
|
||||
t.Helper()
|
||||
hdr := map[string]string{"x-session-id": "probe"}
|
||||
if band != "" {
|
||||
hdr["x-llmgw-band"] = band
|
||||
}
|
||||
|
||||
samples := make([]time.Duration, 0, n)
|
||||
for range n {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
start := time.Now()
|
||||
_, _, err := streamRequest(ctx, h.router.URL, chatBody("ping", 4, true), hdr)
|
||||
cancel()
|
||||
if err != nil {
|
||||
t.Fatalf("probe failed: %v", err)
|
||||
}
|
||||
samples = append(samples, time.Since(start))
|
||||
}
|
||||
slices.Sort(samples)
|
||||
return latency{
|
||||
p50: samples[len(samples)/2],
|
||||
p99: samples[min(int(float64(len(samples))*0.99), len(samples)-1)],
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,414 @@
|
||||
package e2e
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/sim"
|
||||
)
|
||||
|
||||
// simDispatcher is a stand-in for the agent fleet: instead of sending commands
|
||||
// over a stream it starts real sim-engines on this host, on exactly the ports
|
||||
// the reconciler allocated.
|
||||
//
|
||||
// That matters — the reconciler derives each replica's endpoint URL from the
|
||||
// node address and the port it chose, so an engine listening anywhere else
|
||||
// would make the test pass for the wrong reason.
|
||||
type simDispatcher struct {
|
||||
t *testing.T
|
||||
|
||||
mu sync.Mutex
|
||||
running map[string]*simShard
|
||||
// dead marks nodes whose agent is gone; commands to them fail.
|
||||
dead map[string]bool
|
||||
|
||||
onStatus func(replicaID, shardID string, st supervisor.State, errMsg string)
|
||||
}
|
||||
|
||||
type simShard struct {
|
||||
eng *sim.Engine
|
||||
srv *http.Server
|
||||
port int
|
||||
}
|
||||
|
||||
func newSimDispatcher(t *testing.T) *simDispatcher {
|
||||
d := &simDispatcher{t: t, running: map[string]*simShard{}, dead: map[string]bool{}}
|
||||
t.Cleanup(d.shutdown)
|
||||
return d
|
||||
}
|
||||
|
||||
func (d *simDispatcher) StartShard(_ context.Context, cmd reconciler.ShardCommand) error {
|
||||
d.mu.Lock()
|
||||
if d.dead[cmd.NodeID] {
|
||||
d.mu.Unlock()
|
||||
return fmt.Errorf("agent on %s is unreachable", cmd.NodeID)
|
||||
}
|
||||
d.mu.Unlock()
|
||||
|
||||
cfg := simConfig("", int64(cmd.Spec.HTTPPort))
|
||||
cfg.Model = cmd.Spec.Profile.ServedModelName
|
||||
eng := sim.New(cfg)
|
||||
|
||||
ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.Spec.HTTPPort))
|
||||
if err != nil {
|
||||
eng.Close()
|
||||
return fmt.Errorf("binding port %d: %w", cmd.Spec.HTTPPort, err)
|
||||
}
|
||||
srv := &http.Server{Handler: sim.NewServer(eng, slog.New(slog.DiscardHandler)), ReadHeaderTimeout: time.Second}
|
||||
go func() { _ = srv.Serve(ln) }()
|
||||
|
||||
d.mu.Lock()
|
||||
d.running[cmd.ShardID] = &simShard{eng: eng, srv: srv, port: cmd.Spec.HTTPPort}
|
||||
d.mu.Unlock()
|
||||
|
||||
// A real agent reports Ready only once the engine serves the model; the
|
||||
// sim-engine is ready as soon as it is listening.
|
||||
go d.onStatus(cmd.ReplicaID, cmd.ShardID, supervisor.Ready, "")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *simDispatcher) StopShard(_ context.Context, _, shardID string) error {
|
||||
d.mu.Lock()
|
||||
sh, ok := d.running[shardID]
|
||||
delete(d.running, shardID)
|
||||
d.mu.Unlock()
|
||||
if ok {
|
||||
sh.close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *simDispatcher) DrainShard(context.Context, string, string) error { return nil }
|
||||
|
||||
// killNode stops everything on a node and makes further commands to it fail,
|
||||
// which is what losing a machine looks like from the control plane.
|
||||
func (d *simDispatcher) killNode(nodeID string, shardIDs []string) {
|
||||
d.mu.Lock()
|
||||
d.dead[nodeID] = true
|
||||
var doomed []*simShard
|
||||
for _, id := range shardIDs {
|
||||
if sh, ok := d.running[id]; ok {
|
||||
doomed = append(doomed, sh)
|
||||
delete(d.running, id)
|
||||
}
|
||||
}
|
||||
d.mu.Unlock()
|
||||
|
||||
for _, sh := range doomed {
|
||||
sh.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (d *simDispatcher) liveShards() int {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return len(d.running)
|
||||
}
|
||||
|
||||
func (d *simDispatcher) shutdown() {
|
||||
d.mu.Lock()
|
||||
all := make([]*simShard, 0, len(d.running))
|
||||
for _, sh := range d.running {
|
||||
all = append(all, sh)
|
||||
}
|
||||
d.running = map[string]*simShard{}
|
||||
d.mu.Unlock()
|
||||
for _, sh := range all {
|
||||
sh.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *simShard) close() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
_ = s.srv.Shutdown(ctx)
|
||||
s.eng.Close()
|
||||
}
|
||||
|
||||
// controlHarness is a router fed by a real control plane.
|
||||
type controlHarness struct {
|
||||
router *httptest.Server
|
||||
rt *router.Router
|
||||
ctl *control.Controller
|
||||
disp *simDispatcher
|
||||
inv *inventory.Store
|
||||
}
|
||||
|
||||
func startControlPlane(t *testing.T, nodes int, minReplicas int) *controlHarness {
|
||||
t.Helper()
|
||||
|
||||
cat, err := catalog.New([]catalog.Model{{
|
||||
Name: testModel, Source: "/models/test",
|
||||
Scaling: catalog.Scaling{Min: minReplicas, Max: minReplicas + 2},
|
||||
Profiles: []catalog.Profile{{Name: "single", GPUClass: "h100-80g", TP: 1, Weight: 1}},
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cfg := config.Default()
|
||||
cfg.Router.ScrapeInterval = 200 * time.Millisecond
|
||||
cfg.Router.StaleAfter = 3 * time.Second
|
||||
cfg.Router.Admission.Enabled = false // this test is about endpoint churn
|
||||
|
||||
rt, err := router.New(cfg, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
disp := newSimDispatcher(t)
|
||||
|
||||
ccfg := control.DefaultConfig()
|
||||
ccfg.Reconciler.Interval = 100 * time.Millisecond
|
||||
ccfg.Reconciler.BaseBackoff = 200 * time.Millisecond
|
||||
ccfg.Reconciler.MaxBackoff = 500 * time.Millisecond
|
||||
ccfg.Reconciler.UnreachableAfter = 300 * time.Millisecond
|
||||
ccfg.Reconciler.LostAfter = 600 * time.Millisecond
|
||||
ccfg.PublishInterval = 100 * time.Millisecond
|
||||
|
||||
ctl := control.New(ccfg, cat, disp, rt, slog.New(slog.DiscardHandler))
|
||||
disp.onStatus = ctl.Reconciler().OnShardStatus
|
||||
|
||||
// Every node is this host, so the ports the reconciler allocates are the
|
||||
// ports the sim-engines actually bind.
|
||||
for i := range nodes {
|
||||
id := fmt.Sprintf("n%d", i)
|
||||
ctl.Inventory().Upsert(inventory.Node{
|
||||
ID: id, MgmtAddr: "127.0.0.1", State: inventory.Ready,
|
||||
GPUs: []discovery.GPU{{
|
||||
Index: 0, Name: "NVIDIA H100 80GB HBM3",
|
||||
MemoryTotalMB: 81559, ComputeCapability: "9.0",
|
||||
}},
|
||||
})
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
rt.Start(ctx)
|
||||
ctl.Start(ctx)
|
||||
|
||||
h := &controlHarness{
|
||||
router: httptest.NewServer(rt.Handler()),
|
||||
rt: rt, ctl: ctl, disp: disp, inv: ctl.Inventory(),
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
h.router.Close()
|
||||
ctl.Close()
|
||||
cancel()
|
||||
rt.Stop()
|
||||
})
|
||||
return h
|
||||
}
|
||||
|
||||
// heartbeat keeps the named nodes alive.
|
||||
func (h *controlHarness) heartbeat(ids ...string) {
|
||||
for _, id := range ids {
|
||||
h.inv.Touch(id)
|
||||
}
|
||||
}
|
||||
|
||||
// waitForEndpoints blocks until the router sees n endpoints.
|
||||
func (h *controlHarness) waitForEndpoints(t *testing.T, n int, limit time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(limit)
|
||||
for time.Now().Before(deadline) {
|
||||
if len(h.rt.Registry().Snapshots(testModel)) == n {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("router has %d endpoints, expected %d",
|
||||
len(h.rt.Registry().Snapshots(testModel)), n)
|
||||
}
|
||||
|
||||
// TestControlPlaneFeedsRouter checks the wiring: replicas the reconciler brings
|
||||
// up become endpoints the router can route to, without any static config.
|
||||
func TestControlPlaneFeedsRouter(t *testing.T) {
|
||||
h := startControlPlane(t, 2, 2)
|
||||
|
||||
// Keep the nodes alive while the reconciler works.
|
||||
stop := make(chan struct{})
|
||||
go func() {
|
||||
tk := time.NewTicker(50 * time.Millisecond)
|
||||
defer tk.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-tk.C:
|
||||
h.heartbeat("n0", "n1")
|
||||
}
|
||||
}
|
||||
}()
|
||||
defer close(stop)
|
||||
|
||||
h.waitForEndpoints(t, 2, 15*time.Second)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if _, _, err := streamRequest(ctx, h.router.URL, chatBody("hello", 4, true), nil); err != nil {
|
||||
t.Fatalf("the router could not serve from control-plane endpoints: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRouterKeepsServingWhileNodeIsLost is the other half of the phase 3 exit
|
||||
// criterion: a node disappears, its replica is rebuilt elsewhere, and traffic
|
||||
// keeps flowing throughout.
|
||||
func TestRouterKeepsServingWhileNodeIsLost(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("timing-dependent")
|
||||
}
|
||||
// Three nodes, two replicas: losing one still leaves somewhere to rebuild.
|
||||
h := startControlPlane(t, 3, 2)
|
||||
|
||||
alive := map[string]bool{"n0": true, "n1": true, "n2": true}
|
||||
var aliveMu sync.Mutex
|
||||
|
||||
stop := make(chan struct{})
|
||||
go func() {
|
||||
tk := time.NewTicker(50 * time.Millisecond)
|
||||
defer tk.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-tk.C:
|
||||
aliveMu.Lock()
|
||||
for id, ok := range alive {
|
||||
if ok {
|
||||
h.inv.Touch(id)
|
||||
}
|
||||
}
|
||||
aliveMu.Unlock()
|
||||
}
|
||||
}
|
||||
}()
|
||||
defer close(stop)
|
||||
|
||||
h.waitForEndpoints(t, 2, 15*time.Second)
|
||||
|
||||
// Continuous traffic for the whole experiment.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
var ok, failed atomic.Int64
|
||||
var firstErr atomic.Value
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for range 4 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for ctx.Err() == nil {
|
||||
rctx, rcancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
_, _, err := streamRequest(rctx, h.router.URL, chatBody("ping", 4, true), nil)
|
||||
rcancel()
|
||||
switch {
|
||||
case err == nil:
|
||||
ok.Add(1)
|
||||
case ctx.Err() != nil: // shutting down
|
||||
default:
|
||||
failed.Add(1)
|
||||
firstErr.CompareAndSwap(nil, err)
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Let traffic settle before disturbing anything.
|
||||
time.Sleep(600 * time.Millisecond)
|
||||
baseline := ok.Load()
|
||||
if baseline == 0 {
|
||||
cancel()
|
||||
wg.Wait()
|
||||
t.Fatalf("no requests succeeded before the failure: %v", firstErr.Load())
|
||||
}
|
||||
beforeFailures := failed.Load()
|
||||
|
||||
// Find a node hosting a replica and kill it.
|
||||
var victim string
|
||||
var victimShards []string
|
||||
for _, rep := range h.ctl.Reconciler().Replicas() {
|
||||
if rep.State != reconciler.Ready {
|
||||
continue
|
||||
}
|
||||
victim = rep.Shards[0].NodeID
|
||||
for _, s := range rep.Shards {
|
||||
victimShards = append(victimShards, s.ID)
|
||||
}
|
||||
break
|
||||
}
|
||||
if victim == "" {
|
||||
cancel()
|
||||
wg.Wait()
|
||||
t.Fatal("no ready replica to disturb")
|
||||
}
|
||||
|
||||
aliveMu.Lock()
|
||||
alive[victim] = false
|
||||
aliveMu.Unlock()
|
||||
h.disp.killNode(victim, victimShards)
|
||||
t.Logf("killed node %s (shards %v)", victim, victimShards)
|
||||
|
||||
// The control plane should notice, rebuild elsewhere, and get back to two
|
||||
// endpoints. The generous window covers unreachable -> lost -> backoff.
|
||||
deadline := time.Now().Add(25 * time.Second)
|
||||
recovered := false
|
||||
for time.Now().Before(deadline) {
|
||||
if len(h.rt.Registry().Snapshots(testModel)) == 2 && h.disp.liveShards() == 2 {
|
||||
recovered = true
|
||||
break
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
|
||||
// Let traffic run a little longer on the rebuilt topology.
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
cancel()
|
||||
wg.Wait()
|
||||
|
||||
t.Logf("requests ok=%d failed=%d (failures before the kill: %d)",
|
||||
ok.Load(), failed.Load(), beforeFailures)
|
||||
|
||||
if !recovered {
|
||||
var placed []string
|
||||
for _, rep := range h.ctl.Reconciler().Replicas() {
|
||||
placed = append(placed, fmt.Sprintf("%s=%s@%v", rep.ID, rep.State, rep.NodeIDs()))
|
||||
}
|
||||
t.Fatalf("the cluster did not return to 2 replicas: %v", placed)
|
||||
}
|
||||
|
||||
// The victim must not be hosting anything any more.
|
||||
for _, rep := range h.ctl.Reconciler().Replicas() {
|
||||
if rep.OnNode(victim) {
|
||||
t.Errorf("replica %s is still placed on the lost node %s", rep.ID, victim)
|
||||
}
|
||||
}
|
||||
|
||||
// Traffic kept flowing: more succeeded after the failure than before it.
|
||||
if ok.Load() <= baseline {
|
||||
t.Errorf("no requests succeeded after the node was lost (%d before, %d total)",
|
||||
baseline, ok.Load())
|
||||
}
|
||||
// Some failures are expected in the window where a replica is gone and its
|
||||
// replacement has not finished loading. What must not happen is total loss
|
||||
// of service, which is what having a second replica is for.
|
||||
if failed.Load() > ok.Load() {
|
||||
t.Errorf("more requests failed (%d) than succeeded (%d); service was not maintained: %v",
|
||||
failed.Load(), ok.Load(), firstErr.Load())
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user