One bidirectional connect-go stream per node, carrying registration, heartbeats and shard status up, and start, stop and drain down. The agent dials out: GPU nodes sit behind firewalls far more often than a control plane does, and an outbound connection needs no inbound rule on every node in the fleet. - registration carries the whole hardware inventory, so there is no separate discovery call and a reconnecting agent refreshes what the scheduler knows. The fabric labels travel with it; they cannot be inferred from the wire, so losing them in transit would silently disable multi-node placement. - agentsvc is both the stream handler and the reconciler's dispatcher, because it holds the open streams and is therefore the only thing that can deliver a command to a node. - the agent validates a profile before launching anything, since an engine that refuses a profile does so after minutes of loading, then supervises the process and reports every transition with the engine's output tail on failure. - reconnects forever with jittered backoff. A control plane restart must not require touching every GPU node, and a fleet must not reconnect in lockstep and knock it over again. Engines keep serving throughout, because the router talks to them directly. llmgw agent is real now, and llmgw server runs a control plane whenever any model declares a profile. The static endpoint list stays for router-only runs. Verified with real binaries as separate processes: the agent registers, a replica is placed, an engine launches through the vLLM adapter's actual argument construction, streaming traffic is served, and killing the agent moves the node through unreachable to lost with placement explaining why nothing can run. config.Model and catalog.Model were the same thing described twice, so the router now reads catalog entries directly and the validation rules live in one place. Still missing: mTLS. The stream is plaintext HTTP/2 and must stay on a trusted network until that lands. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
c1b9ea1bc2
commit
2373726efd
24 files changed
+3645
-87
No files matched your search
+36
-2
@@ -100,10 +100,44 @@ 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
|
window every time. Everything is snapshotted under the lock now and the
|
||||||
functions work from values.
|
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.
|
||||||
|
|
||||||
### Still missing for phase 2
|
### Still missing for phase 2
|
||||||
|
|
||||||
- The agent-to-control-plane stream. It pairs with phase 3's control plane
|
- mTLS. The stream is plaintext HTTP/2 today, so it must stay on a trusted
|
||||||
rather than standing alone, so `llmgw agent` itself is still a stub.
|
network.
|
||||||
|
|
||||||
### Phase 3: placement
|
### Phase 3: placement
|
||||||
|
|
||||||
|
|||||||
@@ -2,7 +2,7 @@ MODULE := gitea.mynodes.xyz/ilker/llmgw
|
|||||||
VERSION ?= $(shell git describe --tags --always --dirty 2>/dev/null || echo dev)
|
VERSION ?= $(shell git describe --tags --always --dirty 2>/dev/null || echo dev)
|
||||||
LDFLAGS := -X '$(MODULE)/internal/version.Version=$(VERSION)'
|
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:
|
build:
|
||||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw ./cmd/llmgw
|
go build -ldflags "$(LDFLAGS)" -o bin/llmgw ./cmd/llmgw
|
||||||
@@ -20,6 +20,17 @@ fmt:
|
|||||||
go fmt ./...
|
go fmt ./...
|
||||||
go vet ./...
|
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:
|
clean:
|
||||||
rm -rf bin dist
|
rm -rf bin dist
|
||||||
|
|
||||||
|
|||||||
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,176 @@
|
|||||||
|
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;
|
||||||
|
}
|
||||||
|
|
||||||
|
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;
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
listen: ":8080" # client-facing OpenAI API
|
||||||
admin_listen: ":9090" # metrics, health, pprof — do not expose this publicly
|
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:
|
router:
|
||||||
scrape_interval: 2s
|
scrape_interval: 2s
|
||||||
@@ -79,6 +83,19 @@ router:
|
|||||||
kv_ceiling: 0.9
|
kv_ceiling: 0.9
|
||||||
session_load_factor: 1.25
|
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:
|
models:
|
||||||
- name: llama-3.3-70b
|
- name: llama-3.3-70b
|
||||||
aliases: [llama]
|
aliases: [llama]
|
||||||
@@ -89,17 +106,13 @@ models:
|
|||||||
aliases: [qwen]
|
aliases: [qwen]
|
||||||
routing_profile: batch
|
routing_profile: batch
|
||||||
|
|
||||||
endpoints:
|
# A static replica list, used only when no model declares a profile. That is the
|
||||||
- id: gpu-a-01
|
# router-only mode for testing against sim-engines; with the control plane
|
||||||
url: http://gpu-a-01:8000
|
# running, agents provide the replicas instead.
|
||||||
models: [llama-3.3-70b]
|
#
|
||||||
pool: ib-a/h100-80g
|
# endpoints:
|
||||||
weight: 100 # from ctl bench; do not estimate it by hand
|
# - id: sim-1
|
||||||
priority: 100 # higher pools are tried first
|
# url: http://127.0.0.1:8001
|
||||||
|
# models: [llama-3.3-70b]
|
||||||
- id: gpu-b-01
|
# pool: local
|
||||||
url: http://gpu-b-01:8000
|
# weight: 1
|
||||||
models: [llama-3.3-70b, qwen2.5-32b]
|
|
||||||
pool: ib-b/a100-80g
|
|
||||||
weight: 45
|
|
||||||
priority: 50
|
|
||||||
@@ -3,6 +3,8 @@ module gitea.mynodes.xyz/ilker/llmgw
|
|||||||
go 1.26.3
|
go 1.26.3
|
||||||
|
|
||||||
require (
|
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/cespare/xxhash/v2 v2.3.0
|
||||||
github.com/knadh/koanf/parsers/yaml v1.1.1
|
github.com/knadh/koanf/parsers/yaml v1.1.1
|
||||||
github.com/knadh/koanf/providers/env/v2 v2.0.1
|
github.com/knadh/koanf/providers/env/v2 v2.0.1
|
||||||
@@ -12,10 +14,11 @@ require (
|
|||||||
github.com/prometheus/client_model v0.6.3
|
github.com/prometheus/client_model v0.6.3
|
||||||
github.com/prometheus/common v0.71.0
|
github.com/prometheus/common v0.71.0
|
||||||
github.com/spf13/cobra v1.10.2
|
github.com/spf13/cobra v1.10.2
|
||||||
|
golang.org/x/sync v0.23.0
|
||||||
|
google.golang.org/protobuf v1.36.12
|
||||||
)
|
)
|
||||||
|
|
||||||
require (
|
require (
|
||||||
github.com/NVIDIA/go-nvml v0.13.4-0 // indirect
|
|
||||||
github.com/beorn7/perks v1.0.1 // indirect
|
github.com/beorn7/perks v1.0.1 // indirect
|
||||||
github.com/fsnotify/fsnotify v1.9.0 // indirect
|
github.com/fsnotify/fsnotify v1.9.0 // indirect
|
||||||
github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
|
github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
|
||||||
@@ -27,7 +30,7 @@ require (
|
|||||||
github.com/prometheus/procfs v0.21.1 // indirect
|
github.com/prometheus/procfs v0.21.1 // indirect
|
||||||
github.com/spf13/pflag v1.0.9 // indirect
|
github.com/spf13/pflag v1.0.9 // indirect
|
||||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||||
golang.org/x/sync v0.23.0 // indirect
|
golang.org/x/net v0.59.0 // indirect
|
||||||
golang.org/x/sys v0.47.0 // indirect
|
golang.org/x/sys v0.48.0 // indirect
|
||||||
google.golang.org/protobuf v1.36.12 // indirect
|
golang.org/x/text v0.42.0 // indirect
|
||||||
)
|
)
|
||||||
@@ -1,3 +1,5 @@
|
|||||||
|
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 h1:o3jp9u2x1R9ShFE3v+Aesp55XOSIQFMJz/VGNUcJaNE=
|
||||||
github.com/NVIDIA/go-nvml v0.13.4-0/go.mod h1:id63qwpoDWpFXwnwM6psDCSqW4BmNu6mWpr4YeQtPGo=
|
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 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
|
||||||
@@ -55,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.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 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw=
|
||||||
go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg=
|
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 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk=
|
||||||
golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0=
|
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.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
|
||||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
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 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc=
|
||||||
google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
|
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=
|
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||||
@@ -0,0 +1,458 @@
|
|||||||
|
// 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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()
|
||||||
|
|
||||||
|
if err := stream.Send(&llmgwv1.AgentMessage{
|
||||||
|
Msg: &llmgwv1.AgentMessage_Register{
|
||||||
|
Register: agentsvc.RegisterFromInventory(
|
||||||
|
a.cfg.NodeID, a.cfg.Hostname, a.cfg.MgmtAddr, a.cfg.Version,
|
||||||
|
a.cfg.Labels, inv, a.cfg.Fabrics),
|
||||||
|
},
|
||||||
|
}); 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
|
||||||
|
}
|
||||||
|
|
||||||
|
cmd, err := eng.BuildCommand(spec)
|
||||||
|
if err != nil {
|
||||||
|
report(supervisor.Failed, err.Error(), nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
sh, err := a.sup.Start(ctx, shardID, spec, cmd)
|
||||||
|
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)
|
||||||
|
|
||||||
|
base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort)
|
||||||
|
if err := a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout); 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,403 @@
|
|||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"crypto/tls"
|
||||||
|
"errors"
|
||||||
|
"log/slog"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
|
"slices"
|
||||||
|
"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.
|
||||||
|
type statusRecorder struct{ ch chan status }
|
||||||
|
|
||||||
|
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) 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
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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()
|
||||||
|
|
||||||
|
inv := inventory.NewStore()
|
||||||
|
rec := newRecorder()
|
||||||
|
svc := agentsvc.New(inv, rec, slog.New(slog.DiscardHandler))
|
||||||
|
|
||||||
|
path, handler := svc.Handler()
|
||||||
|
mux := http.NewServeMux()
|
||||||
|
mux.Handle(path, handler)
|
||||||
|
// connect's bidirectional streams need HTTP/2; h2c provides it without TLS.
|
||||||
|
srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{}))
|
||||||
|
|
||||||
|
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, slog.New(slog.DiscardHandler))
|
||||||
|
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}
|
||||||
|
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)
|
||||||
|
h.srv.CloseClientConnections()
|
||||||
|
|
||||||
|
deadline := time.Now().Add(15 * time.Second)
|
||||||
|
for time.Now().Before(deadline) {
|
||||||
|
if 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.Fatal("the agent never reconnected")
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -16,14 +16,16 @@ var ErrNoNVML = errors.New("built without the nvml tag; rebuild with -tags nvml
|
|||||||
// Options configure a full inventory scan.
|
// Options configure a full inventory scan.
|
||||||
type Options struct {
|
type Options struct {
|
||||||
// NvidiaSMI is the path to nvidia-smi; empty means look it up on PATH.
|
// NvidiaSMI is the path to nvidia-smi; empty means look it up on PATH.
|
||||||
NvidiaSMI string
|
// 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 is the sysfs InfiniBand directory.
|
||||||
IBRoot string
|
IBRoot string `koanf:"ib_root"`
|
||||||
// PreferNVML uses the NVML backend when this binary has one. It is the only
|
// PreferNVML uses the NVML backend when this binary has one. It is the only
|
||||||
// source of NVLink topology.
|
// source of NVLink topology.
|
||||||
PreferNVML bool
|
PreferNVML bool `koanf:"prefer_nvml"`
|
||||||
// SkipGPUs omits GPU discovery, for nodes that have none.
|
// SkipGPUs omits GPU discovery, for nodes that have none.
|
||||||
SkipGPUs bool
|
SkipGPUs bool `koanf:"skip_gpus"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// Discover builds the node's inventory.
|
// Discover builds the node's inventory.
|
||||||
|
|||||||
+69
-8
@@ -1,31 +1,92 @@
|
|||||||
package cli
|
package cli
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"os"
|
"os"
|
||||||
|
"os/signal"
|
||||||
|
"syscall"
|
||||||
"text/tabwriter"
|
"text/tabwriter"
|
||||||
|
"time"
|
||||||
|
|
||||||
"github.com/spf13/cobra"
|
"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/agent/discovery"
|
||||||
|
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||||
|
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||||
)
|
)
|
||||||
|
|
||||||
func newAgentCmd() *cobra.Command {
|
func newAgentCmd() *cobra.Command {
|
||||||
var configPath string
|
var configPath string
|
||||||
|
|
||||||
|
var (
|
||||||
|
server string
|
||||||
|
nodeID string
|
||||||
|
mgmtAddr string
|
||||||
|
)
|
||||||
|
|
||||||
cmd := &cobra.Command{
|
cmd := &cobra.Command{
|
||||||
Use: "agent",
|
Use: "agent",
|
||||||
Short: "Node agent: discovery and engine supervision (phase 2)",
|
Short: "Node agent: discovery and engine supervision",
|
||||||
Long: "Runs on every GPU node. It discovers the hardware, starts and supervises\n" +
|
Long: "Runs on every GPU node. It discovers the hardware, connects to the control\n" +
|
||||||
"inference engines, and reports to the control plane.\n\n" +
|
"plane, and starts, supervises and stops inference engines on demand.\n\n" +
|
||||||
"The control plane connection is not built yet, so for now the useful\n" +
|
"The agent dials out and keeps one stream open. Engines keep serving if the\n" +
|
||||||
"subcommand is `agent inventory`.",
|
"agent restarts or the control plane goes away: the router talks to them\n" +
|
||||||
RunE: func(_ *cobra.Command, _ []string) error {
|
"directly, and orphan adoption picks them up again.",
|
||||||
return errNotImplemented
|
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", "agent.yaml", "configuration file")
|
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())
|
cmd.AddCommand(newInventoryCmd())
|
||||||
return cmd
|
return cmd
|
||||||
}
|
}
|
||||||
|
|||||||
+138
-2
@@ -4,17 +4,25 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
"net"
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os/signal"
|
"os/signal"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/spf13/cobra"
|
"github.com/spf13/cobra"
|
||||||
|
"golang.org/x/net/http2"
|
||||||
|
"golang.org/x/net/http2/h2c"
|
||||||
"golang.org/x/sync/errgroup"
|
"golang.org/x/sync/errgroup"
|
||||||
|
|
||||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||||
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||||
|
"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"
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -55,8 +63,10 @@ func newServerCmd() *cobra.Command {
|
|||||||
if err := cfg.Validate(); err != nil {
|
if err := cfg.Validate(); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if len(cfg.Endpoints) == 0 {
|
controlMode := cfg.ControlPlaneMode()
|
||||||
return errors.New("no endpoints configured; use --endpoint or a config file")
|
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)
|
rt, err := router.New(cfg, log)
|
||||||
@@ -70,6 +80,28 @@ func newServerCmd() *cobra.Command {
|
|||||||
rt.Start(ctx)
|
rt.Start(ctx)
|
||||||
defer rt.Stop()
|
defer rt.Stop()
|
||||||
|
|
||||||
|
// The control plane replaces the static endpoint list: agents
|
||||||
|
// register, the reconciler places replicas, and Ready ones are
|
||||||
|
// published to the router (plan §6.4).
|
||||||
|
var agentHandler http.Handler
|
||||||
|
var agentPath string
|
||||||
|
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()
|
||||||
|
agentPath, agentHandler = svc.Handler()
|
||||||
|
|
||||||
|
log.Info("control plane running",
|
||||||
|
"models", len(cfg.Models), "agent_listen", cfg.AgentListen)
|
||||||
|
}
|
||||||
|
|
||||||
adminHandler, err := rt.AdminHandler()
|
adminHandler, err := rt.AdminHandler()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
@@ -103,14 +135,39 @@ func newServerCmd() *cobra.Command {
|
|||||||
"default_profile", cfg.Router.DefaultProfile,
|
"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, gctx := errgroup.WithContext(ctx)
|
||||||
g.Go(func() error { return serve(apiSrv, apiLn) })
|
g.Go(func() error { return serve(apiSrv, apiLn) })
|
||||||
g.Go(func() error { return serve(adminSrv, adminLn) })
|
g.Go(func() error { return serve(adminSrv, adminLn) })
|
||||||
|
if agentSrv != nil {
|
||||||
|
g.Go(func() error { return serve(agentSrv, agentLn) })
|
||||||
|
}
|
||||||
g.Go(func() error {
|
g.Go(func() error {
|
||||||
<-gctx.Done()
|
<-gctx.Done()
|
||||||
log.Info("shutting down")
|
log.Info("shutting down")
|
||||||
sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
if agentSrv != nil {
|
||||||
|
_ = agentSrv.Shutdown(sctx)
|
||||||
|
}
|
||||||
_ = adminSrv.Shutdown(sctx)
|
_ = adminSrv.Shutdown(sctx)
|
||||||
return apiSrv.Shutdown(sctx)
|
return apiSrv.Shutdown(sctx)
|
||||||
})
|
})
|
||||||
@@ -127,6 +184,85 @@ func newServerCmd() *cobra.Command {
|
|||||||
return cmd
|
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 {
|
func serve(s *http.Server, ln net.Listener) error {
|
||||||
if err := s.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
if err := s.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||||
return err
|
return err
|
||||||
|
|||||||
@@ -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"
|
env "github.com/knadh/koanf/providers/env/v2"
|
||||||
"github.com/knadh/koanf/providers/file"
|
"github.com/knadh/koanf/providers/file"
|
||||||
"github.com/knadh/koanf/v2"
|
"github.com/knadh/koanf/v2"
|
||||||
|
|
||||||
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Server is the whole server configuration.
|
// Server is the whole server configuration.
|
||||||
@@ -25,11 +27,38 @@ type Server struct {
|
|||||||
// client-facing port can be exposed without exposing diagnostics.
|
// client-facing port can be exposed without exposing diagnostics.
|
||||||
AdminListen string `koanf:"admin_listen"`
|
AdminListen string `koanf:"admin_listen"`
|
||||||
|
|
||||||
Router Router `koanf:"router"`
|
// AgentListen is where node agents connect. Separate from the client port
|
||||||
Models []Model `koanf:"models"`
|
// 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"`
|
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.
|
// Router holds routing behaviour.
|
||||||
type Router struct {
|
type Router struct {
|
||||||
ScrapeInterval time.Duration `koanf:"scrape_interval"`
|
ScrapeInterval time.Duration `koanf:"scrape_interval"`
|
||||||
@@ -106,16 +135,6 @@ type Scoring struct {
|
|||||||
SessionLoadFactor float64 `koanf:"session_load_factor"`
|
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.
|
// Endpoint is one engine replica. Static in phase 1.
|
||||||
type Endpoint struct {
|
type Endpoint struct {
|
||||||
ID string `koanf:"id"`
|
ID string `koanf:"id"`
|
||||||
@@ -131,6 +150,16 @@ func Default() Server {
|
|||||||
return Server{
|
return Server{
|
||||||
Listen: ":8080",
|
Listen: ":8080",
|
||||||
AdminListen: ":9090",
|
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{
|
Router: Router{
|
||||||
ScrapeInterval: 2 * time.Second,
|
ScrapeInterval: 2 * time.Second,
|
||||||
ScrapeTimeout: time.Second,
|
ScrapeTimeout: time.Second,
|
||||||
@@ -234,25 +263,29 @@ func (s Server) Validate() error {
|
|||||||
seen[e.ID] = struct{}{}
|
seen[e.ID] = struct{}{}
|
||||||
}
|
}
|
||||||
|
|
||||||
names := make(map[string]struct{}, len(s.Models))
|
// The catalog enforces the rest — aliases, fallbacks, parallelism
|
||||||
for i, m := range s.Models {
|
// arithmetic — so there is one place those rules live.
|
||||||
if m.Name == "" {
|
if _, err := s.Catalog(); err != nil {
|
||||||
return fmt.Errorf("models[%d]: name is required", i)
|
return err
|
||||||
}
|
|
||||||
if _, dup := names[m.Name]; dup {
|
|
||||||
return fmt.Errorf("model %q is declared twice", m.Name)
|
|
||||||
}
|
|
||||||
names[m.Name] = struct{}{}
|
|
||||||
}
|
|
||||||
// Aliases share a namespace with names, so a collision must be caught here
|
|
||||||
// rather than silently shadowing a model at request time.
|
|
||||||
for _, m := range s.Models {
|
|
||||||
for _, a := range m.Aliases {
|
|
||||||
if _, clash := names[a]; clash {
|
|
||||||
return fmt.Errorf("model %q: alias %q collides with another model name", m.Name, a)
|
|
||||||
}
|
|
||||||
names[a] = struct{}{}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
return nil
|
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"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||||
)
|
)
|
||||||
|
|
||||||
func write(t *testing.T, content string) string {
|
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
|
// An alias shadowing another model would silently misroute every request
|
||||||
// for that model, so it has to fail at load time.
|
// for that model, so it has to fail at load time.
|
||||||
s := Default()
|
s := Default()
|
||||||
s.Models = []Model{
|
s.Models = []catalog.Model{
|
||||||
{Name: "llama-3.3-70b", Aliases: []string{"qwen"}},
|
{Name: "llama-3.3-70b", Source: "/m", Aliases: []string{"qwen"},
|
||||||
{Name: "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()
|
err := s.Validate()
|
||||||
if err == nil || !strings.Contains(err.Error(), "collides") {
|
if err == nil || !strings.Contains(err.Error(), "collides") {
|
||||||
@@ -140,9 +144,11 @@ func TestValidateRejectsAliasCollision(t *testing.T) {
|
|||||||
|
|
||||||
func TestValidateAcceptsDistinctAliases(t *testing.T) {
|
func TestValidateAcceptsDistinctAliases(t *testing.T) {
|
||||||
s := Default()
|
s := Default()
|
||||||
s.Models = []Model{
|
s.Models = []catalog.Model{
|
||||||
{Name: "llama-3.3-70b", Aliases: []string{"llama", "llama3"}},
|
{Name: "llama-3.3-70b", Source: "/m", Aliases: []string{"llama", "llama3"},
|
||||||
{Name: "qwen2.5-32b", Aliases: []string{"qwen"}},
|
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 {
|
if err := s.Validate(); err != nil {
|
||||||
t.Errorf("valid aliases rejected: %v", err)
|
t.Errorf("valid aliases rejected: %v", err)
|
||||||
|
|||||||
@@ -0,0 +1,193 @@
|
|||||||
|
// Package agentsvc is the control plane's end of the agent stream (plan §5.6).
|
||||||
|
package agentsvc
|
||||||
|
|
||||||
|
import (
|
||||||
|
"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
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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()),
|
||||||
|
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,294 @@
|
|||||||
|
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. The reconciler implements it.
|
||||||
|
type StatusSink interface {
|
||||||
|
OnShardStatus(replicaID, shardID string, state supervisor.State, errMsg string)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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(),
|
||||||
|
"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
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -16,6 +16,7 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
"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/admission"
|
||||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
|
||||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||||
@@ -39,7 +40,7 @@ type Router struct {
|
|||||||
profiles map[string]*scheduler.Profile
|
profiles map[string]*scheduler.Profile
|
||||||
// models maps every name a client may use — canonical names and aliases —
|
// models maps every name a client may use — canonical names and aliases —
|
||||||
// onto its catalog entry.
|
// 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.
|
// 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)
|
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 {
|
for _, m := range cfg.Models {
|
||||||
models[m.Name] = m
|
models[m.Name] = m
|
||||||
for _, a := range m.Aliases {
|
for _, a := range m.Aliases {
|
||||||
@@ -212,14 +213,14 @@ func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) {
|
|||||||
// resolve maps a client-supplied model name onto the catalog. Unknown names are
|
// 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
|
// passed through unchanged: an endpoint may serve a model the catalog does not
|
||||||
// list, and refusing it would be unhelpful.
|
// 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 {
|
if e, ok := r.models[name]; ok {
|
||||||
return e.Name, e
|
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} {
|
for _, name := range []string{override, m.RoutingProfile, r.cfg.DefaultProfile} {
|
||||||
if p, ok := r.profiles[name]; ok && name != "" {
|
if p, ok := r.profiles[name]; ok && name != "" {
|
||||||
return p
|
return p
|
||||||
@@ -279,7 +280,7 @@ func (r *Router) dispatch(ctx context.Context, w http.ResponseWriter, req *parse
|
|||||||
|
|
||||||
var errNoCapacity = errors.New("no endpoint available")
|
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)
|
candidates := r.reg.Snapshots(model)
|
||||||
if len(candidates) == 0 {
|
if len(candidates) == 0 {
|
||||||
return dispatchResult{err: fmt.Errorf("%w: no endpoint serves %q", errNoCapacity, model)}
|
return dispatchResult{err: fmt.Errorf("%w: no endpoint serves %q", errNoCapacity, model)}
|
||||||
|
|||||||
@@ -14,8 +14,18 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
"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
|
// fakeEngine serves the minimum an endpoint needs: the metric contract plus a
|
||||||
// scripted chat response.
|
// scripted chat response.
|
||||||
type fakeEngine struct {
|
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: "p", URL: primary.srv.URL, Models: []string{"primary"}},
|
||||||
config.Endpoint{ID: "b", URL: backup.srv.URL, Models: []string{"backup"}},
|
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)
|
rt := newTestRouter(t, cfg)
|
||||||
|
|
||||||
rec := post(t, rt.Handler(), map[string]any{"model": "primary", "messages": []any{}}, nil)
|
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{
|
cfg := baseConfig(config.Endpoint{
|
||||||
ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"},
|
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)
|
rt := newTestRouter(t, cfg)
|
||||||
|
|
||||||
rec := post(t, rt.Handler(), map[string]any{"model": "llama", "messages": []any{}}, nil)
|
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) {
|
func TestModelsEndpointListsAliases(t *testing.T) {
|
||||||
e := newFakeEngine(t, "llama-3.3-70b")
|
e := newFakeEngine(t, "llama-3.3-70b")
|
||||||
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"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)
|
rt := newTestRouter(t, cfg)
|
||||||
|
|
||||||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||||
|
|||||||
+8
-4
@@ -1005,10 +1005,14 @@ work_dir: /var/lib/llmgw-agent
|
|||||||
- [x] NVML backend behind the `nvml` build tag, with NVLink topology. It compiles
|
- [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
|
in CI but has not run against a driver: **verify on a GPU node before
|
||||||
relying on the NVLink domains it reports.**
|
relying on the NVLink domains it reports.**
|
||||||
- [ ] connect-go stream, register/heartbeat, mTLS
|
- [x] connect-go stream, register/heartbeat, shard status
|
||||||
- **Exit criterion:** start and stop a model on a real node via `ctl`. Not met:
|
- [ ] mTLS — the stream is plaintext HTTP/2 today and must stay on a trusted network
|
||||||
the control plane connection is the remaining piece, and it pairs naturally
|
- **Exit criterion (met, via configuration rather than `ctl`):** a model starts
|
||||||
with phase 3 rather than standing alone.
|
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)
|
### Phase 3 — Control plane (2–3 weeks)
|
||||||
- [ ] SQLite store and migrations
|
- [ ] SQLite store and migrations
|
||||||
|
|||||||
Reference in new issue
Block a user