Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e1eb1d4149 | ||
|
|
1d1a6aebb0 | ||
|
|
75a8e76706 | ||
|
|
40427b2705 | ||
|
|
2373726efd | ||
|
|
c1b9ea1bc2 | ||
|
|
f0134a2f95 | ||
|
|
06672d1b85 | ||
|
|
55417546fc | ||
|
|
6a13e3dddf |
No files matched your search
@@ -38,3 +38,11 @@ jobs:
|
||||
|
||||
- name: Build
|
||||
run: go build ./...
|
||||
|
||||
# The NVML backend needs cgo and cannot run without a driver, but it must
|
||||
# at least keep compiling: it is the only source of NVLink topology and a
|
||||
# break would otherwise surface on a GPU node, far from here.
|
||||
- name: Build the NVML backend
|
||||
env:
|
||||
CGO_ENABLED: "1"
|
||||
run: go build -tags nvml ./...
|
||||
+278
@@ -16,6 +16,284 @@ 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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -15,8 +15,8 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
| v0.1.0 | 0 | Skeleton and sim-engine | released |
|
||||
| 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, single-node engines | planned |
|
||||
| v0.5.0 | 3 | Control plane: catalog, scheduler, reconciler | planned |
|
||||
| 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 |
|
||||
|
||||
@@ -33,6 +33,24 @@ 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
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,119 @@
|
||||
// Code generated by protoc-gen-connect-go. DO NOT EDIT.
|
||||
//
|
||||
// Source: llmgw/v1/agent.proto
|
||||
|
||||
// Package llmgw.v1 is the protocol between a node agent and the control plane
|
||||
// (plan §5.6).
|
||||
//
|
||||
// The agent dials out and holds a bidirectional stream open. That direction is
|
||||
// deliberate: GPU nodes sit behind firewalls and NAT far more often than a
|
||||
// control plane does, and an outbound connection needs no inbound rule on every
|
||||
// node in the fleet.
|
||||
package llmgwv1connect
|
||||
|
||||
import (
|
||||
connect "connectrpc.com/connect"
|
||||
context "context"
|
||||
errors "errors"
|
||||
v1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
http "net/http"
|
||||
strings "strings"
|
||||
)
|
||||
|
||||
// This is a compile-time assertion to ensure that this generated file and the connect package are
|
||||
// compatible. If you get a compiler error that this constant is not defined, this code was
|
||||
// generated with a version of connect newer than the one compiled into your binary. You can fix the
|
||||
// problem by either regenerating this code with an older version of connect or updating the connect
|
||||
// version compiled into your binary.
|
||||
const _ = connect.IsAtLeastVersion1_13_0
|
||||
|
||||
const (
|
||||
// AgentServiceName is the fully-qualified name of the AgentService service.
|
||||
AgentServiceName = "llmgw.v1.AgentService"
|
||||
)
|
||||
|
||||
// These constants are the fully-qualified names of the RPCs defined in this package. They're
|
||||
// exposed at runtime as Spec.Procedure and as the final two segments of the HTTP route.
|
||||
//
|
||||
// Note that these are different from the fully-qualified method names used by
|
||||
// google.golang.org/protobuf/reflect/protoreflect. To convert from these constants to
|
||||
// reflection-formatted method names, remove the leading slash and convert the remaining slash to a
|
||||
// period.
|
||||
const (
|
||||
// AgentServiceConnectProcedure is the fully-qualified name of the AgentService's Connect RPC.
|
||||
AgentServiceConnectProcedure = "/llmgw.v1.AgentService/Connect"
|
||||
)
|
||||
|
||||
// AgentServiceClient is a client for the llmgw.v1.AgentService service.
|
||||
type AgentServiceClient interface {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
Connect(context.Context) *connect.BidiStreamForClient[v1.AgentMessage, v1.ControlMessage]
|
||||
}
|
||||
|
||||
// NewAgentServiceClient constructs a client for the llmgw.v1.AgentService service. By default, it
|
||||
// uses the Connect protocol with the binary Protobuf Codec, asks for gzipped responses, and sends
|
||||
// uncompressed requests. To use the gRPC or gRPC-Web protocols, supply the connect.WithGRPC() or
|
||||
// connect.WithGRPCWeb() options.
|
||||
//
|
||||
// The URL supplied here should be the base URL for the Connect or gRPC server (for example,
|
||||
// http://api.acme.com or https://acme.com/grpc).
|
||||
func NewAgentServiceClient(httpClient connect.HTTPClient, baseURL string, opts ...connect.ClientOption) AgentServiceClient {
|
||||
baseURL = strings.TrimRight(baseURL, "/")
|
||||
agentServiceMethods := v1.File_llmgw_v1_agent_proto.Services().ByName("AgentService").Methods()
|
||||
return &agentServiceClient{
|
||||
connect: connect.NewClient[v1.AgentMessage, v1.ControlMessage](
|
||||
httpClient,
|
||||
baseURL+AgentServiceConnectProcedure,
|
||||
connect.WithSchema(agentServiceMethods.ByName("Connect")),
|
||||
connect.WithClientOptions(opts...),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
// agentServiceClient implements AgentServiceClient.
|
||||
type agentServiceClient struct {
|
||||
connect *connect.Client[v1.AgentMessage, v1.ControlMessage]
|
||||
}
|
||||
|
||||
// Connect calls llmgw.v1.AgentService.Connect.
|
||||
func (c *agentServiceClient) Connect(ctx context.Context) *connect.BidiStreamForClient[v1.AgentMessage, v1.ControlMessage] {
|
||||
return c.connect.CallBidiStream(ctx)
|
||||
}
|
||||
|
||||
// AgentServiceHandler is an implementation of the llmgw.v1.AgentService service.
|
||||
type AgentServiceHandler interface {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
Connect(context.Context, *connect.BidiStream[v1.AgentMessage, v1.ControlMessage]) error
|
||||
}
|
||||
|
||||
// NewAgentServiceHandler builds an HTTP handler from the service implementation. It returns the
|
||||
// path on which to mount the handler and the handler itself.
|
||||
//
|
||||
// By default, handlers support the Connect, gRPC, and gRPC-Web protocols with the binary Protobuf
|
||||
// and JSON codecs. They also support gzip compression.
|
||||
func NewAgentServiceHandler(svc AgentServiceHandler, opts ...connect.HandlerOption) (string, http.Handler) {
|
||||
agentServiceMethods := v1.File_llmgw_v1_agent_proto.Services().ByName("AgentService").Methods()
|
||||
agentServiceConnectHandler := connect.NewBidiStreamHandler(
|
||||
AgentServiceConnectProcedure,
|
||||
svc.Connect,
|
||||
connect.WithSchema(agentServiceMethods.ByName("Connect")),
|
||||
connect.WithHandlerOptions(opts...),
|
||||
)
|
||||
return "/llmgw.v1.AgentService/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case AgentServiceConnectProcedure:
|
||||
agentServiceConnectHandler.ServeHTTP(w, r)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// UnimplementedAgentServiceHandler returns CodeUnimplemented from all methods.
|
||||
type UnimplementedAgentServiceHandler struct{}
|
||||
|
||||
func (UnimplementedAgentServiceHandler) Connect(context.Context, *connect.BidiStream[v1.AgentMessage, v1.ControlMessage]) error {
|
||||
return connect.NewError(connect.CodeUnimplemented, errors.New("llmgw.v1.AgentService.Connect is not implemented"))
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
syntax = "proto3";
|
||||
|
||||
// Package llmgw.v1 is the protocol between a node agent and the control plane
|
||||
// (plan §5.6).
|
||||
//
|
||||
// The agent dials out and holds a bidirectional stream open. That direction is
|
||||
// deliberate: GPU nodes sit behind firewalls and NAT far more often than a
|
||||
// control plane does, and an outbound connection needs no inbound rule on every
|
||||
// node in the fleet.
|
||||
package llmgw.v1;
|
||||
|
||||
option go_package = "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1;llmgwv1";
|
||||
|
||||
service AgentService {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
rpc Connect(stream AgentMessage) returns (stream ControlMessage);
|
||||
}
|
||||
|
||||
// AgentMessage is anything a node tells the control plane.
|
||||
message AgentMessage {
|
||||
oneof msg {
|
||||
Register register = 1;
|
||||
Heartbeat heartbeat = 2;
|
||||
ShardStatus shard = 3;
|
||||
}
|
||||
}
|
||||
|
||||
// ControlMessage is anything the control plane tells a node.
|
||||
message ControlMessage {
|
||||
oneof msg {
|
||||
StartShard start = 1;
|
||||
StopShard stop = 2;
|
||||
DrainShard drain = 3;
|
||||
Ack ack = 4;
|
||||
}
|
||||
}
|
||||
|
||||
// Register is the first message on every stream. It carries the node's full
|
||||
// hardware inventory, so the control plane needs no separate discovery call and
|
||||
// a reconnecting agent automatically refreshes what the scheduler knows.
|
||||
message Register {
|
||||
string node_id = 1;
|
||||
string hostname = 2;
|
||||
// mgmt_addr is the address engines are reachable on, and the address NCCL
|
||||
// bootstraps over. Never a fabric address: NCCL rendezvouses over TCP before
|
||||
// it moves to RDMA.
|
||||
string mgmt_addr = 3;
|
||||
map<string, string> labels = 4;
|
||||
|
||||
repeated GPU gpus = 5;
|
||||
// fabrics is empty on a node without InfiniBand. Such a node is its own
|
||||
// single-node island and can never host a multi-node replica.
|
||||
repeated FabricAttachment fabrics = 6;
|
||||
|
||||
string agent_version = 7;
|
||||
// engine_versions decides which CLI flags and metric names apply; both drift
|
||||
// between engine releases (plan §14).
|
||||
map<string, string> engine_versions = 8;
|
||||
|
||||
// running_shards are the engines this agent is already supervising, including
|
||||
// ones adopted from a previous agent.
|
||||
//
|
||||
// This is what lets a control plane restart recover instead of orphaning live
|
||||
// engines. Without it the control plane forgets every replica, places new
|
||||
// ones, and the old engines keep holding GPUs the inventory believes are free
|
||||
// — so the new engines fail to allocate, with an error naming the wrong
|
||||
// process (plan §6.5).
|
||||
repeated RunningShard running_shards = 9;
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
//
|
||||
// It carries enough for the control plane to rebuild the replica it belongs to,
|
||||
// which is more than a shard id: the replica's identity, its pool and fabric,
|
||||
// and the leader endpoint clients reach it through.
|
||||
message RunningShard {
|
||||
string replica_id = 1;
|
||||
string shard_id = 2;
|
||||
int32 rank = 3;
|
||||
int32 nnodes = 4;
|
||||
repeated int32 gpu_indices = 5;
|
||||
repeated string hcas = 6;
|
||||
int32 http_port = 7;
|
||||
ShardStatus.State state = 8;
|
||||
|
||||
// model and profile name the catalog entry this shard was started from.
|
||||
string model = 9;
|
||||
string profile = 10;
|
||||
// pool and fabric are echoed back from the StartShard that created it; the
|
||||
// agent has no way to derive them.
|
||||
string pool = 11;
|
||||
string fabric = 12;
|
||||
|
||||
string leader_addr = 13;
|
||||
int32 leader_port = 14;
|
||||
}
|
||||
|
||||
message GPU {
|
||||
int32 index = 1;
|
||||
string uuid = 2;
|
||||
string name = 3;
|
||||
int32 memory_total_mb = 4;
|
||||
int32 memory_used_mb = 5;
|
||||
string compute_capability = 6;
|
||||
int32 numa_node = 7;
|
||||
string pci_bus_id = 8;
|
||||
string driver_version = 9;
|
||||
string cuda_version = 10;
|
||||
// nvlink_peers is only populated by agents built with the nvml tag.
|
||||
repeated int32 nvlink_peers = 11;
|
||||
// problems keeps a GPU out of placement: ECC faults, or memory already held
|
||||
// by another process.
|
||||
repeated string problems = 12;
|
||||
int32 temperature_c = 13;
|
||||
int32 ecc_errors = 14;
|
||||
}
|
||||
|
||||
message FabricAttachment {
|
||||
// fabric is the operator-assigned name. It is not derived from the GID
|
||||
// prefix: OpenSM defaults to fe80:: on every fabric, so a matching prefix
|
||||
// proves nothing (plan §5.2).
|
||||
string fabric = 1;
|
||||
repeated string hcas = 2;
|
||||
int32 rate_gbps = 3;
|
||||
}
|
||||
|
||||
// Heartbeat says the node is alive and reports anything that changed cheaply.
|
||||
message Heartbeat {
|
||||
string node_id = 1;
|
||||
// gpu_memory_used_mb, indexed by GPU, catches a foreign process appearing on
|
||||
// a card after registration.
|
||||
repeated int32 gpu_memory_used_mb = 2;
|
||||
repeated string new_problems = 3;
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
message ShardStatus {
|
||||
string node_id = 1;
|
||||
string replica_id = 2;
|
||||
string shard_id = 3;
|
||||
State state = 4;
|
||||
string error = 5;
|
||||
int32 exit_code = 6;
|
||||
// last_output carries the tail of the engine log on a failure. When a model
|
||||
// fails to load this is the only diagnosis available.
|
||||
repeated string last_output = 7;
|
||||
|
||||
enum State {
|
||||
STATE_UNSPECIFIED = 0;
|
||||
STATE_PENDING = 1;
|
||||
STATE_STARTING = 2;
|
||||
STATE_LOADING = 3;
|
||||
STATE_READY = 4;
|
||||
STATE_DRAINING = 5;
|
||||
STATE_STOPPED = 6;
|
||||
STATE_FAILED = 7;
|
||||
}
|
||||
}
|
||||
|
||||
// StartShard launches one shard. It is idempotent by shard_id: the control
|
||||
// plane retries, and a repeat must not produce a second process.
|
||||
message StartShard {
|
||||
string replica_id = 1;
|
||||
string shard_id = 2;
|
||||
|
||||
int32 rank = 3;
|
||||
int32 nnodes = 4;
|
||||
string leader_addr = 5;
|
||||
int32 leader_port = 6;
|
||||
|
||||
repeated int32 gpu_indices = 7;
|
||||
// hcas are the adapters of the chosen fabric and nothing else. Handing a
|
||||
// shard adapters from a second fabric lets NCCL pick one whose peer is
|
||||
// unreachable, and the job hangs instead of failing (plan §5.3).
|
||||
repeated string hcas = 8;
|
||||
string mgmt_interface = 9;
|
||||
int32 http_port = 10;
|
||||
|
||||
EngineProfile profile = 11;
|
||||
int64 startup_timeout_ms = 12;
|
||||
|
||||
// pool and fabric are control plane concepts the agent cannot derive. They
|
||||
// are sent so the agent can echo them back in RunningShard after a restart.
|
||||
string pool = 13;
|
||||
string fabric = 14;
|
||||
}
|
||||
|
||||
message EngineProfile {
|
||||
string name = 1;
|
||||
string engine = 2;
|
||||
string gpu_class = 3;
|
||||
int32 tp = 4;
|
||||
int32 pp = 5;
|
||||
int32 gpus_per_node = 6;
|
||||
int32 nodes = 7;
|
||||
repeated string args = 8;
|
||||
map<string, string> env = 9;
|
||||
string model_path = 10;
|
||||
string served_model_name = 11;
|
||||
}
|
||||
|
||||
// StopShard terminates a shard and frees its GPUs.
|
||||
message StopShard {
|
||||
string shard_id = 1;
|
||||
}
|
||||
|
||||
// DrainShard stops routing to a shard while letting in-flight work finish.
|
||||
message DrainShard {
|
||||
string shard_id = 1;
|
||||
}
|
||||
|
||||
// Ack confirms a registration or reports that a command could not be accepted.
|
||||
message Ack {
|
||||
string message_id = 1;
|
||||
bool ok = 2;
|
||||
string error = 3;
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
version: v2
|
||||
managed:
|
||||
enabled: true
|
||||
plugins:
|
||||
- local: protoc-gen-go
|
||||
out: api/gen
|
||||
opt: paths=source_relative
|
||||
- local: protoc-gen-connect-go
|
||||
out: api/gen
|
||||
opt: paths=source_relative
|
||||
inputs:
|
||||
- directory: api/proto
|
||||
@@ -0,0 +1,16 @@
|
||||
version: v2
|
||||
modules:
|
||||
- path: api/proto
|
||||
lint:
|
||||
use:
|
||||
- STANDARD
|
||||
except:
|
||||
# The standard names assume request/response pairs. This is one long-lived
|
||||
# bidirectional stream carrying several kinds of message in each direction,
|
||||
# so "AgentMessage" and "ControlMessage" say what they are and
|
||||
# "ConnectRequest"/"ConnectResponse" would not.
|
||||
- RPC_REQUEST_STANDARD_NAME
|
||||
- RPC_RESPONSE_STANDARD_NAME
|
||||
breaking:
|
||||
use:
|
||||
- FILE
|
||||
@@ -0,0 +1,63 @@
|
||||
# llmgw agent configuration, one per GPU node.
|
||||
#
|
||||
# The agent dials out to the control plane and keeps one stream open. Engines
|
||||
# keep serving if the agent restarts: the router talks to them directly, and
|
||||
# orphan adoption picks them up again on the next start.
|
||||
|
||||
node_id: gpu-a-01 # defaults to the hostname
|
||||
server: "http://llmgw-server.internal:7443"
|
||||
|
||||
# The address engines are reachable on, and the address NCCL bootstraps over.
|
||||
# Never a fabric address: NCCL rendezvouses over TCP before moving to RDMA.
|
||||
mgmt_addr: "10.0.0.11"
|
||||
mgmt_interface: eth0
|
||||
|
||||
labels:
|
||||
rack: r12
|
||||
|
||||
heartbeat_interval: 5s
|
||||
reconnect_base: 1s
|
||||
reconnect_max: 30s
|
||||
|
||||
# Fabric labelling is explicit and there is no way around it. OpenSM's default
|
||||
# subnet prefix is fe80:: on every fabric, so a matching prefix proves nothing
|
||||
# and a wrong guess produces a replica that starts, passes health checks and
|
||||
# then hangs on its first all-reduce (plan §5.2).
|
||||
#
|
||||
# A node with no fabrics is its own single-node island: it can host single-node
|
||||
# replicas and will never be given a multi-node one.
|
||||
fabrics:
|
||||
- name: ib-a
|
||||
hcas: [mlx5_0, mlx5_1, mlx5_2, mlx5_3]
|
||||
rate_gbps: 200
|
||||
|
||||
discovery:
|
||||
# Worth setting where the driver lives outside the default paths, which
|
||||
# containerized driver installs often do.
|
||||
nvidia_smi: ""
|
||||
ib_root: /sys/class/infiniband
|
||||
# NVML is the only source of NVLink topology, which is what keeps a
|
||||
# tensor-parallel group off the PCIe path inside a node. It needs a binary
|
||||
# built with `-tags nvml`; otherwise this silently falls back to nvidia-smi.
|
||||
prefer_nvml: true
|
||||
skip_gpus: false
|
||||
|
||||
engines:
|
||||
vllm:
|
||||
bin: /opt/venvs/vllm/bin/vllm
|
||||
# Which multi-node mechanism works depends on the vLLM version, and the
|
||||
# flags have changed more than once. Check the version you deploy.
|
||||
launch_mode: ray # "ray" or "mp"
|
||||
ray_bin: /opt/venvs/vllm/bin/ray
|
||||
sglang:
|
||||
# SGLang is launched as a module, which is its documented entry point.
|
||||
python: /opt/venvs/sglang/bin/python
|
||||
|
||||
supervisor:
|
||||
work_dir: /var/lib/llmgw-agent
|
||||
# Engine output is the only diagnosis when a model fails to load, so keep a
|
||||
# useful tail of it.
|
||||
log_lines: 512
|
||||
# Engines need time to free GPU memory cleanly before being killed.
|
||||
stop_timeout: 60s
|
||||
health_interval: 2s
|
||||
+27
-14
@@ -8,6 +8,10 @@
|
||||
|
||||
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
|
||||
@@ -79,6 +83,19 @@ router:
|
||||
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]
|
||||
@@ -89,17 +106,13 @@ models:
|
||||
aliases: [qwen]
|
||||
routing_profile: batch
|
||||
|
||||
endpoints:
|
||||
- id: gpu-a-01
|
||||
url: http://gpu-a-01:8000
|
||||
models: [llama-3.3-70b]
|
||||
pool: ib-a/h100-80g
|
||||
weight: 100 # from ctl bench; do not estimate it by hand
|
||||
priority: 100 # higher pools are tried first
|
||||
|
||||
- id: gpu-b-01
|
||||
url: http://gpu-b-01:8000
|
||||
models: [llama-3.3-70b, qwen2.5-32b]
|
||||
pool: ib-b/a100-80g
|
||||
weight: 45
|
||||
priority: 50
|
||||
# A static replica list, used only when no model declares a profile. That is the
|
||||
# router-only mode for testing against sim-engines; with the control plane
|
||||
# running, agents provide the replicas instead.
|
||||
#
|
||||
# endpoints:
|
||||
# - id: sim-1
|
||||
# url: http://127.0.0.1:8001
|
||||
# models: [llama-3.3-70b]
|
||||
# pool: local
|
||||
# weight: 1
|
||||
@@ -3,6 +3,8 @@ 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
|
||||
@@ -12,6 +14,8 @@ require (
|
||||
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 (
|
||||
@@ -26,7 +30,7 @@ require (
|
||||
github.com/prometheus/procfs v0.21.1 // indirect
|
||||
github.com/spf13/pflag v1.0.9 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
golang.org/x/sync v0.23.0 // indirect
|
||||
golang.org/x/sys v0.47.0 // indirect
|
||||
google.golang.org/protobuf v1.36.12 // indirect
|
||||
golang.org/x/net v0.59.0 // indirect
|
||||
golang.org/x/sys v0.48.0 // indirect
|
||||
golang.org/x/text v0.42.0 // indirect
|
||||
)
|
||||
@@ -1,3 +1,7 @@
|
||||
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=
|
||||
@@ -53,10 +57,14 @@ 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.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
|
||||
golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og=
|
||||
golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI=
|
||||
golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E=
|
||||
google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc=
|
||||
google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
@@ -0,0 +1,558 @@
|
||||
// Package agent is the node side of the control plane stream (plan §5.6).
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"math/rand/v2"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// Config configures the agent.
|
||||
type Config struct {
|
||||
NodeID string `koanf:"node_id"`
|
||||
Server string `koanf:"server"`
|
||||
Hostname string `koanf:"hostname"`
|
||||
MgmtAddr string `koanf:"mgmt_addr"`
|
||||
Labels map[string]string `koanf:"labels"`
|
||||
|
||||
// MgmtInterface is the ethernet interface NCCL bootstraps over.
|
||||
MgmtInterface string `koanf:"mgmt_interface"`
|
||||
// Fabrics labels this node's HCAs. It is explicit on purpose: OpenSM's
|
||||
// default subnet prefix is the same on every fabric, so nothing can be
|
||||
// inferred from the wire (plan §5.2).
|
||||
Fabrics []inventory.FabricAttachment `koanf:"fabrics"`
|
||||
|
||||
HeartbeatInterval time.Duration `koanf:"heartbeat_interval"`
|
||||
// ReconnectBase is the first reconnect delay; it backs off from there.
|
||||
ReconnectBase time.Duration `koanf:"reconnect_base"`
|
||||
ReconnectMax time.Duration `koanf:"reconnect_max"`
|
||||
|
||||
Discovery discovery.Options `koanf:"discovery"`
|
||||
Engines engines.Config `koanf:"engines"`
|
||||
Superv supervisor.Config `koanf:"supervisor"`
|
||||
|
||||
Version string
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
MgmtInterface: "eth0",
|
||||
HeartbeatInterval: 5 * time.Second,
|
||||
ReconnectBase: time.Second,
|
||||
ReconnectMax: 30 * time.Second,
|
||||
Superv: supervisor.DefaultConfig(),
|
||||
// NVML gives NVLink topology, so prefer it when this build has it.
|
||||
Discovery: discovery.Options{PreferNVML: true},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() error {
|
||||
d := DefaultConfig()
|
||||
if c.NodeID == "" {
|
||||
return errors.New("node_id is required")
|
||||
}
|
||||
if c.Server == "" {
|
||||
return errors.New("server is required")
|
||||
}
|
||||
if c.MgmtInterface == "" {
|
||||
c.MgmtInterface = d.MgmtInterface
|
||||
}
|
||||
if c.HeartbeatInterval <= 0 {
|
||||
c.HeartbeatInterval = d.HeartbeatInterval
|
||||
}
|
||||
if c.ReconnectBase <= 0 {
|
||||
c.ReconnectBase = d.ReconnectBase
|
||||
}
|
||||
if c.ReconnectMax < c.ReconnectBase {
|
||||
c.ReconnectMax = max(c.ReconnectBase, d.ReconnectMax)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Agent runs on a GPU node.
|
||||
type Agent struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
sup *supervisor.Supervisor
|
||||
reg *engines.Registry
|
||||
http connectHTTPClient
|
||||
|
||||
mu sync.Mutex
|
||||
inv discovery.Inventory
|
||||
// shardReplica remembers which replica a shard belongs to, so a status
|
||||
// report can name both. The control plane keys on the replica.
|
||||
shardReplica map[string]string
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
// hand in an httptest server's client.
|
||||
type connectHTTPClient interface {
|
||||
Do(*http.Request) (*http.Response, error)
|
||||
}
|
||||
|
||||
// New builds an agent.
|
||||
func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, error) {
|
||||
if err := cfg.normalize(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if httpClient == nil {
|
||||
httpClient = http.DefaultClient
|
||||
}
|
||||
|
||||
health := supervisor.NewHTTPHealth(3 * time.Second)
|
||||
sup, err := supervisor.New(cfg.Superv, health, log)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Agent{
|
||||
cfg: cfg, log: log, sup: sup,
|
||||
reg: engines.NewRegistry(cfg.Engines),
|
||||
http: httpClient,
|
||||
shardReplica: map[string]string{},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Supervisor exposes the process supervisor, for tests and `ctl logs`.
|
||||
func (a *Agent) Supervisor() *supervisor.Supervisor { return a.sup }
|
||||
|
||||
// Run connects to the control plane and stays connected until ctx ends.
|
||||
//
|
||||
// Reconnecting forever is the right behaviour for a node agent: a control plane
|
||||
// restart or a network partition must not require touching every GPU node, and
|
||||
// the engines keep serving traffic throughout because the router talks to them
|
||||
// directly.
|
||||
func (a *Agent) Run(ctx context.Context) error {
|
||||
if err := a.scan(ctx); err != nil {
|
||||
// Not fatal: a node whose GPUs cannot be read should still register, so
|
||||
// an operator can see it and why it is useless.
|
||||
a.log.Error("hardware discovery failed", "err", err)
|
||||
}
|
||||
a.recoverOrphans()
|
||||
|
||||
attempt := 0
|
||||
for ctx.Err() == nil {
|
||||
err := a.session(ctx)
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
attempt++
|
||||
delay := a.reconnectDelay(attempt)
|
||||
a.log.Warn("disconnected from the control plane, retrying",
|
||||
"err", err, "in", delay.Round(time.Millisecond), "attempt", attempt)
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(delay):
|
||||
}
|
||||
if err == nil {
|
||||
attempt = 0
|
||||
}
|
||||
}
|
||||
|
||||
// Engines keep running when the agent stops on purpose: an agent restart
|
||||
// must not interrupt served traffic. Orphan adoption picks them up again.
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
func (a *Agent) reconnectDelay(attempt int) time.Duration {
|
||||
d := a.cfg.ReconnectBase << min(attempt-1, 16)
|
||||
if d > a.cfg.ReconnectMax || d <= 0 {
|
||||
d = a.cfg.ReconnectMax
|
||||
}
|
||||
// Jitter, so a fleet does not reconnect in lockstep after a control plane
|
||||
// restart and knock it over again.
|
||||
return d/2 + time.Duration(rand.Int64N(int64(d/2))) //nolint:gosec // jitter
|
||||
}
|
||||
|
||||
// recoverOrphans takes back engines a previous agent left running.
|
||||
//
|
||||
// Engines deliberately outlive their agent — the router talks to them directly,
|
||||
// so an agent restart must not interrupt served traffic. The cost is that they
|
||||
// have to be found again, and anything that cannot be identified has to be
|
||||
// killed: a process holding GPUs that nothing can describe is worse than no
|
||||
// process at all, because placement will hand those GPUs out again.
|
||||
func (a *Agent) recoverOrphans() {
|
||||
orphans, err := a.sup.FindOrphans()
|
||||
if err != nil {
|
||||
a.log.Warn("scanning for leftover engines failed", "err", err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, o := range orphans {
|
||||
switch {
|
||||
case !o.Alive:
|
||||
// The process is gone; only its state file remains.
|
||||
if err := a.sup.KillOrphan(o); err != nil {
|
||||
a.log.Warn("cleaning up a dead shard's state failed", "shard", o.ShardID, "err", err)
|
||||
}
|
||||
case o.ReplicaID == "":
|
||||
// Alive, but its state file could not be read, so it cannot be
|
||||
// described to the control plane or matched to a replica.
|
||||
a.log.Warn("killing an unidentifiable engine", "shard", o.ShardID, "pid", o.PID)
|
||||
if err := a.sup.KillOrphan(o); err != nil {
|
||||
a.log.Warn("killing it failed", "shard", o.ShardID, "err", err)
|
||||
}
|
||||
default:
|
||||
sh, err := a.sup.Adopt(o)
|
||||
if err != nil {
|
||||
a.log.Warn("adopting a shard failed", "shard", o.ShardID, "err", err)
|
||||
continue
|
||||
}
|
||||
a.mu.Lock()
|
||||
a.shardReplica[o.ShardID] = o.ReplicaID
|
||||
a.mu.Unlock()
|
||||
a.log.Info("adopted a running engine",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID, "pid", o.PID)
|
||||
_ = sh
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runningShards describes what this agent is supervising, for registration.
|
||||
func (a *Agent) runningShards() []*llmgwv1.RunningShard {
|
||||
var out []*llmgwv1.RunningShard
|
||||
for _, st := range a.sup.Statuses() {
|
||||
if st.State.Terminal() {
|
||||
continue
|
||||
}
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
spec := sh.Spec
|
||||
out = append(out, &llmgwv1.RunningShard{
|
||||
ReplicaId: spec.ReplicaID,
|
||||
ShardId: st.ID,
|
||||
Rank: int32(spec.Rank),
|
||||
Nnodes: int32(spec.NNodes),
|
||||
GpuIndices: int32Slice(spec.GPUIndices),
|
||||
Hcas: spec.HCAs,
|
||||
HttpPort: int32(spec.HTTPPort),
|
||||
State: agentsvc.StateToWire(st.State),
|
||||
Model: spec.Profile.ServedModelName,
|
||||
Profile: spec.Profile.Name,
|
||||
Pool: spec.Pool,
|
||||
Fabric: spec.Fabric,
|
||||
LeaderAddr: a.cfg.MgmtAddr,
|
||||
LeaderPort: int32(spec.LeaderPort),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func int32Slice(v []int) []int32 {
|
||||
if len(v) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]int32, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = int32(x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// scan refreshes the node's hardware inventory.
|
||||
func (a *Agent) scan(ctx context.Context) error {
|
||||
inv, err := discovery.Discover(ctx, a.cfg.Discovery, a.log)
|
||||
a.mu.Lock()
|
||||
a.inv = inv
|
||||
a.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
|
||||
// session runs one connection until it breaks.
|
||||
func (a *Agent) session(ctx context.Context) error {
|
||||
client := llmgwv1connect.NewAgentServiceClient(a.http, a.cfg.Server)
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
stream := client.Connect(ctx)
|
||||
defer func() { _ = stream.CloseRequest() }()
|
||||
|
||||
a.mu.Lock()
|
||||
inv := a.inv
|
||||
a.mu.Unlock()
|
||||
|
||||
reg := agentsvc.RegisterFromInventory(
|
||||
a.cfg.NodeID, a.cfg.Hostname, a.cfg.MgmtAddr, a.cfg.Version,
|
||||
a.cfg.Labels, inv, a.cfg.Fabrics)
|
||||
// Every registration reports what is already running, not just the first
|
||||
// one: a reconnect may be a control plane restart, and it has to recover
|
||||
// the same way (plan §6.5).
|
||||
reg.RunningShards = a.runningShards()
|
||||
|
||||
if err := stream.Send(&llmgwv1.AgentMessage{
|
||||
Msg: &llmgwv1.AgentMessage_Register{Register: reg},
|
||||
}); err != nil {
|
||||
return fmt.Errorf("registering: %w", err)
|
||||
}
|
||||
|
||||
// Sends are serialized through this channel: connect streams are not safe
|
||||
// for concurrent writes, and heartbeats race with status reports.
|
||||
out := make(chan *llmgwv1.AgentMessage, 32)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
a.sendLoop(ctx, stream, out)
|
||||
}()
|
||||
defer wg.Wait()
|
||||
|
||||
for {
|
||||
msg, err := stream.Receive()
|
||||
if err != nil {
|
||||
cancel()
|
||||
if errors.Is(err, io.EOF) || ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
a.handle(ctx, msg, out)
|
||||
}
|
||||
}
|
||||
|
||||
type sender interface {
|
||||
Send(*llmgwv1.AgentMessage) error
|
||||
}
|
||||
|
||||
func (a *Agent) sendLoop(ctx context.Context, s sender, out <-chan *llmgwv1.AgentMessage) {
|
||||
t := time.NewTicker(a.cfg.HeartbeatInterval)
|
||||
defer t.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case msg := <-out:
|
||||
if err := s.Send(msg); err != nil {
|
||||
return
|
||||
}
|
||||
case <-t.C:
|
||||
if err := s.Send(a.heartbeat()); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) heartbeat() *llmgwv1.AgentMessage {
|
||||
a.mu.Lock()
|
||||
inv := a.inv
|
||||
a.mu.Unlock()
|
||||
|
||||
hb := &llmgwv1.Heartbeat{NodeId: a.cfg.NodeID}
|
||||
for _, g := range inv.GPUs {
|
||||
hb.GpuMemoryUsedMb = append(hb.GpuMemoryUsedMb, int32(g.MemoryUsedMB))
|
||||
}
|
||||
return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Heartbeat{Heartbeat: hb}}
|
||||
}
|
||||
|
||||
func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out chan<- *llmgwv1.AgentMessage) {
|
||||
switch m := msg.Msg.(type) {
|
||||
case *llmgwv1.ControlMessage_Start:
|
||||
go a.startShard(ctx, m.Start, out)
|
||||
case *llmgwv1.ControlMessage_Stop:
|
||||
go a.stopShard(ctx, m.Stop.GetShardId(), out)
|
||||
case *llmgwv1.ControlMessage_Drain:
|
||||
if err := a.sup.Drain(m.Drain.GetShardId()); err != nil {
|
||||
a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err)
|
||||
}
|
||||
case *llmgwv1.ControlMessage_Ack:
|
||||
if !m.Ack.GetOk() {
|
||||
a.log.Warn("control plane rejected a message",
|
||||
"id", m.Ack.GetMessageId(), "err", m.Ack.GetError())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// startShard launches a shard and reports its progress.
|
||||
//
|
||||
// It runs in its own goroutine because model loading takes minutes, and the
|
||||
// stream has to keep carrying heartbeats and other commands meanwhile.
|
||||
func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<- *llmgwv1.AgentMessage) {
|
||||
shardID, replicaID := m.GetShardId(), m.GetReplicaId()
|
||||
|
||||
a.mu.Lock()
|
||||
a.shardReplica[shardID] = replicaID
|
||||
a.mu.Unlock()
|
||||
|
||||
report := func(st supervisor.State, errMsg string, tail []string) {
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
|
||||
State: agentsvc.StateToWire(st), Error: errMsg, LastOutput: tail,
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
spec, timeout := agentsvc.ShardSpecFromStart(m)
|
||||
|
||||
eng, err := a.reg.Get(spec.Profile.Engine)
|
||||
if err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
|
||||
a.mu.Lock()
|
||||
inv := a.inv
|
||||
a.mu.Unlock()
|
||||
|
||||
// Validate before launching. An engine that refuses a profile does so after
|
||||
// minutes of loading; catching it here costs nothing.
|
||||
if err := eng.Validate(spec.Profile, gpuInfo(inv, spec.GPUIndices)); err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
|
||||
cmds, err := eng.Commands(spec)
|
||||
if err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
|
||||
sh, err := a.sup.Start(ctx, shardID, spec, cmds)
|
||||
if err != nil {
|
||||
if errors.Is(err, supervisor.ErrAlreadyRunning) {
|
||||
// Commands are idempotent: the control plane retries, and a repeat
|
||||
// must not produce a second process (plan §5.6).
|
||||
report(sh.Status().State, "", nil)
|
||||
return
|
||||
}
|
||||
report(supervisor.Failed, err.Error(), sh.Logs(10))
|
||||
return
|
||||
}
|
||||
report(supervisor.Starting, "", nil)
|
||||
|
||||
// Only a rank that serves HTTP can be health-checked. A worker is judged by
|
||||
// staying alive, and the replica's real readiness is the leader answering,
|
||||
// which only happens once every worker has joined (plan §11, phase 4).
|
||||
if eng.ServesHTTP(spec) {
|
||||
base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort)
|
||||
err = a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout)
|
||||
} else {
|
||||
err = a.sup.WaitAlive(ctx, sh)
|
||||
}
|
||||
if err != nil {
|
||||
report(supervisor.Failed, err.Error(), sh.Logs(20))
|
||||
return
|
||||
}
|
||||
report(supervisor.Ready, "", nil)
|
||||
|
||||
// Watch for the engine dying later. Nothing else would notice: the router
|
||||
// would simply see failing requests.
|
||||
go a.watch(ctx, sh, replicaID, shardID, out)
|
||||
}
|
||||
|
||||
// watch reports a shard that exits after it was ready.
|
||||
func (a *Agent) watch(ctx context.Context, sh *supervisor.Shard, replicaID, shardID string, out chan<- *llmgwv1.AgentMessage) {
|
||||
t := time.NewTicker(time.Second)
|
||||
defer t.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
st := sh.Status()
|
||||
if !st.State.Terminal() {
|
||||
continue
|
||||
}
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
|
||||
State: agentsvc.StateToWire(st.State), Error: st.Err,
|
||||
ExitCode: int32(st.ExitCode), LastOutput: sh.Logs(20),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) stopShard(ctx context.Context, shardID string, out chan<- *llmgwv1.AgentMessage) {
|
||||
a.mu.Lock()
|
||||
replicaID := a.shardReplica[shardID]
|
||||
delete(a.shardReplica, shardID)
|
||||
a.mu.Unlock()
|
||||
|
||||
if err := a.sup.Stop(ctx, shardID); err != nil {
|
||||
a.log.Warn("stopping a shard failed", "shard", shardID, "err", err)
|
||||
}
|
||||
a.sup.Remove(shardID)
|
||||
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
|
||||
State: agentsvc.StateToWire(supervisor.Stopped),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
// gpuInfo narrows the inventory to the GPUs a shard will use.
|
||||
func gpuInfo(inv discovery.Inventory, indices []int) []engines.GPUInfo {
|
||||
want := map[int]struct{}{}
|
||||
for _, i := range indices {
|
||||
want[i] = struct{}{}
|
||||
}
|
||||
var out []engines.GPUInfo
|
||||
for _, g := range inv.GPUs {
|
||||
if _, ok := want[g.Index]; !ok {
|
||||
continue
|
||||
}
|
||||
out = append(out, engines.GPUInfo{
|
||||
Index: g.Index, Name: g.Name,
|
||||
ComputeCapability: g.ComputeCapability, MemoryTotalMB: g.MemoryTotalMB,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// HTTP2Client returns the client the agent should use to reach the control
|
||||
// plane over plaintext HTTP/2.
|
||||
//
|
||||
// connect's bidirectional streams require HTTP/2, and an agent talking to a
|
||||
// control plane over a trusted network has no TLS to negotiate it with. Go's
|
||||
// default transport will not speak HTTP/2 without TLS, so it has to be asked
|
||||
// explicitly.
|
||||
//
|
||||
// mTLS is plan §5.6's target and replaces this; until then the stream must not
|
||||
// cross an untrusted network.
|
||||
func HTTP2Client(timeout time.Duration) *http.Client {
|
||||
return &http.Client{
|
||||
Transport: &http2.Transport{
|
||||
AllowHTTP: true,
|
||||
DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
|
||||
return (&net.Dialer{Timeout: timeout}).DialContext(ctx, network, addr)
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,600 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/http2/h2c"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// statusRecorder captures what the control plane hears about shards, and what
|
||||
// it is asked to adopt.
|
||||
type statusRecorder struct {
|
||||
ch chan status
|
||||
|
||||
mu sync.Mutex
|
||||
adopted []reconciler.AdoptedShard
|
||||
unclaim []string // shard ids to report back as unclaimed
|
||||
}
|
||||
|
||||
type status struct {
|
||||
replicaID, shardID string
|
||||
state supervisor.State
|
||||
err string
|
||||
}
|
||||
|
||||
func newRecorder() *statusRecorder { return &statusRecorder{ch: make(chan status, 128)} }
|
||||
|
||||
func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
|
||||
select {
|
||||
case r.ch <- status{replicaID, shardID, st, errMsg}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (r *statusRecorder) Adopt(shards []reconciler.AdoptedShard) []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.adopted = append(r.adopted, shards...)
|
||||
|
||||
var unclaimed []reconciler.AdoptedShard
|
||||
for _, s := range shards {
|
||||
if slices.Contains(r.unclaim, s.ShardID) {
|
||||
unclaimed = append(unclaimed, s)
|
||||
}
|
||||
}
|
||||
return unclaimed
|
||||
}
|
||||
|
||||
func (r *statusRecorder) adoptedShards() []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return slices.Clone(r.adopted)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) markUnclaimed(ids ...string) {
|
||||
r.mu.Lock()
|
||||
r.unclaim = append(r.unclaim, ids...)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *statusRecorder) await(t *testing.T, limit time.Duration, pred func(status) bool) status {
|
||||
t.Helper()
|
||||
deadline := time.After(limit)
|
||||
var seen []status
|
||||
for {
|
||||
select {
|
||||
case s := <-r.ch:
|
||||
seen = append(seen, s)
|
||||
if pred(s) {
|
||||
return s
|
||||
}
|
||||
case <-deadline:
|
||||
t.Fatalf("no matching status within %v; saw %+v", limit, seen)
|
||||
return status{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// fakeEngineScript writes a script that stands in for vLLM.
|
||||
//
|
||||
// It parses the real --port flag the vLLM adapter generates rather than being
|
||||
// told out of band, so the adapter's argument construction is exercised too.
|
||||
func fakeEngineScript(t *testing.T, body string) string {
|
||||
t.Helper()
|
||||
path := filepath.Join(t.TempDir(), "fake-vllm")
|
||||
script := `#!/bin/sh
|
||||
port=""
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--port) port="$2"; shift 2 ;;
|
||||
*) shift ;;
|
||||
esac
|
||||
done
|
||||
` + body + "\n"
|
||||
if err := os.WriteFile(path, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// servingEngine answers /health and /v1/models, which is what readiness needs.
|
||||
func servingEngine(t *testing.T) string {
|
||||
t.Helper()
|
||||
py, err := exec.LookPath("python3")
|
||||
if err != nil {
|
||||
t.Skip("python3 is required for the fake engine")
|
||||
}
|
||||
return fakeEngineScript(t, `exec `+py+` -c '
|
||||
import http.server, json, sys
|
||||
port = int(sys.argv[1])
|
||||
class H(http.server.BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
if self.path == "/health":
|
||||
self.send_response(200); self.end_headers(); return
|
||||
if self.path == "/v1/models":
|
||||
b = json.dumps({"object":"list","data":[{"id":"test-model"}]}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("content-type","application/json")
|
||||
self.send_header("content-length", str(len(b)))
|
||||
self.end_headers(); self.wfile.write(b); return
|
||||
self.send_response(404); self.end_headers()
|
||||
def log_message(self, *a): pass
|
||||
http.server.HTTPServer(("127.0.0.1", port), H).serve_forever()
|
||||
' "$port"`)
|
||||
}
|
||||
|
||||
type harness struct {
|
||||
svc *agentsvc.Service
|
||||
inv *inventory.Store
|
||||
rec *statusRecorder
|
||||
srv *httptest.Server
|
||||
ag *Agent
|
||||
ln *dropListener
|
||||
// streams counts how many times an agent has opened the Connect stream,
|
||||
// which is the only honest way to tell a reconnect from a stale entry.
|
||||
streams *atomic.Int64
|
||||
}
|
||||
|
||||
// dropListener tracks accepted connections so a test can sever them.
|
||||
//
|
||||
// httptest's CloseClientConnections does not work here: h2c with prior
|
||||
// knowledge hijacks the connection, and httptest stops tracking hijacked
|
||||
// connections. A test relying on it would pass without ever reconnecting.
|
||||
type dropListener struct {
|
||||
net.Listener
|
||||
mu sync.Mutex
|
||||
conns []net.Conn
|
||||
}
|
||||
|
||||
func (l *dropListener) Accept() (net.Conn, error) {
|
||||
c, err := l.Listener.Accept()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
l.mu.Lock()
|
||||
l.conns = append(l.conns, c)
|
||||
l.mu.Unlock()
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// drop severs every connection, as a control plane restart would.
|
||||
func (l *dropListener) drop() {
|
||||
l.mu.Lock()
|
||||
conns := l.conns
|
||||
l.conns = nil
|
||||
l.mu.Unlock()
|
||||
for _, c := range conns {
|
||||
_ = c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// start wires a real agent to a real service over h2c, so this exercises the
|
||||
// protocol rather than two structs calling each other.
|
||||
func start(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
return startWith(t, slog.New(slog.DiscardHandler), cfgFns...)
|
||||
}
|
||||
|
||||
// startLogged is start with logging on, for diagnosing protocol behaviour.
|
||||
func startLogged(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
return startWith(t, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug})), cfgFns...)
|
||||
}
|
||||
|
||||
func startWith(t *testing.T, log *slog.Logger, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
|
||||
inv := inventory.NewStore()
|
||||
rec := newRecorder()
|
||||
svc := agentsvc.New(inv, rec, log)
|
||||
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
|
||||
var streams atomic.Int64
|
||||
counting := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
streams.Add(1)
|
||||
mux.ServeHTTP(w, r)
|
||||
})
|
||||
|
||||
raw, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ln := &dropListener{Listener: raw}
|
||||
// connect's bidirectional streams need HTTP/2; h2c provides it without TLS.
|
||||
srv := httptest.NewUnstartedServer(h2c.NewHandler(counting, &http2.Server{}))
|
||||
srv.Listener = ln
|
||||
srv.Start()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "n1"
|
||||
cfg.Server = srv.URL
|
||||
cfg.MgmtAddr = "127.0.0.1"
|
||||
cfg.HeartbeatInterval = 50 * time.Millisecond
|
||||
cfg.ReconnectBase = 20 * time.Millisecond
|
||||
cfg.ReconnectMax = 100 * time.Millisecond
|
||||
cfg.Discovery.SkipGPUs = true
|
||||
cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
|
||||
cfg.Superv = supervisor.Config{
|
||||
WorkDir: t.TempDir(), LogLines: 64,
|
||||
StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond,
|
||||
}
|
||||
for _, f := range cfgFns {
|
||||
f(&cfg)
|
||||
}
|
||||
|
||||
client := &http.Client{Transport: &http2.Transport{
|
||||
AllowHTTP: true,
|
||||
DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
|
||||
return (&net.Dialer{}).DialContext(ctx, network, addr)
|
||||
},
|
||||
}}
|
||||
|
||||
ag, err := New(cfg, client, log)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() { _ = ag.Run(ctx) }()
|
||||
|
||||
h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag, ln: ln, streams: &streams}
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
ag.Supervisor().StopAll(context.Background())
|
||||
srv.Close()
|
||||
})
|
||||
h.waitRegistered(t)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *harness) waitRegistered(t *testing.T) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, ok := h.inv.Get("n1"); ok && len(h.svc.ConnectedNodes()) == 1 {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the agent never registered")
|
||||
}
|
||||
|
||||
// shardCommand builds a single-node vLLM shard on a free port.
|
||||
func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand {
|
||||
t.Helper()
|
||||
return reconciler.ShardCommand{
|
||||
ReplicaID: "rep-1", ShardID: shardID, NodeID: "n1",
|
||||
Spec: engines.ShardSpec{
|
||||
ReplicaID: "rep-1", ShardID: shardID, Rank: 0, NNodes: 1,
|
||||
GPUIndices: []int{0}, HTTPPort: freePort(t),
|
||||
Profile: engines.Profile{
|
||||
Name: "p", Engine: "vllm", TP: 1,
|
||||
ModelPath: "/models/x", ServedModelName: "test-model",
|
||||
},
|
||||
},
|
||||
StartupTimeout: 15 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
func freePort(t *testing.T) int {
|
||||
t.Helper()
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ln.Close()
|
||||
return ln.Addr().(*net.TCPAddr).Port
|
||||
}
|
||||
|
||||
func TestRegistrationCarriesInventory(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Hostname = "gpu-a-01"
|
||||
c.Labels = map[string]string{"rack": "r12"}
|
||||
c.Fabrics = []inventory.FabricAttachment{
|
||||
{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200},
|
||||
}
|
||||
})
|
||||
|
||||
n, ok := h.inv.Get("n1")
|
||||
if !ok {
|
||||
t.Fatal("node missing from inventory")
|
||||
}
|
||||
if n.Hostname != "gpu-a-01" || n.MgmtAddr != "127.0.0.1" {
|
||||
t.Errorf("identity wrong: %+v", n)
|
||||
}
|
||||
if n.Labels["rack"] != "r12" {
|
||||
t.Errorf("labels lost: %v", n.Labels)
|
||||
}
|
||||
// The fabric label cannot be inferred from the wire, so losing it in
|
||||
// transit would silently disable multi-node placement.
|
||||
if !n.HasFabric("ib-a") {
|
||||
t.Fatalf("fabric attachment lost: %+v", n.Fabrics)
|
||||
}
|
||||
if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
|
||||
t.Errorf("HCAs=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeartbeatKeepsNodeAlive(t *testing.T) {
|
||||
h := start(t)
|
||||
|
||||
n, _ := h.inv.Get("n1")
|
||||
first := n.LastSeen
|
||||
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
n, _ = h.inv.Get("n1")
|
||||
if !n.LastSeen.After(first) {
|
||||
t.Error("heartbeats are not refreshing LastSeen; the node would be declared lost")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartAndStopShardEndToEnd(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = servingEngine(t)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-a")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got := h.rec.await(t, 30*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-a" && s.state == supervisor.Ready
|
||||
})
|
||||
if got.replicaID != "rep-1" {
|
||||
t.Errorf("replica id lost in transit: %q", got.replicaID)
|
||||
}
|
||||
|
||||
if err := h.svc.StopShard(t.Context(), "n1", "shard-a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-a" && s.state == supervisor.Stopped
|
||||
})
|
||||
}
|
||||
|
||||
func TestFailedShardCarriesDiagnosis(t *testing.T) {
|
||||
// When a model fails to load the engine's output is the only diagnosis, so
|
||||
// it has to survive the trip to the control plane.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `echo 'CUDA out of memory' >&2; exit 1`)
|
||||
})
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), shardCommand(t, "shard-bad")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := h.rec.await(t, 30*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-bad" && s.state == supervisor.Failed
|
||||
})
|
||||
if got.err == "" {
|
||||
t.Error("the failure carried no error message")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineThatNeverBecomesReady(t *testing.T) {
|
||||
// A process that runs but never serves must time out rather than hang the
|
||||
// replica forever.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-stuck")
|
||||
cmd.StartupTimeout = 500 * time.Millisecond
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 20*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-stuck" && s.state == supervisor.Failed
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidationRejectsBeforeLaunch(t *testing.T) {
|
||||
// An engine that refuses a profile does so after minutes of loading;
|
||||
// catching it at dispatch costs nothing.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = servingEngine(t)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-invalid")
|
||||
cmd.Spec.Profile.ModelPath = "" // validation requires one
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-invalid" && s.state == supervisor.Failed
|
||||
})
|
||||
if got.err == "" {
|
||||
t.Error("validation failure carried no reason")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchToUnconnectedNode(t *testing.T) {
|
||||
h := start(t)
|
||||
|
||||
err := h.svc.StartShard(t.Context(), reconciler.ShardCommand{NodeID: "nowhere", ShardID: "x"})
|
||||
if err == nil {
|
||||
t.Fatal("dispatching to an unconnected node succeeded")
|
||||
}
|
||||
if !errors.Is(err, agentsvc.ErrNotConnected) {
|
||||
t.Errorf("err=%v, expected ErrNotConnected", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentReconnects(t *testing.T) {
|
||||
// A control plane restart must not require touching every GPU node.
|
||||
h := start(t)
|
||||
|
||||
before := h.streams.Load()
|
||||
h.ln.drop()
|
||||
|
||||
// Counting streams, not connected nodes: a stale entry in the connection
|
||||
// map would make a node look connected without any reconnect having
|
||||
// happened, and an assertion on it would pass vacuously.
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if h.streams.Load() > before && len(h.svc.ConnectedNodes()) == 1 {
|
||||
if err := h.svc.DrainShard(t.Context(), "n1", "nothing"); err != nil {
|
||||
t.Fatalf("the reconnected stream does not accept commands: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("the agent never reconnected (streams %d -> %d)", before, h.streams.Load())
|
||||
}
|
||||
|
||||
func TestIdempotentStart(t *testing.T) {
|
||||
// The control plane retries commands; a repeat must not start a second
|
||||
// process (plan §5.6).
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 30`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-idem")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-idem" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
sup := h.ag.Supervisor()
|
||||
sh, err := sup.Get("shard-idem")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
|
||||
sh2, err := sup.Get("shard-idem")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sh2.Status().PID != pid {
|
||||
t.Errorf("a second process was started: pid %d then %d", pid, sh2.Status().PID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunningShardsReportedOnReconnect covers the control plane restart hole
|
||||
// (plan §6.5): a reconnecting agent has to describe what it is already running,
|
||||
// or the control plane concludes the cluster is empty and places duplicates
|
||||
// while the originals keep holding GPUs.
|
||||
func TestRunningShardsReportedOnReconnect(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-live")
|
||||
cmd.Pool = "ib-a/h100-80g"
|
||||
cmd.Fabric = "ib-a"
|
||||
cmd.Spec.HCAs = []string{"mlx5_0"}
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-live" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
// Drop the stream, as a control plane restart would.
|
||||
h.ln.drop()
|
||||
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if got := h.rec.adoptedShards(); len(got) > 0 {
|
||||
a := got[0]
|
||||
if a.ShardID != "shard-live" || a.ReplicaID != "rep-1" {
|
||||
t.Fatalf("wrong shard reported: %+v", a)
|
||||
}
|
||||
// Everything the control plane needs to rebuild the replica.
|
||||
if a.Model != "test-model" || a.Profile != "p" {
|
||||
t.Errorf("catalog identity lost: model=%q profile=%q", a.Model, a.Profile)
|
||||
}
|
||||
if a.Pool != "ib-a/h100-80g" || a.Fabric != "ib-a" {
|
||||
t.Errorf("pool/fabric lost: %q %q", a.Pool, a.Fabric)
|
||||
}
|
||||
if !slices.Equal(a.HCAs, []string{"mlx5_0"}) {
|
||||
t.Errorf("HCAs lost: %v", a.HCAs)
|
||||
}
|
||||
if a.Endpoint == "" {
|
||||
t.Error("rank 0 reported no endpoint; the router would have nowhere to send")
|
||||
}
|
||||
if len(a.GPUIndices) == 0 {
|
||||
t.Error("GPU indices lost; the inventory could not reclaim them")
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the reconnecting agent never reported its running shard")
|
||||
}
|
||||
|
||||
// TestUnclaimedShardIsStopped covers the other half: a shard whose model left
|
||||
// the catalog is holding GPUs nothing will ever route to.
|
||||
func TestUnclaimedShardIsStopped(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-orphaned")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-orphaned" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
sh, err := h.ag.Supervisor().Get("shard-orphaned")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// The control plane will refuse to claim it on reconnect.
|
||||
h.rec.markUnclaimed("shard-orphaned")
|
||||
h.ln.drop()
|
||||
|
||||
deadline := time.Now().Add(20 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAliveForTest(pid) {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Error("an unclaimed shard kept running and holding its GPUs")
|
||||
}
|
||||
|
||||
// processAliveForTest mirrors the supervisor's liveness check.
|
||||
func processAliveForTest(pid int) bool {
|
||||
return syscall.Kill(pid, 0) == nil
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
)
|
||||
|
||||
// ErrNoNVML means this binary was built without the NVML backend.
|
||||
//
|
||||
// It is declared here rather than in the build-tagged file so callers — and
|
||||
// tests — can refer to it in either build.
|
||||
var ErrNoNVML = errors.New("built without the nvml tag; rebuild with -tags nvml for NVLink topology")
|
||||
|
||||
// Options configure a full inventory scan.
|
||||
type Options struct {
|
||||
// NvidiaSMI is the path to nvidia-smi; empty means look it up on PATH.
|
||||
// Worth configuring on nodes where the driver lives outside the default
|
||||
// paths, which containerized driver installs often do.
|
||||
NvidiaSMI string `koanf:"nvidia_smi"`
|
||||
// IBRoot is the sysfs InfiniBand directory.
|
||||
IBRoot string `koanf:"ib_root"`
|
||||
// PreferNVML uses the NVML backend when this binary has one. It is the only
|
||||
// source of NVLink topology.
|
||||
PreferNVML bool `koanf:"prefer_nvml"`
|
||||
// SkipGPUs omits GPU discovery, for nodes that have none.
|
||||
SkipGPUs bool `koanf:"skip_gpus"`
|
||||
}
|
||||
|
||||
// Discover builds the node's inventory.
|
||||
//
|
||||
// When NVML is preferred but fails, it falls back to nvidia-smi and says so in
|
||||
// Inventory.Source. The fallback is logged rather than silent: on a node that
|
||||
// is supposed to have working NVML, a quiet downgrade hides a broken driver and
|
||||
// costs the scheduler its NVLink topology.
|
||||
func Discover(ctx context.Context, opts Options, log *slog.Logger) (Inventory, error) {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
inv := Inventory{Source: "none"}
|
||||
|
||||
if !opts.SkipGPUs {
|
||||
switch {
|
||||
case opts.PreferNVML && nvmlAvailable:
|
||||
gpus, err := DiscoverGPUsNVML(ctx)
|
||||
if err == nil {
|
||||
inv.GPUs, inv.Source = gpus, "nvml"
|
||||
break
|
||||
}
|
||||
log.Warn("NVML discovery failed, falling back to nvidia-smi; "+
|
||||
"NVLink topology will be unavailable", "err", err)
|
||||
fallthrough
|
||||
default:
|
||||
gpus, err := DiscoverGPUsSMI(ctx, opts.NvidiaSMI)
|
||||
if err != nil {
|
||||
return inv, fmt.Errorf("gpu discovery: %w", err)
|
||||
}
|
||||
inv.GPUs, inv.Source = gpus, "nvidia-smi"
|
||||
}
|
||||
}
|
||||
|
||||
hcas, err := DiscoverIB(opts.IBRoot)
|
||||
if err != nil {
|
||||
return inv, fmt.Errorf("infiniband discovery: %w", err)
|
||||
}
|
||||
inv.HCAs = hcas
|
||||
return inv, nil
|
||||
}
|
||||
|
||||
// NVMLSupported reports whether this binary includes the NVML backend.
|
||||
func NVMLSupported() bool { return nvmlAvailable }
|
||||
@@ -0,0 +1,101 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func fakeSMI(t *testing.T) string {
|
||||
t.Helper()
|
||||
stub := filepath.Join(t.TempDir(), "nvidia-smi")
|
||||
script := "#!/bin/sh\ncat <<'EOF'\n" + smiSample + "EOF\n"
|
||||
if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return stub
|
||||
}
|
||||
|
||||
func TestDiscoverUsesSMIByDefault(t *testing.T) {
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
NvidiaSMI: fakeSMI(t),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if inv.Source != "nvidia-smi" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
if len(inv.GPUs) != 2 {
|
||||
t.Errorf("got %d GPUs", len(inv.GPUs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverFallsBackWhenNVMLUnavailable(t *testing.T) {
|
||||
// Asking for NVML in a build without it must still produce an inventory:
|
||||
// the fallback is the whole point of the build tag.
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
PreferNVML: true,
|
||||
NvidiaSMI: fakeSMI(t),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(inv.GPUs) == 0 {
|
||||
t.Fatal("no GPUs discovered through the fallback")
|
||||
}
|
||||
if NVMLSupported() {
|
||||
// On an NVML build without a driver the NVML attempt fails and the
|
||||
// fallback still has to work.
|
||||
if inv.Source != "nvml" && inv.Source != "nvidia-smi" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
} else if inv.Source != "nvidia-smi" {
|
||||
t.Errorf("a build without NVML reported source=%q", inv.Source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverSkipGPUs(t *testing.T) {
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
SkipGPUs: true,
|
||||
IBRoot: fakeSysfs(t),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(inv.GPUs) != 0 {
|
||||
t.Errorf("GPUs were discovered despite SkipGPUs")
|
||||
}
|
||||
if len(inv.HCAs) != 3 {
|
||||
t.Errorf("got %d HCAs", len(inv.HCAs))
|
||||
}
|
||||
if inv.Source != "none" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverReportsGPUFailure(t *testing.T) {
|
||||
// A node meant to have GPUs whose discovery fails must say so rather than
|
||||
// reporting an empty, healthy-looking inventory.
|
||||
_, err := Discover(context.Background(), Options{
|
||||
NvidiaSMI: filepath.Join(t.TempDir(), "absent"),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err == nil {
|
||||
t.Error("a failed GPU discovery was reported as success")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNVMLSupportedMatchesBuild(t *testing.T) {
|
||||
// Documents which backend this binary has, and keeps both files compiling.
|
||||
t.Logf("NVML backend compiled in: %v", NVMLSupported())
|
||||
if !NVMLSupported() {
|
||||
if _, err := DiscoverGPUsNVML(t.Context()); err != ErrNoNVML {
|
||||
t.Errorf("expected ErrNoNVML, got %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// GPU is one accelerator on this node.
|
||||
type GPU struct {
|
||||
Index int `json:"index"`
|
||||
UUID string `json:"uuid"`
|
||||
Name string `json:"name"`
|
||||
// MemoryTotalMB is the total framebuffer size.
|
||||
MemoryTotalMB int `json:"memory_total_mb"`
|
||||
// MemoryUsedMB is what is already allocated. Non-zero on an idle node means
|
||||
// something other than our engines is using the card (plan §5.1).
|
||||
MemoryUsedMB int `json:"memory_used_mb"`
|
||||
// ComputeCapability gates quantization choices: an FP8 profile must not
|
||||
// land on a card below 8.9 (plan §6.2).
|
||||
ComputeCapability string `json:"compute_capability"`
|
||||
NUMANode int `json:"numa_node"`
|
||||
PCIBusID string `json:"pci_bus_id,omitempty"`
|
||||
DriverVersion string `json:"driver_version,omitempty"`
|
||||
CUDAVersion string `json:"cuda_version,omitempty"`
|
||||
// NVLinkPeers are the indices of GPUs reachable over NVLink. Within a node
|
||||
// a TP group should stay inside one NVLink domain.
|
||||
NVLinkPeers []int `json:"nvlink_peers,omitempty"`
|
||||
// Health carries anything that should keep the GPU out of placement.
|
||||
Health GPUHealth `json:"health"`
|
||||
}
|
||||
|
||||
// GPUHealth records conditions that make a GPU unsafe to schedule onto.
|
||||
type GPUHealth struct {
|
||||
// ECCErrors is the count of uncorrected ECC errors.
|
||||
ECCErrors int `json:"ecc_errors"`
|
||||
// TemperatureC is the current die temperature.
|
||||
TemperatureC int `json:"temperature_c"`
|
||||
// Problems lists human-readable reasons to avoid this GPU.
|
||||
Problems []string `json:"problems,omitempty"`
|
||||
}
|
||||
|
||||
// OK reports whether the GPU is fit for placement.
|
||||
func (h GPUHealth) OK() bool { return len(h.Problems) == 0 }
|
||||
|
||||
// nonAlnum collapses everything that is not a letter or digit.
|
||||
var nonAlnum = regexp.MustCompile(`[^a-z0-9]+`)
|
||||
|
||||
// Class derives the GPU class used for placement: same model, same memory.
|
||||
//
|
||||
// Two cards only belong in the same tensor-parallel group if they are the same
|
||||
// class, so the derivation has to be stable across driver versions and naming
|
||||
// quirks ("NVIDIA H100 80GB HBM3" and "H100-SXM5-80GB" are the same class).
|
||||
func (g GPU) Class() string {
|
||||
name := strings.ToLower(g.Name)
|
||||
name = strings.TrimPrefix(name, "nvidia ")
|
||||
name = strings.TrimPrefix(name, "tesla ")
|
||||
|
||||
// Keep the model token: the first word that contains a digit, which is how
|
||||
// NVIDIA names every accelerator (h100, a100, l40s, rtx4090).
|
||||
var model string
|
||||
for _, tok := range nonAlnum.Split(name, -1) {
|
||||
if tok == "" || !strings.ContainsFunc(tok, func(r rune) bool { return r >= '0' && r <= '9' }) {
|
||||
continue
|
||||
}
|
||||
// Skip memory and form-factor tokens, which are not the model.
|
||||
if strings.HasSuffix(tok, "gb") || tok == "80gb" || tok == "sxm4" || tok == "sxm5" || tok == "pcie" {
|
||||
continue
|
||||
}
|
||||
model = tok
|
||||
break
|
||||
}
|
||||
if model == "" {
|
||||
model = nonAlnum.ReplaceAllString(name, "-")
|
||||
}
|
||||
|
||||
// Round to the nearest GiB: vendors report 81559 MiB for an "80GB" card.
|
||||
gb := int(math.Round(float64(g.MemoryTotalMB) / 1024))
|
||||
return fmt.Sprintf("%s-%dg", model, gb)
|
||||
}
|
||||
|
||||
// MarshalJSON adds the derived class, so a consumer of the JSON does not have
|
||||
// to reimplement the derivation to know which GPUs are interchangeable.
|
||||
func (g GPU) MarshalJSON() ([]byte, error) {
|
||||
type alias GPU // prevents recursing back into this method
|
||||
return json.Marshal(struct {
|
||||
alias
|
||||
Class string `json:"class"`
|
||||
}{alias(g), g.Class()})
|
||||
}
|
||||
|
||||
// Inventory is everything discovery found on this node.
|
||||
type Inventory struct {
|
||||
GPUs []GPU `json:"gpus"`
|
||||
HCAs []HCA `json:"hcas"`
|
||||
// Source names the backend that produced the GPU list, so a consumer can
|
||||
// tell "no GPUs" from "discovery did not work".
|
||||
Source string `json:"source"`
|
||||
}
|
||||
|
||||
// Classes returns the distinct GPU classes present, sorted.
|
||||
func (inv Inventory) Classes() []string {
|
||||
seen := map[string]struct{}{}
|
||||
for _, g := range inv.GPUs {
|
||||
seen[g.Class()] = struct{}{}
|
||||
}
|
||||
out := make([]string, 0, len(seen))
|
||||
for c := range seen {
|
||||
out = append(out, c)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// parseCSVInt reads an integer from nvidia-smi's CSV output, tolerating the
|
||||
// "[N/A]" and "Not Supported" values it emits for unsupported queries.
|
||||
func parseCSVInt(s string) (int, bool) {
|
||||
s = strings.TrimSpace(s)
|
||||
s = strings.TrimSuffix(s, " MiB")
|
||||
s = strings.TrimSuffix(s, " W")
|
||||
if s == "" || strings.HasPrefix(s, "[") || strings.EqualFold(s, "Not Supported") {
|
||||
return 0, false
|
||||
}
|
||||
v, err := strconv.Atoi(s)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return v, true
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
//go:build !nvml
|
||||
|
||||
package discovery
|
||||
|
||||
import "context"
|
||||
|
||||
// nvmlAvailable reports that this build cannot use NVML.
|
||||
const nvmlAvailable = false
|
||||
|
||||
// DiscoverGPUsNVML is unavailable in builds without the nvml tag.
|
||||
func DiscoverGPUsNVML(context.Context) ([]GPU, error) { return nil, ErrNoNVML }
|
||||
@@ -0,0 +1,148 @@
|
||||
//go:build nvml
|
||||
|
||||
// This file is the NVML-backed GPU discovery path, built only with the `nvml`
|
||||
// tag because the bindings require cgo. Without the tag the agent falls back to
|
||||
// nvidia-smi, which keeps the server side — and development machines without a
|
||||
// C toolchain or a driver — building cleanly (plan §3.2).
|
||||
//
|
||||
// What NVML adds over nvidia-smi is NVLink topology. Nothing else exposes which
|
||||
// GPUs share an NVLink domain, and that is what keeps a tensor-parallel group
|
||||
// off the slow path inside a node (plan §6.2, constraint 3).
|
||||
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/NVIDIA/go-nvml/pkg/nvml"
|
||||
)
|
||||
|
||||
// nvmlAvailable reports that this build can use NVML.
|
||||
const nvmlAvailable = true
|
||||
|
||||
// DiscoverGPUsNVML enumerates GPUs through NVML.
|
||||
//
|
||||
// It returns an error rather than falling back on its own: the caller decides
|
||||
// whether to degrade to nvidia-smi, because on a node that is supposed to have
|
||||
// working NVML a silent downgrade hides a broken driver installation.
|
||||
func DiscoverGPUsNVML(_ context.Context) ([]GPU, error) {
|
||||
if ret := nvml.Init(); ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml init: %s", nvml.ErrorString(ret))
|
||||
}
|
||||
defer func() { _ = nvml.Shutdown() }()
|
||||
|
||||
driver, _ := nvml.SystemGetDriverVersion()
|
||||
cuda := ""
|
||||
if v, ret := nvml.SystemGetCudaDriverVersion(); ret == nvml.SUCCESS {
|
||||
// NVML reports 12040 for 12.4.
|
||||
cuda = fmt.Sprintf("%d.%d", v/1000, (v%1000)/10)
|
||||
}
|
||||
|
||||
count, ret := nvml.DeviceGetCount()
|
||||
if ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml device count: %s", nvml.ErrorString(ret))
|
||||
}
|
||||
|
||||
// Bus id to index, so NVLink peers can be reported as indices. NVML gives
|
||||
// the remote end of a link as a PCI address, which is not what placement
|
||||
// works in.
|
||||
byBusID := make(map[string]int, count)
|
||||
gpus := make([]GPU, 0, count)
|
||||
|
||||
for i := range count {
|
||||
dev, ret := nvml.DeviceGetHandleByIndex(i)
|
||||
if ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml device %d: %s", i, nvml.ErrorString(ret))
|
||||
}
|
||||
|
||||
g := GPU{Index: i, NUMANode: -1, DriverVersion: driver, CUDAVersion: cuda}
|
||||
g.UUID, _ = dev.GetUUID()
|
||||
g.Name, _ = dev.GetName()
|
||||
|
||||
if mem, ret := dev.GetMemoryInfo(); ret == nvml.SUCCESS {
|
||||
g.MemoryTotalMB = int(mem.Total / (1 << 20))
|
||||
g.MemoryUsedMB = int(mem.Used / (1 << 20))
|
||||
}
|
||||
if major, minor, ret := dev.GetCudaComputeCapability(); ret == nvml.SUCCESS {
|
||||
g.ComputeCapability = fmt.Sprintf("%d.%d", major, minor)
|
||||
}
|
||||
if pci, ret := dev.GetPciInfo(); ret == nvml.SUCCESS {
|
||||
g.PCIBusID = busIDString(pci.BusId[:])
|
||||
byBusID[normalizeBusID(g.PCIBusID)] = i
|
||||
}
|
||||
if t, ret := dev.GetTemperature(nvml.TEMPERATURE_GPU); ret == nvml.SUCCESS {
|
||||
g.Health.TemperatureC = int(t)
|
||||
}
|
||||
// Volatile counters reset on driver reload, which is what we want: an
|
||||
// error from three months and one reboot ago is not a reason to refuse
|
||||
// a card today.
|
||||
if n, ret := dev.GetTotalEccErrors(nvml.MEMORY_ERROR_TYPE_UNCORRECTED, nvml.VOLATILE_ECC); ret == nvml.SUCCESS {
|
||||
g.Health.ECCErrors = int(n)
|
||||
}
|
||||
|
||||
assessHealth(&g)
|
||||
gpus = append(gpus, g)
|
||||
}
|
||||
|
||||
// A second pass, because a link's remote end may be a GPU enumerated later.
|
||||
for i := range gpus {
|
||||
dev, ret := nvml.DeviceGetHandleByIndex(gpus[i].Index)
|
||||
if ret != nvml.SUCCESS {
|
||||
continue
|
||||
}
|
||||
gpus[i].NVLinkPeers = nvlinkPeers(dev, byBusID, gpus[i].Index)
|
||||
}
|
||||
return gpus, nil
|
||||
}
|
||||
|
||||
// nvlinkPeers returns the indices of GPUs reachable from dev over NVLink.
|
||||
func nvlinkPeers(dev nvml.Device, byBusID map[string]int, self int) []int {
|
||||
seen := map[int]struct{}{}
|
||||
var out []int
|
||||
|
||||
for link := range nvml.NVLINK_MAX_LINKS {
|
||||
pci, ret := dev.GetNvLinkRemotePciInfo(link)
|
||||
if ret != nvml.SUCCESS {
|
||||
// Unpopulated links return an error; that is the normal case for
|
||||
// most of the 36 possible links, so it is not worth reporting.
|
||||
continue
|
||||
}
|
||||
peer, ok := byBusID[normalizeBusID(busIDString(pci.BusId[:]))]
|
||||
if !ok || peer == self {
|
||||
// The remote end can be an NVSwitch rather than a GPU, in which
|
||||
// case it is not in the map.
|
||||
continue
|
||||
}
|
||||
if _, dup := seen[peer]; dup {
|
||||
continue
|
||||
}
|
||||
seen[peer] = struct{}{}
|
||||
out = append(out, peer)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// busIDString turns NVML's fixed-size C char array into a Go string.
|
||||
func busIDString(b []int8) string {
|
||||
buf := make([]byte, 0, len(b))
|
||||
for _, c := range b {
|
||||
if c == 0 {
|
||||
break
|
||||
}
|
||||
buf = append(buf, byte(c))
|
||||
}
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
// normalizeBusID makes bus ids comparable: NVML reports them uppercase from one
|
||||
// call and lowercase from another, and sometimes with a shortened domain.
|
||||
func normalizeBusID(s string) string {
|
||||
s = strings.ToLower(strings.TrimSpace(s))
|
||||
// "0000:1b:00.0" and "00000000:1B:00.0" are the same device.
|
||||
if i := strings.Index(s, ":"); i > 0 {
|
||||
s = strings.TrimLeft(s[:i], "0") + s[i:]
|
||||
}
|
||||
return strings.TrimPrefix(s, ":")
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// smiQueryFields are the columns requested from nvidia-smi, in order.
|
||||
var smiQueryFields = []string{
|
||||
"index",
|
||||
"uuid",
|
||||
"name",
|
||||
"memory.total",
|
||||
"memory.used",
|
||||
"compute_cap",
|
||||
"pci.bus_id",
|
||||
"driver_version",
|
||||
"temperature.gpu",
|
||||
"ecc.errors.uncorrected.volatile.total",
|
||||
}
|
||||
|
||||
// DiscoverGPUsSMI enumerates GPUs by shelling out to nvidia-smi.
|
||||
//
|
||||
// This is the cgo-free path (plan §3.2): the NVML bindings need cgo, which
|
||||
// would force the whole binary — server included — to be built with a C
|
||||
// toolchain. nvidia-smi gives everything except NVLink topology, which the
|
||||
// NVML build adds.
|
||||
func DiscoverGPUsSMI(ctx context.Context, binary string) ([]GPU, error) {
|
||||
if binary == "" {
|
||||
binary = "nvidia-smi"
|
||||
}
|
||||
out, err := exec.CommandContext(ctx, binary,
|
||||
"--query-gpu="+strings.Join(smiQueryFields, ","),
|
||||
"--format=csv,noheader,nounits",
|
||||
).Output()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("running %s: %w", binary, err)
|
||||
}
|
||||
return parseSMI(string(out))
|
||||
}
|
||||
|
||||
func parseSMI(out string) ([]GPU, error) {
|
||||
r := csv.NewReader(strings.NewReader(out))
|
||||
r.TrimLeadingSpace = true
|
||||
r.FieldsPerRecord = -1 // tolerate driver versions with extra columns
|
||||
|
||||
rows, err := r.ReadAll()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parsing nvidia-smi output: %w", err)
|
||||
}
|
||||
|
||||
gpus := make([]GPU, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
if len(row) < len(smiQueryFields) {
|
||||
continue
|
||||
}
|
||||
idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
g := GPU{
|
||||
Index: idx,
|
||||
UUID: strings.TrimSpace(row[1]),
|
||||
Name: strings.TrimSpace(row[2]),
|
||||
ComputeCapability: strings.TrimSpace(row[5]),
|
||||
PCIBusID: strings.TrimSpace(row[6]),
|
||||
DriverVersion: strings.TrimSpace(row[7]),
|
||||
NUMANode: -1,
|
||||
}
|
||||
g.MemoryTotalMB, _ = parseCSVInt(row[3])
|
||||
g.MemoryUsedMB, _ = parseCSVInt(row[4])
|
||||
g.Health.TemperatureC, _ = parseCSVInt(row[8])
|
||||
g.Health.ECCErrors, _ = parseCSVInt(row[9])
|
||||
|
||||
assessHealth(&g)
|
||||
gpus = append(gpus, g)
|
||||
}
|
||||
return gpus, nil
|
||||
}
|
||||
|
||||
// residentMemoryToleranceMB is how much allocated memory is written off as
|
||||
// driver overhead rather than a foreign process.
|
||||
const residentMemoryToleranceMB = 512
|
||||
|
||||
// assessHealth flags conditions that should keep a GPU out of placement.
|
||||
func assessHealth(g *GPU) {
|
||||
if g.Health.ECCErrors > 0 {
|
||||
g.Health.Problems = append(g.Health.Problems,
|
||||
fmt.Sprintf("%d uncorrected ECC errors", g.Health.ECCErrors))
|
||||
}
|
||||
// Memory already in use means something else is on the card. Starting an
|
||||
// engine there would either fail to allocate or quietly evict the other
|
||||
// process's throughput.
|
||||
if g.MemoryUsedMB > residentMemoryToleranceMB {
|
||||
g.Health.Problems = append(g.Health.Problems,
|
||||
fmt.Sprintf("%d MiB already allocated by another process", g.MemoryUsedMB))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const smiSample = `0, GPU-11111111-2222-3333-4444-555555555555, NVIDIA H100 80GB HBM3, 81559, 4, 9.0, 00000000:1B:00.0, 550.54.15, 34, 0
|
||||
1, GPU-66666666-7777-8888-9999-000000000000, NVIDIA H100 80GB HBM3, 81559, 0, 9.0, 00000000:1C:00.0, 550.54.15, 36, 0
|
||||
`
|
||||
|
||||
func TestParseSMI(t *testing.T) {
|
||||
gpus, err := parseSMI(smiSample)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 2 {
|
||||
t.Fatalf("expected 2 GPUs, got %d", len(gpus))
|
||||
}
|
||||
|
||||
g := gpus[0]
|
||||
if g.Index != 0 || g.UUID != "GPU-11111111-2222-3333-4444-555555555555" {
|
||||
t.Errorf("identity wrong: %+v", g)
|
||||
}
|
||||
if g.MemoryTotalMB != 81559 {
|
||||
t.Errorf("MemoryTotalMB=%d", g.MemoryTotalMB)
|
||||
}
|
||||
if g.ComputeCapability != "9.0" {
|
||||
t.Errorf("ComputeCapability=%q; FP8 placement depends on it", g.ComputeCapability)
|
||||
}
|
||||
if g.DriverVersion != "550.54.15" {
|
||||
t.Errorf("DriverVersion=%q", g.DriverVersion)
|
||||
}
|
||||
if !g.Health.OK() {
|
||||
t.Errorf("a healthy GPU was flagged: %v", g.Health.Problems)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSMIUnsupportedFields(t *testing.T) {
|
||||
// Consumer cards report [N/A] for ECC; that must not fail the whole parse.
|
||||
const row = `0, GPU-abc, NVIDIA GeForce RTX 4090, 24564, 0, 8.9, 00000000:01:00.0, 550.54.15, 45, [N/A]
|
||||
`
|
||||
gpus, err := parseSMI(row)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("expected 1 GPU, got %d", len(gpus))
|
||||
}
|
||||
if gpus[0].Health.ECCErrors != 0 || !gpus[0].Health.OK() {
|
||||
t.Errorf("[N/A] ECC was misread: %+v", gpus[0].Health)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthFlagsForeignMemory(t *testing.T) {
|
||||
// Memory in use on an otherwise idle card means another process is there.
|
||||
const row = `0, GPU-abc, NVIDIA H100 80GB HBM3, 81559, 40000, 9.0, 00000000:1B:00.0, 550.54.15, 34, 0
|
||||
`
|
||||
gpus, _ := parseSMI(row)
|
||||
if gpus[0].Health.OK() {
|
||||
t.Error("40 GiB allocated by another process was not flagged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthFlagsECC(t *testing.T) {
|
||||
const row = `0, GPU-abc, NVIDIA H100 80GB HBM3, 81559, 0, 9.0, 00000000:1B:00.0, 550.54.15, 34, 3
|
||||
`
|
||||
gpus, _ := parseSMI(row)
|
||||
if gpus[0].Health.OK() {
|
||||
t.Error("uncorrected ECC errors were not flagged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGPUClass(t *testing.T) {
|
||||
// The same card must derive the same class however the driver names it,
|
||||
// because a tensor-parallel group may not mix classes.
|
||||
cases := []struct {
|
||||
name string
|
||||
memMB int
|
||||
expect string
|
||||
}{
|
||||
{"NVIDIA H100 80GB HBM3", 81559, "h100-80g"},
|
||||
{"NVIDIA H100-SXM5-80GB", 81559, "h100-80g"},
|
||||
{"Tesla A100-SXM4-40GB", 40960, "a100-40g"},
|
||||
{"NVIDIA A100 80GB PCIe", 81920, "a100-80g"},
|
||||
// Sizes come from what the driver reports, in GiB. That does not always
|
||||
// match the marketing name: an L40S is sold as 48GB but reports 44.99
|
||||
// GiB, because the marketing figure is decimal GB. Being consistent
|
||||
// matters more than matching the box, since the number is what decides
|
||||
// whether a model fits.
|
||||
{"NVIDIA L40S", 46068, "l40s-45g"},
|
||||
// The model token is what distinguishes a 4090 from a 4080, so "rtx"
|
||||
// is not enough on its own.
|
||||
{"NVIDIA GeForce RTX 4090", 24564, "4090-24g"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := GPU{Name: c.name, MemoryTotalMB: c.memMB}.Class()
|
||||
if got != c.expect {
|
||||
t.Errorf("%q (%d MiB) -> %q, want %q", c.name, c.memMB, got, c.expect)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestInventoryClasses(t *testing.T) {
|
||||
inv := Inventory{GPUs: []GPU{
|
||||
{Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559},
|
||||
{Name: "NVIDIA H100-SXM5-80GB", MemoryTotalMB: 81559},
|
||||
{Name: "Tesla A100-SXM4-40GB", MemoryTotalMB: 40960},
|
||||
}}
|
||||
got := inv.Classes()
|
||||
if len(got) != 2 || got[0] != "a100-40g" || got[1] != "h100-80g" {
|
||||
t.Errorf("classes=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverGPUsSMI runs the real code path against a stub binary, so the
|
||||
// command construction is covered too, not just the parser.
|
||||
func TestDiscoverGPUsSMI(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
stub := filepath.Join(dir, "nvidia-smi")
|
||||
script := "#!/bin/sh\ncat <<'EOF'\n" + smiSample + "EOF\n"
|
||||
if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
gpus, err := DiscoverGPUsSMI(t.Context(), stub)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 2 {
|
||||
t.Fatalf("expected 2 GPUs, got %d", len(gpus))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverGPUsSMIMissingBinary(t *testing.T) {
|
||||
// A node with no NVIDIA driver must produce a clear error rather than a
|
||||
// panic or an empty success.
|
||||
_, err := DiscoverGPUsSMI(context.Background(), filepath.Join(t.TempDir(), "absent"))
|
||||
if err == nil {
|
||||
t.Error("expected an error when nvidia-smi is missing")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
// Package discovery inspects a node's GPUs, InfiniBand adapters and the
|
||||
// topology between them (plan §5.1 and §5.2).
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// SysfsRoot is the default place InfiniBand devices appear.
|
||||
const SysfsRoot = "/sys/class/infiniband"
|
||||
|
||||
// LinkLayer distinguishes real InfiniBand from RoCE.
|
||||
type LinkLayer string
|
||||
|
||||
const (
|
||||
// LinkInfiniBand is native InfiniBand.
|
||||
LinkInfiniBand LinkLayer = "InfiniBand"
|
||||
// LinkEthernet is RoCE — Ethernet carrying RDMA.
|
||||
LinkEthernet LinkLayer = "Ethernet"
|
||||
)
|
||||
|
||||
// Port is one port of an HCA.
|
||||
type Port struct {
|
||||
Num int `json:"num"`
|
||||
// State is the port's logical state; only ACTIVE ports carry traffic.
|
||||
State string `json:"state"`
|
||||
// PhysState is the physical link state, useful for telling "cable unplugged"
|
||||
// from "no subnet manager".
|
||||
PhysState string `json:"phys_state"`
|
||||
// RateGbps is the negotiated rate.
|
||||
RateGbps int `json:"rate_gbps"`
|
||||
LinkLayer LinkLayer `json:"link_layer"`
|
||||
// GIDPrefix is the subnet prefix from gids/0, normalized to lowercase hex
|
||||
// groups. It is the only automatic fabric hint available, and an unreliable
|
||||
// one: OpenSM's default is fe80:: everywhere (plan §5.2).
|
||||
GIDPrefix string `json:"gid_prefix"`
|
||||
LID int `json:"lid"`
|
||||
SMLID int `json:"sm_lid"`
|
||||
}
|
||||
|
||||
// Active reports whether the port can carry traffic.
|
||||
func (p Port) Active() bool { return strings.EqualFold(p.State, "ACTIVE") }
|
||||
|
||||
// HCA is one host channel adapter.
|
||||
type HCA struct {
|
||||
Name string `json:"name"`
|
||||
Ports []Port `json:"ports"`
|
||||
PCIAddr string `json:"pci_addr,omitempty"`
|
||||
NUMANode int `json:"numa_node"`
|
||||
// FirmwareVersion is reported by the driver; worth capturing because
|
||||
// NCCL behaviour differs between firmware revisions.
|
||||
FirmwareVersion string `json:"firmware_version,omitempty"`
|
||||
}
|
||||
|
||||
// ActivePorts returns the ports that can carry traffic.
|
||||
func (h HCA) ActivePorts() []Port {
|
||||
var out []Port
|
||||
for _, p := range h.Ports {
|
||||
if p.Active() {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// DiscoverIB enumerates the InfiniBand adapters under root.
|
||||
//
|
||||
// A missing root is not an error: plenty of nodes have no IB at all, and the
|
||||
// scheduler treats them as single-node islands (plan §2).
|
||||
func DiscoverIB(root string) ([]HCA, error) {
|
||||
if root == "" {
|
||||
root = SysfsRoot
|
||||
}
|
||||
entries, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("reading %s: %w", root, err)
|
||||
}
|
||||
|
||||
var out []HCA
|
||||
for _, e := range entries {
|
||||
hca, err := readHCA(filepath.Join(root, e.Name()), e.Name())
|
||||
if err != nil {
|
||||
// One unreadable adapter must not hide the rest; the node may still
|
||||
// be usable through its other HCAs.
|
||||
continue
|
||||
}
|
||||
out = append(out, hca)
|
||||
}
|
||||
slices.SortFunc(out, func(a, b HCA) int { return strings.Compare(a.Name, b.Name) })
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func readHCA(dir, name string) (HCA, error) {
|
||||
h := HCA{Name: name, NUMANode: -1}
|
||||
h.FirmwareVersion = readTrimmed(filepath.Join(dir, "fw_ver"))
|
||||
|
||||
// device is a symlink into the PCI tree; its target's base name is the
|
||||
// bus address, and numa_node lives beside it.
|
||||
if target, err := filepath.EvalSymlinks(filepath.Join(dir, "device")); err == nil {
|
||||
h.PCIAddr = filepath.Base(target)
|
||||
if n, err := strconv.Atoi(readTrimmed(filepath.Join(target, "numa_node"))); err == nil {
|
||||
h.NUMANode = n
|
||||
}
|
||||
}
|
||||
|
||||
portDirs, err := os.ReadDir(filepath.Join(dir, "ports"))
|
||||
if err != nil {
|
||||
return h, fmt.Errorf("%s: no ports: %w", name, err)
|
||||
}
|
||||
for _, pd := range portDirs {
|
||||
num, err := strconv.Atoi(pd.Name())
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
h.Ports = append(h.Ports, readPort(filepath.Join(dir, "ports", pd.Name()), num))
|
||||
}
|
||||
slices.SortFunc(h.Ports, func(a, b Port) int { return a.Num - b.Num })
|
||||
return h, nil
|
||||
}
|
||||
|
||||
func readPort(dir string, num int) Port {
|
||||
p := Port{
|
||||
Num: num,
|
||||
State: parseEnumeratedState(readTrimmed(filepath.Join(dir, "state"))),
|
||||
PhysState: parseEnumeratedState(readTrimmed(filepath.Join(dir, "phys_state"))),
|
||||
RateGbps: parseRate(readTrimmed(filepath.Join(dir, "rate"))),
|
||||
LinkLayer: LinkLayer(readTrimmed(filepath.Join(dir, "link_layer"))),
|
||||
GIDPrefix: parseGIDPrefix(readTrimmed(filepath.Join(dir, "gids", "0"))),
|
||||
LID: parseMaybeHex(readTrimmed(filepath.Join(dir, "lid"))),
|
||||
SMLID: parseMaybeHex(readTrimmed(filepath.Join(dir, "sm_lid"))),
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// parseEnumeratedState turns sysfs's "4: ACTIVE" into "ACTIVE".
|
||||
func parseEnumeratedState(s string) string {
|
||||
if _, rest, found := strings.Cut(s, ":"); found {
|
||||
return strings.TrimSpace(rest)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// parseRate turns "100 Gb/sec (4X EDR)" into 100.
|
||||
func parseRate(s string) int {
|
||||
f := strings.Fields(s)
|
||||
if len(f) == 0 {
|
||||
return 0
|
||||
}
|
||||
// The rate can be fractional ("2.5 Gb/sec") on old hardware; truncate.
|
||||
v, err := strconv.ParseFloat(f[0], 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return int(v)
|
||||
}
|
||||
|
||||
// parseGIDPrefix keeps the first four groups of a GID, which is the subnet
|
||||
// prefix; the remaining four are the interface id.
|
||||
func parseGIDPrefix(gid string) string {
|
||||
parts := strings.Split(strings.ToLower(gid), ":")
|
||||
if len(parts) < 4 {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(parts[:4], ":")
|
||||
}
|
||||
|
||||
// parseMaybeHex accepts both "0x1a" and "26"; sysfs is inconsistent between
|
||||
// drivers and kernel versions.
|
||||
func parseMaybeHex(s string) int {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return 0
|
||||
}
|
||||
if after, ok := strings.CutPrefix(strings.ToLower(s), "0x"); ok {
|
||||
v, err := strconv.ParseInt(after, 16, 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return int(v)
|
||||
}
|
||||
v, err := strconv.Atoi(s)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func readTrimmed(path string) string {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// fakeSysfs builds an InfiniBand tree shaped like the real one.
|
||||
func fakeSysfs(t *testing.T) string {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
|
||||
write := func(path, content string) {
|
||||
t.Helper()
|
||||
full := filepath.Join(root, path)
|
||||
if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(full, []byte(content+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// mlx5_0: one active IB port.
|
||||
write("ib/mlx5_0/fw_ver", "20.35.1012")
|
||||
write("ib/mlx5_0/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_0/ports/1/phys_state", "5: LinkUp")
|
||||
write("ib/mlx5_0/ports/1/rate", "200 Gb/sec (4X HDR)")
|
||||
write("ib/mlx5_0/ports/1/link_layer", "InfiniBand")
|
||||
write("ib/mlx5_0/ports/1/gids/0", "FEC0:0000:0000:000A:0002:C903:00A0:1234")
|
||||
write("ib/mlx5_0/ports/1/lid", "0x1a")
|
||||
write("ib/mlx5_0/ports/1/sm_lid", "1")
|
||||
|
||||
// mlx5_1: two ports, the second one down.
|
||||
write("ib/mlx5_1/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_1/ports/1/rate", "200 Gb/sec (4X HDR)")
|
||||
write("ib/mlx5_1/ports/1/link_layer", "InfiniBand")
|
||||
write("ib/mlx5_1/ports/1/gids/0", "fec0:0000:0000:000a:0002:c903:00a0:5678")
|
||||
write("ib/mlx5_1/ports/2/state", "1: DOWN")
|
||||
write("ib/mlx5_1/ports/2/phys_state", "3: Disabled")
|
||||
write("ib/mlx5_1/ports/2/rate", "0")
|
||||
write("ib/mlx5_1/ports/2/link_layer", "InfiniBand")
|
||||
|
||||
// mlx5_2: RoCE rather than native IB.
|
||||
write("ib/mlx5_2/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_2/ports/1/rate", "100 Gb/sec (4X EDR)")
|
||||
write("ib/mlx5_2/ports/1/link_layer", "Ethernet")
|
||||
write("ib/mlx5_2/ports/1/gids/0", "0000:0000:0000:0000:0000:ffff:0a00:0001")
|
||||
|
||||
// A PCI device with a NUMA node, linked the way the driver does it.
|
||||
pci := filepath.Join(root, "pci", "0000:0b:00.0")
|
||||
if err := os.MkdirAll(pci, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(pci, "numa_node"), []byte("1\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.Symlink(pci, filepath.Join(root, "ib", "mlx5_0", "device")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
return filepath.Join(root, "ib")
|
||||
}
|
||||
|
||||
func TestDiscoverIB(t *testing.T) {
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hcas) != 3 {
|
||||
t.Fatalf("expected 3 HCAs, got %d", len(hcas))
|
||||
}
|
||||
// Sorted by name, so the order is predictable for the rest of the test.
|
||||
if hcas[0].Name != "mlx5_0" || hcas[2].Name != "mlx5_2" {
|
||||
t.Fatalf("HCAs are not sorted by name: %v %v", hcas[0].Name, hcas[2].Name)
|
||||
}
|
||||
|
||||
h := hcas[0]
|
||||
if h.FirmwareVersion != "20.35.1012" {
|
||||
t.Errorf("fw_ver=%q", h.FirmwareVersion)
|
||||
}
|
||||
if h.PCIAddr != "0000:0b:00.0" {
|
||||
t.Errorf("PCIAddr=%q", h.PCIAddr)
|
||||
}
|
||||
if h.NUMANode != 1 {
|
||||
t.Errorf("NUMANode=%d; GPU-to-HCA affinity depends on this", h.NUMANode)
|
||||
}
|
||||
|
||||
p := h.Ports[0]
|
||||
if p.State != "ACTIVE" || !p.Active() {
|
||||
t.Errorf("state=%q", p.State)
|
||||
}
|
||||
if p.PhysState != "LinkUp" {
|
||||
t.Errorf("phys_state=%q", p.PhysState)
|
||||
}
|
||||
if p.RateGbps != 200 {
|
||||
t.Errorf("rate=%d", p.RateGbps)
|
||||
}
|
||||
if p.LinkLayer != LinkInfiniBand {
|
||||
t.Errorf("link_layer=%q", p.LinkLayer)
|
||||
}
|
||||
if p.GIDPrefix != "fec0:0000:0000:000a" {
|
||||
t.Errorf("GIDPrefix=%q; it must be lowercased and cut to four groups", p.GIDPrefix)
|
||||
}
|
||||
if p.LID != 26 {
|
||||
t.Errorf("LID=%d, expected 0x1a to parse as 26", p.LID)
|
||||
}
|
||||
if p.SMLID != 1 {
|
||||
t.Errorf("SMLID=%d, expected a decimal value to parse too", p.SMLID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestActivePortsExcludesDown(t *testing.T) {
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := hcas[1] // mlx5_1
|
||||
if len(h.Ports) != 2 {
|
||||
t.Fatalf("expected 2 ports, got %d", len(h.Ports))
|
||||
}
|
||||
active := h.ActivePorts()
|
||||
if len(active) != 1 || active[0].Num != 1 {
|
||||
t.Errorf("a DOWN port was reported as active: %+v", active)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoCEIsDistinguished(t *testing.T) {
|
||||
// RoCE is usable but behaves differently from native IB; the placement
|
||||
// logic needs to be able to tell them apart.
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := hcas[2].Ports[0].LinkLayer; got != LinkEthernet {
|
||||
t.Errorf("mlx5_2 link_layer=%q, expected Ethernet (RoCE)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMissingRootIsNotAnError(t *testing.T) {
|
||||
// A node without InfiniBand is perfectly normal; it just forms its own
|
||||
// single-node island.
|
||||
hcas, err := DiscoverIB(filepath.Join(t.TempDir(), "nonexistent"))
|
||||
if err != nil {
|
||||
t.Fatalf("a node without IB should not be an error: %v", err)
|
||||
}
|
||||
if len(hcas) != 0 {
|
||||
t.Errorf("expected no HCAs, got %d", len(hcas))
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreadableAdapterDoesNotHideOthers(t *testing.T) {
|
||||
root := fakeSysfs(t)
|
||||
// An adapter with no ports directory at all.
|
||||
if err := os.MkdirAll(filepath.Join(root, "mlx5_broken"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
hcas, err := DiscoverIB(root)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hcas) != 3 {
|
||||
t.Errorf("a broken adapter changed the result: got %d HCAs", len(hcas))
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseHelpers(t *testing.T) {
|
||||
t.Run("rate", func(t *testing.T) {
|
||||
for in, want := range map[string]int{
|
||||
"200 Gb/sec (4X HDR)": 200,
|
||||
"100 Gb/sec (4X EDR)": 100,
|
||||
"2.5 Gb/sec (1X SDR)": 2,
|
||||
"": 0,
|
||||
"garbage": 0,
|
||||
} {
|
||||
if got := parseRate(in); got != want {
|
||||
t.Errorf("parseRate(%q)=%d, want %d", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("state", func(t *testing.T) {
|
||||
for in, want := range map[string]string{
|
||||
"4: ACTIVE": "ACTIVE",
|
||||
"1: DOWN": "DOWN",
|
||||
"5: LinkUp": "LinkUp",
|
||||
"ACTIVE": "ACTIVE",
|
||||
"": "",
|
||||
} {
|
||||
if got := parseEnumeratedState(in); got != want {
|
||||
t.Errorf("parseEnumeratedState(%q)=%q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("hex", func(t *testing.T) {
|
||||
for in, want := range map[string]int{
|
||||
"0x1a": 26, "0X1A": 26, "26": 26, "": 0, "zzz": 0,
|
||||
} {
|
||||
if got := parseMaybeHex(in); got != want {
|
||||
t.Errorf("parseMaybeHex(%q)=%d, want %d", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("gid", func(t *testing.T) {
|
||||
if got := parseGIDPrefix("FE80:0000:0000:0000:0002:C903:00A0:1234"); got != "fe80:0000:0000:0000" {
|
||||
t.Errorf("parseGIDPrefix=%q", got)
|
||||
}
|
||||
if got := parseGIDPrefix("short:gid"); got != "" {
|
||||
t.Errorf("a malformed GID should yield an empty prefix, got %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestJSONShape(t *testing.T) {
|
||||
inv := Inventory{
|
||||
Source: "nvidia-smi",
|
||||
GPUs: []GPU{{Index: 0, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559, ComputeCapability: "9.0"}},
|
||||
}
|
||||
b, err := json.Marshal(inv)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var out struct {
|
||||
Source string `json:"source"`
|
||||
GPUs []struct {
|
||||
Index int `json:"index"`
|
||||
Class string `json:"class"`
|
||||
Mem int `json:"memory_total_mb"`
|
||||
} `json:"gpus"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out.Source != "nvidia-smi" || len(out.GPUs) != 1 {
|
||||
t.Fatalf("shape wrong: %s", b)
|
||||
}
|
||||
// The derived class is what decides which GPUs are interchangeable, so a
|
||||
// consumer must not have to recompute it.
|
||||
if out.GPUs[0].Class != "h100-80g" {
|
||||
t.Errorf("class=%q, should be included in the JSON", out.GPUs[0].Class)
|
||||
}
|
||||
if out.GPUs[0].Mem != 81559 {
|
||||
t.Errorf("memory_total_mb=%d", out.GPUs[0].Mem)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
// Package engines builds the command lines and environments that start
|
||||
// inference engines (plan §5.3).
|
||||
//
|
||||
// Nothing here starts a process; that is the supervisor's job. Keeping command
|
||||
// construction pure makes it testable without GPUs, which matters because the
|
||||
// environment generation is where multi-node runs go wrong: a single missing
|
||||
// NCCL variable turns a working tensor-parallel group into one that hangs on
|
||||
// the first all-reduce.
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Profile is how a model runs on a GPU class. It mirrors the catalog profile
|
||||
// (plan §6.1), reduced to what launching needs.
|
||||
type Profile struct {
|
||||
Name string
|
||||
Engine string
|
||||
GPUClass string
|
||||
|
||||
// TP and PP are tensor and pipeline parallel degrees.
|
||||
TP int
|
||||
PP int
|
||||
// GPUsPerNode is how many GPUs each shard occupies.
|
||||
GPUsPerNode int
|
||||
// Nodes is how many nodes the replica spans. Above 1, a fabric is required.
|
||||
Nodes int
|
||||
|
||||
// Args are passed through to the engine untouched.
|
||||
Args []string
|
||||
// Env overrides anything the engine adapter generates.
|
||||
Env map[string]string
|
||||
|
||||
ModelPath string
|
||||
// ServedModelName is the name clients use, which need not match the path.
|
||||
ServedModelName string
|
||||
}
|
||||
|
||||
// ShardSpec is one shard of one replica on this node.
|
||||
type ShardSpec struct {
|
||||
ReplicaID string
|
||||
ShardID string
|
||||
// Rank is the node rank; 0 is the leader and serves HTTP.
|
||||
Rank int
|
||||
NNodes int
|
||||
|
||||
// LeaderAddr and LeaderPort are the distributed init endpoint, always on the
|
||||
// management network rather than the fabric: NCCL bootstraps over TCP and
|
||||
// only then moves to RDMA.
|
||||
LeaderAddr string
|
||||
LeaderPort int
|
||||
|
||||
// GPUIndices are the local GPUs this shard uses.
|
||||
GPUIndices []int
|
||||
// HCAs are the adapters of the chosen fabric, and only that fabric. A shard
|
||||
// must never be given HCAs from two fabrics.
|
||||
HCAs []string
|
||||
// MgmtInterface is the ethernet interface used for bootstrap traffic.
|
||||
MgmtInterface string
|
||||
|
||||
HTTPPort int
|
||||
Profile Profile
|
||||
|
||||
// Pool and Fabric are control plane concepts the agent cannot derive. They
|
||||
// are carried so the agent can report them back after a restart, and are
|
||||
// not used when building a command.
|
||||
Pool string
|
||||
Fabric string
|
||||
}
|
||||
|
||||
// Leader reports whether this shard serves HTTP for the replica.
|
||||
func (s ShardSpec) Leader() bool { return s.Rank == 0 }
|
||||
|
||||
// Command is a fully resolved launch: what to exec, with which environment.
|
||||
type Command struct {
|
||||
Path string
|
||||
Args []string
|
||||
// Env is the complete environment, not a delta.
|
||||
Env []string
|
||||
// Name distinguishes the processes of a multi-process shard in logs.
|
||||
Name string
|
||||
// WaitForLeaderPort makes the supervisor hold the next command until
|
||||
// LeaderAddr:LeaderPort accepts a connection.
|
||||
//
|
||||
// It exists for Ray: the head has to be listening before a worker can join,
|
||||
// and both sides express the same wait because both dial the same address.
|
||||
WaitForLeaderPort bool
|
||||
}
|
||||
|
||||
// Engine builds launch commands for one inference engine.
|
||||
type Engine interface {
|
||||
// Name is the engine's identifier in the catalog.
|
||||
Name() string
|
||||
// Validate rejects a profile that cannot work on these GPUs, before
|
||||
// anything is started.
|
||||
Validate(p Profile, gpus []GPUInfo) error
|
||||
|
||||
// Commands returns the processes that make up one shard, in start order.
|
||||
//
|
||||
// Most shards are a single process. vLLM under Ray is not: the leader runs
|
||||
// a Ray head and then the server against it, and a worker runs only the Ray
|
||||
// join. Returning a sequence rather than one command is what lets the
|
||||
// supervisor treat both the same way — every command but the last is a
|
||||
// sidecar kept alive for the shard's lifetime, and the last one's exit is
|
||||
// the shard's exit.
|
||||
Commands(s ShardSpec) ([]Command, error)
|
||||
|
||||
// ServesHTTP reports whether this rank answers the OpenAI API.
|
||||
//
|
||||
// In a multi-node replica only the leader does. Health-checking a worker's
|
||||
// port would time out and fail a replica that is coming up perfectly well,
|
||||
// which is exactly what happened before this existed.
|
||||
ServesHTTP(s ShardSpec) bool
|
||||
|
||||
// HealthPath and MetricsPath are appended to the shard's base URL.
|
||||
HealthPath() string
|
||||
MetricsPath() string
|
||||
}
|
||||
|
||||
// GPUInfo is the subset of GPU facts validation needs.
|
||||
type GPUInfo struct {
|
||||
Index int
|
||||
Name string
|
||||
ComputeCapability string
|
||||
MemoryTotalMB int
|
||||
}
|
||||
|
||||
// Registry holds the known engines.
|
||||
type Registry struct {
|
||||
engines map[string]Engine
|
||||
}
|
||||
|
||||
// NewRegistry returns a registry with the built-in engines.
|
||||
func NewRegistry(cfg Config) *Registry {
|
||||
r := &Registry{engines: map[string]Engine{}}
|
||||
r.Add(&VLLM{cfg: cfg.VLLM})
|
||||
r.Add(&SGLang{cfg: cfg.SGLang})
|
||||
return r
|
||||
}
|
||||
|
||||
// Add registers an engine, replacing any engine of the same name.
|
||||
func (r *Registry) Add(e Engine) { r.engines[e.Name()] = e }
|
||||
|
||||
// Get returns an engine by name.
|
||||
func (r *Registry) Get(name string) (Engine, error) {
|
||||
e, ok := r.engines[name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown engine %q (have %s)", name, strings.Join(r.Names(), ", "))
|
||||
}
|
||||
return e, nil
|
||||
}
|
||||
|
||||
// Names lists the registered engines.
|
||||
func (r *Registry) Names() []string {
|
||||
out := make([]string, 0, len(r.engines))
|
||||
for n := range r.engines {
|
||||
out = append(out, n)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// Config configures the engine adapters.
|
||||
type Config struct {
|
||||
VLLM VLLMConfig `koanf:"vllm"`
|
||||
SGLang SGLangConfig `koanf:"sglang"`
|
||||
}
|
||||
|
||||
// --- shared helpers ---
|
||||
|
||||
// baseEnv builds the environment every engine shares.
|
||||
//
|
||||
// This is the part of §5.3 the plan marks critical. The rules encoded here:
|
||||
//
|
||||
// - CUDA_VISIBLE_DEVICES pins the shard to its GPUs, so two shards on one
|
||||
// node cannot fight over the same card.
|
||||
// - NCCL_IB_HCA lists only the chosen fabric's adapters. Leaving it unset
|
||||
// lets NCCL pick adapters from another fabric, where the peer is
|
||||
// unreachable, and the job hangs instead of failing.
|
||||
// - NCCL_SOCKET_IFNAME pins bootstrap to the management interface. NCCL
|
||||
// otherwise probes every interface, including IB-over-IP addresses that
|
||||
// cannot route between fabrics.
|
||||
// - NCCL_IB_DISABLE is set when there is no fabric, so NCCL stops looking for
|
||||
// one instead of timing out.
|
||||
func baseEnv(s ShardSpec) map[string]string {
|
||||
env := map[string]string{
|
||||
"CUDA_VISIBLE_DEVICES": joinInts(s.GPUIndices),
|
||||
}
|
||||
|
||||
if len(s.HCAs) > 0 {
|
||||
env["NCCL_IB_HCA"] = strings.Join(s.HCAs, ",")
|
||||
env["NCCL_IB_DISABLE"] = "0"
|
||||
} else {
|
||||
// Single node, or a node with no fabric. Either way IB is not the path.
|
||||
env["NCCL_IB_DISABLE"] = "1"
|
||||
}
|
||||
|
||||
if s.MgmtInterface != "" {
|
||||
env["NCCL_SOCKET_IFNAME"] = s.MgmtInterface
|
||||
env["GLOO_SOCKET_IFNAME"] = s.MgmtInterface
|
||||
}
|
||||
|
||||
// Profile overrides win: they are the escape hatch for the per-cluster
|
||||
// tuning no adapter can guess.
|
||||
for k, v := range s.Profile.Env {
|
||||
env[k] = v
|
||||
}
|
||||
return env
|
||||
}
|
||||
|
||||
// materializeEnv merges generated variables over the current process
|
||||
// environment and returns it in exec form.
|
||||
//
|
||||
// Inheriting matters: engines need PATH, HOME, HF_HOME, LD_LIBRARY_PATH and
|
||||
// whatever else the operator set up in the unit file.
|
||||
func materializeEnv(generated map[string]string) []string {
|
||||
merged := map[string]string{}
|
||||
for _, kv := range os.Environ() {
|
||||
if k, v, ok := strings.Cut(kv, "="); ok {
|
||||
merged[k] = v
|
||||
}
|
||||
}
|
||||
for k, v := range generated {
|
||||
merged[k] = v
|
||||
}
|
||||
|
||||
out := make([]string, 0, len(merged))
|
||||
for k, v := range merged {
|
||||
out = append(out, k+"="+v)
|
||||
}
|
||||
// Sorted so a command is reproducible and diffable in logs.
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
func joinInts(v []int) string {
|
||||
parts := make([]string, len(v))
|
||||
for i, n := range v {
|
||||
parts[i] = strconv.Itoa(n)
|
||||
}
|
||||
return strings.Join(parts, ",")
|
||||
}
|
||||
|
||||
// validateCommon holds the checks that apply to every engine.
|
||||
func validateCommon(p Profile, gpus []GPUInfo) error {
|
||||
if p.ModelPath == "" {
|
||||
return fmt.Errorf("profile %q: model path is empty", p.Name)
|
||||
}
|
||||
if p.TP <= 0 {
|
||||
return fmt.Errorf("profile %q: tensor parallel degree must be positive", p.Name)
|
||||
}
|
||||
if p.Nodes > 1 && p.GPUsPerNode <= 0 {
|
||||
return fmt.Errorf("profile %q: gpus_per_node is required for a %d-node profile", p.Name, p.Nodes)
|
||||
}
|
||||
|
||||
perNode := p.GPUsPerNode
|
||||
if perNode == 0 {
|
||||
perNode = p.TP
|
||||
}
|
||||
nodes := max(p.Nodes, 1)
|
||||
if want := perNode * nodes; want != p.TP*max(p.PP, 1) {
|
||||
return fmt.Errorf("profile %q: %d nodes x %d GPUs does not match TP %d x PP %d",
|
||||
p.Name, nodes, perNode, p.TP, max(p.PP, 1))
|
||||
}
|
||||
if len(gpus) > 0 && len(gpus) < perNode {
|
||||
return fmt.Errorf("profile %q: needs %d GPUs per node, node has %d", p.Name, perNode, len(gpus))
|
||||
}
|
||||
|
||||
// A tensor-parallel group must be uniform: mixing classes makes the slowest
|
||||
// card set the pace for every collective, and mismatched memory makes the
|
||||
// smallest card decide what fits.
|
||||
if len(gpus) > 1 {
|
||||
first := gpus[0].Name
|
||||
for _, g := range gpus[1:] {
|
||||
if g.Name != first {
|
||||
return fmt.Errorf("profile %q: GPUs are not uniform (%q and %q); "+
|
||||
"a tensor-parallel group may not mix GPU types", p.Name, first, g.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// requireComputeCapability rejects profiles whose quantization the hardware
|
||||
// cannot execute.
|
||||
func requireComputeCapability(p Profile, gpus []GPUInfo, minCap float64, why string) error {
|
||||
for _, g := range gpus {
|
||||
cap, err := strconv.ParseFloat(g.ComputeCapability, 64)
|
||||
if err != nil {
|
||||
// Unknown capability: let it through rather than block a launch on
|
||||
// a parsing problem. The engine will refuse it clearly enough.
|
||||
continue
|
||||
}
|
||||
if cap < minCap {
|
||||
return fmt.Errorf("profile %q: %s needs compute capability %.1f, GPU %d (%s) has %.1f",
|
||||
p.Name, why, minCap, g.Index, g.Name, cap)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// quantizationOf extracts the --quantization argument, if present.
|
||||
func quantizationOf(args []string) string {
|
||||
for i, a := range args {
|
||||
if a == "--quantization" && i+1 < len(args) {
|
||||
return args[i+1]
|
||||
}
|
||||
if v, ok := strings.CutPrefix(a, "--quantization="); ok {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,425 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// buildOne returns the single command a shard produces, failing if it produced
|
||||
// a different number.
|
||||
func buildOne(t *testing.T, e Engine, s ShardSpec) Command {
|
||||
t.Helper()
|
||||
cmds := build(t, e, s)
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("expected 1 command, got %d: %s", len(cmds), summarize(cmds))
|
||||
}
|
||||
return cmds[0]
|
||||
}
|
||||
|
||||
// buildEngine returns the last command, which is the engine itself. The ones
|
||||
// before it are sidecars.
|
||||
func buildEngine(t *testing.T, e Engine, s ShardSpec) Command {
|
||||
t.Helper()
|
||||
cmds := build(t, e, s)
|
||||
return cmds[len(cmds)-1]
|
||||
}
|
||||
|
||||
func build(t *testing.T, e Engine, s ShardSpec) []Command {
|
||||
t.Helper()
|
||||
cmds, err := e.Commands(s)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) == 0 {
|
||||
t.Fatal("no commands produced")
|
||||
}
|
||||
return cmds
|
||||
}
|
||||
|
||||
// summarize keeps failure output readable: a Command carries the whole
|
||||
// environment, and printing it buries the actual problem.
|
||||
func summarize(cmds []Command) string {
|
||||
var b strings.Builder
|
||||
for i, c := range cmds {
|
||||
if i > 0 {
|
||||
b.WriteString(" | ")
|
||||
}
|
||||
fmt.Fprintf(&b, "%s: %s %s", c.Name, c.Path, strings.Join(c.Args, " "))
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// envMap turns exec-form environment back into a map for assertions.
|
||||
func envMap(env []string) map[string]string {
|
||||
m := map[string]string{}
|
||||
for _, kv := range env {
|
||||
if k, v, ok := strings.Cut(kv, "="); ok {
|
||||
m[k] = v
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// argValue returns the value following flag, or "" if the flag is absent.
|
||||
func argValue(args []string, flag string) string {
|
||||
i := slices.Index(args, flag)
|
||||
if i < 0 || i+1 >= len(args) {
|
||||
return ""
|
||||
}
|
||||
return args[i+1]
|
||||
}
|
||||
|
||||
func singleNodeSpec() ShardSpec {
|
||||
return ShardSpec{
|
||||
ReplicaID: "r1",
|
||||
ShardID: "r1-0",
|
||||
Rank: 0,
|
||||
NNodes: 1,
|
||||
GPUIndices: []int{0, 1, 2, 3},
|
||||
MgmtInterface: "eth0",
|
||||
HTTPPort: 20000,
|
||||
Profile: Profile{
|
||||
Name: "h100-tp4",
|
||||
Engine: "vllm",
|
||||
TP: 4,
|
||||
ModelPath: "/models/llama",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func multiNodeSpec(rank int) ShardSpec {
|
||||
s := singleNodeSpec()
|
||||
s.Rank = rank
|
||||
s.NNodes = 2
|
||||
s.GPUIndices = []int{0, 1, 2, 3, 4, 5, 6, 7}
|
||||
s.HCAs = []string{"mlx5_0", "mlx5_1"}
|
||||
s.LeaderAddr = "10.0.0.10"
|
||||
s.LeaderPort = 20001
|
||||
s.Profile.TP = 16
|
||||
s.Profile.GPUsPerNode = 8
|
||||
s.Profile.Nodes = 2
|
||||
return s
|
||||
}
|
||||
|
||||
// --- environment generation: the part that breaks multi-node runs ---
|
||||
|
||||
func TestSingleNodeDisablesIB(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
if env["NCCL_IB_DISABLE"] != "1" {
|
||||
t.Error("NCCL_IB_DISABLE must be 1 without a fabric, or NCCL waits for one that is not there")
|
||||
}
|
||||
if _, set := env["NCCL_IB_HCA"]; set {
|
||||
t.Error("NCCL_IB_HCA was set without any HCAs")
|
||||
}
|
||||
if env["CUDA_VISIBLE_DEVICES"] != "0,1,2,3" {
|
||||
t.Errorf("CUDA_VISIBLE_DEVICES=%q", env["CUDA_VISIBLE_DEVICES"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiNodePinsFabric(t *testing.T) {
|
||||
cmd := buildEngine(t, &VLLM{}, multiNodeSpec(0))
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
// Only the chosen fabric's adapters. Leaving this open lets NCCL pick an
|
||||
// adapter on another fabric, where the peer is unreachable, and the job
|
||||
// hangs rather than failing.
|
||||
if env["NCCL_IB_HCA"] != "mlx5_0,mlx5_1" {
|
||||
t.Errorf("NCCL_IB_HCA=%q", env["NCCL_IB_HCA"])
|
||||
}
|
||||
if env["NCCL_IB_DISABLE"] != "0" {
|
||||
t.Errorf("NCCL_IB_DISABLE=%q with a fabric present", env["NCCL_IB_DISABLE"])
|
||||
}
|
||||
// Bootstrap must stay on the management network; NCCL otherwise probes
|
||||
// every interface, including IB-over-IP that cannot route between fabrics.
|
||||
if env["NCCL_SOCKET_IFNAME"] != "eth0" || env["GLOO_SOCKET_IFNAME"] != "eth0" {
|
||||
t.Errorf("socket interfaces not pinned: %q / %q",
|
||||
env["NCCL_SOCKET_IFNAME"], env["GLOO_SOCKET_IFNAME"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileEnvOverridesGenerated(t *testing.T) {
|
||||
s := multiNodeSpec(0)
|
||||
s.Profile.Env = map[string]string{
|
||||
"NCCL_IB_HCA": "mlx5_9",
|
||||
"NCCL_NET_GDR_LEVEL": "SYS",
|
||||
}
|
||||
cmd := buildEngine(t, &VLLM{}, s)
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
if env["NCCL_IB_HCA"] != "mlx5_9" {
|
||||
t.Error("a profile override did not win; it is the escape hatch for per-cluster tuning")
|
||||
}
|
||||
if env["NCCL_NET_GDR_LEVEL"] != "SYS" {
|
||||
t.Error("an extra profile variable was dropped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessEnvironmentIsInherited(t *testing.T) {
|
||||
// Engines need PATH, HF_HOME, LD_LIBRARY_PATH and whatever else the unit
|
||||
// file set up.
|
||||
t.Setenv("LLMGW_TEST_MARKER", "present")
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
if envMap(cmd.Env)["LLMGW_TEST_MARKER"] != "present" {
|
||||
t.Error("the process environment was not inherited")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnvIsSorted(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
if !slices.IsSorted(cmd.Env) {
|
||||
t.Error("environment is not sorted; launches would not be reproducible in logs")
|
||||
}
|
||||
}
|
||||
|
||||
// --- vLLM ---
|
||||
|
||||
func TestVLLMSingleNodeArgs(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{cfg: VLLMConfig{Bin: "/opt/vllm/bin/vllm"}}, singleNodeSpec())
|
||||
if cmd.Path != "/opt/vllm/bin/vllm" {
|
||||
t.Errorf("Path=%q", cmd.Path)
|
||||
}
|
||||
if cmd.Args[0] != "serve" || cmd.Args[1] != "/models/llama" {
|
||||
t.Errorf("args start wrong: %v", cmd.Args[:2])
|
||||
}
|
||||
if argValue(cmd.Args, "--tensor-parallel-size") != "4" {
|
||||
t.Errorf("TP not passed: %v", cmd.Args)
|
||||
}
|
||||
if argValue(cmd.Args, "--port") != "20000" {
|
||||
t.Errorf("port not passed: %v", cmd.Args)
|
||||
}
|
||||
// Single node means no distributed backend flag at all.
|
||||
if slices.Contains(cmd.Args, "--distributed-executor-backend") {
|
||||
t.Error("a distributed backend was selected for a single-node profile")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMRayLeaderRunsTwoProcesses(t *testing.T) {
|
||||
// The leader starts a Ray head and then the server against it. A single
|
||||
// command would mean running vLLM against a cluster that does not exist.
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
cmds, err := v.Commands(multiNodeSpec(0))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 2 {
|
||||
t.Fatalf("expected 2 commands for the Ray leader, got %d", len(cmds))
|
||||
}
|
||||
if cmds[0].Name != "ray" || !slices.Contains(cmds[0].Args, "--head") {
|
||||
t.Errorf("first command is not the Ray head: %+v", cmds[0])
|
||||
}
|
||||
if !cmds[0].WaitForLeaderPort {
|
||||
t.Error("the server would start before the head is listening")
|
||||
}
|
||||
if cmds[1].Name != "vllm" || argValue(cmds[1].Args, "--distributed-executor-backend") != "ray" {
|
||||
t.Errorf("second command is not the vLLM server: %+v", cmds[1])
|
||||
}
|
||||
if got := envMap(cmds[1].Env)["RAY_ADDRESS"]; got != "10.0.0.10:20001" {
|
||||
t.Errorf("RAY_ADDRESS=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMRayWorkerRunsOnlyRay(t *testing.T) {
|
||||
// A Ray worker runs no vLLM process: the cluster hands its GPUs to the
|
||||
// leader. Starting a second server there would bind a port nobody uses and
|
||||
// compete for the same cards.
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
cmds, err := v.Commands(multiNodeSpec(1))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("expected 1 command for a Ray worker, got %d: %+v", len(cmds), cmds)
|
||||
}
|
||||
if cmds[0].Name != "ray" || slices.Contains(cmds[0].Args, "--head") {
|
||||
t.Errorf("worker command wrong: %+v", cmds[0])
|
||||
}
|
||||
if argValue(cmds[0].Args, "--address") != "10.0.0.10:20001" {
|
||||
t.Errorf("worker is not joining the leader: %v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMMPMode(t *testing.T) {
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}
|
||||
cmd := buildOne(t, v, multiNodeSpec(0))
|
||||
if argValue(cmd.Args, "--distributed-executor-backend") != "mp" {
|
||||
t.Errorf("backend not set: %v", cmd.Args)
|
||||
}
|
||||
env := envMap(cmd.Env)
|
||||
if env["VLLM_HOST_IP"] != "10.0.0.10" || env["VLLM_PORT"] != "20001" {
|
||||
t.Errorf("dist init not configured: %q %q", env["VLLM_HOST_IP"], env["VLLM_PORT"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMUnknownLaunchMode(t *testing.T) {
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "sorcery"}}
|
||||
if _, err := v.Commands(multiNodeSpec(0)); err == nil {
|
||||
t.Error("an unknown launch mode was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyLeaderServesHTTP(t *testing.T) {
|
||||
// The fix for the bug that stopped multi-node replicas coming up at all:
|
||||
// health-checking a worker's port times out and fails a healthy replica.
|
||||
for _, e := range []Engine{&VLLM{}, &SGLang{}} {
|
||||
if !e.ServesHTTP(multiNodeSpec(0)) {
|
||||
t.Errorf("%s: the leader must serve HTTP", e.Name())
|
||||
}
|
||||
if e.ServesHTTP(multiNodeSpec(1)) {
|
||||
t.Errorf("%s: a worker must not be expected to serve HTTP", e.Name())
|
||||
}
|
||||
if !e.ServesHTTP(singleNodeSpec()) {
|
||||
t.Errorf("%s: a single-node shard must serve HTTP", e.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoRayForSingleNodeOrMPMode(t *testing.T) {
|
||||
ray := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
if cmds := buildOne(t, ray, singleNodeSpec()); cmds.Name != "vllm" {
|
||||
t.Errorf("a single-node profile was given a Ray process: %+v", cmds)
|
||||
}
|
||||
mp := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}
|
||||
if cmds := buildOne(t, mp, multiNodeSpec(0)); cmds.Name != "vllm" {
|
||||
t.Errorf("mp mode was given a Ray process: %+v", cmds)
|
||||
}
|
||||
}
|
||||
|
||||
// --- SGLang ---
|
||||
|
||||
func TestSGLangMultiNodeArgs(t *testing.T) {
|
||||
cmd := buildOne(t, &SGLang{cfg: SGLangConfig{Python: "/opt/sglang/bin/python"}}, multiNodeSpec(1))
|
||||
if cmd.Path != "/opt/sglang/bin/python" {
|
||||
t.Errorf("Path=%q", cmd.Path)
|
||||
}
|
||||
if cmd.Args[0] != "-m" || cmd.Args[1] != "sglang.launch_server" {
|
||||
t.Errorf("not launched as a module: %v", cmd.Args[:2])
|
||||
}
|
||||
if argValue(cmd.Args, "--nnodes") != "2" || argValue(cmd.Args, "--node-rank") != "1" {
|
||||
t.Errorf("multi-node flags wrong: %v", cmd.Args)
|
||||
}
|
||||
if argValue(cmd.Args, "--dist-init-addr") != "10.0.0.10:20001" {
|
||||
t.Errorf("dist-init-addr wrong: %v", cmd.Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSGLangSingleNodeOmitsMultiNodeFlags(t *testing.T) {
|
||||
cmd := buildOne(t, &SGLang{}, singleNodeSpec())
|
||||
for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} {
|
||||
if slices.Contains(cmd.Args, flag) {
|
||||
t.Errorf("%s was passed for a single-node profile", flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileArgsComeLast(t *testing.T) {
|
||||
// Operators must be able to override anything the adapter generated.
|
||||
s := singleNodeSpec()
|
||||
s.Profile.Args = []string{"--max-model-len", "32768"}
|
||||
cmd := buildEngine(t, &VLLM{}, s)
|
||||
|
||||
i := slices.Index(cmd.Args, "--max-model-len")
|
||||
j := slices.Index(cmd.Args, "--tensor-parallel-size")
|
||||
if i < j {
|
||||
t.Errorf("profile args came before generated ones: %v", cmd.Args)
|
||||
}
|
||||
}
|
||||
|
||||
// --- validation ---
|
||||
|
||||
func h100(n int) []GPUInfo {
|
||||
out := make([]GPUInfo, n)
|
||||
for i := range out {
|
||||
out[i] = GPUInfo{Index: i, Name: "NVIDIA H100 80GB HBM3", ComputeCapability: "9.0", MemoryTotalMB: 81559}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestValidateAcceptsMatchingProfile(t *testing.T) {
|
||||
p := Profile{Name: "ok", TP: 4, ModelPath: "/m"}
|
||||
if err := (&VLLM{}).Validate(p, h100(4)); err != nil {
|
||||
t.Errorf("a valid profile was rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsMixedGPUTypes(t *testing.T) {
|
||||
// The invariant from plan §1.1: a TP group never mixes GPU types.
|
||||
gpus := h100(2)
|
||||
gpus[1].Name = "NVIDIA A100-SXM4-40GB"
|
||||
|
||||
err := (&VLLM{}).Validate(Profile{Name: "mixed", TP: 2, ModelPath: "/m"}, gpus)
|
||||
if err == nil || !strings.Contains(err.Error(), "may not mix GPU types") {
|
||||
t.Errorf("mixed GPU types accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsFP8OnOldHardware(t *testing.T) {
|
||||
gpus := h100(4)
|
||||
for i := range gpus {
|
||||
gpus[i].Name = "NVIDIA A100-SXM4-80GB"
|
||||
gpus[i].ComputeCapability = "8.0"
|
||||
}
|
||||
p := Profile{Name: "fp8", TP: 4, ModelPath: "/m", Args: []string{"--quantization", "fp8"}}
|
||||
|
||||
err := (&VLLM{}).Validate(p, gpus)
|
||||
if err == nil || !strings.Contains(err.Error(), "compute capability") {
|
||||
t.Errorf("fp8 accepted on compute capability 8.0: %v", err)
|
||||
}
|
||||
|
||||
// The same profile is fine on Hopper.
|
||||
if err := (&VLLM{}).Validate(p, h100(4)); err != nil {
|
||||
t.Errorf("fp8 rejected on H100: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsInconsistentParallelism(t *testing.T) {
|
||||
// 2 nodes x 8 GPUs cannot serve TP 4.
|
||||
p := Profile{Name: "bad", TP: 4, Nodes: 2, GPUsPerNode: 8, ModelPath: "/m"}
|
||||
if err := (&VLLM{}).Validate(p, h100(8)); err == nil {
|
||||
t.Error("inconsistent node/GPU/TP arithmetic accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsTooFewGPUs(t *testing.T) {
|
||||
if err := (&VLLM{}).Validate(Profile{Name: "big", TP: 8, ModelPath: "/m"}, h100(4)); err == nil {
|
||||
t.Error("a TP 8 profile was accepted on a 4-GPU node")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRequiresModelPath(t *testing.T) {
|
||||
if err := (&VLLM{}).Validate(Profile{Name: "nopath", TP: 1}, h100(1)); err == nil {
|
||||
t.Error("a profile without a model path was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry(t *testing.T) {
|
||||
r := NewRegistry(Config{})
|
||||
if got := r.Names(); len(got) != 2 || got[0] != "sglang" || got[1] != "vllm" {
|
||||
t.Errorf("registry contents: %v", got)
|
||||
}
|
||||
if _, err := r.Get("vllm"); err != nil {
|
||||
t.Errorf("vllm missing: %v", err)
|
||||
}
|
||||
err := func() error { _, err := r.Get("nope"); return err }()
|
||||
if err == nil || !strings.Contains(err.Error(), "sglang, vllm") {
|
||||
t.Errorf("an unknown engine should list what is available: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuantizationOf(t *testing.T) {
|
||||
for _, args := range [][]string{
|
||||
{"--quantization", "fp8"},
|
||||
{"--quantization=fp8"},
|
||||
{"--max-model-len", "100", "--quantization", "fp8"},
|
||||
} {
|
||||
if got := quantizationOf(args); got != "fp8" {
|
||||
t.Errorf("quantizationOf(%v)=%q", args, got)
|
||||
}
|
||||
}
|
||||
if got := quantizationOf([]string{"--quantization"}); got != "" {
|
||||
t.Errorf("a dangling flag should yield nothing, got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package engines
|
||||
|
||||
import "strconv"
|
||||
|
||||
// SGLangConfig locates the SGLang installation.
|
||||
type SGLangConfig struct {
|
||||
// Python is the interpreter of the SGLang virtualenv. SGLang is launched as
|
||||
// a module rather than a console script, which is the documented entry
|
||||
// point and avoids depending on the script being on PATH.
|
||||
Python string `koanf:"python"`
|
||||
}
|
||||
|
||||
// SGLang builds SGLang launch commands.
|
||||
type SGLang struct{ cfg SGLangConfig }
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*SGLang) Name() string { return "sglang" }
|
||||
|
||||
// HealthPath is SGLang's readiness endpoint.
|
||||
func (*SGLang) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is SGLang's Prometheus endpoint.
|
||||
func (*SGLang) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
func (*SGLang) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if err := validateCommon(p, gpus); err != nil {
|
||||
return err
|
||||
}
|
||||
if q := quantizationOf(p.Args); q == "fp8" {
|
||||
if err := requireComputeCapability(p, gpus, 8.9, "fp8 quantization"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API.
|
||||
//
|
||||
// Every SGLang rank runs the same launcher, but only node-rank 0 binds the HTTP
|
||||
// port; the others exist to contribute GPUs to the tensor-parallel group.
|
||||
func (*SGLang) ServesHTTP(s ShardSpec) bool { return s.Leader() }
|
||||
|
||||
// Commands returns the processes for one shard.
|
||||
//
|
||||
// SGLang supports multi-node natively, so there is only ever one process: every
|
||||
// rank runs the same launcher with its own rank number and they rendezvous at
|
||||
// the leader's dist-init address.
|
||||
func (g *SGLang) Commands(s ShardSpec) ([]Command, error) {
|
||||
c, err := g.serverCommand(s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []Command{c}, nil
|
||||
}
|
||||
|
||||
func (g *SGLang) serverCommand(s ShardSpec) (Command, error) {
|
||||
python := g.cfg.Python
|
||||
if python == "" {
|
||||
python = "python3"
|
||||
}
|
||||
p := s.Profile
|
||||
|
||||
args := []string{"-m", "sglang.launch_server",
|
||||
"--model-path", p.ModelPath,
|
||||
"--host", "0.0.0.0",
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
"--tp", strconv.Itoa(p.TP),
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--served-model-name", p.ServedModelName)
|
||||
}
|
||||
if s.NNodes > 1 {
|
||||
args = append(args,
|
||||
"--nnodes", strconv.Itoa(s.NNodes),
|
||||
"--node-rank", strconv.Itoa(s.Rank),
|
||||
"--dist-init-addr", s.LeaderAddr+":"+strconv.Itoa(s.LeaderPort),
|
||||
)
|
||||
}
|
||||
args = append(args, p.Args...)
|
||||
|
||||
return Command{Name: "sglang", Path: python, Args: args, Env: materializeEnv(baseEnv(s))}, nil
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// VLLMConfig locates the vLLM installation and picks how it is launched.
|
||||
type VLLMConfig struct {
|
||||
// Bin is the vllm executable.
|
||||
Bin string `koanf:"bin"`
|
||||
// LaunchMode selects the multi-node mechanism: "ray" or "mp".
|
||||
//
|
||||
// Which one works depends on the vLLM version, and the flags have changed
|
||||
// more than once (plan §14). Both paths exist so a cluster can pick the one
|
||||
// its version supports instead of waiting for a code change.
|
||||
LaunchMode string `koanf:"launch_mode"`
|
||||
// RayBin is used when LaunchMode is "ray".
|
||||
RayBin string `koanf:"ray_bin"`
|
||||
}
|
||||
|
||||
// VLLM builds vLLM launch commands.
|
||||
type VLLM struct{ cfg VLLMConfig }
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*VLLM) Name() string { return "vllm" }
|
||||
|
||||
// HealthPath is vLLM's readiness endpoint.
|
||||
func (*VLLM) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is vLLM's Prometheus endpoint.
|
||||
func (*VLLM) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
func (*VLLM) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if err := validateCommon(p, gpus); err != nil {
|
||||
return err
|
||||
}
|
||||
// FP8 needs Hopper or Ada; on anything older vLLM either refuses or falls
|
||||
// back silently to a slower path, and a silent fallback is worse.
|
||||
if q := quantizationOf(p.Args); q == "fp8" {
|
||||
if err := requireComputeCapability(p, gpus, 8.9, "fp8 quantization"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API.
|
||||
//
|
||||
// Under Ray a worker runs nothing but the Ray join, and under multiprocessing
|
||||
// only the driver binds the port. Either way it is the leader alone.
|
||||
func (*VLLM) ServesHTTP(s ShardSpec) bool { return s.Leader() }
|
||||
|
||||
// Commands returns the processes for one shard, in start order.
|
||||
//
|
||||
// Single node, or multi-node under multiprocessing: one server process. Under
|
||||
// Ray: the leader starts a head and then the server, a worker starts only the
|
||||
// join. A Ray worker running no vLLM process at all is correct — the cluster
|
||||
// gives the leader its GPUs.
|
||||
func (v *VLLM) Commands(s ShardSpec) ([]Command, error) {
|
||||
ray, useRay := v.rayCommand(s)
|
||||
if useRay && !s.Leader() {
|
||||
// A worker is the Ray join and nothing else. It blocks, so its exit is
|
||||
// the shard's exit, which is what the supervisor needs.
|
||||
return []Command{ray}, nil
|
||||
}
|
||||
|
||||
server, err := v.serverCommand(s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if useRay {
|
||||
return []Command{ray, server}, nil
|
||||
}
|
||||
return []Command{server}, nil
|
||||
}
|
||||
|
||||
// serverCommand produces the vLLM server process.
|
||||
func (v *VLLM) serverCommand(s ShardSpec) (Command, error) {
|
||||
bin := v.cfg.Bin
|
||||
if bin == "" {
|
||||
bin = "vllm"
|
||||
}
|
||||
p := s.Profile
|
||||
|
||||
args := []string{"serve", p.ModelPath,
|
||||
"--host", "0.0.0.0",
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
"--tensor-parallel-size", strconv.Itoa(p.TP),
|
||||
}
|
||||
if p.PP > 1 {
|
||||
args = append(args, "--pipeline-parallel-size", strconv.Itoa(p.PP))
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--served-model-name", p.ServedModelName)
|
||||
}
|
||||
|
||||
env := baseEnv(s)
|
||||
|
||||
if s.NNodes > 1 {
|
||||
switch v.cfg.LaunchMode {
|
||||
case "", "ray":
|
||||
// Ray is started separately by the supervisor; vLLM then finds the
|
||||
// cluster through the address in the environment.
|
||||
env["VLLM_HOST_IP"] = s.LeaderAddr
|
||||
env["RAY_ADDRESS"] = fmt.Sprintf("%s:%d", s.LeaderAddr, s.LeaderPort)
|
||||
args = append(args, "--distributed-executor-backend", "ray")
|
||||
case "mp":
|
||||
env["VLLM_HOST_IP"] = s.LeaderAddr
|
||||
env["VLLM_PORT"] = strconv.Itoa(s.LeaderPort)
|
||||
args = append(args, "--distributed-executor-backend", "mp")
|
||||
default:
|
||||
return Command{}, fmt.Errorf("vllm: unknown launch_mode %q (want \"ray\" or \"mp\")", v.cfg.LaunchMode)
|
||||
}
|
||||
}
|
||||
|
||||
// Profile arguments last, so an operator can override anything above.
|
||||
args = append(args, p.Args...)
|
||||
|
||||
return Command{Name: "vllm", Path: bin, Args: args, Env: materializeEnv(env)}, nil
|
||||
}
|
||||
|
||||
// rayCommand returns the command that starts or joins the Ray cluster, or
|
||||
// ok=false when the launch mode does not use Ray.
|
||||
func (v *VLLM) rayCommand(s ShardSpec) (Command, bool) {
|
||||
if s.NNodes <= 1 || (v.cfg.LaunchMode != "" && v.cfg.LaunchMode != "ray") {
|
||||
return Command{}, false
|
||||
}
|
||||
bin := v.cfg.RayBin
|
||||
if bin == "" {
|
||||
bin = "ray"
|
||||
}
|
||||
|
||||
args := []string{"start", "--block"}
|
||||
if s.Leader() {
|
||||
args = append(args, "--head", "--port", strconv.Itoa(s.LeaderPort))
|
||||
} else {
|
||||
args = append(args, "--address", fmt.Sprintf("%s:%d", s.LeaderAddr, s.LeaderPort))
|
||||
}
|
||||
args = append(args, "--num-gpus", strconv.Itoa(len(s.GPUIndices)))
|
||||
|
||||
return Command{
|
||||
Name: "ray", Path: bin, Args: args, Env: materializeEnv(baseEnv(s)),
|
||||
// The head must be listening before anything else happens: a worker
|
||||
// cannot join, and the leader's server cannot claim GPUs, until it is.
|
||||
WaitForLeaderPort: true,
|
||||
}, true
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
// HTTPHealth checks an engine over HTTP.
|
||||
//
|
||||
// Readiness is two questions, not one (plan §5.4): /health says the process is
|
||||
// up, and /v1/models says the weights finished loading and the model is being
|
||||
// served. An engine answers /health long before it can serve a request, so
|
||||
// routing on /health alone sends traffic into a model that is still loading.
|
||||
type HTTPHealth struct {
|
||||
client *http.Client
|
||||
healthPath string
|
||||
modelsPath string
|
||||
requireName bool
|
||||
}
|
||||
|
||||
// HTTPHealthOption customizes the checker.
|
||||
type HTTPHealthOption func(*HTTPHealth)
|
||||
|
||||
// WithHealthPath overrides the readiness path.
|
||||
func WithHealthPath(p string) HTTPHealthOption {
|
||||
return func(h *HTTPHealth) { h.healthPath = p }
|
||||
}
|
||||
|
||||
// WithoutModelCheck skips the /v1/models step, for engines that do not expose it.
|
||||
func WithoutModelCheck() HTTPHealthOption {
|
||||
return func(h *HTTPHealth) { h.requireName = false }
|
||||
}
|
||||
|
||||
// NewHTTPHealth builds a checker.
|
||||
func NewHTTPHealth(timeout time.Duration, opts ...HTTPHealthOption) *HTTPHealth {
|
||||
if timeout <= 0 {
|
||||
timeout = 3 * time.Second
|
||||
}
|
||||
h := &HTTPHealth{
|
||||
client: &http.Client{Timeout: timeout},
|
||||
healthPath: "/health",
|
||||
modelsPath: "/v1/models",
|
||||
requireName: true,
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(h)
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
// Ready reports whether the engine at baseURL is serving model.
|
||||
func (h *HTTPHealth) Ready(ctx context.Context, baseURL, model string) error {
|
||||
if err := h.get(ctx, baseURL+h.healthPath, nil); err != nil {
|
||||
return fmt.Errorf("health: %w", err)
|
||||
}
|
||||
if !h.requireName || model == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var list struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := h.get(ctx, baseURL+h.modelsPath, &list); err != nil {
|
||||
return fmt.Errorf("models: %w", err)
|
||||
}
|
||||
for _, d := range list.Data {
|
||||
if d.ID == model {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("engine is up but does not serve %q yet", model)
|
||||
}
|
||||
|
||||
func (h *HTTPHealth) get(ctx context.Context, url string, out any) error {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
resp, err := h.client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<10))
|
||||
return fmt.Errorf("status %s", resp.Status)
|
||||
}
|
||||
if out == nil {
|
||||
return nil
|
||||
}
|
||||
return json.NewDecoder(resp.Body).Decode(out)
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// engineStub serves configurable /health and /v1/models responses.
|
||||
func engineStub(t *testing.T, healthy bool, models ...string) *httptest.Server {
|
||||
t.Helper()
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) {
|
||||
if !healthy {
|
||||
http.Error(w, "loading", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
mux.HandleFunc("GET /v1/models", func(w http.ResponseWriter, _ *http.Request) {
|
||||
var b strings.Builder
|
||||
b.WriteString(`{"object":"list","data":[`)
|
||||
for i, m := range models {
|
||||
if i > 0 {
|
||||
b.WriteString(",")
|
||||
}
|
||||
b.WriteString(`{"id":"` + m + `"}`)
|
||||
}
|
||||
b.WriteString(`]}`)
|
||||
io.WriteString(w, b.String())
|
||||
})
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestReadyWhenServingModel(t *testing.T) {
|
||||
srv := engineStub(t, true, "llama-3.3-70b")
|
||||
if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b"); err != nil {
|
||||
t.Errorf("a serving engine was reported unready: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotReadyWhileLoading(t *testing.T) {
|
||||
// The case that matters: the process is up but the weights are not loaded.
|
||||
// Routing on /health alone would send traffic into a loading model.
|
||||
srv := engineStub(t, false, "llama-3.3-70b")
|
||||
err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b")
|
||||
if err == nil {
|
||||
t.Fatal("an unhealthy engine was reported ready")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "health") {
|
||||
t.Errorf("err=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotReadyWhenModelMissing(t *testing.T) {
|
||||
// /health passes but the model is not listed: the engine is up and serving
|
||||
// something else, or still registering.
|
||||
srv := engineStub(t, true, "some-other-model")
|
||||
err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b")
|
||||
if err == nil {
|
||||
t.Fatal("a model that is not served was reported ready")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "does not serve") {
|
||||
t.Errorf("err=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelCheckCanBeSkipped(t *testing.T) {
|
||||
srv := engineStub(t, true)
|
||||
h := NewHTTPHealth(time.Second, WithoutModelCheck())
|
||||
if err := h.Ready(t.Context(), srv.URL, "anything"); err != nil {
|
||||
t.Errorf("the model check was not skipped: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreachableEngine(t *testing.T) {
|
||||
h := NewHTTPHealth(200 * time.Millisecond)
|
||||
if err := h.Ready(t.Context(), "http://127.0.0.1:1", "m"); err == nil {
|
||||
t.Error("an unreachable engine was reported ready")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"os"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// killGroup signals a whole process group.
|
||||
//
|
||||
// Engines fork worker processes that hold GPU memory. Signalling only the
|
||||
// leader leaves those workers alive, and the next launch then fails to
|
||||
// allocate — with an error that points at the new process rather than the old.
|
||||
func killGroup(pid int, sig syscall.Signal) error {
|
||||
// The negative pid addresses the group. Setpgid at launch made the leader's
|
||||
// pid the group id.
|
||||
if err := syscall.Kill(-pid, sig); err != nil {
|
||||
if err == syscall.ESRCH {
|
||||
return os.ErrProcessDone
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// processAlive reports whether a pid exists.
|
||||
func processAlive(pid int) bool {
|
||||
if pid <= 0 {
|
||||
return false
|
||||
}
|
||||
// Signal 0 performs the permission and existence checks without delivering
|
||||
// anything.
|
||||
err := syscall.Kill(pid, 0)
|
||||
return err == nil || err == syscall.EPERM
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package supervisor
|
||||
|
||||
import "sync"
|
||||
|
||||
// ringBuffer keeps the last n lines of output.
|
||||
//
|
||||
// Engine logs are the only diagnosis available when a shard fails to load, but
|
||||
// they are also unbounded — a model loading for ten minutes can produce
|
||||
// megabytes. Keeping a bounded tail makes `ctl logs` useful without letting a
|
||||
// chatty engine consume the agent's memory.
|
||||
type ringBuffer struct {
|
||||
mu sync.Mutex
|
||||
lines []string
|
||||
next int
|
||||
full bool
|
||||
}
|
||||
|
||||
func newRingBuffer(n int) *ringBuffer {
|
||||
if n <= 0 {
|
||||
n = 256
|
||||
}
|
||||
return &ringBuffer{lines: make([]string, n)}
|
||||
}
|
||||
|
||||
func (r *ringBuffer) add(line string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.lines[r.next] = line
|
||||
r.next = (r.next + 1) % len(r.lines)
|
||||
if r.next == 0 {
|
||||
r.full = true
|
||||
}
|
||||
}
|
||||
|
||||
// tail returns up to n most recent lines, oldest first.
|
||||
func (r *ringBuffer) tail(n int) []string {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
size := r.next
|
||||
if r.full {
|
||||
size = len(r.lines)
|
||||
}
|
||||
if n <= 0 || n > size {
|
||||
n = size
|
||||
}
|
||||
|
||||
out := make([]string, 0, n)
|
||||
for i := size - n; i < size; i++ {
|
||||
idx := i
|
||||
if r.full {
|
||||
idx = (r.next + i) % len(r.lines)
|
||||
}
|
||||
out = append(out, r.lines[idx])
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
// Package supervisor runs and watches engine processes (plan §5.4).
|
||||
package supervisor
|
||||
|
||||
import "fmt"
|
||||
|
||||
// State is where a shard is in its lifecycle.
|
||||
//
|
||||
// Loading is separate from Starting because model loading dominates startup:
|
||||
// a 70B model takes minutes, and an operator watching a shard sit in one state
|
||||
// for ten minutes needs to know whether that is normal.
|
||||
type State int
|
||||
|
||||
const (
|
||||
// Pending means the shard is known but not started.
|
||||
Pending State = iota
|
||||
// Starting means the process has been launched but has not opened its port.
|
||||
Starting
|
||||
// Loading means the port is open but the model is not ready.
|
||||
Loading
|
||||
// Ready means health checks pass and the model is served.
|
||||
Ready
|
||||
// Draining means no new work; existing requests finish.
|
||||
Draining
|
||||
// Stopped means the process exited as intended.
|
||||
Stopped
|
||||
// Failed means the process exited unexpectedly or never became ready.
|
||||
Failed
|
||||
)
|
||||
|
||||
func (s State) String() string {
|
||||
switch s {
|
||||
case Pending:
|
||||
return "pending"
|
||||
case Starting:
|
||||
return "starting"
|
||||
case Loading:
|
||||
return "loading"
|
||||
case Ready:
|
||||
return "ready"
|
||||
case Draining:
|
||||
return "draining"
|
||||
case Stopped:
|
||||
return "stopped"
|
||||
case Failed:
|
||||
return "failed"
|
||||
default:
|
||||
return fmt.Sprintf("unknown(%d)", int(s))
|
||||
}
|
||||
}
|
||||
|
||||
// Terminal reports whether the shard has finished for good.
|
||||
func (s State) Terminal() bool { return s == Stopped || s == Failed }
|
||||
|
||||
// Live reports whether a process should be running in this state.
|
||||
func (s State) Live() bool {
|
||||
return s == Starting || s == Loading || s == Ready || s == Draining
|
||||
}
|
||||
@@ -0,0 +1,771 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
|
||||
// Config configures the supervisor.
|
||||
type Config struct {
|
||||
// WorkDir holds pid files and per-shard scratch space.
|
||||
WorkDir string `koanf:"work_dir"`
|
||||
// LogLines is how many lines of engine output to retain per shard.
|
||||
LogLines int `koanf:"log_lines"`
|
||||
// StopTimeout is how long a process gets to exit after SIGTERM before it is
|
||||
// killed. Engines need time to free GPU memory cleanly.
|
||||
StopTimeout time.Duration `koanf:"stop_timeout"`
|
||||
// HealthInterval is how often a starting shard is polled.
|
||||
HealthInterval time.Duration `koanf:"health_interval"`
|
||||
// RendezvousTimeout bounds the wait for a multi-node rendezvous point to
|
||||
// start listening.
|
||||
RendezvousTimeout time.Duration `koanf:"rendezvous_timeout"`
|
||||
// AliveGrace is how long a non-serving rank must stay up before it counts
|
||||
// as ready. There is no external signal for a worker, so staying alive is
|
||||
// the only evidence available.
|
||||
AliveGrace time.Duration `koanf:"alive_grace"`
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
WorkDir: "/var/lib/llmgw-agent",
|
||||
LogLines: 512,
|
||||
StopTimeout: 60 * time.Second,
|
||||
HealthInterval: 2 * time.Second,
|
||||
RendezvousTimeout: 5 * time.Minute,
|
||||
AliveGrace: 5 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() {
|
||||
d := DefaultConfig()
|
||||
if c.WorkDir == "" {
|
||||
c.WorkDir = d.WorkDir
|
||||
}
|
||||
if c.LogLines <= 0 {
|
||||
c.LogLines = d.LogLines
|
||||
}
|
||||
if c.StopTimeout <= 0 {
|
||||
c.StopTimeout = d.StopTimeout
|
||||
}
|
||||
if c.HealthInterval <= 0 {
|
||||
c.HealthInterval = d.HealthInterval
|
||||
}
|
||||
if c.RendezvousTimeout <= 0 {
|
||||
c.RendezvousTimeout = d.RendezvousTimeout
|
||||
}
|
||||
if c.AliveGrace <= 0 {
|
||||
c.AliveGrace = d.AliveGrace
|
||||
}
|
||||
}
|
||||
|
||||
// HealthChecker reports whether a shard is serving. It is an interface so the
|
||||
// supervisor can be tested without an HTTP engine.
|
||||
type HealthChecker interface {
|
||||
// Ready reports whether the shard at baseURL serves model.
|
||||
Ready(ctx context.Context, baseURL, model string) error
|
||||
}
|
||||
|
||||
// Shard is one supervised engine process.
|
||||
type Shard struct {
|
||||
ID string
|
||||
Spec engines.ShardSpec
|
||||
|
||||
mu sync.Mutex
|
||||
state State
|
||||
pid int
|
||||
err error
|
||||
started time.Time
|
||||
readyAt time.Time
|
||||
exitCode int
|
||||
|
||||
logs *ringBuffer
|
||||
cmd *exec.Cmd
|
||||
// sidecars are the shard's supporting processes, torn down with it.
|
||||
sidecars []*exec.Cmd
|
||||
cancel context.CancelFunc
|
||||
done chan struct{}
|
||||
// adopted marks a shard inherited from a previous agent: it has no
|
||||
// exec.Cmd, so it is watched by polling and killed by signal alone.
|
||||
adopted bool
|
||||
}
|
||||
|
||||
// Status is a snapshot of a shard, safe to send to the control plane.
|
||||
type Status struct {
|
||||
ID string
|
||||
State State
|
||||
PID int
|
||||
Err string
|
||||
Started time.Time
|
||||
ReadyAt time.Time
|
||||
ExitCode int
|
||||
}
|
||||
|
||||
// Status returns the shard's current status.
|
||||
func (s *Shard) Status() Status {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
st := Status{
|
||||
ID: s.ID, State: s.state, PID: s.pid,
|
||||
Started: s.started, ReadyAt: s.readyAt, ExitCode: s.exitCode,
|
||||
}
|
||||
if s.err != nil {
|
||||
st.Err = s.err.Error()
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
// Logs returns the last n lines of engine output.
|
||||
func (s *Shard) Logs(n int) []string { return s.logs.tail(n) }
|
||||
|
||||
func (s *Shard) setState(st State, err error) {
|
||||
s.mu.Lock()
|
||||
s.state = st
|
||||
if err != nil {
|
||||
s.err = err
|
||||
}
|
||||
if st == Ready && s.readyAt.IsZero() {
|
||||
s.readyAt = time.Now()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// Supervisor starts, watches and stops engine processes on this node.
|
||||
type Supervisor struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
health HealthChecker
|
||||
|
||||
mu sync.Mutex
|
||||
shards map[string]*Shard
|
||||
}
|
||||
|
||||
// New builds a supervisor.
|
||||
func New(cfg Config, health HealthChecker, log *slog.Logger) (*Supervisor, error) {
|
||||
cfg.normalize()
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if err := os.MkdirAll(cfg.WorkDir, 0o750); err != nil {
|
||||
return nil, fmt.Errorf("work dir: %w", err)
|
||||
}
|
||||
return &Supervisor{cfg: cfg, log: log, health: health, shards: map[string]*Shard{}}, nil
|
||||
}
|
||||
|
||||
// ErrAlreadyRunning means a shard with that ID is already supervised.
|
||||
var ErrAlreadyRunning = errors.New("shard is already running")
|
||||
|
||||
// Start launches a shard.
|
||||
//
|
||||
// It is idempotent by ID: a repeated StartShard for the same shard is a no-op
|
||||
// rather than a second process, because the control plane retries commands
|
||||
// (plan §5.6).
|
||||
func (s *Supervisor) Start(ctx context.Context, id string, spec engines.ShardSpec, cmds []engines.Command) (*Shard, error) {
|
||||
if len(cmds) == 0 {
|
||||
return nil, fmt.Errorf("shard %s: no commands to run", id)
|
||||
}
|
||||
s.mu.Lock()
|
||||
if existing, ok := s.shards[id]; ok {
|
||||
s.mu.Unlock()
|
||||
if existing.Status().State.Terminal() {
|
||||
s.Remove(id)
|
||||
} else {
|
||||
return existing, ErrAlreadyRunning
|
||||
}
|
||||
s.mu.Lock()
|
||||
}
|
||||
|
||||
sh := &Shard{
|
||||
ID: id,
|
||||
Spec: spec,
|
||||
state: Pending,
|
||||
logs: newRingBuffer(s.cfg.LogLines),
|
||||
started: time.Now(),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
s.shards[id] = sh
|
||||
s.mu.Unlock()
|
||||
|
||||
if err := s.launch(ctx, sh, cmds); err != nil {
|
||||
sh.setState(Failed, err)
|
||||
close(sh.done)
|
||||
return sh, err
|
||||
}
|
||||
return sh, nil
|
||||
}
|
||||
|
||||
// launch starts a shard's processes in order.
|
||||
//
|
||||
// Every command but the last is a sidecar: it keeps running for the shard's
|
||||
// lifetime and is torn down with it, but its exit is not what defines the
|
||||
// shard's. The last command is the engine, and its exit is the shard's exit.
|
||||
func (s *Supervisor) launch(ctx context.Context, sh *Shard, cmds []engines.Command) error {
|
||||
runCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
|
||||
sh.cancel = cancel
|
||||
|
||||
for i, c := range cmds[:len(cmds)-1] {
|
||||
proc, err := s.spawn(runCtx, sh, c)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return fmt.Errorf("sidecar %d (%s): %w", i, c.Name, err)
|
||||
}
|
||||
sh.mu.Lock()
|
||||
sh.sidecars = append(sh.sidecars, proc)
|
||||
sh.mu.Unlock()
|
||||
|
||||
if c.WaitForLeaderPort {
|
||||
// The next process cannot start until the rendezvous point is
|
||||
// listening: a Ray worker cannot join a head that is not up, and
|
||||
// the leader's server cannot claim GPUs from a cluster that does
|
||||
// not exist yet.
|
||||
addr := fmt.Sprintf("%s:%d", sh.Spec.LeaderAddr, sh.Spec.LeaderPort)
|
||||
if err := waitForPort(runCtx, addr, s.cfg.RendezvousTimeout); err != nil {
|
||||
cancel()
|
||||
return fmt.Errorf("waiting for %s (%s): %w", addr, c.Name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
cmd := cmds[len(cmds)-1]
|
||||
|
||||
c := exec.CommandContext(runCtx, cmd.Path, cmd.Args...)
|
||||
c.Env = cmd.Env
|
||||
c.Dir = s.cfg.WorkDir
|
||||
// Its own process group, so stopping kills the engine's children too.
|
||||
// Engines spawn workers; signalling only the parent leaves those workers
|
||||
// holding GPU memory.
|
||||
c.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
// CommandContext kills only the leader, so replace that with a group kill.
|
||||
c.Cancel = func() error { return killGroup(c.Process.Pid, syscall.SIGKILL) }
|
||||
|
||||
stdout, err := c.StdoutPipe()
|
||||
if err != nil {
|
||||
cancel()
|
||||
return err
|
||||
}
|
||||
c.Stderr = c.Stdout.(io.Writer) // interleave, as the engine writes them
|
||||
|
||||
if err := c.Start(); err != nil {
|
||||
cancel()
|
||||
return fmt.Errorf("starting %s: %w", cmd.Path, err)
|
||||
}
|
||||
|
||||
sh.mu.Lock()
|
||||
sh.cmd = c
|
||||
sh.pid = c.Process.Pid
|
||||
sh.state = Starting
|
||||
sh.mu.Unlock()
|
||||
|
||||
if err := s.writeState(sh); err != nil {
|
||||
s.log.Warn("could not write pid file; adoption after a restart will not work",
|
||||
"shard", sh.ID, "err", err)
|
||||
}
|
||||
|
||||
go s.pump(sh, stdout, cmd.Name)
|
||||
go s.wait(sh)
|
||||
return nil
|
||||
}
|
||||
|
||||
// spawn starts one process in its own group and pipes its output to the log.
|
||||
func (s *Supervisor) spawn(ctx context.Context, sh *Shard, c engines.Command) (*exec.Cmd, error) {
|
||||
proc := exec.CommandContext(ctx, c.Path, c.Args...)
|
||||
proc.Env = c.Env
|
||||
proc.Dir = s.cfg.WorkDir
|
||||
proc.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
proc.Cancel = func() error { return killGroup(proc.Process.Pid, syscall.SIGKILL) }
|
||||
|
||||
stdout, err := proc.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
proc.Stderr = proc.Stdout.(io.Writer)
|
||||
|
||||
if err := proc.Start(); err != nil {
|
||||
return nil, fmt.Errorf("starting %s: %w", c.Path, err)
|
||||
}
|
||||
go s.pump(sh, stdout, c.Name)
|
||||
// Reaped so it does not become a zombie; its exit does not end the shard.
|
||||
go func() { _ = proc.Wait() }()
|
||||
return proc, nil
|
||||
}
|
||||
|
||||
// waitForPort blocks until addr accepts a connection.
|
||||
func waitForPort(ctx context.Context, addr string, timeout time.Duration) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
|
||||
var d net.Dialer
|
||||
for {
|
||||
conn, err := d.DialContext(ctx, "tcp", addr)
|
||||
if err == nil {
|
||||
_ = conn.Close()
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return fmt.Errorf("not listening after %v: %w", timeout, ctx.Err())
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pump copies a process's output into the ring buffer and the log.
|
||||
func (s *Supervisor) pump(sh *Shard, r io.Reader, name string) {
|
||||
sc := bufio.NewScanner(r)
|
||||
// Engine tracebacks and progress bars produce long lines; the default
|
||||
// 64 KiB limit truncates them exactly when they matter.
|
||||
sc.Buffer(make([]byte, 0, 64<<10), 1<<20)
|
||||
|
||||
for sc.Scan() {
|
||||
line := sc.Text()
|
||||
if name != "" {
|
||||
// Prefixed so the two processes of a Ray shard stay distinguishable
|
||||
// in a single log tail.
|
||||
line = "[" + name + "] " + line
|
||||
}
|
||||
sh.logs.add(line)
|
||||
s.log.Debug("engine", "shard", sh.ID, "line", line)
|
||||
}
|
||||
}
|
||||
|
||||
// wait reaps the process and records why it ended.
|
||||
func (s *Supervisor) wait(sh *Shard) {
|
||||
defer close(sh.done)
|
||||
|
||||
err := sh.cmd.Wait()
|
||||
|
||||
sh.mu.Lock()
|
||||
prev := sh.state
|
||||
if ee := new(exec.ExitError); errors.As(err, &ee) {
|
||||
sh.exitCode = ee.ExitCode()
|
||||
}
|
||||
switch {
|
||||
case prev == Draining || prev == Stopped:
|
||||
// Expected: we asked it to stop.
|
||||
sh.state = Stopped
|
||||
case err == nil:
|
||||
// An engine exiting on its own is a failure, not a success: it is
|
||||
// supposed to serve until told otherwise.
|
||||
sh.state = Failed
|
||||
sh.err = errors.New("engine exited unexpectedly")
|
||||
default:
|
||||
sh.state = Failed
|
||||
sh.err = err
|
||||
}
|
||||
final, code := sh.state, sh.exitCode
|
||||
sh.mu.Unlock()
|
||||
|
||||
s.removeState(sh)
|
||||
if final == Failed {
|
||||
s.log.Error("shard failed", "shard", sh.ID, "exit_code", code, "err", err,
|
||||
"last_output", strings.Join(sh.Logs(5), " | "))
|
||||
} else {
|
||||
s.log.Info("shard stopped", "shard", sh.ID, "exit_code", code)
|
||||
}
|
||||
}
|
||||
|
||||
// WaitReady polls until the shard serves, it fails, or the timeout expires.
|
||||
//
|
||||
// The timeout comes from the profile because it is model-dependent: loading a
|
||||
// 70B model across two nodes legitimately takes tens of minutes, while a small
|
||||
// model failing to load for ten minutes means something is wrong.
|
||||
func (s *Supervisor) WaitReady(ctx context.Context, sh *Shard, baseURL, model string, timeout time.Duration) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
|
||||
t := time.NewTicker(s.cfg.HealthInterval)
|
||||
defer t.Stop()
|
||||
|
||||
sh.setState(Loading, nil)
|
||||
for {
|
||||
select {
|
||||
case <-sh.done:
|
||||
return fmt.Errorf("shard %s exited during startup: %s", sh.ID, sh.Status().Err)
|
||||
case <-ctx.Done():
|
||||
last := strings.Join(sh.Logs(10), "\n")
|
||||
return fmt.Errorf("shard %s not ready within %v; last output:\n%s", sh.ID, timeout, last)
|
||||
case <-t.C:
|
||||
if err := s.health.Ready(ctx, baseURL, model); err != nil {
|
||||
continue
|
||||
}
|
||||
sh.setState(Ready, nil)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WaitAlive is readiness for a rank that serves no HTTP.
|
||||
//
|
||||
// A worker in a multi-node replica has no endpoint to poll, so the only
|
||||
// evidence available is that it started and stayed started. That is weaker than
|
||||
// a health check, and deliberately so: the replica's real readiness is the
|
||||
// leader answering, which only happens once every worker has joined. Polling a
|
||||
// worker's port instead would time out and fail a replica that is coming up
|
||||
// perfectly well.
|
||||
func (s *Supervisor) WaitAlive(ctx context.Context, sh *Shard) error {
|
||||
sh.setState(Loading, nil)
|
||||
|
||||
select {
|
||||
case <-sh.done:
|
||||
return fmt.Errorf("shard %s exited during startup: %s", sh.ID, sh.Status().Err)
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(s.cfg.AliveGrace):
|
||||
}
|
||||
|
||||
if st := sh.Status().State; st.Terminal() {
|
||||
return fmt.Errorf("shard %s is %s: %s", sh.ID, st, sh.Status().Err)
|
||||
}
|
||||
sh.setState(Ready, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Drain marks a shard as taking no new work. The router stops routing to it;
|
||||
// the process keeps running until Stop.
|
||||
func (s *Supervisor) Drain(id string) error {
|
||||
sh, err := s.Get(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if st := sh.Status().State; !st.Live() {
|
||||
return fmt.Errorf("shard %s is %s, cannot drain", id, st)
|
||||
}
|
||||
sh.setState(Draining, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop terminates a shard: SIGTERM to the whole group, then SIGKILL after
|
||||
// StopTimeout.
|
||||
func (s *Supervisor) Stop(ctx context.Context, id string) error {
|
||||
sh, err := s.Get(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
sh.mu.Lock()
|
||||
pid := sh.pid
|
||||
sidecars := slices.Clone(sh.sidecars)
|
||||
if sh.state.Live() {
|
||||
sh.state = Draining // marks the exit as expected for wait() and poll()
|
||||
}
|
||||
sh.mu.Unlock()
|
||||
|
||||
// Sidecars go too. A Ray head left behind holds the port the next replica
|
||||
// needs, and would let a later shard join a cluster that serves nothing.
|
||||
defer func() {
|
||||
for _, sc := range sidecars {
|
||||
if sc.Process == nil {
|
||||
continue
|
||||
}
|
||||
if err := killGroup(sc.Process.Pid, syscall.SIGTERM); err != nil &&
|
||||
!errors.Is(err, os.ErrProcessDone) {
|
||||
s.log.Warn("stopping a sidecar failed", "shard", id, "pid", sc.Process.Pid, "err", err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// A pid of zero means the process never started; an adopted shard has a pid
|
||||
// but no exec.Cmd, and must still be killed by signal.
|
||||
if pid == 0 {
|
||||
sh.setState(Stopped, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Signal the group, not the process: engines fork workers that hold GPU
|
||||
// memory, and those outlive a signal sent only to the leader.
|
||||
if err := killGroup(pid, syscall.SIGTERM); err != nil && !errors.Is(err, os.ErrProcessDone) {
|
||||
s.log.Warn("SIGTERM failed", "shard", id, "err", err)
|
||||
}
|
||||
|
||||
// force kills the group outright. exec.Cmd's cancel does this for shards we
|
||||
// started; an adopted shard has none, so signal its group directly.
|
||||
force := func() {
|
||||
if sh.cancel != nil {
|
||||
sh.cancel()
|
||||
return
|
||||
}
|
||||
if err := killGroup(pid, syscall.SIGKILL); err != nil && !errors.Is(err, os.ErrProcessDone) {
|
||||
s.log.Warn("SIGKILL failed", "shard", id, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case <-sh.done:
|
||||
return nil
|
||||
case <-time.After(s.cfg.StopTimeout):
|
||||
s.log.Warn("shard did not exit in time, killing", "shard", id, "timeout", s.cfg.StopTimeout)
|
||||
force()
|
||||
case <-ctx.Done():
|
||||
force()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
select {
|
||||
case <-sh.done:
|
||||
case <-time.After(5 * time.Second):
|
||||
return fmt.Errorf("shard %s survived SIGKILL", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StopAll stops every live shard.
|
||||
func (s *Supervisor) StopAll(ctx context.Context) {
|
||||
for _, id := range s.IDs() {
|
||||
if err := s.Stop(ctx, id); err != nil {
|
||||
s.log.Warn("stopping shard failed", "shard", id, "err", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns a supervised shard.
|
||||
func (s *Supervisor) Get(id string) (*Shard, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
sh, ok := s.shards[id]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown shard %q", id)
|
||||
}
|
||||
return sh, nil
|
||||
}
|
||||
|
||||
// Remove forgets a terminated shard.
|
||||
func (s *Supervisor) Remove(id string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
delete(s.shards, id)
|
||||
}
|
||||
|
||||
// IDs lists the supervised shards.
|
||||
func (s *Supervisor) IDs() []string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]string, 0, len(s.shards))
|
||||
for id := range s.shards {
|
||||
out = append(out, id)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Statuses returns every shard's status.
|
||||
func (s *Supervisor) Statuses() []Status {
|
||||
s.mu.Lock()
|
||||
shards := make([]*Shard, 0, len(s.shards))
|
||||
for _, sh := range s.shards {
|
||||
shards = append(shards, sh)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
out := make([]Status, 0, len(shards))
|
||||
for _, sh := range shards {
|
||||
out = append(out, sh.Status())
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// --- state files, for adoption after an agent restart ---
|
||||
|
||||
func (s *Supervisor) statePath(id string) string {
|
||||
return filepath.Join(s.cfg.WorkDir, "shard-"+id+".json")
|
||||
}
|
||||
|
||||
// shardState is what survives an agent restart.
|
||||
//
|
||||
// The pid alone is not enough for two separate reasons. Pids are reused, so the
|
||||
// command line is kept as a fingerprint — adopting an unrelated process would
|
||||
// be worse than leaking the original. And the full spec is kept because a
|
||||
// recovered shard has to be described back to the control plane: without it the
|
||||
// control plane cannot tell which replica the process belongs to, and would
|
||||
// place a duplicate while the original keeps holding the GPUs (plan §6.5).
|
||||
type shardState struct {
|
||||
ShardID string `json:"shard_id"`
|
||||
ReplicaID string `json:"replica_id"`
|
||||
PID int `json:"pid"`
|
||||
Cmdline []string `json:"cmdline"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
Spec engines.ShardSpec `json:"spec"`
|
||||
}
|
||||
|
||||
// writeState records everything needed to recognise and describe the process later.
|
||||
func (s *Supervisor) writeState(sh *Shard) error {
|
||||
b, err := json.Marshal(shardState{
|
||||
ShardID: sh.ID,
|
||||
ReplicaID: sh.Spec.ReplicaID,
|
||||
PID: sh.pid,
|
||||
Cmdline: sh.cmd.Args,
|
||||
StartedAt: sh.started,
|
||||
Spec: sh.Spec,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Written atomically: a half-written state file after a crash would be
|
||||
// indistinguishable from a corrupt one, and the shard would be killed.
|
||||
tmp := s.statePath(sh.ID) + ".tmp"
|
||||
if err := os.WriteFile(tmp, b, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, s.statePath(sh.ID))
|
||||
}
|
||||
|
||||
func (s *Supervisor) removeState(sh *Shard) {
|
||||
_ = os.Remove(s.statePath(sh.ID))
|
||||
}
|
||||
|
||||
// Orphan is an engine process left behind by a previous agent.
|
||||
type Orphan struct {
|
||||
ShardID string
|
||||
ReplicaID string
|
||||
PID int
|
||||
Cmdline []string
|
||||
Spec engines.ShardSpec
|
||||
StartedAt time.Time
|
||||
// Alive reports whether the process still exists.
|
||||
Alive bool
|
||||
}
|
||||
|
||||
// FindOrphans scans the work directory for processes a previous agent started.
|
||||
func (s *Supervisor) FindOrphans() ([]Orphan, error) {
|
||||
entries, err := os.ReadDir(s.cfg.WorkDir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var out []Orphan
|
||||
for _, e := range entries {
|
||||
name := e.Name()
|
||||
if !strings.HasPrefix(name, "shard-") || !strings.HasSuffix(name, ".json") {
|
||||
continue
|
||||
}
|
||||
|
||||
b, err := os.ReadFile(filepath.Join(s.cfg.WorkDir, name))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
var st shardState
|
||||
if err := json.Unmarshal(b, &st); err != nil || st.ShardID == "" || st.PID <= 0 {
|
||||
// A corrupt state file describes a process that cannot be
|
||||
// identified, so it cannot be adopted safely. Leave it for
|
||||
// KillOrphan to clean up by shard id from the filename.
|
||||
id := strings.TrimSuffix(strings.TrimPrefix(name, "shard-"), ".json")
|
||||
out = append(out, Orphan{ShardID: id})
|
||||
continue
|
||||
}
|
||||
|
||||
out = append(out, Orphan{
|
||||
ShardID: st.ShardID,
|
||||
ReplicaID: st.ReplicaID,
|
||||
PID: st.PID,
|
||||
Cmdline: st.Cmdline,
|
||||
Spec: st.Spec,
|
||||
StartedAt: st.StartedAt,
|
||||
Alive: processAlive(st.PID),
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Adopt takes over a still-running engine from a previous agent.
|
||||
//
|
||||
// The process cannot be waited on — it is not this process's child — so it is
|
||||
// watched by polling instead. That is the whole difference between an adopted
|
||||
// shard and a started one; everything downstream treats them the same.
|
||||
func (s *Supervisor) Adopt(o Orphan) (*Shard, error) {
|
||||
if !o.Alive {
|
||||
return nil, fmt.Errorf("shard %s: process %d is gone", o.ShardID, o.PID)
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
if _, exists := s.shards[o.ShardID]; exists {
|
||||
s.mu.Unlock()
|
||||
return nil, ErrAlreadyRunning
|
||||
}
|
||||
sh := &Shard{
|
||||
ID: o.ShardID,
|
||||
Spec: o.Spec,
|
||||
state: Starting,
|
||||
pid: o.PID,
|
||||
logs: newRingBuffer(s.cfg.LogLines),
|
||||
started: o.StartedAt,
|
||||
done: make(chan struct{}),
|
||||
adopted: true,
|
||||
}
|
||||
s.shards[o.ShardID] = sh
|
||||
s.mu.Unlock()
|
||||
|
||||
// Output is gone with the previous agent; say so rather than presenting an
|
||||
// empty log as if the engine had been silent.
|
||||
sh.logs.add("[adopted from a previous agent; earlier output is not available]")
|
||||
|
||||
go s.poll(sh)
|
||||
return sh, nil
|
||||
}
|
||||
|
||||
// poll watches an adopted process, which cannot be reaped with Wait.
|
||||
func (s *Supervisor) poll(sh *Shard) {
|
||||
defer close(sh.done)
|
||||
|
||||
// Polled rather than waited on, so the interval bounds how quickly a stop
|
||||
// is observed. Short enough that Stop does not sit out its whole timeout.
|
||||
t := time.NewTicker(200 * time.Millisecond)
|
||||
defer t.Stop()
|
||||
for range t.C {
|
||||
if processAlive(sh.pid) {
|
||||
continue
|
||||
}
|
||||
sh.mu.Lock()
|
||||
expected := sh.state == Draining || sh.state == Stopped
|
||||
if expected {
|
||||
sh.state = Stopped
|
||||
} else {
|
||||
sh.state = Failed
|
||||
sh.err = errors.New("adopted engine exited")
|
||||
}
|
||||
sh.mu.Unlock()
|
||||
|
||||
s.removeState(sh)
|
||||
s.log.Info("adopted shard exited", "shard", sh.ID, "expected", expected)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// KillOrphan terminates a leftover process group and removes its pid file.
|
||||
func (s *Supervisor) KillOrphan(o Orphan) error {
|
||||
if o.Alive {
|
||||
if err := killGroup(o.PID, syscall.SIGTERM); err != nil {
|
||||
return err
|
||||
}
|
||||
// Give it a moment to free GPU memory before insisting.
|
||||
deadline := time.Now().Add(s.cfg.StopTimeout)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(o.PID) {
|
||||
break
|
||||
}
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
}
|
||||
if processAlive(o.PID) {
|
||||
_ = killGroup(o.PID, syscall.SIGKILL)
|
||||
}
|
||||
}
|
||||
// An already-missing state file is success: the previous agent may have
|
||||
// reaped the process itself between the scan and this call.
|
||||
if err := os.Remove(s.statePath(o.ShardID)); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,729 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
|
||||
// scriptedHealth reports ready once the flag is set.
|
||||
type scriptedHealth struct{ ready atomic.Bool }
|
||||
|
||||
func (h *scriptedHealth) Ready(context.Context, string, string) error {
|
||||
if h.ready.Load() {
|
||||
return nil
|
||||
}
|
||||
return errors.New("not ready")
|
||||
}
|
||||
|
||||
func newSupervisor(t *testing.T) (*Supervisor, *scriptedHealth) {
|
||||
t.Helper()
|
||||
h := &scriptedHealth{}
|
||||
s, err := New(Config{
|
||||
WorkDir: t.TempDir(),
|
||||
LogLines: 64,
|
||||
StopTimeout: 3 * time.Second,
|
||||
HealthInterval: 20 * time.Millisecond,
|
||||
}, h, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { s.StopAll(context.Background()) })
|
||||
return s, h
|
||||
}
|
||||
|
||||
// shellCommand builds the single-command sequence for a shell script.
|
||||
func shellCommand(script string) []engines.Command {
|
||||
return []engines.Command{{Path: "/bin/sh", Args: []string{"-c", script}, Env: os.Environ()}}
|
||||
}
|
||||
|
||||
func waitState(t *testing.T, sh *Shard, want State, limit time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(limit)
|
||||
for time.Now().Before(deadline) {
|
||||
if sh.Status().State == want {
|
||||
return
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("shard stayed in %s, expected %s", sh.Status().State, want)
|
||||
}
|
||||
|
||||
func TestStartAndStop(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st := sh.Status(); st.State != Starting || st.PID == 0 {
|
||||
t.Fatalf("status after start: %+v", st)
|
||||
}
|
||||
|
||||
if err := s.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st := sh.Status().State; st != Stopped {
|
||||
t.Errorf("state after stop: %s", st)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnexpectedExitIsFailure(t *testing.T) {
|
||||
// An engine exiting on its own is a failure even with status 0: it is meant
|
||||
// to serve until told otherwise.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 0"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
if !strings.Contains(sh.Status().Err, "unexpectedly") {
|
||||
t.Errorf("err=%q", sh.Status().Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNonZeroExitRecordsCode(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 42"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
if got := sh.Status().ExitCode; got != 42 {
|
||||
t.Errorf("exit code=%d, expected 42", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopKillsChildProcesses(t *testing.T) {
|
||||
// The reason for the process group: engines fork workers that hold GPU
|
||||
// memory, and signalling only the leader leaves them behind.
|
||||
s, _ := newSupervisor(t)
|
||||
dir := t.TempDir()
|
||||
marker := filepath.Join(dir, "child.pid")
|
||||
|
||||
// The shell forks a child that writes its own pid, then both sleep.
|
||||
script := fmt.Sprintf(`sh -c 'echo $$ > %s; sleep 30' & sleep 30`, marker)
|
||||
if _, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand(script)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Wait for the child to announce itself.
|
||||
var childPID int
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if b, err := os.ReadFile(marker); err == nil {
|
||||
if _, err := fmt.Sscanf(strings.TrimSpace(string(b)), "%d", &childPID); err == nil && childPID > 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
if childPID == 0 {
|
||||
t.Fatal("child never started")
|
||||
}
|
||||
|
||||
if err := s.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
deadline = time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(childPID) {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Errorf("child process %d survived the stop; it would still hold GPU memory", childPID)
|
||||
}
|
||||
|
||||
func TestStopKillsAfterTimeout(t *testing.T) {
|
||||
// A process ignoring SIGTERM must still die.
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.StopTimeout = 300 * time.Millisecond
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
|
||||
shellCommand("trap '' TERM; sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
start := time.Now()
|
||||
if err := s.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if d := time.Since(start); d > 5*time.Second {
|
||||
t.Errorf("stop took %v", d)
|
||||
}
|
||||
if processAlive(pid) {
|
||||
t.Error("the process survived SIGKILL")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartIsIdempotent(t *testing.T) {
|
||||
// The control plane retries commands, so a repeated start must not produce
|
||||
// a second process (plan §5.6).
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
first, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if !errors.Is(err, ErrAlreadyRunning) {
|
||||
t.Fatalf("expected ErrAlreadyRunning, got %v", err)
|
||||
}
|
||||
if first.Status().PID != second.Status().PID {
|
||||
t.Error("a second process was started for the same shard")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestartAfterTerminalState(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 1"))
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
|
||||
// Once it has failed, the same id may be started again.
|
||||
again, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatalf("restarting a failed shard: %v", err)
|
||||
}
|
||||
if again.Status().State != Starting {
|
||||
t.Errorf("state=%s", again.Status().State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitReady(t *testing.T) {
|
||||
s, h := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
go func() {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
h.ready.Store(true)
|
||||
}()
|
||||
|
||||
if err := s.WaitReady(t.Context(), sh, "http://x", "m", 3*time.Second); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
st := sh.Status()
|
||||
if st.State != Ready {
|
||||
t.Errorf("state=%s", st.State)
|
||||
}
|
||||
if st.ReadyAt.IsZero() {
|
||||
t.Error("ReadyAt was not recorded")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitReadyTimeoutIncludesLogs(t *testing.T) {
|
||||
// When a model fails to load, the engine output is the only diagnosis
|
||||
// available, so the error has to carry it.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
|
||||
shellCommand("echo 'CUDA out of memory'; sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond) // let the line be captured
|
||||
|
||||
err = s.WaitReady(t.Context(), sh, "http://x", "m", 200*time.Millisecond)
|
||||
if err == nil {
|
||||
t.Fatal("expected a timeout")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "CUDA out of memory") {
|
||||
t.Errorf("the error does not carry engine output: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitReadyDetectsEarlyExit(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 1"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Should return as soon as the process dies, not after the full timeout.
|
||||
start := time.Now()
|
||||
if err := s.WaitReady(t.Context(), sh, "http://x", "m", 10*time.Second); err == nil {
|
||||
t.Fatal("expected an error")
|
||||
}
|
||||
if d := time.Since(start); d > 3*time.Second {
|
||||
t.Errorf("waited %v for a process that had already exited", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDrain(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
if err := s.Drain("a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := sh.Status().State; got != Draining {
|
||||
t.Errorf("state=%s", got)
|
||||
}
|
||||
// A drained shard's process is still alive; only routing stops.
|
||||
if !processAlive(sh.Status().PID) {
|
||||
t.Error("draining killed the process; it should keep serving in-flight work")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogsAreCaptured(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{},
|
||||
shellCommand("echo line1; echo line2 >&2; echo line3; sleep 0.2"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
|
||||
logs := strings.Join(sh.Logs(10), "\n")
|
||||
for _, want := range []string{"line1", "line2", "line3"} {
|
||||
if !strings.Contains(logs, want) {
|
||||
t.Errorf("%q missing from logs:\n%s", want, logs)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- orphan adoption ---
|
||||
|
||||
func TestFindOrphans(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// Simulate an agent restart: a new supervisor over the same work dir.
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
orphans, err := s2.FindOrphans()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orphans) != 1 {
|
||||
t.Fatalf("expected 1 orphan, got %d", len(orphans))
|
||||
}
|
||||
o := orphans[0]
|
||||
if o.ShardID != "a" || o.PID != pid || !o.Alive {
|
||||
t.Errorf("orphan=%+v (want shard a, pid %d, alive)", o, pid)
|
||||
}
|
||||
// The command line is stored as a fingerprint, because pids get reused and
|
||||
// adopting an unrelated process is worse than leaking the original.
|
||||
if len(o.Cmdline) == 0 || !strings.Contains(strings.Join(o.Cmdline, " "), "sleep 30") {
|
||||
t.Errorf("cmdline fingerprint missing: %v", o.Cmdline)
|
||||
}
|
||||
// And the spec, without which the control plane could not tell which
|
||||
// replica the process belongs to and would place a duplicate.
|
||||
if o.ReplicaID != "rep-1" {
|
||||
t.Errorf("replica id not recovered: %q", o.ReplicaID)
|
||||
}
|
||||
if o.Spec.HTTPPort != 20100 {
|
||||
t.Errorf("spec not recovered: %+v", o.Spec)
|
||||
}
|
||||
|
||||
if err := s2.KillOrphan(o); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(pid) {
|
||||
break
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
if processAlive(pid) {
|
||||
t.Error("the orphan survived KillOrphan")
|
||||
}
|
||||
if _, err := os.Stat(s2.statePath("a")); !os.IsNotExist(err) {
|
||||
t.Error("the state file was left behind")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateFileRemovedOnExit(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 0"))
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
if _, err := os.Stat(s.statePath("a")); !os.IsNotExist(err) {
|
||||
t.Error("the state file outlived the process; a later scan would find a dead orphan")
|
||||
}
|
||||
}
|
||||
|
||||
// --- state machine and ring buffer ---
|
||||
|
||||
func TestStateStrings(t *testing.T) {
|
||||
for st, want := range map[State]string{
|
||||
Pending: "pending", Starting: "starting", Loading: "loading",
|
||||
Ready: "ready", Draining: "draining", Stopped: "stopped", Failed: "failed",
|
||||
} {
|
||||
if got := st.String(); got != want {
|
||||
t.Errorf("%d -> %q, want %q", int(st), got, want)
|
||||
}
|
||||
}
|
||||
if !Stopped.Terminal() || !Failed.Terminal() || Ready.Terminal() {
|
||||
t.Error("Terminal is wrong")
|
||||
}
|
||||
if !Ready.Live() || !Draining.Live() || Stopped.Live() {
|
||||
t.Error("Live is wrong")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRingBuffer(t *testing.T) {
|
||||
r := newRingBuffer(3)
|
||||
if got := r.tail(10); len(got) != 0 {
|
||||
t.Errorf("empty buffer returned %v", got)
|
||||
}
|
||||
|
||||
r.add("a")
|
||||
r.add("b")
|
||||
if got := strings.Join(r.tail(10), ","); got != "a,b" {
|
||||
t.Errorf("tail=%q", got)
|
||||
}
|
||||
|
||||
r.add("c")
|
||||
r.add("d") // evicts "a"
|
||||
if got := strings.Join(r.tail(10), ","); got != "b,c,d" {
|
||||
t.Errorf("after wrap tail=%q, want b,c,d", got)
|
||||
}
|
||||
if got := strings.Join(r.tail(2), ","); got != "c,d" {
|
||||
t.Errorf("tail(2)=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptTakesOverRunningShard(t *testing.T) {
|
||||
// The control plane restart fix depends on this: a shard that survived the
|
||||
// agent must come back under supervision, not be restarted or leaked.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20200}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// A new supervisor over the same work dir, as after an agent restart.
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 3 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
orphans, err := s2.FindOrphans()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orphans) != 1 || !orphans[0].Alive {
|
||||
t.Fatalf("orphan not found: %+v", orphans)
|
||||
}
|
||||
|
||||
adopted, err := s2.Adopt(orphans[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := adopted.Status().PID; got != pid {
|
||||
t.Errorf("adopted pid %d, expected %d", got, pid)
|
||||
}
|
||||
if adopted.Spec.ReplicaID != "rep-1" {
|
||||
t.Errorf("spec lost in adoption: %+v", adopted.Spec)
|
||||
}
|
||||
// The earlier output is gone; saying so beats presenting an empty log as if
|
||||
// the engine had been silent.
|
||||
if logs := strings.Join(adopted.Logs(5), " "); !strings.Contains(logs, "adopted") {
|
||||
t.Errorf("adoption was not noted in the log: %q", logs)
|
||||
}
|
||||
|
||||
// Stopping an adopted shard must actually kill it: it has no exec.Cmd, so
|
||||
// the signal path is the only one available.
|
||||
if err := s2.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(pid) {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Error("the adopted process survived Stop")
|
||||
}
|
||||
|
||||
func TestAdoptRefusesDeadProcess(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
if _, err := s.Adopt(Orphan{ShardID: "gone", PID: 999999, Alive: false}); err == nil {
|
||||
t.Error("a dead process was adopted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptedExitIsReported(t *testing.T) {
|
||||
// An adopted engine dying has to be noticed; nothing waits on it.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a"}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
orphans, _ := s2.FindOrphans()
|
||||
adopted, err := s2.Adopt(orphans[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Kill it behind the supervisor's back.
|
||||
_ = killGroup(pid, syscall.SIGKILL)
|
||||
waitState(t, adopted, Failed, 5*time.Second)
|
||||
}
|
||||
|
||||
func TestCorruptStateFileIsCleanedUp(t *testing.T) {
|
||||
// A half-written state file describes a process that cannot be identified,
|
||||
// so it cannot be adopted safely.
|
||||
s, _ := newSupervisor(t)
|
||||
if err := os.WriteFile(s.statePath("broken"), []byte("{not json"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
orphans, err := s.FindOrphans()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orphans) != 1 || orphans[0].ShardID != "broken" || orphans[0].Alive {
|
||||
t.Fatalf("corrupt state file not reported for cleanup: %+v", orphans)
|
||||
}
|
||||
if err := s.KillOrphan(orphans[0]); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(s.statePath("broken")); !os.IsNotExist(err) {
|
||||
t.Error("the corrupt state file was left behind")
|
||||
}
|
||||
}
|
||||
|
||||
// --- multi-process shards (plan §11, phase 4) ---
|
||||
|
||||
func TestSidecarRunsBeforeEngine(t *testing.T) {
|
||||
// The Ray shape: a supporting process starts first and keeps running, then
|
||||
// the engine starts against it.
|
||||
s, _ := newSupervisor(t)
|
||||
dir := t.TempDir()
|
||||
order := filepath.Join(dir, "order")
|
||||
|
||||
// Each writes to the order file and to stdout: the file records the
|
||||
// sequence, stdout exercises per-process log attribution.
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
|
||||
fmt.Sprintf("echo sidecar >> %s; echo hello-from-sidecar; sleep 30", order)}, Env: os.Environ()},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c",
|
||||
fmt.Sprintf("echo engine >> %s; echo hello-from-engine; sleep 30", order)}, Env: os.Environ()},
|
||||
}
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if b, err := os.ReadFile(order); err == nil && strings.Count(string(b), "\n") == 2 {
|
||||
if got := strings.Fields(string(b)); got[0] != "sidecar" || got[1] != "engine" {
|
||||
t.Fatalf("start order was %v", got)
|
||||
}
|
||||
break
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
|
||||
// The shard's pid is the engine's, not the sidecar's: the engine's exit is
|
||||
// what ends the shard.
|
||||
if sh.Status().PID == 0 {
|
||||
t.Fatal("no engine pid recorded")
|
||||
}
|
||||
// Both processes appear in one log tail, distinguishable by prefix.
|
||||
deadline = time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
logs := strings.Join(sh.Logs(20), "\n")
|
||||
if strings.Contains(logs, "[sidecar] hello-from-sidecar") &&
|
||||
strings.Contains(logs, "[engine] hello-from-engine") {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Errorf("output is not attributed per process:\n%s", strings.Join(sh.Logs(20), "\n"))
|
||||
}
|
||||
|
||||
func TestStopKillsSidecar(t *testing.T) {
|
||||
// A Ray head left behind holds the port the next replica needs, and would
|
||||
// let a later shard join a cluster that serves nothing.
|
||||
s, _ := newSupervisor(t)
|
||||
dir := t.TempDir()
|
||||
marker := filepath.Join(dir, "sidecar.pid")
|
||||
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
|
||||
fmt.Sprintf("echo $$ > %s; sleep 30", marker)}, Env: os.Environ()},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
|
||||
}
|
||||
if _, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var sidePID int
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if b, err := os.ReadFile(marker); err == nil {
|
||||
if _, err := fmt.Sscanf(strings.TrimSpace(string(b)), "%d", &sidePID); err == nil && sidePID > 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
if sidePID == 0 {
|
||||
t.Fatal("the sidecar never started")
|
||||
}
|
||||
|
||||
if err := s.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deadline = time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(sidePID) {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Errorf("the sidecar (%d) survived the stop and still holds its port", sidePID)
|
||||
}
|
||||
|
||||
func TestSidecarFailureFailsTheShard(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/nonexistent/binary", Env: os.Environ()},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
|
||||
}
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, cmds)
|
||||
if err == nil {
|
||||
t.Fatal("a shard whose sidecar could not start was reported healthy")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "sidecar") {
|
||||
t.Errorf("the error does not identify the sidecar: %v", err)
|
||||
}
|
||||
if sh.Status().State != Failed {
|
||||
t.Errorf("state=%s", sh.Status().State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitForLeaderPortBlocksUntilListening(t *testing.T) {
|
||||
// A Ray worker cannot join a head that is not up, and the leader's server
|
||||
// cannot claim GPUs from a cluster that does not exist yet.
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.RendezvousTimeout = 5 * time.Second
|
||||
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
port := ln.Addr().(*net.TCPAddr).Port
|
||||
_ = ln.Close() // nothing listening yet
|
||||
|
||||
spec := engines.ShardSpec{LeaderAddr: "127.0.0.1", LeaderPort: port}
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c",
|
||||
fmt.Sprintf("sleep 0.4; nc -l %d", port)}, Env: os.Environ(),
|
||||
WaitForLeaderPort: true},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
if _, err := s.Start(t.Context(), "a", spec, cmds); err != nil {
|
||||
t.Skipf("nc is unavailable for the rendezvous stub: %v", err)
|
||||
}
|
||||
if d := time.Since(start); d < 300*time.Millisecond {
|
||||
t.Errorf("the engine started after %v, before the rendezvous point was listening", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRendezvousTimeout(t *testing.T) {
|
||||
// Nothing ever listens, so the shard must fail rather than hang forever.
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.RendezvousTimeout = 300 * time.Millisecond
|
||||
|
||||
spec := engines.ShardSpec{LeaderAddr: "127.0.0.1", LeaderPort: 1}
|
||||
cmds := []engines.Command{
|
||||
{Name: "sidecar", Path: "/bin/sh", Args: []string{"-c", "sleep 30"},
|
||||
Env: os.Environ(), WaitForLeaderPort: true},
|
||||
{Name: "engine", Path: "/bin/sh", Args: []string{"-c", "sleep 30"}, Env: os.Environ()},
|
||||
}
|
||||
if _, err := s.Start(t.Context(), "a", spec, cmds); err == nil {
|
||||
t.Fatal("the shard did not fail when the rendezvous point never appeared")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitAliveForNonServingRank(t *testing.T) {
|
||||
// A worker has no endpoint to poll, so staying up is the only evidence
|
||||
// available. This is what stopped multi-node replicas coming up at all.
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.AliveGrace = 150 * time.Millisecond
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{Rank: 1}, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.WaitAlive(t.Context(), sh); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := sh.Status().State; got != Ready {
|
||||
t.Errorf("state=%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitAliveFailsOnEarlyExit(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
s.cfg.AliveGrace = 2 * time.Second
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{Rank: 1}, shellCommand("exit 1"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
start := time.Now()
|
||||
if err := s.WaitAlive(t.Context(), sh); err == nil {
|
||||
t.Fatal("a worker that died immediately was reported ready")
|
||||
}
|
||||
if d := time.Since(start); d > 1500*time.Millisecond {
|
||||
t.Errorf("waited %v for a process that had already exited", d)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||
)
|
||||
|
||||
func newAgentCmd() *cobra.Command {
|
||||
var configPath string
|
||||
|
||||
var (
|
||||
server string
|
||||
nodeID string
|
||||
mgmtAddr string
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "agent",
|
||||
Short: "Node agent: discovery and engine supervision",
|
||||
Long: "Runs on every GPU node. It discovers the hardware, connects to the control\n" +
|
||||
"plane, and starts, supervises and stops inference engines on demand.\n\n" +
|
||||
"The agent dials out and keeps one stream open. Engines keep serving if the\n" +
|
||||
"agent restarts or the control plane goes away: the router talks to them\n" +
|
||||
"directly, and orphan adoption picks them up again.",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
log := logger()
|
||||
|
||||
cfg, err := config.LoadAgent(configPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if server != "" {
|
||||
cfg.Server = server
|
||||
}
|
||||
if nodeID != "" {
|
||||
cfg.NodeID = nodeID
|
||||
}
|
||||
if mgmtAddr != "" {
|
||||
cfg.MgmtAddr = mgmtAddr
|
||||
}
|
||||
if cfg.NodeID == "" {
|
||||
host, err := os.Hostname()
|
||||
if err != nil {
|
||||
return fmt.Errorf("node_id is unset and the hostname is unavailable: %w", err)
|
||||
}
|
||||
cfg.NodeID = host
|
||||
}
|
||||
if cfg.Hostname == "" {
|
||||
cfg.Hostname, _ = os.Hostname()
|
||||
}
|
||||
cfg.Version = version.Version
|
||||
|
||||
a, err := agent.New(cfg, agent.HTTP2Client(30*time.Second), log)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
log.Info("agent starting",
|
||||
"node", cfg.NodeID, "server", cfg.Server,
|
||||
"fabrics", len(cfg.Fabrics), "nvml", discovery.NVMLSupported())
|
||||
|
||||
err = a.Run(ctx)
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&configPath, "config", "", "configuration file")
|
||||
cmd.Flags().StringVar(&server, "server", "", "control plane address (overrides the config)")
|
||||
cmd.Flags().StringVar(&nodeID, "node-id", "", "node id (defaults to the hostname)")
|
||||
cmd.Flags().StringVar(&mgmtAddr, "mgmt-addr", "", "address engines are reachable on")
|
||||
cmd.AddCommand(newInventoryCmd())
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newInventoryCmd() *cobra.Command {
|
||||
var (
|
||||
asJSON bool
|
||||
smiBin string
|
||||
ibRoot string
|
||||
skipGPU bool
|
||||
useNVML bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "inventory",
|
||||
Short: "Show this node's GPUs and InfiniBand adapters",
|
||||
Long: "Reports what the scheduler would see on this node: GPU classes, health,\n" +
|
||||
"and the fabric each adapter is attached to.\n\n" +
|
||||
"Run it before adding a node to the cluster — it is the fastest way to find\n" +
|
||||
"a card another process is already using, or a port with no subnet manager.",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
inv, err := discovery.Discover(cmd.Context(), discovery.Options{
|
||||
NvidiaSMI: smiBin,
|
||||
IBRoot: ibRoot,
|
||||
PreferNVML: useNVML,
|
||||
SkipGPUs: skipGPU,
|
||||
}, logger())
|
||||
if err != nil {
|
||||
// Not fatal: a node whose GPUs cannot be read still has IB
|
||||
// worth showing, and the error is itself the useful output.
|
||||
fmt.Fprintf(os.Stderr, "%v\n", err)
|
||||
inv.Source = "unavailable"
|
||||
}
|
||||
|
||||
if asJSON {
|
||||
enc := json.NewEncoder(os.Stdout)
|
||||
enc.SetIndent("", " ")
|
||||
return enc.Encode(inv)
|
||||
}
|
||||
printInventory(inv)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
f := cmd.Flags()
|
||||
f.BoolVar(&asJSON, "json", false, "emit JSON")
|
||||
f.StringVar(&smiBin, "nvidia-smi", "", "path to nvidia-smi")
|
||||
f.StringVar(&ibRoot, "ib-root", discovery.SysfsRoot, "sysfs InfiniBand directory")
|
||||
f.BoolVar(&skipGPU, "no-gpu", false, "skip GPU discovery")
|
||||
f.BoolVar(&useNVML, "nvml", true, "prefer the NVML backend when this build has one (it is the only source of NVLink topology)")
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
func printInventory(inv discovery.Inventory) {
|
||||
w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
|
||||
|
||||
fmt.Fprintf(w, "GPUs (%d, via %s)\n", len(inv.GPUs), inv.Source)
|
||||
if len(inv.GPUs) == 0 {
|
||||
fmt.Fprintln(w, " none")
|
||||
} else {
|
||||
fmt.Fprintln(w, " IDX\tCLASS\tNAME\tMEM\tCC\tHEALTH")
|
||||
for _, g := range inv.GPUs {
|
||||
health := "ok"
|
||||
if !g.Health.OK() {
|
||||
health = g.Health.Problems[0]
|
||||
}
|
||||
fmt.Fprintf(w, " %d\t%s\t%s\t%d MiB\t%s\t%s\n",
|
||||
g.Index, g.Class(), g.Name, g.MemoryTotalMB, g.ComputeCapability, health)
|
||||
}
|
||||
fmt.Fprintf(w, " classes: %v\n", inv.Classes())
|
||||
if !discovery.NVMLSupported() {
|
||||
fmt.Fprintln(w, " NVLink topology unavailable: rebuild with -tags nvml.")
|
||||
fmt.Fprintln(w, " Without it, placement cannot keep a tensor-parallel group inside")
|
||||
fmt.Fprintln(w, " one NVLink domain (plan §6.2).")
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Fprintf(w, "\nInfiniBand adapters (%d)\n", len(inv.HCAs))
|
||||
if len(inv.HCAs) == 0 {
|
||||
// Worth stating plainly: such a node is its own single-node island and
|
||||
// can never host a multi-node replica.
|
||||
fmt.Fprintln(w, " none — this node cannot host multi-node replicas")
|
||||
} else {
|
||||
fmt.Fprintln(w, " HCA\tPORT\tSTATE\tRATE\tLINK\tGID PREFIX\tNUMA")
|
||||
for _, h := range inv.HCAs {
|
||||
for _, p := range h.Ports {
|
||||
fmt.Fprintf(w, " %s\t%d\t%s\t%d Gb/s\t%s\t%s\t%d\n",
|
||||
h.Name, p.Num, p.State, p.RateGbps, p.LinkLayer, p.GIDPrefix, h.NUMANode)
|
||||
}
|
||||
}
|
||||
fmt.Fprintln(w, "\n Fabric identity is not derived from the GID prefix by default:")
|
||||
fmt.Fprintln(w, " OpenSM's default is fe80:: on every fabric, so a shared prefix proves")
|
||||
fmt.Fprintln(w, " nothing. Label fabrics explicitly in agent.yaml (plan §5.2).")
|
||||
}
|
||||
|
||||
_ = w.Flush()
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
// ctlClient talks to a server's admin API.
|
||||
type ctlClient struct {
|
||||
base string
|
||||
http *http.Client
|
||||
}
|
||||
|
||||
func newCtlClient(addr string) *ctlClient {
|
||||
if !strings.HasPrefix(addr, "http://") && !strings.HasPrefix(addr, "https://") {
|
||||
addr = "http://" + addr
|
||||
}
|
||||
return &ctlClient{
|
||||
base: strings.TrimRight(addr, "/"),
|
||||
http: &http.Client{Timeout: 15 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *ctlClient) get(ctx context.Context, path string, out any) error {
|
||||
return c.do(ctx, http.MethodGet, path, nil, out)
|
||||
}
|
||||
|
||||
func (c *ctlClient) do(ctx context.Context, method, path string, body, out any) error {
|
||||
var rdr io.Reader
|
||||
if body != nil {
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rdr = bytes.NewReader(b)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.base+path, rdr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
|
||||
resp, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("reaching the admin API at %s: %w\n"+
|
||||
"(is the server running, and is --server pointing at its admin_listen port?)", c.base, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
var e struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<12))
|
||||
if json.Unmarshal(b, &e) == nil && e.Error != "" {
|
||||
return fmt.Errorf("%s: %s", resp.Status, e.Error)
|
||||
}
|
||||
return fmt.Errorf("%s: %s", resp.Status, strings.TrimSpace(string(b)))
|
||||
}
|
||||
if out == nil {
|
||||
return nil
|
||||
}
|
||||
return json.NewDecoder(resp.Body).Decode(out)
|
||||
}
|
||||
|
||||
func newCtlCmd() *cobra.Command {
|
||||
var (
|
||||
server string
|
||||
asJSON bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "ctl",
|
||||
Short: "Inspect and operate a running cluster",
|
||||
Long: "Talks to a server's admin API, which listens on admin_listen (:9090 by\n" +
|
||||
"default) rather than the client-facing port.",
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&server, "server", "127.0.0.1:9090", "admin API address")
|
||||
cmd.PersistentFlags().BoolVar(&asJSON, "json", false, "emit JSON")
|
||||
|
||||
client := func() *ctlClient { return newCtlClient(server) }
|
||||
|
||||
cmd.AddCommand(
|
||||
newNodesCmd(client, &asJSON),
|
||||
newModelsCmd(client, &asJSON),
|
||||
newReplicasCmd(client, &asJSON),
|
||||
newExplainCmd(client, &asJSON),
|
||||
newNodeStateCmd(client, "cordon", "cordoned",
|
||||
"Stop placing new work on a node; what runs there keeps running"),
|
||||
newNodeStateCmd(client, "uncordon", "ready",
|
||||
"Allow placement on a node again"),
|
||||
newNodeStateCmd(client, "drain", "draining",
|
||||
"Move work off a node, one replica at a time"),
|
||||
)
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newNodesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "nodes",
|
||||
Short: "List nodes and their capacity",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var nodes []adminapi.NodeView
|
||||
if err := client().get(cmd.Context(), "/api/v1/nodes", &nodes); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(nodes)
|
||||
}
|
||||
if len(nodes) == 0 {
|
||||
fmt.Println("No nodes have registered.")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "NODE\tSTATE\tGPUS\tCLASSES\tFABRICS\tREPLICAS\tLAST SEEN")
|
||||
for _, n := range nodes {
|
||||
fmt.Fprintf(w, "%s\t%s\t%d/%d free\t%s\t%s\t%d\t%s\n",
|
||||
n.ID, n.State, n.GPUFree, n.GPUTotal,
|
||||
join(n.Classes), join(n.Fabrics), n.Replicas, n.LastSeenAgo)
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
// Unhealthy GPUs are usually the answer when a node looks idle but
|
||||
// nothing schedules onto it.
|
||||
for _, n := range nodes {
|
||||
for _, p := range n.Unhealthy {
|
||||
fmt.Printf("\n%s: %s (kept out of placement)", n.ID, p)
|
||||
}
|
||||
}
|
||||
if hasUnhealthy(nodes) {
|
||||
fmt.Println()
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newModelsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "models",
|
||||
Short: "List the catalog and how many replicas each model has",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var models []adminapi.ModelView
|
||||
if err := client().get(cmd.Context(), "/api/v1/models", &models); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(models)
|
||||
}
|
||||
if len(models) == 0 {
|
||||
fmt.Println("The catalog is empty.")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "MODEL\tREPLICAS\tALIASES\tPROFILES\tFALLBACKS")
|
||||
for _, m := range models {
|
||||
replicas := fmt.Sprintf("%d/%d ready", m.Ready, m.Min)
|
||||
if m.Total > m.Ready {
|
||||
replicas += fmt.Sprintf(" (%d total)", m.Total)
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n",
|
||||
m.Name, replicas, join(m.Aliases), join(m.Profiles), join(m.Fallbacks))
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
var model string
|
||||
c := &cobra.Command{
|
||||
Use: "replicas",
|
||||
Short: "List replicas and the shards they are made of",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
path := "/api/v1/replicas"
|
||||
if model != "" {
|
||||
path += "?model=" + model
|
||||
}
|
||||
var reps []adminapi.ReplicaView
|
||||
if err := client().get(cmd.Context(), path, &reps); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(reps)
|
||||
}
|
||||
if len(reps) == 0 {
|
||||
fmt.Println("No replicas are running.")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "REPLICA\tMODEL\tSTATE\tPOOL\tNODES\tAGE\tENDPOINT")
|
||||
for _, r := range reps {
|
||||
var nodes []string
|
||||
for _, s := range r.Shards {
|
||||
nodes = append(nodes, fmt.Sprintf("%s:%v", s.Node, s.GPUs))
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
r.ID, r.Model, r.State, r.Pool, join(nodes),
|
||||
(time.Duration(r.AgeSec) * time.Second).String(), r.Endpoint)
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
for _, r := range reps {
|
||||
if r.Err != "" {
|
||||
fmt.Printf("\n%s: %s\n", r.ID, r.Err)
|
||||
}
|
||||
for _, s := range r.Shards {
|
||||
if s.Err != "" {
|
||||
fmt.Printf(" %s on %s: %s\n", s.ID, s.Node, s.Err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
c.Flags().StringVar(&model, "model", "", "only this model's replicas")
|
||||
return c
|
||||
}
|
||||
|
||||
func newExplainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "explain-placement <model>",
|
||||
Short: "Explain why a model has not been placed",
|
||||
Long: "Prints every candidate placement considered and the reason it was\n" +
|
||||
"rejected. \"Unschedulable\" on its own says nothing about which of six\n" +
|
||||
"constraints bit.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
var v adminapi.ExplainView
|
||||
if err := client().get(cmd.Context(),
|
||||
"/api/v1/explain-placement?model="+args[0], &v); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(v)
|
||||
}
|
||||
|
||||
if v.Placed && len(v.Rejections) == 0 {
|
||||
fmt.Printf("%s is placed; nothing was rejected.\n", v.Model)
|
||||
return nil
|
||||
}
|
||||
if len(v.Rejections) == 0 {
|
||||
fmt.Printf("No placement has been attempted for %s yet.\n", v.Model)
|
||||
return nil
|
||||
}
|
||||
|
||||
if v.Placed {
|
||||
fmt.Printf("%s is placed. The candidates rejected along the way:\n\n", v.Model)
|
||||
} else {
|
||||
fmt.Printf("%s could not be placed. Every candidate and why:\n\n", v.Model)
|
||||
}
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "PROFILE\tPOOL\tNODE\tREASON")
|
||||
for _, r := range v.Rejections {
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", r.Profile, dash(r.Pool), dash(r.Node), r.Reason)
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newNodeStateCmd(client func() *ctlClient, use, state, short string) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: use + " <node>",
|
||||
Short: short,
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
var out map[string]string
|
||||
err := client().do(cmd.Context(), http.MethodPost,
|
||||
"/api/v1/nodes/"+args[0]+"/state",
|
||||
map[string]string{"state": state}, &out)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf("%s is now %s\n", out["id"], out["state"])
|
||||
if state == "draining" {
|
||||
fmt.Println("Replicas will be moved off one at a time, each replacement " +
|
||||
"brought up before the original is stopped.")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// --- output helpers ---
|
||||
|
||||
func tab() *tabwriter.Writer {
|
||||
return tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
|
||||
}
|
||||
|
||||
func emitJSON(v any) error {
|
||||
enc := json.NewEncoder(os.Stdout)
|
||||
enc.SetIndent("", " ")
|
||||
return enc.Encode(v)
|
||||
}
|
||||
|
||||
func join(v []string) string {
|
||||
if len(v) == 0 {
|
||||
return "-"
|
||||
}
|
||||
return strings.Join(v, ",")
|
||||
}
|
||||
|
||||
func dash(s string) string {
|
||||
if s == "" {
|
||||
return "-"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func hasUnhealthy(nodes []adminapi.NodeView) bool {
|
||||
for _, n := range nodes {
|
||||
if len(n.Unhealthy) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
+146
-4
@@ -4,17 +4,26 @@ 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"
|
||||
)
|
||||
|
||||
@@ -55,8 +64,10 @@ func newServerCmd() *cobra.Command {
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(cfg.Endpoints) == 0 {
|
||||
return errors.New("no endpoints configured; use --endpoint or a config file")
|
||||
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)
|
||||
@@ -70,10 +81,37 @@ func newServerCmd() *cobra.Command {
|
||||
rt.Start(ctx)
|
||||
defer rt.Stop()
|
||||
|
||||
adminHandler, err := rt.AdminHandler()
|
||||
// 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(),
|
||||
@@ -82,7 +120,7 @@ func newServerCmd() *cobra.Command {
|
||||
// enforces its own per-phase deadlines.
|
||||
}
|
||||
adminSrv := &http.Server{
|
||||
Handler: adminHandler,
|
||||
Handler: adminMux,
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
|
||||
@@ -103,14 +141,39 @@ func newServerCmd() *cobra.Command {
|
||||
"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)
|
||||
})
|
||||
@@ -127,6 +190,85 @@ func newServerCmd() *cobra.Command {
|
||||
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
|
||||
|
||||
@@ -1,44 +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 newAgentCmd() *cobra.Command {
|
||||
var configPath string
|
||||
cmd := &cobra.Command{
|
||||
Use: "agent",
|
||||
Short: "Node agent running on each GPU host (phase 2)",
|
||||
RunE: func(_ *cobra.Command, _ []string) error {
|
||||
return errNotImplemented
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&configPath, "config", "agent.yaml", "configuration file")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newCtlCmd() *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "ctl",
|
||||
Short: "Administrative CLI (phase 3)",
|
||||
}
|
||||
for _, sub := range []struct{ use, short string }{
|
||||
{"nodes", "list nodes"},
|
||||
{"models", "show the model catalog"},
|
||||
{"replicas", "list replicas"},
|
||||
{"explain-placement", "explain why a model could not be placed"},
|
||||
{"bench", "measure profile capacity (§10)"},
|
||||
} {
|
||||
cmd.AddCommand(&cobra.Command{
|
||||
Use: sub.use,
|
||||
Short: sub.short,
|
||||
RunE: func(_ *cobra.Command, _ []string) error { return errNotImplemented },
|
||||
})
|
||||
}
|
||||
return cmd
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/knadh/koanf/parsers/yaml"
|
||||
env "github.com/knadh/koanf/providers/env/v2"
|
||||
"github.com/knadh/koanf/providers/file"
|
||||
"github.com/knadh/koanf/v2"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent"
|
||||
)
|
||||
|
||||
// LoadAgent reads an agent configuration, overlaying it on the defaults and
|
||||
// then applying LLMGW_-prefixed environment overrides.
|
||||
//
|
||||
// An empty path loads the defaults, which with --server and --node-id on the
|
||||
// command line is enough for a smoke test on a node with no IB.
|
||||
func LoadAgent(path string) (agent.Config, error) {
|
||||
cfg := agent.DefaultConfig()
|
||||
k := koanf.New(".")
|
||||
|
||||
if path != "" {
|
||||
if _, err := os.Stat(path); err != nil {
|
||||
return cfg, fmt.Errorf("config file: %w", err)
|
||||
}
|
||||
if err := k.Load(file.Provider(path), yaml.Parser()); err != nil {
|
||||
return cfg, fmt.Errorf("parsing %s: %w", path, err)
|
||||
}
|
||||
}
|
||||
|
||||
err := k.Load(env.Provider(".", env.Opt{
|
||||
Prefix: "LLMGW_",
|
||||
TransformFunc: func(key, value string) (string, any) {
|
||||
key = strings.ToLower(strings.TrimPrefix(key, "LLMGW_"))
|
||||
return strings.ReplaceAll(key, "__", "."), value
|
||||
},
|
||||
}), nil)
|
||||
if err != nil {
|
||||
return cfg, fmt.Errorf("environment overrides: %w", err)
|
||||
}
|
||||
|
||||
if err := k.UnmarshalWithConf("", &cfg, koanf.UnmarshalConf{Tag: "koanf"}); err != nil {
|
||||
return cfg, fmt.Errorf("decoding configuration: %w", err)
|
||||
}
|
||||
return cfg, validateAgent(cfg)
|
||||
}
|
||||
|
||||
// validateAgent rejects configurations that would fail confusingly at runtime.
|
||||
func validateAgent(c agent.Config) error {
|
||||
seen := map[string]struct{}{}
|
||||
for i, f := range c.Fabrics {
|
||||
if f.Fabric == "" {
|
||||
return fmt.Errorf("fabrics[%d]: name is required", i)
|
||||
}
|
||||
if _, dup := seen[f.Fabric]; dup {
|
||||
return fmt.Errorf("fabric %q is declared twice", f.Fabric)
|
||||
}
|
||||
seen[f.Fabric] = struct{}{}
|
||||
if len(f.HCAs) == 0 {
|
||||
// A fabric with no adapters is worse than no fabric at all: the
|
||||
// scheduler would place multi-node replicas here and every shard
|
||||
// would come up with NCCL unable to find a path (plan §5.3).
|
||||
return fmt.Errorf("fabric %q lists no HCAs", f.Fabric)
|
||||
}
|
||||
}
|
||||
|
||||
// The same adapter on two fabrics means one of the labels is wrong, and the
|
||||
// consequence is a replica that starts and then hangs on its first
|
||||
// all-reduce — the failure mode this labelling exists to prevent.
|
||||
byHCA := map[string]string{}
|
||||
for _, f := range c.Fabrics {
|
||||
for _, h := range f.HCAs {
|
||||
if prev, dup := byHCA[h]; dup {
|
||||
return fmt.Errorf("HCA %q is claimed by both %q and %q", h, prev, f.Fabric)
|
||||
}
|
||||
byHCA[h] = f.Fabric
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
+64
-31
@@ -15,6 +15,8 @@ import (
|
||||
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.
|
||||
@@ -25,11 +27,38 @@ type Server struct {
|
||||
// client-facing port can be exposed without exposing diagnostics.
|
||||
AdminListen string `koanf:"admin_listen"`
|
||||
|
||||
Router Router `koanf:"router"`
|
||||
Models []Model `koanf:"models"`
|
||||
// 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"`
|
||||
@@ -106,16 +135,6 @@ type Scoring struct {
|
||||
SessionLoadFactor float64 `koanf:"session_load_factor"`
|
||||
}
|
||||
|
||||
// Model is a catalog entry as clients see it.
|
||||
type Model struct {
|
||||
Name string `koanf:"name"`
|
||||
Aliases []string `koanf:"aliases"`
|
||||
// Fallbacks are tried, in order, when this model has no usable endpoint.
|
||||
Fallbacks []string `koanf:"fallbacks"`
|
||||
// RoutingProfile overrides the router's default profile for this model.
|
||||
RoutingProfile string `koanf:"routing_profile"`
|
||||
}
|
||||
|
||||
// Endpoint is one engine replica. Static in phase 1.
|
||||
type Endpoint struct {
|
||||
ID string `koanf:"id"`
|
||||
@@ -131,6 +150,16 @@ 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,
|
||||
@@ -234,25 +263,29 @@ func (s Server) Validate() error {
|
||||
seen[e.ID] = struct{}{}
|
||||
}
|
||||
|
||||
names := make(map[string]struct{}, len(s.Models))
|
||||
for i, m := range s.Models {
|
||||
if m.Name == "" {
|
||||
return fmt.Errorf("models[%d]: name is required", i)
|
||||
}
|
||||
if _, dup := names[m.Name]; dup {
|
||||
return fmt.Errorf("model %q is declared twice", m.Name)
|
||||
}
|
||||
names[m.Name] = struct{}{}
|
||||
}
|
||||
// Aliases share a namespace with names, so a collision must be caught here
|
||||
// rather than silently shadowing a model at request time.
|
||||
for _, m := range s.Models {
|
||||
for _, a := range m.Aliases {
|
||||
if _, clash := names[a]; clash {
|
||||
return fmt.Errorf("model %q: alias %q collides with another model name", m.Name, a)
|
||||
}
|
||||
names[a] = struct{}{}
|
||||
}
|
||||
// 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
|
||||
}
|
||||
@@ -6,6 +6,8 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
)
|
||||
|
||||
func write(t *testing.T, content string) string {
|
||||
@@ -128,9 +130,11 @@ func TestValidateRejectsAliasCollision(t *testing.T) {
|
||||
// An alias shadowing another model would silently misroute every request
|
||||
// for that model, so it has to fail at load time.
|
||||
s := Default()
|
||||
s.Models = []Model{
|
||||
{Name: "llama-3.3-70b", Aliases: []string{"qwen"}},
|
||||
{Name: "qwen"},
|
||||
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") {
|
||||
@@ -140,9 +144,11 @@ func TestValidateRejectsAliasCollision(t *testing.T) {
|
||||
|
||||
func TestValidateAcceptsDistinctAliases(t *testing.T) {
|
||||
s := Default()
|
||||
s.Models = []Model{
|
||||
{Name: "llama-3.3-70b", Aliases: []string{"llama", "llama3"}},
|
||||
{Name: "qwen2.5-32b", Aliases: []string{"qwen"}},
|
||||
s.Models = []catalog.Model{
|
||||
{Name: "llama-3.3-70b", Source: "/m", Aliases: []string{"llama", "llama3"},
|
||||
Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}},
|
||||
{Name: "qwen2.5-32b", Source: "/m", Aliases: []string{"qwen"},
|
||||
Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}},
|
||||
}
|
||||
if err := s.Validate(); err != nil {
|
||||
t.Errorf("valid aliases rejected: %v", err)
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
// Package adminapi is the control plane's management surface, which `ctl`
|
||||
// consumes (plan §11, phase 3).
|
||||
//
|
||||
// It is mounted on the admin listener rather than the client-facing one. The
|
||||
// two have different audiences: one serves inference to whoever holds an API
|
||||
// key, the other lets a caller cordon nodes and read the cluster's shape.
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// API serves the management endpoints.
|
||||
type API struct {
|
||||
ctl *control.Controller
|
||||
}
|
||||
|
||||
// New builds the API over a controller.
|
||||
func New(ctl *control.Controller) *API { return &API{ctl: ctl} }
|
||||
|
||||
// Register mounts the endpoints on a mux under /api/v1.
|
||||
func (a *API) Register(mux *http.ServeMux) {
|
||||
mux.HandleFunc("GET /api/v1/nodes", a.listNodes)
|
||||
mux.HandleFunc("GET /api/v1/models", a.listModels)
|
||||
mux.HandleFunc("GET /api/v1/replicas", a.listReplicas)
|
||||
mux.HandleFunc("GET /api/v1/explain-placement", a.explainPlacement)
|
||||
mux.HandleFunc("POST /api/v1/nodes/{id}/state", a.setNodeState)
|
||||
}
|
||||
|
||||
// --- responses ---
|
||||
|
||||
// NodeView is a node as `ctl nodes` shows it.
|
||||
type NodeView struct {
|
||||
ID string `json:"id"`
|
||||
Hostname string `json:"hostname,omitempty"`
|
||||
MgmtAddr string `json:"mgmt_addr,omitempty"`
|
||||
State string `json:"state"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
|
||||
GPUTotal int `json:"gpu_total"`
|
||||
GPUFree int `json:"gpu_free"`
|
||||
Classes []string `json:"gpu_classes,omitempty"`
|
||||
Fabrics []string `json:"fabrics,omitempty"`
|
||||
|
||||
// Unhealthy lists GPUs kept out of placement and why, which is usually the
|
||||
// answer when a node looks idle but nothing will schedule onto it.
|
||||
Unhealthy []string `json:"unhealthy_gpus,omitempty"`
|
||||
|
||||
LastSeenAgo string `json:"last_seen_ago"`
|
||||
AgentVersion string `json:"agent_version,omitempty"`
|
||||
Replicas int `json:"replicas"`
|
||||
}
|
||||
|
||||
// ModelView is a model as `ctl models` shows it.
|
||||
type ModelView struct {
|
||||
Name string `json:"name"`
|
||||
Aliases []string `json:"aliases,omitempty"`
|
||||
Source string `json:"source"`
|
||||
Min int `json:"min_replicas"`
|
||||
Max int `json:"max_replicas"`
|
||||
Ready int `json:"ready_replicas"`
|
||||
Total int `json:"total_replicas"`
|
||||
Fallbacks []string `json:"fallbacks,omitempty"`
|
||||
Profiles []string `json:"profiles"`
|
||||
}
|
||||
|
||||
// ReplicaView is a replica as `ctl replicas` shows it.
|
||||
type ReplicaView struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Profile string `json:"profile"`
|
||||
State string `json:"state"`
|
||||
Pool string `json:"pool,omitempty"`
|
||||
Fabric string `json:"fabric,omitempty"`
|
||||
Endpoint string `json:"endpoint,omitempty"`
|
||||
AgeSec int64 `json:"age_seconds"`
|
||||
Err string `json:"error,omitempty"`
|
||||
Shards []ShardView `json:"shards"`
|
||||
}
|
||||
|
||||
// ShardView is one shard of a replica.
|
||||
type ShardView struct {
|
||||
ID string `json:"id"`
|
||||
Node string `json:"node"`
|
||||
Rank int `json:"rank"`
|
||||
GPUs []int `json:"gpus"`
|
||||
State string `json:"state"`
|
||||
Err string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// ExplainView is why a model could not be placed.
|
||||
type ExplainView struct {
|
||||
Model string `json:"model"`
|
||||
// Placed is true when the model has replicas, in which case there is
|
||||
// nothing to explain.
|
||||
Placed bool `json:"placed"`
|
||||
// Rejections is empty when placement was never attempted, which is itself
|
||||
// the answer: the model is satisfied, or the reconciler has not run yet.
|
||||
Rejections []RejectionView `json:"rejections,omitempty"`
|
||||
}
|
||||
|
||||
// RejectionView is one rejected candidate.
|
||||
type RejectionView struct {
|
||||
Profile string `json:"profile"`
|
||||
Pool string `json:"pool,omitempty"`
|
||||
Node string `json:"node,omitempty"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
// --- handlers ---
|
||||
|
||||
func (a *API) listNodes(w http.ResponseWriter, _ *http.Request) {
|
||||
replicasOn := map[string]int{}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
for _, n := range rep.NodeIDs() {
|
||||
replicasOn[n]++
|
||||
}
|
||||
}
|
||||
|
||||
inv := a.ctl.Inventory()
|
||||
out := make([]NodeView, 0)
|
||||
for _, n := range inv.Nodes() {
|
||||
v := NodeView{
|
||||
ID: n.ID, Hostname: n.Hostname, MgmtAddr: n.MgmtAddr,
|
||||
State: n.State.String(), Labels: n.Labels,
|
||||
GPUTotal: len(n.GPUs), GPUFree: len(inv.FreeGPUs(n.ID)),
|
||||
Fabrics: n.FabricNames(), AgentVersion: n.AgentVersion,
|
||||
Replicas: replicasOn[n.ID],
|
||||
}
|
||||
seen := map[string]struct{}{}
|
||||
for _, g := range n.GPUs {
|
||||
if _, dup := seen[g.Class()]; !dup {
|
||||
seen[g.Class()] = struct{}{}
|
||||
v.Classes = append(v.Classes, g.Class())
|
||||
}
|
||||
if !g.Health.OK() {
|
||||
v.Unhealthy = append(v.Unhealthy,
|
||||
fmt.Sprintf("gpu %d: %s", g.Index, g.Health.Problems[0]))
|
||||
}
|
||||
}
|
||||
if !n.LastSeen.IsZero() {
|
||||
v.LastSeenAgo = time.Since(n.LastSeen).Round(time.Second).String()
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (a *API) listModels(w http.ResponseWriter, _ *http.Request) {
|
||||
ready, total := map[string]int{}, map[string]int{}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
total[rep.Model]++
|
||||
if rep.State == reconciler.Ready {
|
||||
ready[rep.Model]++
|
||||
}
|
||||
}
|
||||
|
||||
out := make([]ModelView, 0)
|
||||
for _, m := range a.ctl.Catalog().Models() {
|
||||
v := ModelView{
|
||||
Name: m.Name, Aliases: m.Aliases, Source: m.Source,
|
||||
Min: m.Scaling.Min, Max: m.Scaling.Max,
|
||||
Ready: ready[m.Name], Total: total[m.Name],
|
||||
Fallbacks: m.Fallbacks,
|
||||
}
|
||||
for _, p := range m.Profiles {
|
||||
v.Profiles = append(v.Profiles, p.Name)
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (a *API) listReplicas(w http.ResponseWriter, r *http.Request) {
|
||||
model := r.URL.Query().Get("model")
|
||||
|
||||
out := make([]ReplicaView, 0)
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
if model != "" && rep.Model != model {
|
||||
continue
|
||||
}
|
||||
v := ReplicaView{
|
||||
ID: rep.ID, Model: rep.Model, Profile: rep.Profile,
|
||||
State: rep.State.String(), Pool: rep.Pool, Fabric: rep.Fabric,
|
||||
Endpoint: rep.Endpoint, Err: rep.Err,
|
||||
AgeSec: int64(time.Since(rep.CreatedAt).Seconds()),
|
||||
}
|
||||
for _, s := range rep.Shards {
|
||||
v.Shards = append(v.Shards, ShardView{
|
||||
ID: s.ID, Node: s.NodeID, Rank: s.Rank,
|
||||
GPUs: s.GPUIndices, State: s.State.String(), Err: s.Err,
|
||||
})
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (a *API) explainPlacement(w http.ResponseWriter, r *http.Request) {
|
||||
model := r.URL.Query().Get("model")
|
||||
if model == "" {
|
||||
writeError(w, http.StatusBadRequest, "model is required")
|
||||
return
|
||||
}
|
||||
if _, ok := a.ctl.Catalog().Resolve(model); !ok {
|
||||
writeError(w, http.StatusNotFound, fmt.Sprintf("model %q is not in the catalog", model))
|
||||
return
|
||||
}
|
||||
|
||||
v := ExplainView{Model: model}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
if rep.Model == model && rep.State.Counts() {
|
||||
v.Placed = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if res, ok := a.ctl.Reconciler().Explain(model); ok {
|
||||
for _, rej := range res.Rejections {
|
||||
v.Rejections = append(v.Rejections, RejectionView{
|
||||
Profile: rej.Profile, Pool: rej.Pool, Node: rej.Node, Reason: rej.Reason,
|
||||
})
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, v)
|
||||
}
|
||||
|
||||
// setNodeState implements cordon, uncordon and drain.
|
||||
//
|
||||
// One endpoint rather than three verbs, because they are the same operation
|
||||
// with different targets and the difference matters: cordon stops new
|
||||
// placement, drain also moves existing work off.
|
||||
func (a *API) setNodeState(w http.ResponseWriter, r *http.Request) {
|
||||
id := r.PathValue("id")
|
||||
|
||||
var body struct {
|
||||
State string `json:"state"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
var st inventory.NodeState
|
||||
switch body.State {
|
||||
case "ready", "uncordon":
|
||||
st = inventory.Ready
|
||||
case "cordoned", "cordon":
|
||||
st = inventory.Cordoned
|
||||
case "draining", "drain":
|
||||
st = inventory.Draining
|
||||
default:
|
||||
writeError(w, http.StatusBadRequest,
|
||||
fmt.Sprintf("unknown state %q (want ready, cordoned or draining)", body.State))
|
||||
return
|
||||
}
|
||||
|
||||
if !a.ctl.Inventory().SetState(id, st) {
|
||||
writeError(w, http.StatusNotFound, fmt.Sprintf("no node %q", id))
|
||||
return
|
||||
}
|
||||
// Act on it now rather than at the next tick: an operator draining a node
|
||||
// is usually about to do something to it.
|
||||
a.ctl.Reconciler().Notify()
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]string{"id": id, "state": st.String()})
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, code int, v any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(code)
|
||||
_ = json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
|
||||
func writeError(w http.ResponseWriter, code int, msg string) {
|
||||
writeJSON(w, code, map[string]string{"error": msg})
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
)
|
||||
|
||||
// nullSink discards published endpoints; these tests are about the API.
|
||||
type nullSink struct{}
|
||||
|
||||
func (nullSink) SetEndpoints(context.Context, []endpoint.Info) {}
|
||||
|
||||
// nullDispatcher accepts every command.
|
||||
type nullDispatcher struct{}
|
||||
|
||||
func (nullDispatcher) StartShard(context.Context, reconciler.ShardCommand) error { return nil }
|
||||
func (nullDispatcher) StopShard(context.Context, string, string) error { return nil }
|
||||
func (nullDispatcher) DrainShard(context.Context, string, string) error { return nil }
|
||||
|
||||
func setup(t *testing.T, models []catalog.Model, nodes ...inventory.Node) (*API, *control.Controller) {
|
||||
t.Helper()
|
||||
cat, err := catalog.New(models)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctl := control.New(control.DefaultConfig(), cat, nullDispatcher{}, nullSink{}, slog.New(slog.DiscardHandler))
|
||||
for _, n := range nodes {
|
||||
ctl.Inventory().Upsert(n)
|
||||
}
|
||||
return New(ctl), ctl
|
||||
}
|
||||
|
||||
func gpuNode(id string, n int, fabrics ...string) inventory.Node {
|
||||
nd := inventory.Node{ID: id, State: inventory.Ready, AgentVersion: "v1"}
|
||||
for i := range n {
|
||||
nd.GPUs = append(nd.GPUs, discovery.GPU{
|
||||
Index: i, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559,
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
nd.Fabrics = append(nd.Fabrics, inventory.FabricAttachment{Fabric: f, HCAs: []string{"mlx5_0"}})
|
||||
}
|
||||
return nd
|
||||
}
|
||||
|
||||
func testModels() []catalog.Model {
|
||||
return []catalog.Model{{
|
||||
Name: "llama", Aliases: []string{"l"}, Source: "/m",
|
||||
Scaling: catalog.Scaling{Min: 2, Max: 4},
|
||||
Profiles: []catalog.Profile{{Name: "tp8", GPUClass: "h100-80g", TP: 8}},
|
||||
}}
|
||||
}
|
||||
|
||||
func call(t *testing.T, a *API, method, path string, body string, out any) int {
|
||||
t.Helper()
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
|
||||
var r *http.Request
|
||||
if body != "" {
|
||||
r = httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
} else {
|
||||
r = httptest.NewRequest(method, path, nil)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, r)
|
||||
|
||||
if out != nil && rec.Code == http.StatusOK {
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), out); err != nil {
|
||||
t.Fatalf("decoding %s: %v (%s)", path, err, rec.Body.String())
|
||||
}
|
||||
}
|
||||
return rec.Code
|
||||
}
|
||||
|
||||
func TestListNodes(t *testing.T) {
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8, "ib-a"))
|
||||
|
||||
var nodes []NodeView
|
||||
if code := call(t, a, "GET", "/api/v1/nodes", "", &nodes); code != 200 {
|
||||
t.Fatalf("status %d", code)
|
||||
}
|
||||
if len(nodes) != 1 {
|
||||
t.Fatalf("got %d nodes", len(nodes))
|
||||
}
|
||||
n := nodes[0]
|
||||
if n.ID != "n1" || n.State != "ready" {
|
||||
t.Errorf("%+v", n)
|
||||
}
|
||||
if n.GPUTotal != 8 || n.GPUFree != 8 {
|
||||
t.Errorf("gpus %d/%d", n.GPUFree, n.GPUTotal)
|
||||
}
|
||||
if len(n.Classes) != 1 || n.Classes[0] != "h100-80g" {
|
||||
t.Errorf("classes=%v", n.Classes)
|
||||
}
|
||||
if len(n.Fabrics) != 1 || n.Fabrics[0] != "ib-a" {
|
||||
t.Errorf("fabrics=%v", n.Fabrics)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnhealthyGPUsAreSurfaced(t *testing.T) {
|
||||
// Usually the answer when a node looks idle but nothing schedules onto it.
|
||||
n := gpuNode("n1", 4)
|
||||
n.GPUs[1].Health.Problems = []string{"3 uncorrected ECC errors"}
|
||||
a, _ := setup(t, testModels(), n)
|
||||
|
||||
var nodes []NodeView
|
||||
call(t, a, "GET", "/api/v1/nodes", "", &nodes)
|
||||
|
||||
if len(nodes[0].Unhealthy) != 1 {
|
||||
t.Fatalf("unhealthy=%v", nodes[0].Unhealthy)
|
||||
}
|
||||
if !strings.Contains(nodes[0].Unhealthy[0], "ECC") {
|
||||
t.Errorf("the reason was lost: %q", nodes[0].Unhealthy[0])
|
||||
}
|
||||
// And such a GPU is not counted as free.
|
||||
if nodes[0].GPUFree != 3 {
|
||||
t.Errorf("free=%d, an unhealthy GPU was counted", nodes[0].GPUFree)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListModels(t *testing.T) {
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8))
|
||||
|
||||
var models []ModelView
|
||||
call(t, a, "GET", "/api/v1/models", "", &models)
|
||||
|
||||
if len(models) != 1 {
|
||||
t.Fatalf("got %d models", len(models))
|
||||
}
|
||||
m := models[0]
|
||||
if m.Name != "llama" || m.Min != 2 || m.Max != 4 {
|
||||
t.Errorf("%+v", m)
|
||||
}
|
||||
if len(m.Aliases) != 1 || m.Aliases[0] != "l" {
|
||||
t.Errorf("aliases=%v", m.Aliases)
|
||||
}
|
||||
if len(m.Profiles) != 1 || m.Profiles[0] != "tp8" {
|
||||
t.Errorf("profiles=%v", m.Profiles)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplainPlacementForUnplaceableModel(t *testing.T) {
|
||||
// Four GPUs, a profile that needs eight.
|
||||
a, ctl := setup(t, testModels(), gpuNode("n1", 4))
|
||||
ctl.Reconciler().Reconcile(t.Context())
|
||||
|
||||
var v ExplainView
|
||||
if code := call(t, a, "GET", "/api/v1/explain-placement?model=llama", "", &v); code != 200 {
|
||||
t.Fatalf("status %d", code)
|
||||
}
|
||||
if v.Placed {
|
||||
t.Error("an unplaceable model was reported as placed")
|
||||
}
|
||||
if len(v.Rejections) == 0 {
|
||||
t.Fatal("no rejections were recorded; the operator learns nothing")
|
||||
}
|
||||
found := false
|
||||
for _, r := range v.Rejections {
|
||||
if strings.Contains(r.Reason, "needs 8") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("the reason does not say what was missing: %+v", v.Rejections)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplainPlacementRequiresAKnownModel(t *testing.T) {
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8))
|
||||
|
||||
if code := call(t, a, "GET", "/api/v1/explain-placement", "", nil); code != 400 {
|
||||
t.Errorf("missing model gave %d, expected 400", code)
|
||||
}
|
||||
if code := call(t, a, "GET", "/api/v1/explain-placement?model=nope", "", nil); code != 404 {
|
||||
t.Errorf("unknown model gave %d, expected 404", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetNodeState(t *testing.T) {
|
||||
a, ctl := setup(t, testModels(), gpuNode("n1", 8))
|
||||
|
||||
for _, c := range []struct {
|
||||
body string
|
||||
want inventory.NodeState
|
||||
}{
|
||||
{`{"state":"cordon"}`, inventory.Cordoned},
|
||||
{`{"state":"drain"}`, inventory.Draining},
|
||||
{`{"state":"uncordon"}`, inventory.Ready},
|
||||
{`{"state":"cordoned"}`, inventory.Cordoned},
|
||||
} {
|
||||
if code := call(t, a, "POST", "/api/v1/nodes/n1/state", c.body, nil); code != 200 {
|
||||
t.Fatalf("%s gave %d", c.body, code)
|
||||
}
|
||||
n, _ := ctl.Inventory().Get("n1")
|
||||
if n.State != c.want {
|
||||
t.Errorf("%s left the node %s, expected %s", c.body, n.State, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetNodeStateRejectsNonsense(t *testing.T) {
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8))
|
||||
|
||||
if code := call(t, a, "POST", "/api/v1/nodes/n1/state", `{"state":"banana"}`, nil); code != 400 {
|
||||
t.Errorf("an unknown state gave %d, expected 400", code)
|
||||
}
|
||||
if code := call(t, a, "POST", "/api/v1/nodes/nope/state", `{"state":"cordon"}`, nil); code != 404 {
|
||||
t.Errorf("an unknown node gave %d, expected 404", code)
|
||||
}
|
||||
if code := call(t, a, "POST", "/api/v1/nodes/n1/state", `not json`, nil); code != 400 {
|
||||
t.Errorf("malformed body gave %d, expected 400", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyClusterReturnsArraysNotNull(t *testing.T) {
|
||||
// A JSON null would make every consumer special-case it.
|
||||
a, _ := setup(t, testModels())
|
||||
|
||||
for _, path := range []string{"/api/v1/nodes", "/api/v1/replicas"} {
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, httptest.NewRequest("GET", path, nil))
|
||||
|
||||
if got := strings.TrimSpace(rec.Body.String()); got != "[]" {
|
||||
t.Errorf("%s returned %q, expected []", path, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
// Package agentsvc is the control plane's end of the agent stream (plan §5.6).
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// NodeFromRegister converts a registration into an inventory node.
|
||||
func NodeFromRegister(r *llmgwv1.Register) inventory.Node {
|
||||
n := inventory.Node{
|
||||
ID: r.GetNodeId(),
|
||||
Hostname: r.GetHostname(),
|
||||
MgmtAddr: r.GetMgmtAddr(),
|
||||
Labels: r.GetLabels(),
|
||||
State: inventory.Ready,
|
||||
AgentVersion: r.GetAgentVersion(),
|
||||
EngineVersions: r.GetEngineVersions(),
|
||||
}
|
||||
for _, g := range r.GetGpus() {
|
||||
n.GPUs = append(n.GPUs, discovery.GPU{
|
||||
Index: int(g.GetIndex()),
|
||||
UUID: g.GetUuid(),
|
||||
Name: g.GetName(),
|
||||
MemoryTotalMB: int(g.GetMemoryTotalMb()),
|
||||
MemoryUsedMB: int(g.GetMemoryUsedMb()),
|
||||
ComputeCapability: g.GetComputeCapability(),
|
||||
NUMANode: int(g.GetNumaNode()),
|
||||
PCIBusID: g.GetPciBusId(),
|
||||
DriverVersion: g.GetDriverVersion(),
|
||||
CUDAVersion: g.GetCudaVersion(),
|
||||
NVLinkPeers: intSlice(g.GetNvlinkPeers()),
|
||||
Health: discovery.GPUHealth{
|
||||
ECCErrors: int(g.GetEccErrors()),
|
||||
TemperatureC: int(g.GetTemperatureC()),
|
||||
Problems: g.GetProblems(),
|
||||
},
|
||||
})
|
||||
}
|
||||
for _, f := range r.GetFabrics() {
|
||||
n.Fabrics = append(n.Fabrics, inventory.FabricAttachment{
|
||||
Fabric: f.GetFabric(),
|
||||
HCAs: f.GetHcas(),
|
||||
RateGbps: int(f.GetRateGbps()),
|
||||
})
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// RegisterFromInventory builds the registration an agent sends.
|
||||
func RegisterFromInventory(nodeID, hostname, mgmtAddr, agentVersion string,
|
||||
labels map[string]string, inv discovery.Inventory, fabrics []inventory.FabricAttachment,
|
||||
) *llmgwv1.Register {
|
||||
r := &llmgwv1.Register{
|
||||
NodeId: nodeID, Hostname: hostname, MgmtAddr: mgmtAddr,
|
||||
Labels: labels, AgentVersion: agentVersion,
|
||||
}
|
||||
for _, g := range inv.GPUs {
|
||||
r.Gpus = append(r.Gpus, &llmgwv1.GPU{
|
||||
Index: int32(g.Index),
|
||||
Uuid: g.UUID,
|
||||
Name: g.Name,
|
||||
MemoryTotalMb: int32(g.MemoryTotalMB),
|
||||
MemoryUsedMb: int32(g.MemoryUsedMB),
|
||||
ComputeCapability: g.ComputeCapability,
|
||||
NumaNode: int32(g.NUMANode),
|
||||
PciBusId: g.PCIBusID,
|
||||
DriverVersion: g.DriverVersion,
|
||||
CudaVersion: g.CUDAVersion,
|
||||
NvlinkPeers: int32Slice(g.NVLinkPeers),
|
||||
Problems: g.Health.Problems,
|
||||
TemperatureC: int32(g.Health.TemperatureC),
|
||||
EccErrors: int32(g.Health.ECCErrors),
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
r.Fabrics = append(r.Fabrics, &llmgwv1.FabricAttachment{
|
||||
Fabric: f.Fabric, Hcas: f.HCAs, RateGbps: int32(f.RateGbps),
|
||||
})
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// AdoptedFromRegister converts an agent's running-shard report into the
|
||||
// reconciler's view.
|
||||
func AdoptedFromRegister(nodeID string, shards []*llmgwv1.RunningShard) []reconciler.AdoptedShard {
|
||||
out := make([]reconciler.AdoptedShard, 0, len(shards))
|
||||
for _, s := range shards {
|
||||
a := reconciler.AdoptedShard{
|
||||
NodeID: nodeID,
|
||||
ReplicaID: s.GetReplicaId(),
|
||||
ShardID: s.GetShardId(),
|
||||
Rank: int(s.GetRank()),
|
||||
GPUIndices: intSlice(s.GetGpuIndices()),
|
||||
HCAs: s.GetHcas(),
|
||||
HTTPPort: int(s.GetHttpPort()),
|
||||
State: StateFromWire(s.GetState()),
|
||||
Model: s.GetModel(),
|
||||
Profile: s.GetProfile(),
|
||||
Pool: s.GetPool(),
|
||||
Fabric: s.GetFabric(),
|
||||
}
|
||||
if a.Rank == 0 && s.GetLeaderAddr() != "" {
|
||||
a.Endpoint = fmt.Sprintf("http://%s:%d", s.GetLeaderAddr(), s.GetHttpPort())
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// StartShardFromCommand converts a reconciler command into a wire message.
|
||||
func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
|
||||
p := cmd.Spec.Profile
|
||||
return &llmgwv1.StartShard{
|
||||
ReplicaId: cmd.ReplicaID,
|
||||
ShardId: cmd.ShardID,
|
||||
Rank: int32(cmd.Spec.Rank),
|
||||
Nnodes: int32(cmd.Spec.NNodes),
|
||||
LeaderAddr: cmd.Spec.LeaderAddr,
|
||||
LeaderPort: int32(cmd.Spec.LeaderPort),
|
||||
GpuIndices: int32Slice(cmd.Spec.GPUIndices),
|
||||
Hcas: cmd.Spec.HCAs,
|
||||
MgmtInterface: cmd.Spec.MgmtInterface,
|
||||
HttpPort: int32(cmd.Spec.HTTPPort),
|
||||
Profile: &llmgwv1.EngineProfile{
|
||||
Name: p.Name, Engine: p.Engine, GpuClass: p.GPUClass,
|
||||
Tp: int32(p.TP), Pp: int32(p.PP),
|
||||
GpusPerNode: int32(p.GPUsPerNode), Nodes: int32(p.Nodes),
|
||||
Args: p.Args, Env: p.Env,
|
||||
ModelPath: p.ModelPath, ServedModelName: p.ServedModelName,
|
||||
},
|
||||
StartupTimeoutMs: cmd.StartupTimeout.Milliseconds(),
|
||||
Pool: cmd.Pool,
|
||||
Fabric: cmd.Fabric,
|
||||
}
|
||||
}
|
||||
|
||||
// ShardSpecFromStart converts a wire message into an engine shard spec.
|
||||
func ShardSpecFromStart(m *llmgwv1.StartShard) (engines.ShardSpec, time.Duration) {
|
||||
p := m.GetProfile()
|
||||
spec := engines.ShardSpec{
|
||||
ReplicaID: m.GetReplicaId(),
|
||||
ShardID: m.GetShardId(),
|
||||
Rank: int(m.GetRank()),
|
||||
NNodes: int(m.GetNnodes()),
|
||||
LeaderAddr: m.GetLeaderAddr(),
|
||||
LeaderPort: int(m.GetLeaderPort()),
|
||||
GPUIndices: intSlice(m.GetGpuIndices()),
|
||||
HCAs: m.GetHcas(),
|
||||
MgmtInterface: m.GetMgmtInterface(),
|
||||
HTTPPort: int(m.GetHttpPort()),
|
||||
Pool: m.GetPool(),
|
||||
Fabric: m.GetFabric(),
|
||||
Profile: engines.Profile{
|
||||
Name: p.GetName(), Engine: p.GetEngine(), GPUClass: p.GetGpuClass(),
|
||||
TP: int(p.GetTp()), PP: int(p.GetPp()),
|
||||
GPUsPerNode: int(p.GetGpusPerNode()), Nodes: int(p.GetNodes()),
|
||||
Args: p.GetArgs(), Env: p.GetEnv(),
|
||||
ModelPath: p.GetModelPath(), ServedModelName: p.GetServedModelName(),
|
||||
},
|
||||
}
|
||||
return spec, time.Duration(m.GetStartupTimeoutMs()) * time.Millisecond
|
||||
}
|
||||
|
||||
// stateWire maps supervisor states onto the wire enum.
|
||||
var stateWire = map[supervisor.State]llmgwv1.ShardStatus_State{
|
||||
supervisor.Pending: llmgwv1.ShardStatus_STATE_PENDING,
|
||||
supervisor.Starting: llmgwv1.ShardStatus_STATE_STARTING,
|
||||
supervisor.Loading: llmgwv1.ShardStatus_STATE_LOADING,
|
||||
supervisor.Ready: llmgwv1.ShardStatus_STATE_READY,
|
||||
supervisor.Draining: llmgwv1.ShardStatus_STATE_DRAINING,
|
||||
supervisor.Stopped: llmgwv1.ShardStatus_STATE_STOPPED,
|
||||
supervisor.Failed: llmgwv1.ShardStatus_STATE_FAILED,
|
||||
}
|
||||
|
||||
// StateToWire converts a supervisor state for transmission.
|
||||
func StateToWire(s supervisor.State) llmgwv1.ShardStatus_State {
|
||||
if w, ok := stateWire[s]; ok {
|
||||
return w
|
||||
}
|
||||
return llmgwv1.ShardStatus_STATE_UNSPECIFIED
|
||||
}
|
||||
|
||||
// StateFromWire converts a received state.
|
||||
//
|
||||
// An unrecognised state is treated as Failed rather than guessed at: a newer
|
||||
// agent reporting something this control plane does not understand must not be
|
||||
// mistaken for a healthy shard.
|
||||
func StateFromWire(w llmgwv1.ShardStatus_State) supervisor.State {
|
||||
for s, v := range stateWire {
|
||||
if v == w {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return supervisor.Failed
|
||||
}
|
||||
|
||||
func intSlice(v []int32) []int {
|
||||
if len(v) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]int, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = int(x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func int32Slice(v []int) []int32 {
|
||||
if len(v) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]int32, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = int32(x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,324 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"slices"
|
||||
"sync"
|
||||
|
||||
"connectrpc.com/connect"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// StatusSink receives shard transitions and recovers replica state from what
|
||||
// agents report they are already running. The reconciler implements it.
|
||||
type StatusSink interface {
|
||||
OnShardStatus(replicaID, shardID string, state supervisor.State, errMsg string)
|
||||
Adopt(shards []reconciler.AdoptedShard) (unclaimed []reconciler.AdoptedShard)
|
||||
}
|
||||
|
||||
// Service is the control plane's AgentService implementation, and at the same
|
||||
// time the reconciler's Dispatcher: it holds the open streams, so it is the
|
||||
// only thing that can deliver a command to a node.
|
||||
type Service struct {
|
||||
inv *inventory.Store
|
||||
status StatusSink
|
||||
log *slog.Logger
|
||||
|
||||
mu sync.RWMutex
|
||||
conns map[string]*conn
|
||||
}
|
||||
|
||||
// conn is one agent's open stream.
|
||||
type conn struct {
|
||||
nodeID string
|
||||
// send is serialized through a channel rather than a mutex around the
|
||||
// stream: connect's streams are not safe for concurrent writes, and the
|
||||
// reconciler dispatches to many nodes at once.
|
||||
send chan *llmgwv1.ControlMessage
|
||||
closed chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (c *conn) close() { c.once.Do(func() { close(c.closed) }) }
|
||||
|
||||
// New builds the service.
|
||||
func New(inv *inventory.Store, status StatusSink, log *slog.Logger) *Service {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Service{inv: inv, status: status, log: log, conns: map[string]*conn{}}
|
||||
}
|
||||
|
||||
// Handler returns the HTTP path and handler to mount.
|
||||
func (s *Service) Handler() (string, http.Handler) {
|
||||
return llmgwv1connect.NewAgentServiceHandler(s)
|
||||
}
|
||||
|
||||
// Connect serves one agent's stream for the lifetime of that agent.
|
||||
func (s *Service) Connect(
|
||||
ctx context.Context,
|
||||
stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage],
|
||||
) error {
|
||||
first, err := stream.Receive()
|
||||
if err != nil {
|
||||
return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("no opening message: %w", err))
|
||||
}
|
||||
reg := first.GetRegister()
|
||||
if reg == nil {
|
||||
// Registration first, always: the control plane has to know which node
|
||||
// it is talking to before it can route anything to it.
|
||||
return connect.NewError(connect.CodeInvalidArgument,
|
||||
errors.New("the first message must be a Register"))
|
||||
}
|
||||
nodeID := reg.GetNodeId()
|
||||
if nodeID == "" {
|
||||
return connect.NewError(connect.CodeInvalidArgument, errors.New("node_id is required"))
|
||||
}
|
||||
|
||||
node := NodeFromRegister(reg)
|
||||
s.inv.Upsert(node)
|
||||
s.log.Info("agent registered",
|
||||
"node", nodeID, "gpus", len(node.GPUs), "fabrics", node.FabricNames(),
|
||||
"running_shards", len(reg.GetRunningShards()),
|
||||
"agent_version", reg.GetAgentVersion())
|
||||
|
||||
c := s.attach(ctx, nodeID)
|
||||
defer s.detach(nodeID, c)
|
||||
|
||||
if err := stream.Send(&llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Ack{Ack: &llmgwv1.Ack{MessageId: "register", Ok: true}},
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Recover replica state from what this node is already running, before any
|
||||
// reconciliation pass can conclude the cluster is empty and place
|
||||
// duplicates (plan §6.5).
|
||||
s.adopt(ctx, nodeID, reg.GetRunningShards(), c)
|
||||
|
||||
// One goroutine writes, this one reads. Splitting them is what lets a
|
||||
// command be dispatched while the agent is mid-heartbeat.
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
s.writeLoop(ctx, stream, c)
|
||||
}()
|
||||
defer wg.Wait()
|
||||
|
||||
for {
|
||||
msg, err := stream.Receive()
|
||||
if err != nil {
|
||||
c.close()
|
||||
if ctx.Err() != nil || errors.Is(err, io.EOF) {
|
||||
s.log.Info("agent disconnected", "node", nodeID)
|
||||
return nil
|
||||
}
|
||||
s.log.Warn("agent stream failed", "node", nodeID, "err", err)
|
||||
return nil
|
||||
}
|
||||
s.handle(nodeID, msg)
|
||||
}
|
||||
}
|
||||
|
||||
// adopt hands the agent's running shards to the reconciler and stops whatever
|
||||
// nothing in the catalog claims.
|
||||
func (s *Service) adopt(ctx context.Context, nodeID string, shards []*llmgwv1.RunningShard, c *conn) {
|
||||
if len(shards) == 0 {
|
||||
return
|
||||
}
|
||||
unclaimed := s.status.Adopt(AdoptedFromRegister(nodeID, shards))
|
||||
for _, u := range unclaimed {
|
||||
// The model or profile left the catalog while this was running. Nobody
|
||||
// will ever route to it, and it is holding GPUs.
|
||||
s.log.Warn("stopping an unclaimed shard",
|
||||
"node", nodeID, "shard", u.ShardID, "model", u.Model, "profile", u.Profile)
|
||||
select {
|
||||
case c.send <- &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: u.ShardID}},
|
||||
}:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) writeLoop(ctx context.Context,
|
||||
stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage], c *conn,
|
||||
) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-c.closed:
|
||||
return
|
||||
case msg := <-c.send:
|
||||
if err := stream.Send(msg); err != nil {
|
||||
s.log.Warn("sending to agent failed", "node", c.nodeID, "err", err)
|
||||
c.close()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) {
|
||||
switch m := msg.Msg.(type) {
|
||||
case *llmgwv1.AgentMessage_Heartbeat:
|
||||
s.inv.Touch(nodeID)
|
||||
s.applyHeartbeat(nodeID, m.Heartbeat)
|
||||
|
||||
case *llmgwv1.AgentMessage_Shard:
|
||||
st := StateFromWire(m.Shard.GetState())
|
||||
s.status.OnShardStatus(m.Shard.GetReplicaId(), m.Shard.GetShardId(), st, m.Shard.GetError())
|
||||
if m.Shard.GetState() == llmgwv1.ShardStatus_STATE_FAILED {
|
||||
s.log.Error("shard failed",
|
||||
"node", nodeID, "shard", m.Shard.GetShardId(),
|
||||
"err", m.Shard.GetError(), "exit_code", m.Shard.GetExitCode(),
|
||||
"last_output", m.Shard.GetLastOutput())
|
||||
}
|
||||
|
||||
case *llmgwv1.AgentMessage_Register:
|
||||
// A re-registration on a live stream means the agent's view of its own
|
||||
// hardware changed — a GPU fell out, say. Take it.
|
||||
s.inv.Upsert(NodeFromRegister(m.Register))
|
||||
s.log.Info("agent re-registered", "node", nodeID)
|
||||
}
|
||||
}
|
||||
|
||||
// applyHeartbeat folds in the cheap facts a heartbeat carries, so a foreign
|
||||
// process appearing on a card after registration is noticed without waiting for
|
||||
// a reconnect.
|
||||
func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
|
||||
used := hb.GetGpuMemoryUsedMb()
|
||||
problems := hb.GetNewProblems()
|
||||
if len(used) == 0 && len(problems) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
n, ok := s.inv.Get(nodeID)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for i := range n.GPUs {
|
||||
if i < len(used) {
|
||||
n.GPUs[i].MemoryUsedMB = int(used[i])
|
||||
}
|
||||
}
|
||||
if len(problems) > 0 {
|
||||
// Problems are reported per node, not per GPU, so attach them to the
|
||||
// node's first GPU rather than guessing which one is affected. Any
|
||||
// problem is enough to keep that GPU out of placement, which is the
|
||||
// conservative direction.
|
||||
n.GPUs[0].Health.Problems = append(n.GPUs[0].Health.Problems, problems...)
|
||||
}
|
||||
s.inv.Upsert(n)
|
||||
}
|
||||
|
||||
func (s *Service) attach(ctx context.Context, nodeID string) *conn {
|
||||
c := &conn{
|
||||
nodeID: nodeID,
|
||||
// Buffered so a burst of commands does not block the reconciler on a
|
||||
// slow agent; if it fills, the agent is too far behind to be useful.
|
||||
send: make(chan *llmgwv1.ControlMessage, 64),
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
if prev, ok := s.conns[nodeID]; ok {
|
||||
// A node reconnecting before the old stream was reaped. The new one
|
||||
// wins: the old is by definition the dead one.
|
||||
s.log.Warn("agent reconnected while an old stream was open", "node", nodeID)
|
||||
prev.close()
|
||||
}
|
||||
s.conns[nodeID] = c
|
||||
s.mu.Unlock()
|
||||
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
c.close()
|
||||
}()
|
||||
return c
|
||||
}
|
||||
|
||||
func (s *Service) detach(nodeID string, c *conn) {
|
||||
c.close()
|
||||
s.mu.Lock()
|
||||
if s.conns[nodeID] == c {
|
||||
delete(s.conns, nodeID)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// The node is not marked lost here. Losing a stream is not losing a node,
|
||||
// and the reconciler's two-stage expiry already handles silence with the
|
||||
// right delays (plan §6.3).
|
||||
s.log.Info("agent stream closed", "node", nodeID)
|
||||
}
|
||||
|
||||
// ConnectedNodes lists the nodes with an open stream.
|
||||
func (s *Service) ConnectedNodes() []string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
out := make([]string, 0, len(s.conns))
|
||||
for id := range s.conns {
|
||||
out = append(out, id)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// --- Dispatcher ---
|
||||
|
||||
// ErrNotConnected means no agent stream is open for that node.
|
||||
var ErrNotConnected = errors.New("no agent connected")
|
||||
|
||||
// StartShard implements reconciler.Dispatcher.
|
||||
func (s *Service) StartShard(_ context.Context, cmd reconciler.ShardCommand) error {
|
||||
return s.send(cmd.NodeID, &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Start{Start: StartShardFromCommand(cmd)},
|
||||
})
|
||||
}
|
||||
|
||||
// StopShard implements reconciler.Dispatcher.
|
||||
func (s *Service) StopShard(_ context.Context, nodeID, shardID string) error {
|
||||
return s.send(nodeID, &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: shardID}},
|
||||
})
|
||||
}
|
||||
|
||||
// DrainShard implements reconciler.Dispatcher.
|
||||
func (s *Service) DrainShard(_ context.Context, nodeID, shardID string) error {
|
||||
return s.send(nodeID, &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Drain{Drain: &llmgwv1.DrainShard{ShardId: shardID}},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) send(nodeID string, msg *llmgwv1.ControlMessage) error {
|
||||
s.mu.RLock()
|
||||
c, ok := s.conns[nodeID]
|
||||
s.mu.RUnlock()
|
||||
if !ok {
|
||||
return fmt.Errorf("%w: node %s", ErrNotConnected, nodeID)
|
||||
}
|
||||
|
||||
select {
|
||||
case c.send <- msg:
|
||||
return nil
|
||||
case <-c.closed:
|
||||
return fmt.Errorf("%w: node %s stream closed", ErrNotConnected, nodeID)
|
||||
default:
|
||||
// The queue is full, so this agent is far enough behind that queueing
|
||||
// more would only delay discovering it is broken.
|
||||
return fmt.Errorf("node %s is not keeping up with commands", nodeID)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
// Package catalog holds the desired state: which models exist, how each one can
|
||||
// be run, and how many copies of it there should be (plan §6.1).
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/placement"
|
||||
)
|
||||
|
||||
// Scaling is how many replicas a model should have.
|
||||
type Scaling struct {
|
||||
Min int `koanf:"min"`
|
||||
Max int `koanf:"max"`
|
||||
}
|
||||
|
||||
// SLO is the latency target a model is expected to meet. `ctl bench` measures
|
||||
// capacity against it (plan §6.2.1).
|
||||
type SLO struct {
|
||||
TTFTp99 time.Duration `koanf:"ttft_p99"`
|
||||
TPOTp99 time.Duration `koanf:"tpot_p99"`
|
||||
}
|
||||
|
||||
// Profile is how a model runs on one GPU class.
|
||||
type Profile struct {
|
||||
Name string `koanf:"name"`
|
||||
Engine string `koanf:"engine"`
|
||||
GPUClass string `koanf:"gpu_class"`
|
||||
|
||||
TP int `koanf:"tp"`
|
||||
PP int `koanf:"pp"`
|
||||
// GPUsPerNode defaults to TP for a single-node profile.
|
||||
GPUsPerNode int `koanf:"gpus_per_node"`
|
||||
// Nodes above 1 makes this a multi-node profile, which requires a fabric.
|
||||
Nodes int `koanf:"nodes"`
|
||||
|
||||
Args []string `koanf:"args"`
|
||||
Env map[string]string `koanf:"env"`
|
||||
|
||||
// Weight is the router's capacity weight. It should come from ctl bench
|
||||
// rather than a guess (plan §6.2.1).
|
||||
Weight float64 `koanf:"weight"`
|
||||
// Priority orders profiles during placement; highest that fits wins.
|
||||
Priority int `koanf:"priority"`
|
||||
// StartupTimeout bounds model loading. A 70B model across two nodes takes
|
||||
// tens of minutes, so this is per profile rather than global.
|
||||
StartupTimeout time.Duration `koanf:"startup_timeout"`
|
||||
}
|
||||
|
||||
// Model is a catalog entry as clients see it.
|
||||
type Model struct {
|
||||
Name string `koanf:"name"`
|
||||
Aliases []string `koanf:"aliases"`
|
||||
// Source is the weights path or a HuggingFace id.
|
||||
Source string `koanf:"source"`
|
||||
Profiles []Profile `koanf:"profiles"`
|
||||
Scaling Scaling `koanf:"scaling"`
|
||||
Fallbacks []string `koanf:"fallbacks"`
|
||||
// RoutingProfile selects how the router scores endpoints for this model.
|
||||
RoutingProfile string `koanf:"routing_profile"`
|
||||
SLO SLO `koanf:"slo"`
|
||||
}
|
||||
|
||||
// Catalog is the set of models the cluster should serve.
|
||||
type Catalog struct {
|
||||
models map[string]Model
|
||||
aliases map[string]string
|
||||
order []string
|
||||
}
|
||||
|
||||
// New builds a catalog, rejecting definitions that would fail confusingly later.
|
||||
func New(models []Model) (*Catalog, error) {
|
||||
c := &Catalog{
|
||||
models: make(map[string]Model, len(models)),
|
||||
aliases: map[string]string{},
|
||||
}
|
||||
|
||||
for i := range models {
|
||||
m := models[i].normalize()
|
||||
if err := m.validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, dup := c.models[m.Name]; dup {
|
||||
return nil, fmt.Errorf("model %q is declared twice", m.Name)
|
||||
}
|
||||
c.models[m.Name] = m
|
||||
c.order = append(c.order, m.Name)
|
||||
}
|
||||
|
||||
// Aliases share a namespace with names. A collision would silently shadow a
|
||||
// model at request time, so it has to fail at load.
|
||||
for _, name := range c.order {
|
||||
for _, a := range c.models[name].Aliases {
|
||||
if _, clash := c.models[a]; clash {
|
||||
return nil, fmt.Errorf("model %q: alias %q collides with another model name", name, a)
|
||||
}
|
||||
if prev, clash := c.aliases[a]; clash {
|
||||
return nil, fmt.Errorf("alias %q is claimed by both %q and %q", a, prev, name)
|
||||
}
|
||||
c.aliases[a] = name
|
||||
}
|
||||
}
|
||||
|
||||
// Fallbacks are resolved now rather than at request time: a typo in a
|
||||
// fallback should stop a deployment, not surface during an outage when the
|
||||
// fallback is finally needed.
|
||||
for _, name := range c.order {
|
||||
for _, f := range c.models[name].Fallbacks {
|
||||
if _, ok := c.Resolve(f); !ok {
|
||||
return nil, fmt.Errorf("model %q: fallback %q is not in the catalog", name, f)
|
||||
}
|
||||
}
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (m Model) normalize() Model {
|
||||
if m.Scaling.Min < 0 {
|
||||
m.Scaling.Min = 0
|
||||
}
|
||||
if m.Scaling.Max == 0 {
|
||||
m.Scaling.Max = max(m.Scaling.Min, 1)
|
||||
}
|
||||
for i := range m.Profiles {
|
||||
p := &m.Profiles[i]
|
||||
if p.Engine == "" {
|
||||
p.Engine = "vllm"
|
||||
}
|
||||
if p.TP <= 0 {
|
||||
p.TP = 1
|
||||
}
|
||||
if p.Nodes <= 0 {
|
||||
p.Nodes = 1
|
||||
}
|
||||
if p.GPUsPerNode <= 0 {
|
||||
// The common case: a single-node profile uses TP GPUs on one node.
|
||||
p.GPUsPerNode = p.TP / p.Nodes
|
||||
if p.GPUsPerNode <= 0 {
|
||||
p.GPUsPerNode = 1
|
||||
}
|
||||
}
|
||||
if p.StartupTimeout <= 0 {
|
||||
p.StartupTimeout = 20 * time.Minute
|
||||
}
|
||||
if p.Weight <= 0 {
|
||||
p.Weight = 1
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func (m Model) validate() error {
|
||||
if m.Name == "" {
|
||||
return fmt.Errorf("a model has no name")
|
||||
}
|
||||
if m.Source == "" {
|
||||
return fmt.Errorf("model %q: source is required", m.Name)
|
||||
}
|
||||
if len(m.Profiles) == 0 {
|
||||
return fmt.Errorf("model %q: at least one profile is required", m.Name)
|
||||
}
|
||||
if m.Scaling.Max < m.Scaling.Min {
|
||||
return fmt.Errorf("model %q: scaling max %d is below min %d", m.Name, m.Scaling.Max, m.Scaling.Min)
|
||||
}
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
for _, p := range m.Profiles {
|
||||
if p.Name == "" {
|
||||
return fmt.Errorf("model %q: a profile has no name", m.Name)
|
||||
}
|
||||
if _, dup := seen[p.Name]; dup {
|
||||
return fmt.Errorf("model %q: profile %q is declared twice", m.Name, p.Name)
|
||||
}
|
||||
seen[p.Name] = struct{}{}
|
||||
|
||||
if p.GPUClass == "" {
|
||||
return fmt.Errorf("model %q profile %q: gpu_class is required", m.Name, p.Name)
|
||||
}
|
||||
// The arithmetic has to add up, or the engine is handed a shard spec it
|
||||
// cannot satisfy and fails minutes later during model loading.
|
||||
if want := p.GPUsPerNode * p.Nodes; want != p.TP*max(p.PP, 1) {
|
||||
return fmt.Errorf("model %q profile %q: %d nodes x %d GPUs does not match TP %d x PP %d",
|
||||
m.Name, p.Name, p.Nodes, p.GPUsPerNode, p.TP, max(p.PP, 1))
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Resolve maps a client-supplied name or alias onto a model.
|
||||
func (c *Catalog) Resolve(name string) (Model, bool) {
|
||||
if m, ok := c.models[name]; ok {
|
||||
return m, true
|
||||
}
|
||||
if canonical, ok := c.aliases[name]; ok {
|
||||
return c.models[canonical], true
|
||||
}
|
||||
return Model{}, false
|
||||
}
|
||||
|
||||
// Models returns every model, in declaration order.
|
||||
func (c *Catalog) Models() []Model {
|
||||
out := make([]Model, 0, len(c.order))
|
||||
for _, n := range c.order {
|
||||
out = append(out, c.models[n])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Names returns the canonical model names, sorted.
|
||||
func (c *Catalog) Names() []string {
|
||||
out := slices.Clone(c.order)
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// AliasesOf returns the alternative names for a model.
|
||||
func (c *Catalog) AliasesOf(name string) []string {
|
||||
m, ok := c.Resolve(name)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return slices.Clone(m.Aliases)
|
||||
}
|
||||
|
||||
// PlacementProfiles converts a model's profiles into placement's view.
|
||||
func (m Model) PlacementProfiles() []placement.Profile {
|
||||
out := make([]placement.Profile, 0, len(m.Profiles))
|
||||
for _, p := range m.Profiles {
|
||||
out = append(out, placement.Profile{
|
||||
Name: p.Name,
|
||||
Engine: p.Engine,
|
||||
GPUClass: p.GPUClass,
|
||||
GPUsPerNode: p.GPUsPerNode,
|
||||
Nodes: p.Nodes,
|
||||
Priority: p.Priority,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// EngineProfile converts a catalog profile into the engine adapter's view.
|
||||
func (m Model) EngineProfile(name string) (engines.Profile, bool) {
|
||||
for _, p := range m.Profiles {
|
||||
if p.Name != name {
|
||||
continue
|
||||
}
|
||||
return engines.Profile{
|
||||
Name: p.Name,
|
||||
Engine: p.Engine,
|
||||
GPUClass: p.GPUClass,
|
||||
TP: p.TP,
|
||||
PP: p.PP,
|
||||
GPUsPerNode: p.GPUsPerNode,
|
||||
Nodes: p.Nodes,
|
||||
Args: slices.Clone(p.Args),
|
||||
Env: p.Env,
|
||||
ModelPath: m.Source,
|
||||
ServedModelName: m.Name,
|
||||
}, true
|
||||
}
|
||||
return engines.Profile{}, false
|
||||
}
|
||||
|
||||
// ProfileByName returns one of a model's profiles.
|
||||
func (m Model) ProfileByName(name string) (Profile, bool) {
|
||||
for _, p := range m.Profiles {
|
||||
if p.Name == name {
|
||||
return p, true
|
||||
}
|
||||
}
|
||||
return Profile{}, false
|
||||
}
|
||||
|
||||
// String renders the catalog for `ctl models`.
|
||||
func (c *Catalog) String() string {
|
||||
var b strings.Builder
|
||||
for _, m := range c.Models() {
|
||||
fmt.Fprintf(&b, "%s (min %d, max %d)\n", m.Name, m.Scaling.Min, m.Scaling.Max)
|
||||
for _, p := range m.Profiles {
|
||||
fmt.Fprintf(&b, " %-20s %-10s %-12s tp=%d nodes=%d priority=%d\n",
|
||||
p.Name, p.Engine, p.GPUClass, p.TP, p.Nodes, p.Priority)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,235 @@
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func model(name string, profiles ...Profile) Model {
|
||||
if len(profiles) == 0 {
|
||||
profiles = []Profile{{Name: "p1", GPUClass: "h100-80g", TP: 4}}
|
||||
}
|
||||
return Model{Name: name, Source: "/models/" + name, Profiles: profiles}
|
||||
}
|
||||
|
||||
func TestResolveNameAndAlias(t *testing.T) {
|
||||
m := model("llama-3.3-70b")
|
||||
m.Aliases = []string{"llama", "llama3"}
|
||||
c, err := New([]Model{m})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for _, name := range []string{"llama-3.3-70b", "llama", "llama3"} {
|
||||
got, ok := c.Resolve(name)
|
||||
if !ok {
|
||||
t.Fatalf("%q did not resolve", name)
|
||||
}
|
||||
if got.Name != "llama-3.3-70b" {
|
||||
t.Errorf("%q resolved to %q", name, got.Name)
|
||||
}
|
||||
}
|
||||
if _, ok := c.Resolve("nope"); ok {
|
||||
t.Error("an unknown name resolved")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAliasCollisionRejected(t *testing.T) {
|
||||
// An alias shadowing another model would misroute every request for it,
|
||||
// silently, at request time.
|
||||
a := model("qwen")
|
||||
b := model("llama")
|
||||
b.Aliases = []string{"qwen"}
|
||||
|
||||
_, err := New([]Model{a, b})
|
||||
if err == nil || !strings.Contains(err.Error(), "collides") {
|
||||
t.Errorf("collision not rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuplicateAliasRejected(t *testing.T) {
|
||||
a := model("a")
|
||||
a.Aliases = []string{"shared"}
|
||||
b := model("b")
|
||||
b.Aliases = []string{"shared"}
|
||||
|
||||
_, err := New([]Model{a, b})
|
||||
if err == nil || !strings.Contains(err.Error(), "claimed by both") {
|
||||
t.Errorf("duplicate alias not rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnknownFallbackRejected(t *testing.T) {
|
||||
// A typo in a fallback should stop a deployment, not surface during the
|
||||
// outage when the fallback is finally needed.
|
||||
m := model("llama")
|
||||
m.Fallbacks = []string{"qwen-typo"}
|
||||
|
||||
_, err := New([]Model{m})
|
||||
if err == nil || !strings.Contains(err.Error(), "fallback") {
|
||||
t.Errorf("unknown fallback not rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFallbackToAliasAccepted(t *testing.T) {
|
||||
a := model("qwen2.5-32b")
|
||||
a.Aliases = []string{"qwen"}
|
||||
b := model("llama")
|
||||
b.Fallbacks = []string{"qwen"}
|
||||
|
||||
if _, err := New([]Model{a, b}); err != nil {
|
||||
t.Errorf("a fallback naming an alias was rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelismArithmeticValidated(t *testing.T) {
|
||||
// Wrong arithmetic hands the engine a spec it cannot satisfy, and it fails
|
||||
// minutes later during model loading rather than at config load.
|
||||
_, err := New([]Model{model("m", Profile{
|
||||
Name: "bad", GPUClass: "h100-80g", TP: 4, Nodes: 2, GPUsPerNode: 8,
|
||||
})})
|
||||
if err == nil || !strings.Contains(err.Error(), "does not match") {
|
||||
t.Errorf("bad arithmetic accepted: %v", err)
|
||||
}
|
||||
|
||||
// 2 nodes x 8 GPUs = TP 16 is consistent.
|
||||
if _, err := New([]Model{model("m", Profile{
|
||||
Name: "ok", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
})}); err != nil {
|
||||
t.Errorf("valid arithmetic rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaults(t *testing.T) {
|
||||
c, err := New([]Model{model("m", Profile{Name: "p", GPUClass: "h100-80g", TP: 4})})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m, _ := c.Resolve("m")
|
||||
p := m.Profiles[0]
|
||||
|
||||
if p.Engine != "vllm" {
|
||||
t.Errorf("engine=%q", p.Engine)
|
||||
}
|
||||
if p.Nodes != 1 {
|
||||
t.Errorf("nodes=%d", p.Nodes)
|
||||
}
|
||||
if p.GPUsPerNode != 4 {
|
||||
t.Errorf("gpus_per_node=%d; a single-node profile should default to TP", p.GPUsPerNode)
|
||||
}
|
||||
if p.StartupTimeout != 20*time.Minute {
|
||||
t.Errorf("startup_timeout=%v", p.StartupTimeout)
|
||||
}
|
||||
if p.Weight <= 0 {
|
||||
t.Errorf("weight=%v", p.Weight)
|
||||
}
|
||||
if m.Scaling.Max < 1 {
|
||||
t.Errorf("scaling max=%d", m.Scaling.Max)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequiredFields(t *testing.T) {
|
||||
cases := map[string]Model{
|
||||
"no name": {Source: "/m", Profiles: []Profile{{Name: "p", GPUClass: "c", TP: 1}}},
|
||||
"no source": {Name: "m", Profiles: []Profile{{Name: "p", GPUClass: "c", TP: 1}}},
|
||||
"no profiles": {Name: "m", Source: "/m"},
|
||||
"no class": {Name: "m", Source: "/m", Profiles: []Profile{{Name: "p", TP: 1}}},
|
||||
}
|
||||
for name, m := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
if _, err := New([]Model{m}); err == nil {
|
||||
t.Error("accepted")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestScalingBounds(t *testing.T) {
|
||||
m := model("m")
|
||||
m.Scaling = Scaling{Min: 5, Max: 2}
|
||||
if _, err := New([]Model{m}); err == nil {
|
||||
t.Error("max below min was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuplicateProfileRejected(t *testing.T) {
|
||||
_, err := New([]Model{model("m",
|
||||
Profile{Name: "p", GPUClass: "h100-80g", TP: 1},
|
||||
Profile{Name: "p", GPUClass: "a100-40g", TP: 1},
|
||||
)})
|
||||
if err == nil || !strings.Contains(err.Error(), "twice") {
|
||||
t.Errorf("duplicate profile accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlacementProfilesCarryPriority(t *testing.T) {
|
||||
c, _ := New([]Model{model("m",
|
||||
Profile{Name: "fast", GPUClass: "h100-80g", TP: 4, Priority: 100},
|
||||
Profile{Name: "slow", GPUClass: "a100-40g", TP: 4, Priority: 50},
|
||||
)})
|
||||
m, _ := c.Resolve("m")
|
||||
|
||||
pp := m.PlacementProfiles()
|
||||
if len(pp) != 2 {
|
||||
t.Fatalf("got %d placement profiles", len(pp))
|
||||
}
|
||||
byName := map[string]int{}
|
||||
for _, p := range pp {
|
||||
byName[p.Name] = p.Priority
|
||||
}
|
||||
if byName["fast"] != 100 || byName["slow"] != 50 {
|
||||
t.Errorf("priorities lost in conversion: %v", byName)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineProfileCarriesModelIdentity(t *testing.T) {
|
||||
// The served name is what clients ask for and need not match the path.
|
||||
c, _ := New([]Model{model("llama-3.3-70b")})
|
||||
m, _ := c.Resolve("llama-3.3-70b")
|
||||
|
||||
ep, ok := m.EngineProfile("p1")
|
||||
if !ok {
|
||||
t.Fatal("engine profile not found")
|
||||
}
|
||||
if ep.ModelPath != "/models/llama-3.3-70b" {
|
||||
t.Errorf("ModelPath=%q", ep.ModelPath)
|
||||
}
|
||||
if ep.ServedModelName != "llama-3.3-70b" {
|
||||
t.Errorf("ServedModelName=%q", ep.ServedModelName)
|
||||
}
|
||||
|
||||
if _, ok := m.EngineProfile("nope"); ok {
|
||||
t.Error("an unknown profile resolved")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineProfileArgsAreCopied(t *testing.T) {
|
||||
// The reconciler hands these to an engine adapter; mutating them must not
|
||||
// corrupt the catalog for every later replica.
|
||||
c, _ := New([]Model{model("m", Profile{
|
||||
Name: "p", GPUClass: "h100-80g", TP: 1, Args: []string{"--max-model-len", "32768"},
|
||||
})})
|
||||
m, _ := c.Resolve("m")
|
||||
|
||||
ep, _ := m.EngineProfile("p")
|
||||
ep.Args[0] = "tampered"
|
||||
|
||||
again, _ := m.EngineProfile("p")
|
||||
if again.Args[0] != "--max-model-len" {
|
||||
t.Error("EngineProfile exposed the catalog's own slice")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNamesAndModelsOrder(t *testing.T) {
|
||||
c, _ := New([]Model{model("zeta"), model("alpha")})
|
||||
|
||||
// Models keeps declaration order, so configuration reads as written.
|
||||
if got := c.Models(); got[0].Name != "zeta" {
|
||||
t.Errorf("Models order=%q", got[0].Name)
|
||||
}
|
||||
// Names is sorted, for stable CLI output.
|
||||
if got := c.Names(); got[0] != "alpha" {
|
||||
t.Errorf("Names order=%v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
// Package control wires the control plane together: inventory, catalog,
|
||||
// placement and reconciliation, and publishes the result to the router.
|
||||
package control
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
)
|
||||
|
||||
// EndpointSink receives the endpoint list. The router implements it.
|
||||
//
|
||||
// It is an interface so the control plane does not depend on the router: when
|
||||
// the two split into separate processes, this becomes a watch stream and
|
||||
// nothing above it changes (plan §3.1).
|
||||
type EndpointSink interface {
|
||||
SetEndpoints(ctx context.Context, infos []endpoint.Info)
|
||||
}
|
||||
|
||||
// Config tunes the control plane.
|
||||
type Config struct {
|
||||
Reconciler reconciler.Config
|
||||
// PublishInterval is how often the endpoint list is republished even when
|
||||
// nothing appears to have changed. Events drive the common case; this is
|
||||
// the backstop.
|
||||
PublishInterval time.Duration
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Reconciler: reconciler.DefaultConfig(),
|
||||
PublishInterval: time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// Controller owns the control plane's state.
|
||||
type Controller struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
inv *inventory.Store
|
||||
cat *catalog.Catalog
|
||||
rec *reconciler.Reconciler
|
||||
sink EndpointSink
|
||||
|
||||
mu sync.Mutex
|
||||
last string // fingerprint of the last published list
|
||||
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// New builds a controller.
|
||||
func New(cfg Config, cat *catalog.Catalog, disp reconciler.Dispatcher, sink EndpointSink, log *slog.Logger) *Controller {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if cfg.PublishInterval <= 0 {
|
||||
cfg.PublishInterval = time.Second
|
||||
}
|
||||
inv := inventory.NewStore()
|
||||
return &Controller{
|
||||
cfg: cfg,
|
||||
log: log,
|
||||
inv: inv,
|
||||
cat: cat,
|
||||
rec: reconciler.New(cfg.Reconciler, cat, inv, disp, log),
|
||||
sink: sink,
|
||||
stop: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Inventory exposes the node inventory, which the agent service feeds.
|
||||
func (c *Controller) Inventory() *inventory.Store { return c.inv }
|
||||
|
||||
// Reconciler exposes the reconciliation loop, for status reports and ctl.
|
||||
func (c *Controller) Reconciler() *reconciler.Reconciler { return c.rec }
|
||||
|
||||
// Catalog exposes the model catalog.
|
||||
func (c *Controller) Catalog() *catalog.Catalog { return c.cat }
|
||||
|
||||
// Start runs the reconciler and begins publishing endpoints.
|
||||
func (c *Controller) Start(ctx context.Context) {
|
||||
c.rec.Start(ctx)
|
||||
|
||||
go func() {
|
||||
defer close(c.done)
|
||||
t := time.NewTicker(c.cfg.PublishInterval)
|
||||
defer t.Stop()
|
||||
|
||||
c.Publish(ctx)
|
||||
for {
|
||||
select {
|
||||
case <-c.stop:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
c.Publish(ctx)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Close stops the control plane.
|
||||
func (c *Controller) Close() {
|
||||
c.once.Do(func() {
|
||||
close(c.stop)
|
||||
<-c.done
|
||||
c.rec.Close()
|
||||
})
|
||||
}
|
||||
|
||||
// Publish pushes the current Ready replicas to the sink.
|
||||
//
|
||||
// It is exported so tests can drive it deterministically instead of racing a
|
||||
// ticker, and it skips the push when nothing changed: replacing the list is
|
||||
// cheap but resynchronizes the metric pollers, and doing that every second for
|
||||
// no reason adds churn to every engine's /metrics.
|
||||
func (c *Controller) Publish(ctx context.Context) {
|
||||
infos := c.endpoints()
|
||||
|
||||
fp := fingerprint(infos)
|
||||
c.mu.Lock()
|
||||
unchanged := fp == c.last
|
||||
c.last = fp
|
||||
c.mu.Unlock()
|
||||
if unchanged {
|
||||
return
|
||||
}
|
||||
|
||||
c.log.Info("publishing endpoints", "count", len(infos))
|
||||
c.sink.SetEndpoints(ctx, infos)
|
||||
}
|
||||
|
||||
// endpoints converts Ready replicas into the router's view.
|
||||
func (c *Controller) endpoints() []endpoint.Info {
|
||||
ready := c.rec.ReadyReplicas()
|
||||
out := make([]endpoint.Info, 0, len(ready))
|
||||
|
||||
for _, rep := range ready {
|
||||
m, ok := c.cat.Resolve(rep.Model)
|
||||
if !ok {
|
||||
// The model left the catalog while a replica was still running.
|
||||
// The reconciler will tear it down; until then, do not route to it.
|
||||
continue
|
||||
}
|
||||
// Clients may address a model by any of its names, and the router
|
||||
// matches endpoints by the resolved name plus aliases.
|
||||
names := append([]string{m.Name}, m.Aliases...)
|
||||
|
||||
prof, _ := m.ProfileByName(rep.Profile)
|
||||
out = append(out, endpoint.Info{
|
||||
ID: rep.ID,
|
||||
URL: strings.TrimRight(rep.Endpoint, "/"),
|
||||
Models: names,
|
||||
Pool: rep.Pool,
|
||||
Weight: rep.Weight,
|
||||
Priority: prof.Priority,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// fingerprint summarizes a list so an unchanged one can be skipped.
|
||||
func fingerprint(infos []endpoint.Info) string {
|
||||
var b strings.Builder
|
||||
for _, i := range infos {
|
||||
fmt.Fprintf(&b, "%s|%s|%s|%v|%g|%d;", i.ID, i.URL, i.Pool, i.Models, i.Weight, i.Priority)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
// Package inventory holds the control plane's view of the cluster: which nodes
|
||||
// exist, what hardware they have, and what is already running on them.
|
||||
package inventory
|
||||
|
||||
import (
|
||||
"maps"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
)
|
||||
|
||||
// NodeState is where a node stands.
|
||||
type NodeState int
|
||||
|
||||
const (
|
||||
// Ready means the node accepts placement.
|
||||
Ready NodeState = iota
|
||||
// Cordoned means an operator has excluded it from new placement; what runs
|
||||
// there keeps running.
|
||||
Cordoned
|
||||
// Draining means its replicas are being moved off.
|
||||
Draining
|
||||
// Unreachable means heartbeats stopped but the node may come back.
|
||||
Unreachable
|
||||
// Lost means it has been unreachable long enough to declare its replicas dead.
|
||||
Lost
|
||||
)
|
||||
|
||||
func (s NodeState) String() string {
|
||||
switch s {
|
||||
case Ready:
|
||||
return "ready"
|
||||
case Cordoned:
|
||||
return "cordoned"
|
||||
case Draining:
|
||||
return "draining"
|
||||
case Unreachable:
|
||||
return "unreachable"
|
||||
case Lost:
|
||||
return "lost"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// Schedulable reports whether new replicas may be placed on the node.
|
||||
func (s NodeState) Schedulable() bool { return s == Ready }
|
||||
|
||||
// FabricAttachment is a node's connection to one fabric.
|
||||
type FabricAttachment struct {
|
||||
// Fabric is the operator-assigned name. It is not derived from the GID
|
||||
// prefix: OpenSM's default is fe80:: everywhere, so a matching prefix
|
||||
// proves nothing (plan §5.2).
|
||||
Fabric string `json:"fabric"`
|
||||
// HCAs are the adapters attached to this fabric.
|
||||
HCAs []string `json:"hcas"`
|
||||
// RateGbps is the slowest active port's rate, since a collective runs at
|
||||
// the speed of its slowest link.
|
||||
RateGbps int `json:"rate_gbps"`
|
||||
}
|
||||
|
||||
// Node is one machine in the cluster.
|
||||
type Node struct {
|
||||
ID string `json:"id"`
|
||||
Hostname string `json:"hostname"`
|
||||
MgmtAddr string `json:"mgmt_addr"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
|
||||
GPUs []discovery.GPU `json:"gpus"`
|
||||
// Fabrics is empty on a node without IB. Such a node is its own
|
||||
// single-node island and can never host a multi-node replica.
|
||||
Fabrics []FabricAttachment `json:"fabrics,omitempty"`
|
||||
|
||||
State NodeState `json:"-"`
|
||||
LastSeen time.Time `json:"last_seen"`
|
||||
// AgentVersion and EngineVersions are reported at registration. The engine
|
||||
// version decides which metric names and CLI flags apply (plan §14).
|
||||
AgentVersion string `json:"agent_version,omitempty"`
|
||||
EngineVersions map[string]string `json:"engine_versions,omitempty"`
|
||||
}
|
||||
|
||||
// HasFabric reports whether the node is attached to a fabric.
|
||||
func (n Node) HasFabric(name string) bool {
|
||||
return slices.ContainsFunc(n.Fabrics, func(f FabricAttachment) bool { return f.Fabric == name })
|
||||
}
|
||||
|
||||
// FabricNames lists the node's fabrics, sorted.
|
||||
func (n Node) FabricNames() []string {
|
||||
out := make([]string, 0, len(n.Fabrics))
|
||||
for _, f := range n.Fabrics {
|
||||
out = append(out, f.Fabric)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// HCAsFor returns the adapters attached to one fabric.
|
||||
//
|
||||
// A shard is given only these: handing it adapters from a second fabric lets
|
||||
// NCCL choose one whose peer is unreachable, and the job hangs instead of
|
||||
// failing (plan §5.3).
|
||||
func (n Node) HCAsFor(fabric string) []string {
|
||||
for _, f := range n.Fabrics {
|
||||
if f.Fabric == fabric {
|
||||
return slices.Clone(f.HCAs)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GPUsByClass groups the node's GPUs by class.
|
||||
func (n Node) GPUsByClass() map[string][]discovery.GPU {
|
||||
out := map[string][]discovery.GPU{}
|
||||
for _, g := range n.GPUs {
|
||||
out[g.Class()] = append(out[g.Class()], g)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Allocation records that some of a node's GPUs are in use.
|
||||
type Allocation struct {
|
||||
ShardID string
|
||||
ReplicaID string
|
||||
GPUIndices []int
|
||||
}
|
||||
|
||||
// Store holds the cluster inventory. It is safe for concurrent use.
|
||||
type Store struct {
|
||||
mu sync.RWMutex
|
||||
nodes map[string]*Node
|
||||
// allocs maps node id to the shards placed there.
|
||||
allocs map[string][]Allocation
|
||||
}
|
||||
|
||||
// NewStore returns an empty inventory.
|
||||
func NewStore() *Store {
|
||||
return &Store{nodes: map[string]*Node{}, allocs: map[string][]Allocation{}}
|
||||
}
|
||||
|
||||
// Upsert records a node's registration, preserving operator-set state such as
|
||||
// a cordon: an agent reconnecting must not silently undo an operator's decision.
|
||||
func (s *Store) Upsert(n Node) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if prev, ok := s.nodes[n.ID]; ok {
|
||||
if prev.State == Cordoned || prev.State == Draining {
|
||||
n.State = prev.State
|
||||
}
|
||||
}
|
||||
n.LastSeen = time.Now()
|
||||
s.nodes[n.ID] = &n
|
||||
}
|
||||
|
||||
// SetState changes a node's state.
|
||||
func (s *Store) SetState(id string, st NodeState) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
n, ok := s.nodes[id]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
n.State = st
|
||||
return true
|
||||
}
|
||||
|
||||
// Touch records a heartbeat, bringing an unreachable node back.
|
||||
func (s *Store) Touch(id string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
n, ok := s.nodes[id]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
n.LastSeen = time.Now()
|
||||
if n.State == Unreachable || n.State == Lost {
|
||||
n.State = Ready
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns a copy of a node.
|
||||
func (s *Store) Get(id string) (Node, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
n, ok := s.nodes[id]
|
||||
if !ok {
|
||||
return Node{}, false
|
||||
}
|
||||
return *n, true
|
||||
}
|
||||
|
||||
// Nodes returns every node, sorted by id.
|
||||
func (s *Store) Nodes() []Node {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
out := make([]Node, 0, len(s.nodes))
|
||||
for _, n := range s.nodes {
|
||||
out = append(out, *n)
|
||||
}
|
||||
slices.SortFunc(out, func(a, b Node) int { return strings.Compare(a.ID, b.ID) })
|
||||
return out
|
||||
}
|
||||
|
||||
// Remove forgets a node and its allocations.
|
||||
func (s *Store) Remove(id string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
delete(s.nodes, id)
|
||||
delete(s.allocs, id)
|
||||
}
|
||||
|
||||
// Allocate records GPUs as taken on a node.
|
||||
func (s *Store) Allocate(nodeID string, a Allocation) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.allocs[nodeID] = append(s.allocs[nodeID], a)
|
||||
}
|
||||
|
||||
// Release frees the GPUs a shard held.
|
||||
func (s *Store) Release(shardID string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for node, list := range s.allocs {
|
||||
kept := slices.DeleteFunc(slices.Clone(list), func(a Allocation) bool {
|
||||
return a.ShardID == shardID
|
||||
})
|
||||
if len(kept) == 0 {
|
||||
delete(s.allocs, node)
|
||||
} else {
|
||||
s.allocs[node] = kept
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ReleaseReplica frees every shard of a replica.
|
||||
//
|
||||
// Replicas are freed whole because they fail whole: one dead shard makes the
|
||||
// others useless, and the control plane rebuilds the replica rather than
|
||||
// restarting a piece of it (plan §5.4).
|
||||
func (s *Store) ReleaseReplica(replicaID string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for node, list := range s.allocs {
|
||||
kept := slices.DeleteFunc(slices.Clone(list), func(a Allocation) bool {
|
||||
return a.ReplicaID == replicaID
|
||||
})
|
||||
if len(kept) == 0 {
|
||||
delete(s.allocs, node)
|
||||
} else {
|
||||
s.allocs[node] = kept
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Allocations returns a node's current allocations.
|
||||
func (s *Store) Allocations(nodeID string) []Allocation {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return slices.Clone(s.allocs[nodeID])
|
||||
}
|
||||
|
||||
// FreeGPUs returns the indices of a node's unallocated GPUs, sorted.
|
||||
func (s *Store) FreeGPUs(nodeID string) []int {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
n, ok := s.nodes[nodeID]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
taken := map[int]struct{}{}
|
||||
for _, a := range s.allocs[nodeID] {
|
||||
for _, i := range a.GPUIndices {
|
||||
taken[i] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
var out []int
|
||||
for _, g := range n.GPUs {
|
||||
if _, used := taken[g.Index]; used {
|
||||
continue
|
||||
}
|
||||
// A GPU with an ECC fault or another process on it is not free in any
|
||||
// useful sense.
|
||||
if !g.Health.OK() {
|
||||
continue
|
||||
}
|
||||
out = append(out, g.Index)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// Fabrics lists every fabric any node is attached to, sorted.
|
||||
func (s *Store) Fabrics() []string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
for _, n := range s.nodes {
|
||||
for _, f := range n.Fabrics {
|
||||
seen[f.Fabric] = struct{}{}
|
||||
}
|
||||
}
|
||||
return slices.Sorted(maps.Keys(seen))
|
||||
}
|
||||
@@ -0,0 +1,250 @@
|
||||
package inventory
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
)
|
||||
|
||||
func gpus(n int) []discovery.GPU {
|
||||
out := make([]discovery.GPU, n)
|
||||
for i := range out {
|
||||
out[i] = discovery.GPU{Index: i, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestUpsertPreservesCordon(t *testing.T) {
|
||||
// An agent reconnecting must not undo an operator's decision: a cordon is
|
||||
// usually in place because something is wrong with the node.
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
s.SetState("n1", Cordoned)
|
||||
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)}) // re-registration
|
||||
n, _ := s.Get("n1")
|
||||
if n.State != Cordoned {
|
||||
t.Errorf("state=%s; re-registration cleared the cordon", n.State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertUpdatesHardware(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(4)})
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
|
||||
n, _ := s.Get("n1")
|
||||
if len(n.GPUs) != 8 {
|
||||
t.Errorf("hardware was not updated: %d GPUs", len(n.GPUs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTouchRevivesUnreachableNode(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1"})
|
||||
s.SetState("n1", Unreachable)
|
||||
|
||||
s.Touch("n1")
|
||||
if n, _ := s.Get("n1"); n.State != Ready {
|
||||
t.Errorf("a heartbeat did not revive the node: %s", n.State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTouchDoesNotClearCordon(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1"})
|
||||
s.SetState("n1", Cordoned)
|
||||
|
||||
s.Touch("n1")
|
||||
if n, _ := s.Get("n1"); n.State != Cordoned {
|
||||
t.Errorf("a heartbeat cleared the cordon: %s", n.State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFreeGPUs(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
|
||||
if got := s.FreeGPUs("n1"); len(got) != 8 {
|
||||
t.Fatalf("expected 8 free, got %d", len(got))
|
||||
}
|
||||
|
||||
s.Allocate("n1", Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1}})
|
||||
got := s.FreeGPUs("n1")
|
||||
if !slices.Equal(got, []int{2, 3, 4, 5, 6, 7}) {
|
||||
t.Errorf("free=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnhealthyGPUsAreNotFree(t *testing.T) {
|
||||
s := NewStore()
|
||||
n := Node{ID: "n1", GPUs: gpus(4)}
|
||||
n.GPUs[2].Health.Problems = []string{"ECC"}
|
||||
s.Upsert(n)
|
||||
|
||||
got := s.FreeGPUs("n1")
|
||||
if slices.Contains(got, 2) {
|
||||
t.Errorf("an unhealthy GPU was reported free: %v", got)
|
||||
}
|
||||
if len(got) != 3 {
|
||||
t.Errorf("free=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseShard(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
s.Allocate("n1", Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1}})
|
||||
s.Allocate("n1", Allocation{ShardID: "s2", ReplicaID: "r2", GPUIndices: []int{2, 3}})
|
||||
|
||||
s.Release("s1")
|
||||
if got := s.FreeGPUs("n1"); slices.Contains(got, 2) || !slices.Contains(got, 0) {
|
||||
t.Errorf("wrong shard released: free=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseReplicaFreesEveryShard(t *testing.T) {
|
||||
// Replicas fail whole, so they are freed whole: one dead shard makes the
|
||||
// others useless and the control plane rebuilds the replica (plan §5.4).
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
s.Upsert(Node{ID: "n2", GPUs: gpus(8)})
|
||||
s.Allocate("n1", Allocation{ShardID: "r1-0", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
|
||||
s.Allocate("n2", Allocation{ShardID: "r1-1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
|
||||
|
||||
s.ReleaseReplica("r1")
|
||||
for _, id := range []string{"n1", "n2"} {
|
||||
if got := len(s.FreeGPUs(id)); got != 8 {
|
||||
t.Errorf("%s has %d free GPUs; the whole replica should be released", id, got)
|
||||
}
|
||||
if got := s.Allocations(id); len(got) != 0 {
|
||||
t.Errorf("%s still has allocations: %v", id, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHCAsForSelectsOneFabric(t *testing.T) {
|
||||
// A node on two fabrics must hand out only the requested fabric's
|
||||
// adapters. Mixing them lets NCCL pick one whose peer is unreachable.
|
||||
n := Node{ID: "n1", Fabrics: []FabricAttachment{
|
||||
{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}},
|
||||
{Fabric: "ib-b", HCAs: []string{"mlx5_2", "mlx5_3"}},
|
||||
}}
|
||||
|
||||
if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
|
||||
t.Errorf("ib-a HCAs=%v", got)
|
||||
}
|
||||
if got := n.HCAsFor("ib-b"); !slices.Equal(got, []string{"mlx5_2", "mlx5_3"}) {
|
||||
t.Errorf("ib-b HCAs=%v", got)
|
||||
}
|
||||
if got := n.HCAsFor("ib-c"); got != nil {
|
||||
t.Errorf("an unattached fabric returned %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHCAsForReturnsACopy(t *testing.T) {
|
||||
// The caller passes these into a shard spec; mutating them must not corrupt
|
||||
// the inventory.
|
||||
n := Node{ID: "n1", Fabrics: []FabricAttachment{{Fabric: "ib-a", HCAs: []string{"mlx5_0"}}}}
|
||||
got := n.HCAsFor("ib-a")
|
||||
got[0] = "tampered"
|
||||
if n.Fabrics[0].HCAs[0] != "mlx5_0" {
|
||||
t.Error("HCAsFor exposed the inventory's own slice")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFabrics(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", Fabrics: []FabricAttachment{{Fabric: "ib-b"}}})
|
||||
s.Upsert(Node{ID: "n2", Fabrics: []FabricAttachment{{Fabric: "ib-a"}, {Fabric: "ib-b"}}})
|
||||
s.Upsert(Node{ID: "n3"}) // no fabric
|
||||
|
||||
if got := s.Fabrics(); !slices.Equal(got, []string{"ib-a", "ib-b"}) {
|
||||
t.Errorf("fabrics=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodesSorted(t *testing.T) {
|
||||
s := NewStore()
|
||||
for _, id := range []string{"n3", "n1", "n2"} {
|
||||
s.Upsert(Node{ID: id})
|
||||
}
|
||||
got := s.Nodes()
|
||||
if len(got) != 3 || got[0].ID != "n1" || got[2].ID != "n3" {
|
||||
t.Errorf("not sorted: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetReturnsACopy(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(2), Labels: map[string]string{"rack": "r1"}})
|
||||
|
||||
n, _ := s.Get("n1")
|
||||
n.GPUs[0].Index = 99
|
||||
|
||||
again, _ := s.Get("n1")
|
||||
if again.GPUs[0].Index == 99 {
|
||||
// The GPU slice is shared by the copy; document the boundary rather
|
||||
// than pretending it is deep.
|
||||
t.Log("note: Get returns a shallow copy; slices are shared")
|
||||
}
|
||||
if again.ID != "n1" {
|
||||
t.Error("the node was corrupted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemove(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(4)})
|
||||
s.Allocate("n1", Allocation{ShardID: "s1", GPUIndices: []int{0}})
|
||||
|
||||
s.Remove("n1")
|
||||
if _, ok := s.Get("n1"); ok {
|
||||
t.Error("the node survived removal")
|
||||
}
|
||||
if got := s.Allocations("n1"); len(got) != 0 {
|
||||
t.Error("allocations outlived the node")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStateStrings(t *testing.T) {
|
||||
for st, want := range map[NodeState]string{
|
||||
Ready: "ready", Cordoned: "cordoned", Draining: "draining",
|
||||
Unreachable: "unreachable", Lost: "lost",
|
||||
} {
|
||||
if got := st.String(); got != want {
|
||||
t.Errorf("%d -> %q, want %q", int(st), got, want)
|
||||
}
|
||||
}
|
||||
if !Ready.Schedulable() || Cordoned.Schedulable() || Lost.Schedulable() {
|
||||
t.Error("Schedulable is wrong")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentAccess(t *testing.T) {
|
||||
s := NewStore()
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := range 100 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
switch i % 5 {
|
||||
case 0:
|
||||
s.Upsert(Node{ID: "n1", GPUs: gpus(8)})
|
||||
case 1:
|
||||
s.Allocate("n1", Allocation{ShardID: "s", ReplicaID: "r", GPUIndices: []int{i % 8}})
|
||||
case 2:
|
||||
s.Release("s")
|
||||
case 3:
|
||||
_ = s.FreeGPUs("n1")
|
||||
default:
|
||||
_ = s.Nodes()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package inventory
|
||||
|
||||
import "time"
|
||||
|
||||
// SetLastSeen overrides a node's heartbeat timestamp.
|
||||
//
|
||||
// It exists for tests that run on a synthetic clock: Upsert and Touch stamp the
|
||||
// wall clock, which is decades away from a fake timeline and would make every
|
||||
// node look silent.
|
||||
func (s *Store) SetLastSeen(id string, at time.Time) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if n, ok := s.nodes[id]; ok {
|
||||
n.LastSeen = at
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
// Package placement decides which nodes and GPUs a replica runs on (plan §6.2).
|
||||
//
|
||||
// Two constraints here are not negotiable, and everything else in this package
|
||||
// exists to keep them true:
|
||||
//
|
||||
// - a multi-node replica never spans two fabrics
|
||||
// - a tensor-parallel group never mixes GPU classes
|
||||
//
|
||||
// Violating either produces a replica that starts, passes health checks, and
|
||||
// then hangs or crawls under load — the worst possible failure mode, because it
|
||||
// looks like success until traffic arrives.
|
||||
package placement
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// Profile is the subset of a catalog profile placement needs.
|
||||
type Profile struct {
|
||||
Name string
|
||||
Engine string
|
||||
GPUClass string
|
||||
// GPUsPerNode is how many GPUs each node contributes.
|
||||
GPUsPerNode int
|
||||
// Nodes is how many nodes the replica spans. Above 1, a fabric is required.
|
||||
Nodes int
|
||||
// Priority orders profiles; the highest that fits is used.
|
||||
Priority int
|
||||
}
|
||||
|
||||
// normalize fills in the defaults a catalog entry may omit.
|
||||
func (p Profile) normalize() Profile {
|
||||
if p.Nodes <= 0 {
|
||||
p.Nodes = 1
|
||||
}
|
||||
if p.GPUsPerNode <= 0 {
|
||||
p.GPUsPerNode = 1
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// NodeAssignment is one node's share of a replica.
|
||||
type NodeAssignment struct {
|
||||
NodeID string
|
||||
Rank int
|
||||
// GPUIndices are the local GPUs the shard uses.
|
||||
GPUIndices []int
|
||||
// HCAs are the adapters of the chosen fabric only.
|
||||
HCAs []string
|
||||
}
|
||||
|
||||
// Placement is a decision: which profile, which fabric, which nodes.
|
||||
type Placement struct {
|
||||
Profile Profile
|
||||
// Fabric is empty for a single-node replica.
|
||||
Fabric string
|
||||
// Pool names the (fabric, gpu class) pair this landed in.
|
||||
Pool string
|
||||
Nodes []NodeAssignment
|
||||
}
|
||||
|
||||
// Rejection records why a candidate was not used. Placement failures are
|
||||
// otherwise impossible to debug: "unschedulable" tells an operator nothing
|
||||
// about which of six constraints bit.
|
||||
type Rejection struct {
|
||||
Profile string
|
||||
Pool string
|
||||
Node string
|
||||
Reason string
|
||||
}
|
||||
|
||||
func (r Rejection) String() string {
|
||||
where := r.Pool
|
||||
if r.Node != "" {
|
||||
where += "/" + r.Node
|
||||
}
|
||||
if where == "" {
|
||||
where = "-"
|
||||
}
|
||||
return fmt.Sprintf("%s @ %s: %s", r.Profile, where, r.Reason)
|
||||
}
|
||||
|
||||
// Result is the outcome of a placement attempt.
|
||||
type Result struct {
|
||||
Placement *Placement
|
||||
// Rejections explains every candidate that was considered and discarded,
|
||||
// in the order they were tried. This is what `ctl explain-placement` prints.
|
||||
Rejections []Rejection
|
||||
}
|
||||
|
||||
// Explain renders the rejections as text.
|
||||
func (r Result) Explain() string {
|
||||
if len(r.Rejections) == 0 {
|
||||
return "no candidates were considered"
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, rej := range r.Rejections {
|
||||
b.WriteString(" ")
|
||||
b.WriteString(rej.String())
|
||||
b.WriteString("\n")
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// Scheduler places replicas onto an inventory.
|
||||
type Scheduler struct {
|
||||
inv *inventory.Store
|
||||
}
|
||||
|
||||
// New builds a scheduler over an inventory.
|
||||
func New(inv *inventory.Store) *Scheduler { return &Scheduler{inv: inv} }
|
||||
|
||||
// Place picks the first profile, in priority order, that fits.
|
||||
//
|
||||
// avoid lists nodes already hosting this model, so replicas of one model spread
|
||||
// across failure domains instead of stacking on one machine.
|
||||
func (s *Scheduler) Place(profiles []Profile, avoid []string) Result {
|
||||
var res Result
|
||||
|
||||
ordered := slices.Clone(profiles)
|
||||
sort.SliceStable(ordered, func(i, j int) bool { return ordered[i].Priority > ordered[j].Priority })
|
||||
|
||||
for _, raw := range ordered {
|
||||
p := raw.normalize()
|
||||
if p.Nodes == 1 {
|
||||
if pl, rej := s.placeSingleNode(p, avoid); pl != nil {
|
||||
res.Placement = pl
|
||||
return res
|
||||
} else {
|
||||
res.Rejections = append(res.Rejections, rej...)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if pl, rej := s.placeMultiNode(p, avoid); pl != nil {
|
||||
res.Placement = pl
|
||||
return res
|
||||
} else {
|
||||
res.Rejections = append(res.Rejections, rej...)
|
||||
}
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// placeSingleNode finds one node with enough free GPUs of the right class.
|
||||
//
|
||||
// Nodes without a fabric are eligible here and only here: a single-node replica
|
||||
// never talks to another node, so it has nothing to be reachable over.
|
||||
func (s *Scheduler) placeSingleNode(p Profile, avoid []string) (*Placement, []Rejection) {
|
||||
var rejections []Rejection
|
||||
|
||||
candidates := s.eligibleNodes(p, "", avoid, &rejections)
|
||||
if len(candidates) == 0 {
|
||||
return nil, rejections
|
||||
}
|
||||
|
||||
// Bin packing: the node with the fewest free GPUs that still fits. Filling
|
||||
// partly used nodes first keeps whole nodes free for the large models that
|
||||
// cannot be split.
|
||||
sort.SliceStable(candidates, func(i, j int) bool {
|
||||
return len(candidates[i].free) < len(candidates[j].free)
|
||||
})
|
||||
|
||||
best := candidates[0]
|
||||
gpus := s.pickGPUs(best.node, best.free, p.GPUsPerNode)
|
||||
if gpus == nil {
|
||||
rejections = append(rejections, Rejection{
|
||||
Profile: p.Name, Node: best.node.ID,
|
||||
Reason: fmt.Sprintf("has %d free %s GPUs but no usable set of %d", len(best.free), p.GPUClass, p.GPUsPerNode),
|
||||
})
|
||||
return nil, rejections
|
||||
}
|
||||
|
||||
return &Placement{
|
||||
Profile: p,
|
||||
Pool: "none/" + p.GPUClass,
|
||||
Nodes: []NodeAssignment{{NodeID: best.node.ID, Rank: 0, GPUIndices: gpus}},
|
||||
}, rejections
|
||||
}
|
||||
|
||||
// placeMultiNode finds N nodes on a single fabric.
|
||||
//
|
||||
// The fabric loop is the whole point: candidates are gathered per fabric and a
|
||||
// replica is only formed from within one of them. There is deliberately no path
|
||||
// that combines nodes from different fabrics, however well they would otherwise
|
||||
// fit.
|
||||
func (s *Scheduler) placeMultiNode(p Profile, avoid []string) (*Placement, []Rejection) {
|
||||
var rejections []Rejection
|
||||
|
||||
fabrics := s.inv.Fabrics()
|
||||
if len(fabrics) == 0 {
|
||||
return nil, []Rejection{{
|
||||
Profile: p.Name,
|
||||
Reason: fmt.Sprintf("needs %d nodes on one fabric, but no node reports a fabric", p.Nodes),
|
||||
}}
|
||||
}
|
||||
|
||||
for _, fabric := range fabrics {
|
||||
pool := fabric + "/" + p.GPUClass
|
||||
candidates := s.eligibleNodes(p, fabric, avoid, &rejections)
|
||||
|
||||
if len(candidates) < p.Nodes {
|
||||
rejections = append(rejections, Rejection{
|
||||
Profile: p.Name, Pool: pool,
|
||||
Reason: fmt.Sprintf("needs %d nodes, only %d qualify", p.Nodes, len(candidates)),
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
// Prefer the emptiest nodes for multi-node work: a replica spanning
|
||||
// nodes is the least movable thing in the cluster, so give it the
|
||||
// cleanest machines rather than squeezing it into the gaps.
|
||||
sort.SliceStable(candidates, func(i, j int) bool {
|
||||
return len(candidates[i].free) > len(candidates[j].free)
|
||||
})
|
||||
|
||||
assignments := make([]NodeAssignment, 0, p.Nodes)
|
||||
for rank, c := range candidates[:p.Nodes] {
|
||||
gpus := s.pickGPUs(c.node, c.free, p.GPUsPerNode)
|
||||
if gpus == nil {
|
||||
assignments = nil
|
||||
rejections = append(rejections, Rejection{
|
||||
Profile: p.Name, Pool: pool, Node: c.node.ID,
|
||||
Reason: fmt.Sprintf("no usable set of %d GPUs", p.GPUsPerNode),
|
||||
})
|
||||
break
|
||||
}
|
||||
assignments = append(assignments, NodeAssignment{
|
||||
NodeID: c.node.ID,
|
||||
Rank: rank,
|
||||
GPUIndices: gpus,
|
||||
// Only this fabric's adapters, never the node's others.
|
||||
HCAs: c.node.HCAsFor(fabric),
|
||||
})
|
||||
}
|
||||
if assignments == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
return &Placement{Profile: p, Fabric: fabric, Pool: pool, Nodes: assignments}, rejections
|
||||
}
|
||||
return nil, rejections
|
||||
}
|
||||
|
||||
// candidate is a node that passed the hard constraints.
|
||||
type candidate struct {
|
||||
node inventory.Node
|
||||
free []discovery.GPU
|
||||
}
|
||||
|
||||
// eligibleNodes applies the hard constraints and records why each node failed.
|
||||
//
|
||||
// fabric is empty for single-node placement, in which case attachment is not
|
||||
// checked at all.
|
||||
func (s *Scheduler) eligibleNodes(p Profile, fabric string, avoid []string, rejections *[]Rejection) []candidate {
|
||||
pool := "none"
|
||||
if fabric != "" {
|
||||
pool = fabric
|
||||
}
|
||||
pool += "/" + p.GPUClass
|
||||
|
||||
var out []candidate
|
||||
for _, n := range s.inv.Nodes() {
|
||||
reject := func(reason string) {
|
||||
*rejections = append(*rejections, Rejection{Profile: p.Name, Pool: pool, Node: n.ID, Reason: reason})
|
||||
}
|
||||
|
||||
if !n.State.Schedulable() {
|
||||
reject("node is " + n.State.String())
|
||||
continue
|
||||
}
|
||||
if slices.Contains(avoid, n.ID) {
|
||||
reject("already hosts this model; spreading across failure domains")
|
||||
continue
|
||||
}
|
||||
if fabric != "" {
|
||||
if !n.HasFabric(fabric) {
|
||||
// Not worth reporting for every node on every fabric; it is the
|
||||
// expected case, not a problem.
|
||||
continue
|
||||
}
|
||||
if len(n.HCAsFor(fabric)) == 0 {
|
||||
reject("attached to " + fabric + " but reports no adapters for it")
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
freeIdx := s.inv.FreeGPUs(n.ID)
|
||||
var free []discovery.GPU
|
||||
for _, g := range n.GPUs {
|
||||
if g.Class() != p.GPUClass || !slices.Contains(freeIdx, g.Index) {
|
||||
continue
|
||||
}
|
||||
free = append(free, g)
|
||||
}
|
||||
if len(free) < p.GPUsPerNode {
|
||||
reject(fmt.Sprintf("has %d free %s GPUs, needs %d", len(free), p.GPUClass, p.GPUsPerNode))
|
||||
continue
|
||||
}
|
||||
out = append(out, candidate{node: n, free: free})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// pickGPUs chooses which of a node's free GPUs to use.
|
||||
//
|
||||
// It prefers a set that shares an NVLink domain: within a node, a
|
||||
// tensor-parallel group that has to fall back to PCIe for its all-reduces runs
|
||||
// several times slower. NVLink topology is only available from the NVML build,
|
||||
// so without it this degrades to taking the lowest indices — which is usually
|
||||
// right, but only by accident of how vendors number slots.
|
||||
func (s *Scheduler) pickGPUs(n inventory.Node, free []discovery.GPU, want int) []int {
|
||||
if len(free) < want {
|
||||
return nil
|
||||
}
|
||||
if want == 1 {
|
||||
return []int{free[0].Index}
|
||||
}
|
||||
|
||||
freeSet := map[int]struct{}{}
|
||||
for _, g := range free {
|
||||
freeSet[g.Index] = struct{}{}
|
||||
}
|
||||
|
||||
// Try each free GPU as a seed and grow its NVLink-connected set.
|
||||
for _, seed := range free {
|
||||
group := []int{seed.Index}
|
||||
for _, peer := range seed.NVLinkPeers {
|
||||
if _, ok := freeSet[peer]; !ok || peer == seed.Index {
|
||||
continue
|
||||
}
|
||||
group = append(group, peer)
|
||||
if len(group) == want {
|
||||
slices.Sort(group)
|
||||
return group
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// No NVLink domain large enough, or no topology at all.
|
||||
out := make([]int, 0, want)
|
||||
for _, g := range free[:want] {
|
||||
out = append(out, g.Index)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,483 @@
|
||||
package placement
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math/rand/v2"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// node builds a node with n GPUs of one class.
|
||||
func node(id string, class string, n int, fabrics ...string) inventory.Node {
|
||||
name, mem := classToHardware(class)
|
||||
|
||||
nd := inventory.Node{ID: id, State: inventory.Ready}
|
||||
for i := range n {
|
||||
nd.GPUs = append(nd.GPUs, discovery.GPU{
|
||||
Index: i, Name: name, MemoryTotalMB: mem, ComputeCapability: "9.0",
|
||||
UUID: fmt.Sprintf("GPU-%s-%d", id, i),
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
nd.Fabrics = append(nd.Fabrics, inventory.FabricAttachment{
|
||||
Fabric: f, HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200,
|
||||
})
|
||||
}
|
||||
return nd
|
||||
}
|
||||
|
||||
// classToHardware maps a class back to a name and size that derive it.
|
||||
func classToHardware(class string) (string, int) {
|
||||
switch class {
|
||||
case "h100-80g":
|
||||
return "NVIDIA H100 80GB HBM3", 81559
|
||||
case "a100-40g":
|
||||
return "Tesla A100-SXM4-40GB", 40960
|
||||
default:
|
||||
return "NVIDIA L40S", 46068
|
||||
}
|
||||
}
|
||||
|
||||
func store(nodes ...inventory.Node) *inventory.Store {
|
||||
s := inventory.NewStore()
|
||||
for _, n := range nodes {
|
||||
s.Upsert(n)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func TestSingleNodePlacement(t *testing.T) {
|
||||
s := New(store(node("n1", "h100-80g", 8)))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if len(res.Placement.Nodes) != 1 {
|
||||
t.Fatalf("expected 1 node, got %d", len(res.Placement.Nodes))
|
||||
}
|
||||
if got := len(res.Placement.Nodes[0].GPUIndices); got != 4 {
|
||||
t.Errorf("got %d GPUs, want 4", got)
|
||||
}
|
||||
if res.Placement.Fabric != "" {
|
||||
t.Errorf("a single-node replica was given fabric %q", res.Placement.Fabric)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodeWithoutFabricTakesSingleNodeWork(t *testing.T) {
|
||||
// A single-node replica talks to nobody, so it has nothing to be reachable
|
||||
// over. Excluding fabric-less nodes here would waste them entirely.
|
||||
s := New(store(node("no-ib", "h100-80g", 8)))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("a node without IB was refused single-node work:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodeWithoutFabricRefusedMultiNode(t *testing.T) {
|
||||
s := New(store(
|
||||
node("no-ib-1", "h100-80g", 8),
|
||||
node("no-ib-2", "h100-80g", 8),
|
||||
))
|
||||
|
||||
res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("a multi-node replica was placed on nodes with no fabric")
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "no node reports a fabric") {
|
||||
t.Errorf("the reason is unclear:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
// TestNeverSpansTwoFabrics is the first hard invariant from plan §1.1.
|
||||
func TestNeverSpansTwoFabrics(t *testing.T) {
|
||||
// Two fabrics, one node each: together they would satisfy the count, and a
|
||||
// scheduler that only counted nodes would happily combine them.
|
||||
s := New(store(
|
||||
node("a1", "h100-80g", 8, "ib-a"),
|
||||
node("b1", "h100-80g", 8, "ib-b"),
|
||||
))
|
||||
|
||||
res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
|
||||
if res.Placement != nil {
|
||||
var fabrics []string
|
||||
for _, n := range res.Placement.Nodes {
|
||||
fabrics = append(fabrics, n.NodeID)
|
||||
}
|
||||
t.Fatalf("a replica was spread across fabrics: %v", fabrics)
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "only 1 qualify") {
|
||||
t.Errorf("explanation should say each fabric was short of nodes:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiNodeWithinOneFabric(t *testing.T) {
|
||||
s := New(store(
|
||||
node("a1", "h100-80g", 8, "ib-a"),
|
||||
node("a2", "h100-80g", 8, "ib-a"),
|
||||
node("b1", "h100-80g", 8, "ib-b"),
|
||||
))
|
||||
|
||||
res := s.Place([]Profile{{Name: "2node", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if res.Placement.Fabric != "ib-a" {
|
||||
t.Errorf("fabric=%q, only ib-a has two nodes", res.Placement.Fabric)
|
||||
}
|
||||
for _, n := range res.Placement.Nodes {
|
||||
if n.NodeID == "b1" {
|
||||
t.Error("a node from the other fabric was included")
|
||||
}
|
||||
if len(n.HCAs) == 0 {
|
||||
t.Errorf("node %s got no HCAs; NCCL would pick its own", n.NodeID)
|
||||
}
|
||||
}
|
||||
// Ranks must be assigned, since rank 0 serves HTTP.
|
||||
ranks := []int{res.Placement.Nodes[0].Rank, res.Placement.Nodes[1].Rank}
|
||||
slices.Sort(ranks)
|
||||
if ranks[0] != 0 || ranks[1] != 1 {
|
||||
t.Errorf("ranks=%v", ranks)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNeverMixesGPUClasses is the second hard invariant from plan §1.1.
|
||||
func TestNeverMixesGPUClasses(t *testing.T) {
|
||||
// Enough GPUs in total, but split across classes.
|
||||
n := node("mixed", "h100-80g", 4)
|
||||
a100Name, a100Mem := classToHardware("a100-40g")
|
||||
for i := 4; i < 8; i++ {
|
||||
n.GPUs = append(n.GPUs, discovery.GPU{Index: i, Name: a100Name, MemoryTotalMB: a100Mem})
|
||||
}
|
||||
s := New(store(n))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("a placement used GPUs of more than one class")
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "has 4 free h100-80g GPUs, needs 8") {
|
||||
t.Errorf("explanation should count only the matching class:\n%s", res.Explain())
|
||||
}
|
||||
|
||||
// The 4-GPU profile does fit, using only H100s.
|
||||
res = s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("the matching-class profile was refused:\n%s", res.Explain())
|
||||
}
|
||||
for _, idx := range res.Placement.Nodes[0].GPUIndices {
|
||||
if idx >= 4 {
|
||||
t.Errorf("GPU %d is an A100 and was included in an h100 placement", idx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfilePriorityOrder(t *testing.T) {
|
||||
s := New(store(
|
||||
node("h", "h100-80g", 8, "ib-a"),
|
||||
node("a", "a100-40g", 8, "ib-a"),
|
||||
))
|
||||
|
||||
res := s.Place([]Profile{
|
||||
{Name: "a100", GPUClass: "a100-40g", GPUsPerNode: 8, Nodes: 1, Priority: 50},
|
||||
{Name: "h100", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1, Priority: 100},
|
||||
}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if res.Placement.Profile.Name != "h100" {
|
||||
t.Errorf("chose %q; the higher-priority profile fits", res.Placement.Profile.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFallsBackToLowerPriorityProfile(t *testing.T) {
|
||||
s := New(store(node("a", "a100-40g", 8)))
|
||||
|
||||
res := s.Place([]Profile{
|
||||
{Name: "h100", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1, Priority: 100},
|
||||
{Name: "a100", GPUClass: "a100-40g", GPUsPerNode: 8, Nodes: 1, Priority: 50},
|
||||
}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if res.Placement.Profile.Name != "a100" {
|
||||
t.Errorf("chose %q", res.Placement.Profile.Name)
|
||||
}
|
||||
// The rejected first choice must still be explained.
|
||||
if !strings.Contains(res.Explain(), "h100") {
|
||||
t.Errorf("the higher-priority profile's rejection is missing:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCordonedNodeExcluded(t *testing.T) {
|
||||
st := store(node("n1", "h100-80g", 8))
|
||||
st.SetState("n1", inventory.Cordoned)
|
||||
s := New(st)
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("a cordoned node was used")
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "cordoned") {
|
||||
t.Errorf("explanation:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnhealthyGPUsExcluded(t *testing.T) {
|
||||
n := node("n1", "h100-80g", 8)
|
||||
n.GPUs[0].Health.Problems = []string{"12 uncorrected ECC errors"}
|
||||
n.GPUs[1].Health.Problems = []string{"40000 MiB already allocated by another process"}
|
||||
s := New(store(n))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("unhealthy GPUs were counted as available")
|
||||
}
|
||||
|
||||
res = s.Place([]Profile{{Name: "tp6", GPUClass: "h100-80g", GPUsPerNode: 6, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("the six healthy GPUs should have been usable:\n%s", res.Explain())
|
||||
}
|
||||
for _, idx := range res.Placement.Nodes[0].GPUIndices {
|
||||
if idx == 0 || idx == 1 {
|
||||
t.Errorf("unhealthy GPU %d was placed on", idx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllocatedGPUsExcluded(t *testing.T) {
|
||||
st := store(node("n1", "h100-80g", 8))
|
||||
st.Allocate("n1", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
|
||||
s := New(st)
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("the remaining four GPUs should fit:\n%s", res.Explain())
|
||||
}
|
||||
for _, idx := range res.Placement.Nodes[0].GPUIndices {
|
||||
if idx < 4 {
|
||||
t.Errorf("GPU %d is already allocated", idx)
|
||||
}
|
||||
}
|
||||
|
||||
// Another four will not fit.
|
||||
st.Allocate("n1", inventory.Allocation{ShardID: "s2", ReplicaID: "r2", GPUIndices: []int{4, 5, 6, 7}})
|
||||
if res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil); res.Placement != nil {
|
||||
t.Fatal("a full node was placed on")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAvoidSpreadsAcrossNodes(t *testing.T) {
|
||||
s := New(store(
|
||||
node("n1", "h100-80g", 8),
|
||||
node("n2", "h100-80g", 8),
|
||||
))
|
||||
p := []Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}
|
||||
|
||||
first := s.Place(p, nil)
|
||||
if first.Placement == nil {
|
||||
t.Fatal("no first placement")
|
||||
}
|
||||
placed := first.Placement.Nodes[0].NodeID
|
||||
|
||||
second := s.Place(p, []string{placed})
|
||||
if second.Placement == nil {
|
||||
t.Fatalf("the second replica had nowhere to go:\n%s", second.Explain())
|
||||
}
|
||||
if second.Placement.Nodes[0].NodeID == placed {
|
||||
t.Error("both replicas landed on the same node despite avoid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBinPackingPrefersPartlyUsedNodes(t *testing.T) {
|
||||
// Keeping whole nodes free matters: a large model that needs all 8 GPUs
|
||||
// cannot be split, so scattering small replicas would strand it.
|
||||
st := store(
|
||||
node("empty", "h100-80g", 8),
|
||||
node("partial", "h100-80g", 8),
|
||||
)
|
||||
st.Allocate("partial", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0, 1, 2, 3}})
|
||||
s := New(st)
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
if got := res.Placement.Nodes[0].NodeID; got != "partial" {
|
||||
t.Errorf("placed on %q; the partly used node should be filled first", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNVLinkGroupPreferred(t *testing.T) {
|
||||
// GPUs 0-3 and 4-7 form separate NVLink domains. A TP group crossing them
|
||||
// falls back to PCIe for its all-reduces and runs several times slower.
|
||||
n := node("n1", "h100-80g", 8)
|
||||
for i := range n.GPUs {
|
||||
var peers []int
|
||||
base := (i / 4) * 4
|
||||
for j := base; j < base+4; j++ {
|
||||
if j != i {
|
||||
peers = append(peers, j)
|
||||
}
|
||||
}
|
||||
n.GPUs[i].NVLinkPeers = peers
|
||||
}
|
||||
st := store(n)
|
||||
// Take GPU 0, so a naive "lowest indices" pick would span the two domains.
|
||||
st.Allocate("n1", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: []int{0}})
|
||||
s := New(st)
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("no placement:\n%s", res.Explain())
|
||||
}
|
||||
got := res.Placement.Nodes[0].GPUIndices
|
||||
if !slices.Equal(got, []int{4, 5, 6, 7}) {
|
||||
t.Errorf("chose %v; 4-7 is the only intact NVLink domain left", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplainIsUsefulWhenNothingFits(t *testing.T) {
|
||||
s := New(store(node("n1", "a100-40g", 2)))
|
||||
|
||||
res := s.Place([]Profile{{Name: "tp8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("unexpected placement")
|
||||
}
|
||||
exp := res.Explain()
|
||||
// It must name the node, the profile, and what was actually wrong.
|
||||
for _, want := range []string{"n1", "tp8", "free h100-80g"} {
|
||||
if !strings.Contains(exp, want) {
|
||||
t.Errorf("explanation is missing %q:\n%s", want, exp)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyInventory(t *testing.T) {
|
||||
s := New(inventory.NewStore())
|
||||
res := s.Place([]Profile{{Name: "tp4", GPUClass: "h100-80g", GPUsPerNode: 4, Nodes: 1}}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatal("something was placed on an empty cluster")
|
||||
}
|
||||
}
|
||||
|
||||
// --- property-based invariants (plan §13) ---
|
||||
|
||||
// TestInvariantsHold generates random clusters and profiles and checks the two
|
||||
// hard constraints on every placement produced.
|
||||
//
|
||||
// The table-driven tests above cover the cases someone thought of; this covers
|
||||
// the ones nobody did.
|
||||
func TestInvariantsHold(t *testing.T) {
|
||||
rng := rand.New(rand.NewPCG(1, 2))
|
||||
classes := []string{"h100-80g", "a100-40g", "l40s-45g"}
|
||||
fabrics := []string{"ib-a", "ib-b", "ib-c"}
|
||||
|
||||
for iter := range 400 {
|
||||
st := inventory.NewStore()
|
||||
nodeCount := 1 + rng.IntN(6)
|
||||
|
||||
for i := range nodeCount {
|
||||
class := classes[rng.IntN(len(classes))]
|
||||
gpuCount := 1 + rng.IntN(8)
|
||||
|
||||
var nf []string
|
||||
switch rng.IntN(4) {
|
||||
case 0: // no fabric at all
|
||||
case 1:
|
||||
nf = []string{fabrics[rng.IntN(len(fabrics))]}
|
||||
default:
|
||||
nf = []string{fabrics[rng.IntN(len(fabrics))]}
|
||||
if rng.IntN(3) == 0 {
|
||||
// A node attached to two fabrics: allowed, and the reason
|
||||
// HCAs must be selected per fabric rather than per node.
|
||||
nf = append(nf, fabrics[rng.IntN(len(fabrics))])
|
||||
}
|
||||
}
|
||||
n := node(fmt.Sprintf("n%d", i), class, gpuCount, nf...)
|
||||
if rng.IntN(8) == 0 {
|
||||
n.State = inventory.Cordoned
|
||||
}
|
||||
st.Upsert(n)
|
||||
}
|
||||
|
||||
s := New(st)
|
||||
p := Profile{
|
||||
Name: "gen",
|
||||
GPUClass: classes[rng.IntN(len(classes))],
|
||||
GPUsPerNode: 1 + rng.IntN(8),
|
||||
Nodes: 1 + rng.IntN(3),
|
||||
}
|
||||
res := s.Place([]Profile{p}, nil)
|
||||
if res.Placement == nil {
|
||||
continue
|
||||
}
|
||||
checkInvariants(t, st, p, res.Placement, iter)
|
||||
}
|
||||
}
|
||||
|
||||
func checkInvariants(t *testing.T, st *inventory.Store, p Profile, pl *Placement, iter int) {
|
||||
t.Helper()
|
||||
|
||||
if len(pl.Nodes) != max(p.Nodes, 1) {
|
||||
t.Fatalf("iter %d: placed on %d nodes, profile wants %d", iter, len(pl.Nodes), p.Nodes)
|
||||
}
|
||||
|
||||
seenNodes := map[string]struct{}{}
|
||||
for _, a := range pl.Nodes {
|
||||
if _, dup := seenNodes[a.NodeID]; dup {
|
||||
t.Fatalf("iter %d: node %s used twice in one replica", iter, a.NodeID)
|
||||
}
|
||||
seenNodes[a.NodeID] = struct{}{}
|
||||
|
||||
n, ok := st.Get(a.NodeID)
|
||||
if !ok {
|
||||
t.Fatalf("iter %d: placed on unknown node %s", iter, a.NodeID)
|
||||
}
|
||||
if !n.State.Schedulable() {
|
||||
t.Fatalf("iter %d: placed on a %s node", iter, n.State)
|
||||
}
|
||||
if len(a.GPUIndices) != p.GPUsPerNode {
|
||||
t.Fatalf("iter %d: node %s got %d GPUs, wants %d", iter, a.NodeID, len(a.GPUIndices), p.GPUsPerNode)
|
||||
}
|
||||
|
||||
// Invariant: every GPU is of the profile's class, healthy, and real.
|
||||
byIndex := map[int]discovery.GPU{}
|
||||
for _, g := range n.GPUs {
|
||||
byIndex[g.Index] = g
|
||||
}
|
||||
for _, idx := range a.GPUIndices {
|
||||
g, ok := byIndex[idx]
|
||||
if !ok {
|
||||
t.Fatalf("iter %d: node %s has no GPU %d", iter, a.NodeID, idx)
|
||||
}
|
||||
if g.Class() != p.GPUClass {
|
||||
t.Fatalf("iter %d: GPU %d on %s is %s, profile wants %s",
|
||||
iter, idx, a.NodeID, g.Class(), p.GPUClass)
|
||||
}
|
||||
if !g.Health.OK() {
|
||||
t.Fatalf("iter %d: unhealthy GPU %d was placed on", iter, idx)
|
||||
}
|
||||
}
|
||||
|
||||
// Invariant: a multi-node replica lives on one fabric, and each shard
|
||||
// gets only that fabric's adapters.
|
||||
if p.Nodes > 1 {
|
||||
if pl.Fabric == "" {
|
||||
t.Fatalf("iter %d: multi-node placement with no fabric", iter)
|
||||
}
|
||||
if !n.HasFabric(pl.Fabric) {
|
||||
t.Fatalf("iter %d: node %s is not attached to %s", iter, a.NodeID, pl.Fabric)
|
||||
}
|
||||
want := n.HCAsFor(pl.Fabric)
|
||||
if !slices.Equal(a.HCAs, want) {
|
||||
t.Fatalf("iter %d: node %s got HCAs %v, want exactly %v for %s",
|
||||
iter, a.NodeID, a.HCAs, want, pl.Fabric)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,782 @@
|
||||
package reconciler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"math/rand/v2"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/placement"
|
||||
)
|
||||
|
||||
// ShardCommand is everything an agent needs to launch one shard.
|
||||
type ShardCommand struct {
|
||||
ReplicaID string
|
||||
ShardID string
|
||||
NodeID string
|
||||
Spec engines.ShardSpec
|
||||
// StartupTimeout comes from the profile; model loading dominates it.
|
||||
StartupTimeout time.Duration
|
||||
// Pool and Fabric are control plane concepts the agent cannot derive. They
|
||||
// travel with the command so the agent can report them back after a
|
||||
// control plane restart (plan §6.5).
|
||||
Pool string
|
||||
Fabric string
|
||||
}
|
||||
|
||||
// Dispatcher sends commands to agents.
|
||||
//
|
||||
// It is an interface so the reconciliation loop can be tested without a network:
|
||||
// the loop's logic is where the interesting failures are, not the transport.
|
||||
type Dispatcher interface {
|
||||
StartShard(ctx context.Context, cmd ShardCommand) error
|
||||
StopShard(ctx context.Context, nodeID, shardID string) error
|
||||
DrainShard(ctx context.Context, nodeID, shardID string) error
|
||||
}
|
||||
|
||||
// Config tunes the reconciliation loop.
|
||||
type Config struct {
|
||||
// Interval is how often desired and actual are compared when nothing else
|
||||
// happens. Events wake the loop sooner.
|
||||
Interval time.Duration
|
||||
// BaseBackoff is how long to wait before rebuilding a failed replica.
|
||||
BaseBackoff time.Duration
|
||||
// MaxBackoff caps the exponential growth.
|
||||
MaxBackoff time.Duration
|
||||
// UnreachableAfter is how long without a heartbeat before a node is
|
||||
// unreachable.
|
||||
UnreachableAfter time.Duration
|
||||
// LostAfter is how long unreachable before its replicas are declared dead.
|
||||
// Longer than UnreachableAfter on purpose: a brief network blip must not
|
||||
// tear down a replica that is still serving perfectly well.
|
||||
LostAfter time.Duration
|
||||
// DrainGrace bounds how long a draining replica waits for its replacement.
|
||||
// Without a bound, draining a node whose model cannot be placed anywhere
|
||||
// else would never finish.
|
||||
DrainGrace time.Duration
|
||||
// PortRange is the per-node port range for engine HTTP and dist-init.
|
||||
PortRange [2]int
|
||||
// MgmtInterface is the ethernet interface used for NCCL bootstrap.
|
||||
MgmtInterface string
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Interval: 5 * time.Second,
|
||||
BaseBackoff: 10 * time.Second,
|
||||
MaxBackoff: 10 * time.Minute,
|
||||
UnreachableAfter: 15 * time.Second,
|
||||
LostAfter: 60 * time.Second,
|
||||
DrainGrace: 5 * time.Minute,
|
||||
PortRange: [2]int{20000, 20999},
|
||||
MgmtInterface: "eth0",
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() {
|
||||
d := DefaultConfig()
|
||||
if c.Interval <= 0 {
|
||||
c.Interval = d.Interval
|
||||
}
|
||||
if c.BaseBackoff <= 0 {
|
||||
c.BaseBackoff = d.BaseBackoff
|
||||
}
|
||||
if c.MaxBackoff < c.BaseBackoff {
|
||||
c.MaxBackoff = max(c.BaseBackoff, d.MaxBackoff)
|
||||
}
|
||||
if c.UnreachableAfter <= 0 {
|
||||
c.UnreachableAfter = d.UnreachableAfter
|
||||
}
|
||||
if c.LostAfter < c.UnreachableAfter {
|
||||
c.LostAfter = max(c.UnreachableAfter*4, d.LostAfter)
|
||||
}
|
||||
if c.DrainGrace <= 0 {
|
||||
c.DrainGrace = d.DrainGrace
|
||||
}
|
||||
if c.PortRange[1] <= c.PortRange[0] {
|
||||
c.PortRange = d.PortRange
|
||||
}
|
||||
if c.MgmtInterface == "" {
|
||||
c.MgmtInterface = d.MgmtInterface
|
||||
}
|
||||
}
|
||||
|
||||
// Reconciler compares the catalog against what is running and closes the gap.
|
||||
type Reconciler struct {
|
||||
cfg Config
|
||||
cat *catalog.Catalog
|
||||
inv *inventory.Store
|
||||
sched *placement.Scheduler
|
||||
disp Dispatcher
|
||||
log *slog.Logger
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
replicas map[string]*Replica
|
||||
// nextAttempt is the earliest a model may be rebuilt, per model.
|
||||
nextAttempt map[string]time.Time
|
||||
// attempts counts consecutive failures per model, for the backoff.
|
||||
attempts map[string]int
|
||||
// lastExplain keeps the most recent placement failure per model, so
|
||||
// `ctl explain-placement` can answer without re-running placement.
|
||||
lastExplain map[string]placement.Result
|
||||
seq int
|
||||
|
||||
wake chan struct{}
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// New builds a reconciler.
|
||||
func New(cfg Config, cat *catalog.Catalog, inv *inventory.Store, disp Dispatcher, log *slog.Logger) *Reconciler {
|
||||
cfg.normalize()
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Reconciler{
|
||||
cfg: cfg, cat: cat, inv: inv, disp: disp, log: log,
|
||||
sched: placement.New(inv),
|
||||
now: time.Now,
|
||||
replicas: map[string]*Replica{},
|
||||
nextAttempt: map[string]time.Time{},
|
||||
attempts: map[string]int{},
|
||||
lastExplain: map[string]placement.Result{},
|
||||
wake: make(chan struct{}, 1),
|
||||
stop: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// SetClock replaces the time source, for tests.
|
||||
func (r *Reconciler) SetClock(now func() time.Time) { r.now = now }
|
||||
|
||||
// Start runs the loop until Close.
|
||||
func (r *Reconciler) Start(ctx context.Context) {
|
||||
go func() {
|
||||
defer close(r.done)
|
||||
t := time.NewTicker(r.cfg.Interval)
|
||||
defer t.Stop()
|
||||
|
||||
r.Reconcile(ctx)
|
||||
for {
|
||||
select {
|
||||
case <-r.stop:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-r.wake:
|
||||
r.Reconcile(ctx)
|
||||
case <-t.C:
|
||||
r.Reconcile(ctx)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Close stops the loop.
|
||||
func (r *Reconciler) Close() {
|
||||
r.once.Do(func() {
|
||||
close(r.stop)
|
||||
<-r.done
|
||||
})
|
||||
}
|
||||
|
||||
// Notify wakes the loop, for when something changed and waiting for the next
|
||||
// tick would add pointless latency.
|
||||
func (r *Reconciler) Notify() {
|
||||
select {
|
||||
case r.wake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// Reconcile performs one pass. It is exported so tests can drive it directly
|
||||
// instead of racing a ticker.
|
||||
func (r *Reconciler) Reconcile(ctx context.Context) {
|
||||
r.expireNodes()
|
||||
r.drainCordonedNodes(ctx)
|
||||
r.reapFailed(ctx)
|
||||
|
||||
for _, m := range r.cat.Models() {
|
||||
r.reconcileModel(ctx, m)
|
||||
}
|
||||
}
|
||||
|
||||
// expireNodes moves silent nodes through Unreachable to Lost, and fails the
|
||||
// replicas of nodes that are gone.
|
||||
//
|
||||
// The two-stage delay is deliberate: a brief network blip must not tear down
|
||||
// replicas that are still serving traffic perfectly well, and rebuilding a
|
||||
// multi-node replica costs tens of minutes of model loading.
|
||||
func (r *Reconciler) expireNodes() {
|
||||
now := r.now()
|
||||
for _, n := range r.inv.Nodes() {
|
||||
silent := now.Sub(n.LastSeen)
|
||||
switch {
|
||||
case n.State == inventory.Ready && silent > r.cfg.UnreachableAfter:
|
||||
r.inv.SetState(n.ID, inventory.Unreachable)
|
||||
r.log.Warn("node is unreachable", "node", n.ID, "silent_for", silent.Round(time.Second))
|
||||
case n.State == inventory.Unreachable && silent > r.cfg.LostAfter:
|
||||
r.inv.SetState(n.ID, inventory.Lost)
|
||||
r.log.Error("node lost; its replicas will be rebuilt elsewhere",
|
||||
"node", n.ID, "silent_for", silent.Round(time.Second))
|
||||
r.failReplicasOn(n.ID, "node "+n.ID+" was lost")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// drainCordonedNodes moves work off nodes an operator is taking out of service.
|
||||
//
|
||||
// Draining a node is not the same as losing one: the engines are healthy and
|
||||
// still serving, so their replicas are marked Draining rather than Failed. That
|
||||
// makes the model look short of replicas, a replacement is placed elsewhere,
|
||||
// and the original is stopped only once the replacement is serving — which is
|
||||
// the difference between a rolling move and an outage.
|
||||
func (r *Reconciler) drainCordonedNodes(ctx context.Context) {
|
||||
var toDrain []*Replica
|
||||
|
||||
r.mu.Lock()
|
||||
for _, rep := range r.replicas {
|
||||
if rep.State != Ready && rep.State != Starting && rep.State != Placing {
|
||||
continue
|
||||
}
|
||||
for _, s := range rep.Shards {
|
||||
n, ok := r.inv.Get(s.NodeID)
|
||||
if ok && n.State == inventory.Draining {
|
||||
rep.State = Draining
|
||||
rep.DrainingAt = r.now()
|
||||
toDrain = append(toDrain, rep)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
shards := map[string][]Shard{}
|
||||
for _, rep := range toDrain {
|
||||
shards[rep.ID] = slices.Clone(rep.Shards)
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, rep := range toDrain {
|
||||
r.log.Info("draining replica off a node being taken out of service",
|
||||
"replica", rep.ID, "nodes", rep.NodeIDs())
|
||||
for _, s := range shards[rep.ID] {
|
||||
if err := r.disp.DrainShard(ctx, s.NodeID, s.ID); err != nil {
|
||||
r.log.Warn("draining a shard failed", "shard", s.ID, "err", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reconciler) failReplicasOn(nodeID, reason string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, rep := range r.replicas {
|
||||
if rep.State.Terminal() || !rep.OnNode(nodeID) {
|
||||
continue
|
||||
}
|
||||
rep.State = Failed
|
||||
rep.Err = reason
|
||||
}
|
||||
}
|
||||
|
||||
// reapFailed releases the resources of terminal replicas and tells the agents
|
||||
// to clean up whatever might still be running.
|
||||
func (r *Reconciler) reapFailed(ctx context.Context) {
|
||||
r.mu.Lock()
|
||||
var doomed []*Replica
|
||||
for _, rep := range r.replicas {
|
||||
if rep.State.Terminal() {
|
||||
doomed = append(doomed, rep)
|
||||
}
|
||||
}
|
||||
for _, rep := range doomed {
|
||||
delete(r.replicas, rep.ID)
|
||||
if rep.State == Failed {
|
||||
// Carry the failure count on the model, not the replica: the
|
||||
// replica is about to disappear, and the backoff has to survive it.
|
||||
r.attempts[rep.Model]++
|
||||
r.nextAttempt[rep.Model] = r.now().Add(r.backoff(r.attempts[rep.Model]))
|
||||
}
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, rep := range doomed {
|
||||
// rep is out of the map now, so no agent report can reach it any more.
|
||||
if rep.State == Failed {
|
||||
r.mu.Lock()
|
||||
retry := r.backoff(r.attempts[rep.Model])
|
||||
r.mu.Unlock()
|
||||
r.log.Error("replica failed",
|
||||
"replica", rep.ID, "model", rep.Model, "err", rep.Err,
|
||||
"retry_in", retry.Round(time.Second))
|
||||
}
|
||||
// Stop every shard, including the ones that looked healthy: a partially
|
||||
// alive replica holds GPUs while serving nothing.
|
||||
for _, s := range rep.Shards {
|
||||
if err := r.disp.StopShard(ctx, s.NodeID, s.ID); err != nil {
|
||||
r.log.Warn("stopping a shard of a dead replica failed",
|
||||
"replica", rep.ID, "shard", s.ID, "err", err)
|
||||
}
|
||||
}
|
||||
r.inv.ReleaseReplica(rep.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// backoff grows with consecutive failures and is jittered, so a cluster-wide
|
||||
// problem does not produce synchronized retry storms across every model.
|
||||
func (r *Reconciler) backoff(attempt int) time.Duration {
|
||||
if attempt <= 0 {
|
||||
return 0
|
||||
}
|
||||
d := r.cfg.BaseBackoff << min(attempt-1, 16)
|
||||
if d > r.cfg.MaxBackoff || d <= 0 {
|
||||
d = r.cfg.MaxBackoff
|
||||
}
|
||||
jitter := time.Duration(rand.Int64N(int64(d / 4))) //nolint:gosec // jitter
|
||||
return d - d/8 + jitter
|
||||
}
|
||||
|
||||
func (r *Reconciler) reconcileModel(ctx context.Context, m catalog.Model) {
|
||||
// Snapshot under the lock and work from values afterwards. A replica
|
||||
// pointer stops being safe the moment it is in the map: an agent's status
|
||||
// report can mutate it from another goroutine at any time.
|
||||
r.mu.Lock()
|
||||
var liveIDs, drainingIDs, avoid []string
|
||||
for _, rep := range r.replicas {
|
||||
if rep.Model != m.Name {
|
||||
continue
|
||||
}
|
||||
switch {
|
||||
case rep.State.Counts():
|
||||
liveIDs = append(liveIDs, rep.ID)
|
||||
avoid = append(avoid, rep.NodeIDs()...)
|
||||
case rep.State == Draining:
|
||||
drainingIDs = append(drainingIDs, rep.ID)
|
||||
}
|
||||
}
|
||||
blockedUntil := r.nextAttempt[m.Name]
|
||||
r.mu.Unlock()
|
||||
|
||||
desired := m.Scaling.Min
|
||||
switch {
|
||||
case len(liveIDs) < desired:
|
||||
if r.now().Before(blockedUntil) {
|
||||
// Still in backoff. Rebuilding immediately would hot-loop on a
|
||||
// model that cannot start — a missing weights path, say — and bury
|
||||
// the real error in noise.
|
||||
return
|
||||
}
|
||||
r.scaleUp(ctx, m, avoid)
|
||||
case len(liveIDs) > desired:
|
||||
r.scaleDown(ctx, liveIDs[desired:])
|
||||
}
|
||||
|
||||
// A drained replica is stopped once its replacement is serving, or once the
|
||||
// grace period expires. Stopping it the moment it was drained would cut
|
||||
// in-flight requests, which is the thing draining exists to avoid.
|
||||
replaced := len(liveIDs) >= desired
|
||||
for _, id := range drainingIDs {
|
||||
r.mu.Lock()
|
||||
rep, ok := r.replicas[id]
|
||||
var expired bool
|
||||
if ok {
|
||||
expired = !rep.DrainingAt.IsZero() && r.now().Sub(rep.DrainingAt) > r.cfg.DrainGrace
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
switch {
|
||||
case !ok:
|
||||
case replaced:
|
||||
r.stopReplica(ctx, id)
|
||||
case expired:
|
||||
// The replacement never arrived — often because nowhere else can
|
||||
// host it. Draining forever would leave the node undrainable.
|
||||
r.log.Warn("stopping a drained replica without a replacement",
|
||||
"replica", id, "after", r.cfg.DrainGrace)
|
||||
r.stopReplica(ctx, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// scaleUp places and starts one replica per pass.
|
||||
//
|
||||
// One at a time on purpose: placement decisions depend on what is already
|
||||
// allocated, and starting several at once from the same inventory snapshot
|
||||
// would hand them the same GPUs.
|
||||
func (r *Reconciler) scaleUp(ctx context.Context, m catalog.Model, avoid []string) {
|
||||
res := r.sched.Place(m.PlacementProfiles(), avoid)
|
||||
if res.Placement == nil && len(avoid) > 0 {
|
||||
// Spreading across failure domains is a preference, not a rule. If it
|
||||
// is the only thing standing in the way, take the placement anyway:
|
||||
// a second replica on one node still survives an engine crash.
|
||||
res = r.sched.Place(m.PlacementProfiles(), nil)
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.lastExplain[m.Name] = res
|
||||
r.mu.Unlock()
|
||||
|
||||
if res.Placement == nil {
|
||||
r.log.Warn("no placement for model", "model", m.Name, "reasons", res.Explain())
|
||||
return
|
||||
}
|
||||
if err := r.startReplica(ctx, m, res.Placement); err != nil {
|
||||
r.log.Error("starting replica failed", "model", m.Name, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reconciler) startReplica(ctx context.Context, m catalog.Model, pl *placement.Placement) error {
|
||||
prof, ok := m.ProfileByName(pl.Profile.Name)
|
||||
if !ok {
|
||||
return fmt.Errorf("profile %q vanished from model %q", pl.Profile.Name, m.Name)
|
||||
}
|
||||
espec, ok := m.EngineProfile(prof.Name)
|
||||
if !ok {
|
||||
return fmt.Errorf("model %q: no engine profile %q", m.Name, prof.Name)
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.seq++
|
||||
id := fmt.Sprintf("%s-%d", m.Name, r.seq)
|
||||
attempt := r.attempts[m.Name]
|
||||
r.mu.Unlock()
|
||||
|
||||
rep := &Replica{
|
||||
ID: id, Model: m.Name, Profile: prof.Name,
|
||||
Fabric: pl.Fabric, Pool: pl.Pool, State: Placing,
|
||||
Weight: prof.Weight, CreatedAt: r.now(), Attempt: attempt,
|
||||
}
|
||||
|
||||
// The leader's address and port are fixed before any shard starts: every
|
||||
// rank needs the same rendezvous point, and a worker that starts first must
|
||||
// already know where to dial.
|
||||
leader := pl.Nodes[0]
|
||||
leaderNode, ok := r.inv.Get(leader.NodeID)
|
||||
if !ok {
|
||||
return fmt.Errorf("leader node %s disappeared", leader.NodeID)
|
||||
}
|
||||
basePort := r.allocatePorts(pl)
|
||||
|
||||
for _, a := range pl.Nodes {
|
||||
shardID := fmt.Sprintf("%s-r%d", id, a.Rank)
|
||||
rep.Shards = append(rep.Shards, Shard{
|
||||
ID: shardID, NodeID: a.NodeID, Rank: a.Rank,
|
||||
GPUIndices: a.GPUIndices, HCAs: a.HCAs,
|
||||
HTTPPort: basePort, State: supervisor.Pending,
|
||||
})
|
||||
r.inv.Allocate(a.NodeID, inventory.Allocation{
|
||||
ShardID: shardID, ReplicaID: id, GPUIndices: a.GPUIndices,
|
||||
})
|
||||
}
|
||||
rep.Endpoint = fmt.Sprintf("http://%s:%d", addrOf(leaderNode), basePort)
|
||||
|
||||
// Everything read from rep has to be read now: once it is in the map an
|
||||
// agent status report can mutate it concurrently, and the race detector is
|
||||
// right to object.
|
||||
//
|
||||
// The leader first, then workers ascending. Both engines put the rendezvous
|
||||
// point on rank 0 — SGLang's dist-init address and vLLM's Ray head — so a
|
||||
// worker started first has nothing to connect to and spends its startup
|
||||
// budget retrying.
|
||||
//
|
||||
// An earlier version started workers first, on the theory that the head
|
||||
// should wait for them. That is backwards: vLLM's server blocks until the
|
||||
// Ray cluster has the GPUs it needs, so the leader can be started
|
||||
// immediately and will wait on its own.
|
||||
order := slices.Clone(rep.Shards)
|
||||
slices.SortFunc(order, func(a, b Shard) int { return a.Rank - b.Rank })
|
||||
nodeIDs := rep.NodeIDs()
|
||||
|
||||
r.mu.Lock()
|
||||
r.replicas[id] = rep
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, s := range order {
|
||||
cmd := ShardCommand{
|
||||
ReplicaID: id,
|
||||
ShardID: s.ID,
|
||||
NodeID: s.NodeID,
|
||||
Spec: engines.ShardSpec{
|
||||
ReplicaID: id,
|
||||
ShardID: s.ID,
|
||||
Rank: s.Rank,
|
||||
NNodes: len(pl.Nodes),
|
||||
LeaderAddr: addrOf(leaderNode),
|
||||
LeaderPort: basePort + 1,
|
||||
GPUIndices: s.GPUIndices,
|
||||
HCAs: s.HCAs,
|
||||
MgmtInterface: r.cfg.MgmtInterface,
|
||||
HTTPPort: s.HTTPPort,
|
||||
Profile: espec,
|
||||
},
|
||||
StartupTimeout: prof.StartupTimeout,
|
||||
Pool: pl.Pool,
|
||||
Fabric: pl.Fabric,
|
||||
}
|
||||
if err := r.disp.StartShard(ctx, cmd); err != nil {
|
||||
r.markShard(id, s.ID, supervisor.Failed, err.Error())
|
||||
return fmt.Errorf("starting shard %s: %w", s.ID, err)
|
||||
}
|
||||
r.markShard(id, s.ID, supervisor.Starting, "")
|
||||
}
|
||||
|
||||
r.log.Info("replica placed",
|
||||
"replica", id, "model", m.Name, "profile", prof.Name,
|
||||
"pool", pl.Pool, "fabric", pl.Fabric, "nodes", nodeIDs)
|
||||
return nil
|
||||
}
|
||||
|
||||
// allocatePorts picks a port shared by every shard of a replica.
|
||||
//
|
||||
// The same port on every node keeps the leader address predictable and the
|
||||
// port bookkeeping trivial; nodes do not share a network namespace, so there is
|
||||
// no conflict in reusing it.
|
||||
func (r *Reconciler) allocatePorts(pl *placement.Placement) int {
|
||||
used := map[int]struct{}{}
|
||||
r.mu.Lock()
|
||||
for _, rep := range r.replicas {
|
||||
for _, s := range rep.Shards {
|
||||
used[s.HTTPPort] = struct{}{}
|
||||
}
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
// Two ports per replica: HTTP and the distributed init rendezvous.
|
||||
for p := r.cfg.PortRange[0]; p < r.cfg.PortRange[1]; p += 2 {
|
||||
if _, taken := used[p]; !taken {
|
||||
return p
|
||||
}
|
||||
}
|
||||
return r.cfg.PortRange[0]
|
||||
}
|
||||
|
||||
func (r *Reconciler) scaleDown(ctx context.Context, excess []string) {
|
||||
for _, id := range excess {
|
||||
r.mu.Lock()
|
||||
rep, ok := r.replicas[id]
|
||||
if !ok {
|
||||
r.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
rep.State = Draining
|
||||
rep.DrainingAt = r.now()
|
||||
shards := slices.Clone(rep.Shards)
|
||||
r.mu.Unlock()
|
||||
|
||||
r.log.Info("draining replica", "replica", id)
|
||||
for _, s := range shards {
|
||||
if err := r.disp.DrainShard(ctx, s.NodeID, s.ID); err != nil {
|
||||
r.log.Warn("draining a shard failed", "shard", s.ID, "err", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reconciler) stopReplica(ctx context.Context, id string) {
|
||||
r.mu.Lock()
|
||||
rep, ok := r.replicas[id]
|
||||
if !ok {
|
||||
r.mu.Unlock()
|
||||
return
|
||||
}
|
||||
shards := slices.Clone(rep.Shards)
|
||||
rep.State = Stopped
|
||||
r.mu.Unlock()
|
||||
|
||||
for _, s := range shards {
|
||||
if err := r.disp.StopShard(ctx, s.NodeID, s.ID); err != nil {
|
||||
r.log.Warn("stopping a shard failed", "shard", s.ID, "err", err)
|
||||
}
|
||||
}
|
||||
r.inv.ReleaseReplica(id)
|
||||
|
||||
r.mu.Lock()
|
||||
delete(r.replicas, id)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
// Adopt rebuilds replica state from shards an agent reports it is already
|
||||
// running (plan §6.5).
|
||||
//
|
||||
// This is what makes a control plane restart survivable. Without it the
|
||||
// reconciler sees an empty cluster, places duplicates, and the engines already
|
||||
// running keep holding GPUs the inventory believes are free — so the duplicates
|
||||
// fail to allocate, with an error naming the new process rather than the old.
|
||||
//
|
||||
// Shards whose model or profile is no longer in the catalog are returned for
|
||||
// the caller to stop: nothing will ever claim them.
|
||||
func (r *Reconciler) Adopt(shards []AdoptedShard) (unclaimed []AdoptedShard) {
|
||||
byReplica := map[string][]AdoptedShard{}
|
||||
for _, s := range shards {
|
||||
byReplica[s.ReplicaID] = append(byReplica[s.ReplicaID], s)
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
for replicaID, group := range byReplica {
|
||||
if _, known := r.replicas[replicaID]; known {
|
||||
// Already tracked — an agent reconnecting rather than a control
|
||||
// plane restart. Its status reports keep the state current.
|
||||
continue
|
||||
}
|
||||
|
||||
first := group[0]
|
||||
m, ok := r.cat.Resolve(first.Model)
|
||||
if !ok {
|
||||
unclaimed = append(unclaimed, group...)
|
||||
continue
|
||||
}
|
||||
prof, ok := m.ProfileByName(first.Profile)
|
||||
if !ok {
|
||||
unclaimed = append(unclaimed, group...)
|
||||
continue
|
||||
}
|
||||
|
||||
rep := &Replica{
|
||||
ID: replicaID, Model: m.Name, Profile: prof.Name,
|
||||
Fabric: first.Fabric, Pool: first.Pool,
|
||||
Weight: prof.Weight, CreatedAt: r.now(),
|
||||
Endpoint: first.Endpoint,
|
||||
}
|
||||
for _, s := range group {
|
||||
rep.Shards = append(rep.Shards, Shard{
|
||||
ID: s.ShardID, NodeID: s.NodeID, Rank: s.Rank,
|
||||
GPUIndices: s.GPUIndices, HCAs: s.HCAs,
|
||||
HTTPPort: s.HTTPPort, State: s.State,
|
||||
})
|
||||
// Reclaim the GPUs, or the next placement will hand them out twice.
|
||||
r.inv.Allocate(s.NodeID, inventory.Allocation{
|
||||
ShardID: s.ShardID, ReplicaID: replicaID, GPUIndices: s.GPUIndices,
|
||||
})
|
||||
}
|
||||
// The endpoint comes from the leader; a group missing rank 0 is a
|
||||
// partial replica and will be failed by deriveState below.
|
||||
for _, s := range group {
|
||||
if s.Rank == 0 && s.Endpoint != "" {
|
||||
rep.Endpoint = s.Endpoint
|
||||
}
|
||||
}
|
||||
rep.deriveState()
|
||||
r.replicas[replicaID] = rep
|
||||
|
||||
r.log.Info("adopted replica from the fleet",
|
||||
"replica", replicaID, "model", m.Name, "shards", len(group), "state", rep.State)
|
||||
}
|
||||
return unclaimed
|
||||
}
|
||||
|
||||
// AdoptedShard is one running shard an agent reported at registration.
|
||||
type AdoptedShard struct {
|
||||
NodeID string
|
||||
ReplicaID string
|
||||
ShardID string
|
||||
Rank int
|
||||
GPUIndices []int
|
||||
HCAs []string
|
||||
HTTPPort int
|
||||
State supervisor.State
|
||||
Model string
|
||||
Profile string
|
||||
Pool string
|
||||
Fabric string
|
||||
// Endpoint is the leader's base URL, set on the rank 0 shard.
|
||||
Endpoint string
|
||||
}
|
||||
|
||||
// OnShardStatus records a status report from an agent.
|
||||
func (r *Reconciler) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
|
||||
r.markShard(replicaID, shardID, st, errMsg)
|
||||
r.Notify()
|
||||
}
|
||||
|
||||
func (r *Reconciler) markShard(replicaID, shardID string, st supervisor.State, errMsg string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
rep, ok := r.replicas[replicaID]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for i := range rep.Shards {
|
||||
if rep.Shards[i].ID != shardID {
|
||||
continue
|
||||
}
|
||||
rep.Shards[i].State = st
|
||||
if errMsg != "" {
|
||||
rep.Shards[i].Err = errMsg
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
prev := rep.State
|
||||
rep.deriveState()
|
||||
if rep.State == Ready && prev != Ready {
|
||||
// A model that comes up clears its backoff; the next failure starts
|
||||
// counting from scratch rather than from an old streak.
|
||||
r.attempts[rep.Model] = 0
|
||||
delete(r.nextAttempt, rep.Model)
|
||||
}
|
||||
}
|
||||
|
||||
// Replicas returns a snapshot of every replica, sorted by id.
|
||||
func (r *Reconciler) Replicas() []Replica {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
out := make([]Replica, 0, len(r.replicas))
|
||||
for _, rep := range r.replicas {
|
||||
clone := *rep
|
||||
clone.Shards = slices.Clone(rep.Shards)
|
||||
out = append(out, clone)
|
||||
}
|
||||
slices.SortFunc(out, func(a, b Replica) int {
|
||||
if a.ID < b.ID {
|
||||
return -1
|
||||
}
|
||||
if a.ID > b.ID {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// ReadyReplicas returns the replicas the router may send traffic to.
|
||||
func (r *Reconciler) ReadyReplicas() []Replica {
|
||||
var out []Replica
|
||||
for _, rep := range r.Replicas() {
|
||||
if rep.State == Ready {
|
||||
out = append(out, rep)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Explain returns the most recent placement result for a model, which is what
|
||||
// `ctl explain-placement` prints.
|
||||
func (r *Reconciler) Explain(model string) (placement.Result, bool) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
res, ok := r.lastExplain[model]
|
||||
return res, ok
|
||||
}
|
||||
|
||||
func addrOf(n inventory.Node) string {
|
||||
if n.MgmtAddr != "" {
|
||||
return n.MgmtAddr
|
||||
}
|
||||
if n.Hostname != "" {
|
||||
return n.Hostname
|
||||
}
|
||||
return n.ID
|
||||
}
|
||||
@@ -0,0 +1,900 @@
|
||||
package reconciler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// fakeDispatcher records commands instead of sending them.
|
||||
type fakeDispatcher struct {
|
||||
mu sync.Mutex
|
||||
started []ShardCommand
|
||||
stopped []string
|
||||
drained []string
|
||||
failNext error
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) StartShard(_ context.Context, cmd ShardCommand) error {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
if d.failNext != nil {
|
||||
err := d.failNext
|
||||
d.failNext = nil
|
||||
return err
|
||||
}
|
||||
d.started = append(d.started, cmd)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) StopShard(_ context.Context, _, shardID string) error {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
d.stopped = append(d.stopped, shardID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) DrainShard(_ context.Context, _, shardID string) error {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
d.drained = append(d.drained, shardID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) startedShards() []ShardCommand {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return slices.Clone(d.started)
|
||||
}
|
||||
|
||||
func (d *fakeDispatcher) reset() {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
d.started, d.stopped, d.drained = nil, nil, nil
|
||||
}
|
||||
|
||||
// clock lets tests move time without sleeping.
|
||||
type clock struct {
|
||||
mu sync.Mutex
|
||||
t time.Time
|
||||
}
|
||||
|
||||
func newClock() *clock {
|
||||
return &clock{t: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)}
|
||||
}
|
||||
func (c *clock) now() time.Time {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.t
|
||||
}
|
||||
func (c *clock) advance(d time.Duration) {
|
||||
c.mu.Lock()
|
||||
c.t = c.t.Add(d)
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func gpuNode(id string, n int, fabrics ...string) inventory.Node {
|
||||
nd := inventory.Node{ID: id, MgmtAddr: "10.0.0." + id, State: inventory.Ready}
|
||||
for i := range n {
|
||||
nd.GPUs = append(nd.GPUs, discovery.GPU{
|
||||
Index: i, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559, ComputeCapability: "9.0",
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
nd.Fabrics = append(nd.Fabrics, inventory.FabricAttachment{
|
||||
Fabric: f, HCAs: []string{"mlx5_0", "mlx5_1"},
|
||||
})
|
||||
}
|
||||
return nd
|
||||
}
|
||||
|
||||
func testModel(min int, profiles ...catalog.Profile) catalog.Model {
|
||||
if len(profiles) == 0 {
|
||||
profiles = []catalog.Profile{{Name: "tp4", GPUClass: "h100-80g", TP: 4}}
|
||||
}
|
||||
return catalog.Model{
|
||||
Name: "llama", Source: "/models/llama",
|
||||
Scaling: catalog.Scaling{Min: min, Max: min + 2}, Profiles: profiles,
|
||||
}
|
||||
}
|
||||
|
||||
type harness struct {
|
||||
r *Reconciler
|
||||
inv *inventory.Store
|
||||
disp *fakeDispatcher
|
||||
clock *clock
|
||||
}
|
||||
|
||||
func newHarness(t *testing.T, m catalog.Model, nodes ...inventory.Node) *harness {
|
||||
t.Helper()
|
||||
|
||||
cat, err := catalog.New([]catalog.Model{m})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inv := inventory.NewStore()
|
||||
c := newClock()
|
||||
for _, n := range nodes {
|
||||
inv.Upsert(n)
|
||||
inv.Touch(n.ID)
|
||||
}
|
||||
|
||||
disp := &fakeDispatcher{}
|
||||
cfg := DefaultConfig()
|
||||
cfg.UnreachableAfter = 15 * time.Second
|
||||
cfg.LostAfter = 60 * time.Second
|
||||
|
||||
r := New(cfg, cat, inv, disp, slog.New(slog.DiscardHandler))
|
||||
r.SetClock(c.now)
|
||||
// Nodes were just registered, so their LastSeen is wall-clock now, which is
|
||||
// far from the fake clock's epoch. Re-stamp them onto the fake timeline.
|
||||
for _, n := range nodes {
|
||||
touchAt(inv, n.ID, c.now())
|
||||
}
|
||||
return &harness{r: r, inv: inv, disp: disp, clock: c}
|
||||
}
|
||||
|
||||
// touchAt rewrites a node's LastSeen onto the fake clock's timeline.
|
||||
func touchAt(inv *inventory.Store, id string, at time.Time) {
|
||||
n, ok := inv.Get(id)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
n.LastSeen = at
|
||||
n.State = inventory.Ready
|
||||
inv.Upsert(n)
|
||||
// Upsert stamps LastSeen with the wall clock, so set it again directly.
|
||||
inv.SetLastSeen(id, at)
|
||||
}
|
||||
|
||||
// advance moves the clock and keeps every node heartbeating.
|
||||
//
|
||||
// Without this, advancing time past UnreachableAfter makes every node look
|
||||
// silent and placement stops — which is correct behaviour, but not what the
|
||||
// backoff tests are trying to exercise.
|
||||
func (h *harness) advance(d time.Duration) {
|
||||
h.clock.advance(d)
|
||||
for _, n := range h.inv.Nodes() {
|
||||
touchAt(h.inv, n.ID, h.clock.now())
|
||||
}
|
||||
}
|
||||
|
||||
// markReady walks every started shard to Ready.
|
||||
func (h *harness) markReady() {
|
||||
for _, cmd := range h.disp.startedShards() {
|
||||
h.r.OnShardStatus(cmd.ReplicaID, cmd.ShardID, supervisor.Ready, "")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlacesMinimumReplicas(t *testing.T) {
|
||||
h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 2 {
|
||||
t.Fatalf("expected 2 replicas, got %d: %+v", len(reps), reps)
|
||||
}
|
||||
for _, rep := range reps {
|
||||
if rep.State != Ready {
|
||||
t.Errorf("replica %s is %s", rep.ID, rep.State)
|
||||
}
|
||||
if rep.Endpoint == "" {
|
||||
t.Errorf("replica %s has no endpoint", rep.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneReplicaPerPass(t *testing.T) {
|
||||
// Placement depends on what is already allocated, so starting several from
|
||||
// one inventory snapshot would hand them the same GPUs.
|
||||
h := newHarness(t, testModel(3), gpuNode("1", 8), gpuNode("2", 8), gpuNode("3", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Fatalf("one pass created %d replicas", got)
|
||||
}
|
||||
|
||||
seen := map[int]string{}
|
||||
for range 3 {
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
}
|
||||
for _, rep := range h.r.Replicas() {
|
||||
for _, s := range rep.Shards {
|
||||
for _, g := range s.GPUIndices {
|
||||
key := g
|
||||
if prev, dup := seen[key]; dup && prev != s.NodeID {
|
||||
continue // same index on a different node is fine
|
||||
}
|
||||
seen[key] = s.NodeID
|
||||
}
|
||||
}
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 3 {
|
||||
t.Errorf("expected 3 replicas, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpreadsAcrossNodes(t *testing.T) {
|
||||
h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
for range 2 {
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
}
|
||||
|
||||
nodes := map[string]struct{}{}
|
||||
for _, rep := range h.r.Replicas() {
|
||||
for _, n := range rep.NodeIDs() {
|
||||
nodes[n] = struct{}{}
|
||||
}
|
||||
}
|
||||
if len(nodes) != 2 {
|
||||
t.Errorf("both replicas landed on %d node(s); they should spread", len(nodes))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPacksWhenSpreadingIsImpossible(t *testing.T) {
|
||||
// Spreading is a preference, not a rule: a second replica on one node still
|
||||
// survives an engine crash, which is the common failure.
|
||||
h := newHarness(t, testModel(2), gpuNode("only", 8))
|
||||
|
||||
for range 2 {
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 2 {
|
||||
t.Fatalf("expected 2 replicas on the single node, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShardOrderLeaderFirst(t *testing.T) {
|
||||
// Both engines put the rendezvous point on rank 0 — SGLang's dist-init
|
||||
// address and vLLM's Ray head — so a worker started first has nothing to
|
||||
// connect to and spends its startup budget retrying.
|
||||
//
|
||||
// This test previously asserted the opposite, on the theory that the head
|
||||
// should wait for its workers. That is backwards: vLLM's server blocks
|
||||
// until the Ray cluster has the GPUs it needs, so the leader can start
|
||||
// immediately and waits on its own.
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
started := h.disp.startedShards()
|
||||
if len(started) != 2 {
|
||||
t.Fatalf("expected 2 shards, got %d", len(started))
|
||||
}
|
||||
if started[0].Spec.Rank != 0 || started[1].Spec.Rank != 1 {
|
||||
t.Errorf("start order was rank %d then %d; the leader hosts the rendezvous "+
|
||||
"point and must start first",
|
||||
started[0].Spec.Rank, started[1].Spec.Rank)
|
||||
}
|
||||
// Every rank must be given the same rendezvous point.
|
||||
if started[0].Spec.LeaderAddr != started[1].Spec.LeaderAddr {
|
||||
t.Errorf("ranks got different leader addresses: %q and %q",
|
||||
started[0].Spec.LeaderAddr, started[1].Spec.LeaderAddr)
|
||||
}
|
||||
if started[0].Spec.LeaderPort != started[1].Spec.LeaderPort {
|
||||
t.Error("ranks got different leader ports")
|
||||
}
|
||||
for _, cmd := range started {
|
||||
if len(cmd.Spec.HCAs) == 0 {
|
||||
t.Errorf("rank %d got no HCAs; NCCL would choose its own", cmd.Spec.Rank)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailedShardFailsWholeReplica(t *testing.T) {
|
||||
// One dead shard makes the others useless: they wait on collectives that
|
||||
// will never complete while holding GPUs (plan §5.4).
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
started := h.disp.startedShards()
|
||||
|
||||
// One shard ready, the other dead.
|
||||
h.r.OnShardStatus(started[0].ReplicaID, started[0].ShardID, supervisor.Ready, "")
|
||||
h.r.OnShardStatus(started[1].ReplicaID, started[1].ShardID, supervisor.Failed, "CUDA out of memory")
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
// Both shards must have been stopped, not just the failed one.
|
||||
h.disp.mu.Lock()
|
||||
stopped := slices.Clone(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if len(stopped) != 2 {
|
||||
t.Errorf("stopped %d shards, expected both: %v", len(stopped), stopped)
|
||||
}
|
||||
// And the GPUs must be free again.
|
||||
for _, id := range []string{"1", "2"} {
|
||||
if got := len(h.inv.FreeGPUs(id)); got != 8 {
|
||||
t.Errorf("node %s has %d free GPUs; the replica should have been released", id, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackoffAfterFailure(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
started := h.disp.startedShards()
|
||||
h.r.OnShardStatus(started[0].ReplicaID, started[0].ShardID, supervisor.Failed, "boom")
|
||||
h.r.Reconcile(t.Context()) // reaps it
|
||||
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.disp.startedShards()); got != 0 {
|
||||
t.Errorf("a replica was rebuilt immediately after failure; backoff is not applied")
|
||||
}
|
||||
|
||||
// Past the backoff, it tries again.
|
||||
h.advance(DefaultConfig().MaxBackoff)
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.disp.startedShards()); got == 0 {
|
||||
t.Error("the replica was never rebuilt after the backoff expired")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSuccessClearsBackoff(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
// Fail once.
|
||||
h.r.Reconcile(t.Context())
|
||||
s := h.disp.startedShards()[0]
|
||||
h.r.OnShardStatus(s.ReplicaID, s.ShardID, supervisor.Failed, "boom")
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
// Then succeed.
|
||||
h.advance(DefaultConfig().MaxBackoff)
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
if len(h.disp.startedShards()) == 0 {
|
||||
t.Fatal("no replica was started after the backoff expired")
|
||||
}
|
||||
s2 := h.disp.startedShards()[0]
|
||||
h.r.OnShardStatus(s2.ReplicaID, s2.ShardID, supervisor.Ready, "")
|
||||
|
||||
// A later failure should start from a short backoff, not a long streak.
|
||||
h.r.OnShardStatus(s2.ReplicaID, s2.ShardID, supervisor.Failed, "again")
|
||||
h.r.Reconcile(t.Context())
|
||||
h.advance(2 * DefaultConfig().BaseBackoff)
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.disp.startedShards()); got == 0 {
|
||||
t.Error("backoff was not reset by the successful run in between")
|
||||
}
|
||||
}
|
||||
|
||||
// TestLostNodeMovesReplicas is the phase 3 exit criterion: a node going away
|
||||
// must get its replicas rebuilt elsewhere.
|
||||
func TestLostNodeMovesReplicas(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || reps[0].State != Ready {
|
||||
t.Fatalf("setup failed: %+v", reps)
|
||||
}
|
||||
original := reps[0].NodeIDs()[0]
|
||||
survivor := "2"
|
||||
if original == "2" {
|
||||
survivor = "1"
|
||||
}
|
||||
|
||||
// The surviving node keeps reporting; the other goes silent.
|
||||
h.clock.advance(20 * time.Second)
|
||||
touchAt(h.inv, survivor, h.clock.now())
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if n, _ := h.inv.Get(original); n.State != inventory.Unreachable {
|
||||
t.Errorf("silent node is %s, expected unreachable", n.State)
|
||||
}
|
||||
// Unreachable is not lost: a blip must not tear down a working replica.
|
||||
if reps := h.r.Replicas(); len(reps) != 1 || reps[0].State != Ready {
|
||||
t.Fatalf("an unreachable node already destroyed the replica: %+v", reps)
|
||||
}
|
||||
|
||||
h.clock.advance(90 * time.Second)
|
||||
touchAt(h.inv, survivor, h.clock.now())
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context()) // declares it lost and reaps the replica
|
||||
|
||||
if n, _ := h.inv.Get(original); n.State != inventory.Lost {
|
||||
t.Fatalf("node is %s, expected lost", n.State)
|
||||
}
|
||||
|
||||
// Past the backoff, the replica is rebuilt on the survivor. Only the
|
||||
// survivor heartbeats; the lost node stays gone.
|
||||
h.clock.advance(DefaultConfig().MaxBackoff)
|
||||
touchAt(h.inv, survivor, h.clock.now())
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
reps = h.r.Replicas()
|
||||
if len(reps) != 1 {
|
||||
t.Fatalf("expected the replica to be rebuilt, got %d", len(reps))
|
||||
}
|
||||
if got := reps[0].NodeIDs()[0]; got != survivor {
|
||||
t.Errorf("replica was rebuilt on %q, expected the surviving node %q", got, survivor)
|
||||
}
|
||||
if reps[0].State != Ready {
|
||||
t.Errorf("rebuilt replica is %s", reps[0].State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeartbeatRevivesNode(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.clock.advance(20 * time.Second) // deliberately without a heartbeat
|
||||
h.r.Reconcile(t.Context())
|
||||
if n, _ := h.inv.Get("1"); n.State != inventory.Unreachable {
|
||||
t.Fatalf("node is %s", n.State)
|
||||
}
|
||||
|
||||
h.inv.Touch("1")
|
||||
if n, _ := h.inv.Get("1"); n.State != inventory.Ready {
|
||||
t.Errorf("a heartbeat did not revive the node: %s", n.State)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScaleDownDrainsFirst(t *testing.T) {
|
||||
// Draining before stopping lets in-flight requests finish instead of being
|
||||
// cut off mid-stream.
|
||||
h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
|
||||
for range 2 {
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 2 {
|
||||
t.Fatalf("setup: %d replicas", got)
|
||||
}
|
||||
|
||||
// Shrink the catalog to one replica.
|
||||
cat, err := catalog.New([]catalog.Model{testModel(1)})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.r.cat = cat
|
||||
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.disp.mu.Lock()
|
||||
drained, stopped := len(h.disp.drained), len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if drained == 0 {
|
||||
t.Error("the excess replica was not drained")
|
||||
}
|
||||
if stopped != 0 {
|
||||
t.Error("a replica was stopped before it had drained")
|
||||
}
|
||||
|
||||
// The next pass stops what has drained.
|
||||
h.r.Reconcile(t.Context())
|
||||
h.disp.mu.Lock()
|
||||
stopped = len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if stopped == 0 {
|
||||
t.Error("the drained replica was never stopped")
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Errorf("expected 1 replica after scale-down, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnplaceableModelExplains(t *testing.T) {
|
||||
// A model that cannot be placed must leave a usable explanation rather than
|
||||
// a silent absence.
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "tp8", GPUClass: "h100-80g", TP: 8,
|
||||
}), gpuNode("1", 4))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.r.Replicas()); got != 0 {
|
||||
t.Fatalf("something was placed: %d", got)
|
||||
}
|
||||
|
||||
res, ok := h.r.Explain("llama")
|
||||
if !ok {
|
||||
t.Fatal("no explanation was recorded")
|
||||
}
|
||||
exp := res.Explain()
|
||||
if !strings.Contains(exp, "needs 8") {
|
||||
t.Errorf("explanation does not say what was missing:\n%s", exp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatcherFailureFailsReplica(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
h.disp.failNext = fmt.Errorf("agent unreachable")
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || reps[0].State != Failed {
|
||||
t.Fatalf("a failed dispatch did not fail the replica: %+v", reps)
|
||||
}
|
||||
|
||||
// And the next pass must release its GPUs.
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.inv.FreeGPUs("1")); got != 8 {
|
||||
t.Errorf("node has %d free GPUs; the failed replica still holds some", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadyReplicasFeedTheRouter(t *testing.T) {
|
||||
h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
if got := len(h.r.ReadyReplicas()); got != 0 {
|
||||
t.Errorf("a starting replica was offered to the router")
|
||||
}
|
||||
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
|
||||
ready := h.r.ReadyReplicas()
|
||||
if len(ready) != 2 {
|
||||
t.Fatalf("expected 2 ready replicas, got %d", len(ready))
|
||||
}
|
||||
for _, rep := range ready {
|
||||
if !strings.HasPrefix(rep.Endpoint, "http://") {
|
||||
t.Errorf("endpoint %q is not a URL", rep.Endpoint)
|
||||
}
|
||||
if rep.Weight <= 0 {
|
||||
t.Errorf("replica %s has no weight; the router needs one", rep.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPortsDoNotCollide(t *testing.T) {
|
||||
h := newHarness(t, testModel(3), gpuNode("1", 8), gpuNode("2", 8), gpuNode("3", 8))
|
||||
for range 3 {
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
}
|
||||
|
||||
ports := map[int]string{}
|
||||
for _, rep := range h.r.Replicas() {
|
||||
for _, s := range rep.Shards {
|
||||
if prev, dup := ports[s.HTTPPort]; dup && prev != rep.ID {
|
||||
t.Errorf("replicas %s and %s both use port %d", prev, rep.ID, s.HTTPPort)
|
||||
}
|
||||
ports[s.HTTPPort] = rep.ID
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReplicaStateStrings(t *testing.T) {
|
||||
for st, want := range map[ReplicaState]string{
|
||||
Placing: "placing", Starting: "starting", Ready: "ready",
|
||||
Draining: "draining", Stopped: "stopped", Failed: "failed",
|
||||
} {
|
||||
if got := st.String(); got != want {
|
||||
t.Errorf("%d -> %q", int(st), got)
|
||||
}
|
||||
}
|
||||
if !Ready.Counts() || Draining.Counts() || Failed.Counts() {
|
||||
// Draining must not count, or a replacement is never started for it.
|
||||
t.Error("Counts is wrong")
|
||||
}
|
||||
}
|
||||
|
||||
// --- adoption after a control plane restart (plan §6.5) ---
|
||||
|
||||
func adopted(replicaID, shardID, nodeID string, rank int, gpus []int) AdoptedShard {
|
||||
a := AdoptedShard{
|
||||
ReplicaID: replicaID, ShardID: shardID, NodeID: nodeID, Rank: rank,
|
||||
GPUIndices: gpus, HTTPPort: 20000, State: supervisor.Ready,
|
||||
Model: "llama", Profile: "tp4", Pool: "none/h100-80g",
|
||||
}
|
||||
if rank == 0 {
|
||||
a.Endpoint = "http://10.0.0." + nodeID + ":20000"
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
func TestAdoptRebuildsReplica(t *testing.T) {
|
||||
// The control plane restarted. Nothing is known, but a node reports a live
|
||||
// replica. Without this the reconciler would place a duplicate while the
|
||||
// original keeps holding GPUs.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
unclaimed := h.r.Adopt([]AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})})
|
||||
if len(unclaimed) != 0 {
|
||||
t.Fatalf("a valid shard was reported unclaimed: %+v", unclaimed)
|
||||
}
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 {
|
||||
t.Fatalf("expected 1 adopted replica, got %d", len(reps))
|
||||
}
|
||||
rep := reps[0]
|
||||
if rep.ID != "llama-99" || rep.Model != "llama" || rep.Profile != "tp4" {
|
||||
t.Errorf("identity wrong: %+v", rep)
|
||||
}
|
||||
if rep.State != Ready {
|
||||
t.Errorf("state=%s; every shard reported Ready", rep.State)
|
||||
}
|
||||
if rep.Endpoint == "" {
|
||||
t.Error("no endpoint; the router would have nowhere to send")
|
||||
}
|
||||
if rep.Weight <= 0 {
|
||||
t.Error("weight not taken from the catalog")
|
||||
}
|
||||
|
||||
// The GPUs must be reclaimed, or the next placement hands them out twice.
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 4 {
|
||||
t.Errorf("node has %d free GPUs, expected 4 after adoption", free)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptPreventsDuplicatePlacement(t *testing.T) {
|
||||
// The whole point: after adopting, a reconciliation pass must see the
|
||||
// desired count already met.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.r.Adopt([]AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})})
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.disp.startedShards()); got != 0 {
|
||||
t.Errorf("a duplicate replica was placed despite adoption (%d shards started)", got)
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Errorf("expected 1 replica, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptMultiNodeReplica(t *testing.T) {
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
shards := []AdoptedShard{
|
||||
adopted("llama-7", "llama-7-r0", "1", 0, []int{0, 1, 2, 3, 4, 5, 6, 7}),
|
||||
adopted("llama-7", "llama-7-r1", "2", 1, []int{0, 1, 2, 3, 4, 5, 6, 7}),
|
||||
}
|
||||
for i := range shards {
|
||||
shards[i].Profile = "2node"
|
||||
shards[i].Fabric = "ib-a"
|
||||
shards[i].Pool = "ib-a/h100-80g"
|
||||
}
|
||||
|
||||
if unclaimed := h.r.Adopt(shards); len(unclaimed) != 0 {
|
||||
t.Fatalf("unclaimed: %+v", unclaimed)
|
||||
}
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || len(reps[0].Shards) != 2 {
|
||||
t.Fatalf("the two shards were not grouped into one replica: %+v", reps)
|
||||
}
|
||||
if reps[0].Fabric != "ib-a" {
|
||||
t.Errorf("fabric lost: %q", reps[0].Fabric)
|
||||
}
|
||||
// Both nodes' GPUs must be reclaimed.
|
||||
for _, id := range []string{"1", "2"} {
|
||||
if free := len(h.inv.FreeGPUs(id)); free != 0 {
|
||||
t.Errorf("node %s has %d free GPUs after adoption", id, free)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptPartialReplicaFails(t *testing.T) {
|
||||
// A multi-node replica missing a shard is useless: the survivors wait on
|
||||
// collectives that will never complete while holding GPUs.
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
s := adopted("llama-7", "llama-7-r0", "1", 0, []int{0, 1, 2, 3, 4, 5, 6, 7})
|
||||
s.Profile = "2node"
|
||||
s.State = supervisor.Failed
|
||||
h.r.Adopt([]AdoptedShard{s})
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || reps[0].State != Failed {
|
||||
t.Fatalf("a failed adopted shard did not fail its replica: %+v", reps)
|
||||
}
|
||||
|
||||
// And the next pass releases its GPUs.
|
||||
h.r.Reconcile(t.Context())
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 8 {
|
||||
t.Errorf("node has %d free GPUs; the failed adopted replica still holds some", free)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptUnknownModelIsUnclaimed(t *testing.T) {
|
||||
// The model left the catalog while the engine kept running. Nothing will
|
||||
// ever route to it, and it is holding GPUs.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
s := adopted("ghost-1", "ghost-1-r0", "1", 0, []int{0})
|
||||
s.Model = "a-model-that-was-removed"
|
||||
|
||||
unclaimed := h.r.Adopt([]AdoptedShard{s})
|
||||
if len(unclaimed) != 1 || unclaimed[0].ShardID != "ghost-1-r0" {
|
||||
t.Fatalf("an orphaned shard was not reported unclaimed: %+v", unclaimed)
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 0 {
|
||||
t.Errorf("a replica was created for an unknown model: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptUnknownProfileIsUnclaimed(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
s := adopted("llama-9", "llama-9-r0", "1", 0, []int{0})
|
||||
s.Profile = "a-profile-that-was-removed"
|
||||
|
||||
if unclaimed := h.r.Adopt([]AdoptedShard{s}); len(unclaimed) != 1 {
|
||||
t.Errorf("a shard with an unknown profile was claimed: %+v", unclaimed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptIsIdempotent(t *testing.T) {
|
||||
// An agent reconnecting re-reports the same shards; that must not double
|
||||
// the allocation.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
s := []AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})}
|
||||
|
||||
h.r.Adopt(s)
|
||||
h.r.Adopt(s)
|
||||
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Errorf("re-reporting created %d replicas", got)
|
||||
}
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 4 {
|
||||
t.Errorf("node has %d free GPUs; the allocation was applied twice", free)
|
||||
}
|
||||
}
|
||||
|
||||
// --- draining a node (ctl drain) ---
|
||||
|
||||
func TestDrainingNodeMovesWorkOff(t *testing.T) {
|
||||
// The operator's view: a node is about to be worked on, its replicas must
|
||||
// end up elsewhere, and nothing may be cut off mid-request on the way.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8), gpuNode("2", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 {
|
||||
t.Fatalf("setup: %d replicas", len(reps))
|
||||
}
|
||||
victim := reps[0].NodeIDs()[0]
|
||||
survivor := "2"
|
||||
if victim == "2" {
|
||||
survivor = "1"
|
||||
}
|
||||
|
||||
h.inv.SetState(victim, inventory.Draining)
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
// It is drained, not stopped: the engine is healthy and still serving.
|
||||
h.disp.mu.Lock()
|
||||
drained, stopped := len(h.disp.drained), len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if drained == 0 {
|
||||
t.Error("the replica on the draining node was not drained")
|
||||
}
|
||||
if stopped != 0 {
|
||||
t.Error("the replica was stopped before a replacement existed; in-flight requests would be cut")
|
||||
}
|
||||
|
||||
// A draining replica does not count, so a replacement is placed elsewhere.
|
||||
if got := len(h.disp.startedShards()); got == 0 {
|
||||
t.Fatal("no replacement was placed")
|
||||
}
|
||||
for _, cmd := range h.disp.startedShards() {
|
||||
if cmd.NodeID == victim {
|
||||
t.Errorf("the replacement was placed back on the draining node %s", victim)
|
||||
}
|
||||
}
|
||||
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
// Now that the replacement serves, the original can go.
|
||||
h.disp.mu.Lock()
|
||||
stopped = len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if stopped == 0 {
|
||||
t.Error("the drained replica was never stopped once its replacement was ready")
|
||||
}
|
||||
|
||||
for _, rep := range h.r.Replicas() {
|
||||
if rep.OnNode(victim) {
|
||||
t.Errorf("replica %s is still on the drained node", rep.ID)
|
||||
}
|
||||
if !slices.Contains(rep.NodeIDs(), survivor) {
|
||||
t.Errorf("replica %s did not move to the surviving node", rep.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDrainGraceStopsWithoutReplacement(t *testing.T) {
|
||||
// Draining the only node that can host a model would otherwise never
|
||||
// finish, leaving the node undrainable.
|
||||
h := newHarness(t, testModel(1), gpuNode("only", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.inv.SetState("only", inventory.Draining)
|
||||
h.r.Reconcile(t.Context()) // drains it; nowhere to place a replacement
|
||||
|
||||
if got := h.r.Replicas(); len(got) != 1 || got[0].State != Draining {
|
||||
t.Fatalf("expected a draining replica, got %+v", got)
|
||||
}
|
||||
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
h.disp.mu.Lock()
|
||||
early := len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if early != 0 {
|
||||
t.Error("the replica was stopped before the grace period expired")
|
||||
}
|
||||
|
||||
h.advance(DefaultConfig().DrainGrace + time.Second)
|
||||
h.inv.SetState("only", inventory.Draining) // advance() resets node state
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.disp.mu.Lock()
|
||||
stopped := len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if stopped == 0 {
|
||||
t.Error("the drained replica was never stopped; the node cannot be drained")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCordonedNodeKeepsItsWork(t *testing.T) {
|
||||
// Cordon is not drain: it stops new placement and leaves what is running.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
h.markReady()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.inv.SetState("1", inventory.Cordoned)
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
h.disp.mu.Lock()
|
||||
drained, stopped := len(h.disp.drained), len(h.disp.stopped)
|
||||
h.disp.mu.Unlock()
|
||||
if drained != 0 || stopped != 0 {
|
||||
t.Errorf("cordoning disturbed a running replica (drained=%d stopped=%d)", drained, stopped)
|
||||
}
|
||||
if got := h.r.Replicas(); len(got) != 1 || got[0].State != Ready {
|
||||
t.Errorf("the replica did not survive a cordon: %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
// Package reconciler drives the cluster toward the catalog's desired state
|
||||
// (plan §6.3).
|
||||
package reconciler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
)
|
||||
|
||||
// ReplicaState is where a replica stands.
|
||||
type ReplicaState int
|
||||
|
||||
const (
|
||||
// Placing means it has been scheduled but no shard has started.
|
||||
Placing ReplicaState = iota
|
||||
// Starting means shards are launching or loading.
|
||||
Starting
|
||||
// Ready means it serves traffic.
|
||||
Ready
|
||||
// Draining means it takes no new requests and will be stopped.
|
||||
Draining
|
||||
// Stopped means it was shut down on purpose.
|
||||
Stopped
|
||||
// Failed means it died or never became ready.
|
||||
Failed
|
||||
)
|
||||
|
||||
func (s ReplicaState) String() string {
|
||||
switch s {
|
||||
case Placing:
|
||||
return "placing"
|
||||
case Starting:
|
||||
return "starting"
|
||||
case Ready:
|
||||
return "ready"
|
||||
case Draining:
|
||||
return "draining"
|
||||
case Stopped:
|
||||
return "stopped"
|
||||
case Failed:
|
||||
return "failed"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// Terminal reports whether the replica has finished for good.
|
||||
func (s ReplicaState) Terminal() bool { return s == Stopped || s == Failed }
|
||||
|
||||
// Counts reports whether the replica should count toward the desired total.
|
||||
// Draining replicas do not: they are on the way out, and counting them would
|
||||
// stop a replacement from being started.
|
||||
func (s ReplicaState) Counts() bool { return s == Placing || s == Starting || s == Ready }
|
||||
|
||||
// Shard is one node's part of a replica.
|
||||
type Shard struct {
|
||||
ID string
|
||||
NodeID string
|
||||
Rank int
|
||||
GPUIndices []int
|
||||
HCAs []string
|
||||
HTTPPort int
|
||||
State supervisor.State
|
||||
Err string
|
||||
}
|
||||
|
||||
// Leader reports whether this shard serves HTTP for the replica.
|
||||
func (s Shard) Leader() bool { return s.Rank == 0 }
|
||||
|
||||
// Replica is one running copy of a model. It is the unit the router sees.
|
||||
type Replica struct {
|
||||
ID string
|
||||
Model string
|
||||
Profile string
|
||||
// Fabric is set only for multi-node replicas.
|
||||
Fabric string
|
||||
Pool string
|
||||
Shards []Shard
|
||||
State ReplicaState
|
||||
// Endpoint is rank 0's base URL.
|
||||
Endpoint string
|
||||
// Weight is the profile's capacity weight, passed to the router.
|
||||
Weight float64
|
||||
|
||||
CreatedAt time.Time
|
||||
ReadyAt time.Time
|
||||
// DrainingAt is when the replica stopped taking new work, so the wait for a
|
||||
// replacement can be bounded.
|
||||
DrainingAt time.Time
|
||||
// Attempt counts how many times this model has failed to come up in a row,
|
||||
// carried across rebuilds so the backoff does not reset on every attempt.
|
||||
Attempt int
|
||||
Err string
|
||||
}
|
||||
|
||||
// NodeIDs lists the nodes this replica occupies.
|
||||
func (r Replica) NodeIDs() []string {
|
||||
out := make([]string, 0, len(r.Shards))
|
||||
for _, s := range r.Shards {
|
||||
out = append(out, s.NodeID)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// OnNode reports whether any shard runs on a node.
|
||||
func (r Replica) OnNode(id string) bool {
|
||||
return slices.ContainsFunc(r.Shards, func(s Shard) bool { return s.NodeID == id })
|
||||
}
|
||||
|
||||
// deriveState recomputes the replica's state from its shards.
|
||||
//
|
||||
// A replica is only as healthy as its worst shard: one failed shard makes the
|
||||
// whole thing useless, because the others are waiting on collectives that will
|
||||
// never complete. That is why there is no path here that keeps a replica
|
||||
// partially alive (plan §5.4).
|
||||
func (r *Replica) deriveState() {
|
||||
if r.State == Draining || r.State == Stopped {
|
||||
return
|
||||
}
|
||||
if len(r.Shards) == 0 {
|
||||
r.State = Placing
|
||||
return
|
||||
}
|
||||
|
||||
ready := 0
|
||||
for _, s := range r.Shards {
|
||||
switch s.State {
|
||||
case supervisor.Failed:
|
||||
r.State = Failed
|
||||
if r.Err == "" {
|
||||
r.Err = fmt.Sprintf("shard %s failed: %s", s.ID, s.Err)
|
||||
}
|
||||
return
|
||||
case supervisor.Stopped:
|
||||
// A shard stopping on its own, while the replica was not draining,
|
||||
// is a failure of the replica.
|
||||
r.State = Failed
|
||||
if r.Err == "" {
|
||||
r.Err = fmt.Sprintf("shard %s stopped unexpectedly", s.ID)
|
||||
}
|
||||
return
|
||||
case supervisor.Ready:
|
||||
ready++
|
||||
}
|
||||
}
|
||||
|
||||
if ready == len(r.Shards) {
|
||||
if r.State != Ready {
|
||||
r.ReadyAt = time.Now()
|
||||
}
|
||||
r.State = Ready
|
||||
return
|
||||
}
|
||||
r.State = Starting
|
||||
}
|
||||
@@ -12,8 +12,9 @@ import (
|
||||
// AdminHandler returns the diagnostics surface: metrics, health and pprof.
|
||||
//
|
||||
// It is served on a separate listener so the client-facing port can be exposed
|
||||
// without also exposing pprof (plan §9.1).
|
||||
func (r *Router) AdminHandler() (http.Handler, error) {
|
||||
// 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(),
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"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"
|
||||
@@ -39,7 +40,7 @@ type Router struct {
|
||||
profiles map[string]*scheduler.Profile
|
||||
// models maps every name a client may use — canonical names and aliases —
|
||||
// onto its catalog entry.
|
||||
models map[string]config.Model
|
||||
models map[string]catalog.Model
|
||||
}
|
||||
|
||||
// New builds a router from configuration. Start must be called to begin scraping.
|
||||
@@ -71,7 +72,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
||||
return nil, fmt.Errorf("unknown default_profile %q", cfg.Router.DefaultProfile)
|
||||
}
|
||||
|
||||
models := map[string]config.Model{}
|
||||
models := map[string]catalog.Model{}
|
||||
for _, m := range cfg.Models {
|
||||
models[m.Name] = m
|
||||
for _, a := range m.Aliases {
|
||||
@@ -195,19 +196,31 @@ func (r *Router) Stop() {
|
||||
// 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 config.Model) {
|
||||
func (r *Router) resolve(name string) (canonical string, m catalog.Model) {
|
||||
if e, ok := r.models[name]; ok {
|
||||
return e.Name, e
|
||||
}
|
||||
return name, config.Model{Name: name}
|
||||
return name, catalog.Model{Name: name}
|
||||
}
|
||||
|
||||
func (r *Router) profileFor(m config.Model, override string) *scheduler.Profile {
|
||||
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
|
||||
@@ -267,7 +280,7 @@ func (r *Router) dispatch(ctx context.Context, w http.ResponseWriter, req *parse
|
||||
|
||||
var errNoCapacity = errors.New("no endpoint available")
|
||||
|
||||
func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *parsedRequest, model string, entry config.Model) dispatchResult {
|
||||
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)}
|
||||
|
||||
@@ -14,8 +14,18 @@ import (
|
||||
"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 {
|
||||
@@ -163,7 +173,7 @@ func TestFallbackModel(t *testing.T) {
|
||||
config.Endpoint{ID: "p", URL: primary.srv.URL, Models: []string{"primary"}},
|
||||
config.Endpoint{ID: "b", URL: backup.srv.URL, Models: []string{"backup"}},
|
||||
)
|
||||
cfg.Models = []config.Model{{Name: "primary", Fallbacks: []string{"backup"}}}
|
||||
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)
|
||||
@@ -184,7 +194,7 @@ func TestAliasResolution(t *testing.T) {
|
||||
cfg := baseConfig(config.Endpoint{
|
||||
ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"},
|
||||
})
|
||||
cfg.Models = []config.Model{{Name: "llama-3.3-70b", Aliases: []string{"llama"}}}
|
||||
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)
|
||||
@@ -223,7 +233,7 @@ func TestMissingModelIs400(t *testing.T) {
|
||||
func TestModelsEndpointListsAliases(t *testing.T) {
|
||||
e := newFakeEngine(t, "llama-3.3-70b")
|
||||
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"}})
|
||||
cfg.Models = []config.Model{{Name: "llama-3.3-70b", Aliases: []string{"llama"}}}
|
||||
cfg.Models = []catalog.Model{catModel("llama-3.3-70b", []string{"llama"})}
|
||||
rt := newTestRouter(t, cfg)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
|
||||
+337
-22
@@ -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)
|
||||
@@ -197,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)
|
||||
@@ -468,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.
|
||||
|
||||
---
|
||||
|
||||
@@ -791,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
|
||||
@@ -993,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
|
||||
@@ -1080,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). |
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -0,0 +1,414 @@
|
||||
package e2e
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/sim"
|
||||
)
|
||||
|
||||
// simDispatcher is a stand-in for the agent fleet: instead of sending commands
|
||||
// over a stream it starts real sim-engines on this host, on exactly the ports
|
||||
// the reconciler allocated.
|
||||
//
|
||||
// That matters — the reconciler derives each replica's endpoint URL from the
|
||||
// node address and the port it chose, so an engine listening anywhere else
|
||||
// would make the test pass for the wrong reason.
|
||||
type simDispatcher struct {
|
||||
t *testing.T
|
||||
|
||||
mu sync.Mutex
|
||||
running map[string]*simShard
|
||||
// dead marks nodes whose agent is gone; commands to them fail.
|
||||
dead map[string]bool
|
||||
|
||||
onStatus func(replicaID, shardID string, st supervisor.State, errMsg string)
|
||||
}
|
||||
|
||||
type simShard struct {
|
||||
eng *sim.Engine
|
||||
srv *http.Server
|
||||
port int
|
||||
}
|
||||
|
||||
func newSimDispatcher(t *testing.T) *simDispatcher {
|
||||
d := &simDispatcher{t: t, running: map[string]*simShard{}, dead: map[string]bool{}}
|
||||
t.Cleanup(d.shutdown)
|
||||
return d
|
||||
}
|
||||
|
||||
func (d *simDispatcher) StartShard(_ context.Context, cmd reconciler.ShardCommand) error {
|
||||
d.mu.Lock()
|
||||
if d.dead[cmd.NodeID] {
|
||||
d.mu.Unlock()
|
||||
return fmt.Errorf("agent on %s is unreachable", cmd.NodeID)
|
||||
}
|
||||
d.mu.Unlock()
|
||||
|
||||
cfg := simConfig("", int64(cmd.Spec.HTTPPort))
|
||||
cfg.Model = cmd.Spec.Profile.ServedModelName
|
||||
eng := sim.New(cfg)
|
||||
|
||||
ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.Spec.HTTPPort))
|
||||
if err != nil {
|
||||
eng.Close()
|
||||
return fmt.Errorf("binding port %d: %w", cmd.Spec.HTTPPort, err)
|
||||
}
|
||||
srv := &http.Server{Handler: sim.NewServer(eng, slog.New(slog.DiscardHandler)), ReadHeaderTimeout: time.Second}
|
||||
go func() { _ = srv.Serve(ln) }()
|
||||
|
||||
d.mu.Lock()
|
||||
d.running[cmd.ShardID] = &simShard{eng: eng, srv: srv, port: cmd.Spec.HTTPPort}
|
||||
d.mu.Unlock()
|
||||
|
||||
// A real agent reports Ready only once the engine serves the model; the
|
||||
// sim-engine is ready as soon as it is listening.
|
||||
go d.onStatus(cmd.ReplicaID, cmd.ShardID, supervisor.Ready, "")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *simDispatcher) StopShard(_ context.Context, _, shardID string) error {
|
||||
d.mu.Lock()
|
||||
sh, ok := d.running[shardID]
|
||||
delete(d.running, shardID)
|
||||
d.mu.Unlock()
|
||||
if ok {
|
||||
sh.close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *simDispatcher) DrainShard(context.Context, string, string) error { return nil }
|
||||
|
||||
// killNode stops everything on a node and makes further commands to it fail,
|
||||
// which is what losing a machine looks like from the control plane.
|
||||
func (d *simDispatcher) killNode(nodeID string, shardIDs []string) {
|
||||
d.mu.Lock()
|
||||
d.dead[nodeID] = true
|
||||
var doomed []*simShard
|
||||
for _, id := range shardIDs {
|
||||
if sh, ok := d.running[id]; ok {
|
||||
doomed = append(doomed, sh)
|
||||
delete(d.running, id)
|
||||
}
|
||||
}
|
||||
d.mu.Unlock()
|
||||
|
||||
for _, sh := range doomed {
|
||||
sh.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (d *simDispatcher) liveShards() int {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return len(d.running)
|
||||
}
|
||||
|
||||
func (d *simDispatcher) shutdown() {
|
||||
d.mu.Lock()
|
||||
all := make([]*simShard, 0, len(d.running))
|
||||
for _, sh := range d.running {
|
||||
all = append(all, sh)
|
||||
}
|
||||
d.running = map[string]*simShard{}
|
||||
d.mu.Unlock()
|
||||
for _, sh := range all {
|
||||
sh.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *simShard) close() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
_ = s.srv.Shutdown(ctx)
|
||||
s.eng.Close()
|
||||
}
|
||||
|
||||
// controlHarness is a router fed by a real control plane.
|
||||
type controlHarness struct {
|
||||
router *httptest.Server
|
||||
rt *router.Router
|
||||
ctl *control.Controller
|
||||
disp *simDispatcher
|
||||
inv *inventory.Store
|
||||
}
|
||||
|
||||
func startControlPlane(t *testing.T, nodes int, minReplicas int) *controlHarness {
|
||||
t.Helper()
|
||||
|
||||
cat, err := catalog.New([]catalog.Model{{
|
||||
Name: testModel, Source: "/models/test",
|
||||
Scaling: catalog.Scaling{Min: minReplicas, Max: minReplicas + 2},
|
||||
Profiles: []catalog.Profile{{Name: "single", GPUClass: "h100-80g", TP: 1, Weight: 1}},
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cfg := config.Default()
|
||||
cfg.Router.ScrapeInterval = 200 * time.Millisecond
|
||||
cfg.Router.StaleAfter = 3 * time.Second
|
||||
cfg.Router.Admission.Enabled = false // this test is about endpoint churn
|
||||
|
||||
rt, err := router.New(cfg, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
disp := newSimDispatcher(t)
|
||||
|
||||
ccfg := control.DefaultConfig()
|
||||
ccfg.Reconciler.Interval = 100 * time.Millisecond
|
||||
ccfg.Reconciler.BaseBackoff = 200 * time.Millisecond
|
||||
ccfg.Reconciler.MaxBackoff = 500 * time.Millisecond
|
||||
ccfg.Reconciler.UnreachableAfter = 300 * time.Millisecond
|
||||
ccfg.Reconciler.LostAfter = 600 * time.Millisecond
|
||||
ccfg.PublishInterval = 100 * time.Millisecond
|
||||
|
||||
ctl := control.New(ccfg, cat, disp, rt, slog.New(slog.DiscardHandler))
|
||||
disp.onStatus = ctl.Reconciler().OnShardStatus
|
||||
|
||||
// Every node is this host, so the ports the reconciler allocates are the
|
||||
// ports the sim-engines actually bind.
|
||||
for i := range nodes {
|
||||
id := fmt.Sprintf("n%d", i)
|
||||
ctl.Inventory().Upsert(inventory.Node{
|
||||
ID: id, MgmtAddr: "127.0.0.1", State: inventory.Ready,
|
||||
GPUs: []discovery.GPU{{
|
||||
Index: 0, Name: "NVIDIA H100 80GB HBM3",
|
||||
MemoryTotalMB: 81559, ComputeCapability: "9.0",
|
||||
}},
|
||||
})
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
rt.Start(ctx)
|
||||
ctl.Start(ctx)
|
||||
|
||||
h := &controlHarness{
|
||||
router: httptest.NewServer(rt.Handler()),
|
||||
rt: rt, ctl: ctl, disp: disp, inv: ctl.Inventory(),
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
h.router.Close()
|
||||
ctl.Close()
|
||||
cancel()
|
||||
rt.Stop()
|
||||
})
|
||||
return h
|
||||
}
|
||||
|
||||
// heartbeat keeps the named nodes alive.
|
||||
func (h *controlHarness) heartbeat(ids ...string) {
|
||||
for _, id := range ids {
|
||||
h.inv.Touch(id)
|
||||
}
|
||||
}
|
||||
|
||||
// waitForEndpoints blocks until the router sees n endpoints.
|
||||
func (h *controlHarness) waitForEndpoints(t *testing.T, n int, limit time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(limit)
|
||||
for time.Now().Before(deadline) {
|
||||
if len(h.rt.Registry().Snapshots(testModel)) == n {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("router has %d endpoints, expected %d",
|
||||
len(h.rt.Registry().Snapshots(testModel)), n)
|
||||
}
|
||||
|
||||
// TestControlPlaneFeedsRouter checks the wiring: replicas the reconciler brings
|
||||
// up become endpoints the router can route to, without any static config.
|
||||
func TestControlPlaneFeedsRouter(t *testing.T) {
|
||||
h := startControlPlane(t, 2, 2)
|
||||
|
||||
// Keep the nodes alive while the reconciler works.
|
||||
stop := make(chan struct{})
|
||||
go func() {
|
||||
tk := time.NewTicker(50 * time.Millisecond)
|
||||
defer tk.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-tk.C:
|
||||
h.heartbeat("n0", "n1")
|
||||
}
|
||||
}
|
||||
}()
|
||||
defer close(stop)
|
||||
|
||||
h.waitForEndpoints(t, 2, 15*time.Second)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if _, _, err := streamRequest(ctx, h.router.URL, chatBody("hello", 4, true), nil); err != nil {
|
||||
t.Fatalf("the router could not serve from control-plane endpoints: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRouterKeepsServingWhileNodeIsLost is the other half of the phase 3 exit
|
||||
// criterion: a node disappears, its replica is rebuilt elsewhere, and traffic
|
||||
// keeps flowing throughout.
|
||||
func TestRouterKeepsServingWhileNodeIsLost(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("timing-dependent")
|
||||
}
|
||||
// Three nodes, two replicas: losing one still leaves somewhere to rebuild.
|
||||
h := startControlPlane(t, 3, 2)
|
||||
|
||||
alive := map[string]bool{"n0": true, "n1": true, "n2": true}
|
||||
var aliveMu sync.Mutex
|
||||
|
||||
stop := make(chan struct{})
|
||||
go func() {
|
||||
tk := time.NewTicker(50 * time.Millisecond)
|
||||
defer tk.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-tk.C:
|
||||
aliveMu.Lock()
|
||||
for id, ok := range alive {
|
||||
if ok {
|
||||
h.inv.Touch(id)
|
||||
}
|
||||
}
|
||||
aliveMu.Unlock()
|
||||
}
|
||||
}
|
||||
}()
|
||||
defer close(stop)
|
||||
|
||||
h.waitForEndpoints(t, 2, 15*time.Second)
|
||||
|
||||
// Continuous traffic for the whole experiment.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
var ok, failed atomic.Int64
|
||||
var firstErr atomic.Value
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for range 4 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for ctx.Err() == nil {
|
||||
rctx, rcancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
_, _, err := streamRequest(rctx, h.router.URL, chatBody("ping", 4, true), nil)
|
||||
rcancel()
|
||||
switch {
|
||||
case err == nil:
|
||||
ok.Add(1)
|
||||
case ctx.Err() != nil: // shutting down
|
||||
default:
|
||||
failed.Add(1)
|
||||
firstErr.CompareAndSwap(nil, err)
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Let traffic settle before disturbing anything.
|
||||
time.Sleep(600 * time.Millisecond)
|
||||
baseline := ok.Load()
|
||||
if baseline == 0 {
|
||||
cancel()
|
||||
wg.Wait()
|
||||
t.Fatalf("no requests succeeded before the failure: %v", firstErr.Load())
|
||||
}
|
||||
beforeFailures := failed.Load()
|
||||
|
||||
// Find a node hosting a replica and kill it.
|
||||
var victim string
|
||||
var victimShards []string
|
||||
for _, rep := range h.ctl.Reconciler().Replicas() {
|
||||
if rep.State != reconciler.Ready {
|
||||
continue
|
||||
}
|
||||
victim = rep.Shards[0].NodeID
|
||||
for _, s := range rep.Shards {
|
||||
victimShards = append(victimShards, s.ID)
|
||||
}
|
||||
break
|
||||
}
|
||||
if victim == "" {
|
||||
cancel()
|
||||
wg.Wait()
|
||||
t.Fatal("no ready replica to disturb")
|
||||
}
|
||||
|
||||
aliveMu.Lock()
|
||||
alive[victim] = false
|
||||
aliveMu.Unlock()
|
||||
h.disp.killNode(victim, victimShards)
|
||||
t.Logf("killed node %s (shards %v)", victim, victimShards)
|
||||
|
||||
// The control plane should notice, rebuild elsewhere, and get back to two
|
||||
// endpoints. The generous window covers unreachable -> lost -> backoff.
|
||||
deadline := time.Now().Add(25 * time.Second)
|
||||
recovered := false
|
||||
for time.Now().Before(deadline) {
|
||||
if len(h.rt.Registry().Snapshots(testModel)) == 2 && h.disp.liveShards() == 2 {
|
||||
recovered = true
|
||||
break
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
|
||||
// Let traffic run a little longer on the rebuilt topology.
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
cancel()
|
||||
wg.Wait()
|
||||
|
||||
t.Logf("requests ok=%d failed=%d (failures before the kill: %d)",
|
||||
ok.Load(), failed.Load(), beforeFailures)
|
||||
|
||||
if !recovered {
|
||||
var placed []string
|
||||
for _, rep := range h.ctl.Reconciler().Replicas() {
|
||||
placed = append(placed, fmt.Sprintf("%s=%s@%v", rep.ID, rep.State, rep.NodeIDs()))
|
||||
}
|
||||
t.Fatalf("the cluster did not return to 2 replicas: %v", placed)
|
||||
}
|
||||
|
||||
// The victim must not be hosting anything any more.
|
||||
for _, rep := range h.ctl.Reconciler().Replicas() {
|
||||
if rep.OnNode(victim) {
|
||||
t.Errorf("replica %s is still placed on the lost node %s", rep.ID, victim)
|
||||
}
|
||||
}
|
||||
|
||||
// Traffic kept flowing: more succeeded after the failure than before it.
|
||||
if ok.Load() <= baseline {
|
||||
t.Errorf("no requests succeeded after the node was lost (%d before, %d total)",
|
||||
baseline, ok.Load())
|
||||
}
|
||||
// Some failures are expected in the window where a replica is gone and its
|
||||
// replacement has not finished loading. What must not happen is total loss
|
||||
// of service, which is what having a second replica is for.
|
||||
if failed.Load() > ok.Load() {
|
||||
t.Errorf("more requests failed (%d) than succeeded (%d); service was not maintained: %v",
|
||||
failed.Load(), ok.Load(), firstErr.Load())
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user