Compare commits

...
12 Commits
Author SHA1 Message Date
Ilker ManapandClaude Opus 5.5 611fddd236 Describe v0.10.3
Release / build (push) Successful in 1m4s
Release / notes (push) Successful in 3s
CI / test (push) Successful in 3m39s
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-30 00:38:55 +02:00
Ilker ManapandClaude Opus 5.5 9853936e0b Give every engine a key only the router holds
Engines listen on every address their node has, with no key. Checked on the
live cluster: from a neighbouring node, a request straight to an engine's
port was answered — past the gateway's authentication, quotas and
accounting, by anyone who could reach the node.

Each replica's engine is now started with a key and refuses requests without
it. The key is derived from the replica id with an HMAC over a secret kept
in the store, rather than stored per replica: the router has to send the
right key to replicas this process never started — adopted after a restart,
or placed by the other control plane — and a function of a shared secret
gives that without a table to keep in step. A key taken from one node opens
that replica only. The secret is made on first use and copied with the rest
of the store when moving to PostgreSQL; a new one would lock the router out
of every running engine at the moment of the move.

The key reaches llama.cpp and vLLM through their environment
(LLAMA_API_KEY, VLLM_API_KEY) rather than their command line, which every
user on the node can read; SGLang only takes it as an argument. The agent's
readiness check, the router's proxy and its metrics scrape send it, the
client's own Authorization still never reaches the engine, and `ctl bench`
fetches the key of the replica it is pointed at from the admin API.

Engines started before this keep running without a key until their replica
is next started. With no store there is nowhere to keep the secret, and
engines start without one, with a warning.

Also formats two files that went in unformatted.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-30 00:38:29 +02:00
Ilker ManapandClaude Opus 5.5 d167686cf1 Stop a browser page, or a model definition, from running code on the nodes
The admin listener has no login. Reaching it is meant to mean holding the
control plane, so it binds to loopback and an operator reaches it over an
ssh tunnel — and so does every page open in that operator's browser. The
API decoded JSON whatever the Content-Type said, so a page could POST a
text/plain model definition, no preflight needed. A model's profile env
reached an engine the agent runs as root, unfiltered: LD_PRELOAD in a
definition was code execution on every node. A page that rebound its own
name to 127.0.0.1 could also read pprof's heap, which holds API keys.

The admin listener now answers only when Host is an address or localhost,
refuses requests a browser marks as cross-site (Origin, Sec-Fetch-Site),
and takes changes only as application/json, which no page can send across
origins without a preflight nobody grants. ctl sends that on every change,
bodiless DELETEs included; curl and ctl are otherwise unaffected.

Profile env that changes what runs rather than how — the dynamic loader's
LD_*, PATH, Python's PYTHON* start-up variables, BASH_ENV and their like —
is refused when a model is defined, and dropped again by the engine adapter
before the process starts, whatever sent the profile. None of the live
catalog's models sets any.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-30 00:29:52 +02:00
Ilker ManapandClaude Opus 5.5 97dc49cbeb Keep the mirror's management API and its downloads off the network behind it
Checked on the live cluster: from any node, the mirror's management API on
:8082 answered — list, fetch, delete — with no credentials. Unset
--api-allow meant anyone. It now means this machine only; the control plane
is let in with --api-allow, which `server init` already wrote for the
mirrors it sets up, or by mutual TLS.

A fetch went to whatever URL it was given, and the file server on :8081
then handed the result to anyone: "fetch http://169.254.169.254/..." read a
cloud's instance metadata through the mirror, and a URL on the LAN read
whatever answered there. Downloads started through the API now go to http
and https only, and are refused at dial time — after resolution, on every
redirect — when the address is loopback, private, link-local, CGNAT or
multicast. --allow-private-sources lets them through for a mirror that
really does copy from a server on its own network. `llmgw mirror add`, run
by someone on the mirror itself, is unchanged.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-30 00:27:27 +02:00
Ilker ManapandClaude Opus 5.5 3a26277c32 Refuse names that leave their directory, on the mirror and the agent
Two places built a path from a name somebody else supplied and checked only
for slashes.

The mirror refused "a/b" as a model name but took "..", which is a single
path element too: a fetch for model ".." wrote beside the mirror root rather
than inside it. Model, version and file names now have to be one element
that does not start with a dot — the mirror's own checksum files are dotted
— on the fetch path and the API alike.

The agent took the engine version from the control plane as given and made
a directory of it under /opt/llmgw/engines, and a RemoveAll before replacing
one. A version of "../x" would have done both outside the install root, as
root. The agent now installs only versions that look like one — letters,
digits and ._+- — whatever the control plane sent.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-30 00:24:48 +02:00
Ilker ManapandClaude Opus 5.5 7062645bcb Describe v0.10.2
Release / notes (push) Successful in 3s
CI / test (push) Failing after 4m10s
Release / build (push) Successful in 3m45s
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-29 22:37:28 +02:00
Ilker ManapandClaude Opus 5.5 bb90a927aa Print only what must outlive init's screen, and drop bubbles
Once server init's screen closed, the whole run was printed again under the
frame that had just shown it. Now the password and the next steps are, the
whole run goes to /var/lib/llmgw/init.log (0600: it holds the password), and
it is printed in full only when init fails. Seen in a draft of the install
video, where it scrolled the screen away.

The spinner and the text field are written here instead of taken from
bubbles, whose packages query the terminal's colours as they load. bubbletea
does the same deliberately, so the query stays; the dependency goes.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-29 22:37:28 +02:00
Ilker ManapandClaude Opus 5.5 49c446271c Describe v0.10.1
Release / build (push) Successful in 1m1s
Release / notes (push) Successful in 3s
CI / test (push) Successful in 3m39s
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-29 22:17:33 +02:00
Ilker ManapandClaude Opus 5.5 b5f412f9df Put the mirror on its own machine in two commands
mirror init sets a machine up as the model mirror — the user, /srv/models,
and a unit serving files on :8081 and the management API on :8082 — and
server init --mirror points the control plane at it instead of running one.
Before, that took a unit written by hand and an edit to server.yaml.

The management API can fetch and delete models and answered anyone who
could reach it. mirror serve --api-allow limits it to the addresses given
and the machine itself; mirror init passes the control plane's. The unit
text is shared with the one server init writes for a mirror beside it.

With --no-node the starter model is still fetched and catalogued, unpinned,
so it starts on the first node added, and the init form shows the mirror in
use instead of offering to run one.

Tried on three empty VMs: the mirror answered the control plane and refused
the node, and the starter model came up on the node when it was added.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-29 22:17:24 +02:00
Ilker ManapandClaude Opus 5.5 a1b5daceb8 Let server init ask for what it needs, and node add take several nodes
In a terminal and without --admin, server init shows a short form: the
administrator's email, the other names nodes will use (with the addresses
it found), SQLite or PostgreSQL, and the mirror, node and starter model.
Its own lines then fill in as a list — read from the writer init already
uses, so nothing in init changed — with downloads updating in place and the
password in a box until the end. The whole run, password included, is
printed again when the screen closes, so it is in the terminal's history.
The last screen says how to add a node and to back up the authority's key.

node add takes several addresses and adds them in turn, with a summary.
Suggested fabrics carry what the adapters report them to be, so a name
given for one node is used for the next on the same fabric without asking.

Tried on empty VMs: init from the form through to the starter model
serving, and two nodes added to that cluster in one command. The first
attempt at the form lost its keystrokes to the terminal colour query in a
pseudo-terminal that never answers it; a real terminal does.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-29 21:54:18 +02:00
Ilker ManapandClaude Opus 5.5 7071c67487 Show which requests each node's agent answers
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-29 21:31:05 +02:00
Ilker ManapandClaude Opus 5.5 7bb360fc78 Measure two nodes' fabric rail by rail: llmgw ctl fabric-test
doctor checks a node alone, and every check can pass while two nodes still
cannot reach each other. ctl fabric-test A B asks B's agent for an
ib_write_bw server and A's for the client, adapter by adapter as NCCL pairs
them, one rail at a time, and grades each against its link's rate. It goes
through the agents' streams, so there is no SSH between nodes; the agent
installs perftest the first time, and names the RoCE v2 GID itself.

The result is read from perftest's table by its heading, not a column
number. Agents now list the requests they answer beyond the usual
(Register.features), and the control plane refuses a test on one too old to
answer instead of waiting minutes for it. The forwarding dispatcher carries
the new capability, and its test says so: a missing forward would have made
the command quietly unavailable in production.

Not yet run over real InfiniBand or RoCE; this cluster has neither.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-29 21:29:23 +02:00
78 changed files with 3318 additions and 199 deletions

No files matched your search

+106
View File
@@ -20,6 +20,112 @@ is only released once its **exit criterion** is met, not when its code lands.
Until v1.0.0 the minor version is the breaking one: configuration formats and Go
APIs change between minors without a deprecation period.
## v0.10.3 — 2026-09-30
Security fixes from a review of v0.10.2. Three of them were confirmed on the
live cluster before being fixed.
### Engines demand a key only the router holds
Engines listened on every address of their node with no key. A request from a
neighbouring node reached the model directly, past the gateway's
authentication, quotas and accounting. Each replica's engine is now started
with a key derived from the replica id and a secret in the store. The router,
its metrics scrape, the agent's readiness check and `ctl bench` send it.
llama.cpp and vLLM receive it through their environment, not their command
line. Engines started before this keep running without a key until their
replica is next started.
### The mirror's management API answers only who it was told to
Without `--api-allow`, the management API answered anyone who could reach it,
including every node. It now answers this machine only; the control plane is
let in with `--api-allow` or mutual TLS. **Upgrading a mirror set up by hand:
add `--api-allow <control plane address>` to its unit.** Downloads started
through the API go to http and https only, and never to loopback, private or
link-local addresses (`--allow-private-sources` to permit that). Model,
version and file names can no longer be `..` or start with a dot.
### The admin listener refuses what a browser page could send
The admin API has no login and trusts whoever reaches it. A page open in an
operator's browser, reaching it over an ssh tunnel, could post a model
definition. The listener now requires Host to be an address or localhost,
refuses cross-site requests, and takes changes only as `application/json`.
A profile's `env` may no longer set variables that make an engine run other
code (`LD_*`, `PATH`, `PYTHON*` start-up variables, `BASH_ENV` and similar).
They are refused where a model is defined and dropped again before an engine
starts.
### Fixes
- The agent installs only engine versions that look like one, whatever the
control plane sent; a version of `../x` could have installed outside
`/opt/llmgw/engines`.
## v0.10.2 — 2026-09-29
- After `server init`'s screen closes, only the administrator's password and
the next steps are printed, not the whole run again under the screen that
just showed it. The whole run is kept in `/var/lib/llmgw/init.log`, root's
only, and printed in full when init fails.
- The screens' spinner and text field are llmgw's own rather than from
bubbles, one dependency fewer. Every llmgw command still asks the terminal
for its colours as it starts, which bubbletea v1 does on purpose; a terminal
that never answers waits for its timeout.
## v0.10.1 — 2026-09-29
### `llmgw ctl fabric-test`: two nodes, rail by rail
`llmgw doctor` checks a node on its own, and every check can pass while two
nodes still cannot reach each other: a cable into the wrong switch, a
partition key, a RoCE network with no route. `llmgw ctl fabric-test A B` runs
`ib_write_bw` from each of A's adapters on a fabric to the same adapter on B,
one rail at a time, through the agents — no SSH between nodes, and perftest is
installed the first time — and grades each against its link's rate: ok at 80%,
slow at 50%, fail below or when nothing got through.
Agents now say which requests they answer beyond the usual (`features`), so
one too old for the test is refused at once instead of waited for.
Written against perftest's real output and tested with it, but not yet run
over real InfiniBand or RoCE: this cluster has neither.
### The mirror on its own machine, in two commands
`llmgw mirror init --control-plane <address>` sets a machine up as the model
mirror: the user, `/srv/models`, and a service serving files to every node on
:8081 and the management API on :8082. `server init --mirror
http://<mirror>:8082` uses it instead of running one beside the control plane.
It took a unit written by hand and an edit to `server.yaml` before.
The management API can fetch and delete models; `mirror serve --api-allow`
makes it answer only the addresses given and the machine itself, and `mirror
init` gives it the control plane's. With `--no-node`, the starter model is
fetched and catalogued anyway, unpinned, and starts on the first node added.
Tried on three empty VMs: mirror, control plane and a node, and a request
answered by the starter model on the node.
### `server init` asks, and `node add` takes several nodes
- **`sudo llmgw server init`** in a terminal, without `--admin`, asks for the
administrator's email, the other names nodes will reach the machine by
(showing the addresses it already found), SQLite or PostgreSQL, and whether
to run the mirror, a node and the starter model here. Then its steps fill
in as a list, downloads updating in place, and the administrator's password
stays in a box until the end — and is printed again after, with the whole
run, so it is in the terminal's history. The last screen says how to add a
node and to back up the authority's key. `--plain`, or no terminal, prints
lines as before.
- **`llmgw node add a b c`** adds each in turn and ends with which were added.
A fabric named for one node's adapters is recognised on the next by what the
adapters report, and not asked again.
Both tried on empty VMs: init from the form to a starter model answering, and
two nodes added to it in one command in 40 seconds.
## v0.10.0 — 2026-09-29
Independent InfiniBand fabrics and NVLink domains, told apart, drawn and
+328 -63
View File
@@ -103,6 +103,7 @@ type AgentMessage struct {
// *AgentMessage_Shard
// *AgentMessage_Logs
// *AgentMessage_FabricScan
// *AgentMessage_FabricTest
Msg isAgentMessage_Msg `protobuf_oneof:"msg"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
@@ -190,6 +191,15 @@ func (x *AgentMessage) GetFabricScan() *FabricScanReply {
return nil
}
func (x *AgentMessage) GetFabricTest() *FabricTestReply {
if x != nil {
if x, ok := x.Msg.(*AgentMessage_FabricTest); ok {
return x.FabricTest
}
}
return nil
}
type isAgentMessage_Msg interface {
isAgentMessage_Msg()
}
@@ -214,6 +224,10 @@ type AgentMessage_FabricScan struct {
FabricScan *FabricScanReply `protobuf:"bytes,5,opt,name=fabric_scan,json=fabricScan,proto3,oneof"`
}
type AgentMessage_FabricTest struct {
FabricTest *FabricTestReply `protobuf:"bytes,6,opt,name=fabric_test,json=fabricTest,proto3,oneof"`
}
func (*AgentMessage_Register) isAgentMessage_Msg() {}
func (*AgentMessage_Heartbeat) isAgentMessage_Msg() {}
@@ -224,6 +238,8 @@ func (*AgentMessage_Logs) isAgentMessage_Msg() {}
func (*AgentMessage_FabricScan) isAgentMessage_Msg() {}
func (*AgentMessage_FabricTest) isAgentMessage_Msg() {}
// ControlMessage is anything the control plane tells a node.
type ControlMessage struct {
state protoimpl.MessageState `protogen:"open.v1"`
@@ -237,6 +253,7 @@ type ControlMessage struct {
// *ControlMessage_Provision
// *ControlMessage_Logs
// *ControlMessage_FabricScan
// *ControlMessage_FabricTest
Msg isControlMessage_Msg `protobuf_oneof:"msg"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
@@ -351,6 +368,15 @@ func (x *ControlMessage) GetFabricScan() *FabricScanRequest {
return nil
}
func (x *ControlMessage) GetFabricTest() *FabricTestRequest {
if x != nil {
if x, ok := x.Msg.(*ControlMessage_FabricTest); ok {
return x.FabricTest
}
}
return nil
}
type isControlMessage_Msg interface {
isControlMessage_Msg()
}
@@ -387,6 +413,10 @@ type ControlMessage_FabricScan struct {
FabricScan *FabricScanRequest `protobuf:"bytes,8,opt,name=fabric_scan,json=fabricScan,proto3,oneof"`
}
type ControlMessage_FabricTest struct {
FabricTest *FabricTestRequest `protobuf:"bytes,9,opt,name=fabric_test,json=fabricTest,proto3,oneof"`
}
func (*ControlMessage_Start) isControlMessage_Msg() {}
func (*ControlMessage_Stop) isControlMessage_Msg() {}
@@ -403,6 +433,8 @@ func (*ControlMessage_Logs) isControlMessage_Msg() {}
func (*ControlMessage_FabricScan) isControlMessage_Msg() {}
func (*ControlMessage_FabricTest) isControlMessage_Msg() {}
// 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.
@@ -478,8 +510,13 @@ type Register struct {
// vLLM, it answered "unsupported", and every node of a cluster half way
// through an upgrade showed as failing.
InstallableEngines []string `protobuf:"bytes,16,rep,name=installable_engines,json=installableEngines,proto3" json:"installable_engines,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
// features are the requests this agent answers beyond the ones every agent
// does: "fabric-test". An agent that does not know a request ignores it, so
// without this the control plane would wait minutes for an answer that is
// never coming.
Features []string `protobuf:"bytes,17,rep,name=features,proto3" json:"features,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *Register) Reset() {
@@ -624,6 +661,13 @@ func (x *Register) GetInstallableEngines() []string {
return nil
}
func (x *Register) GetFeatures() []string {
if x != nil {
return x.Features
}
return nil
}
// Host describes the operating system and processor of a node.
type Host struct {
state protoimpl.MessageState `protogen:"open.v1"`
@@ -2481,7 +2525,11 @@ type StartShard struct {
Fabric string `protobuf:"bytes,14,opt,name=fabric,proto3" json:"fabric,omitempty"`
// profile_fp identifies the version of the profile, so the agent can say
// which one a running shard was started from. The agent never interprets it.
ProfileFp string `protobuf:"bytes,15,opt,name=profile_fp,json=profileFp,proto3" json:"profile_fp,omitempty"`
ProfileFp string `protobuf:"bytes,15,opt,name=profile_fp,json=profileFp,proto3" json:"profile_fp,omitempty"`
// api_key is what the engine demands of every request, derived by the
// control plane from the replica id. The router sends it; nothing else
// that can reach the node's ports has it.
ApiKey string `protobuf:"bytes,16,opt,name=api_key,json=apiKey,proto3" json:"api_key,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
@@ -2621,6 +2669,13 @@ func (x *StartShard) GetProfileFp() string {
return ""
}
func (x *StartShard) GetApiKey() string {
if x != nil {
return x.ApiKey
}
return ""
}
type EngineProfile struct {
state protoimpl.MessageState `protogen:"open.v1"`
Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"`
@@ -3223,19 +3278,203 @@ func (x *FabricScanReply) GetTookMs() int64 {
return 0
}
// FabricTestRequest asks a node for one side of an RDMA bandwidth test
// between two nodes: ib_write_bw, from perftest, over one adapter. The
// server side is asked first and answers when the test is over; the client
// dials it. What it proves is what a replica spanning nodes needs and nothing
// on one node can show: that these two adapters reach each other through the
// fabric, and how fast.
type FabricTestRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
RequestId string `protobuf:"bytes,1,opt,name=request_id,json=requestId,proto3" json:"request_id,omitempty"`
// role is "server" or "client".
Role string `protobuf:"bytes,2,opt,name=role,proto3" json:"role,omitempty"`
Hca string `protobuf:"bytes,3,opt,name=hca,proto3" json:"hca,omitempty"`
Port int32 `protobuf:"varint,4,opt,name=port,proto3" json:"port,omitempty"`
// peer is the server's management address, for the client. perftest
// exchanges its connection details over TCP before any RDMA.
Peer string `protobuf:"bytes,5,opt,name=peer,proto3" json:"peer,omitempty"`
// tcp_port is that exchange's port, the same on both sides.
TcpPort int32 `protobuf:"varint,6,opt,name=tcp_port,json=tcpPort,proto3" json:"tcp_port,omitempty"`
Seconds int32 `protobuf:"varint,7,opt,name=seconds,proto3" json:"seconds,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *FabricTestRequest) Reset() {
*x = FabricTestRequest{}
mi := &file_llmgw_v1_agent_proto_msgTypes[31]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *FabricTestRequest) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*FabricTestRequest) ProtoMessage() {}
func (x *FabricTestRequest) ProtoReflect() protoreflect.Message {
mi := &file_llmgw_v1_agent_proto_msgTypes[31]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use FabricTestRequest.ProtoReflect.Descriptor instead.
func (*FabricTestRequest) Descriptor() ([]byte, []int) {
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{31}
}
func (x *FabricTestRequest) GetRequestId() string {
if x != nil {
return x.RequestId
}
return ""
}
func (x *FabricTestRequest) GetRole() string {
if x != nil {
return x.Role
}
return ""
}
func (x *FabricTestRequest) GetHca() string {
if x != nil {
return x.Hca
}
return ""
}
func (x *FabricTestRequest) GetPort() int32 {
if x != nil {
return x.Port
}
return 0
}
func (x *FabricTestRequest) GetPeer() string {
if x != nil {
return x.Peer
}
return ""
}
func (x *FabricTestRequest) GetTcpPort() int32 {
if x != nil {
return x.TcpPort
}
return 0
}
func (x *FabricTestRequest) GetSeconds() int32 {
if x != nil {
return x.Seconds
}
return 0
}
// FabricTestReply is what one side measured.
type FabricTestReply struct {
state protoimpl.MessageState `protogen:"open.v1"`
RequestId string `protobuf:"bytes,1,opt,name=request_id,json=requestId,proto3" json:"request_id,omitempty"`
// gbps is the average bandwidth, in gigabits a second.
Gbps float64 `protobuf:"fixed64,2,opt,name=gbps,proto3" json:"gbps,omitempty"`
// output is the end of the tool's output, for when the number is not
// enough.
Output string `protobuf:"bytes,3,opt,name=output,proto3" json:"output,omitempty"`
Error string `protobuf:"bytes,4,opt,name=error,proto3" json:"error,omitempty"`
TookMs int64 `protobuf:"varint,5,opt,name=took_ms,json=tookMs,proto3" json:"took_ms,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *FabricTestReply) Reset() {
*x = FabricTestReply{}
mi := &file_llmgw_v1_agent_proto_msgTypes[32]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *FabricTestReply) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*FabricTestReply) ProtoMessage() {}
func (x *FabricTestReply) ProtoReflect() protoreflect.Message {
mi := &file_llmgw_v1_agent_proto_msgTypes[32]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use FabricTestReply.ProtoReflect.Descriptor instead.
func (*FabricTestReply) Descriptor() ([]byte, []int) {
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{32}
}
func (x *FabricTestReply) GetRequestId() string {
if x != nil {
return x.RequestId
}
return ""
}
func (x *FabricTestReply) GetGbps() float64 {
if x != nil {
return x.Gbps
}
return 0
}
func (x *FabricTestReply) GetOutput() string {
if x != nil {
return x.Output
}
return ""
}
func (x *FabricTestReply) GetError() string {
if x != nil {
return x.Error
}
return ""
}
func (x *FabricTestReply) GetTookMs() int64 {
if x != nil {
return x.TookMs
}
return 0
}
var File_llmgw_v1_agent_proto protoreflect.FileDescriptor
const file_llmgw_v1_agent_proto_rawDesc = "" +
"\n" +
"\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\x94\x02\n" +
"\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\xd2\x02\n" +
"\fAgentMessage\x120\n" +
"\bregister\x18\x01 \x01(\v2\x12.llmgw.v1.RegisterH\x00R\bregister\x123\n" +
"\theartbeat\x18\x02 \x01(\v2\x13.llmgw.v1.HeartbeatH\x00R\theartbeat\x12-\n" +
"\x05shard\x18\x03 \x01(\v2\x15.llmgw.v1.ShardStatusH\x00R\x05shard\x12)\n" +
"\x04logs\x18\x04 \x01(\v2\x13.llmgw.v1.LogsReplyH\x00R\x04logs\x12<\n" +
"\vfabric_scan\x18\x05 \x01(\v2\x19.llmgw.v1.FabricScanReplyH\x00R\n" +
"fabricScanB\x05\n" +
"\x03msg\"\xa5\x03\n" +
"fabricScan\x12<\n" +
"\vfabric_test\x18\x06 \x01(\v2\x19.llmgw.v1.FabricTestReplyH\x00R\n" +
"fabricTestB\x05\n" +
"\x03msg\"\xe5\x03\n" +
"\x0eControlMessage\x12,\n" +
"\x05start\x18\x01 \x01(\v2\x14.llmgw.v1.StartShardH\x00R\x05start\x12)\n" +
"\x04stop\x18\x02 \x01(\v2\x13.llmgw.v1.StopShardH\x00R\x04stop\x12,\n" +
@@ -3245,8 +3484,10 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
"\tprovision\x18\x06 \x01(\v2\x13.llmgw.v1.ProvisionH\x00R\tprovision\x12+\n" +
"\x04logs\x18\a \x01(\v2\x15.llmgw.v1.LogsRequestH\x00R\x04logs\x12>\n" +
"\vfabric_scan\x18\b \x01(\v2\x1b.llmgw.v1.FabricScanRequestH\x00R\n" +
"fabricScanB\x05\n" +
"\x03msg\"\xc7\x06\n" +
"fabricScan\x12>\n" +
"\vfabric_test\x18\t \x01(\v2\x1b.llmgw.v1.FabricTestRequestH\x00R\n" +
"fabricTestB\x05\n" +
"\x03msg\"\xe3\x06\n" +
"\bRegister\x12\x17\n" +
"\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1a\n" +
"\bhostname\x18\x02 \x01(\tR\bhostname\x12\x1b\n" +
@@ -3264,7 +3505,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
"\x04host\x18\r \x01(\v2\x0e.llmgw.v1.HostR\x04host\x12%\n" +
"\x0eengines_probed\x18\x0e \x01(\bR\renginesProbed\x12\"\n" +
"\x04imex\x18\x0f \x01(\v2\x0e.llmgw.v1.IMEXR\x04imex\x12/\n" +
"\x13installable_engines\x18\x10 \x03(\tR\x12installableEngines\x1a9\n" +
"\x13installable_engines\x18\x10 \x03(\tR\x12installableEngines\x12\x1a\n" +
"\bfeatures\x18\x11 \x03(\tR\bfeatures\x1a9\n" +
"\vLabelsEntry\x12\x10\n" +
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\x1aA\n" +
@@ -3462,7 +3704,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
"\n" +
"bytes_done\x18\x02 \x01(\x03R\tbytesDone\x12\x1f\n" +
"\vbytes_total\x18\x03 \x01(\x03R\n" +
"bytesTotal\"\xd9\x03\n" +
"bytesTotal\"\xf2\x03\n" +
"\n" +
"StartShard\x12\x1d\n" +
"\n" +
@@ -3485,7 +3727,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
"\x04pool\x18\r \x01(\tR\x04pool\x12\x16\n" +
"\x06fabric\x18\x0e \x01(\tR\x06fabric\x12\x1d\n" +
"\n" +
"profile_fp\x18\x0f \x01(\tR\tprofileFp\"\xfd\x02\n" +
"profile_fp\x18\x0f \x01(\tR\tprofileFp\x12\x17\n" +
"\aapi_key\x18\x10 \x01(\tR\x06apiKey\"\xfd\x02\n" +
"\rEngineProfile\x12\x12\n" +
"\x04name\x18\x01 \x01(\tR\x04name\x12\x16\n" +
"\x06engine\x18\x02 \x01(\tR\x06engine\x12\x1b\n" +
@@ -3536,7 +3779,23 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
"request_id\x18\x01 \x01(\tR\trequestId\x12\x16\n" +
"\x06output\x18\x02 \x01(\tR\x06output\x12\x14\n" +
"\x05error\x18\x03 \x01(\tR\x05error\x12\x17\n" +
"\atook_ms\x18\x04 \x01(\x03R\x06tookMs2O\n" +
"\atook_ms\x18\x04 \x01(\x03R\x06tookMs\"\xb5\x01\n" +
"\x11FabricTestRequest\x12\x1d\n" +
"\n" +
"request_id\x18\x01 \x01(\tR\trequestId\x12\x12\n" +
"\x04role\x18\x02 \x01(\tR\x04role\x12\x10\n" +
"\x03hca\x18\x03 \x01(\tR\x03hca\x12\x12\n" +
"\x04port\x18\x04 \x01(\x05R\x04port\x12\x12\n" +
"\x04peer\x18\x05 \x01(\tR\x04peer\x12\x19\n" +
"\btcp_port\x18\x06 \x01(\x05R\atcpPort\x12\x18\n" +
"\aseconds\x18\a \x01(\x05R\aseconds\"\x8b\x01\n" +
"\x0fFabricTestReply\x12\x1d\n" +
"\n" +
"request_id\x18\x01 \x01(\tR\trequestId\x12\x12\n" +
"\x04gbps\x18\x02 \x01(\x01R\x04gbps\x12\x16\n" +
"\x06output\x18\x03 \x01(\tR\x06output\x12\x14\n" +
"\x05error\x18\x04 \x01(\tR\x05error\x12\x17\n" +
"\atook_ms\x18\x05 \x01(\x03R\x06tookMs2O\n" +
"\fAgentService\x12?\n" +
"\aConnect\x12\x16.llmgw.v1.AgentMessage\x1a\x18.llmgw.v1.ControlMessage(\x010\x01B\x93\x01\n" +
"\fcom.llmgw.v1B\n" +
@@ -3555,7 +3814,7 @@ func file_llmgw_v1_agent_proto_rawDescGZIP() []byte {
}
var file_llmgw_v1_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 37)
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 39)
var file_llmgw_v1_agent_proto_goTypes = []any{
(ShardStatus_State)(0), // 0: llmgw.v1.ShardStatus.State
(*AgentMessage)(nil), // 1: llmgw.v1.AgentMessage
@@ -3589,12 +3848,14 @@ var file_llmgw_v1_agent_proto_goTypes = []any{
(*LogsReply)(nil), // 29: llmgw.v1.LogsReply
(*FabricScanRequest)(nil), // 30: llmgw.v1.FabricScanRequest
(*FabricScanReply)(nil), // 31: llmgw.v1.FabricScanReply
nil, // 32: llmgw.v1.Register.LabelsEntry
nil, // 33: llmgw.v1.Register.EngineVersionsEntry
nil, // 34: llmgw.v1.GPU.NicPathsEntry
nil, // 35: llmgw.v1.IBPort.CountersEntry
nil, // 36: llmgw.v1.Heartbeat.EngineVersionsEntry
nil, // 37: llmgw.v1.EngineProfile.EnvEntry
(*FabricTestRequest)(nil), // 32: llmgw.v1.FabricTestRequest
(*FabricTestReply)(nil), // 33: llmgw.v1.FabricTestReply
nil, // 34: llmgw.v1.Register.LabelsEntry
nil, // 35: llmgw.v1.Register.EngineVersionsEntry
nil, // 36: llmgw.v1.GPU.NicPathsEntry
nil, // 37: llmgw.v1.IBPort.CountersEntry
nil, // 38: llmgw.v1.Heartbeat.EngineVersionsEntry
nil, // 39: llmgw.v1.EngineProfile.EnvEntry
}
var file_llmgw_v1_agent_proto_depIdxs = []int32{
3, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register
@@ -3602,49 +3863,51 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{
20, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus
29, // 3: llmgw.v1.AgentMessage.logs:type_name -> llmgw.v1.LogsReply
31, // 4: llmgw.v1.AgentMessage.fabric_scan:type_name -> llmgw.v1.FabricScanReply
22, // 5: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
25, // 6: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
26, // 7: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
27, // 8: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
24, // 9: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights
17, // 10: llmgw.v1.ControlMessage.provision:type_name -> llmgw.v1.Provision
28, // 11: llmgw.v1.ControlMessage.logs:type_name -> llmgw.v1.LogsRequest
30, // 12: llmgw.v1.ControlMessage.fabric_scan:type_name -> llmgw.v1.FabricScanRequest
32, // 13: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
7, // 14: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
8, // 15: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
33, // 16: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
6, // 17: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
9, // 18: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights
5, // 19: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage
4, // 20: llmgw.v1.Register.host:type_name -> llmgw.v1.Host
14, // 21: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX
0, // 22: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
34, // 23: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry
11, // 24: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad
35, // 25: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry
13, // 26: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager
5, // 27: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage
9, // 28: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights
36, // 29: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry
19, // 30: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus
10, // 31: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad
12, // 32: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort
8, // 33: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment
15, // 34: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState
14, // 35: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX
18, // 36: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant
0, // 37: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
21, // 38: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress
23, // 39: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
37, // 40: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
1, // 41: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
2, // 42: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
42, // [42:43] is the sub-list for method output_type
41, // [41:42] is the sub-list for method input_type
41, // [41:41] is the sub-list for extension type_name
41, // [41:41] is the sub-list for extension extendee
0, // [0:41] is the sub-list for field type_name
33, // 5: llmgw.v1.AgentMessage.fabric_test:type_name -> llmgw.v1.FabricTestReply
22, // 6: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
25, // 7: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
26, // 8: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
27, // 9: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
24, // 10: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights
17, // 11: llmgw.v1.ControlMessage.provision:type_name -> llmgw.v1.Provision
28, // 12: llmgw.v1.ControlMessage.logs:type_name -> llmgw.v1.LogsRequest
30, // 13: llmgw.v1.ControlMessage.fabric_scan:type_name -> llmgw.v1.FabricScanRequest
32, // 14: llmgw.v1.ControlMessage.fabric_test:type_name -> llmgw.v1.FabricTestRequest
34, // 15: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
7, // 16: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
8, // 17: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
35, // 18: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
6, // 19: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
9, // 20: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights
5, // 21: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage
4, // 22: llmgw.v1.Register.host:type_name -> llmgw.v1.Host
14, // 23: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX
0, // 24: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
36, // 25: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry
11, // 26: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad
37, // 27: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry
13, // 28: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager
5, // 29: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage
9, // 30: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights
38, // 31: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry
19, // 32: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus
10, // 33: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad
12, // 34: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort
8, // 35: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment
15, // 36: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState
14, // 37: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX
18, // 38: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant
0, // 39: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
21, // 40: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress
23, // 41: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
39, // 42: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
1, // 43: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
2, // 44: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
44, // [44:45] is the sub-list for method output_type
43, // [43:44] is the sub-list for method input_type
43, // [43:43] is the sub-list for extension type_name
43, // [43:43] is the sub-list for extension extendee
0, // [0:43] is the sub-list for field type_name
}
func init() { file_llmgw_v1_agent_proto_init() }
@@ -3658,6 +3921,7 @@ func file_llmgw_v1_agent_proto_init() {
(*AgentMessage_Shard)(nil),
(*AgentMessage_Logs)(nil),
(*AgentMessage_FabricScan)(nil),
(*AgentMessage_FabricTest)(nil),
}
file_llmgw_v1_agent_proto_msgTypes[1].OneofWrappers = []any{
(*ControlMessage_Start)(nil),
@@ -3668,6 +3932,7 @@ func file_llmgw_v1_agent_proto_init() {
(*ControlMessage_Provision)(nil),
(*ControlMessage_Logs)(nil),
(*ControlMessage_FabricScan)(nil),
(*ControlMessage_FabricTest)(nil),
}
type x struct{}
out := protoimpl.TypeBuilder{
@@ -3675,7 +3940,7 @@ func file_llmgw_v1_agent_proto_init() {
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: unsafe.Slice(unsafe.StringData(file_llmgw_v1_agent_proto_rawDesc), len(file_llmgw_v1_agent_proto_rawDesc)),
NumEnums: 1,
NumMessages: 37,
NumMessages: 39,
NumExtensions: 0,
NumServices: 1,
},
+44
View File
@@ -25,6 +25,7 @@ message AgentMessage {
ShardStatus shard = 3;
LogsReply logs = 4;
FabricScanReply fabric_scan = 5;
FabricTestReply fabric_test = 6;
}
}
@@ -39,6 +40,7 @@ message ControlMessage {
Provision provision = 6;
LogsRequest logs = 7;
FabricScanRequest fabric_scan = 8;
FabricTestRequest fabric_test = 9;
}
}
@@ -125,6 +127,12 @@ message Register {
// vLLM, it answered "unsupported", and every node of a cluster half way
// through an upgrade showed as failing.
repeated string installable_engines = 16;
// features are the requests this agent answers beyond the ones every agent
// does: "fabric-test". An agent that does not know a request ignores it, so
// without this the control plane would wait minutes for an answer that is
// never coming.
repeated string features = 17;
}
// Host describes the operating system and processor of a node.
@@ -516,6 +524,10 @@ message StartShard {
// profile_fp identifies the version of the profile, so the agent can say
// which one a running shard was started from. The agent never interprets it.
string profile_fp = 15;
// api_key is what the engine demands of every request, derived by the
// control plane from the replica id. The router sends it; nothing else
// that can reach the node's ports has it.
string api_key = 16;
}
message EngineProfile {
@@ -605,3 +617,35 @@ message FabricScanReply {
string error = 3;
int64 took_ms = 4;
}
// FabricTestRequest asks a node for one side of an RDMA bandwidth test
// between two nodes: ib_write_bw, from perftest, over one adapter. The
// server side is asked first and answers when the test is over; the client
// dials it. What it proves is what a replica spanning nodes needs and nothing
// on one node can show: that these two adapters reach each other through the
// fabric, and how fast.
message FabricTestRequest {
string request_id = 1;
// role is "server" or "client".
string role = 2;
string hca = 3;
int32 port = 4;
// peer is the server's management address, for the client. perftest
// exchanges its connection details over TCP before any RDMA.
string peer = 5;
// tcp_port is that exchange's port, the same on both sides.
int32 tcp_port = 6;
int32 seconds = 7;
}
// FabricTestReply is what one side measured.
message FabricTestReply {
string request_id = 1;
// gbps is the average bandwidth, in gigabits a second.
double gbps = 2;
// output is the end of the tool's output, for when the number is not
// enough.
string output = 3;
string error = 4;
int64 took_ms = 5;
}
+34 -1
View File
@@ -121,6 +121,39 @@ rate and errors, and the NCCL environment a shard on each fabric would be given.
`--json` writes every fact it read as well as what it concluded — send that
file when the report is not enough.
## Checking two nodes: `llmgw ctl fabric-test`
Every check on one node can pass while two nodes still cannot reach each other
— a cable into the wrong switch, a partition key, a RoCE network with no route
— and a replica spanning them then hangs on its first all-reduce. Before
placing one, on the control plane:
```
llmgw ctl fabric-test gpu01 gpu02
```
```
gpu01 → gpu02 on ib-compute, 41s
RAIL LINK MEASURED VERDICT DETAIL
mlx5_0 → mlx5_0 400 Gb/s 388.2 Gb/s ok 97% of 400 Gb/s
mlx5_1 → mlx5_1 400 Gb/s 241.0 Gb/s slow 60% of 400 Gb/s: check the cable, and the card's PCIe width and NUMA node
mlx5_2 → mlx5_2 400 Gb/s 0.0 Gb/s fail gpu01: ib_write_bw: Couldn't connect to 10.1.0.22:18517
```
It runs `ib_write_bw` (perftest, installed on a node the first time) from each
of the first node's adapters on the fabric to the same adapter on the second —
mlx5_0 to mlx5_0, the way NCCL pairs them — one rail at a time, through the
agents, so no SSH between nodes is needed. Each rail is graded against its
link's rate: **ok** at 80% or more, **slow** at 50%, **fail** below that or when
nothing got through. `--fabric` picks one when the nodes share several,
`--seconds` the length of each rail's run. The traffic is real: run it when the
fabric is quiet.
It measures RDMA between adapters, not GPUDirect or NCCL itself: a pass says
the fabric carries traffic between these two nodes at its rate, which is what
a first all-reduce needs and what nothing on one node can show.
## Watching them
- **Cluster page, Fabrics.** Each declared fabric and what it physically is;
@@ -172,6 +205,6 @@ workers headless. `launch_mode: ray` is there for a vLLM too old for that.
| It | Look at |
|---|---|
| is not placed | `llmgw ctl explain-placement <model>`: a fabric split by identity, a node whose IMEX is down, a node refused because its adapters are cabled to two fabrics |
| starts and hangs on its first collective | the cluster page's fabric problems; `doctor` on each node — a port down, the wrong interface, RoCE without a v2 GID |
| starts and hangs on its first collective | `llmgw ctl fabric-test` between two of its nodes; the cluster page's fabric problems; `doctor` on each node — a port down, the wrong interface, RoCE without a v2 GID |
| runs, slowly | `nvidia_peermem`; the map for a degraded cable; the node page's NVLink matrix for a group that crossed PCIe |
| worked yesterday | the subnet manager table for a failover; `link_downed` on the node's ports |
+27 -3
View File
@@ -35,10 +35,12 @@ install it.
## 2. Set up the control plane
```
sudo llmgw server init --admin you@example.com
sudo llmgw server init
```
This one command sets up everything on this machine, and takes a few minutes:
It asks for your email and a few choices, then sets up everything on this
machine, in a few minutes. In a script, give them as flags instead:
`sudo llmgw server init --admin you@example.com`.
- the service, its certificates and its configuration;
- your administrator account. **The password is printed once. Copy it now.**
@@ -54,18 +56,33 @@ sudo llmgw server init --admin you@example.com --host llm.example.com
It is safe to run again. It keeps what exists and finishes what is missing.
**The mirror on its own machine**, when it wants a larger disk than the
control plane has: first on that machine, with llmgw installed as in step 1,
```
sudo llmgw mirror init --control-plane <control plane address>
```
then on the control plane `sudo llmgw server init --mirror http://<mirror
address>:8082`. The mirror serves model files to every node, and takes
orders only from the control plane. Untick "This machine as a node" in the
form (or add `--no-node`) to keep models off the control plane; the starter
model then starts on the first node you add.
**Check:** open `http://<control-plane>:8080/chat`, sign in with the printed
password, and ask something. The small model answers. Change the password on
the account page.
## 3. Add the GPU nodes
On the control plane, for each node:
On the control plane:
```
sudo llmgw node add ubuntu@10.0.0.21
```
Several at once: `sudo llmgw node add ubuntu@10.0.0.21 ubuntu@10.0.0.22 ubuntu@10.0.0.23`.
Replace `ubuntu` with the user you log in to the node as, and `10.0.0.21` with
its address. That user needs `sudo` without a password on the node.
@@ -89,6 +106,13 @@ the same fabric the same name. In a script, give it instead:
sudo llmgw node add ubuntu@10.0.0.21 --fabric ib-a:mlx5_0,mlx5_1,mlx5_2,mlx5_3
```
Once two nodes on a fabric are in, check they reach each other at the link's
rate before placing anything across them:
```
llmgw ctl fabric-test gpu01 gpu02
```
**If it cannot log in.** `sudo` runs SSH as root, with root's keys. To use your
own keys instead, log in to the control plane with `ssh -A` and keep your
agent through sudo:
+8
View File
@@ -71,6 +71,7 @@ rather than repeated.
| `lease_ttl` | `time.Duration` | [Control](#control) |
| `id` | `string` | [Endpoint](#endpoint) |
| `url` | `string` | [Endpoint](#endpoint) |
| `api_key` | `string` | [Endpoint](#endpoint) |
| `models` | `[]string` | [Endpoint](#endpoint) |
| `engine` | `string` | [Endpoint](#endpoint) |
| `pool` | `string` | [Endpoint](#endpoint) |
@@ -566,6 +567,13 @@ reading the alert. Zero takes the default.
`string`
### `api_key`
`string`
APIKey is sent to the engine as a bearer token, for an engine started
with one.
### `models`
`[]string`
-2
View File
@@ -6,7 +6,6 @@ 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/charmbracelet/bubbles v1.0.0
github.com/charmbracelet/bubbletea v1.3.10
github.com/charmbracelet/lipgloss v1.1.0
github.com/jackc/pgx/v5 v5.11.0
@@ -28,7 +27,6 @@ require (
)
require (
github.com/atotto/clipboard v0.1.4 // indirect
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/charmbracelet/colorprofile v0.4.1 // indirect
-4
View File
@@ -2,16 +2,12 @@ connectrpc.com/connect v1.21.0 h1:LhqSJt7jHf5NJBo9Jq/t/9FjcYAideif0mg+qe2jCUs=
connectrpc.com/connect v1.21.0/go.mod h1:A2ygJrukXwWy32vkCAAHNVguZrqZ+jeZ9rGRnGR4dN4=
github.com/NVIDIA/go-nvml v0.13.4-0 h1:o3jp9u2x1R9ShFE3v+Aesp55XOSIQFMJz/VGNUcJaNE=
github.com/NVIDIA/go-nvml v0.13.4-0/go.mod h1:id63qwpoDWpFXwnwM6psDCSqW4BmNu6mWpr4YeQtPGo=
github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc=
github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E=
github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw=
github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
github.com/charmbracelet/colorprofile v0.4.1 h1:a1lO03qTrSIRaK8c3JRxJDZOvhvIeSco3ej+ngLk1kk=
+75
View File
@@ -844,6 +844,7 @@ func (a *Agent) session(ctx context.Context, server string) error {
reg.EngineVersions = a.engines(parent)
reg.EnginesProbed = true
reg.InstallableEngines = provision.Installable()
reg.Features = []string{FeatureFabricTest}
reg.RunningShards = a.runningShards()
a.mu.Lock()
h := a.prov.Host
@@ -1145,6 +1146,16 @@ func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out cha
case <-ctx.Done():
}
}(m.FabricScan)
case *llmgwv1.ControlMessage_FabricTest:
// Seconds of traffic, and a package install the first time: off the
// receive loop too.
go func(req *llmgwv1.FabricTestRequest) {
reply := a.fabricTest(ctx, req)
select {
case out <- &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_FabricTest{FabricTest: reply}}:
case <-ctx.Done():
}
}(m.FabricTest)
case *llmgwv1.ControlMessage_Ack:
if !m.Ack.GetOk() {
a.log.Warn("control plane rejected a message",
@@ -1758,6 +1769,70 @@ func (a *Agent) ensureDiags(ctx context.Context) {
}
}
// FeatureFabricTest is the agent answering FabricTestRequest.
const FeatureFabricTest = "fabric-test"
// runWriteBW is discovery.RunWriteBW, and haveWriteBW whether it is
// installed; both replaced in tests.
var (
runWriteBW = discovery.RunWriteBW
haveWriteBW = func() bool { _, err := exec.LookPath("ib_write_bw"); return err == nil }
)
// fabricTest runs this node's side of a bandwidth test with another.
func (a *Agent) fabricTest(ctx context.Context, req *llmgwv1.FabricTestRequest) *llmgwv1.FabricTestReply {
start := time.Now()
reply := &llmgwv1.FabricTestReply{RequestId: req.GetRequestId()}
defer func() { reply.TookMs = time.Since(start).Milliseconds() }()
w := discovery.WriteBW{HCA: req.GetHca(), Port: int(req.GetPort()), GIDIndex: -1,
TCPPort: int(req.GetTcpPort()), Seconds: int(req.GetSeconds())}
if req.GetRole() == "client" {
w.Peer = req.GetPeer()
}
a.mu.Lock()
found := false
for _, h := range a.inv.HCAs {
for _, p := range h.Ports {
if h.Name == w.HCA && p.Num == w.Port {
found = true
// RoCE needs its v2 GID named, as NCCL does; InfiniBand uses the LID.
if p.LinkLayer == discovery.LinkEthernet && p.RoCEv2GID >= 0 {
w.GIDIndex = p.RoCEv2GID
}
}
}
}
a.mu.Unlock()
if !found {
reply.Error = fmt.Sprintf("this node has no port %s/%d", w.HCA, w.Port)
return reply
}
if !haveWriteBW() && a.prov != nil {
a.log.Info("installing perftest, for a bandwidth test with another node")
pkgs := map[string][]string{"deb": {"perftest"}, "rpm": {"perftest"}}
if err := a.prov.InstallPackages(ctx, pkgs); err != nil {
reply.Error = "installing perftest: " + err.Error()
return reply
}
}
gbps, out, err := runWriteBW(ctx, w)
reply.Gbps = gbps
reply.Output = tailLines(out, 12)
if err != nil {
reply.Error = err.Error()
}
return reply
}
func tailLines(s string, n int) string {
lines := strings.Split(strings.TrimRight(s, "\n"), "\n")
if len(lines) > n {
lines = lines[len(lines)-n:]
}
return strings.Join(lines, "\n")
}
// discoverGPUFabric and discoverIMEX are replaced in tests.
var (
discoverGPUFabric = discovery.DiscoverGPUFabric
+32
View File
@@ -23,6 +23,7 @@ import (
"golang.org/x/net/http2"
"golang.org/x/net/http2/h2c"
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"
@@ -881,3 +882,34 @@ func TestAFabricScanIsAskedForAndAnswered(t *testing.T) {
t.Errorf("asked %q, got %+v", asked, r)
}
}
func TestABandwidthTestIsRunOnTheAdapterAskedFor(t *testing.T) {
oldRun, oldHave := runWriteBW, haveWriteBW
t.Cleanup(func() { runWriteBW, haveWriteBW = oldRun, oldHave })
haveWriteBW = func() bool { return true }
var got discovery.WriteBW
runWriteBW = func(_ context.Context, w discovery.WriteBW) (float64, string, error) {
got = w
return 388.2, " #bytes #iterations BW peak[Gb/sec] BW average[Gb/sec]\n", nil
}
h := start(t, func(c *Config) { c.Discovery.IBRoot = ibTreeWithActivePorts(t, "mlx5_0") })
r, err := h.svc.FabricTest(t.Context(), "n1", &llmgwv1.FabricTestRequest{
Role: "client", Hca: "mlx5_0", Port: 1, Peer: "10.0.0.2", TcpPort: 18515, Seconds: 3})
if err != nil {
t.Fatal(err)
}
if r.GetGbps() != 388.2 || r.GetError() != "" {
t.Errorf("reply %+v", r)
}
if got.HCA != "mlx5_0" || got.Peer != "10.0.0.2" || got.TCPPort != 18515 || got.GIDIndex != -1 {
t.Errorf("ran %+v", got)
}
// An adapter the node does not have is said so, without running anything.
got = discovery.WriteBW{}
r, _ = h.svc.FabricTest(t.Context(), "n1", &llmgwv1.FabricTestRequest{Role: "server", Hca: "mlx5_9", Port: 1, Seconds: 1})
if !strings.Contains(r.GetError(), "no port mlx5_9/1") || got.HCA != "" {
t.Errorf("reply %+v, ran %+v", r, got)
}
}
+133
View File
@@ -0,0 +1,133 @@
package discovery
import (
"context"
"errors"
"fmt"
"os/exec"
"regexp"
"strconv"
"strings"
"time"
)
// WriteBW is one side of an ib_write_bw run, from perftest: an RDMA write
// bandwidth test between two adapters, which is what a collective spanning
// nodes rides on.
type WriteBW struct {
HCA string
Port int
// GIDIndex is the RoCE v2 GID to use; negative on InfiniBand, where the
// port's LID is used.
GIDIndex int
// Peer is the server's address, for the client; empty makes the server.
Peer string
TCPPort int
Seconds int
}
// Args are the command line: a duration rather than a count, so the test
// takes as long on a slow link as a fast one, and results in Gb/s.
func (w WriteBW) Args() []string {
args := []string{"-d", w.HCA, "-i", strconv.Itoa(max(w.Port, 1)),
"-p", strconv.Itoa(w.TCPPort), "-D", strconv.Itoa(max(w.Seconds, 1)),
"-F", "--report_gbits"}
if w.GIDIndex >= 0 {
args = append(args, "-x", strconv.Itoa(w.GIDIndex))
}
if w.Peer != "" {
args = append(args, w.Peer)
}
return args
}
// RunWriteBW runs one side and returns the average bandwidth the output
// reports, in Gb/s, and the output. The client retries its connection for a
// while: the server is started first, but a few seconds may pass before it
// listens.
func RunWriteBW(ctx context.Context, w WriteBW) (float64, string, error) {
bin, err := exec.LookPath("ib_write_bw")
if err != nil {
return 0, "", errors.New("ib_write_bw is not installed (the perftest package)")
}
limit := time.Duration(w.Seconds+60) * time.Second
var out []byte
for attempt := 0; ; attempt++ {
cctx, cancel := context.WithTimeout(ctx, limit)
out, err = exec.CommandContext(cctx, bin, w.Args()...).CombinedOutput() //nolint:gosec // fixed tool
cancel()
if err == nil || w.Peer == "" || !refused(string(out)) || attempt >= 10 {
break
}
select {
case <-ctx.Done():
return 0, string(out), ctx.Err()
case <-time.After(time.Second):
}
}
text := string(out)
if err != nil {
return 0, text, fmt.Errorf("ib_write_bw: %v: %s", err, lastNonEmpty(text))
}
gbps, perr := ParseWriteBW(text)
if perr != nil && w.Peer != "" {
// Only the client is sure to print the result; a server that does not
// is not a failure.
return 0, text, perr
}
return gbps, text, nil
}
func refused(out string) bool {
for _, s := range []string{"Couldn't connect", "Unable to init the socket connection", "Connection refused"} {
if strings.Contains(out, s) {
return true
}
}
return false
}
var columnSplit = regexp.MustCompile(`\s{2,}`)
// ParseWriteBW finds the average bandwidth in perftest's result table:
//
// #bytes #iterations BW peak[Gb/sec] BW average[Gb/sec] MsgRate[Mpps]
// 65536 887036 0.00 96.93 0.184877
//
// by the heading's position rather than a fixed column, since versions add
// columns and the peak reads 0.00 when the test runs for a duration.
func ParseWriteBW(out string) (float64, error) {
lines := strings.Split(out, "\n")
for i, l := range lines {
if !strings.Contains(l, "BW average") {
continue
}
heads := columnSplit.Split(strings.TrimSpace(l), -1)
col := -1
for j, h := range heads {
if strings.HasPrefix(h, "BW average") {
col = j
}
}
for _, d := range lines[i+1:] {
f := strings.Fields(d)
if len(f) == 0 || strings.HasPrefix(f[0], "-") {
continue
}
if len(f) != len(heads) || col < 0 {
return 0, fmt.Errorf("perftest's result table has a shape this does not know: %q / %q", l, d)
}
v, err := strconv.ParseFloat(f[col], 64)
if err != nil {
return 0, fmt.Errorf("reading the average bandwidth %q: %w", f[col], err)
}
return v, nil
}
}
return 0, errors.New("perftest printed no result")
}
func lastNonEmpty(s string) string {
lines := strings.Split(strings.TrimSpace(s), "\n")
return strings.TrimSpace(lines[len(lines)-1])
}
+49
View File
@@ -0,0 +1,49 @@
package discovery
import (
"slices"
"testing"
)
// The client's output of perftest 24.x, with a duration.
const writeBWClient = `---------------------------------------------------------------------------------------
RDMA_Write BW Test
Dual-port : OFF Device : mlx5_0
Number of qps : 1 Transport type : IB
Connection type : RC Using SRQ : OFF
TX depth : 128
CQ Moderation : 100
Mtu : 4096[B]
Link type : IB
Max inline data : 0[B]
rdma_cm QPs : OFF
Data ex. method : Ethernet
---------------------------------------------------------------------------------------
local address: LID 0x05 QPN 0x0112 PSN 0x7b1f02 RKey 0x1fffbd VAddr 0x007f1d2a5f0000
remote address: LID 0x09 QPN 0x0113 PSN 0x2c1e44 RKey 0x1ffeaf VAddr 0x007fa4e1a60000
---------------------------------------------------------------------------------------
#bytes #iterations BW peak[Gb/sec] BW average[Gb/sec] MsgRate[Mpps]
65536 1851214 0.00 388.21 0.740422
---------------------------------------------------------------------------------------
`
func TestTheAverageBandwidthIsReadFromPerftest(t *testing.T) {
got, err := ParseWriteBW(writeBWClient)
if err != nil || got != 388.21 {
t.Errorf("got %v, %v", got, err)
}
if _, err := ParseWriteBW(" Couldn't connect to 10.0.0.2:18515\n"); err == nil {
t.Error("output with no result table gave a number")
}
}
func TestWriteBWArguments(t *testing.T) {
ib := WriteBW{HCA: "mlx5_0", Port: 1, GIDIndex: -1, TCPPort: 18515, Seconds: 3}.Args()
if slices.Contains(ib, "-x") {
t.Errorf("InfiniBand was given a GID index: %v", ib)
}
roce := WriteBW{HCA: "mlx5_0", Port: 1, GIDIndex: 3, Peer: "10.0.0.2", TCPPort: 18515, Seconds: 3}.Args()
if !slices.Contains(roce, "-x") || roce[len(roce)-1] != "10.0.0.2" {
t.Errorf("RoCE client: %v", roce)
}
}
+12
View File
@@ -70,6 +70,13 @@ type ShardSpec struct {
HTTPPort int
Profile Profile
// APIKey is what the engine demands of every request. Engines listen on
// all of a node's addresses; without a key anybody who can reach the port
// talks to the model directly, past the gateway's authentication, quotas
// and accounting. Given to the engine in its environment rather than on
// its command line, which every user on the node can read.
APIKey string `json:"api_key,omitempty"`
// Pool and Fabric are control plane concepts the agent cannot derive. They
// are carried so the agent can report them back after a restart; Fabric
// also says, with an "nvl:" prefix, that the replica is inside a
@@ -249,6 +256,11 @@ func baseEnv(s ShardSpec) map[string]string {
// Profile overrides win: they are the escape hatch for the per-cluster
// tuning no adapter can guess.
for k, v := range s.Profile.Env {
// Checked again here, not only where the model was defined: this is
// the process about to run as root, whatever sent the profile.
if bad, _ := UnsafeEnv(k); bad {
continue
}
env[k] = v
}
return env
+5 -1
View File
@@ -120,8 +120,12 @@ func (l *LlamaCpp) Commands(s ShardSpec) ([]Command, error) {
}
args = append(args, p.Args...)
env := baseEnv(s)
if s.APIKey != "" {
env["LLAMA_API_KEY"] = s.APIKey
}
return []Command{{
Name: "llama-server", Path: bin, Args: args, Env: materializeEnv(baseEnv(s)),
Name: "llama-server", Path: bin, Args: args, Env: materializeEnv(env),
}}, nil
}
+16
View File
@@ -245,3 +245,19 @@ func TestAMissingPathIsLeftToTheEngine(t *testing.T) {
t.Errorf("expected it passed through, got %q %v", got, err)
}
}
func TestLlamaCppIsGivenItsKeyOutOfSightOfTheProcessList(t *testing.T) {
// The command line is readable by every user on the node; the
// environment of a root process is not.
cmds, err := (&LlamaCpp{}).Commands(ShardSpec{HTTPPort: 8080, APIKey: "k1",
Profile: Profile{ModelPath: "/m.gguf"}})
if err != nil {
t.Fatal(err)
}
if slices.Contains(cmds[0].Args, "k1") || slices.Contains(cmds[0].Args, "--api-key") {
t.Errorf("the key is on the command line: %v", cmds[0].Args)
}
if !slices.Contains(cmds[0].Env, "LLAMA_API_KEY=k1") {
t.Error("llama-server was not given its key")
}
}
+5
View File
@@ -96,6 +96,11 @@ func (g *SGLang) serverCommand(s ShardSpec) (Command, error) {
)
}
args = append(args, p.Args...)
// SGLang reads its key only from the command line, so on a node with
// other users it is visible to them in the process list.
if s.APIKey != "" && s.Leader() {
args = append(args, "--api-key", s.APIKey)
}
return Command{Name: "sglang", Path: python, Args: args, Env: materializeEnv(baseEnv(s))}, nil
}
+30
View File
@@ -0,0 +1,30 @@
package engines
import "strings"
// UnsafeEnv reports whether a profile may not set an environment variable,
// and why.
//
// A profile's env is the escape hatch for tuning no adapter can guess —
// CUDA_*, NCCL_*, VLLM_*, HF_* — and it reaches an engine the agent starts as
// root. Some variables are not tuning but a way to run other code in that
// process: the dynamic loader's LD_PRELOAD, Python's PYTHONSTARTUP, a shell's
// BASH_ENV. Set from a model definition, which an admin API or a config file
// supplies, they would turn "add a model" into "run this as root on every
// node". They are refused where a model is defined and again where an engine
// starts.
func UnsafeEnv(key string) (bool, string) {
k := strings.ToUpper(key)
switch {
case strings.HasPrefix(k, "LD_"), strings.HasPrefix(k, "DYLD_"):
return true, "it changes what the dynamic loader loads"
case k == "PATH":
return true, "it changes which binaries run"
case strings.HasPrefix(k, "PYTHON") && k != "PYTHONUNBUFFERED" && k != "PYTHONHASHSEED":
return true, "it changes what Python imports or runs at start"
case k == "BASH_ENV" || k == "ENV" || k == "PERL5OPT" || k == "PERL5LIB" ||
k == "NODE_OPTIONS" || k == "RUBYOPT" || k == "GCONV_PATH" || k == "OPENSSL_CONF":
return true, "it makes an interpreter or library load code"
}
return false, ""
}
+3
View File
@@ -150,6 +150,9 @@ func (v *VLLM) serverCommand(s ShardSpec) (Command, error) {
// Profile arguments last, so an operator can override anything above.
args = append(args, p.Args...)
if s.APIKey != "" {
env["VLLM_API_KEY"] = s.APIKey
}
return Command{Name: "vllm", Path: bin, Args: args, Env: materializeEnv(env)}, nil
}
+14
View File
@@ -16,6 +16,7 @@ import (
"os"
"os/exec"
"path/filepath"
"regexp"
"strings"
"time"
@@ -77,6 +78,10 @@ type Installer struct {
PyPIBase string
}
// validVersion is what a version may look like: "b10991", "0.29.0",
// "0.11.0+cu128". No separators, and not starting with a dot.
var validVersion = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._+-]{0,63}$`)
// DefaultRoot is where engines are installed when nothing says otherwise.
const DefaultRoot = "/opt/llmgw/engines"
@@ -114,6 +119,15 @@ func (in *Installer) Converge(ctx context.Context, wants []Want) []Status {
func (in *Installer) one(ctx context.Context, w Want) Status {
st := Status{Engine: w.Engine, Version: w.Version}
// Checked here, on the node, whatever the control plane checked: the
// version names a directory under the install root and is passed to a
// package manager, and "../x" would put an install — and the RemoveAll
// that replaces one — outside it, as root.
if !validVersion.MatchString(w.Version) {
st.State = Unsupported
st.Detail = fmt.Sprintf("%q is not a version this agent will install", w.Version)
return st
}
recipe, ok := Recipes[w.Engine]
if !ok {
st.State = Unsupported
+12
View File
@@ -330,3 +330,15 @@ func TestAPackageManagerBusyWithItsOwnUpdatesIsWaitedFor(t *testing.T) {
t.Error("the wait was not reported")
}
}
func TestAVersionThatIsAPathIsRefused(t *testing.T) {
// The version names a directory under the install root; "../" would put
// an install, and the RemoveAll that replaces one, outside it, as root.
in := &Installer{Root: t.TempDir(), Host: ubuntu}
for _, v := range []string{"../../etc", "..", ".hidden", "b1/../../x", "", "b1 && rm -rf /"} {
got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: v}})
if got[0].State != Unsupported {
t.Errorf("version %q: %+v", v, got[0])
}
}
}
+7 -4
View File
@@ -53,8 +53,8 @@ func NewHTTPHealth(timeout time.Duration, opts ...HTTPHealthOption) *HTTPHealth
}
// Ready reports whether the engine at baseURL is serving model.
func (h *HTTPHealth) Ready(ctx context.Context, baseURL, model string) error {
if err := h.get(ctx, baseURL+h.healthPath, nil); err != nil {
func (h *HTTPHealth) Ready(ctx context.Context, baseURL, model, apiKey string) error {
if err := h.get(ctx, baseURL+h.healthPath, apiKey, nil); err != nil {
return fmt.Errorf("health: %w", err)
}
if !h.requireName || model == "" {
@@ -66,7 +66,7 @@ func (h *HTTPHealth) Ready(ctx context.Context, baseURL, model string) error {
ID string `json:"id"`
} `json:"data"`
}
if err := h.get(ctx, baseURL+h.modelsPath, &list); err != nil {
if err := h.get(ctx, baseURL+h.modelsPath, apiKey, &list); err != nil {
return fmt.Errorf("models: %w", err)
}
for _, d := range list.Data {
@@ -77,11 +77,14 @@ func (h *HTTPHealth) Ready(ctx context.Context, baseURL, model string) error {
return fmt.Errorf("engine is up but does not serve %q yet", model)
}
func (h *HTTPHealth) get(ctx context.Context, url string, out any) error {
func (h *HTTPHealth) get(ctx context.Context, url, apiKey string, out any) error {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return err
}
if apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := h.client.Do(req)
if err != nil {
return err
+24 -5
View File
@@ -39,7 +39,7 @@ func engineStub(t *testing.T, healthy bool, models ...string) *httptest.Server {
func TestReadyWhenServingModel(t *testing.T) {
srv := engineStub(t, true, "llama-3.3-70b")
if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b"); err != nil {
if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b", ""); err != nil {
t.Errorf("a serving engine was reported unready: %v", err)
}
}
@@ -48,7 +48,7 @@ func TestNotReadyWhileLoading(t *testing.T) {
// The case that matters: the process is up but the weights are not loaded.
// Routing on /health alone would send traffic into a loading model.
srv := engineStub(t, false, "llama-3.3-70b")
err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b")
err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b", "")
if err == nil {
t.Fatal("an unhealthy engine was reported ready")
}
@@ -61,7 +61,7 @@ func TestNotReadyWhenModelMissing(t *testing.T) {
// /health passes but the model is not listed: the engine is up and serving
// something else, or still registering.
srv := engineStub(t, true, "some-other-model")
err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b")
err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b", "")
if err == nil {
t.Fatal("a model that is not served was reported ready")
}
@@ -73,14 +73,33 @@ func TestNotReadyWhenModelMissing(t *testing.T) {
func TestModelCheckCanBeSkipped(t *testing.T) {
srv := engineStub(t, true)
h := NewHTTPHealth(time.Second, WithoutModelCheck())
if err := h.Ready(t.Context(), srv.URL, "anything"); err != nil {
if err := h.Ready(t.Context(), srv.URL, "anything", ""); err != nil {
t.Errorf("the model check was not skipped: %v", err)
}
}
func TestUnreachableEngine(t *testing.T) {
h := NewHTTPHealth(200 * time.Millisecond)
if err := h.Ready(t.Context(), "http://127.0.0.1:1", "m"); err == nil {
if err := h.Ready(t.Context(), "http://127.0.0.1:1", "m", ""); err == nil {
t.Error("an unreachable engine was reported ready")
}
}
func TestTheHealthCheckSendsTheEnginesKey(t *testing.T) {
// An engine started with a key refuses /v1/models without it, and a
// shard that is serving would read as never ready.
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("Authorization") != "Bearer k1" {
w.WriteHeader(http.StatusUnauthorized)
return
}
_, _ = w.Write([]byte(`{"data":[{"id":"m"}]}`))
}))
defer srv.Close()
if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "m", "k1"); err != nil {
t.Errorf("with the key: %v", err)
}
if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "m", ""); err == nil {
t.Error("without the key the engine refused, and that read as ready")
}
}
+2 -2
View File
@@ -79,7 +79,7 @@ func (c *Config) normalize() {
// supervisor can be tested without an HTTP engine.
type HealthChecker interface {
// Ready reports whether the shard at baseURL serves model.
Ready(ctx context.Context, baseURL, model string) error
Ready(ctx context.Context, baseURL, model, apiKey string) error
}
// Shard is one supervised engine process.
@@ -411,7 +411,7 @@ func (s *Supervisor) WaitReady(ctx context.Context, sh *Shard, baseURL, model st
last := strings.Join(sh.Logs(10), "\n")
return fmt.Errorf("shard %s not ready within %v; last output:\n%s", sh.ID, timeout, last)
case <-t.C:
if err := s.health.Ready(ctx, baseURL, model); err != nil {
if err := s.health.Ready(ctx, baseURL, model, sh.Spec.APIKey); err != nil {
continue
}
sh.setState(Ready, nil)
+1 -1
View File
@@ -21,7 +21,7 @@ import (
// scriptedHealth reports ready once the flag is set.
type scriptedHealth struct{ ready atomic.Bool }
func (h *scriptedHealth) Ready(context.Context, string, string) error {
func (h *scriptedHealth) Ready(context.Context, string, string, string) error {
if h.ready.Load() {
return nil
}
+1
View File
@@ -225,6 +225,7 @@ func runStep(ctx context.Context, client *http.Client, cfg Config, concurrency i
)
sampler := newSampler(cfg.MetricsURL, cfg.Engine)
sampler.apiKey = cfg.APIKey
sampler.start(ctx)
start := time.Now()
+4
View File
@@ -19,6 +19,7 @@ import (
// Optional, because a run pointed at the gateway has no single engine to read.
type sampler struct {
url string
apiKey string
engine string
client *http.Client
cancel context.CancelFunc
@@ -65,6 +66,9 @@ func (s *sampler) sample(ctx context.Context) {
if err != nil {
return
}
if s.apiKey != "" {
req.Header.Set("Authorization", "Bearer "+s.apiKey)
}
resp, err := s.client.Do(req)
if err != nil {
return
+20
View File
@@ -41,6 +41,11 @@ func newBenchCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
cfg.MetricsURL = metrics
cfg.Engine = engine
cfg.APIKey = key
// A replica's engine demands its own key. Asked for rather than
// typed: the admin API knows which replica answers at the target.
if cfg.APIKey == "" {
cfg.APIKey = replicaKey(cmd, client, target)
}
if len(steps) > 0 {
cfg.Concurrency = steps
}
@@ -227,3 +232,18 @@ func ms(d time.Duration) string {
return d.Round(10 * time.Millisecond).String()
}
}
// replicaKey is the engine key of the replica at target, or "" when target is
// not a replica (the gateway, say) or the admin API cannot be asked.
func replicaKey(cmd *cobra.Command, client func() *ctlClient, target string) string {
var reps []adminapi.ReplicaView
if err := client().get(cmd.Context(), "/api/v1/replicas", &reps); err != nil {
return ""
}
for _, rep := range reps {
if sameEndpoint(rep.Endpoint, target) {
return rep.EngineKey
}
}
return ""
}
+4 -1
View File
@@ -67,7 +67,9 @@ func (c *ctlClient) do(ctx context.Context, method, path string, body, out any)
if err != nil {
return err
}
if body != nil {
// On every change, body or not: the admin listener refuses anything else,
// because a browser can send a bodiless POST across origins unasked.
if method != http.MethodGet {
req.Header.Set("Content-Type", "application/json")
}
// Who made the change, kept with a stored model so a catalog that grew
@@ -120,6 +122,7 @@ func newCtlCmd() *cobra.Command {
newModelsCmd(client, &asJSON),
newReplicasCmd(client, &asJSON),
newExplainCmd(client, &asJSON),
newFabricTestCmd(client, &asJSON),
newNodeStateCmd(client, "cordon", "cordoned",
"Stop placing new work on a node; what runs there keeps running"),
newNodeStateCmd(client, "uncordon", "ready",
+75
View File
@@ -0,0 +1,75 @@
package cli
import (
"encoding/json"
"fmt"
"text/tabwriter"
"time"
"github.com/spf13/cobra"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
)
func newFabricTestCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
var (
fabric string
seconds int
)
c := &cobra.Command{
Use: "fabric-test NODE-A NODE-B",
Short: "Measure RDMA bandwidth between two nodes, rail by rail",
Long: `Runs ib_write_bw from each of NODE-A's adapters on a fabric to the same
adapter on NODE-B — mlx5_0 to mlx5_0, and so on, the way NCCL pairs them —
one rail at a time, and grades each against its link's rate: ok at 80% or
more, slow at 50%, fail below or when nothing got through.
Every check on one node can pass while two nodes still cannot reach each
other: a cable into the wrong switch, a partition key, a RoCE network with no
route. A replica spanning them then hangs on its first all-reduce. Run this
before placing one.
It goes through the agents, so it needs no ssh between nodes. perftest is
installed on a node the first time it is asked. The traffic is real and
takes a few seconds a rail; run it when the fabric is quiet.`,
Example: ` llmgw ctl fabric-test gpu01 gpu02
llmgw ctl fabric-test gpu01 gpu02 --fabric ib-compute --seconds 10`,
Args: cobra.ExactArgs(2),
RunE: func(cmd *cobra.Command, args []string) error {
body, _ := json.Marshal(map[string]any{"a": args[0], "b": args[1], "fabric": fabric, "seconds": seconds})
cl := client()
// Rails one after another, each for its seconds and a settle; the
// client's usual timeout is for questions, not measurements.
cl.http.Timeout = 30 * time.Minute
var v adminapi.FabricTestView
if err := cl.post(cmd.Context(), "/api/v1/fabric-test", body, &v); err != nil {
return err
}
if *asJSON {
return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
}
out := cmd.OutOrStdout()
fmt.Fprintf(out, "%s → %s on %s, %s\n\n", v.A, v.B, v.Fabric, v.Took)
w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
fmt.Fprintln(w, "RAIL\tLINK\tMEASURED\tVERDICT\tDETAIL")
for _, r := range v.Rails {
link := "-"
if r.RateGbps > 0 {
link = fmt.Sprintf("%d Gb/s", r.RateGbps)
}
fmt.Fprintf(w, "%s → %s\t%s\t%.1f Gb/s\t%s\t%s\n", r.A, r.B, link, r.Gbps, r.Verdict, r.Detail)
}
_ = w.Flush()
for _, n := range v.Notes {
fmt.Fprintf(out, "\nnote: %s\n", n)
}
if !v.OK {
return fmt.Errorf("not every rail passed")
}
return nil
},
}
c.Flags().StringVar(&fabric, "fabric", "", "the fabric to test (default: the one the two nodes share)")
c.Flags().IntVar(&seconds, "seconds", 3, "how long each rail runs, up to 30")
return c
}
+8
View File
@@ -5,6 +5,7 @@ import (
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fabrictest"
)
func TestTheForwardingDispatcherCarriesEveryCapability(t *testing.T) {
@@ -23,10 +24,17 @@ func TestTheForwardingDispatcherCarriesEveryCapability(t *testing.T) {
t.Error("dispatcherRef drops WeightRemover, so the control plane cannot reclaim disk")
}
if _, ok := ref.(fabrictest.Asker); !ok {
t.Error("dispatcherRef drops the bandwidth test, so ctl fabric-test cannot reach the agents")
}
// And the thing it forwards to must have it, or the forwarding has nothing
// to reach.
var svc any = &agentsvc.Service{}
if _, ok := svc.(control.WeightRemover); !ok {
t.Error("the agent service does not implement WeightRemover")
}
if _, ok := svc.(fabrictest.Asker); !ok {
t.Error("the agent service does not implement the bandwidth test")
}
}
+85 -8
View File
@@ -4,7 +4,9 @@ import (
"encoding/json"
"fmt"
"log/slog"
"net"
"net/http"
"net/netip"
"os"
"strings"
"time"
@@ -44,6 +46,7 @@ func newMirrorCmd() *cobra.Command {
newMirrorUnusedCmd(&dir),
newMirrorRemoveCmd(&dir),
newMirrorServeCmd(&dir),
newMirrorInitCmd(&dir),
)
return c
}
@@ -229,6 +232,8 @@ func newMirrorRemoveCmd(dir *string) *cobra.Command {
func newMirrorServeCmd(dir *string) *cobra.Command {
var listen, apiListen, base string
var tlsc tlsconf.Config
var allow []string
var allowPrivate bool
c := &cobra.Command{
Use: "serve",
@@ -258,6 +263,9 @@ func newMirrorServeCmd(dir *string) *cobra.Command {
out := cmd.OutOrStdout()
log := slog.Default()
mgr := mirror.NewManager(*dir, log)
if allowPrivate {
mgr.AllowPrivateSources()
}
api := mirror.NewAPI(mgr, strings.TrimSuffix(base, "/"))
// http.Dir refuses paths that escape the root, which is the whole
@@ -268,9 +276,13 @@ func newMirrorServeCmd(dir *string) *cobra.Command {
Handler: readOnly(withManifest(*dir, http.FileServer(http.Dir(*dir)))),
ReadHeaderTimeout: 10 * time.Second,
}
allowed, err := parseAllow(allow)
if err != nil {
return err
}
admin := &http.Server{
Addr: apiListen,
Handler: api.Handler(),
Handler: apiGate(tlsc.Enabled(), allowed, api.Handler()),
ReadHeaderTimeout: 10 * time.Second,
}
@@ -279,14 +291,17 @@ func newMirrorServeCmd(dir *string) *cobra.Command {
fmt.Fprintf(out, "serving %s on %s\n", *dir, listen)
go func() { errs <- fmt.Errorf("serving the API: %w", serveAPI(admin, tlsc)) }()
if tlsc.Enabled() {
switch {
case tlsc.Enabled():
fmt.Fprintf(out, "management API on %s, mutual TLS against %s\n", apiListen, tlsc.CA)
} else {
// Said plainly rather than left to be discovered. An API that
// can delete models and start downloads, open to anyone who can
// reach the port, is a choice the operator should have made
// rather than inherited from a default.
fmt.Fprintf(out, "management API on %s, UNAUTHENTICATED — keep it on a trusted network\n", apiListen)
case len(allowed) > 0:
fmt.Fprintf(out, "management API on %s, for %s and this machine only\n", apiListen, strings.Join(allow, ", "))
default:
// This machine only, which was once "anyone": an API that can
// delete models and start downloads, open to every node that
// could reach the port, is a choice the operator should make
// with --api-allow rather than inherit.
fmt.Fprintf(out, "management API on %s, for this machine only (--api-allow to add the control plane)\n", apiListen)
}
fmt.Fprintf(out, "catalog sources should point at %s\n", base)
@@ -300,6 +315,10 @@ func newMirrorServeCmd(dir *string) *cobra.Command {
c.Flags().StringVar(&apiListen, "api-listen", ":8082", "address to serve the management API on")
c.Flags().StringVar(&base, "base", "http://localhost:8081",
"the address nodes use to reach this mirror, as it should appear in catalog sources")
c.Flags().StringArrayVar(&allow, "api-allow", nil,
"an address or network the management API answers, besides this machine; repeatable (default: this machine only)")
c.Flags().BoolVar(&allowPrivate, "allow-private-sources", false,
"let downloads reach loopback, private and link-local addresses")
c.Flags().StringVar(&tlsc.CA, "tls-ca", "", "CA bundle that management clients must present a certificate from")
c.Flags().StringVar(&tlsc.Cert, "tls-cert", "", "this mirror's certificate")
c.Flags().StringVar(&tlsc.Key, "tls-key", "", "this mirror's private key")
@@ -398,3 +417,61 @@ func versionPath(p string) (model, version string, ok bool) {
}
return parts[0], parts[1], true
}
// parseAllow reads addresses and networks: 10.0.0.140, 10.0.0.0/24.
func parseAllow(list []string) ([]netip.Prefix, error) {
var out []netip.Prefix
for _, s := range list {
for _, one := range strings.Split(s, ",") {
one = strings.TrimSpace(one)
if one == "" {
continue
}
if p, err := netip.ParsePrefix(one); err == nil {
out = append(out, p.Masked())
continue
}
a, err := netip.ParseAddr(one)
if err != nil {
return nil, fmt.Errorf("--api-allow %q is not an address or a network", one)
}
out = append(out, netip.PrefixFrom(a.Unmap(), a.Unmap().BitLen()))
}
}
return out, nil
}
// apiGate decides who reaches the management API. With mutual TLS the
// certificate decides, narrowed by --api-allow when given. Without it, only
// this machine and the addresses allowed: nothing allowed used to mean anyone,
// and every node that could reach the port could delete models.
func apiGate(mtls bool, allowed []netip.Prefix, next http.Handler) http.Handler {
if mtls && len(allowed) == 0 {
return next
}
return onlyFrom(allowed, next)
}
// onlyFrom answers only this machine and the addresses allowed: the management
// API can delete models and start downloads, and a mirror on its own machine
// is reached over the network by one control plane and no one else.
func onlyFrom(allowed []netip.Prefix, next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
host, _, err := net.SplitHostPort(r.RemoteAddr)
a, perr := netip.ParseAddr(host)
if err != nil || perr != nil {
http.Error(w, "forbidden", http.StatusForbidden)
return
}
a = a.Unmap()
ok := a.IsLoopback()
for _, p := range allowed {
ok = ok || p.Contains(a)
}
if !ok {
http.Error(w, "this mirror is managed by its control plane only", http.StatusForbidden)
return
}
next.ServeHTTP(w, r)
})
}
+39
View File
@@ -0,0 +1,39 @@
package cli
import (
"net/http"
"net/http/httptest"
"net/netip"
"testing"
)
func TestTheMirrorAPIAnswersOnlyThisMachineUnlessTold(t *testing.T) {
ok := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {})
call := func(h http.Handler, from string) int {
r := httptest.NewRequest("DELETE", "/v1/models/m/v", nil)
r.RemoteAddr = from + ":5000"
w := httptest.NewRecorder()
h.ServeHTTP(w, r)
return w.Code
}
cp := []netip.Prefix{netip.MustParsePrefix("10.0.0.140/32")}
// Nothing allowed and no TLS used to mean anyone: every node could delete
// models.
if got := call(apiGate(false, nil, ok), "10.0.0.141"); got != http.StatusForbidden {
t.Errorf("a node reached an unconfigured mirror API: %d", got)
}
if got := call(apiGate(false, nil, ok), "127.0.0.1"); got != http.StatusOK {
t.Errorf("the mirror's own machine was refused: %d", got)
}
if got := call(apiGate(false, cp, ok), "10.0.0.140"); got != http.StatusOK {
t.Errorf("the allowed control plane was refused: %d", got)
}
if got := call(apiGate(false, cp, ok), "10.0.0.141"); got != http.StatusForbidden {
t.Errorf("a node outside the allowlist got through: %d", got)
}
// With mutual TLS the certificate decides.
if got := call(apiGate(true, nil, ok), "10.0.0.141"); got != http.StatusOK {
t.Errorf("mTLS without an allowlist refused a client: %d", got)
}
}
+125
View File
@@ -0,0 +1,125 @@
package cli
import (
"context"
"errors"
"fmt"
"io"
"net"
"os"
"runtime"
"strings"
"time"
"github.com/spf13/cobra"
"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)
// mirror init sets a machine up as the cluster's model mirror, on its own:
// the bytes for every node on :8081, the management API on :8082 for the
// control plane alone. server init does the same beside the control plane;
// this is for when the mirror wants its own disk.
func newMirrorInitCmd(dir *string) *cobra.Command {
var (
controlPlanes []string
base string
)
cmd := &cobra.Command{
Use: "init",
Short: "Set this machine up as the cluster's model mirror and start it",
Long: `Creates the llmgw user and the mirror's directory, and a systemd unit that
serves the model files to every node on :8081 and the management API on
:8082, and starts it.
The management API can fetch and delete models, so it answers this machine
and the control planes given with --control-plane, and nobody else. Then on
the control plane:
sudo llmgw server init --mirror http://<this machine>:8082
Safe to run again: it changes only what differs.`,
Example: ` sudo llmgw mirror init --control-plane 10.0.0.10`,
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, _ []string) error {
if runtime.GOOS != "linux" {
return fmt.Errorf("mirror init sets up a Linux machine; this is %s", runtime.GOOS)
}
if os.Geteuid() != 0 {
return errors.New("mirror init creates a user, writes /etc and starts a service: run it as root")
}
if len(controlPlanes) == 0 {
return errors.New("--control-plane: which control plane manages this mirror (its address)")
}
if _, err := parseAllow(controlPlanes); err != nil {
return err
}
return mirrorInit(cmd.Context(), cmd.OutOrStdout(), *dir, base, controlPlanes)
},
}
cmd.Flags().StringArrayVar(&controlPlanes, "control-plane", nil,
"address of the control plane that manages this mirror; repeatable")
cmd.Flags().StringVar(&base, "base", "",
"the address nodes fetch from, http://host:8081 (default: this machine's first address)")
return cmd
}
func mirrorInit(ctx context.Context, out io.Writer, dir, base string, controlPlanes []string) error {
in := &serverInit{user: "llmgw", stateDir: "/var/lib/llmgw", out: out}
if err := in.ensureUser(ctx); err != nil {
return err
}
if err := os.MkdirAll(dir, 0o755); err != nil {
return err
}
if err := in.chown(dir, true); err != nil {
return err
}
addr := ""
if ips := localAddrs(); len(ips) > 0 {
addr = ips[0].String()
}
if base == "" {
if addr == "" {
return errors.New("this machine has no address to give nodes; give one with --base")
}
base = "http://" + net.JoinHostPort(addr, "8081")
}
unit := mirrorUnitFor(dir, ":8082", base, controlPlanes,
"'llmgw mirror init'. Bytes on :8081 for every node; the\n# management API on :8082, for the control plane only.")
changed, err := in.writeUnit(mirrorUnitPath, unit)
if err != nil {
return err
}
for _, args := range [][]string{{"daemon-reload"}, {"enable", "--now", "llmgw-mirror"}} {
if err := in.systemctl(ctx, args...); err != nil {
return err
}
}
if changed {
if err := in.systemctl(ctx, "restart", "llmgw-mirror"); err != nil {
return err
}
in.done("model mirror in %s: files on :8081 at %s, management on :8082 for %s", dir, base, strings.Join(controlPlanes, ", "))
} else {
in.kept("model mirror in %s", dir)
}
m, err := mirror.NewClient("http://127.0.0.1:8082", tlsconf.Config{}, 5*time.Second)
if err != nil {
return err
}
if err := waitFor(ctx, 30*time.Second, "the mirror to answer", func() (bool, error) {
return m.Health(ctx) == nil, nil
}); err != nil {
return err
}
in.kept("answering")
api := "http://<this machine>:8082"
if addr != "" {
api = "http://" + net.JoinHostPort(addr, "8082")
}
fmt.Fprintf(out, "\nThe mirror is up. On the control plane:\n\n sudo llmgw server init --mirror %s\n\n", api)
return nil
}
+118 -27
View File
@@ -6,6 +6,7 @@ import (
"fmt"
"io"
"os"
"slices"
"strings"
"time"
@@ -64,12 +65,15 @@ type nodeAdd struct {
noWait bool
plain bool
serverCfg string
// named are the fabric names given in this run, by what the adapters
// report the fabric to be.
named map[string]string
}
func newNodeAddCmd() *cobra.Command {
var na nodeAdd
cmd := &cobra.Command{
Use: "add [user@]address",
Use: "add [user@]address...",
Short: "Make a machine a node: check it, give it an identity, start its agent",
Long: `Run on the machine holding the certificate authority — the control plane,
after 'llmgw server init' — as root. It logs in to the machine with ssh,
@@ -87,14 +91,18 @@ Nothing else needs giving, except --fabric for InfiniBand or RoCE: which
fabric an adapter is on has a name only you can choose, and it must be the
same on every node attached to it. Doctor suggests one when it sees adapters.
Several addresses add several nodes, one after another; a fabric named for
one is recognised on the next and not asked again.
The user needs sudo on the machine without a password. Run it again at any
time: the node gets a fresh certificate and configuration, and its running
engines are kept.`,
Example: ` sudo llmgw node add 10.0.0.21
sudo llmgw node add ubuntu@10.0.0.21 ubuntu@10.0.0.22 ubuntu@10.0.0.23
sudo llmgw node add ubuntu@gpu01 --fabric ib-a:mlx5_0,mlx5_1,mlx5_2,mlx5_3`,
Args: cobra.ExactArgs(1),
Args: cobra.MinimumNArgs(1),
RunE: func(cmd *cobra.Command, args []string) error {
return na.run(cmd.Context(), args[0], cmd.OutOrStdout())
return na.run(cmd.Context(), args, cmd.OutOrStdout())
},
}
f := cmd.Flags()
@@ -120,18 +128,9 @@ engines are kept.`,
return cmd
}
func (na *nodeAdd) prepare(addr string) error {
// prepare works out what every node added in this run shares.
func (na *nodeAdd) prepare() error {
o := &na.opts
o.Addr = addr
if user, host, ok := strings.Cut(addr, "@"); ok {
o.User, o.Addr = user, host
}
if o.User == "" {
// Under sudo, the person who ran it — their account is usually the
// one that exists on the nodes too. root's own keys still do the
// logging in, since that is who is running ssh.
o.User = os.Getenv("SUDO_USER")
}
for _, s := range na.fabrics {
fab, err := onboard.ParseFabric(s)
if err != nil {
@@ -158,32 +157,124 @@ func (na *nodeAdd) prepare(addr string) error {
o.Binary = exe
}
onboard.AgentUnit = agentUnit
na.named = map[string]string{}
return nil
}
func (na *nodeAdd) run(ctx context.Context, addr string, out io.Writer) error {
if err := na.prepare(addr); err != nil {
// target is the options for one address, [user@]host.
func (na *nodeAdd) target(addr string) onboard.Options {
o := na.opts
o.Fabrics = slices.Clone(na.opts.Fabrics)
o.Addr = addr
if user, host, ok := strings.Cut(addr, "@"); ok {
o.User, o.Addr = user, host
}
if o.User == "" {
// Under sudo, the person who ran it — their account is usually the
// one that exists on the nodes too. root's own keys still do the
// logging in, since that is who is running ssh.
o.User = os.Getenv("SUDO_USER")
}
return o
}
// rememberFabrics wraps the question of which fabric a node's adapters are
// on, so a name given for one node is used for the next whose adapters
// report the same fabric, without asking again. ask nil is a script, which
// is asked nothing and gets only what an earlier answer said.
func (na *nodeAdd) rememberFabrics(ask func([]onboard.Fabric) []onboard.Fabric) func([]onboard.Fabric) []onboard.Fabric {
return func(s []onboard.Fabric) []onboard.Fabric {
var known []onboard.Fabric
all := true
for i, f := range s {
if name, ok := na.named[f.Key]; ok && f.Key != "" && f.Key != "unknown" {
f.Name = name
s[i].Name = name
known = append(known, f)
} else {
all = false
}
}
if all || ask == nil {
return known
}
got := ask(s)
for _, f := range got {
if f.Key != "" && f.Key != "unknown" {
na.named[f.Key] = f.Name
}
}
return got
}
}
func (na *nodeAdd) run(ctx context.Context, addrs []string, out io.Writer) error {
if len(addrs) > 1 && na.opts.Node != "" {
return errors.New("--node names one node; give one address with it")
}
if err := na.prepare(); err != nil {
return err
}
if !na.plain && isTerminal(out) {
return na.runTUI(ctx)
tui := !na.plain && isTerminal(out)
type outcome struct {
addr, node string
err error
}
var results []outcome
for i, addr := range addrs {
o := na.target(addr)
title := "llmgw · adding " + o.Addr
if len(addrs) > 1 {
title += fmt.Sprintf(" (%d of %d)", i+1, len(addrs))
}
var node string
var err error
if tui {
node, err = na.runTUI(ctx, o, title)
} else {
node, err = na.runPlain(ctx, o, title, out)
}
results = append(results, outcome{addr, node, err})
if ctx.Err() != nil {
break
}
}
if len(addrs) == 1 {
return results[0].err
}
fmt.Fprintln(out)
failed := 0
for _, r := range results {
if r.err != nil {
failed++
fmt.Fprintf(out, " FAIL %-20s %s\n", r.addr, r.err)
} else {
fmt.Fprintf(out, " ok %-20s %s\n", r.addr, r.node)
}
}
if failed > 0 {
return fmt.Errorf("%d of %d not added; running the same command again finishes them and leaves the rest as they are", failed, len(addrs))
}
return nil
}
func (na *nodeAdd) runPlain(ctx context.Context, o onboard.Options, title string, out io.Writer) (string, error) {
p := &plainProgress{out: out}
fmt.Fprintf(out, "adding %s\n", na.opts.Addr)
r := onboard.New(na.opts, p.event)
res, err := r.Run(ctx)
fmt.Fprintf(out, "%s\n", title)
o.AskFabrics = na.rememberFabrics(nil)
res, err := onboard.New(o, p.event).Run(ctx)
if err != nil {
return err
return "", err
}
if na.noWait {
fmt.Fprintf(out, "\n%s is added. 'llmgw ctl nodes' shows it once it joins.\n", res.Node)
return nil
fmt.Fprintf(out, "\n%s is added. 'llmgw ctl nodes' shows it once it joins.\n\n", res.Node)
return res.Node, nil
}
if err := na.watch(ctx, res.Node, p.event); err != nil {
return err
return res.Node, err
}
fmt.Fprintf(out, "\n%s is in the cluster.\n", res.Node)
return nil
fmt.Fprintf(out, "\n%s is in the cluster.\n\n", res.Node)
return res.Node, nil
}
// adminAPI is the control plane's admin API, from its configuration when
+29 -1
View File
@@ -21,11 +21,13 @@ import (
"golang.org/x/net/http2/h2c"
"golang.org/x/sync/errgroup"
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fabrictest"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fleetmetrics"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
@@ -244,7 +246,7 @@ func newServerCmd() *cobra.Command {
// enforces its own per-phase deadlines.
}
adminSrv := &http.Server{
Handler: adminMux,
Handler: router.AdminGuard(adminMux),
ReadHeaderTimeout: 10 * time.Second,
}
@@ -383,6 +385,20 @@ func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Live
ccfg.Reconciler.PortRange = [2]int{cfg.Control.PortRange[0], cfg.Control.PortRange[1]}
}
// Each replica's engine demands a key only the router has. The secret it
// is derived from lives in the store, so every control plane against it —
// and this one after a restart — computes the same key for a replica.
if h, ok := sink.(interface{ Store() *storepkg.DB }); ok && h.Store() != nil {
secret, err := h.Store().ClusterSecret(ctx, "engine-key")
if err != nil {
return nil, nil, fmt.Errorf("reading the engine key secret: %w", err)
}
ccfg.Reconciler.EngineKey = control.EngineKeys([]byte(secret))
} else {
log.Warn("engines start without a key: there is no store to keep the secret " +
"they are derived from, so anyone who can reach a node's ports reaches its models")
}
// 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
@@ -493,6 +509,18 @@ func (r *dispatcherRef) Logs(ctx context.Context, nodeID, source string, lines i
return lr.Logs(ctx, nodeID, source, lines)
}
func (r *dispatcherRef) FabricTest(ctx context.Context, nodeID string, req *llmgwv1.FabricTestRequest) (*llmgwv1.FabricTestReply, error) {
d, err := r.get()
if err != nil {
return nil, err
}
ft, ok := d.(fabrictest.Asker)
if !ok {
return nil, errors.New("this dispatcher cannot run a bandwidth test")
}
return ft.FabricTest(ctx, nodeID, req)
}
func (r *dispatcherRef) StartShard(ctx context.Context, cmd reconciler.ShardCommand) error {
d, err := r.get()
if err != nil {
+23 -9
View File
@@ -73,6 +73,8 @@ type serverInit struct {
// make it serve something: a mirror, this machine as a node, and a small
// model on it.
noMirror, noNode, noModel bool
// mirrorURL is a mirror on another machine to use instead of running one.
mirrorURL string
// pgSetup prepares PostgreSQL and returns the store address; a seam for
// tests. Nil does it for real.
pgSetup func(context.Context) (string, error)
@@ -92,6 +94,7 @@ func newServerInitCmd() *cobra.Command {
unitPath: "/etc/systemd/system/llmgw-server.service",
user: "llmgw",
}
var plain bool
cmd := &cobra.Command{
Use: "init",
Short: "Set this machine up as a control plane and start it",
@@ -122,15 +125,21 @@ func newServerInitCmd() *cobra.Command {
return errors.New("init creates a user, writes /etc and starts a service: run it as root")
}
in.out = cmd.OutOrStdout()
if !plain && isTerminal(in.out) {
return in.runInitTUI(cmd.Context(), in.out)
}
return in.run(cmd.Context())
},
}
f := cmd.Flags()
f.BoolVar(&plain, "plain", false, "plain output, even in a terminal")
f.StringArrayVar(&in.hosts, "host", nil,
"an address agents dial besides the ones found on this machine; repeatable")
f.StringVar(&in.admin, "admin", "", "email of the first administrator; a password is generated and printed once")
f.StringVar(&in.pkiDir, "pki", in.pkiDir, "directory holding the certificate authority")
f.BoolVar(&in.noMirror, "no-mirror", false, "do not run a model mirror on this machine")
f.StringVar(&in.mirrorURL, "mirror", "",
"use the mirror on another machine, set up there by 'llmgw mirror init': its management API, http://host:8082")
f.BoolVar(&in.noNode, "no-node", false, "do not make this machine a node (and serve no model on it)")
f.BoolVar(&in.noModel, "no-model", false, "do not fetch and serve the starter model ("+starterName+")")
f.BoolVar(&in.postgres, "postgres", false,
@@ -163,7 +172,12 @@ func (in *serverInit) run(ctx context.Context) error {
hasSystemd = false
}
// Before the service starts, because the mirror goes into server.yaml.
if hasSystemd && !in.noMirror {
switch {
case in.mirrorURL != "":
if err := in.useMirror(ctx); err != nil {
return fmt.Errorf("mirror: %w", err)
}
case hasSystemd && !in.noMirror:
if err := in.ensureMirror(ctx); err != nil {
return fmt.Errorf("mirror: %w", err)
}
@@ -175,14 +189,14 @@ func (in *serverInit) run(ctx context.Context) error {
if _, err := in.ensureNode(ctx); err != nil {
return fmt.Errorf("node: %w", err)
}
switch {
case in.noModel:
case in.noMirror:
fmt.Fprintln(in.out, " skip no starter model without the mirror")
default:
if err := in.ensureStarter(ctx); err != nil {
return fmt.Errorf("starter model: %w", err)
}
}
switch {
case !hasSystemd || in.noModel:
case in.noMirror && in.mirrorURL == "":
fmt.Fprintln(in.out, " skip no starter model without the mirror")
default:
if err := in.ensureStarter(ctx); err != nil {
return fmt.Errorf("starter model: %w", err)
}
}
in.nextSteps()
+79 -16
View File
@@ -111,15 +111,25 @@ func (in *serverInit) writeUnit(path, body string) (bool, error) {
func (in *serverInit) mirrorUnit() string {
base := "http://" + net.JoinHostPort(in.reachableAddr(), "8081")
return mirrorUnitFor(mirrorDir, mirrorAPI, base, nil,
"'llmgw server init'. Bytes on :8081 for every node; the\n# management API on loopback only, for the control plane beside it.")
}
// mirrorUnitFor is a mirror's systemd unit. allow is who, besides the
// machine itself, may use the management API.
func mirrorUnitFor(dir, apiListen, base string, allow []string, about string) string {
exec := "/usr/local/bin/llmgw mirror serve --dir " + dir + " --listen :8081 --api-listen " + apiListen + " --base " + base
for _, a := range allow {
exec += " --api-allow " + a
}
return `[Unit]
Description=llmgw model mirror
After=network-online.target
Wants=network-online.target
[Service]
# Written by 'llmgw server init'. Bytes on :8081 for every node; the
# management API on loopback only, for the control plane beside it.
ExecStart=/usr/local/bin/llmgw mirror serve --dir ` + mirrorDir + ` --listen :8081 --api-listen ` + mirrorAPI + ` --base ` + base + `
# Written by ` + about + `
ExecStart=` + exec + `
User=llmgw
Group=llmgw
Restart=on-failure
@@ -128,13 +138,53 @@ NoNewPrivileges=true
ProtectSystem=strict
ProtectHome=true
PrivateTmp=true
ReadWritePaths=` + mirrorDir + `
ReadWritePaths=` + dir + `
[Install]
WantedBy=multi-user.target
`
}
// mirrorAPIURL is the mirror's management API: one on its own machine, or
// the one beside this control plane.
func (in *serverInit) mirrorAPIURL() string {
if in.mirrorURL != "" {
return in.mirrorURL
}
return "http://" + mirrorAPI
}
// useMirror points the control plane at a mirror on another machine, set up
// there by 'llmgw mirror init'.
func (in *serverInit) useMirror(ctx context.Context) error {
m, err := mirror.NewClient(in.mirrorURL, tlsconf.Config{}, 5*time.Second)
if err != nil {
return err
}
if err := m.Health(ctx); err != nil {
return fmt.Errorf("the mirror at %s does not answer: %v. On that machine: sudo llmgw mirror init --control-plane <this machine's address>", in.mirrorURL, err)
}
cfg, err := config.Load(in.configPath())
if err != nil {
return err
}
switch cfg.Mirror.API {
case in.mirrorURL:
in.kept("mirror at %s", in.mirrorURL)
return nil
case "":
default:
in.kept("mirror at %s, as server.yaml says; change it there to use %s", cfg.Mirror.API, in.mirrorURL)
return nil
}
if err := appendMirrorBlockFor(in.configPath(), in.mirrorURL,
"# The model mirror, on its own machine, added by 'llmgw server init'."); err != nil {
return err
}
in.done("server.yaml points at the mirror at %s", in.mirrorURL)
return nil
}
// ensureMirror runs a mirror on this machine and points the control plane at
// it. Its management API is on loopback, which is why it needs no certificate:
// nothing but this machine can reach it.
@@ -174,7 +224,9 @@ func (in *serverInit) ensureMirror(ctx context.Context) error {
}
if cfg.Mirror.API == "" {
if err := appendMirrorBlock(in.configPath()); err != nil {
if err := appendMirrorBlockFor(in.configPath(), "http://"+mirrorAPI,
"# The model mirror on this machine, added by 'llmgw server init'. Its\n"+
"# management API is on loopback, so it needs no certificate."); err != nil {
return err
}
in.done("server.yaml points at the mirror")
@@ -189,9 +241,9 @@ func (in *serverInit) ensureMirror(ctx context.Context) error {
})
}
// appendMirrorBlock adds the mirror to server.yaml without touching the rest,
// keeping the file as it was beside it.
func appendMirrorBlock(path string) error {
// appendMirrorBlockFor adds the mirror to server.yaml without touching the
// rest, keeping the file as it was beside it.
func appendMirrorBlockFor(path, api, comment string) error {
b, err := os.ReadFile(path)
if err != nil {
return err
@@ -200,9 +252,7 @@ func appendMirrorBlock(path string) error {
if err := os.WriteFile(backup, b, 0o640); err != nil {
return err
}
block := "\n# The model mirror on this machine, added by 'llmgw server init'. Its\n" +
"# management API is on loopback, so it needs no certificate.\n" +
"mirror:\n api: http://" + mirrorAPI + "\n"
block := "\n" + comment + "\nmirror:\n api: " + api + "\n"
info, err := os.Stat(path)
if err != nil {
return err
@@ -356,8 +406,11 @@ func hasNVIDIADevice() bool {
// starterModel is the catalog entry: a GPU profile first when there is a card,
// and a CPU profile always, both pinned to this machine.
func starterModel(source string, gpuClasses []string) catalog.Model {
pin := map[string]string{localLabelKey: localLabelValue}
func starterModel(source string, gpuClasses []string, pinned bool) catalog.Model {
var pin map[string]string
if pinned {
pin = map[string]string{localLabelKey: localLabelValue}
}
var profiles []catalog.Profile
for _, class := range gpuClasses {
profiles = append(profiles, catalog.Profile{
@@ -401,7 +454,7 @@ func (in *serverInit) ensureStarter(ctx context.Context) error {
if !ok {
return fmt.Errorf("cannot turn %s into a download", starterSource)
}
m, err := mirror.NewClient("http://"+mirrorAPI, tlsconf.Config{}, 15*time.Second)
m, err := mirror.NewClient(in.mirrorAPIURL(), tlsconf.Config{}, 15*time.Second)
if err != nil {
return err
}
@@ -444,14 +497,24 @@ func (in *serverInit) ensureStarter(ctx context.Context) error {
}
in.done("%s in the mirror", starterName)
classes, _ := gpuClasses(ctx)
body, err := json.Marshal(starterModel(st.URL, classes))
// Served here when this machine is a node, pinned so it does not wander
// to the first GPU node that joins; otherwise on whichever node joins.
pinned := !in.noNode
var classes []string
if pinned {
classes, _ = gpuClasses(ctx)
}
body, err := json.Marshal(starterModel(st.URL, classes, pinned))
if err != nil {
return err
}
if err := client.post(ctx, "/api/v1/models", body, &map[string]string{}); err != nil {
return fmt.Errorf("adding %s to the catalog: %w", starterName, err)
}
if !pinned {
in.done("%s in the catalog; it starts on the first node that joins", starterName)
return nil
}
where := "on CPU"
if len(classes) > 0 {
where = "on " + strings.Join(classes, ", ") + ", or CPU if the card is busy"
+37 -3
View File
@@ -1,6 +1,8 @@
package cli
import (
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
@@ -21,7 +23,7 @@ func TestTheInitAgentUnitIsTheShippedUnit(t *testing.T) {
}
func TestTheStarterUsesTheGPUAndFallsBackToCPUOnThisMachine(t *testing.T) {
m := starterModel("http://10.0.0.5:8081/qwen2.5-0.5b-instruct/starter/q.gguf", []string{"4090-24g"})
m := starterModel("http://10.0.0.5:8081/qwen2.5-0.5b-instruct/starter/q.gguf", []string{"4090-24g"}, true)
c, err := catalog.New([]catalog.Model{m})
if err != nil {
t.Fatalf("the starter is not a valid catalog entry: %v", err)
@@ -46,7 +48,7 @@ func TestTheStarterUsesTheGPUAndFallsBackToCPUOnThisMachine(t *testing.T) {
t.Error("the starter should be running, not waiting to be asked for")
}
if m := starterModel("http://x/q.gguf", nil); len(m.Profiles) != 1 || m.Profiles[0].Name != "cpu" {
if m := starterModel("http://x/q.gguf", nil, true); len(m.Profiles) != 1 || m.Profiles[0].Name != "cpu" {
t.Errorf("without a GPU only the CPU profile: %+v", m.Profiles)
}
}
@@ -56,7 +58,7 @@ func TestTheMirrorBlockIsAddedAndReads(t *testing.T) {
if err := in.run(t.Context()); err != nil {
t.Fatalf("%v\n%s", err, out)
}
if err := appendMirrorBlock(in.configPath()); err != nil {
if err := appendMirrorBlockFor(in.configPath(), "http://"+mirrorAPI, "# the mirror"); err != nil {
t.Fatal(err)
}
cfg, err := config.Load(in.configPath())
@@ -87,3 +89,35 @@ func TestTheMirrorUnitAdvertisesAnAddressNodesCanReach(t *testing.T) {
t.Errorf("the management API is not on loopback: %s", u)
}
}
func TestWithNoNodeHereTheStarterGoesToTheFirstNodeThatJoins(t *testing.T) {
m := starterModel("http://x/q.gguf", nil, false)
for _, p := range m.Profiles {
if len(p.NodeSelector) > 0 {
t.Errorf("profile %s is pinned to a machine that is not a node: %v", p.Name, p.NodeSelector)
}
}
}
func TestTheMirrorAPIOnlyAnswersItsControlPlane(t *testing.T) {
allowed, err := parseAllow([]string{"10.0.0.140", "10.1.0.0/24"})
if err != nil {
t.Fatal(err)
}
h := onlyFrom(allowed, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {}))
for addr, want := range map[string]int{
"10.0.0.140:5000": 200, "10.1.0.7:5000": 200, "127.0.0.1:5000": 200, "[::1]:5000": 200,
"10.0.0.141:5000": 403, "192.168.1.9:5000": 403,
} {
r := httptest.NewRequest("GET", "/api/v1/health", nil)
r.RemoteAddr = addr
w := httptest.NewRecorder()
h.ServeHTTP(w, r)
if w.Code != want {
t.Errorf("%s: %d, want %d", addr, w.Code, want)
}
}
if _, err := parseAllow([]string{"not-an-address"}); err == nil {
t.Error("nonsense was accepted")
}
}
+421
View File
@@ -0,0 +1,421 @@
package cli
import (
"context"
"fmt"
"io"
"net/mail"
"os"
"path/filepath"
"strings"
"sync"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
// The full-screen way through 'llmgw server init': a short form for what
// init would otherwise take as flags, then init's own lines as a list that
// fills in, with the administrator's password kept in view — and printed
// again once the screen closes, so it is in the terminal's history.
// initLineMsg is one line init wrote.
type initLineMsg string
// initDoneMsg is init finishing.
type initDoneMsg struct{ err error }
// lineWriter turns init's output into one message per line.
type lineWriter struct {
mu sync.Mutex
buf []byte
send func(string)
}
func (w *lineWriter) Write(p []byte) (int, error) {
w.mu.Lock()
defer w.mu.Unlock()
w.buf = append(w.buf, p...)
for {
i := strings.IndexByte(string(w.buf), '\n')
if i < 0 {
return len(p), nil
}
w.send(string(w.buf[:i]))
w.buf = w.buf[i+1:]
}
}
// Form fields, in order.
const (
fAdmin = iota
fHosts
fPostgres
fMirror
fNode
fModel
fStart
fCount
)
type initModel struct {
in *serverInit
found []string // this machine's addresses
// The form, until it is submitted.
form bool
focus int
admin textField
hosts textField
toggles [fCount]bool
formError string
// Then the run.
running bool
lines []string // init's lines, as it wrote them
password string
next []string // the closing "Next:" text
spin spinnerModel
width int
done bool
err error
quit bool
start func()
}
func newInitModel(in *serverInit, found []string, askForm bool) initModel {
admin := newTextField()
admin.Placeholder = "you@example.com"
admin.SetValue(in.admin)
admin.CharLimit = 200
admin.Width = 40
hosts := newTextField()
hosts.Placeholder = "llm.example.com, 203.0.113.10"
hosts.SetValue(strings.Join(in.hosts, ", "))
hosts.CharLimit = 400
hosts.Width = 40
sp := newSpinner()
m := initModel{in: in, found: found, form: askForm, admin: admin, hosts: hosts, spin: sp, width: 100}
m.toggles[fPostgres] = in.postgres
m.toggles[fMirror] = !in.noMirror
m.toggles[fNode] = !in.noNode
m.toggles[fModel] = !in.noModel
if askForm {
m.admin.Focus()
}
return m
}
func (m initModel) Init() tea.Cmd {
if m.form {
return nil
}
return tea.Batch(m.spin.Tick, func() tea.Msg { return startInitMsg{} })
}
type startInitMsg struct{}
func (m initModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.WindowSizeMsg:
m.width = msg.Width
case spinTickMsg:
var cmd tea.Cmd
m.spin, cmd = m.spin.Update(msg)
return m, cmd
case startInitMsg:
m.running = true
if m.start != nil {
m.start()
}
case initLineMsg:
m.addLine(string(msg))
case initDoneMsg:
m.done, m.err = true, msg.err
return m, tea.Quit
case tea.KeyMsg:
if msg.Type == tea.KeyCtrlC {
m.quit = true
return m, tea.Quit
}
if m.form {
return m.formKey(msg)
}
}
return m, nil
}
// addLine files one of init's lines: the password aside, the closing text
// aside, a progress line replacing the one before it.
func (m *initModel) addLine(l string) {
t := strings.TrimSpace(l)
switch {
case t == "":
return
case strings.HasPrefix(t, "password: "):
m.password = strings.TrimPrefix(t, "password: ")
return
case strings.HasPrefix(t, "shown once"):
return
case !strings.HasPrefix(l, " ") || len(m.next) > 0:
// Unindented lines are the closing text, and everything after its
// first line belongs to it.
m.next = append(m.next, l)
return
}
if strings.HasPrefix(t, "..") && len(m.lines) > 0 && strings.HasPrefix(strings.TrimSpace(m.lines[len(m.lines)-1]), "..") {
m.lines[len(m.lines)-1] = l
return
}
m.lines = append(m.lines, l)
}
func (m initModel) formKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
switch k.Type {
case tea.KeyTab, tea.KeyDown:
m.moveFocus(1)
return m, nil
case tea.KeyShiftTab, tea.KeyUp:
m.moveFocus(-1)
return m, nil
case tea.KeyEnter:
if m.focus != fStart {
m.moveFocus(1)
return m, nil
}
admin := strings.TrimSpace(m.admin.Value())
if admin != "" {
if _, err := mail.ParseAddress(admin); err != nil {
m.formError = "the administrator's email is not an address"
m.focus = fAdmin
m.admin.Focus()
return m, nil
}
}
m.in.admin = admin
m.in.hosts = nil
for _, h := range strings.Split(m.hosts.Value(), ",") {
if h = strings.TrimSpace(h); h != "" {
m.in.hosts = append(m.in.hosts, h)
}
}
m.in.postgres = m.toggles[fPostgres]
m.in.noMirror = !m.toggles[fMirror]
m.in.noNode = !m.toggles[fNode]
m.in.noModel = !m.toggles[fModel]
m.form = false
return m, tea.Batch(m.spin.Tick, func() tea.Msg { return startInitMsg{} })
}
switch m.focus {
case fAdmin:
var cmd tea.Cmd
m.admin, cmd = m.admin.Update(k)
return m, cmd
case fHosts:
var cmd tea.Cmd
m.hosts, cmd = m.hosts.Update(k)
return m, cmd
case fPostgres, fMirror, fNode, fModel:
if k.String() == " " || k.String() == "x" {
m.toggles[m.focus] = !m.toggles[m.focus]
}
}
return m, nil
}
func (m *initModel) moveFocus(d int) {
m.focus = (m.focus + d + fCount) % fCount
// With --mirror there is no mirror to run here, so nothing to tick.
if m.focus == fMirror && m.in.mirrorURL != "" {
m.focus = (m.focus + d + fCount) % fCount
}
m.admin.Blur()
m.hosts.Blur()
switch m.focus {
case fAdmin:
m.admin.Focus()
case fHosts:
m.hosts.Focus()
}
}
func (m initModel) View() string {
var b strings.Builder
b.WriteString(sTitle.Render("llmgw · setting up this machine as a control plane") + "\n\n")
if m.form {
return b.String() + m.formView()
}
for _, l := range m.visibleLines() {
b.WriteString(m.lineView(l) + "\n")
}
if m.password != "" {
box := sBox.BorderForeground(tuiWarn).Render(
fmt.Sprintf("Administrator %s\npassword %s\n%s", m.in.admin,
lipgloss.NewStyle().Bold(true).Render(m.password),
sDim.Render("Shown once. Copy it now; change it on the account page after signing in.")))
b.WriteString("\n" + box + "\n")
}
switch {
case m.done && m.err != nil:
b.WriteString("\n" + sBad.Render("✗ ") + wrap(m.err.Error(), m.width-4, " ") + "\n" +
sDim.Render(" It is safe to run init again once that is fixed: it finishes what is missing.") + "\n")
case m.done:
b.WriteString("\n" + sOK.Render("✓ ") + "The control plane is up. Add a node: " +
lipgloss.NewStyle().Bold(true).Render("sudo llmgw node add <user>@<address>") + "\n")
b.WriteString(" " + sDim.Render("Back up "+m.in.pkiDir+": its ca-key.pem signs every node's identity, and nothing recovers it.") + "\n")
case m.quit:
b.WriteString("\n" + sDim.Render("stopped; run init again to finish") + "\n")
case m.running:
b.WriteString("\n" + sDim.Render("ctrl+c to stop") + "\n")
}
return b.String()
}
// visibleLines are the last lines, enough to follow without the screen
// outgrowing the terminal: the whole log is printed when it closes.
func (m initModel) visibleLines() []string {
const keep = 16
if len(m.lines) > keep {
return m.lines[len(m.lines)-keep:]
}
return m.lines
}
func (m initModel) lineView(l string) string {
t := strings.TrimSpace(l)
word, rest, _ := strings.Cut(t, " ")
rest = truncate(strings.TrimSpace(rest), max(20, m.width-6))
switch word {
case "ok":
return " " + sOK.Render("✓") + " " + rest
case "kept":
return " " + sDim.Render("· "+rest)
case "skip":
return " " + sDim.Render("- "+rest)
case "!!":
return " " + sWarn.Render("! ") + rest
case "..":
if m.done {
return " " + sDim.Render("· "+rest)
}
return " " + m.spin.View() + " " + rest
}
return " " + t
}
func (m initModel) formView() string {
var b strings.Builder
label := func(i int, s string) string {
if m.focus == i {
return sTitle.Render("› " + s)
}
return " " + s
}
check := func(i int, s, hint string) string {
box := "[ ]"
if m.toggles[i] {
box = "[" + sOK.Render("x") + "]"
}
return fmt.Sprintf("%s %s %s\n", label(i, box), s, sDim.Render(hint))
}
fmt.Fprintf(&b, "%s\n %s\n", label(fAdmin, "Administrator's email"), m.admin.View())
b.WriteString(" " + sDim.Render("A password is generated and shown once. Leave empty to create accounts later.") + "\n\n")
fmt.Fprintf(&b, "%s\n %s\n", label(fHosts, "Other names nodes will reach this machine by"), m.hosts.View())
found := "none found"
if len(m.found) > 0 {
found = strings.Join(m.found, ", ")
}
b.WriteString(" " + sDim.Render("Already in the certificate: "+found+". Add a DNS name or NAT address here.") + "\n\n")
b.WriteString(check(fPostgres, "PostgreSQL", "rather than SQLite; installed here, for several control planes or a larger store"))
if m.in.mirrorURL != "" {
b.WriteString(" Model mirror " + sDim.Render("at "+m.in.mirrorURL+", from --mirror") + "\n")
} else {
b.WriteString(check(fMirror, "Model mirror", "downloads model files once for every node"))
}
b.WriteString(check(fNode, "This machine as a node", "serves models here too, on its GPUs or CPU"))
b.WriteString(check(fModel, "Starter model", "a small Qwen, so there is something to ask at once"))
b.WriteString("\n" + label(fStart, lipgloss.NewStyle().Bold(true).Render("Set up")) + "\n")
if m.formError != "" {
b.WriteString("\n" + sBad.Render(m.formError) + "\n")
}
b.WriteString("\n" + sDim.Render("tab or ↑↓ to move · space to tick · enter on Set up to start · ctrl+c to leave") + "\n")
return b.String()
}
// runInitTUI runs init with the screen above, asking the form first when
// the flags left the administrator unset.
func (in *serverInit) runInitTUI(ctx context.Context, out io.Writer) error {
var found []string
for _, ip := range localAddrs() {
found = append(found, ip.String())
}
m := newInitModel(in, found, in.admin == "")
var p *tea.Program
var (
logMu sync.Mutex
log []string
)
w := &lineWriter{send: func(l string) {
logMu.Lock()
log = append(log, l)
logMu.Unlock()
p.Send(initLineMsg(l))
}}
ctx, cancel := context.WithCancel(ctx)
defer cancel()
started := make(chan struct{})
m.start = func() {
close(started)
go func() {
in.out = w
err := in.run(ctx)
p.Send(initDoneMsg{err})
}()
}
p = tea.NewProgram(m, tea.WithOutput(out))
final, err := p.Run()
if err != nil {
return err
}
fm := final.(initModel)
select {
case <-started:
default:
return fmt.Errorf("left before anything was done")
}
if fm.quit && !fm.done {
cancel()
return fmt.Errorf("stopped; run init again to finish what was started")
}
// The whole run is kept, root's only, since it holds the password. On a
// failure it is printed too, for the lines the screen had no room for;
// otherwise the screen's last frame says it, and only what must outlive
// the screen is printed again: the password, and what to do next.
logMu.Lock()
defer logMu.Unlock()
whole := strings.Join(log, "\n") + "\n"
saved := ""
if in.stateDir != "" {
p := filepath.Join(in.stateDir, "init.log")
if err := os.WriteFile(p, []byte(whole), 0o600); err == nil {
saved = p
}
}
fmt.Fprintln(out)
if fm.err != nil {
fmt.Fprint(out, whole)
return fm.err
}
if fm.password != "" {
fmt.Fprintf(out, "Administrator %s, password %s (shown once)\n\n", in.admin, fm.password)
}
for _, l := range fm.next {
fmt.Fprintln(out, l)
}
if saved != "" {
fmt.Fprintf(out, "\nThe whole run is in %s.\n", saved)
}
return fm.err
}
+103
View File
@@ -0,0 +1,103 @@
package cli
import (
"strings"
"testing"
tea "github.com/charmbracelet/bubbletea"
)
func typeIn(m tea.Model, s string) tea.Model {
return send(m, tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(s)})
}
func TestTheInitFormBecomesInitsSettings(t *testing.T) {
in := &serverInit{pkiDir: "/etc/llmgw/pki"}
var m tea.Model = newInitModel(in, []string{"10.0.0.140"}, true)
if v := m.View(); !strings.Contains(v, "Already in the certificate: 10.0.0.140") {
t.Fatalf("the addresses found are not shown:\n%s", v)
}
m = typeIn(m, "ops@example.com")
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
m = typeIn(m, "llm.example.com, 203.0.113.10")
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(" ")}) // PostgreSQL on
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab},
tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(" ")}) // starter model off
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyEnter})
if in.admin != "ops@example.com" || strings.Join(in.hosts, "|") != "llm.example.com|203.0.113.10" {
t.Errorf("admin %q hosts %q", in.admin, in.hosts)
}
if !in.postgres || in.noMirror || in.noNode || !in.noModel {
t.Errorf("toggles: postgres %v mirror %v node %v model %v", in.postgres, !in.noMirror, !in.noNode, !in.noModel)
}
if m.(initModel).form {
t.Error("the form stayed up after Set up")
}
}
func TestAnEmailThatIsNotOneIsSaidSo(t *testing.T) {
in := &serverInit{}
var m tea.Model = newInitModel(in, nil, true)
m = typeIn(m, "not an address")
for range fStart {
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
}
m = send(m, tea.KeyMsg{Type: tea.KeyEnter})
if !m.(initModel).form || !strings.Contains(m.View(), "not an address") {
t.Errorf("a bad email was accepted:\n%s", m.View())
}
}
func TestInitsLinesAreFiledAndThePasswordKeptInView(t *testing.T) {
in := &serverInit{admin: "ops@example.com", pkiDir: "/etc/llmgw/pki"}
var m tea.Model = newInitModel(in, nil, false)
for _, l := range []string{
" ok system user llmgw",
" kept certificate authority in /etc/llmgw/pki",
" ok administrator ops@example.com (u1)",
"",
" password: Tr0ub4dor-3",
" shown once — sign in and change it from /profile",
"",
" .. downloading qwen2.5-0.5b-instruct into the mirror",
" .. 40%",
" .. 80%",
" !! an NVIDIA card is on the PCI bus but nvidia-smi does not answer",
"The control plane is up. Next:",
" llmgw ctl status",
} {
m = send(m, initLineMsg(l))
}
m = send(m, initDoneMsg{})
im := m.(initModel)
if im.password != "Tr0ub4dor-3" {
t.Errorf("password %q", im.password)
}
v := m.View()
for _, want := range []string{"✓ system user llmgw", "Tr0ub4dor-3", "Shown once", "80%", "! an NVIDIA card", "Back up /etc/llmgw/pki", "node add"} {
if !strings.Contains(v, want) {
t.Errorf("screen lacks %q:\n%s", want, v)
}
}
if strings.Contains(v, "40%") {
t.Error("progress lines piled up instead of replacing each other")
}
if strings.Contains(v, "llmgw ctl status") {
t.Error("the closing text was listed as a step")
}
}
func TestWithAMirrorElsewhereTheFormDoesNotOfferOne(t *testing.T) {
in := &serverInit{mirrorURL: "http://10.0.0.150:8082"}
var m tea.Model = newInitModel(in, nil, true)
if v := m.View(); !strings.Contains(v, "at http://10.0.0.150:8082, from --mirror") {
t.Errorf("the mirror in use is not shown:\n%s", v)
}
// Tab from the addresses goes past the mirror line to PostgreSQL, then
// the node.
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab})
if f := m.(initModel).focus; f != fNode {
t.Errorf("focus %d, want the node toggle %d", f, fNode)
}
}
+17 -20
View File
@@ -7,8 +7,6 @@ import (
"os"
"strings"
"github.com/charmbracelet/bubbles/spinner"
"github.com/charmbracelet/bubbles/textinput"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
"github.com/mattn/go-isatty"
@@ -67,7 +65,7 @@ type nodeModel struct {
title string
steps map[string]stepState
notes []onboard.Event // warnings and failures, with their fixes
spin spinner.Model
spin spinnerModel
cancel context.CancelFunc
width int
done bool
@@ -77,16 +75,14 @@ type nodeModel struct {
// The question being asked, if any.
fabrics *askFabricsMsg
inputs []textinput.Model
inputs []textField
focus int
confirm *askContinueMsg
quitAsked bool
}
func newNodeModel(title string, cancel context.CancelFunc) nodeModel {
sp := spinner.New()
sp.Spinner = spinner.MiniDot
sp.Style = lipgloss.NewStyle().Foreground(tuiAccent)
sp := newSpinner()
return nodeModel{title: title, steps: map[string]stepState{}, spin: sp, cancel: cancel, width: 100}
}
@@ -96,7 +92,7 @@ func (m nodeModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.WindowSizeMsg:
m.width = msg.Width
case spinner.TickMsg:
case spinTickMsg:
var cmd tea.Cmd
m.spin, cmd = m.spin.Update(msg)
return m, cmd
@@ -121,7 +117,7 @@ func (m nodeModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.fabrics = &msg
m.inputs = nil
for i, f := range msg.suggested {
in := textinput.New()
in := newTextField()
in.Prompt = ""
in.SetValue(f.Name)
in.CharLimit = 40
@@ -176,7 +172,7 @@ func (m nodeModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
var out []onboard.Fabric
for i, in := range m.inputs {
if name := strings.TrimSpace(in.Value()); name != "" {
out = append(out, onboard.Fabric{Name: name, HCAs: m.fabrics.suggested[i].HCAs})
out = append(out, onboard.Fabric{Name: name, HCAs: m.fabrics.suggested[i].HCAs, Key: m.fabrics.suggested[i].Key})
}
}
m.fabrics.reply <- out
@@ -297,15 +293,15 @@ func wrap(s string, width int, indent string) string {
return strings.Join(lines, "\n"+indent)
}
// runTUI adds the node with the screen above.
func (na *nodeAdd) runTUI(ctx context.Context) error {
// runTUI adds one node with the screen above.
func (na *nodeAdd) runTUI(ctx context.Context, o onboard.Options, title string) (string, error) {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
m := newNodeModel("llmgw · adding "+na.opts.Addr, cancel)
m := newNodeModel(title, cancel)
p := tea.NewProgram(m)
emit := func(e onboard.Event) { p.Send(evMsg(e)) }
na.opts.AskFabrics = func(s []onboard.Fabric) []onboard.Fabric {
o.AskFabrics = na.rememberFabrics(func(s []onboard.Fabric) []onboard.Fabric {
reply := make(chan []onboard.Fabric)
p.Send(askFabricsMsg{s, reply})
select {
@@ -314,8 +310,8 @@ func (na *nodeAdd) runTUI(ctx context.Context) error {
case <-ctx.Done():
return nil
}
}
na.opts.AskContinue = func(fails int) bool {
})
o.AskContinue = func(fails int) bool {
reply := make(chan bool)
p.Send(askContinueMsg{fails, reply})
select {
@@ -327,7 +323,7 @@ func (na *nodeAdd) runTUI(ctx context.Context) error {
}
go func() {
res, err := onboard.New(na.opts, emit).Run(ctx)
res, err := onboard.New(o, emit).Run(ctx)
node := ""
if res != nil {
node = res.Node
@@ -344,11 +340,12 @@ func (na *nodeAdd) runTUI(ctx context.Context) error {
final, err := p.Run()
if err != nil {
return err
return "", err
}
fm := final.(nodeModel)
if fm.quitAsked && !fm.done {
return fmt.Errorf("stopped; the node may be half set up, and running this again finishes it")
cancel()
return fm.node, fmt.Errorf("stopped; the node may be half set up, and running this again finishes it")
}
return fm.err
return fm.node, fm.err
}
+21
View File
@@ -94,3 +94,24 @@ func TestANodeCountsAsJoinedOnlyWhenHeardSinceItsAgentStarted(t *testing.T) {
}
}
}
func TestAFabricNamedForOneNodeIsNotAskedAgainForTheNext(t *testing.T) {
na := &nodeAdd{named: map[string]string{}}
asked := 0
ask := func(s []onboard.Fabric) []onboard.Fabric {
asked++
return []onboard.Fabric{{Name: "ib-compute", HCAs: s[0].HCAs, Key: s[0].Key}}
}
first := na.rememberFabrics(ask)([]onboard.Fabric{{Name: "fabric", HCAs: []string{"mlx5_0"}, Key: "sm:0xa"}})
second := na.rememberFabrics(ask)([]onboard.Fabric{{Name: "fabric", HCAs: []string{"mlx5_0", "mlx5_1"}, Key: "sm:0xa"}})
if asked != 1 {
t.Errorf("asked %d times", asked)
}
if len(second) != 1 || second[0].Name != "ib-compute" || len(second[0].HCAs) != 2 || first[0].Name != "ib-compute" {
t.Errorf("first %+v second %+v", first, second)
}
// A fabric not seen before is asked about; in a script nothing is asked.
if got := na.rememberFabrics(nil)([]onboard.Fabric{{Name: "fabric", Key: "sm:0xb"}}); len(got) != 0 {
t.Errorf("a script was given a fabric nobody named: %+v", got)
}
}
+98
View File
@@ -0,0 +1,98 @@
package cli
import (
"strings"
"time"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
// Two small widgets, written here rather than taken from bubbles: its
// packages ask the terminal for its colours as they load, so every llmgw
// command — llmgw --version included — sent a query and, in a terminal that
// never answers one, waited five seconds for it. Only the screens need to
// know the colours, and lipgloss asks when they first draw.
// spinTickMsg advances a spinner.
type spinTickMsg struct{}
var spinFrames = []string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}
type spinnerModel struct {
frame int
Style lipgloss.Style
}
func newSpinner() spinnerModel { return spinnerModel{Style: lipgloss.NewStyle().Foreground(tuiAccent)} }
// Tick starts the spinner turning; it is a tea.Cmd.
func (s spinnerModel) Tick() tea.Msg { return spinTickMsg{} }
func (s spinnerModel) Update(msg tea.Msg) (spinnerModel, tea.Cmd) {
if _, ok := msg.(spinTickMsg); !ok {
return s, nil
}
s.frame = (s.frame + 1) % len(spinFrames)
return s, tea.Tick(100*time.Millisecond, func(time.Time) tea.Msg { return spinTickMsg{} })
}
func (s spinnerModel) View() string { return s.Style.Render(spinFrames[s.frame]) }
// textField is one line of text typed at its end: all a form here needs.
type textField struct {
value []rune
focused bool
Prompt string
Placeholder string
CharLimit int
Width int
}
func newTextField() textField { return textField{Prompt: "> "} }
func (f *textField) SetValue(s string) { f.value = []rune(s) }
func (f textField) Value() string { return string(f.value) }
func (f *textField) Focus() { f.focused = true }
func (f *textField) Blur() { f.focused = false }
func (f textField) Update(msg tea.Msg) (textField, tea.Cmd) {
k, ok := msg.(tea.KeyMsg)
if !ok || !f.focused {
return f, nil
}
switch k.Type {
case tea.KeyRunes, tea.KeySpace:
r := k.Runes
if k.Type == tea.KeySpace {
r = []rune{' '}
}
for _, c := range r {
if f.CharLimit == 0 || len(f.value) < f.CharLimit {
f.value = append(f.value, c)
}
}
case tea.KeyBackspace:
if len(f.value) > 0 {
f.value = f.value[:len(f.value)-1]
}
case tea.KeyCtrlU:
f.value = nil
}
return f, nil
}
func (f textField) View() string {
text := string(f.value)
shown := text
if text == "" && f.Placeholder != "" {
shown = sDim.Render(f.Placeholder)
}
if f.focused {
shown += "█"
}
if pad := f.Width - len([]rune(text)); pad > 0 && text != "" {
shown += strings.Repeat(" ", pad)
}
return f.Prompt + shown
}
+5 -2
View File
@@ -404,8 +404,11 @@ type Scoring struct {
// Endpoint is one engine replica. Static in phase 1.
type Endpoint struct {
ID string `koanf:"id"`
URL string `koanf:"url"`
ID string `koanf:"id"`
URL string `koanf:"url"`
// APIKey is sent to the engine as a bearer token, for an engine started
// with one.
APIKey string `koanf:"api_key"`
Models []string `koanf:"models"`
// Engine names the inference engine, which decides what its metrics are
// called. Empty means vLLM's names.
+17 -8
View File
@@ -75,6 +75,7 @@ func (a *API) Register(mux *http.ServeMux) {
mux.HandleFunc("GET /api/v1/explain-placement", a.explainPlacement)
mux.HandleFunc("POST /api/v1/nodes/{id}/state", a.setNodeState)
mux.HandleFunc("DELETE /api/v1/nodes/{id}", a.removeNode)
mux.HandleFunc("POST /api/v1/fabric-test", a.fabricTest)
mux.HandleFunc("GET /api/v1/nodes/{id}/logs", a.nodeLogs)
}
@@ -133,7 +134,10 @@ type NodeView struct {
LastSeenAgo string `json:"last_seen_ago"`
AgentVersion string `json:"agent_version,omitempty"`
Replicas int `json:"replicas"`
// Features are the requests the agent answers beyond the usual:
// "fabric-test". Absent is an agent too old to say.
Features []string `json:"features,omitempty"`
Replicas int `json:"replicas"`
// Load is how busy the machine is, absent when its agent has not
// measured it lately — which is unknown, not idle.
@@ -234,9 +238,13 @@ type ModelView struct {
// ReplicaView is a replica as `ctl replicas` shows it.
type ReplicaView struct {
ID string `json:"id"`
Model string `json:"model"`
Profile string `json:"profile"`
ID string `json:"id"`
// EngineKey is the key this replica's engine demands, for `ctl bench`
// and anyone else holding the admin listener who needs to reach an engine
// directly. The client-facing pages never render it.
EngineKey string `json:"engine_key,omitempty"`
Model string `json:"model"`
Profile string `json:"profile"`
// Stale is true when the profile has been edited since this replica was
// placed, so it is running a definition the catalog no longer holds.
//
@@ -334,7 +342,7 @@ func (a *API) Nodes() []NodeView {
State: n.State.String(), Labels: n.Labels,
GPUTotal: len(n.GPUs), GPUFree: len(inv.FreeGPUs(n.ID)),
MemoryGB: n.MemoryTotalMB / 1024,
Fabrics: n.FabricNames(), AgentVersion: n.AgentVersion,
Fabrics: n.FabricNames(), AgentVersion: n.AgentVersion, Features: n.Features,
Replicas: replicasOn[n.ID],
}
if free, total := n.LocalDisk(); total > 0 {
@@ -467,9 +475,10 @@ func (a *API) Replicas(model string) []ReplicaView {
ID: rep.ID, Model: rep.Model, Profile: rep.Profile,
State: rep.State.String(), Pool: rep.Pool, Fabric: rep.Fabric,
Endpoint: rep.Endpoint, Err: rep.Err,
Stale: a.profileChanged(rep.Model, rep.Profile, rep.ProfileFP),
AgeSec: int64(time.Since(rep.CreatedAt).Seconds()),
IdleSec: -1,
EngineKey: a.ctl.EngineKey(rep.ID),
Stale: a.profileChanged(rep.Model, rep.Profile, rep.ProfileFP),
AgeSec: int64(time.Since(rep.CreatedAt).Seconds()),
IdleSec: -1,
}
if rep.LastUsedAt.IsZero() {
v.IdleSec = -1
+71
View File
@@ -0,0 +1,71 @@
package adminapi
import (
"encoding/json"
"fmt"
"net/http"
"time"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fabrictest"
)
// FabricTestView is a bandwidth test between two nodes.
type FabricTestView struct {
A string `json:"a"`
B string `json:"b"`
Fabric string `json:"fabric"`
OK bool `json:"ok"`
Rails []RailView `json:"rails"`
Notes []string `json:"notes,omitempty"`
Took string `json:"took"`
}
// RailView is one rail's result.
type RailView struct {
A string `json:"a"`
B string `json:"b"`
RateGbps int `json:"rate_gbps"`
Gbps float64 `json:"gbps"`
Verdict string `json:"verdict"`
Detail string `json:"detail,omitempty"`
}
// fabricTest measures RDMA bandwidth between two nodes, rail by rail. It
// sends traffic over the fabric for seconds a rail, so it is a POST and only
// ever run when asked.
func (a *API) fabricTest(w http.ResponseWriter, r *http.Request) {
var body struct {
A, B, Fabric string
Seconds int
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeError(w, http.StatusBadRequest, err.Error())
return
}
ask := a.ctl.FabricTester()
if ask == nil {
writeError(w, http.StatusNotImplemented, "this control plane cannot reach its nodes' agents")
return
}
inv := a.ctl.Inventory()
na, okA := inv.Get(body.A)
nb, okB := inv.Get(body.B)
if !okA || !okB {
writeError(w, http.StatusNotFound, fmt.Sprintf("no node %q", map[bool]string{true: body.B, false: body.A}[okA]))
return
}
fabric, rails, notes, err := fabrictest.Plan(na, nb, body.Fabric)
if err != nil {
writeError(w, http.StatusBadRequest, err.Error())
return
}
start := time.Now()
rep := fabrictest.Run(r.Context(), ask, na, nb, fabric, rails,
fabrictest.Options{Seconds: min(max(body.Seconds, 1), 30), Settle: 2 * time.Second})
v := FabricTestView{A: na.ID, B: nb.ID, Fabric: fabric, OK: rep.OK(), Notes: append(notes, rep.Notes...),
Took: time.Since(start).Round(time.Second).String()}
for _, x := range rep.Results {
v.Rails = append(v.Rails, RailView{x.A, x.B, x.RateGbps, x.Gbps, x.Verdict, x.Detail})
}
writeJSON(w, http.StatusOK, v)
}
+3
View File
@@ -23,6 +23,7 @@ func NodeFromRegister(r *llmgwv1.Register) inventory.Node {
Labels: r.GetLabels(),
State: inventory.Ready,
AgentVersion: r.GetAgentVersion(),
Features: r.GetFeatures(),
EngineVersions: r.GetEngineVersions(),
EnginesProbed: r.GetEnginesProbed(),
MemoryTotalMB: int(r.GetMemoryTotalMb()),
@@ -190,6 +191,7 @@ func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
Pool: cmd.Pool,
Fabric: cmd.Fabric,
ProfileFp: cmd.ProfileFP,
ApiKey: cmd.Spec.APIKey,
}
}
@@ -210,6 +212,7 @@ func ShardSpecFromStart(m *llmgwv1.StartShard) (engines.ShardSpec, time.Duration
Pool: m.GetPool(),
Fabric: m.GetFabric(),
ProfileFP: m.GetProfileFp(),
APIKey: m.GetApiKey(),
Profile: engines.Profile{
Name: p.GetName(), Engine: p.GetEngine(), GPUClass: p.GetGpuClass(),
TP: int(p.GetTp()), PP: int(p.GetPp()),
+47
View File
@@ -19,6 +19,7 @@ import (
"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/supervisor"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fabrictest"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
)
@@ -61,6 +62,7 @@ type Service struct {
logSeq uint64
// scanWaits are the fabric scans awaiting a reply, by request id.
scanWaits map[string]chan *llmgwv1.FabricScanReply
testWaits map[string]chan *llmgwv1.FabricTestReply
}
// SetEngines sets the engines every node installs, by name and version.
@@ -289,6 +291,15 @@ func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) {
"last_output", m.Shard.GetLastOutput())
}
case *llmgwv1.AgentMessage_FabricTest:
s.logMu.Lock()
ch, ok := s.testWaits[m.FabricTest.GetRequestId()]
delete(s.testWaits, m.FabricTest.GetRequestId())
s.logMu.Unlock()
if ok {
ch <- m.FabricTest
}
case *llmgwv1.AgentMessage_FabricScan:
s.logMu.Lock()
ch, ok := s.scanWaits[m.FabricScan.GetRequestId()]
@@ -717,6 +728,38 @@ func (s *Service) ScanFabric(ctx context.Context, nodeID, hca string, port int)
}
}
// FabricTest asks a node for one side of a bandwidth test, and waits for
// what it measured. The request's id is filled in here.
func (s *Service) FabricTest(ctx context.Context, nodeID string, req *llmgwv1.FabricTestRequest) (*llmgwv1.FabricTestReply, error) {
s.logMu.Lock()
if s.testWaits == nil {
s.testWaits = map[string]chan *llmgwv1.FabricTestReply{}
}
s.logSeq++
id := fmt.Sprintf("test-%d", s.logSeq)
ch := make(chan *llmgwv1.FabricTestReply, 1)
s.testWaits[id] = ch
s.logMu.Unlock()
defer func() {
s.logMu.Lock()
delete(s.testWaits, id)
s.logMu.Unlock()
}()
req.RequestId = id
if err := s.send(nodeID, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_FabricTest{FabricTest: req}}); err != nil {
return nil, err
}
// The test's own length, and a package install the first time.
ctx, cancel := context.WithTimeout(ctx, time.Duration(req.GetSeconds())*time.Second+5*time.Minute)
defer cancel()
select {
case r := <-ch:
return r, nil
case <-ctx.Done():
return nil, fmt.Errorf("node %s did not answer the bandwidth test", nodeID)
}
}
// ScanFabrics maps every physical InfiniBand fabric every so often, asking
// one node on each: the first, by name, whose port on it is up. A fabric
// looks the same from all its ports, so one is enough, and the rest of its
@@ -795,3 +838,7 @@ func installable(reg *llmgwv1.Register) []string {
}
return []string{"llamacpp"}
}
// The admin API finds the bandwidth test by asserting this; a signature that
// drifted would make it quietly nil.
var _ fabrictest.Asker = (*Service)(nil)
+7
View File
@@ -302,6 +302,13 @@ func (m Model) validate() error {
}
seen[p.Name] = struct{}{}
for k := range p.Env {
if bad, why := engines.UnsafeEnv(k); bad {
return fmt.Errorf("model %q profile %q: env %s is not allowed: %s",
m.Name, p.Name, k, why)
}
}
if p.GPUClass == "" && !cpuCapable(p.Engine) {
return fmt.Errorf("model %q profile %q: gpu_class is required for engine %q",
m.Name, p.Name, p.Engine)
+19
View File
@@ -302,3 +302,22 @@ func TestSimulatedNeedsEveryProfileToBeSimulated(t *testing.T) {
})
}
}
func TestAProfileCannotSetWhatRunsCodeInTheEngine(t *testing.T) {
// A profile's env reaches an engine started as root. LD_PRELOAD from a
// model definition would make "add a model" into "run this everywhere".
for _, k := range []string{"LD_PRELOAD", "ld_library_path", "PYTHONSTARTUP", "PATH", "BASH_ENV"} {
m := Model{Name: "m", Source: "http://x/m.gguf", Profiles: []Profile{{
Name: "cpu", Engine: "llamacpp", Env: map[string]string{k: "/tmp/x"},
}}}
if _, err := New([]Model{m}); err == nil || !strings.Contains(err.Error(), "not allowed") {
t.Errorf("%s: %v", k, err)
}
}
ok := Model{Name: "m", Source: "http://x/m.gguf", Profiles: []Profile{{
Name: "cpu", Engine: "llamacpp", Env: map[string]string{"OMP_NUM_THREADS": "4", "PYTHONUNBUFFERED": "1"},
}}}
if _, err := New([]Model{ok}); err != nil {
t.Errorf("ordinary tuning refused: %v", err)
}
}
+37
View File
@@ -4,6 +4,9 @@ package control
import (
"context"
"crypto/hmac"
"crypto/sha256"
"encoding/hex"
"fmt"
"log/slog"
"strings"
@@ -11,6 +14,7 @@ import (
"time"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fabrictest"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
@@ -215,6 +219,12 @@ func (c *Controller) LogReader() LogReader {
return r
}
// FabricTester returns the dispatcher's bandwidth test capability, or nil.
func (c *Controller) FabricTester() fabrictest.Asker {
r, _ := c.disp.(fabrictest.Asker)
return r
}
// Start runs the reconciler and begins publishing endpoints.
func (c *Controller) Start(ctx context.Context) {
if c.lead != nil {
@@ -297,6 +307,7 @@ func (c *Controller) endpoints() []endpoint.Info {
ID: rep.ID,
URL: strings.TrimRight(rep.Endpoint, "/"),
Models: names,
APIKey: c.engineKey(rep.ID),
Engine: prof.Engine,
Pool: rep.Pool,
Weight: rep.Weight,
@@ -314,3 +325,29 @@ func fingerprint(infos []endpoint.Info) string {
}
return b.String()
}
// EngineKey is the key a replica's engine demands, empty when engines run
// without one.
func (c *Controller) EngineKey(replicaID string) string { return c.engineKey(replicaID) }
func (c *Controller) engineKey(replicaID string) string {
if c.cfg.Reconciler.EngineKey == nil {
return ""
}
return c.cfg.Reconciler.EngineKey(replicaID)
}
// EngineKeys derives each replica's engine key from one secret.
//
// HMAC rather than a stored random key per replica: the router needs the key
// for every replica it sends to, including ones adopted after a restart that
// this process never started, and a second control plane has to agree. A
// function of a shared secret and the replica id gives all three without a
// table to keep in step. A key taken from one node opens that replica only.
func EngineKeys(secret []byte) func(replicaID string) string {
return func(replicaID string) string {
m := hmac.New(sha256.New, secret)
m.Write([]byte("llmgw engine key\x00" + replicaID))
return "llmgw-" + hex.EncodeToString(m.Sum(nil))[:40]
}
}
+19
View File
@@ -0,0 +1,19 @@
package control
import "testing"
func TestEngineKeysAreTheSameForTheSameReplicaAndNoOther(t *testing.T) {
// A control plane that restarts, or a second one, must compute the key a
// running engine was started with; a key taken from one node must open
// that replica only.
a, b := EngineKeys([]byte("s")), EngineKeys([]byte("s"))
if a("r1") != b("r1") {
t.Error("the same secret gave two keys for one replica")
}
if a("r1") == a("r2") {
t.Error("two replicas share a key")
}
if a("r1") == EngineKeys([]byte("other"))("r1") {
t.Error("the key does not depend on the secret")
}
}
+235
View File
@@ -0,0 +1,235 @@
// Package fabrictest measures the fabric between two nodes, rail by rail.
//
// Every check on one node — ports up, a subnet manager, the right modules —
// can pass while the two nodes still cannot reach each other: a cable into the
// wrong switch, a partition key, a RoCE network with no route. A replica
// spanning them then starts, and hangs on its first all-reduce with nothing to
// say why. This runs what that all-reduce would ride on — an RDMA write from
// each adapter to its counterpart — and says how it went against the link's
// rate, before a replica does.
package fabrictest
import (
"context"
"errors"
"fmt"
"slices"
"strings"
"time"
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)
// Asker sends one side of a test to a node and waits for its answer.
type Asker interface {
FabricTest(ctx context.Context, node string, req *llmgwv1.FabricTestRequest) (*llmgwv1.FabricTestReply, error)
}
// Rail is one adapter on each node, paired by their position in the fabric's
// declaration: mlx5_0 with mlx5_0, the way NCCL pairs them.
type Rail struct {
A, B string // adapters
PortA, PortB int
// RateGbps is the slower of the two links.
RateGbps int
}
// Result is one rail's test.
type Result struct {
Rail
Gbps float64 `json:"gbps"`
// Verdict is "ok" at 80% of the link's rate or more, "slow" at 50%,
// "fail" below that or when it did not run.
Verdict string `json:"verdict"`
Detail string `json:"detail,omitempty"`
}
// Report is a whole test.
type Report struct {
Fabric string `json:"fabric"`
A, B string `json:"-"`
Results []Result `json:"results"`
// Notes say what the test could not cover.
Notes []string `json:"notes,omitempty"`
}
// OK reports whether every rail passed.
func (r Report) OK() bool {
for _, x := range r.Results {
if x.Verdict != "ok" {
return false
}
}
return len(r.Results) > 0
}
// Plan pairs the adapters two nodes have on a fabric. With no fabric named,
// the one fabric they share is used.
func Plan(a, b inventory.Node, fabric string) (string, []Rail, []string, error) {
if a.ID == b.ID {
return "", nil, nil, errors.New("give two different nodes")
}
if fabric == "" {
var common []string
for _, f := range a.FabricNames() {
if b.HasFabric(f) {
common = append(common, f)
}
}
switch len(common) {
case 0:
return "", nil, nil, fmt.Errorf("%s (%s) and %s (%s) declare no fabric in common",
a.ID, orNone(a.FabricNames()), b.ID, orNone(b.FabricNames()))
case 1:
fabric = common[0]
default:
return "", nil, nil, fmt.Errorf("%s and %s share %s: name one with --fabric", a.ID, b.ID, strings.Join(common, ", "))
}
}
for _, n := range []inventory.Node{a, b} {
if !slices.Contains(n.Features, "fabric-test") {
return "", nil, nil, fmt.Errorf("%s's agent (%s) is too old to run a bandwidth test: upgrade it", n.ID, orNone([]string{n.AgentVersion}))
}
}
ha, hb := a.HCAsFor(fabric), b.HCAsFor(fabric)
if len(ha) == 0 || len(hb) == 0 {
return "", nil, nil, fmt.Errorf("%s is not declared on both %s and %s", fabric, a.ID, b.ID)
}
var notes []string
if len(ha) != len(hb) {
notes = append(notes, fmt.Sprintf("%s has %d adapters on %s and %s has %d: only the first %d are paired",
a.ID, len(ha), fabric, b.ID, len(hb), min(len(ha), len(hb))))
}
var rails []Rail
for i := range min(len(ha), len(hb)) {
pa, ra := port(a, ha[i])
pb, rb := port(b, hb[i])
rails = append(rails, Rail{A: ha[i], B: hb[i], PortA: pa, PortB: pb, RateGbps: minPositive(ra, rb)})
}
return fabric, rails, notes, nil
}
// port is an adapter's first active port and its rate, or port 1 when the
// node reported none.
func port(n inventory.Node, hca string) (int, int) {
for _, p := range n.IBPorts {
if p.HCA == hca && strings.EqualFold(p.State, "ACTIVE") {
return p.Port, p.RateGbps
}
}
return 1, 0
}
// Options for a run.
type Options struct {
Seconds int
// BasePort is perftest's TCP exchange port for the first rail; each rail
// takes the next, so a rail whose server is still closing never meets the
// next one's.
BasePort int
// Settle is how long the server is given to listen before the client
// dials; the client retries besides.
Settle time.Duration
}
// Run tests every rail, one after another: rails at once would share the
// PCIe and memory bandwidth they are measuring.
func Run(ctx context.Context, ask Asker, a, b inventory.Node, fabric string, rails []Rail, o Options) Report {
if o.Seconds <= 0 {
o.Seconds = 3
}
if o.BasePort == 0 {
o.BasePort = 18515
}
rep := Report{Fabric: fabric, A: a.ID, B: b.ID}
for i, r := range rails {
rep.Results = append(rep.Results, runRail(ctx, ask, a, b, r, o, o.BasePort+i))
}
return rep
}
func runRail(ctx context.Context, ask Asker, a, b inventory.Node, r Rail, o Options, tcp int) Result {
res := Result{Rail: r, Verdict: "fail"}
if b.MgmtAddr == "" {
res.Detail = b.ID + " reported no management address for the client to dial"
return res
}
type answer struct {
r *llmgwv1.FabricTestReply
err error
}
server := make(chan answer, 1)
go func() {
rep, err := ask.FabricTest(ctx, b.ID, &llmgwv1.FabricTestRequest{Role: "server", Hca: r.B,
Port: int32(r.PortB), TcpPort: int32(tcp), Seconds: int32(o.Seconds)}) //nolint:gosec // small numbers
server <- answer{rep, err}
}()
select {
case <-ctx.Done():
res.Detail = ctx.Err().Error()
return res
case <-time.After(o.Settle):
}
client, err := ask.FabricTest(ctx, a.ID, &llmgwv1.FabricTestRequest{Role: "client", Hca: r.A,
Port: int32(r.PortA), Peer: b.MgmtAddr, TcpPort: int32(tcp), Seconds: int32(o.Seconds)}) //nolint:gosec // small numbers
var srv answer
select {
case srv = <-server:
case <-time.After(time.Duration(o.Seconds)*time.Second + time.Minute):
srv.err = errors.New("the server side did not finish")
}
switch {
case err != nil:
res.Detail = a.ID + ": " + err.Error()
case client.GetError() != "":
res.Detail = a.ID + ": " + client.GetError()
// The server's reason is often the better one: it saw no one arrive.
if srv.r != nil && srv.r.GetError() != "" {
res.Detail += "; " + b.ID + ": " + srv.r.GetError()
}
default:
res.Gbps = client.GetGbps()
res.Verdict, res.Detail = grade(res.Gbps, r.RateGbps)
}
return res
}
// grade compares a measurement with the link's rate. RDMA writes of 64 KiB
// reach 90–97% of the signalling rate on a healthy link; well under that is
// a degraded cable, a card in a narrow PCIe slot, or traffic crossing a
// socket.
func grade(gbps float64, rate int) (string, string) {
if gbps <= 0 {
return "fail", "no traffic got through"
}
if rate <= 0 {
return "ok", "the link's rate is not known to compare with"
}
pct := gbps / float64(rate) * 100
switch {
case pct >= 80:
return "ok", fmt.Sprintf("%.0f%% of %d Gb/s", pct, rate)
case pct >= 50:
return "slow", fmt.Sprintf("%.0f%% of %d Gb/s: check the cable, and the card's PCIe width and NUMA node", pct, rate)
default:
return "fail", fmt.Sprintf("%.0f%% of %d Gb/s", pct, rate)
}
}
func minPositive(a, b int) int {
switch {
case a <= 0:
return b
case b <= 0:
return a
}
return min(a, b)
}
func orNone(s []string) string {
if len(s) == 0 {
return "none"
}
return strings.Join(slices.Clone(s), ", ")
}
@@ -0,0 +1,100 @@
package fabrictest
import (
"context"
"strings"
"sync"
"testing"
"time"
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)
func node(id string, hcas ...string) inventory.Node {
n := inventory.Node{ID: id, MgmtAddr: "10.0.0." + id[len(id)-1:], Features: []string{"fabric-test"},
Fabrics: []inventory.FabricAttachment{{Fabric: "ib-a", HCAs: hcas}}}
for _, h := range hcas {
n.IBPorts = append(n.IBPorts, inventory.IBPort{HCA: h, Port: 1, State: "ACTIVE", RateGbps: 400})
}
return n
}
// fakeFabric answers as two agents would: the server waits for its client.
type fakeFabric struct {
mu sync.Mutex
gbps map[string]float64 // by client adapter
fails map[string]string
asked []string
}
func (f *fakeFabric) FabricTest(_ context.Context, n string, req *llmgwv1.FabricTestRequest) (*llmgwv1.FabricTestReply, error) {
f.mu.Lock()
defer f.mu.Unlock()
f.asked = append(f.asked, n+" "+req.GetRole()+" "+req.GetHca())
if req.GetRole() == "server" {
return &llmgwv1.FabricTestReply{}, nil
}
if e := f.fails[req.GetHca()]; e != "" {
return &llmgwv1.FabricTestReply{Error: e}, nil
}
return &llmgwv1.FabricTestReply{Gbps: f.gbps[req.GetHca()]}, nil
}
func TestRailsArePairedByPositionOnTheSharedFabric(t *testing.T) {
a, b := node("n1", "mlx5_0", "mlx5_1"), node("n2", "mlx5_0", "mlx5_1", "mlx5_2")
fabric, rails, notes, err := Plan(a, b, "")
if err != nil {
t.Fatal(err)
}
if fabric != "ib-a" || len(rails) != 2 || rails[1].A != "mlx5_1" || rails[1].B != "mlx5_1" || rails[0].RateGbps != 400 {
t.Errorf("got %s %+v", fabric, rails)
}
if len(notes) != 1 || !strings.Contains(notes[0], "only the first 2") {
t.Errorf("an unequal count went unmentioned: %v", notes)
}
}
func TestNodesWithNoFabricInCommonAreSaidSo(t *testing.T) {
b := node("n2", "mlx5_0")
b.Fabrics[0].Fabric = "ib-b"
if _, _, _, err := Plan(node("n1", "mlx5_0"), b, ""); err == nil || !strings.Contains(err.Error(), "no fabric in common") {
t.Errorf("got %v", err)
}
}
func TestEachRailIsGradedAgainstItsLink(t *testing.T) {
a, b := node("n1", "mlx5_0", "mlx5_1", "mlx5_2"), node("n2", "mlx5_0", "mlx5_1", "mlx5_2")
_, rails, _, _ := Plan(a, b, "ib-a")
f := &fakeFabric{
gbps: map[string]float64{"mlx5_0": 388, "mlx5_1": 240},
fails: map[string]string{"mlx5_2": "ib_write_bw: Couldn't connect to 10.0.0.2:18517"},
}
rep := Run(context.Background(), f, a, b, "ib-a", rails, Options{Seconds: 1, Settle: time.Millisecond})
want := []string{"ok", "slow", "fail"}
for i, r := range rep.Results {
if r.Verdict != want[i] {
t.Errorf("rail %d: %s (%s), want %s", i, r.Verdict, r.Detail, want[i])
}
}
if rep.OK() {
t.Error("a report with a failed rail is OK")
}
if !strings.Contains(rep.Results[2].Detail, "Couldn't connect") {
t.Errorf("the reason was lost: %q", rep.Results[2].Detail)
}
// Servers on the second node, clients on the first, one rail at a time.
if f.asked[0] != "n2 server mlx5_0" || f.asked[1] != "n1 client mlx5_0" {
t.Errorf("asked %v", f.asked)
}
}
func TestAnAgentTooOldToAnswerIsRefusedAtOnce(t *testing.T) {
// It would ignore the request, and the control plane would wait minutes.
old := node("n2", "mlx5_0")
old.Features, old.AgentVersion = nil, "v0.10.0"
_, _, _, err := Plan(node("n1", "mlx5_0"), old, "")
if err == nil || !strings.Contains(err.Error(), "n2's agent (v0.10.0) is too old") {
t.Errorf("got %v", err)
}
}
+4 -1
View File
@@ -100,7 +100,10 @@ type Node struct {
LastSeen time.Time `json:"last_seen"`
// AgentVersion and EngineVersions are reported at registration. The engine
// version decides which metric names and CLI flags apply (plan §14).
AgentVersion string `json:"agent_version,omitempty"`
AgentVersion string `json:"agent_version,omitempty"`
// Features are the requests the agent answers beyond the ones every
// agent does.
Features []string `json:"features,omitempty"`
EngineVersions map[string]string `json:"engine_versions,omitempty"`
// EnginesProbed says EngineVersions is an answer, so an empty map means
// "has none" rather than "an agent too old to say".
+13
View File
@@ -69,6 +69,11 @@ type Config struct {
PortRange [2]int
// MgmtInterface is the ethernet interface used for NCCL bootstrap.
MgmtInterface string
// EngineKey gives each replica the key its engine demands. Derived from
// the replica id rather than stored, so a control plane that restarts —
// or a second one — computes the same key for a replica it adopts. Nil
// starts engines without one.
EngineKey func(replicaID string) string
// EvictionGrace is how long a replica must have gone unasked-for before it
// may be stopped to make room for another model.
//
@@ -883,6 +888,7 @@ func (r *Reconciler) startReplica(ctx context.Context, m catalog.Model, pl *plac
MgmtInterface: r.cfg.MgmtInterface,
HTTPPort: s.HTTPPort,
Profile: espec,
APIKey: r.engineKey(id),
},
StartupTimeout: prof.StartupTimeout,
Pool: pl.Pool,
@@ -1380,3 +1386,10 @@ func (r *Reconciler) profileMemory(model, profile string) int {
}
return 0
}
func (r *Reconciler) engineKey(replicaID string) string {
if r.cfg.EngineKey == nil {
return ""
}
return r.cfg.EngineKey(replicaID)
}
@@ -1721,3 +1721,23 @@ func TestAnAgentTooOldToReportTheProfileVersionIsNotGuessedAt(t *testing.T) {
t.Errorf("an unreported fingerprint should stay empty, got %q", fp)
}
}
func TestAnEngineIsStartedWithItsReplicasKey(t *testing.T) {
// Engines listen on every address a node has. Started without a key, any
// machine that could reach the port reached the model past the gateway's
// authentication, quotas and accounting — checked from a neighbouring node
// on the live cluster.
h := newHarness(t, testModel(1), gpuNode("1", 8))
h.r.cfg.EngineKey = func(id string) string { return "key-for-" + id }
h.r.Reconcile(t.Context())
started := h.disp.startedShards()
if len(started) == 0 {
t.Fatal("nothing was started")
}
for _, c := range started {
if c.Spec.APIKey != "key-for-"+c.ReplicaID {
t.Errorf("shard %s started with key %q", c.ShardID, c.Spec.APIKey)
}
}
}
+14 -6
View File
@@ -7,7 +7,6 @@ import (
"net/http"
"os"
"path/filepath"
"strings"
"time"
)
@@ -29,13 +28,22 @@ type Added struct {
// the layout — a version that can be overwritten takes away the ability to
// change the catalog back and still have the old weights.
func Add(ctx context.Context, root, source string, e Entry, progress func(done, total int64)) (Added, error) {
return addWith(ctx, http.DefaultClient, root, source, e, progress)
}
// addWith is Add through a given client. The management API passes one that
// will not reach private addresses; `llmgw mirror add`, run by somebody on
// the mirror, passes the ordinary one.
func addWith(ctx context.Context, client *http.Client, root, source string, e Entry,
progress func(done, total int64),
) (Added, error) {
if e.Model == "" || e.Version == "" || e.File == "" {
return Added{}, fmt.Errorf("a model, a version and a filename are all required")
}
if strings.ContainsAny(e.Model+e.Version+e.File, `/\`) {
// Each is one path element. A slash would put the file somewhere the
// layout does not describe, and List would not find it again.
return Added{}, fmt.Errorf("model, version and filename are single names, not paths")
for _, n := range []string{e.Model, e.Version, e.File} {
if err := ValidName(n); err != nil {
return Added{}, err
}
}
final := e.Path(root)
@@ -50,7 +58,7 @@ func Add(ctx context.Context, root, source string, e Entry, progress func(done,
if err != nil {
return Added{}, err
}
resp, err := http.DefaultClient.Do(req)
resp, err := client.Do(req)
if err != nil {
return Added{}, fmt.Errorf("fetching %s: %w", source, err)
}
+2 -2
View File
@@ -130,8 +130,8 @@ func (a *API) status(w http.ResponseWriter, r *http.Request) {
func (a *API) remove(w http.ResponseWriter, r *http.Request) {
model, version := r.PathValue("model"), r.PathValue("version")
if strings.ContainsAny(model+version, `/\`) {
fail(w, http.StatusBadRequest, errors.New("model and version are single names, not paths"))
if err := errors.Join(ValidName(model), ValidName(version)); err != nil {
fail(w, http.StatusBadRequest, err)
return
}
+34 -1
View File
@@ -17,7 +17,10 @@ import (
func testAPI(t *testing.T) (*API, string) {
t.Helper()
root := t.TempDir()
a := NewAPI(NewManager(root, discardLog()), "http://mirror:8081")
m := NewManager(root, discardLog())
// The upstreams in these tests are httptest servers on loopback.
m.AllowPrivateSources()
a := NewAPI(m, "http://mirror:8081")
// A download outlives the request that started it — that is the point of
// the manager. So the temp directory must not be torn down while one is
@@ -434,3 +437,33 @@ func TestRemovingAVersionIsRefusedWhileAnyShardIsArriving(t *testing.T) {
t.Errorf("expected 409 while a shard is arriving, got %d", code)
}
}
func TestTheAPIWillNotFetchFromThePrivateNetwork(t *testing.T) {
// The file server hands whatever was fetched to anyone, so a fetch that
// could reach 169.254.169.254 or a database on the LAN was a way to read
// them through the mirror.
m := NewManager(t.TempDir(), discardLog())
for _, src := range []string{
"http://127.0.0.1:1/x.gguf", "http://169.254.169.254/latest/meta-data",
"http://10.0.0.140:9090/api/v1/nodes", "http://[::1]:1/x",
} {
st, err := m.Start(src, Entry{Model: "m", Version: "v", File: "x.gguf"})
if err != nil {
t.Fatalf("%s: %v", src, err)
}
deadline := time.Now().Add(5 * time.Second)
for st.State == Fetching && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
st, _ = m.Status("m", "v")
}
if st.State != Failed || !strings.Contains(st.Err, "refusing to fetch") {
t.Errorf("%s: %+v", src, st)
}
m.Forget("m", "v")
}
for _, src := range []string{"file:///etc/passwd", "gopher://x/", "not a url"} {
if _, err := m.Start(src, Entry{Model: "m", Version: "v", File: "x"}); err == nil {
t.Errorf("%s was accepted as a source", src)
}
}
}
+16 -2
View File
@@ -5,6 +5,7 @@ import (
"errors"
"fmt"
"log/slog"
"net/http"
"os"
"sync"
"time"
@@ -72,8 +73,18 @@ type Manager struct {
// until the first had finished. A fifteen-shard model had to be pulled one
// shard at a time, each waiting on the last.
running map[string]*Status
// client fetches sources. It will not reach private addresses unless
// AllowPrivateSources was called.
client *http.Client
}
// AllowPrivateSources lets downloads reach loopback, private and link-local
// addresses — a model server on the same network, say. Off by default: the
// file server hands whatever was fetched to anyone, so a fetch that can reach
// the private network is a way to read it.
func (m *Manager) AllowPrivateSources() { m.client = fetchClient(true) }
// runKey identifies one download.
func runKey(e Entry) string { return e.Model + "/" + e.Version + "/" + e.File }
@@ -82,7 +93,7 @@ func NewManager(root string, log *slog.Logger) *Manager {
if log == nil {
log = slog.Default()
}
return &Manager{root: root, log: log, running: map[string]*Status{}}
return &Manager{root: root, log: log, running: map[string]*Status{}, client: fetchClient(false)}
}
// ErrAlreadyFetching means a download for that model and version is running.
@@ -97,6 +108,9 @@ func (m *Manager) Start(source string, e Entry) (Status, error) {
if source == "" {
return Status{}, fmt.Errorf("a source URL is required")
}
if err := validSource(source); err != nil {
return Status{}, err
}
key := runKey(e)
m.mu.Lock()
@@ -125,7 +139,7 @@ func (m *Manager) run(source string, e Entry, key string) {
// minutes to hours and the thing that asked for it is long gone.
ctx := context.Background()
added, err := Add(ctx, m.root, source, e, func(done, total int64) {
added, err := addWith(ctx, m.client, m.root, source, e, func(done, total int64) {
m.mu.Lock()
if st, ok := m.running[key]; ok {
st.BytesDone, st.BytesTotal = done, total
+15
View File
@@ -37,6 +37,21 @@ type Entry struct {
Bytes int64
}
// ValidName checks one element of an entry — a model, version or file name.
//
// Each is one path element, and a name that is not one writes somewhere the
// layout does not describe. Rejecting slashes was not enough: ".." is a
// single element too, and a model called ".." put its file beside the mirror
// root rather than inside it. A name starting with a dot is refused as well,
// because that is where the mirror keeps its own bookkeeping.
func ValidName(n string) error {
if n == "" || n[0] == '.' || strings.ContainsAny(n, "/\\\x00") {
return fmt.Errorf("%q is not a usable name: it must be a single path element "+
"that does not start with a dot", n)
}
return nil
}
// Path is where this entry lives under a root.
func (e Entry) Path(root string) string {
return filepath.Join(root, e.Model, e.Version, e.File)
+15
View File
@@ -162,3 +162,18 @@ func TestProgressIsReportedForALongDownload(t *testing.T) {
t.Error("a download reported no progress at all")
}
}
func TestANameThatLeavesTheMirrorIsRefused(t *testing.T) {
root := t.TempDir()
for _, e := range []Entry{
{Model: "..", Version: "v", File: "f.gguf"},
{Model: "m", Version: "..", File: "f.gguf"},
{Model: "m", Version: "v", File: ".f.sha256"},
{Model: "a/b", Version: "v", File: "f.gguf"},
} {
if _, err := Add(t.Context(), root, "http://127.0.0.1:1/x", e, nil); err == nil ||
!strings.Contains(err.Error(), "not a usable name") {
t.Errorf("%+v: %v", e, err)
}
}
}
+86
View File
@@ -0,0 +1,86 @@
package mirror
import (
"errors"
"fmt"
"net"
"net/http"
"net/netip"
"net/url"
"syscall"
"time"
)
// errPrivateSource is a fetch refused because it would reach this machine or
// the network behind it rather than the internet.
var errPrivateSource = errors.New("refusing to fetch from a loopback, private or link-local address")
// fetchClient downloads sources for the management API.
//
// The API fetches whatever URL it is given and the file server then hands the
// result to anyone. Without this, "fetch http://169.254.169.254/..." was a way
// to read a cloud's instance metadata — or anything else on the network the
// mirror sits in — through the mirror. The check is made on the address
// actually dialled, after resolution and on every redirect, so a name that
// resolves somewhere public and then somewhere private is caught too.
func fetchClient(allowPrivate bool) *http.Client {
d := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}
if !allowPrivate {
d.Control = func(_, address string, _ syscall.RawConn) error {
host, _, err := net.SplitHostPort(address)
if err != nil {
return err
}
a, err := netip.ParseAddr(host)
if err != nil {
return err
}
if blockedAddr(a.Unmap()) {
return fmt.Errorf("%w (%s)", errPrivateSource, a)
}
return nil
}
}
return &http.Client{
Transport: &http.Transport{
Proxy: http.ProxyFromEnvironment,
DialContext: d.DialContext,
TLSHandshakeTimeout: 30 * time.Second,
ResponseHeaderTimeout: 2 * time.Minute,
ForceAttemptHTTP2: true,
},
// No overall timeout: a model is tens of gigabytes. The dial, the
// handshake and the first byte are bounded above; after that a stall
// is the context's to end.
CheckRedirect: func(req *http.Request, via []*http.Request) error {
if len(via) >= 10 {
return errors.New("too many redirects")
}
return checkScheme(req.URL)
},
}
}
func blockedAddr(a netip.Addr) bool {
cgnat := netip.MustParsePrefix("100.64.0.0/10")
return a.IsLoopback() || a.IsPrivate() || a.IsLinkLocalUnicast() ||
a.IsLinkLocalMulticast() || a.IsMulticast() || a.IsUnspecified() || cgnat.Contains(a)
}
// checkScheme allows the two schemes a model is published over. file:// and
// the rest would read this machine's own disk.
func checkScheme(u *url.URL) error {
if u.Scheme != "http" && u.Scheme != "https" {
return fmt.Errorf("%q is not an http or https source", u.Redacted())
}
return nil
}
// validSource checks a source before a download is started for it.
func validSource(source string) error {
u, err := url.Parse(source)
if err != nil || u.Host == "" {
return fmt.Errorf("%q is not a URL", source)
}
return checkScheme(u)
}
+5 -1
View File
@@ -35,6 +35,10 @@ import (
type Fabric struct {
Name string
HCAs []string
// Key is what the adapters report the fabric to be — its subnet manager,
// or a RoCE network — on a suggestion, so a name given for one node can
// be given again to the next on the same fabric.
Key string
}
// ParseFabric reads "name:hca1,hca2", or "name" alone.
@@ -462,7 +466,7 @@ func suggestFabrics(hcas []discovery.HCA) []Fabric {
if len(keys) > 1 {
name = fmt.Sprintf("fabric%d", i+1)
}
out = append(out, Fabric{Name: name, HCAs: byKey[k]})
out = append(out, Fabric{Name: name, HCAs: byKey[k], Key: k})
}
return out
}
+75
View File
@@ -0,0 +1,75 @@
package router
import (
"net"
"net/http"
"net/netip"
"net/url"
"strings"
)
// AdminGuard keeps a browser from being used against the admin listener.
//
// The listener has no login: reaching it is meant to mean holding the machine,
// so it binds to loopback and an operator reaches it over an ssh tunnel. But a
// browser on that operator's machine reaches it too, from any page it has
// open. A page could POST a model whose profile runs as root on every node,
// or — by pointing its own domain at 127.0.0.1 — read pprof's heap, which
// holds API keys in the clear. Three checks close that without touching what
// ctl and curl do:
//
// - Host must be an address or localhost. A page that rebinds its own name
// to 127.0.0.1 still sends that name.
// - A request carrying Origin or Sec-Fetch-Site from another site is
// refused. Browsers send them; ctl and curl do not.
// - Anything but GET and HEAD must be application/json. A page can send a
// form or text/plain anywhere without asking; JSON across origins needs a
// preflight this listener never grants.
func AdminGuard(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !addressHost(r.Host) {
http.Error(w, "the admin listener answers to an address or localhost, not a name", http.StatusMisdirectedRequest)
return
}
if crossSite(r) {
http.Error(w, "cross-site request refused", http.StatusForbidden)
return
}
if r.Method != http.MethodGet && r.Method != http.MethodHead {
ct := r.Header.Get("Content-Type")
if mt, _, _ := strings.Cut(ct, ";"); strings.TrimSpace(strings.ToLower(mt)) != "application/json" {
http.Error(w, "a change to the admin API must be sent as application/json", http.StatusUnsupportedMediaType)
return
}
}
next.ServeHTTP(w, r)
})
}
func addressHost(hostport string) bool {
host := hostport
if h, _, err := net.SplitHostPort(hostport); err == nil {
host = h
}
host = strings.Trim(host, "[]")
if strings.EqualFold(host, "localhost") {
return true
}
_, err := netip.ParseAddr(host)
return err == nil
}
func crossSite(r *http.Request) bool {
switch r.Header.Get("Sec-Fetch-Site") {
case "", "same-origin", "none":
default:
return true
}
if o := r.Header.Get("Origin"); o != "" {
u, err := url.Parse(o)
if err != nil || !strings.EqualFold(u.Host, r.Host) {
return true
}
}
return false
}
+57
View File
@@ -0,0 +1,57 @@
package router
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestTheAdminListenerRefusesWhatABrowserPageCouldSend(t *testing.T) {
h := AdminGuard(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {}))
do := func(method, host, ct string, hdr map[string]string) int {
r := httptest.NewRequest(method, "/api/v1/models", strings.NewReader("{}"))
r.Host = host
if ct != "" {
r.Header.Set("Content-Type", ct)
}
for k, v := range hdr {
r.Header.Set(k, v)
}
w := httptest.NewRecorder()
h.ServeHTTP(w, r)
return w.Code
}
// What ctl and curl send.
if got := do("POST", "127.0.0.1:9090", "application/json", nil); got != 200 {
t.Errorf("ctl's request was refused: %d", got)
}
if got := do("GET", "localhost:9090", "", nil); got != 200 {
t.Errorf("a GET through an ssh tunnel was refused: %d", got)
}
if got := do("DELETE", "[::1]:9090", "application/json", nil); got != 200 {
t.Errorf("IPv6 loopback refused: %d", got)
}
// A page posting a form, or text/plain, needs no preflight.
if got := do("POST", "127.0.0.1:9090", "text/plain", nil); got != http.StatusUnsupportedMediaType {
t.Errorf("text/plain POST accepted: %d", got)
}
if got := do("POST", "127.0.0.1:9090", "", nil); got != http.StatusUnsupportedMediaType {
t.Errorf("a bodiless POST accepted: %d", got)
}
// DNS rebinding: the page's own name, now resolving to 127.0.0.1.
if got := do("GET", "evil.example:9090", "", nil); got != http.StatusMisdirectedRequest {
t.Errorf("a rebound name was answered: %d", got)
}
// And anything a browser marks as coming from elsewhere.
if got := do("POST", "127.0.0.1:9090", "application/json",
map[string]string{"Sec-Fetch-Site": "cross-site"}); got != http.StatusForbidden {
t.Errorf("cross-site fetch accepted: %d", got)
}
if got := do("POST", "127.0.0.1:9090", "application/json",
map[string]string{"Origin": "https://evil.example"}); got != http.StatusForbidden {
t.Errorf("foreign Origin accepted: %d", got)
}
}
+3
View File
@@ -22,6 +22,9 @@ type Info struct {
ID string
// URL is the engine's root address (e.g. http://gpu-a-01:8000), no trailing slash.
URL string
// APIKey is what the engine demands of every request, sent as a bearer
// token. Empty for an engine that asks for none.
APIKey string
// Models lists the model names this replica serves.
Models []string
// Engine names the inference engine serving here, which decides what its
+6 -2
View File
@@ -148,6 +148,8 @@ type Request struct {
Body []byte
// Header carries headers to forward (authorization is deliberately not).
Header http.Header
// APIKey is the engine's own key, sent in place of the client's.
APIKey string
// Stream says whether an SSE response is expected.
Stream bool
// SuppressUsageChunk drops the final usage-only chunk from the client while
@@ -218,6 +220,9 @@ func (p *Proxy) Do(ctx context.Context, w http.ResponseWriter, r Request) (Outco
}
req.Header = cloneForUpstream(r.Header)
req.Header.Set("Content-Type", "application/json")
if r.APIKey != "" {
req.Header.Set("Authorization", "Bearer "+r.APIKey)
}
req.ContentLength = int64(len(r.Body))
// A non-streaming request waits for the whole completion before a single
@@ -518,8 +523,7 @@ func copyHeaders(dst, src http.Header) {
// cloneForUpstream forwards only what the engine needs. The client's
// Authorization header is deliberately dropped: it authenticates the caller to
// llmgw, not llmgw to the engine, and engines are reached over a trusted
// network (plan §14).
// llmgw, not llmgw to the engine. llmgw's own key for the engine is set by Do.
func cloneForUpstream(h http.Header) http.Header {
out := make(http.Header, 4)
for _, k := range []string{"Accept", "User-Agent", "Traceparent", "Tracestate"} {
+19
View File
@@ -355,3 +355,22 @@ func TestAStreamingRequestKeepsTheShortLivenessTimeout(t *testing.T) {
t.Errorf("took %s; the streaming liveness timeout is not being applied", elapsed)
}
}
func TestTheEngineGetsLlmgwsKeyNotTheClients(t *testing.T) {
var got string
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
got = r.Header.Get("Authorization")
io.WriteString(w, `{}`)
}))
defer up.Close()
client := http.Header{"Authorization": {"Bearer sk-the-callers-own"}}
_, err := testProxy().Do(t.Context(), httptest.NewRecorder(),
Request{URL: up.URL, Body: []byte(`{}`), Header: client, APIKey: "llmgw-engine"})
if err != nil {
t.Fatal(err)
}
if got != "Bearer llmgw-engine" {
t.Errorf("the engine was sent %q", got)
}
}
+2
View File
@@ -373,6 +373,7 @@ func toInfos(eps []config.Endpoint) []endpoint.Info {
out = append(out, endpoint.Info{
ID: e.ID,
URL: strings.TrimRight(e.URL, "/"),
APIKey: e.APIKey,
Models: e.Models,
Engine: e.Engine,
Pool: e.Pool,
@@ -673,6 +674,7 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa
out, perr := r.px.Do(ctx, w, proxy.Request{
URL: dec.Endpoint.URL + req.Path,
APIKey: dec.Endpoint.APIKey,
Body: req.UpstreamBody,
Header: req.Header,
Stream: req.Stream,
+6 -2
View File
@@ -157,7 +157,7 @@ func (p *Poller) loop(ctx context.Context, in endpoint.Info) {
}
func (p *Poller) once(ctx context.Context, in endpoint.Info, reported *string) {
stats, err := p.fetch(ctx, in.URL, p.mappingFor(in))
stats, err := p.fetch(ctx, in.URL, in.APIKey, p.mappingFor(in))
if err != nil {
if ctx.Err() != nil {
return
@@ -201,7 +201,7 @@ func (p *Poller) report(in endpoint.Info, err error, reported *string) {
*reported = "unreachable"
}
func (p *Poller) fetch(ctx context.Context, base string, mapping Mapping) (endpoint.Stats, error) {
func (p *Poller) fetch(ctx context.Context, base, apiKey string, mapping Mapping) (endpoint.Stats, error) {
ctx, cancel := context.WithTimeout(ctx, p.cfg.Timeout)
defer cancel()
@@ -209,6 +209,10 @@ func (p *Poller) fetch(ctx context.Context, base string, mapping Mapping) (endpo
if err != nil {
return endpoint.Stats{}, err
}
// llama.cpp guards /metrics with its key; vLLM does not, and ignores one.
if apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := p.http.Do(req)
if err != nil {
return endpoint.Stats{}, err
+4
View File
@@ -38,6 +38,10 @@ var copyOrder = []string{
"threads",
"thread_messages",
"memory_notes",
// Copied, not regenerated: every running engine was started with a key
// derived from this, and a new secret would lock the router out of all
// of them at the moment of the move.
"cluster_secrets",
}
// notCopied are tables deliberately left behind.
@@ -0,0 +1,12 @@
-- Secrets the control planes share and nobody types: generated on first use,
-- read by every control plane against this store afterwards.
--
-- The first is the one each replica's engine key is derived from. It has to be
-- the same for every control plane — the router sends a replica the key it was
-- started with, whichever control plane started it — and the same after a
-- restart, when running replicas are adopted rather than restarted.
CREATE TABLE cluster_secrets (
name TEXT PRIMARY KEY,
value TEXT NOT NULL,
created_at BIGINT NOT NULL
);
@@ -0,0 +1,12 @@
-- Secrets the control planes share and nobody types: generated on first use,
-- read by every control plane against this store afterwards.
--
-- The first is the one each replica's engine key is derived from. It has to be
-- the same for every control plane — the router sends a replica the key it was
-- started with, whichever control plane started it — and the same after a
-- restart, when running replicas are adopted rather than restarted.
CREATE TABLE cluster_secrets (
name TEXT PRIMARY KEY,
value TEXT NOT NULL,
created_at BIGINT NOT NULL
);
+31
View File
@@ -0,0 +1,31 @@
package store
import (
"context"
"crypto/rand"
"encoding/hex"
"time"
)
// ClusterSecret returns the named secret, creating a random one the first time
// it is asked for.
//
// Insert-if-absent and then read, so two control planes asking at once end up
// with the same value: the loser's insert changes nothing and it reads the
// winner's.
func (s *DB) ClusterSecret(ctx context.Context, name string) (string, error) {
b := make([]byte, 32)
if _, err := rand.Read(b); err != nil {
return "", err
}
if _, err := s.db.ExecContext(ctx,
`INSERT INTO cluster_secrets (name, value, created_at) VALUES (?, ?, ?)
ON CONFLICT (name) DO NOTHING`,
name, hex.EncodeToString(b), time.Now().UnixMilli()); err != nil {
return "", err
}
var v string
err := s.db.QueryRowContext(ctx,
`SELECT value FROM cluster_secrets WHERE name = ?`, name).Scan(&v)
return v, err
}
+25
View File
@@ -0,0 +1,25 @@
package store
import (
"path/filepath"
"testing"
)
func TestAClusterSecretIsMadeOnceAndKept(t *testing.T) {
db, err := Open(t.Context(), filepath.Join(t.TempDir(), "s.db"))
if err != nil {
t.Fatal(err)
}
defer db.Close()
a, err := db.ClusterSecret(t.Context(), "engine-key")
if err != nil || len(a) != 64 {
t.Fatalf("%q %v", a, err)
}
b, err := db.ClusterSecret(t.Context(), "engine-key")
if err != nil || a != b {
t.Errorf("a second read made a new secret: %q then %q", a, b)
}
if c, _ := db.ClusterSecret(t.Context(), "other"); c == a {
t.Error("two names share a secret")
}
}
+8 -1
View File
@@ -303,7 +303,7 @@ func TestThePageRendersThroughTheRealClient(t *testing.T) {
log := slog.New(slog.NewTextHandler(io.Discard, nil))
srv := httptest.NewServer(
mirror.NewAPI(mirror.NewManager(dir, log), "http://mirror:8081").Handler())
mirror.NewAPI(privateOK(mirror.NewManager(dir, log)), "http://mirror:8081").Handler())
defer srv.Close()
c, err := mirror.NewClient(srv.URL, tlsconf.Config{}, 0)
@@ -648,3 +648,10 @@ func TestAModelsFamilyIsHowItsNameStarts(t *testing.T) {
}
}
}
// privateOK lets a test's manager fetch from the httptest upstreams on
// loopback, which a real one refuses.
func privateOK(m *mirror.Manager) *mirror.Manager {
m.AllowPrivateSources()
return m
}