Compare commits
299
Commits
No files matched your search
@@ -38,3 +38,11 @@ jobs:
|
||||
|
||||
- name: Build
|
||||
run: go build ./...
|
||||
|
||||
# The NVML backend needs cgo and cannot run without a driver, but it must
|
||||
# at least keep compiling: it is the only source of NVLink topology and a
|
||||
# break would otherwise surface on a GPU node, far from here.
|
||||
- name: Build the NVML backend
|
||||
env:
|
||||
CGO_ENABLED: "1"
|
||||
run: go build -tags nvml ./...
|
||||
@@ -0,0 +1,205 @@
|
||||
# Attaches binaries to a release when a version tag is pushed.
|
||||
#
|
||||
# Without this a release is a git tag, and installing llmgw on a node means
|
||||
# having a Go toolchain and the source. The onboarding script copies a binary
|
||||
# from wherever it is run, which works and quietly makes whoever runs it the
|
||||
# build server.
|
||||
name: Release
|
||||
|
||||
on:
|
||||
push:
|
||||
tags: ["v*"]
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
# The built-in token every job is given carries whatever the instance's
|
||||
# default mode allows, and under the restricted default that is read-only
|
||||
# on releases — so the build would succeed and the publish would be
|
||||
# refused. `contents` covers code and releases together, which is what
|
||||
# creating a release needs: `releases: write` alone was refused, because
|
||||
# creating one can create its tag.
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
# Tags are what the version string is built from, and a shallow clone
|
||||
# has none of them.
|
||||
fetch-depth: 0
|
||||
|
||||
- uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version-file: go.mod
|
||||
cache: true
|
||||
|
||||
- name: Build
|
||||
run: |
|
||||
set -euo pipefail
|
||||
version="${GITHUB_REF_NAME:-$(git describe --tags --always)}"
|
||||
mkdir -p dist
|
||||
|
||||
# CGO_ENABLED=0 is the whole point: the binary has to run on a node
|
||||
# whose libc is older than the builder's, and llmgw has no C
|
||||
# dependency — modernc.org/sqlite is pure Go precisely so this stays
|
||||
# true.
|
||||
for target in linux/amd64 linux/arm64 darwin/amd64 darwin/arm64; do
|
||||
os="${target%/*}"; arch="${target#*/}"
|
||||
name="llmgw-${version}-${os}-${arch}"
|
||||
echo "building $name"
|
||||
CGO_ENABLED=0 GOOS="$os" GOARCH="$arch" go build \
|
||||
-trimpath \
|
||||
-ldflags "-s -w -X gitea.mynodes.xyz/ilker/llmgw/internal/version.Version=${version}" \
|
||||
-o "dist/${name}" ./cmd/llmgw
|
||||
(cd dist && sha256sum "${name}" >> checksums.txt)
|
||||
done
|
||||
|
||||
# The reverse proxy configurations, so a release page carries
|
||||
# everything install.md asks for and the repository is not needed.
|
||||
for f in deploy/proxy/Caddyfile deploy/proxy/nginx.conf; do
|
||||
cp "$f" dist/
|
||||
(cd dist && sha256sum "$(basename "$f")" >> checksums.txt)
|
||||
done
|
||||
|
||||
# An NVML build for nodes with NVLink, which the portable binary
|
||||
# cannot serve.
|
||||
#
|
||||
# NVML is the only source of NVLink topology, its bindings need cgo,
|
||||
# and cgo is exactly what CGO_ENABLED=0 above exists to avoid. So the
|
||||
# portable binary silently falls back to taking the lowest free GPU
|
||||
# indices — right on an SXM box with NVSwitch, and wrong on a PCIe
|
||||
# box where NVLink bridges pair specific cards: a tensor group
|
||||
# straddling two bridges runs several times slower and nothing says
|
||||
# so. Shipping only the portable build meant nobody could have that
|
||||
# topology without building it themselves.
|
||||
#
|
||||
# Dynamically linked against the host's libc by necessity, so it is
|
||||
# published beside the portable one rather than instead of it.
|
||||
if command -v gcc >/dev/null; then
|
||||
nvml_name="llmgw-${version}-linux-amd64-nvml"
|
||||
CGO_ENABLED=1 GOOS=linux GOARCH=amd64 go build \
|
||||
-tags nvml -trimpath \
|
||||
-ldflags "-s -w -X gitea.mynodes.xyz/ilker/llmgw/internal/version.Version=${version}" \
|
||||
-o "dist/${nvml_name}" ./cmd/llmgw
|
||||
(cd dist && sha256sum "${nvml_name}" >> checksums.txt)
|
||||
else
|
||||
echo "no C toolchain; skipping the nvml build" >&2
|
||||
fi
|
||||
|
||||
# Proof the artifact runs at all before it is published. It cannot
|
||||
# test the cross-compiled ones, which is why the native build is
|
||||
# checked rather than assumed.
|
||||
./dist/llmgw-${version}-linux-amd64 --version
|
||||
cat dist/checksums.txt
|
||||
|
||||
# Signed when the repository has a signing key, so a machine that trusts
|
||||
# it installs nothing the server alone could have put there. Without the
|
||||
# secret nothing is signed and releases are checked on their checksums,
|
||||
# as before; a build that trusts a key refuses those, so the secret has
|
||||
# to be in place before trusted.pub is filled in.
|
||||
- name: Sign the checksums
|
||||
env:
|
||||
RELEASE_SIGNING_KEY: ${{ secrets.RELEASE_SIGNING_KEY }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
version="${GITHUB_REF_NAME:-$(git describe --tags --always)}"
|
||||
if [ -z "${RELEASE_SIGNING_KEY:-}" ]; then
|
||||
echo "no RELEASE_SIGNING_KEY secret; releases are not signed"
|
||||
exit 0
|
||||
fi
|
||||
./dist/llmgw-${version}-linux-amd64 release sign dist/checksums.txt
|
||||
echo "signed: dist/checksums.txt.sig"
|
||||
|
||||
- name: Attach to the release
|
||||
env:
|
||||
TOKEN: ${{ secrets.GITEA_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
# Gitea gives every job a GITEA_TOKEN, so this is not the usual
|
||||
# "you forgot to add a secret" — an empty one means the job was denied
|
||||
# a token, which is an instance setting rather than a repository one.
|
||||
if [ -z "${TOKEN:-}" ]; then
|
||||
echo "The job was given no GITEA_TOKEN, so nothing can be uploaded."
|
||||
echo "Check Settings -> Actions -> General; the binaries are in dist/."
|
||||
exit 1
|
||||
fi
|
||||
# workflow_dispatch can be fired from a branch, where GITHUB_REF_NAME
|
||||
# is the branch name. Uploading then would create a release called
|
||||
# "main". The binaries are still built and kept as a check that the
|
||||
# release build works; only the publishing is a tag's business.
|
||||
if [ "${GITHUB_REF_TYPE:-}" != "tag" ]; then
|
||||
echo "not a tag (ref=${GITHUB_REF_NAME}); built, not published"
|
||||
exit 0
|
||||
fi
|
||||
tag="${GITHUB_REF_NAME}"
|
||||
api="${GITHUB_SERVER_URL}/api/v1/repos/${GITHUB_REPOSITORY}"
|
||||
|
||||
# Gitea materialises a release row for every pushed tag, so this
|
||||
# nearly always finds one and the create below is for a tag pushed
|
||||
# some other way.
|
||||
#
|
||||
# The id is read with grep rather than sed because sed's .* is greedy:
|
||||
# "id" appears again inside the nested author object, so the old
|
||||
# pattern took the author's id when it matched at all — and it did not
|
||||
# match "id": 13 with a space, which is what actually happened. The
|
||||
# first "id" in a release object is the release's own.
|
||||
release_id() {
|
||||
curl -sS -H "Authorization: token $TOKEN" "$api/releases/tags/$tag" \
|
||||
| grep -o '"id"[[:space:]]*:[[:space:]]*[0-9]*' \
|
||||
| head -1 | grep -o '[0-9]*$'
|
||||
}
|
||||
|
||||
id=$(release_id || true)
|
||||
if [ -z "$id" ]; then
|
||||
echo "no release for $tag yet; creating one"
|
||||
out=$(curl -sS -w '\n%{http_code}' -X POST \
|
||||
-H "Authorization: token $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "{\"tag_name\":\"$tag\",\"name\":\"$tag\"}" \
|
||||
"$api/releases")
|
||||
echo "create returned $(printf '%s' "$out" | tail -1)"
|
||||
# Re-read rather than parsing the create response: a 409 means it
|
||||
# was there all along, which is a success for our purposes.
|
||||
id=$(release_id || true)
|
||||
[ -n "$id" ] || printf '%s\n' "$out" | head -5 >&2
|
||||
fi
|
||||
[ -n "$id" ] || { echo "could not find or create the release" >&2; exit 1; }
|
||||
echo "release id $id"
|
||||
|
||||
for f in dist/*; do
|
||||
echo "uploading $(basename "$f")"
|
||||
curl -sS -f -X POST -H "Authorization: token $TOKEN" \
|
||||
-F "attachment=@$f" \
|
||||
"$api/releases/$id/assets?name=$(basename "$f")" >/dev/null \
|
||||
|| { echo "upload of $(basename "$f") was refused" >&2; exit 1; }
|
||||
done
|
||||
echo "attached $(ls dist | wc -l) files to $tag"
|
||||
|
||||
# The release page says what changed, copied from CHANGELOG.md. Gitea makes
|
||||
# an empty release for every tag and the job above only attaches files, so
|
||||
# every release until this existed was a list of binaries and nothing else.
|
||||
#
|
||||
# After the build, because that is what guarantees the release exists. Run by
|
||||
# hand from a branch it syncs every version in the changelog, which is how
|
||||
# the releases made before this got their notes.
|
||||
notes:
|
||||
needs: build
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Copy the changelog onto the release
|
||||
env:
|
||||
GITEA_TOKEN: ${{ secrets.GITEA_TOKEN }}
|
||||
GITEA_INSTANCE: ${{ github.server_url }}
|
||||
GITEA_REPO: ${{ github.repository }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
if [ "${GITHUB_REF_TYPE:-}" = "tag" ]; then
|
||||
python3 scripts/gitea/release_notes.py --tag "${GITHUB_REF_NAME}"
|
||||
else
|
||||
python3 scripts/gitea/release_notes.py
|
||||
fi
|
||||
+1906
-4
File diff suppressed because it is too large.
Load diff
@@ -2,7 +2,7 @@ MODULE := gitea.mynodes.xyz/ilker/llmgw
|
||||
VERSION ?= $(shell git describe --tags --always --dirty 2>/dev/null || echo dev)
|
||||
LDFLAGS := -X '$(MODULE)/internal/version.Version=$(VERSION)'
|
||||
|
||||
.PHONY: build test race lint fmt clean rename
|
||||
.PHONY: build test race lint fmt clean rename proto proto-tools
|
||||
|
||||
build:
|
||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw ./cmd/llmgw
|
||||
@@ -20,6 +20,17 @@ fmt:
|
||||
go fmt ./...
|
||||
go vet ./...
|
||||
|
||||
# Regenerate the agent protocol. Generated code is committed, so this only
|
||||
# needs running when api/proto changes.
|
||||
proto: proto-tools
|
||||
buf lint
|
||||
buf generate
|
||||
|
||||
proto-tools:
|
||||
go install github.com/bufbuild/buf/cmd/buf@latest
|
||||
go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
|
||||
go install connectrpc.com/connect/cmd/protoc-gen-connect-go@latest
|
||||
|
||||
clean:
|
||||
rm -rf bin dist
|
||||
|
||||
|
||||
@@ -1,11 +1,23 @@
|
||||
# llmgw
|
||||
|
||||
A single-binary LLM gateway for clusters with multiple InfiniBand fabrics and
|
||||
heterogeneous GPUs.
|
||||
A single-binary LLM gateway for clusters with multiple InfiniBand fabrics,
|
||||
NVLink domains and heterogeneous GPUs.
|
||||
|
||||
It serves LLMs running on your own GPU servers (vLLM / SGLang) behind one
|
||||
OpenAI-compatible endpoint. Design document: [llmgw-plan.md](llmgw-plan.md).
|
||||
|
||||
## Getting started
|
||||
|
||||
Starting from nothing — no packages installed, no drivers, no nodes — follow
|
||||
[docs/install.md](docs/install.md): it builds a control plane, adds a node and
|
||||
serves a model, and says plainly which parts llmgw does not install for you.
|
||||
|
||||
| Guide | For |
|
||||
|---|---|
|
||||
| [docs/install.md](docs/install.md) | Standing up a cluster from scratch |
|
||||
| [docs/node-onboarding.md](docs/node-onboarding.md) | Adding nodes to one that exists |
|
||||
| [docs/onboarding.md](docs/onboarding.md) | Giving people accounts and keys |
|
||||
|
||||
## Status
|
||||
|
||||
Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
@@ -14,14 +26,61 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
|---|---|---|---|
|
||||
| v0.1.0 | 0 | Skeleton and sim-engine | released |
|
||||
| v0.2.0 | 1 | Static router (filter/score/pick, streaming proxy) | released |
|
||||
| v0.3.0 | 1.5 | Admission control (priority bands) | planned |
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | planned |
|
||||
| v0.5.0 | 3 | Control plane: catalog, scheduler, reconciler | planned |
|
||||
| v0.6.0 | 4 | Multi-node placement and fabric awareness | planned |
|
||||
| v0.3.0 | 1.5 | Admission control (priority bands) | released |
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, engine supervision | released |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler, ctl | released |
|
||||
| v0.6.0 | 4 | Multi-node placement and fabric awareness | released |
|
||||
| v0.7.0 | — | Conversations, agent memory, disk management | released |
|
||||
| v0.8.0 | — | PostgreSQL, leader election, agent failover | released |
|
||||
| v0.9.0 | — | Managing people and the tenancy tree | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
As of v0.2.0 the router serves traffic. There is still no admission control, so
|
||||
a batch workload can starve interactive traffic — that is what v0.3.0 is for.
|
||||
As of v0.6.0 a cluster runs unattended: agents report their hardware, the
|
||||
control plane places replicas — including ones larger than a single machine —
|
||||
models load when someone asks for them and stop when nobody does, and keys,
|
||||
budgets and rate limits decide who may use any of it.
|
||||
|
||||
Before v1.0.0: tracing, reference documentation, and filtering on `ctl nodes`
|
||||
and `ctl replicas` for fleets too large to read a whole table of. Two items
|
||||
that were on this list are not gaps: a model falls through to its fallbacks
|
||||
when there is no capacity as well as when a budget is spent, and routing
|
||||
profiles are chosen by the key rather than by a request header on purpose —
|
||||
a header would let any client promote itself into the top priority band.
|
||||
|
||||
One thing is worth saying plainly:
|
||||
multi-node replicas are proven on a simulated cluster and on real hardware
|
||||
without a fabric. The InfiniBand, NVLink and NCCL paths — telling fabrics apart
|
||||
by their subnet managers, mapping them, finding NVLink domains, launching across
|
||||
nodes — are written against the formats and failure modes those tools actually
|
||||
have and tested against them, but they have not yet run on a machine that has
|
||||
them. On such a machine, run `llmgw doctor --json` first and keep the report;
|
||||
[docs/fabrics.md](docs/fabrics.md) says what to look at after that.
|
||||
|
||||
## Documentation
|
||||
|
||||
- [docs/quick-install.md](docs/quick-install.md) — installing on servers, in
|
||||
six steps: start here
|
||||
- [docs/install.md](docs/install.md) — the same from nothing, with every option
|
||||
explained: control plane, nodes, engines, weights, first request
|
||||
- [docs/quickstart.md](docs/quickstart.md) — one workstation, from source
|
||||
- [docs/fabrics.md](docs/fabrics.md) — InfiniBand fabrics and NVLink: what
|
||||
llmgw needs of them, checking a node with `llmgw doctor`, and what to look at
|
||||
when a replica spanning nodes will not run
|
||||
- [docs/choosing-models.md](docs/choosing-models.md) — which models to run,
|
||||
how many of each, and where the weights should live
|
||||
- [docs/data.md](docs/data.md) — what llmgw stores, what it does not, and
|
||||
what leaves the machine; written for a security review
|
||||
- [docs/reference.md](docs/reference.md) — every configuration key, generated
|
||||
from the structs and checked by a test, so it cannot describe a setting that
|
||||
does not exist or miss one that does
|
||||
- [deploy/mirror/](deploy/mirror/) — a model mirror, which is the cheapest
|
||||
and most useful thing to build around llmgw
|
||||
- [deploy/proxy/](deploy/proxy/) — Caddy and nginx in front of the client
|
||||
port, with streaming and the forwarded headers right
|
||||
- [deploy/ha/](deploy/ha/) — running two control planes, and the one health
|
||||
check that has to be right
|
||||
- [docs/node-onboarding.md](docs/node-onboarding.md),
|
||||
[docs/onboarding.md](docs/onboarding.md) — adding machines, and adding people
|
||||
|
||||
## Building
|
||||
|
||||
@@ -31,6 +90,25 @@ make test # unit tests
|
||||
make race # with the race detector
|
||||
```
|
||||
|
||||
## Operating a cluster
|
||||
|
||||
```sh
|
||||
llmgw ctl nodes # capacity, health, what is running where
|
||||
llmgw ctl models # catalog and replica counts
|
||||
llmgw ctl replicas --model llama # replicas and their shards
|
||||
llmgw ctl explain-placement llama # why a model has not been placed
|
||||
llmgw ctl drain gpu-a-01 # move work off a node, one replica at a time
|
||||
llmgw ctl cordon gpu-a-01 # stop new placement, leave what runs
|
||||
sudo llmgw doctor # on a node: what its GPUs, NVLink and fabric need
|
||||
```
|
||||
|
||||
`ctl` talks to the admin port (`:9090` by default), not the client-facing one.
|
||||
Add `--json` to any command for machine-readable output.
|
||||
|
||||
`explain-placement` is the one worth knowing about: "unschedulable" on its own
|
||||
says nothing about which of six constraints bit, so it prints every candidate
|
||||
considered and the reason it was rejected.
|
||||
|
||||
## Running the router
|
||||
|
||||
```sh
|
||||
@@ -51,6 +129,12 @@ For anything real use a config file; [deploy/examples/server.yaml](deploy/exampl
|
||||
is commented throughout. Metrics, health and pprof are on the admin port (`:9090`
|
||||
by default), deliberately separate from the client-facing one.
|
||||
|
||||
Requests carry a priority band (`x-llmgw-band`). Under load, low bands wait at
|
||||
the gateway while high bands keep flowing, so the cluster degrades from the
|
||||
bottom up instead of everywhere at once. Measured effect: interactive p99 stays
|
||||
within 1.24x of idle while 500 concurrent batch requests run, against 58x with
|
||||
admission off.
|
||||
|
||||
Routing profiles select how endpoints are scored: `chat` favours session
|
||||
affinity, `batch` favours the shortest queue, `long-context` favours free KV
|
||||
cache. Pick one per model in the config, or per request with the
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,119 @@
|
||||
// Code generated by protoc-gen-connect-go. DO NOT EDIT.
|
||||
//
|
||||
// Source: llmgw/v1/agent.proto
|
||||
|
||||
// Package llmgw.v1 is the protocol between a node agent and the control plane
|
||||
// (plan §5.6).
|
||||
//
|
||||
// The agent dials out and holds a bidirectional stream open. That direction is
|
||||
// deliberate: GPU nodes sit behind firewalls and NAT far more often than a
|
||||
// control plane does, and an outbound connection needs no inbound rule on every
|
||||
// node in the fleet.
|
||||
package llmgwv1connect
|
||||
|
||||
import (
|
||||
connect "connectrpc.com/connect"
|
||||
context "context"
|
||||
errors "errors"
|
||||
v1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
http "net/http"
|
||||
strings "strings"
|
||||
)
|
||||
|
||||
// This is a compile-time assertion to ensure that this generated file and the connect package are
|
||||
// compatible. If you get a compiler error that this constant is not defined, this code was
|
||||
// generated with a version of connect newer than the one compiled into your binary. You can fix the
|
||||
// problem by either regenerating this code with an older version of connect or updating the connect
|
||||
// version compiled into your binary.
|
||||
const _ = connect.IsAtLeastVersion1_13_0
|
||||
|
||||
const (
|
||||
// AgentServiceName is the fully-qualified name of the AgentService service.
|
||||
AgentServiceName = "llmgw.v1.AgentService"
|
||||
)
|
||||
|
||||
// These constants are the fully-qualified names of the RPCs defined in this package. They're
|
||||
// exposed at runtime as Spec.Procedure and as the final two segments of the HTTP route.
|
||||
//
|
||||
// Note that these are different from the fully-qualified method names used by
|
||||
// google.golang.org/protobuf/reflect/protoreflect. To convert from these constants to
|
||||
// reflection-formatted method names, remove the leading slash and convert the remaining slash to a
|
||||
// period.
|
||||
const (
|
||||
// AgentServiceConnectProcedure is the fully-qualified name of the AgentService's Connect RPC.
|
||||
AgentServiceConnectProcedure = "/llmgw.v1.AgentService/Connect"
|
||||
)
|
||||
|
||||
// AgentServiceClient is a client for the llmgw.v1.AgentService service.
|
||||
type AgentServiceClient interface {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
Connect(context.Context) *connect.BidiStreamForClient[v1.AgentMessage, v1.ControlMessage]
|
||||
}
|
||||
|
||||
// NewAgentServiceClient constructs a client for the llmgw.v1.AgentService service. By default, it
|
||||
// uses the Connect protocol with the binary Protobuf Codec, asks for gzipped responses, and sends
|
||||
// uncompressed requests. To use the gRPC or gRPC-Web protocols, supply the connect.WithGRPC() or
|
||||
// connect.WithGRPCWeb() options.
|
||||
//
|
||||
// The URL supplied here should be the base URL for the Connect or gRPC server (for example,
|
||||
// http://api.acme.com or https://acme.com/grpc).
|
||||
func NewAgentServiceClient(httpClient connect.HTTPClient, baseURL string, opts ...connect.ClientOption) AgentServiceClient {
|
||||
baseURL = strings.TrimRight(baseURL, "/")
|
||||
agentServiceMethods := v1.File_llmgw_v1_agent_proto.Services().ByName("AgentService").Methods()
|
||||
return &agentServiceClient{
|
||||
connect: connect.NewClient[v1.AgentMessage, v1.ControlMessage](
|
||||
httpClient,
|
||||
baseURL+AgentServiceConnectProcedure,
|
||||
connect.WithSchema(agentServiceMethods.ByName("Connect")),
|
||||
connect.WithClientOptions(opts...),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
// agentServiceClient implements AgentServiceClient.
|
||||
type agentServiceClient struct {
|
||||
connect *connect.Client[v1.AgentMessage, v1.ControlMessage]
|
||||
}
|
||||
|
||||
// Connect calls llmgw.v1.AgentService.Connect.
|
||||
func (c *agentServiceClient) Connect(ctx context.Context) *connect.BidiStreamForClient[v1.AgentMessage, v1.ControlMessage] {
|
||||
return c.connect.CallBidiStream(ctx)
|
||||
}
|
||||
|
||||
// AgentServiceHandler is an implementation of the llmgw.v1.AgentService service.
|
||||
type AgentServiceHandler interface {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
Connect(context.Context, *connect.BidiStream[v1.AgentMessage, v1.ControlMessage]) error
|
||||
}
|
||||
|
||||
// NewAgentServiceHandler builds an HTTP handler from the service implementation. It returns the
|
||||
// path on which to mount the handler and the handler itself.
|
||||
//
|
||||
// By default, handlers support the Connect, gRPC, and gRPC-Web protocols with the binary Protobuf
|
||||
// and JSON codecs. They also support gzip compression.
|
||||
func NewAgentServiceHandler(svc AgentServiceHandler, opts ...connect.HandlerOption) (string, http.Handler) {
|
||||
agentServiceMethods := v1.File_llmgw_v1_agent_proto.Services().ByName("AgentService").Methods()
|
||||
agentServiceConnectHandler := connect.NewBidiStreamHandler(
|
||||
AgentServiceConnectProcedure,
|
||||
svc.Connect,
|
||||
connect.WithSchema(agentServiceMethods.ByName("Connect")),
|
||||
connect.WithHandlerOptions(opts...),
|
||||
)
|
||||
return "/llmgw.v1.AgentService/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case AgentServiceConnectProcedure:
|
||||
agentServiceConnectHandler.ServeHTTP(w, r)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// UnimplementedAgentServiceHandler returns CodeUnimplemented from all methods.
|
||||
type UnimplementedAgentServiceHandler struct{}
|
||||
|
||||
func (UnimplementedAgentServiceHandler) Connect(context.Context, *connect.BidiStream[v1.AgentMessage, v1.ControlMessage]) error {
|
||||
return connect.NewError(connect.CodeUnimplemented, errors.New("llmgw.v1.AgentService.Connect is not implemented"))
|
||||
}
|
||||
@@ -0,0 +1,716 @@
|
||||
syntax = "proto3";
|
||||
|
||||
// Package llmgw.v1 is the protocol between a node agent and the control plane
|
||||
// (plan §5.6).
|
||||
//
|
||||
// The agent dials out and holds a bidirectional stream open. That direction is
|
||||
// deliberate: GPU nodes sit behind firewalls and NAT far more often than a
|
||||
// control plane does, and an outbound connection needs no inbound rule on every
|
||||
// node in the fleet.
|
||||
package llmgw.v1;
|
||||
|
||||
option go_package = "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1;llmgwv1";
|
||||
|
||||
service AgentService {
|
||||
// Connect opens the agent's stream. It stays open for the agent's lifetime;
|
||||
// losing it is how each side learns the other is gone.
|
||||
rpc Connect(stream AgentMessage) returns (stream ControlMessage);
|
||||
}
|
||||
|
||||
// AgentMessage is anything a node tells the control plane.
|
||||
message AgentMessage {
|
||||
oneof msg {
|
||||
Register register = 1;
|
||||
Heartbeat heartbeat = 2;
|
||||
ShardStatus shard = 3;
|
||||
LogsReply logs = 4;
|
||||
FabricScanReply fabric_scan = 5;
|
||||
FabricTestReply fabric_test = 6;
|
||||
UpgradeReply upgrade = 7;
|
||||
}
|
||||
}
|
||||
|
||||
// ControlMessage is anything the control plane tells a node.
|
||||
message ControlMessage {
|
||||
oneof msg {
|
||||
StartShard start = 1;
|
||||
StopShard stop = 2;
|
||||
DrainShard drain = 3;
|
||||
Ack ack = 4;
|
||||
RemoveWeights remove_weights = 5;
|
||||
Provision provision = 6;
|
||||
LogsRequest logs = 7;
|
||||
FabricScanRequest fabric_scan = 8;
|
||||
FabricTestRequest fabric_test = 9;
|
||||
UpgradeAgent upgrade = 10;
|
||||
}
|
||||
}
|
||||
|
||||
// Register is the first message on every stream. It carries the node's full
|
||||
// hardware inventory, so the control plane needs no separate discovery call and
|
||||
// a reconnecting agent automatically refreshes what the scheduler knows.
|
||||
message Register {
|
||||
string node_id = 1;
|
||||
string hostname = 2;
|
||||
// mgmt_addr is the address engines are reachable on, and the address NCCL
|
||||
// bootstraps over. Never a fabric address: NCCL rendezvouses over TCP before
|
||||
// it moves to RDMA.
|
||||
string mgmt_addr = 3;
|
||||
map<string, string> labels = 4;
|
||||
|
||||
repeated GPU gpus = 5;
|
||||
// fabrics is empty on a node without InfiniBand. Such a node is its own
|
||||
// single-node island and can never host a multi-node replica.
|
||||
repeated FabricAttachment fabrics = 6;
|
||||
|
||||
string agent_version = 7;
|
||||
// engine_versions decides which CLI flags and metric names apply; both drift
|
||||
// between engine releases (plan §14).
|
||||
map<string, string> engine_versions = 8;
|
||||
|
||||
// running_shards are the engines this agent is already supervising, including
|
||||
// ones adopted from a previous agent.
|
||||
//
|
||||
// This is what lets a control plane restart recover instead of orphaning live
|
||||
// engines. Without it the control plane forgets every replica, places new
|
||||
// ones, and the old engines keep holding GPUs the inventory believes are free
|
||||
// — so the new engines fail to allocate, with an error naming the wrong
|
||||
// process (plan §6.5).
|
||||
repeated RunningShard running_shards = 9;
|
||||
|
||||
// memory_total_mb is the machine's RAM, and 0 means the agent did not say —
|
||||
// which is what every agent older than this field reports, and is treated as
|
||||
// no constraint rather than as a node with no memory.
|
||||
//
|
||||
// Placement needs it because a CPU profile has no card to be free, so every
|
||||
// ready node looks like a candidate: a 14B model whose working set is 24 GB
|
||||
// would be placed just as happily on a node with 8 GB.
|
||||
int64 memory_total_mb = 10;
|
||||
|
||||
// weights are the model files this node can see, from the directories its
|
||||
// configuration names.
|
||||
//
|
||||
// How they got there is the operator's business and deliberately outside
|
||||
// llmgw's: a shared filesystem, an image exported over the fabric, rsync, or
|
||||
// a copy made by hand all look identical here. What matters to placement is
|
||||
// only whether they are present, because a node without them either cannot
|
||||
// start the replica at all or must first move tens of gigabytes.
|
||||
//
|
||||
// Empty means the node was not asked to look, which is not the same as having
|
||||
// nothing — an agent with no model directories configured behaves exactly as
|
||||
// every agent did before this existed, and constrains nothing.
|
||||
repeated LocalWeights weights = 11;
|
||||
|
||||
// storage is the disk behind those weights, one entry per filesystem.
|
||||
//
|
||||
// Reported so an operator can see a node filling up before a fetch dies at
|
||||
// ninety percent, and so placement can refuse a node that cannot hold a
|
||||
// model rather than discovering it during the download.
|
||||
//
|
||||
// Empty means the node was not asked to look, exactly as for weights, and
|
||||
// constrains nothing.
|
||||
repeated Storage storage = 12;
|
||||
|
||||
// host is what the machine is, so the control plane can have the mirror
|
||||
// fetch software built for it rather than assuming every node matches the
|
||||
// first.
|
||||
Host host = 13;
|
||||
|
||||
// engines_probed says engine_versions is an answer. Without it an empty map
|
||||
// meant "an agent too old to say" and constrained nothing — so a freshly
|
||||
// built node with no engine installed yet was a candidate for every model.
|
||||
bool engines_probed = 14;
|
||||
// cpu_count is the CPUs the agent may run on. Placement reserves cores
|
||||
// for CPU replicas out of it, the way it reserves memory out of
|
||||
// memory_total_mb; zero, from an older agent, reserves nothing.
|
||||
int32 cpu_count = 18;
|
||||
// imex is the service that lets a multi-node NVLink domain share memory.
|
||||
IMEX imex = 15;
|
||||
|
||||
// installable_engines are the engines this agent knows how to install
|
||||
// itself. The control plane asks it for no others. Empty is an agent older
|
||||
// than the field, which could install llama.cpp and nothing else: asked for
|
||||
// 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.
|
||||
message Host {
|
||||
// From /etc/os-release: "ubuntu", "24.04", "noble".
|
||||
string os_id = 1;
|
||||
string os_version = 2;
|
||||
string os_codename = 3;
|
||||
// Debian's spelling, which is what package indexes are keyed by: "amd64",
|
||||
// "arm64".
|
||||
string arch = 4;
|
||||
// The vector extensions inference cares about, of those the CPU has. A VM
|
||||
// given QEMU's default CPU model has none of them, and runs, and is five
|
||||
// times slower, and nothing says why.
|
||||
repeated string cpu_flags = 5;
|
||||
// nvidia_driver is the driver version, empty when there is none. A card
|
||||
// without a driver is not a GPU anything can use.
|
||||
string nvidia_driver = 6;
|
||||
}
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// Per filesystem rather than per configured directory: free space belongs to
|
||||
// the disk, and two model directories on one disk are one pool of bytes.
|
||||
message Storage {
|
||||
// path is one directory on this filesystem, so the operator recognises it.
|
||||
string path = 1;
|
||||
// free_bytes is what a non-root process can write — not the root reserve,
|
||||
// which is space the agent cannot have.
|
||||
int64 free_bytes = 2;
|
||||
int64 total_bytes = 3;
|
||||
// shared marks a filesystem the node did not fill. Free space on one is
|
||||
// shared with every other node using it, so two nodes each reporting 500 GB
|
||||
// free on the same filer have 500 GB between them and not a terabyte.
|
||||
bool shared = 4;
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
//
|
||||
// It carries enough for the control plane to rebuild the replica it belongs to,
|
||||
// which is more than a shard id: the replica's identity, its pool and fabric,
|
||||
// and the leader endpoint clients reach it through.
|
||||
message RunningShard {
|
||||
string replica_id = 1;
|
||||
string shard_id = 2;
|
||||
int32 rank = 3;
|
||||
int32 nnodes = 4;
|
||||
repeated int32 gpu_indices = 5;
|
||||
repeated string hcas = 6;
|
||||
int32 http_port = 7;
|
||||
ShardStatus.State state = 8;
|
||||
|
||||
// model and profile name the catalog entry this shard was started from.
|
||||
string model = 9;
|
||||
string profile = 10;
|
||||
// pool and fabric are echoed back from the StartShard that created it; the
|
||||
// agent has no way to derive them.
|
||||
string pool = 11;
|
||||
string fabric = 12;
|
||||
|
||||
string leader_addr = 13;
|
||||
int32 leader_port = 14;
|
||||
|
||||
// profile_fp is the fingerprint of the profile this shard was started from,
|
||||
// echoed back from StartShard. The control plane keeps no record of a running
|
||||
// replica across its own restart, so without this a shard adopted after a
|
||||
// restart could not be compared with the catalog — which is exactly when the
|
||||
// catalog is most likely to have been edited.
|
||||
string profile_fp = 15;
|
||||
}
|
||||
|
||||
message GPU {
|
||||
int32 index = 1;
|
||||
string uuid = 2;
|
||||
string name = 3;
|
||||
int32 memory_total_mb = 4;
|
||||
int32 memory_used_mb = 5;
|
||||
string compute_capability = 6;
|
||||
int32 numa_node = 7;
|
||||
string pci_bus_id = 8;
|
||||
string driver_version = 9;
|
||||
string cuda_version = 10;
|
||||
// nvlink_peers is only populated by agents built with the nvml tag.
|
||||
repeated int32 nvlink_peers = 11;
|
||||
// problems keeps a GPU out of placement: ECC faults, or memory already held
|
||||
// by another process.
|
||||
repeated string problems = 12;
|
||||
int32 temperature_c = 13;
|
||||
int32 ecc_errors = 14;
|
||||
// utilization_percent is how busy the card was over the last sample period.
|
||||
// Zero when the driver is too old to report it, which is not the same as
|
||||
// idle: the metric is simply absent rather than wrong.
|
||||
int32 utilization_percent = 15;
|
||||
// nic_paths is how far each network adapter is from this card, by HCA
|
||||
// name: PIX, PXB, PHB, NODE or SYS, from nvidia-smi topo.
|
||||
map<string, string> nic_paths = 16;
|
||||
// The card's multi-node NVLink membership, from nvidia-smi -q: state is
|
||||
// Completed once the fabric manager has joined it; cluster_uuid and
|
||||
// clique_id name the domain it joined.
|
||||
string fabric_state = 17;
|
||||
string fabric_status = 18;
|
||||
string cluster_uuid = 19;
|
||||
string clique_id = 20;
|
||||
string fabric_bandwidth = 21;
|
||||
}
|
||||
|
||||
message FabricAttachment {
|
||||
// fabric is the operator-assigned name. It is not derived from the GID
|
||||
// prefix: OpenSM defaults to fe80:: on every fabric, so a matching prefix
|
||||
// proves nothing (plan §5.2).
|
||||
string fabric = 1;
|
||||
repeated string hcas = 2;
|
||||
int32 rate_gbps = 3;
|
||||
}
|
||||
|
||||
// Heartbeat says the node is alive and reports anything that changed cheaply.
|
||||
// LocalWeights is one model file or directory a node can see.
|
||||
message LocalWeights {
|
||||
// path is what an engine would be given, so it can be compared against a
|
||||
// model's source without interpretation.
|
||||
string path = 1;
|
||||
int64 size_bytes = 2;
|
||||
// shared marks weights on a filesystem the node did not have to fill —
|
||||
// reported so a fabric with a shared filer is not mistaken for one where
|
||||
// every node happens to have fetched its own copy.
|
||||
bool shared = 3;
|
||||
}
|
||||
|
||||
// NodeLoad is a node's use of its own CPU and memory, and each engine's share.
|
||||
//
|
||||
// Carried on the heartbeat rather than scraped from the node: the control
|
||||
// plane already exports the fleet from one place, and a node exporter per
|
||||
// machine would be a scrape target, a port and a firewall rule each.
|
||||
message NodeLoad {
|
||||
// cpu_percent is the whole machine, 0 to 100 however many cores it has.
|
||||
double cpu_percent = 1;
|
||||
int32 cpu_count = 2;
|
||||
// The kernel's run-queue averages over 1, 5 and 15 minutes.
|
||||
double load1 = 3;
|
||||
double load5 = 4;
|
||||
double load15 = 5;
|
||||
// memory_used_bytes is total less available: what the machine could not
|
||||
// hand to a new process without evicting something.
|
||||
int64 memory_used_bytes = 6;
|
||||
int64 memory_available_bytes = 7;
|
||||
repeated ShardLoad shards = 8;
|
||||
// What the machine's own sensors say; zero is not measured. cpu_temp_c is
|
||||
// the hottest CPU sensor, cpu_power_w the CPU packages' draw from RAPL,
|
||||
// inlet_temp_c the board's inlet or ambient sensor where it has one.
|
||||
double cpu_temp_c = 9;
|
||||
double cpu_power_w = 10;
|
||||
double inlet_temp_c = 11;
|
||||
repeated GPUSensor gpus = 12;
|
||||
}
|
||||
|
||||
// GPUSensor is one card's heat, power, and what is holding it back.
|
||||
message GPUSensor {
|
||||
int32 index = 1;
|
||||
double temp_c = 2;
|
||||
double memory_temp_c = 3;
|
||||
double power_w = 4;
|
||||
double power_limit_w = 5;
|
||||
// throttle names what is holding the clocks down now: sw_power_cap,
|
||||
// hw_slowdown, sw_thermal, hw_thermal, power_brake.
|
||||
repeated string throttle = 6;
|
||||
}
|
||||
|
||||
// ShardLoad is one engine process group's use.
|
||||
message ShardLoad {
|
||||
string shard_id = 1;
|
||||
// cpu_cores is in cores, the way top counts a process: 2.5 is two and a
|
||||
// half cores busy, which a percentage of the machine would hide on a large
|
||||
// one.
|
||||
double cpu_cores = 2;
|
||||
int64 memory_rss_bytes = 3;
|
||||
}
|
||||
|
||||
// IBPort is one InfiniBand or RoCE port's state and counters.
|
||||
message IBPort {
|
||||
string hca = 1;
|
||||
int32 port = 2;
|
||||
// state is ACTIVE, INIT, DOWN...; phys_state tells a cable unplugged
|
||||
// (Polling, Disabled) from a link up with no subnet manager (LinkUp and INIT).
|
||||
string state = 3;
|
||||
string phys_state = 4;
|
||||
int32 rate_gbps = 5;
|
||||
// link_layer is InfiniBand or Ethernet (RoCE).
|
||||
string link_layer = 6;
|
||||
// counters are cumulative, by sysfs name; data counters in bytes.
|
||||
map<string, uint64> counters = 7;
|
||||
|
||||
// Identity: the adapter's node GUID and the port's own, as the fabric
|
||||
// knows them.
|
||||
string node_guid = 8;
|
||||
string port_guid = 9;
|
||||
// The subnet manager this port sees as master, on InfiniBand. Every port on
|
||||
// one physical fabric sees the same one, which is what tells separate
|
||||
// fabrics apart when their names are only what somebody typed.
|
||||
string sm_guid = 10;
|
||||
int32 sm_lid = 11;
|
||||
int32 sm_priority = 12;
|
||||
string sm_state = 13;
|
||||
// sm_error says why sm_guid is empty on an active InfiniBand port:
|
||||
// infiniband-diags missing, or no subnet manager answering.
|
||||
string sm_error = 14;
|
||||
// network is, on RoCE, the IP network of the RoCE v2 address: RoCE has no
|
||||
// subnet manager, and the network is what its fabric is.
|
||||
string network = 15;
|
||||
string gid_prefix = 16;
|
||||
// subnet_managers is every subnet manager the fabric's subnet
|
||||
// administrator lists — master and standbys — asked through one port per
|
||||
// fabric, so it is on that port only. It is what keeps a fabric one fabric
|
||||
// through a failover: the master changes, the list does not.
|
||||
repeated SubnetManager subnet_managers = 17;
|
||||
string subnet_managers_error = 18;
|
||||
}
|
||||
|
||||
// SubnetManager is one subnet manager on a fabric.
|
||||
message SubnetManager {
|
||||
string guid = 1;
|
||||
int32 lid = 2;
|
||||
int32 priority = 3;
|
||||
// state is master, standby, discovering or not active.
|
||||
string state = 4;
|
||||
}
|
||||
|
||||
// IMEX is the nvidia-imex service and its channels.
|
||||
message IMEX {
|
||||
// state is systemd's word for it, empty where it is not installed.
|
||||
string state = 1;
|
||||
int32 channels = 2;
|
||||
}
|
||||
|
||||
// GPUFabricState is one GPU's NVLink domain membership.
|
||||
message GPUFabricState {
|
||||
int32 index = 1;
|
||||
string state = 2;
|
||||
string status = 3;
|
||||
string cluster_uuid = 4;
|
||||
string clique_id = 5;
|
||||
string bandwidth = 6;
|
||||
}
|
||||
|
||||
message Heartbeat {
|
||||
string node_id = 1;
|
||||
// gpu_memory_used_mb, indexed by GPU, catches a foreign process appearing on
|
||||
// a card after registration.
|
||||
repeated int32 gpu_memory_used_mb = 2;
|
||||
repeated string new_problems = 3;
|
||||
|
||||
// storage for the same reason the GPU memory is here: it changes while the
|
||||
// agent runs, and a figure sampled once at registration is a figure that is
|
||||
// wrong on exactly the node worth watching — the one busy filling up.
|
||||
//
|
||||
// Always sent, because it is a few numbers per filesystem.
|
||||
repeated Storage storage = 4;
|
||||
|
||||
// weights only when they changed, because the list is long enough that
|
||||
// sending it every few seconds would be paying for nothing most of the time.
|
||||
//
|
||||
// weights_changed is what says the list is meant: without it an empty list
|
||||
// would be ambiguous between "this node now holds nothing" and "nothing to
|
||||
// report", and the first would quietly empty a node's inventory.
|
||||
bool weights_changed = 5;
|
||||
repeated LocalWeights weights = 6;
|
||||
|
||||
// memory_total_mb for the same reason again: RAM is added to a virtual
|
||||
// machine while it runs, or across a reboot the agent does not notice, and a
|
||||
// figure read once at registration kept a node that had grown to 32 GB
|
||||
// looking like the 7 GB machine it used to be. 0 means not read.
|
||||
int64 memory_total_mb = 7;
|
||||
|
||||
// engine_versions only when they changed, like weights. An engine installed
|
||||
// after the agent started is otherwise invisible to placement until
|
||||
// somebody restarts the agent — which is a step nobody knows to take.
|
||||
bool engines_changed = 8;
|
||||
map<string, string> engine_versions = 9;
|
||||
|
||||
// provision is where installing the node's software stands, sent when it
|
||||
// changed. The flag makes an empty list mean "nothing to install" rather
|
||||
// than "nothing to say", as for weights.
|
||||
bool provision_changed = 10;
|
||||
repeated ProvisionStatus provision = 11;
|
||||
|
||||
// load is how busy the machine is, sampled over the interval since the
|
||||
// previous heartbeat. Absent on the first one, which has nothing to take a
|
||||
// difference from, and from an agent on a system without /proc.
|
||||
NodeLoad load = 12;
|
||||
|
||||
// ib_ports is every InfiniBand or RoCE port as it stands now, re-read on
|
||||
// each heartbeat: a port that goes down after the agent started used to go
|
||||
// unnoticed until the agent restarted.
|
||||
repeated IBPort ib_ports = 13;
|
||||
// fabrics is the fabrics this node can now carry traffic on, sent only when
|
||||
// that changed since registration or the last heartbeat that said so — a
|
||||
// port going down takes its fabric away, and one coming back returns it.
|
||||
bool fabrics_changed = 14;
|
||||
repeated FabricAttachment fabrics = 15;
|
||||
|
||||
// gpu_fabric is each GPU's NVLink domain membership as it stands now, sent
|
||||
// once a minute: a fabric manager restarting drops every card out of its
|
||||
// domain, and the node has to stop being offered work that spans it.
|
||||
repeated GPUFabricState gpu_fabric = 16;
|
||||
IMEX imex = 17;
|
||||
// cpu_count as in Register, for a VM given more CPUs while it runs.
|
||||
int32 cpu_count = 18;
|
||||
}
|
||||
|
||||
// Provision is the software a node should have: which engines, at which
|
||||
// versions. Nothing more.
|
||||
//
|
||||
// Everything else is the node's to work out, because only the node knows it:
|
||||
// which build fits its hardware, where that build is published, what its
|
||||
// operating system is missing, and whether the result actually runs. The agent
|
||||
// converges on this — installs what is absent, checks what is there — and
|
||||
// reports each engine through ProvisionStatus. Nothing is taken on trust from
|
||||
// what happens to be on the machine: an engine is offered to placement only
|
||||
// once the agent has installed it and seen it run.
|
||||
message Provision {
|
||||
repeated EngineWant engines = 1;
|
||||
}
|
||||
|
||||
// EngineWant names one engine at one version.
|
||||
message EngineWant {
|
||||
string engine = 1; // "llamacpp"
|
||||
string version = 2; // "b10991"
|
||||
}
|
||||
|
||||
// ProvisionStatus is where one engine's installation stands on a node.
|
||||
message ProvisionStatus {
|
||||
string engine = 1;
|
||||
string version = 2;
|
||||
string variant = 3;
|
||||
// "installing", "ready", "failed", or "unsupported" when no build fits.
|
||||
string state = 4;
|
||||
// detail says what is happening or what went wrong, in words for a person.
|
||||
string detail = 5;
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
message ShardStatus {
|
||||
string node_id = 1;
|
||||
string replica_id = 2;
|
||||
string shard_id = 3;
|
||||
State state = 4;
|
||||
string error = 5;
|
||||
int32 exit_code = 6;
|
||||
// last_output carries the tail of the engine log on a failure. When a model
|
||||
// fails to load this is the only diagnosis available.
|
||||
repeated string last_output = 7;
|
||||
|
||||
// fetch is how far a download has got, when the shard is waiting on one.
|
||||
//
|
||||
// The agent logs this already, on the node, where nobody is looking. From
|
||||
// the control plane a replica pulling a hundred gigabytes and one that is
|
||||
// wedged are the same thing — a shard that has been starting for ten
|
||||
// minutes — and telling them apart is the difference between waiting and
|
||||
// going to look.
|
||||
FetchProgress fetch = 8;
|
||||
|
||||
enum State {
|
||||
STATE_UNSPECIFIED = 0;
|
||||
STATE_PENDING = 1;
|
||||
STATE_STARTING = 2;
|
||||
STATE_LOADING = 3;
|
||||
STATE_READY = 4;
|
||||
STATE_DRAINING = 5;
|
||||
STATE_STOPPED = 6;
|
||||
STATE_FAILED = 7;
|
||||
}
|
||||
}
|
||||
|
||||
// FetchProgress is a download in flight.
|
||||
message FetchProgress {
|
||||
string source = 1;
|
||||
int64 bytes_done = 2;
|
||||
// bytes_total is 0 when the server did not say how large the file is, which
|
||||
// is a real answer: a percentage would then be invented, and an invented
|
||||
// percentage that stops moving is worse than a byte count that keeps rising.
|
||||
int64 bytes_total = 3;
|
||||
}
|
||||
|
||||
// StartShard launches one shard. It is idempotent by shard_id: the control
|
||||
// plane retries, and a repeat must not produce a second process.
|
||||
message StartShard {
|
||||
string replica_id = 1;
|
||||
string shard_id = 2;
|
||||
|
||||
int32 rank = 3;
|
||||
int32 nnodes = 4;
|
||||
string leader_addr = 5;
|
||||
int32 leader_port = 6;
|
||||
|
||||
repeated int32 gpu_indices = 7;
|
||||
// hcas are the adapters of the chosen fabric and nothing else. Handing a
|
||||
// shard adapters from a second fabric lets NCCL pick one whose peer is
|
||||
// unreachable, and the job hangs instead of failing (plan §5.3).
|
||||
repeated string hcas = 8;
|
||||
string mgmt_interface = 9;
|
||||
int32 http_port = 10;
|
||||
|
||||
EngineProfile profile = 11;
|
||||
int64 startup_timeout_ms = 12;
|
||||
|
||||
// pool and fabric are control plane concepts the agent cannot derive. They
|
||||
// are sent so the agent can echo them back in RunningShard after a restart.
|
||||
string pool = 13;
|
||||
string fabric = 14;
|
||||
|
||||
// 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;
|
||||
// cpu_threads is how many CPU threads a shard with no GPU computes on:
|
||||
// the cores placement reserved for it. Zero leaves it to the agent.
|
||||
int32 cpu_threads = 17;
|
||||
}
|
||||
|
||||
message EngineProfile {
|
||||
string name = 1;
|
||||
string engine = 2;
|
||||
string gpu_class = 3;
|
||||
int32 tp = 4;
|
||||
int32 pp = 5;
|
||||
int32 gpus_per_node = 6;
|
||||
int32 nodes = 7;
|
||||
repeated string args = 8;
|
||||
map<string, string> env = 9;
|
||||
string model_path = 10;
|
||||
string served_model_name = 11;
|
||||
// adapters are LoRA adapters loaded beside the model, in order: the index
|
||||
// is the id a request applies one by.
|
||||
repeated LoraAdapter adapters = 12;
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
//
|
||||
// The paths are computed by the control plane from the catalog, but the agent
|
||||
// checks them again before unlinking anything: it deletes only files inside its
|
||||
// own cache directory whose names this system generated. A control plane that
|
||||
// is wrong, compromised or simply newer than the agent must not be able to
|
||||
// talk it into deleting a model an operator spent a week copying there.
|
||||
message RemoveWeights {
|
||||
repeated string paths = 1;
|
||||
}
|
||||
|
||||
// StopShard terminates a shard and frees its GPUs.
|
||||
message StopShard {
|
||||
string shard_id = 1;
|
||||
}
|
||||
|
||||
// DrainShard stops routing to a shard while letting in-flight work finish.
|
||||
message DrainShard {
|
||||
string shard_id = 1;
|
||||
}
|
||||
|
||||
// Ack confirms a registration or reports that a command could not be accepted.
|
||||
message Ack {
|
||||
string message_id = 1;
|
||||
bool ok = 2;
|
||||
string error = 3;
|
||||
}
|
||||
|
||||
// LogsRequest asks a node for the recent lines of one of its logs.
|
||||
//
|
||||
// Over the stream the agent already holds, because that is the one way into a
|
||||
// node every deployment has: the admin page otherwise means ssh through a jump
|
||||
// host to a machine whose address nobody remembers, to read a journal only
|
||||
// root may open.
|
||||
message LogsRequest {
|
||||
// request_id pairs the reply with its request; several may be in flight.
|
||||
string request_id = 1;
|
||||
// source is "agent" for the agent's own log, or a shard id for that
|
||||
// engine's output.
|
||||
string source = 2;
|
||||
int32 lines = 3;
|
||||
}
|
||||
|
||||
// LogsReply answers a LogsRequest.
|
||||
message LogsReply {
|
||||
string request_id = 1;
|
||||
repeated string lines = 2;
|
||||
// sources are every log this node can show, so a page can offer them
|
||||
// without knowing what runs there.
|
||||
repeated string sources = 3;
|
||||
// error says why the lines are absent: an unknown source, say.
|
||||
string error = 4;
|
||||
}
|
||||
|
||||
// FabricScanRequest asks a node to map the fabric one of its ports is on,
|
||||
// with ibnetdiscover: every switch, every adapter and every link between
|
||||
// them. One node per fabric is asked, every few minutes; a fabric looks the
|
||||
// same from all of its ports.
|
||||
message FabricScanRequest {
|
||||
string request_id = 1;
|
||||
string hca = 2;
|
||||
int32 port = 3;
|
||||
}
|
||||
|
||||
// FabricScanReply is ibnetdiscover's output as it printed it. Parsed by the
|
||||
// control plane, so a format the parser does not yet know is a control plane
|
||||
// upgrade rather than one on every node.
|
||||
message FabricScanReply {
|
||||
string request_id = 1;
|
||||
string output = 2;
|
||||
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;
|
||||
}
|
||||
|
||||
// UpgradeAgent tells an agent to replace its own binary with a release and
|
||||
// restart.
|
||||
//
|
||||
// The agent fetches the release for its own platform and checks it against
|
||||
// the release's checksums before replacing anything, keeps the old binary
|
||||
// beside the new, then exits so systemd starts the new one. Engines keep
|
||||
// running throughout, as for any agent restart, and are adopted on the way
|
||||
// back. Over the stream the agent already holds, so upgrading a fleet needs
|
||||
// no ssh into it.
|
||||
message UpgradeAgent {
|
||||
string request_id = 1;
|
||||
// version must be a release: v0.10.4 or v0.10.4-rc.2.
|
||||
string version = 2;
|
||||
// releases is the base URL release files are downloaded from.
|
||||
string releases = 3;
|
||||
}
|
||||
|
||||
// UpgradeReply says whether the new binary is in place. Sent before the
|
||||
// agent exits; an empty error means it will be back on the new version.
|
||||
message UpgradeReply {
|
||||
string request_id = 1;
|
||||
string version = 2;
|
||||
string error = 3;
|
||||
}
|
||||
|
||||
// LoraAdapter is one adapter: the name clients ask for it by, and its source,
|
||||
// which the agent fetches as it does the model's weights.
|
||||
message LoraAdapter {
|
||||
string name = 1;
|
||||
string source = 2;
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
version: v2
|
||||
managed:
|
||||
enabled: true
|
||||
plugins:
|
||||
- local: protoc-gen-go
|
||||
out: api/gen
|
||||
opt: paths=source_relative
|
||||
- local: protoc-gen-connect-go
|
||||
out: api/gen
|
||||
opt: paths=source_relative
|
||||
inputs:
|
||||
- directory: api/proto
|
||||
@@ -0,0 +1,16 @@
|
||||
version: v2
|
||||
modules:
|
||||
- path: api/proto
|
||||
lint:
|
||||
use:
|
||||
- STANDARD
|
||||
except:
|
||||
# The standard names assume request/response pairs. This is one long-lived
|
||||
# bidirectional stream carrying several kinds of message in each direction,
|
||||
# so "AgentMessage" and "ControlMessage" say what they are and
|
||||
# "ConnectRequest"/"ConnectResponse" would not.
|
||||
- RPC_REQUEST_STANDARD_NAME
|
||||
- RPC_RESPONSE_STANDARD_NAME
|
||||
breaking:
|
||||
use:
|
||||
- FILE
|
||||
Executable
+112
@@ -0,0 +1,112 @@
|
||||
#!/bin/sh
|
||||
# add-node.sh — build a node and put it in the fleet, in one command.
|
||||
#
|
||||
# provision-node.sh and onboard-node.sh stay separate because the two halves
|
||||
# need different access: one the hypervisor, one the certificate authority.
|
||||
# That is a real boundary and this script does not erase it — it drives both
|
||||
# where one person holds both, which is the common case and was the case here.
|
||||
#
|
||||
# The reason to have it is not keystrokes. The VM name, the hostname, the node
|
||||
# id and the certificate's common name must all be the same string, and the
|
||||
# address appears in three places; split across two invocations run minutes
|
||||
# apart, that agreement is maintained by whoever is typing. Here it is stated
|
||||
# once.
|
||||
set -eu
|
||||
|
||||
usage() {
|
||||
cat >&2 <<'USAGE'
|
||||
usage: add-node.sh --name NODE --ip ADDR/CIDR --server URL [options]
|
||||
|
||||
--name NODE node id, hostname and VM name; also the certificate's CN
|
||||
--ip ADDR/CIDR e.g. 10.0.0.145/24
|
||||
--server URL control plane, e.g. https://10.0.0.140:7443
|
||||
|
||||
--vmid N Proxmox VMID (default: the last octet of --ip)
|
||||
--cores N (default 4)
|
||||
--cpu TYPE QEMU CPU model (default host; see provision-node.sh)
|
||||
--memory MB (default 8192)
|
||||
--disk GB (default 20)
|
||||
--fabric NAME InfiniBand fabric; repeatable as NAME:hca1,hca2.
|
||||
Omit for a node with no fabric.
|
||||
--llamacpp-bin P path to llama-server on the node
|
||||
--vllm-bin P path to vllm
|
||||
--sglang-python P path to the python that has sglang
|
||||
--binary PATH llmgw binary to install (default: the one on PATH)
|
||||
--pki DIR certificate authority (default /etc/llmgw/pki)
|
||||
--ssh-user U (default ilker)
|
||||
--provision-only build the VM and stop
|
||||
--onboard-only skip provisioning; the machine already exists
|
||||
USAGE
|
||||
exit 2
|
||||
}
|
||||
|
||||
NAME= IP= SERVER= VMID= CORES=4 MEMORY=8192 DISK=20 CPU=
|
||||
BINARY= PKI=/etc/llmgw/pki SSH_USER=ilker
|
||||
PROVISION=1 ONBOARD=1
|
||||
set -- "$@"
|
||||
PASS_THROUGH=
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--name) NAME="$2"; shift 2;;
|
||||
--ip) IP="$2"; shift 2;;
|
||||
--server) SERVER="$2"; shift 2;;
|
||||
--vmid) VMID="$2"; shift 2;;
|
||||
--cores) CORES="$2"; shift 2;;
|
||||
--cpu) CPU="$2"; shift 2;;
|
||||
--memory) MEMORY="$2"; shift 2;;
|
||||
--disk) DISK="$2"; shift 2;;
|
||||
--binary) BINARY="$2"; shift 2;;
|
||||
--pki) PKI="$2"; shift 2;;
|
||||
--ssh-user) SSH_USER="$2"; shift 2;;
|
||||
--provision-only) ONBOARD=0; shift;;
|
||||
--onboard-only) PROVISION=0; shift;;
|
||||
# Everything the onboarding step understands and this one does not need to.
|
||||
--fabric|--llamacpp-bin|--vllm-bin|--sglang-python|--ib-root|--nvidia-smi|--mgmt-addr|--mgmt-if)
|
||||
PASS_THROUGH="$PASS_THROUGH $1 $2"; shift 2;;
|
||||
-h|--help) usage;;
|
||||
*) echo "unknown option: $1" >&2; usage;;
|
||||
esac
|
||||
done
|
||||
[ -n "$NAME" ] && [ -n "$IP" ] && [ -n "$SERVER" ] || usage
|
||||
|
||||
# The address without its prefix, which is what onboarding and ssh want.
|
||||
ADDR="${IP%%/*}"
|
||||
# The last octet as the VMID by default: it keeps `qm list` and `ctl nodes` in
|
||||
# the same order, and one fewer number to choose is one fewer to get wrong.
|
||||
[ -n "$VMID" ] || VMID="${ADDR##*.}"
|
||||
|
||||
here() { dirname "$0"; }
|
||||
PROVISION_SH="$(here)/provision-node.sh"
|
||||
ONBOARD_SH="$(here)/onboard-node.sh"
|
||||
[ -x "$PROVISION_SH" ] || PROVISION_SH="$(command -v llmgw-provision-node || true)"
|
||||
[ -x "$ONBOARD_SH" ] || ONBOARD_SH="$(command -v llmgw-onboard-node || true)"
|
||||
|
||||
if [ "$PROVISION" = 1 ]; then
|
||||
[ -x "$PROVISION_SH" ] || { echo "cannot find provision-node.sh" >&2; exit 1; }
|
||||
echo "== 1/2 building $NAME as VM $VMID"
|
||||
# shellcheck disable=SC2086
|
||||
"$PROVISION_SH" --vmid "$VMID" --name "$NAME" --ip "$IP" \
|
||||
--cores "$CORES" --memory "$MEMORY" --disk "$DISK" ${CPU:+--cpu "$CPU"}
|
||||
fi
|
||||
|
||||
if [ "$ONBOARD" = 0 ]; then
|
||||
echo "$NAME built. Onboard it from the CA host when you are ready."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
[ -x "$ONBOARD_SH" ] || { echo "cannot find onboard-node.sh" >&2; exit 1; }
|
||||
echo
|
||||
echo "== 2/2 onboarding $NAME"
|
||||
# Onboarding waits for ssh itself, which is what makes running the two halves
|
||||
# back to back safe: cloud-init is still regenerating host keys and restarting
|
||||
# sshd for the first minute of a new machine's life.
|
||||
# shellcheck disable=SC2086
|
||||
"$ONBOARD_SH" --node "$NAME" --addr "$ADDR" --server "$SERVER" \
|
||||
--pki "$PKI" --ssh-user "$SSH_USER" \
|
||||
${BINARY:+--binary "$BINARY"} $PASS_THROUGH
|
||||
|
||||
echo
|
||||
echo "Check it with: llmgw ctl nodes"
|
||||
echo
|
||||
echo "The agent installs its engines itself; the node takes no model until they"
|
||||
echo "are installed and seen to run. 'llmgw ctl nodes' shows where that stands."
|
||||
Executable
+297
@@ -0,0 +1,297 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Give a machine an identity in an llmgw cluster.
|
||||
#
|
||||
# Run on the host that holds the certificate authority, because that is the one
|
||||
# thing that cannot be delegated: this script issues the credential by which the
|
||||
# control plane will believe the node is who it says it is.
|
||||
#
|
||||
# What it does, in the order it matters:
|
||||
# 1. issue a client certificate whose common name is the node id
|
||||
# 2. copy the CA root, the certificate and its key to the machine
|
||||
# 3. trust the root system-wide, so other tooling there stops complaining
|
||||
# 4. install the binary, the configuration and the unit
|
||||
# 5. start the agent
|
||||
#
|
||||
# The node id, the VM name, the hostname and the certificate common name are all
|
||||
# the same string. They do not have to be, and every time they are not, someone
|
||||
# spends an afternoon working out which of the four a message is talking about.
|
||||
set -euo pipefail
|
||||
|
||||
usage() {
|
||||
cat >&2 <<USAGE
|
||||
usage: $0 --node ID --addr HOST --server URL [options]
|
||||
|
||||
--node ID node id, hostname and certificate common name
|
||||
--addr HOST how to reach the machine over ssh
|
||||
--server URL the control plane's agent endpoint, e.g. https://10.0.0.140:7443;
|
||||
several, comma separated, for a node that fails over
|
||||
--fabric NAME InfiniBand fabric this node is attached to; repeatable as
|
||||
--fabric NAME:hca1,hca2 . Omit for a node with no fabric,
|
||||
which can host single-node replicas and will never be given
|
||||
a multi-node one.
|
||||
--mgmt-addr ADDR address engines are reachable on (default: the node's
|
||||
address on the route to --server, found by the agent)
|
||||
--mgmt-if NAME interface NCCL bootstraps over (default: the one holding
|
||||
that address)
|
||||
--llamacpp-bin P use this llama-server instead of the one the agent installs
|
||||
--vllm-bin P use this vllm instead of the one the agent installs on a
|
||||
GPU node (also set engines.vllm: "" in server.yaml)
|
||||
--sglang-python P python interpreter of an SGLang virtualenv you installed
|
||||
--nvidia-smi PATH where nvidia-smi lives, when it is not on PATH. Containerized
|
||||
driver installs routinely put it somewhere else.
|
||||
--ib-root DIR the sysfs InfiniBand directory (default /sys/class/infiniband)
|
||||
--pki DIR where the authority lives (default /etc/llmgw/pki)
|
||||
--binary PATH llmgw binary to install (default: the one running this).
|
||||
Releases carry static linux/amd64 and linux/arm64 builds;
|
||||
downloading one is better than copying whatever happens to
|
||||
be on the machine you are running this from.
|
||||
--releases URL where release binaries are downloaded from (default the
|
||||
project's Gitea releases)
|
||||
--ssh-user NAME (default ilker)
|
||||
USAGE
|
||||
exit 2
|
||||
}
|
||||
|
||||
NODE= ADDR= SERVER= MGMT_OVERRIDE= MGMT_IF_OVERRIDE=
|
||||
NVIDIA_SMI= IB_ROOT=/sys/class/infiniband
|
||||
LLAMACPP_BIN= VLLM_BIN= SGLANG_PYTHON=
|
||||
SSH_WAIT=180
|
||||
PKI=/etc/llmgw/pki BINARY= SSH_USER=ilker
|
||||
RELEASES=https://gitea.mynodes.xyz/ilker/llmgw/releases/download
|
||||
FABRIC_NAME= FABRIC_HCAS="mlx5_0,mlx5_1"
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--node) NODE="$2"; shift 2;;
|
||||
--addr) ADDR="$2"; shift 2;;
|
||||
--server) SERVER="$2"; shift 2;;
|
||||
--fabric) FABRIC_NAME="${2%%:*}"; [ "${2#*:}" != "$2" ] && FABRIC_HCAS="${2#*:}"; shift 2;;
|
||||
--mgmt-addr) MGMT_OVERRIDE="$2"; shift 2;;
|
||||
--mgmt-if) MGMT_IF_OVERRIDE="$2"; shift 2;;
|
||||
--llamacpp-bin) LLAMACPP_BIN="$2"; shift 2;;
|
||||
--vllm-bin) VLLM_BIN="$2"; shift 2;;
|
||||
--sglang-python) SGLANG_PYTHON="$2"; shift 2;;
|
||||
--nvidia-smi) NVIDIA_SMI="$2"; shift 2;;
|
||||
--ib-root) IB_ROOT="$2"; shift 2;;
|
||||
--pki) PKI="$2"; shift 2;;
|
||||
--binary) BINARY="$2"; shift 2;;
|
||||
--releases) RELEASES="$2"; shift 2;;
|
||||
--ssh-user) SSH_USER="$2"; shift 2;;
|
||||
*) usage;;
|
||||
esac
|
||||
done
|
||||
[ -n "$NODE" ] && [ -n "$ADDR" ] && [ -n "$SERVER" ] || usage
|
||||
BINARY="${BINARY:-$(command -v llmgw)}"
|
||||
[ -x "$BINARY" ] || { echo "no llmgw binary; pass --binary" >&2; exit 1; }
|
||||
[ -f "$PKI/ca-key.pem" ] || { echo "no authority in $PKI; run 'llmgw ctl cert init-ca' first" >&2; exit 1; }
|
||||
|
||||
# accept-new, not "no": a machine provisioned a minute ago has a host key
|
||||
# nobody has seen, and refusing it stopped onboarding dead with "Host key
|
||||
# verification failed" — an error about ssh, in the middle of a script whose
|
||||
# whole job is to talk to a new machine. accept-new learns an unknown key and
|
||||
# still refuses a changed one, which is the distinction that matters: the risk
|
||||
# being guarded against is a substituted machine, not a new one.
|
||||
SSHOPTS="-o BatchMode=yes -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new"
|
||||
SSH="ssh $SSHOPTS $SSH_USER@$ADDR"
|
||||
say() { printf ' %s\n' "$*"; }
|
||||
|
||||
echo "onboarding $NODE at $ADDR"
|
||||
|
||||
# Waiting rather than failing: cloud-init restarts sshd after it regenerates
|
||||
# the host keys, so a machine that answered a moment ago refuses connections
|
||||
# for a few seconds. Observed while onboarding node05 — the agent was already
|
||||
# registered and healthy while ssh was refusing.
|
||||
wait_for_ssh() {
|
||||
deadline=$(( $(date +%s) + SSH_WAIT ))
|
||||
until $SSH true 2>/dev/null; do
|
||||
[ "$(date +%s)" -lt "$deadline" ] || {
|
||||
echo "cannot reach $SSH_USER@$ADDR over ssh after ${SSH_WAIT}s" >&2
|
||||
exit 1
|
||||
}
|
||||
sleep 5
|
||||
done
|
||||
}
|
||||
wait_for_ssh
|
||||
|
||||
# 1. The certificate. Issued fresh every time: re-onboarding a machine should
|
||||
# not depend on finding an old key, and the old one keeps working until it
|
||||
# is revoked, which is a separate decision.
|
||||
say "issuing a certificate"
|
||||
TMP="$(mktemp -d)"
|
||||
trap 'rm -rf "$TMP"' EXIT
|
||||
"$BINARY" ctl cert issue-client "$NODE" --dir "$PKI" --out "$TMP" >/dev/null
|
||||
|
||||
# 2 and 3. Identity. The root is trusted by the agent's configuration only,
|
||||
# never system-wide: it signs whatever the control plane asks, and every
|
||||
# TLS client on the machine would take it for any site.
|
||||
say "installing the certificate"
|
||||
scp -q $SSHOPTS "$PKI/ca.pem" "$TMP/$NODE.pem" "$TMP/$NODE-key.pem" "$SSH_USER@$ADDR:/tmp/"
|
||||
$SSH "set -e
|
||||
sudo install -d -m 0755 /etc/llmgw
|
||||
sudo install -m 0644 /tmp/ca.pem /etc/llmgw/ca.pem
|
||||
sudo install -m 0644 /tmp/$NODE.pem /etc/llmgw/agent.pem
|
||||
sudo install -m 0600 /tmp/$NODE-key.pem /etc/llmgw/agent-key.pem
|
||||
sudo chown root:root /etc/llmgw/*
|
||||
if [ -f /usr/local/share/ca-certificates/llmgw-ca.crt ]; then
|
||||
sudo rm -f /usr/local/share/ca-certificates/llmgw-ca.crt
|
||||
sudo update-ca-certificates --fresh >/dev/null 2>&1 || true
|
||||
fi
|
||||
rm -f /tmp/ca.pem /tmp/$NODE.pem /tmp/$NODE-key.pem"
|
||||
|
||||
# 4. The binary, the configuration and the unit.
|
||||
# The node fetches the release CI built for this version, for its own
|
||||
# architecture, and checks it against the release's checksums — rather than
|
||||
# being handed whatever binary is on the machine running this script. A build
|
||||
# that is not a release (a developer's, with commits past the tag) has no such
|
||||
# file, and is copied as before.
|
||||
TAG="$("$BINARY" --version 2>/dev/null | awk '{print $NF}')"
|
||||
if printf '%s' "$TAG" | grep -Eq '^v[0-9]+\.[0-9]+\.[0-9]+(-rc\.[0-9]+)?$'; then
|
||||
say "installing llmgw $TAG from the release"
|
||||
$SSH "set -e
|
||||
arch=\$(dpkg --print-architecture)
|
||||
name=llmgw-$TAG-linux-\$arch
|
||||
cd /tmp && rm -f \$name checksums.txt
|
||||
curl -fsSL -o \$name '$RELEASES/$TAG/'\$name
|
||||
curl -fsSL -o checksums.txt '$RELEASES/$TAG/checksums.txt'
|
||||
grep \" \$name\\\$\" checksums.txt | sha256sum -c - >/dev/null
|
||||
sudo install -m 0755 \$name /usr/local/bin/llmgw
|
||||
rm -f \$name checksums.txt"
|
||||
else
|
||||
say "installing the binary (${TAG:-unversioned}, not a release)"
|
||||
scp -q $SSHOPTS "$BINARY" "$SSH_USER@$ADDR:/tmp/llmgw"
|
||||
$SSH 'sudo install -m 0755 /tmp/llmgw /usr/local/bin/llmgw && rm -f /tmp/llmgw'
|
||||
fi
|
||||
|
||||
say "writing the configuration"
|
||||
NODE_IS_HOSTNAME=no
|
||||
[ "$($SSH hostname 2>/dev/null)" = "$NODE" ] && NODE_IS_HOSTNAME=yes
|
||||
# Only what the node cannot find out for itself. Its name is its hostname, its
|
||||
# address is the one on the route to the control plane, llama-server is looked
|
||||
# for on PATH and where release archives are unpacked, GPUs are there when
|
||||
# nvidia-smi is, and memory, disk and engines are re-read while it runs. Every
|
||||
# setting written here is one that has to be kept true by hand — so an option
|
||||
# below is written only when it was given, as an override.
|
||||
{
|
||||
echo "# $NODE, written by onboard-node.sh."
|
||||
# Several addresses, comma separated, are every control plane the node may
|
||||
# fail over to; a standby refuses it and it moves on to the next.
|
||||
if [ "${SERVER#*,}" != "$SERVER" ]; then
|
||||
echo "servers:"
|
||||
printf '%s\n' "$SERVER" | tr ',' '\n' | sed 's/^ *//; s/^/ - "/; s/$/"/'
|
||||
else
|
||||
echo "server: \"$SERVER\""
|
||||
fi
|
||||
echo
|
||||
echo "tls:"
|
||||
echo " ca: /etc/llmgw/ca.pem"
|
||||
echo " cert: /etc/llmgw/agent.pem"
|
||||
echo " key: /etc/llmgw/agent-key.pem"
|
||||
# The certificate's common name is --node, and the control plane refuses a
|
||||
# stream whose node id differs from it. Said explicitly only when the
|
||||
# hostname would say something else.
|
||||
[ "$NODE_IS_HOSTNAME" = yes ] || echo "node_id: $NODE"
|
||||
[ -n "$MGMT_OVERRIDE" ] && echo "mgmt_addr: \"$MGMT_OVERRIDE\""
|
||||
[ -n "$MGMT_IF_OVERRIDE" ] && echo "mgmt_interface: $MGMT_IF_OVERRIDE"
|
||||
if [ -n "$FABRIC_NAME" ]; then
|
||||
# Fabric membership is the one thing that cannot be found: OpenSM's default
|
||||
# subnet prefix is the same on every fabric, so a matching prefix proves
|
||||
# nothing and a wrong guess hangs the first all-reduce.
|
||||
echo "fabrics:"
|
||||
echo " - name: $FABRIC_NAME"
|
||||
echo " hcas: [$(echo "$FABRIC_HCAS" | sed 's/,/, /g')]"
|
||||
fi
|
||||
if [ -n "$NVIDIA_SMI" ] || [ "$IB_ROOT" != /sys/class/infiniband ]; then
|
||||
echo "discovery:"
|
||||
[ -n "$NVIDIA_SMI" ] && echo " nvidia_smi: \"$NVIDIA_SMI\""
|
||||
[ "$IB_ROOT" != /sys/class/infiniband ] && echo " ib_root: $IB_ROOT"
|
||||
fi
|
||||
if [ -n "$LLAMACPP_BIN$VLLM_BIN$SGLANG_PYTHON" ]; then
|
||||
echo "engines:"
|
||||
[ -n "$LLAMACPP_BIN" ] && printf ' llamacpp:\n bin: "%s"\n' "$LLAMACPP_BIN"
|
||||
[ -n "$VLLM_BIN" ] && printf ' vllm:\n bin: "%s"\n' "$VLLM_BIN"
|
||||
[ -n "$SGLANG_PYTHON" ] && printf ' sglang:\n python: "%s"\n' "$SGLANG_PYTHON"
|
||||
fi
|
||||
true
|
||||
} > "$TMP/agent.yaml"
|
||||
scp -q "$TMP/agent.yaml" "$SSH_USER@$ADDR:/tmp/agent.yaml"
|
||||
|
||||
cat > "$TMP/llmgw-agent.service" <<'UNIT'
|
||||
[Unit]
|
||||
Description=llmgw node agent
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
ExecStart=/usr/local/bin/llmgw agent --config /etc/llmgw/agent.yaml
|
||||
# Root, because the agent launches engines and needs device access. Engines
|
||||
# outlive a restart of this unit and are adopted again on the next start, so
|
||||
# KillMode must not take the whole process group with it.
|
||||
KillMode=process
|
||||
Restart=on-failure
|
||||
RestartSec=2s
|
||||
# RDMA pins the buffers it registers, and NCCL registers a great deal: with
|
||||
# systemd's default 8 MB lock limit, an engine that drops the capability that
|
||||
# lets root ignore it fails its first collective over InfiniBand.
|
||||
LimitMEMLOCK=infinity
|
||||
StateDirectory=llmgw-agent
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
UNIT
|
||||
scp -q "$TMP/llmgw-agent.service" "$SSH_USER@$ADDR:/tmp/"
|
||||
|
||||
# An engine path that is wrong, or right but unrunnable, is the failure this
|
||||
# node type is most likely to have, and it is invisible: the agent registers
|
||||
# healthy and every shard placed on it dies with "executable file not found"
|
||||
# or a missing shared library. Checked here because it is cheap here and
|
||||
# expensive later — llama-server on node05 was present, executable, and could
|
||||
# not start for want of libgomp1, and the node reported ready throughout.
|
||||
check_engine() {
|
||||
[ -n "$2" ] || return 0
|
||||
if $SSH "test -x '$2'" 2>/dev/null; then
|
||||
if $SSH "'$2' --version >/dev/null 2>&1 || '$2' --help >/dev/null 2>&1"; then
|
||||
say "$1: $2 runs"
|
||||
return 0
|
||||
fi
|
||||
ENGINE_PROBLEMS="$ENGINE_PROBLEMS
|
||||
$1 at $2 exists but will not run. Usually a missing shared library:
|
||||
ssh $SSH_USER@$ADDR ldd $2 | grep 'not found'"
|
||||
return 0
|
||||
fi
|
||||
ENGINE_PROBLEMS="$ENGINE_PROBLEMS
|
||||
$1 is configured at $2, which is not an executable on $NODE."
|
||||
}
|
||||
ENGINE_PROBLEMS=
|
||||
if [ -n "$LLAMACPP_BIN$VLLM_BIN$SGLANG_PYTHON" ]; then
|
||||
say "checking the engines"
|
||||
check_engine llamacpp "$LLAMACPP_BIN"
|
||||
check_engine vllm "$VLLM_BIN"
|
||||
check_engine sglang "$SGLANG_PYTHON"
|
||||
fi
|
||||
|
||||
say "starting the agent"
|
||||
$SSH 'set -e
|
||||
sudo install -m 0644 -o root -g root /tmp/agent.yaml /etc/llmgw/agent.yaml
|
||||
sudo install -m 0644 /tmp/llmgw-agent.service /etc/systemd/system/llmgw-agent.service
|
||||
rm -f /tmp/agent.yaml /tmp/llmgw-agent.service
|
||||
sudo systemctl daemon-reload
|
||||
sudo systemctl enable --now llmgw-agent >/dev/null
|
||||
sudo systemctl restart llmgw-agent'
|
||||
|
||||
echo
|
||||
if [ -z "$LLAMACPP_BIN$VLLM_BIN$SGLANG_PYTHON" ]; then
|
||||
say "the agent now installs the engines the control plane asks for: the"
|
||||
say "build for this machine, checked against its publisher's checksum and run"
|
||||
say "before the node is offered any model. 'llmgw ctl nodes' shows progress."
|
||||
fi
|
||||
if [ -n "$ENGINE_PROBLEMS" ]; then
|
||||
echo "$NODE is onboarded but an engine will not run:$ENGINE_PROBLEMS"
|
||||
echo
|
||||
echo "The node will register as ready regardless. Fix this before a replica"
|
||||
echo "is placed on it, or the shard fails and the reason is three logs away."
|
||||
echo
|
||||
fi
|
||||
echo "$NODE onboarded. It should appear in 'llmgw ctl nodes' within a few seconds."
|
||||
echo "If it does not, the control plane's log says why — a certificate whose"
|
||||
echo "common name does not match the claimed node id is refused by design."
|
||||
Executable
+96
@@ -0,0 +1,96 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Create a VM on Proxmox for a new llmgw node.
|
||||
#
|
||||
# Run on the Proxmox host. It does the half of onboarding that is about
|
||||
# hardware; onboard-node.sh does the half that is about identity, and the two
|
||||
# are separate because they are done by different people with different access.
|
||||
#
|
||||
# Idempotent only in the sense that it refuses to touch an existing VMID. A node
|
||||
# that already exists is not something a provisioning script should have an
|
||||
# opinion about.
|
||||
set -euo pipefail
|
||||
|
||||
usage() {
|
||||
cat >&2 <<USAGE
|
||||
usage: $0 --vmid N --name NAME --ip ADDR [options]
|
||||
|
||||
--vmid N Proxmox VM id; must be free
|
||||
--name NAME VM name; use the node id, so the machine, the hostname and
|
||||
the certificate common name are the same string
|
||||
--ip ADDR address with prefix, e.g. 10.0.0.144/24
|
||||
--gw ADDR gateway (default 10.0.0.1)
|
||||
--cores N (default 4)
|
||||
--cpu TYPE QEMU CPU model (default host)
|
||||
--memory MB (default 8192)
|
||||
--disk GB (default 20)
|
||||
--bridge NAME (default vmbr0)
|
||||
--storage NAME (default local-lvm)
|
||||
--image PATH cloud image (default Ubuntu 24.04 in the ISO store)
|
||||
--user NAME cloud-init user (default ilker)
|
||||
--sshkeys PATH authorized_keys to install (default /root/.ssh/authorized_keys)
|
||||
USAGE
|
||||
exit 2
|
||||
}
|
||||
|
||||
VMID= NAME= IP= GW=10.0.0.1 CORES=4 MEMORY=8192 DISK=20
|
||||
# host, not Proxmox's kvm64 default. kvm64 hides AVX, AVX2, FMA and even
|
||||
# SSE4.2 from the guest, and an inference node without vector instructions is
|
||||
# not slightly slower, it is unusable: llama.cpp falls back to its scalar
|
||||
# backend and a 1.5B model dropped from 39 tokens a second to 5.8, with prompt
|
||||
# processing 19x worse. It cost an afternoon to find because nothing is broken
|
||||
# — every number is simply small.
|
||||
#
|
||||
# The trade is live migration: a VM pinned to the host CPU cannot move to a
|
||||
# machine with a different one. For a node whose job is arithmetic, that is
|
||||
# the right way round.
|
||||
CPU=host
|
||||
BRIDGE=vmbr0 STORAGE=local-lvm USER=ilker
|
||||
IMAGE=/var/lib/vz/template/iso/noble-server-cloudimg-amd64.img
|
||||
SSHKEYS=/root/.ssh/authorized_keys
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--vmid) VMID="$2"; shift 2;;
|
||||
--name) NAME="$2"; shift 2;;
|
||||
--ip) IP="$2"; shift 2;;
|
||||
--gw) GW="$2"; shift 2;;
|
||||
--cores) CORES="$2"; shift 2;;
|
||||
--cpu) CPU="$2"; shift 2;;
|
||||
--memory) MEMORY="$2"; shift 2;;
|
||||
--disk) DISK="$2"; shift 2;;
|
||||
--bridge) BRIDGE="$2"; shift 2;;
|
||||
--storage) STORAGE="$2"; shift 2;;
|
||||
--image) IMAGE="$2"; shift 2;;
|
||||
--user) USER="$2"; shift 2;;
|
||||
--sshkeys) SSHKEYS="$2"; shift 2;;
|
||||
*) usage;;
|
||||
esac
|
||||
done
|
||||
[ -n "$VMID" ] && [ -n "$NAME" ] && [ -n "$IP" ] || usage
|
||||
|
||||
if qm status "$VMID" >/dev/null 2>&1; then
|
||||
echo "VM $VMID already exists. Pick a free id; this script will not adopt one." >&2
|
||||
exit 1
|
||||
fi
|
||||
[ -f "$IMAGE" ] || { echo "no cloud image at $IMAGE" >&2; exit 1; }
|
||||
|
||||
echo "creating $NAME as VM $VMID at $IP"
|
||||
qm create "$VMID" \
|
||||
--name "$NAME" --cores "$CORES" --memory "$MEMORY" --cpu "$CPU" \
|
||||
--net0 "virtio,bridge=$BRIDGE" --ostype l26 --agent enabled=1
|
||||
|
||||
# Imported rather than installed: a cloud image plus cloud-init is the whole
|
||||
# reason a node takes a minute to build instead of an afternoon.
|
||||
qm importdisk "$VMID" "$IMAGE" "$STORAGE" >/dev/null
|
||||
qm set "$VMID" --scsihw virtio-scsi-pci --scsi0 "$STORAGE:vm-$VMID-disk-0,discard=on,ssd=1" >/dev/null
|
||||
qm disk resize "$VMID" scsi0 "${DISK}G" >/dev/null
|
||||
qm set "$VMID" --ide2 "$STORAGE:cloudinit" --boot order=scsi0 >/dev/null
|
||||
qm set "$VMID" \
|
||||
--ipconfig0 "ip=$IP,gw=$GW" \
|
||||
--nameserver "$GW 8.8.8.8" \
|
||||
--ciuser "$USER" \
|
||||
--sshkeys "$SSHKEYS" >/dev/null
|
||||
|
||||
qm start "$VMID"
|
||||
echo "started. cloud-init takes a minute; then run onboard-node.sh from the CA host."
|
||||
@@ -0,0 +1,113 @@
|
||||
# llmgw agent configuration, one per GPU node.
|
||||
#
|
||||
# The agent dials out to the control plane and keeps one stream open. Engines
|
||||
# keep serving if the agent restarts: the router talks to them directly, and
|
||||
# orphan adoption picks them up again on the next start.
|
||||
|
||||
# Everything a node can find out about itself, it does: its id is its hostname,
|
||||
# its address is the one on the route to the control plane, llama-server is
|
||||
# looked for on PATH and in /opt/llama.cpp, GPUs are there when nvidia-smi is,
|
||||
# weights are fetched into /var/lib/llmgw-weights and looked for there, and
|
||||
# memory, disk and installed engines are re-read while the agent runs. So a
|
||||
# plain CPU node needs only the two blocks below, and the same file on every
|
||||
# node. What follows them are overrides, for when the defaults are wrong.
|
||||
|
||||
server: "http://llmgw-server.internal:7443"
|
||||
|
||||
# The agent's half of mutual TLS. The certificate's common name must be this
|
||||
# node's id — its hostname, unless node_id says otherwise — because the control
|
||||
# plane compares the two and refuses the stream when they disagree:
|
||||
#
|
||||
# llmgw ctl cert issue-client gpu-a-01 --dir /etc/llmgw/pki
|
||||
#
|
||||
# With this set, `server` above must be an https:// address — an agent holding
|
||||
# certificates and dialling http:// is refused at startup rather than connecting
|
||||
# in the clear without noticing.
|
||||
#
|
||||
# tls:
|
||||
# ca: /etc/llmgw/ca.pem
|
||||
# cert: /etc/llmgw/agent.pem
|
||||
# key: /etc/llmgw/agent-key.pem
|
||||
|
||||
# ---- overrides ---------------------------------------------------------------
|
||||
|
||||
# node_id: gpu-a-01 # when it is not the hostname
|
||||
|
||||
# The address engines are reachable on, and the address NCCL bootstraps over.
|
||||
# Never a fabric address: NCCL rendezvouses over TCP before moving to RDMA.
|
||||
# mgmt_addr: "10.0.0.11"
|
||||
# mgmt_interface: eth0
|
||||
|
||||
# Free-form, for a profile's node_selector. Nothing needs them; do not restate
|
||||
# what is measured (memory, GPUs, engines), because a label goes stale.
|
||||
# labels:
|
||||
# rack: r12
|
||||
|
||||
heartbeat_interval: 5s
|
||||
reconnect_base: 1s
|
||||
reconnect_max: 30s
|
||||
|
||||
# Each fabric this node's adapters are on, by name. The name is checked against
|
||||
# what the adapters report — the subnet managers every port on one fabric sees —
|
||||
# so two independent fabrics given one name are placed as two and the cluster
|
||||
# page says so; give each its own. The rate is measured, not read from here.
|
||||
#
|
||||
# A node with no fabrics is its own single-node island: it can host single-node
|
||||
# replicas and will never be given a multi-node one. docs/fabrics.md has more.
|
||||
fabrics:
|
||||
- name: ib-a
|
||||
hcas: [mlx5_0, mlx5_1, mlx5_2, mlx5_3]
|
||||
|
||||
discovery:
|
||||
# Worth setting where the driver lives outside the default paths, which
|
||||
# containerized driver installs often do.
|
||||
nvidia_smi: ""
|
||||
ib_root: /sys/class/infiniband
|
||||
# NVLink topology comes from `nvidia-smi topo -m` on every build; a binary
|
||||
# built with `-tags nvml` reads it from NVML instead, when this is set.
|
||||
prefer_nvml: true
|
||||
|
||||
# Where this node keeps model weights besides /var/lib/llmgw-weights, which is
|
||||
# where fetched weights land and is always looked in. Storage differs at
|
||||
# every installation and llmgw deliberately does not own it: a Lustre mount, an
|
||||
# image exported over the fabric with SRP or NVMe-oF, local NVMe, or a copy
|
||||
# made by hand all look identical from here.
|
||||
#
|
||||
# What it buys is placement preferring a node that already holds the model over
|
||||
# one that must first move tens of gigabytes. That matters most where fabrics
|
||||
# are isolated from one another, because weights in one fabric are not weights
|
||||
# in another and the only path between them is the management network.
|
||||
#
|
||||
# model_dirs:
|
||||
# - /var/lib/llmgw-models
|
||||
# - /mnt/lustre/models
|
||||
|
||||
# The agent installs the engines the control plane names (see `engines:` in
|
||||
# server.yaml) under /opt/llmgw/engines: the build for this machine — CUDA when
|
||||
# an NVIDIA driver is loaded — checked against its publisher's checksum, with
|
||||
# the OS packages it needs, and run before the node takes any model. llama.cpp
|
||||
# is installed this way today. Naming a binary below opts that engine out and
|
||||
# uses yours instead; vLLM and SGLang still have to be named until the agent
|
||||
# installs them too.
|
||||
# engines_dir: /opt/llmgw/engines
|
||||
|
||||
engines:
|
||||
vllm:
|
||||
bin: /opt/venvs/vllm/bin/vllm
|
||||
# How a replica spanning nodes is launched. mp, the default, is what vLLM
|
||||
# documents today and needs nothing but vLLM; ray is for a vLLM too old
|
||||
# for mp's multi-node flags, and needs Ray on every node.
|
||||
launch_mode: mp # "mp" or "ray"
|
||||
# ray_bin: /opt/venvs/vllm/bin/ray
|
||||
sglang:
|
||||
# SGLang is launched as a module, which is its documented entry point.
|
||||
python: /opt/venvs/sglang/bin/python
|
||||
|
||||
supervisor:
|
||||
work_dir: /var/lib/llmgw-agent
|
||||
# Engine output is the only diagnosis when a model fails to load, so keep a
|
||||
# useful tail of it.
|
||||
log_lines: 512
|
||||
# Engines need time to free GPU memory cleanly before being killed.
|
||||
stop_timeout: 60s
|
||||
health_interval: 2s
|
||||
+210
-15
@@ -7,7 +7,91 @@
|
||||
# __ for nesting: LLMGW_ROUTER__RETRIES=3
|
||||
|
||||
listen: ":8080" # client-facing OpenAI API
|
||||
admin_listen: ":9090" # metrics, health, pprof — do not expose this publicly
|
||||
admin_listen: "127.0.0.1:9090" # dashboard, management API, pprof — keep it local
|
||||
# Metrics only, on their own port, so Prometheus can reach them without pprof
|
||||
# being reachable too. Empty means off; metrics stay on admin_listen either way.
|
||||
# metrics_listen: "0.0.0.0:9091"
|
||||
# Where node agents connect. Separate from the client port so the fleet-facing
|
||||
# surface can be firewalled independently.
|
||||
agent_listen: ":7443"
|
||||
|
||||
# Mutual TLS on the agent port (plan §5.6). With it, a node's identity is taken
|
||||
# from its certificate's common name and an agent claiming a different node id
|
||||
# is refused. Without it the port is plaintext HTTP/2 and anyone who can reach
|
||||
# it can claim to be any node, so it must stay on a trusted network.
|
||||
#
|
||||
# These certificates come from llmgw's own authority, not a public one: public
|
||||
# authorities stopped issuing the clientAuth usage, so no Let's Encrypt
|
||||
# certificate can authenticate an agent.
|
||||
#
|
||||
# llmgw ctl cert init-ca --dir /etc/llmgw/pki
|
||||
# llmgw ctl cert issue-server llmgw-ctl --host 10.0.0.140 --dir /etc/llmgw/pki
|
||||
#
|
||||
# Every address an agent might dial has to appear in --host, bare IP addresses
|
||||
# included: hostname verification reads subject alternative names and ignores
|
||||
# the common name entirely.
|
||||
#
|
||||
# This covers the fleet-facing port only. The client-facing API above is public
|
||||
# and belongs behind a publicly trusted certificate.
|
||||
#
|
||||
# agent_tls:
|
||||
# ca: /etc/llmgw/ca.pem
|
||||
# cert: /etc/llmgw/server.pem
|
||||
# key: /etc/llmgw/server-key.pem
|
||||
|
||||
# The public page at "/". Someone who reaches this address without a key needs a
|
||||
# route to getting one, and the gateway cannot guess who to send them to.
|
||||
#
|
||||
# landing:
|
||||
# contact: "platform@example.com" # or an https:// URL, or "#platform on Slack"
|
||||
# notice: "Internal use only."
|
||||
|
||||
# API keys on the client-facing port (plan §7.9). Off by default, because a
|
||||
# gateway on a trusted network with no keys issued still has to serve — but
|
||||
# nothing else protects that port, so turning it on is what makes it safe to
|
||||
# expose.
|
||||
#
|
||||
# llmgw keys init --store /var/lib/llmgw/llmgw.db
|
||||
# llmgw keys create --store /var/lib/llmgw/llmgw.db --project default "laptop"
|
||||
#
|
||||
# The key is printed once. Only its hash is stored, so there is nothing to
|
||||
# recover it from.
|
||||
#
|
||||
# auth:
|
||||
# enabled: true
|
||||
# store: /var/lib/llmgw/llmgw.db
|
||||
# # How long a resolved key is trusted before being read again. This is the
|
||||
# # price of keeping the database off the request path: it bounds how long a
|
||||
# # revoked key keeps working.
|
||||
# cache_ttl: 30s
|
||||
# # How long a rejected key is remembered as rejected, so a client looping
|
||||
# # with a bad key cannot become a database read per request.
|
||||
# negative_ttl: 5s
|
||||
#
|
||||
# # Serve the signup form. What it creates is a pending account that can do
|
||||
# # nothing until `llmgw users approve` places it in a project, so this is not
|
||||
# # a way in — it only moves key creation from you to the person who needs one.
|
||||
# # Off, there is no form and you issue keys with `llmgw keys create --claim`.
|
||||
# signup_open: true
|
||||
#
|
||||
# # What each model costs, in micro-units per thousand tokens. A unit is
|
||||
# # whatever you decide: self-hosted models have no per-token price, so it can
|
||||
# # be GPU-seconds, a notional rate, or 1000 everywhere — which turns budgets
|
||||
# # into plain token quotas.
|
||||
# #
|
||||
# # A model priced at zero is exempt from budget checks, which is what lets a
|
||||
# # key with nothing left still be served a cheap fallback.
|
||||
# default_price: { prompt_per_1k: 0, completion_per_1k: 0 }
|
||||
# pricing:
|
||||
# llama-3.3-70b: { prompt_per_1k: 1000, completion_per_1k: 3000 }
|
||||
# qwen2.5-32b: { prompt_per_1k: 400, completion_per_1k: 1200 }
|
||||
#
|
||||
# # Usage rows are written in batches off the request path. Beyond buffer_size
|
||||
# # they are dropped rather than blocking a request: losing some accounting is
|
||||
# # bad, turning a stalled database into an outage is worse.
|
||||
# flush_interval: 200ms
|
||||
# batch_size: 500
|
||||
# buffer_size: 8192
|
||||
|
||||
router:
|
||||
scrape_interval: 2s
|
||||
@@ -29,6 +113,39 @@ router:
|
||||
failure_threshold: 3
|
||||
|
||||
default_profile: chat # chat | batch | long-context | debug
|
||||
|
||||
# Approximate prefix routing (plan §7.5 tier 1): send a request to the replica
|
||||
# that has probably already prefilled its prefix.
|
||||
#
|
||||
# It earns its keep when the working set does not fit in one replica's cache.
|
||||
# With a handful of prompts and a large cache, scattering is free — every
|
||||
# replica ends up holding everything — and this only adds constraint. Measured
|
||||
# with sim-engines whose caches were too small for the working set: cache hits
|
||||
# went from 22.5% to 81.7%.
|
||||
#
|
||||
# Watch llmgw_prefix_decisions_total to tune it. A high
|
||||
# load_balance_override count means the guard is doing its job and affinity is
|
||||
# not getting a chance; a high cache_miss_min_load means match_ratio or
|
||||
# min_match is too strict for your traffic.
|
||||
prefix:
|
||||
enabled: false
|
||||
# The shortest match worth acting on, in bytes. Below this the overlap is
|
||||
# coincidence — two English prompts share "The".
|
||||
min_match: 64
|
||||
# The share of the prompt that must match. A short shared header on a very
|
||||
# long prompt saves almost no prefill.
|
||||
match_ratio: 0.3
|
||||
# The imbalance guard, and both must trip before affinity is abandoned.
|
||||
# Absolute alone fires on an idle cluster where one replica has two requests
|
||||
# and another none; relative alone fires on a busy one where 200 against 180
|
||||
# is a rounding error. Without the guard the cache hit rate rises while p99
|
||||
# collapses.
|
||||
balance_abs: 8
|
||||
balance_rel: 1.5
|
||||
# An engine evicts on its own schedule, so an old record may describe a
|
||||
# cache that is gone.
|
||||
ttl: 10m
|
||||
max_nodes: 100000
|
||||
session_header: x-session-id # affinity key; falls back to the `user` field
|
||||
# Ask engines for usage even when the client did not, so token accounting
|
||||
# works regardless of client behaviour. The extra chunk is dropped on the way
|
||||
@@ -43,6 +160,32 @@ router:
|
||||
idle: 60s # between streamed tokens; catches a wedged engine
|
||||
total: 30m
|
||||
|
||||
# Flow control: decides *when* a request may proceed, before scheduling
|
||||
# decides where it goes. Without it a batch workload starves interactive
|
||||
# traffic, because an engine queue is FIFO and knows nothing about tenants.
|
||||
admission:
|
||||
enabled: true
|
||||
# With auth on, the key decides the band and this header is ignored. With
|
||||
# auth off the header is trusted, so any client can promote itself into the
|
||||
# top band — do not expose the port in that configuration.
|
||||
band_header: x-llmgw-band
|
||||
default_band: gold
|
||||
# A band waits once pool saturation passes its ceiling, so the cluster
|
||||
# degrades from the bottom up. A ceiling of 1.0 means "flow up to the
|
||||
# healthy limit"; use a value above 1 for a band that must never wait.
|
||||
bands:
|
||||
- { name: platinum, priority: 400, saturation_ceiling: 1.00 }
|
||||
- { name: gold, priority: 300, saturation_ceiling: 0.85 }
|
||||
- { name: silver, priority: 200, saturation_ceiling: 0.60 }
|
||||
- { name: bronze, priority: 100, saturation_ceiling: 0.30 }
|
||||
# Drop a request that has waited this long: the client has usually given up,
|
||||
# and sending the GPU work nobody is waiting for helps no one.
|
||||
queue_ttl: 30s
|
||||
# Reject immediately once a band's queue is this deep, rather than letting
|
||||
# it accumulate silently.
|
||||
queue_capacity_per_band: 2000
|
||||
interval: 10ms
|
||||
|
||||
# What the scorers normalize against. Measure these with `ctl bench` rather
|
||||
# than guessing: target_waiting in particular governs how the whole system
|
||||
# behaves under load. Too low starves the GPU, too high brings back queueing
|
||||
@@ -53,27 +196,79 @@ router:
|
||||
kv_ceiling: 0.9
|
||||
session_load_factor: 1.25
|
||||
|
||||
# How aggressively the control plane reconciles.
|
||||
# The engines every node installs, and at which version. Unset keeps the
|
||||
# versions this release was checked against (llama.cpp b10991). Each agent
|
||||
# picks the build for its own hardware and installs it itself; changing a
|
||||
# version here reaches every node on its next connect.
|
||||
# engines:
|
||||
# llamacpp: b10991
|
||||
|
||||
control:
|
||||
interval: 5s
|
||||
base_backoff: 10s
|
||||
max_backoff: 10m
|
||||
# Node expiry is staged deliberately: a network blip must not tear down
|
||||
# replicas that are still serving, and rebuilding a multi-node replica costs
|
||||
# tens of minutes of model loading.
|
||||
unreachable_after: 15s
|
||||
lost_after: 60s
|
||||
port_range: [20000, 20999]
|
||||
mgmt_interface: eth0
|
||||
|
||||
# Each model needs a source — a weights path or a HuggingFace id — and at least
|
||||
# one profile if the control plane is to place it. Without profiles the model is
|
||||
# catalogued but never started, which is the router-only mode below.
|
||||
models:
|
||||
- name: llama-3.3-70b
|
||||
source: meta-llama/Llama-3.3-70B-Instruct
|
||||
aliases: [llama]
|
||||
fallbacks: [qwen2.5-32b]
|
||||
routing_profile: chat
|
||||
scaling: { min: 1, max: 4 }
|
||||
profiles:
|
||||
# Two nodes of four H100s each. A multi-node profile requires a fabric,
|
||||
# and a replica never spans two of them.
|
||||
- name: tp8
|
||||
engine: vllm
|
||||
gpu_class: h100-80g
|
||||
tp: 8
|
||||
nodes: 2
|
||||
gpus_per_node: 4
|
||||
priority: 100
|
||||
weight: 8
|
||||
startup_timeout: 30m
|
||||
# Where a whole node is free, one node is cheaper to schedule and faster
|
||||
# to start than two.
|
||||
- name: tp4-a100
|
||||
engine: vllm
|
||||
gpu_class: a100-40g
|
||||
tp: 4
|
||||
priority: 50
|
||||
weight: 3
|
||||
startup_timeout: 20m
|
||||
|
||||
- name: qwen2.5-32b
|
||||
source: Qwen/Qwen2.5-32B-Instruct
|
||||
aliases: [qwen]
|
||||
routing_profile: batch
|
||||
scaling: { min: 1, max: 2 }
|
||||
profiles:
|
||||
- name: tp2
|
||||
engine: sglang
|
||||
gpu_class: h100-80g
|
||||
tp: 2
|
||||
priority: 100
|
||||
weight: 2
|
||||
startup_timeout: 15m
|
||||
|
||||
endpoints:
|
||||
- id: gpu-a-01
|
||||
url: http://gpu-a-01:8000
|
||||
models: [llama-3.3-70b]
|
||||
pool: ib-a/h100-80g
|
||||
weight: 100 # from ctl bench; do not estimate it by hand
|
||||
priority: 100 # higher pools are tried first
|
||||
|
||||
- id: gpu-b-01
|
||||
url: http://gpu-b-01:8000
|
||||
models: [llama-3.3-70b, qwen2.5-32b]
|
||||
pool: ib-b/a100-80g
|
||||
weight: 45
|
||||
priority: 50
|
||||
# A static replica list, used only when no model declares a profile. That is the
|
||||
# router-only mode for testing against sim-engines; with the control plane
|
||||
# running, agents provide the replicas instead.
|
||||
#
|
||||
# endpoints:
|
||||
# - id: sim-1
|
||||
# url: http://127.0.0.1:8001
|
||||
# models: [llama-3.3-70b]
|
||||
# pool: local
|
||||
# weight: 1
|
||||
@@ -0,0 +1,87 @@
|
||||
# Two control planes
|
||||
|
||||
llmgw handles the hard half of this itself: only one control plane places work,
|
||||
and agents find whichever one is leading. What it does not do is move client
|
||||
traffic, because that happens in front of both and llmgw is not in front of
|
||||
anything.
|
||||
|
||||
This directory is the missing half. It is a pattern rather than a product:
|
||||
substitute your own load balancer if you have one.
|
||||
|
||||
## What llmgw does on its own
|
||||
|
||||
**One leader.** Leadership is a lease in the shared store. A second control
|
||||
plane starts, sees the lease, and stands by — it places nothing, evicts
|
||||
nothing and drains nothing. Start one by accident and your cluster is
|
||||
unharmed.
|
||||
|
||||
**Agents follow.** A standby refuses agent streams, so agents move to the
|
||||
leader by themselves. Give each agent both addresses:
|
||||
|
||||
```yaml
|
||||
servers:
|
||||
- "https://ctl-1.internal:7443"
|
||||
- "https://ctl-2.internal:7443"
|
||||
```
|
||||
|
||||
Measured on a cluster: leadership moved three seconds after the leader was
|
||||
stopped, every agent had arrived six seconds later, and every replica was
|
||||
adopted rather than restarted — the engines never stopped serving.
|
||||
|
||||
## What it does not do
|
||||
|
||||
Client traffic. During that same failover the standby served inference
|
||||
correctly on its own port while the public address returned 502, because
|
||||
nothing in front of the two had moved.
|
||||
|
||||
So something has to. The options, in the order most people should consider
|
||||
them:
|
||||
|
||||
**A load balancer you already run.** Health-check `/ready` on each control
|
||||
plane and send traffic to whichever answers. This is the right answer if you
|
||||
have one, and it needs nothing else from this directory.
|
||||
|
||||
**A floating address.** `keepalived.conf.example` here moves one IP between
|
||||
the two machines. Fewer moving parts than a load balancer and no extra hop,
|
||||
at the cost of one machine carrying all the traffic.
|
||||
|
||||
**DNS.** Works, eventually. A record with a short TTL is a failover measured
|
||||
in whatever your resolvers actually honour rather than what you set, which on
|
||||
a bad day is minutes. Fine for a cluster people use interactively, wrong for
|
||||
one an agent is driving.
|
||||
|
||||
## Check `/ready`, not `/health`
|
||||
|
||||
This is the one thing to get right, and the obvious choice is the wrong one.
|
||||
|
||||
`/health` says the process is up. It says nothing about whether it can serve,
|
||||
and a standby is exactly that case: alive, answering, and holding no endpoints
|
||||
at all — its reconciler placed nothing and it has no agent streams to adopt
|
||||
from. A load balancer checking `/health` would send it half the traffic, and
|
||||
half the traffic would be a 503.
|
||||
|
||||
`/ready` says this process has somewhere to send a request. It is 503 on a
|
||||
standby, 503 on a leader whose replicas have all failed, and 200 when there is
|
||||
something to serve — which is the question a load balancer is actually asking.
|
||||
Both are open, because a load balancer has no key.
|
||||
|
||||
During a failover neither passes for about the length of the lease handover:
|
||||
three seconds when the leader is stopped cleanly, up to the lease TTL when it
|
||||
is not. That gap is real and there is no way to have it both ways — traffic
|
||||
cannot go to a control plane that has nothing to route to.
|
||||
|
||||
Not the admin port, for a different reason: it serves pprof, and pprof can
|
||||
dump the process heap, where API keys sit in plaintext.
|
||||
|
||||
## Both control planes need the same store
|
||||
|
||||
Leadership lives in it, and so do keys, budgets and conversations. Two control
|
||||
planes on two SQLite files are two clusters that happen to share a name.
|
||||
|
||||
```yaml
|
||||
auth:
|
||||
store: "@/etc/llmgw/secrets/store-dsn"
|
||||
```
|
||||
|
||||
See [docs/install.md](../../docs/install.md) for moving an existing
|
||||
installation onto PostgreSQL without reissuing every key.
|
||||
@@ -0,0 +1,65 @@
|
||||
# A floating address for two control planes.
|
||||
#
|
||||
# One IP that lives on whichever machine is up. Clients and the reverse proxy
|
||||
# in front of llmgw point at the address, not at either host, so a failover is
|
||||
# invisible to them.
|
||||
#
|
||||
# This is not a load balancer: all traffic goes to the holder. That is usually
|
||||
# the right trade for a control plane, which serves far less than the nodes
|
||||
# it schedules.
|
||||
#
|
||||
# Copy to /etc/keepalived/keepalived.conf on both machines, changing state,
|
||||
# priority and interface.
|
||||
|
||||
global_defs {
|
||||
enable_script_security
|
||||
script_user nobody
|
||||
}
|
||||
|
||||
# The address moves only while llmgw can actually serve, not merely while the
|
||||
# machine is powered on.
|
||||
#
|
||||
# /ready, not /health. /health is 200 on a standby, which is up, answering and
|
||||
# holding no endpoints at all — so the address would sit on a host that refuses
|
||||
# every request. /ready is 200 only when there is somewhere to send one.
|
||||
#
|
||||
# Not the admin port either: it serves pprof, which can dump the heap that
|
||||
# holds every API key.
|
||||
vrrp_script check_llmgw {
|
||||
script "/usr/bin/curl -fsS --max-time 2 -o /dev/null http://127.0.0.1:8080/ready"
|
||||
interval 2
|
||||
timeout 3
|
||||
# Two failures before giving up the address, so one slow moment does not
|
||||
# move it. Two successes before taking it back, so a flapping gateway does
|
||||
# not drag the address with it.
|
||||
fall 2
|
||||
rise 2
|
||||
weight -40
|
||||
}
|
||||
|
||||
vrrp_instance llmgw {
|
||||
# MASTER on the machine you would rather serve; BACKUP on the other.
|
||||
state MASTER
|
||||
interface eth0
|
||||
# Any number both machines agree on and nothing else on the segment uses.
|
||||
virtual_router_id 51
|
||||
# The higher one wins while both are healthy. The weight above is enough
|
||||
# to invert this when the preferred machine's gateway is not answering.
|
||||
priority 100
|
||||
advert_int 1
|
||||
|
||||
authentication {
|
||||
auth_type PASS
|
||||
# Change this. It stops another host on the segment claiming the
|
||||
# address, which is otherwise a single unicast packet away.
|
||||
auth_pass CHANGEME
|
||||
}
|
||||
|
||||
virtual_ipaddress {
|
||||
10.0.0.150/24
|
||||
}
|
||||
|
||||
track_script {
|
||||
check_llmgw
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
# A model mirror
|
||||
|
||||
One machine that holds the model files, so nodes fetch over your own network
|
||||
instead of over the internet.
|
||||
|
||||
This is the highest-value piece of infrastructure around llmgw and the cheapest.
|
||||
Eight nodes each pulling forty gigabytes from Hugging Face is a day's work, and
|
||||
on plenty of corporate networks it is not possible at all. From a mirror on the
|
||||
same switch it is minutes.
|
||||
|
||||
llmgw needs nothing special for this: a model's `source` is a URL, and the URL
|
||||
can be yours.
|
||||
|
||||
## What it is
|
||||
|
||||
An HTTP server serving a directory. That is the whole design.
|
||||
|
||||
```
|
||||
/srv/models/
|
||||
qwen2.5-14b-instruct/
|
||||
2026-03/
|
||||
q4_k_m.gguf
|
||||
bge-m3/
|
||||
2026-01/
|
||||
model.safetensors
|
||||
```
|
||||
|
||||
The date in the path is the point of the layout, not decoration. A model
|
||||
repository that changes upstream changes what you are running, and the first
|
||||
sign is a prompt that worked last week answering differently this week with
|
||||
nobody able to say why. Pin it here and that cannot happen.
|
||||
|
||||
## Running it
|
||||
|
||||
`llmgw` is the mirror. Same binary as the control plane and the agents, one
|
||||
subcommand:
|
||||
|
||||
```
|
||||
sudo mkdir -p /srv/models && sudo chown llmgw:llmgw /srv/models
|
||||
llmgw mirror serve --dir /srv/models --base http://models.internal:8081
|
||||
```
|
||||
|
||||
There is a unit in `deploy/systemd/llmgw-mirror.service`. Set `--base` to the
|
||||
address nodes reach this machine on — not `localhost`, and not a name only this
|
||||
host knows itself by. It is the string that ends up in catalog entries, and a
|
||||
wrong one makes `mirror unused` report files as unreferenced when they are in
|
||||
use. The next step after that list is deletion.
|
||||
|
||||
### Two listeners
|
||||
|
||||
`--listen :8081` serves the weights. Read-only, unauthenticated, Range requests
|
||||
supported so an interrupted fetch on a node resumes rather than restarts. This
|
||||
is the path eight nodes pull forty gigabytes from.
|
||||
|
||||
`--api-listen :8082` is the management API. Separate on purpose: putting a
|
||||
client certificate in front of the byte path would mean re-certifying every
|
||||
agent's fetcher in order to change how models are administered.
|
||||
|
||||
The API is unauthenticated unless you give it the fleet's CA, and it says so at
|
||||
startup rather than leaving you to find out. Anyone who can reach that port can
|
||||
start downloads and delete models, so either keep it on a trusted network or
|
||||
hand it certificates:
|
||||
|
||||
```
|
||||
llmgw mirror serve --dir /srv/models --base http://models.internal:8081 \
|
||||
--tls-ca /etc/llmgw/ca.pem \
|
||||
--tls-cert /etc/llmgw/mirror.pem \
|
||||
--tls-key /etc/llmgw/mirror-key.pem
|
||||
```
|
||||
|
||||
That is the same CA the agents already use, so the certificates come from the
|
||||
authority you have:
|
||||
|
||||
```
|
||||
llmgw ctl cert issue-server mirror --host models.internal --host 10.0.0.146
|
||||
llmgw ctl cert issue-client control-plane
|
||||
```
|
||||
|
||||
`--host` has to cover every address the API is called on, because that is what
|
||||
a client verifies against. Both sides are checked — the mirror will not accept a
|
||||
caller without a certificate from this CA, the same rule the agent stream
|
||||
follows.
|
||||
|
||||
### nginx instead
|
||||
|
||||
If you already run nginx, it serves a directory perfectly well and you do not
|
||||
need `mirror serve` for the bytes:
|
||||
|
||||
```nginx
|
||||
server {
|
||||
listen 80;
|
||||
server_name models.internal;
|
||||
|
||||
root /srv/models;
|
||||
|
||||
# Read-only. This is deliberate: a mirror that can be written to over HTTP
|
||||
# is a way to replace a model with another one, and the nodes would fetch
|
||||
# it without complaint.
|
||||
location / {
|
||||
autoindex on;
|
||||
# Large files, and the client is a machine with a resumable fetcher.
|
||||
# Sendfile and no buffering keep memory flat regardless of size.
|
||||
sendfile on;
|
||||
tcp_nopush on;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
`autoindex` is there so a person can browse what is available; llmgw does not
|
||||
need it.
|
||||
|
||||
## Filling it
|
||||
|
||||
From the mirror itself:
|
||||
|
||||
```
|
||||
llmgw mirror add \
|
||||
--model qwen2.5-14b-instruct --version 2026-03 \
|
||||
"https://huggingface.co/<repo>/resolve/<revision>/q4_k_m.gguf"
|
||||
```
|
||||
|
||||
Use a revision rather than `main` in the source URL. `main` is what moves.
|
||||
|
||||
It downloads to a temporary name and renames when complete, so an interrupted
|
||||
download never leaves something a catalog could be pointed at, and it prints the
|
||||
sha256 when it lands. Compare that against what upstream publishes — Hugging
|
||||
Face gives it as `lfs.oid` from
|
||||
`https://huggingface.co/api/models/<repo>/tree/<revision>`. A mirror is the one
|
||||
place in this system where a corrupt file is copied to every node instead of
|
||||
being caught once.
|
||||
|
||||
It refuses to overwrite a version that is already there. Add a new one beside
|
||||
it: that refusal is what makes the old weights still available when the new ones
|
||||
turn out to be worse.
|
||||
|
||||
## Driving it from elsewhere
|
||||
|
||||
The same fetch over the API, so the control plane — or anything else — can ask
|
||||
for a model without anyone logging in:
|
||||
|
||||
```
|
||||
curl -X POST http://models.internal:8082/v1/models -d '{
|
||||
"source": "https://huggingface.co/<repo>/resolve/<revision>/q4_k_m.gguf",
|
||||
"model": "qwen2.5-14b-instruct",
|
||||
"version": "2026-03"
|
||||
}'
|
||||
```
|
||||
|
||||
The mirror does the downloading. A control plane that pulled the bytes itself
|
||||
and pushed them here would carry every model twice over the network, through the
|
||||
one machine least able to spare the bandwidth.
|
||||
|
||||
It answers `202` with the URL a catalog entry should use, and the bytes are not
|
||||
here yet. Poll for them:
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| `GET /v1/models` | everything held or arriving, and the base to build URLs from |
|
||||
| `POST /v1/models` | fetch this source into model/version |
|
||||
| `GET /v1/models/{model}/{version}` | `fetching` with a byte count, `present`, or `failed` with the reason |
|
||||
| `DELETE /v1/models/{model}/{version}` | remove a version; refuses while it is downloading |
|
||||
| `GET /v1/health` | up *and* able to read its own directory |
|
||||
|
||||
A failed fetch stays `failed` with what went wrong, rather than disappearing:
|
||||
otherwise the thing waiting on that model waits forever with nothing to say why.
|
||||
Asking twice for the same download joins the one already running instead of
|
||||
starting a rival writer on the same file.
|
||||
|
||||
`DELETE` takes the whole version, never one file of it. A model split across
|
||||
several files is useless with one missing, and half a version on disk looks
|
||||
present and will not load.
|
||||
|
||||
## Retiring a model
|
||||
|
||||
```
|
||||
llmgw mirror unused # what no catalog entry refers to
|
||||
llmgw mirror remove --model <m> --version <v> --file <f>
|
||||
```
|
||||
|
||||
`unused` compares against every `source` in the running catalog, so the answer
|
||||
is what can go rather than what looks old. It refuses to answer when it cannot
|
||||
reach the control plane: an empty catalog and an unreachable one produce the
|
||||
same list — everything — and the obvious next step after that list is deletion.
|
||||
|
||||
Retire the catalog entry first, let the fleet settle, then remove. A node that
|
||||
has to refetch a file you deleted cannot, and the failure arrives later, on the
|
||||
node, looking like something else.
|
||||
|
||||
## From the control plane's admin page
|
||||
|
||||
Configure the mirror on the control plane and the same operations appear at
|
||||
`/admin/mirror`, behind the admin session the other admin pages use:
|
||||
|
||||
```yaml
|
||||
mirror:
|
||||
api: https://models.internal:8082
|
||||
tls:
|
||||
ca: /etc/llmgw/ca.crt
|
||||
cert: /etc/llmgw/control-plane.crt
|
||||
key: /etc/llmgw/control-plane.key
|
||||
```
|
||||
|
||||
Leave the block out and the page is not mounted at all, which is the right page
|
||||
for a deployment with no mirror — not an empty one pointing at nothing.
|
||||
|
||||
The page lists what is held, what is arriving with its byte count, and what
|
||||
failed with the reason; it starts a fetch, retries a failed one, and removes a
|
||||
version. It reloads itself only while something is downloading, because a reload
|
||||
at any other moment takes with it whatever is half-typed in the form.
|
||||
|
||||
Two things it deliberately will not do. It cannot point the control plane at a
|
||||
different mirror — that is configuration, and changing it needs access to the
|
||||
host rather than a session. And it does not write the catalog for you: it shows
|
||||
the URL an entry should use, and putting it there is still a deliberate act,
|
||||
because that is the moment a model starts being served.
|
||||
|
||||
The address is a mistake worth naming: `api` is the **management** port, not the
|
||||
one nodes fetch bytes from. They are two listeners so this one can be locked
|
||||
down without re-certifying every agent's fetcher.
|
||||
|
||||
## Pointing llmgw at it
|
||||
|
||||
```yaml
|
||||
- name: qwen2.5-14b-instruct
|
||||
source: http://models.internal/qwen2.5-14b-instruct/2026-03/q4_k_m.gguf
|
||||
scaling: { min: 2, max: 2 }
|
||||
profiles:
|
||||
- name: gpu
|
||||
engine: vllm
|
||||
gpu_class: h100-80g
|
||||
tp: 1
|
||||
```
|
||||
|
||||
Each node fetches once, in the background, before its engine starts — and
|
||||
`llmgw ctl replicas` shows how far along it is while it does. A killed download
|
||||
resumes rather than restarting, and a partial file is never visible under the
|
||||
model's name, so an interrupted fetch cannot be handed to an engine.
|
||||
|
||||
## What this is not
|
||||
|
||||
It is not a cache. Nothing populates it on its own: a model is there because
|
||||
somebody asked for it, by hand or over the API, and one nobody asked for is not
|
||||
there. That is on purpose. A mirror that fetched on demand would reach the
|
||||
internet at the least convenient moment, and in a network where that is blocked
|
||||
it would fail in a way that looks like llmgw being broken.
|
||||
|
||||
It is also not a replacement for a shared filesystem. If you already run one
|
||||
that every node mounts, put the models on it and use a path as the `source` —
|
||||
nothing is fetched at all, and llmgw notices the filesystem is shared and will
|
||||
not count the same free space once per node.
|
||||
|
||||
## When this stops being enough
|
||||
|
||||
Somewhere above ten or twenty nodes, one mirror serving a forty gigabyte file
|
||||
to all of them at once becomes the bottleneck instead of the fix. At that point
|
||||
the options are a filer, a second mirror per rack, or distribution between the
|
||||
nodes themselves.
|
||||
|
||||
llmgw knows which nodes already hold which weights — it is in `ctl status` as
|
||||
"N models on X/Y nodes" — so distributing between them is a feature it is
|
||||
positioned to add. It has not been built, because a mirror is enough for
|
||||
everyone who has needed it so far.
|
||||
@@ -0,0 +1,184 @@
|
||||
# Watching a cluster
|
||||
|
||||
Prometheus, Alertmanager and Grafana, configured for what llmgw actually
|
||||
publishes. Every alert here is written against a series this system emits;
|
||||
nothing is aimed at a metric that might exist.
|
||||
|
||||
## What it is
|
||||
|
||||
One scrape target for the whole fleet.
|
||||
|
||||
Agents publish no metrics of their own. The control plane gathers what they
|
||||
report — node state, memory, disk, engine queue depth, KV occupancy — and
|
||||
republishes all of it, so a single endpoint describes every machine. That is
|
||||
worth knowing before going looking for an exporter on each node.
|
||||
|
||||
## The port it scrapes
|
||||
|
||||
`metrics_listen`, not `admin_listen`.
|
||||
|
||||
The admin listener is bound to loopback for one reason: it also serves pprof,
|
||||
and pprof can dump the process heap, where API keys sit in plaintext. Metrics
|
||||
are the one thing on that listener something else genuinely needs to reach, so
|
||||
they get a port of their own instead of becoming an argument for opening the
|
||||
other one.
|
||||
|
||||
```yaml
|
||||
# /etc/llmgw/server.yaml
|
||||
metrics_listen: ":9091"
|
||||
```
|
||||
|
||||
Confirm it carries nothing else:
|
||||
|
||||
```
|
||||
curl -s -o /dev/null -w '%{http_code}\n' http://<control-plane>:9091/debug/pprof/ # 404
|
||||
curl -s -o /dev/null -w '%{http_code}\n' http://<control-plane>:9091/api/v1/nodes # 404
|
||||
```
|
||||
|
||||
## Installing
|
||||
|
||||
On Ubuntu 24.04 the first two come from the archive; Grafana needs its own
|
||||
repository.
|
||||
|
||||
```
|
||||
apt-get install prometheus prometheus-alertmanager
|
||||
|
||||
install -m 0644 prometheus.yml /etc/prometheus/prometheus.yml
|
||||
install -m 0644 llmgw-rules.yml /etc/prometheus/rules/llmgw.yml
|
||||
install -m 0644 alertmanager.yml /etc/prometheus/alertmanager.yml
|
||||
# Alertmanager answers on loopback only. Its API takes silences from anyone who
|
||||
# can reach it, with no login, and silencing BackupTooOld or ModelParked is how
|
||||
# something that should be noticed is not. Prometheus beside it is its only
|
||||
# client; one Alertmanager needs no cluster port.
|
||||
echo 'ARGS="--web.listen-address=127.0.0.1:9093 --cluster.listen-address="' \
|
||||
> /etc/default/prometheus-alertmanager
|
||||
# Prometheus itself is read-only as installed: its admin and lifecycle APIs
|
||||
# stay off (no --web.enable-admin-api, no --web.enable-lifecycle). It still
|
||||
# answers any query from anyone who reaches :9090, which says which models and
|
||||
# nodes there are and how busy each is. Its readers are the control plane
|
||||
# (monitoring.prometheus) and Grafana beside it; to keep it to those, allow
|
||||
# the control plane and drop the rest, e.g. with nftables:
|
||||
# nft add rule inet filter input tcp dport 9090 ip saddr != { 127.0.0.1, CONTROL_PLANE_IP } drop
|
||||
promtool check config /etc/prometheus/prometheus.yml
|
||||
systemctl restart prometheus prometheus-alertmanager
|
||||
|
||||
install -m 0644 grafana/datasource.yaml /etc/grafana/provisioning/datasources/llmgw.yaml
|
||||
install -m 0644 grafana/dashboards.yaml /etc/grafana/provisioning/dashboards/llmgw.yaml
|
||||
install -m 0644 -o grafana -g grafana grafana/llmgw-overview.json \
|
||||
/var/lib/grafana/dashboards/llmgw.json
|
||||
systemctl restart grafana-server
|
||||
```
|
||||
|
||||
The dashboard is provisioned from a file rather than clicked together, so it
|
||||
survives a reinstall and can be reviewed in a diff.
|
||||
|
||||
## Where to run it
|
||||
|
||||
Not on the control plane. Monitoring that dies with the thing it monitors tells
|
||||
you nothing at the moment you need it, and the control plane is also the busiest
|
||||
machine in a small cluster.
|
||||
|
||||
Anywhere else on the management network will do. The mirror is a reasonable
|
||||
host: it has the disk, and serving weights is bursty rather than constant.
|
||||
|
||||
## Nothing is sent anywhere
|
||||
|
||||
Alertmanager groups and holds alerts and has no receiver.
|
||||
|
||||
Where an alert should go — an address, a Slack workspace, a webhook — is a
|
||||
decision with credentials attached, and a monitoring system that quietly starts
|
||||
mailing people is worse than one that is visibly not wired up. Alerts are
|
||||
visible in Prometheus, and at `:9093` on the monitoring host itself (`ssh -L
|
||||
9093:127.0.0.1:9093 <host>` to open its page); add a receiver to `alertmanager.yml` to
|
||||
make them travel. `critical` already routes separately so it can be pointed
|
||||
somewhere else without touching the rest.
|
||||
|
||||
## Check that it fires
|
||||
|
||||
An alerting system nobody has seen fire is an assumption.
|
||||
|
||||
```
|
||||
llmgw ctl cordon <node> # work already there keeps running
|
||||
# ~2 minutes later
|
||||
ssh <host> curl -s http://127.0.0.1:9093/api/v2/alerts | jq '.[].labels'
|
||||
llmgw ctl uncordon <node> # the alert resolves on its own
|
||||
```
|
||||
|
||||
Cordoning is the cheapest real test: it changes a metric the rules read, stops
|
||||
nothing, and reverses in one command.
|
||||
|
||||
## What the alerts assume
|
||||
|
||||
**A stopped model publishes no series.** `llmgw_model_replicas_ready` exists
|
||||
only for models with a replica placed, so `== 0` cannot fire for one that is
|
||||
deliberately off. That is what makes an on-demand catalogue quiet here.
|
||||
|
||||
**Saturation of 1 is the healthy limit, not failure.** It is the worse of queue
|
||||
depth and KV occupancy, measured by `ctl bench` rather than guessed. The alert
|
||||
allows five minutes above it before saying anything.
|
||||
|
||||
**Latencies are in seconds.** These thresholds are written for CPU inference. A
|
||||
p99 first token of thirty seconds is the alerting point here and would be
|
||||
absurd on a GPU; change it when the hardware changes.
|
||||
|
||||
**Disk is only reported by nodes that hold weights.** A node with no model
|
||||
directory reports nothing rather than zero, and the rule divides free by total
|
||||
so a node that reports neither never fires.
|
||||
|
||||
## Where the metrics live
|
||||
|
||||
On the machine that scrapes them, in Prometheus's own store, and that is the
|
||||
right answer at this size. Measured on this cluster: 2,359 series at 41 samples
|
||||
a second, which compresses to roughly 7 MB a day — about 3 GB for a year.
|
||||
|
||||
```
|
||||
# /etc/default/prometheus
|
||||
ARGS="--storage.tsdb.retention.time=365d --storage.tsdb.retention.size=20GB"
|
||||
```
|
||||
|
||||
A year because the questions worth asking of this data are seasonal: what did
|
||||
capacity look like last quarter, was it always this slow, how much did the
|
||||
fleet grow. Fifteen days answers none of them.
|
||||
|
||||
Both limits, not one. The time limit is the intent; the size cap is what stops
|
||||
Prometheus becoming the thing that fills the disk if the series count grows
|
||||
tenfold after somebody adds a fleet.
|
||||
|
||||
No remote write, no object store, no second tier. At two thousand series that
|
||||
is infrastructure for a problem this cluster does not have, and every one of
|
||||
those has its own failure modes to learn.
|
||||
|
||||
### The part that is not about capacity
|
||||
|
||||
The disk is not the constraint — 900 GB free against 3 GB a year. What matters
|
||||
is what shares it.
|
||||
|
||||
On this cluster the model mirror and the metrics database are on the same
|
||||
filesystem. A mirror that fills it stops Prometheus writing, so monitoring goes
|
||||
quiet at exactly the moment the disk is the problem, and the alert that would
|
||||
have said so is the one that died.
|
||||
|
||||
Two ways out, in order of how much they cost:
|
||||
|
||||
- Watch it, which is now done. `prometheus-node-exporter` arrives with the
|
||||
Prometheus package and was running, unscraped, on this machine — so the disk
|
||||
holding both was watched by nothing. It is a scrape target now, and
|
||||
`MonitoringDiskFilling` fires at fifteen percent rather than five, because
|
||||
the thing being protected is the ability to be told.
|
||||
- Separate them: a volume of its own for `/var/lib/prometheus`, or monitoring
|
||||
on its own machine. Worth doing before the mirror holds a few hundred
|
||||
gigabytes; not worth doing today at seven percent used.
|
||||
|
||||
## What is not here
|
||||
|
||||
No exporter on the llmgw nodes, so no host CPU, load average or network for
|
||||
them. The control plane publishes memory, disk and liveness for every node,
|
||||
which is what placement decides on; the rest is a `node_exporter` away if
|
||||
somebody wants it. The monitoring host itself is scraped, for the reason above.
|
||||
|
||||
No probe of the mirror. Its `/v1/health` is JSON rather than Prometheus text,
|
||||
and checking it properly means `blackbox_exporter`.
|
||||
|
||||
No recording rules. At this size every query in the dashboard runs fast enough
|
||||
raw, and a recording rule is a second definition of a number to keep in step
|
||||
with the first.
|
||||
@@ -0,0 +1,41 @@
|
||||
# Alertmanager for the llmgw cluster.
|
||||
#
|
||||
# Nothing is sent anywhere yet, on purpose. Where an alert should go — an
|
||||
# address, a Slack workspace, a webhook — is a decision with credentials
|
||||
# attached, and a monitoring system that quietly starts mailing people is worse
|
||||
# than one that is visibly not wired up. Alerts group and fire here and are
|
||||
# visible at :9093 and in Prometheus; add a receiver below to make them travel.
|
||||
global:
|
||||
resolve_timeout: 5m
|
||||
|
||||
route:
|
||||
# By alert and by what it is about, so five nodes going quiet in the same
|
||||
# network blip arrive as one notification rather than five.
|
||||
group_by: ['alertname', 'node', 'model']
|
||||
group_wait: 30s
|
||||
group_interval: 5m
|
||||
# Long, because a repeat is a reminder and not news. A cluster with one thing
|
||||
# wrong should not produce a message an hour.
|
||||
repeat_interval: 12h
|
||||
receiver: none
|
||||
|
||||
routes:
|
||||
# Critical gets its own branch so it can be pointed somewhere else later
|
||||
# without touching the rest: the difference between "a node is draining"
|
||||
# and "nothing is serving" is who should be woken.
|
||||
- matchers: [severity="critical"]
|
||||
receiver: none
|
||||
group_wait: 10s
|
||||
repeat_interval: 4h
|
||||
|
||||
receivers:
|
||||
- name: none
|
||||
|
||||
inhibit_rules:
|
||||
# A node that is down explains every replica on it. Reporting both says the
|
||||
# same thing twice and buries whichever is the actual cause.
|
||||
- source_matchers: [alertname="NodeNotReady"]
|
||||
target_matchers: [alertname="EndpointMetricsStale"]
|
||||
equal: ['node']
|
||||
- source_matchers: [alertname="ControlPlaneUnreachable"]
|
||||
target_matchers: [severity="warning"]
|
||||
@@ -0,0 +1,11 @@
|
||||
apiVersion: 1
|
||||
providers:
|
||||
- name: llmgw
|
||||
orgId: 1
|
||||
folder: llmgw
|
||||
type: file
|
||||
disableDeletion: false
|
||||
updateIntervalSeconds: 30
|
||||
allowUiUpdates: true
|
||||
options:
|
||||
path: /var/lib/grafana/dashboards
|
||||
@@ -0,0 +1,11 @@
|
||||
# Provisioned rather than clicked in, so the dashboard below can reference it by
|
||||
# a fixed uid and the whole setup survives a reinstall of the machine.
|
||||
apiVersion: 1
|
||||
datasources:
|
||||
- name: Prometheus
|
||||
uid: llmgw-prom
|
||||
type: prometheus
|
||||
access: proxy
|
||||
url: http://localhost:9090
|
||||
isDefault: true
|
||||
editable: false
|
||||
@@ -0,0 +1,492 @@
|
||||
{
|
||||
"uid": "llmgw-overview",
|
||||
"title": "llmgw \u2014 cluster",
|
||||
"tags": [
|
||||
"llmgw"
|
||||
],
|
||||
"timezone": "browser",
|
||||
"schemaVersion": 39,
|
||||
"version": 1,
|
||||
"refresh": "30s",
|
||||
"time": {
|
||||
"from": "now-6h",
|
||||
"to": "now"
|
||||
},
|
||||
"panels": [
|
||||
{
|
||||
"id": 1,
|
||||
"title": "Nodes ready",
|
||||
"type": "stat",
|
||||
"description": "Nodes taking work. A node that is known but draining or cordoned is not counted.",
|
||||
"gridPos": {
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 4,
|
||||
"h": 4
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "count(llmgw_node_up == 1)",
|
||||
"legendFormat": "ready",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"reduceOptions": {
|
||||
"calcs": [
|
||||
"lastNotNull"
|
||||
]
|
||||
},
|
||||
"textMode": "auto",
|
||||
"colorMode": "value"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"title": "Replicas serving",
|
||||
"type": "stat",
|
||||
"description": "Across every model. A model deliberately stopped contributes nothing and is not a fault.",
|
||||
"gridPos": {
|
||||
"x": 4,
|
||||
"y": 0,
|
||||
"w": 4,
|
||||
"h": 4
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "sum(llmgw_model_replicas_ready)",
|
||||
"legendFormat": "ready",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"reduceOptions": {
|
||||
"calcs": [
|
||||
"lastNotNull"
|
||||
]
|
||||
},
|
||||
"textMode": "auto",
|
||||
"colorMode": "value"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"title": "Models with something serving",
|
||||
"type": "stat",
|
||||
"description": "",
|
||||
"gridPos": {
|
||||
"x": 8,
|
||||
"y": 0,
|
||||
"w": 4,
|
||||
"h": 4
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "count(llmgw_model_replicas_ready > 0)",
|
||||
"legendFormat": "models",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"reduceOptions": {
|
||||
"calcs": [
|
||||
"lastNotNull"
|
||||
]
|
||||
},
|
||||
"textMode": "auto",
|
||||
"colorMode": "value"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"title": "Requests per second",
|
||||
"type": "timeseries",
|
||||
"description": "Failing is counted separately rather than as a percentage, so a quiet cluster reads as quiet instead of as 0% of nothing.",
|
||||
"gridPos": {
|
||||
"x": 12,
|
||||
"y": 0,
|
||||
"w": 12,
|
||||
"h": 4
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "sum(rate(llmgw_requests_total[5m]))",
|
||||
"legendFormat": "all",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
},
|
||||
{
|
||||
"refId": "B",
|
||||
"expr": "sum(rate(llmgw_requests_total{status=~\"5..\"}[5m]))",
|
||||
"legendFormat": "failing",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "reqps"
|
||||
},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"title": "Replicas ready, by model",
|
||||
"type": "timeseries",
|
||||
"description": "Only models with a replica placed appear. One that is stopped on purpose publishes no series.",
|
||||
"gridPos": {
|
||||
"x": 0,
|
||||
"y": 4,
|
||||
"w": 12,
|
||||
"h": 7
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "llmgw_model_replicas_ready",
|
||||
"legendFormat": "{{model}}",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"title": "Pool saturation",
|
||||
"type": "timeseries",
|
||||
"description": "The worse of queue depth and KV occupancy. 1 is the limit measured as healthy, not the point of failure.",
|
||||
"gridPos": {
|
||||
"x": 12,
|
||||
"y": 4,
|
||||
"w": 12,
|
||||
"h": 7
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "llmgw_pool_saturation",
|
||||
"legendFormat": "{{model}}",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"max": 2,
|
||||
"thresholds": {
|
||||
"mode": "absolute",
|
||||
"steps": [
|
||||
{
|
||||
"color": "green",
|
||||
"value": null
|
||||
},
|
||||
{
|
||||
"color": "red",
|
||||
"value": 1
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 7,
|
||||
"title": "Time to first token",
|
||||
"type": "timeseries",
|
||||
"description": "As the router sees it. On CPU these are seconds, not milliseconds; compare against what ctl bench measured.",
|
||||
"gridPos": {
|
||||
"x": 0,
|
||||
"y": 11,
|
||||
"w": 12,
|
||||
"h": 7
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "histogram_quantile(0.5, sum by (model, le) (rate(llmgw_ttft_seconds_bucket[5m])))",
|
||||
"legendFormat": "{{model}} p50",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
},
|
||||
{
|
||||
"refId": "B",
|
||||
"expr": "histogram_quantile(0.99, sum by (model, le) (rate(llmgw_ttft_seconds_bucket[5m])))",
|
||||
"legendFormat": "{{model}} p99",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "s"
|
||||
},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"title": "Tokens per second",
|
||||
"type": "timeseries",
|
||||
"description": "",
|
||||
"gridPos": {
|
||||
"x": 12,
|
||||
"y": 11,
|
||||
"w": 12,
|
||||
"h": 7
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "sum by (direction) (rate(llmgw_tokens_total[5m]))",
|
||||
"legendFormat": "{{direction}}",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "short"
|
||||
},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 9,
|
||||
"title": "Engine queue depth",
|
||||
"type": "timeseries",
|
||||
"description": "Requests the engine has accepted and not started. Rising here is the first sign of saturation.",
|
||||
"gridPos": {
|
||||
"x": 0,
|
||||
"y": 18,
|
||||
"w": 8,
|
||||
"h": 6
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "llmgw_endpoint_waiting",
|
||||
"legendFormat": "{{replica}}",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 10,
|
||||
"title": "KV cache occupancy",
|
||||
"type": "timeseries",
|
||||
"description": "",
|
||||
"gridPos": {
|
||||
"x": 8,
|
||||
"y": 18,
|
||||
"w": 8,
|
||||
"h": 6
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "llmgw_endpoint_kv_cache_usage",
|
||||
"legendFormat": "{{replica}}",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "percentunit"
|
||||
},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 11,
|
||||
"title": "Model disk free",
|
||||
"type": "timeseries",
|
||||
"description": "The filesystems holding weights. A node that cannot hold the next model is skipped by placement rather than failing.",
|
||||
"gridPos": {
|
||||
"x": 16,
|
||||
"y": 18,
|
||||
"w": 8,
|
||||
"h": 6
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "llmgw_node_disk_free_bytes",
|
||||
"legendFormat": "{{node}}",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "bytes"
|
||||
},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 12,
|
||||
"title": "Seconds since each node last reported",
|
||||
"type": "timeseries",
|
||||
"description": "Heartbeats are every 5s. A line that climbs is an agent going quiet; the control plane stops placing there at 15s.",
|
||||
"gridPos": {
|
||||
"x": 0,
|
||||
"y": 24,
|
||||
"w": 12,
|
||||
"h": 6
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "llmgw_node_last_seen_seconds",
|
||||
"legendFormat": "{{node}}",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "s"
|
||||
},
|
||||
"overrides": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 13,
|
||||
"title": "Replicas by state",
|
||||
"type": "timeseries",
|
||||
"description": "",
|
||||
"gridPos": {
|
||||
"x": 12,
|
||||
"y": 24,
|
||||
"w": 12,
|
||||
"h": 6
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"expr": "sum by (state) (llmgw_replicas)",
|
||||
"legendFormat": "{{state}}",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "llmgw-prom"
|
||||
}
|
||||
}
|
||||
],
|
||||
"fieldConfig": {
|
||||
"defaults": {},
|
||||
"overrides": []
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
# Alerts for llmgw. Every expression here is against a series this cluster
|
||||
# actually publishes; nothing is written against a metric that might exist.
|
||||
groups:
|
||||
- name: llmgw-fleet
|
||||
rules:
|
||||
- alert: ControlPlaneUnreachable
|
||||
expr: up{job="llmgw"} == 0
|
||||
for: 1m
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "The control plane is not answering"
|
||||
description: >-
|
||||
Nothing has scraped 10.0.0.140:9091 for a minute. Engines already
|
||||
running keep serving — the router is in the same process, so if
|
||||
this is down, so is the gateway.
|
||||
|
||||
- alert: NodeNotReady
|
||||
expr: llmgw_node_up == 0
|
||||
for: 2m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.node }} is {{ $labels.state }}"
|
||||
description: >-
|
||||
The node is known but not taking work. Two minutes, because a node
|
||||
restarting its agent is briefly not ready and is not a problem.
|
||||
|
||||
- alert: AgentGoneQuiet
|
||||
expr: llmgw_node_last_seen_seconds > 60
|
||||
for: 1m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.node }} has not reported for {{ $value | humanizeDuration }}"
|
||||
description: >-
|
||||
Heartbeats are every 5s. A minute of silence is the network or the
|
||||
agent, and the control plane stages expiry at 15s unreachable and
|
||||
60s lost — so this fires around the point it stops placing there.
|
||||
|
||||
- alert: ModelHasNothingServing
|
||||
expr: llmgw_model_replicas_ready == 0
|
||||
for: 10m
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} has no replica ready"
|
||||
description: >-
|
||||
The model should be running and nothing is serving it. Ten minutes,
|
||||
because a large model on CPU legitimately takes minutes to load. The
|
||||
series is published at 0 for every model that should be running —
|
||||
before v0.10.4 it was published only while a replica was ready, so
|
||||
this could not fire at all. A model deliberately stopped, or an
|
||||
on-demand one nobody is asking for, publishes nothing.
|
||||
|
||||
- alert: ModelParked
|
||||
expr: llmgw_model_parked == 1
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} failed every time it started and is no longer retried"
|
||||
description: >-
|
||||
Five replicas in a row failed at start. Before this existed a model
|
||||
whose weights its engine could not load was started 155 times over
|
||||
eight hours. `llmgw ctl status` has the last error; once it is dealt
|
||||
with, `llmgw ctl retry <model>`. Changing the model's definition
|
||||
also retries it.
|
||||
|
||||
- alert: ModelFailingToStart
|
||||
expr: llmgw_model_start_failures >= 2
|
||||
for: 5m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} has failed to start {{ $value }} times in a row"
|
||||
description: >-
|
||||
On its way to being parked. The engine's own words are in
|
||||
`llmgw ctl logs <node> --source <shard>` for the next attempt, or in
|
||||
the control plane's log for the last one.
|
||||
|
||||
- alert: ModelCannotBePlaced
|
||||
expr: llmgw_model_unplaceable == 1
|
||||
for: 10m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} is wanted and fits on no node"
|
||||
description: >-
|
||||
Short of replicas, and the last attempt found nowhere with the
|
||||
memory, cards or engine it needs. `llmgw ctl explain-placement
|
||||
<model>` says what each node lacks. For an on-demand model, requests
|
||||
for it are told this rather than asked to retry.
|
||||
|
||||
- alert: ModelPreempted
|
||||
expr: increase(llmgw_model_preempted_total[15m]) > 0
|
||||
labels: { severity: info }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} ({{ $labels.class }}) gave up replicas to a higher class"
|
||||
description: >-
|
||||
Capacity ran short and a model of a higher class took this one's
|
||||
room. Expected now and then; if it keeps firing for the same model,
|
||||
the cluster is too small for what is asked of it at these classes.
|
||||
`llmgw ctl models` shows who it gave way to.
|
||||
|
||||
- alert: CapacityAtLimit
|
||||
expr: llmgw_capacity_holds >= on(scope, resource) llmgw_capacity_limit
|
||||
for: 15m
|
||||
labels: { severity: info }
|
||||
annotations:
|
||||
summary: "{{ $labels.scope }} holds all the {{ $labels.resource }} it may"
|
||||
description: >-
|
||||
Its models get no further replicas, under load or not, until
|
||||
something of theirs stops or the limit is raised: `llmgw ctl quota
|
||||
list`, `llmgw ctl quota set`.
|
||||
|
||||
- alert: EngineInstallFailed
|
||||
expr: llmgw_node_engine_install{state="failed"} == 1
|
||||
for: 15m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.engine }} {{ $labels.version }} will not install on {{ $labels.node }}"
|
||||
description: >-
|
||||
The node takes no model needing it until it does. The agent retries
|
||||
on its own; `llmgw ctl nodes` and the node's logs on the cluster
|
||||
page say why it failed.
|
||||
|
||||
- alert: ControlPlaneLeadership
|
||||
expr: sum(llmgw_control_plane_leader) != 1
|
||||
for: 2m
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "{{ $value }} control planes believe they are placing work"
|
||||
description: >-
|
||||
None means nothing is placed or replaced; two means both are, each
|
||||
into GPUs the other believes free. The lease is in the store: check
|
||||
it is reachable from every control plane.
|
||||
|
||||
- alert: CertificateExpiresSoon
|
||||
expr: (llmgw_certificate_expiry_timestamp_seconds - time()) < 30 * 86400
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.cert }} ({{ $labels.subject }}) expires in {{ $value | humanizeDuration }}"
|
||||
description: >-
|
||||
Certificates last a year and nothing renews them. When the agent
|
||||
server's runs out every agent refuses to connect at once. Reissue
|
||||
with `llmgw ctl cert` and restart the control plane.
|
||||
|
||||
- alert: CertificateExpiresVerySoon
|
||||
expr: (llmgw_certificate_expiry_timestamp_seconds - time()) < 7 * 86400
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "{{ $labels.cert }} ({{ $labels.subject }}) expires in {{ $value | humanizeDuration }}"
|
||||
description: See CertificateExpiresSoon; a week is not long.
|
||||
|
||||
- alert: BackupTooOld
|
||||
expr: time() - llmgw_backup_last_success_timestamp_seconds > 26 * 3600
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "The newest llmgw backup is {{ $value | humanizeDuration }} old"
|
||||
description: >-
|
||||
llmgw-backup.timer runs nightly. Also fires when there has never
|
||||
been a backup (the series is 0 then). `journalctl -u
|
||||
llmgw-backup.service` on the control plane says why the last one
|
||||
failed; docs/install.md, "Backing up", says how to set it up.
|
||||
|
||||
- alert: PoolSaturated
|
||||
expr: llmgw_pool_saturation > 1
|
||||
for: 5m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} is over its healthy limit"
|
||||
description: >-
|
||||
Saturation is the worse of queue depth and KV occupancy; 1 is the
|
||||
limit measured as healthy, not the point of failure. Sustained above
|
||||
it means requests are queueing — add a replica or lower the load.
|
||||
|
||||
- alert: EndpointBreakerOpen
|
||||
expr: llmgw_endpoint_breaker_open == 1
|
||||
for: 2m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.replica }} has been taken out of rotation"
|
||||
description: >-
|
||||
The router stopped sending to this replica after repeated failures.
|
||||
Traffic is going elsewhere; if it is the only replica, the model is
|
||||
down.
|
||||
|
||||
- alert: EndpointMetricsStale
|
||||
expr: llmgw_endpoint_metrics_stale == 1
|
||||
for: 5m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.replica }} is being routed to blind"
|
||||
description: >-
|
||||
Its queue depth and KV usage are old, so the scheduler is choosing
|
||||
with stale numbers. The replica may be serving perfectly well; the
|
||||
decisions about it are the problem.
|
||||
|
||||
- alert: NodeDiskFilling
|
||||
expr: llmgw_node_disk_free_bytes / llmgw_node_disk_total_bytes < 0.10
|
||||
for: 15m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.node }} has {{ $value | humanizePercentage }} of its model disk free"
|
||||
description: >-
|
||||
Weights are tens of gigabytes. A node that cannot hold the next one
|
||||
is skipped by placement rather than failing, so this is quiet until
|
||||
somebody wonders why a model will not land there.
|
||||
|
||||
- name: monitoring-host
|
||||
rules:
|
||||
- alert: MonitoringDiskFilling
|
||||
expr: >-
|
||||
node_filesystem_avail_bytes{job="monitoring-host",fstype!~"tmpfs|overlay|squashfs"}
|
||||
/ node_filesystem_size_bytes{job="monitoring-host",fstype!~"tmpfs|overlay|squashfs"} < 0.15
|
||||
for: 10m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.mountpoint }} on the monitoring host is {{ $value | humanizePercentage }} free"
|
||||
description: >-
|
||||
This filesystem holds the model mirror and the metrics database at
|
||||
once. Filling it stops new weights arriving and stops Prometheus
|
||||
writing — so monitoring goes quiet at the moment it is most needed,
|
||||
which is why this fires at fifteen percent rather than five.
|
||||
|
||||
- alert: MonitoringDiskAlmostFull
|
||||
expr: >-
|
||||
node_filesystem_avail_bytes{job="monitoring-host",fstype!~"tmpfs|overlay|squashfs"}
|
||||
/ node_filesystem_size_bytes{job="monitoring-host",fstype!~"tmpfs|overlay|squashfs"} < 0.05
|
||||
for: 5m
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "{{ $labels.mountpoint }} on the monitoring host has {{ $value | humanizePercentage }} left"
|
||||
|
||||
- name: llmgw-traffic
|
||||
rules:
|
||||
# These need traffic through the gateway to have any series at all. A
|
||||
# counter with no observations publishes nothing, so a quiet cluster
|
||||
# simply has no data here rather than a rule that is always false.
|
||||
- alert: HighErrorRate
|
||||
expr: >-
|
||||
sum by (model) (rate(llmgw_requests_total{status=~"5.."}[5m]))
|
||||
/ sum by (model) (rate(llmgw_requests_total[5m])) > 0.05
|
||||
for: 5m
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} is failing {{ $value | humanizePercentage }} of requests"
|
||||
|
||||
- alert: SlowFirstToken
|
||||
expr: >-
|
||||
histogram_quantile(0.99,
|
||||
sum by (model, le) (rate(llmgw_ttft_seconds_bucket[10m]))) > 30
|
||||
for: 10m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} takes {{ $value | humanizeDuration }} to first token at p99"
|
||||
description: >-
|
||||
On CPU this is expected under load and is why the threshold is
|
||||
thirty seconds rather than one. Compare with what ctl bench measured
|
||||
before treating it as a fault.
|
||||
|
||||
- alert: RequestsBeingRejected
|
||||
expr: sum by (band) (rate(llmgw_admission_rejected_total[5m])) > 0
|
||||
for: 5m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "Admission is turning away {{ $labels.band }} traffic"
|
||||
description: >-
|
||||
Working as designed when the cluster is full: the lowest band is
|
||||
held first. Worth knowing about because the callers being refused
|
||||
do not know why.
|
||||
@@ -0,0 +1,42 @@
|
||||
# Prometheus for the llmgw cluster.
|
||||
#
|
||||
# One scrape target for the whole fleet. Agents publish nothing of their own —
|
||||
# the control plane gathers what they report and republishes it, so node health,
|
||||
# disk and endpoint state all arrive here from 10.0.0.140.
|
||||
global:
|
||||
scrape_interval: 15s
|
||||
evaluation_interval: 15s
|
||||
external_labels:
|
||||
cluster: mynodes
|
||||
|
||||
alerting:
|
||||
alertmanagers:
|
||||
- static_configs:
|
||||
- targets: ['localhost:9093']
|
||||
|
||||
rule_files:
|
||||
- /etc/prometheus/rules/llmgw.yml
|
||||
|
||||
scrape_configs:
|
||||
# The control plane's dedicated metrics port. Not its admin listener: that one
|
||||
# is loopback because it also serves pprof, and pprof can dump the heap where
|
||||
# API keys sit in plaintext.
|
||||
- job_name: llmgw
|
||||
static_configs:
|
||||
- targets: ['10.0.0.140:9091']
|
||||
labels:
|
||||
instance: control-plane
|
||||
|
||||
- job_name: prometheus
|
||||
static_configs:
|
||||
- targets: ['localhost:9090']
|
||||
|
||||
# The machine this runs on. It came with the Prometheus package and was not
|
||||
# being scraped, which left the disk holding both the model mirror and the
|
||||
# metrics themselves watched by nothing — so a mirror that filled it would
|
||||
# take monitoring down and the alert that would have warned did not exist.
|
||||
- job_name: monitoring-host
|
||||
static_configs:
|
||||
- targets: ['localhost:9100']
|
||||
labels:
|
||||
instance: monitoring
|
||||
@@ -0,0 +1,41 @@
|
||||
# llmgw behind Caddy: the client-facing port (8080) on a public name, with a
|
||||
# certificate Caddy obtains and renews by itself.
|
||||
#
|
||||
# Replace llm.example.com, point its DNS at this machine, open 80 and 443, and
|
||||
# run: caddy run --config /etc/caddy/Caddyfile
|
||||
#
|
||||
# Only 8080 goes through here. 7443 (agents) is mutual TLS with the cluster's
|
||||
# own authority and must reach llmgw directly; 9090 and 9091 stay private.
|
||||
|
||||
llm.example.com {
|
||||
# The router reads at most 32 MiB of body (router.max_body_bytes). Refusing
|
||||
# more here gives the client a 413 before the upload, not after it.
|
||||
request_body {
|
||||
max_size 32MB
|
||||
}
|
||||
|
||||
reverse_proxy 127.0.0.1:8080 {
|
||||
# Tokens are written the moment they are generated. Without this a
|
||||
# streamed reply can arrive in bursts, and a client measuring time to
|
||||
# first token measures the proxy instead.
|
||||
flush_interval -1
|
||||
|
||||
# Caddy sets X-Forwarded-Proto and X-Forwarded-Host itself and, with
|
||||
# no trusted_proxies, discards whatever the client sent. llmgw relies
|
||||
# on that: the proto decides whether session cookies are Secure, and
|
||||
# the host is the base URL the landing page tells people to use.
|
||||
# Behind a load balancer of your own, name it in trusted_proxies.
|
||||
|
||||
transport http {
|
||||
# A long generation can be silent for a while before the first
|
||||
# token. llmgw has its own limits (timeouts.*); this must not be
|
||||
# the shorter one.
|
||||
response_header_timeout 15m
|
||||
}
|
||||
}
|
||||
|
||||
# No `encode` on purpose: compressing a stream buffers it.
|
||||
|
||||
# Access log to stderr, which under systemd is the journal.
|
||||
log
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
# A reverse proxy in front of llmgw
|
||||
|
||||
The client-facing port (8080) is plain HTTP. On anything but a trusted network,
|
||||
put it behind a proxy that holds a publicly trusted certificate.
|
||||
|
||||
| File | For |
|
||||
| --- | --- |
|
||||
| [Caddyfile](Caddyfile) | Caddy — obtains and renews its own certificate; the least work |
|
||||
| [nginx.conf](nginx.conf) | nginx, with a certificate from certbot |
|
||||
|
||||
Both are also attached to every release.
|
||||
|
||||
The `/about` page shows both, from copies in `internal/web/proxy`; a test fails
|
||||
when those differ from these, so change them here and copy them over.
|
||||
|
||||
Four things both files get right, and any other proxy has to as well:
|
||||
|
||||
- **Streaming is not buffered.** Replies are server-sent events; a buffering
|
||||
proxy delivers them in one burst at the end.
|
||||
- **The read timeout is longer than a model's silence.** nginx's default of 60
|
||||
seconds turns a slow first token into a 504.
|
||||
- **`X-Forwarded-Proto` and `X-Forwarded-Host` are set by the proxy**, replacing
|
||||
whatever the client sent. llmgw uses the first to mark session cookies
|
||||
`Secure` and the second for the base URL on its landing page.
|
||||
- **The body limit matches llmgw's** (32 MiB, `router.max_body_bytes`), so an
|
||||
oversized request is refused before it is uploaded.
|
||||
|
||||
Only port 8080 is proxied. Agents reach 7443 directly — it is mutual TLS with
|
||||
the cluster's own authority, and a proxy terminating it would take away the
|
||||
node identity it carries. 9090 (admin, pprof) stays on loopback; 9091
|
||||
(metrics) goes to your monitoring host only.
|
||||
|
||||
With a proxy on the same machine, bind llmgw to loopback so nothing reaches it
|
||||
around the proxy:
|
||||
|
||||
```yaml
|
||||
# /etc/llmgw/server.yaml
|
||||
listen: "127.0.0.1:8080"
|
||||
```
|
||||
@@ -0,0 +1,63 @@
|
||||
# llmgw behind nginx: the client-facing port (8080) on a public name.
|
||||
#
|
||||
# Put this in /etc/nginx/sites-available/llmgw (or conf.d/llmgw.conf), replace
|
||||
# llm.example.com, and get a certificate — certbot --nginx -d llm.example.com
|
||||
# fills in the ssl_ lines below. Then: nginx -t && systemctl reload nginx
|
||||
#
|
||||
# Only 8080 goes through here. 7443 (agents) is mutual TLS with the cluster's
|
||||
# own authority and must reach llmgw directly; 9090 and 9091 stay private.
|
||||
|
||||
upstream llmgw {
|
||||
server 127.0.0.1:8080;
|
||||
keepalive 32;
|
||||
}
|
||||
|
||||
server {
|
||||
listen 80;
|
||||
listen [::]:80;
|
||||
server_name llm.example.com;
|
||||
return 301 https://$host$request_uri;
|
||||
}
|
||||
|
||||
server {
|
||||
# `listen ... http2` rather than `http2 on;`: the latter needs nginx
|
||||
# 1.25.1, and Ubuntu 22.04 and Debian 12 ship older ones. Newer nginx
|
||||
# accepts this form with a deprecation warning.
|
||||
listen 443 ssl http2;
|
||||
listen [::]:443 ssl http2;
|
||||
server_name llm.example.com;
|
||||
|
||||
ssl_certificate /etc/letsencrypt/live/llm.example.com/fullchain.pem;
|
||||
ssl_certificate_key /etc/letsencrypt/live/llm.example.com/privkey.pem;
|
||||
|
||||
# The router reads at most 32 MiB of body (router.max_body_bytes).
|
||||
client_max_body_size 32m;
|
||||
|
||||
location / {
|
||||
proxy_pass http://llmgw;
|
||||
proxy_http_version 1.1;
|
||||
proxy_set_header Connection "";
|
||||
|
||||
# Set, never appended: llmgw decides from X-Forwarded-Proto whether
|
||||
# session cookies are Secure, and builds the base URL the landing page
|
||||
# shows from X-Forwarded-Host. A value passed through from the client
|
||||
# would be the client's choice.
|
||||
proxy_set_header Host $host;
|
||||
proxy_set_header X-Forwarded-Host $host;
|
||||
proxy_set_header X-Forwarded-Proto $scheme;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
|
||||
# Streamed replies: tokens go out as they are generated. Buffering
|
||||
# turns a stream into one late burst, and gzip buffers too.
|
||||
proxy_buffering off;
|
||||
proxy_cache off;
|
||||
gzip off;
|
||||
|
||||
# nginx's default is 60 seconds between two reads from upstream. A
|
||||
# large model can think longer than that before its first token, and
|
||||
# the client then sees a 504 for a request that was about to succeed.
|
||||
# llmgw has its own limits (timeouts.*); this must not be the shorter.
|
||||
proxy_read_timeout 15m;
|
||||
proxy_send_timeout 15m;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
# systemd units
|
||||
|
||||
Two roles, two units. Copy the one you need to `/etc/systemd/system/`, put the
|
||||
matching configuration in `/etc/llmgw/`, and `systemctl enable --now` it.
|
||||
|
||||
## The control plane — `llmgw-server.service`
|
||||
|
||||
Runs unprivileged. It needs to read its own TLS key, so that file must belong to
|
||||
the `llmgw` user and nothing else:
|
||||
|
||||
```
|
||||
useradd --system --home-dir /var/lib/llmgw --shell /usr/sbin/nologin llmgw
|
||||
chown llmgw:llmgw /etc/llmgw/server-key.pem
|
||||
chmod 0400 /etc/llmgw/server-key.pem
|
||||
```
|
||||
|
||||
`StateDirectory=llmgw` creates `/var/lib/llmgw` with the right owner, which is
|
||||
where the store belongs. `ProtectSystem=strict` makes the rest of the filesystem
|
||||
read-only, so anything else it writes to has to be named explicitly.
|
||||
|
||||
## The agent — `llmgw-agent.service`
|
||||
|
||||
Runs as root, because it launches engines and needs device access.
|
||||
|
||||
`KillMode=process` is the line that matters and the reason this file exists
|
||||
rather than being generated. Without it, systemd kills the whole control group
|
||||
on stop — including the engines, which hold a loaded model and are serving
|
||||
traffic. Engines are meant to outlive an agent restart: the router talks to them
|
||||
directly, and the agent adopts them again when it comes back. Taking them down
|
||||
with the agent turns a thirty-second upgrade into a full model reload on every
|
||||
node.
|
||||
|
||||
systemd will log `Found left-over process ... in control group while starting
|
||||
unit. Ignoring.` on the next start. That is the mechanism working, not a fault.
|
||||
@@ -0,0 +1,21 @@
|
||||
[Unit]
|
||||
Description=llmgw node agent
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
ExecStart=/usr/local/bin/llmgw agent --config /etc/llmgw/agent.yaml
|
||||
# Root, because the agent launches engines and needs device access. Engines
|
||||
# outlive a restart of this unit and are adopted again on the next start, so
|
||||
# KillMode must not take the whole process group with it.
|
||||
KillMode=process
|
||||
Restart=on-failure
|
||||
RestartSec=2s
|
||||
# RDMA pins the buffers it registers, and NCCL registers a great deal: with
|
||||
# systemd's default 8 MB lock limit, an engine that drops the capability that
|
||||
# lets root ignore it fails its first collective over InfiniBand.
|
||||
LimitMEMLOCK=infinity
|
||||
StateDirectory=llmgw-agent
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -0,0 +1,14 @@
|
||||
[Unit]
|
||||
Description=llmgw backup: the store and /etc/llmgw, to /var/backups/llmgw-server
|
||||
After=llmgw-server.service postgresql.service
|
||||
|
||||
[Service]
|
||||
Type=oneshot
|
||||
# Nothing it writes is for anyone else: the archive holds the CA's key and a
|
||||
# copy of the store.
|
||||
UMask=0077
|
||||
# Root, because the archive includes /etc/llmgw/secrets and the certificate
|
||||
# authority's key, which the llmgw user deliberately cannot read all of.
|
||||
ExecStart=/usr/local/bin/llmgw server backup --config /etc/llmgw/server.yaml --keep 14
|
||||
Nice=10
|
||||
IOSchedulingClass=idle
|
||||
@@ -0,0 +1,12 @@
|
||||
[Unit]
|
||||
Description=Back up llmgw every night
|
||||
|
||||
[Timer]
|
||||
OnCalendar=*-*-* 03:15:00
|
||||
# A machine that was off at 03:15 backs up when it next starts, rather than
|
||||
# waiting another day with a gap nobody notices until the alert.
|
||||
Persistent=true
|
||||
RandomizedDelaySec=10m
|
||||
|
||||
[Install]
|
||||
WantedBy=timers.target
|
||||
@@ -0,0 +1,29 @@
|
||||
[Unit]
|
||||
Description=llmgw model mirror
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
# --base is what a catalog entry will say, so it has to be the address nodes
|
||||
# reach this machine on — not localhost, and not the name only this host knows
|
||||
# itself by. Getting it wrong makes `mirror unused` report files as unreferenced
|
||||
# when they are in use, and the next step after that list is deletion.
|
||||
ExecStart=/usr/local/bin/llmgw mirror serve \
|
||||
--dir /srv/models \
|
||||
--listen :8081 \
|
||||
--api-listen :8082 \
|
||||
--base http://models.internal:8081
|
||||
User=llmgw
|
||||
Group=llmgw
|
||||
Restart=on-failure
|
||||
RestartSec=2s
|
||||
NoNewPrivileges=true
|
||||
ProtectSystem=strict
|
||||
ProtectHome=true
|
||||
PrivateTmp=true
|
||||
# The one writable path. The mirror serves reads to the whole fleet and writes
|
||||
# only what its own API was asked to fetch, so nothing else needs to be.
|
||||
ReadWritePaths=/srv/models
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -0,0 +1,19 @@
|
||||
[Unit]
|
||||
Description=llmgw control plane
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
ExecStart=/usr/local/bin/llmgw server --config /etc/llmgw/server.yaml
|
||||
User=llmgw
|
||||
Group=llmgw
|
||||
Restart=on-failure
|
||||
RestartSec=2s
|
||||
NoNewPrivileges=true
|
||||
ProtectSystem=strict
|
||||
ProtectHome=true
|
||||
PrivateTmp=true
|
||||
StateDirectory=llmgw
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -0,0 +1,149 @@
|
||||
# Choosing what to run
|
||||
|
||||
llmgw serves whatever you give it and has no opinion about which models those
|
||||
are. This page has an opinion, because "we have GPUs, now what" is the question
|
||||
that stalls most of these projects — and the usual first answer, one large
|
||||
model, is the wrong shape.
|
||||
|
||||
Model releases move faster than documentation. Treat the families named here as
|
||||
a starting point and check what is current before committing; the reasoning is
|
||||
what should outlive the names.
|
||||
|
||||
## Three jobs, not one model
|
||||
|
||||
The requests a company actually makes fall into three groups with different
|
||||
needs, and no single model serves all three well.
|
||||
|
||||
**Interactive** — chat, autocomplete, short questions. Most of your traffic.
|
||||
Here latency beats quality: a person waits two seconds and does not wait
|
||||
twenty. A 7–14B class model, always resident.
|
||||
|
||||
**Heavy** — long document analysis, code generation, anything that needs
|
||||
reasoning. Quality beats latency, and a queue is acceptable. A 30–70B class
|
||||
model, loaded on demand:
|
||||
|
||||
```yaml
|
||||
scaling: { min: 0, max: 1, idle_timeout: 30m }
|
||||
```
|
||||
|
||||
llmgw starts it when somebody asks and stops it when nobody has, so it costs
|
||||
nothing while it is not in use. The first request waits for the load, and the
|
||||
client is told that rather than refused.
|
||||
|
||||
**Embeddings** — and this is the one that gets forgotten. Most "let us chat
|
||||
with our documents" projects are retrieval projects, and retrieval runs on an
|
||||
embedding model. It is small, it is cheap, and it carries the highest request
|
||||
volume of the three: every document chunk once, every query forever. llmgw
|
||||
routes `/v1/embeddings` like any other endpoint.
|
||||
|
||||
A deployment with the first two and not the third can answer questions from
|
||||
what the model already knew. It cannot answer questions about your company.
|
||||
|
||||
## Swedish and English
|
||||
|
||||
The first thing to know is that this combination is harder than either alone.
|
||||
|
||||
Open models are strongest in English by a wide margin. Swedish is served —
|
||||
better than most languages its size, because the Nordic languages appear in
|
||||
multilingual training sets — but the gap is real, and it widens as models get
|
||||
smaller. **A 7B model that feels fine in English can feel noticeably worse in
|
||||
Swedish**, which means the interactive tier may need to be larger than an
|
||||
English-only deployment would choose. Test before deciding; do not assume the
|
||||
English benchmark transfers.
|
||||
|
||||
Two things matter more than the model's headline multilingual score:
|
||||
|
||||
**Cross-lingual retrieval.** A Swedish company writes internal documents in
|
||||
Swedish and reads technical documentation in English. Somebody will ask a
|
||||
question in Swedish whose answer is in an English document. An embedding model
|
||||
that handles both languages separately may still fail to connect them, and it
|
||||
fails silently — the model gives a confident answer from the wrong source, and
|
||||
nobody notices because the answer reads well. Test this explicitly: ask in one
|
||||
language, check that documents in the other are retrieved.
|
||||
|
||||
**Swedish compounds.** *Kundtjänstmedarbetare* is one token-hostile word where
|
||||
English would use three. This hurts keyword search more than dense retrieval,
|
||||
so if you run hybrid search, expect the lexical half to underperform on Swedish
|
||||
and weight accordingly.
|
||||
|
||||
## How to actually decide
|
||||
|
||||
Benchmarks will not settle this. Take twenty real documents and twenty real
|
||||
questions from the company — in both languages, including at least a few where
|
||||
the question and the answer are in different languages — and compare candidates
|
||||
on those. An afternoon of this is worth more than a week of reading leaderboards,
|
||||
and it is the only way to find the failure modes that matter to you.
|
||||
|
||||
Do the same for the embedding model, separately. A good generation model on top
|
||||
of bad retrieval produces fluent, confident, wrong answers, which is the worst
|
||||
failure this stack has: it does not look like a failure.
|
||||
|
||||
## How many replicas
|
||||
|
||||
Not "how many employees". **Concurrency**, which is far lower than headcount
|
||||
because people spend most of their time reading rather than waiting on a model.
|
||||
Two hundred people generate perhaps five to fifteen concurrent requests at
|
||||
peak, for interactive use.
|
||||
|
||||
vLLM batches continuously, so one replica on a capable card serves many
|
||||
concurrent requests. Which means the replica count is usually decided by
|
||||
**redundancy rather than throughput**:
|
||||
|
||||
| Tier | Replicas | Why |
|
||||
|---|---|---|
|
||||
| Embeddings | 2 | Not for capacity. If retrieval stops, every document feature stops. |
|
||||
| Interactive | 2 | So a node can be rebooted during working hours. |
|
||||
| Heavy | 0–1, on demand | Loaded when asked for, stopped when idle. |
|
||||
|
||||
Then measure rather than guess:
|
||||
|
||||
```
|
||||
llmgw ctl bench <model>
|
||||
```
|
||||
|
||||
It climbs concurrency until the SLO breaks and reports the last level that met
|
||||
it. Peak throughput arrives long after latency has become unusable, so
|
||||
throughput alone overstates what the hardware can serve.
|
||||
|
||||
## Where the weights live
|
||||
|
||||
The single highest-value piece of infrastructure here is **a local mirror**: one
|
||||
internal HTTP server holding the model files. [deploy/mirror/](../deploy/mirror/)
|
||||
has the nginx config and the directory layout. Hugging Face is slow, rate
|
||||
limited, and blocked outright on plenty of corporate networks. Eight nodes each
|
||||
pulling forty gigabytes from the internet is a day; from a mirror on the same
|
||||
switch it is minutes. Point a model's `source` at it and llmgw fetches from
|
||||
there.
|
||||
|
||||
**Pin the version.** Put it in the path — `/models/qwen2.5-14b/2026-03/q4_k_m.gguf`
|
||||
— so a repository that changes upstream cannot change what you are running.
|
||||
Without this, a prompt that worked last week gives a different answer this week
|
||||
and nobody can explain why.
|
||||
|
||||
**Quantisation.** For llama.cpp, `q4_k_m` is the practical default: the quality
|
||||
difference from `q8_0` is smaller than most people expect and the memory
|
||||
difference is a factor of two. For vLLM on modern cards, FP8 or AWQ. Run
|
||||
`llmgw ctl plan <model>` to see whether a given precision fits before
|
||||
downloading anything.
|
||||
|
||||
**Disk.** Budget three to five times what you first estimate. Several
|
||||
quantisations, several versions, and partial downloads all accumulate.
|
||||
`llmgw ctl nodes` shows free space per node once `model_dirs` is configured,
|
||||
and `llmgw ctl weights` reclaims files llmgw downloaded and no model needs —
|
||||
never files you placed yourself.
|
||||
|
||||
Under ten nodes, let each keep its own copy; it is simpler and the disk is
|
||||
cheap. Above that, consider a shared filesystem — llmgw notices one and will
|
||||
not count the same free terabyte once per node.
|
||||
|
||||
## Before anyone starts using it
|
||||
|
||||
Set budgets and limits on day one, not later:
|
||||
|
||||
```
|
||||
llmgw keys limits --scope project --id <project> --rpm 60 --budget 50
|
||||
```
|
||||
|
||||
The reason is not technical. A resource that appears free and unlimited becomes,
|
||||
three months later, a cost nobody can account for — and introducing limits at
|
||||
that point is a much harder conversation than having had them from the start.
|
||||
+133
@@ -0,0 +1,133 @@
|
||||
# What llmgw stores, and what leaves the machine
|
||||
|
||||
Written for the review that happens before a gateway is allowed near anyone's
|
||||
documents. It describes what the software does. Mapping that onto an
|
||||
obligation — GDPR, a sector rule, an internal policy — is work for whoever
|
||||
holds the obligation; this page exists so that work starts from facts rather
|
||||
than from a vendor's summary.
|
||||
|
||||
Everything below is checkable in the source, and the places to look are named.
|
||||
|
||||
## The short version
|
||||
|
||||
By default llmgw stores **no prompt or completion text anywhere**. It stores
|
||||
who made a request, when, how many tokens, and what it cost.
|
||||
|
||||
Two features change that, and both are off until somebody turns them on for a
|
||||
named project, and neither can be turned on without also saying for how long.
|
||||
|
||||
Nothing on the serving path leaves the cluster.
|
||||
|
||||
## What is stored, exactly
|
||||
|
||||
**Usage records**, one per request: the key, the project, the model asked for,
|
||||
the model that answered, prompt and completion token counts, cost, HTTP status
|
||||
and latency. No message content. The row is in `internal/store/store.go` as
|
||||
`Usage` and there is no field for text in it.
|
||||
|
||||
**API keys**, as SHA-256 hashes. The key itself is shown once, at creation, and
|
||||
never stored — so a leaked database hands over no working credential, and
|
||||
neither llmgw nor an administrator can recover a key somebody lost. A prefix is
|
||||
kept so a key can be identified in a listing without revealing it.
|
||||
|
||||
**People and tenancy**: accounts, the organisation/team/project tree, and the
|
||||
limits attached to it. Passwords are hashed.
|
||||
|
||||
**Sessions** for the browser pages, as tokens.
|
||||
|
||||
**Conversations and agent notes** — only for a project that asked. See below.
|
||||
|
||||
## What is not stored
|
||||
|
||||
Prompts. Completions. Anything the model was shown or produced, unless the
|
||||
project turned on conversations or notes.
|
||||
|
||||
The chat page in llmgw's own web interface keeps the conversation in a
|
||||
JavaScript variable in the browser. Refresh it and the conversation is gone,
|
||||
because it was never anywhere else.
|
||||
|
||||
## Conversations and notes, when they are turned on
|
||||
|
||||
Two features keep text, and both work the same way:
|
||||
|
||||
- **Conversations** (`x-llmgw-thread`) let a client send only its new turn while
|
||||
llmgw replays the history — including to a different model than the one the
|
||||
conversation started on.
|
||||
- **Agent notes** (`/v1/memory`, and MCP tools) let an agent keep something for
|
||||
itself across conversations.
|
||||
|
||||
Both are **off for every project** until an operator runs `llmgw keys threads`
|
||||
or `llmgw keys memory` for that project by name. A request carrying a thread id
|
||||
against a project that has not is refused with an explanation, not quietly
|
||||
served without its history.
|
||||
|
||||
Both **require a retention window**. There is no way to turn either on without
|
||||
saying how many hours, because a store with no expiry is a transcript of
|
||||
everything a team ever typed, kept forever because nobody chose a number. The
|
||||
window is measured from the last turn, so an active conversation does not
|
||||
expire underneath it, and a background sweep removes what has passed it.
|
||||
|
||||
Both are **scoped to a project**, and the scope is part of every lookup rather
|
||||
than a check afterwards. Reading a conversation that belongs to another project
|
||||
returns the same answer as reading one that does not exist, so a leaked id
|
||||
cannot be used to discover what exists elsewhere.
|
||||
|
||||
`DELETE /v1/threads/{id}` removes a conversation and its messages immediately.
|
||||
It is a client endpoint, not an administrative one: the person who said the
|
||||
thing can unsay it without asking somebody else to read it first.
|
||||
|
||||
Turning retention off stops new storage, and existing conversations go at the
|
||||
next hourly sweep.
|
||||
|
||||
## What leaves the machine
|
||||
|
||||
**Nothing on the serving path.** A request arrives, is routed to an engine on
|
||||
your own hardware, and the answer comes back. The router talks to engines, the
|
||||
engines talk to nothing, the agents talk to the control plane. All of it is
|
||||
inside the cluster.
|
||||
|
||||
Two things can reach outside, both started by an operator rather than by a
|
||||
request:
|
||||
|
||||
**Fetching weights.** When a model's `source` is a URL, the agent downloads it.
|
||||
Point that at an internal mirror and nothing leaves at all — which is also the
|
||||
faster arrangement, because eight nodes pulling forty gigabytes from the
|
||||
internet is a day's work and from a local mirror it is minutes. A `source` that
|
||||
is a path on disk is never fetched.
|
||||
|
||||
**`llmgw ctl plan`.** Reads a model's architecture from huggingface.co to work
|
||||
out whether it fits your hardware. An operator command run at planning time,
|
||||
not part of serving, and not run by anything automatically.
|
||||
|
||||
There is no telemetry, no licence check, and no call home. The binary is
|
||||
static, built with `CGO_ENABLED=0`, and runs in a network with no route to the
|
||||
internet as long as the weights are already on disk.
|
||||
|
||||
## What the software cannot do for you
|
||||
|
||||
**The admin port serves pprof**, and pprof can dump the process heap — where
|
||||
API keys sit in plaintext while they are being checked. Bind it to loopback and
|
||||
reach it over SSH. The default does this; a deployment that changes it should
|
||||
know why.
|
||||
|
||||
**The store holds everything.** With conversations turned on it holds text as
|
||||
well as accounting, so where it lives, who can read it and how it is backed up
|
||||
are decisions with the same weight as the ones above. The PostgreSQL DSN
|
||||
belongs in a file readable only by the user llmgw runs as, which is why the
|
||||
configuration takes `@/path/to/file` rather than the address itself.
|
||||
|
||||
**Model weights are yours.** llmgw reclaims only files it downloaded, and it
|
||||
identifies those by a name it generated. Anything you put on a node it will not
|
||||
touch, whatever the catalog says.
|
||||
|
||||
## Where to look
|
||||
|
||||
| Claim | Where |
|
||||
|---|---|
|
||||
| Usage rows carry no text | `internal/store/store.go`, type `Usage` |
|
||||
| Keys are stored hashed | `internal/store/keys.go` |
|
||||
| Conversations are off by default | `internal/store/threads.go`, `ThreadRetention` |
|
||||
| Retention is required | `llmgw keys threads --help` |
|
||||
| Scope is part of the lookup | `internal/store/threads.go`, `Thread` |
|
||||
| Deletion is a client endpoint | `internal/router/threads.go`, `handleDeleteThread` |
|
||||
| Weight fetching is the only serving-path egress | `internal/agent/weights/fetch.go` |
|
||||
+210
@@ -0,0 +1,210 @@
|
||||
# InfiniBand fabrics and NVLink
|
||||
|
||||
For a site whose GPU nodes are joined by InfiniBand or RoCE, by NVLink, or both
|
||||
— several independent fabrics, NVL72 racks, eight-GPU HGX boards. It says what
|
||||
llmgw does with them, what it needs from you before it can, how to check a node
|
||||
in one command, and where to look when a replica spanning nodes will not run.
|
||||
|
||||
Read [install.md](install.md) first; this is the part of it that only matters
|
||||
once there is a fabric.
|
||||
|
||||
> **Where this stands.** Every piece here was written against the real formats
|
||||
> of `sminfo`, `saquery`, `ibnetdiscover`, `nvidia-smi topo -m` and
|
||||
> `nvidia-smi -q`, and tested against them — but none of it has yet run on a
|
||||
> site with the hardware. Run `llmgw doctor --json` on two or three nodes first
|
||||
> and keep the reports: if a parser meets a format it does not know, that report
|
||||
> is what fixes it, and the fix is a control plane upgrade, not one per node.
|
||||
|
||||
## What llmgw does with them
|
||||
|
||||
**Two kinds of interconnect.** InfiniBand or RoCE joins nodes (scale-out).
|
||||
NVLink joins GPUs — inside one machine on an HGX board, or across the nodes of a
|
||||
rack on an NVL72 or an NVLink Switch system (scale-up). A replica spanning nodes
|
||||
is placed, in order of preference:
|
||||
|
||||
1. inside one **multi-node NVLink domain**, if any has enough usable nodes:
|
||||
tensor parallelism between its nodes then runs over NVLink;
|
||||
2. on one **physical InfiniBand or RoCE fabric**, preferring nodes under one
|
||||
leaf switch and putting nodes with a degraded cable last.
|
||||
|
||||
It is never split across two fabrics or two NVLink domains. Within a node, the
|
||||
GPUs a shard gets are the set with the most NVLink pairs, then the nearest to
|
||||
the fabric's adapters, then on one NUMA node.
|
||||
|
||||
**A fabric has a name and an identity.** The name is what `agent.yaml` says; the
|
||||
identity is what the adapters report — on InfiniBand the subnet managers every
|
||||
port on one fabric sees, on RoCE the IP network of its RoCE v2 address. The name
|
||||
is how you address a fabric; the identity is how llmgw catches the name being
|
||||
wrong: two independent fabrics given one name are placed as two, and the
|
||||
cluster page says so.
|
||||
|
||||
## Before a node joins
|
||||
|
||||
On every GPU node:
|
||||
|
||||
| What | Why | Check |
|
||||
|---|---|---|
|
||||
| NVIDIA driver | everything | `nvidia-smi` |
|
||||
| `nvidia-fabricmanager`, on an NVSwitch machine | without it CUDA does not start on those GPUs | `systemctl is-active nvidia-fabricmanager` |
|
||||
| `rdma-core`, `ib_uverbs` loaded | user-space RDMA at all | `ls /sys/class/infiniband`, `lsmod \| grep ib_uverbs` |
|
||||
| `nvidia_peermem` loaded | GPUDirect RDMA; without it NCCL copies every transfer through host memory | `lsmod \| grep nvidia_peermem` |
|
||||
| `infiniband-diags` | tells llmgw which fabric a port is on; the agent installs it if missing | `which sminfo` |
|
||||
|
||||
On an NVL72 or NVLink Switch system, also:
|
||||
|
||||
| What | Why | Check |
|
||||
|---|---|---|
|
||||
| `nvidia-imex` running, every node of the domain in `/etc/nvidia-imex/nodes_config.cfg` | GPUs in different nodes share memory through it | `systemctl is-active nvidia-imex` |
|
||||
| an IMEX channel (`options nvidia NVreg_CreateImexChannel0=1`) | an engine opens one | `ls /dev/nvidia-caps-imex-channels` |
|
||||
| every GPU joined to its domain | the fabric manager does this | `nvidia-smi -q \| grep -A6 Fabric` shows `State: Completed` |
|
||||
|
||||
On each InfiniBand fabric:
|
||||
|
||||
- **A subnet manager, and a standby.** Without one, ports come up to `INIT` and
|
||||
stay there; without a standby, the fabric keeps its routes if the master
|
||||
stops but cannot recover from any change after. llmgw watches this but does
|
||||
not run one.
|
||||
- **Its own name** in `agent.yaml` on every node attached to it. Optionally, a
|
||||
distinct `subnet_prefix` in each subnet manager's configuration, which makes
|
||||
GIDs unique across fabrics; llmgw does not need it.
|
||||
|
||||
## Declaring the fabrics
|
||||
|
||||
In `/etc/llmgw/agent.yaml`, the adapters on each fabric, by name:
|
||||
|
||||
```yaml
|
||||
fabrics:
|
||||
- name: ib-compute
|
||||
hcas: [mlx5_0, mlx5_1, mlx5_2, mlx5_3, mlx5_6, mlx5_7, mlx5_8, mlx5_9]
|
||||
- name: ib-storage
|
||||
hcas: [mlx5_4]
|
||||
```
|
||||
|
||||
or `--fabric ib-compute:mlx5_0,mlx5_1,…` when onboarding ([node-onboarding.md](node-onboarding.md)).
|
||||
An adapter belongs to one fabric; a node may be on several. A replica on
|
||||
`ib-compute` is given only `ib-compute`'s adapters (`NCCL_IB_HCA`).
|
||||
|
||||
The rate is measured, not taken from the file: a fabric's rate on a node is the
|
||||
slowest of its active adapters, because a collective runs at the speed of its
|
||||
slowest link.
|
||||
|
||||
The agent checks the declaration against the adapters on every heartbeat. An
|
||||
adapter that is missing, or whose port is down, takes that fabric off the node
|
||||
within seconds — the node stops being offered multi-node work on it — and
|
||||
returns it when the port comes back.
|
||||
|
||||
## Checking a node: `llmgw doctor`
|
||||
|
||||
```
|
||||
sudo llmgw doctor # uses /etc/llmgw/agent.yaml when present
|
||||
sudo llmgw doctor --json > $(hostname)-doctor.json
|
||||
```
|
||||
|
||||
It reads the node and says what is wrong and how to fix it, most serious first,
|
||||
and exits non-zero when anything fails:
|
||||
|
||||
```
|
||||
FAIL nvlink this machine has NVSwitch but nvidia-fabricmanager is inactive: CUDA will not initialise on its GPUs
|
||||
→ systemctl enable --now nvidia-fabricmanager (its version must match the driver's)
|
||||
WARN fabric nvidia_peermem is not loaded: GPUDirect RDMA is off, and NCCL copies every transfer through host memory
|
||||
→ modprobe nvidia_peermem
|
||||
OK gpu 8 GPUs: 8× NVIDIA H100 80GB HBM3
|
||||
OK nvlink NVLink groups: [0 1 2 3 4 5 6 7]
|
||||
OK fabric mlx5_0/1 ACTIVE at 400 Gb/s (InfiniBand), subnet manager 0x248a07030030f0c0 (LID 1, SMINFO_MASTER)
|
||||
OK fabric subnet managers of 0x248a07030030f0c0's fabric: 0x248a07030030f0c0 (master, priority 15), 0x248a07030030f1d4 (standby, priority 14)
|
||||
OK agent fabric ib-compute: mlx5_0,mlx5_1,… usable at 400 Gb/s
|
||||
OK agent NCCL bootstraps over ens3 (10.1.0.21)
|
||||
```
|
||||
|
||||
followed by each GPU's NVLink peers and nearest adapter, each port's state,
|
||||
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;
|
||||
a table of subnet managers per fabric — master, priority, standbys, the last
|
||||
change of master — and every disagreement: one name on two fabrics, two names
|
||||
on one, an adapter up but undeclared, ports up with no subnet manager, a
|
||||
fabric with no standby, a recent failover.
|
||||
- **`/admin/fabrics`.** Each fabric's map, from `ibnetdiscover` run by one of
|
||||
its nodes every five minutes: switches in their tiers, adapters under their
|
||||
leaves, cables coloured by rate, degraded ones in red, rail alignment, and
|
||||
the cluster's own adapters live with their traffic. Click anything for its
|
||||
detail.
|
||||
- **`/admin/nvlink`.** Each multi-node NVLink domain as a rack of trays: which
|
||||
GPUs are joined, which node's IMEX is down, which link runs degraded.
|
||||
- **A node's page.** Its ports with state, rate, errors and what they are
|
||||
physically on; its GPU-by-GPU NVLink matrix; its load.
|
||||
- **Metrics**, for alerts:
|
||||
|
||||
| Metric | Alert when |
|
||||
|---|---|
|
||||
| `llmgw_ib_port_active` | a declared port is 0 |
|
||||
| `llmgw_ib_port_counter_total{counter="symbol_error"}` and `link_downed` | rising |
|
||||
| `llmgw_ib_port_rate_gbps` | below its usual rate |
|
||||
| `llmgw_fabric_subnet_managers{state="standby"}` | absent or 0 |
|
||||
| `llmgw_fabric_master_sm` | its `sm_guid` changes (a failover) |
|
||||
| `llmgw_fabric_problems{level="fail"}` | above 0 |
|
||||
| `llmgw_gpu_nvlink_joined` | 0 |
|
||||
| `llmgw_imex_up` | 0 |
|
||||
|
||||
## What an engine is given
|
||||
|
||||
For a shard on an InfiniBand or RoCE fabric:
|
||||
|
||||
| Variable | Value |
|
||||
|---|---|
|
||||
| `NCCL_IB_HCA` | that fabric's adapters on this node, and no others |
|
||||
| `NCCL_IB_DISABLE` | `0` (`1` for a single-node shard) |
|
||||
| `NCCL_SOCKET_IFNAME`, `GLOO_SOCKET_IFNAME` | the node's management interface: `mgmt_interface`, or the one holding its management address |
|
||||
| `NCCL_IB_GID_INDEX` | on RoCE only: the RoCE v2 entry with an IPv4 address |
|
||||
| `NCCL_MNNVL_ENABLE`, `NCCL_CUMEM_ENABLE` | `1`, for a replica inside a multi-node NVLink domain |
|
||||
|
||||
A profile's `env` overrides any of them. vLLM launches a multi-node replica
|
||||
with its own multiprocessing (`engines.vllm.launch_mode: mp`, the default):
|
||||
every node runs `vllm serve` with its rank and the leader's address, the
|
||||
workers headless. `launch_mode: ray` is there for a vLLM too old for that.
|
||||
|
||||
## When a replica spanning nodes will not run
|
||||
|
||||
| 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 | `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 |
|
||||
@@ -0,0 +1,303 @@
|
||||
# Facility-aware placement: rack power, temperature and cooling
|
||||
|
||||
**Status: step 0, reading the nodes' own sensors, is built; steps 1–5 are
|
||||
deferred** until the rack, PDU and power information of a real site is at
|
||||
hand — without it they would be designed against guesses.
|
||||
This is the design for step 5 of resource allocation, after reserved CPU cores
|
||||
and memory (1), capacity classes and preemption (2), tenant capacity quotas (3)
|
||||
and sizing by measured throughput (4), all of which shipped in v0.10.5. Of
|
||||
what follows, only the node sensors (rollout step 0) exist in the code: they
|
||||
are read, published and shown, and change no decision.
|
||||
|
||||
## Why
|
||||
|
||||
In a large installation the scarce thing is often not a GPU but the room it
|
||||
sits in. Three failures are invisible to placement today:
|
||||
|
||||
- **A rack goes over its power budget.** Every node in it reports healthy;
|
||||
placement sees free cards and fills them; the breaker, or the facility's
|
||||
power cap, decides the rest. A rack with A and B feeds is worse: it runs
|
||||
fine until one feed fails and the other cannot carry the load.
|
||||
- **Hot hardware is slow hardware.** GPUs throttle long before they fail. A
|
||||
replica placed in a rack whose inlet air is too warm serves fewer tokens per
|
||||
second than its measured capacity says, and nothing connects the two.
|
||||
- **Correlated failure.** Two replicas of a model on two nodes in the same rack
|
||||
or on the same PDU are one failure away from zero. Placement spreads replicas
|
||||
across nodes, not across racks.
|
||||
|
||||
The goal is for placement to treat facility limits the way it already treats
|
||||
memory and cores: reserve against a budget, refuse what does not fit, say why,
|
||||
and never let missing data stop the cluster.
|
||||
|
||||
## What it builds on
|
||||
|
||||
| Already in place | Used for |
|
||||
|---|---|
|
||||
| Node labels, reported by agents, matched by `node_selector` | Where a node is: its rack, row, PDU |
|
||||
| `inventory.Allocation` carries GPUs, memory and cores per shard | Adding watts, summed per rack and PDU |
|
||||
| Placement rejections (`explain-placement`, headroom's "why not") | Naming the rack that is full or hot |
|
||||
| Headroom: the scheduler dry-run on a copy of the inventory | Room that respects power budgets, for free |
|
||||
| Preemption with classes, minimums and a dry run (step 2) | Shedding load from a hot rack in the right order |
|
||||
| Measured throughput per profile (step 4) | Tokens per watt as a placement preference |
|
||||
| Node sensors, read by every agent (step 0, built) | Heat and power with no facility integration: see below |
|
||||
| `internal/promq`, a Prometheus client | Reading facility metrics that DCIM exporters publish |
|
||||
|
||||
## Concepts
|
||||
|
||||
**Location.** Each node has a place in a tree: site, room, row, rack. Power has
|
||||
its own tree beside it: a rack draws from one or more PDUs, and each PDU from a
|
||||
feed. Both are read from node labels with well-known keys:
|
||||
|
||||
```yaml
|
||||
# agent.yaml
|
||||
labels:
|
||||
llmgw.io/site: sto1
|
||||
llmgw.io/room: hall-b
|
||||
llmgw.io/row: "12"
|
||||
llmgw.io/rack: r12-07
|
||||
llmgw.io/pdu: r12-07-a,r12-07-b # both feeds, for a dual-fed rack
|
||||
```
|
||||
|
||||
A node without these labels is its own rack of one: placement behaves as it
|
||||
does today.
|
||||
|
||||
**Domain.** A rack, a row, a PDU, a room: anything with a budget or a
|
||||
condition. Domains are declared in a facility file (below) or discovered from
|
||||
node labels; a label naming a domain the file does not declare creates it with
|
||||
no budget and no condition.
|
||||
|
||||
**Budget.** A domain may have a power cap in watts. Placement reserves each
|
||||
replica's estimated draw against every domain its node belongs to and refuses a
|
||||
placement that would take any of them past its cap less a safety margin. For a
|
||||
dual-fed rack the cap that matters is what one feed can carry alone, since that
|
||||
is what the rack must survive.
|
||||
|
||||
**Condition.** A domain has a thermal and cooling condition: `ok`, `warm`, `hot`
|
||||
or `unknown`, derived from readings against thresholds with hysteresis.
|
||||
Conditions change what placement prefers and allows; they are not amounts to
|
||||
fill.
|
||||
|
||||
**Reading.** A measured value with a time: inlet temperature, power draw, a
|
||||
cooling unit's alarm state. Readings older than their staleness limit are
|
||||
treated as absent.
|
||||
|
||||
## Estimating a replica's draw
|
||||
|
||||
What placement reserves has to be known before the replica runs, so it is an
|
||||
estimate first and a measurement later:
|
||||
|
||||
1. **Declared.** A profile may say `power_watts`. The operator's figure wins.
|
||||
2. **Derived.** Otherwise GPU profiles take each card's power limit from NVML
|
||||
(the agent already reads the cards; the power limit is one more field) plus
|
||||
a per-node share for CPU, memory and fans; CPU profiles take their reserved
|
||||
cores times a per-core figure for the node's CPU model.
|
||||
3. **Measured.** Once replicas run, the agent reports their cards' actual draw
|
||||
(NVML) and, where the BMC allows, the node's (Redfish or IPMI DCMI). A
|
||||
domain's headroom is computed from the larger of what is reserved and what
|
||||
is measured, so an estimate that was too low is corrected by the room it
|
||||
stops offering, not by an outage.
|
||||
|
||||
## Decisions
|
||||
|
||||
### Power
|
||||
|
||||
- **Hard limit.** A placement that would take any domain past `cap × (1 −
|
||||
margin)` is rejected, with the domain and the figures: *rack r12-07: 14.2 of
|
||||
16.0 kW reserved, replica needs ~2.8, margin 10%*.
|
||||
- **Feeds.** For a domain with two feeds the cap is the smaller feed's rating
|
||||
times its derating (typically 80%), not the sum.
|
||||
- **Already over.** A domain found over its cap, because the estimate was low
|
||||
or the cap was lowered, gets no new placements, and an alert. It is not
|
||||
emptied automatically; shedding power is the facility's decision, and
|
||||
stopping inference to do it is ours only if asked (see shedding below).
|
||||
|
||||
### Temperature and cooling
|
||||
|
||||
Thresholds are per domain with defaults, and each condition needs to hold for a
|
||||
while before it changes, in both directions, so a single reading neither
|
||||
starts nor ends anything:
|
||||
|
||||
| Condition | Enters when | Placement |
|
||||
|---|---|---|
|
||||
| `ok` | below the warm threshold for 5 min | normal |
|
||||
| `warm` | inlet above the warm threshold for 2 min, GPU temperatures near throttling, or any GPU thermally throttled | allowed, scored down: other domains first |
|
||||
| `hot` | inlet above the hot threshold for 2 min, a cooling unit in alarm, or a sustained share of the domain's GPUs thermally throttled | no new placements, as a cordon would, but automatic and reversible |
|
||||
| `unknown` | no fresh reading | as `ok`, and reported as unknown |
|
||||
|
||||
**Shedding** (optional, off by default): a domain `hot` for longer than its
|
||||
shed delay has load moved off it using step 2's machinery. The lowest class goes
|
||||
first, and replicas above their model's minimum go before any at it. A critical
|
||||
minimum is never touched, and moves are rate limited to one replica per domain
|
||||
per interval. The replica is placed elsewhere first where there is room, so
|
||||
shedding moves work rather than dropping it. GPU power capping (`nvidia-smi
|
||||
-pl`) is a gentler alternative worth measuring against shedding: less
|
||||
throughput, no restart.
|
||||
|
||||
### Spreading
|
||||
|
||||
Replicas of one model are spread across racks, and across PDUs, the way they
|
||||
are spread across nodes today. This is a preference, dropped when it is the only
|
||||
thing standing in the way, as node spreading already is.
|
||||
|
||||
### Scoring
|
||||
|
||||
Placement today filters, then takes the first candidate in a fixed order
|
||||
(fewest free GPUs for GPU profiles, fewest replicas for CPU ones, nodes holding
|
||||
the weights first). Facility signals are better as a score than as more
|
||||
filters. Candidates that pass the hard limits are ordered by:
|
||||
|
||||
1. condition: `ok` before `warm`;
|
||||
2. power headroom left in the rack after this replica, more first, so load
|
||||
levels across racks;
|
||||
3. the existing order.
|
||||
|
||||
Once measurements include power, step 4's profile preference can use tokens per
|
||||
watt instead of tokens per GPU where both are known.
|
||||
|
||||
## Where the data comes from
|
||||
|
||||
A `FacilitySource` interface, implemented several ways, chosen in
|
||||
`server.yaml`:
|
||||
|
||||
```go
|
||||
type FacilitySource interface {
|
||||
// Domains is what the source knows: budgets, thresholds, members.
|
||||
Domains(ctx) ([]Domain, error)
|
||||
// Readings is the latest value of each signal, per domain.
|
||||
Readings(ctx) ([]Reading, error)
|
||||
}
|
||||
```
|
||||
|
||||
- **Static file.** Domains, caps and thresholds, written by hand or generated
|
||||
from a DCIM export. Enough for power budgets with no live data at all.
|
||||
- **Prometheus.** DCIM, PDU (SNMP exporter) and BMC (Redfish exporter)
|
||||
readings usually end up in Prometheus. The source runs configured instant
|
||||
queries and maps their labels onto domains:
|
||||
|
||||
```yaml
|
||||
facility:
|
||||
source: prometheus
|
||||
prometheus: http://monitoring:9090
|
||||
signals:
|
||||
inlet_temp_c:
|
||||
query: max by (rack) (dcim_rack_inlet_temperature_celsius)
|
||||
domain_label: rack
|
||||
stale_after: 2m
|
||||
power_watts:
|
||||
query: sum by (pdu) (pdu_outlet_power_watts)
|
||||
domain_label: pdu
|
||||
cooling_alarm:
|
||||
query: max by (row) (crac_alarm)
|
||||
domain_label: row
|
||||
```
|
||||
- **Agents (built, step 0).** What each node reads about itself, every ten
|
||||
seconds at most, sent with its heartbeat:
|
||||
- **GPUs, from nvidia-smi:** temperature, memory (HBM) temperature where the
|
||||
driver reports it, power draw, the enforced power limit, and what is
|
||||
holding the clocks down now — `sw_power_cap`, `hw_slowdown`, `sw_thermal`,
|
||||
`hw_thermal`, `power_brake`.
|
||||
- **CPUs, from hwmon** (what lm_sensors reads, without it installed): the
|
||||
package temperature, or the hottest core; and the board's inlet or ambient
|
||||
sensor where it has one (`SYSTIN`, `Inlet`, `Ambient`).
|
||||
- **CPU power, from RAPL:** the packages' energy counters over time,
|
||||
subdomains not counted twice, counter wrap handled.
|
||||
|
||||
A machine without a sensor reports nothing for it, never a zero; a VM has
|
||||
none of them. Published as `llmgw_gpu_temperature_celsius`,
|
||||
`llmgw_gpu_memory_temperature_celsius`, `llmgw_gpu_power_watts`,
|
||||
`llmgw_gpu_power_limit_watts`, `llmgw_gpu_throttled{reason}`,
|
||||
`llmgw_node_cpu_temperature_celsius`, `llmgw_node_cpu_power_watts` and
|
||||
`llmgw_node_inlet_temperature_celsius`, and on each node's page.
|
||||
|
||||
What they can and cannot say:
|
||||
- **Throttling is the strongest signal there is.** A temperature is a proxy;
|
||||
a card holding its clocks down is the fact that invalidates its measured
|
||||
capacity, and it needs no threshold to be chosen.
|
||||
- **Measured power** is GPU draw plus CPU packages; memory, disks, fans and
|
||||
power-supply losses are not in it, so a node's draw is that plus an
|
||||
estimated remainder. A rack's is its llmgw nodes' plus a declared figure
|
||||
for everything else in it: switches and storage report nothing here.
|
||||
- **CPU and GPU temperatures follow the load, not the room.** A hot rack is
|
||||
inferred from many nodes running hot at comparable load, or from
|
||||
throttling, not from one node's CPU. Only a board inlet sensor reads the
|
||||
air, and most boards have none.
|
||||
- **No budget comes from a sensor.** A reading is a draw; the cap it is
|
||||
measured against still comes from the facility.
|
||||
- **Simulated.** Readings from a script, for tests and for rehearsing a
|
||||
thermal event on a cluster without facility data.
|
||||
|
||||
## Safety
|
||||
|
||||
- **Missing data never blocks.** A source that is down, a reading that is stale,
|
||||
a node without location labels: placement behaves as it does today, and says
|
||||
the data is missing (a metric, a line on the cluster page). Facility
|
||||
awareness can make placement more careful, never impossible.
|
||||
- **Hysteresis and hold times** on every condition, so a noisy sensor does not
|
||||
flap a rack between allowed and forbidden.
|
||||
- **Observe-only first.** A mode where every decision is computed, logged and
|
||||
shown ("would have refused: rack r12-07 hot") but nothing is enforced, run
|
||||
for long enough to trust the thresholds before turning enforcement on.
|
||||
- **Manual override.** `llmgw ctl facility hold RACK` and `release RACK`, for the
|
||||
facility team to mark a domain unavailable, or to clear a condition they know
|
||||
to be a sensor fault.
|
||||
- **Shedding is opt-in**, per domain or cluster-wide, and rate limited.
|
||||
|
||||
## What people see
|
||||
|
||||
- `explain-placement` and headroom name the domain: *rack r12-07 is hot (inlet
|
||||
31°C for 6 min)*, *pdu r12-07-a: 7.1 of 7.4 kW*.
|
||||
- `llmgw ctl facility` lists domains with budget, reserved, measured and
|
||||
condition. The cluster page gets a rack summary beside the nodes table.
|
||||
- Metrics: `llmgw_domain_power_budget_watts`, `_reserved_watts`,
|
||||
`_measured_watts`, `llmgw_domain_condition{condition}`,
|
||||
`llmgw_domain_reading_age_seconds`; alerts for a domain near its cap, a
|
||||
domain hot, a domain over its cap, and facility data gone stale.
|
||||
|
||||
## Rollout
|
||||
|
||||
Each step ships on its own and is useful without the next:
|
||||
|
||||
0. **Node sensors, observe-only (built).** GPU heat, power and throttling,
|
||||
CPU heat and power, inlet air where the board has a sensor; published and
|
||||
shown, no decision changed.
|
||||
1. **Location and spreading.** Well-known labels, domains from labels, replicas
|
||||
spread across racks and PDUs, locations on the node views. No facility data
|
||||
needed.
|
||||
2. **Power budgets from a static file.** Caps and margins, estimated draw per
|
||||
replica (declared or derived), hard limit in placement and headroom.
|
||||
3. **Readings, observe-only.** The Prometheus source beside the agents',
|
||||
conditions with hysteresis (throttling first among them), measured power
|
||||
folded into headroom; decisions shown, not enforced.
|
||||
4. **Thermal enforcement.** `warm` scored down, `hot` refused, manual hold and
|
||||
release.
|
||||
5. **Shedding and tokens per watt.** Opt-in shedding through step 2's
|
||||
preemption; measured power in profile preference.
|
||||
|
||||
## Testing
|
||||
|
||||
- The simulated source drives unit tests of every condition transition, the
|
||||
hysteresis, staleness, and the placement outcomes, with mutation checks as
|
||||
for steps 1–4.
|
||||
- Property tests: whatever the readings, placement never exceeds a declared
|
||||
cap with fresh data, and never refuses everything because data is missing.
|
||||
- On the PoC cluster, with no facility data: step 1 with labelled VMs, step 2
|
||||
with a static file of made-up caps, steps 3–4 with the simulated source
|
||||
replaying a thermal event.
|
||||
- Before enforcement anywhere real: a period in observe-only mode against the
|
||||
site's actual data.
|
||||
|
||||
## Open questions
|
||||
|
||||
- **What the sites have.** Which DCIM, which PDUs, which BMCs, and whether any
|
||||
of it already reaches a Prometheus. The source design assumes Prometheus is
|
||||
the common ground; a site without it needs another source.
|
||||
- **Who owns the numbers.** Caps, deratings and thresholds belong to the
|
||||
facility team. Whether they edit llmgw's file, or llmgw reads theirs.
|
||||
- **Phase balance.** Three-phase PDUs have per-phase limits; whether placement
|
||||
should see phases or only the PDU total.
|
||||
- **Cooling granularity.** Cooling is often per row or per aisle, not per rack;
|
||||
the domain tree allows either, but the defaults should match what sites
|
||||
actually measure.
|
||||
- **Power capping instead of shedding.** Whether to lower GPU power limits in a
|
||||
warm rack rather than move replicas, and how that interacts with measured
|
||||
capacity from step 4, which would no longer hold at the lower limit.
|
||||
+1239
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,209 @@
|
||||
# Adding a node
|
||||
|
||||
Written for whoever runs the cluster.
|
||||
|
||||
## The short version
|
||||
|
||||
On the control plane:
|
||||
|
||||
```
|
||||
sudo llmgw node add ubuntu@10.0.0.145
|
||||
```
|
||||
|
||||
That is all a machine that already exists needs — bare metal, or a VM you
|
||||
built. [install.md](install.md#5-adding-the-node) says what it does. On
|
||||
Proxmox, the machine itself can be built first by the script below, on the
|
||||
hypervisor.
|
||||
|
||||
## The two halves
|
||||
|
||||
Two scripts, because the two halves of onboarding are done by different people
|
||||
with different access: one needs the hypervisor, the other needs the
|
||||
certificate authority.
|
||||
|
||||
## 1. The machine
|
||||
|
||||
On the Proxmox host:
|
||||
|
||||
```
|
||||
llmgw-provision-node --vmid 144 --name llmgw-node04 --ip 10.0.0.144/24
|
||||
```
|
||||
|
||||
Name it after the node id. **The VM name, the hostname, the node id and the
|
||||
certificate's common name are all the same string.** They do not have to be, and
|
||||
every time they are not, somebody spends an afternoon working out which of the
|
||||
four a message is talking about.
|
||||
|
||||
It imports a cloud image and lets cloud-init do the rest, which is why a node
|
||||
takes a minute rather than an afternoon. It refuses to touch an existing VMID:
|
||||
adopting a machine is not a decision a provisioning script should make.
|
||||
|
||||
Releases carry static `linux/amd64` and `linux/arm64` binaries. Download one and
|
||||
pass it with `--binary` rather than copying whatever happens to be on the
|
||||
machine you are running onboarding from — otherwise whoever runs the script
|
||||
quietly becomes the build server, and the fleet ends up running whatever was in
|
||||
their working tree.
|
||||
|
||||
## 2. The identity
|
||||
|
||||
On the control plane, which holds the certificate authority — this is the
|
||||
part that cannot be delegated, because it issues the credential by which the
|
||||
control plane believes the node is who it says it is:
|
||||
|
||||
```
|
||||
sudo llmgw node add ilker@10.0.0.144 --fabric ib-b:mlx5_0,mlx5_1
|
||||
```
|
||||
|
||||
It checks the machine with `llmgw doctor`, issues a client certificate whose
|
||||
common name is the node id (the hostname), installs it along with the CA root,
|
||||
trusts that root system-wide, installs the binary and the unit, starts the
|
||||
agent, and waits for the node to join.
|
||||
|
||||
If it does not join, the agent's log on the node says why. A certificate whose
|
||||
common name does not match the claimed node id is refused by design — that
|
||||
check is what makes a node's identity something it proves rather than
|
||||
something it picks.
|
||||
|
||||
**It needs SSH from the control plane to the node.** Provisioning copies the
|
||||
hypervisor's `authorized_keys` into the new machine; with `ssh -A` to the
|
||||
control plane and `sudo --preserve-env=SSH_AUTH_SOCK`, the same key that
|
||||
reaches the hypervisor reaches the node.
|
||||
|
||||
`deploy/bootstrap/onboard-node.sh` is the script this replaced, for a machine
|
||||
holding an authority but not a current llmgw.
|
||||
|
||||
## The CPU model matters, and its default is wrong
|
||||
|
||||
Nodes are provisioned with `--cpu host`. Proxmox's default, `kvm64`, hides
|
||||
AVX, AVX2, FMA and even SSE4.2 from the guest, and llama.cpp then loads its
|
||||
scalar backend instead of the one built for the host's instruction set.
|
||||
Measured on node05, an 8-core VM on a Xeon E5-2680 v4, same weights, same
|
||||
threads, the only change being the CPU model:
|
||||
|
||||
| Qwen2.5-Coder | kvm64 tok/s | host tok/s |
|
||||
| -------------- | ----------: | ---------: |
|
||||
| 1.5B q4_K_M | 5.80 | 39.09 |
|
||||
| 3B q4_K_M | 2.90 | 20.51 |
|
||||
| 7B q4_K_M | 1.33 | 9.88 |
|
||||
|
||||
Prompt processing is worse still: 7.6 tokens a second against 143.8, a factor
|
||||
of nineteen.
|
||||
|
||||
It is worth calling out because nothing looks broken. The node registers, the
|
||||
engine starts, replies are correct — every number is simply small, and the
|
||||
obvious conclusion is that CPU inference is hopeless rather than that the VM
|
||||
cannot multiply efficiently.
|
||||
|
||||
The trade is live migration: a VM pinned to the host CPU cannot move to a
|
||||
machine with a different one. For a node whose job is arithmetic that is the
|
||||
right way round, and `--cpu` overrides it where it is not.
|
||||
|
||||
## Engines
|
||||
|
||||
The script installs none, and does not need to: once the agent is connected,
|
||||
the control plane tells it which engines to have and it installs them itself
|
||||
— llama.cpp on every node, vLLM on every node with an NVIDIA GPU it can run
|
||||
on, SGLang where the control plane names a version. Each is checked to run
|
||||
before the node offers it, so a node never advertises an engine that cannot
|
||||
start. [install.md](install.md#engines) says what vLLM needs of the node.
|
||||
|
||||
`--llamacpp-bin`, `--vllm-bin` and `--sglang-python` are for an engine you
|
||||
installed yourself. For the agent to use yours rather than install its own,
|
||||
also set that engine's version to `""` under `engines` in `server.yaml`.
|
||||
|
||||
Weights are the other thing a node needs. llmgw fetches them for a model whose
|
||||
`source` is a URL, one copy per node, which covers the simple case; anything
|
||||
larger — a shared filesystem, a cache — stays the operator's arrangement, and
|
||||
llmgw is content to find the files already there.
|
||||
|
||||
The `sim` engine needs neither, which is what makes it useful for exercising
|
||||
everything else.
|
||||
|
||||
## Checking an engine after an upgrade
|
||||
|
||||
Engines rename and remove flags between versions, and when that happens the
|
||||
failure arrives minutes into a model load, in the engine's own error message, on
|
||||
a node nobody was watching.
|
||||
|
||||
```
|
||||
llmgw ctl engine-check vllm
|
||||
llmgw ctl engine-check llamacpp --help-output "$(llama-server --help 2>&1)"
|
||||
```
|
||||
|
||||
It builds the widest command llmgw would ever produce — a multi-node one, since
|
||||
a flag that only appears there is exactly the one nobody notices has gone — and
|
||||
checks every flag against the engine's own `--help`. No GPU required: a CLI
|
||||
prints its help on any machine.
|
||||
|
||||
Some programs do not give up their whole help for `--help`. vLLM prints a
|
||||
summary of section names and expects `--help=all`; checking against the summary
|
||||
reports every flag as missing, confidently and wrongly. `engine-check` knows
|
||||
that one. If you add an engine, check that what you get back is the full list.
|
||||
|
||||
vLLM also refuses to do anything at all on a machine with no accelerator,
|
||||
including print its help, unless `VLLM_TARGET_DEVICE` says what to assume.
|
||||
`engine-check` sets it.
|
||||
|
||||
Give it the **whole** help. Truncating llama.cpp's 725 lines to 400 reports five
|
||||
flags as missing that the running cluster uses; the command warns when the text
|
||||
it was given does not document `--help` itself, which is the cheapest sign it
|
||||
was cut short.
|
||||
|
||||
## Fabrics
|
||||
|
||||
`--fabric` is how a node says which InfiniBand fabric it is attached to. The
|
||||
name is checked against what the adapters report — every port on one fabric
|
||||
sees the same subnet managers — so two independent fabrics given one name are
|
||||
placed as two, and the cluster page says so; give each fabric its own name.
|
||||
The agent installs `infiniband-diags`, which that check needs, if it is missing.
|
||||
|
||||
Before a node with a fabric joins, run `sudo llmgw doctor` on it: it checks the
|
||||
subnet manager, GPUDirect RDMA, the fabric manager and IMEX on NVLink machines,
|
||||
and the declared fabrics against the adapters. [fabrics.md](fabrics.md) has the
|
||||
rest.
|
||||
|
||||
Omit it for a node with no fabric. Such a node is its own island: it can host
|
||||
single-node replicas and will never be given a multi-node one.
|
||||
|
||||
Two adapters per fabric is the default; `--fabric ib-b:mlx5_2,mlx5_3` names them
|
||||
explicitly. An adapter must never appear under two fabrics — the agent refuses
|
||||
that configuration: a replica given adapters on two fabrics starts, passes its
|
||||
health checks, and hangs on its first all-reduce.
|
||||
|
||||
## Where the driver is
|
||||
|
||||
```
|
||||
--nvidia-smi /opt/nvidia/bin/nvidia-smi
|
||||
--ib-root /sys/class/infiniband
|
||||
```
|
||||
|
||||
Both default to the usual places. Containerized driver installs routinely put
|
||||
`nvidia-smi` somewhere else, and an agent that cannot find it registers with no
|
||||
GPUs — which looks exactly like a node that has none.
|
||||
|
||||
## Removing a node
|
||||
|
||||
```
|
||||
llmgw ctl drain llmgw-node04 # move its replicas off, one at a time
|
||||
llmgw ctl nodes # wait until it holds none
|
||||
```
|
||||
|
||||
Then stop the agent. Draining brings each replacement up before stopping the
|
||||
original, which is the difference between a rolling move and an outage.
|
||||
|
||||
The node's certificate stays valid until it expires. There is no revocation list
|
||||
— if a machine is decommissioned rather than retired, reissue the authority or
|
||||
accept that a certificate naming a node that no longer exists can still connect
|
||||
and claim to be it.
|
||||
|
||||
## Testing without hardware
|
||||
|
||||
Every part of this works on machines with no GPUs and no InfiniBand. Point
|
||||
`--nvidia-smi` at a script that prints the CSV discovery asks for, point
|
||||
`--ib-root` at a sysfs-shaped directory, and use `engine: sim` in the catalog.
|
||||
|
||||
The simulated engine rendezvouses for real: rank 0 binds the dist-init port and
|
||||
does not serve until every worker has joined, and workers serve nothing. That
|
||||
exercises the multi-node launch sequence, the fabric constraint, and failure and
|
||||
rebuild — which is most of what goes wrong — without exercising NCCL, which is
|
||||
the rest.
|
||||
@@ -0,0 +1,443 @@
|
||||
# Onboarding a user
|
||||
|
||||
Written for whoever operates an llmgw deployment: the person with a shell on the
|
||||
control plane machine.
|
||||
|
||||
There are two ways a person gets a key. Pick one and put it in the config.
|
||||
|
||||
| | Self-service | Operator-issued |
|
||||
|---|---|---|
|
||||
| Config | `auth.signup_open: true` | `auth.signup_open: false` |
|
||||
| They do | sign up, wait, then make their own keys | wait for a link |
|
||||
| You do | approve once | issue every key |
|
||||
| Good for | more than a handful of people | a few, or machines |
|
||||
|
||||
Self-service does not mean open. A signup produces an account that can do
|
||||
nothing at all; **approval is what makes it work**, and approval needs a shell
|
||||
here rather than a session cookie, because it is the step that decides who gets
|
||||
to spend the cluster.
|
||||
|
||||
Everything below runs as the user the server runs as, so the files it creates
|
||||
stay readable by it:
|
||||
|
||||
```
|
||||
sudo -u llmgw llmgw <command> --store /var/lib/llmgw/llmgw.db
|
||||
```
|
||||
|
||||
## Self-service
|
||||
|
||||
Turn on the form:
|
||||
|
||||
```yaml
|
||||
auth:
|
||||
enabled: true
|
||||
store: /var/lib/llmgw/llmgw.db
|
||||
signup_open: true
|
||||
```
|
||||
|
||||
The landing page then links to it, and the rest is:
|
||||
|
||||
```
|
||||
llmgw users list --pending --store …
|
||||
ID EMAIL STATE ADMIN MEMBER OF SIGNED UP
|
||||
usr-k4mq… ayse@example.com pending - on their own 4m12s ago
|
||||
|
||||
llmgw users approve ayse@example.com --by ilker --store …
|
||||
llmgw users approve ayse@example.com --project Acme/ml/training --store … # or straight into a project
|
||||
```
|
||||
|
||||
That is the whole operator loop. From then on they sign in at `/login`, and
|
||||
their account page (`/profile`) is where they create keys, see when each was
|
||||
last used, revoke the ones they are done with, and see their own use over the
|
||||
last thirty days against every budget that applies to them. A key is shown once
|
||||
on the page that creates it and never again.
|
||||
|
||||
### Connecting an editor
|
||||
|
||||
The account page's **Connect your editor** makes a key for Continue, Codex or
|
||||
Cline and shows what the tool needs with that key, this gateway's address and a
|
||||
model already in it:
|
||||
|
||||
- **Continue** — an assistant file of its own, `~/.continue/assistants/llmgw.yaml`,
|
||||
so their `config.yaml` is not touched.
|
||||
- **Codex** — a provider and a profile in `~/.codex/config.toml`, run with
|
||||
`codex --profile llmgw`; their default model stays as it was. Codex speaks the
|
||||
Responses API, which the gateway translates at `/v1/responses`.
|
||||
- **Cline** — four settings to paste into its own screen, since it keeps them
|
||||
in VS Code's secret storage where no file reaches.
|
||||
|
||||
For Continue and Codex there is a one-line command,
|
||||
`curl -fsSL <gateway>/setup/<tool>?model=… | sh`. Its script is served without
|
||||
signing in and holds no secret — the key comes from the environment — and it
|
||||
keeps a copy of anything it replaces. Revoking the key disconnects the tool.
|
||||
|
||||
How well an editor's agent mode works depends on the model: it needs one that
|
||||
calls tools well, which the coder models do and small general ones often do not.
|
||||
|
||||
**This gateway sends no email.** There is no "your account was approved" message
|
||||
and no password reset, so tell people what to expect: they sign up, they tell
|
||||
you, you approve, they sign in. If that loop is too manual for your size, it is
|
||||
the thing to automate — not the approval step.
|
||||
|
||||
To close an account:
|
||||
|
||||
```
|
||||
llmgw users suspend usr-k4mq… --store …
|
||||
```
|
||||
|
||||
which closes their sessions immediately. Their **API keys keep working** — those
|
||||
are separate credentials, and revoking them is a separate step:
|
||||
|
||||
```
|
||||
llmgw keys list --project default --store …
|
||||
llmgw keys revoke key-… --store …
|
||||
```
|
||||
|
||||
That separation is deliberate: a suspended person's running jobs should not all
|
||||
fail at once without someone deciding they should.
|
||||
|
||||
## Operator-issued
|
||||
|
||||
With `signup_open: false` there is no form, and you issue keys yourself. This is
|
||||
also the path for things that are not people — CI, a notebook server, a script.
|
||||
|
||||
### Decide where the key belongs
|
||||
|
||||
Keys hang off a project, projects off a team, teams off an organization. Limits
|
||||
can be set at any of those levels and compose downward by taking the **most
|
||||
restrictive** value — a key cannot raise what its team set, and a team cannot
|
||||
raise what its organization set. That is the only rule under which a budget
|
||||
means anything.
|
||||
|
||||
If you have no structure yet, one of each is fine:
|
||||
|
||||
```
|
||||
llmgw org add org mynodes
|
||||
llmgw org add team default --in mynodes
|
||||
llmgw org add project default --in mynodes/default
|
||||
llmgw org tree --people
|
||||
```
|
||||
|
||||
Everything that takes an organization, team or project — here, in `users` and
|
||||
in `keys` — accepts its id, its path (`mynodes/default/default`) or its name
|
||||
when the name is unique, so an id never has to be copied off a page. When there
|
||||
is only one project, `--project` can be left out.
|
||||
|
||||
The tree is not fixed once made:
|
||||
|
||||
```
|
||||
llmgw org rename project Acme/platform/gateway api-gateway
|
||||
llmgw org move project Acme/platform/gateway --to Scania/ml
|
||||
llmgw org remove team Acme/old # only when nothing is in it
|
||||
llmgw users delete ayse@example.com --yes # signs out, revokes their keys
|
||||
llmgw keys list --all # every key, where it spends, and whose
|
||||
```
|
||||
|
||||
A project that ever had a key cannot be removed, because usage history refers
|
||||
to it — rename it instead.
|
||||
|
||||
### Where people belong
|
||||
|
||||
A person can be a member of any number of organizations, teams and projects, or
|
||||
of none:
|
||||
|
||||
```
|
||||
llmgw users join ayse@example.com --team Acme/ml # every project in ml
|
||||
llmgw users join ayse@example.com --project Scania/research/vision
|
||||
llmgw users leave ayse@example.com --project Acme/ml/training
|
||||
```
|
||||
|
||||
A member of a team or organization may use every project in it. Leaving revokes
|
||||
the keys the person made in projects they can no longer reach; a project they
|
||||
still reach through another membership keeps them.
|
||||
|
||||
Someone who is a member of nothing still uses the gateway, on their own:
|
||||
personal keys, and the chat page, within their personal limits. When they make
|
||||
a key on `/profile` they choose what it spends against — one of their projects,
|
||||
or "Personal". An operator can do the same:
|
||||
|
||||
```
|
||||
llmgw keys create laptop --for ayse@example.com # personal
|
||||
llmgw keys create ci --for ayse@example.com --project Acme/ml/training
|
||||
```
|
||||
|
||||
Personal use can be limited by default, or turned off so people use the gateway
|
||||
only through their projects:
|
||||
|
||||
```yaml
|
||||
auth:
|
||||
personal:
|
||||
rpm: 20
|
||||
tpm: 50000
|
||||
budget: 10 # whole units
|
||||
period: month
|
||||
# disabled: true
|
||||
```
|
||||
|
||||
A personal key has no project, so it keeps no conversations or agent notes;
|
||||
those are kept per project.
|
||||
|
||||
### Organization managers
|
||||
|
||||
A member of an organization can be made its manager:
|
||||
|
||||
```
|
||||
llmgw users join boss@acme.com --org Acme --role manager
|
||||
```
|
||||
|
||||
or with the star on their organization membership on `/admin/users`. A manager
|
||||
gets an **Organization** link and a page at `/org`: usage over time by team,
|
||||
project or person, the teams and projects with their limits and spend, and every
|
||||
member. Only what was spent in that organization's projects is on it; what its
|
||||
people do elsewhere is not. The page is read-only — a manager is not an
|
||||
administrator of the cluster, and limits and memberships stay with those who
|
||||
are.
|
||||
|
||||
`llmgw keys init` still works for provisioning scripts and is idempotent, but it
|
||||
will no longer move a team or project that already exists under another parent;
|
||||
it says so and stops.
|
||||
|
||||
Split it when you want separate limits or separate reporting, not before.
|
||||
|
||||
### Issue the key as a one-time link
|
||||
|
||||
```
|
||||
llmgw keys create --store … \
|
||||
--project default \
|
||||
--band gold \
|
||||
--base-url https://llmgw.example \
|
||||
--claim \
|
||||
"ayse-laptop"
|
||||
|
||||
https://llmgw.example/claim/3mKq…
|
||||
```
|
||||
|
||||
Send the link, not the key.
|
||||
|
||||
A key pasted into a chat window stays in that chat window forever, searchable by
|
||||
everyone who joins the channel later. A link works once and then expires, so the
|
||||
worst case is someone opening it before the intended person — which you find out
|
||||
about, because they will tell you the link is dead.
|
||||
|
||||
The name is for you, not the user: it appears in `keys list` and is how you will
|
||||
know which key to revoke.
|
||||
|
||||
**Why the link is safe to store.** The key is sealed with the link's own token,
|
||||
and the token is never written down — only its hash is. A database leak yields
|
||||
two useless hashes and a blob nothing in the database can open.
|
||||
|
||||
Without `--claim` the key is printed to stdout instead, for when you are
|
||||
provisioning something that is not a person:
|
||||
|
||||
```
|
||||
llmgw keys create --store … --project ci "github-runner" > /run/secrets/llmgw-key
|
||||
```
|
||||
|
||||
## Limits
|
||||
|
||||
Limits are optional. A key with none is unlimited, which is the right default on
|
||||
a private cluster and the wrong one the moment the gateway faces anything else.
|
||||
|
||||
```
|
||||
# Everyone in this project
|
||||
llmgw keys limits --store … --scope project --id default \
|
||||
--rpm 60 --tpm 200000
|
||||
|
||||
# One key, tighter
|
||||
llmgw keys limits --store … --scope key --id key-mpmip77twpt7qaxc \
|
||||
--rpm 10
|
||||
|
||||
# Restrict which models
|
||||
llmgw keys limits --store … --scope project --id default \
|
||||
--models qwen-0.5b,llama-3.3-70b
|
||||
```
|
||||
|
||||
An **absent** model list permits everything. An **empty** one permits nothing —
|
||||
`--models ""` is a way to suspend a project without revoking its keys.
|
||||
|
||||
Allowlists compose by intersection: a key cannot add a model its team withheld.
|
||||
A model's canonical name and any of its aliases name the same thing, so either
|
||||
form works.
|
||||
|
||||
Budgets are given in whole units, as the quotas page shows them, and what a unit
|
||||
means is yours to decide. Self-hosted models have no per-token price, so it can
|
||||
be GPU-seconds, a notional rate, or one unit per thousand tokens everywhere —
|
||||
a price of 1 000 000 in `auth.pricing`, which is in micro-units — and a budget
|
||||
becomes a plain quota of thousands of tokens.
|
||||
Models are priced in the server's `auth.pricing` section; a model priced at zero
|
||||
is exempt from budget checks, which is what lets a key with nothing left still be
|
||||
served a cheap fallback.
|
||||
|
||||
```
|
||||
llmgw keys limits --store … --scope project --id Acme/platform/gateway \
|
||||
--budget 50 --period month # 50 units a month
|
||||
```
|
||||
|
||||
Two more kinds of limit follow people rather than the tree:
|
||||
|
||||
```
|
||||
# A person, everywhere: every key they hold, every project, the chat page
|
||||
llmgw keys limits --scope user --id ayse@example.com --budget 25
|
||||
|
||||
# Each member of a project separately — a fair share, not one shared pool
|
||||
llmgw keys limits --scope project --id Acme/ml/training --per-member --rpm 20
|
||||
```
|
||||
|
||||
A request is held to all of them at once: the key's own, its project's, team's
|
||||
and organization's, each per-member limit on those for its owner, and the
|
||||
owner's own. The strictest wins, and each is counted in its own pool.
|
||||
|
||||
## The landing page
|
||||
|
||||
Someone who reaches the gateway without a key sees the landing page. Tell it who
|
||||
to send them to, or it can only shrug:
|
||||
|
||||
```yaml
|
||||
landing:
|
||||
contact: "platform@example.com" # or an https:// URL, or "#platform on Slack"
|
||||
notice: "Internal use only."
|
||||
```
|
||||
|
||||
An address becomes a mailto link, an `http(s)` URL stays a link, anything else is
|
||||
shown as text. Other schemes are never linked — configuration is not always
|
||||
written by whoever runs the gateway, and a `javascript:` href on a page served to
|
||||
every visitor is a short path to a stored XSS.
|
||||
|
||||
## What the user sees
|
||||
|
||||
They open the link, press the button, and get the key once. The page shows them
|
||||
the base URL and a working `curl`, and points at two endpoints:
|
||||
|
||||
- `GET /v1/models` — what this key may use
|
||||
- `GET /me` — its limits and what it has spent
|
||||
|
||||
`/me` is there so nobody has to ask you "why was I rate limited" or "how much
|
||||
have I used". It is outside `/v1` because it is not part of the OpenAI API.
|
||||
|
||||
## Rotation and revocation
|
||||
|
||||
```
|
||||
llmgw keys list --store … --project default
|
||||
```
|
||||
|
||||
`LAST USED` is written in batches rather than per request, so it lags by up to a
|
||||
minute. It is the column that tells you whether a key is still in use, which is
|
||||
what makes rotation safe:
|
||||
|
||||
**To rotate**, issue the new key, let both work while the user switches, watch
|
||||
`LAST USED` on the old one stop moving, then revoke it.
|
||||
|
||||
```
|
||||
llmgw keys revoke --store … key-mpmip77twpt7qaxc
|
||||
```
|
||||
|
||||
Revocation is a soft delete: the key stops working, the row stays, and its usage
|
||||
history keeps its subject. A running server keeps accepting a revoked key until
|
||||
its cache expires — `auth.cache_ttl`, 30 seconds by default. That is the price of
|
||||
keeping the database off the request path, and it is the number to lower if you
|
||||
need revocation to bite faster.
|
||||
|
||||
## Passwords
|
||||
|
||||
People change their own from `/profile`. It asks for the current one — a session
|
||||
is what someone gets from a borrowed laptop, and the password is what they do not
|
||||
have — and it signs every other browser out, which is the point if they are
|
||||
changing it because something went wrong.
|
||||
|
||||
There is no self-service reset, because this gateway sends no email and a reset
|
||||
link with nowhere to send it is not a reset. When someone loses theirs:
|
||||
|
||||
```
|
||||
llmgw users reset-password usr-k4mq… --password "a temporary passphrase" --store …
|
||||
```
|
||||
|
||||
Tell them out of band and have them change it from their profile. Their sessions
|
||||
are closed — if they lost the password, the reason might be that someone else has
|
||||
it. Their API keys keep working; those are separate credentials.
|
||||
|
||||
## Trying a model in the browser
|
||||
|
||||
`/chat` is a playground for anyone signed in. It exists so somebody can see
|
||||
whether a model works without setting up a client, and — more often — so they
|
||||
can tell whether a problem is the gateway or their code.
|
||||
|
||||
**The browser never holds a key.** The page carries a session; the server turns
|
||||
that into the same identity an API key would produce, resolved from the person's
|
||||
project, and the request goes down the ordinary path: admission, model
|
||||
allowlist, rate limits, budget, accounting. Anything else would make the page a
|
||||
hole in every limit the rest of the gateway enforces.
|
||||
|
||||
Usage from the page is attributed to `session:<user-id>` rather than a key, so
|
||||
it is billed to the right project and visibly did not come from a credential
|
||||
somebody holds. One turn is capped at 512 tokens: a playground should not be the
|
||||
thing that ties up a replica for ten minutes, and anyone who wants a longer
|
||||
answer has an API key for that.
|
||||
|
||||
Nothing is stored. The conversation lives in the page and is gone when it is
|
||||
closed.
|
||||
|
||||
## Seeing the cluster
|
||||
|
||||
The dashboard — nodes, replicas, models, why something did not get placed, usage
|
||||
per key — is served twice, and the difference matters.
|
||||
|
||||
On the **admin listener** (`127.0.0.1:9090` by default) it needs no
|
||||
authentication: having a shell on the machine is the credential.
|
||||
|
||||
```
|
||||
ssh -L 9090:127.0.0.1:9090 <control-plane> # then http://127.0.0.1:9090/
|
||||
```
|
||||
|
||||
That listener also serves pprof, which can dump the process heap — where API
|
||||
keys sit in plaintext. **Never expose that port.** It is bound to loopback for
|
||||
that reason.
|
||||
|
||||
On the **client-facing port** the same page is at `/admin`, behind a session and
|
||||
an admin flag, so it can be reached from a browser without the port carrying
|
||||
pprof along with it:
|
||||
|
||||
```
|
||||
llmgw users promote usr-k4mq… --store …
|
||||
```
|
||||
|
||||
Someone without the flag gets a 404 rather than a 403, because 403 confirms the
|
||||
page exists.
|
||||
|
||||
## Managing accounts from the browser
|
||||
|
||||
`/admin/users` is the tree: organizations, teams and projects, each project
|
||||
opening onto its members and its keys. From there an admin approves people
|
||||
(into a project or on their own), adds them to and takes them out of
|
||||
organizations, teams and projects, suspends and removes them, revokes keys, and
|
||||
adds, renames, moves and removes organizations, teams and projects. `/admin/quotas`
|
||||
sets limits on the tree, per member of any part of it, and per person. A searchable list of every account is at the
|
||||
bottom. The page works as plain forms; with scripts on, a change updates it in
|
||||
place and the projects you had open stay open.
|
||||
|
||||
Three things are deliberately not there:
|
||||
|
||||
| | Why |
|
||||
|---|---|
|
||||
| `llmgw users promote` | An admin who can make admins turns a stolen session into permanent access |
|
||||
| `llmgw users reset-password` | Setting someone's password is taking over their account |
|
||||
| `llmgw keys create` | Issuing a credential on another person's behalf |
|
||||
|
||||
Each of those turns a stolen admin cookie into more than it was, and each is
|
||||
rare enough that a shell is no hardship. Cordon, drain and everything else that
|
||||
touches the cluster are also CLI-only for the same reason: the dashboard shows
|
||||
the cluster, it does not steer it.
|
||||
|
||||
Every administrative action is logged with who did it, including the ones that
|
||||
were refused.
|
||||
|
||||
## What this does not do
|
||||
|
||||
- **No email.** No approval notice, no self-service password reset, no address
|
||||
verification. The loop is manual on purpose at this size; if it stops being
|
||||
the right size, that is the thing to automate.
|
||||
- **No SSO.** Accounts are local, with local passwords.
|
||||
- **No roles.** There is one flag, `is_admin`, and it grants looking at the
|
||||
dashboard. Everything else an administrator does is the CLI, and having a
|
||||
shell here is what makes someone one.
|
||||
- **Rate limits are per gateway instance.** With two instances each counts its own,
|
||||
so the effective limit doubles. v1 is single-instance by design.
|
||||
@@ -0,0 +1,202 @@
|
||||
# Quick install
|
||||
|
||||
From empty servers to a model answering requests, in six steps. This page
|
||||
has only the commands and what to check after each; [install.md](install.md)
|
||||
explains every choice.
|
||||
|
||||
## What you need
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| **One control plane machine** | Linux, x86-64 or ARM64. Debian 12, Ubuntu 22.04+ or RHEL 10. No GPU needed; 2 GB RAM, 20 GB disk. |
|
||||
| **GPU nodes** (optional to start) | The same Linux, with the NVIDIA driver installed. vLLM needs driver 580 or newer and a Turing GPU or newer (not V100). |
|
||||
| **root** | on the control plane, and `sudo` without a password on every node. |
|
||||
| **SSH** | from the control plane to each node, as a user with passwordless `sudo` there. |
|
||||
| **Network** | nodes reach the control plane on 7443 and 8081; users reach it on 8080. The machines need internet access to download engines and models. |
|
||||
|
||||
You do not install llama.cpp, vLLM, Python or CUDA. The nodes do that
|
||||
themselves once they join.
|
||||
|
||||
## 1. Put llmgw on the control plane
|
||||
|
||||
```
|
||||
VERSION=$(curl -fsSL https://gitea.mynodes.xyz/api/v1/repos/ilker/llmgw/releases/latest \
|
||||
| grep -o '"tag_name":"[^"]*"' | cut -d'"' -f4)
|
||||
U=https://gitea.mynodes.xyz/ilker/llmgw/releases/download/$VERSION
|
||||
curl -fsSLO $U/llmgw-$VERSION-linux-amd64 # -linux-arm64 on ARM
|
||||
curl -fsSL $U/checksums.txt | sha256sum -c --ignore-missing
|
||||
sudo install -m0755 llmgw-$VERSION-linux-amd64 /usr/local/bin/llmgw
|
||||
llmgw --version
|
||||
```
|
||||
|
||||
`sha256sum` must print `OK`. If it does not, the download is damaged; do not
|
||||
install it.
|
||||
|
||||
## 2. Set up the control plane
|
||||
|
||||
```
|
||||
sudo llmgw server init
|
||||
```
|
||||
|
||||
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.**
|
||||
- a model mirror, which downloads model files once for the whole cluster;
|
||||
- this machine as a first node, and a small model (Qwen2.5 0.5B) on it.
|
||||
|
||||
If nodes will reach this machine by a name or address it does not know itself
|
||||
(a DNS name, a NAT address), add `--host` for each:
|
||||
|
||||
```
|
||||
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:
|
||||
|
||||
```
|
||||
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.
|
||||
|
||||
It logs in over SSH and, while you watch:
|
||||
|
||||
1. installs llmgw on the node, the release for its processor, checksum checked;
|
||||
2. runs `llmgw doctor` there, and stops if the node has a problem. Each problem
|
||||
comes with the command that fixes it. Fix it and run the same command again;
|
||||
3. gives the node its certificate and configuration, and starts its agent;
|
||||
4. waits for the node to join and install llama.cpp.
|
||||
|
||||
The node's name in the cluster is its hostname. vLLM installs by itself after
|
||||
that, in a few more minutes on a GPU node (about 8 GB). The cluster page
|
||||
(`/admin`) shows its progress, or why the node cannot run it.
|
||||
|
||||
**InfiniBand.** When the node's adapters are up, `node add` shows them grouped
|
||||
by the fabric they are on and asks for each fabric's name. Give every node on
|
||||
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:
|
||||
|
||||
```
|
||||
sudo --preserve-env=SSH_AUTH_SOCK llmgw node add ubuntu@10.0.0.21
|
||||
```
|
||||
|
||||
**Check:** `llmgw ctl nodes` lists the node. Running `node add` again on the
|
||||
same node is safe: it renews its certificate and configuration, and keeps what
|
||||
runs there.
|
||||
|
||||
A machine retired for good is removed with `llmgw node remove <name>` once its
|
||||
agent is stopped.
|
||||
|
||||
## 4. Add a model
|
||||
|
||||
Open `/admin/mirror` and choose **Explore models**. It lists what is popular on
|
||||
Hugging Face, showing only what your nodes can run.
|
||||
|
||||
1. Click a model to read its card.
|
||||
2. Pick one of its files, usually a `q4_k_m` for a first try, and click
|
||||
**Fetch**. It is downloaded into the mirror once, for every node.
|
||||
3. When the download finishes, click **serve** on its row in the table above.
|
||||
4. In the form, choose how it runs: **always on**, or **on demand**, which
|
||||
loads it on the first request and stops it after 15 minutes unused. Click
|
||||
**Add to the catalog**.
|
||||
|
||||
The model is placed on a node that fits it, and its status on `/admin` turns
|
||||
`ready`.
|
||||
|
||||
**Check:** the model appears in the list on `/chat`. Ask it something.
|
||||
|
||||
## 5. Give people access
|
||||
|
||||
**A key for a program:** each person makes their own on the account page
|
||||
(`/profile`). The key is shown once.
|
||||
|
||||
```
|
||||
curl http://<control-plane>:8080/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "<model-name>", "messages": [{"role": "user", "content": "hello"}]}'
|
||||
```
|
||||
|
||||
Any OpenAI client works. Only the base URL changes:
|
||||
`http://<control-plane>:8080/v1`.
|
||||
|
||||
**An editor:** the account page connects Continue, Codex or Cline. Pick the
|
||||
tool and a model, and it gives a one-line command that writes the
|
||||
configuration.
|
||||
|
||||
**Other people's accounts:** they sign up at `/signup`, and you approve them on
|
||||
the people page (`/admin/users`).
|
||||
|
||||
## 6. Put it behind HTTPS
|
||||
|
||||
Port 8080 is plain HTTP. Before users outside a trusted network use it, put a
|
||||
reverse proxy in front with a real certificate. With Caddy, which gets the
|
||||
certificate by itself:
|
||||
|
||||
```
|
||||
sudo apt-get install -y caddy # Debian and Ubuntu
|
||||
curl -fsSL $U/Caddyfile | sed 's/llm.example.com/YOUR.DOMAIN/' | sudo tee /etc/caddy/Caddyfile
|
||||
sudo systemctl reload caddy
|
||||
```
|
||||
|
||||
Keep 7443 (nodes) and 8081 (mirror) on the internal network. They never go
|
||||
through the proxy. There is an nginx configuration too:
|
||||
[deploy/proxy/](../deploy/proxy/).
|
||||
|
||||
## When something does not work
|
||||
|
||||
| What you see | What to do |
|
||||
|---|---|
|
||||
| `node add` stops at `join` | `sudo journalctl -u llmgw-agent -n 50` on the node. Usually it cannot reach port 7443 on the control plane. |
|
||||
| `node add` says the address is not in the certificate | Run `sudo llmgw server init --host <that address>`, then `node add` again. |
|
||||
| A node says it cannot run vLLM | The reason is on `/admin`: the driver is older than 580, the GPU is older than Turing, or there is less than 25 GB free. |
|
||||
| A model shows `not placed` | `llmgw ctl explain-placement <model>` names each node and why it was not used. |
|
||||
| A reply stops after 60 seconds behind nginx | Use the nginx configuration from the release. Its timeouts are set for slow models. |
|
||||
| Anything else about a node | `sudo llmgw doctor` on it. `--json` writes everything it read, to send to whoever helps you. |
|
||||
|
||||
## Next
|
||||
|
||||
- [install.md](install.md): the same steps with every option explained.
|
||||
- [fabrics.md](fabrics.md): InfiniBand and NVLink.
|
||||
- [onboarding.md](onboarding.md): accounts, projects, budgets and keys.
|
||||
@@ -0,0 +1,375 @@
|
||||
# Quickstart: one machine, one model, ten minutes
|
||||
|
||||
Written for a developer who wants to see llmgw answer a request on a single
|
||||
workstation, built from source, with no systemd and no `sudo`. **To install
|
||||
llmgw on servers, read [quick-install.md](quick-install.md) instead.**
|
||||
|
||||
What you end up with: an OpenAI-compatible endpoint on `localhost:8080` in front
|
||||
of a real model served by llama.cpp.
|
||||
|
||||
## Two ways to run
|
||||
|
||||
Steps 1 to 4 use the simple one, **router-only**: you start the engines yourself
|
||||
and tell the router where they are with `--endpoint`. No certificates, no node
|
||||
agent.
|
||||
|
||||
Step 5 uses the **control plane and node agent**, the same architecture as a real
|
||||
cluster on one machine: the agent starts and supervises the engine, and the
|
||||
control plane places the model.
|
||||
|
||||
## 1. Build
|
||||
|
||||
```
|
||||
git clone https://gitea.mynodes.xyz/ilker/llmgw && cd llmgw
|
||||
make build
|
||||
make test
|
||||
./bin/llmgw --version
|
||||
```
|
||||
|
||||
`make build` writes `./bin/llmgw`. Go 1.26 or later is the only requirement.
|
||||
|
||||
## 2. Check the path with a fake engine
|
||||
|
||||
Before involving a GPU, prove the router works. `sim-engine` is an
|
||||
OpenAI-compatible fake that streams placeholder words.
|
||||
|
||||
```
|
||||
./bin/llmgw sim-engine --listen :8001 --model llama &
|
||||
./bin/llmgw server --endpoint e1=http://127.0.0.1:8001=llama &
|
||||
|
||||
curl -N localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-H 'x-session-id: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}'
|
||||
```
|
||||
|
||||
Chunks of `alpha beta gamma ...` ending in `data: [DONE]` mean the router is
|
||||
fine. Stop both before continuing (see [Stopping](#stopping)).
|
||||
|
||||
## 3. Serve a real model
|
||||
|
||||
Two processes, two ports. Clients only ever talk to the router; the router
|
||||
knows where the engine is because you tell it with `--endpoint`:
|
||||
|
||||
```
|
||||
curl :8080 -> llmgw router :8080 -> llama serve :8082 -> model
|
||||
(client) (the gateway) (the engine)
|
||||
```
|
||||
|
||||
Requests go to 8080. Port 8082 is the engine's own, and you only touch it to
|
||||
check that the engine is up.
|
||||
|
||||
llmgw does not install engines or weights. You need:
|
||||
|
||||
- a `llama-server` build from llama.cpp (CUDA build if you want the GPU), and
|
||||
- a `.gguf` file.
|
||||
|
||||
Check the card first with `nvidia-smi`. An 8 GB card holds a 1B model at bf16
|
||||
(2.4 GB) or a 9B model at `q4_k_m` (5.4 GB).
|
||||
|
||||
Start the engine on its own port. The model name given to `--alias` is the name
|
||||
clients will use:
|
||||
|
||||
```
|
||||
llama serve \
|
||||
-m /path/to/Llama-3.2-1B-Instruct-BF16.gguf \
|
||||
--host 0.0.0.0 --port 8082 -ngl 99 --alias llama
|
||||
```
|
||||
|
||||
`-ngl 99` offloads every layer to the GPU. Confirm it is up and knows the model:
|
||||
|
||||
```
|
||||
curl -s localhost:8082/v1/models
|
||||
```
|
||||
|
||||
Then point the router at it. The format is `id=URL=model`:
|
||||
|
||||
```
|
||||
./bin/llmgw server --endpoint e1=http://127.0.0.1:8082=llama
|
||||
```
|
||||
|
||||
## 4. Ask it something
|
||||
|
||||
Streaming:
|
||||
|
||||
```
|
||||
curl -N localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-H 'x-session-id: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}'
|
||||
```
|
||||
|
||||
One reply at once:
|
||||
|
||||
```
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"max_tokens":100}'
|
||||
```
|
||||
|
||||
Also useful: `curl -s localhost:8080/v1/models` and
|
||||
`curl -s localhost:8080/health`. `llmgw ctl` does not work in this mode: it
|
||||
talks to the control plane's admin API, which router-only mode does not have,
|
||||
so `ctl status` answers `404`.
|
||||
|
||||
On the machine this was written on, a 1B bf16 model on an RTX 4070 laptop GPU
|
||||
answered at roughly 85 tokens per second.
|
||||
|
||||
## Several replicas of one model
|
||||
|
||||
Give the router more than one `--endpoint` with the same model name and it
|
||||
spreads requests across them. `x-session-id` keeps a conversation on the same
|
||||
replica so its cache is reused:
|
||||
|
||||
```
|
||||
./bin/llmgw server \
|
||||
--endpoint e1=http://127.0.0.1:8082=llama \
|
||||
--endpoint e2=http://127.0.0.1:8083=llama
|
||||
```
|
||||
|
||||
## Stopping
|
||||
|
||||
```
|
||||
pkill -f "llmgw server"
|
||||
pkill -f "llmgw sim-engine"
|
||||
```
|
||||
|
||||
Stop the engine the way you started it (Ctrl-C in its terminal).
|
||||
|
||||
## What goes wrong
|
||||
|
||||
**`listening on :8080: bind: address already in use`.** An earlier router is
|
||||
still running. `ss -ltnp | grep 8080` names the process; stop it first. This is
|
||||
the usual result of moving from step 2 to step 3 without stopping the fake
|
||||
engine's router.
|
||||
|
||||
**`/v1/models` on the engine returns an empty list.** The engine was started
|
||||
without `-m`, so it is a router with nothing to load. Restart it with a model.
|
||||
|
||||
**`error: nothing to serve`.** `llmgw server` with no `--endpoint` and no model
|
||||
carrying `profiles` has nothing to do. In router-only mode, always pass at least
|
||||
one `--endpoint`.
|
||||
|
||||
**Requests fail with a model-not-found error.** The name in the request must
|
||||
match the name after the last `=` in `--endpoint`, and it must be a name the
|
||||
engine actually serves (its `--alias`).
|
||||
|
||||
**Answers are poor.** A 1B model is a plumbing test, not a chat assistant. Swap
|
||||
in a larger `.gguf`; nothing else changes.
|
||||
|
||||
## 5. Let llmgw run the engine (control plane and node agent)
|
||||
|
||||
Steps 2 and 3 start the engine by hand. Here a **node agent** starts and
|
||||
supervises llama.cpp for you, and a control plane decides where the model runs.
|
||||
Still one machine, still no `sudo`. Stop everything from steps 2 and 3 first:
|
||||
the GPU has to be free and port 8080 unused.
|
||||
|
||||
```
|
||||
curl :8080 -> llmgw server (control plane + router) <- llmgw agent (this node)
|
||||
|
|
||||
starts llama-server :20000
|
||||
```
|
||||
|
||||
The agent dials the control plane over mutual TLS, so this needs a small private
|
||||
certificate authority. Everything goes in a scratch directory, `devcluster/`,
|
||||
which is in `.gitignore`.
|
||||
|
||||
### Certificates
|
||||
|
||||
```
|
||||
mkdir -p devcluster/pki
|
||||
./bin/llmgw ctl cert init-ca --dir devcluster/pki
|
||||
./bin/llmgw ctl cert issue-server llmgw-ctl --host 127.0.0.1 --host localhost \
|
||||
--dir devcluster/pki --out devcluster
|
||||
./bin/llmgw ctl cert issue-client llmgw-node01 --dir devcluster/pki --out devcluster
|
||||
```
|
||||
|
||||
The name given to `issue-client` becomes the certificate's common name and must
|
||||
equal `node_id` in the agent config. The control plane refuses an agent whose
|
||||
claimed id differs from its certificate.
|
||||
|
||||
### The engine wrapper
|
||||
|
||||
The agent starts the engine as `<bin> --model ... --port ...`. A `llama` binary
|
||||
that wants `llama serve ...` refuses every flag, so put a wrapper in front of it:
|
||||
|
||||
```
|
||||
cat > devcluster/llama-server <<'EOF'
|
||||
#!/bin/sh
|
||||
exec /path/to/llama serve "$@"
|
||||
EOF
|
||||
chmod +x devcluster/llama-server
|
||||
|
||||
./bin/llmgw ctl engine-check llamacpp --bin devcluster/llama-server
|
||||
```
|
||||
|
||||
`engine-check` should end with `every flag llmgw builds is accepted`. If your
|
||||
build already ships a `llama-server` binary, point `bin` at that and skip the
|
||||
wrapper.
|
||||
|
||||
### Control plane config: `devcluster/server.yaml`
|
||||
|
||||
```yaml
|
||||
listen: "127.0.0.1:8080"
|
||||
admin_listen: "127.0.0.1:9090"
|
||||
agent_listen: "127.0.0.1:7443"
|
||||
|
||||
agent_tls:
|
||||
ca: devcluster/pki/ca.pem
|
||||
cert: devcluster/server.pem
|
||||
key: devcluster/server-key.pem
|
||||
|
||||
models:
|
||||
- name: llama
|
||||
source: /path/to/Llama-3.2-1B-Instruct-BF16.gguf
|
||||
routing_profile: chat
|
||||
context_length: 8192
|
||||
scaling: { min: 1, max: 1 }
|
||||
profiles:
|
||||
- name: gpu
|
||||
engine: llamacpp
|
||||
gpu_class: 4070-8g
|
||||
weight: 1
|
||||
priority: 100
|
||||
startup_timeout: 5m
|
||||
```
|
||||
|
||||
Three things in there are easy to get wrong:
|
||||
|
||||
- **The model must have at least one profile.** That is what switches the control
|
||||
plane on. With `models: []` the server starts as a plain router and exits with
|
||||
`nothing to serve`.
|
||||
- **`gpu_class` decides CPU or GPU.** Without it the replica is placed in the CPU
|
||||
pool and runs on the CPU. Get the class name from `./bin/llmgw agent inventory`
|
||||
(the `CLASS` column).
|
||||
- **There is no `auth:` block.** Leave it out for a local trial; with auth on,
|
||||
requests need an API key.
|
||||
|
||||
### Agent config: `devcluster/agent.yaml`
|
||||
|
||||
```yaml
|
||||
node_id: llmgw-node01
|
||||
server: "https://127.0.0.1:7443"
|
||||
mgmt_addr: "127.0.0.1"
|
||||
|
||||
tls:
|
||||
ca: devcluster/pki/ca.pem
|
||||
cert: devcluster/llmgw-node01.pem
|
||||
key: devcluster/llmgw-node01-key.pem
|
||||
|
||||
engines:
|
||||
llamacpp:
|
||||
bin: /absolute/path/to/devcluster/llama-server
|
||||
|
||||
supervisor:
|
||||
work_dir: devcluster/agent-work
|
||||
```
|
||||
|
||||
### Run it
|
||||
|
||||
Two terminals:
|
||||
|
||||
```
|
||||
./bin/llmgw server --config devcluster/server.yaml
|
||||
./bin/llmgw agent --config devcluster/agent.yaml
|
||||
```
|
||||
|
||||
The model is not placed at once. The control plane waits about 35 seconds for the
|
||||
fleet to settle after it starts, then places the replica, and the agent then
|
||||
loads the weights. Watch it with:
|
||||
|
||||
```
|
||||
./bin/llmgw ctl nodes
|
||||
./bin/llmgw ctl models
|
||||
./bin/llmgw ctl replicas
|
||||
./bin/llmgw ctl explain-placement llama
|
||||
```
|
||||
|
||||
`STATUS ready` in `ctl models` means it is serving. `nvidia-smi` should now show
|
||||
the model in GPU memory. The same curl from step 4 works against port 8080.
|
||||
|
||||
`ctl` works here, unlike in router-only mode, because there is a control plane
|
||||
for it to talk to.
|
||||
|
||||
Stop it with Ctrl-C in each terminal. The agent leaves the engine running when
|
||||
the agent itself stops, so if you restart the agent and the GPU stays full,
|
||||
`pgrep -af "llama serve"` will show the engine still up.
|
||||
|
||||
## 6. Memory: conversations and notes
|
||||
|
||||
Off unless you ask. llmgw otherwise stores nothing anyone typed. On a single
|
||||
machine with no API keys, one block in the server config turns on two things:
|
||||
|
||||
```yaml
|
||||
threads:
|
||||
store: devcluster/data/threads.db
|
||||
retain_hours: 24
|
||||
```
|
||||
|
||||
Restart the server after adding it. `retain_hours` is required: everything kept
|
||||
expires that long after it was last written.
|
||||
|
||||
Everything kept this way sits in one shared pool with no owner, and anyone who
|
||||
can reach the port and knows an id can read it. That is fine on one machine and
|
||||
wrong on a shared network. With auth on, threads and notes are per project
|
||||
instead and the block is ignored; see [install.md](install.md).
|
||||
|
||||
### Conversations (threads)
|
||||
|
||||
Send the same `x-llmgw-thread` id and only the new turn. llmgw puts the stored
|
||||
history in front of it:
|
||||
|
||||
```
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' -H 'x-llmgw-thread: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"My name is Deniz."}]}'
|
||||
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' -H 'x-llmgw-thread: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"What is my name?"}]}'
|
||||
|
||||
curl -X DELETE localhost:8080/v1/threads/demo # forget it
|
||||
```
|
||||
|
||||
Without the block, a request carrying `x-llmgw-thread` is refused with
|
||||
`this gateway keeps no conversations` rather than answered without its history.
|
||||
|
||||
### Notes an agent keeps for itself
|
||||
|
||||
Notes outlive any one conversation. Each is a key, some text and optional tags,
|
||||
kept per agent; name the agent with `x-llmgw-agent` (default `default`), so two
|
||||
agents using the same key do not read each other's notes.
|
||||
|
||||
```
|
||||
curl -X PUT localhost:8080/v1/memory/favorite-color \
|
||||
-H 'content-type: application/json' \
|
||||
-d '{"text":"The user likes dark green.","tags":["preference"]}'
|
||||
|
||||
curl 'localhost:8080/v1/memory?q=green' # search by words
|
||||
curl 'localhost:8080/v1/memory?tag=preference' # or by tag
|
||||
curl localhost:8080/v1/memory/favorite-color
|
||||
curl -X DELETE localhost:8080/v1/memory/favorite-color
|
||||
```
|
||||
|
||||
The same notes are offered to an agent as tools over MCP at
|
||||
`http://localhost:8080/mcp`: `remember`, `recall` and `forget`. Add that URL to
|
||||
your agent's MCP server list (HTTP transport; how you set request headers such
|
||||
as `x-llmgw-agent` depends on the client). llmgw does not put these tools into
|
||||
chat requests on its own, so an agent that is not given the server never sees
|
||||
memory.
|
||||
|
||||
Limits: a key is up to 200 bytes, a note's text up to 16 KiB, at most 16 tags,
|
||||
1000 notes per agent. Tags are lowercase letters, digits, `-` and `_`.
|
||||
|
||||
## Where this stops
|
||||
|
||||
Router-only mode (steps 2 to 4) has no authentication, no placement and no model
|
||||
management: you start and restart engines yourself. Step 5 adds placement, model
|
||||
management and a node agent, and can add authentication as `install.md`
|
||||
describes. For a real cluster, with several nodes, systemd and a persistent
|
||||
certificate authority, follow [install.md](install.md).
|
||||
|
||||
One caveat found while writing this page: `install.md` recommends starting with
|
||||
an empty `models: []` and adding models later with `ctl model add`. The control
|
||||
plane only turns on when the config declares a model with `profiles`, so that
|
||||
first start does not produce a control plane.
|
||||
+1555
File diff suppressed because it is too large.
Load diff
@@ -3,30 +3,83 @@ module gitea.mynodes.xyz/ilker/llmgw
|
||||
go 1.26.3
|
||||
|
||||
require (
|
||||
connectrpc.com/connect v1.21.0
|
||||
github.com/NVIDIA/go-nvml v0.13.4-0
|
||||
github.com/cespare/xxhash/v2 v2.3.0
|
||||
github.com/charmbracelet/bubbletea v1.3.10
|
||||
github.com/charmbracelet/lipgloss v1.1.0
|
||||
github.com/jackc/pgx/v5 v5.11.0
|
||||
github.com/knadh/koanf/parsers/yaml v1.1.1
|
||||
github.com/knadh/koanf/providers/env/v2 v2.0.1
|
||||
github.com/knadh/koanf/providers/file v1.2.1
|
||||
github.com/knadh/koanf/v2 v2.3.6
|
||||
github.com/mattn/go-isatty v0.0.24
|
||||
github.com/prometheus/client_golang v1.24.1
|
||||
github.com/prometheus/client_model v0.6.3
|
||||
github.com/prometheus/common v0.71.0
|
||||
github.com/spf13/cobra v1.10.2
|
||||
github.com/yuin/goldmark v1.8.6
|
||||
go.opentelemetry.io/otel v1.46.0
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.46.0
|
||||
go.opentelemetry.io/otel/sdk v1.46.0
|
||||
go.opentelemetry.io/otel/trace v1.46.0
|
||||
golang.org/x/crypto v0.57.0
|
||||
golang.org/x/net v0.59.0
|
||||
golang.org/x/sync v0.23.0
|
||||
google.golang.org/protobuf v1.36.12
|
||||
modernc.org/sqlite v1.59.0
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
|
||||
github.com/beorn7/perks v1.0.1 // indirect
|
||||
github.com/cenkalti/backoff/v5 v5.0.3 // indirect
|
||||
github.com/charmbracelet/colorprofile v0.4.1 // indirect
|
||||
github.com/charmbracelet/x/ansi v0.11.6 // indirect
|
||||
github.com/charmbracelet/x/cellbuf v0.0.15 // indirect
|
||||
github.com/charmbracelet/x/term v0.2.2 // indirect
|
||||
github.com/clipperhouse/displaywidth v0.9.0 // indirect
|
||||
github.com/clipperhouse/stringish v0.1.1 // indirect
|
||||
github.com/clipperhouse/uax29/v2 v2.5.0 // indirect
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
|
||||
github.com/fsnotify/fsnotify v1.9.0 // indirect
|
||||
github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
|
||||
github.com/go-logr/logr v1.4.4 // indirect
|
||||
github.com/go-logr/stdr v1.2.2 // indirect
|
||||
github.com/go-viper/mapstructure/v2 v2.5.0 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/knadh/koanf/maps v0.1.2 // indirect
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
|
||||
github.com/mattn/go-localereader v0.0.1 // indirect
|
||||
github.com/mattn/go-runewidth v0.0.19 // indirect
|
||||
github.com/mitchellh/copystructure v1.2.0 // indirect
|
||||
github.com/mitchellh/reflectwalk v1.0.2 // indirect
|
||||
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
|
||||
github.com/muesli/cancelreader v0.2.2 // indirect
|
||||
github.com/muesli/termenv v0.16.0 // indirect
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
|
||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||
github.com/prometheus/procfs v0.21.1 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/spf13/pflag v1.0.9 // indirect
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.46.0 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.46.0 // indirect
|
||||
go.opentelemetry.io/proto/otlp v1.11.0 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
golang.org/x/sync v0.23.0 // indirect
|
||||
golang.org/x/sys v0.47.0 // indirect
|
||||
google.golang.org/protobuf v1.36.12 // indirect
|
||||
golang.org/x/sys v0.48.0 // indirect
|
||||
golang.org/x/text v0.42.0 // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20260819154853-08b0e4226688 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260819154853-08b0e4226688 // indirect
|
||||
google.golang.org/grpc v1.83.1 // indirect
|
||||
modernc.org/libc v1.75.7 // indirect
|
||||
modernc.org/mathutil v1.7.1 // indirect
|
||||
modernc.org/memory v1.12.1 // indirect
|
||||
)
|
||||
@@ -1,16 +1,70 @@
|
||||
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/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/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM=
|
||||
github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=
|
||||
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/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=
|
||||
github.com/charmbracelet/colorprofile v0.4.1/go.mod h1:U1d9Dljmdf9DLegaJ0nGZNJvoXAhayhmidOdcBwAvKk=
|
||||
github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
|
||||
github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
|
||||
github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8=
|
||||
github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ=
|
||||
github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI=
|
||||
github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q=
|
||||
github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk=
|
||||
github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI=
|
||||
github.com/clipperhouse/displaywidth v0.9.0 h1:Qb4KOhYwRiN3viMv1v/3cTBlz3AcAZX3+y9OLhMtAtA=
|
||||
github.com/clipperhouse/displaywidth v0.9.0/go.mod h1:aCAAqTlh4GIVkhQnJpbL0T/WfcrJXHcj8C0yjYcjOZA=
|
||||
github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfatpWHKCs=
|
||||
github.com/clipperhouse/stringish v0.1.1/go.mod h1:v/WhFtE1q0ovMta2+m+UbpZ+2/HEXNWYXQgCt4hdOzA=
|
||||
github.com/clipperhouse/uax29/v2 v2.5.0 h1:x7T0T4eTHDONxFJsL94uKNKPHrclyFI0lm7+w94cO8U=
|
||||
github.com/clipperhouse/uax29/v2 v2.5.0/go.mod h1:Wn1g7MK6OoeDT0vL+Q0SQLDz/KpfsVRgg6W7ihQeh4g=
|
||||
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
|
||||
github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k=
|
||||
github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
|
||||
github.com/go-viper/mapstructure/v2 v2.4.0 h1:EBsztssimR/CONLSZZ04E8qAkxNYq4Qp9LvH92wZUgs=
|
||||
github.com/go-viper/mapstructure/v2 v2.4.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
|
||||
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
|
||||
github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8=
|
||||
github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
|
||||
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
|
||||
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
|
||||
github.com/go-viper/mapstructure/v2 v2.5.0 h1:vM5IJoUAy3d7zRSVtIwQgBj7BiWtMPfmPEgAXnvj1Ro=
|
||||
github.com/go-viper/mapstructure/v2 v2.5.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
|
||||
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
|
||||
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3 h1:LMLX+LgTNWpfvCBdFebv6EsYotImrt/Ppc5cXIriCSo=
|
||||
github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3/go.mod h1:jl5iWTm0/hd5PjEYEOuwAJ57L/CibdZfrqZ5XA5GrCk=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0 h1:/Tnpcb2E0Pz/tN9s3bfEY2Q8ePCEX9iuS+cneUwncnw=
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0/go.mod h1:zOBXOsUaBSjKgmH4OGzV1esUpR3oUSCPYVd2cUBjKYY=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
|
||||
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.11.0 h1:IzBBtyK9AHqf98cctWFifYSci2hgQR/cd56wB4p+ogg=
|
||||
github.com/jackc/pgx/v5 v5.11.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk=
|
||||
github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
|
||||
github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo=
|
||||
@@ -25,12 +79,29 @@ github.com/knadh/koanf/v2 v2.3.6 h1:JoQPSJmvS4aP0xNc8xMDr5tcrkSEInL23/Il7pITAKo=
|
||||
github.com/knadh/koanf/v2 v2.3.6/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28=
|
||||
github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
|
||||
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
|
||||
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
|
||||
github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
|
||||
github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4=
|
||||
github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88=
|
||||
github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byFGLdw=
|
||||
github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
|
||||
github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw=
|
||||
github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s=
|
||||
github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ=
|
||||
github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw=
|
||||
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
|
||||
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
|
||||
github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
|
||||
github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
|
||||
github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc=
|
||||
github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU=
|
||||
github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE=
|
||||
github.com/prometheus/client_model v0.6.3 h1:O0jaTVAYNxTHYInEPFJt5I3+sN8zqBtVMPTB1qyxiEo=
|
||||
@@ -39,13 +110,42 @@ github.com/prometheus/common v0.71.0 h1:9KDAKb7Mj3HEVKyFCK6Dc/HIwlBzZIN2l7/lrHl3
|
||||
github.com/prometheus/common v0.71.0/go.mod h1:CLJ5H8TEsGX8bl31BdMkfhIZ+QmZ9tBPPotUxUbfcmk=
|
||||
github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI=
|
||||
github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
|
||||
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
|
||||
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
|
||||
github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
|
||||
github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
|
||||
github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY=
|
||||
github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE=
|
||||
github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg=
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
|
||||
github.com/yuin/goldmark v1.8.6 h1:d0VcaP1sx9GkFVkoW+KtggpGi2KZ965i14b0+bDQST4=
|
||||
github.com/yuin/goldmark v1.8.6/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
|
||||
go.opentelemetry.io/otel v1.46.0 h1:FHt5/CDyVxi/8IM1CH7VE/rRgq3kLHa2mSTVMO8AWyc=
|
||||
go.opentelemetry.io/otel v1.46.0/go.mod h1:Gj3SEScelsNC45tp4nSxRYlS+f5iez7W8XPMCt905kE=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.46.0 h1:OFnwLJr+pF3iHrlGSzbxyuo6/6HyBlnlN1CWEJmBVcw=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.46.0/go.mod h1:716wFneO0ov19A2beH5hjfh9AK5z/VWNAtDijp1Y0/g=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.46.0 h1:KrC1YrQeSt46ITMWAbgQx1M1eV1/1TKzttrBzymPmss=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.46.0/go.mod h1:zDSEzoEqsOrgBeGvH66KRgxh90VonFyJqBHA0Pk3+rM=
|
||||
go.opentelemetry.io/otel/metric v1.46.0 h1:yBnkXvgV7AXFILZc5K6IZe/CBFF3OS7BJ8ov6/lj0K8=
|
||||
go.opentelemetry.io/otel/metric v1.46.0/go.mod h1:iPmdWqifKUdzziPkvvzIJXITl56fQx2mGM/DHLB3/2o=
|
||||
go.opentelemetry.io/otel/sdk v1.46.0 h1:h5CNQQjEbuQXY/JfZtgt3i7HVFV3aHPO2OAwO2eTYPI=
|
||||
go.opentelemetry.io/otel/sdk v1.46.0/go.mod h1:GAERFXFt5SYCEB+YiKUbMBeza6UaDH7GmGOZEfh2gSM=
|
||||
go.opentelemetry.io/otel/sdk/metric v1.46.0 h1:0piZ26EG4RBfebb2jhDH6ERCYHoVWduc3kLgPCwSnSE=
|
||||
go.opentelemetry.io/otel/sdk/metric v1.46.0/go.mod h1:I1PbKrdVc8Qu8HYVDNtqVIwLwjNrhsV/uFuxfwg8mO4=
|
||||
go.opentelemetry.io/otel/trace v1.46.0 h1:OULy7ccdJnZtJ0UDYFOIGaCmiWzJ8Vi2G/Rsu60qs1c=
|
||||
go.opentelemetry.io/otel/trace v1.46.0/go.mod h1:J7GAXweO77XSFkB/rmAqk9D6ihszhFjLU+d9WuUxDLI=
|
||||
go.opentelemetry.io/proto/otlp v1.11.0 h1:5rrYs0Ykyj50sdU/JU0x8etU+LubXWb+gED6TbEdMIk=
|
||||
go.opentelemetry.io/proto/otlp v1.11.0/go.mod h1:SmVizdCOAm3XBtG1g1NnOdhW6jtddT72hLMhv8VwA8E=
|
||||
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
|
||||
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
|
||||
go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ=
|
||||
@@ -53,10 +153,60 @@ go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ=
|
||||
go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
|
||||
go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw=
|
||||
go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg=
|
||||
golang.org/x/crypto v0.57.0 h1:3ZVCjf8Ggz7zneR/EHRVx68Ctf+2pmIMP2UFhh9cC6M=
|
||||
golang.org/x/crypto v0.57.0/go.mod h1:Fdz0i5U6CoizGwLda9DttjSk6qlZo25zYNtR+ycvuZA=
|
||||
golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
|
||||
golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
|
||||
golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c=
|
||||
golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o=
|
||||
golang.org/x/net v0.59.0 h1:5zfYln+w5XCxwrnMMJPufRgNoXEaGxl0wo5GqPXyues=
|
||||
golang.org/x/net v0.59.0/go.mod h1:2DA/G1UfVbCpQPeWTmMPGY7Cs2PkBkwu743bVX5PIVg=
|
||||
golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk=
|
||||
golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0=
|
||||
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
|
||||
golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og=
|
||||
golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI=
|
||||
golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E=
|
||||
golang.org/x/tools v0.49.0 h1:3NI7VXzL9+1WZD52Dx2ttoPwD5DWrFGpl9mFZDlmisI=
|
||||
golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo=
|
||||
gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4=
|
||||
gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E=
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20260819154853-08b0e4226688 h1:ax2KzoSRIZU/M0cIxri3pKxy99vniH1PVxWC6si/eZI=
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20260819154853-08b0e4226688/go.mod h1:1RJ9BQGyNdZwkGc1eTqkErfRZ6RJyYPHZo73BZ1vQqI=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260819154853-08b0e4226688 h1:cYNAzI2sUwhmCcoj9TxvihSrqsxt6uIkj3rDRhSDmW4=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20260819154853-08b0e4226688/go.mod h1:DjtHYE8FKJLivXcBEjGwndXfIC23G0VpXiXKqG179uA=
|
||||
google.golang.org/grpc v1.83.1 h1:HIO0+BEtBP6soyqvqC8sNUjZ7bTs+0hFQuFF+RAy++Y=
|
||||
google.golang.org/grpc v1.83.1/go.mod h1:kDyl6SKsiHKt0uylY5gtn5cEjkrIOhQOGDgIc4JGwzQ=
|
||||
google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc=
|
||||
google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
modernc.org/cc/v4 v4.29.2 h1:h6+9ciCnPKutf4I03CvheAvDLX7+IHlqR6Iy6J+cgd8=
|
||||
modernc.org/cc/v4 v4.29.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
|
||||
modernc.org/ccgo/v4 v4.35.0 h1:F+TUsmw09QxLzmi3aeYYGxjAXarmZaKgj3mKQHNaA8w=
|
||||
modernc.org/ccgo/v4 v4.35.0/go.mod h1:qrVGs9S3Sr2Ztcg9ve+kTAYMp5a3YvWjo+SoN06kJ5I=
|
||||
modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
|
||||
modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
|
||||
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
|
||||
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
|
||||
modernc.org/gc/v3 v3.1.5 h1:21ldfPfRYE31Tb7B3mwAK8gy1AxP4+dKjrOQPfqakoc=
|
||||
modernc.org/gc/v3 v3.1.5/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
|
||||
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
|
||||
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
|
||||
modernc.org/libc v1.75.7 h1:o3DTP9/0p9pKmY2WCKQaySW6wIiZhNM7wc2lUoyhfew=
|
||||
modernc.org/libc v1.75.7/go.mod h1:bO5o2ztHxBb2rjz0PgdHN0sSMw57CgxGFLZ3Qd/QpVQ=
|
||||
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||
modernc.org/memory v1.12.1 h1:nFMiWrpStgZczNl6XI9GnIk/rWhYIyHGUaR04pGbp9g=
|
||||
modernc.org/memory v1.12.1/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
|
||||
modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg=
|
||||
modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
|
||||
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
|
||||
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
|
||||
modernc.org/sqlite v1.59.0 h1:X1es1GpqBlS/5T+vbM4HLUdaa8OtQx468DF2vrx+38A=
|
||||
modernc.org/sqlite v1.59.0/go.mod h1:+paeT2A3iPRHkQDwG7oA6Tk0zQd5woMEI8q7orfry8k=
|
||||
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
||||
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
|
||||
@@ -0,0 +1,400 @@
|
||||
package account
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func newStore(t *testing.T) *store.DB {
|
||||
t.Helper()
|
||||
s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "a.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = s.Close() })
|
||||
|
||||
ctx := context.Background()
|
||||
for _, err := range []error{
|
||||
s.PutOrg(ctx, store.Org{ID: "org", Name: "Org"}),
|
||||
s.PutTeam(ctx, store.Team{ID: "team", OrgID: "org", Name: "Team"}),
|
||||
s.PutProject(ctx, store.Project{ID: "proj", TeamID: "team", Name: "Proj"}),
|
||||
} {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func newService(t *testing.T) (*Service, *store.DB) {
|
||||
t.Helper()
|
||||
s := newStore(t)
|
||||
return New(s, DefaultConfig(), slog.New(slog.DiscardHandler)), s
|
||||
}
|
||||
|
||||
const goodPassword = "correct horse battery staple"
|
||||
|
||||
func TestSignupLeavesAnAccountWaiting(t *testing.T) {
|
||||
// A public signup form that produced a working account would be a way in.
|
||||
svc, st := newService(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := svc.SignUp(ctx, "Ayse@Example.COM", "Ayşe", goodPassword); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u, err := st.UserByEmail(ctx, "ayse@example.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if u.State != store.UserPending {
|
||||
t.Errorf("a new account is %q", u.State)
|
||||
}
|
||||
if u.Email != "ayse@example.com" {
|
||||
t.Errorf("the address was not normalized: %q", u.Email)
|
||||
}
|
||||
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); !errors.Is(err, ErrNotApproved) {
|
||||
t.Errorf("a pending account signed in: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSignupDoesNotRevealExistingAccounts(t *testing.T) {
|
||||
// A form that says "this address already has an account" is an account
|
||||
// enumerator anyone can run, and the person signing up learns nothing they
|
||||
// need: either way they wait for approval.
|
||||
svc, _ := newService(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := svc.SignUp(ctx, "ayse@example.com", "A", goodPassword); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := svc.SignUp(ctx, "ayse@example.com", "Someone Else", goodPassword); err != nil {
|
||||
t.Errorf("the second signup reported %v, which confirms the address is taken", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalIsWhatMakesAnAccountWork(t *testing.T) {
|
||||
svc, st := newService(t)
|
||||
ctx := context.Background()
|
||||
must(t, svc.SignUp(ctx, "ayse@example.com", "A", goodPassword))
|
||||
|
||||
u, _ := st.UserByEmail(ctx, "ayse@example.com")
|
||||
must(t, st.ApproveUser(ctx, u.ID, "proj", "ilker", time.Now()))
|
||||
|
||||
token, csrf, err := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if token == "" || csrf == "" {
|
||||
t.Fatal("sign-in produced no session")
|
||||
}
|
||||
_, got, err := svc.Resolve(ctx, token)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ok, err := st.CanUseProject(ctx, got.ID, "proj"); err != nil || !ok {
|
||||
t.Errorf("the approved account is not a member of the project it was approved into: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalNeedNotNameAProject(t *testing.T) {
|
||||
// A person need not belong to anything to use the gateway: approved into
|
||||
// nothing, they work on their own within their personal limits.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
if err := st.ApproveUser(ctx, "usr-1", "", "ilker", time.Now()); err != nil {
|
||||
t.Fatalf("approving with no project: %v", err)
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); err != nil {
|
||||
t.Errorf("an account approved into nothing cannot sign in: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func newStoreAndUser(t *testing.T) (*Service, *store.DB) {
|
||||
t.Helper()
|
||||
svc, st := newService(t)
|
||||
hash, _ := HashPassword(goodPassword)
|
||||
must(t, st.CreateUser(context.Background(), store.User{
|
||||
ID: "usr-1", Email: "ayse@example.com", Name: "A",
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: time.Now(),
|
||||
}))
|
||||
return svc, st
|
||||
}
|
||||
|
||||
func TestSuspensionTakesEffectOnOpenSessions(t *testing.T) {
|
||||
// Leaving existing sessions open would make suspension advisory.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
token, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := svc.Resolve(ctx, token); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
must(t, st.SetUserState(ctx, "usr-1", store.UserSuspended))
|
||||
if _, _, err := svc.Resolve(ctx, token); !errors.Is(err, store.ErrNoSession) {
|
||||
t.Errorf("a suspended account kept its session: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSignOutEndsTheSession(t *testing.T) {
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
token, _, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
must(t, svc.SignOut(ctx, token))
|
||||
if _, _, err := svc.Resolve(ctx, token); !errors.Is(err, store.ErrNoSession) {
|
||||
t.Errorf("the session survived sign-out: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWrongPasswordAndAMissingAccountAreIndistinguishable(t *testing.T) {
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
_, _, wrong := svc.SignIn(ctx, "ayse@example.com", "wrong password here")
|
||||
_, _, missing := svc.SignIn(ctx, "nobody@example.com", "wrong password here")
|
||||
if !errors.Is(wrong, ErrBadCredentials) || !errors.Is(missing, ErrBadCredentials) {
|
||||
t.Fatalf("wrong=%v missing=%v", wrong, missing)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGuessingIsBounded(t *testing.T) {
|
||||
// Without this a public login form is an open brute-force target.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.MaxAttempts = 3
|
||||
svc = New(st, cfg, slog.New(slog.DiscardHandler))
|
||||
|
||||
for i := range 3 {
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", "nope"); !errors.Is(err, ErrBadCredentials) {
|
||||
t.Fatalf("attempt %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", "nope"); !errors.Is(err, ErrTooManyTries) {
|
||||
t.Errorf("guessing was not bounded: %v", err)
|
||||
}
|
||||
// And a correct password is refused too while the window is open, or the
|
||||
// limit would be trivially bypassed by interleaving.
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); !errors.Is(err, ErrTooManyTries) {
|
||||
t.Errorf("the window did not apply to a correct password: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASuccessfulSignInClearsTheWindow(t *testing.T) {
|
||||
// Otherwise a person who mistypes their password a few times is locked out
|
||||
// for the rest of the window even after getting it right.
|
||||
st := newStore(t)
|
||||
hash, _ := HashPassword(goodPassword)
|
||||
ctx := context.Background()
|
||||
must(t, st.CreateUser(ctx, store.User{
|
||||
ID: "usr-1", Email: "a@example.com", Name: "A",
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: time.Now(),
|
||||
}))
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.MaxAttempts = 3
|
||||
svc := New(st, cfg, slog.New(slog.DiscardHandler))
|
||||
|
||||
for range 2 {
|
||||
_, _, _ = svc.SignIn(ctx, "a@example.com", "nope")
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "a@example.com", goodPassword); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for range 3 {
|
||||
if _, _, err := svc.SignIn(ctx, "a@example.com", "nope"); errors.Is(err, ErrTooManyTries) {
|
||||
t.Fatal("the window was not cleared by a successful sign-in")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCSRFTokensAreCheckedAgainstTheSession(t *testing.T) {
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
token, csrf, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
sess, _, _ := svc.Resolve(ctx, token)
|
||||
|
||||
if !svc.CheckCSRF(sess, csrf) {
|
||||
t.Error("the session's own token was rejected")
|
||||
}
|
||||
for _, bad := range []string{"", "wrong", csrf + "x", csrf[:len(csrf)-1]} {
|
||||
if svc.CheckCSRF(sess, bad) {
|
||||
t.Errorf("%q was accepted as a csrf token", bad)
|
||||
}
|
||||
}
|
||||
// A token from another session must not work in this one.
|
||||
other, otherCSRF, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
_ = other
|
||||
if svc.CheckCSRF(sess, otherCSRF) {
|
||||
t.Error("another session's token was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasswordsAreHashedNotStored(t *testing.T) {
|
||||
svc, st := newService(t)
|
||||
ctx := context.Background()
|
||||
must(t, svc.SignUp(ctx, "ayse@example.com", "A", goodPassword))
|
||||
|
||||
u, _ := st.UserByEmail(ctx, "ayse@example.com")
|
||||
if strings.Contains(u.PasswordHash, goodPassword) {
|
||||
t.Fatal("the password is in the database")
|
||||
}
|
||||
if !strings.HasPrefix(u.PasswordHash, "$argon2id$") {
|
||||
t.Errorf("hash format: %q", u.PasswordHash)
|
||||
}
|
||||
ok, err := VerifyPassword(goodPassword, u.PasswordHash)
|
||||
if err != nil || !ok {
|
||||
t.Errorf("the hash does not verify: %v %v", ok, err)
|
||||
}
|
||||
if ok, _ := VerifyPassword("something else entirely", u.PasswordHash); ok {
|
||||
t.Error("a wrong password verified")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTwoHashesOfOnePasswordDiffer(t *testing.T) {
|
||||
// A salt that did not vary would make the table a rainbow-table lookup.
|
||||
a, _ := HashPassword(goodPassword)
|
||||
b, _ := HashPassword(goodPassword)
|
||||
if a == b {
|
||||
t.Fatal("hashing is not salted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestShortPasswordsAreRefused(t *testing.T) {
|
||||
if _, err := HashPassword("short"); !errors.Is(err, ErrWeakPassword) {
|
||||
t.Errorf("a five-character password was accepted: %v", err)
|
||||
}
|
||||
svc, _ := newService(t)
|
||||
if err := svc.SignUp(context.Background(), "a@example.com", "A", "short"); !errors.Is(err, ErrWeakPassword) {
|
||||
t.Errorf("signup accepted it: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMalformedHashesAreRejectedRatherThanMatched(t *testing.T) {
|
||||
for _, bad := range []string{"", "plaintext", "$argon2id$", "$bcrypt$v=19$m=1,t=1,p=1$aa$bb"} {
|
||||
if ok, err := VerifyPassword(goodPassword, bad); ok || err == nil {
|
||||
t.Errorf("VerifyPassword(%q) = %v, %v", bad, ok, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestObviousNonAddressesAreRefused(t *testing.T) {
|
||||
svc, _ := newService(t)
|
||||
ctx := context.Background()
|
||||
for _, bad := range []string{"", "ayse", "ayse@", "@example.com", "a b@example.com", "a@b", "a@@b.com"} {
|
||||
if err := svc.SignUp(ctx, bad, "A", goodPassword); !errors.Is(err, ErrInvalidEmail) {
|
||||
t.Errorf("SignUp(%q) = %v", bad, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func must(t *testing.T, err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChangingAPasswordNeedsTheOldOne(t *testing.T) {
|
||||
// A session is what an attacker gets from a borrowed laptop or a stolen
|
||||
// cookie. The password is what they do not have, and without this check the
|
||||
// first thing they would do is take the account.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
u, _ := st.UserByID(ctx, "usr-1")
|
||||
|
||||
if err := svc.ChangePassword(ctx, st, u, "not the password", "a brand new passphrase"); !errors.Is(err, ErrBadCredentials) {
|
||||
t.Errorf("a wrong current password was accepted: %v", err)
|
||||
}
|
||||
// And the old one still works.
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); err != nil {
|
||||
t.Errorf("the password changed anyway: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChangingAPasswordSignsOutEverywhere(t *testing.T) {
|
||||
// Leaving other sessions open would mean a password change does nothing
|
||||
// about the reason someone is changing it.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
laptop, _, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
phone, _, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
|
||||
u, _ := st.UserByID(ctx, "usr-1")
|
||||
must(t, svc.ChangePassword(ctx, st, u, goodPassword, "a brand new passphrase"))
|
||||
|
||||
for name, token := range map[string]string{"laptop": laptop, "phone": phone} {
|
||||
if _, _, err := svc.Resolve(ctx, token); !errors.Is(err, store.ErrNoSession) {
|
||||
t.Errorf("the %s session survived the change: %v", name, err)
|
||||
}
|
||||
}
|
||||
// The new password works and the old one does not.
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", "a brand new passphrase"); err != nil {
|
||||
t.Errorf("the new password does not work: %v", err)
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); !errors.Is(err, ErrBadCredentials) {
|
||||
t.Errorf("the old password still works: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANewPasswordMustActuallyBeNew(t *testing.T) {
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
u, _ := st.UserByID(ctx, "usr-1")
|
||||
|
||||
if err := svc.ChangePassword(ctx, st, u, goodPassword, goodPassword); !errors.Is(err, ErrSamePassword) {
|
||||
t.Errorf("changing a password to itself: %v", err)
|
||||
}
|
||||
if err := svc.ChangePassword(ctx, st, u, goodPassword, "short"); !errors.Is(err, ErrWeakPassword) {
|
||||
t.Errorf("a short new password was accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGuessingTheCurrentPasswordIsBounded(t *testing.T) {
|
||||
// Someone with a stolen session should not get unlimited attempts at the
|
||||
// password that would let them keep it.
|
||||
st := newStore(t)
|
||||
ctx := context.Background()
|
||||
hash, _ := HashPassword(goodPassword)
|
||||
must(t, st.CreateUser(ctx, store.User{
|
||||
ID: "usr-1", Email: "a@example.com", Name: "A",
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: time.Now(),
|
||||
}))
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.MaxAttempts = 3
|
||||
svc := New(st, cfg, slog.New(slog.DiscardHandler))
|
||||
u, _ := st.UserByID(ctx, "usr-1")
|
||||
|
||||
for range 3 {
|
||||
_ = svc.ChangePassword(ctx, st, u, "wrong", "a brand new passphrase")
|
||||
}
|
||||
if err := svc.ChangePassword(ctx, st, u, "wrong", "a brand new passphrase"); !errors.Is(err, ErrTooManyTries) {
|
||||
t.Errorf("guessing was unbounded: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
// Package account is the human half of authentication: people sign up, an
|
||||
// operator approves them, and they issue their own API keys from a profile page
|
||||
// (plan §7.9).
|
||||
//
|
||||
// This is deliberately a different mechanism from the one in internal/auth.
|
||||
// That one authenticates machines holding 256-bit random keys; this one
|
||||
// authenticates people holding passwords they chose. The two have opposite
|
||||
// requirements — a fast hash is correct there and wrong here — and keeping them
|
||||
// apart is what stops one set of assumptions leaking into the other.
|
||||
package account
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/crypto/argon2"
|
||||
)
|
||||
|
||||
// Argon2 parameters. OWASP's current floor for argon2id, which costs roughly
|
||||
// 50ms per attempt on ordinary hardware — slow enough to make an offline
|
||||
// dictionary attack expensive, fast enough that a login does not feel broken.
|
||||
//
|
||||
// They are stored with each hash, so raising them later leaves existing
|
||||
// passwords working until their owners next sign in.
|
||||
const (
|
||||
argonTime = 3
|
||||
argonMemory = 64 * 1024 // KiB
|
||||
argonThreads = 4
|
||||
argonKeyLen = 32
|
||||
saltLen = 16
|
||||
)
|
||||
|
||||
// MinPasswordLen is the only password rule.
|
||||
//
|
||||
// Length is the one requirement that reliably buys entropy. Composition rules —
|
||||
// a digit, a symbol, a capital — push people towards Password1! and produce
|
||||
// less entropy than they cost, so there are none.
|
||||
const MinPasswordLen = 12
|
||||
|
||||
// ErrWeakPassword says a password is too short to accept.
|
||||
var ErrWeakPassword = fmt.Errorf("a password must be at least %d characters", MinPasswordLen)
|
||||
|
||||
// ErrBadHash says a stored hash could not be parsed.
|
||||
var ErrBadHash = errors.New("unreadable password hash")
|
||||
|
||||
// HashPassword derives a storable hash.
|
||||
func HashPassword(password string) (string, error) {
|
||||
if len([]rune(password)) < MinPasswordLen {
|
||||
return "", ErrWeakPassword
|
||||
}
|
||||
salt := make([]byte, saltLen)
|
||||
if _, err := rand.Read(salt); err != nil {
|
||||
return "", err
|
||||
}
|
||||
sum := idKey([]byte(password), salt, argonTime, argonMemory, argonThreads, argonKeyLen)
|
||||
|
||||
b64 := base64.RawStdEncoding.EncodeToString
|
||||
return fmt.Sprintf("$argon2id$v=%d$m=%d,t=%d,p=%d$%s$%s",
|
||||
argon2.Version, argonMemory, argonTime, argonThreads, b64(salt), b64(sum)), nil
|
||||
}
|
||||
|
||||
// VerifyPassword reports whether a password matches a stored hash.
|
||||
//
|
||||
// The comparison is constant time. The parameters come from the hash rather
|
||||
// than from the constants above, so a hash written under older settings still
|
||||
// verifies.
|
||||
func VerifyPassword(password, encoded string) (bool, error) {
|
||||
parts := strings.Split(encoded, "$")
|
||||
if len(parts) != 6 || parts[1] != "argon2id" {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
|
||||
var version int
|
||||
if _, err := fmt.Sscanf(parts[2], "v=%d", &version); err != nil || version != argon2.Version {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
var memory, time uint32
|
||||
var threads uint8
|
||||
if _, err := fmt.Sscanf(parts[3], "m=%d,t=%d,p=%d", &memory, &time, &threads); err != nil {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
|
||||
salt, err := base64.RawStdEncoding.DecodeString(parts[4])
|
||||
if err != nil {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
want, err := base64.RawStdEncoding.DecodeString(parts[5])
|
||||
if err != nil {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
|
||||
got := idKey([]byte(password), salt, time, memory, threads, uint32(len(want)))
|
||||
return subtle.ConstantTimeCompare(got, want) == 1, nil
|
||||
}
|
||||
|
||||
// dummyHash is verified against when no account matches, so a login attempt for
|
||||
// an address that does not exist costs the same time as one that does.
|
||||
//
|
||||
// Without it, a stopwatch enumerates every account on the system.
|
||||
var dummyHash string
|
||||
|
||||
func init() {
|
||||
h, err := HashPassword(strings.Repeat("x", MinPasswordLen))
|
||||
if err != nil {
|
||||
panic("account: cannot build the dummy hash: " + err.Error())
|
||||
}
|
||||
dummyHash = h
|
||||
}
|
||||
|
||||
// argonSlots bounds how many hashes run at once. Each takes 64 MiB, and the
|
||||
// login and signup forms are on the client-facing port: without a bound, a
|
||||
// few hundred parallel attempts at random addresses asked for gigabytes at
|
||||
// once, and the per-address throttle does not see them — each address is new.
|
||||
// Beyond the bound, attempts wait their turn instead.
|
||||
var argonSlots = make(chan struct{}, max(2, runtime.NumCPU()/2))
|
||||
|
||||
func idKey(password, salt []byte, time, memory uint32, threads uint8, keyLen uint32) []byte {
|
||||
argonSlots <- struct{}{}
|
||||
defer func() { <-argonSlots }()
|
||||
return argon2.IDKey(password, salt, time, memory, threads, keyLen)
|
||||
}
|
||||
@@ -0,0 +1,350 @@
|
||||
package account
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/base32"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// What a sign-in or sign-up can fail with.
|
||||
var (
|
||||
ErrBadCredentials = errors.New("that email and password do not match an account")
|
||||
ErrNotApproved = errors.New("this account is waiting for approval")
|
||||
ErrSuspended = errors.New("this account has been suspended")
|
||||
ErrTooManyTries = errors.New("too many attempts; try again shortly")
|
||||
ErrInvalidEmail = errors.New("that does not look like an email address")
|
||||
)
|
||||
|
||||
// Store is what the service needs from persistence.
|
||||
type Store interface {
|
||||
CreateUser(ctx context.Context, u store.User) error
|
||||
UserByEmail(ctx context.Context, email string) (store.User, error)
|
||||
UserByID(ctx context.Context, id string) (store.User, error)
|
||||
TouchLogin(ctx context.Context, id string, at time.Time) error
|
||||
|
||||
CreateSession(ctx context.Context, userID, token, csrf string, ttl time.Duration) error
|
||||
SessionByToken(ctx context.Context, token string) (store.Session, store.User, error)
|
||||
DeleteSession(ctx context.Context, token string) error
|
||||
TouchSession(ctx context.Context, token string, at time.Time) error
|
||||
}
|
||||
|
||||
// Config tunes the service.
|
||||
type Config struct {
|
||||
// SignupOpen decides whether anyone may create an account. Off, the form is
|
||||
// not served at all; accounts are made with `llmgw users create`.
|
||||
SignupOpen bool
|
||||
// SessionTTL is how long a sign-in lasts.
|
||||
SessionTTL time.Duration
|
||||
// MaxAttempts and AttemptWindow bound password guessing per address.
|
||||
MaxAttempts int
|
||||
AttemptWindow time.Duration
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible values.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
SignupOpen: true,
|
||||
SessionTTL: store.SessionTTL,
|
||||
MaxAttempts: 10,
|
||||
AttemptWindow: 15 * time.Minute,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() {
|
||||
d := DefaultConfig()
|
||||
if c.SessionTTL <= 0 {
|
||||
c.SessionTTL = d.SessionTTL
|
||||
}
|
||||
if c.MaxAttempts <= 0 {
|
||||
c.MaxAttempts = d.MaxAttempts
|
||||
}
|
||||
if c.AttemptWindow <= 0 {
|
||||
c.AttemptWindow = d.AttemptWindow
|
||||
}
|
||||
}
|
||||
|
||||
// Service is signup, sign-in and sessions.
|
||||
type Service struct {
|
||||
store Store
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
attempts map[string]*attemptWindow
|
||||
}
|
||||
|
||||
type attemptWindow struct {
|
||||
count int
|
||||
since time.Time
|
||||
}
|
||||
|
||||
// New builds the service.
|
||||
func New(s Store, cfg Config, log *slog.Logger) *Service {
|
||||
cfg.normalize()
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Service{store: s, cfg: cfg, log: log, now: time.Now, attempts: map[string]*attemptWindow{}}
|
||||
}
|
||||
|
||||
// SignupOpen reports whether the form should be served.
|
||||
func (s *Service) SignupOpen() bool { return s.cfg.SignupOpen }
|
||||
|
||||
// SignUp creates a pending account.
|
||||
//
|
||||
// It returns nil whether or not the address was already taken. A signup form
|
||||
// that says "this address already has an account" is an account enumerator that
|
||||
// anyone can run, and the person signing up learns nothing they need: either
|
||||
// way, what happens next is that they wait for approval.
|
||||
func (s *Service) SignUp(ctx context.Context, email, name, password string) error {
|
||||
email = store.NormalizeEmail(email)
|
||||
if !validEmail(email) {
|
||||
return ErrInvalidEmail
|
||||
}
|
||||
hash, err := HashPassword(password)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if name = strings.TrimSpace(name); name == "" {
|
||||
name = email
|
||||
}
|
||||
|
||||
err = s.store.CreateUser(ctx, store.User{
|
||||
ID: randomID("usr"), Email: email, Name: name,
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: s.now(),
|
||||
})
|
||||
if errors.Is(err, store.ErrEmailTaken) {
|
||||
// Logged so an operator can see it; not told to the caller.
|
||||
s.log.Info("signup for an address that already has an account")
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.log.Info("a new account is waiting for approval", "email", email)
|
||||
return nil
|
||||
}
|
||||
|
||||
// SignIn checks a password and opens a session, returning the session token and
|
||||
// the CSRF token that must accompany the forms in it.
|
||||
func (s *Service) SignIn(ctx context.Context, email, password string) (token, csrf string, err error) {
|
||||
email = store.NormalizeEmail(email)
|
||||
if !s.allowAttempt(email) {
|
||||
return "", "", ErrTooManyTries
|
||||
}
|
||||
|
||||
u, err := s.store.UserByEmail(ctx, email)
|
||||
if errors.Is(err, store.ErrNotFound) {
|
||||
// Verify anyway, against a hash that cannot match. Returning early here
|
||||
// would let a stopwatch enumerate every account on the system.
|
||||
_, _ = VerifyPassword(password, dummyHash)
|
||||
return "", "", ErrBadCredentials
|
||||
}
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
|
||||
ok, err := VerifyPassword(password, u.PasswordHash)
|
||||
if err != nil || !ok {
|
||||
return "", "", ErrBadCredentials
|
||||
}
|
||||
// The password was right, so this was not an attack.
|
||||
s.clearAttempts(email)
|
||||
|
||||
// State is checked after the password, on purpose. Checking first would
|
||||
// tell anyone who guesses an address whether it has a pending account.
|
||||
switch u.State {
|
||||
case store.UserPending:
|
||||
return "", "", ErrNotApproved
|
||||
case store.UserSuspended:
|
||||
return "", "", ErrSuspended
|
||||
}
|
||||
|
||||
if token, err = store.NewSessionToken(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if csrf, err = store.NewSessionToken(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if err := s.store.CreateSession(ctx, u.ID, token, csrf, s.cfg.SessionTTL); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if err := s.store.TouchLogin(ctx, u.ID, s.now()); err != nil {
|
||||
s.log.Warn("could not record a sign-in", "err", err)
|
||||
}
|
||||
return token, csrf, nil
|
||||
}
|
||||
|
||||
// Resolve turns a session cookie into the person holding it.
|
||||
func (s *Service) Resolve(ctx context.Context, token string) (store.Session, store.User, error) {
|
||||
if token == "" {
|
||||
return store.Session{}, store.User{}, store.ErrNoSession
|
||||
}
|
||||
sess, u, err := s.store.SessionByToken(ctx, token)
|
||||
if err != nil {
|
||||
return store.Session{}, store.User{}, err
|
||||
}
|
||||
if err := s.store.TouchSession(ctx, token, s.now()); err != nil {
|
||||
s.log.Debug("could not touch a session", "err", err)
|
||||
}
|
||||
return sess, u, nil
|
||||
}
|
||||
|
||||
// SignOut closes one session.
|
||||
func (s *Service) SignOut(ctx context.Context, token string) error {
|
||||
if token == "" {
|
||||
return nil
|
||||
}
|
||||
return s.store.DeleteSession(ctx, token)
|
||||
}
|
||||
|
||||
// CheckCSRF compares a submitted token against the session's.
|
||||
//
|
||||
// A cookie authenticates the browser, not the page that asked. Without this, any
|
||||
// site the person visits can post a form to ours and have the browser attach the
|
||||
// session automatically.
|
||||
func (s *Service) CheckCSRF(sess store.Session, submitted string) bool {
|
||||
if submitted == "" || len(sess.CSRFHash) == 0 {
|
||||
return false
|
||||
}
|
||||
return subtle.ConstantTimeCompare(store.HashKey(submitted), sess.CSRFHash) == 1
|
||||
}
|
||||
|
||||
// allowAttempt bounds password guessing per address.
|
||||
//
|
||||
// Per address rather than per IP: an attacker has many addresses to come from
|
||||
// and only one to guess against, and limiting by IP would lock out everyone
|
||||
// behind one office NAT the moment somebody fat-fingers their password.
|
||||
func (s *Service) allowAttempt(email string) bool {
|
||||
now := s.now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
w, ok := s.attempts[email]
|
||||
if !ok && len(s.attempts) >= maxTrackedAddresses {
|
||||
// Full of addresses nobody has signed in as for a while: forget the
|
||||
// ones whose window has passed. Only expiring, never evicting early —
|
||||
// dropping a live window is how an attacker spraying random addresses
|
||||
// would buy unlimited guesses at the one they care about.
|
||||
for k, v := range s.attempts {
|
||||
if now.Sub(v.since) > s.cfg.AttemptWindow {
|
||||
delete(s.attempts, k)
|
||||
}
|
||||
}
|
||||
if len(s.attempts) >= maxTrackedAddresses {
|
||||
// Still full, so somebody is spraying addresses faster than they
|
||||
// expire. New ones wait; those already tracked are unaffected.
|
||||
return false
|
||||
}
|
||||
}
|
||||
if !ok || now.Sub(w.since) > s.cfg.AttemptWindow {
|
||||
s.attempts[email] = &attemptWindow{count: 1, since: now}
|
||||
return true
|
||||
}
|
||||
w.count++
|
||||
return w.count <= s.cfg.MaxAttempts
|
||||
}
|
||||
|
||||
// maxTrackedAddresses bounds the throttle's memory. Unbounded, it grew by one
|
||||
// entry per address tried and shrank only on a successful sign-in.
|
||||
const maxTrackedAddresses = 50000
|
||||
|
||||
func (s *Service) clearAttempts(email string) {
|
||||
s.mu.Lock()
|
||||
delete(s.attempts, email)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// validEmail is deliberately crude. A real check is delivery, which this does
|
||||
// not do; this only rejects what is obviously not an address so the database
|
||||
// does not fill with them.
|
||||
func validEmail(s string) bool {
|
||||
at := strings.IndexByte(s, '@')
|
||||
if at <= 0 || at != strings.LastIndexByte(s, '@') || at == len(s)-1 {
|
||||
return false
|
||||
}
|
||||
if strings.ContainsAny(s, " \t\r\n<>\"") {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(s[at+1:], ".")
|
||||
}
|
||||
|
||||
func randomID(prefix string) string {
|
||||
var buf [10]byte
|
||||
_, _ = rand.Read(buf[:])
|
||||
return prefix + "-" + strings.ToLower(
|
||||
base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:]))
|
||||
}
|
||||
|
||||
// ErrSamePassword says the new password is the old one.
|
||||
var ErrSamePassword = errors.New("that is the password you already have")
|
||||
|
||||
// PasswordChanger is the extra storage a password change needs.
|
||||
type PasswordChanger interface {
|
||||
SetPassword(ctx context.Context, userID, hash string) error
|
||||
DeleteUserSessions(ctx context.Context, userID string) error
|
||||
}
|
||||
|
||||
// ChangePassword replaces a person's password and signs them out everywhere.
|
||||
//
|
||||
// The current password is required even though the caller already holds a
|
||||
// session. A session is what an attacker gets from a borrowed laptop or a stolen
|
||||
// cookie; the password is what they do not have, and without this check the
|
||||
// first thing they would do is take the account.
|
||||
//
|
||||
// Every session is closed, this browser's included. Leaving the others open
|
||||
// would mean a password change does nothing about the reason someone is
|
||||
// changing it. The caller re-issues a session for the browser in front of them.
|
||||
func (s *Service) ChangePassword(ctx context.Context, pw PasswordChanger,
|
||||
u store.User, current, next string,
|
||||
) error {
|
||||
if !s.allowAttempt("pw:" + u.ID) {
|
||||
return ErrTooManyTries
|
||||
}
|
||||
ok, err := VerifyPassword(current, u.PasswordHash)
|
||||
if err != nil || !ok {
|
||||
return ErrBadCredentials
|
||||
}
|
||||
s.clearAttempts("pw:" + u.ID)
|
||||
|
||||
if same, _ := VerifyPassword(next, u.PasswordHash); same {
|
||||
return ErrSamePassword
|
||||
}
|
||||
hash, err := HashPassword(next)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := pw.SetPassword(ctx, u.ID, hash); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := pw.DeleteUserSessions(ctx, u.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
s.log.Info("a password was changed", "user", u.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// OpenSession starts a session for a person who has just proved who they are by
|
||||
// some other means — immediately after changing their password, which closed
|
||||
// the one they were using.
|
||||
func (s *Service) OpenSession(ctx context.Context, userID string) (token, csrf string, err error) {
|
||||
if token, err = store.NewSessionToken(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if csrf, err = store.NewSessionToken(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if err := s.store.CreateSession(ctx, userID, token, csrf, s.cfg.SessionTTL); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return token, csrf, nil
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package account
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSprayingAddressesCannotGrowTheThrottleOrEscapeIt(t *testing.T) {
|
||||
now := time.Unix(0, 0)
|
||||
s := &Service{cfg: DefaultConfig(), attempts: map[string]*attemptWindow{}, now: func() time.Time { return now }}
|
||||
|
||||
// The target is being guessed at.
|
||||
for i := 0; i < s.cfg.MaxAttempts; i++ {
|
||||
s.allowAttempt("victim@example.com")
|
||||
}
|
||||
// Then the table is flooded with made-up addresses.
|
||||
for i := 0; i < maxTrackedAddresses+100; i++ {
|
||||
s.allowAttempt(fmt.Sprintf("x%d@example.com", i))
|
||||
}
|
||||
if n := len(s.attempts); n > maxTrackedAddresses {
|
||||
t.Errorf("the table grew to %d", n)
|
||||
}
|
||||
if s.allowAttempt("victim@example.com") {
|
||||
t.Error("flooding the table reset the target's throttle")
|
||||
}
|
||||
|
||||
// Once the window passes, everything expires and sign-ins work again.
|
||||
now = now.Add(s.cfg.AttemptWindow + time.Second)
|
||||
if !s.allowAttempt("new@example.com") {
|
||||
t.Error("a new address was refused after the flood expired")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,421 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// Refusals a key can hit. They are distinguished so the client is told which
|
||||
// limit it hit, and so retrying is only suggested where retrying helps.
|
||||
var (
|
||||
ErrRateLimited = errors.New("request rate limit exceeded")
|
||||
ErrTokenLimited = errors.New("token rate limit exceeded")
|
||||
ErrBudgetSpent = errors.New("budget exhausted")
|
||||
)
|
||||
|
||||
// LimitError carries which limit was hit and, where it makes sense, when to
|
||||
// come back.
|
||||
type LimitError struct {
|
||||
Err error
|
||||
// RetryAfter is zero when waiting does not help, which is the case for a
|
||||
// spent budget: it will not refill until the period rolls over, and telling
|
||||
// a client to retry in a second would be a lie.
|
||||
RetryAfter time.Duration
|
||||
}
|
||||
|
||||
func (e *LimitError) Error() string { return e.Err.Error() }
|
||||
func (e *LimitError) Unwrap() error { return e.Err }
|
||||
|
||||
// Accountant enforces limits and records what was used.
|
||||
//
|
||||
// One object because the three halves are not independent: a budget check needs
|
||||
// the price, the price needs the model, and the record closes the loop by
|
||||
// feeding the counters that the next check reads.
|
||||
type Accountant struct {
|
||||
pricing *Pricing
|
||||
budgets *Budgets
|
||||
requests *Limiter
|
||||
tokens *Limiter
|
||||
recorder *Recorder
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// New builds an accountant.
|
||||
func New(pricing *Pricing, recorder *Recorder, log *slog.Logger) *Accountant {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Accountant{
|
||||
pricing: pricing,
|
||||
budgets: NewBudgets(),
|
||||
requests: NewLimiter(),
|
||||
tokens: NewLimiter(),
|
||||
recorder: recorder,
|
||||
log: log,
|
||||
}
|
||||
}
|
||||
|
||||
// RestoreBudgets seeds the in-memory counters from the store.
|
||||
//
|
||||
// Called at startup. A restart must not reset a budget, which it would if the
|
||||
// counters only ever lived in memory.
|
||||
func (a *Accountant) RestoreBudgets(ctx context.Context, s store.Store, period store.Period) error {
|
||||
start := PeriodStart(period, time.Now())
|
||||
if start.IsZero() {
|
||||
return nil
|
||||
}
|
||||
rows, err := s.SpendSinceByScope(ctx, start)
|
||||
if err != nil {
|
||||
return fmt.Errorf("restoring budgets: %w", err)
|
||||
}
|
||||
// Every scope a key sits under, not just the key. A budget owned by a team
|
||||
// is one counter fed by every key beneath it, and restoring only per-key
|
||||
// totals would bring the control plane back up believing the team had spent
|
||||
// nothing — which is a budget that resets on every restart, for exactly the
|
||||
// scopes where it matters most.
|
||||
// And every pool a person is in: their own, and each per-member limit of
|
||||
// the project, team and organization they spent it in.
|
||||
spend := make(map[string]int64, len(rows)*8)
|
||||
for _, r := range rows {
|
||||
owners := []store.LimitOwner{
|
||||
{Scope: store.ScopeKey, ID: r.KeyID},
|
||||
{Scope: store.ScopeProject, ID: r.ProjectID},
|
||||
{Scope: store.ScopeTeam, ID: r.TeamID},
|
||||
{Scope: store.ScopeOrg, ID: r.OrgID},
|
||||
{Scope: store.ScopeUser, ID: r.UserID},
|
||||
}
|
||||
if r.UserID != "" && r.ProjectID != "" {
|
||||
owners = append(owners,
|
||||
store.MemberOwner(store.ScopeProjectMember, r.ProjectID, r.UserID),
|
||||
store.MemberOwner(store.ScopeTeamMember, r.TeamID, r.UserID),
|
||||
store.MemberOwner(store.ScopeOrgMember, r.OrgID, r.UserID))
|
||||
}
|
||||
for _, o := range owners {
|
||||
if o.ID == "" || strings.HasPrefix(o.ID, "/") {
|
||||
continue
|
||||
}
|
||||
if k := o.CounterKey(); k != "" {
|
||||
spend[k] += r.Micros
|
||||
}
|
||||
}
|
||||
}
|
||||
a.budgets.Restore(period, spend)
|
||||
a.log.Info("restored budgets from the store",
|
||||
"keys", len(rows), "counters", len(spend), "since", start)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Reservation is what Check returns: the estimate it charged, so Record can
|
||||
// settle the difference against the real figure.
|
||||
// pool is one level's counter and the period its budget runs over.
|
||||
type pool struct {
|
||||
Counter string
|
||||
Period store.Period
|
||||
}
|
||||
|
||||
type Reservation struct {
|
||||
KeyID string
|
||||
// Budgets and Tokens are every counter the check went through, carried so
|
||||
// the settlement lands on the same ones. Every, not the tightest: an
|
||||
// organization's pool only fills up if each team's spend is added to it
|
||||
// as well as to the team's own.
|
||||
//
|
||||
// Derived again at settlement would mean re-resolving the limits, and a
|
||||
// budget checked against the team and settled against the key is a budget
|
||||
// that never fills up.
|
||||
Budgets []pool
|
||||
Tokens []string
|
||||
EstimatedTokens int64
|
||||
Price Price
|
||||
// Period is the key's budget period, carried from the check so the record
|
||||
// counts against the same window the check read. Deriving it later would
|
||||
// mean guessing, and a monthly budget credited to a daily counter is a
|
||||
// budget that never binds.
|
||||
Period store.Period
|
||||
}
|
||||
|
||||
// Check decides whether a request may proceed.
|
||||
//
|
||||
// It runs before scheduling and touches no database: every counter it reads is
|
||||
// in memory. The token limit is checked against an estimate, because the real
|
||||
// figure only exists once the request has run; Record settles the difference.
|
||||
func (a *Accountant) Check(p auth.Principal, model string, estimatedTokens int64) (Reservation, error) {
|
||||
price := a.pricing.For(model)
|
||||
lim := p.Limits
|
||||
// Every level that set a limit, each against its own pool. A quota written
|
||||
// on a team is one pool the team shares, and the organization above it is a
|
||||
// second pool that the team's spending also fills — both have to hold, or
|
||||
// an organization given 1000 with four teams given 500 spends 2000.
|
||||
chain := levels(lim, p.KeyID)
|
||||
|
||||
res := Reservation{
|
||||
KeyID: p.KeyID, EstimatedTokens: estimatedTokens,
|
||||
Price: price, Period: lim.BudgetPeriod,
|
||||
}
|
||||
for _, l := range chain {
|
||||
if l.Budget != nil {
|
||||
res.Budgets = append(res.Budgets, pool{l.Counter, l.Period})
|
||||
}
|
||||
if l.TPM != nil {
|
||||
res.Tokens = append(res.Tokens, l.Counter)
|
||||
}
|
||||
}
|
||||
if p.KeyID == "" {
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// Budget first. It is the limit whose refusal cannot be retried away, so
|
||||
// telling the client about it first is more useful than reporting a rate
|
||||
// limit it would hit again in a second.
|
||||
//
|
||||
// A model that costs nothing is exempt, which is what lets a key with no
|
||||
// budget left still be served a cheap fallback (plan §7.7).
|
||||
// Budgets first, and all of them. Exhausted only reads, so the order costs
|
||||
// nothing and a refusal here charges no counter.
|
||||
if !price.Zero() {
|
||||
for _, l := range chain {
|
||||
if l.Budget != nil && a.budgets.Exhausted(l.Counter, l.Period, *l.Budget) {
|
||||
return res, &LimitError{Err: ErrBudgetSpent}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rate limits charge as they check — deliberately, so two concurrent
|
||||
// requests cannot both see room. With several levels that means a refusal
|
||||
// at one has already taken from the levels before it, so those are handed
|
||||
// back. Without the refund a team at its ceiling would quietly eat its
|
||||
// organization's rate limit with requests that were never served.
|
||||
if err := a.allowAll(chain, func(l level) (string, int, int64, bool) {
|
||||
return l.Counter, valInt(l.RPM), 1, l.RPM != nil
|
||||
}, a.requests); err != nil {
|
||||
return res, &LimitError{Err: ErrRateLimited, RetryAfter: time.Second}
|
||||
}
|
||||
if err := a.allowAll(chain, func(l level) (string, int, int64, bool) {
|
||||
return l.Counter, valInt(l.TPM), estimatedTokens, l.TPM != nil
|
||||
}, a.tokens); err != nil {
|
||||
// The request was not charged against the request-rate limit in vain:
|
||||
// it was, and that is correct — it did arrive.
|
||||
return res, &LimitError{Err: ErrTokenLimited, RetryAfter: time.Second}
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// level is one scope's limits with the counter they are enforced against.
|
||||
type level struct {
|
||||
Counter string
|
||||
Period store.Period
|
||||
RPM *int
|
||||
TPM *int
|
||||
Budget *int64
|
||||
}
|
||||
|
||||
// levels turns a resolved chain into the pools a request must pass.
|
||||
//
|
||||
// Falls back to the collapsed values when the chain is empty, which is what a
|
||||
// principal assembled by hand looks like — in tests, and anywhere that builds
|
||||
// limits without walking the tree.
|
||||
func levels(e store.Effective, keyID string) []level {
|
||||
if len(e.Chain) == 0 {
|
||||
var out []level
|
||||
if e.RPM != nil || e.TPM != nil || e.Budget != nil {
|
||||
out = append(out, level{
|
||||
Counter: counterFor(e.BudgetOwner, keyID),
|
||||
Period: e.BudgetPeriod, Budget: e.Budget,
|
||||
})
|
||||
out[0].RPM, out[0].TPM = e.RPM, e.TPM
|
||||
// One entry, so a hand-built principal with different owners per
|
||||
// limit still counts each where it said to.
|
||||
if r, b := counterFor(e.RPMOwner, keyID), out[0].Counter; r != b && e.RPM != nil {
|
||||
out[0].RPM = nil
|
||||
out = append(out, level{Counter: r, RPM: e.RPM})
|
||||
}
|
||||
if tp, b := counterFor(e.TPMOwner, keyID), out[0].Counter; tp != b && e.TPM != nil {
|
||||
out[0].TPM = nil
|
||||
out = append(out, level{Counter: tp, TPM: e.TPM})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
out := make([]level, 0, len(e.Chain))
|
||||
for _, l := range e.Chain {
|
||||
out = append(out, level{
|
||||
Counter: store.LimitOwner{Scope: l.Scope, ID: l.ScopeID}.CounterKey(),
|
||||
Period: l.BudgetPeriod, RPM: l.RPM, TPM: l.TPM, Budget: l.Budget,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// allowAll charges every level that has this limit, refunding on a refusal.
|
||||
func (a *Accountant) allowAll(chain []level,
|
||||
pick func(level) (key string, limit int, n int64, has bool), lim *Limiter,
|
||||
) error {
|
||||
var charged []struct {
|
||||
key string
|
||||
n int64
|
||||
}
|
||||
for _, l := range chain {
|
||||
key, limit, n, has := pick(l)
|
||||
if !has {
|
||||
continue
|
||||
}
|
||||
if !lim.Allow(key, limit, n) {
|
||||
for _, c := range charged {
|
||||
lim.Charge(c.key, -c.n)
|
||||
}
|
||||
return errRefused
|
||||
}
|
||||
charged = append(charged, struct {
|
||||
key string
|
||||
n int64
|
||||
}{key, n})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var errRefused = errors.New("refused")
|
||||
|
||||
func valInt(v *int) int {
|
||||
if v == nil {
|
||||
return 0
|
||||
}
|
||||
return *v
|
||||
}
|
||||
|
||||
// Record settles a finished request: it reconciles the token estimate, adds to
|
||||
// the budget, and queues the usage row.
|
||||
func (a *Accountant) Record(res Reservation, u store.Usage) {
|
||||
if res.KeyID == "" {
|
||||
return
|
||||
}
|
||||
actual := u.PromptTokens + u.CompletionTokens
|
||||
|
||||
// The estimate was charged up front; charge the difference now. A caller
|
||||
// whose requests are consistently larger than estimated is still bounded,
|
||||
// one window later than with a perfect estimate.
|
||||
// Every level the check went through, so an organization's pool fills up
|
||||
// from its teams' spending rather than only from requests that named it.
|
||||
for _, c := range res.Tokens {
|
||||
a.tokens.Charge(c, actual-res.EstimatedTokens)
|
||||
}
|
||||
|
||||
u.CostMicros = res.Price.Cost(u.PromptTokens, u.CompletionTokens)
|
||||
for _, b := range res.Budgets {
|
||||
a.budgets.Add(b.Counter, b.Period, u.CostMicros)
|
||||
}
|
||||
|
||||
if a.recorder != nil {
|
||||
a.recorder.Record(u)
|
||||
}
|
||||
}
|
||||
|
||||
// Status returns what a key has used, for headers and `ctl`.
|
||||
func (a *Accountant) Status(p auth.Principal) (spent int64, rpm, tpm float64) {
|
||||
if p.KeyID == "" {
|
||||
return 0, 0, 0
|
||||
}
|
||||
// The pool that will refuse them, which is not always the one with the
|
||||
// smallest number on it. A team given 2000 inside an organization given the
|
||||
// same is bounded by whichever is nearer its ceiling, and that is the
|
||||
// organization as soon as a second team spends anything. A header that
|
||||
// reported the team's would disagree with the refusal that follows it.
|
||||
for _, l := range levels(p.Limits, p.KeyID) {
|
||||
if l.Budget != nil {
|
||||
if v := a.budgets.Spent(l.Counter, l.Period); v > spent {
|
||||
spent = v
|
||||
}
|
||||
}
|
||||
if l.RPM != nil {
|
||||
if v := a.requests.Rate(l.Counter); v > rpm {
|
||||
rpm = v
|
||||
}
|
||||
}
|
||||
if l.TPM != nil {
|
||||
if v := a.tokens.Rate(l.Counter); v > tpm {
|
||||
tpm = v
|
||||
}
|
||||
}
|
||||
}
|
||||
return spent, rpm, tpm
|
||||
}
|
||||
|
||||
// counterFor is the counter a limit owned by this scope should use.
|
||||
//
|
||||
// Falls back to the key's own counter when no scope owns the limit, which
|
||||
// happens when nothing in the chain set one.
|
||||
//
|
||||
// Namespaced even then. A bare key id here and a "key:"-prefixed one in the
|
||||
// restore path would be two different counters for the same key, so a budget
|
||||
// would check against an empty one and survive every restart — which is how
|
||||
// this was wrong when it was first written.
|
||||
func counterFor(o store.LimitOwner, keyID string) string {
|
||||
if k := o.CounterKey(); k != "" {
|
||||
return k
|
||||
}
|
||||
return keyCounter(keyID)
|
||||
}
|
||||
|
||||
// keyCounter is a key's own counter, in the same namespace as every other.
|
||||
func or(v, fallback string) string {
|
||||
if v == "" {
|
||||
return fallback
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func keyCounter(keyID string) string {
|
||||
if keyID == "" {
|
||||
return ""
|
||||
}
|
||||
return store.LimitOwner{Scope: store.ScopeKey, ID: keyID}.CounterKey()
|
||||
}
|
||||
|
||||
// AnythingPriced reports whether a budget can ever be spent against.
|
||||
//
|
||||
// On the accountant rather than making callers reach for the price table: what
|
||||
// they are asking is whether the budgets they are about to display mean
|
||||
// anything, and the accountant is what would enforce them.
|
||||
func (a *Accountant) AnythingPriced() bool { return a.pricing.AnythingPriced() }
|
||||
|
||||
// Headroom is the tightest of a principal's request and token limits: how
|
||||
// much is left of each, and when the minute it is counted over rolls on.
|
||||
type Headroom struct {
|
||||
RequestLimit, RequestsLeft int
|
||||
TokenLimit, TokensLeft int
|
||||
Reset time.Duration
|
||||
}
|
||||
|
||||
// Headroom reports what a principal has left, after whatever Check just
|
||||
// charged. Of several levels the one with least left is the one that will
|
||||
// refuse first, so that is the one reported. A limit of 0 means none is set.
|
||||
func (a *Accountant) Headroom(p auth.Principal) Headroom {
|
||||
var h Headroom
|
||||
firstReq, firstTok := true, true
|
||||
for _, l := range levels(p.Limits, p.KeyID) {
|
||||
if l.RPM != nil && *l.RPM > 0 {
|
||||
left := max(0, *l.RPM-int(a.requests.Rate(l.Counter)))
|
||||
if firstReq || left < h.RequestsLeft {
|
||||
h.RequestLimit, h.RequestsLeft, firstReq = *l.RPM, left, false
|
||||
}
|
||||
}
|
||||
if l.TPM != nil && *l.TPM > 0 {
|
||||
left := max(0, *l.TPM-int(a.tokens.Rate(l.Counter)))
|
||||
if firstTok || left < h.TokensLeft {
|
||||
h.TokenLimit, h.TokensLeft, firstTok = *l.TPM, left, false
|
||||
}
|
||||
}
|
||||
}
|
||||
// The windows slide, so room comes back continuously; the next minute
|
||||
// boundary is when a whole window's worth has moved on, and is the figure
|
||||
// a client pacing itself can act on.
|
||||
now := time.Now()
|
||||
h.Reset = now.Truncate(windowLen).Add(windowLen).Sub(now)
|
||||
return h
|
||||
}
|
||||
@@ -0,0 +1,660 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func intp(v int) *int { return &v }
|
||||
func i64p(v int64) *int64 { return &v }
|
||||
|
||||
func principal(id string, tune func(*store.Effective)) auth.Principal {
|
||||
e := store.Effective{KeyID: id, ProjectID: "proj"}
|
||||
if tune != nil {
|
||||
tune(&e)
|
||||
}
|
||||
return auth.Principal{KeyID: id, ProjectID: "proj", Limits: e}
|
||||
}
|
||||
|
||||
func pricing(t *testing.T) *Pricing {
|
||||
t.Helper()
|
||||
return NewPricing(map[string]Price{
|
||||
"paid": {PromptPer1k: 1000, CompletionPer1k: 3000},
|
||||
"free": {},
|
||||
}, Price{})
|
||||
}
|
||||
|
||||
func TestCostRoundsUpSoSmallRequestsAreNotFree(t *testing.T) {
|
||||
// Rounding down would make a stream of tiny requests cost nothing, which is
|
||||
// exactly the traffic shape a quota exists to bound.
|
||||
p := Price{PromptPer1k: 1000, CompletionPer1k: 3000}
|
||||
if got := p.Cost(1, 0); got != 1 {
|
||||
t.Errorf("one prompt token cost %d", got)
|
||||
}
|
||||
if got := p.Cost(1000, 1000); got != 4000 {
|
||||
t.Errorf("a thousand of each cost %d, want 4000", got)
|
||||
}
|
||||
if got := p.Cost(0, 0); got != 0 {
|
||||
t.Errorf("an empty request cost %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestZeroPricedModelsStayAvailableWhenTheBudgetIsSpent(t *testing.T) {
|
||||
// This is what makes budget fallback work: a key with nothing left can
|
||||
// still be served something that costs nothing.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) {
|
||||
e.Budget = i64p(100)
|
||||
e.BudgetPeriod = store.PeriodMonth
|
||||
})
|
||||
|
||||
res, err := a.Check(p, "paid", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 0})
|
||||
|
||||
if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("a spent budget still served a paid model: %v", err)
|
||||
}
|
||||
if _, err := a.Check(p, "free", 10); err != nil {
|
||||
t.Errorf("a spent budget refused a free model: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetRefusalDoesNotSuggestRetrying(t *testing.T) {
|
||||
// A budget does not refill until the period rolls over, so a Retry-After of
|
||||
// a second would be a lie that turns one client into a retry loop.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) {
|
||||
e.Budget = i64p(1)
|
||||
e.BudgetPeriod = store.PeriodDay
|
||||
})
|
||||
res, _ := a.Check(p, "paid", 10)
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 5000})
|
||||
|
||||
_, err := a.Check(p, "paid", 10)
|
||||
var le *LimitError
|
||||
if !errors.As(err, &le) {
|
||||
t.Fatalf("expected a limit error, got %v", err)
|
||||
}
|
||||
if le.RetryAfter != 0 {
|
||||
t.Errorf("a spent budget suggested retrying in %v", le.RetryAfter)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitsBind(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) { e.RPM = intp(3) })
|
||||
|
||||
for i := range 3 {
|
||||
if _, err := a.Check(p, "free", 1); err != nil {
|
||||
t.Fatalf("request %d was refused: %v", i, err)
|
||||
}
|
||||
}
|
||||
if _, err := a.Check(p, "free", 1); !errors.Is(err, ErrRateLimited) {
|
||||
t.Errorf("the fourth request was allowed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenLimitsUseTheEstimateAndThenSettle(t *testing.T) {
|
||||
// The real token count only exists once the request has run, so the check
|
||||
// has to work on an estimate. A caller that consistently underestimates
|
||||
// must still end up bounded.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) { e.TPM = intp(1000) })
|
||||
|
||||
// Claims 10 tokens, actually uses 900.
|
||||
res, err := a.Check(p, "free", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 400, CompletionTokens: 500})
|
||||
|
||||
// The settled figure is what the next check sees, so a second request of
|
||||
// the same shape cannot fit.
|
||||
if _, err := a.Check(p, "free", 200); !errors.Is(err, ErrTokenLimited) {
|
||||
t.Errorf("the underestimate was never settled: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoLimitIsNotALimitOfZero(t *testing.T) {
|
||||
// The distinction the store keeps with a nil pointer. Getting it wrong here
|
||||
// would silently forbid every request for every key with no limits set.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", nil)
|
||||
|
||||
for range 50 {
|
||||
if _, err := a.Check(p, "paid", 10_000); err != nil {
|
||||
t.Fatalf("a key with no limits was refused: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnUnauthenticatedRequestIsNotAccounted(t *testing.T) {
|
||||
// With auth off there is no key to charge, and Check must not refuse.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
res, err := a.Check(auth.Principal{}, "paid", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{PromptTokens: 1000})
|
||||
if spent, _, _ := a.Status(auth.Principal{}); spent != 0 {
|
||||
t.Errorf("an anonymous request was charged %d", spent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetSurvivesARestart(t *testing.T) {
|
||||
// The counters live in memory; the store is what makes them survive.
|
||||
ctx := context.Background()
|
||||
path := filepath.Join(t.TempDir(), "acct.db")
|
||||
s, err := store.OpenSQLite(ctx, path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
must(t, s.PutOrg(ctx, store.Org{ID: "o", Name: "o"}))
|
||||
must(t, s.PutTeam(ctx, store.Team{ID: "t", OrgID: "o", Name: "t"}))
|
||||
must(t, s.PutProject(ctx, store.Project{ID: "proj", TeamID: "t", Name: "p"}))
|
||||
_, k, _ := store.NewKey("k1", "n", "proj")
|
||||
k.CreatedAt = time.Now()
|
||||
must(t, s.CreateKey(ctx, k))
|
||||
|
||||
must(t, s.WriteUsage(ctx, []store.Usage{
|
||||
{At: time.Now(), KeyID: "k1", ProjectID: "proj", Model: "paid",
|
||||
ServedModel: "paid", CostMicros: 5000, Status: 200},
|
||||
}))
|
||||
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
must(t, a.RestoreBudgets(ctx, s, store.PeriodMonth))
|
||||
|
||||
p := principal("k1", func(e *store.Effective) {
|
||||
e.Budget = i64p(5000)
|
||||
e.BudgetPeriod = store.PeriodMonth
|
||||
})
|
||||
if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("the budget was reset by the restart: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetsResetAtThePeriodBoundary(t *testing.T) {
|
||||
b := NewBudgets()
|
||||
now := time.Date(2026, 1, 31, 23, 59, 0, 0, time.UTC)
|
||||
b.now = func() time.Time { return now }
|
||||
|
||||
b.Add("k1", store.PeriodMonth, 900)
|
||||
if got := b.Spent("k1", store.PeriodMonth); got != 900 {
|
||||
t.Fatalf("spent %d", got)
|
||||
}
|
||||
|
||||
now = now.Add(2 * time.Minute) // into February
|
||||
if got := b.Spent("k1", store.PeriodMonth); got != 0 {
|
||||
t.Errorf("the counter carried %d across the month boundary", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDailyAndMonthlyBoundariesDiffer(t *testing.T) {
|
||||
b := NewBudgets()
|
||||
now := time.Date(2026, 3, 15, 12, 0, 0, 0, time.UTC)
|
||||
b.now = func() time.Time { return now }
|
||||
|
||||
b.Add("daily", store.PeriodDay, 100)
|
||||
b.Add("monthly", store.PeriodMonth, 100)
|
||||
|
||||
now = now.Add(24 * time.Hour)
|
||||
if got := b.Spent("daily", store.PeriodDay); got != 0 {
|
||||
t.Errorf("the daily counter carried %d across a day", got)
|
||||
}
|
||||
if got := b.Spent("monthly", store.PeriodMonth); got != 100 {
|
||||
t.Errorf("the monthly counter lost %d within a month", 100-got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheWindowDoesNotAllowTwiceTheLimit(t *testing.T) {
|
||||
// The reason for a sliding window. With a fixed one, a caller sends its
|
||||
// whole allowance at the end of a window and again at the start of the
|
||||
// next, so the observed rate is double the limit exactly when the system is
|
||||
// least able to take it.
|
||||
l := NewLimiter()
|
||||
base := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC)
|
||||
now := base.Add(59 * time.Second)
|
||||
l.now = func() time.Time { return now }
|
||||
|
||||
const limit = 60
|
||||
sent := 0
|
||||
for range limit {
|
||||
if l.Allow("k", limit, 1) {
|
||||
sent++
|
||||
}
|
||||
}
|
||||
if sent != limit {
|
||||
t.Fatalf("the first burst only managed %d", sent)
|
||||
}
|
||||
|
||||
// One second later, in the next fixed window.
|
||||
now = base.Add(61 * time.Second)
|
||||
second := 0
|
||||
for range limit {
|
||||
if l.Allow("k", limit, 1) {
|
||||
second++
|
||||
}
|
||||
}
|
||||
// Some allowance returns as the old window slides out of view, but nowhere
|
||||
// near the whole of it.
|
||||
if second > limit/2 {
|
||||
t.Errorf("a second burst of %d passed immediately after %d; the window is not sliding",
|
||||
second, sent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterIsSafeUnderConcurrency(t *testing.T) {
|
||||
// Checking and charging have to be one operation, or two concurrent
|
||||
// requests both see room and both take it.
|
||||
l := NewLimiter()
|
||||
const limit = 100
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
allowed := 0
|
||||
for range 500 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if l.Allow("k", limit, 1) {
|
||||
mu.Lock()
|
||||
allowed++
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if allowed != limit {
|
||||
t.Errorf("%d requests passed a limit of %d", allowed, limit)
|
||||
}
|
||||
}
|
||||
|
||||
func must(t *testing.T, err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// teamMember is a key whose budget and rate limit were set on its team, which
|
||||
// is the arrangement an enterprise actually asks for: a group gets a quota and
|
||||
// the people in it share it.
|
||||
func teamMember(keyID, teamID string, tune func(*store.Effective)) auth.Principal {
|
||||
e := store.Effective{KeyID: keyID, ProjectID: "proj", TeamID: teamID}
|
||||
owner := store.LimitOwner{Scope: store.ScopeTeam, ID: teamID}
|
||||
e.BudgetOwner, e.RPMOwner, e.TPMOwner = owner, owner, owner
|
||||
if tune != nil {
|
||||
tune(&e)
|
||||
}
|
||||
return auth.Principal{KeyID: keyID, ProjectID: "proj", Limits: e}
|
||||
}
|
||||
|
||||
func TestATeamBudgetIsOnePoolRatherThanOnePerMember(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
budget := func(e *store.Effective) {
|
||||
e.Budget = i64p(3000)
|
||||
e.BudgetPeriod = store.PeriodMonth
|
||||
}
|
||||
first := teamMember("k1", "t1", budget)
|
||||
second := teamMember("k2", "t1", budget)
|
||||
|
||||
// One member spends the lot.
|
||||
res, err := a.Check(first, "paid", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 1000})
|
||||
|
||||
// The second must now be refused. Counted per key it would have its own
|
||||
// untouched 3000 — a team budget of 3000 with ten keys would permit thirty
|
||||
// thousand, which is the bug this change exists for.
|
||||
if _, err := a.Check(second, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("the team's budget should be shared; the second member got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestATeamRateLimitIsSharedToo(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
rpm := func(e *store.Effective) { e.RPM = intp(1) }
|
||||
|
||||
if _, err := a.Check(teamMember("k1", "t1", rpm), "free", 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := a.Check(teamMember("k2", "t1", rpm), "free", 1); !errors.Is(err, ErrRateLimited) {
|
||||
t.Error("a rate limit set on the team should be shared by its keys")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTwoTeamsDoNotShareACounter(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
rpm := func(e *store.Effective) { e.RPM = intp(1) }
|
||||
|
||||
if _, err := a.Check(teamMember("k1", "t1", rpm), "free", 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Pooling is the point; pooling the wrong things is the danger. One team
|
||||
// exhausting another's quota would be worse than the bug being fixed.
|
||||
if _, err := a.Check(teamMember("k2", "t2", rpm), "free", 1); err != nil {
|
||||
t.Errorf("a second team should have its own counter, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAKeyLimitStaysPerKey(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
perKey := func(id string) auth.Principal {
|
||||
return principal(id, func(e *store.Effective) {
|
||||
e.RPM = intp(1)
|
||||
e.RPMOwner = store.LimitOwner{Scope: store.ScopeKey, ID: id}
|
||||
})
|
||||
}
|
||||
|
||||
if _, err := a.Check(perKey("k1"), "free", 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The tree still has to be able to say "each of you may do this much".
|
||||
// Sharing a group pool must not become the only thing expressible.
|
||||
if _, err := a.Check(perKey("k2"), "free", 1); err != nil {
|
||||
t.Errorf("a limit owned by the key should not be shared, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheDifferentLimitsCanBeOwnedByDifferentScopes(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
// A team budget with a per-key rate limit is the ordinary arrangement: the
|
||||
// group's money, the individual's throughput.
|
||||
mk := func(id string) auth.Principal {
|
||||
e := store.Effective{KeyID: id, TeamID: "t1"}
|
||||
e.Budget, e.BudgetPeriod = i64p(3000), store.PeriodMonth
|
||||
e.BudgetOwner = store.LimitOwner{Scope: store.ScopeTeam, ID: "t1"}
|
||||
e.RPM = intp(1)
|
||||
e.RPMOwner = store.LimitOwner{Scope: store.ScopeKey, ID: id}
|
||||
return auth.Principal{KeyID: id, Limits: e}
|
||||
}
|
||||
|
||||
res, err := a.Check(mk("k1"), "paid", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 1000})
|
||||
|
||||
// k2's rate limit is its own and untouched; the budget is the team's and
|
||||
// gone. The budget is checked first, so that is what comes back.
|
||||
if _, err := a.Check(mk("k2"), "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("expected the shared budget to bite, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusReportsTheCounterThatWillRefuseYou(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
budget := func(e *store.Effective) {
|
||||
e.Budget = i64p(9000)
|
||||
e.BudgetPeriod = store.PeriodMonth
|
||||
}
|
||||
res, err := a.Check(teamMember("k1", "t1", budget), "paid", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 1000})
|
||||
|
||||
// A member who has spent nothing themselves still sees the team's spend,
|
||||
// because that is the number that will stop them. A header disagreeing with
|
||||
// the refusal it precedes is worse than no header.
|
||||
spent, _, _ := a.Status(teamMember("k2", "t1", budget))
|
||||
if spent == 0 {
|
||||
t.Error("a member should see the pool's spend, not their own")
|
||||
}
|
||||
}
|
||||
|
||||
func TestATeamBudgetSurvivesARestart(t *testing.T) {
|
||||
// The per-key restore could not do this: it rebuilt a counter named after
|
||||
// the key, while the check reads one named after the team. The control
|
||||
// plane came back believing the team had spent nothing — a budget that
|
||||
// resets on every restart, for exactly the scope where it matters most.
|
||||
ctx := context.Background()
|
||||
s, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "team.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
must(t, s.PutOrg(ctx, store.Org{ID: "o", Name: "o"}))
|
||||
must(t, s.PutTeam(ctx, store.Team{ID: "t1", OrgID: "o", Name: "t"}))
|
||||
must(t, s.PutProject(ctx, store.Project{ID: "proj", TeamID: "t1", Name: "p"}))
|
||||
for _, id := range []string{"k1", "k2"} {
|
||||
_, k, _ := store.NewKey(id, "n", "proj")
|
||||
k.CreatedAt = time.Now()
|
||||
must(t, s.CreateKey(ctx, k))
|
||||
}
|
||||
|
||||
// Two members of one team, 2500 each: neither alone exceeds a 4000 budget,
|
||||
// together they do.
|
||||
must(t, s.WriteUsage(ctx, []store.Usage{
|
||||
{At: time.Now(), KeyID: "k1", ProjectID: "proj", Model: "paid",
|
||||
ServedModel: "paid", CostMicros: 2500, Status: 200},
|
||||
{At: time.Now(), KeyID: "k2", ProjectID: "proj", Model: "paid",
|
||||
ServedModel: "paid", CostMicros: 2500, Status: 200},
|
||||
}))
|
||||
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
must(t, a.RestoreBudgets(ctx, s, store.PeriodMonth))
|
||||
|
||||
p := teamMember("k1", "t1", func(e *store.Effective) {
|
||||
e.Budget = i64p(4000)
|
||||
e.BudgetPeriod = store.PeriodMonth
|
||||
})
|
||||
if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("the team's spend should have been restored and pooled: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// inTree is a key under org/team/project, with whatever each level sets.
|
||||
func inTree(keyID string, chain ...store.Limits) auth.Principal {
|
||||
e := store.Effective{KeyID: keyID, Chain: chain}
|
||||
// The collapsed values still exist for headers and the profile page; they
|
||||
// are not what enforcement reads any more.
|
||||
for _, l := range chain {
|
||||
if l.Budget != nil && (e.Budget == nil || *l.Budget < *e.Budget) {
|
||||
e.Budget, e.BudgetPeriod = l.Budget, l.BudgetPeriod
|
||||
e.BudgetOwner = store.LimitOwner{Scope: l.Scope, ID: l.ScopeID}
|
||||
}
|
||||
}
|
||||
return auth.Principal{KeyID: keyID, Limits: e}
|
||||
}
|
||||
|
||||
func orgBudget(id string, v int64) store.Limits {
|
||||
return store.Limits{Scope: store.ScopeOrg, ScopeID: id,
|
||||
Budget: i64p(v), BudgetPeriod: store.PeriodMonth}
|
||||
}
|
||||
|
||||
func teamBudget(id string, v int64) store.Limits {
|
||||
return store.Limits{Scope: store.ScopeTeam, ScopeID: id,
|
||||
Budget: i64p(v), BudgetPeriod: store.PeriodMonth}
|
||||
}
|
||||
|
||||
func TestTheSumOfTeamsCannotExceedTheOrganization(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
|
||||
// Each request costs 800. The organization allows 2000 and each of its four
|
||||
// teams allows 2000, so no team comes close to its own ceiling — but the
|
||||
// organization's fills after three. Taking only the most restrictive value,
|
||||
// the organization's number was never looked at and all four went through.
|
||||
//
|
||||
// A budget stops spending once it is reached rather than refusing the
|
||||
// request that would cross it, so the third is served and the fourth is not.
|
||||
spend := func(team string) error {
|
||||
p := inTree(team+"-k", orgBudget("o", 2000), teamBudget(team, 2000))
|
||||
res, err := a.Check(p, "paid", 1000)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: team + "-k", PromptTokens: 200, CompletionTokens: 200})
|
||||
return nil
|
||||
}
|
||||
|
||||
total := 0
|
||||
for _, team := range []string{"t1", "t2", "t3", "t4"} {
|
||||
if err := spend(team); err == nil {
|
||||
total++
|
||||
}
|
||||
}
|
||||
if total != 3 {
|
||||
t.Errorf("expected the organization to stop the fourth team, %d got through", total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestATeamStillStopsBeforeItsOrganizationDoes(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := func() auth.Principal {
|
||||
return inTree("k1", orgBudget("o", 100_000), teamBudget("t1", 800))
|
||||
}
|
||||
|
||||
res, err := a.Check(p(), "paid", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 200, CompletionTokens: 200})
|
||||
|
||||
// The narrower pool binds first, which is what it is for. Enforcing every
|
||||
// level must not stop the tighter one working.
|
||||
if _, err := a.Check(p(), "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("the team's own budget should still bite, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpendingFillsEveryPoolAboveIt(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := inTree("k1", orgBudget("o", 5000), teamBudget("t1", 5000))
|
||||
|
||||
res, err := a.Check(p, "paid", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 200, CompletionTokens: 200})
|
||||
|
||||
// Settled onto both, or the organization's pool would only ever fill from
|
||||
// requests that named it directly — which is none of them.
|
||||
if got := a.budgets.Spent("org:o", store.PeriodMonth); got == 0 {
|
||||
t.Error("the organization's pool did not fill")
|
||||
}
|
||||
if got := a.budgets.Spent("team:t1", store.PeriodMonth); got == 0 {
|
||||
t.Error("the team's pool did not fill")
|
||||
}
|
||||
}
|
||||
|
||||
func TestARefusalHandsBackWhatTheNarrowerLevelsTook(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
// The team allows plenty; the organization allows one request a minute.
|
||||
chain := []store.Limits{
|
||||
{Scope: store.ScopeOrg, ScopeID: "o", RPM: intp(1)},
|
||||
{Scope: store.ScopeTeam, ScopeID: "t1", RPM: intp(100)},
|
||||
}
|
||||
if _, err := a.Check(inTree("k1", chain...), "free", 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
before := a.requests.Rate("team:t1")
|
||||
|
||||
if _, err := a.Check(inTree("k2", chain...), "free", 1); !errors.Is(err, ErrRateLimited) {
|
||||
t.Fatalf("the organization should have refused the second, got %v", err)
|
||||
}
|
||||
|
||||
// A rate limiter charges as it checks, so a refusal further along has
|
||||
// already taken from the levels before it. Left there, a team at its
|
||||
// organization's ceiling would quietly eat its own rate limit with requests
|
||||
// that were never served.
|
||||
if after := a.requests.Rate("team:t1"); after != before {
|
||||
t.Errorf("the refused request stayed charged to the team: %v -> %v", before, after)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAChainWithNoLimitsAnywhereAllowsEverything(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
for i := 0; i < 5; i++ {
|
||||
if _, err := a.Check(inTree("k1"), "paid", 1000); err != nil {
|
||||
t.Fatalf("nothing is limited, but request %d was refused: %v", i, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusReportsTheFullestPoolNotTheSmallestNumber(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
chain := func(team string) auth.Principal {
|
||||
return inTree(team+"-k", orgBudget("o", 9000), teamBudget(team, 9000))
|
||||
}
|
||||
|
||||
// Another team spends; this one has spent nothing of its own.
|
||||
res, err := a.Check(chain("t1"), "paid", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "t1-k", PromptTokens: 200, CompletionTokens: 200})
|
||||
|
||||
spent, _, _ := a.Status(chain("t2"))
|
||||
// The organization's pool is the one nearer its ceiling, so it is the one
|
||||
// that will refuse them. Reporting the team's nil would disagree with the
|
||||
// refusal that eventually follows it.
|
||||
if spent == 0 {
|
||||
t.Error("a member should see the fullest pool that bounds them")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEmptyPriceTableMeansNoBudgetCanEverBind(t *testing.T) {
|
||||
// The ordinary case for self-hosted models, and the reason a budget on a
|
||||
// team can read as a limit while binding on nothing.
|
||||
if NewPricing(nil, Price{}).AnythingPriced() {
|
||||
t.Error("nothing priced should say so")
|
||||
}
|
||||
if NewPricing(map[string]Price{"m": {}}, Price{}).AnythingPriced() {
|
||||
t.Error("a model priced at zero is not priced")
|
||||
}
|
||||
if !NewPricing(map[string]Price{"m": {PromptPer1k: 1}}, Price{}).AnythingPriced() {
|
||||
t.Error("one priced model is enough for a budget to move")
|
||||
}
|
||||
if !NewPricing(nil, Price{CompletionPer1k: 1}).AnythingPriced() {
|
||||
t.Error("a fallback price prices everything")
|
||||
}
|
||||
var nilPricing *Pricing
|
||||
if nilPricing.AnythingPriced() {
|
||||
t.Error("no price table at all is not priced")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadroomIsTheTightestLimitAfterThisRequest(t *testing.T) {
|
||||
// A project allowing 100 a minute under a team allowing 3: the team is
|
||||
// what will refuse, so the team is what a client is told about.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) {
|
||||
e.Chain = []store.Limits{
|
||||
{Scope: store.ScopeProject, ScopeID: "proj", RPM: intp(100), TPM: intp(50000)},
|
||||
{Scope: store.ScopeTeam, ScopeID: "team", RPM: intp(3)},
|
||||
}
|
||||
})
|
||||
if _, err := a.Check(p, "m", 10); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := a.Headroom(p)
|
||||
if h.RequestLimit != 3 || h.RequestsLeft != 2 {
|
||||
t.Errorf("requests %d left of %d; expected 2 of the team's 3", h.RequestsLeft, h.RequestLimit)
|
||||
}
|
||||
if h.TokenLimit != 50000 || h.TokensLeft != 50000-10 {
|
||||
t.Errorf("tokens %d left of %d", h.TokensLeft, h.TokenLimit)
|
||||
}
|
||||
if h.Reset <= 0 || h.Reset > time.Minute {
|
||||
t.Errorf("reset %v", h.Reset)
|
||||
}
|
||||
if none := a.Headroom(principal("k2", nil)); none.RequestLimit != 0 || none.TokenLimit != 0 {
|
||||
t.Errorf("a key with no limits advertised some: %+v", none)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// PeriodStart is when the current budget period began.
|
||||
//
|
||||
// Periods are anchored to UTC rather than to a per-key start date, so every
|
||||
// key's day and month boundary is the same instant. A key whose period started
|
||||
// when it was created would make "spend this month" mean something different
|
||||
// for every key, and make the restore query unanswerable in one pass.
|
||||
func PeriodStart(p store.Period, now time.Time) time.Time {
|
||||
u := now.UTC()
|
||||
switch p {
|
||||
case store.PeriodDay:
|
||||
return time.Date(u.Year(), u.Month(), u.Day(), 0, 0, 0, 0, time.UTC)
|
||||
case store.PeriodMonth:
|
||||
return time.Date(u.Year(), u.Month(), 1, 0, 0, 0, 0, time.UTC)
|
||||
default:
|
||||
return time.Time{}
|
||||
}
|
||||
}
|
||||
|
||||
// Budgets tracks spend per key in memory, so a budget check costs no query.
|
||||
//
|
||||
// The database catches up through the batch writer. That means the in-memory
|
||||
// figure leads the stored one, which is the right direction: a budget that only
|
||||
// counted what had already been written would let a burst overspend by however
|
||||
// much the writer was behind.
|
||||
type Budgets struct {
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
spend map[string]counter
|
||||
}
|
||||
|
||||
type counter struct {
|
||||
micros int64
|
||||
since time.Time
|
||||
}
|
||||
|
||||
// NewBudgets builds an empty tracker.
|
||||
func NewBudgets() *Budgets {
|
||||
return &Budgets{now: time.Now, spend: map[string]counter{}}
|
||||
}
|
||||
|
||||
// Restore seeds the counters from what the store already recorded.
|
||||
//
|
||||
// A restart must not reset a budget, which it would if the counter lived only
|
||||
// in memory. The caller passes the spend since each period's start.
|
||||
func (b *Budgets) Restore(period store.Period, spend map[string]int64) {
|
||||
start := PeriodStart(period, b.now())
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
for key, micros := range spend {
|
||||
b.spend[key] = counter{micros: micros, since: start}
|
||||
}
|
||||
}
|
||||
|
||||
// Spent returns a key's spend in the current period, resetting it if the period
|
||||
// has rolled over.
|
||||
func (b *Budgets) Spent(keyID string, period store.Period) int64 {
|
||||
start := PeriodStart(period, b.now())
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
return b.currentLocked(keyID, start).micros
|
||||
}
|
||||
|
||||
// Add records spend.
|
||||
func (b *Budgets) Add(keyID string, period store.Period, micros int64) {
|
||||
if micros == 0 {
|
||||
return
|
||||
}
|
||||
start := PeriodStart(period, b.now())
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
c := b.currentLocked(keyID, start)
|
||||
c.micros += micros
|
||||
b.spend[keyID] = c
|
||||
}
|
||||
|
||||
// Exhausted reports whether a key has spent its budget.
|
||||
func (b *Budgets) Exhausted(keyID string, period store.Period, budget int64) bool {
|
||||
return b.Spent(keyID, period) >= budget
|
||||
}
|
||||
|
||||
// currentLocked returns the counter, rolled over if the period changed.
|
||||
func (b *Budgets) currentLocked(keyID string, start time.Time) counter {
|
||||
c, ok := b.spend[keyID]
|
||||
if !ok || c.since.Before(start) {
|
||||
// Either never seen, or the period rolled over while the process was
|
||||
// running. Both start from zero.
|
||||
c = counter{since: start}
|
||||
b.spend[keyID] = c
|
||||
}
|
||||
return c
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
// Package accounting turns finished requests into spend, enforces what a key
|
||||
// is allowed, and gets both into the store without touching the request path
|
||||
// (plan §7.9.3, §7.9.4).
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Price is what one model costs, in micro-units per thousand tokens.
|
||||
//
|
||||
// Micro-units rather than a float because spend is accumulated by repeated
|
||||
// addition, and repeated addition of floats drifts. What a unit means is the
|
||||
// operator's: self-hosted models have no per-token price, so it may be
|
||||
// GPU-seconds, a notional rate, or 1000 everywhere — which turns a budget into
|
||||
// a token quota.
|
||||
type Price struct {
|
||||
PromptPer1k int64 `koanf:"prompt_per_1k"`
|
||||
CompletionPer1k int64 `koanf:"completion_per_1k"`
|
||||
}
|
||||
|
||||
// Zero reports whether this model spends nothing.
|
||||
//
|
||||
// It matters beyond arithmetic: a zero-cost model stays available to a key
|
||||
// whose budget is spent, which is what makes the plan's budget fallback work
|
||||
// (§7.7).
|
||||
func (p Price) Zero() bool { return p.PromptPer1k == 0 && p.CompletionPer1k == 0 }
|
||||
|
||||
// Cost computes what a request spent, rounding up.
|
||||
//
|
||||
// Up rather than to nearest: rounding down would make a stream of tiny requests
|
||||
// free, which is exactly the shape of traffic a quota exists to bound.
|
||||
func (p Price) Cost(promptTokens, completionTokens int64) int64 {
|
||||
return divUp(promptTokens*p.PromptPer1k, 1000) + divUp(completionTokens*p.CompletionPer1k, 1000)
|
||||
}
|
||||
|
||||
func divUp(n, d int64) int64 {
|
||||
if n <= 0 {
|
||||
return 0
|
||||
}
|
||||
return (n + d - 1) / d
|
||||
}
|
||||
|
||||
// Pricing resolves a model name to its price.
|
||||
type Pricing struct {
|
||||
mu sync.RWMutex
|
||||
byModel map[string]Price
|
||||
fallback Price
|
||||
}
|
||||
|
||||
// NewPricing builds a price table. Names are matched case-insensitively,
|
||||
// because a model name reaches us from whatever the client typed.
|
||||
func NewPricing(byModel map[string]Price, fallback Price) *Pricing {
|
||||
table := make(map[string]Price, len(byModel))
|
||||
for name, p := range byModel {
|
||||
table[strings.ToLower(name)] = p
|
||||
}
|
||||
return &Pricing{byModel: table, fallback: fallback}
|
||||
}
|
||||
|
||||
// For returns a model's price, or the fallback.
|
||||
func (p *Pricing) For(model string) Price {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
if v, ok := p.byModel[strings.ToLower(model)]; ok {
|
||||
return v
|
||||
}
|
||||
return p.fallback
|
||||
}
|
||||
|
||||
// AnythingPriced reports whether any model costs anything.
|
||||
//
|
||||
// A budget is spend per period, and spend is computed from a price table. With
|
||||
// no prices every request costs nothing, so a budget written on a team is a
|
||||
// number that can never move — displayed as a limit, reported as zero percent
|
||||
// used, and binding on nothing.
|
||||
//
|
||||
// Self-hosted models have no per-token price, so an empty table is the ordinary
|
||||
// case rather than a misconfiguration. What is not ordinary is showing somebody
|
||||
// a budget without saying so.
|
||||
func (p *Pricing) AnythingPriced() bool {
|
||||
if p == nil {
|
||||
return false
|
||||
}
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
if !p.fallback.Zero() {
|
||||
return true
|
||||
}
|
||||
for _, price := range p.byModel {
|
||||
if !price.Zero() {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Limiter enforces per-key rate limits.
|
||||
//
|
||||
// A sliding window rather than a fixed one. A fixed window lets a caller send
|
||||
// its whole minute's allowance in the last second of one window and again in
|
||||
// the first second of the next, so the observed rate is twice the configured
|
||||
// limit exactly when the system is least able to take it.
|
||||
//
|
||||
// The window is approximated from two buckets, which is the usual compromise:
|
||||
// exact enough that the doubling cannot happen, and bounded memory per key
|
||||
// rather than a timestamp per request.
|
||||
type Limiter struct {
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
keys map[string]*window
|
||||
}
|
||||
|
||||
type window struct {
|
||||
// current and previous are counts in this minute and the one before.
|
||||
current, previous int64
|
||||
start time.Time
|
||||
}
|
||||
|
||||
// NewLimiter builds an empty limiter.
|
||||
func NewLimiter() *Limiter {
|
||||
return &Limiter{now: time.Now, keys: map[string]*window{}}
|
||||
}
|
||||
|
||||
const windowLen = time.Minute
|
||||
|
||||
// Allow charges n against a key's limit and reports whether it fits.
|
||||
//
|
||||
// Charging and checking are one operation on purpose: doing them separately
|
||||
// lets two concurrent requests both see room and both take it.
|
||||
func (l *Limiter) Allow(key string, limit int, n int64) bool {
|
||||
if limit <= 0 {
|
||||
// No limit configured. Zero would be a limit of zero, which is why
|
||||
// callers pass a nil pointer rather than 0 for "unset".
|
||||
return true
|
||||
}
|
||||
now := l.now()
|
||||
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
w := l.windowLocked(key, now)
|
||||
if w.rate(now)+float64(n) > float64(limit) {
|
||||
return false
|
||||
}
|
||||
w.current += n
|
||||
return true
|
||||
}
|
||||
|
||||
// Charge records usage that was already allowed, for reconciliation.
|
||||
//
|
||||
// Token limits are checked against an estimate before the request runs, because
|
||||
// the real figure only exists once it has. The difference is charged here, so a
|
||||
// caller whose requests are consistently larger than estimated is still bounded
|
||||
// — one window later than it would be with a perfect estimate.
|
||||
func (l *Limiter) Charge(key string, n int64) {
|
||||
if n == 0 {
|
||||
return
|
||||
}
|
||||
now := l.now()
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
w := l.windowLocked(key, now)
|
||||
w.current += n
|
||||
if w.current < 0 {
|
||||
w.current = 0
|
||||
}
|
||||
}
|
||||
|
||||
// Rate returns a key's current observed rate.
|
||||
func (l *Limiter) Rate(key string) float64 {
|
||||
now := l.now()
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.windowLocked(key, now).rate(now)
|
||||
}
|
||||
|
||||
func (l *Limiter) windowLocked(key string, now time.Time) *window {
|
||||
w, ok := l.keys[key]
|
||||
if !ok {
|
||||
w = &window{start: now.Truncate(windowLen)}
|
||||
l.keys[key] = w
|
||||
}
|
||||
w.roll(now)
|
||||
return w
|
||||
}
|
||||
|
||||
// roll advances the window to now.
|
||||
func (w *window) roll(now time.Time) {
|
||||
start := now.Truncate(windowLen)
|
||||
switch {
|
||||
case start.Equal(w.start):
|
||||
return
|
||||
case start.Sub(w.start) == windowLen:
|
||||
w.previous, w.current = w.current, 0
|
||||
default:
|
||||
// More than one window has passed with no traffic.
|
||||
w.previous, w.current = 0, 0
|
||||
}
|
||||
w.start = start
|
||||
}
|
||||
|
||||
// rate weights the previous window by how much of it is still in view.
|
||||
func (w *window) rate(now time.Time) float64 {
|
||||
elapsed := now.Sub(w.start)
|
||||
if elapsed >= windowLen {
|
||||
return float64(w.current)
|
||||
}
|
||||
remaining := 1 - float64(elapsed)/float64(windowLen)
|
||||
return float64(w.previous)*remaining + float64(w.current)
|
||||
}
|
||||
|
||||
// Forget drops a key's window. Used when a key is revoked; keeping it would
|
||||
// leak memory for every key that ever called.
|
||||
func (l *Limiter) Forget(key string) {
|
||||
l.mu.Lock()
|
||||
delete(l.keys, key)
|
||||
l.mu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// peopleStore has Acme/ml/training with two members, ayse and deniz.
|
||||
func peopleStore(t *testing.T) (*store.DB, string) {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
s, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "p.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = s.Close() })
|
||||
org, _ := s.CreateOrg(ctx, "Acme")
|
||||
team, _ := s.CreateTeam(ctx, org, "ml")
|
||||
proj, _ := s.CreateProject(ctx, team, "training")
|
||||
for _, id := range []string{"ayse", "deniz"} {
|
||||
if err := s.CreateUser(ctx, store.User{ID: id, Email: id + "@x", Name: id,
|
||||
PasswordHash: "h", State: store.UserPending, CreatedAt: time.Now()}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.ApproveUser(ctx, id, proj, "test", time.Now()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
return s, proj
|
||||
}
|
||||
|
||||
// keyFor makes a key for a person and resolves it the way authentication does.
|
||||
func keyFor(t *testing.T, s *store.DB, id, user, project string) auth.Principal {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
_, k, _ := store.NewKey(id, id, project)
|
||||
if err := s.CreateUserKey(ctx, user, k); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
e, err := s.Resolve(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatalf("resolving %s: %v", id, err)
|
||||
}
|
||||
return auth.Principal{KeyID: id, UserID: e.UserID, ProjectID: e.ProjectID, Limits: e}
|
||||
}
|
||||
|
||||
func TestAPersonsLimitIsOnePoolAcrossTheirKeys(t *testing.T) {
|
||||
s, proj := peopleStore(t)
|
||||
if err := s.PutLimits(context.Background(),
|
||||
store.Limits{Scope: store.ScopeUser, ScopeID: "ayse", RPM: intp(2)}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
personal := keyFor(t, s, "k-personal", "ayse", "")
|
||||
inProject := keyFor(t, s, "k-project", "ayse", proj)
|
||||
|
||||
// Two requests, one on each key, use up the person's two a minute.
|
||||
for _, p := range []auth.Principal{personal, inProject} {
|
||||
if _, err := a.Check(p, "free", 1); err != nil {
|
||||
t.Fatalf("%s: %v", p.KeyID, err)
|
||||
}
|
||||
}
|
||||
if _, err := a.Check(personal, "free", 1); !errors.Is(err, ErrRateLimited) {
|
||||
t.Errorf("a third request across her keys was allowed: %v", err)
|
||||
}
|
||||
// Deniz is somebody else.
|
||||
if _, err := a.Check(keyFor(t, s, "k-deniz", "deniz", proj), "free", 1); err != nil {
|
||||
t.Errorf("another person was limited by hers: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPerMemberLimitCountsEachMemberApart(t *testing.T) {
|
||||
s, proj := peopleStore(t)
|
||||
if err := s.PutLimits(context.Background(),
|
||||
store.Limits{Scope: store.ScopeProjectMember, ScopeID: proj, RPM: intp(1)}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
ayse := keyFor(t, s, "k-a", "ayse", proj)
|
||||
deniz := keyFor(t, s, "k-d", "deniz", proj)
|
||||
|
||||
if _, err := a.Check(ayse, "free", 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := a.Check(ayse, "free", 1); !errors.Is(err, ErrRateLimited) {
|
||||
t.Errorf("her second request passed a one-a-minute per-member limit: %v", err)
|
||||
}
|
||||
// Her share is hers: Deniz still has his.
|
||||
if _, err := a.Check(deniz, "free", 1); err != nil {
|
||||
t.Errorf("one member's use counted against another: %v", err)
|
||||
}
|
||||
// And outside the project it does not apply at all.
|
||||
if _, err := a.Check(keyFor(t, s, "k-a-own", "ayse", ""), "free", 1); err != nil {
|
||||
t.Errorf("a project's per-member limit held her personal key: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPersonsBudgetSurvivesARestart(t *testing.T) {
|
||||
s, _ := peopleStore(t)
|
||||
ctx := context.Background()
|
||||
if err := s.PutLimits(ctx, store.Limits{Scope: store.ScopeUser, ScopeID: "ayse",
|
||||
Budget: i64p(1000), BudgetPeriod: store.PeriodMonth}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Spent before the restart, personally and on the chat page.
|
||||
if err := s.WriteUsage(ctx, []store.Usage{
|
||||
{At: time.Now(), KeyID: "k-old", UserID: "ayse", CostMicros: 600, Model: "paid", ServedModel: "paid"},
|
||||
{At: time.Now(), KeyID: "session:ayse", UserID: "ayse", CostMicros: 500, Model: "paid", ServedModel: "paid"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
if err := a.RestoreBudgets(ctx, s, store.PeriodMonth); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := a.Check(keyFor(t, s, "k-new", "ayse", ""), "paid", 1); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("a restart forgot what she had spent: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// RecorderConfig tunes the batch writer.
|
||||
type RecorderConfig struct {
|
||||
// FlushInterval and BatchSize decide when a batch is written. Whichever
|
||||
// comes first.
|
||||
FlushInterval time.Duration
|
||||
BatchSize int
|
||||
// Capacity bounds the buffer. Beyond it, records are dropped rather than
|
||||
// blocking a request — see Record.
|
||||
Capacity int
|
||||
}
|
||||
|
||||
// DefaultRecorderConfig returns sensible values.
|
||||
func DefaultRecorderConfig() RecorderConfig {
|
||||
return RecorderConfig{FlushInterval: 200 * time.Millisecond, BatchSize: 500, Capacity: 8192}
|
||||
}
|
||||
|
||||
func (c *RecorderConfig) normalize() {
|
||||
d := DefaultRecorderConfig()
|
||||
if c.FlushInterval <= 0 {
|
||||
c.FlushInterval = d.FlushInterval
|
||||
}
|
||||
if c.BatchSize <= 0 {
|
||||
c.BatchSize = d.BatchSize
|
||||
}
|
||||
if c.Capacity <= 0 {
|
||||
c.Capacity = d.Capacity
|
||||
}
|
||||
}
|
||||
|
||||
// Toucher is the part of the authenticator the recorder flushes alongside
|
||||
// usage, so last-used timestamps ride the same transaction schedule.
|
||||
type Toucher interface {
|
||||
FlushTouches(ctx context.Context) error
|
||||
}
|
||||
|
||||
// Recorder batches usage records into the store.
|
||||
//
|
||||
// The buffered channel is the whole design. A request must never wait on a
|
||||
// database write, so Record is a non-blocking send and the writer drains it on
|
||||
// its own schedule.
|
||||
type Recorder struct {
|
||||
cfg RecorderConfig
|
||||
store store.Store
|
||||
touch Toucher
|
||||
log *slog.Logger
|
||||
|
||||
ch chan store.Usage
|
||||
done chan struct{}
|
||||
wg sync.WaitGroup
|
||||
|
||||
mu sync.Mutex
|
||||
dropped int64
|
||||
written int64
|
||||
}
|
||||
|
||||
// NewRecorder builds a recorder. Start begins writing.
|
||||
func NewRecorder(s store.Store, touch Toucher, cfg RecorderConfig, log *slog.Logger) *Recorder {
|
||||
cfg.normalize()
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Recorder{
|
||||
cfg: cfg, store: s, touch: touch, log: log,
|
||||
ch: make(chan store.Usage, cfg.Capacity),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Record queues one request's accounting.
|
||||
//
|
||||
// Never blocks. If the buffer is full — a database stalled long enough to fill
|
||||
// eight thousand records — the record is dropped and counted. Losing some
|
||||
// accounting is bad; making every request wait on a stalled database is worse,
|
||||
// because it converts a reporting problem into an outage.
|
||||
func (r *Recorder) Record(u store.Usage) {
|
||||
select {
|
||||
case r.ch <- u:
|
||||
default:
|
||||
r.mu.Lock()
|
||||
r.dropped++
|
||||
n := r.dropped
|
||||
r.mu.Unlock()
|
||||
if n == 1 || n%1000 == 0 {
|
||||
r.log.Warn("dropping usage records: the store is not keeping up",
|
||||
"dropped_total", n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Start runs the writer until ctx ends.
|
||||
func (r *Recorder) Start(ctx context.Context) {
|
||||
r.wg.Add(1)
|
||||
go func() {
|
||||
defer r.wg.Done()
|
||||
defer close(r.done)
|
||||
r.run(ctx)
|
||||
}()
|
||||
}
|
||||
|
||||
// Stop waits for the writer to drain and finish.
|
||||
func (r *Recorder) Stop() { r.wg.Wait() }
|
||||
|
||||
func (r *Recorder) run(ctx context.Context) {
|
||||
t := time.NewTicker(r.cfg.FlushInterval)
|
||||
defer t.Stop()
|
||||
|
||||
batch := make([]store.Usage, 0, r.cfg.BatchSize)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// Drain what is already queued before leaving. These are completed
|
||||
// requests that have already been served; throwing them away would
|
||||
// mean a clean shutdown silently loses the last moments of billing.
|
||||
r.drain(&batch)
|
||||
r.flush(context.WithoutCancel(ctx), batch)
|
||||
return
|
||||
|
||||
case u := <-r.ch:
|
||||
batch = append(batch, u)
|
||||
if len(batch) >= r.cfg.BatchSize {
|
||||
r.flush(ctx, batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
|
||||
case <-t.C:
|
||||
r.drain(&batch)
|
||||
if len(batch) > 0 {
|
||||
r.flush(ctx, batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
if err := r.touchFlush(ctx); err != nil {
|
||||
r.log.Warn("failed to write last-used timestamps", "err", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// drain moves everything queued into the batch without blocking.
|
||||
func (r *Recorder) drain(batch *[]store.Usage) {
|
||||
for len(*batch) < r.cfg.BatchSize {
|
||||
select {
|
||||
case u := <-r.ch:
|
||||
*batch = append(*batch, u)
|
||||
default:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) flush(ctx context.Context, batch []store.Usage) {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
if err := r.store.WriteUsage(ctx, batch); err != nil {
|
||||
// The records are gone either way; what matters is that this is loud.
|
||||
// Silently losing billing data is the kind of thing discovered at the
|
||||
// end of a month.
|
||||
r.log.Error("failed to write usage records", "count", len(batch), "err", err)
|
||||
return
|
||||
}
|
||||
r.mu.Lock()
|
||||
r.written += int64(len(batch))
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *Recorder) touchFlush(ctx context.Context) error {
|
||||
if r.touch == nil {
|
||||
return nil
|
||||
}
|
||||
return r.touch.FlushTouches(ctx)
|
||||
}
|
||||
|
||||
// Stats reports what the recorder has done, for metrics and tests.
|
||||
func (r *Recorder) Stats() (written, dropped int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.written, r.dropped
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// fakeStore records what was written and can be made slow or broken.
|
||||
type fakeStore struct {
|
||||
store.Store
|
||||
|
||||
mu sync.Mutex
|
||||
batches [][]store.Usage
|
||||
|
||||
block chan struct{}
|
||||
fail atomic.Bool
|
||||
touched atomic.Int64
|
||||
}
|
||||
|
||||
func (f *fakeStore) WriteUsage(_ context.Context, rows []store.Usage) error {
|
||||
if f.block != nil {
|
||||
<-f.block
|
||||
}
|
||||
if f.fail.Load() {
|
||||
return errors.New("store is down")
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.batches = append(f.batches, append([]store.Usage(nil), rows...))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeStore) rows() []store.Usage {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
var out []store.Usage
|
||||
for _, b := range f.batches {
|
||||
out = append(out, b...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (f *fakeStore) batchCount() int {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return len(f.batches)
|
||||
}
|
||||
|
||||
type fakeToucher struct{ n *atomic.Int64 }
|
||||
|
||||
func (f fakeToucher) FlushTouches(context.Context) error {
|
||||
f.n.Add(1)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestRecordsAreBatchedNotWrittenOneByOne(t *testing.T) {
|
||||
// The whole point of the buffer: a hundred requests must not be a hundred
|
||||
// transactions.
|
||||
f := &fakeStore{}
|
||||
r := NewRecorder(f, nil, RecorderConfig{FlushInterval: 20 * time.Millisecond, BatchSize: 500},
|
||||
slog.New(slog.DiscardHandler))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
|
||||
for i := range 100 {
|
||||
r.Record(store.Usage{KeyID: "k", PromptTokens: int64(i)})
|
||||
}
|
||||
waitFor(t, time.Second, func() bool { return len(f.rows()) == 100 })
|
||||
|
||||
if n := f.batchCount(); n > 5 {
|
||||
t.Errorf("100 records took %d transactions", n)
|
||||
}
|
||||
cancel()
|
||||
r.Stop()
|
||||
}
|
||||
|
||||
func TestRecordNeverBlocksARequest(t *testing.T) {
|
||||
// A stalled database must not become an outage. Losing accounting is bad;
|
||||
// making every request wait on the store is worse.
|
||||
f := &fakeStore{block: make(chan struct{})}
|
||||
r := NewRecorder(f, nil, RecorderConfig{
|
||||
FlushInterval: time.Millisecond, BatchSize: 1, Capacity: 8,
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
defer func() { close(f.block); cancel(); r.Stop() }()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
for range 1000 {
|
||||
r.Record(store.Usage{KeyID: "k"})
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Record blocked while the store was stalled")
|
||||
}
|
||||
|
||||
_, dropped := r.Stats()
|
||||
if dropped == 0 {
|
||||
t.Error("the overflow was not counted, so nobody would ever know")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAShutdownDoesNotLoseWhatIsAlreadyQueued(t *testing.T) {
|
||||
// Those are completed requests that were already served. Throwing them away
|
||||
// means a clean shutdown silently loses the last moments of billing.
|
||||
f := &fakeStore{}
|
||||
r := NewRecorder(f, nil, RecorderConfig{FlushInterval: time.Hour, BatchSize: 10_000},
|
||||
slog.New(slog.DiscardHandler))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
|
||||
for range 50 {
|
||||
r.Record(store.Usage{KeyID: "k"})
|
||||
}
|
||||
cancel()
|
||||
r.Stop()
|
||||
|
||||
if got := len(f.rows()); got != 50 {
|
||||
t.Errorf("a clean shutdown wrote %d of 50 queued records", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFailedWriteIsLoudRatherThanSilent(t *testing.T) {
|
||||
// Silently losing billing data is the kind of thing discovered at the end
|
||||
// of a month.
|
||||
f := &fakeStore{}
|
||||
f.fail.Store(true)
|
||||
logged := &countingHandler{}
|
||||
r := NewRecorder(f, nil, RecorderConfig{FlushInterval: 10 * time.Millisecond, BatchSize: 1},
|
||||
slog.New(logged))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
r.Record(store.Usage{KeyID: "k"})
|
||||
|
||||
waitFor(t, time.Second, func() bool { return logged.count(slog.LevelError) > 0 })
|
||||
cancel()
|
||||
r.Stop()
|
||||
|
||||
if written, _ := r.Stats(); written != 0 {
|
||||
t.Errorf("a failed write was counted as %d written", written)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLastUsedRidesTheSameSchedule(t *testing.T) {
|
||||
var n atomic.Int64
|
||||
f := &fakeStore{}
|
||||
r := NewRecorder(f, fakeToucher{n: &n}, RecorderConfig{FlushInterval: 10 * time.Millisecond},
|
||||
slog.New(slog.DiscardHandler))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
waitFor(t, time.Second, func() bool { return n.Load() > 0 })
|
||||
cancel()
|
||||
r.Stop()
|
||||
}
|
||||
|
||||
type countingHandler struct {
|
||||
mu sync.Mutex
|
||||
levels []slog.Level
|
||||
}
|
||||
|
||||
func (c *countingHandler) Enabled(context.Context, slog.Level) bool { return true }
|
||||
func (c *countingHandler) WithAttrs([]slog.Attr) slog.Handler { return c }
|
||||
func (c *countingHandler) WithGroup(string) slog.Handler { return c }
|
||||
func (c *countingHandler) Handle(_ context.Context, r slog.Record) error {
|
||||
c.mu.Lock()
|
||||
c.levels = append(c.levels, r.Level)
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *countingHandler) count(l slog.Level) int {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
n := 0
|
||||
for _, got := range c.levels {
|
||||
if got == l {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func waitFor(t *testing.T, limit time.Duration, pred func() bool) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(limit)
|
||||
for time.Now().Before(deadline) {
|
||||
if pred() {
|
||||
return
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("condition never became true")
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,920 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"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"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// statusRecorder captures what the control plane hears about shards, and what
|
||||
// it is asked to adopt.
|
||||
type statusRecorder struct {
|
||||
ch chan status
|
||||
|
||||
mu sync.Mutex
|
||||
calls []adoptCall
|
||||
adopted []reconciler.AdoptedShard
|
||||
unclaim []string // shard ids to report back as unclaimed
|
||||
fetches []fetchReport
|
||||
}
|
||||
|
||||
type status struct {
|
||||
replicaID, shardID string
|
||||
state supervisor.State
|
||||
err string
|
||||
}
|
||||
|
||||
func newRecorder() *statusRecorder { return &statusRecorder{ch: make(chan status, 128)} }
|
||||
|
||||
func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
|
||||
select {
|
||||
case r.ch <- status{replicaID, shardID, st, errMsg}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// fetches records download progress, so a test can assert that a fetch is
|
||||
// visible from the control plane rather than only in the node's log.
|
||||
func (r *statusRecorder) OnFetchProgress(replicaID, shardID, source string, done, total int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.fetches = append(r.fetches, fetchReport{replicaID, shardID, source, done, total})
|
||||
}
|
||||
|
||||
func (r *statusRecorder) fetchReports() []fetchReport {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return append([]fetchReport(nil), r.fetches...)
|
||||
}
|
||||
|
||||
type fetchReport struct {
|
||||
replicaID, shardID, source string
|
||||
done, total int64
|
||||
}
|
||||
|
||||
func (r *statusRecorder) Adopt(nodeID string, shards []reconciler.AdoptedShard) []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.calls = append(r.calls, adoptCall{nodeID: nodeID, shards: shards})
|
||||
r.adopted = append(r.adopted, shards...)
|
||||
|
||||
var unclaimed []reconciler.AdoptedShard
|
||||
for _, s := range shards {
|
||||
if slices.Contains(r.unclaim, s.ShardID) {
|
||||
unclaimed = append(unclaimed, s)
|
||||
}
|
||||
}
|
||||
return unclaimed
|
||||
}
|
||||
|
||||
// adoptCall records that the service reached the reconciler at all, which an
|
||||
// empty shard list used to prevent.
|
||||
type adoptCall struct {
|
||||
nodeID string
|
||||
shards []reconciler.AdoptedShard
|
||||
}
|
||||
|
||||
func (r *statusRecorder) adoptCalls() []adoptCall {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return slices.Clone(r.calls)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) adoptedShards() []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return slices.Clone(r.adopted)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) markUnclaimed(ids ...string) {
|
||||
r.mu.Lock()
|
||||
r.unclaim = append(r.unclaim, ids...)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *statusRecorder) await(t *testing.T, limit time.Duration, pred func(status) bool) status {
|
||||
t.Helper()
|
||||
deadline := time.After(limit)
|
||||
var seen []status
|
||||
for {
|
||||
select {
|
||||
case s := <-r.ch:
|
||||
seen = append(seen, s)
|
||||
if pred(s) {
|
||||
return s
|
||||
}
|
||||
case <-deadline:
|
||||
t.Fatalf("no matching status within %v; saw %+v", limit, seen)
|
||||
return status{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// fakeEngineScript writes a script that stands in for vLLM.
|
||||
//
|
||||
// It parses the real --port flag the vLLM adapter generates rather than being
|
||||
// told out of band, so the adapter's argument construction is exercised too.
|
||||
func fakeEngineScript(t *testing.T, body string) string {
|
||||
t.Helper()
|
||||
path := filepath.Join(t.TempDir(), "fake-vllm")
|
||||
script := `#!/bin/sh
|
||||
port=""
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--port) port="$2"; shift 2 ;;
|
||||
*) shift ;;
|
||||
esac
|
||||
done
|
||||
` + body + "\n"
|
||||
if err := os.WriteFile(path, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// servingEngine answers /health and /v1/models, which is what readiness needs.
|
||||
func servingEngine(t *testing.T) string {
|
||||
t.Helper()
|
||||
py, err := exec.LookPath("python3")
|
||||
if err != nil {
|
||||
t.Skip("python3 is required for the fake engine")
|
||||
}
|
||||
return fakeEngineScript(t, `exec `+py+` -c '
|
||||
import http.server, json, sys
|
||||
port = int(sys.argv[1])
|
||||
class H(http.server.BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
if self.path == "/health":
|
||||
self.send_response(200); self.end_headers(); return
|
||||
if self.path == "/v1/models":
|
||||
b = json.dumps({"object":"list","data":[{"id":"test-model"}]}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("content-type","application/json")
|
||||
self.send_header("content-length", str(len(b)))
|
||||
self.end_headers(); self.wfile.write(b); return
|
||||
self.send_response(404); self.end_headers()
|
||||
def log_message(self, *a): pass
|
||||
http.server.HTTPServer(("127.0.0.1", port), H).serve_forever()
|
||||
' "$port"`)
|
||||
}
|
||||
|
||||
type harness struct {
|
||||
svc *agentsvc.Service
|
||||
inv *inventory.Store
|
||||
rec *statusRecorder
|
||||
srv *httptest.Server
|
||||
ag *Agent
|
||||
ln *dropListener
|
||||
// streams counts how many times an agent has opened the Connect stream,
|
||||
// which is the only honest way to tell a reconnect from a stale entry.
|
||||
streams *atomic.Int64
|
||||
}
|
||||
|
||||
// dropListener tracks accepted connections so a test can sever them.
|
||||
//
|
||||
// httptest's CloseClientConnections does not work here: h2c with prior
|
||||
// knowledge hijacks the connection, and httptest stops tracking hijacked
|
||||
// connections. A test relying on it would pass without ever reconnecting.
|
||||
type dropListener struct {
|
||||
net.Listener
|
||||
mu sync.Mutex
|
||||
conns []net.Conn
|
||||
}
|
||||
|
||||
func (l *dropListener) Accept() (net.Conn, error) {
|
||||
c, err := l.Listener.Accept()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
l.mu.Lock()
|
||||
l.conns = append(l.conns, c)
|
||||
l.mu.Unlock()
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// drop severs every connection, as a control plane restart would.
|
||||
func (l *dropListener) drop() {
|
||||
l.mu.Lock()
|
||||
conns := l.conns
|
||||
l.conns = nil
|
||||
l.mu.Unlock()
|
||||
for _, c := range conns {
|
||||
_ = c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// start wires a real agent to a real service over h2c, so this exercises the
|
||||
// protocol rather than two structs calling each other.
|
||||
func start(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
return startWith(t, slog.New(slog.DiscardHandler), cfgFns...)
|
||||
}
|
||||
|
||||
// startLogged is start with logging on, for diagnosing protocol behaviour.
|
||||
func startLogged(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
return startWith(t, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug})), cfgFns...)
|
||||
}
|
||||
|
||||
func startWith(t *testing.T, log *slog.Logger, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
|
||||
inv := inventory.NewStore()
|
||||
rec := newRecorder()
|
||||
svc := agentsvc.New(inv, rec, false, log)
|
||||
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
|
||||
var streams atomic.Int64
|
||||
counting := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
streams.Add(1)
|
||||
mux.ServeHTTP(w, r)
|
||||
})
|
||||
|
||||
raw, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ln := &dropListener{Listener: raw}
|
||||
// connect's bidirectional streams need HTTP/2; h2c provides it without TLS.
|
||||
srv := httptest.NewUnstartedServer(h2c.NewHandler(counting, &http2.Server{}))
|
||||
srv.Listener = ln
|
||||
srv.Start()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "n1"
|
||||
cfg.Server = srv.URL
|
||||
cfg.MgmtAddr = "127.0.0.1"
|
||||
cfg.HeartbeatInterval = 50 * time.Millisecond
|
||||
cfg.ReconnectBase = 20 * time.Millisecond
|
||||
cfg.ReconnectMax = 100 * time.Millisecond
|
||||
cfg.Discovery.SkipGPUs = true
|
||||
cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
|
||||
cfg.Superv = supervisor.Config{
|
||||
WorkDir: t.TempDir(), LogLines: 64,
|
||||
StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond,
|
||||
}
|
||||
for _, f := range cfgFns {
|
||||
f(&cfg)
|
||||
}
|
||||
|
||||
client := &http.Client{Transport: &http2.Transport{
|
||||
AllowHTTP: true,
|
||||
DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
|
||||
return (&net.Dialer{}).DialContext(ctx, network, addr)
|
||||
},
|
||||
}}
|
||||
|
||||
ag, err := New(cfg, client, log)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() { _ = ag.Run(ctx) }()
|
||||
|
||||
h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag, ln: ln, streams: &streams}
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
ag.Supervisor().StopAll(context.Background())
|
||||
srv.Close()
|
||||
})
|
||||
h.waitRegistered(t)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *harness) waitRegistered(t *testing.T) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, ok := h.inv.Get("n1"); ok && len(h.svc.ConnectedNodes()) == 1 {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the agent never registered")
|
||||
}
|
||||
|
||||
// shardCommand builds a single-node vLLM shard on a free port.
|
||||
//
|
||||
// The model path is a real file. It used to be "/models/x", which stopped
|
||||
// working once the agent began checking that weights exist before launching an
|
||||
// engine — the check doing exactly what it is for, on a test that was relying
|
||||
// on nobody looking.
|
||||
func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand {
|
||||
t.Helper()
|
||||
modelPath := filepath.Join(t.TempDir(), "model.gguf")
|
||||
if err := os.WriteFile(modelPath, []byte("weights"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return reconciler.ShardCommand{
|
||||
ReplicaID: "rep-1", ShardID: shardID, NodeID: "n1",
|
||||
Spec: engines.ShardSpec{
|
||||
ReplicaID: "rep-1", ShardID: shardID, Rank: 0, NNodes: 1,
|
||||
GPUIndices: []int{0}, HTTPPort: freePort(t),
|
||||
Profile: engines.Profile{
|
||||
Name: "p", Engine: "vllm", TP: 1,
|
||||
ModelPath: modelPath, ServedModelName: "test-model",
|
||||
},
|
||||
},
|
||||
StartupTimeout: 15 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
func freePort(t *testing.T) int {
|
||||
t.Helper()
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ln.Close()
|
||||
return ln.Addr().(*net.TCPAddr).Port
|
||||
}
|
||||
|
||||
func TestRegistrationCarriesInventory(t *testing.T) {
|
||||
// The node must actually have the adapters it declares. It used to declare
|
||||
// mlx5_0 and mlx5_1 on a node whose sysfs held neither, which passed
|
||||
// because nothing compared the two — and that is the gap a real cluster
|
||||
// would have found as a replica hanging on its first all-reduce.
|
||||
ibRoot := ibTreeWithActivePorts(t, "mlx5_0", "mlx5_1")
|
||||
h := start(t, func(c *Config) {
|
||||
c.Hostname = "gpu-a-01"
|
||||
c.Labels = map[string]string{"rack": "r12"}
|
||||
c.Discovery.IBRoot = ibRoot
|
||||
c.Fabrics = []inventory.FabricAttachment{
|
||||
{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200},
|
||||
}
|
||||
})
|
||||
|
||||
n, ok := h.inv.Get("n1")
|
||||
if !ok {
|
||||
t.Fatal("node missing from inventory")
|
||||
}
|
||||
if n.Hostname != "gpu-a-01" || n.MgmtAddr != "127.0.0.1" {
|
||||
t.Errorf("identity wrong: %+v", n)
|
||||
}
|
||||
if n.Labels["rack"] != "r12" {
|
||||
t.Errorf("labels lost: %v", n.Labels)
|
||||
}
|
||||
// The fabric label cannot be inferred from the wire, so losing it in
|
||||
// transit would silently disable multi-node placement.
|
||||
if !n.HasFabric("ib-a") {
|
||||
t.Fatalf("fabric attachment lost: %+v", n.Fabrics)
|
||||
}
|
||||
if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
|
||||
t.Errorf("HCAs=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeartbeatKeepsNodeAlive(t *testing.T) {
|
||||
h := start(t)
|
||||
|
||||
n, _ := h.inv.Get("n1")
|
||||
first := n.LastSeen
|
||||
|
||||
// Waited for rather than slept through: four heartbeat intervals was not
|
||||
// always enough on a loaded runner for one to have landed.
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if n, _ = h.inv.Get("n1"); n.LastSeen.After(first) {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Error("heartbeats are not refreshing LastSeen; the node would be declared lost")
|
||||
}
|
||||
|
||||
func TestStartAndStopShardEndToEnd(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = servingEngine(t)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-a")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got := h.rec.await(t, 30*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-a" && s.state == supervisor.Ready
|
||||
})
|
||||
if got.replicaID != "rep-1" {
|
||||
t.Errorf("replica id lost in transit: %q", got.replicaID)
|
||||
}
|
||||
|
||||
if err := h.svc.StopShard(t.Context(), "n1", "shard-a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-a" && s.state == supervisor.Stopped
|
||||
})
|
||||
}
|
||||
|
||||
func TestFailedShardCarriesDiagnosis(t *testing.T) {
|
||||
// When a model fails to load the engine's output is the only diagnosis, so
|
||||
// it has to survive the trip to the control plane.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `echo 'CUDA out of memory' >&2; exit 1`)
|
||||
})
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), shardCommand(t, "shard-bad")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := h.rec.await(t, 30*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-bad" && s.state == supervisor.Failed
|
||||
})
|
||||
if got.err == "" {
|
||||
t.Error("the failure carried no error message")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineThatNeverBecomesReady(t *testing.T) {
|
||||
// A process that runs but never serves must time out rather than hang the
|
||||
// replica forever.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-stuck")
|
||||
cmd.StartupTimeout = 500 * time.Millisecond
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 20*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-stuck" && s.state == supervisor.Failed
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidationRejectsBeforeLaunch(t *testing.T) {
|
||||
// An engine that refuses a profile does so after minutes of loading;
|
||||
// catching it at dispatch costs nothing.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = servingEngine(t)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-invalid")
|
||||
cmd.Spec.Profile.ModelPath = "" // validation requires one
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-invalid" && s.state == supervisor.Failed
|
||||
})
|
||||
if got.err == "" {
|
||||
t.Error("validation failure carried no reason")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchToUnconnectedNode(t *testing.T) {
|
||||
h := start(t)
|
||||
|
||||
err := h.svc.StartShard(t.Context(), reconciler.ShardCommand{NodeID: "nowhere", ShardID: "x"})
|
||||
if err == nil {
|
||||
t.Fatal("dispatching to an unconnected node succeeded")
|
||||
}
|
||||
if !errors.Is(err, agentsvc.ErrNotConnected) {
|
||||
t.Errorf("err=%v, expected ErrNotConnected", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentReconnects(t *testing.T) {
|
||||
// A control plane restart must not require touching every GPU node.
|
||||
h := start(t)
|
||||
|
||||
before := h.streams.Load()
|
||||
h.ln.drop()
|
||||
|
||||
// Counting streams, not connected nodes: a stale entry in the connection
|
||||
// map would make a node look connected without any reconnect having
|
||||
// happened, and an assertion on it would pass vacuously.
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if h.streams.Load() > before && len(h.svc.ConnectedNodes()) == 1 {
|
||||
if err := h.svc.DrainShard(t.Context(), "n1", "nothing"); err != nil {
|
||||
t.Fatalf("the reconnected stream does not accept commands: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("the agent never reconnected (streams %d -> %d)", before, h.streams.Load())
|
||||
}
|
||||
|
||||
func TestIdempotentStart(t *testing.T) {
|
||||
// The control plane retries commands; a repeat must not start a second
|
||||
// process (plan §5.6).
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 30`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-idem")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-idem" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
sup := h.ag.Supervisor()
|
||||
sh, err := sup.Get("shard-idem")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
|
||||
sh2, err := sup.Get("shard-idem")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sh2.Status().PID != pid {
|
||||
t.Errorf("a second process was started: pid %d then %d", pid, sh2.Status().PID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunningShardsReportedOnReconnect covers the control plane restart hole
|
||||
// (plan §6.5): a reconnecting agent has to describe what it is already running,
|
||||
// or the control plane concludes the cluster is empty and places duplicates
|
||||
// while the originals keep holding GPUs.
|
||||
func TestRunningShardsReportedOnReconnect(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-live")
|
||||
cmd.Pool = "ib-a/h100-80g"
|
||||
cmd.Fabric = "ib-a"
|
||||
cmd.Spec.HCAs = []string{"mlx5_0"}
|
||||
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-live" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
// Drop the stream, as a control plane restart would.
|
||||
h.ln.drop()
|
||||
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if got := h.rec.adoptedShards(); len(got) > 0 {
|
||||
a := got[0]
|
||||
if a.ShardID != "shard-live" || a.ReplicaID != "rep-1" {
|
||||
t.Fatalf("wrong shard reported: %+v", a)
|
||||
}
|
||||
// Everything the control plane needs to rebuild the replica.
|
||||
if a.Model != "test-model" || a.Profile != "p" {
|
||||
t.Errorf("catalog identity lost: model=%q profile=%q", a.Model, a.Profile)
|
||||
}
|
||||
if a.Pool != "ib-a/h100-80g" || a.Fabric != "ib-a" {
|
||||
t.Errorf("pool/fabric lost: %q %q", a.Pool, a.Fabric)
|
||||
}
|
||||
if !slices.Equal(a.HCAs, []string{"mlx5_0"}) {
|
||||
t.Errorf("HCAs lost: %v", a.HCAs)
|
||||
}
|
||||
if a.Endpoint == "" {
|
||||
t.Error("rank 0 reported no endpoint; the router would have nowhere to send")
|
||||
}
|
||||
if len(a.GPUIndices) == 0 {
|
||||
t.Error("GPU indices lost; the inventory could not reclaim them")
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the reconnecting agent never reported its running shard")
|
||||
}
|
||||
|
||||
// TestUnclaimedShardIsStopped covers the other half: a shard whose model left
|
||||
// the catalog is holding GPUs nothing will ever route to.
|
||||
func TestUnclaimedShardIsStopped(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-orphaned")
|
||||
if err := h.svc.StartShard(t.Context(), cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.rec.await(t, 15*time.Second, func(s status) bool {
|
||||
return s.shardID == "shard-orphaned" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
sh, err := h.ag.Supervisor().Get("shard-orphaned")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// The control plane will refuse to claim it on reconnect.
|
||||
h.rec.markUnclaimed("shard-orphaned")
|
||||
h.ln.drop()
|
||||
|
||||
deadline := time.Now().Add(20 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAliveForTest(pid) {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Error("an unclaimed shard kept running and holding its GPUs")
|
||||
}
|
||||
|
||||
// processAliveForTest mirrors the supervisor's liveness check.
|
||||
func processAliveForTest(pid int) bool {
|
||||
return syscall.Kill(pid, 0) == nil
|
||||
}
|
||||
|
||||
func TestAShardWithMissingWeightsFailsBeforeLaunching(t *testing.T) {
|
||||
// The commonest setup mistake on a fresh node. Left to the engine it
|
||||
// surfaces minutes into a load, in the engine's own words, on a machine
|
||||
// nobody is watching — so it is caught here, where the message can say
|
||||
// what is actually wrong.
|
||||
a := &Agent{cfg: Config{WeightsCache: t.TempDir()}, log: slog.New(slog.DiscardHandler)}
|
||||
_, err := a.weights().Ensure(t.Context(), filepath.Join(t.TempDir(), "not-there.gguf"))
|
||||
if err == nil {
|
||||
t.Fatal("a missing weights file was accepted")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not-there.gguf") {
|
||||
t.Errorf("the error does not name the file: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheCacheDefaultsToTheFirstModelDirectory(t *testing.T) {
|
||||
// A node that has said where its models live has already answered this;
|
||||
// asking twice is a second thing to get out of step.
|
||||
dir := t.TempDir()
|
||||
a := &Agent{cfg: Config{ModelDirs: []string{dir, "/elsewhere"}}, log: slog.New(slog.DiscardHandler)}
|
||||
if got := a.weights().CacheDir; got != dir {
|
||||
t.Errorf("cache = %q, want the first model directory %q", got, dir)
|
||||
}
|
||||
}
|
||||
|
||||
// ibTreeWithActivePorts builds a sysfs InfiniBand tree with the named adapters,
|
||||
// each with one active port, shaped like the real thing.
|
||||
func ibTreeWithActivePorts(t *testing.T, names ...string) string {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
for _, name := range names {
|
||||
dir := filepath.Join(root, name, "ports", "1")
|
||||
if err := os.MkdirAll(filepath.Join(dir, "gids"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for file, content := range map[string]string{
|
||||
"state": "4: ACTIVE",
|
||||
"phys_state": "5: LinkUp",
|
||||
"rate": "200 Gb/sec (4X HDR)",
|
||||
"link_layer": "InfiniBand",
|
||||
"lid": "0x1a",
|
||||
"gids/0": "fec0:0000:0000:000a:0002:c903:00a0:1234",
|
||||
} {
|
||||
if err := os.WriteFile(filepath.Join(dir, file), []byte(content+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return root
|
||||
}
|
||||
|
||||
func TestAFabricWhoseAdaptersAreAbsentIsNotAdvertised(t *testing.T) {
|
||||
// The failure this prevents: a typo, or a card that enumerated differently
|
||||
// after a reboot, produced NCCL_IB_HCA naming an adapter that is not there.
|
||||
// The replica starts, passes its health checks, and hangs on its first
|
||||
// all-reduce — the exact case the fabric invariant exists to prevent,
|
||||
// arriving through a door nobody was watching.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Discovery.IBRoot = ibTreeWithActivePorts(t, "mlx5_0")
|
||||
c.Fabrics = []inventory.FabricAttachment{
|
||||
{Fabric: "ib-typo", HCAs: []string{"mlx5_9"}, RateGbps: 200},
|
||||
}
|
||||
})
|
||||
|
||||
n, ok := h.inv.Get("n1")
|
||||
if !ok {
|
||||
t.Fatal("node missing from inventory")
|
||||
}
|
||||
if n.HasFabric("ib-typo") {
|
||||
t.Errorf("a fabric with no usable adapter was advertised: %+v", n.Fabrics)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFabricRateIsTheSlowestActiveLink(t *testing.T) {
|
||||
// A collective runs at the speed of its slowest link, so reporting the
|
||||
// fastest would have the scheduler believe a fabric is quicker than it is.
|
||||
root := t.TempDir()
|
||||
for name, rate := range map[string]string{"mlx5_0": "200", "mlx5_1": "100"} {
|
||||
dir := filepath.Join(root, name, "ports", "1")
|
||||
if err := os.MkdirAll(filepath.Join(dir, "gids"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for file, content := range map[string]string{
|
||||
"state": "4: ACTIVE", "phys_state": "5: LinkUp",
|
||||
"rate": rate + " Gb/sec (4X HDR)", "link_layer": "InfiniBand",
|
||||
"lid": "0x1a", "gids/0": "fec0::1",
|
||||
} {
|
||||
if err := os.WriteFile(filepath.Join(dir, file), []byte(content+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
h := start(t, func(c *Config) {
|
||||
c.Discovery.IBRoot = root
|
||||
c.Fabrics = []inventory.FabricAttachment{
|
||||
{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}},
|
||||
}
|
||||
})
|
||||
|
||||
n, _ := h.inv.Get("n1")
|
||||
for _, f := range n.Fabrics {
|
||||
if f.Fabric == "ib-a" && f.RateGbps != 100 {
|
||||
t.Errorf("fabric rate = %d, want the slower link's 100", f.RateGbps)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPortGoingDownIsToldToTheControlPlaneOnce(t *testing.T) {
|
||||
// Checked only at registration, a port going down went unnoticed until
|
||||
// the agent restarted, and the node went on being offered multi-node
|
||||
// work that would hang on its first collective.
|
||||
root := ibTreeWithActivePorts(t, "mlx5_0")
|
||||
h := start(t, func(c *Config) {
|
||||
c.Discovery.IBRoot = root
|
||||
c.Fabrics = []inventory.FabricAttachment{{Fabric: "ib-a", HCAs: []string{"mlx5_0"}}}
|
||||
})
|
||||
if n, _ := h.inv.Get("n1"); !n.HasFabric("ib-a") {
|
||||
t.Fatalf("the fabric should be advertised at registration: %+v", n.Fabrics)
|
||||
}
|
||||
|
||||
hb := h.ag.heartbeat().GetHeartbeat()
|
||||
if hb.GetFabricsChanged() || len(hb.GetIbPorts()) != 1 || hb.GetIbPorts()[0].GetState() != "ACTIVE" {
|
||||
t.Fatalf("a quiet heartbeat should carry the port and no change: %+v", hb)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(filepath.Join(root, "mlx5_0", "ports", "1", "state"), []byte("1: DOWN\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hb = h.ag.heartbeat().GetHeartbeat()
|
||||
if !hb.GetFabricsChanged() || len(hb.GetFabrics()) != 0 {
|
||||
t.Errorf("the port going down should withdraw the fabric: changed=%v fabrics=%v",
|
||||
hb.GetFabricsChanged(), hb.GetFabrics())
|
||||
}
|
||||
if hb = h.ag.heartbeat().GetHeartbeat(); hb.GetFabricsChanged() {
|
||||
t.Error("the change should be sent once, not on every heartbeat")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAShardUsesThisNodesInterfaceAndRoCEIndex(t *testing.T) {
|
||||
// The control plane sends one interface name for every node; a node whose
|
||||
// management network is on another had NCCL bind nothing. And which RoCE
|
||||
// GID entry routes is a property of this node's adapters.
|
||||
a := &Agent{cfg: Config{MgmtInterface: "ens3"}}
|
||||
a.inv.HCAs = []discovery.HCA{{Name: "mlx5_0", Ports: []discovery.Port{
|
||||
{Num: 1, State: "ACTIVE", LinkLayer: discovery.LinkEthernet, RoCEv2GID: 3}}}}
|
||||
spec := engines.ShardSpec{MgmtInterface: "eth0", HCAs: []string{"mlx5_0"}}
|
||||
a.localizeSpec(&spec)
|
||||
if spec.MgmtInterface != "ens3" {
|
||||
t.Errorf("interface = %q", spec.MgmtInterface)
|
||||
}
|
||||
if spec.GIDIndex == nil || *spec.GIDIndex != 3 {
|
||||
t.Errorf("GID index = %v", spec.GIDIndex)
|
||||
}
|
||||
|
||||
ib := engines.ShardSpec{HCAs: []string{"mlx5_1"}}
|
||||
a.inv.HCAs = append(a.inv.HCAs, discovery.HCA{Name: "mlx5_1", Ports: []discovery.Port{
|
||||
{Num: 1, State: "ACTIVE", LinkLayer: discovery.LinkInfiniBand, RoCEv2GID: -1}}})
|
||||
a.localizeSpec(&ib)
|
||||
if ib.GIDIndex != nil {
|
||||
t.Error("an InfiniBand shard was given a GID index")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheManagementInterfaceIsFoundNotAssumed(t *testing.T) {
|
||||
// It was defaulted to eth0 before it could be looked for, so it never
|
||||
// was: a node on ens3 or bond0 had NCCL bootstrap over nothing.
|
||||
cfg := DefaultConfig()
|
||||
cfg.Server = "http://127.0.0.1:7443"
|
||||
cfg.NodeID = "n1"
|
||||
cfg.MgmtAddr = "127.0.0.1"
|
||||
if err := cfg.normalize(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.MgmtInterface == "" || cfg.MgmtInterface == "eth0" {
|
||||
t.Errorf("interface = %q, want the loopback interface holding 127.0.0.1", cfg.MgmtInterface)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEachFabricsSubnetManagersAreAskedOnce(t *testing.T) {
|
||||
// Two adapters on one fabric see one master. Its managers are listed
|
||||
// through one of them — the same list twice is a query too many — and
|
||||
// every port says which master it sees.
|
||||
oldSM, oldSMs := querySM, querySMs
|
||||
t.Cleanup(func() { querySM, querySMs = oldSM, oldSMs })
|
||||
querySM = func(string, int) (*discovery.SMInfo, error) {
|
||||
return &discovery.SMInfo{GUID: "0x01", LID: 1, State: "SMINFO_MASTER"}, nil
|
||||
}
|
||||
lists := 0
|
||||
querySMs = func(string, int) ([]discovery.SMRecord, error) {
|
||||
lists++
|
||||
return []discovery.SMRecord{{GUID: "0x01", State: "master"}, {GUID: "0x02", State: "standby"}}, nil
|
||||
}
|
||||
h := start(t, func(c *Config) {
|
||||
c.Discovery.IBRoot = ibTreeWithActivePorts(t, "mlx5_0", "mlx5_1")
|
||||
c.Fabrics = []inventory.FabricAttachment{{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}}}
|
||||
})
|
||||
hb := h.ag.heartbeat().GetHeartbeat()
|
||||
withList := 0
|
||||
for _, p := range hb.GetIbPorts() {
|
||||
if p.GetSmGuid() != "0x01" {
|
||||
t.Errorf("%s sees master %q", p.GetHca(), p.GetSmGuid())
|
||||
}
|
||||
if len(p.GetSubnetManagers()) > 0 {
|
||||
withList++
|
||||
}
|
||||
}
|
||||
if lists != 1 || withList != 1 {
|
||||
t.Errorf("listed %d times, carried on %d ports; want once, on one", lists, withList)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFabricScanIsAskedForAndAnswered(t *testing.T) {
|
||||
// One node per fabric maps it; the control plane asks and parses.
|
||||
old := scanFabric
|
||||
t.Cleanup(func() { scanFabric = old })
|
||||
var asked string
|
||||
scanFabric = func(_ context.Context, hca string, port int) (string, error) {
|
||||
asked = fmt.Sprintf("%s/%d", hca, port)
|
||||
return "Switch\t36 \"S-00000000000000a1\"\t\t# \"MF0;leaf-a:QM8700/U1\" enhanced port 0 lid 1 lmc 0\n", nil
|
||||
}
|
||||
h := start(t)
|
||||
r, err := h.svc.ScanFabric(t.Context(), "n1", "mlx5_0", 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if asked != "mlx5_0/1" || !strings.Contains(r.GetOutput(), "leaf-a") || r.GetError() != "" {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
// Config configures the agent.
|
||||
type Config struct {
|
||||
// NodeID defaults to the hostname. It must equal the common name of the
|
||||
// agent's certificate, so set it only where the two differ.
|
||||
NodeID string `koanf:"node_id"`
|
||||
// Server is the control plane, and Servers is the same thing when there is
|
||||
// more than one. Both are accepted: every configuration written before
|
||||
// there could be a second control plane names one address, and none of
|
||||
// them should have to change.
|
||||
//
|
||||
// The order is tried in order, and a refusal moves to the next. Only the
|
||||
// leader accepts an agent, so a standby is skipped in the time it takes to
|
||||
// be told so — which is what makes failover a reconnect rather than an
|
||||
// edit to every node's configuration.
|
||||
Server string `koanf:"server"`
|
||||
Servers []string `koanf:"servers"`
|
||||
// Hostname defaults to the machine's.
|
||||
Hostname string `koanf:"hostname"`
|
||||
// MgmtAddr is the address engines are reached on. It defaults to this
|
||||
// machine's address on the route to the control plane, which is the one
|
||||
// the control plane and the router can reach.
|
||||
MgmtAddr string `koanf:"mgmt_addr"`
|
||||
// Labels are free-form facts about a node for a profile's node_selector.
|
||||
// Nothing needs them: memory, GPUs, engines and disk are measured, and a
|
||||
// label restating one of those goes stale the day the machine changes.
|
||||
Labels map[string]string `koanf:"labels"`
|
||||
|
||||
// MgmtInterface is the ethernet interface NCCL bootstraps over. It
|
||||
// defaults to the interface holding MgmtAddr, then eth0.
|
||||
MgmtInterface string `koanf:"mgmt_interface"`
|
||||
// Fabrics labels this node's HCAs. It is explicit on purpose: OpenSM's
|
||||
// default subnet prefix is the same on every fabric, so nothing can be
|
||||
// inferred from the wire (plan §5.2).
|
||||
Fabrics []inventory.FabricAttachment `koanf:"fabrics"`
|
||||
|
||||
HeartbeatInterval time.Duration `koanf:"heartbeat_interval"`
|
||||
// ReconnectBase is the first reconnect delay; it backs off from there.
|
||||
ReconnectBase time.Duration `koanf:"reconnect_base"`
|
||||
ReconnectMax time.Duration `koanf:"reconnect_max"`
|
||||
|
||||
// AdoptReadyTimeout bounds how long an adopted engine is given to answer a
|
||||
// health check before the agent stops waiting on it.
|
||||
//
|
||||
// Shorter than a startup timeout on purpose: this engine is already
|
||||
// running, so it is not loading a model, it is either serving or it is not.
|
||||
AdoptReadyTimeout time.Duration `koanf:"adopt_ready_timeout"`
|
||||
|
||||
// TLS authenticates both ends of the control-plane stream. Leaving it unset
|
||||
// keeps the plaintext behaviour, which only belongs on a trusted network.
|
||||
TLS tlsconf.Config `koanf:"tls"`
|
||||
|
||||
// ModelDirs are the directories this node keeps model weights in, beyond
|
||||
// WeightsCache, which is always looked in. Storage differs at every installation,
|
||||
// so llmgw observes it rather than owning it: a Lustre mount, an image
|
||||
// exported over the fabric, local NVMe and a copy made by hand all look
|
||||
// identical from here, and placement only needs to know whether the file
|
||||
// is present.
|
||||
ModelDirs []string `koanf:"model_dirs"`
|
||||
|
||||
// EnginesDir is where the agent installs engines, one directory per
|
||||
// build. Defaults to /opt/llmgw/engines.
|
||||
EnginesDir string `koanf:"engines_dir"`
|
||||
|
||||
// WeightsCache is where a remote source is downloaded to. Defaults to the
|
||||
// first entry of ModelDirs, because a node that says where its models live
|
||||
// has already answered this question, and otherwise to
|
||||
// /var/lib/llmgw-weights.
|
||||
//
|
||||
// Fetching is the simple arrangement: every node keeps its own copy. The
|
||||
// exotic ones need no code — weights on a shared filer or an image exported
|
||||
// over the fabric are simply present, and present means nothing is fetched.
|
||||
WeightsCache string `koanf:"weights_cache"`
|
||||
|
||||
Discovery discovery.Options `koanf:"discovery"`
|
||||
Engines engines.Config `koanf:"engines"`
|
||||
Superv supervisor.Config `koanf:"supervisor"`
|
||||
|
||||
Version string
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
// No MgmtInterface: normalize finds the one holding the management
|
||||
// address, and falls back to eth0 only when it cannot. Defaulted here, it
|
||||
// was never found — every node said eth0 unless told otherwise, and a node
|
||||
// on ens3 or bond0 had NCCL bootstrap over an interface it does not have.
|
||||
return Config{
|
||||
HeartbeatInterval: 5 * time.Second,
|
||||
ReconnectBase: time.Second,
|
||||
ReconnectMax: 30 * time.Second,
|
||||
AdoptReadyTimeout: 60 * time.Second,
|
||||
Superv: supervisor.DefaultConfig(),
|
||||
// NVML gives NVLink topology, so prefer it when this build has it.
|
||||
Discovery: discovery.Options{PreferNVML: true},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() error {
|
||||
d := DefaultConfig()
|
||||
if err := c.TLS.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := c.normalizeServerURL(); err != nil {
|
||||
return err
|
||||
}
|
||||
// Everything below is found rather than written. A file that has to say
|
||||
// which machine it is on, and at what address, is a file that differs on
|
||||
// every node — and a copy taken from one node to the next registers the
|
||||
// second as the first.
|
||||
host, _ := os.Hostname()
|
||||
if c.NodeID == "" {
|
||||
c.NodeID = host
|
||||
}
|
||||
if c.NodeID == "" {
|
||||
return errors.New("node_id is not set and the hostname could not be read")
|
||||
}
|
||||
if c.Hostname == "" {
|
||||
c.Hostname = host
|
||||
}
|
||||
if c.MgmtAddr == "" {
|
||||
c.MgmtAddr = addrToward(c.Server)
|
||||
}
|
||||
if c.MgmtInterface == "" {
|
||||
c.MgmtInterface = interfaceOf(c.MgmtAddr)
|
||||
}
|
||||
if c.MgmtInterface == "" {
|
||||
c.MgmtInterface = "eth0"
|
||||
}
|
||||
// The directory downloads land in is always one that is looked in. When
|
||||
// they were two settings, a node could fill its disk with fetched weights
|
||||
// that placement never saw — which is what every node here was doing.
|
||||
if c.WeightsCache == "" && len(c.ModelDirs) > 0 {
|
||||
c.WeightsCache = c.ModelDirs[0]
|
||||
}
|
||||
if c.WeightsCache == "" {
|
||||
c.WeightsCache = DefaultWeightsCache
|
||||
}
|
||||
if !slices.Contains(c.ModelDirs, c.WeightsCache) {
|
||||
c.ModelDirs = append(c.ModelDirs, c.WeightsCache)
|
||||
}
|
||||
if c.HeartbeatInterval <= 0 {
|
||||
c.HeartbeatInterval = d.HeartbeatInterval
|
||||
}
|
||||
if c.ReconnectBase <= 0 {
|
||||
c.ReconnectBase = d.ReconnectBase
|
||||
}
|
||||
if c.AdoptReadyTimeout <= 0 {
|
||||
c.AdoptReadyTimeout = d.AdoptReadyTimeout
|
||||
}
|
||||
if c.ReconnectMax < c.ReconnectBase {
|
||||
c.ReconnectMax = max(c.ReconnectBase, d.ReconnectMax)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// normalizeServerURL makes the configured address a URL connect can use, and
|
||||
// refuses the combination that would quietly undo the certificates: a scheme
|
||||
// that says plaintext while TLS is configured. Silently connecting anyway would
|
||||
// mean a cluster believing it has mutual authentication when it has none.
|
||||
func (c *Config) normalizeServerURL() error {
|
||||
// One list, however it was written. Server first when both are given,
|
||||
// because an operator adding a standby to an existing file expects the
|
||||
// one already there to stay preferred.
|
||||
all := make([]string, 0, 1+len(c.Servers))
|
||||
if c.Server != "" {
|
||||
all = append(all, c.Server)
|
||||
}
|
||||
all = append(all, c.Servers...)
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
out := make([]string, 0, len(all))
|
||||
for _, s := range all {
|
||||
fixed, err := normalizeOne(s, c.TLS.Enabled())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// A duplicate would be tried twice in a row on failover, which reads
|
||||
// as the rotation not working.
|
||||
if _, dup := seen[fixed]; dup {
|
||||
continue
|
||||
}
|
||||
seen[fixed] = struct{}{}
|
||||
out = append(out, fixed)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return errors.New("server is required")
|
||||
}
|
||||
c.Servers = out
|
||||
c.Server = out[0]
|
||||
return nil
|
||||
}
|
||||
|
||||
// DefaultWeightsCache is where fetched weights go when nothing says otherwise.
|
||||
const DefaultWeightsCache = "/var/lib/llmgw-weights"
|
||||
|
||||
// addrToward is this machine's address on the route to the control plane.
|
||||
//
|
||||
// The address engines are reached on is whichever one the control plane can
|
||||
// reach, and the kernel already knows which that is: a UDP "connection" picks
|
||||
// the source address for the route without sending a packet. Asking the
|
||||
// interfaces instead would have to guess among docker bridges, VPNs and
|
||||
// link-local addresses.
|
||||
func addrToward(server string) string {
|
||||
u, err := url.Parse(server)
|
||||
if err != nil || u.Host == "" {
|
||||
return ""
|
||||
}
|
||||
host := u.Hostname()
|
||||
port := u.Port()
|
||||
if port == "" {
|
||||
port = "443"
|
||||
}
|
||||
conn, err := net.Dial("udp", net.JoinHostPort(host, port))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer conn.Close()
|
||||
if a, ok := conn.LocalAddr().(*net.UDPAddr); ok {
|
||||
return a.IP.String()
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// interfaceOf names the interface that holds an address, or "" when none does.
|
||||
func interfaceOf(addr string) string {
|
||||
ip := net.ParseIP(addr)
|
||||
if ip == nil {
|
||||
return ""
|
||||
}
|
||||
ifaces, err := net.Interfaces()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
for _, ifc := range ifaces {
|
||||
addrs, err := ifc.Addrs()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, a := range addrs {
|
||||
if n, ok := a.(*net.IPNet); ok && n.IP.Equal(ip) {
|
||||
return ifc.Name
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// normalizeOne applies the scheme rules to a single address.
|
||||
func normalizeOne(addr string, secure bool) (string, error) {
|
||||
switch {
|
||||
case strings.HasPrefix(addr, "https://"):
|
||||
if !secure {
|
||||
return "", fmt.Errorf("%s is an https:// address but no tls certificates are configured", addr)
|
||||
}
|
||||
case strings.HasPrefix(addr, "http://"):
|
||||
if secure {
|
||||
return "", fmt.Errorf("tls is configured but %s is an http:// address; use https://", addr)
|
||||
}
|
||||
case strings.Contains(addr, "://"):
|
||||
return "", fmt.Errorf("server %q must be an http:// or https:// address", addr)
|
||||
case secure:
|
||||
return "https://" + addr, nil
|
||||
default:
|
||||
return "http://" + addr, nil
|
||||
}
|
||||
return addr, nil
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
)
|
||||
|
||||
// ErrNoNVML means this binary was built without the NVML backend.
|
||||
//
|
||||
// It is declared here rather than in the build-tagged file so callers — and
|
||||
// tests — can refer to it in either build.
|
||||
var ErrNoNVML = errors.New("built without the nvml tag; rebuild with -tags nvml for NVLink topology")
|
||||
|
||||
// Options configure a full inventory scan.
|
||||
type Options struct {
|
||||
// NvidiaSMI is the path to nvidia-smi; empty means look it up on PATH.
|
||||
// Worth configuring on nodes where the driver lives outside the default
|
||||
// paths, which containerized driver installs often do.
|
||||
NvidiaSMI string `koanf:"nvidia_smi"`
|
||||
// IBRoot is the sysfs InfiniBand directory.
|
||||
IBRoot string `koanf:"ib_root"`
|
||||
// PreferNVML uses the NVML backend when this binary has one. NVLink
|
||||
// topology is read from nvidia-smi topo -m on every build; NVML, where it
|
||||
// is used, is the more exact source of the peers.
|
||||
PreferNVML bool `koanf:"prefer_nvml"`
|
||||
// SkipGPUs omits GPU discovery. A node without nvidia-smi is already
|
||||
// found to have none, so this is for one whose driver is broken.
|
||||
SkipGPUs bool `koanf:"skip_gpus"`
|
||||
// ModelDirs are scanned for weights. Empty means do not look.
|
||||
ModelDirs []string `koanf:"-"`
|
||||
}
|
||||
|
||||
// Discover builds the node's inventory.
|
||||
//
|
||||
// When NVML is preferred but fails, it falls back to nvidia-smi and says so in
|
||||
// Inventory.Source. The fallback is logged rather than silent: on a node that
|
||||
// is supposed to have working NVML, a quiet downgrade hides a broken driver and
|
||||
// costs the scheduler its NVLink topology.
|
||||
func Discover(ctx context.Context, opts Options, log *slog.Logger) (Inventory, error) {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
inv := Inventory{Source: "none"}
|
||||
// First, because it has nothing to do with GPUs and every path below can
|
||||
// return early. Read after them, a node with no nvidia-smi reported no
|
||||
// memory at all: GPU discovery fails, Discover returns, and the fact that
|
||||
// has nothing to do with cards goes with it. That is exactly the node whose
|
||||
// memory matters most, since a CPU profile is what it will be asked to run.
|
||||
inv.MemoryTotalMB = MemoryTotalMB()
|
||||
// Also before the GPU paths, and for the same reason: which weights a node
|
||||
// holds has nothing to do with its cards, and a node with none is exactly
|
||||
// where a CPU profile will look for them.
|
||||
inv.Weights = ScanWeights(opts.ModelDirs)
|
||||
// And the disk they sit on. A node can hold a model today and be unable to
|
||||
// take the next one, which the file list alone never shows.
|
||||
inv.Storage = ScanStorage(opts.ModelDirs)
|
||||
|
||||
if !opts.SkipGPUs {
|
||||
switch {
|
||||
case opts.PreferNVML && nvmlAvailable:
|
||||
gpus, err := DiscoverGPUsNVML(ctx)
|
||||
if err == nil {
|
||||
inv.GPUs, inv.Source = gpus, "nvml"
|
||||
break
|
||||
}
|
||||
log.Warn("NVML discovery failed, falling back to nvidia-smi", "err", err)
|
||||
fallthrough
|
||||
default:
|
||||
gpus, err := DiscoverGPUsSMI(ctx, opts.NvidiaSMI)
|
||||
if err != nil {
|
||||
return inv, fmt.Errorf("gpu discovery: %w", err)
|
||||
}
|
||||
inv.GPUs, inv.Source = gpus, "nvidia-smi"
|
||||
}
|
||||
// NVLink and which adapter sits nearest each card, for either path:
|
||||
// NVML knows NVLink but not the adapters, nvidia-smi's query knows
|
||||
// neither. A failure costs the topology and nothing else — the cards
|
||||
// are still cards.
|
||||
if len(inv.GPUs) > 0 {
|
||||
topo, err := DiscoverTopology(ctx, opts.NvidiaSMI)
|
||||
if err != nil {
|
||||
log.Warn("could not read the GPU topology; tensor-parallel groups "+
|
||||
"will be chosen without NVLink", "err", err)
|
||||
} else {
|
||||
topo.Apply(inv.GPUs)
|
||||
}
|
||||
// Multi-node NVLink: which domain each card has joined.
|
||||
if fab, err := DiscoverGPUFabric(ctx, opts.NvidiaSMI); err != nil {
|
||||
log.Warn("could not read the GPUs' NVLink fabric state", "err", err)
|
||||
} else {
|
||||
ApplyFabric(inv.GPUs, fab)
|
||||
}
|
||||
inv.IMEX = DiscoverIMEX(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
hcas, err := DiscoverIB(opts.IBRoot)
|
||||
if err != nil {
|
||||
return inv, fmt.Errorf("infiniband discovery: %w", err)
|
||||
}
|
||||
inv.HCAs = hcas
|
||||
return inv, nil
|
||||
}
|
||||
|
||||
// NVMLSupported reports whether this binary includes the NVML backend.
|
||||
func NVMLSupported() bool { return nvmlAvailable }
|
||||
@@ -0,0 +1,119 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func fakeSMI(t *testing.T) string {
|
||||
t.Helper()
|
||||
stub := filepath.Join(t.TempDir(), "nvidia-smi")
|
||||
script := "#!/bin/sh\ncat <<'EOF'\n" + smiSample + "EOF\n"
|
||||
if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return stub
|
||||
}
|
||||
|
||||
func TestDiscoverUsesSMIByDefault(t *testing.T) {
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
NvidiaSMI: fakeSMI(t),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if inv.Source != "nvidia-smi" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
if len(inv.GPUs) != 2 {
|
||||
t.Errorf("got %d GPUs", len(inv.GPUs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverFallsBackWhenNVMLUnavailable(t *testing.T) {
|
||||
// Asking for NVML in a build without it must still produce an inventory:
|
||||
// the fallback is the whole point of the build tag.
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
PreferNVML: true,
|
||||
NvidiaSMI: fakeSMI(t),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(inv.GPUs) == 0 {
|
||||
t.Fatal("no GPUs discovered through the fallback")
|
||||
}
|
||||
if NVMLSupported() {
|
||||
// On an NVML build without a driver the NVML attempt fails and the
|
||||
// fallback still has to work.
|
||||
if inv.Source != "nvml" && inv.Source != "nvidia-smi" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
} else if inv.Source != "nvidia-smi" {
|
||||
t.Errorf("a build without NVML reported source=%q", inv.Source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverSkipGPUs(t *testing.T) {
|
||||
inv, err := Discover(t.Context(), Options{
|
||||
SkipGPUs: true,
|
||||
IBRoot: fakeSysfs(t),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(inv.GPUs) != 0 {
|
||||
t.Errorf("GPUs were discovered despite SkipGPUs")
|
||||
}
|
||||
if len(inv.HCAs) != 3 {
|
||||
t.Errorf("got %d HCAs", len(inv.HCAs))
|
||||
}
|
||||
if inv.Source != "none" {
|
||||
t.Errorf("source=%q", inv.Source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverReportsGPUFailure(t *testing.T) {
|
||||
// A node meant to have GPUs whose discovery fails must say so rather than
|
||||
// reporting an empty, healthy-looking inventory.
|
||||
_, err := Discover(context.Background(), Options{
|
||||
NvidiaSMI: filepath.Join(t.TempDir(), "absent"),
|
||||
IBRoot: filepath.Join(t.TempDir(), "absent"),
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
if err == nil {
|
||||
t.Error("a failed GPU discovery was reported as success")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNVMLSupportedMatchesBuild(t *testing.T) {
|
||||
// Documents which backend this binary has, and keeps both files compiling.
|
||||
t.Logf("NVML backend compiled in: %v", NVMLSupported())
|
||||
if !NVMLSupported() {
|
||||
if _, err := DiscoverGPUsNVML(t.Context()); err != ErrNoNVML {
|
||||
t.Errorf("expected ErrNoNVML, got %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMemoryIsReportedEvenWhenGPUDiscoveryFails(t *testing.T) {
|
||||
// A node with no nvidia-smi is exactly the node whose memory matters, since
|
||||
// a CPU profile is what it will be asked to run. Reading memory after GPU
|
||||
// discovery meant that node reported none: the GPU error returns early and
|
||||
// takes a fact that has nothing to do with cards with it. Seen on node05,
|
||||
// which has 32 GB and was reporting "-".
|
||||
inv, err := Discover(t.Context(), Options{NvidiaSMI: "/nonexistent/nvidia-smi"}, nil)
|
||||
if err == nil {
|
||||
t.Fatal("GPU discovery was expected to fail with no nvidia-smi")
|
||||
}
|
||||
// On Linux this is a real number; elsewhere there is no /proc/meminfo and 0
|
||||
// is the honest answer. Either way it must not be lost to the GPU error.
|
||||
if runtime.GOOS == "linux" && inv.MemoryTotalMB <= 0 {
|
||||
t.Errorf("memory was not reported alongside a GPU failure: %d", inv.MemoryTotalMB)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// FabricCheck is one declared fabric measured against what the node actually
|
||||
// has.
|
||||
type FabricCheck struct {
|
||||
Fabric string
|
||||
// Declared are the adapters the configuration names for it.
|
||||
Declared []string
|
||||
// Usable are the declared adapters that exist and have an active port.
|
||||
Usable []string
|
||||
// RateGbps is the slowest active port among them, because a collective
|
||||
// runs at the speed of its slowest link.
|
||||
RateGbps int
|
||||
// Problems are the declared adapters that cannot carry traffic, and why.
|
||||
Problems []string
|
||||
}
|
||||
|
||||
// OK reports whether this fabric can carry a collective at all.
|
||||
func (c FabricCheck) OK() bool { return len(c.Usable) > 0 }
|
||||
|
||||
// CheckFabrics measures declared fabric membership against discovered adapters.
|
||||
//
|
||||
// Fabric membership has to be configuration — OpenSM's default subnet prefix is
|
||||
// the same on every fabric, so a matching prefix proves nothing — but taking it
|
||||
// entirely on trust was a mistake of a different kind. A typo, or a card that
|
||||
// enumerated differently after a reboot, produced NCCL_IB_HCA naming an adapter
|
||||
// that does not exist; and an adapter whose cable is out, or whose subnet
|
||||
// manager is down, is present in sysfs and carries nothing.
|
||||
//
|
||||
// Either way the replica starts, passes its health checks, and hangs on its
|
||||
// first all-reduce — which is the failure the whole fabric invariant exists to
|
||||
// prevent, arriving through a door nobody was watching. Declared and real are
|
||||
// now compared where both are visible, which is only here.
|
||||
func CheckFabrics(declared map[string][]string, discovered []HCA) []FabricCheck {
|
||||
byName := make(map[string]HCA, len(discovered))
|
||||
for _, h := range discovered {
|
||||
byName[h.Name] = h
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(declared))
|
||||
for f := range declared {
|
||||
names = append(names, f)
|
||||
}
|
||||
sort.Strings(names)
|
||||
|
||||
out := make([]FabricCheck, 0, len(names))
|
||||
for _, fabric := range names {
|
||||
c := FabricCheck{Fabric: fabric, Declared: declared[fabric]}
|
||||
for _, want := range declared[fabric] {
|
||||
hca, found := byName[want]
|
||||
if !found {
|
||||
c.Problems = append(c.Problems, fmt.Sprintf(
|
||||
"%s is not on this node (found: %s)", want, listNames(discovered)))
|
||||
continue
|
||||
}
|
||||
rate, active := activeRate(hca)
|
||||
if !active {
|
||||
c.Problems = append(c.Problems, fmt.Sprintf(
|
||||
"%s has no active port (%s)", want, portStates(hca)))
|
||||
continue
|
||||
}
|
||||
c.Usable = append(c.Usable, want)
|
||||
if c.RateGbps == 0 || rate < c.RateGbps {
|
||||
c.RateGbps = rate
|
||||
}
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// activeRate returns the fastest active port's rate on an adapter.
|
||||
//
|
||||
// Only ACTIVE ports carry traffic. A port that is physically linked but has no
|
||||
// subnet manager sits in INIT, looks healthy to anything counting adapters, and
|
||||
// passes no data — which is precisely the case worth catching, because it
|
||||
// survives a reboot and looks like a software problem.
|
||||
func activeRate(h HCA) (int, bool) {
|
||||
best, ok := 0, false
|
||||
for _, p := range h.Ports {
|
||||
if !strings.EqualFold(p.State, "ACTIVE") {
|
||||
continue
|
||||
}
|
||||
ok = true
|
||||
if p.RateGbps > best {
|
||||
best = p.RateGbps
|
||||
}
|
||||
}
|
||||
return best, ok
|
||||
}
|
||||
|
||||
func portStates(h HCA) string {
|
||||
if len(h.Ports) == 0 {
|
||||
return "no ports"
|
||||
}
|
||||
var parts []string
|
||||
for _, p := range h.Ports {
|
||||
parts = append(parts, fmt.Sprintf("port %d %s/%s", p.Num, p.State, p.PhysState))
|
||||
}
|
||||
return strings.Join(parts, ", ")
|
||||
}
|
||||
|
||||
func listNames(hcas []HCA) string {
|
||||
if len(hcas) == 0 {
|
||||
return "none"
|
||||
}
|
||||
var names []string
|
||||
for _, h := range hcas {
|
||||
names = append(names, h.Name)
|
||||
}
|
||||
return strings.Join(names, ", ")
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func hca(name string, ports ...Port) HCA { return HCA{Name: name, Ports: ports} }
|
||||
|
||||
func active(rate int) Port {
|
||||
return Port{Num: 1, State: "ACTIVE", PhysState: "LinkUp", RateGbps: rate, LinkLayer: LinkInfiniBand}
|
||||
}
|
||||
|
||||
func TestADeclaredAdapterThatIsNotThereIsReported(t *testing.T) {
|
||||
// A typo, or a card that enumerated differently after a reboot. Taken on
|
||||
// trust it becomes NCCL_IB_HCA naming an adapter that does not exist, and
|
||||
// the replica hangs on its first all-reduce.
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-a": {"mlx5_0", "mlx5_9"}},
|
||||
[]HCA{hca("mlx5_0", active(200))},
|
||||
)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("got %d checks", len(got))
|
||||
}
|
||||
c := got[0]
|
||||
if len(c.Usable) != 1 || c.Usable[0] != "mlx5_0" {
|
||||
t.Errorf("usable = %v, want only the adapter that exists", c.Usable)
|
||||
}
|
||||
if len(c.Problems) != 1 || !strings.Contains(c.Problems[0], "mlx5_9") {
|
||||
t.Errorf("problems = %v", c.Problems)
|
||||
}
|
||||
// The message names what is actually present, because the operator is
|
||||
// looking at a typo and cannot see the node.
|
||||
if !strings.Contains(c.Problems[0], "mlx5_0") {
|
||||
t.Errorf("the problem does not say what was found: %q", c.Problems[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAdapterWithNoActivePortCannotCarryTraffic(t *testing.T) {
|
||||
// A port that is physically linked but has no subnet manager sits in INIT.
|
||||
// It looks healthy to anything counting adapters and passes no data, which
|
||||
// is the case worth catching: it survives a reboot and looks like software.
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-a": {"mlx5_0"}},
|
||||
[]HCA{hca("mlx5_0", Port{Num: 1, State: "INIT", PhysState: "LinkUp", RateGbps: 200})},
|
||||
)
|
||||
c := got[0]
|
||||
if c.OK() {
|
||||
t.Error("a fabric with no active port was reported usable")
|
||||
}
|
||||
if len(c.Problems) != 1 || !strings.Contains(c.Problems[0], "INIT") {
|
||||
t.Errorf("the problem does not say what state the port is in: %v", c.Problems)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheRateIsTheSlowestActiveLink(t *testing.T) {
|
||||
// A collective runs at the speed of its slowest link, so reporting the
|
||||
// fastest would have the scheduler believe a fabric is quicker than it is.
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-a": {"fast", "slow"}},
|
||||
[]HCA{hca("fast", active(400)), hca("slow", active(100))},
|
||||
)
|
||||
if got[0].RateGbps != 100 {
|
||||
t.Errorf("rate = %d, want 100", got[0].RateGbps)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAdapterWithSeveralPortsUsesItsBestActiveOne(t *testing.T) {
|
||||
// A dual-port card with one cable in is a perfectly good adapter, and
|
||||
// judging it by the dead port would take a working fabric out of service.
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-a": {"mlx5_0"}},
|
||||
[]HCA{hca("mlx5_0",
|
||||
Port{Num: 1, State: "DOWN", PhysState: "Disabled"},
|
||||
active(200),
|
||||
)},
|
||||
)
|
||||
if !got[0].OK() {
|
||||
t.Fatalf("a card with one live port was rejected: %+v", got[0])
|
||||
}
|
||||
if got[0].RateGbps != 200 {
|
||||
t.Errorf("rate = %d, want the active port's 200", got[0].RateGbps)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFabricWithNothingUsableIsNotOK(t *testing.T) {
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-dead": {"mlx5_0"}},
|
||||
[]HCA{hca("mlx5_0", Port{Num: 1, State: "DOWN", PhysState: "Polling"})},
|
||||
)
|
||||
if got[0].OK() {
|
||||
t.Error("a fabric whose only adapter is down was reported usable")
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodeWithNoFabricsDeclaredProducesNoChecks(t *testing.T) {
|
||||
// Every node without InfiniBand, which must stay entirely unaffected.
|
||||
if got := CheckFabrics(nil, []HCA{hca("mlx5_0", active(200))}); len(got) != 0 {
|
||||
t.Errorf("got %d checks for a node that declared none", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestChecksComeBackInAStableOrder(t *testing.T) {
|
||||
// Map iteration would otherwise reorder log lines between restarts, which
|
||||
// makes a diff of two registrations unreadable.
|
||||
declared := map[string][]string{"ib-c": {"m"}, "ib-a": {"m"}, "ib-b": {"m"}}
|
||||
got := CheckFabrics(declared, nil)
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("got %d", len(got))
|
||||
}
|
||||
if got[0].Fabric != "ib-a" || got[1].Fabric != "ib-b" || got[2].Fabric != "ib-c" {
|
||||
t.Errorf("order = %s, %s, %s", got[0].Fabric, got[1].Fabric, got[2].Fabric)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// GPU is one accelerator on this node.
|
||||
type GPU struct {
|
||||
Index int `json:"index"`
|
||||
UUID string `json:"uuid"`
|
||||
Name string `json:"name"`
|
||||
// MemoryTotalMB is the total framebuffer size.
|
||||
MemoryTotalMB int `json:"memory_total_mb"`
|
||||
// MemoryUsedMB is what is already allocated. Non-zero on an idle node means
|
||||
// something other than our engines is using the card (plan §5.1).
|
||||
MemoryUsedMB int `json:"memory_used_mb"`
|
||||
// UtilizationPercent is how busy the card was over the last sample period.
|
||||
//
|
||||
// The number people actually watch, and the one nothing collected: memory
|
||||
// used says a model is loaded, and this says whether it is doing anything.
|
||||
// A fleet of cards at ninety percent occupancy and two percent utilisation
|
||||
// is a fleet holding models nobody is asking for, which is invisible
|
||||
// without it.
|
||||
UtilizationPercent int `json:"utilization_percent,omitempty"`
|
||||
// ComputeCapability gates quantization choices: an FP8 profile must not
|
||||
// land on a card below 8.9 (plan §6.2).
|
||||
ComputeCapability string `json:"compute_capability"`
|
||||
NUMANode int `json:"numa_node"`
|
||||
PCIBusID string `json:"pci_bus_id,omitempty"`
|
||||
DriverVersion string `json:"driver_version,omitempty"`
|
||||
CUDAVersion string `json:"cuda_version,omitempty"`
|
||||
// NVLinkPeers are the indices of GPUs reachable over NVLink. Within a node
|
||||
// a TP group should stay inside one NVLink domain.
|
||||
NVLinkPeers []int `json:"nvlink_peers,omitempty"`
|
||||
// NICPaths is how far each network adapter is from this GPU, by HCA name:
|
||||
// PIX, PXB, PHB, NODE or SYS, as nvidia-smi topo reports it. What decides
|
||||
// which adapter a GPU's share of a multi-node job should go through.
|
||||
NICPaths map[string]string `json:"nic_paths,omitempty"`
|
||||
// Fabric is the GPU's multi-node NVLink membership, nil where the
|
||||
// driver reports none.
|
||||
Fabric *GPUFabric `json:"fabric,omitempty"`
|
||||
// Health carries anything that should keep the GPU out of placement.
|
||||
Health GPUHealth `json:"health"`
|
||||
}
|
||||
|
||||
// GPUHealth records conditions that make a GPU unsafe to schedule onto.
|
||||
type GPUHealth struct {
|
||||
// ECCErrors is the count of uncorrected ECC errors.
|
||||
ECCErrors int `json:"ecc_errors"`
|
||||
// TemperatureC is the current die temperature.
|
||||
TemperatureC int `json:"temperature_c"`
|
||||
// Problems lists human-readable reasons to avoid this GPU.
|
||||
Problems []string `json:"problems,omitempty"`
|
||||
}
|
||||
|
||||
// OK reports whether the GPU is fit for placement.
|
||||
func (h GPUHealth) OK() bool { return len(h.Problems) == 0 }
|
||||
|
||||
// nonAlnum collapses everything that is not a letter or digit.
|
||||
var nonAlnum = regexp.MustCompile(`[^a-z0-9]+`)
|
||||
|
||||
// Class derives the GPU class used for placement: same model, same memory.
|
||||
//
|
||||
// Two cards only belong in the same tensor-parallel group if they are the same
|
||||
// class, so the derivation has to be stable across driver versions and naming
|
||||
// quirks ("NVIDIA H100 80GB HBM3" and "H100-SXM5-80GB" are the same class).
|
||||
func (g GPU) Class() string {
|
||||
name := strings.ToLower(g.Name)
|
||||
name = strings.TrimPrefix(name, "nvidia ")
|
||||
name = strings.TrimPrefix(name, "tesla ")
|
||||
|
||||
// Keep the model token: the first word that contains a digit, which is how
|
||||
// NVIDIA names every accelerator (h100, a100, l40s, rtx4090).
|
||||
var model string
|
||||
for _, tok := range nonAlnum.Split(name, -1) {
|
||||
if tok == "" || !strings.ContainsFunc(tok, func(r rune) bool { return r >= '0' && r <= '9' }) {
|
||||
continue
|
||||
}
|
||||
// Skip memory and form-factor tokens, which are not the model.
|
||||
if strings.HasSuffix(tok, "gb") || tok == "80gb" || tok == "sxm4" || tok == "sxm5" || tok == "pcie" {
|
||||
continue
|
||||
}
|
||||
model = tok
|
||||
break
|
||||
}
|
||||
if model == "" {
|
||||
model = nonAlnum.ReplaceAllString(name, "-")
|
||||
}
|
||||
|
||||
// Round to the nearest GiB: vendors report 81559 MiB for an "80GB" card.
|
||||
gb := int(math.Round(float64(g.MemoryTotalMB) / 1024))
|
||||
return fmt.Sprintf("%s-%dg", model, gb)
|
||||
}
|
||||
|
||||
// MarshalJSON adds the derived class, so a consumer of the JSON does not have
|
||||
// to reimplement the derivation to know which GPUs are interchangeable.
|
||||
func (g GPU) MarshalJSON() ([]byte, error) {
|
||||
type alias GPU // prevents recursing back into this method
|
||||
return json.Marshal(struct {
|
||||
alias
|
||||
Class string `json:"class"`
|
||||
}{alias(g), g.Class()})
|
||||
}
|
||||
|
||||
// Inventory is everything discovery found on this node.
|
||||
type Inventory struct {
|
||||
GPUs []GPU `json:"gpus"`
|
||||
HCAs []HCA `json:"hcas"`
|
||||
// Source names the backend that produced the GPU list, so a consumer can
|
||||
// tell "no GPUs" from "discovery did not work".
|
||||
Source string `json:"source"`
|
||||
// MemoryTotalMB is the machine's RAM. Zero means unknown, which is what an
|
||||
// agent older than this field reports, and is treated as no constraint.
|
||||
MemoryTotalMB int `json:"memory_total_mb,omitempty"`
|
||||
// Weights are the model files this node can see. Empty means it was not
|
||||
// asked to look.
|
||||
Weights []LocalWeights `json:"weights,omitempty"`
|
||||
// Storage is the disk behind those weights, one entry per filesystem.
|
||||
// Empty means the same thing it means for Weights: not asked.
|
||||
Storage []Storage `json:"storage,omitempty"`
|
||||
// IMEX is the service a multi-node NVLink domain needs, read where the
|
||||
// node has GPUs.
|
||||
IMEX IMEX `json:"imex"`
|
||||
}
|
||||
|
||||
// Classes returns the distinct GPU classes present, sorted.
|
||||
func (inv Inventory) Classes() []string {
|
||||
seen := map[string]struct{}{}
|
||||
for _, g := range inv.GPUs {
|
||||
seen[g.Class()] = struct{}{}
|
||||
}
|
||||
out := make([]string, 0, len(seen))
|
||||
for c := range seen {
|
||||
out = append(out, c)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// parseCSVInt reads an integer from nvidia-smi's CSV output, tolerating the
|
||||
// "[N/A]" and "Not Supported" values it emits for unsupported queries.
|
||||
func parseCSVInt(s string) (int, bool) {
|
||||
s = strings.TrimSpace(s)
|
||||
s = strings.TrimSuffix(s, " MiB")
|
||||
s = strings.TrimSuffix(s, " W")
|
||||
if s == "" || strings.HasPrefix(s, "[") || strings.EqualFold(s, "Not Supported") {
|
||||
return 0, false
|
||||
}
|
||||
v, err := strconv.Atoi(s)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return v, true
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// GPUFabric is a GPU's place in a multi-node NVLink domain, as the driver
|
||||
// reports it: the NVL72 rack, or the NVLink Switch system, it belongs to.
|
||||
type GPUFabric struct {
|
||||
// State is Completed once the fabric manager has joined the GPU to its
|
||||
// domain; In Progress, Not Started or Not Supported otherwise.
|
||||
State string `json:"state"`
|
||||
// Status is Success, or the reason the join failed.
|
||||
Status string `json:"status,omitempty"`
|
||||
// ClusterUUID and CliqueID name the domain: GPUs sharing both can reach
|
||||
// each other over NVLink whatever node they are in.
|
||||
ClusterUUID string `json:"cluster_uuid,omitempty"`
|
||||
CliqueID string `json:"clique_id,omitempty"`
|
||||
// Bandwidth is Full, or how the domain's NVLink is degraded.
|
||||
Bandwidth string `json:"bandwidth,omitempty"`
|
||||
}
|
||||
|
||||
// Domain is the NVLink domain the GPU is joined to, empty when it is in none
|
||||
// that crosses nodes: a card with no fabric, one not joined, or one whose
|
||||
// cluster is all zeros, which is what a single machine's NVSwitch reports.
|
||||
func (f *GPUFabric) Domain() string {
|
||||
if f == nil || !strings.EqualFold(f.State, "Completed") {
|
||||
return ""
|
||||
}
|
||||
u := strings.ToLower(strings.TrimSpace(f.ClusterUUID))
|
||||
if u == "" || strings.Trim(u, "0-") == "" {
|
||||
return ""
|
||||
}
|
||||
return u + "/" + f.CliqueID
|
||||
}
|
||||
|
||||
// ParseGPUFabric reads the Fabric sections of `nvidia-smi -q`, by the PCI
|
||||
// address each GPU's section starts with.
|
||||
func ParseGPUFabric(out string) map[string]GPUFabric {
|
||||
res := map[string]GPUFabric{}
|
||||
var bus string
|
||||
fabricIndent := -1
|
||||
var cur *GPUFabric
|
||||
flush := func() {
|
||||
if bus != "" && cur != nil && cur.State != "" {
|
||||
res[bus] = *cur
|
||||
}
|
||||
cur, fabricIndent = nil, -1
|
||||
}
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
trim := strings.TrimSpace(line)
|
||||
indent := len(line) - len(strings.TrimLeft(line, " \t"))
|
||||
if strings.HasPrefix(line, "GPU ") && indent == 0 {
|
||||
flush()
|
||||
bus = strings.ToLower(strings.TrimSpace(strings.TrimPrefix(line, "GPU ")))
|
||||
continue
|
||||
}
|
||||
if trim == "Fabric" {
|
||||
flush()
|
||||
cur, fabricIndent = &GPUFabric{}, indent
|
||||
continue
|
||||
}
|
||||
if cur == nil {
|
||||
continue
|
||||
}
|
||||
if trim != "" && indent <= fabricIndent {
|
||||
flush()
|
||||
continue
|
||||
}
|
||||
k, v, ok := strings.Cut(trim, ":")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
k, v = strings.TrimSpace(k), strings.TrimSpace(v)
|
||||
if strings.EqualFold(v, "N/A") {
|
||||
v = ""
|
||||
}
|
||||
switch k {
|
||||
case "State":
|
||||
cur.State = v
|
||||
case "Status":
|
||||
cur.Status = v
|
||||
case "ClusterUUID", "Cluster UUID":
|
||||
cur.ClusterUUID = v
|
||||
case "CliqueId", "Clique ID":
|
||||
cur.CliqueID = v
|
||||
case "Bandwidth":
|
||||
cur.Bandwidth = v
|
||||
}
|
||||
}
|
||||
flush()
|
||||
return res
|
||||
}
|
||||
|
||||
// DiscoverGPUFabric runs `nvidia-smi -q` for the Fabric sections. Nothing
|
||||
// is reported on a driver or a card without them.
|
||||
func DiscoverGPUFabric(ctx context.Context, binary string) (map[string]GPUFabric, error) {
|
||||
if binary == "" {
|
||||
if _, err := exec.LookPath("nvidia-smi"); err != nil {
|
||||
return nil, nil
|
||||
}
|
||||
binary = "nvidia-smi"
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
|
||||
defer cancel()
|
||||
out, err := exec.CommandContext(ctx, binary, "-q").Output()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ParseGPUFabric(string(out)), nil
|
||||
}
|
||||
|
||||
// ApplyFabric puts each GPU's fabric on it, matched by PCI address.
|
||||
func ApplyFabric(gpus []GPU, fab map[string]GPUFabric) {
|
||||
for i := range gpus {
|
||||
if f, ok := fab[strings.ToLower(gpus[i].PCIBusID)]; ok {
|
||||
f := f
|
||||
gpus[i].Fabric = &f
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// IMEX is the service that lets GPUs in different nodes of one NVLink
|
||||
// domain share memory: without it, multi-node NVLink carries nothing.
|
||||
type IMEX struct {
|
||||
// State is systemd's word for nvidia-imex: active, inactive, or empty
|
||||
// where it is not installed.
|
||||
State string `json:"state,omitempty"`
|
||||
// Channels is how many IMEX channels exist; a process needs one.
|
||||
Channels int `json:"channels"`
|
||||
}
|
||||
|
||||
// DiscoverIMEX reads the IMEX service and its channels.
|
||||
func DiscoverIMEX(ctx context.Context) IMEX {
|
||||
var m IMEX
|
||||
if entries, err := os.ReadDir("/dev/nvidia-caps-imex-channels"); err == nil {
|
||||
m.Channels = len(entries)
|
||||
}
|
||||
if _, err := exec.LookPath("systemctl"); err == nil {
|
||||
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
out, _ := exec.CommandContext(ctx, "systemctl", "is-active", "nvidia-imex").Output()
|
||||
m.State = strings.TrimSpace(string(out))
|
||||
if m.State == "inactive" {
|
||||
// "inactive" is also what systemd says of a unit that does not
|
||||
// exist; the binary says which.
|
||||
if _, err := exec.LookPath("nvidia-imex"); err != nil {
|
||||
if _, err := os.Stat(filepath.Join("/usr/bin", "nvidia-imex")); err != nil {
|
||||
m.State = ""
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package discovery
|
||||
|
||||
import "testing"
|
||||
|
||||
// smiQ is an excerpt of `nvidia-smi -q` on a GB200 tray: two GPUs joined to
|
||||
// one NVL72 domain, then a PCIe card that has no fabric at all.
|
||||
const smiQ = `
|
||||
==============NVSMI LOG==============
|
||||
|
||||
Timestamp : Mon Sep 29 10:00:00 2026
|
||||
Driver Version : 570.86.15
|
||||
Attached GPUs : 3
|
||||
GPU 00000008:01:00.0
|
||||
Product Name : NVIDIA GB200
|
||||
Minor Number : 0
|
||||
Fabric
|
||||
State : Completed
|
||||
Status : Success
|
||||
CliqueId : 32766
|
||||
ClusterUUID : 76B1AD76-7c24-4b6b-9f6c-2a1e0d8c11aa
|
||||
Health
|
||||
Bandwidth : Full
|
||||
Processes : None
|
||||
|
||||
GPU 00000009:01:00.0
|
||||
Product Name : NVIDIA GB200
|
||||
Fabric
|
||||
State : In Progress
|
||||
Status : N/A
|
||||
CliqueId : N/A
|
||||
ClusterUUID : N/A
|
||||
Processes : None
|
||||
|
||||
GPU 0000000A:01:00.0
|
||||
Product Name : NVIDIA L40S
|
||||
Fabric
|
||||
State : N/A
|
||||
Status : N/A
|
||||
`
|
||||
|
||||
func TestAGPUsNVLinkDomainIsRead(t *testing.T) {
|
||||
fab := ParseGPUFabric(smiQ)
|
||||
a, ok := fab["00000008:01:00.0"]
|
||||
if !ok {
|
||||
t.Fatalf("fabric = %+v", fab)
|
||||
}
|
||||
if a.State != "Completed" || a.CliqueID != "32766" || a.Bandwidth != "Full" {
|
||||
t.Errorf("GB200 = %+v", a)
|
||||
}
|
||||
if got := a.Domain(); got != "76b1ad76-7c24-4b6b-9f6c-2a1e0d8c11aa/32766" {
|
||||
t.Errorf("domain = %q", got)
|
||||
}
|
||||
// Still joining: no domain yet, and said as such.
|
||||
b := fab["00000009:01:00.0"]
|
||||
if b.State != "In Progress" || b.Domain() != "" {
|
||||
t.Errorf("joining GPU = %+v", b)
|
||||
}
|
||||
// A card with no fabric reports nothing rather than an empty one.
|
||||
if _, ok := fab["0000000a:01:00.0"]; ok {
|
||||
t.Error("a card without a fabric was given one")
|
||||
}
|
||||
|
||||
gpus := []GPU{{Index: 0, PCIBusID: "00000008:01:00.0"}, {Index: 1, PCIBusID: "0000000A:01:00.0"}}
|
||||
ApplyFabric(gpus, fab)
|
||||
if gpus[0].Fabric == nil || gpus[1].Fabric != nil {
|
||||
t.Errorf("applied = %+v %+v", gpus[0].Fabric, gpus[1].Fabric)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASingleMachinesNVSwitchIsNoDomain(t *testing.T) {
|
||||
// An eight-GPU HGX reports its fabric joined, with a cluster of zeros: its
|
||||
// NVLink ends at the chassis.
|
||||
f := &GPUFabric{State: "Completed", ClusterUUID: "00000000-0000-0000-0000-000000000000", CliqueID: "0"}
|
||||
if f.Domain() != "" {
|
||||
t.Errorf("a single machine read as a domain: %q", f.Domain())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
//go:build !nvml
|
||||
|
||||
package discovery
|
||||
|
||||
import "context"
|
||||
|
||||
// nvmlAvailable reports that this build cannot use NVML.
|
||||
const nvmlAvailable = false
|
||||
|
||||
// DiscoverGPUsNVML is unavailable in builds without the nvml tag.
|
||||
func DiscoverGPUsNVML(context.Context) ([]GPU, error) { return nil, ErrNoNVML }
|
||||
@@ -0,0 +1,148 @@
|
||||
//go:build nvml
|
||||
|
||||
// This file is the NVML-backed GPU discovery path, built only with the `nvml`
|
||||
// tag because the bindings require cgo. Without the tag the agent falls back to
|
||||
// nvidia-smi, which keeps the server side — and development machines without a
|
||||
// C toolchain or a driver — building cleanly (plan §3.2).
|
||||
//
|
||||
// What NVML adds over nvidia-smi is NVLink topology. Nothing else exposes which
|
||||
// GPUs share an NVLink domain, and that is what keeps a tensor-parallel group
|
||||
// off the slow path inside a node (plan §6.2, constraint 3).
|
||||
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/NVIDIA/go-nvml/pkg/nvml"
|
||||
)
|
||||
|
||||
// nvmlAvailable reports that this build can use NVML.
|
||||
const nvmlAvailable = true
|
||||
|
||||
// DiscoverGPUsNVML enumerates GPUs through NVML.
|
||||
//
|
||||
// It returns an error rather than falling back on its own: the caller decides
|
||||
// whether to degrade to nvidia-smi, because on a node that is supposed to have
|
||||
// working NVML a silent downgrade hides a broken driver installation.
|
||||
func DiscoverGPUsNVML(_ context.Context) ([]GPU, error) {
|
||||
if ret := nvml.Init(); ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml init: %s", nvml.ErrorString(ret))
|
||||
}
|
||||
defer func() { _ = nvml.Shutdown() }()
|
||||
|
||||
driver, _ := nvml.SystemGetDriverVersion()
|
||||
cuda := ""
|
||||
if v, ret := nvml.SystemGetCudaDriverVersion(); ret == nvml.SUCCESS {
|
||||
// NVML reports 12040 for 12.4.
|
||||
cuda = fmt.Sprintf("%d.%d", v/1000, (v%1000)/10)
|
||||
}
|
||||
|
||||
count, ret := nvml.DeviceGetCount()
|
||||
if ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml device count: %s", nvml.ErrorString(ret))
|
||||
}
|
||||
|
||||
// Bus id to index, so NVLink peers can be reported as indices. NVML gives
|
||||
// the remote end of a link as a PCI address, which is not what placement
|
||||
// works in.
|
||||
byBusID := make(map[string]int, count)
|
||||
gpus := make([]GPU, 0, count)
|
||||
|
||||
for i := range count {
|
||||
dev, ret := nvml.DeviceGetHandleByIndex(i)
|
||||
if ret != nvml.SUCCESS {
|
||||
return nil, fmt.Errorf("nvml device %d: %s", i, nvml.ErrorString(ret))
|
||||
}
|
||||
|
||||
g := GPU{Index: i, NUMANode: -1, DriverVersion: driver, CUDAVersion: cuda}
|
||||
g.UUID, _ = dev.GetUUID()
|
||||
g.Name, _ = dev.GetName()
|
||||
|
||||
if mem, ret := dev.GetMemoryInfo(); ret == nvml.SUCCESS {
|
||||
g.MemoryTotalMB = int(mem.Total / (1 << 20))
|
||||
g.MemoryUsedMB = int(mem.Used / (1 << 20))
|
||||
}
|
||||
if major, minor, ret := dev.GetCudaComputeCapability(); ret == nvml.SUCCESS {
|
||||
g.ComputeCapability = fmt.Sprintf("%d.%d", major, minor)
|
||||
}
|
||||
if pci, ret := dev.GetPciInfo(); ret == nvml.SUCCESS {
|
||||
g.PCIBusID = busIDString(pci.BusId[:])
|
||||
byBusID[normalizeBusID(g.PCIBusID)] = i
|
||||
}
|
||||
if t, ret := dev.GetTemperature(nvml.TEMPERATURE_GPU); ret == nvml.SUCCESS {
|
||||
g.Health.TemperatureC = int(t)
|
||||
}
|
||||
// Volatile counters reset on driver reload, which is what we want: an
|
||||
// error from three months and one reboot ago is not a reason to refuse
|
||||
// a card today.
|
||||
if n, ret := dev.GetTotalEccErrors(nvml.MEMORY_ERROR_TYPE_UNCORRECTED, nvml.VOLATILE_ECC); ret == nvml.SUCCESS {
|
||||
g.Health.ECCErrors = int(n)
|
||||
}
|
||||
|
||||
assessHealth(&g)
|
||||
gpus = append(gpus, g)
|
||||
}
|
||||
|
||||
// A second pass, because a link's remote end may be a GPU enumerated later.
|
||||
for i := range gpus {
|
||||
dev, ret := nvml.DeviceGetHandleByIndex(gpus[i].Index)
|
||||
if ret != nvml.SUCCESS {
|
||||
continue
|
||||
}
|
||||
gpus[i].NVLinkPeers = nvlinkPeers(dev, byBusID, gpus[i].Index)
|
||||
}
|
||||
return gpus, nil
|
||||
}
|
||||
|
||||
// nvlinkPeers returns the indices of GPUs reachable from dev over NVLink.
|
||||
func nvlinkPeers(dev nvml.Device, byBusID map[string]int, self int) []int {
|
||||
seen := map[int]struct{}{}
|
||||
var out []int
|
||||
|
||||
for link := range nvml.NVLINK_MAX_LINKS {
|
||||
pci, ret := dev.GetNvLinkRemotePciInfo(link)
|
||||
if ret != nvml.SUCCESS {
|
||||
// Unpopulated links return an error; that is the normal case for
|
||||
// most of the 36 possible links, so it is not worth reporting.
|
||||
continue
|
||||
}
|
||||
peer, ok := byBusID[normalizeBusID(busIDString(pci.BusId[:]))]
|
||||
if !ok || peer == self {
|
||||
// The remote end can be an NVSwitch rather than a GPU, in which
|
||||
// case it is not in the map.
|
||||
continue
|
||||
}
|
||||
if _, dup := seen[peer]; dup {
|
||||
continue
|
||||
}
|
||||
seen[peer] = struct{}{}
|
||||
out = append(out, peer)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// busIDString turns NVML's fixed-size C char array into a Go string.
|
||||
func busIDString(b []int8) string {
|
||||
buf := make([]byte, 0, len(b))
|
||||
for _, c := range b {
|
||||
if c == 0 {
|
||||
break
|
||||
}
|
||||
buf = append(buf, byte(c))
|
||||
}
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
// normalizeBusID makes bus ids comparable: NVML reports them uppercase from one
|
||||
// call and lowercase from another, and sometimes with a shortened domain.
|
||||
func normalizeBusID(s string) string {
|
||||
s = strings.ToLower(strings.TrimSpace(s))
|
||||
// "0000:1b:00.0" and "00000000:1B:00.0" are the same device.
|
||||
if i := strings.Index(s, ":"); i > 0 {
|
||||
s = strings.TrimLeft(s[:i], "0") + s[i:]
|
||||
}
|
||||
return strings.TrimPrefix(s, ":")
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// smiQueryFields are the columns requested from nvidia-smi, in order.
|
||||
var smiQueryFields = []string{
|
||||
"index",
|
||||
"uuid",
|
||||
"name",
|
||||
"memory.total",
|
||||
"memory.used",
|
||||
"compute_cap",
|
||||
"pci.bus_id",
|
||||
"driver_version",
|
||||
"temperature.gpu",
|
||||
"ecc.errors.uncorrected.volatile.total",
|
||||
// Last, and optional. An nvidia-smi too old to know this field returns a
|
||||
// shorter row, and every column before it still lines up — which is why it
|
||||
// goes on the end rather than anywhere else, and why the parser requires
|
||||
// only the mandatory prefix.
|
||||
"utilization.gpu",
|
||||
}
|
||||
|
||||
// mandatorySMIFields is how many columns a row must have to be usable.
|
||||
//
|
||||
// Everything before utilization.gpu. A driver that does not know the last field
|
||||
// gives a short row, and rejecting those would make an llmgw upgrade silently
|
||||
// stop discovering GPUs on an older node — a far worse outcome than not knowing
|
||||
// how busy a card is.
|
||||
const mandatorySMIFields = 10
|
||||
|
||||
// DiscoverGPUsSMI enumerates GPUs by shelling out to nvidia-smi.
|
||||
//
|
||||
// This is the cgo-free path (plan §3.2): the NVML bindings need cgo, which
|
||||
// would force the whole binary — server included — to be built with a C
|
||||
// toolchain. The query gives everything except topology, which Discover
|
||||
// reads from `nvidia-smi topo -m` afterwards.
|
||||
func DiscoverGPUsSMI(ctx context.Context, binary string) ([]GPU, error) {
|
||||
if binary == "" {
|
||||
// Not configured and not installed is a node without an NVIDIA card,
|
||||
// not a failed discovery. Reported as an error it also cut discovery
|
||||
// short before InfiniBand was looked at, on every CPU node that had
|
||||
// not been told by hand to skip GPUs.
|
||||
if _, err := exec.LookPath("nvidia-smi"); err != nil {
|
||||
return nil, nil
|
||||
}
|
||||
binary = "nvidia-smi"
|
||||
}
|
||||
out, err := exec.CommandContext(ctx, binary,
|
||||
"--query-gpu="+strings.Join(smiQueryFields, ","),
|
||||
"--format=csv,noheader,nounits",
|
||||
).Output()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("running %s: %w", binary, err)
|
||||
}
|
||||
return parseSMI(string(out))
|
||||
}
|
||||
|
||||
func parseSMI(out string) ([]GPU, error) {
|
||||
r := csv.NewReader(strings.NewReader(out))
|
||||
r.TrimLeadingSpace = true
|
||||
r.FieldsPerRecord = -1 // tolerate driver versions with extra columns
|
||||
|
||||
rows, err := r.ReadAll()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parsing nvidia-smi output: %w", err)
|
||||
}
|
||||
|
||||
gpus := make([]GPU, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
if len(row) < mandatorySMIFields {
|
||||
continue
|
||||
}
|
||||
idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
g := GPU{
|
||||
Index: idx,
|
||||
UUID: strings.TrimSpace(row[1]),
|
||||
Name: strings.TrimSpace(row[2]),
|
||||
ComputeCapability: strings.TrimSpace(row[5]),
|
||||
PCIBusID: strings.TrimSpace(row[6]),
|
||||
DriverVersion: strings.TrimSpace(row[7]),
|
||||
NUMANode: -1,
|
||||
}
|
||||
g.MemoryTotalMB, _ = parseCSVInt(row[3])
|
||||
g.MemoryUsedMB, _ = parseCSVInt(row[4])
|
||||
g.Health.TemperatureC, _ = parseCSVInt(row[8])
|
||||
g.Health.ECCErrors, _ = parseCSVInt(row[9])
|
||||
if len(row) > 10 {
|
||||
g.UtilizationPercent, _ = parseCSVInt(row[10])
|
||||
}
|
||||
|
||||
assessHealth(&g)
|
||||
gpus = append(gpus, g)
|
||||
}
|
||||
return gpus, nil
|
||||
}
|
||||
|
||||
// residentMemoryToleranceMB is how much allocated memory is written off as
|
||||
// driver overhead rather than a foreign process.
|
||||
const residentMemoryToleranceMB = 512
|
||||
|
||||
// assessHealth flags conditions that should keep a GPU out of placement.
|
||||
func assessHealth(g *GPU) {
|
||||
if g.Health.ECCErrors > 0 {
|
||||
g.Health.Problems = append(g.Health.Problems,
|
||||
fmt.Sprintf("%d uncorrected ECC errors", g.Health.ECCErrors))
|
||||
}
|
||||
// Memory already in use means something else is on the card. Starting an
|
||||
// engine there would either fail to allocate or quietly evict the other
|
||||
// process's throughput.
|
||||
if g.MemoryUsedMB > residentMemoryToleranceMB {
|
||||
g.Health.Problems = append(g.Health.Problems,
|
||||
fmt.Sprintf("%d MiB already allocated by another process", g.MemoryUsedMB))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const smiSample = `0, GPU-11111111-2222-3333-4444-555555555555, NVIDIA H100 80GB HBM3, 81559, 4, 9.0, 00000000:1B:00.0, 550.54.15, 34, 0
|
||||
1, GPU-66666666-7777-8888-9999-000000000000, NVIDIA H100 80GB HBM3, 81559, 0, 9.0, 00000000:1C:00.0, 550.54.15, 36, 0
|
||||
`
|
||||
|
||||
func TestParseSMI(t *testing.T) {
|
||||
gpus, err := parseSMI(smiSample)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 2 {
|
||||
t.Fatalf("expected 2 GPUs, got %d", len(gpus))
|
||||
}
|
||||
|
||||
g := gpus[0]
|
||||
if g.Index != 0 || g.UUID != "GPU-11111111-2222-3333-4444-555555555555" {
|
||||
t.Errorf("identity wrong: %+v", g)
|
||||
}
|
||||
if g.MemoryTotalMB != 81559 {
|
||||
t.Errorf("MemoryTotalMB=%d", g.MemoryTotalMB)
|
||||
}
|
||||
if g.ComputeCapability != "9.0" {
|
||||
t.Errorf("ComputeCapability=%q; FP8 placement depends on it", g.ComputeCapability)
|
||||
}
|
||||
if g.DriverVersion != "550.54.15" {
|
||||
t.Errorf("DriverVersion=%q", g.DriverVersion)
|
||||
}
|
||||
if !g.Health.OK() {
|
||||
t.Errorf("a healthy GPU was flagged: %v", g.Health.Problems)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSMIUnsupportedFields(t *testing.T) {
|
||||
// Consumer cards report [N/A] for ECC; that must not fail the whole parse.
|
||||
const row = `0, GPU-abc, NVIDIA GeForce RTX 4090, 24564, 0, 8.9, 00000000:01:00.0, 550.54.15, 45, [N/A]
|
||||
`
|
||||
gpus, err := parseSMI(row)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("expected 1 GPU, got %d", len(gpus))
|
||||
}
|
||||
if gpus[0].Health.ECCErrors != 0 || !gpus[0].Health.OK() {
|
||||
t.Errorf("[N/A] ECC was misread: %+v", gpus[0].Health)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthFlagsForeignMemory(t *testing.T) {
|
||||
// Memory in use on an otherwise idle card means another process is there.
|
||||
const row = `0, GPU-abc, NVIDIA H100 80GB HBM3, 81559, 40000, 9.0, 00000000:1B:00.0, 550.54.15, 34, 0
|
||||
`
|
||||
gpus, _ := parseSMI(row)
|
||||
if gpus[0].Health.OK() {
|
||||
t.Error("40 GiB allocated by another process was not flagged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthFlagsECC(t *testing.T) {
|
||||
const row = `0, GPU-abc, NVIDIA H100 80GB HBM3, 81559, 0, 9.0, 00000000:1B:00.0, 550.54.15, 34, 3
|
||||
`
|
||||
gpus, _ := parseSMI(row)
|
||||
if gpus[0].Health.OK() {
|
||||
t.Error("uncorrected ECC errors were not flagged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGPUClass(t *testing.T) {
|
||||
// The same card must derive the same class however the driver names it,
|
||||
// because a tensor-parallel group may not mix classes.
|
||||
cases := []struct {
|
||||
name string
|
||||
memMB int
|
||||
expect string
|
||||
}{
|
||||
{"NVIDIA H100 80GB HBM3", 81559, "h100-80g"},
|
||||
{"NVIDIA H100-SXM5-80GB", 81559, "h100-80g"},
|
||||
{"Tesla A100-SXM4-40GB", 40960, "a100-40g"},
|
||||
{"NVIDIA A100 80GB PCIe", 81920, "a100-80g"},
|
||||
// Sizes come from what the driver reports, in GiB. That does not always
|
||||
// match the marketing name: an L40S is sold as 48GB but reports 44.99
|
||||
// GiB, because the marketing figure is decimal GB. Being consistent
|
||||
// matters more than matching the box, since the number is what decides
|
||||
// whether a model fits.
|
||||
{"NVIDIA L40S", 46068, "l40s-45g"},
|
||||
// The model token is what distinguishes a 4090 from a 4080, so "rtx"
|
||||
// is not enough on its own.
|
||||
{"NVIDIA GeForce RTX 4090", 24564, "4090-24g"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := GPU{Name: c.name, MemoryTotalMB: c.memMB}.Class()
|
||||
if got != c.expect {
|
||||
t.Errorf("%q (%d MiB) -> %q, want %q", c.name, c.memMB, got, c.expect)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestInventoryClasses(t *testing.T) {
|
||||
inv := Inventory{GPUs: []GPU{
|
||||
{Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559},
|
||||
{Name: "NVIDIA H100-SXM5-80GB", MemoryTotalMB: 81559},
|
||||
{Name: "Tesla A100-SXM4-40GB", MemoryTotalMB: 40960},
|
||||
}}
|
||||
got := inv.Classes()
|
||||
if len(got) != 2 || got[0] != "a100-40g" || got[1] != "h100-80g" {
|
||||
t.Errorf("classes=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverGPUsSMI runs the real code path against a stub binary, so the
|
||||
// command construction is covered too, not just the parser.
|
||||
func TestDiscoverGPUsSMI(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
stub := filepath.Join(dir, "nvidia-smi")
|
||||
script := "#!/bin/sh\ncat <<'EOF'\n" + smiSample + "EOF\n"
|
||||
if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
gpus, err := DiscoverGPUsSMI(t.Context(), stub)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 2 {
|
||||
t.Fatalf("expected 2 GPUs, got %d", len(gpus))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverGPUsSMIMissingBinary(t *testing.T) {
|
||||
// A node with no NVIDIA driver must produce a clear error rather than a
|
||||
// panic or an empty success.
|
||||
_, err := DiscoverGPUsSMI(context.Background(), filepath.Join(t.TempDir(), "absent"))
|
||||
if err == nil {
|
||||
t.Error("expected an error when nvidia-smi is missing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnOlderDriverWithoutUtilisationStillDiscovers(t *testing.T) {
|
||||
// An nvidia-smi too old for utilization.gpu returns a shorter row. Rejecting
|
||||
// those would make an llmgw upgrade silently stop discovering GPUs on an
|
||||
// older node, which is far worse than not knowing how busy a card is — and
|
||||
// is what the first version of this change did.
|
||||
old := "0, GPU-aaa, NVIDIA H100 80GB HBM3, 81559, 1024, 9.0, 0000:01:00.0, 550.54, 41, 0"
|
||||
gpus, err := parseSMI(old)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("a row without utilisation produced %d GPUs", len(gpus))
|
||||
}
|
||||
if gpus[0].MemoryTotalMB != 81559 || gpus[0].Health.TemperatureC != 41 {
|
||||
t.Errorf("columns shifted: %+v", gpus[0])
|
||||
}
|
||||
if gpus[0].UtilizationPercent != 0 {
|
||||
t.Errorf("utilisation = %d, want 0 when unreported", gpus[0].UtilizationPercent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUtilisationIsReadWhenPresent(t *testing.T) {
|
||||
row := "0, GPU-aaa, NVIDIA H100 80GB HBM3, 81559, 1024, 9.0, 0000:01:00.0, 550.54, 41, 0, 87"
|
||||
gpus, err := parseSMI(row)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("got %d GPUs", len(gpus))
|
||||
}
|
||||
if gpus[0].UtilizationPercent != 87 {
|
||||
t.Errorf("utilisation = %d, want 87", gpus[0].UtilizationPercent)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"os"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Host is what a node is, as far as installing software on it is concerned.
|
||||
type Host struct {
|
||||
// OSID, OSVersion and OSCodename come from /etc/os-release: "ubuntu",
|
||||
// "24.04", "noble".
|
||||
OSID string `json:"os_id,omitempty"`
|
||||
OSVersion string `json:"os_version,omitempty"`
|
||||
OSCodename string `json:"os_codename,omitempty"`
|
||||
// OSLike is ID_LIKE: the families this system says it follows, "debian"
|
||||
// for Ubuntu, "rhel centos fedora" for Rocky. It is what decides the
|
||||
// package manager, which the name alone does not for every derivative.
|
||||
OSLike []string `json:"os_like,omitempty"`
|
||||
// Arch is Debian's spelling, "amd64" or "arm64", because that is what
|
||||
// package names and release archives are keyed by.
|
||||
Arch string `json:"arch,omitempty"`
|
||||
// CPUFlags are the vector extensions inference cares about, of those this
|
||||
// processor has.
|
||||
CPUFlags []string `json:"cpu_flags,omitempty"`
|
||||
// NVIDIADriver is the loaded driver's version, empty when there is none.
|
||||
NVIDIADriver string `json:"nvidia_driver,omitempty"`
|
||||
// MinComputeCapability is the oldest GPU's, 0 when there are none or it
|
||||
// is not known. Not read here but from the GPU inventory, by the agent.
|
||||
MinComputeCapability float64 `json:"min_compute_capability,omitempty"`
|
||||
}
|
||||
|
||||
// Family is the packaging family: "deb", "rpm", or "" when neither is known.
|
||||
func (h Host) Family() string {
|
||||
for _, id := range append([]string{h.OSID}, h.OSLike...) {
|
||||
switch id {
|
||||
case "debian", "ubuntu":
|
||||
return "deb"
|
||||
case "rhel", "fedora", "centos", "rocky", "almalinux", "suse", "opensuse":
|
||||
return "rpm"
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// HasCPUFlag reports whether the processor has an extension.
|
||||
func (h Host) HasCPUFlag(f string) bool {
|
||||
for _, have := range h.CPUFlags {
|
||||
if have == f {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// interestingFlags are the extensions that decide which build of an engine
|
||||
// runs well. A VM on QEMU's default CPU model has none of them — and still
|
||||
// runs, five times slower, with nothing saying why.
|
||||
var interestingFlags = []string{"avx", "avx2", "fma", "f16c", "avx512f", "avx512_bf16", "amx_tile", "asimd", "sve"}
|
||||
|
||||
// DiscoverHost reads what the machine is.
|
||||
func DiscoverHost() Host {
|
||||
h := hostFromOSRelease("/etc/os-release")
|
||||
h.Arch = debianArch(runtime.GOARCH)
|
||||
h.CPUFlags = cpuFlagsFrom("/proc/cpuinfo")
|
||||
h.NVIDIADriver = nvidiaDriverFrom("/proc/driver/nvidia/version")
|
||||
return h
|
||||
}
|
||||
|
||||
func hostFromOSRelease(path string) Host {
|
||||
var h Host
|
||||
f, err := os.Open(path) //nolint:gosec // a fixed system path
|
||||
if err != nil {
|
||||
return h
|
||||
}
|
||||
defer f.Close()
|
||||
sc := bufio.NewScanner(f)
|
||||
for sc.Scan() {
|
||||
k, v, ok := strings.Cut(sc.Text(), "=")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
v = strings.Trim(v, `"'`)
|
||||
switch k {
|
||||
case "ID":
|
||||
h.OSID = v
|
||||
case "VERSION_ID":
|
||||
h.OSVersion = v
|
||||
case "VERSION_CODENAME":
|
||||
h.OSCodename = v
|
||||
case "ID_LIKE":
|
||||
h.OSLike = strings.Fields(v)
|
||||
}
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
func debianArch(goarch string) string {
|
||||
switch goarch {
|
||||
case "amd64", "arm64":
|
||||
return goarch
|
||||
case "386":
|
||||
return "i386"
|
||||
case "arm":
|
||||
return "armhf"
|
||||
}
|
||||
return goarch
|
||||
}
|
||||
|
||||
func cpuFlagsFrom(path string) []string {
|
||||
b, err := os.ReadFile(path) //nolint:gosec // a fixed system path
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
have := map[string]bool{}
|
||||
for _, line := range strings.Split(string(b), "\n") {
|
||||
k, v, ok := strings.Cut(line, ":")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
// "flags" on x86, "Features" on ARM. The first processor is enough:
|
||||
// a machine does not mix instruction sets across its cores.
|
||||
if k = strings.TrimSpace(k); k != "flags" && k != "Features" {
|
||||
continue
|
||||
}
|
||||
for _, f := range strings.Fields(v) {
|
||||
have[f] = true
|
||||
}
|
||||
break
|
||||
}
|
||||
var out []string
|
||||
for _, f := range interestingFlags {
|
||||
if have[f] {
|
||||
out = append(out, f)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// nvidiaDriverFrom reads the loaded kernel module's version, which exists
|
||||
// exactly when a driver is loaded — unlike nvidia-smi, which can be installed
|
||||
// on a machine whose module failed to load.
|
||||
func nvidiaDriverFrom(path string) string {
|
||||
b, err := os.ReadFile(path) //nolint:gosec // a fixed system path
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
// "NVRM version: NVIDIA UNIX x86_64 Kernel Module 570.86.15 Thu ..."
|
||||
fields := strings.Fields(string(b))
|
||||
for i, f := range fields {
|
||||
if f == "Module" && i+1 < len(fields) {
|
||||
return fields[i+1]
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// MinComputeCapability is the oldest GPU's compute capability, 0 when none
|
||||
// is known: an engine built for newer GPUs is refused by the oldest.
|
||||
func MinComputeCapability(gpus []GPU) float64 {
|
||||
low := 0.0
|
||||
for _, g := range gpus {
|
||||
cc, err := strconv.ParseFloat(strings.TrimSpace(g.ComputeCapability), 64)
|
||||
if err != nil || cc <= 0 {
|
||||
continue
|
||||
}
|
||||
if low == 0 || cc < low {
|
||||
low = cc
|
||||
}
|
||||
}
|
||||
return low
|
||||
}
|
||||
@@ -0,0 +1,378 @@
|
||||
// Package discovery inspects a node's GPUs, InfiniBand adapters and the
|
||||
// topology between them (plan §5.1 and §5.2).
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// SysfsRoot is the default place InfiniBand devices appear.
|
||||
const SysfsRoot = "/sys/class/infiniband"
|
||||
|
||||
// LinkLayer distinguishes real InfiniBand from RoCE.
|
||||
type LinkLayer string
|
||||
|
||||
const (
|
||||
// LinkInfiniBand is native InfiniBand.
|
||||
LinkInfiniBand LinkLayer = "InfiniBand"
|
||||
// LinkEthernet is RoCE — Ethernet carrying RDMA.
|
||||
LinkEthernet LinkLayer = "Ethernet"
|
||||
)
|
||||
|
||||
// Port is one port of an HCA.
|
||||
type Port struct {
|
||||
Num int `json:"num"`
|
||||
// State is the port's logical state; only ACTIVE ports carry traffic.
|
||||
State string `json:"state"`
|
||||
// PhysState is the physical link state, useful for telling "cable unplugged"
|
||||
// from "no subnet manager".
|
||||
PhysState string `json:"phys_state"`
|
||||
// RateGbps is the negotiated rate.
|
||||
RateGbps int `json:"rate_gbps"`
|
||||
LinkLayer LinkLayer `json:"link_layer"`
|
||||
// GIDPrefix is the subnet prefix from gids/0, normalized to lowercase hex
|
||||
// groups. It is the only automatic fabric hint available, and an unreliable
|
||||
// one: OpenSM's default is fe80:: everywhere (plan §5.2).
|
||||
GIDPrefix string `json:"gid_prefix"`
|
||||
LID int `json:"lid"`
|
||||
SMLID int `json:"sm_lid"`
|
||||
// Counters are the port's error and traffic counters, by their sysfs
|
||||
// names, as far as the driver has them. The data counters are converted
|
||||
// to bytes; sysfs counts them in four-byte words.
|
||||
Counters map[string]uint64 `json:"counters,omitempty"`
|
||||
// RoCEv2GID is the GID table index NCCL should use on a RoCE port: the
|
||||
// first RoCE v2 entry holding an IPv4 address, else the first RoCE v2
|
||||
// entry. -1 on InfiniBand, and on a RoCE port with no v2 entry. NCCL's
|
||||
// default of 0 is a RoCE v1 link-local entry on most setups, which does
|
||||
// not route, and a job on it hangs.
|
||||
RoCEv2GID int `json:"rocev2_gid"`
|
||||
// PortGUID is the port's own identifier, the interface half of its GID.
|
||||
PortGUID string `json:"port_guid,omitempty"`
|
||||
// SM is the subnet manager this port sees as master, from sminfo. Nil on
|
||||
// RoCE, which has none, and when infiniband-diags is not installed.
|
||||
SM *SMInfo `json:"sm,omitempty"`
|
||||
// Network is, on RoCE, the IP network of the RoCE v2 address: the only
|
||||
// thing that says two RoCE ports are on the same fabric.
|
||||
Network string `json:"network,omitempty"`
|
||||
}
|
||||
|
||||
// SMInfo is a subnet manager as a port sees it.
|
||||
type SMInfo struct {
|
||||
GUID string `json:"guid"`
|
||||
LID int `json:"lid"`
|
||||
Priority int `json:"priority"`
|
||||
// State is SMINFO_MASTER, SMINFO_STANDBY, ...
|
||||
State string `json:"state"`
|
||||
}
|
||||
|
||||
// FabricKey is what identifies the physical fabric this port is on: the
|
||||
// master subnet manager on InfiniBand, since every port on one fabric sees
|
||||
// the same one and ports on separate fabrics cannot; the IP network on RoCE.
|
||||
// Empty when neither is known, and for a port that is not up.
|
||||
func (p Port) FabricKey() string {
|
||||
if !p.Active() {
|
||||
return ""
|
||||
}
|
||||
switch {
|
||||
case p.LinkLayer == LinkEthernet && p.Network != "":
|
||||
return "roce:" + p.Network
|
||||
case p.SM != nil && p.SM.GUID != "":
|
||||
return "sm:" + p.SM.GUID
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// PortCounters are the counters worth watching, from each port's counters/
|
||||
// directory. The errors are the ones that say a cable, a transceiver or a
|
||||
// switch port is going bad before a job fails on it; the data counters say
|
||||
// whether the fabric is carrying anything at all.
|
||||
var PortCounters = []string{
|
||||
"symbol_error", "link_error_recovery", "link_downed", "port_rcv_errors",
|
||||
"port_rcv_remote_physical_errors", "port_xmit_discards",
|
||||
"local_link_integrity_errors", "excessive_buffer_overrun_errors",
|
||||
"port_xmit_data", "port_rcv_data",
|
||||
}
|
||||
|
||||
// dataCounter reports whether a counter counts data in four-byte words.
|
||||
func dataCounter(name string) bool { return name == "port_xmit_data" || name == "port_rcv_data" }
|
||||
|
||||
// Active reports whether the port can carry traffic.
|
||||
func (p Port) Active() bool { return strings.EqualFold(p.State, "ACTIVE") }
|
||||
|
||||
// HCA is one host channel adapter.
|
||||
type HCA struct {
|
||||
Name string `json:"name"`
|
||||
Ports []Port `json:"ports"`
|
||||
PCIAddr string `json:"pci_addr,omitempty"`
|
||||
NUMANode int `json:"numa_node"`
|
||||
// FirmwareVersion is reported by the driver; worth capturing because
|
||||
// NCCL behaviour differs between firmware revisions.
|
||||
FirmwareVersion string `json:"firmware_version,omitempty"`
|
||||
// NodeGUID identifies the adapter on its fabric: it is the name a
|
||||
// switch's view of the fabric, and every other tool, knows it by.
|
||||
NodeGUID string `json:"node_guid,omitempty"`
|
||||
}
|
||||
|
||||
// ActivePorts returns the ports that can carry traffic.
|
||||
func (h HCA) ActivePorts() []Port {
|
||||
var out []Port
|
||||
for _, p := range h.Ports {
|
||||
if p.Active() {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// DiscoverIB enumerates the InfiniBand adapters under root.
|
||||
//
|
||||
// A missing root is not an error: plenty of nodes have no IB at all, and the
|
||||
// scheduler treats them as single-node islands (plan §2).
|
||||
func DiscoverIB(root string) ([]HCA, error) {
|
||||
if root == "" {
|
||||
root = SysfsRoot
|
||||
}
|
||||
entries, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("reading %s: %w", root, err)
|
||||
}
|
||||
|
||||
var out []HCA
|
||||
for _, e := range entries {
|
||||
hca, err := readHCA(filepath.Join(root, e.Name()), e.Name())
|
||||
if err != nil {
|
||||
// One unreadable adapter must not hide the rest; the node may still
|
||||
// be usable through its other HCAs.
|
||||
continue
|
||||
}
|
||||
out = append(out, hca)
|
||||
}
|
||||
slices.SortFunc(out, func(a, b HCA) int { return strings.Compare(a.Name, b.Name) })
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func readHCA(dir, name string) (HCA, error) {
|
||||
h := HCA{Name: name, NUMANode: -1}
|
||||
h.FirmwareVersion = readTrimmed(filepath.Join(dir, "fw_ver"))
|
||||
h.NodeGUID = normalizeGUID(readTrimmed(filepath.Join(dir, "node_guid")))
|
||||
|
||||
// device is a symlink into the PCI tree; its target's base name is the
|
||||
// bus address, and numa_node lives beside it.
|
||||
if target, err := filepath.EvalSymlinks(filepath.Join(dir, "device")); err == nil {
|
||||
h.PCIAddr = filepath.Base(target)
|
||||
if n, err := strconv.Atoi(readTrimmed(filepath.Join(target, "numa_node"))); err == nil {
|
||||
h.NUMANode = n
|
||||
}
|
||||
}
|
||||
|
||||
portDirs, err := os.ReadDir(filepath.Join(dir, "ports"))
|
||||
if err != nil {
|
||||
return h, fmt.Errorf("%s: no ports: %w", name, err)
|
||||
}
|
||||
for _, pd := range portDirs {
|
||||
num, err := strconv.Atoi(pd.Name())
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
h.Ports = append(h.Ports, readPort(filepath.Join(dir, "ports", pd.Name()), num))
|
||||
}
|
||||
slices.SortFunc(h.Ports, func(a, b Port) int { return a.Num - b.Num })
|
||||
return h, nil
|
||||
}
|
||||
|
||||
func readPort(dir string, num int) Port {
|
||||
p := Port{
|
||||
Num: num,
|
||||
State: parseEnumeratedState(readTrimmed(filepath.Join(dir, "state"))),
|
||||
PhysState: parseEnumeratedState(readTrimmed(filepath.Join(dir, "phys_state"))),
|
||||
RateGbps: parseRate(readTrimmed(filepath.Join(dir, "rate"))),
|
||||
LinkLayer: LinkLayer(readTrimmed(filepath.Join(dir, "link_layer"))),
|
||||
GIDPrefix: parseGIDPrefix(readTrimmed(filepath.Join(dir, "gids", "0"))),
|
||||
LID: parseMaybeHex(readTrimmed(filepath.Join(dir, "lid"))),
|
||||
SMLID: parseMaybeHex(readTrimmed(filepath.Join(dir, "sm_lid"))),
|
||||
}
|
||||
p.RoCEv2GID = -1
|
||||
p.PortGUID = portGUID(readTrimmed(filepath.Join(dir, "gids", "0")))
|
||||
if p.LinkLayer == LinkEthernet {
|
||||
p.RoCEv2GID = rocev2GID(dir)
|
||||
if p.RoCEv2GID >= 0 {
|
||||
idx := strconv.Itoa(p.RoCEv2GID)
|
||||
p.Network = RoCENetwork(readTrimmed(filepath.Join(dir, "gid_attrs", "ndevs", idx)),
|
||||
readTrimmed(filepath.Join(dir, "gids", idx)))
|
||||
}
|
||||
}
|
||||
for _, name := range PortCounters {
|
||||
raw := readTrimmed(filepath.Join(dir, "counters", name))
|
||||
v, err := strconv.ParseUint(raw, 10, 64)
|
||||
if err != nil {
|
||||
continue // not every driver has every counter
|
||||
}
|
||||
if dataCounter(name) {
|
||||
v *= 4
|
||||
}
|
||||
if p.Counters == nil {
|
||||
p.Counters = map[string]uint64{}
|
||||
}
|
||||
p.Counters[name] = v
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// parseEnumeratedState turns sysfs's "4: ACTIVE" into "ACTIVE".
|
||||
func parseEnumeratedState(s string) string {
|
||||
if _, rest, found := strings.Cut(s, ":"); found {
|
||||
return strings.TrimSpace(rest)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// parseRate turns "100 Gb/sec (4X EDR)" into 100.
|
||||
func parseRate(s string) int {
|
||||
f := strings.Fields(s)
|
||||
if len(f) == 0 {
|
||||
return 0
|
||||
}
|
||||
// The rate can be fractional ("2.5 Gb/sec") on old hardware; truncate.
|
||||
v, err := strconv.ParseFloat(f[0], 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return int(v)
|
||||
}
|
||||
|
||||
// parseGIDPrefix keeps the first four groups of a GID, which is the subnet
|
||||
// prefix; the remaining four are the interface id.
|
||||
func parseGIDPrefix(gid string) string {
|
||||
parts := strings.Split(strings.ToLower(gid), ":")
|
||||
if len(parts) < 4 {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(parts[:4], ":")
|
||||
}
|
||||
|
||||
// parseMaybeHex accepts both "0x1a" and "26"; sysfs is inconsistent between
|
||||
// drivers and kernel versions.
|
||||
func parseMaybeHex(s string) int {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return 0
|
||||
}
|
||||
if after, ok := strings.CutPrefix(strings.ToLower(s), "0x"); ok {
|
||||
v, err := strconv.ParseInt(after, 16, 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return int(v)
|
||||
}
|
||||
v, err := strconv.Atoi(s)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func readTrimmed(path string) string {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
|
||||
// rocev2GID finds the GID index a RoCE port's traffic should use.
|
||||
func rocev2GID(dir string) int {
|
||||
first := -1
|
||||
for i := 0; i < 64; i++ {
|
||||
typ := readTrimmed(filepath.Join(dir, "gid_attrs", "types", strconv.Itoa(i)))
|
||||
if !strings.EqualFold(typ, "RoCE v2") {
|
||||
continue
|
||||
}
|
||||
gid := strings.ToLower(readTrimmed(filepath.Join(dir, "gids", strconv.Itoa(i))))
|
||||
if strings.HasPrefix(gid, "0000:0000:0000:0000:0000:ffff:") {
|
||||
return i
|
||||
}
|
||||
if first < 0 {
|
||||
first = i
|
||||
}
|
||||
}
|
||||
return first
|
||||
}
|
||||
|
||||
// normalizeGUID writes a GUID one way: sysfs says "0c42:a103:0045:bd6c",
|
||||
// sminfo and ibnetdiscover say "0x0c42a1030045bd6c".
|
||||
func normalizeGUID(g string) string {
|
||||
g = strings.ToLower(strings.TrimSpace(g))
|
||||
g = strings.TrimPrefix(g, "0x")
|
||||
g = strings.ReplaceAll(g, ":", "")
|
||||
if g == "" || strings.Trim(g, "0") == "" {
|
||||
return ""
|
||||
}
|
||||
for len(g) < 16 {
|
||||
g = "0" + g
|
||||
}
|
||||
return "0x" + g
|
||||
}
|
||||
|
||||
// NormalizeGUID is normalizeGUID for other packages.
|
||||
func NormalizeGUID(g string) string { return normalizeGUID(g) }
|
||||
|
||||
// portGUID is the interface identifier of a GID: its last four groups.
|
||||
func portGUID(gid string) string {
|
||||
parts := strings.Split(gid, ":")
|
||||
if len(parts) != 8 {
|
||||
return ""
|
||||
}
|
||||
return normalizeGUID(strings.Join(parts[4:], ""))
|
||||
}
|
||||
|
||||
// RoCENetwork is the IP network an IPv4-mapped RoCE GID's address sits in,
|
||||
// read from the interface that carries it. Replaced in tests.
|
||||
var RoCENetwork = func(ndev, gid string) string {
|
||||
ip := gidIPv4(gid)
|
||||
if ip == nil || ndev == "" {
|
||||
return ""
|
||||
}
|
||||
ifc, err := net.InterfaceByName(ndev)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
addrs, _ := ifc.Addrs()
|
||||
for _, a := range addrs {
|
||||
if n, ok := a.(*net.IPNet); ok && n.IP.Equal(ip) {
|
||||
return (&net.IPNet{IP: ip.Mask(n.Mask), Mask: n.Mask}).String()
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// gidIPv4 is the IPv4 address in an IPv4-mapped GID, nil otherwise.
|
||||
func gidIPv4(gid string) net.IP {
|
||||
g := strings.ToLower(gid)
|
||||
const mapped = "0000:0000:0000:0000:0000:ffff:"
|
||||
if !strings.HasPrefix(g, mapped) {
|
||||
return nil
|
||||
}
|
||||
hex := strings.ReplaceAll(strings.TrimPrefix(g, mapped), ":", "")
|
||||
if len(hex) != 8 {
|
||||
return nil
|
||||
}
|
||||
b := make(net.IP, 4)
|
||||
for i := range 4 {
|
||||
v, err := strconv.ParseUint(hex[i*2:i*2+2], 16, 8)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
b[i] = byte(v)
|
||||
}
|
||||
return b
|
||||
}
|
||||
@@ -0,0 +1,353 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// fakeSysfs builds an InfiniBand tree shaped like the real one.
|
||||
func fakeSysfs(t *testing.T) string {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
|
||||
write := func(path, content string) {
|
||||
t.Helper()
|
||||
full := filepath.Join(root, path)
|
||||
if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(full, []byte(content+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// mlx5_0: one active IB port.
|
||||
write("ib/mlx5_0/fw_ver", "20.35.1012")
|
||||
write("ib/mlx5_0/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_0/ports/1/phys_state", "5: LinkUp")
|
||||
write("ib/mlx5_0/ports/1/rate", "200 Gb/sec (4X HDR)")
|
||||
write("ib/mlx5_0/ports/1/link_layer", "InfiniBand")
|
||||
write("ib/mlx5_0/ports/1/gids/0", "FEC0:0000:0000:000A:0002:C903:00A0:1234")
|
||||
write("ib/mlx5_0/ports/1/lid", "0x1a")
|
||||
write("ib/mlx5_0/ports/1/sm_lid", "1")
|
||||
write("ib/mlx5_0/ports/1/counters/symbol_error", "3")
|
||||
write("ib/mlx5_0/ports/1/counters/link_downed", "1")
|
||||
write("ib/mlx5_0/ports/1/counters/port_xmit_data", "1000")
|
||||
|
||||
// mlx5_1: two ports, the second one down.
|
||||
write("ib/mlx5_1/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_1/ports/1/rate", "200 Gb/sec (4X HDR)")
|
||||
write("ib/mlx5_1/ports/1/link_layer", "InfiniBand")
|
||||
write("ib/mlx5_1/ports/1/gids/0", "fec0:0000:0000:000a:0002:c903:00a0:5678")
|
||||
write("ib/mlx5_1/ports/2/state", "1: DOWN")
|
||||
write("ib/mlx5_1/ports/2/phys_state", "3: Disabled")
|
||||
write("ib/mlx5_1/ports/2/rate", "0")
|
||||
write("ib/mlx5_1/ports/2/link_layer", "InfiniBand")
|
||||
|
||||
// mlx5_2: RoCE rather than native IB.
|
||||
write("ib/mlx5_2/ports/1/state", "4: ACTIVE")
|
||||
write("ib/mlx5_2/ports/1/rate", "100 Gb/sec (4X EDR)")
|
||||
write("ib/mlx5_2/ports/1/link_layer", "Ethernet")
|
||||
write("ib/mlx5_2/ports/1/gids/0", "fe80:0000:0000:0000:0000:0000:0000:0001")
|
||||
write("ib/mlx5_2/ports/1/gid_attrs/types/0", "IB/RoCE v1")
|
||||
write("ib/mlx5_2/ports/1/gids/1", "fe80:0000:0000:0000:0000:0000:0000:0001")
|
||||
write("ib/mlx5_2/ports/1/gid_attrs/types/1", "RoCE v2")
|
||||
write("ib/mlx5_2/ports/1/gids/3", "0000:0000:0000:0000:0000:ffff:0a00:0001")
|
||||
write("ib/mlx5_2/ports/1/gid_attrs/types/3", "RoCE v2")
|
||||
write("ib/mlx5_2/ports/1/gid_attrs/ndevs/3", "ens1f0")
|
||||
|
||||
// A PCI device with a NUMA node, linked the way the driver does it.
|
||||
pci := filepath.Join(root, "pci", "0000:0b:00.0")
|
||||
if err := os.MkdirAll(pci, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(pci, "numa_node"), []byte("1\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.Symlink(pci, filepath.Join(root, "ib", "mlx5_0", "device")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
return filepath.Join(root, "ib")
|
||||
}
|
||||
|
||||
func TestDiscoverIB(t *testing.T) {
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hcas) != 3 {
|
||||
t.Fatalf("expected 3 HCAs, got %d", len(hcas))
|
||||
}
|
||||
// Sorted by name, so the order is predictable for the rest of the test.
|
||||
if hcas[0].Name != "mlx5_0" || hcas[2].Name != "mlx5_2" {
|
||||
t.Fatalf("HCAs are not sorted by name: %v %v", hcas[0].Name, hcas[2].Name)
|
||||
}
|
||||
|
||||
h := hcas[0]
|
||||
if h.FirmwareVersion != "20.35.1012" {
|
||||
t.Errorf("fw_ver=%q", h.FirmwareVersion)
|
||||
}
|
||||
if h.PCIAddr != "0000:0b:00.0" {
|
||||
t.Errorf("PCIAddr=%q", h.PCIAddr)
|
||||
}
|
||||
if h.NUMANode != 1 {
|
||||
t.Errorf("NUMANode=%d; GPU-to-HCA affinity depends on this", h.NUMANode)
|
||||
}
|
||||
|
||||
p := h.Ports[0]
|
||||
if p.State != "ACTIVE" || !p.Active() {
|
||||
t.Errorf("state=%q", p.State)
|
||||
}
|
||||
if p.PhysState != "LinkUp" {
|
||||
t.Errorf("phys_state=%q", p.PhysState)
|
||||
}
|
||||
if p.RateGbps != 200 {
|
||||
t.Errorf("rate=%d", p.RateGbps)
|
||||
}
|
||||
if p.LinkLayer != LinkInfiniBand {
|
||||
t.Errorf("link_layer=%q", p.LinkLayer)
|
||||
}
|
||||
if p.GIDPrefix != "fec0:0000:0000:000a" {
|
||||
t.Errorf("GIDPrefix=%q; it must be lowercased and cut to four groups", p.GIDPrefix)
|
||||
}
|
||||
if p.LID != 26 {
|
||||
t.Errorf("LID=%d, expected 0x1a to parse as 26", p.LID)
|
||||
}
|
||||
if p.SMLID != 1 {
|
||||
t.Errorf("SMLID=%d, expected a decimal value to parse too", p.SMLID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestActivePortsExcludesDown(t *testing.T) {
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := hcas[1] // mlx5_1
|
||||
if len(h.Ports) != 2 {
|
||||
t.Fatalf("expected 2 ports, got %d", len(h.Ports))
|
||||
}
|
||||
active := h.ActivePorts()
|
||||
if len(active) != 1 || active[0].Num != 1 {
|
||||
t.Errorf("a DOWN port was reported as active: %+v", active)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoCEIsDistinguished(t *testing.T) {
|
||||
// RoCE is usable but behaves differently from native IB; the placement
|
||||
// logic needs to be able to tell them apart.
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := hcas[2].Ports[0].LinkLayer; got != LinkEthernet {
|
||||
t.Errorf("mlx5_2 link_layer=%q, expected Ethernet (RoCE)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMissingRootIsNotAnError(t *testing.T) {
|
||||
// A node without InfiniBand is perfectly normal; it just forms its own
|
||||
// single-node island.
|
||||
hcas, err := DiscoverIB(filepath.Join(t.TempDir(), "nonexistent"))
|
||||
if err != nil {
|
||||
t.Fatalf("a node without IB should not be an error: %v", err)
|
||||
}
|
||||
if len(hcas) != 0 {
|
||||
t.Errorf("expected no HCAs, got %d", len(hcas))
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreadableAdapterDoesNotHideOthers(t *testing.T) {
|
||||
root := fakeSysfs(t)
|
||||
// An adapter with no ports directory at all.
|
||||
if err := os.MkdirAll(filepath.Join(root, "mlx5_broken"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
hcas, err := DiscoverIB(root)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hcas) != 3 {
|
||||
t.Errorf("a broken adapter changed the result: got %d HCAs", len(hcas))
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseHelpers(t *testing.T) {
|
||||
t.Run("rate", func(t *testing.T) {
|
||||
for in, want := range map[string]int{
|
||||
"200 Gb/sec (4X HDR)": 200,
|
||||
"100 Gb/sec (4X EDR)": 100,
|
||||
"2.5 Gb/sec (1X SDR)": 2,
|
||||
"": 0,
|
||||
"garbage": 0,
|
||||
} {
|
||||
if got := parseRate(in); got != want {
|
||||
t.Errorf("parseRate(%q)=%d, want %d", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("state", func(t *testing.T) {
|
||||
for in, want := range map[string]string{
|
||||
"4: ACTIVE": "ACTIVE",
|
||||
"1: DOWN": "DOWN",
|
||||
"5: LinkUp": "LinkUp",
|
||||
"ACTIVE": "ACTIVE",
|
||||
"": "",
|
||||
} {
|
||||
if got := parseEnumeratedState(in); got != want {
|
||||
t.Errorf("parseEnumeratedState(%q)=%q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("hex", func(t *testing.T) {
|
||||
for in, want := range map[string]int{
|
||||
"0x1a": 26, "0X1A": 26, "26": 26, "": 0, "zzz": 0,
|
||||
} {
|
||||
if got := parseMaybeHex(in); got != want {
|
||||
t.Errorf("parseMaybeHex(%q)=%d, want %d", in, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("gid", func(t *testing.T) {
|
||||
if got := parseGIDPrefix("FE80:0000:0000:0000:0002:C903:00A0:1234"); got != "fe80:0000:0000:0000" {
|
||||
t.Errorf("parseGIDPrefix=%q", got)
|
||||
}
|
||||
if got := parseGIDPrefix("short:gid"); got != "" {
|
||||
t.Errorf("a malformed GID should yield an empty prefix, got %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestAPortsCountersAreRead(t *testing.T) {
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c := hcas[0].Ports[0].Counters
|
||||
// Data in bytes: sysfs counts four-byte words.
|
||||
if c["symbol_error"] != 3 || c["link_downed"] != 1 || c["port_xmit_data"] != 4000 {
|
||||
t.Errorf("counters = %v", c)
|
||||
}
|
||||
if _, ok := c["port_rcv_errors"]; ok {
|
||||
t.Error("a counter the driver does not have should be absent, not zero")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoCEFindsItsRoutableGID(t *testing.T) {
|
||||
// NCCL's default GID index is a v1 link-local entry on most hosts, which
|
||||
// does not route: the job hangs. The v2 entry holding an IPv4 address is
|
||||
// the one to use, over an earlier v2 link-local one.
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, h := range hcas {
|
||||
p := h.Ports[0]
|
||||
switch h.Name {
|
||||
case "mlx5_2":
|
||||
if p.RoCEv2GID != 3 {
|
||||
t.Errorf("RoCE GID index = %d, want 3", p.RoCEv2GID)
|
||||
}
|
||||
default:
|
||||
if p.RoCEv2GID != -1 {
|
||||
t.Errorf("%s is InfiniBand and has a GID index %d", h.Name, p.RoCEv2GID)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheSubnetManagerAPortSeesIsRead(t *testing.T) {
|
||||
sm, err := ParseSMInfo("sminfo: sm lid 1 sm guid 0x98039B03009fcfd6, activity count 5452 priority 15 state 3 SMINFO_MASTER\n")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sm.GUID != "0x98039b03009fcfd6" || sm.LID != 1 || sm.Priority != 15 || sm.State != "SMINFO_MASTER" {
|
||||
t.Errorf("sm = %+v", sm)
|
||||
}
|
||||
if _, err := ParseSMInfo("ibwarn: [1234] mad_rpc_open_port: can't open UMAD port"); err == nil {
|
||||
t.Error("an error from sminfo should not parse")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPortsFabricIsItsSubnetManagerOrItsNetwork(t *testing.T) {
|
||||
// Two ports on one InfiniBand fabric see one master SM, and on separate
|
||||
// fabrics cannot; RoCE has no SM, and is its IP network.
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if hcas[0].NodeGUID != "" {
|
||||
t.Errorf("no node_guid file, yet %q", hcas[0].NodeGUID)
|
||||
}
|
||||
err = AttachSM(hcas, func(hca string, port int) (*SMInfo, error) {
|
||||
return &SMInfo{GUID: "0x00000000000000aa", LID: 1, State: "SMINFO_MASTER"}, nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := hcas[0].Ports[0].FabricKey(); got != "sm:0x00000000000000aa" {
|
||||
t.Errorf("IB key = %q", got)
|
||||
}
|
||||
if got := hcas[1].Ports[1].FabricKey(); got != "" {
|
||||
t.Errorf("a port that is down is on no fabric, got %q", got)
|
||||
}
|
||||
if hcas[2].Ports[0].SM != nil {
|
||||
t.Error("RoCE was asked for a subnet manager")
|
||||
}
|
||||
if got := hcas[0].Ports[0].PortGUID; got != "0x0002c90300a01234" {
|
||||
t.Errorf("port GUID = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestARoCEPortsNetworkIsItsFabric(t *testing.T) {
|
||||
old := RoCENetwork
|
||||
t.Cleanup(func() { RoCENetwork = old })
|
||||
RoCENetwork = func(ndev, gid string) string {
|
||||
if ip := gidIPv4(gid); ip != nil && ip.String() == "10.0.0.1" {
|
||||
return "10.0.0.0/24"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
hcas, err := DiscoverIB(fakeSysfs(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
roce := hcas[2].Ports[0]
|
||||
if roce.Network != "10.0.0.0/24" || roce.FabricKey() != "roce:10.0.0.0/24" {
|
||||
t.Errorf("RoCE network %q key %q", roce.Network, roce.FabricKey())
|
||||
}
|
||||
}
|
||||
|
||||
// smirDump is `saquery SMIR` on a fabric with a master and a standby, in the
|
||||
// layout infiniband-diags prints.
|
||||
const smirDump = "SMInfoRecord dump:\n\t\tRID\n\t\tLID...................1\n\t\tSMInfo dump:\n" +
|
||||
"\t\tGUID..................0x248a07030030f0c0\n\t\tSM_Key................0x0000000000000000\n" +
|
||||
"\t\tActCount..............386773\n\t\tPriority..............15\n\t\tSMState...............3\n" +
|
||||
"SMInfoRecord dump:\n\t\tRID\n\t\tLID...................12\n\t\tSMInfo dump:\n" +
|
||||
"\t\tGUID..................0x248a07030030f1d4\n\t\tSM_Key................0x0000000000000000\n" +
|
||||
"\t\tActCount..............12\n\t\tPriority..............14\n\t\tSMState...............2\n"
|
||||
|
||||
func TestEverySubnetManagerOnAFabricIsListed(t *testing.T) {
|
||||
recs, err := ParseSMIR(smirDump)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(recs) != 2 {
|
||||
t.Fatalf("records = %+v", recs)
|
||||
}
|
||||
if recs[0].GUID != "0x248a07030030f0c0" || recs[0].State != "master" || recs[0].Priority != 15 || recs[0].LID != 1 {
|
||||
t.Errorf("master = %+v", recs[0])
|
||||
}
|
||||
if recs[1].State != "standby" || recs[1].LID != 12 {
|
||||
t.Errorf("standby = %+v", recs[1])
|
||||
}
|
||||
if _, err := ParseSMIR("ibwarn: can't open UMAD port"); err == nil {
|
||||
t.Error("an error should not parse")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestJSONShape(t *testing.T) {
|
||||
inv := Inventory{
|
||||
Source: "nvidia-smi",
|
||||
GPUs: []GPU{{Index: 0, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559, ComputeCapability: "9.0"}},
|
||||
}
|
||||
b, err := json.Marshal(inv)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var out struct {
|
||||
Source string `json:"source"`
|
||||
GPUs []struct {
|
||||
Index int `json:"index"`
|
||||
Class string `json:"class"`
|
||||
Mem int `json:"memory_total_mb"`
|
||||
} `json:"gpus"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out.Source != "nvidia-smi" || len(out.GPUs) != 1 {
|
||||
t.Fatalf("shape wrong: %s", b)
|
||||
}
|
||||
// The derived class is what decides which GPUs are interchangeable, so a
|
||||
// consumer must not have to recompute it.
|
||||
if out.GPUs[0].Class != "h100-80g" {
|
||||
t.Errorf("class=%q, should be included in the JSON", out.GPUs[0].Class)
|
||||
}
|
||||
if out.GPUs[0].Mem != 81559 {
|
||||
t.Errorf("memory_total_mb=%d", out.GPUs[0].Mem)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,260 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Load is how busy a node is: the machine as a whole, and each engine's share.
|
||||
type Load struct {
|
||||
CPUPercent float64
|
||||
CPUCount int
|
||||
Load1 float64
|
||||
Load5 float64
|
||||
Load15 float64
|
||||
MemUsed int64
|
||||
MemAvail int64
|
||||
Shards []ShardLoad
|
||||
}
|
||||
|
||||
// ShardLoad is one engine process group's use.
|
||||
type ShardLoad struct {
|
||||
ID string
|
||||
CPUCores float64
|
||||
RSSBytes int64
|
||||
}
|
||||
|
||||
// userHZ is the unit /proc counts CPU time in. Fixed at 100 by the kernel's
|
||||
// ABI on every architecture a node runs, whatever the internal tick rate.
|
||||
const userHZ = 100
|
||||
|
||||
// LoadSampler measures load as the difference between two readings of /proc.
|
||||
//
|
||||
// CPU time is a counter, so a percentage needs two samples and the time
|
||||
// between them. The sampler keeps the previous one; the first call has
|
||||
// nothing to compare with and reports no load at all, rather than a zero that
|
||||
// would read as an idle machine.
|
||||
//
|
||||
// Linux only, like every node. Elsewhere the files are missing and it reports
|
||||
// nothing, as an older agent does.
|
||||
type LoadSampler struct {
|
||||
root string
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
prevAt time.Time
|
||||
prevCPU cpuTimes
|
||||
prevPG map[int]groupUse // process group -> what it has used
|
||||
}
|
||||
|
||||
// groupUse is a process group's CPU time so far and its resident memory now.
|
||||
type groupUse struct {
|
||||
jiffies uint64
|
||||
rssPages int64
|
||||
}
|
||||
|
||||
type cpuTimes struct{ total, idle uint64 }
|
||||
|
||||
// NewLoadSampler reads the real /proc.
|
||||
func NewLoadSampler() *LoadSampler { return newLoadSampler("/proc", time.Now) }
|
||||
|
||||
func newLoadSampler(root string, now func() time.Time) *LoadSampler {
|
||||
return &LoadSampler{root: root, now: now}
|
||||
}
|
||||
|
||||
// Sample reads the machine and the process groups led by the given shards'
|
||||
// PIDs, keyed by shard ID. False when there is nothing to report yet.
|
||||
func (s *LoadSampler) Sample(shards map[string]int) (Load, bool) {
|
||||
cpu, ok := s.cpu()
|
||||
if !ok {
|
||||
return Load{}, false
|
||||
}
|
||||
groups := s.groups()
|
||||
at := s.now()
|
||||
|
||||
s.mu.Lock()
|
||||
prevAt, prevCPU, prevPG := s.prevAt, s.prevCPU, s.prevPG
|
||||
s.prevAt, s.prevCPU, s.prevPG = at, cpu, groups
|
||||
s.mu.Unlock()
|
||||
|
||||
if prevAt.IsZero() || cpu.total <= prevCPU.total {
|
||||
return Load{}, false
|
||||
}
|
||||
l := Load{CPUCount: s.cpuCount()}
|
||||
dTotal := float64(cpu.total - prevCPU.total)
|
||||
dIdle := float64(cpu.idle - prevCPU.idle)
|
||||
l.CPUPercent = clamp((dTotal-dIdle)/dTotal*100, 0, 100)
|
||||
l.Load1, l.Load5, l.Load15 = s.loadavg()
|
||||
total, avail := s.meminfo()
|
||||
if total > 0 {
|
||||
l.MemUsed, l.MemAvail = total-avail, avail
|
||||
}
|
||||
|
||||
secs := at.Sub(prevAt).Seconds()
|
||||
for id, pid := range shards {
|
||||
if pid <= 0 {
|
||||
continue
|
||||
}
|
||||
now := groups[pid]
|
||||
sl := ShardLoad{ID: id, RSSBytes: now.rssPages * int64(os.Getpagesize())}
|
||||
if before, ok := prevPG[pid]; ok && secs > 0 && now.jiffies >= before.jiffies {
|
||||
sl.CPUCores = float64(now.jiffies-before.jiffies) / userHZ / secs
|
||||
}
|
||||
l.Shards = append(l.Shards, sl)
|
||||
}
|
||||
return l, true
|
||||
}
|
||||
|
||||
func clamp(v, lo, hi float64) float64 { return max(lo, min(v, hi)) }
|
||||
|
||||
// cpu reads the aggregate line of /proc/stat. Idle includes iowait: a core
|
||||
// waiting on the disk is not doing work.
|
||||
func (s *LoadSampler) cpu() (cpuTimes, bool) {
|
||||
b, err := os.ReadFile(filepath.Join(s.root, "stat"))
|
||||
if err != nil {
|
||||
return cpuTimes{}, false
|
||||
}
|
||||
line, _, _ := bytes.Cut(b, []byte("\n"))
|
||||
f := strings.Fields(string(line))
|
||||
if len(f) < 5 || f[0] != "cpu" {
|
||||
return cpuTimes{}, false
|
||||
}
|
||||
var t cpuTimes
|
||||
for i, v := range f[1:] {
|
||||
// guest and guest_nice are already counted in user and nice.
|
||||
if i >= 8 {
|
||||
break
|
||||
}
|
||||
n, _ := strconv.ParseUint(v, 10, 64)
|
||||
t.total += n
|
||||
if i == 3 || i == 4 {
|
||||
t.idle += n
|
||||
}
|
||||
}
|
||||
return t, true
|
||||
}
|
||||
|
||||
func (s *LoadSampler) cpuCount() int {
|
||||
b, err := os.ReadFile(filepath.Join(s.root, "stat"))
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
n := 0
|
||||
for _, line := range strings.Split(string(b), "\n") {
|
||||
if strings.HasPrefix(line, "cpu") && len(line) > 3 && line[3] >= '0' && line[3] <= '9' {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (s *LoadSampler) loadavg() (l1, l5, l15 float64) {
|
||||
b, err := os.ReadFile(filepath.Join(s.root, "loadavg"))
|
||||
if err != nil {
|
||||
return 0, 0, 0
|
||||
}
|
||||
f := strings.Fields(string(b))
|
||||
if len(f) < 3 {
|
||||
return 0, 0, 0
|
||||
}
|
||||
l1, _ = strconv.ParseFloat(f[0], 64)
|
||||
l5, _ = strconv.ParseFloat(f[1], 64)
|
||||
l15, _ = strconv.ParseFloat(f[2], 64)
|
||||
return l1, l5, l15
|
||||
}
|
||||
|
||||
// meminfo returns MemTotal and MemAvailable in bytes.
|
||||
func (s *LoadSampler) meminfo() (total, avail int64) {
|
||||
f, err := os.Open(filepath.Join(s.root, "meminfo"))
|
||||
if err != nil {
|
||||
return 0, 0
|
||||
}
|
||||
defer f.Close()
|
||||
sc := bufio.NewScanner(f)
|
||||
for sc.Scan() {
|
||||
k, v, ok := strings.Cut(sc.Text(), ":")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fields := strings.Fields(v)
|
||||
if len(fields) == 0 {
|
||||
continue
|
||||
}
|
||||
kb, _ := strconv.ParseInt(fields[0], 10, 64)
|
||||
switch k {
|
||||
case "MemTotal":
|
||||
total = kb * 1024
|
||||
case "MemAvailable":
|
||||
avail = kb * 1024
|
||||
}
|
||||
}
|
||||
return total, avail
|
||||
}
|
||||
|
||||
// procStat is what one /proc/<pid>/stat says that load needs.
|
||||
type procStat struct {
|
||||
pgrp int
|
||||
jiffies uint64
|
||||
rssPages int64
|
||||
}
|
||||
|
||||
// readProcStat parses /proc/<pid>/stat. The command name is in parentheses
|
||||
// and may itself hold spaces and parentheses, so fields are counted from the
|
||||
// last closing one.
|
||||
func readProcStat(path string) (procStat, bool) {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return procStat{}, false
|
||||
}
|
||||
i := bytes.LastIndexByte(b, ')')
|
||||
if i < 0 {
|
||||
return procStat{}, false
|
||||
}
|
||||
f := strings.Fields(string(b[i+1:]))
|
||||
// After the name: state(0) ppid(1) pgrp(2) ... utime(11) stime(12) ... rss(21)
|
||||
if len(f) < 22 {
|
||||
return procStat{}, false
|
||||
}
|
||||
var st procStat
|
||||
st.pgrp, _ = strconv.Atoi(f[2])
|
||||
ut, _ := strconv.ParseUint(f[11], 10, 64)
|
||||
stt, _ := strconv.ParseUint(f[12], 10, 64)
|
||||
st.jiffies = ut + stt
|
||||
st.rssPages, _ = strconv.ParseInt(f[21], 10, 64)
|
||||
return st, true
|
||||
}
|
||||
|
||||
// groups totals CPU time and memory per process group, in one pass over
|
||||
// /proc. An engine is started as the leader of its own group, so its workers
|
||||
// — vLLM starts several — are counted with it.
|
||||
func (s *LoadSampler) groups() map[int]groupUse {
|
||||
out := map[int]groupUse{}
|
||||
s.eachProc(func(st procStat) {
|
||||
g := out[st.pgrp]
|
||||
g.jiffies += st.jiffies
|
||||
g.rssPages += st.rssPages
|
||||
out[st.pgrp] = g
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *LoadSampler) eachProc(fn func(procStat)) {
|
||||
entries, err := os.ReadDir(s.root)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, e := range entries {
|
||||
if _, err := strconv.Atoi(e.Name()); err != nil {
|
||||
continue
|
||||
}
|
||||
if st, ok := readProcStat(filepath.Join(s.root, e.Name(), "stat")); ok {
|
||||
fn(st)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// fakeProc is a /proc with two CPUs, a memory picture and processes.
|
||||
type fakeProc struct {
|
||||
t *testing.T
|
||||
root string
|
||||
}
|
||||
|
||||
func newFakeProc(t *testing.T) *fakeProc {
|
||||
f := &fakeProc{t: t, root: t.TempDir()}
|
||||
f.write("loadavg", "1.50 0.75 0.25 2/300 4242\n")
|
||||
f.write("meminfo", "MemTotal: 8000000 kB\nMemFree: 1000000 kB\nMemAvailable: 3000000 kB\n")
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *fakeProc) write(name, body string) {
|
||||
f.t.Helper()
|
||||
p := filepath.Join(f.root, name)
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
|
||||
f.t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(p, []byte(body), 0o644); err != nil {
|
||||
f.t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// cpu sets the aggregate counters: user and idle, the rest zero.
|
||||
func (f *fakeProc) cpu(user, idle uint64) {
|
||||
f.write("stat", fmt.Sprintf("cpu %d 0 0 %d 0 0 0 0 0 0\ncpu0 0 0 0 0 0 0 0 0 0 0\ncpu1 0 0 0 0 0 0 0 0 0 0\nintr 1\n", user, idle))
|
||||
}
|
||||
|
||||
// proc sets one process: its group, CPU jiffies (as utime) and RSS in pages.
|
||||
func (f *fakeProc) proc(pid, pgrp int, jiffies uint64, rss int64) {
|
||||
// The name holds a space and a parenthesis, as real ones can.
|
||||
f.write(fmt.Sprintf("%d/stat", pid), fmt.Sprintf(
|
||||
"%d (llama-server (x)) S 1 %d %d 0 -1 0 0 0 0 0 %d 0 0 0 20 0 1 0 100 1000 %d 0\n",
|
||||
pid, pgrp, pgrp, jiffies, rss))
|
||||
}
|
||||
|
||||
func TestLoadIsTheDifferenceBetweenTwoReadings(t *testing.T) {
|
||||
p := newFakeProc(t)
|
||||
at := time.Unix(1000, 0)
|
||||
s := newLoadSampler(p.root, func() time.Time { return at })
|
||||
|
||||
p.cpu(100, 900)
|
||||
p.proc(42, 42, 1000, 10)
|
||||
p.proc(43, 42, 500, 5) // a worker in the engine's group
|
||||
p.proc(7, 7, 99999, 1) // something else entirely
|
||||
if _, ok := s.Sample(map[string]int{"shard-a": 42}); ok {
|
||||
t.Fatal("the first reading has nothing to compare with and should report nothing")
|
||||
}
|
||||
|
||||
// Two seconds later: 300 of the 400 jiffies that passed were busy, and
|
||||
// the engine's group used 400 of them.
|
||||
at = at.Add(2 * time.Second)
|
||||
p.cpu(400, 1000)
|
||||
p.proc(42, 42, 1300, 10)
|
||||
p.proc(43, 42, 600, 6)
|
||||
l, ok := s.Sample(map[string]int{"shard-a": 42})
|
||||
if !ok {
|
||||
t.Fatal("the second reading should report")
|
||||
}
|
||||
if l.CPUPercent != 75 || l.CPUCount != 2 {
|
||||
t.Errorf("cpu %.1f%% of %d, want 75%% of 2", l.CPUPercent, l.CPUCount)
|
||||
}
|
||||
if l.Load1 != 1.5 || l.Load15 != 0.25 {
|
||||
t.Errorf("load %v %v", l.Load1, l.Load15)
|
||||
}
|
||||
if l.MemUsed != 5000000*1024 || l.MemAvail != 3000000*1024 {
|
||||
t.Errorf("memory used %d available %d", l.MemUsed, l.MemAvail)
|
||||
}
|
||||
if len(l.Shards) != 1 {
|
||||
t.Fatalf("shards = %+v", l.Shards)
|
||||
}
|
||||
sh := l.Shards[0]
|
||||
// 400 jiffies over 2 seconds at 100 a second is two cores.
|
||||
if sh.ID != "shard-a" || sh.CPUCores != 2 || sh.RSSBytes != 16*int64(os.Getpagesize()) {
|
||||
t.Errorf("shard = %+v", sh)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoProcMeansNoLoad(t *testing.T) {
|
||||
s := newLoadSampler(filepath.Join(t.TempDir(), "missing"), time.Now)
|
||||
s.Sample(nil)
|
||||
if _, ok := s.Sample(nil); ok {
|
||||
t.Error("a system without /proc should report nothing, not zeros")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// MemoryTotalMB reports the machine's total RAM.
|
||||
//
|
||||
// It exists because placement had no idea how much a node could hold. A CPU
|
||||
// profile has no card to be free, so every ready node looked like a candidate:
|
||||
// a 14B model whose working set is 24 GB would have been placed just as
|
||||
// happily on a node with 8 GB, which reports itself perfectly healthy right up
|
||||
// until the OOM killer arrives. Labels were the stopgap; this is the fact they
|
||||
// were standing in for.
|
||||
//
|
||||
// Read from /proc/meminfo, which is Linux-only — every node is. On anything
|
||||
// else it returns 0, which callers treat as "unknown" and therefore no
|
||||
// constraint, exactly as an older agent that does not report it at all.
|
||||
func MemoryTotalMB() int { return memTotalMBFrom("/proc/meminfo") }
|
||||
|
||||
func memTotalMBFrom(path string) int {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
sc := bufio.NewScanner(f)
|
||||
for sc.Scan() {
|
||||
rest, ok := strings.CutPrefix(sc.Text(), "MemTotal:")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fields := strings.Fields(rest)
|
||||
if len(fields) == 0 {
|
||||
return 0
|
||||
}
|
||||
kb, err := strconv.Atoi(fields[0])
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
// MemTotal is what the kernel can hand out, already short of the
|
||||
// installed total by the firmware's reservations. That is the right
|
||||
// number for deciding what fits.
|
||||
return kb / 1024
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMemTotalIsReadFromMeminfo(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "meminfo")
|
||||
// The real shape, including the fields before and after, so the parser is
|
||||
// not tested against something tidier than it will see.
|
||||
if err := os.WriteFile(p, []byte(
|
||||
"MemTotal: 32876244 kB\nMemFree: 1234 kB\nMemAvailable: 5678 kB\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := memTotalMBFrom(p), 32876244/1024; got != want {
|
||||
t.Errorf("memTotalMBFrom = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreadableMeminfoReportsUnknownRatherThanZeroMemory(t *testing.T) {
|
||||
// Zero means "did not say" everywhere downstream, and placement treats it
|
||||
// as no constraint. A node with genuinely no memory is not a case.
|
||||
if got := memTotalMBFrom(filepath.Join(t.TempDir(), "absent")); got != 0 {
|
||||
t.Errorf("got %d for a missing file", got)
|
||||
}
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "meminfo")
|
||||
for _, junk := range []string{"", "MemTotal:\n", "MemTotal: not-a-number kB\n", "Other: 1 kB\n"} {
|
||||
if err := os.WriteFile(p, []byte(junk), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := memTotalMBFrom(p); got != 0 {
|
||||
t.Errorf("got %d for %q", got, junk)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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])
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Sensors is what a node's own hardware says about heat and power: the CPUs,
|
||||
// the board where it has an inlet sensor, and each GPU. A zero is not
|
||||
// measured, never a reading of zero — a VM has none of these, and must not
|
||||
// look like a machine running at 0 °C on no power.
|
||||
type Sensors struct {
|
||||
// CPUTempC is the hottest CPU sensor: the package, or the hottest core
|
||||
// where no package reading exists.
|
||||
CPUTempC float64
|
||||
// CPUPowerW is the CPU packages' draw, from RAPL: the change in their
|
||||
// energy counters over the time between two samples.
|
||||
CPUPowerW float64
|
||||
// InletTempC is the air coming in, where the board has a sensor it calls
|
||||
// that. Most do not, and a CPU's temperature is not a substitute: it
|
||||
// follows the load, not the room.
|
||||
InletTempC float64
|
||||
GPUs []GPUSensor
|
||||
}
|
||||
|
||||
// GPUSensor is one card's heat, power and whether it is slowing itself down.
|
||||
type GPUSensor struct {
|
||||
Index int
|
||||
TempC float64
|
||||
MemoryTempC float64
|
||||
PowerW float64
|
||||
PowerLimitW float64
|
||||
// Throttle names what is holding the card's clocks down now:
|
||||
// sw_power_cap, hw_slowdown, sw_thermal, hw_thermal, power_brake. Empty
|
||||
// is running free. Idle and application clock settings are left out;
|
||||
// they are not the card being held back.
|
||||
Throttle []string
|
||||
}
|
||||
|
||||
// sensorEvery bounds how often nvidia-smi is run. Heartbeats come more often
|
||||
// than temperatures change, and each run costs a few hundred milliseconds.
|
||||
const sensorEvery = 10 * time.Second
|
||||
|
||||
// SensorSampler reads Sensors. The first sample has no CPU power, which needs
|
||||
// two readings of the energy counters.
|
||||
type SensorSampler struct {
|
||||
sys string // /sys, or a test's copy of it
|
||||
smi func(ctx context.Context, fields string) ([]byte, error)
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
last Sensors
|
||||
lastAt time.Time
|
||||
prevEnergy map[string]uint64
|
||||
prevAt time.Time
|
||||
}
|
||||
|
||||
// NewSensorSampler reads this machine.
|
||||
func NewSensorSampler() *SensorSampler {
|
||||
return &SensorSampler{sys: "/sys", smi: runSMI, now: time.Now}
|
||||
}
|
||||
|
||||
func runSMI(ctx context.Context, fields string) ([]byte, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
//nolint:gosec // fixed binary, fields from this file
|
||||
return exec.CommandContext(ctx, "nvidia-smi", "--query-gpu="+fields,
|
||||
"--format=csv,noheader,nounits").Output()
|
||||
}
|
||||
|
||||
// Sample reads the sensors, or returns the last reading when that is recent.
|
||||
func (s *SensorSampler) Sample(ctx context.Context) Sensors {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
now := s.now()
|
||||
if !s.lastAt.IsZero() && now.Sub(s.lastAt) < sensorEvery {
|
||||
return s.last
|
||||
}
|
||||
out := Sensors{}
|
||||
out.CPUTempC, out.InletTempC = s.hwmon()
|
||||
out.CPUPowerW = s.rapl(now)
|
||||
out.GPUs = s.gpus(ctx)
|
||||
s.last, s.lastAt = out, now
|
||||
return out
|
||||
}
|
||||
|
||||
// hwmon reads the hottest CPU sensor and an inlet sensor, from every hwmon
|
||||
// device the kernel exposes — what lm_sensors reads, without it installed.
|
||||
func (s *SensorSampler) hwmon() (cpu, inlet float64) {
|
||||
devs, _ := filepath.Glob(filepath.Join(s.sys, "class/hwmon/hwmon*"))
|
||||
var pkg, core float64
|
||||
for _, d := range devs {
|
||||
name := readTrim(filepath.Join(d, "name"))
|
||||
inputs, _ := filepath.Glob(filepath.Join(d, "temp*_input"))
|
||||
for _, in := range inputs {
|
||||
milli, err := strconv.ParseFloat(readTrim(in), 64)
|
||||
if err != nil || milli <= 0 {
|
||||
continue
|
||||
}
|
||||
c := milli / 1000
|
||||
label := strings.ToLower(readTrim(strings.TrimSuffix(in, "_input") + "_label"))
|
||||
switch {
|
||||
case isCPUSensor(name) && (strings.HasPrefix(label, "package") || label == "tctl" || label == "tdie"):
|
||||
pkg = max(pkg, c)
|
||||
case isCPUSensor(name):
|
||||
core = max(core, c)
|
||||
case label == "cpu":
|
||||
// A board's own reading of the CPU, as Dell's SMM gives.
|
||||
core = max(core, c)
|
||||
case strings.Contains(label, "inlet") || strings.Contains(label, "ambient") ||
|
||||
strings.Contains(label, "intake") || label == "systin":
|
||||
inlet = max(inlet, c)
|
||||
}
|
||||
}
|
||||
}
|
||||
if pkg > 0 {
|
||||
return pkg, inlet
|
||||
}
|
||||
return core, inlet
|
||||
}
|
||||
|
||||
func isCPUSensor(name string) bool {
|
||||
switch name {
|
||||
case "coretemp", "k10temp", "zenpower", "cpu_thermal":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// rapl is the CPU packages' power from their energy counters. Package
|
||||
// domains only: the dram and core subdomains are parts of a package, and
|
||||
// adding them would count the same energy twice.
|
||||
func (s *SensorSampler) rapl(now time.Time) float64 {
|
||||
doms, _ := filepath.Glob(filepath.Join(s.sys, "class/powercap/intel-rapl:*"))
|
||||
energy := map[string]uint64{}
|
||||
wrap := map[string]uint64{}
|
||||
for _, d := range doms {
|
||||
if strings.Count(filepath.Base(d), ":") != 1 {
|
||||
continue // intel-rapl:0:0 is a subdomain of intel-rapl:0
|
||||
}
|
||||
e, err := strconv.ParseUint(readTrim(filepath.Join(d, "energy_uj")), 10, 64)
|
||||
if err != nil {
|
||||
continue // unreadable without root on newer kernels
|
||||
}
|
||||
energy[d] = e
|
||||
wrap[d], _ = strconv.ParseUint(readTrim(filepath.Join(d, "max_energy_range_uj")), 10, 64)
|
||||
}
|
||||
prev, prevAt := s.prevEnergy, s.prevAt
|
||||
s.prevEnergy, s.prevAt = energy, now
|
||||
dt := now.Sub(prevAt).Seconds()
|
||||
if prev == nil || dt <= 0 || len(energy) == 0 {
|
||||
return 0
|
||||
}
|
||||
var joules float64
|
||||
for d, e := range energy {
|
||||
p, ok := prev[d]
|
||||
if !ok {
|
||||
return 0 // a domain appeared: no fair difference this time
|
||||
}
|
||||
if e < p { // the counter wrapped
|
||||
e += wrap[d]
|
||||
}
|
||||
joules += float64(e-p) / 1e6
|
||||
}
|
||||
return joules / dt
|
||||
}
|
||||
|
||||
// The throttle reasons nvidia-smi reports, as bits. Only those that mean the
|
||||
// card is being held back are named.
|
||||
var throttleBits = []struct {
|
||||
bit uint64
|
||||
name string
|
||||
}{
|
||||
{0x4, "sw_power_cap"},
|
||||
{0x8, "hw_slowdown"},
|
||||
{0x20, "sw_thermal"},
|
||||
{0x40, "hw_thermal"},
|
||||
{0x80, "power_brake"},
|
||||
}
|
||||
|
||||
// gpus asks nvidia-smi for each card's heat and power. The memory temperature
|
||||
// is not known to every driver, and a field nvidia-smi does not know fails the
|
||||
// whole query, so it is asked for first and dropped on failure.
|
||||
func (s *SensorSampler) gpus(ctx context.Context) []GPUSensor {
|
||||
withMem := "index,temperature.gpu,power.draw,enforced.power.limit,clocks_throttle_reasons.active,temperature.memory"
|
||||
out, err := s.smi(ctx, withMem)
|
||||
if err != nil {
|
||||
if out, err = s.smi(ctx, strings.TrimSuffix(withMem, ",temperature.memory")); err != nil {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
var gs []GPUSensor
|
||||
for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
|
||||
f := strings.Split(line, ",")
|
||||
if len(f) < 5 {
|
||||
continue
|
||||
}
|
||||
for i := range f {
|
||||
f[i] = strings.TrimSpace(f[i])
|
||||
}
|
||||
idx, err := strconv.Atoi(f[0])
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
g := GPUSensor{Index: idx, TempC: num(f[1]), PowerW: num(f[2]), PowerLimitW: num(f[3])}
|
||||
if bits, err := strconv.ParseUint(strings.TrimPrefix(f[4], "0x"), 16, 64); err == nil {
|
||||
for _, b := range throttleBits {
|
||||
if bits&b.bit != 0 {
|
||||
g.Throttle = append(g.Throttle, b.name)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(f) > 5 {
|
||||
g.MemoryTempC = num(f[5])
|
||||
}
|
||||
gs = append(gs, g)
|
||||
}
|
||||
return gs
|
||||
}
|
||||
|
||||
// num reads a figure nvidia-smi may give as "[N/A]" or "N/A".
|
||||
func num(s string) float64 {
|
||||
v, err := strconv.ParseFloat(s, 64)
|
||||
if err != nil || v < 0 {
|
||||
return 0
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func readTrim(path string) string {
|
||||
b, err := os.ReadFile(path) //nolint:gosec // sysfs
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func put(t *testing.T, root, rel, body string) {
|
||||
t.Helper()
|
||||
p := filepath.Join(root, rel)
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(p, []byte(body+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// dellHost is the layout of the Proxmox host these tests were written
|
||||
// against: Dell's SMM with a CPU and SODIMM sensors, coretemp with a package
|
||||
// and cores, and RAPL with a package and its dram subdomain.
|
||||
func dellHost(t *testing.T) string {
|
||||
root := t.TempDir()
|
||||
put(t, root, "class/hwmon/hwmon0/name", "dell_smm")
|
||||
put(t, root, "class/hwmon/hwmon0/temp1_input", "55000")
|
||||
put(t, root, "class/hwmon/hwmon0/temp1_label", "CPU")
|
||||
put(t, root, "class/hwmon/hwmon0/temp2_input", "47000")
|
||||
put(t, root, "class/hwmon/hwmon0/temp2_label", "SODIMM")
|
||||
put(t, root, "class/hwmon/hwmon1/name", "coretemp")
|
||||
put(t, root, "class/hwmon/hwmon1/temp1_input", "58000")
|
||||
put(t, root, "class/hwmon/hwmon1/temp1_label", "Package id 0")
|
||||
put(t, root, "class/hwmon/hwmon1/temp10_input", "61000")
|
||||
put(t, root, "class/hwmon/hwmon1/temp10_label", "Core 8")
|
||||
put(t, root, "class/powercap/intel-rapl:0/energy_uj", "1000000000")
|
||||
put(t, root, "class/powercap/intel-rapl:0/max_energy_range_uj", "262143328850")
|
||||
put(t, root, "class/powercap/intel-rapl:0:0/energy_uj", "5000000")
|
||||
return root
|
||||
}
|
||||
|
||||
func sampler(sys string, smi func(context.Context, string) ([]byte, error), clock *time.Time) *SensorSampler {
|
||||
return &SensorSampler{sys: sys, smi: smi, now: func() time.Time { return *clock }}
|
||||
}
|
||||
|
||||
func noSMI(context.Context, string) ([]byte, error) { return nil, errors.New("not found") }
|
||||
|
||||
func TestAHostsCPUHeatAndPowerAreReadAsLmSensorsWould(t *testing.T) {
|
||||
root := dellHost(t)
|
||||
now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
s := sampler(root, noSMI, &now)
|
||||
|
||||
first := s.Sample(t.Context())
|
||||
// The package reading, not the hotter core: the package is what the
|
||||
// chip itself reports as its temperature.
|
||||
if first.CPUTempC != 58 {
|
||||
t.Errorf("CPU %.1f °C, expected the package's 58", first.CPUTempC)
|
||||
}
|
||||
if first.CPUPowerW != 0 {
|
||||
t.Error("power from one reading of a counter is invented")
|
||||
}
|
||||
if first.InletTempC != 0 || first.GPUs != nil {
|
||||
t.Errorf("sensors this host does not have were reported: %+v", first)
|
||||
}
|
||||
|
||||
// 10 s later the package has used 650 J, and its dram subdomain some of
|
||||
// that — counted once.
|
||||
put(t, root, "class/powercap/intel-rapl:0/energy_uj", "1650000000")
|
||||
put(t, root, "class/powercap/intel-rapl:0:0/energy_uj", "99000000")
|
||||
now = now.Add(10 * time.Second)
|
||||
if got := s.Sample(t.Context()).CPUPowerW; got < 64.9 || got > 65.1 {
|
||||
t.Errorf("CPU power %.2f W, expected 65", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestACounterThatWrapsIsNotANegativeDraw(t *testing.T) {
|
||||
root := dellHost(t)
|
||||
now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
put(t, root, "class/powercap/intel-rapl:0/energy_uj", "262143000000")
|
||||
s := sampler(root, noSMI, &now)
|
||||
s.Sample(t.Context())
|
||||
// 328,850 µJ to the wrap, then 499,671,150 past it: 500 J in 10 s.
|
||||
put(t, root, "class/powercap/intel-rapl:0/energy_uj", "499671150")
|
||||
now = now.Add(10 * time.Second)
|
||||
if got := s.Sample(t.Context()).CPUPowerW; got < 49.9 || got > 50.1 {
|
||||
t.Errorf("across a wrap: %.2f W, expected 50", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAVMHasNoSensorsAndSaysNothing(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
put(t, root, "class/thermal/cooling_device0/type", "Processor")
|
||||
now := time.Now()
|
||||
got := sampler(root, noSMI, &now).Sample(t.Context())
|
||||
if got.CPUTempC != 0 || got.CPUPowerW != 0 || got.InletTempC != 0 || got.GPUs != nil {
|
||||
t.Errorf("a machine with no sensors reported %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestABoardsInletSensorIsTheRoomNotTheCPU(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
put(t, root, "class/hwmon/hwmon2/name", "nct6798")
|
||||
put(t, root, "class/hwmon/hwmon2/temp1_input", "24500")
|
||||
put(t, root, "class/hwmon/hwmon2/temp1_label", "SYSTIN")
|
||||
put(t, root, "class/hwmon/hwmon3/name", "k10temp")
|
||||
put(t, root, "class/hwmon/hwmon3/temp1_input", "71250")
|
||||
put(t, root, "class/hwmon/hwmon3/temp1_label", "Tctl")
|
||||
now := time.Now()
|
||||
got := sampler(root, noSMI, &now).Sample(t.Context())
|
||||
if got.InletTempC != 24.5 || got.CPUTempC != 71.25 {
|
||||
t.Errorf("inlet %.2f, CPU %.2f", got.InletTempC, got.CPUTempC)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEachGPUSaysHowHotHowHungryAndWhatHoldsItBack(t *testing.T) {
|
||||
var asked []string
|
||||
smi := func(_ context.Context, fields string) ([]byte, error) {
|
||||
asked = append(asked, fields)
|
||||
return []byte(
|
||||
"0, 71, 612.45, 700.00, 0x0000000000000000, 84\n" +
|
||||
"1, 88, 698.10, 700.00, 0x0000000000000024, 95\n" + // sw_power_cap + sw_thermal
|
||||
"2, 45, [N/A], 700.00, 0x0000000000000001, [N/A]\n"), nil // idle only
|
||||
}
|
||||
now := time.Now()
|
||||
ss := sampler(t.TempDir(), smi, &now)
|
||||
got := ss.Sample(t.Context())
|
||||
if len(got.GPUs) != 3 {
|
||||
t.Fatalf("%+v", got.GPUs)
|
||||
}
|
||||
g0, g1, g2 := got.GPUs[0], got.GPUs[1], got.GPUs[2]
|
||||
if g0.TempC != 71 || g0.PowerW != 612.45 || g0.PowerLimitW != 700 || g0.MemoryTempC != 84 || g0.Throttle != nil {
|
||||
t.Errorf("card 0 %+v", g0)
|
||||
}
|
||||
if strings.Join(g1.Throttle, ",") != "sw_power_cap,sw_thermal" {
|
||||
t.Errorf("card 1 throttled by %v", g1.Throttle)
|
||||
}
|
||||
if g2.PowerW != 0 || g2.MemoryTempC != 0 || g2.Throttle != nil {
|
||||
t.Errorf("card 2: N/A is not a reading, and idle is not being held back: %+v", g2)
|
||||
}
|
||||
// Within ten seconds the last reading is reused rather than running
|
||||
// nvidia-smi on every heartbeat; after, it is asked again.
|
||||
now = now.Add(5 * time.Second)
|
||||
ss.Sample(t.Context())
|
||||
if len(asked) != 1 {
|
||||
t.Errorf("nvidia-smi was asked %d times within ten seconds", len(asked))
|
||||
}
|
||||
now = now.Add(6 * time.Second)
|
||||
ss.Sample(t.Context())
|
||||
if len(asked) != 2 {
|
||||
t.Errorf("nvidia-smi was asked %d times after eleven seconds, expected 2", len(asked))
|
||||
}
|
||||
}
|
||||
|
||||
func TestADriverThatDoesNotKnowMemoryTemperatureStillReportsTheRest(t *testing.T) {
|
||||
smi := func(_ context.Context, fields string) ([]byte, error) {
|
||||
if strings.Contains(fields, "temperature.memory") {
|
||||
return nil, errors.New(`Field "temperature.memory" is not a valid field to query.`)
|
||||
}
|
||||
return []byte("0, 66, 250.0, 300.0, 0x0000000000000040\n"), nil
|
||||
}
|
||||
now := time.Now()
|
||||
got := sampler(t.TempDir(), smi, &now).Sample(t.Context())
|
||||
if len(got.GPUs) != 1 || got.GPUs[0].TempC != 66 || strings.Join(got.GPUs[0].Throttle, ",") != "hw_thermal" {
|
||||
t.Errorf("%+v", got.GPUs)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrNoDiags means infiniband-diags is not installed, so the subnet manager
|
||||
// cannot be asked who it is.
|
||||
var ErrNoDiags = errors.New("sminfo not found: install infiniband-diags")
|
||||
|
||||
// smLine is sminfo's one line of answer:
|
||||
//
|
||||
// sminfo: sm lid 1 sm guid 0x98039b03009fcfd6, activity count 5452 priority 15 state 3 SMINFO_MASTER
|
||||
var smLine = regexp.MustCompile(`sm lid (\d+) sm guid (0x[0-9a-fA-F]+), activity count \d+ priority (\d+) state \d+ (\S+)`)
|
||||
|
||||
// ParseSMInfo reads sminfo's output.
|
||||
func ParseSMInfo(out string) (*SMInfo, error) {
|
||||
m := smLine.FindStringSubmatch(out)
|
||||
if m == nil {
|
||||
return nil, fmt.Errorf("unrecognised sminfo output: %q", strings.TrimSpace(out))
|
||||
}
|
||||
lid, _ := strconv.Atoi(m[1])
|
||||
prio, _ := strconv.Atoi(m[3])
|
||||
return &SMInfo{GUID: normalizeGUID(m[2]), LID: lid, Priority: prio, State: m[4]}, nil
|
||||
}
|
||||
|
||||
// QuerySM asks, through one port, which subnet manager is master. It is a
|
||||
// management datagram to the SM, not a scan: cheap, but not free, which is
|
||||
// why the agent asks once a minute rather than every heartbeat.
|
||||
func QuerySM(ctx context.Context, hca string, port int) (*SMInfo, error) {
|
||||
bin, err := exec.LookPath("sminfo")
|
||||
if err != nil {
|
||||
return nil, ErrNoDiags
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
out, err := exec.CommandContext(ctx, bin, "-C", hca, "-P", strconv.Itoa(port)).CombinedOutput()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("sminfo on %s/%d: %w: %s", hca, port, err, strings.TrimSpace(string(out)))
|
||||
}
|
||||
return ParseSMInfo(string(out))
|
||||
}
|
||||
|
||||
// AttachSM fills in the subnet manager each active InfiniBand port sees.
|
||||
// RoCE has none. The first error is returned for reporting; the ports it
|
||||
// could answer for keep their answer.
|
||||
func AttachSM(hcas []HCA, query func(hca string, port int) (*SMInfo, error)) error {
|
||||
var first error
|
||||
for i := range hcas {
|
||||
for j := range hcas[i].Ports {
|
||||
p := &hcas[i].Ports[j]
|
||||
if !p.Active() || p.LinkLayer == LinkEthernet {
|
||||
continue
|
||||
}
|
||||
sm, err := query(hcas[i].Name, p.Num)
|
||||
if err != nil {
|
||||
if first == nil {
|
||||
first = err
|
||||
}
|
||||
continue
|
||||
}
|
||||
p.SM = sm
|
||||
}
|
||||
}
|
||||
return first
|
||||
}
|
||||
|
||||
// SMRecord is one subnet manager on a fabric, as the subnet administrator
|
||||
// lists them: master and standbys alike.
|
||||
type SMRecord struct {
|
||||
GUID string `json:"guid"`
|
||||
LID int `json:"lid"`
|
||||
Priority int `json:"priority"`
|
||||
// State is master, standby, discovering or not active.
|
||||
State string `json:"state"`
|
||||
}
|
||||
|
||||
// smStates are SMInfo's state numbers.
|
||||
var smStates = map[string]string{"0": "not active", "1": "discovering", "2": "standby", "3": "master"}
|
||||
|
||||
// smirField is one line of saquery's record dump: "LID.........1".
|
||||
var smirField = regexp.MustCompile(`^\s*([A-Za-z_]+)\.{2,}\s*(\S+)`)
|
||||
|
||||
// ParseSMIR reads `saquery SMIR`: every subnet manager the fabric's subnet
|
||||
// administrator knows of, one "SMInfoRecord dump:" block each.
|
||||
func ParseSMIR(out string) ([]SMRecord, error) {
|
||||
var recs []SMRecord
|
||||
var cur *SMRecord
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
if strings.Contains(line, "SMInfoRecord dump") {
|
||||
recs = append(recs, SMRecord{})
|
||||
cur = &recs[len(recs)-1]
|
||||
continue
|
||||
}
|
||||
m := smirField.FindStringSubmatch(line)
|
||||
if m == nil || cur == nil {
|
||||
continue
|
||||
}
|
||||
switch strings.ToLower(m[1]) {
|
||||
case "lid":
|
||||
cur.LID, _ = strconv.Atoi(m[2])
|
||||
case "guid":
|
||||
cur.GUID = normalizeGUID(m[2])
|
||||
case "priority":
|
||||
cur.Priority, _ = strconv.Atoi(m[2])
|
||||
case "smstate":
|
||||
cur.State = smStates[m[2]]
|
||||
if cur.State == "" {
|
||||
cur.State = m[2]
|
||||
}
|
||||
}
|
||||
}
|
||||
out2 := recs[:0]
|
||||
for _, r := range recs {
|
||||
if r.GUID != "" {
|
||||
out2 = append(out2, r)
|
||||
}
|
||||
}
|
||||
if len(out2) == 0 {
|
||||
return nil, fmt.Errorf("no subnet manager records in saquery output")
|
||||
}
|
||||
return out2, nil
|
||||
}
|
||||
|
||||
// QuerySMs asks the subnet administrator, through one port, for every subnet
|
||||
// manager on the fabric.
|
||||
func QuerySMs(ctx context.Context, hca string, port int) ([]SMRecord, error) {
|
||||
bin, err := exec.LookPath("saquery")
|
||||
if err != nil {
|
||||
return nil, ErrNoDiags
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
out, err := exec.CommandContext(ctx, bin, "-C", hca, "-P", strconv.Itoa(port), "SMIR").CombinedOutput()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("saquery on %s/%d: %w: %s", hca, port, err, strings.TrimSpace(string(out)))
|
||||
}
|
||||
return ParseSMIR(string(out))
|
||||
}
|
||||
|
||||
// ScanFabric maps the fabric a port is on with ibnetdiscover: every switch,
|
||||
// adapter and link. It can take a while on a large fabric, so it is bounded,
|
||||
// and its output is capped: a fabric of thousands of ports is a few
|
||||
// megabytes of text.
|
||||
func ScanFabric(ctx context.Context, hca string, port int) (string, error) {
|
||||
bin, err := exec.LookPath("ibnetdiscover")
|
||||
if err != nil {
|
||||
return "", ErrNoDiags
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
|
||||
defer cancel()
|
||||
out, err := exec.CommandContext(ctx, bin, "-C", hca, "-P", strconv.Itoa(port)).Output()
|
||||
if err != nil {
|
||||
var ee *exec.ExitError
|
||||
if errors.As(err, &ee) {
|
||||
return "", fmt.Errorf("ibnetdiscover on %s/%d: %w: %s", hca, port, err, strings.TrimSpace(string(ee.Stderr)))
|
||||
}
|
||||
return "", fmt.Errorf("ibnetdiscover on %s/%d: %w", hca, port, err)
|
||||
}
|
||||
const limit = 16 << 20
|
||||
if len(out) > limit {
|
||||
return "", fmt.Errorf("ibnetdiscover printed %d MB, more than the %d MB a scan may send", len(out)>>20, limit>>20)
|
||||
}
|
||||
return string(out), nil
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// One entry per filesystem rather than per configured directory, because free
|
||||
// space belongs to the disk and not to the path. Two model directories on the
|
||||
// same disk are one pool of bytes, and reporting them separately would let the
|
||||
// control plane add 24 GB to 24 GB and place a 40 GB model on a node with 24.
|
||||
type Storage struct {
|
||||
// Path is one directory on this filesystem — the first one configured that
|
||||
// landed here, so the operator recognises it.
|
||||
Path string `json:"path"`
|
||||
// FreeBytes is what a non-root process can still write. Deliberately not
|
||||
// the root reserve: the agent does not run as root, so the reserve is space
|
||||
// it cannot have.
|
||||
FreeBytes int64 `json:"free_bytes"`
|
||||
// TotalBytes is the filesystem's size, so a reading of "8 GB free" can be
|
||||
// told apart on a 10 TB filer and on a 10 GB root partition.
|
||||
TotalBytes int64 `json:"total_bytes"`
|
||||
// Shared marks a filesystem the node did not have to fill. Free space on
|
||||
// one is shared with every other node using it, so two nodes reporting
|
||||
// 500 GB free on the same filer have 500 GB between them, not a terabyte.
|
||||
Shared bool `json:"shared,omitempty"`
|
||||
}
|
||||
|
||||
// ScanStorage reports the filesystems behind the given model directories.
|
||||
//
|
||||
// Empty in, empty out, and a directory that cannot be read is skipped rather
|
||||
// than reported as full: a mount that has not come back yet looks exactly like
|
||||
// a missing directory, and claiming zero free bytes would take the node out of
|
||||
// placement for a reason that is not true. Absent is not zero here, as
|
||||
// everywhere else in this system.
|
||||
func ScanStorage(dirs []string) []Storage {
|
||||
var out []Storage
|
||||
seen := map[uint64]int{} // filesystem id -> index in out
|
||||
|
||||
for _, dir := range dirs {
|
||||
if dir == "" {
|
||||
continue
|
||||
}
|
||||
dev, ok := filesystemID(dir)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if i, dup := seen[dev]; dup {
|
||||
// Same disk, already counted. Keep the first path: it is the one
|
||||
// the operator wrote first and so the one they will recognise.
|
||||
_ = i
|
||||
continue
|
||||
}
|
||||
|
||||
var st syscall.Statfs_t
|
||||
if err := syscall.Statfs(dir, &st); err != nil {
|
||||
continue
|
||||
}
|
||||
// Bsize is int32 on some platforms and int64 on others, which is why
|
||||
// every field here is widened explicitly rather than assumed.
|
||||
block := int64(st.Bsize)
|
||||
seen[dev] = len(out)
|
||||
out = append(out, Storage{
|
||||
Path: dir,
|
||||
FreeBytes: int64(st.Bavail) * block, //nolint:gosec,unconvert // widened on purpose
|
||||
TotalBytes: int64(st.Blocks) * block, //nolint:gosec,unconvert // widened on purpose
|
||||
Shared: isNetworkFS(dir),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// filesystemID identifies the disk a path is on.
|
||||
//
|
||||
// By device number rather than by the statfs filesystem id, because the field
|
||||
// holding that id is spelled differently on Linux and on macOS and this has to
|
||||
// compile on both — the agent runs on one and is developed on the other.
|
||||
func filesystemID(path string) (uint64, bool) {
|
||||
fi, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
st, ok := fi.Sys().(*syscall.Stat_t)
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
return uint64(st.Dev), true //nolint:gosec,unconvert // int32 on macOS, uint64 on Linux
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTwoDirectoriesOnOneDiskAreOnePoolOfBytes(t *testing.T) {
|
||||
// The reason this matters: free space belongs to the disk, not to the path.
|
||||
// Counting two directories separately lets the control plane add 24 GB to
|
||||
// 24 GB and place a 40 GB model on a node that has 24.
|
||||
base := t.TempDir()
|
||||
a := filepath.Join(base, "models-a")
|
||||
b := filepath.Join(base, "models-b")
|
||||
for _, d := range []string{a, b} {
|
||||
if err := os.Mkdir(d, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
got := ScanStorage([]string{a, b})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("two directories on one filesystem reported as %d pools: %+v", len(got), got)
|
||||
}
|
||||
if got[0].Path != a {
|
||||
t.Errorf("kept %q rather than the first directory configured", got[0].Path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStorageReportsSomethingPlausible(t *testing.T) {
|
||||
got := ScanStorage([]string{t.TempDir()})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("a real directory reported %d filesystems", len(got))
|
||||
}
|
||||
s := got[0]
|
||||
if s.TotalBytes <= 0 {
|
||||
t.Error("a filesystem with no size at all")
|
||||
}
|
||||
if s.FreeBytes < 0 || s.FreeBytes > s.TotalBytes {
|
||||
t.Errorf("free %d of total %d", s.FreeBytes, s.TotalBytes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingDirectoryIsNotAFullDisk(t *testing.T) {
|
||||
// A mount that has not come back yet looks exactly like a missing
|
||||
// directory. Reporting it as zero free bytes would take the node out of
|
||||
// placement for a reason that is not true; reporting nothing says only
|
||||
// that nothing is known, which is the honest answer.
|
||||
got := ScanStorage([]string{filepath.Join(t.TempDir(), "never-mounted")})
|
||||
if len(got) != 0 {
|
||||
t.Errorf("a missing directory was reported as a filesystem: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoDirectoriesConstrainNothing(t *testing.T) {
|
||||
if got := ScanStorage(nil); len(got) != 0 {
|
||||
t.Errorf("scanning nothing found %+v", got)
|
||||
}
|
||||
if got := ScanStorage([]string{""}); len(got) != 0 {
|
||||
t.Errorf("an empty path found %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Topology is what `nvidia-smi topo -m` says about how a node's GPUs and
|
||||
// network adapters are wired together.
|
||||
//
|
||||
// Read from nvidia-smi rather than NVML so the portable, cgo-free binary knows
|
||||
// it too. Without it that binary saw no NVLink at all: every GPU reported no
|
||||
// peers, and a tensor-parallel group was chosen by index on machines where
|
||||
// NVLink decides whether it runs at full speed.
|
||||
type Topology struct {
|
||||
// NVLink is, per GPU, the GPUs it reaches over NVLink and how many links
|
||||
// are bonded to each.
|
||||
NVLink map[int]map[int]int
|
||||
// NICPath is, per GPU, how far each network adapter is, by its HCA name:
|
||||
// PIX, PXB, PHB, NODE or SYS, nearest first. A GPU should talk to the
|
||||
// fabric through the adapter nearest it; one across the socket costs a
|
||||
// trip over the processor interconnect on every transfer.
|
||||
NICPath map[int]map[string]string
|
||||
// NUMA is the NUMA node each GPU is attached to, where reported.
|
||||
NUMA map[int]int
|
||||
}
|
||||
|
||||
// PathRank orders the relations nvidia-smi reports, nearest first. Lower is
|
||||
// closer; an unknown relation ranks last.
|
||||
func PathRank(rel string) int {
|
||||
switch strings.ToUpper(rel) {
|
||||
case "PIX":
|
||||
return 0
|
||||
case "PXB":
|
||||
return 1
|
||||
case "PHB":
|
||||
return 2
|
||||
case "NODE":
|
||||
return 3
|
||||
case "SYS":
|
||||
return 4
|
||||
}
|
||||
return 5
|
||||
}
|
||||
|
||||
// DiscoverTopology runs `nvidia-smi topo -m`. A node without nvidia-smi has
|
||||
// no topology to report, which is not an error.
|
||||
func DiscoverTopology(ctx context.Context, binary string) (Topology, error) {
|
||||
if binary == "" {
|
||||
if _, err := exec.LookPath("nvidia-smi"); err != nil {
|
||||
return Topology{}, nil
|
||||
}
|
||||
binary = "nvidia-smi"
|
||||
}
|
||||
out, err := exec.CommandContext(ctx, binary, "topo", "-m").Output()
|
||||
if err != nil {
|
||||
return Topology{}, fmt.Errorf("running %s topo -m: %w", binary, err)
|
||||
}
|
||||
return parseTopo(string(out))
|
||||
}
|
||||
|
||||
// ParseTopology reads `nvidia-smi topo -m` output.
|
||||
func ParseTopology(out string) (Topology, error) { return parseTopo(out) }
|
||||
|
||||
var (
|
||||
ansi = regexp.MustCompile(`\x1b\[[0-9;]*m`)
|
||||
nvLinks = regexp.MustCompile(`^NV(\d+)$`)
|
||||
nicLegend = regexp.MustCompile(`^\s*(NIC\d+)\s*:\s*(\S+)\s*$`)
|
||||
gpuColumn = regexp.MustCompile(`^GPU(\d+)$`)
|
||||
numaColumn = "NUMA Affinity"
|
||||
)
|
||||
|
||||
// parseTopo reads the matrix. Columns are tab-separated; the header names
|
||||
// each one — GPUn, then the adapters, as NICn with a legend below on newer
|
||||
// drivers or by HCA name on older ones — then CPU and NUMA affinity.
|
||||
func parseTopo(out string) (Topology, error) {
|
||||
out = ansi.ReplaceAllString(out, "")
|
||||
lines := strings.Split(out, "\n")
|
||||
|
||||
legend := map[string]string{}
|
||||
for _, l := range lines {
|
||||
if m := nicLegend.FindStringSubmatch(l); m != nil {
|
||||
legend[m[1]] = m[2]
|
||||
}
|
||||
}
|
||||
|
||||
var header []string
|
||||
start := -1
|
||||
for i, l := range lines {
|
||||
cells := splitTabs(l)
|
||||
if len(cells) > 1 && cells[0] == "" && gpuColumn.MatchString(cells[1]) {
|
||||
header, start = cells, i
|
||||
break
|
||||
}
|
||||
}
|
||||
if start < 0 {
|
||||
return Topology{}, fmt.Errorf("no GPU matrix in nvidia-smi topo output")
|
||||
}
|
||||
|
||||
t := Topology{NVLink: map[int]map[int]int{}, NICPath: map[int]map[string]string{}, NUMA: map[int]int{}}
|
||||
for _, l := range lines[start+1:] {
|
||||
cells := splitTabs(l)
|
||||
if len(cells) < 2 {
|
||||
// The matrix ends at the first blank line, before the legend.
|
||||
if strings.TrimSpace(l) == "" {
|
||||
break
|
||||
}
|
||||
continue
|
||||
}
|
||||
m := gpuColumn.FindStringSubmatch(cells[0])
|
||||
if m == nil {
|
||||
continue // an adapter's row: its view of the GPUs is the same matrix
|
||||
}
|
||||
gpu, _ := strconv.Atoi(m[1])
|
||||
for c := 1; c < len(cells) && c < len(header); c++ {
|
||||
col, v := header[c], strings.TrimSpace(cells[c])
|
||||
switch {
|
||||
case gpuColumn.MatchString(col):
|
||||
peer, _ := strconv.Atoi(gpuColumn.FindStringSubmatch(col)[1])
|
||||
if n := nvLinks.FindStringSubmatch(v); n != nil && peer != gpu {
|
||||
links, _ := strconv.Atoi(n[1])
|
||||
if t.NVLink[gpu] == nil {
|
||||
t.NVLink[gpu] = map[int]int{}
|
||||
}
|
||||
t.NVLink[gpu][peer] = links
|
||||
}
|
||||
case col == numaColumn:
|
||||
if n, err := strconv.Atoi(v); err == nil {
|
||||
t.NUMA[gpu] = n
|
||||
}
|
||||
case col == "CPU Affinity" || col == "GPU NUMA ID" || col == "":
|
||||
default:
|
||||
// An adapter: NICn named by the legend, or an HCA by name.
|
||||
hca := col
|
||||
if name, ok := legend[col]; ok {
|
||||
hca = name
|
||||
}
|
||||
if PathRank(v) < 5 {
|
||||
if t.NICPath[gpu] == nil {
|
||||
t.NICPath[gpu] = map[string]string{}
|
||||
}
|
||||
t.NICPath[gpu][hca] = strings.ToUpper(v)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(t.NVLink) == 0 && len(t.NICPath) == 0 && len(t.NUMA) == 0 {
|
||||
return t, fmt.Errorf("the nvidia-smi topo matrix had no rows")
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
func splitTabs(l string) []string {
|
||||
l = strings.TrimRight(l, "\r")
|
||||
cells := strings.Split(l, "\t")
|
||||
for i := range cells {
|
||||
cells[i] = strings.TrimSpace(cells[i])
|
||||
}
|
||||
return cells
|
||||
}
|
||||
|
||||
// Apply puts the topology on the GPUs it describes. What a GPU already knows
|
||||
// — NVLink peers from NVML, which is the more exact source — is kept.
|
||||
func (t Topology) Apply(gpus []GPU) {
|
||||
for i := range gpus {
|
||||
g := &gpus[i]
|
||||
if len(g.NVLinkPeers) == 0 {
|
||||
for peer := range t.NVLink[g.Index] {
|
||||
g.NVLinkPeers = append(g.NVLinkPeers, peer)
|
||||
}
|
||||
sort.Ints(g.NVLinkPeers)
|
||||
}
|
||||
if n, ok := t.NUMA[g.Index]; ok {
|
||||
g.NUMANode = n
|
||||
}
|
||||
if paths := t.NICPath[g.Index]; len(paths) > 0 {
|
||||
g.NICPaths = paths
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// row joins cells the way nvidia-smi prints them: tab-separated, each padded.
|
||||
func row(cells ...string) string { return strings.Join(cells, "\t") }
|
||||
|
||||
// dgxH100 is `nvidia-smi topo -m` on an eight-GPU NVSwitch machine with two
|
||||
// adapters, as newer drivers print it: NICn columns and a legend naming them.
|
||||
var dgxH100 = strings.Join([]string{
|
||||
row("\x1b[4m", "GPU0", "GPU1", "GPU2", "GPU3", "NIC0", "NIC1", "CPU Affinity", "NUMA Affinity", "GPU NUMA ID\x1b[0m"),
|
||||
row("GPU0", " X ", "NV18", "NV18", "NV18", "PXB", "SYS", "0-55,112-167", "0", "N/A"),
|
||||
row("GPU1", "NV18", " X ", "NV18", "NV18", "PXB", "SYS", "0-55,112-167", "0", "N/A"),
|
||||
row("GPU2", "NV18", "NV18", " X ", "NV18", "SYS", "PXB", "56-111,168-223", "1", "N/A"),
|
||||
row("GPU3", "NV18", "NV18", "NV18", " X ", "SYS", "PIX", "56-111,168-223", "1", "N/A"),
|
||||
row("NIC0", "PXB", "PXB", "SYS", "SYS", " X ", "SYS"),
|
||||
row("NIC1", "SYS", "SYS", "PXB", "PIX", "SYS", " X "),
|
||||
"",
|
||||
"Legend:",
|
||||
"",
|
||||
" X = Self",
|
||||
" SYS = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)",
|
||||
" NV# = Connection traversing a bonded set of # NVLinks",
|
||||
"",
|
||||
"NIC Legend:",
|
||||
"",
|
||||
" NIC0: mlx5_0",
|
||||
" NIC1: mlx5_1",
|
||||
"",
|
||||
}, "\n")
|
||||
|
||||
func TestNVLinkAndAdapterDistanceAreReadFromTopo(t *testing.T) {
|
||||
topo, err := parseTopo(dgxH100)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := topo.NVLink[0]; !reflect.DeepEqual(got, map[int]int{1: 18, 2: 18, 3: 18}) {
|
||||
t.Errorf("GPU0's NVLink peers = %v", got)
|
||||
}
|
||||
// The legend turns NIC1 into the adapter's own name.
|
||||
if got := topo.NICPath[3]; !reflect.DeepEqual(got, map[string]string{"mlx5_0": "SYS", "mlx5_1": "PIX"}) {
|
||||
t.Errorf("GPU3's adapters = %v", got)
|
||||
}
|
||||
if topo.NUMA[2] != 1 || topo.NUMA[0] != 0 {
|
||||
t.Errorf("NUMA = %v", topo.NUMA)
|
||||
}
|
||||
|
||||
gpus := []GPU{{Index: 0}, {Index: 3}}
|
||||
topo.Apply(gpus)
|
||||
if !reflect.DeepEqual(gpus[0].NVLinkPeers, []int{1, 2, 3}) || gpus[1].NUMANode != 1 ||
|
||||
gpus[1].NICPaths["mlx5_1"] != "PIX" {
|
||||
t.Errorf("applied = %+v", gpus)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOlderDriversNameAdaptersInTheHeader(t *testing.T) {
|
||||
// No legend: the columns are the adapters' own names, and NVLink between
|
||||
// pairs only, as on an A100 PCIe box with bridges.
|
||||
out := strings.Join([]string{
|
||||
row("", "GPU0", "GPU1", "GPU2", "GPU3", "mlx5_0", "CPU Affinity", "NUMA Affinity"),
|
||||
row("GPU0", " X ", "NV12", "SYS", "SYS", "NODE", "0-31", "0"),
|
||||
row("GPU1", "NV12", " X ", "SYS", "SYS", "NODE", "0-31", "0"),
|
||||
row("GPU2", "SYS", "SYS", " X ", "NV12", "SYS", "32-63", "1"),
|
||||
row("GPU3", "SYS", "SYS", "NV12", " X ", "SYS", "32-63", "1"),
|
||||
row("mlx5_0", "NODE", "NODE", "SYS", "SYS", " X "),
|
||||
"",
|
||||
}, "\n")
|
||||
topo, err := parseTopo(out)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reflect.DeepEqual(topo.NVLink[2], map[int]int{3: 12}) || len(topo.NVLink[0]) != 1 {
|
||||
t.Errorf("NVLink = %v", topo.NVLink)
|
||||
}
|
||||
if topo.NICPath[0]["mlx5_0"] != "NODE" {
|
||||
t.Errorf("adapters = %v", topo.NICPath)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPCIeMachineHasNoNVLink(t *testing.T) {
|
||||
out := strings.Join([]string{
|
||||
row("", "GPU0", "GPU1", "CPU Affinity", "NUMA Affinity"),
|
||||
row("GPU0", " X ", "PHB", "0-15", "0"),
|
||||
row("GPU1", "PHB", " X ", "0-15", "0"),
|
||||
"",
|
||||
}, "\n")
|
||||
topo, err := parseTopo(out)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(topo.NVLink) != 0 {
|
||||
t.Errorf("NVLink on a PCIe machine: %v", topo.NVLink)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNVMLPeersAreKeptOverTopo(t *testing.T) {
|
||||
topo, _ := parseTopo(dgxH100)
|
||||
gpus := []GPU{{Index: 0, NVLinkPeers: []int{1}}}
|
||||
topo.Apply(gpus)
|
||||
if !reflect.DeepEqual(gpus[0].NVLinkPeers, []int{1}) {
|
||||
t.Errorf("NVML's answer was replaced: %v", gpus[0].NVLinkPeers)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSomethingThatIsNotAMatrixIsAnError(t *testing.T) {
|
||||
if _, err := parseTopo("NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver."); err == nil {
|
||||
t.Error("expected an error")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"io/fs"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// LocalWeights is one model file or directory a node can see.
|
||||
type LocalWeights struct {
|
||||
Path string `json:"path"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
// Shared marks weights on a filesystem this node did not have to fill.
|
||||
//
|
||||
// Worth distinguishing: twenty nodes each holding their own copy and twenty
|
||||
// nodes seeing one filer are the same list of paths and completely
|
||||
// different situations — one has twenty copies of a 141 GB model and the
|
||||
// other has one.
|
||||
Shared bool `json:"shared,omitempty"`
|
||||
}
|
||||
|
||||
// ScanWeights lists the model files under the given directories.
|
||||
//
|
||||
// Storage is the part of this system that will differ at every installation —
|
||||
// a Lustre mount, an image exported over the fabric, local NVMe, an rsync from
|
||||
// a laptop — so llmgw does not try to own it. It observes the result instead:
|
||||
// whatever put the file there, the question placement needs answered is only
|
||||
// whether it is present.
|
||||
//
|
||||
// Scanning nothing is the default and constrains nothing. A node with no
|
||||
// directories configured behaves exactly as every node did before this existed,
|
||||
// which is what makes the whole feature optional.
|
||||
func ScanWeights(dirs []string) []LocalWeights {
|
||||
var out []LocalWeights
|
||||
seen := map[string]struct{}{}
|
||||
|
||||
for _, dir := range dirs {
|
||||
if dir == "" {
|
||||
continue
|
||||
}
|
||||
shared := isNetworkFS(dir)
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
// A directory that is not there is not an error to report upward:
|
||||
// a mount that has not come back yet looks exactly like this, and
|
||||
// the node is simply one without those weights until it does.
|
||||
continue
|
||||
}
|
||||
for _, e := range entries {
|
||||
path := filepath.Join(dir, e.Name())
|
||||
if _, dup := seen[path]; dup {
|
||||
continue
|
||||
}
|
||||
size, ok := weightSize(path, e)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
seen[path] = struct{}{}
|
||||
out = append(out, LocalWeights{Path: path, SizeBytes: size, Shared: shared})
|
||||
}
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path })
|
||||
return out
|
||||
}
|
||||
|
||||
// weightSize measures a candidate and says whether it is one at all.
|
||||
//
|
||||
// A model is either a single file — a .gguf, usually — or a directory of
|
||||
// shards and a config, which is what every HuggingFace checkout looks like.
|
||||
// Anything else in the directory is somebody's notes and is not reported.
|
||||
func weightSize(path string, e fs.DirEntry) (int64, bool) {
|
||||
if e.IsDir() {
|
||||
// A directory counts only if it actually holds weights. An empty one
|
||||
// is a half-finished download, and reporting it as present would have
|
||||
// placement choose a node whose engine then fails to load.
|
||||
var total int64
|
||||
var found bool
|
||||
_ = filepath.WalkDir(path, func(p string, d fs.DirEntry, err error) error {
|
||||
if err != nil || d.IsDir() {
|
||||
return nil //nolint:nilerr // an unreadable subtree is not fatal
|
||||
}
|
||||
if isWeightFile(p) {
|
||||
found = true
|
||||
}
|
||||
if info, err := d.Info(); err == nil {
|
||||
total += info.Size()
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return total, found
|
||||
}
|
||||
if !isWeightFile(path) {
|
||||
return 0, false
|
||||
}
|
||||
info, err := e.Info()
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return info.Size(), true
|
||||
}
|
||||
|
||||
func isWeightFile(path string) bool {
|
||||
switch strings.ToLower(filepath.Ext(path)) {
|
||||
case ".gguf", ".safetensors", ".bin", ".pt", ".pth":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isNetworkFS reports whether a path is on a filesystem the node did not fill.
|
||||
//
|
||||
// By magic number rather than by parsing mounts, because the mount table says
|
||||
// what was asked for and this says what the kernel actually has. The list is
|
||||
// the network filesystems that turn up in front of GPUs: NFS, Lustre, GPFS,
|
||||
// CephFS, GlusterFS and SMB. Anything unrecognised is treated as local, which
|
||||
// is the conservative direction — a shared filer mistaken for local storage
|
||||
// costs an unnecessary copy, while local storage mistaken for shared would have
|
||||
// the cluster believe a model is everywhere when it is on one node.
|
||||
func isNetworkFS(path string) bool {
|
||||
var st syscall.Statfs_t
|
||||
if err := syscall.Statfs(path, &st); err != nil {
|
||||
return false
|
||||
}
|
||||
switch uint32(st.Type) { //nolint:unconvert // Type is int64 on some platforms
|
||||
case 0x6969, // NFS
|
||||
0x0BD00BD0, // Lustre
|
||||
0x47504653, // GPFS
|
||||
0x00C36400, // CephFS
|
||||
0x65735546, // FUSE, which is how GlusterFS and many object mounts appear
|
||||
0xFF534D42: // SMB/CIFS
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func write(t *testing.T, path string, size int) {
|
||||
t.Helper()
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, make([]byte, size), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanFindsFilesAndDirectoriesOfWeights(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
write(t, filepath.Join(dir, "qwen-3b-q4.gguf"), 1024)
|
||||
// A HuggingFace checkout: a directory of shards rather than one file.
|
||||
write(t, filepath.Join(dir, "Llama-3.3-70B", "model-00001.safetensors"), 2048)
|
||||
write(t, filepath.Join(dir, "Llama-3.3-70B", "config.json"), 64)
|
||||
|
||||
got := ScanWeights([]string{dir})
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("found %d entries: %+v", len(got), got)
|
||||
}
|
||||
byName := map[string]int64{}
|
||||
for _, w := range got {
|
||||
byName[filepath.Base(w.Path)] = w.SizeBytes
|
||||
}
|
||||
if byName["qwen-3b-q4.gguf"] != 1024 {
|
||||
t.Errorf("gguf size = %d", byName["qwen-3b-q4.gguf"])
|
||||
}
|
||||
// The directory's size is everything in it, config included: what matters
|
||||
// is how much would have to be moved to put it somewhere else.
|
||||
if byName["Llama-3.3-70B"] != 2048+64 {
|
||||
t.Errorf("checkout size = %d, want %d", byName["Llama-3.3-70B"], 2048+64)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEmptyDirectoryIsNotWeights(t *testing.T) {
|
||||
// A half-finished download looks exactly like this, and reporting it as
|
||||
// present would have placement choose a node whose engine then fails to
|
||||
// load — the failure this whole feature exists to avoid.
|
||||
dir := t.TempDir()
|
||||
if err := os.MkdirAll(filepath.Join(dir, "downloading"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
write(t, filepath.Join(dir, "downloading", "README.md"), 10)
|
||||
|
||||
if got := ScanWeights([]string{dir}); len(got) != 0 {
|
||||
t.Errorf("a directory with no weight files was reported: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotesAndScriptsAreNotReportedAsModels(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
write(t, filepath.Join(dir, "notes.txt"), 10)
|
||||
write(t, filepath.Join(dir, "fetch.sh"), 10)
|
||||
write(t, filepath.Join(dir, "real.gguf"), 10)
|
||||
|
||||
got := ScanWeights([]string{dir})
|
||||
if len(got) != 1 || filepath.Base(got[0].Path) != "real.gguf" {
|
||||
t.Errorf("scan returned %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingDirectoryIsNotAnError(t *testing.T) {
|
||||
// A mount that has not come back yet looks exactly like a directory that is
|
||||
// not there, and the node is simply one without those weights until it is.
|
||||
got := ScanWeights([]string{filepath.Join(t.TempDir(), "never-existed")})
|
||||
if len(got) != 0 {
|
||||
t.Errorf("got %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanningNothingIsTheDefaultAndConstrainsNothing(t *testing.T) {
|
||||
// What makes the feature optional: an agent with no directories configured
|
||||
// behaves exactly as every agent did before this existed.
|
||||
if got := ScanWeights(nil); len(got) != 0 {
|
||||
t.Errorf("got %+v", got)
|
||||
}
|
||||
if got := ScanWeights([]string{""}); len(got) != 0 {
|
||||
t.Errorf("an empty directory name produced %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheSamePathIsReportedOnce(t *testing.T) {
|
||||
// Two configured directories can overlap, or be the same mount reached two
|
||||
// ways. Counting a model twice would overstate what a fabric holds.
|
||||
dir := t.TempDir()
|
||||
write(t, filepath.Join(dir, "m.gguf"), 10)
|
||||
if got := ScanWeights([]string{dir, dir}); len(got) != 1 {
|
||||
t.Errorf("the same directory twice produced %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalStorageIsNeverMistakenForShared(t *testing.T) {
|
||||
// The conservative direction. A shared filer read as local costs an
|
||||
// unnecessary copy; local read as shared would have the cluster believe a
|
||||
// model is everywhere when it is on one node.
|
||||
dir := t.TempDir()
|
||||
write(t, filepath.Join(dir, "m.gguf"), 10)
|
||||
got := ScanWeights([]string{dir})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("got %+v", got)
|
||||
}
|
||||
if got[0].Shared {
|
||||
t.Error("a temporary directory on local disk was reported as shared")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,385 @@
|
||||
// Package engines builds the command lines and environments that start
|
||||
// inference engines (plan §5.3).
|
||||
//
|
||||
// Nothing here starts a process; that is the supervisor's job. Keeping command
|
||||
// construction pure makes it testable without GPUs, which matters because the
|
||||
// environment generation is where multi-node runs go wrong: a single missing
|
||||
// NCCL variable turns a working tensor-parallel group into one that hangs on
|
||||
// the first all-reduce.
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Profile is how a model runs on a GPU class. It mirrors the catalog profile
|
||||
// (plan §6.1), reduced to what launching needs.
|
||||
type Profile struct {
|
||||
Name string
|
||||
Engine string
|
||||
GPUClass string
|
||||
|
||||
// TP and PP are tensor and pipeline parallel degrees.
|
||||
TP int
|
||||
PP int
|
||||
// GPUsPerNode is how many GPUs each shard occupies.
|
||||
GPUsPerNode int
|
||||
// Nodes is how many nodes the replica spans. Above 1, a fabric is required.
|
||||
Nodes int
|
||||
|
||||
// Args are passed through to the engine untouched.
|
||||
Args []string
|
||||
// Env overrides anything the engine adapter generates.
|
||||
Env map[string]string
|
||||
|
||||
ModelPath string
|
||||
// ServedModelName is the name clients use, which need not match the path.
|
||||
ServedModelName string
|
||||
// Adapters are LoRA adapters to load beside the model, in order: an
|
||||
// adapter's index is the id requests apply it by.
|
||||
Adapters []Adapter
|
||||
}
|
||||
|
||||
// ShardSpec is one shard of one replica on this node.
|
||||
type ShardSpec struct {
|
||||
ReplicaID string
|
||||
ShardID string
|
||||
// Rank is the node rank; 0 is the leader and serves HTTP.
|
||||
Rank int
|
||||
NNodes int
|
||||
|
||||
// LeaderAddr and LeaderPort are the distributed init endpoint, always on the
|
||||
// management network rather than the fabric: NCCL bootstraps over TCP and
|
||||
// only then moves to RDMA.
|
||||
LeaderAddr string
|
||||
LeaderPort int
|
||||
|
||||
// GPUIndices are the local GPUs this shard uses.
|
||||
GPUIndices []int
|
||||
// HCAs are the adapters of the chosen fabric, and only that fabric. A shard
|
||||
// must never be given HCAs from two fabrics.
|
||||
HCAs []string
|
||||
// MgmtInterface is the ethernet interface used for bootstrap traffic.
|
||||
MgmtInterface string
|
||||
// GIDIndex is the RoCE v2 GID index for the HCAs, set by the agent when
|
||||
// they are RoCE; nil on InfiniBand, where NCCL needs none.
|
||||
GIDIndex *int
|
||||
|
||||
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"`
|
||||
// Threads is the CPU cores placement reserved for a shard with no GPU,
|
||||
// and so what its engine computes on. Zero leaves it to the engine's own
|
||||
// configuration.
|
||||
Threads int `json:"threads,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
|
||||
// multi-node NVLink domain, which changes NCCL's environment.
|
||||
Pool string
|
||||
Fabric string
|
||||
// ProfileFP identifies the version of the profile this shard was started
|
||||
// from. Carried for the same reason as Pool and Fabric, and never read: it
|
||||
// lets the control plane tell, after its own restart, whether a shard that
|
||||
// is still running matches the profile the catalog holds now.
|
||||
ProfileFP string
|
||||
}
|
||||
|
||||
// Leader reports whether this shard serves HTTP for the replica.
|
||||
func (s ShardSpec) Leader() bool { return s.Rank == 0 }
|
||||
|
||||
// Command is a fully resolved launch: what to exec, with which environment.
|
||||
type Command struct {
|
||||
Path string
|
||||
Args []string
|
||||
// Env is the complete environment, not a delta.
|
||||
Env []string
|
||||
// Name distinguishes the processes of a multi-process shard in logs.
|
||||
Name string
|
||||
// WaitForLeaderPort makes the supervisor hold the next command until
|
||||
// LeaderAddr:LeaderPort accepts a connection.
|
||||
//
|
||||
// It exists for Ray: the head has to be listening before a worker can join,
|
||||
// and both sides express the same wait because both dial the same address.
|
||||
WaitForLeaderPort bool
|
||||
}
|
||||
|
||||
// Engine builds launch commands for one inference engine.
|
||||
type Engine interface {
|
||||
// Name is the engine's identifier in the catalog.
|
||||
Name() string
|
||||
// Validate rejects a profile that cannot work on these GPUs, before
|
||||
// anything is started.
|
||||
Validate(p Profile, gpus []GPUInfo) error
|
||||
|
||||
// Commands returns the processes that make up one shard, in start order.
|
||||
//
|
||||
// Most shards are a single process. vLLM under Ray is not: the leader runs
|
||||
// a Ray head and then the server against it, and a worker runs only the Ray
|
||||
// join. Returning a sequence rather than one command is what lets the
|
||||
// supervisor treat both the same way — every command but the last is a
|
||||
// sidecar kept alive for the shard's lifetime, and the last one's exit is
|
||||
// the shard's exit.
|
||||
Commands(s ShardSpec) ([]Command, error)
|
||||
|
||||
// ServesHTTP reports whether this rank answers the OpenAI API.
|
||||
//
|
||||
// In a multi-node replica only the leader does. Health-checking a worker's
|
||||
// port would time out and fail a replica that is coming up perfectly well,
|
||||
// which is exactly what happened before this existed.
|
||||
ServesHTTP(s ShardSpec) bool
|
||||
|
||||
// HealthPath and MetricsPath are appended to the shard's base URL.
|
||||
HealthPath() string
|
||||
MetricsPath() string
|
||||
}
|
||||
|
||||
// GPUInfo is the subset of GPU facts validation needs.
|
||||
type GPUInfo struct {
|
||||
Index int
|
||||
Name string
|
||||
ComputeCapability string
|
||||
MemoryTotalMB int
|
||||
}
|
||||
|
||||
// Registry holds the known engines.
|
||||
type Registry struct {
|
||||
engines map[string]Engine
|
||||
}
|
||||
|
||||
// NewRegistry returns a registry with the built-in engines.
|
||||
func NewRegistry(cfg Config) *Registry {
|
||||
r := &Registry{engines: map[string]Engine{}}
|
||||
r.Add(&VLLM{cfg: cfg.VLLM})
|
||||
r.Add(&SGLang{cfg: cfg.SGLang})
|
||||
r.Add(&LlamaCpp{cfg: cfg.LlamaCpp})
|
||||
r.Add(&Sim{cfg: cfg.Sim})
|
||||
return r
|
||||
}
|
||||
|
||||
// UseInstalled points an engine at a build the agent installed, and reports
|
||||
// whether the engine takes one.
|
||||
func (r *Registry) UseInstalled(name, path string) bool {
|
||||
u, ok := r.engines[name].(interface{ UseInstalled(string) })
|
||||
if ok {
|
||||
u.UseInstalled(path)
|
||||
}
|
||||
return ok
|
||||
}
|
||||
|
||||
// Add registers an engine, replacing any engine of the same name.
|
||||
func (r *Registry) Add(e Engine) { r.engines[e.Name()] = e }
|
||||
|
||||
// Get returns an engine by name.
|
||||
func (r *Registry) Get(name string) (Engine, error) {
|
||||
e, ok := r.engines[name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown engine %q (have %s)", name, strings.Join(r.Names(), ", "))
|
||||
}
|
||||
return e, nil
|
||||
}
|
||||
|
||||
// Names lists the registered engines.
|
||||
func (r *Registry) Names() []string {
|
||||
out := make([]string, 0, len(r.engines))
|
||||
for n := range r.engines {
|
||||
out = append(out, n)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// Config configures the engine adapters.
|
||||
type Config struct {
|
||||
VLLM VLLMConfig `koanf:"vllm"`
|
||||
SGLang SGLangConfig `koanf:"sglang"`
|
||||
LlamaCpp LlamaCppConfig `koanf:"llamacpp"`
|
||||
Sim SimConfig `koanf:"sim"`
|
||||
}
|
||||
|
||||
// --- shared helpers ---
|
||||
|
||||
// NCCLEnv is the environment a shard with this spec is given, for a tool that
|
||||
// shows an operator what an engine would see.
|
||||
func NCCLEnv(s ShardSpec) map[string]string { return baseEnv(s) }
|
||||
|
||||
// baseEnv builds the environment every engine shares.
|
||||
//
|
||||
// This is the part of §5.3 the plan marks critical. The rules encoded here:
|
||||
//
|
||||
// - CUDA_VISIBLE_DEVICES pins the shard to its GPUs, so two shards on one
|
||||
// node cannot fight over the same card.
|
||||
// - NCCL_IB_HCA lists only the chosen fabric's adapters. Leaving it unset
|
||||
// lets NCCL pick adapters from another fabric, where the peer is
|
||||
// unreachable, and the job hangs instead of failing.
|
||||
// - NCCL_SOCKET_IFNAME pins bootstrap to the management interface. NCCL
|
||||
// otherwise probes every interface, including IB-over-IP addresses that
|
||||
// cannot route between fabrics.
|
||||
// - NCCL_IB_DISABLE is set when there is no fabric, so NCCL stops looking for
|
||||
// one instead of timing out.
|
||||
func baseEnv(s ShardSpec) map[string]string {
|
||||
env := map[string]string{
|
||||
"CUDA_VISIBLE_DEVICES": joinInts(s.GPUIndices),
|
||||
}
|
||||
|
||||
if len(s.HCAs) > 0 {
|
||||
env["NCCL_IB_HCA"] = strings.Join(s.HCAs, ",")
|
||||
env["NCCL_IB_DISABLE"] = "0"
|
||||
// RoCE: the index of the routable RoCE v2 entry. NCCL's default is a
|
||||
// v1 link-local one on most hosts, and a job on it hangs.
|
||||
if s.GIDIndex != nil {
|
||||
env["NCCL_IB_GID_INDEX"] = strconv.Itoa(*s.GIDIndex)
|
||||
}
|
||||
} else {
|
||||
// Single node, or a node with no fabric. Either way IB is not the path.
|
||||
env["NCCL_IB_DISABLE"] = "1"
|
||||
}
|
||||
|
||||
// Placed inside a multi-node NVLink domain: NCCL is to carry the
|
||||
// replica's collectives over NVLink between its nodes, which it only
|
||||
// does when told to, and needs IMEX for.
|
||||
if s.NNodes > 1 && strings.HasPrefix(s.Fabric, "nvl:") {
|
||||
env["NCCL_MNNVL_ENABLE"] = "1"
|
||||
env["NCCL_CUMEM_ENABLE"] = "1"
|
||||
}
|
||||
|
||||
if s.MgmtInterface != "" {
|
||||
env["NCCL_SOCKET_IFNAME"] = s.MgmtInterface
|
||||
env["GLOO_SOCKET_IFNAME"] = s.MgmtInterface
|
||||
}
|
||||
|
||||
// Profile overrides win: they are the escape hatch for the per-cluster
|
||||
// tuning no adapter can guess.
|
||||
for k, v := range s.Profile.Env {
|
||||
// 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
|
||||
}
|
||||
|
||||
// materializeEnv merges generated variables over the current process
|
||||
// environment and returns it in exec form.
|
||||
//
|
||||
// Inheriting matters: engines need PATH, HOME, HF_HOME, LD_LIBRARY_PATH and
|
||||
// whatever else the operator set up in the unit file.
|
||||
func materializeEnv(generated map[string]string) []string {
|
||||
merged := map[string]string{}
|
||||
for _, kv := range os.Environ() {
|
||||
if k, v, ok := strings.Cut(kv, "="); ok {
|
||||
merged[k] = v
|
||||
}
|
||||
}
|
||||
for k, v := range generated {
|
||||
merged[k] = v
|
||||
}
|
||||
|
||||
out := make([]string, 0, len(merged))
|
||||
for k, v := range merged {
|
||||
out = append(out, k+"="+v)
|
||||
}
|
||||
// Sorted so a command is reproducible and diffable in logs.
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
func joinInts(v []int) string {
|
||||
parts := make([]string, len(v))
|
||||
for i, n := range v {
|
||||
parts[i] = strconv.Itoa(n)
|
||||
}
|
||||
return strings.Join(parts, ",")
|
||||
}
|
||||
|
||||
// validateCommon holds the checks that apply to every engine.
|
||||
func validateCommon(p Profile, gpus []GPUInfo) error {
|
||||
if p.ModelPath == "" {
|
||||
return fmt.Errorf("profile %q: model path is empty", p.Name)
|
||||
}
|
||||
if p.TP <= 0 {
|
||||
return fmt.Errorf("profile %q: tensor parallel degree must be positive", p.Name)
|
||||
}
|
||||
if p.Nodes > 1 && p.GPUsPerNode <= 0 {
|
||||
return fmt.Errorf("profile %q: gpus_per_node is required for a %d-node profile", p.Name, p.Nodes)
|
||||
}
|
||||
|
||||
perNode := p.GPUsPerNode
|
||||
if perNode == 0 {
|
||||
perNode = p.TP
|
||||
}
|
||||
nodes := max(p.Nodes, 1)
|
||||
if want := perNode * nodes; want != p.TP*max(p.PP, 1) {
|
||||
return fmt.Errorf("profile %q: %d nodes x %d GPUs does not match TP %d x PP %d",
|
||||
p.Name, nodes, perNode, p.TP, max(p.PP, 1))
|
||||
}
|
||||
if len(gpus) > 0 && len(gpus) < perNode {
|
||||
return fmt.Errorf("profile %q: needs %d GPUs per node, node has %d", p.Name, perNode, len(gpus))
|
||||
}
|
||||
|
||||
// A tensor-parallel group must be uniform: mixing classes makes the slowest
|
||||
// card set the pace for every collective, and mismatched memory makes the
|
||||
// smallest card decide what fits.
|
||||
if len(gpus) > 1 {
|
||||
first := gpus[0].Name
|
||||
for _, g := range gpus[1:] {
|
||||
if g.Name != first {
|
||||
return fmt.Errorf("profile %q: GPUs are not uniform (%q and %q); "+
|
||||
"a tensor-parallel group may not mix GPU types", p.Name, first, g.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// requireComputeCapability rejects profiles whose quantization the hardware
|
||||
// cannot execute.
|
||||
func requireComputeCapability(p Profile, gpus []GPUInfo, minCap float64, why string) error {
|
||||
for _, g := range gpus {
|
||||
cap, err := strconv.ParseFloat(g.ComputeCapability, 64)
|
||||
if err != nil {
|
||||
// Unknown capability: let it through rather than block a launch on
|
||||
// a parsing problem. The engine will refuse it clearly enough.
|
||||
continue
|
||||
}
|
||||
if cap < minCap {
|
||||
return fmt.Errorf("profile %q: %s needs compute capability %.1f, GPU %d (%s) has %.1f",
|
||||
p.Name, why, minCap, g.Index, g.Name, cap)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// quantizationOf extracts the --quantization argument, if present.
|
||||
func quantizationOf(args []string) string {
|
||||
for i, a := range args {
|
||||
if a == "--quantization" && i+1 < len(args) {
|
||||
return args[i+1]
|
||||
}
|
||||
if v, ok := strings.CutPrefix(a, "--quantization="); ok {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// Adapter is a LoRA adapter an engine loads beside its model. Path is a
|
||||
// source until the agent has fetched it, a local file after.
|
||||
type Adapter struct {
|
||||
Name string `json:"name"`
|
||||
Path string `json:"path"`
|
||||
}
|
||||
@@ -0,0 +1,478 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// buildOne returns the single command a shard produces, failing if it produced
|
||||
// a different number.
|
||||
func buildOne(t *testing.T, e Engine, s ShardSpec) Command {
|
||||
t.Helper()
|
||||
cmds := build(t, e, s)
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("expected 1 command, got %d: %s", len(cmds), summarize(cmds))
|
||||
}
|
||||
return cmds[0]
|
||||
}
|
||||
|
||||
// buildEngine returns the last command, which is the engine itself. The ones
|
||||
// before it are sidecars.
|
||||
func buildEngine(t *testing.T, e Engine, s ShardSpec) Command {
|
||||
t.Helper()
|
||||
cmds := build(t, e, s)
|
||||
return cmds[len(cmds)-1]
|
||||
}
|
||||
|
||||
func build(t *testing.T, e Engine, s ShardSpec) []Command {
|
||||
t.Helper()
|
||||
cmds, err := e.Commands(s)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) == 0 {
|
||||
t.Fatal("no commands produced")
|
||||
}
|
||||
return cmds
|
||||
}
|
||||
|
||||
// summarize keeps failure output readable: a Command carries the whole
|
||||
// environment, and printing it buries the actual problem.
|
||||
func summarize(cmds []Command) string {
|
||||
var b strings.Builder
|
||||
for i, c := range cmds {
|
||||
if i > 0 {
|
||||
b.WriteString(" | ")
|
||||
}
|
||||
fmt.Fprintf(&b, "%s: %s %s", c.Name, c.Path, strings.Join(c.Args, " "))
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// envMap turns exec-form environment back into a map for assertions.
|
||||
func envMap(env []string) map[string]string {
|
||||
m := map[string]string{}
|
||||
for _, kv := range env {
|
||||
if k, v, ok := strings.Cut(kv, "="); ok {
|
||||
m[k] = v
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// argValue returns the value following flag, or "" if the flag is absent.
|
||||
func argValue(args []string, flag string) string {
|
||||
i := slices.Index(args, flag)
|
||||
if i < 0 || i+1 >= len(args) {
|
||||
return ""
|
||||
}
|
||||
return args[i+1]
|
||||
}
|
||||
|
||||
func singleNodeSpec() ShardSpec {
|
||||
return ShardSpec{
|
||||
ReplicaID: "r1",
|
||||
ShardID: "r1-0",
|
||||
Rank: 0,
|
||||
NNodes: 1,
|
||||
GPUIndices: []int{0, 1, 2, 3},
|
||||
MgmtInterface: "eth0",
|
||||
HTTPPort: 20000,
|
||||
Profile: Profile{
|
||||
Name: "h100-tp4",
|
||||
Engine: "vllm",
|
||||
TP: 4,
|
||||
ModelPath: "/models/llama",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func multiNodeSpec(rank int) ShardSpec {
|
||||
s := singleNodeSpec()
|
||||
s.Rank = rank
|
||||
s.NNodes = 2
|
||||
s.GPUIndices = []int{0, 1, 2, 3, 4, 5, 6, 7}
|
||||
s.HCAs = []string{"mlx5_0", "mlx5_1"}
|
||||
s.LeaderAddr = "10.0.0.10"
|
||||
s.LeaderPort = 20001
|
||||
s.Profile.TP = 16
|
||||
s.Profile.GPUsPerNode = 8
|
||||
s.Profile.Nodes = 2
|
||||
return s
|
||||
}
|
||||
|
||||
// --- environment generation: the part that breaks multi-node runs ---
|
||||
|
||||
func TestSingleNodeDisablesIB(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
if env["NCCL_IB_DISABLE"] != "1" {
|
||||
t.Error("NCCL_IB_DISABLE must be 1 without a fabric, or NCCL waits for one that is not there")
|
||||
}
|
||||
if _, set := env["NCCL_IB_HCA"]; set {
|
||||
t.Error("NCCL_IB_HCA was set without any HCAs")
|
||||
}
|
||||
if env["CUDA_VISIBLE_DEVICES"] != "0,1,2,3" {
|
||||
t.Errorf("CUDA_VISIBLE_DEVICES=%q", env["CUDA_VISIBLE_DEVICES"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiNodePinsFabric(t *testing.T) {
|
||||
cmd := buildEngine(t, &VLLM{}, multiNodeSpec(0))
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
// Only the chosen fabric's adapters. Leaving this open lets NCCL pick an
|
||||
// adapter on another fabric, where the peer is unreachable, and the job
|
||||
// hangs rather than failing.
|
||||
if env["NCCL_IB_HCA"] != "mlx5_0,mlx5_1" {
|
||||
t.Errorf("NCCL_IB_HCA=%q", env["NCCL_IB_HCA"])
|
||||
}
|
||||
if env["NCCL_IB_DISABLE"] != "0" {
|
||||
t.Errorf("NCCL_IB_DISABLE=%q with a fabric present", env["NCCL_IB_DISABLE"])
|
||||
}
|
||||
// Bootstrap must stay on the management network; NCCL otherwise probes
|
||||
// every interface, including IB-over-IP that cannot route between fabrics.
|
||||
if env["NCCL_SOCKET_IFNAME"] != "eth0" || env["GLOO_SOCKET_IFNAME"] != "eth0" {
|
||||
t.Errorf("socket interfaces not pinned: %q / %q",
|
||||
env["NCCL_SOCKET_IFNAME"], env["GLOO_SOCKET_IFNAME"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileEnvOverridesGenerated(t *testing.T) {
|
||||
s := multiNodeSpec(0)
|
||||
s.Profile.Env = map[string]string{
|
||||
"NCCL_IB_HCA": "mlx5_9",
|
||||
"NCCL_NET_GDR_LEVEL": "SYS",
|
||||
}
|
||||
cmd := buildEngine(t, &VLLM{}, s)
|
||||
env := envMap(cmd.Env)
|
||||
|
||||
if env["NCCL_IB_HCA"] != "mlx5_9" {
|
||||
t.Error("a profile override did not win; it is the escape hatch for per-cluster tuning")
|
||||
}
|
||||
if env["NCCL_NET_GDR_LEVEL"] != "SYS" {
|
||||
t.Error("an extra profile variable was dropped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessEnvironmentIsInherited(t *testing.T) {
|
||||
// Engines need PATH, HF_HOME, LD_LIBRARY_PATH and whatever else the unit
|
||||
// file set up.
|
||||
t.Setenv("LLMGW_TEST_MARKER", "present")
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
if envMap(cmd.Env)["LLMGW_TEST_MARKER"] != "present" {
|
||||
t.Error("the process environment was not inherited")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnvIsSorted(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{}, singleNodeSpec())
|
||||
if !slices.IsSorted(cmd.Env) {
|
||||
t.Error("environment is not sorted; launches would not be reproducible in logs")
|
||||
}
|
||||
}
|
||||
|
||||
// --- vLLM ---
|
||||
|
||||
func TestVLLMSingleNodeArgs(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{cfg: VLLMConfig{Bin: "/opt/vllm/bin/vllm"}}, singleNodeSpec())
|
||||
if cmd.Path != "/opt/vllm/bin/vllm" {
|
||||
t.Errorf("Path=%q", cmd.Path)
|
||||
}
|
||||
if cmd.Args[0] != "serve" || cmd.Args[1] != "/models/llama" {
|
||||
t.Errorf("args start wrong: %v", cmd.Args[:2])
|
||||
}
|
||||
if argValue(cmd.Args, "--tensor-parallel-size") != "4" {
|
||||
t.Errorf("TP not passed: %v", cmd.Args)
|
||||
}
|
||||
if argValue(cmd.Args, "--port") != "20000" {
|
||||
t.Errorf("port not passed: %v", cmd.Args)
|
||||
}
|
||||
// Single node means no distributed backend flag at all.
|
||||
if slices.Contains(cmd.Args, "--distributed-executor-backend") {
|
||||
t.Error("a distributed backend was selected for a single-node profile")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMRayLeaderRunsTwoProcesses(t *testing.T) {
|
||||
// The leader starts a Ray head and then the server against it. A single
|
||||
// command would mean running vLLM against a cluster that does not exist.
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
cmds, err := v.Commands(multiNodeSpec(0))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 2 {
|
||||
t.Fatalf("expected 2 commands for the Ray leader, got %d", len(cmds))
|
||||
}
|
||||
if cmds[0].Name != "ray" || !slices.Contains(cmds[0].Args, "--head") {
|
||||
t.Errorf("first command is not the Ray head: %+v", cmds[0])
|
||||
}
|
||||
if !cmds[0].WaitForLeaderPort {
|
||||
t.Error("the server would start before the head is listening")
|
||||
}
|
||||
if cmds[1].Name != "vllm" || argValue(cmds[1].Args, "--distributed-executor-backend") != "ray" {
|
||||
t.Errorf("second command is not the vLLM server: %+v", cmds[1])
|
||||
}
|
||||
if got := envMap(cmds[1].Env)["RAY_ADDRESS"]; got != "10.0.0.10:20001" {
|
||||
t.Errorf("RAY_ADDRESS=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMRayWorkerRunsOnlyRay(t *testing.T) {
|
||||
// A Ray worker runs no vLLM process: the cluster hands its GPUs to the
|
||||
// leader. Starting a second server there would bind a port nobody uses and
|
||||
// compete for the same cards.
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
cmds, err := v.Commands(multiNodeSpec(1))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("expected 1 command for a Ray worker, got %d: %+v", len(cmds), cmds)
|
||||
}
|
||||
if cmds[0].Name != "ray" || slices.Contains(cmds[0].Args, "--head") {
|
||||
t.Errorf("worker command wrong: %+v", cmds[0])
|
||||
}
|
||||
if argValue(cmds[0].Args, "--address") != "10.0.0.10:20001" {
|
||||
t.Errorf("worker is not joining the leader: %v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMMPMode(t *testing.T) {
|
||||
// vLLM's multi-node launch: every node the same serve with its rank and
|
||||
// the leader's address, workers headless. Workers used to be given the
|
||||
// leader's address as their own VLLM_HOST_IP, and no rank at all.
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}
|
||||
for _, rank := range []int{0, 1} {
|
||||
cmd := buildOne(t, v, multiNodeSpec(rank))
|
||||
for flag, want := range map[string]string{
|
||||
"--distributed-executor-backend": "mp", "--nnodes": "2",
|
||||
"--node-rank": strconv.Itoa(rank), "--master-addr": "10.0.0.10", "--master-port": "20001",
|
||||
} {
|
||||
if got := argValue(cmd.Args, flag); got != want {
|
||||
t.Errorf("rank %d: %s = %q, want %q (%v)", rank, flag, got, want, cmd.Args)
|
||||
}
|
||||
}
|
||||
if headless := slices.Contains(cmd.Args, "--headless"); headless != (rank > 0) {
|
||||
t.Errorf("rank %d: headless = %v", rank, headless)
|
||||
}
|
||||
if ip, ok := envMap(cmd.Env)["VLLM_HOST_IP"]; ok && ip == "10.0.0.10" && rank > 0 {
|
||||
t.Errorf("a worker was told the leader's address is its own")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoCEGetsItsGIDIndex(t *testing.T) {
|
||||
s := multiNodeSpec(0)
|
||||
idx := 3
|
||||
s.GIDIndex = &idx
|
||||
env := envMap(buildOne(t, &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}, s).Env)
|
||||
if env["NCCL_IB_GID_INDEX"] != "3" {
|
||||
t.Errorf("NCCL_IB_GID_INDEX = %q", env["NCCL_IB_GID_INDEX"])
|
||||
}
|
||||
// InfiniBand needs none, and must not be given one.
|
||||
env = envMap(buildOne(t, &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}, multiNodeSpec(0)).Env)
|
||||
if _, ok := env["NCCL_IB_GID_INDEX"]; ok {
|
||||
t.Error("an InfiniBand shard was given a GID index")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMUnknownLaunchMode(t *testing.T) {
|
||||
v := &VLLM{cfg: VLLMConfig{LaunchMode: "sorcery"}}
|
||||
if _, err := v.Commands(multiNodeSpec(0)); err == nil {
|
||||
t.Error("an unknown launch mode was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyLeaderServesHTTP(t *testing.T) {
|
||||
// The fix for the bug that stopped multi-node replicas coming up at all:
|
||||
// health-checking a worker's port times out and fails a healthy replica.
|
||||
for _, e := range []Engine{&VLLM{}, &SGLang{}} {
|
||||
if !e.ServesHTTP(multiNodeSpec(0)) {
|
||||
t.Errorf("%s: the leader must serve HTTP", e.Name())
|
||||
}
|
||||
if e.ServesHTTP(multiNodeSpec(1)) {
|
||||
t.Errorf("%s: a worker must not be expected to serve HTTP", e.Name())
|
||||
}
|
||||
if !e.ServesHTTP(singleNodeSpec()) {
|
||||
t.Errorf("%s: a single-node shard must serve HTTP", e.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoRayForSingleNodeOrMPMode(t *testing.T) {
|
||||
ray := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
|
||||
if cmds := buildOne(t, ray, singleNodeSpec()); cmds.Name != "vllm" {
|
||||
t.Errorf("a single-node profile was given a Ray process: %+v", cmds)
|
||||
}
|
||||
mp := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}
|
||||
if cmds := buildOne(t, mp, multiNodeSpec(0)); cmds.Name != "vllm" {
|
||||
t.Errorf("mp mode was given a Ray process: %+v", cmds)
|
||||
}
|
||||
}
|
||||
|
||||
// --- SGLang ---
|
||||
|
||||
func TestSGLangMultiNodeArgs(t *testing.T) {
|
||||
cmd := buildOne(t, &SGLang{cfg: SGLangConfig{Python: "/opt/sglang/bin/python"}}, multiNodeSpec(1))
|
||||
if cmd.Path != "/opt/sglang/bin/python" {
|
||||
t.Errorf("Path=%q", cmd.Path)
|
||||
}
|
||||
if cmd.Args[0] != "-m" || cmd.Args[1] != "sglang.launch_server" {
|
||||
t.Errorf("not launched as a module: %v", cmd.Args[:2])
|
||||
}
|
||||
if argValue(cmd.Args, "--nnodes") != "2" || argValue(cmd.Args, "--node-rank") != "1" {
|
||||
t.Errorf("multi-node flags wrong: %v", cmd.Args)
|
||||
}
|
||||
if argValue(cmd.Args, "--dist-init-addr") != "10.0.0.10:20001" {
|
||||
t.Errorf("dist-init-addr wrong: %v", cmd.Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSGLangSingleNodeOmitsMultiNodeFlags(t *testing.T) {
|
||||
cmd := buildOne(t, &SGLang{}, singleNodeSpec())
|
||||
for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} {
|
||||
if slices.Contains(cmd.Args, flag) {
|
||||
t.Errorf("%s was passed for a single-node profile", flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileArgsComeLast(t *testing.T) {
|
||||
// Operators must be able to override anything the adapter generated.
|
||||
s := singleNodeSpec()
|
||||
s.Profile.Args = []string{"--max-model-len", "32768"}
|
||||
cmd := buildEngine(t, &VLLM{}, s)
|
||||
|
||||
i := slices.Index(cmd.Args, "--max-model-len")
|
||||
j := slices.Index(cmd.Args, "--tensor-parallel-size")
|
||||
if i < j {
|
||||
t.Errorf("profile args came before generated ones: %v", cmd.Args)
|
||||
}
|
||||
}
|
||||
|
||||
// --- validation ---
|
||||
|
||||
func h100(n int) []GPUInfo {
|
||||
out := make([]GPUInfo, n)
|
||||
for i := range out {
|
||||
out[i] = GPUInfo{Index: i, Name: "NVIDIA H100 80GB HBM3", ComputeCapability: "9.0", MemoryTotalMB: 81559}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestValidateAcceptsMatchingProfile(t *testing.T) {
|
||||
p := Profile{Name: "ok", TP: 4, ModelPath: "/m"}
|
||||
if err := (&VLLM{}).Validate(p, h100(4)); err != nil {
|
||||
t.Errorf("a valid profile was rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsMixedGPUTypes(t *testing.T) {
|
||||
// The invariant from plan §1.1: a TP group never mixes GPU types.
|
||||
gpus := h100(2)
|
||||
gpus[1].Name = "NVIDIA A100-SXM4-40GB"
|
||||
|
||||
err := (&VLLM{}).Validate(Profile{Name: "mixed", TP: 2, ModelPath: "/m"}, gpus)
|
||||
if err == nil || !strings.Contains(err.Error(), "may not mix GPU types") {
|
||||
t.Errorf("mixed GPU types accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsFP8OnOldHardware(t *testing.T) {
|
||||
gpus := h100(4)
|
||||
for i := range gpus {
|
||||
gpus[i].Name = "NVIDIA A100-SXM4-80GB"
|
||||
gpus[i].ComputeCapability = "8.0"
|
||||
}
|
||||
p := Profile{Name: "fp8", TP: 4, ModelPath: "/m", Args: []string{"--quantization", "fp8"}}
|
||||
|
||||
err := (&VLLM{}).Validate(p, gpus)
|
||||
if err == nil || !strings.Contains(err.Error(), "compute capability") {
|
||||
t.Errorf("fp8 accepted on compute capability 8.0: %v", err)
|
||||
}
|
||||
|
||||
// The same profile is fine on Hopper.
|
||||
if err := (&VLLM{}).Validate(p, h100(4)); err != nil {
|
||||
t.Errorf("fp8 rejected on H100: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsInconsistentParallelism(t *testing.T) {
|
||||
// 2 nodes x 8 GPUs cannot serve TP 4.
|
||||
p := Profile{Name: "bad", TP: 4, Nodes: 2, GPUsPerNode: 8, ModelPath: "/m"}
|
||||
if err := (&VLLM{}).Validate(p, h100(8)); err == nil {
|
||||
t.Error("inconsistent node/GPU/TP arithmetic accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRejectsTooFewGPUs(t *testing.T) {
|
||||
if err := (&VLLM{}).Validate(Profile{Name: "big", TP: 8, ModelPath: "/m"}, h100(4)); err == nil {
|
||||
t.Error("a TP 8 profile was accepted on a 4-GPU node")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRequiresModelPath(t *testing.T) {
|
||||
if err := (&VLLM{}).Validate(Profile{Name: "nopath", TP: 1}, h100(1)); err == nil {
|
||||
t.Error("a profile without a model path was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry(t *testing.T) {
|
||||
r := NewRegistry(Config{})
|
||||
if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "sim", "vllm"}) {
|
||||
t.Errorf("registry contents: %v", got)
|
||||
}
|
||||
if _, err := r.Get("vllm"); err != nil {
|
||||
t.Errorf("vllm missing: %v", err)
|
||||
}
|
||||
err := func() error { _, err := r.Get("nope"); return err }()
|
||||
if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, vllm") {
|
||||
t.Errorf("an unknown engine should list what is available: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuantizationOf(t *testing.T) {
|
||||
for _, args := range [][]string{
|
||||
{"--quantization", "fp8"},
|
||||
{"--quantization=fp8"},
|
||||
{"--max-model-len", "100", "--quantization", "fp8"},
|
||||
} {
|
||||
if got := quantizationOf(args); got != "fp8" {
|
||||
t.Errorf("quantizationOf(%v)=%q", args, got)
|
||||
}
|
||||
}
|
||||
if got := quantizationOf([]string{"--quantization"}); got != "" {
|
||||
t.Errorf("a dangling flag should yield nothing, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAReplicaInAnNVLinkDomainIsToldToUseIt(t *testing.T) {
|
||||
s := multiNodeSpec(0)
|
||||
s.Fabric = "nvl:76b1ad76-7c24-4b6b-9f6c-2a1e0d8c11aa/7"
|
||||
env := envMap(buildOne(t, &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}, s).Env)
|
||||
if env["NCCL_MNNVL_ENABLE"] != "1" {
|
||||
t.Errorf("NCCL_MNNVL_ENABLE = %q", env["NCCL_MNNVL_ENABLE"])
|
||||
}
|
||||
if env := envMap(buildOne(t, &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}, multiNodeSpec(0)).Env); env["NCCL_MNNVL_ENABLE"] != "" {
|
||||
t.Error("a replica on InfiniBand was told to use multi-node NVLink")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiNodeVLLMDefaultsToItsOwnLaunch(t *testing.T) {
|
||||
// mp is what vLLM documents and needs nothing but vLLM; Ray was the
|
||||
// default and was never checked against a real cluster.
|
||||
cmds, err := (&VLLM{}).Commands(multiNodeSpec(1))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 1 || cmds[0].Name != "vllm" || argValue(cmds[0].Args, "--distributed-executor-backend") != "mp" ||
|
||||
!slices.Contains(cmds[0].Args, "--headless") {
|
||||
t.Errorf("default worker = %+v", cmds)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// LlamaCppConfig locates the llama.cpp installation.
|
||||
type LlamaCppConfig struct {
|
||||
// Bin is the llama-server binary. It is the server, not the CLI: the
|
||||
// gateway talks OpenAI over HTTP and needs the one that serves it.
|
||||
// Empty means the agent installs llama.cpp itself, at the version the
|
||||
// control plane asks for. Setting it opts out of that.
|
||||
Bin string `koanf:"bin"`
|
||||
// Parallel is how many requests a server handles at once, when a profile
|
||||
// does not say. llama.cpp defaults to one, which would make every second
|
||||
// request wait behind the first no matter what the router decided.
|
||||
Parallel int `koanf:"parallel"`
|
||||
// Threads is how many CPU threads a shard with no GPU computes on. Zero
|
||||
// takes half this machine's CPUs.
|
||||
//
|
||||
// Not all of them, which is llama.cpp's own default: two busy servers
|
||||
// that each spin on every CPU fight for them, and on an 8-vCPU node
|
||||
// that made 7B generation fall from 10 tokens a second to under 1 each.
|
||||
// At half, one server alone is about a third slower and two together
|
||||
// keep 5.5 each. A node that only ever runs one CPU replica does better
|
||||
// with all of them; say so here.
|
||||
Threads int `koanf:"threads"`
|
||||
}
|
||||
|
||||
// numCPU is the CPUs this process may run on; a variable for tests.
|
||||
var numCPU = runtime.NumCPU
|
||||
|
||||
// LlamaCpp builds llama.cpp launch commands.
|
||||
//
|
||||
// It is the engine that makes a cluster without GPUs useful: llama.cpp runs a
|
||||
// small model on CPU at a speed that is fine for testing the whole path, and it
|
||||
// serves the same OpenAI API as the others, so nothing downstream can tell the
|
||||
// difference.
|
||||
type LlamaCpp struct {
|
||||
cfg LlamaCppConfig
|
||||
installed atomic.Pointer[string]
|
||||
}
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*LlamaCpp) Name() string { return "llamacpp" }
|
||||
|
||||
// HealthPath is llama.cpp's readiness endpoint. It returns 503 while the model
|
||||
// is still loading, which is exactly what the supervisor needs.
|
||||
func (*LlamaCpp) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is llama.cpp's Prometheus endpoint, which only exists when the
|
||||
// server was started with --metrics. Commands always passes it.
|
||||
func (*LlamaCpp) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API. Every llama.cpp shard
|
||||
// does, because a llama.cpp replica is always exactly one shard.
|
||||
func (*LlamaCpp) ServesHTTP(ShardSpec) bool { return true }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
//
|
||||
// The two refusals here are about not pretending: llama.cpp has no equivalent
|
||||
// of a tensor-parallel group spanning machines, and its multi-GPU support is a
|
||||
// different mechanism (--split-mode, --tensor-split) that this adapter does not
|
||||
// drive. Accepting either would produce a replica that starts and then serves
|
||||
// from one machine while the catalog believes it has several.
|
||||
func (*LlamaCpp) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if err := validateCommon(p, gpus); err != nil {
|
||||
return err
|
||||
}
|
||||
if p.Nodes > 1 {
|
||||
return fmt.Errorf("profile %q: llamacpp does not span nodes; use vllm or sglang for a %d-node replica",
|
||||
p.Name, p.Nodes)
|
||||
}
|
||||
if p.TP > 1 {
|
||||
return fmt.Errorf("profile %q: llamacpp does not do tensor parallelism; "+
|
||||
"it splits a model across local GPUs with --tensor-split, which this adapter does not drive",
|
||||
p.Name)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Commands returns the processes for one shard, which is always one process.
|
||||
func (l *LlamaCpp) Commands(s ShardSpec) ([]Command, error) {
|
||||
bin := l.bin()
|
||||
p := s.Profile
|
||||
|
||||
model, err := ggufIn(p.ModelPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
args := []string{
|
||||
"--model", model,
|
||||
"--host", "0.0.0.0",
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
// Without this there is no /metrics, and an endpoint the router cannot
|
||||
// measure is one it has to route to blind.
|
||||
"--metrics",
|
||||
}
|
||||
// The model's own chat template rather than a generic one, which is what
|
||||
// renders tool definitions in the form the model was trained to read.
|
||||
//
|
||||
// Pinned rather than relied on: build 10991 has it on by default, older
|
||||
// builds do not, and an engine upgrade should not be able to change how
|
||||
// prompts are built underneath a running fleet. Passing it explicitly was
|
||||
// a no-op on the build here — it was added while chasing tool calls that
|
||||
// came back as text, and did not fix that, which turned out to be the
|
||||
// model rather than the template.
|
||||
if !hasFlag(p.Args, "--jinja", "--no-jinja") {
|
||||
args = append(args, "--jinja")
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--alias", p.ServedModelName)
|
||||
}
|
||||
if !hasFlag(p.Args, "--parallel", "-np") {
|
||||
n := l.cfg.Parallel
|
||||
if n <= 0 {
|
||||
n = 4
|
||||
}
|
||||
args = append(args, "--parallel", strconv.Itoa(n))
|
||||
}
|
||||
// With GPUs assigned, offload everything; llama.cpp keeps the model on the
|
||||
// CPU otherwise, which is not what a GPU placement asked for. Without them
|
||||
// the flag is omitted entirely rather than set to zero, so a build with no
|
||||
// GPU support does not see an argument it rejects.
|
||||
if len(s.GPUIndices) > 0 && !hasFlag(p.Args, "--n-gpu-layers", "-ngl") {
|
||||
args = append(args, "--n-gpu-layers", "999")
|
||||
}
|
||||
if len(s.GPUIndices) == 0 && !hasFlag(p.Args, "--threads", "-t") {
|
||||
// What placement reserved first: it is what the node's other
|
||||
// replicas were placed around.
|
||||
n := s.Threads
|
||||
if n <= 0 {
|
||||
n = l.cfg.Threads
|
||||
}
|
||||
if n <= 0 {
|
||||
n = max(numCPU()/2, 1)
|
||||
}
|
||||
args = append(args, "--threads", strconv.Itoa(n))
|
||||
}
|
||||
// Loaded at start, applied per request: a request that names an adapter
|
||||
// carries "lora": [{"id": i, "scale": s}], and one that names none is the
|
||||
// base model, since every adapter starts at scale zero.
|
||||
if len(p.Adapters) > 0 {
|
||||
args = append(args, "--lora-init-without-apply")
|
||||
for _, a := range p.Adapters {
|
||||
args = append(args, "--lora", a.Path)
|
||||
}
|
||||
}
|
||||
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(env),
|
||||
}}, nil
|
||||
}
|
||||
|
||||
// hasFlag reports whether the profile already sets one of these flags, in
|
||||
// either "--flag value" or "--flag=value" form.
|
||||
func hasFlag(args []string, names ...string) bool {
|
||||
for _, a := range args {
|
||||
for _, n := range names {
|
||||
if a == n || strings.HasPrefix(a, n+"=") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ggufIn resolves a model path that may be a directory.
|
||||
//
|
||||
// A multi-file version arrives as the directory it was downloaded into, because
|
||||
// that is the only shape a set can take. llama.cpp wants a file: given the
|
||||
// first part of a split GGUF it finds the rest beside it, so the first part is
|
||||
// what it gets.
|
||||
//
|
||||
// A path that is already a file is returned untouched, which is every
|
||||
// single-file model and every hand-staged one.
|
||||
func ggufIn(path string) (string, error) {
|
||||
st, err := os.Stat(path)
|
||||
if err != nil || !st.IsDir() {
|
||||
// Not a directory, or not there at all — the second is the engine's to
|
||||
// report, with the message it would have given anyway.
|
||||
return path, nil //nolint:nilerr // a missing path is the engine's error to give
|
||||
}
|
||||
entries, err := os.ReadDir(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var names []string
|
||||
for _, e := range entries {
|
||||
if !e.IsDir() && strings.HasSuffix(strings.ToLower(e.Name()), ".gguf") {
|
||||
names = append(names, e.Name())
|
||||
}
|
||||
}
|
||||
if len(names) == 0 {
|
||||
return "", fmt.Errorf("%s holds no .gguf file", path)
|
||||
}
|
||||
// Sorted, so a split model's first part comes first: the parts are named
|
||||
// "-00001-of-00004" and lexical order is shard order.
|
||||
sort.Strings(names)
|
||||
return filepath.Join(path, names[0]), nil
|
||||
}
|
||||
|
||||
// UseInstalled points the engine at a build the agent installed and checked.
|
||||
// It takes precedence over everything else, because it is the only one that
|
||||
// was seen to run.
|
||||
func (l *LlamaCpp) UseInstalled(path string) { l.installed.Store(&path) }
|
||||
|
||||
// Configured reports whether the agent's configuration names a binary. A
|
||||
// named binary is the operator's deliberate choice; one merely found on PATH
|
||||
// is not, and is not reported as installed.
|
||||
func (l *LlamaCpp) Configured() bool { return l.cfg.Bin != "" }
|
||||
|
||||
// bin is the llama-server to run: the installed build, else the configured
|
||||
// one, else whatever PATH has — which is only ever reached by a node not
|
||||
// provisioned by the agent and not told where its engine is.
|
||||
func (l *LlamaCpp) bin() string {
|
||||
if p := l.installed.Load(); p != nil {
|
||||
return *p
|
||||
}
|
||||
if l.cfg.Bin != "" {
|
||||
return l.cfg.Bin
|
||||
}
|
||||
return "llama-server"
|
||||
}
|
||||
@@ -0,0 +1,325 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func llamaSpec(gpus []int) ShardSpec {
|
||||
return ShardSpec{
|
||||
ReplicaID: "rep-1", ShardID: "sh-1", Rank: 0, NNodes: 1,
|
||||
GPUIndices: gpus, HTTPPort: 8000, MgmtInterface: "eth0",
|
||||
Profile: Profile{
|
||||
Name: "cpu", Engine: "llamacpp", TP: 1,
|
||||
ModelPath: "/models/qwen.gguf", ServedModelName: "qwen-0.5b",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppCommand(t *testing.T) {
|
||||
cmds, err := (&LlamaCpp{}).Commands(llamaSpec(nil))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("a llama.cpp shard is one process, got %d", len(cmds))
|
||||
}
|
||||
c := cmds[0]
|
||||
if c.Path != "llama-server" {
|
||||
t.Errorf("path %q", c.Path)
|
||||
}
|
||||
|
||||
// --metrics is not optional: llama-server serves no /metrics without it, and
|
||||
// an endpoint the router cannot measure is one it routes to blind.
|
||||
if !slices.Contains(c.Args, "--metrics") {
|
||||
t.Error("the server was started without --metrics")
|
||||
}
|
||||
if v, _ := flagValue(c.Args, "--model"); v != "/models/qwen.gguf" {
|
||||
t.Errorf("--model %q", v)
|
||||
}
|
||||
if v, _ := flagValue(c.Args, "--port"); v != "8000" {
|
||||
t.Errorf("--port %q", v)
|
||||
}
|
||||
if v, _ := flagValue(c.Args, "--alias"); v != "qwen-0.5b" {
|
||||
t.Errorf("--alias %q", v)
|
||||
}
|
||||
// llama.cpp serves one request at a time unless told otherwise, which would
|
||||
// make every second request wait no matter what the router decided.
|
||||
if v, ok := flagValue(c.Args, "--parallel"); !ok || v == "1" {
|
||||
t.Errorf("--parallel %q", v)
|
||||
}
|
||||
if slices.Contains(c.Args, "--n-gpu-layers") {
|
||||
t.Error("a CPU placement was given GPU layers")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppOffloadsWhenGivenGPUs(t *testing.T) {
|
||||
cmds, _ := (&LlamaCpp{}).Commands(llamaSpec([]int{2}))
|
||||
if v, ok := flagValue(cmds[0].Args, "--n-gpu-layers"); !ok || v != "999" {
|
||||
t.Errorf("a GPU placement should offload the model, got %q", v)
|
||||
}
|
||||
// The shard must still be pinned to its own card.
|
||||
if !slices.Contains(cmds[0].Env, "CUDA_VISIBLE_DEVICES=2") {
|
||||
t.Error("the shard was not pinned to its GPU")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileArgumentsWinOverDefaults(t *testing.T) {
|
||||
// The escape hatch has to actually work, or an operator tuning slot count
|
||||
// ends up with two --parallel flags and whichever llama.cpp happens to
|
||||
// prefer.
|
||||
s := llamaSpec(nil)
|
||||
s.Profile.Args = []string{"--parallel", "16", "--ctx-size", "4096"}
|
||||
|
||||
cmds, _ := (&LlamaCpp{}).Commands(s)
|
||||
n := 0
|
||||
for _, a := range cmds[0].Args {
|
||||
if a == "--parallel" {
|
||||
n++
|
||||
}
|
||||
}
|
||||
if n != 1 {
|
||||
t.Errorf("--parallel appears %d times", n)
|
||||
}
|
||||
if v, _ := flagValue(cmds[0].Args, "--parallel"); v != "16" {
|
||||
t.Errorf("the profile's value lost: %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShortFormArgumentsAreRecognized(t *testing.T) {
|
||||
s := llamaSpec([]int{0})
|
||||
s.Profile.Args = []string{"-np", "8", "-ngl", "20"}
|
||||
|
||||
cmds, _ := (&LlamaCpp{}).Commands(s)
|
||||
for _, flag := range []string{"--parallel", "--n-gpu-layers"} {
|
||||
if slices.Contains(cmds[0].Args, flag) {
|
||||
t.Errorf("%s was added although the profile set its short form", flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppRefusesWhatItCannotDo(t *testing.T) {
|
||||
// Accepting either of these would produce a replica that starts and then
|
||||
// serves from one machine while the catalog believes it has several.
|
||||
multi := Profile{Name: "multi", Engine: "llamacpp", TP: 1, Nodes: 2, GPUsPerNode: 1, ModelPath: "/m"}
|
||||
if err := (&LlamaCpp{}).Validate(multi, nil); err == nil {
|
||||
t.Error("a multi-node llama.cpp profile was accepted")
|
||||
}
|
||||
|
||||
tp := Profile{Name: "tp", Engine: "llamacpp", TP: 2, GPUsPerNode: 2, ModelPath: "/m"}
|
||||
if err := (&LlamaCpp{}).Validate(tp, h100(2)); err == nil {
|
||||
t.Error("a tensor-parallel llama.cpp profile was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCPUProfileValidates(t *testing.T) {
|
||||
p := Profile{Name: "cpu", Engine: "llamacpp", TP: 1, ModelPath: "/models/qwen.gguf"}
|
||||
if err := (&LlamaCpp{}).Validate(p, nil); err != nil {
|
||||
t.Errorf("a CPU profile on a node with no GPUs was rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppAlwaysServesHTTP(t *testing.T) {
|
||||
// Unlike the others there is no worker rank to get this wrong for.
|
||||
if !(&LlamaCpp{}).ServesHTTP(llamaSpec(nil)) {
|
||||
t.Error("a llama.cpp shard must answer the API")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBinaryIsConfigurable(t *testing.T) {
|
||||
l := &LlamaCpp{cfg: LlamaCppConfig{Bin: "/opt/llama.cpp/bin/llama-server", Parallel: 2}}
|
||||
cmds, _ := l.Commands(llamaSpec(nil))
|
||||
if cmds[0].Path != "/opt/llama.cpp/bin/llama-server" {
|
||||
t.Errorf("path %q", cmds[0].Path)
|
||||
}
|
||||
if v, _ := flagValue(cmds[0].Args, "--parallel"); v != "2" {
|
||||
t.Errorf("--parallel %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnknownEngineNamesLlamaCpp(t *testing.T) {
|
||||
_, err := NewRegistry(Config{}).Get("llama.cpp")
|
||||
if err == nil || !strings.Contains(err.Error(), "llamacpp") {
|
||||
t.Errorf("the suggestion should name the real engine: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// flagValue reads a "--flag value" pair out of a command line.
|
||||
func flagValue(args []string, flag string) (string, bool) {
|
||||
for i, a := range args {
|
||||
if a == flag && i+1 < len(args) {
|
||||
return args[i+1], true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func TestLlamaCppUsesTheModelsChatTemplate(t *testing.T) {
|
||||
// Pinned so an engine upgrade cannot change how prompts are built under a
|
||||
// running fleet: build 10991 defaults to jinja, older builds do not.
|
||||
//
|
||||
// Not a fix for tool calling, though it was added while chasing that. On
|
||||
// this build the flag is a no-op and the behaviour is identical with and
|
||||
// without it.
|
||||
l := &LlamaCpp{}
|
||||
cmds, err := l.Commands(ShardSpec{HTTPPort: 8080, Profile: Profile{ModelPath: "/m.gguf"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !slices.Contains(cmds[0].Args, "--jinja") {
|
||||
t.Errorf("--jinja is missing; tool calls will come back as text: %v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppLetsAProfileTurnTheTemplateOff(t *testing.T) {
|
||||
// A model whose bundled template is broken has to be serviceable without
|
||||
// editing the adapter.
|
||||
l := &LlamaCpp{}
|
||||
cmds, err := l.Commands(ShardSpec{
|
||||
HTTPPort: 8080,
|
||||
Profile: Profile{ModelPath: "/m.gguf", Args: []string{"--no-jinja"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if slices.Contains(cmds[0].Args, "--jinja") {
|
||||
t.Errorf("the profile said --no-jinja and got both: %v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASetDirectoryResolvesToItsFirstPart(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
// Deliberately out of order on disk, because readdir order is not shard
|
||||
// order and the first part is the one llama.cpp must be given.
|
||||
for _, n := range []string{
|
||||
"qwen-q4_k_m-00003-of-00003.gguf",
|
||||
"qwen-q4_k_m-00001-of-00003.gguf",
|
||||
"qwen-q4_k_m-00002-of-00003.gguf",
|
||||
} {
|
||||
if err := os.WriteFile(filepath.Join(dir, n), []byte("x"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
got, err := ggufIn(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Given the first part llama.cpp finds the rest beside it; given the third
|
||||
// it loads a third of a model.
|
||||
if filepath.Base(got) != "qwen-q4_k_m-00001-of-00003.gguf" {
|
||||
t.Errorf("expected the first shard, got %s", filepath.Base(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPlainFilePathIsUntouched(t *testing.T) {
|
||||
// Every single-file model and every hand-staged one takes this path.
|
||||
p := filepath.Join(t.TempDir(), "model.gguf")
|
||||
if err := os.WriteFile(p, []byte("x"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, _ := ggufIn(p); got != p {
|
||||
t.Errorf("expected the path unchanged, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestADirectoryWithNoWeightsSaysSo(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte("{}"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := ggufIn(dir); err == nil {
|
||||
t.Error("a directory with no .gguf in it is not a model llama.cpp can load")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingPathIsLeftToTheEngine(t *testing.T) {
|
||||
// The engine's own "no such file" names the path and the flag it came
|
||||
// from, which is more use than anything this could say.
|
||||
p := "/nonexistent/model.gguf"
|
||||
got, err := ggufIn(p)
|
||||
if err != nil || got != p {
|
||||
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")
|
||||
}
|
||||
}
|
||||
|
||||
func TestACPUShardLeavesHalfTheMachineForAnother(t *testing.T) {
|
||||
prev := numCPU
|
||||
numCPU = func() int { return 8 }
|
||||
t.Cleanup(func() { numCPU = prev })
|
||||
|
||||
threads := func(l *LlamaCpp, s ShardSpec) (string, bool) {
|
||||
cmds, err := l.Commands(s)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return flagValue(cmds[0].Args, "--threads")
|
||||
}
|
||||
if v, _ := threads(&LlamaCpp{}, llamaSpec(nil)); v != "4" {
|
||||
t.Errorf("--threads %q on 8 CPUs, expected 4", v)
|
||||
}
|
||||
if v, _ := threads(&LlamaCpp{cfg: LlamaCppConfig{Threads: 8}}, llamaSpec(nil)); v != "8" {
|
||||
t.Errorf("--threads %q with 8 configured", v)
|
||||
}
|
||||
if _, ok := threads(&LlamaCpp{}, llamaSpec([]int{0})); ok {
|
||||
t.Error("a GPU shard was given a thread count")
|
||||
}
|
||||
s := llamaSpec(nil)
|
||||
s.Profile.Args = []string{"-t", "2"}
|
||||
cmds, _ := (&LlamaCpp{}).Commands(s)
|
||||
if slices.Contains(cmds[0].Args, "--threads") {
|
||||
t.Error("the profile's own -t was overridden")
|
||||
}
|
||||
numCPU = func() int { return 1 }
|
||||
if v, _ := threads(&LlamaCpp{}, llamaSpec(nil)); v != "1" {
|
||||
t.Errorf("--threads %q on one CPU", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheCoresPlacementReservedAreTheThreads(t *testing.T) {
|
||||
s := llamaSpec(nil)
|
||||
s.Threads = 3
|
||||
cmds, err := (&LlamaCpp{cfg: LlamaCppConfig{Threads: 8}}).Commands(s)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if v, _ := flagValue(cmds[0].Args, "--threads"); v != "3" {
|
||||
t.Errorf("--threads %q; what placement reserved should win over the node's setting", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdaptersAreLoadedAtStartAndAppliedToNoneByDefault(t *testing.T) {
|
||||
s := llamaSpec(nil)
|
||||
s.Profile.Adapters = []Adapter{{Name: "sql", Path: "/w/sql.gguf"}, {Name: "docs", Path: "/w/docs.gguf"}}
|
||||
cmds, err := (&LlamaCpp{}).Commands(s)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
args := strings.Join(cmds[0].Args, " ")
|
||||
if !strings.Contains(args, "--lora-init-without-apply --lora /w/sql.gguf --lora /w/docs.gguf") {
|
||||
t.Errorf("args %s: both adapters, in order, loaded unapplied", args)
|
||||
}
|
||||
plain, _ := (&LlamaCpp{}).Commands(llamaSpec(nil))
|
||||
if strings.Contains(strings.Join(plain[0].Args, " "), "--lora") {
|
||||
t.Error("a model with no adapters was given lora flags")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Unknown is the version reported for an engine that is installed and would
|
||||
// not say which version it is.
|
||||
//
|
||||
// A value, not an absent key, because the two mean different things and
|
||||
// placement reads the keys: absent is "this node cannot run it", Unknown is
|
||||
// "it is here and the probe could not get a number out of it". Collapsing them
|
||||
// would make a slow engine look like a missing one and take a node out of
|
||||
// scheduling for being slow to answer --version.
|
||||
const Unknown = "unknown"
|
||||
|
||||
// Prober is an engine that can say whether it is installed on this machine.
|
||||
//
|
||||
// Optional, and asserted for rather than part of Engine: an adapter that has no
|
||||
// way to check should report nothing rather than be forced to invent an answer.
|
||||
type Prober interface {
|
||||
// Probe returns the engine's version, or an error when it is not installed.
|
||||
Probe(ctx context.Context) (string, error)
|
||||
}
|
||||
|
||||
// Probe asks every registered engine whether it is here.
|
||||
//
|
||||
// The result is what the node reports at registration, and its key set is what
|
||||
// placement filters on. Until this existed nobody knew: the registry registers
|
||||
// every adapter whether or not its binary exists, so a node with no vLLM
|
||||
// accepted a vLLM shard, registered healthy, and failed the shard minutes later
|
||||
// with "executable file not found" — in the engine's own words, on a machine
|
||||
// nobody was watching.
|
||||
//
|
||||
// Bounded, because it runs on the path to registering and an agent that cannot
|
||||
// reach the control plane is worse than one that reports an incomplete list.
|
||||
func (r *Registry) Probe(ctx context.Context) map[string]string {
|
||||
out := map[string]string{}
|
||||
for name, e := range r.engines {
|
||||
p, ok := e.(Prober)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
v, err := p.Probe(ctx)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if v == "" {
|
||||
v = Unknown
|
||||
}
|
||||
out[name] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Names of the engines in a probe result, sorted.
|
||||
func ProbedNames(m map[string]string) []string {
|
||||
out := make([]string, 0, len(m))
|
||||
for n := range m {
|
||||
out = append(out, n)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// probeTimeout bounds one engine's version check.
|
||||
//
|
||||
// Generous, because vLLM's --version imports torch and that is seconds on a
|
||||
// cold page cache, and mean, because this is startup: an engine that has not
|
||||
// answered in this long is reported as present-with-unknown-version rather
|
||||
// than waited for.
|
||||
const probeTimeout = 8 * time.Second
|
||||
|
||||
// run executes a version command, saying whether the program was even there.
|
||||
//
|
||||
// The two failures are different answers and the caller has to tell them
|
||||
// apart. A binary that is absent is an engine this node cannot run. A binary
|
||||
// that is present and exits non-zero is an engine it probably can — unless the
|
||||
// command was a test of whether a module imports, where a non-zero exit is the
|
||||
// whole answer.
|
||||
func run(ctx context.Context, bin string, args ...string) (out string, found bool, ok bool) {
|
||||
path, err := exec.LookPath(bin)
|
||||
if err != nil {
|
||||
return "", false, false
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, probeTimeout)
|
||||
defer cancel()
|
||||
|
||||
// Combined, not just stdout: engines print their version to either, and one
|
||||
// that chose stderr would otherwise read as having no version at all.
|
||||
b, err := exec.CommandContext(ctx, path, args...).CombinedOutput()
|
||||
return string(b), true, err == nil
|
||||
}
|
||||
|
||||
// installedBinary reports an engine that is its own executable.
|
||||
//
|
||||
// It resolves the path and runs nothing. The first version of this ran the
|
||||
// engine with --version, which is fine for an engine and not fine for the
|
||||
// wrapper scripts people put in front of one: a script that shifts past flags
|
||||
// it does not recognise and then execs the server will start an inference
|
||||
// server, at agent startup, on every node. The agent's own test suite installs
|
||||
// exactly such a script, and caught it.
|
||||
//
|
||||
// So the version is not read here. Presence is what placement needs, and what
|
||||
// the version is for — telling whether an engine still accepts the flags llmgw
|
||||
// builds — is `llmgw ctl engine-check`, which is run deliberately by somebody
|
||||
// who meant to run the engine.
|
||||
func installedBinary(bin string) (string, error) {
|
||||
if _, err := exec.LookPath(bin); err != nil {
|
||||
return "", errNotInstalled
|
||||
}
|
||||
return Unknown, nil
|
||||
}
|
||||
|
||||
// versionOfModule probes an engine that is a module inside an interpreter.
|
||||
//
|
||||
// This one does execute, and safely: the command is llmgw's own one-line
|
||||
// import, not the engine's launcher, so there is no wrapper to shift past
|
||||
// flags and start a server. There is also no alternative — a module has no
|
||||
// path to look up.
|
||||
//
|
||||
// A non-zero exit is the answer: python3 exists on almost every machine, so
|
||||
// presence of the interpreter says nothing. Treating a failed import as
|
||||
// "installed" reported SGLang on every node in a fleet that had none, which
|
||||
// would take nodes into scheduling for an engine they cannot run — the exact
|
||||
// failure this whole change exists to stop.
|
||||
func versionOfModule(ctx context.Context, python string, args ...string) (string, error) {
|
||||
out, found, ok := run(ctx, python, args...)
|
||||
if !found || !ok {
|
||||
return "", errNotInstalled
|
||||
}
|
||||
return firstVersion(out), nil
|
||||
}
|
||||
|
||||
var errNotInstalled = errors.New("not installed on this node")
|
||||
|
||||
var versionPattern = regexp.MustCompile(`\b\d+\.\d+(\.\d+)?\S*`)
|
||||
|
||||
// firstVersion is the first thing in the output that looks like a version.
|
||||
//
|
||||
// Pattern rather than parsing each engine's format: they print "vllm 0.6.3",
|
||||
// "version: b4123 (…)", and a paragraph with the number in the middle, and all
|
||||
// of those change between releases. What this is for is telling one build from
|
||||
// another, which the first number does.
|
||||
func firstVersion(s string) string {
|
||||
for _, line := range strings.Split(s, "\n") {
|
||||
if m := versionPattern.FindString(line); m != "" {
|
||||
return strings.TrimRight(m, ",;)")
|
||||
}
|
||||
}
|
||||
return Unknown
|
||||
}
|
||||
|
||||
// Configured is an engine whose location the agent's configuration names.
|
||||
type Configured interface {
|
||||
Configured() bool
|
||||
}
|
||||
|
||||
// ProbeConfigured probes only the engines the configuration names.
|
||||
//
|
||||
// Probe looks everywhere — PATH, the default interpreter — and reports what
|
||||
// it finds, which is exactly the trust in whatever happens to be on a machine
|
||||
// that installing engines replaces. A node the agent provisions reports what
|
||||
// it installed, plus what the operator named on purpose; nothing else.
|
||||
func (r *Registry) ProbeConfigured(ctx context.Context) map[string]string {
|
||||
out := map[string]string{}
|
||||
for name, e := range r.engines {
|
||||
c, ok := e.(Configured)
|
||||
if !ok || !c.Configured() {
|
||||
continue
|
||||
}
|
||||
p, ok := e.(Prober)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if v, err := p.Probe(ctx); err == nil {
|
||||
if v == "" {
|
||||
v = Unknown
|
||||
}
|
||||
out[name] = v
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Probe reports whether the llama.cpp server is installed.
|
||||
func (l *LlamaCpp) Probe(context.Context) (string, error) {
|
||||
return installedBinary(l.bin())
|
||||
}
|
||||
|
||||
// Configured reports whether the configuration names vLLM.
|
||||
func (v *VLLM) Configured() bool { return v.cfg.Bin != "" }
|
||||
|
||||
// Configured reports whether the configuration names SGLang's interpreter.
|
||||
func (g *SGLang) Configured() bool { return g.cfg.Python != "" }
|
||||
|
||||
// Probe reports whether vLLM is installed.
|
||||
func (v *VLLM) Probe(context.Context) (string, error) {
|
||||
return installedBinary(v.bin())
|
||||
}
|
||||
|
||||
// Probe reports the SGLang version.
|
||||
//
|
||||
// Through the interpreter, because SGLang is a module rather than a binary:
|
||||
// asking the shell for "sglang" finds nothing on a machine that can run it
|
||||
// perfectly well.
|
||||
func (g *SGLang) Probe(ctx context.Context) (string, error) {
|
||||
return versionOfModule(ctx, g.python(), "-c",
|
||||
"import sglang, sys; sys.stdout.write(getattr(sglang, '__version__', ''))")
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// script writes an executable that stands in for an engine.
|
||||
func script(t *testing.T, name, body string) string {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, name)
|
||||
if err := os.WriteFile(p, []byte("#!/bin/sh\n"+body+"\n"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func TestAnEngineThatIsNotInstalledIsAbsentFromTheReport(t *testing.T) {
|
||||
r := NewRegistry(Config{})
|
||||
r.engines = map[string]Engine{
|
||||
"vllm": &VLLM{cfg: VLLMConfig{Bin: "/nonexistent/vllm"}},
|
||||
}
|
||||
// Absent, not present-with-unknown-version. Placement reads the keys, and
|
||||
// a node listed for an engine it cannot run is the failure being fixed.
|
||||
if got := r.Probe(t.Context()); len(got) != 0 {
|
||||
t.Errorf("expected nothing, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbingNeverRunsTheEngine(t *testing.T) {
|
||||
// A wrapper that shifts past flags it does not recognise and then starts
|
||||
// the server. People put these in front of engines, and the agent's own
|
||||
// tests install one. Probing with --version would launch an inference
|
||||
// server at agent startup, on every node.
|
||||
marker := filepath.Join(t.TempDir(), "ran")
|
||||
bin := script(t, "fake-vllm", "touch "+marker+"; sleep 60")
|
||||
|
||||
r := NewRegistry(Config{})
|
||||
r.engines = map[string]Engine{"vllm": &VLLM{cfg: VLLMConfig{Bin: bin}}}
|
||||
|
||||
got := r.Probe(t.Context())
|
||||
if _, ok := got["vllm"]; !ok {
|
||||
t.Error("a binary that is there should be reported present")
|
||||
}
|
||||
if _, err := os.Stat(marker); err == nil {
|
||||
t.Fatal("the engine was executed by a probe")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAModuleThatWillNotImportIsNotInstalled(t *testing.T) {
|
||||
// python3 is on almost every machine, so the interpreter being there says
|
||||
// nothing about SGLang. Reporting it present put every node in a fleet
|
||||
// forward for an engine none of them had.
|
||||
bin := script(t, "python3", "exit 1")
|
||||
r := NewRegistry(Config{})
|
||||
r.engines = map[string]Engine{"sglang": &SGLang{cfg: SGLangConfig{Python: bin}}}
|
||||
|
||||
if got := r.Probe(t.Context()); len(got) != 0 {
|
||||
t.Errorf("a failed import means not installed, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAModuleThatImportsIsReportedWithItsVersion(t *testing.T) {
|
||||
bin := script(t, "python3", `echo "0.4.1"`)
|
||||
r := NewRegistry(Config{})
|
||||
r.engines = map[string]Engine{"sglang": &SGLang{cfg: SGLangConfig{Python: bin}}}
|
||||
|
||||
got := r.Probe(t.Context())
|
||||
if got["sglang"] != "0.4.1" {
|
||||
t.Errorf("expected the version, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAVersionIsPickedOutOfWhateverIsPrinted(t *testing.T) {
|
||||
// Only the module probe reaches this now, and a Python package version is
|
||||
// dotted. A bare build number is deliberately not matched: it would turn a
|
||||
// stray number in a line of prose into a version.
|
||||
for in, want := range map[string]string{
|
||||
"0.4.1": "0.4.1",
|
||||
"sglang 0.4.post1": "0.4.post1",
|
||||
"warning: noise\n0.4.1": "0.4.1",
|
||||
"": Unknown,
|
||||
"no numbers here": Unknown,
|
||||
"just 4123": Unknown,
|
||||
} {
|
||||
if got := firstVersion(in); got != want {
|
||||
t.Errorf("firstVersion(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbingIsBoundedByTheContext(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
bin := script(t, "python3", "sleep 60")
|
||||
r := NewRegistry(Config{})
|
||||
r.engines = map[string]Engine{"sglang": &SGLang{cfg: SGLangConfig{Python: bin}}}
|
||||
|
||||
// Registration waits on this. An agent that cannot reach the control plane
|
||||
// is worse than one reporting an incomplete list.
|
||||
if got := r.Probe(ctx); len(got) != 0 {
|
||||
t.Errorf("a cancelled probe should report nothing, got %v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// SGLangConfig locates the SGLang installation.
|
||||
type SGLangConfig struct {
|
||||
// Python is the interpreter of the SGLang virtualenv. SGLang is launched as
|
||||
// a module rather than a console script, which is the documented entry
|
||||
// point and avoids depending on the script being on PATH.
|
||||
Python string `koanf:"python"`
|
||||
}
|
||||
|
||||
// SGLang builds SGLang launch commands.
|
||||
type SGLang struct {
|
||||
cfg SGLangConfig
|
||||
installed atomic.Pointer[string]
|
||||
}
|
||||
|
||||
// UseInstalled points the engine at the interpreter of the environment the
|
||||
// agent installed and checked, ahead of the configuration.
|
||||
func (g *SGLang) UseInstalled(path string) { g.installed.Store(&path) }
|
||||
|
||||
// python is the interpreter to launch SGLang with.
|
||||
func (g *SGLang) python() string {
|
||||
if p := g.installed.Load(); p != nil {
|
||||
return *p
|
||||
}
|
||||
if g.cfg.Python != "" {
|
||||
return g.cfg.Python
|
||||
}
|
||||
return "python3"
|
||||
}
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*SGLang) Name() string { return "sglang" }
|
||||
|
||||
// HealthPath is SGLang's readiness endpoint.
|
||||
func (*SGLang) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is SGLang's Prometheus endpoint.
|
||||
func (*SGLang) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
func (*SGLang) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if err := validateCommon(p, gpus); err != nil {
|
||||
return err
|
||||
}
|
||||
if q := quantizationOf(p.Args); q == "fp8" {
|
||||
if err := requireComputeCapability(p, gpus, 8.9, "fp8 quantization"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API.
|
||||
//
|
||||
// Every SGLang rank runs the same launcher, but only node-rank 0 binds the HTTP
|
||||
// port; the others exist to contribute GPUs to the tensor-parallel group.
|
||||
func (*SGLang) ServesHTTP(s ShardSpec) bool { return s.Leader() }
|
||||
|
||||
// Commands returns the processes for one shard.
|
||||
//
|
||||
// SGLang supports multi-node natively, so there is only ever one process: every
|
||||
// rank runs the same launcher with its own rank number and they rendezvous at
|
||||
// the leader's dist-init address.
|
||||
func (g *SGLang) Commands(s ShardSpec) ([]Command, error) {
|
||||
c, err := g.serverCommand(s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []Command{c}, nil
|
||||
}
|
||||
|
||||
func (g *SGLang) serverCommand(s ShardSpec) (Command, error) {
|
||||
python := g.python()
|
||||
p := s.Profile
|
||||
|
||||
args := []string{"-m", "sglang.launch_server",
|
||||
"--model-path", p.ModelPath,
|
||||
"--host", "0.0.0.0",
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
"--tp", strconv.Itoa(p.TP),
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--served-model-name", p.ServedModelName)
|
||||
}
|
||||
if s.NNodes > 1 {
|
||||
args = append(args,
|
||||
"--nnodes", strconv.Itoa(s.NNodes),
|
||||
"--node-rank", strconv.Itoa(s.Rank),
|
||||
"--dist-init-addr", s.LeaderAddr+":"+strconv.Itoa(s.LeaderPort),
|
||||
)
|
||||
}
|
||||
args = append(args, p.Args...)
|
||||
// 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
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// SimConfig locates the simulated engine.
|
||||
type SimConfig struct {
|
||||
// Bin is the llmgw binary itself; the engine is one of its subcommands.
|
||||
// Empty means whatever started this agent, which is right almost always.
|
||||
Bin string `koanf:"bin"`
|
||||
}
|
||||
|
||||
// Sim launches llmgw's own simulated engine.
|
||||
//
|
||||
// It is in the production registry on purpose. A cluster can be brought up and
|
||||
// exercised end to end — placement, fabrics, multi-node sequencing, failure and
|
||||
// rebuild — on machines that have no GPUs, which is the only way most of that
|
||||
// path gets tested at all. The multi-node launch sequence went four phases
|
||||
// without ever running, because the only things exercising it were stub
|
||||
// processes that exercised the command line and nothing else.
|
||||
//
|
||||
// Nothing about it pretends to be real: the model name it serves is whatever it
|
||||
// is told, the tokens are invented, and it says so in its own logs.
|
||||
type Sim struct{ cfg SimConfig }
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*Sim) Name() string { return "sim" }
|
||||
|
||||
// HealthPath is the simulated engine's readiness endpoint.
|
||||
func (*Sim) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is its Prometheus endpoint, in vLLM's vocabulary.
|
||||
func (*Sim) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API.
|
||||
//
|
||||
// Only the leader, exactly as vLLM and SGLang behave — which is the behaviour
|
||||
// this engine exists to let us test.
|
||||
func (*Sim) ServesHTTP(s ShardSpec) bool { return s.Leader() }
|
||||
|
||||
// Validate accepts anything the arithmetic allows.
|
||||
//
|
||||
// No compute-capability checks and no quantization rules: there is no hardware
|
||||
// to be wrong about, and inventing refusals would make this a worse stand-in
|
||||
// for the engines it stands in for.
|
||||
func (*Sim) Validate(p Profile, gpus []GPUInfo) error {
|
||||
return validateCommon(p, gpus)
|
||||
}
|
||||
|
||||
// Commands returns the single process that makes up one shard.
|
||||
func (s *Sim) Commands(spec ShardSpec) ([]Command, error) {
|
||||
bin := s.cfg.Bin
|
||||
if bin == "" {
|
||||
exe, err := os.Executable()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("locating the llmgw binary for the sim engine: %w", err)
|
||||
}
|
||||
bin = exe
|
||||
}
|
||||
p := spec.Profile
|
||||
|
||||
args := []string{"sim-engine",
|
||||
"--listen", "0.0.0.0:" + strconv.Itoa(spec.HTTPPort),
|
||||
"--model", p.ServedModelName,
|
||||
}
|
||||
if spec.NNodes > 1 {
|
||||
args = append(args,
|
||||
"--nnodes", strconv.Itoa(spec.NNodes),
|
||||
"--node-rank", strconv.Itoa(spec.Rank),
|
||||
"--dist-init-addr", spec.LeaderAddr+":"+strconv.Itoa(spec.LeaderPort),
|
||||
)
|
||||
}
|
||||
args = append(args, p.Args...)
|
||||
|
||||
return []Command{{
|
||||
Name: "sim-engine", Path: bin, Args: args,
|
||||
Env: materializeEnv(baseEnv(spec)),
|
||||
// Workers dial the leader, so the leader has to be listening first.
|
||||
// Both ranks express the same wait because both know the same address.
|
||||
WaitForLeaderPort: spec.NNodes > 1 && !spec.Leader(),
|
||||
}}, nil
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func simSpec(rank, nnodes int) ShardSpec {
|
||||
return ShardSpec{
|
||||
ReplicaID: "rep-1", ShardID: "sh-1", Rank: rank, NNodes: nnodes,
|
||||
LeaderAddr: "10.0.0.5", LeaderPort: 20001,
|
||||
GPUIndices: []int{0, 1}, HTTPPort: 20000, MgmtInterface: "eth0",
|
||||
HCAs: []string{"mlx5_0", "mlx5_1"},
|
||||
Profile: Profile{
|
||||
Name: "tp4", Engine: "sim", TP: 4, GPUsPerNode: 2, Nodes: nnodes,
|
||||
ModelPath: "/models/x", ServedModelName: "big-model",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyTheLeaderServes(t *testing.T) {
|
||||
// The behaviour that made replicas fail while they were coming up perfectly
|
||||
// well: health checking a worker's port, which never opens.
|
||||
if !(&Sim{}).ServesHTTP(simSpec(0, 2)) {
|
||||
t.Error("the leader does not serve")
|
||||
}
|
||||
if (&Sim{}).ServesHTTP(simSpec(1, 2)) {
|
||||
t.Error("a worker claims to serve")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMultiNodeShardCarriesItsRendezvous(t *testing.T) {
|
||||
cmds, err := (&Sim{cfg: SimConfig{Bin: "/usr/local/bin/llmgw"}}).Commands(simSpec(1, 2))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("%d commands", len(cmds))
|
||||
}
|
||||
c := cmds[0]
|
||||
|
||||
if c.Args[0] != "sim-engine" {
|
||||
t.Errorf("the engine is a subcommand of llmgw: %v", c.Args)
|
||||
}
|
||||
for flag, want := range map[string]string{
|
||||
"--nnodes": "2", "--node-rank": "1", "--dist-init-addr": "10.0.0.5:20001",
|
||||
} {
|
||||
if got, _ := flagValue(c.Args, flag); got != want {
|
||||
t.Errorf("%s = %q, want %q", flag, got, want)
|
||||
}
|
||||
}
|
||||
// A worker must wait for the leader's port, or the launch is a race.
|
||||
if !c.WaitForLeaderPort {
|
||||
t.Error("a worker does not wait for the leader")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheLeaderDoesNotWaitForItself(t *testing.T) {
|
||||
cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(0, 2))
|
||||
if cmds[0].WaitForLeaderPort {
|
||||
t.Error("the leader waits for a port it is about to open")
|
||||
}
|
||||
}
|
||||
|
||||
func TestASingleNodeShardCarriesNoRendezvous(t *testing.T) {
|
||||
cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(0, 1))
|
||||
for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} {
|
||||
if slices.Contains(cmds[0].Args, flag) {
|
||||
t.Errorf("a single-node shard was given %s", flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheShardIsStillPinnedToItsFabric(t *testing.T) {
|
||||
// The simulated engine does not use NCCL, but the environment is generated
|
||||
// by the same code that generates it for the ones that do — so a cluster
|
||||
// exercised with this engine is exercising that too.
|
||||
cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(1, 2))
|
||||
env := cmds[0].Env
|
||||
for _, want := range []string{
|
||||
"NCCL_IB_HCA=mlx5_0,mlx5_1", "NCCL_SOCKET_IFNAME=eth0",
|
||||
"CUDA_VISIBLE_DEVICES=0,1", "NCCL_IB_DISABLE=0",
|
||||
} {
|
||||
if !slices.Contains(env, want) {
|
||||
t.Errorf("the environment is missing %s", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loaded 100 of 582 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user