Compare commits

...
13 Commits
Author SHA1 Message Date
ilkerandClaude Opus 5 e1eb1d4149 Phase 3: ctl, and make draining a rolling move
ctl talks to a management API 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 and read the cluster's shape.

- nodes surfaces unhealthy GPUs and why they are excluded. That is usually the
  answer when a node looks idle but nothing will schedule onto it, and it was
  invisible before.
- explain-placement prints every candidate considered and why it was rejected.
- replicas shows the shards a replica is made of, with per-shard errors.
- cordon, uncordon and drain are one endpoint with different targets, because
  the difference between cordon and drain is exactly what an operator needs to
  be sure of before running either.

Draining had to be fixed first to make ctl drain worth having. A drained replica
was stopped on the next pass, which cuts in-flight requests — the thing draining
exists to avoid. Now a node marked Draining has its replicas drained, the model
falls below its desired count, a replacement is placed elsewhere, and the
original is stopped only once the replacement is serving. A grace period bounds
it, because draining the only node that can host a model would otherwise never
finish and the node could never be taken out of service.

Verified against a two-node cluster of real binaries: draining moved a replica
to the surviving node and traffic kept returning 200 throughout.

ctl logs is deliberately absent. It needs a request/response pair the agent
protocol does not have, and the failure path already ships an engine log tail
with every status report — so the case that matters most is covered, and
reading a healthy shard's output on demand is not worth a protocol change yet.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 06:27:21 +02:00
ilkerandClaude Opus 5 1d1a6aebb0 Phase 4: make a multi-node replica able to start
Placement and environment generation had landed early, which made multi-node
look further along than it was. A replica spanning nodes could not come up at
all, and all three reasons came down to assumptions that only hold on one node.

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 now says which ranks serve, and a
worker is judged by staying alive. That is deliberately weaker evidence: 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 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. vLLM under Ray is why: the
leader runs a Ray head and then the server against it, while a worker runs only
the join and 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. Sidecars are killed with the shard, because a Ray head left
behind holds the port the next replica needs.

Start order was backwards. Workers went 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. vLLM's server blocks until the
Ray cluster has the GPUs it needs, so the leader can start immediately and waits
on its own. A command can now also require that rendezvous point to be listening
before the next starts, bounded by a timeout so a head that never appears fails
the shard instead of hanging it.

Stop was also killing only the engine, leaving sidecars behind.

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.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 06:20:03 +02:00
ilkerandClaude Opus 5 75a8e76706 Survive a control plane restart without orphaning live engines
Not persisting replica state left one hole: restarting the control plane made it
forget every replica while the engines kept running. They held 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 belongs on the agent, not in a database. A row can say "running" about a
process that died; the fleet cannot.

- the supervisor 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 an engine a previous agent left running. It is not this
  process's child, so it cannot be waited on: 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 reports what the agent is running, not only the first, because
  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 is killed. A process holding GPUs that
  nothing can describe is worse than no process, because placement hands those
  GPUs out again.

Stop was also wrong for adopted shards: it checked for an exec.Cmd, which they
do not have, and reported them stopped without killing anything.

TestAgentReconnects had been passing for no reason. httptest's
CloseClientConnections does nothing to an h2c stream — prior-knowledge h2c
hijacks the connection and httptest stops tracking hijacked ones — so the agent
never reconnected, and the assertion on the connection map passed on a stale
entry. Tests now sever connections through their own listener and count opened
streams.

Plan §11 also records what multi-node still needs, since placement landing early
makes it look further along than it is: workers never become Ready because the
agent health-checks every shard while only rank 0 serves HTTP; RayCommand is
written and tested but never called; and startup is fire-and-forget.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 06:15:04 +02:00
ilkerandClaude Opus 5 40427b2705 plan: correct the store design and specify how tenancy is actually stored
Two problems with what v0.1 said about persistence.

First, it listed eight tables, most of which were a mistake. A database that
describes the cluster can be wrong; the cluster cannot. An engine that dies
during a restart leaves a row saying "running", and the scheduler believes it.
Nodes, replicas, shards and the catalog are all rebuilt from the fleet or from
configuration on every reconnect, so they need no rows at all. What genuinely
needs persisting — API keys, usage, capacity profiles — is entirely phase 5
work, so the store moves there and gets a schema shaped by real requirements
instead of guesses.

Not persisting replica state has one consequence that must be closed first:
restarting the control plane currently 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 new process rather than the old one. The fix
belongs on the agent, not in a database — Register carries the shards the agent
is supervising, and the supervisor's orphan scan (written, tested, never called)
kills whatever nothing claims. That cannot disagree with reality, and it covers
agent restarts and control plane restarts with one mechanism. Added to phase 3.

Second, §7.9 described tenancy behaviour but never said how any of it is stored,
which left the most security-sensitive part of the system unspecified. It now
covers the schema for keys and the tenancy tree, why a fast hash is the correct
choice for a 256-bit random key and what assumption that rests on, how limits
resolve to the most restrictive value along the chain, why authentication is a
map lookup rather than four queries, the batching pipeline that keeps 86 million
rows a day off the request path, what "units" mean when self-hosted models have
no per-token price, and how budgets survive a restart.

Two tradeoffs are now stated rather than left implicit: usage buffered for the
batch writer is lost on a hard kill, which is fine for budgets and not fine for
external billing; and rate limit counters are per instance, so N routers allow N
times the limit. Both are in the risk table.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 06:08:16 +02:00
ilkerandClaude Opus 5 2373726efd Agent stream: connect the fleet to the control plane
One bidirectional connect-go stream per node, carrying registration, heartbeats
and shard status up, and start, stop and drain down. The agent dials out: 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 hardware 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.
- agentsvc is both the stream handler and the reconciler's dispatcher, because
  it holds the open streams and is therefore the only thing that can deliver a
  command to a node.
- the agent validates a profile before launching anything, since an engine that
  refuses a profile does so after minutes of loading, then supervises the
  process and reports every transition with 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 it over again. Engines keep serving throughout, because the router
  talks to them directly.

llmgw agent is real now, and llmgw server runs a control plane whenever any
model declares a profile. The static endpoint list stays for router-only runs.

Verified with real binaries as separate processes: the agent registers, a
replica is placed, an engine launches through the vLLM adapter's actual argument
construction, streaming traffic is served, and killing the agent moves the node
through unreachable to lost with placement explaining why nothing can run.

config.Model and catalog.Model were the same thing described twice, so the
router now reads catalog entries directly and the validation rules live in one
place.

Still missing: mTLS. The stream is plaintext HTTP/2 and must stay on a trusted
network until that lands.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 02:48:21 +02:00
ilkerandClaude Opus 5 c1b9ea1bc2 Phase 3 (3/3): wire the control plane to the router
Ready replicas become router endpoints with no static configuration involved.
Endpoints that survive an update keep their live state — 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 it becomes a watch stream
and nothing above it changes.

Phase 3's exit criterion now holds end to end. 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 with zero failures and no replica was
left 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.

The race detector caught a genuine bug while this went in. startReplica read the
replica's fields after publishing it to the map, and scaleUp walked replica
pointers outside the lock; both raced with agent status reports, which arrive
from another goroutine and mutate exactly those fields. Against the simulated
fleet the reports came back fast enough to hit the window on every run. State is
snapshotted under the lock now and those functions work from values.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 02:16:35 +02:00
ilkerandClaude Opus 5 f0134a2f95 Phase 3 (2/3): catalog and reconciler
The desired-versus-actual loop, with a Dispatcher interface where the agent
stream will go. The loop's logic is where the interesting failures live, not the
transport, so it is tested against a fake dispatcher and a synthetic clock.

- catalog: models, profiles, aliases, fallbacks, scaling bounds. Aliases and
  fallbacks resolve at load rather than at request time, because a typo in a
  fallback should stop a deployment instead of surfacing 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 rather than minutes into loading.
- reconciler: one replica placed per pass, since placement depends on what is
  already allocated and several from one snapshot would be handed the same GPUs.
  Workers start before rank 0: SGLang's ranks rendezvous and vLLM under Ray
  needs its workers attached before the head serves.
- a failed shard fails the whole replica. The others are waiting on collectives
  that will never complete while holding GPUs, so every shard is stopped and the
  replica rebuilt, with a jittered exponential backoff carried on the model
  rather than on the replica that was just destroyed.
- node expiry has two stages, Unreachable then Lost, deliberately far apart: 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.
- scale-down drains before stopping, so in-flight requests finish.

The phase 3 exit criterion is half met. TestLostNodeMovesReplicas covers the
reconciliation half end to end. The router continuing to serve throughout needs
the EndpointSource wiring and the agent stream.

Writing the tests caught a real property of the design worth recording: advancing
the clock makes every node look silent, so placement correctly stops. The test
harness now heartbeats nodes as it advances time, except where a node going
silent is the point.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 02:12:06 +02:00
ilkerandClaude Opus 5 06672d1b85 Phase 3 (1/3): cluster inventory and placement
The two hard constraints from plan §1.1 are enforced structurally rather than
checked after the fact: there is no code path that can combine nodes from two
fabrics, and only GPUs of the profile's class are ever considered. Getting
either wrong produces a replica that starts, passes health checks, and then
hangs or crawls under load — which looks like success until traffic arrives.

- inventory: nodes, hardware, fabric attachments, allocations. A re-registering
  agent cannot clear an operator's cordon, since a cordon is usually in place
  because something is wrong with the node. A GPU with an ECC fault or another
  process on it is not free in any useful sense. Replicas are released whole,
  because they fail whole.
- placement: highest-priority profile that fits, with every rejected candidate
  explained. "Unschedulable" on its own tells an operator nothing about which
  of six constraints bit.
- bin packing fills partly used nodes first so whole nodes stay free for large
  models that cannot be split; multi-node replicas get the emptiest machines
  instead, being the least movable thing in the cluster.
- NVLink-aware GPU selection keeps a tensor-parallel group inside one domain
  where the topology is known, falling back to lowest free indices otherwise.

Property-based tests generate 400 random clusters and assert the invariants on
every placement produced: one fabric per multi-node replica, only that fabric's
HCAs per shard, one GPU class, no unhealthy or already-allocated GPU. The
table-driven tests cover the cases someone thought of; these cover the rest.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 02:07:29 +02:00
ilkerandClaude Opus 5 55417546fc Phase 2 (2/2): NVML backend behind a build tag
NVML is the only source of NVLink topology, and that topology is what keeps a
tensor-parallel group off the slow path inside a node. It needs cgo, so it sits
behind the nvml tag and the default build keeps using nvidia-smi.

Discover() picks the backend and falls back to nvidia-smi when NVML fails,
logging the fallback rather than performing it silently: on a node that is meant
to have working NVML, a quiet downgrade hides a broken driver installation and
costs the scheduler its topology.

CI now compiles the tagged build. That immediately paid for itself: ErrNoNVML
was declared only in the untagged file, so the tests did not compile under
-tags nvml. It is now declared where both builds can see it.

Not verified: the NVML path has never run against a driver, because there is no
GPU machine available. The NVLink domains it reports should be treated as
unverified until someone runs `llmgw agent inventory` on a real node.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 02:05:25 +02:00
ilkerandClaude Opus 5 6a13e3dddf Phase 2 (1/2): node discovery, engine adapters, process supervisor
Everything here works without GPUs and is covered by tests. The control plane
connection is not built yet, so llmgw agent itself is still a stub and the phase
exit criterion is not met.

- discovery: IB from /sys/class/infiniband, parsing the formats sysfs actually
  uses and telling native InfiniBand from RoCE; GPUs via nvidia-smi, the
  cgo-free path, with health flags for uncorrected ECC and memory already held
  by another process. A node without IB is not an error, it is a single-node
  island.
- engines: vLLM and SGLang command and environment construction, kept free of
  process handling so it stays testable. The NCCL environment is the part the
  plan marks critical: NCCL_IB_HCA lists only the chosen fabric's adapters,
  because leaving it open lets NCCL pick an adapter on a fabric where the peer
  is unreachable, and the job then hangs instead of failing.
- validation before launch: no mixed GPU types in a tensor-parallel group, the
  node/GPU/TP arithmetic has to add up, and fp8 is refused below compute
  capability 8.9 rather than silently falling back to a slower path.
- supervisor: process groups, because engines fork workers that hold GPU memory
  and outlive a signal sent only to the leader; a bounded log tail, since engine
  output is the only diagnosis when loading fails; Starting separate from
  Loading, because a 70B model legitimately takes minutes; and orphan adoption
  through pid files fingerprinted with the command line, since 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 loading model.
- llmgw agent inventory reports what the scheduler would see on a node.

One plan correction: GPU class 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, not l40s-48g. Consistency matters more than matching
the box, because that number decides whether a model fits. Plan §2 and §5.1
updated.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 02:01:59 +02:00
ilkerandClaude Opus 5 dc38dd6096 Phase 1.5: admission control
Inserted between quota checks and scheduling: admission 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 endpoint selection rescues the request behind it.

- saturation detector: one number per pool, the worse of queue depth and KV
  occupancy, since either alone misleads.
- priority bands with per-band saturation ceilings, so the cluster degrades from
  the bottom up.
- tenant round-robin within a band: one caller sending a thousand requests takes
  one slot in the rotation, not a thousand.
- bounded queues with capacity (503) and TTL (504); cancelled requests leave the
  queue instead of holding capacity.
- metrics for queue depth, wait time, rejections and pool saturation.

Exit criterion, measured in test/e2e: with 500 concurrent bronze requests
running continuously, interactive p99 in the platinum band stays at 1.04-1.24x
idle against a 1.5x budget. With admission off the same load pushes it to
55-60x. Both directions are asserted, so the test fails if admission regresses
and also if the load generator stops being a real load.

Two lags had to be closed, and they needed different answers. Engine metrics are
scraped and lag an interval; router inflight does not, but it counts running
requests, and treating those as queue makes a healthy engine look overloaded
every interval. What is actually queued is inflight minus batch capacity, and
capacity is not in the metric contract, so it is learned as the high-water mark
of observed running. Separately, between admission granting a request and the
router dispatching it the request is invisible everywhere, so a simultaneous
burst passed the check en masse; the controller now counts its own outstanding
leases against dispatched inflight.

One plan correction: a saturation ceiling of 1.0 was documented as "never held
back". It is not — it means "flow up to the healthy limit". A band that must
never wait needs a ceiling above 1. Plan §7.2.2 updated.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 01:53:38 +02:00
ilkerandClaude Opus 5 a3f71aa5ee Phase 1: the router serves traffic
llmgw server is real now. It fronts a static set of engine replicas with an
OpenAI-compatible API, scoring endpoints through the pipeline added in the
previous commit.

- proxy: SSE relayed event by event with a flush after each, cancellation
  propagated upstream, and separate deadlines for dial, headers, first token,
  inter-token idle and total. The idle watchdog is what distinguishes a wedged
  engine from a slow one.
- router: dispatch with retry and model fallback. Retry only happens while
  nothing has reached the client; after that, splicing two responses together
  would be worse than the short stream the client already sees.
- config: YAML via koanf overlaid on defaults, with LLMGW_ env overrides.
  Validation rejects duplicate endpoint ids and aliases that shadow a model,
  because both misroute silently at request time rather than failing loudly.
- metrics: requests, duration, TTFT, tokens, plus per-scorer values and
  per-filter drops on a separate admin listener with health and pprof.

Phase 1 exit criterion, measured in test/e2e: 1000 concurrent streams across
3 sim-engines complete at ~6800 req/s; every client disconnect frees the KV
blocks on every engine; router overhead p99 is about 0.9 ms against the plan's
2 ms budget.

The handlers live in internal/router rather than an internal/router/openai
package as the plan's tree sketched: they need the dispatch loop, and splitting
them apart would have meant an interface with exactly one implementation. Plan
§4 updated to match.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 01:41:13 +02:00
ilkerandClaude Opus 5 f0c1ca87f1 ci: set milestone due dates at noon UTC
23:59:59Z renders as the following day in any positive-offset timezone, so
every due date showed up a day late in the UI.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 01:34:05 +02:00
77 changed files with 18910 additions and 117 deletions

No files matched your search

+8
View File
@@ -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
View File
@@ -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
+12 -1
View File
@@ -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
+57 -6
View File
@@ -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"))
}
+218
View File
@@ -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;
}
+12
View File
@@ -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
+16
View File
@@ -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
+63
View 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
+118
View File
@@ -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
+22 -3
View File
@@ -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
)
+42 -2
View File
@@ -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=
+558
View File
@@ -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)
},
},
}
}
+600
View File
@@ -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
}
+72
View File
@@ -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 }
+101
View File
@@ -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)
}
}
}
+133
View File
@@ -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
}
+11
View File
@@ -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 }
+148
View File
@@ -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, ":")
}
+103
View File
@@ -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))
}
}
+144
View File
@@ -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")
}
}
+203
View File
@@ -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))
}
+216
View File
@@ -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)
}
})
}
+40
View File
@@ -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)
}
}
+319
View File
@@ -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 ""
}
+425
View File
@@ -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)
}
}
+83
View File
@@ -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
}
+149
View File
@@ -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
}
+99
View File
@@ -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)
}
+86
View File
@@ -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")
}
}
+34
View File
@@ -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
}
+57
View File
@@ -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
}
+57
View File
@@ -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
}
+771
View File
@@ -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)
}
}
+187
View File
@@ -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()
}
+336
View File
@@ -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
}
+299
View File
@@ -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
}
-57
View File
@@ -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
}
+83
View File
@@ -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
}
+291
View File
@@ -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
}
+156
View File
@@ -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)
}
}
+283
View File
@@ -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})
}
+243
View File
@@ -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)
}
}
}
+225
View File
@@ -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
}
+324
View File
@@ -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)
}
}
+289
View File
@@ -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()
}
+235
View File
@@ -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)
}
}
+182
View File
@@ -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()
}
+309
View File
@@ -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()
}
+16
View File
@@ -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
}
}
+352
View File
@@ -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)
}
}
}
}
+782
View File
@@ -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)
}
}
+159
View File
@@ -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
}
+128
View File
@@ -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
}
+436
View File
@@ -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)
}
}
+412
View File
@@ -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"
}
+168
View File
@@ -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
}
}
+98
View File
@@ -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)
}
+22 -11
View File
@@ -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) })
+361
View File
@@ -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,
}})
}
+137
View File
@@ -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))
}
+468
View File
@@ -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)
}
+302
View File
@@ -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")
}
}
+39
View File
@@ -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
}
+354
View File
@@ -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 })
}
+325
View File
@@ -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
View File
@@ -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). |
---
+7 -7
View File
@@ -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",
+242
View File
@@ -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
}
+203
View File
@@ -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)],
}
}
+197
View File
@@ -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)
}
}
+414
View File
@@ -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())
}
}