Compare commits
237
Commits
v0.11.0
...
settings-e
No files matched your search
@@ -39,6 +39,24 @@ jobs:
|
||||
- name: Build
|
||||
run: go build ./...
|
||||
|
||||
# llmgw-kb runs on people's own machines, Windows included, which no
|
||||
# test here runs on: building it for each catches a Unix-only call.
|
||||
- name: Build llmgw-kb for every system
|
||||
run: |
|
||||
for target in windows/amd64 windows/arm64 darwin/amd64 darwin/arm64 linux/amd64 linux/arm64; do
|
||||
CGO_ENABLED=0 GOOS="${target%/*}" GOARCH="${target#*/}" go build -o /dev/null ./cmd/llmgw-kb
|
||||
done
|
||||
|
||||
# The Terraform provider is a module of its own, which ./... above does
|
||||
# not reach. Its tests check its client against docs/admin-api.openapi.json,
|
||||
# so a change to the admin API that breaks the provider fails here. The
|
||||
# acceptance tests (TF_ACC=1) are left out: they need Terraform itself.
|
||||
- name: Test the Terraform provider
|
||||
working-directory: terraform-provider-llmgw
|
||||
run: |
|
||||
go vet ./...
|
||||
go test ./...
|
||||
|
||||
# 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.
|
||||
|
||||
@@ -55,6 +55,21 @@ jobs:
|
||||
(cd dist && sha256sum "${name}" >> checksums.txt)
|
||||
done
|
||||
|
||||
# llmgw-kb, which people run on their own machines to put their
|
||||
# documents into a knowledge base: Windows too, which llmgw itself
|
||||
# does not build for.
|
||||
for target in windows/amd64 windows/arm64 darwin/amd64 darwin/arm64 linux/amd64 linux/arm64; do
|
||||
os="${target%/*}"; arch="${target#*/}"
|
||||
ext=""; [ "$os" = windows ] && ext=".exe"
|
||||
name="llmgw-kb-${version}-${os}-${arch}${ext}"
|
||||
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-kb
|
||||
(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
|
||||
@@ -91,13 +106,13 @@ jobs:
|
||||
# test the cross-compiled ones, which is why the native build is
|
||||
# checked rather than assumed.
|
||||
./dist/llmgw-${version}-linux-amd64 --version
|
||||
./dist/llmgw-kb-${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.
|
||||
# Signed, so a machine that trusts the key built into it installs
|
||||
# nothing the server alone could have put there. `release sign` checks
|
||||
# the signature against that key before writing it, so a secret that
|
||||
# does not match trusted.pub stops the release too.
|
||||
- name: Sign the checksums
|
||||
env:
|
||||
RELEASE_SIGNING_KEY: ${{ secrets.RELEASE_SIGNING_KEY }}
|
||||
@@ -105,7 +120,15 @@ jobs:
|
||||
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"
|
||||
# Builds that carry a trusted key install only signed releases:
|
||||
# publishing an unsigned one would be a release no upgraded
|
||||
# machine accepts, so it stops here instead.
|
||||
if [ -s internal/release/trusted.pub ]; then
|
||||
echo "internal/release/trusted.pub holds a key, and the RELEASE_SIGNING_KEY secret is not set:"
|
||||
echo "add it under Settings -> Secrets, or this release could not be installed anywhere."
|
||||
exit 1
|
||||
fi
|
||||
echo "no RELEASE_SIGNING_KEY secret and no trusted key; releases are not signed"
|
||||
exit 0
|
||||
fi
|
||||
./dist/llmgw-${version}-linux-amd64 release sign dist/checksums.txt
|
||||
|
||||
@@ -5,3 +5,7 @@
|
||||
*.db-shm
|
||||
.DS_Store
|
||||
coverage.out
|
||||
# The provider built in place, and what Terraform leaves beside an example.
|
||||
/terraform-provider-llmgw/terraform-provider-llmgw
|
||||
/terraform-provider-llmgw/examples/.terraform*
|
||||
/terraform-provider-llmgw/examples/terraform.tfstate*
|
||||
+1033
File diff suppressed because it is too large.
Load diff
+1460
File diff suppressed because it is too large.
Load diff
@@ -2,10 +2,11 @@ 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 proto proto-tools
|
||||
.PHONY: build test race lint fmt vuln check clean rename proto proto-tools
|
||||
|
||||
build:
|
||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw ./cmd/llmgw
|
||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw-kb ./cmd/llmgw-kb
|
||||
|
||||
test:
|
||||
go test ./...
|
||||
@@ -20,6 +21,17 @@ fmt:
|
||||
go fmt ./...
|
||||
go vet ./...
|
||||
|
||||
# Known vulnerabilities reachable from this code, in the standard library and
|
||||
# in dependencies. A toolchain left behind is the commonest cause: the go
|
||||
# directive in go.mod is what fixes it.
|
||||
vuln:
|
||||
go run golang.org/x/vuln/cmd/govulncheck@latest ./...
|
||||
|
||||
# What has to pass before a change is merged.
|
||||
check: test
|
||||
go vet ./...
|
||||
$(MAKE) vuln
|
||||
|
||||
# Regenerate the agent protocol. Generated code is committed, so this only
|
||||
# needs running when api/proto changes.
|
||||
proto: proto-tools
|
||||
|
||||
@@ -3,20 +3,23 @@
|
||||
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).
|
||||
It serves LLMs running on your own GPU and CPU servers (vLLM, SGLang,
|
||||
llama.cpp, TensorRT-LLM) behind one OpenAI- and Anthropic-compatible
|
||||
endpoint, and answers from your own documents. 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.
|
||||
On the machine that will be the control plane, with the release binary
|
||||
installed:
|
||||
|
||||
| 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 |
|
||||
```sh
|
||||
sudo llmgw
|
||||
```
|
||||
|
||||
The setup screens install one machine or three (high availability), check
|
||||
them first, and leave a model answering. After that the same command opens
|
||||
the console. [docs/quick-install.md](docs/quick-install.md) has the six steps
|
||||
from empty servers to a first request; everything else is below.
|
||||
|
||||
## Status
|
||||
|
||||
@@ -33,6 +36,11 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
| 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 |
|
||||
| v0.10.0 | — | InfiniBand fabrics and NVLink, doctor, node add, engines installed by the nodes | released |
|
||||
| v0.11.0 | — | High availability from setup, the setup and console screens, usage sums and per-tenant alerts | released |
|
||||
| v0.12.0 | — | The cluster's settings in the store, in force when changed; the firewall switch; wipe | released |
|
||||
| v0.13.0 | — | The whole system stopped and started from one command; Qwen3Guard judges, on a GPU and CPUs together | released |
|
||||
| v0.14.0 | — | Knowledge bases: documents from people's own machines, kept apart per tenant, answered from | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
As of v0.6.0 a cluster runs unattended: agents report their hardware, the
|
||||
@@ -40,12 +48,69 @@ 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.
|
||||
Since v0.11.0 (each release's detail is in [CHANGELOG.md](CHANGELOG.md)):
|
||||
|
||||
- **Knowledge bases** (v0.14.0). People put a company's documents in from
|
||||
their own machines with `llmgw-kb`, a separate download for Windows,
|
||||
macOS and Linux that reads PDF, Office, OpenDocument, mail and many more
|
||||
formats and sends text and vectors, never the files. Collections are kept
|
||||
apart per tenant, down to row-level security in PostgreSQL. A chat
|
||||
request, or a part of the tree by default, is answered from them with
|
||||
its sources: [docs/knowledge.md](docs/knowledge.md).
|
||||
- **One command for the whole system** (v0.13.0): `llmgw stop`, `start` and
|
||||
`restart`. Qwen3Guard is the guardrails' judge, on a GPU and on CPUs
|
||||
together.
|
||||
- **Settings that apply when changed** (v0.12.0). The cluster's
|
||||
configuration lives in the store, the same on every control plane and in
|
||||
force within seconds (`llmgw ctl config`, the console's Settings). The
|
||||
firewall can be switched on and off for the whole cluster.
|
||||
- **Security hardening** (v0.11.2): signed releases only, a token on the
|
||||
admin API, engines not run as root, TLS for the agent port, etcd and
|
||||
Patroni, a firewall for llmgw's internal ports, encrypted backups,
|
||||
bounded connections, sign-in throttling, a Content-Security-Policy.
|
||||
- **Guardrails** (v0.11.2). Policies check requests and answers for PII,
|
||||
secrets, prompt injection, deny patterns, and with a judge model; they
|
||||
observe or enforce, and attach to any part of the tenancy tree:
|
||||
[docs/guardrails.md](docs/guardrails.md).
|
||||
- **Administration** (v0.11.2):
|
||||
- an audit log of every change (`ctl audit`);
|
||||
- staff roles (viewer, operator, billing, admin);
|
||||
- single sign-on through OpenID Connect and provisioning through SCIM
|
||||
2.0;
|
||||
- key lifetime limits, and rotation with overlap (`keys ttl`, `keys
|
||||
rotate`);
|
||||
- usage export as CSV (`ctl usage export`).
|
||||
- **Notifications** (v0.11.2): webhooks, signed or Slack-shaped, for a
|
||||
model parked, a node lost, a budget spent and new weights published, and
|
||||
Alertmanager receiver templates.
|
||||
- **Traffic and models** (v0.11.2):
|
||||
- request tags, canaries of new weights with automatic rollback, and
|
||||
quality suites gating updates;
|
||||
- a response cache, kept and replayed requests, and batches;
|
||||
- Anthropic's Messages API, and more of OpenAI's (structured output,
|
||||
images, rerank, audio, moderations);
|
||||
- scaling by the clock, and SLO-aware admission;
|
||||
- TensorRT-LLM, MIG, and models imported from Ollama;
|
||||
- routing by what each vLLM replica's cache holds, and LoRA adapters
|
||||
loaded on demand;
|
||||
- other llmgw clusters as providers, and placement that knows racks,
|
||||
PDUs and power caps.
|
||||
- **Operations** (v0.11.2):
|
||||
- `server restore`, off-machine backup copies, and secrets in Vault;
|
||||
- forgetting a person, and keeping a project's requests in the cluster;
|
||||
- offline installs from a bundle, and a Helm chart;
|
||||
- tracing from the control plane to the engines;
|
||||
- the admin API in OpenAPI, with a Terraform provider;
|
||||
- the portal in English and Turkish, with its docs at /docs.
|
||||
|
||||
Every one of these is in the console as well as on the command line:
|
||||
[docs/console.md](docs/console.md).
|
||||
|
||||
What is still ahead is in [BACKLOG.md](BACKLOG.md). Two things that look like
|
||||
gaps are not: 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
|
||||
@@ -58,34 +123,65 @@ them. On such a machine, run `llmgw doctor --json` first and keep the report;
|
||||
|
||||
## 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
|
||||
Every page is also in the binary: `llmgw docs`, or **Manual** in the console.
|
||||
|
||||
**Installing**
|
||||
|
||||
- [docs/quick-install.md](docs/quick-install.md) — from empty servers to a
|
||||
first request, in six steps: start here
|
||||
- [docs/install.md](docs/install.md) — the same with every option explained,
|
||||
and running the cluster after: monitoring, notifications, single sign-on,
|
||||
the firewall, upgrades, backups
|
||||
- [docs/node-onboarding.md](docs/node-onboarding.md) — adding nodes to a
|
||||
cluster that exists
|
||||
- [docs/kubernetes.md](docs/kubernetes.md) — the control plane on
|
||||
Kubernetes with the Helm chart, and why GPU nodes stay outside it
|
||||
- [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
|
||||
|
||||
**Running it**
|
||||
|
||||
- [docs/console.md](docs/console.md) — the console that bare `llmgw` opens:
|
||||
every view, and the command each of its actions runs
|
||||
- [docs/onboarding.md](docs/onboarding.md) — accounts, keys, limits, staff
|
||||
roles and single sign-on
|
||||
- [docs/clients.md](docs/clients.md) — using the gateway from code with the
|
||||
official OpenAI and Anthropic SDKs, and llmgw's extensions
|
||||
- [docs/guardrails.md](docs/guardrails.md) — checking requests and answers:
|
||||
policies, detectors, a judge model, and where each applies
|
||||
- [docs/knowledge.md](docs/knowledge.md) — knowledge bases: collections,
|
||||
who may read them, `llmgw-kb` on people's machines, and answers from
|
||||
documents
|
||||
- [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
|
||||
|
||||
**Reference**
|
||||
|
||||
- [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
|
||||
- [docs/data.md](docs/data.md) — what llmgw stores, what it does not, and
|
||||
what leaves the machine; written for a security review
|
||||
- [docs/admin-api.md](docs/admin-api.md) and its OpenAPI document — the
|
||||
admin API, for scripts; [terraform-provider-llmgw/](terraform-provider-llmgw/)
|
||||
manages models, quotas and guardrails with Terraform or OpenTofu
|
||||
- [deploy/](deploy/) — [a mirror](deploy/mirror/) on its own machine,
|
||||
[Caddy and nginx](deploy/proxy/) in front of the client port, [two control
|
||||
planes](deploy/ha/) behind a balancer, and [Prometheus, Alertmanager and
|
||||
Grafana](deploy/monitoring/)
|
||||
|
||||
**Working on llmgw**
|
||||
|
||||
- [docs/development.md](docs/development.md) — built from source on one
|
||||
workstation, no systemd and no `sudo`
|
||||
- [llmgw-plan.md](llmgw-plan.md) — the design; [BACKLOG.md](BACKLOG.md) —
|
||||
what is next
|
||||
|
||||
## Building
|
||||
|
||||
```sh
|
||||
make build # bin/llmgw
|
||||
make build # bin/llmgw and bin/llmgw-kb
|
||||
make test # unit tests
|
||||
make race # with the race detector
|
||||
```
|
||||
@@ -100,8 +196,13 @@ 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
|
||||
llmgw ctl audit --since 24h # who changed what
|
||||
llmgw ctl usage export --from 2026-09-01 --to 2026-10-01 --by project,model
|
||||
llmgw ctl kb list # knowledge-base collections, grants, defaults
|
||||
```
|
||||
|
||||
Bare `llmgw` opens the console, which does all of this from one screen.
|
||||
|
||||
`ctl` talks to the admin port (`:9090` by default), not the client-facing one.
|
||||
Add `--json` to any command for machine-readable output.
|
||||
|
||||
|
||||
+893
-235
File diff suppressed because it is too large.
Load diff
@@ -27,9 +27,36 @@ message AgentMessage {
|
||||
FabricScanReply fabric_scan = 5;
|
||||
FabricTestReply fabric_test = 6;
|
||||
UpgradeReply upgrade = 7;
|
||||
KVEvents kv_events = 8;
|
||||
}
|
||||
}
|
||||
|
||||
// KVEvents is a change to what one vLLM replica's prefix cache holds, read
|
||||
// by the agent from the engine's KV cache events and sent on for the router's
|
||||
// precise index (plan §7.5 tier 2).
|
||||
//
|
||||
// A net effect since the last one, not the events themselves: a block stored
|
||||
// and evicted between two sends is not worth the bytes. On every new stream
|
||||
// the agent first sends each replica in full, with clear_first set, because a
|
||||
// control plane that has just started or taken over knows nothing.
|
||||
message KVEvents {
|
||||
string replica_id = 1;
|
||||
// block_size is how many tokens make one of the replica's blocks.
|
||||
int32 block_size = 2;
|
||||
// clear_first says to forget what was known of the replica before applying the
|
||||
// rest: its cache was cleared, its engine restarted, events were lost, or
|
||||
// this is the first piece of a full copy.
|
||||
bool clear_first = 3;
|
||||
// stored and removed are block hashes as vLLM publishes them: the low 64
|
||||
// bits of the sha256_cbor digest.
|
||||
repeated fixed64 stored = 4;
|
||||
repeated fixed64 removed = 5;
|
||||
// problem, when set, is why the replica's events cannot be used — most
|
||||
// likely an engine hashing another way than llmgw configured. Its index is
|
||||
// empty while it stands.
|
||||
string problem = 6;
|
||||
}
|
||||
|
||||
// ControlMessage is anything the control plane tells a node.
|
||||
message ControlMessage {
|
||||
oneof msg {
|
||||
@@ -44,7 +71,48 @@ message ControlMessage {
|
||||
FabricTestRequest fabric_test = 9;
|
||||
UpgradeAgent upgrade = 10;
|
||||
ControlPlanes control_planes = 11;
|
||||
Firewall firewall = 13;
|
||||
RestartAgent restart_agent = 14;
|
||||
AgentSettings agent_settings = 15;
|
||||
}
|
||||
// trace is the W3C trace context (traceparent, and tracestate when there
|
||||
// is one) of the control plane's span that sent this command, so what the
|
||||
// agent does for it — fetching weights, launching the engine, waiting for
|
||||
// it to answer — lands in the same trace. Empty when the control plane is
|
||||
// not tracing, and ignored by an agent that is not.
|
||||
map<string, string> trace = 12;
|
||||
}
|
||||
|
||||
// RestartAgent has the agent start again, as llmgw restart asks of every
|
||||
// node: its engines keep serving and are taken over again.
|
||||
message RestartAgent {
|
||||
string reason = 1;
|
||||
}
|
||||
|
||||
// Firewall is whether llmgw's firewall is kept on the node — its nftables
|
||||
// table, which lets the cluster's networks alone reach the engines' ports and
|
||||
// the node exporter — sent when an agent registers and whenever the
|
||||
// cluster's settings change it (firewall.enabled).
|
||||
message Firewall {
|
||||
bool enabled = 1;
|
||||
// networks are let in besides those the node reaches the control planes
|
||||
// on: firewall_networks, networks or single addresses.
|
||||
repeated string networks = 2;
|
||||
// engine_ports_from and engine_ports_to are the engines' ports
|
||||
// (control.port_range).
|
||||
int32 engine_ports_from = 3;
|
||||
int32 engine_ports_to = 4;
|
||||
}
|
||||
|
||||
// AgentSettings is what the cluster's settings say of this node's agent
|
||||
// (agents.default, with agents.nodes.<id> over it), sent when it registers
|
||||
// and whenever they change. The agent keeps it beside agent.yaml, which
|
||||
// holds only what it needs to start; over that file, it is what the agent
|
||||
// runs on. A change from what it runs on has the agent start again on it,
|
||||
// its engines serving throughout.
|
||||
message AgentSettings {
|
||||
// yaml is agent.yaml's keys, as YAML; empty for none.
|
||||
string yaml = 1;
|
||||
}
|
||||
|
||||
// ControlPlanes is every control plane an agent may dial, sent by the leader
|
||||
@@ -217,6 +285,11 @@ message RunningShard {
|
||||
// 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;
|
||||
|
||||
// compute is where its engine said it put the model when it became ready,
|
||||
// kept by the agent across its own restarts: a control plane that takes
|
||||
// over, or restarts, learns it here.
|
||||
Compute compute = 16;
|
||||
}
|
||||
|
||||
message GPU {
|
||||
@@ -252,6 +325,25 @@ message GPU {
|
||||
string cluster_uuid = 19;
|
||||
string clique_id = 20;
|
||||
string fabric_bandwidth = 21;
|
||||
// mig is set when this is a MIG instance rather than a whole card. index
|
||||
// is then 1000 + 100 × card + device, uuid the instance's (MIG-…), and
|
||||
// memory the instance's share; utilization_percent stays zero, because no
|
||||
// driver measures it per instance.
|
||||
MIGDevice mig = 22;
|
||||
}
|
||||
|
||||
// MIGDevice is which card a MIG instance is carved from, and how.
|
||||
message MIGDevice {
|
||||
int32 parent_index = 1;
|
||||
string parent_uuid = 2;
|
||||
// parent_memory_total_mb is the whole card's: the class is named from it.
|
||||
int32 parent_memory_total_mb = 3;
|
||||
// device is the instance's number on its card, as nvidia-smi -L lists it.
|
||||
int32 device = 4;
|
||||
// profile is NVIDIA's name for the slice, like 3g.40gb.
|
||||
string profile = 5;
|
||||
int32 gpu_instance = 6;
|
||||
int32 compute_instance = 7;
|
||||
}
|
||||
|
||||
message FabricAttachment {
|
||||
@@ -515,6 +607,15 @@ message EngineWant {
|
||||
// variant is the build to install, when one is chosen for the node rather
|
||||
// than left to its hardware: "cu13", "cu129". Empty is the agent's choice.
|
||||
string variant = 3;
|
||||
// offline_base, on a cluster with no way to the internet, is where this
|
||||
// engine version's files are on its mirror:
|
||||
// http://mirror:8081/_llmgw/engines/llamacpp/b10991. Set, the node never
|
||||
// asks the engine's publisher.
|
||||
string offline_base = 4;
|
||||
// offline_sha256 is each of those files' checksum by name, from the
|
||||
// bundle the control plane trusts — not from the mirror, which serves
|
||||
// every node over plain HTTP.
|
||||
map<string, string> offline_sha256 = 5;
|
||||
}
|
||||
|
||||
// ProvisionStatus is where one engine's installation stands on a node.
|
||||
@@ -528,6 +629,18 @@ message ProvisionStatus {
|
||||
string detail = 5;
|
||||
}
|
||||
|
||||
// Compute is where an engine runs its model, as it said when it loaded it.
|
||||
message Compute {
|
||||
// backend is "CUDA", "Vulkan", "Metal" or "CPU".
|
||||
string backend = 1;
|
||||
// devices are the devices it named: "CUDA0: Quadro RTX 6000 (7.5)".
|
||||
repeated string devices = 2;
|
||||
// layers_offloaded of layers_total went to the devices; both 0 when the
|
||||
// engine did not say.
|
||||
int32 layers_offloaded = 3;
|
||||
int32 layers_total = 4;
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
message ShardStatus {
|
||||
string node_id = 1;
|
||||
@@ -549,6 +662,10 @@ message ShardStatus {
|
||||
// going to look.
|
||||
FetchProgress fetch = 8;
|
||||
|
||||
// compute is what the engine says it computes on, read from its own output
|
||||
// as it became ready: sent with STATE_READY when the engine says.
|
||||
Compute compute = 9;
|
||||
|
||||
enum State {
|
||||
STATE_UNSPECIFIED = 0;
|
||||
STATE_PENDING = 1;
|
||||
@@ -625,6 +742,9 @@ message EngineProfile {
|
||||
// adapters are LoRA adapters loaded beside the model, in order: the index
|
||||
// is the id a request applies one by.
|
||||
repeated LoraAdapter adapters = 12;
|
||||
// kv_events has a vLLM engine publish its KV cache events to the agent, for
|
||||
// the router's precise index.
|
||||
bool kv_events = 13;
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
|
||||
@@ -0,0 +1,467 @@
|
||||
// Command llmgw-kb puts a person's documents into an llmgw knowledge base
|
||||
// from their own machine (BACKLOG F2): it reads the files there, cuts them
|
||||
// into chunks, has the gateway embed them with the person's key, and sends
|
||||
// the text and vectors — only for what changed since the last run, which a
|
||||
// ledger on the machine remembers.
|
||||
//
|
||||
// Its own binary, for Windows, macOS and Linux: it carries none of the
|
||||
// server, agent or setup code that llmgw does.
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"golang.org/x/term"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/extract"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||
)
|
||||
|
||||
func main() {
|
||||
if err := root().Execute(); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "error:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
// config is what `login` keeps: the gateway, the key, and for a personal key
|
||||
// the project it acts in.
|
||||
type config struct {
|
||||
Gateway string `json:"gateway"`
|
||||
Key string `json:"key"`
|
||||
Project string `json:"project,omitempty"`
|
||||
CA string `json:"ca,omitempty"`
|
||||
}
|
||||
|
||||
func configPath() (string, error) {
|
||||
dir, err := os.UserConfigDir()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Join(dir, "llmgw-kb", "config.json"), nil
|
||||
}
|
||||
|
||||
// loadConfig reads the saved configuration, with the environment over it:
|
||||
// LLMGW_URL, LLMGW_API_KEY, LLMGW_PROJECT.
|
||||
func loadConfig() (config, error) {
|
||||
var c config
|
||||
if p, err := configPath(); err == nil {
|
||||
if b, err := os.ReadFile(p); err == nil {
|
||||
if err := json.Unmarshal(b, &c); err != nil {
|
||||
return c, fmt.Errorf("%s: %w", p, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
for env, dst := range map[string]*string{"LLMGW_URL": &c.Gateway, "LLMGW_API_KEY": &c.Key, "LLMGW_PROJECT": &c.Project} {
|
||||
if v := strings.TrimSpace(os.Getenv(env)); v != "" {
|
||||
*dst = v
|
||||
}
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func client(flags *globalFlags) (*kbsync.Client, error) {
|
||||
c, err := loadConfig()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if flags.gateway != "" {
|
||||
c.Gateway = flags.gateway
|
||||
}
|
||||
if flags.project != "" {
|
||||
c.Project = flags.project
|
||||
}
|
||||
if flags.ca != "" {
|
||||
c.CA = flags.ca
|
||||
}
|
||||
if c.Gateway == "" || c.Key == "" {
|
||||
return nil, errors.New("no gateway or key: run `llmgw-kb login`, or set LLMGW_URL and LLMGW_API_KEY")
|
||||
}
|
||||
h, err := kbsync.NewHTTP(c.CA)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &kbsync.Client{Base: c.Gateway, Key: c.Key, Project: c.Project, HTTP: h}, nil
|
||||
}
|
||||
|
||||
type globalFlags struct {
|
||||
gateway, project, ca, ledger string
|
||||
}
|
||||
|
||||
func (g *globalFlags) openLedger(ctx context.Context) (*kbsync.Ledger, error) {
|
||||
p := g.ledger
|
||||
if p == "" {
|
||||
var err error
|
||||
if p, err = kbsync.DefaultLedgerPath(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return kbsync.OpenLedger(ctx, p)
|
||||
}
|
||||
|
||||
func root() *cobra.Command {
|
||||
g := &globalFlags{}
|
||||
cmd := &cobra.Command{
|
||||
Use: "llmgw-kb",
|
||||
Short: "Put your documents into an llmgw knowledge base, from this machine",
|
||||
Long: `llmgw-kb reads documents on this machine, cuts them into chunks, has the
|
||||
gateway embed them with your key, and sends the text and vectors to a
|
||||
knowledge base. Files are never uploaded. Run it as often as you like: a
|
||||
ledger on this machine remembers what was sent, and only what changed is
|
||||
read and sent again.`,
|
||||
Version: version.Version,
|
||||
SilenceUsage: true,
|
||||
SilenceErrors: true,
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&g.gateway, "gateway", "", "the gateway's address (default: from login, or LLMGW_URL)")
|
||||
cmd.PersistentFlags().StringVar(&g.project, "project", "", "the project a personal key acts in (default: from login, or LLMGW_PROJECT)")
|
||||
cmd.PersistentFlags().StringVar(&g.ca, "ca", "", "a PEM file of certificates to trust besides the system's")
|
||||
cmd.PersistentFlags().StringVar(&g.ledger, "ledger", "", "the ledger file (default: in your configuration directory)")
|
||||
cmd.AddCommand(loginCmd(g), collectionsCmd(g), syncCmd(g), statusCmd(g), historyCmd(g), formatsCmd())
|
||||
return cmd
|
||||
}
|
||||
|
||||
func loginCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "login",
|
||||
Short: "Save the gateway's address and your key",
|
||||
Long: `Asks for the gateway's address and your API key, checks them, and saves them
|
||||
in your configuration directory, readable by you only. The key is typed
|
||||
without being shown, so it does not end up in your shell's history.`,
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
c, _ := loadConfig()
|
||||
in := bufio.NewReader(os.Stdin)
|
||||
ask := func(q, def string) string {
|
||||
if def != "" {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s [%s]: ", q, def)
|
||||
} else {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: ", q)
|
||||
}
|
||||
s, _ := in.ReadString('\n')
|
||||
if s = strings.TrimSpace(s); s == "" {
|
||||
return def
|
||||
}
|
||||
return s
|
||||
}
|
||||
c.Gateway = strings.TrimRight(ask("Gateway address (https://…)", or(g.gateway, c.Gateway)), "/")
|
||||
fmt.Fprint(cmd.OutOrStdout(), "API key: ")
|
||||
var key string
|
||||
if term.IsTerminal(int(os.Stdin.Fd())) {
|
||||
b, err := term.ReadPassword(int(os.Stdin.Fd()))
|
||||
fmt.Fprintln(cmd.OutOrStdout())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key = strings.TrimSpace(string(b))
|
||||
} else {
|
||||
key, _ = in.ReadString('\n')
|
||||
key = strings.TrimSpace(key)
|
||||
}
|
||||
if key != "" {
|
||||
c.Key = key
|
||||
}
|
||||
c.Project = ask("Project, for a personal key (empty for a project's key)", or(g.project, c.Project))
|
||||
if g.ca != "" {
|
||||
c.CA = g.ca
|
||||
}
|
||||
h, err := kbsync.NewHTTP(c.CA)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cl := &kbsync.Client{Base: c.Gateway, Key: c.Key, Project: c.Project, HTTP: h, Wait: time.Second}
|
||||
ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
|
||||
defer cancel()
|
||||
cs, err := cl.Collections(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("the gateway refused: %w", err)
|
||||
}
|
||||
p, err := configPath()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0o700); err != nil {
|
||||
return err
|
||||
}
|
||||
b, _ := json.MarshalIndent(c, "", " ")
|
||||
if err := os.WriteFile(p, b, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "Saved in %s. This key reaches %d collection(s); `llmgw-kb collections` lists them.\n", p, len(cs))
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func or(a, b string) string {
|
||||
if a != "" {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func collectionsCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "collections",
|
||||
Short: "List the collections your key may read, and which it may add to",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cs, err := c.Collections(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(cs) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), "No collections. An administrator creates them and grants access (llmgw ctl kb).")
|
||||
return nil
|
||||
}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "NAME\tACCESS\tOWNER\tEMBEDDED BY")
|
||||
for _, x := range cs {
|
||||
access := "read"
|
||||
if x.Writable {
|
||||
access = "read, write"
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%s\t%s (%d)\n", x.Name, access, or(x.OwnerName, x.Owner), x.EmbedModel, x.Dims)
|
||||
}
|
||||
return tw.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func syncCmd(g *globalFlags) *cobra.Command {
|
||||
var o kbsync.Options
|
||||
var verbose bool
|
||||
var ocr string
|
||||
var maxMiB int64
|
||||
cmd := &cobra.Command{
|
||||
Use: "sync COLLECTION PATH...",
|
||||
Short: "Send a collection what changed in these folders and files",
|
||||
Long: `Reads the files under each PATH that the collection does not already have as
|
||||
they are, and sends their text and vectors. Run it again whenever you like;
|
||||
unchanged files are not even read.
|
||||
|
||||
Documents are named by their path under PATH, with forward slashes; with
|
||||
several PATHs, each folder's own name comes first.
|
||||
|
||||
--dry-run reads and cuts everything, sends nothing, and says what would
|
||||
leave this machine
|
||||
--delete removes documents this machine sent whose files are gone
|
||||
(only under these PATHs; nothing anybody else sent)
|
||||
--retry reads again files skipped before (after installing LibreOffice
|
||||
or Tesseract, say)
|
||||
|
||||
Without --include, every file of a known format is read (llmgw-kb formats);
|
||||
hidden files, version control folders and Office lock files are not.`,
|
||||
Args: cobra.MinimumNArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), os.Interrupt)
|
||||
defer stop()
|
||||
l, err := g.openLedger(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
o.Collection, o.Roots = args[0], args[1:]
|
||||
o.Extract = extract.Options{Helpers: extract.FindHelpers(), OCRLanguages: ocr, MaxBytes: maxMiB << 20}
|
||||
out := cmd.OutOrStdout()
|
||||
c.Waiting = func(why string, d time.Duration) {
|
||||
fmt.Fprintf(out, " .. waiting %s: %s\n", d, why)
|
||||
}
|
||||
start := time.Now()
|
||||
rep, err := kbsync.Sync(ctx, c, l, o, func(e kbsync.Event) {
|
||||
switch e.Kind {
|
||||
case "unchanged":
|
||||
if !verbose {
|
||||
return
|
||||
}
|
||||
fmt.Fprintf(out, " = %s\n", e.Name)
|
||||
case "sent", "would-send":
|
||||
fmt.Fprintf(out, " %-8s %s (%d chunks, %s)\n", label(e.Kind), e.Name, e.Chunks, size(e.Bytes))
|
||||
default:
|
||||
line := fmt.Sprintf(" %-8s %s", label(e.Kind), e.Name)
|
||||
if e.Detail != "" {
|
||||
line += ": " + e.Detail
|
||||
}
|
||||
fmt.Fprintln(out, line)
|
||||
}
|
||||
})
|
||||
if rep.Collection.Name != "" {
|
||||
verb := "sent"
|
||||
if o.DryRun {
|
||||
verb = "would send"
|
||||
}
|
||||
fmt.Fprintf(out, "\n%s: %d files; %s %d (%d chunks, %s), %d unchanged, %d skipped, %d failed, %d deleted, in %s\n",
|
||||
rep.Collection.Name, rep.Files, verb, rep.Sent, rep.Chunks, size(rep.Bytes), rep.Unchanged,
|
||||
rep.Skipped, rep.Failed, rep.Deleted, time.Since(start).Round(time.Second))
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rep.Failed > 0 {
|
||||
return fmt.Errorf("%d document(s) failed; run again to retry them", rep.Failed)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
f := cmd.Flags()
|
||||
f.StringSliceVar(&o.Include, "include", nil, "read only files matching these globs (\"*.pdf,docs/**/*.md\")")
|
||||
f.StringSliceVar(&o.Exclude, "exclude", nil, "never read files matching these globs")
|
||||
f.BoolVar(&o.Hidden, "hidden", false, "read hidden files and folders too")
|
||||
f.BoolVar(&o.DryRun, "dry-run", false, "send nothing; say what would be sent")
|
||||
f.BoolVar(&o.Delete, "delete", false, "remove documents this machine sent whose files are gone")
|
||||
f.BoolVar(&o.Retry, "retry", false, "read again files skipped before")
|
||||
f.IntVar(&o.Workers, "workers", 4, "files read and sent at once")
|
||||
f.StringVar(&ocr, "ocr-languages", "tur+eng", "Tesseract's languages for images and scanned pages")
|
||||
f.Int64Var(&maxMiB, "max-size", 256, "skip files larger than this many MiB")
|
||||
f.BoolVarP(&verbose, "verbose", "v", false, "list unchanged files too")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func label(kind string) string {
|
||||
switch kind {
|
||||
case "would-send":
|
||||
return "send"
|
||||
case "would-delete":
|
||||
return "delete"
|
||||
}
|
||||
return kind
|
||||
}
|
||||
|
||||
func size(b int64) string {
|
||||
switch {
|
||||
case b >= 1<<20:
|
||||
return fmt.Sprintf("%.1f MiB", float64(b)/(1<<20))
|
||||
case b >= 1<<10:
|
||||
return fmt.Sprintf("%.0f KiB", float64(b)/(1<<10))
|
||||
}
|
||||
return fmt.Sprintf("%d B", b)
|
||||
}
|
||||
|
||||
func statusCmd(g *globalFlags) *cobra.Command {
|
||||
var all bool
|
||||
return &cobra.Command{
|
||||
Use: "status COLLECTION",
|
||||
Short: "Show what this machine has sent to a collection, and what it skipped",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
m, err := c.Manifest(cmd.Context(), args[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
l, err := g.openLedger(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
es, err := l.Entries(cmd.Context(), strings.TrimRight(c.Base, "/"), m.Collection.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
counts := map[kbsync.State]int{}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "STATE\tDOCUMENT\tCHUNKS\tWHEN\tWHY")
|
||||
for _, e := range es {
|
||||
counts[e.State]++
|
||||
if e.State == kbsync.Sent && !all {
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%d\t%s\t%s\n", e.State, e.Name, e.Chunks, e.At.Format("2006-01-02 15:04"), e.Reason)
|
||||
}
|
||||
_ = tw.Flush()
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n%s holds %d documents; from this machine %d sent, %d skipped, %d failed.",
|
||||
m.Collection.Name, len(m.Documents), counts[kbsync.Sent], counts[kbsync.Skipped], counts[kbsync.Failed])
|
||||
if !all {
|
||||
fmt.Fprint(cmd.OutOrStdout(), " (--all lists the sent ones too)")
|
||||
}
|
||||
fmt.Fprintln(cmd.OutOrStdout())
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func historyCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "history",
|
||||
Short: "List the latest syncs from this machine",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
l, err := g.openLedger(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
runs, err := l.Runs(cmd.Context(), 20)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "STARTED\tTOOK\tCOLLECTION\tSENT\tUNCHANGED\tSKIPPED\tFAILED\tDELETED")
|
||||
for _, r := range runs {
|
||||
took := "unfinished"
|
||||
if !r.Finished.IsZero() {
|
||||
took = r.Finished.Sub(r.Started).Round(time.Second).String()
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%s\t%d\t%d\t%d\t%d\t%d\n", r.Started.Format("2006-01-02 15:04"), took,
|
||||
r.Collection, r.Sent, r.Unchanged, r.Skipped, r.Failed, r.Delete)
|
||||
}
|
||||
return tw.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func formatsCmd() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "formats",
|
||||
Short: "List the formats read, and the helpers found on this machine",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
builtIn, helper := extract.Supported()
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintf(out, "Read here:\n%s\n\n", wrap(builtIn))
|
||||
fmt.Fprintf(out, "Read with a helper program:\n%s\n\n", wrap(helper))
|
||||
fmt.Fprintln(out, "Helpers:")
|
||||
for _, l := range extract.FindHelpers().Describe() {
|
||||
fmt.Fprintln(out, " "+l)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func wrap(words []string) string {
|
||||
var b strings.Builder
|
||||
line := " "
|
||||
for _, w := range words {
|
||||
if len(line)+len(w) > 78 {
|
||||
b.WriteString(strings.TrimRight(line, " ") + "\n")
|
||||
line = " "
|
||||
}
|
||||
line += w + " "
|
||||
}
|
||||
b.WriteString(strings.TrimRight(line, " "))
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The client goes onto employees' laptops: none of the server, agent, store
|
||||
// or setup code may come with it, and it must build where llmgw does not.
|
||||
func TestTheClientCarriesNoServerCode(t *testing.T) {
|
||||
out, err := exec.Command("go", "list", "-deps", ".").Output()
|
||||
if err != nil {
|
||||
t.Skipf("go list: %v", err)
|
||||
}
|
||||
for _, dep := range strings.Fields(string(out)) {
|
||||
for _, banned := range []string{"/internal/store", "/internal/router", "/internal/agent", "/internal/control",
|
||||
"/internal/cli", "/internal/guard", "github.com/jackc/pgx"} {
|
||||
if strings.Contains(dep, banned) {
|
||||
t.Errorf("llmgw-kb depends on %s", dep)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,8 @@ package main
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
// Time zones for scaling schedules, on a machine without the tz database.
|
||||
_ "time/tzdata"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/cli"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
# The llmgw control plane as a container image, for deploy/helm/llmgw.
|
||||
#
|
||||
# The project publishes binaries, not an image, so this is built where the
|
||||
# image is to live:
|
||||
#
|
||||
# docker build -f deploy/container/Dockerfile --build-arg VERSION=v0.11.1 \
|
||||
# -t registry.internal/llmgw:v0.11.1 .
|
||||
# docker push registry.internal/llmgw:v0.11.1
|
||||
#
|
||||
# The binary is static (CGO_ENABLED=0, as every release is), so the image is
|
||||
# that one file on a base with nothing else in it: no shell, no package
|
||||
# manager. `kubectl exec ... -- llmgw ctl status` still works, because it
|
||||
# runs the binary rather than a shell.
|
||||
#
|
||||
# Only the control plane runs from it. An agent installs engines on its
|
||||
# machine, manages its GPUs and supervises processes there as root, which is
|
||||
# what a container is built to prevent: docs/kubernetes.md.
|
||||
|
||||
FROM golang:1.26 AS build
|
||||
WORKDIR /src
|
||||
COPY go.mod go.sum ./
|
||||
RUN go mod download
|
||||
COPY . .
|
||||
ARG VERSION=dev
|
||||
RUN CGO_ENABLED=0 go build -trimpath \
|
||||
-ldflags "-s -w -X gitea.mynodes.xyz/ilker/llmgw/internal/version.Version=${VERSION}" \
|
||||
-o /out/llmgw ./cmd/llmgw && \
|
||||
mkdir -p /out/state
|
||||
|
||||
FROM gcr.io/distroless/static-debian12:nonroot
|
||||
COPY --from=build /out/llmgw /usr/local/bin/llmgw
|
||||
# The chart mounts server.yaml, the certificates and the token under
|
||||
# /etc/llmgw, and the SQLite store, when it is one, under /var/lib/llmgw —
|
||||
# owned by the user the server runs as, so a volume Docker fills from the
|
||||
# image is writable without Kubernetes' fsGroup.
|
||||
COPY --from=build --chown=65532:65532 /out/state /var/lib/llmgw
|
||||
USER 65532:65532
|
||||
EXPOSE 8080 7443 9091
|
||||
ENTRYPOINT ["/usr/local/bin/llmgw"]
|
||||
CMD ["server", "--config", "/etc/llmgw/server.yaml"]
|
||||
@@ -0,0 +1,58 @@
|
||||
# Arithmetic a small instruction model gets right: sums, products and a
|
||||
# little reading, each answer scored by the last number in it. A quantisation
|
||||
# that has gone wrong, or a chat template that no longer matches the
|
||||
# weights, shows up here as a falling score long before anyone notices it in
|
||||
# conversation.
|
||||
#
|
||||
# llmgw ctl eval MODEL --suite arithmetic
|
||||
# llmgw ctl model update MODEL --eval arithmetic --min-score 0.8
|
||||
#
|
||||
# Copy it into control.eval_dir (/etc/llmgw/eval) to change it; a suite
|
||||
# there of the same name replaces this one.
|
||||
name: arithmetic
|
||||
description: Sums, products and word problems, each scored by the number in its answer.
|
||||
threshold: 0.7
|
||||
system: You are a careful calculator. Answer with the number only, without words or units.
|
||||
max_tokens: 64
|
||||
cases:
|
||||
- id: add-small
|
||||
prompt: What is 17 + 25?
|
||||
number: 42
|
||||
- id: add-carry
|
||||
prompt: What is 389 + 476?
|
||||
number: 865
|
||||
- id: subtract
|
||||
prompt: What is 1000 - 357?
|
||||
number: 643
|
||||
- id: subtract-negative
|
||||
prompt: What is 12 - 30?
|
||||
number: -18
|
||||
- id: multiply
|
||||
prompt: What is 12 * 12?
|
||||
number: 144
|
||||
- id: multiply-larger
|
||||
prompt: What is 23 * 17?
|
||||
number: 391
|
||||
weight: 2
|
||||
- id: divide
|
||||
prompt: What is 144 / 6?
|
||||
number: 24
|
||||
- id: decimal
|
||||
prompt: What is 2.5 * 4?
|
||||
number: 10
|
||||
- id: percent
|
||||
prompt: What is 15% of 200?
|
||||
number: 30
|
||||
- id: words-apples
|
||||
prompt: Ana has 7 apples and buys 5 more, then gives 3 away. How many apples does she have?
|
||||
number: 9
|
||||
weight: 2
|
||||
- id: words-time
|
||||
prompt: A train leaves at 14:00 and the trip takes 150 minutes. How many hours does the trip take?
|
||||
number: 2.5
|
||||
tolerance: 0.01
|
||||
weight: 2
|
||||
- id: order-of-operations
|
||||
prompt: What is 2 + 3 * 4?
|
||||
number: 14
|
||||
weight: 2
|
||||
@@ -0,0 +1,94 @@
|
||||
# Following a format: asked for JSON of a given shape, does the model give
|
||||
# it — the thing an application parsing its answers depends on, and the
|
||||
# first to go when a chat template drifts. Each answer is read as JSON (a
|
||||
# code fence or a sentence around it is forgiven) and held to a schema.
|
||||
#
|
||||
# llmgw ctl eval MODEL --suite json-format
|
||||
#
|
||||
# Copy it into control.eval_dir (/etc/llmgw/eval) to change it; a suite
|
||||
# there of the same name replaces this one.
|
||||
name: json-format
|
||||
description: Answers asked for as JSON of a given shape, held to a schema.
|
||||
threshold: 0.7
|
||||
system: You answer with a single JSON value and nothing else.
|
||||
max_tokens: 200
|
||||
cases:
|
||||
- id: person
|
||||
prompt: >-
|
||||
Give a JSON object for a person named Ada Lovelace, born in 1815, with the
|
||||
keys "name" (a string) and "born" (a number).
|
||||
json_schema:
|
||||
type: object
|
||||
required: [name, born]
|
||||
properties:
|
||||
name: {type: string, pattern: "Ada"}
|
||||
born: {type: integer, minimum: 1815, maximum: 1815}
|
||||
|
||||
- id: list-of-colours
|
||||
prompt: Give a JSON array of exactly three colour names, as strings.
|
||||
json_schema:
|
||||
type: array
|
||||
minItems: 3
|
||||
maxItems: 3
|
||||
items: {type: string, minLength: 1}
|
||||
|
||||
- id: boolean-answer
|
||||
prompt: >-
|
||||
Is water wet? Answer as a JSON object with one key, "answer", whose
|
||||
value is true or false.
|
||||
json_schema:
|
||||
type: object
|
||||
required: [answer]
|
||||
properties:
|
||||
answer: {type: boolean}
|
||||
|
||||
- id: nested
|
||||
prompt: >-
|
||||
Give a JSON object describing a book: "title" (a string), "author" (an
|
||||
object with "first" and "last", both strings) and "pages" (a whole number
|
||||
above zero). Use any book.
|
||||
json_schema:
|
||||
type: object
|
||||
required: [title, author, pages]
|
||||
properties:
|
||||
title: {type: string, minLength: 1}
|
||||
author:
|
||||
type: object
|
||||
required: [first, last]
|
||||
properties:
|
||||
first: {type: string}
|
||||
last: {type: string}
|
||||
pages: {type: integer, exclusiveMinimum: 0}
|
||||
weight: 2
|
||||
|
||||
- id: enum-choice
|
||||
prompt: >-
|
||||
Classify the sentence "I love this product!" as positive, negative or
|
||||
neutral. Answer as a JSON object {"sentiment": "..."} using one of those
|
||||
three words in lower case.
|
||||
json_schema:
|
||||
type: object
|
||||
required: [sentiment]
|
||||
properties:
|
||||
sentiment: {enum: [positive]}
|
||||
weight: 2
|
||||
|
||||
- id: numbers
|
||||
prompt: 'Give a JSON object {"sum": N} where N is the sum of 4, 5 and 6.'
|
||||
json_schema:
|
||||
type: object
|
||||
required: [sum]
|
||||
properties:
|
||||
sum: {const: 15}
|
||||
|
||||
- id: no-extra-keys
|
||||
prompt: >-
|
||||
Give a JSON object with exactly two keys, "city" and "country", for
|
||||
Paris. No other keys.
|
||||
json_schema:
|
||||
type: object
|
||||
required: [city, country]
|
||||
additionalProperties: false
|
||||
properties:
|
||||
city: {type: string}
|
||||
country: {type: string}
|
||||
@@ -17,8 +17,10 @@ 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.
|
||||
# is refused. Without it the port would be plaintext HTTP/2 and anyone who can
|
||||
# reach it could claim to be any node, so the control plane refuses to start
|
||||
# unless `agent_insecure: true` is set — and then listens on loopback unless
|
||||
# agent_listen names an address. That is for one machine or a test.
|
||||
#
|
||||
# These certificates come from llmgw's own authority, not a public one: public
|
||||
# authorities stopped issuing the clientAuth usage, so no Let's Encrypt
|
||||
@@ -157,6 +159,12 @@ router:
|
||||
# back to a client that did not ask for it.
|
||||
inject_usage: true
|
||||
max_body_bytes: 33554432
|
||||
# Memory for the answers of projects that reuse them (llmgw keys cache),
|
||||
# all together; the least recently used go first. Nothing is reused until a
|
||||
# project turns it on.
|
||||
# cache:
|
||||
# max_bytes: 134217728 # 128 MB
|
||||
# max_entry_bytes: 1048576 # 1 MB
|
||||
|
||||
timeouts:
|
||||
dial: 5s
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
apiVersion: v2
|
||||
name: llmgw
|
||||
description: >-
|
||||
The llmgw control plane — the OpenAI-compatible gateway, the router and the
|
||||
scheduler that places models on GPU and CPU nodes. The nodes themselves run
|
||||
outside Kubernetes; see docs/kubernetes.md.
|
||||
type: application
|
||||
# The chart's own version; appVersion is the llmgw release the image holds.
|
||||
version: 0.1.0
|
||||
appVersion: "v0.11.1"
|
||||
home: https://gitea.mynodes.xyz/ilker/llmgw
|
||||
sources:
|
||||
- https://gitea.mynodes.xyz/ilker/llmgw
|
||||
keywords:
|
||||
- llm
|
||||
- gateway
|
||||
- inference
|
||||
@@ -0,0 +1,27 @@
|
||||
The llmgw control plane is starting as {{ include "llmgw.fullname" . }}.
|
||||
|
||||
The admin API listens on the pod's loopback only. To reach it from here:
|
||||
|
||||
kubectl -n {{ .Release.Namespace }} port-forward statefulset/{{ include "llmgw.fullname" . }} 9090:9090 &
|
||||
export LLMGW_ADMIN_TOKEN=$(kubectl -n {{ .Release.Namespace }} get secret {{ include "llmgw.adminSecret" . }} \
|
||||
-o jsonpath='{.data.{{ if .Values.adminToken.existingSecret }}{{ .Values.adminToken.secretKey }}{{ else }}admin-token{{ end }}}' | base64 -d)
|
||||
llmgw ctl status
|
||||
|
||||
The first administrator:
|
||||
|
||||
kubectl -n {{ .Release.Namespace }} exec -it statefulset/{{ include "llmgw.fullname" . }} -- \
|
||||
llmgw users create you@example.com --password '<a password>'
|
||||
kubectl -n {{ .Release.Namespace }} exec -it statefulset/{{ include "llmgw.fullname" . }} -- \
|
||||
llmgw users role you@example.com admin
|
||||
|
||||
Agents connect to the {{ include "llmgw.fullname" . }}-agents Service on port {{ .Values.agentService.port }}:
|
||||
|
||||
kubectl -n {{ .Release.Namespace }} get service {{ include "llmgw.fullname" . }}-agents
|
||||
|
||||
Its address must be in the certificate in {{ .Values.agentTLS.existingSecret }}. Nodes join from
|
||||
the machine holding llmgw's authority:
|
||||
|
||||
sudo llmgw node add ubuntu@gpu01 --server https://<that address>:{{ .Values.agentService.port }} \
|
||||
--pki ./pki --admin 127.0.0.1:9090
|
||||
|
||||
docs/kubernetes.md has the rest, and why the nodes stay outside Kubernetes.
|
||||
@@ -0,0 +1,74 @@
|
||||
{{/* The release's name for everything the chart makes. */}}
|
||||
{{- define "llmgw.fullname" -}}
|
||||
{{- if contains .Chart.Name .Release.Name -}}
|
||||
{{- .Release.Name | trunc 63 | trimSuffix "-" -}}
|
||||
{{- else -}}
|
||||
{{- printf "%s-%s" .Release.Name .Chart.Name | trunc 63 | trimSuffix "-" -}}
|
||||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "llmgw.selectorLabels" -}}
|
||||
app.kubernetes.io/name: {{ .Chart.Name }}
|
||||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "llmgw.labels" -}}
|
||||
{{ include "llmgw.selectorLabels" . }}
|
||||
app.kubernetes.io/version: {{ .Values.image.tag | default .Chart.AppVersion | quote }}
|
||||
app.kubernetes.io/managed-by: {{ .Release.Service }}
|
||||
helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "llmgw.image" -}}
|
||||
{{- if not .Values.image.repository -}}
|
||||
{{- fail "image.repository: the image built from deploy/container/Dockerfile, pushed where the cluster can pull it (the project publishes binaries, not an image)" -}}
|
||||
{{- end -}}
|
||||
{{- printf "%s:%s" .Values.image.repository (.Values.image.tag | default .Chart.AppVersion) -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "llmgw.adminSecret" -}}
|
||||
{{- .Values.adminToken.existingSecret | default (printf "%s-admin" (include "llmgw.fullname" .)) -}}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
server.yaml: what the chart decides — listeners, where the certificates,
|
||||
the token and the store are mounted — laid over .Values.config. The chart's
|
||||
keys win, because each names a port or a path the pod is built around.
|
||||
*/}}
|
||||
{{- define "llmgw.serverConfig" -}}
|
||||
{{- $store := "/var/lib/llmgw/llmgw.db" -}}
|
||||
{{- if eq .Values.store.type "postgres" -}}
|
||||
{{- $store = "@/etc/llmgw/store/dsn" -}}
|
||||
{{- end -}}
|
||||
{{- $base := dict
|
||||
"listen" ":8080"
|
||||
"admin_listen" "127.0.0.1:9090"
|
||||
"metrics_listen" ":9091"
|
||||
"agent_listen" ":7443"
|
||||
"admin_token_file" "/etc/llmgw/admin/admin-token"
|
||||
"agent_tls" (dict "ca" "/etc/llmgw/tls/ca.pem" "cert" "/etc/llmgw/tls/server.pem" "key" "/etc/llmgw/tls/server-key.pem")
|
||||
"auth" (dict "store" $store)
|
||||
-}}
|
||||
{{- if .Values.offlineChecksums.configMap -}}
|
||||
{{- $_ := set $base "offline" (dict "checksums" "/etc/llmgw/offline/offline.sha256") -}}
|
||||
{{- end -}}
|
||||
{{- $cfg := mustMergeOverwrite (deepCopy (.Values.config | default dict)) $base -}}
|
||||
# Made by the llmgw Helm chart from its values. Every key: docs/reference.md
|
||||
{{ toYaml $cfg }}
|
||||
{{- end -}}
|
||||
|
||||
{{/* Refuses values that would make a pod that cannot work. */}}
|
||||
{{- define "llmgw.validate" -}}
|
||||
{{- if not (has .Values.store.type (list "sqlite" "postgres")) -}}
|
||||
{{- fail (printf "store.type %q: sqlite or postgres" .Values.store.type) -}}
|
||||
{{- end -}}
|
||||
{{- if ne (int .Values.replicas) 1 -}}
|
||||
{{- fail "replicas: this chart runs one control plane; values.yaml says why" -}}
|
||||
{{- end -}}
|
||||
{{- if and (eq .Values.store.type "postgres") (not .Values.store.postgres.existingSecret) -}}
|
||||
{{- fail "store.postgres.existingSecret: a Secret holding the PostgreSQL address" -}}
|
||||
{{- end -}}
|
||||
{{- if not .Values.agentTLS.existingSecret -}}
|
||||
{{- fail "agentTLS.existingSecret: a Secret with ca.pem, server.pem and server-key.pem from llmgw's authority (see values.yaml)" -}}
|
||||
{{- end -}}
|
||||
{{- end -}}
|
||||
@@ -0,0 +1,10 @@
|
||||
{{- include "llmgw.validate" . -}}
|
||||
apiVersion: v1
|
||||
kind: ConfigMap
|
||||
metadata:
|
||||
name: {{ include "llmgw.fullname" . }}
|
||||
labels:
|
||||
{{- include "llmgw.labels" . | nindent 4 }}
|
||||
data:
|
||||
server.yaml: |
|
||||
{{- include "llmgw.serverConfig" . | nindent 4 }}
|
||||
@@ -0,0 +1,37 @@
|
||||
{{- if .Values.ingress.enabled }}
|
||||
# The OpenAI-compatible API, the portal and the chat page on port 8080. Never
|
||||
# the agent port: its mutual TLS cannot pass through an HTTP proxy.
|
||||
apiVersion: networking.k8s.io/v1
|
||||
kind: Ingress
|
||||
metadata:
|
||||
name: {{ include "llmgw.fullname" . }}
|
||||
labels:
|
||||
{{- include "llmgw.labels" . | nindent 4 }}
|
||||
{{- with .Values.ingress.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
{{- with .Values.ingress.className }}
|
||||
ingressClassName: {{ . }}
|
||||
{{- end }}
|
||||
{{- with .Values.ingress.tls }}
|
||||
tls:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
rules:
|
||||
{{- range .Values.ingress.hosts }}
|
||||
- host: {{ .host | quote }}
|
||||
http:
|
||||
paths:
|
||||
{{- range .paths }}
|
||||
- path: {{ .path }}
|
||||
pathType: {{ .pathType }}
|
||||
backend:
|
||||
service:
|
||||
name: {{ include "llmgw.fullname" $ }}
|
||||
port:
|
||||
name: api
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,19 @@
|
||||
{{- if not .Values.adminToken.existingSecret }}
|
||||
{{- $name := include "llmgw.adminSecret" . }}
|
||||
{{- $token := randAlphaNum 64 }}
|
||||
{{- /* Kept across upgrades: a new token on every upgrade would lock out every console holding the old one. */}}
|
||||
{{- with lookup "v1" "Secret" .Release.Namespace $name }}
|
||||
{{- $token = index .data "admin-token" | b64dec }}
|
||||
{{- end }}
|
||||
apiVersion: v1
|
||||
kind: Secret
|
||||
metadata:
|
||||
name: {{ $name }}
|
||||
labels:
|
||||
{{- include "llmgw.labels" . | nindent 4 }}
|
||||
annotations:
|
||||
helm.sh/resource-policy: keep
|
||||
type: Opaque
|
||||
data:
|
||||
admin-token: {{ $token | b64enc | quote }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,21 @@
|
||||
{{- if .Values.serviceMonitor.enabled }}
|
||||
# For the Prometheus Operator: the metrics port, which serves metrics and
|
||||
# nothing else. deploy/monitoring holds the rules and dashboards that read them.
|
||||
apiVersion: monitoring.coreos.com/v1
|
||||
kind: ServiceMonitor
|
||||
metadata:
|
||||
name: {{ include "llmgw.fullname" . }}
|
||||
labels:
|
||||
{{- include "llmgw.labels" . | nindent 4 }}
|
||||
{{- with .Values.serviceMonitor.labels }}
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "llmgw.selectorLabels" . | nindent 6 }}
|
||||
endpoints:
|
||||
- port: metrics
|
||||
path: /metrics
|
||||
interval: {{ .Values.serviceMonitor.interval }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,70 @@
|
||||
# The API and the metrics, for the cluster and an Ingress.
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "llmgw.fullname" . }}
|
||||
labels:
|
||||
{{- include "llmgw.labels" . | nindent 4 }}
|
||||
{{- with .Values.service.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
type: {{ .Values.service.type }}
|
||||
selector:
|
||||
{{- include "llmgw.selectorLabels" . | nindent 4 }}
|
||||
ports:
|
||||
- name: api
|
||||
port: {{ .Values.service.port }}
|
||||
targetPort: api
|
||||
- name: metrics
|
||||
port: {{ .Values.service.metricsPort }}
|
||||
targetPort: metrics
|
||||
---
|
||||
# Where agents on the nodes connect. Mutual TLS ends in the pod: whatever is
|
||||
# in between must pass TCP through untouched.
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "llmgw.fullname" . }}-agents
|
||||
labels:
|
||||
{{- include "llmgw.labels" . | nindent 4 }}
|
||||
{{- with .Values.agentService.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
type: {{ .Values.agentService.type }}
|
||||
{{- if and (eq .Values.agentService.type "LoadBalancer") .Values.agentService.loadBalancerIP }}
|
||||
loadBalancerIP: {{ .Values.agentService.loadBalancerIP }}
|
||||
{{- end }}
|
||||
{{- with .Values.agentService.loadBalancerSourceRanges }}
|
||||
loadBalancerSourceRanges:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
{{- if has .Values.agentService.type (list "LoadBalancer" "NodePort") }}
|
||||
externalTrafficPolicy: {{ .Values.agentService.externalTrafficPolicy }}
|
||||
{{- end }}
|
||||
selector:
|
||||
{{- include "llmgw.selectorLabels" . | nindent 4 }}
|
||||
ports:
|
||||
- name: agents
|
||||
port: {{ .Values.agentService.port }}
|
||||
targetPort: agents
|
||||
---
|
||||
# The StatefulSet's own, giving each pod a stable name.
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "llmgw.fullname" . }}-headless
|
||||
labels:
|
||||
{{- include "llmgw.labels" . | nindent 4 }}
|
||||
spec:
|
||||
clusterIP: None
|
||||
publishNotReadyAddresses: true
|
||||
selector:
|
||||
{{- include "llmgw.selectorLabels" . | nindent 4 }}
|
||||
ports:
|
||||
- name: api
|
||||
port: 8080
|
||||
targetPort: api
|
||||
@@ -0,0 +1,181 @@
|
||||
apiVersion: apps/v1
|
||||
kind: StatefulSet
|
||||
metadata:
|
||||
name: {{ include "llmgw.fullname" . }}
|
||||
labels:
|
||||
{{- include "llmgw.labels" . | nindent 4 }}
|
||||
spec:
|
||||
# One control plane: Kubernetes replaces it within seconds when it goes,
|
||||
# and the agents reconnect on their own. See replicas in values.yaml.
|
||||
replicas: 1
|
||||
serviceName: {{ include "llmgw.fullname" . }}-headless
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "llmgw.selectorLabels" . | nindent 6 }}
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
{{- include "llmgw.selectorLabels" . | nindent 8 }}
|
||||
{{- with .Values.podLabels }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
annotations:
|
||||
# server.yaml is read at start: a changed one restarts the pod.
|
||||
checksum/config: {{ include "llmgw.serverConfig" . | sha256sum }}
|
||||
{{- with .Values.podAnnotations }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
{{- with .Values.imagePullSecrets }}
|
||||
imagePullSecrets:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
# Agents reconnect within seconds of a control plane going; a long
|
||||
# grace period only delays the one that replaces it.
|
||||
terminationGracePeriodSeconds: 30
|
||||
securityContext:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65532
|
||||
runAsGroup: 65532
|
||||
fsGroup: 65532
|
||||
seccompProfile:
|
||||
type: RuntimeDefault
|
||||
containers:
|
||||
- name: llmgw
|
||||
image: {{ include "llmgw.image" . }}
|
||||
imagePullPolicy: {{ .Values.image.pullPolicy }}
|
||||
args: ["server", "--config", "/etc/llmgw/server.yaml"]
|
||||
env:
|
||||
# For `kubectl exec ... -- llmgw ctl`, which reads the token
|
||||
# from here; the server reads admin_token_file.
|
||||
- name: LLMGW_ADMIN_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ include "llmgw.adminSecret" . }}
|
||||
key: {{ if .Values.adminToken.existingSecret }}{{ .Values.adminToken.secretKey }}{{ else }}admin-token{{ end }}
|
||||
ports:
|
||||
- name: api
|
||||
containerPort: 8080
|
||||
- name: agents
|
||||
containerPort: 7443
|
||||
- name: metrics
|
||||
containerPort: 9091
|
||||
# Up is /health. Ready is /leader: a control plane holds the
|
||||
# agents' connections and places models only once it leads — on
|
||||
# PostgreSQL, once the lease of the one it replaced has run out. A
|
||||
# cluster with no node yet is still ready: /ready would wait for a
|
||||
# replica, and nodes could then never connect.
|
||||
startupProbe:
|
||||
httpGet: {path: /health, port: api}
|
||||
periodSeconds: 5
|
||||
failureThreshold: 60
|
||||
livenessProbe:
|
||||
httpGet: {path: /health, port: api}
|
||||
periodSeconds: 10
|
||||
failureThreshold: 6
|
||||
readinessProbe:
|
||||
httpGet: {path: /leader, port: api}
|
||||
periodSeconds: 5
|
||||
failureThreshold: 2
|
||||
securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop: ["ALL"]
|
||||
resources:
|
||||
{{- toYaml .Values.resources | nindent 12 }}
|
||||
volumeMounts:
|
||||
- name: config
|
||||
mountPath: /etc/llmgw/server.yaml
|
||||
subPath: server.yaml
|
||||
readOnly: true
|
||||
- name: agent-tls
|
||||
mountPath: /etc/llmgw/tls
|
||||
readOnly: true
|
||||
- name: admin-token
|
||||
mountPath: /etc/llmgw/admin
|
||||
readOnly: true
|
||||
{{- if eq .Values.store.type "postgres" }}
|
||||
- name: store
|
||||
mountPath: /etc/llmgw/store
|
||||
readOnly: true
|
||||
{{- end }}
|
||||
{{- if .Values.offlineChecksums.configMap }}
|
||||
- name: offline
|
||||
mountPath: /etc/llmgw/offline
|
||||
readOnly: true
|
||||
{{- end }}
|
||||
- name: data
|
||||
mountPath: /var/lib/llmgw
|
||||
- name: tmp
|
||||
mountPath: /tmp
|
||||
volumes:
|
||||
- name: config
|
||||
configMap:
|
||||
name: {{ include "llmgw.fullname" . }}
|
||||
- name: agent-tls
|
||||
secret:
|
||||
secretName: {{ .Values.agentTLS.existingSecret }}
|
||||
defaultMode: 0440
|
||||
items:
|
||||
- {key: ca.pem, path: ca.pem}
|
||||
- {key: server.pem, path: server.pem}
|
||||
- {key: server-key.pem, path: server-key.pem}
|
||||
- name: admin-token
|
||||
secret:
|
||||
secretName: {{ include "llmgw.adminSecret" . }}
|
||||
defaultMode: 0440
|
||||
items:
|
||||
- key: {{ if .Values.adminToken.existingSecret }}{{ .Values.adminToken.secretKey }}{{ else }}admin-token{{ end }}
|
||||
path: admin-token
|
||||
{{- if eq .Values.store.type "postgres" }}
|
||||
- name: store
|
||||
secret:
|
||||
secretName: {{ .Values.store.postgres.existingSecret }}
|
||||
defaultMode: 0440
|
||||
items:
|
||||
- key: {{ .Values.store.postgres.secretKey }}
|
||||
path: dsn
|
||||
{{- end }}
|
||||
{{- if .Values.offlineChecksums.configMap }}
|
||||
- name: offline
|
||||
configMap:
|
||||
name: {{ .Values.offlineChecksums.configMap }}
|
||||
items:
|
||||
- key: {{ .Values.offlineChecksums.key }}
|
||||
path: offline.sha256
|
||||
{{- end }}
|
||||
- name: tmp
|
||||
emptyDir: {}
|
||||
{{- if eq .Values.store.type "postgres" }}
|
||||
# The store is PostgreSQL; nothing here outlives the pod.
|
||||
- name: data
|
||||
emptyDir: {}
|
||||
{{- end }}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.tolerations }}
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- if eq .Values.store.type "sqlite" }}
|
||||
volumeClaimTemplates:
|
||||
- metadata:
|
||||
name: data
|
||||
labels:
|
||||
{{- include "llmgw.selectorLabels" . | nindent 10 }}
|
||||
spec:
|
||||
accessModes: ["ReadWriteOnce"]
|
||||
{{- with .Values.store.sqlite.storageClass }}
|
||||
storageClassName: {{ . }}
|
||||
{{- end }}
|
||||
resources:
|
||||
requests:
|
||||
storage: {{ .Values.store.sqlite.size }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,132 @@
|
||||
# The llmgw control plane in Kubernetes. docs/kubernetes.md says what this
|
||||
# runs, what it does not, and how nodes outside the cluster join it.
|
||||
|
||||
image:
|
||||
# Where you pushed the image built from deploy/container/Dockerfile. The
|
||||
# project publishes binaries, not an image, so this has no default.
|
||||
repository: ""
|
||||
# Empty takes the chart's appVersion.
|
||||
tag: ""
|
||||
pullPolicy: IfNotPresent
|
||||
imagePullSecrets: []
|
||||
|
||||
# One control plane, always. Two against PostgreSQL would elect a leader and
|
||||
# keep a standby, but a standby serves nothing and is never ready, so a
|
||||
# rollout would wait on it for ever. Until the chart handles that, it runs
|
||||
# one: Kubernetes replaces it within seconds when it goes, and the agents
|
||||
# reconnect on their own.
|
||||
replicas: 1
|
||||
|
||||
store:
|
||||
# sqlite: a file on a PersistentVolumeClaim, for one replica.
|
||||
# postgres: a PostgreSQL you run, its address in a Secret.
|
||||
type: sqlite
|
||||
sqlite:
|
||||
size: 10Gi
|
||||
storageClass: ""
|
||||
postgres:
|
||||
# A Secret holding the address, postgres://user:password@host:5432/llmgw?sslmode=verify-full
|
||||
existingSecret: ""
|
||||
secretKey: dsn
|
||||
|
||||
# The agent port's certificate. It comes from llmgw's own authority, not a
|
||||
# public one — agents present certificates from the same authority — so it
|
||||
# is made outside the cluster, where the authority's key stays:
|
||||
#
|
||||
# llmgw ctl cert init-ca --dir ./pki
|
||||
# llmgw ctl cert issue-server llmgw-ctl --host <agent address> --dir ./pki --out ./pki
|
||||
# kubectl create secret generic llmgw-agent-tls \
|
||||
# --from-file=ca.pem=pki/ca.pem \
|
||||
# --from-file=server.pem=pki/server.pem \
|
||||
# --from-file=server-key.pem=pki/server-key.pem
|
||||
#
|
||||
# --host is every address agents dial: agentService's load balancer address
|
||||
# or DNS name.
|
||||
agentTLS:
|
||||
existingSecret: ""
|
||||
|
||||
# The admin API's token (llmgw ctl, the console). Empty makes one, kept in a
|
||||
# Secret across upgrades.
|
||||
adminToken:
|
||||
existingSecret: ""
|
||||
secretKey: admin-token
|
||||
|
||||
# server.yaml beyond what the chart sets (listeners, TLS, store, admin
|
||||
# token): any key of docs/reference.md, laid over the chart's. An offline
|
||||
# cluster sets offline.mirror here, and offlineChecksums below.
|
||||
config:
|
||||
control:
|
||||
enabled: true
|
||||
auth:
|
||||
enabled: true
|
||||
signup_open: false
|
||||
default_rpm: 60
|
||||
default_tpm: 200000
|
||||
default_price:
|
||||
prompt_per_1k: 1000
|
||||
completion_per_1k: 1000
|
||||
max_key_ttl: 2160h
|
||||
# mirror:
|
||||
# api: https://models.internal:8082
|
||||
# offline:
|
||||
# mirror: http://models.internal:8081
|
||||
# firewall_networks: ["10.42.0.0/16"] # the pods' network, so nodes' engines answer the router
|
||||
|
||||
# On an offline cluster, a ConfigMap or Secret holding the engine checksums
|
||||
# 'llmgw bundle trust' writes (offline.sha256), mounted where offline.checksums
|
||||
# points. Empty mounts none.
|
||||
offlineChecksums:
|
||||
configMap: ""
|
||||
key: offline.sha256
|
||||
|
||||
# The OpenAI-compatible API and the metrics, inside the cluster.
|
||||
service:
|
||||
type: ClusterIP
|
||||
port: 8080
|
||||
metricsPort: 9091
|
||||
annotations: {}
|
||||
|
||||
# Where agents connect, with mutual TLS. Nodes run outside the cluster, so
|
||||
# this is a load balancer by default; the TLS ends in the pod, never at a
|
||||
# proxy in between.
|
||||
agentService:
|
||||
type: LoadBalancer
|
||||
port: 7443
|
||||
annotations: {}
|
||||
loadBalancerIP: ""
|
||||
loadBalancerSourceRanges: []
|
||||
# Keeps the agents' addresses as they are, which the node's identity does
|
||||
# not need but the logs and the audit read better with.
|
||||
externalTrafficPolicy: Local
|
||||
|
||||
ingress:
|
||||
enabled: false
|
||||
className: ""
|
||||
annotations: {}
|
||||
hosts:
|
||||
- host: llm.example.com
|
||||
paths:
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
tls: []
|
||||
# - secretName: llm-example-com
|
||||
# hosts: [llm.example.com]
|
||||
|
||||
# A ServiceMonitor for the Prometheus Operator, scraping the metrics port.
|
||||
serviceMonitor:
|
||||
enabled: false
|
||||
interval: 30s
|
||||
labels: {}
|
||||
|
||||
resources:
|
||||
requests:
|
||||
cpu: 500m
|
||||
memory: 512Mi
|
||||
limits:
|
||||
memory: 2Gi
|
||||
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
podAnnotations: {}
|
||||
podLabels: {}
|
||||
+12
-1
@@ -25,6 +25,13 @@ An HTTP server serving a directory. That is the whole design.
|
||||
model.safetensors
|
||||
```
|
||||
|
||||
At an offline site the mirror also holds what the nodes would otherwise fetch
|
||||
from the internet — llmgw's releases and the llama.cpp builds — under
|
||||
`_llmgw/releases/` and `_llmgw/engines/`, put there by `llmgw bundle import`
|
||||
(docs/install.md, "Offline sites"). They are not models: `mirror list` and
|
||||
`mirror unused` pass over them, and a node checks each against what the
|
||||
control plane says it must be, not against anything the mirror says.
|
||||
|
||||
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
|
||||
@@ -65,9 +72,13 @@ 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
|
||||
--tls-key /etc/llmgw/mirror-key.pem \
|
||||
--api-client control-plane
|
||||
```
|
||||
|
||||
`--api-client` is required with `--tls-ca`: it names the one certificate the
|
||||
API answers, because the CA signs every node's certificate as well.
|
||||
|
||||
That is the same CA the agents already use, so the certificates come from the
|
||||
authority you have:
|
||||
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
# Alertmanager for the llmgw cluster.
|
||||
#
|
||||
# Nothing is sent anywhere yet, on purpose. Where an alert should go — an
|
||||
# Nothing is sent anywhere until a receiver is chosen, 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.
|
||||
# visible at :9093 and in Prometheus; uncomment a receiver template below and
|
||||
# name it on the route to make them travel.
|
||||
global:
|
||||
resolve_timeout: 5m
|
||||
|
||||
@@ -31,6 +32,53 @@ route:
|
||||
receivers:
|
||||
- name: none
|
||||
|
||||
# Receiver templates. Uncomment one, fill in where it goes, and name it in
|
||||
# `receiver:` above — on the route, the critical branch, or both. Secrets go
|
||||
# in files Alertmanager reads (the *_file settings), mode 0640 owned by
|
||||
# root:prometheus, not in this file. `amtool check-config` before a restart.
|
||||
#
|
||||
# Slack (an incoming webhook):
|
||||
# - name: slack
|
||||
# slack_configs:
|
||||
# - api_url_file: /etc/prometheus/slack-webhook-url
|
||||
# channel: '#llmgw'
|
||||
# send_resolved: true
|
||||
# title: '{{ .Status | toUpper }}: {{ .CommonLabels.alertname }}'
|
||||
# text: >-
|
||||
# {{ range .Alerts }}{{ .Annotations.summary }}
|
||||
# {{ .Annotations.description }}
|
||||
# {{ end }}
|
||||
#
|
||||
# E-mail (through a relay that takes authenticated submission):
|
||||
# - name: email
|
||||
# email_configs:
|
||||
# - to: ops@example.com
|
||||
# from: llmgw-alerts@example.com
|
||||
# smarthost: smtp.example.com:587
|
||||
# auth_username: llmgw-alerts@example.com
|
||||
# auth_password_file: /etc/prometheus/smtp-password
|
||||
# send_resolved: true
|
||||
#
|
||||
# PagerDuty (Events API v2), for the critical branch — what should wake
|
||||
# someone:
|
||||
# - name: pagerduty
|
||||
# pagerduty_configs:
|
||||
# - routing_key_file: /etc/prometheus/pagerduty-routing-key
|
||||
# severity: '{{ .CommonLabels.severity }}'
|
||||
# description: '{{ .CommonAnnotations.summary }}'
|
||||
#
|
||||
# A webhook of your own (Alertmanager's JSON, version 4):
|
||||
# - name: webhook
|
||||
# webhook_configs:
|
||||
# - url_file: /etc/prometheus/alert-webhook-url
|
||||
# send_resolved: true
|
||||
# max_alerts: 20
|
||||
#
|
||||
# The gateway itself also sends a few events straight to webhooks, with no
|
||||
# Prometheus in between — a model parked, a node lost, a budget spent, new
|
||||
# weights published: `notify.webhooks` in server.yaml (docs/install.md,
|
||||
# "Being told").
|
||||
|
||||
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.
|
||||
|
||||
@@ -156,6 +156,18 @@ groups:
|
||||
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: ModelUpdateAvailable
|
||||
expr: llmgw_model_update_available == 1
|
||||
for: 1h
|
||||
labels: { severity: info }
|
||||
annotations:
|
||||
summary: "{{ $labels.repo }} has new weights for {{ $labels.model }}"
|
||||
description: >-
|
||||
The registry repository the model was served from has published
|
||||
different files under the same names. `llmgw ctl model updates`
|
||||
lists what changed; `llmgw ctl model update {{ $labels.model }}`
|
||||
moves it to them, and `llmgw ctl model rollback` goes back.
|
||||
|
||||
- alert: PoolSaturated
|
||||
expr: llmgw_pool_saturation > 1
|
||||
for: 5m
|
||||
@@ -477,3 +489,98 @@ groups:
|
||||
description: >-
|
||||
Patroni will not promote a standby more than 1 MB behind, so a
|
||||
failover now would find none fit. -1 is a lag Patroni cannot read.
|
||||
|
||||
- name: llmgw-guard
|
||||
rules:
|
||||
- alert: GuardFailingOpen
|
||||
expr: >-
|
||||
sum(rate(llmgw_guard_fail_open_total[5m]))
|
||||
/ clamp_min(sum(rate(llmgw_guard_checks_total[5m])), 1e-9) > 0.01
|
||||
for: 5m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $value | humanizePercentage }} of judged traffic is served unjudged"
|
||||
description: >-
|
||||
A guard policy's judge model is not answering within its latency
|
||||
budget, and the policy fails open: requests go through flagged
|
||||
"judge unavailable" instead of being judged.
|
||||
llmgw_guard_judge_errors_total says whether it is timing out or
|
||||
not there; `llmgw ctl replicas` whether the judge has replicas.
|
||||
|
||||
- alert: GuardJudgeSlow
|
||||
expr: >-
|
||||
histogram_quantile(0.95, sum by (le, judge_model) (rate(llmgw_guard_judge_latency_seconds_bucket[5m]))) > 0.25
|
||||
for: 10m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "Judge {{ $labels.judge_model }} answers in {{ $value | humanizeDuration }} at p95"
|
||||
description: >-
|
||||
Over the default latency budget: every judged request waits this
|
||||
long, and those over their policy's budget are served unjudged or
|
||||
refused. Give the judge more replicas, or faster ones.
|
||||
|
||||
- alert: GuardJudgeUnavailable
|
||||
expr: sum by (reason) (rate(llmgw_guard_judge_errors_total{reason!="timeout"}[5m])) > 0
|
||||
for: 5m
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "The guardrails cannot ask their judge ({{ $labels.reason }})"
|
||||
description: >-
|
||||
no_judge: a policy names a model that is not in the catalog with
|
||||
role: guard. unavailable: it has no replica answering. format: it
|
||||
answers in a form its guard.format does not read.
|
||||
|
||||
- alert: GuardBlocksSpiking
|
||||
expr: >-
|
||||
sum by (tenant) (rate(llmgw_tenant_guard_blocks_total[1h]))
|
||||
> 3 * sum by (tenant) (rate(llmgw_tenant_guard_blocks_total[7d]))
|
||||
and sum by (tenant) (increase(llmgw_tenant_guard_blocks_total[1h])) > 10
|
||||
for: 15m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.tenant }} is being blocked three times as often as usual"
|
||||
description: >-
|
||||
A tenant's requests or answers are blocked by the guardrails
|
||||
more than three times their weekly rate. Someone probing, a
|
||||
client sending data it should not, or a policy that changed.
|
||||
The /admin/guard page lists the recent ones by request.
|
||||
|
||||
- alert: GuardSecretInAnswer
|
||||
expr: sum(increase(llmgw_guard_checks_total{stage="output",category="secret",action!="allowed"}[5m])) > 0
|
||||
labels: { severity: critical }
|
||||
annotations:
|
||||
summary: "A model put a credential in an answer"
|
||||
description: >-
|
||||
An answer contained something shaped like a key or a private key.
|
||||
It was stopped or rewritten if the policy enforces; either way a
|
||||
model has seen it, in its prompt or its training. The usage rows
|
||||
with guard_category secret say which request.
|
||||
|
||||
- name: llmgw-slo
|
||||
rules:
|
||||
- alert: ModelTTFTOverSLO
|
||||
expr: >-
|
||||
histogram_quantile(0.99, sum by (le, model) (rate(llmgw_ttft_seconds_bucket[10m])))
|
||||
> on (model) group_left llmgw_model_slo_seconds{kind="ttft_p99"}
|
||||
for: 15m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} answers its first token past its SLO at p99"
|
||||
description: >-
|
||||
The p99 time to the first token over ten minutes is past the
|
||||
model's slo.ttft_p99. Admission already holds the lower bands back
|
||||
while it is; this says it has lasted. More replicas, or less
|
||||
traffic in the bands above.
|
||||
|
||||
- alert: ModelTPOTOverSLO
|
||||
expr: >-
|
||||
histogram_quantile(0.99, sum by (le, model) (rate(llmgw_tpot_seconds_bucket[10m])))
|
||||
> on (model) group_left llmgw_model_slo_seconds{kind="tpot_p99"}
|
||||
for: 15m
|
||||
labels: { severity: warning }
|
||||
annotations:
|
||||
summary: "{{ $labels.model }} produces tokens slower than its SLO at p99"
|
||||
description: >-
|
||||
The p99 time per output token is past slo.tpot_p99: replicas are
|
||||
batching more than they can decode quickly, or a slower card is
|
||||
serving. llmgw_output_tokens_per_second by pool says which.
|
||||
@@ -27,6 +27,17 @@ scrape_configs:
|
||||
labels:
|
||||
instance: control-plane
|
||||
|
||||
# Every machine of the cluster, for its node exporter: network and disk
|
||||
# traffic, filesystems, CPU time. The control plane lists them — nodes,
|
||||
# itself, the store's members — so a node added is a node scraped with
|
||||
# nothing edited here, and each carries its llmgw id as the node label.
|
||||
# `llmgw node add` installs the exporter; on a machine added before that,
|
||||
# apt-get install prometheus-node-exporter.
|
||||
- job_name: nodes
|
||||
http_sd_configs:
|
||||
- url: http://10.0.0.140:9091/prometheus/targets
|
||||
refresh_interval: 1m
|
||||
|
||||
- job_name: prometheus
|
||||
static_configs:
|
||||
- targets: ['localhost:9090']
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
// Package llmgw holds what the binary carries besides code: its manual, the
|
||||
// description of its admin API, and the example quality suites.
|
||||
package llmgw
|
||||
|
||||
import (
|
||||
"embed"
|
||||
"io/fs"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Docs is the documentation, read on any machine by `llmgw docs` and the
|
||||
// console's manual without a checkout or the internet.
|
||||
//
|
||||
//go:embed README.md CHANGELOG.md docs/*.md
|
||||
var Docs embed.FS
|
||||
|
||||
// Doc is one page of the manual.
|
||||
type Doc struct {
|
||||
// Key is what `llmgw docs KEY` takes, and the page's address under /docs/
|
||||
// on the gateway.
|
||||
Key string
|
||||
// Path is the file in Docs.
|
||||
Path string
|
||||
// Title is the document's first heading.
|
||||
Title string
|
||||
}
|
||||
|
||||
// Source is the document's markdown.
|
||||
func (d Doc) Source() string {
|
||||
b, err := fs.ReadFile(Docs, d.Path)
|
||||
if err != nil {
|
||||
return "# " + d.Path + "\n\nNot in this build."
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// docOrder is the documents in the order someone new reads them. Here rather
|
||||
// than with either reader, because the console's manual and the gateway's
|
||||
// docs portal both list them, and two lists drift.
|
||||
var docOrder = []string{"console", "quick-install", "install", "node-onboarding", "onboarding",
|
||||
"clients", "guardrails", "knowledge", "choosing-models", "fabrics", "facility-aware-placement", "data",
|
||||
"reference", "development"}
|
||||
|
||||
// Manual is every document, in reading order: those docOrder names, then any
|
||||
// other under docs/ alphabetically, so a new one is never missing, then the
|
||||
// README and the changes.
|
||||
func Manual() []Doc {
|
||||
var out []Doc
|
||||
seen := map[string]bool{}
|
||||
add := func(key, path string) {
|
||||
if seen[key] {
|
||||
return
|
||||
}
|
||||
b, err := fs.ReadFile(Docs, path)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
seen[key] = true
|
||||
out = append(out, Doc{Key: key, Path: path, Title: titleOf(string(b), path)})
|
||||
}
|
||||
for _, k := range docOrder {
|
||||
add(k, "docs/"+k+".md")
|
||||
}
|
||||
entries, _ := fs.ReadDir(Docs, "docs")
|
||||
var rest []string
|
||||
for _, e := range entries {
|
||||
if strings.HasSuffix(e.Name(), ".md") {
|
||||
rest = append(rest, strings.TrimSuffix(e.Name(), ".md"))
|
||||
}
|
||||
}
|
||||
sort.Strings(rest)
|
||||
for _, k := range rest {
|
||||
add(k, "docs/"+k+".md")
|
||||
}
|
||||
add("readme", "README.md")
|
||||
add("changes", "CHANGELOG.md")
|
||||
return out
|
||||
}
|
||||
|
||||
// titleOf is a document's first top-level heading, or its path when it has
|
||||
// none.
|
||||
func titleOf(src, path string) string {
|
||||
for _, l := range strings.Split(src, "\n") {
|
||||
if strings.HasPrefix(l, "# ") {
|
||||
return strings.TrimSpace(strings.TrimPrefix(l, "# "))
|
||||
}
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// AdminOpenAPI is the admin API described in OpenAPI 3.1, generated from its
|
||||
// route table by `go generate` in internal/docgen and served at
|
||||
// /api/v1/openapi.json, so a client can be pointed at the server it talks to.
|
||||
//
|
||||
//go:embed docs/admin-api.openapi.json
|
||||
var AdminOpenAPI []byte
|
||||
|
||||
// EvalSuites are the example quality suites under deploy/eval, which the
|
||||
// control plane runs by name when its eval directory has none of the name.
|
||||
//
|
||||
//go:embed deploy/eval/*.yaml
|
||||
var EvalSuites embed.FS
|
||||
@@ -0,0 +1,160 @@
|
||||
# The admin API
|
||||
|
||||
Written for whoever automates a cluster rather than typing at it: how to reach
|
||||
the control plane's management API, what it asks of a request, and where its
|
||||
description is.
|
||||
|
||||
`llmgw ctl` and the console are clients of this API and nothing more, so
|
||||
anything they do can be done with `curl`, a generated client or the
|
||||
[Terraform provider](../terraform-provider-llmgw/README.md). It adds and
|
||||
removes models, scales them, cordons and upgrades nodes, sets capacity
|
||||
quotas, attaches guardrail policies, and reads the fleet, usage and the audit
|
||||
log.
|
||||
|
||||
## Reaching it
|
||||
|
||||
The API is under `/api/v1` on the admin listener, `admin_listen` in
|
||||
`/etc/llmgw/server.yaml`: `127.0.0.1:9090` unless changed. It is on loopback
|
||||
on purpose. The same listener serves pprof when it is turned on, and the
|
||||
server refuses to start with it on any other address unless
|
||||
`admin_remote: true` — the token would cross the network in clear.
|
||||
|
||||
From another machine, tunnel to it rather than moving it:
|
||||
|
||||
```sh
|
||||
ssh -N -L 9090:127.0.0.1:9090 admin@control-plane &
|
||||
curl -s -H "Authorization: Bearer $LLMGW_ADMIN_TOKEN" http://127.0.0.1:9090/api/v1/summary
|
||||
```
|
||||
|
||||
With `admin_remote: true` and `admin_listen: 10.0.0.5:9090`, call that address
|
||||
directly — on a network you trust, or behind a TLS proxy that keeps the
|
||||
`Host` header an address.
|
||||
|
||||
## The token
|
||||
|
||||
Every path under `/api` needs the admin token as a bearer token:
|
||||
|
||||
```
|
||||
Authorization: Bearer 6f1c…
|
||||
```
|
||||
|
||||
The server writes it the first time it starts, 0600, to
|
||||
`/var/lib/llmgw/admin-token` under its systemd unit (`admin_token_file`
|
||||
moves it; a server run by hand keeps it in `secrets/admin-token` beside its
|
||||
configuration). `llmgw ctl` reads it from there, which is why it runs with
|
||||
sudo, or from `LLMGW_ADMIN_TOKEN`. To give it to automation:
|
||||
|
||||
```sh
|
||||
sudo cat /var/lib/llmgw/admin-token
|
||||
```
|
||||
|
||||
The token is the whole of the API's authorisation: whoever holds it can do
|
||||
everything the API does. Keep it where you keep other root credentials.
|
||||
|
||||
## What a request must look like
|
||||
|
||||
Besides the token, the listener refuses three things before a handler sees
|
||||
them, as defences against a browser being turned against it:
|
||||
|
||||
- **A name in `Host`.** The listener answers to an address or `localhost`,
|
||||
not a DNS name, so a page on some other site cannot reach it by rebinding
|
||||
a name to 127.0.0.1. Call it as `http://127.0.0.1:9090`, not
|
||||
`http://cp.example.com:9090`; behind a proxy, have the proxy send an
|
||||
address.
|
||||
- **A cross-site request.** One whose `Origin` is another host, or whose
|
||||
`Sec-Fetch-Site` says another site sent it.
|
||||
- **A change that is not JSON.** Every method but `GET` and `HEAD` must carry
|
||||
`Content-Type: application/json`, even a `DELETE` or a `POST` with no body.
|
||||
A form post from a page cannot set that header.
|
||||
|
||||
A change is recorded in the audit log with who made it, as the request says
|
||||
in `X-Llmgw-Actor` (letters, digits and `._-@`, at most 64; `unknown` when
|
||||
absent). `llmgw ctl` sends the name of whoever ran it; automation should send
|
||||
its own, `terraform` or `ci@deploy`, so `llmgw ctl audit` can tell them apart.
|
||||
|
||||
A refusal is a 4xx or 5xx with `{"error": "…"}` saying why. 501 means this
|
||||
control plane lacks what the endpoint needs — a store, a mirror, guard
|
||||
policies — and says which.
|
||||
|
||||
## Some calls
|
||||
|
||||
```sh
|
||||
H=(-H "Authorization: Bearer $LLMGW_ADMIN_TOKEN" -H "Content-Type: application/json" -H "X-Llmgw-Actor: $USER")
|
||||
API=http://127.0.0.1:9090/api/v1
|
||||
|
||||
# The fleet in one answer, and every model.
|
||||
curl -s "${H[@]}" $API/summary
|
||||
curl -s "${H[@]}" $API/models
|
||||
|
||||
# Add a model, read back its definition with the defaults filled in, scale it.
|
||||
curl -s "${H[@]}" -X POST $API/models -d '{
|
||||
"name": "qwen-small",
|
||||
"source": "/srv/models/qwen2.5-1.5b-instruct-q4_k_m.gguf",
|
||||
"profiles": [{"name": "cpu", "engine": "llamacpp", "min_memory_gb": 4, "cpu_cores": 4}],
|
||||
"scaling": {"min": 1, "max": 2}
|
||||
}'
|
||||
curl -s "${H[@]}" $API/models/qwen-small
|
||||
curl -s "${H[@]}" -X POST $API/models/qwen-small/scale -d '{"min": 0, "max": 2, "idle_timeout": 1800000000000}'
|
||||
|
||||
# A team's capacity quota.
|
||||
curl -s "${H[@]}" -X PUT $API/quotas -d '{"scope": "team:Acme/ml", "limit": {"gpus": 8}}'
|
||||
|
||||
# Remove the model; a DELETE carries the content type too.
|
||||
curl -s "${H[@]}" -X DELETE $API/models/qwen-small
|
||||
```
|
||||
|
||||
Durations in bodies are integer nanoseconds, as Go encodes them: half an hour
|
||||
is `1800000000000`.
|
||||
|
||||
## The description
|
||||
|
||||
The API is described in OpenAPI 3.1:
|
||||
|
||||
- served by every control plane at `GET /api/v1/openapi.json` (with the
|
||||
token, like the rest), as built into that binary, so it describes the
|
||||
server you are talking to;
|
||||
- committed as [admin-api.openapi.json](admin-api.openapi.json) for reading
|
||||
without a cluster.
|
||||
|
||||
```sh
|
||||
curl -s "${H[@]}" $API/openapi.json > llmgw-admin.json
|
||||
```
|
||||
|
||||
It lists every endpoint, its path and query parameters, and the schema of
|
||||
each request and answer, with the meaning of each field. Load it into any
|
||||
OpenAPI tool — Swagger UI, Redoc, an editor — or generate a client from it:
|
||||
|
||||
```sh
|
||||
docker run --rm -v "$PWD:/w" openapitools/openapi-generator-cli generate \
|
||||
-i /w/llmgw-admin.json -g python -o /w/llmgw-admin-client
|
||||
```
|
||||
|
||||
### How it is kept true
|
||||
|
||||
The document is generated, not written. `internal/control/adminapi/routes.go`
|
||||
holds a table of every endpoint: its method and path, a summary, its query
|
||||
parameters, and Go values of the types it reads and answers with. The
|
||||
generator in `internal/docgen/openapi` reads those types by reflection, the
|
||||
way `encoding/json` encodes them, and takes each field's description from its
|
||||
comment in the source.
|
||||
|
||||
Two tests keep it from drifting. One registers the API on a recording mux and
|
||||
fails when an endpoint is served and not in the table, or in the table and not
|
||||
served. The other regenerates the document and fails when it differs from the
|
||||
committed one. After changing an endpoint or a type it carries:
|
||||
|
||||
```sh
|
||||
go generate ./internal/docgen
|
||||
```
|
||||
|
||||
## What it does not cover
|
||||
|
||||
- **Organizations, teams, projects, people and API keys.** These live in the
|
||||
store and are managed with `llmgw org`, `llmgw users` and `llmgw keys` on the
|
||||
control plane, or in the web console's `/admin` pages on the client port
|
||||
behind a session. The admin API reads the tree only where a quota or a
|
||||
policy names a part of it.
|
||||
- **The client-facing API** (`/v1/chat/completions` and the rest), which is
|
||||
OpenAI's and Anthropic's, on `listen`, with API keys.
|
||||
- **The rest of the admin listener**: `/metrics`, `/health`, the dashboard
|
||||
pages and `/debug/pprof`, which are not under `/api`.
|
||||
File diff suppressed because it is too large.
Load diff
+557
@@ -0,0 +1,557 @@
|
||||
# Using llmgw from code
|
||||
|
||||
Written for a developer whose program talks to a model through llmgw. There
|
||||
is no llmgw SDK, and there will not be one: the gateway speaks OpenAI's API
|
||||
and Anthropic's Messages API, so the official SDKs for both work as they are.
|
||||
Only two things change — the base URL and the key. Everything llmgw adds on
|
||||
top is an HTTP header, which every one of those SDKs lets you set.
|
||||
|
||||
The examples use `https://llm.example.com` for the gateway; use the address
|
||||
your operator gave you, or the one on the gateway's own front page. Every
|
||||
example reads the key from `LLMGW_API_KEY`. `<model-name>` is a model this
|
||||
gateway serves: `GET /v1/models` lists the ones your key may use, and sending
|
||||
the placeholder unchanged is answered with that list rather than with an
|
||||
error.
|
||||
|
||||
## Getting a key
|
||||
|
||||
Keys are made on the gateway's **Account** page after signing in, or issued
|
||||
by an operator. A key is shown once and never again — only its hash is kept —
|
||||
so put it in a password manager straight away. `GET /me` with the key says
|
||||
what it may do and what it has spent.
|
||||
|
||||
## The base URL and the key
|
||||
|
||||
| SDK | Base URL | Key |
|
||||
|---|---|---|
|
||||
| OpenAI (any language) | `https://llm.example.com/v1` | `Authorization: Bearer …`, as the SDK sends it |
|
||||
| Anthropic (any language) | `https://llm.example.com` — the SDK adds `/v1/messages` | `x-api-key: …`, as the SDK sends it |
|
||||
|
||||
Both SDKs also read their usual environment variables, so a program that
|
||||
already uses them needs no change at all:
|
||||
|
||||
```
|
||||
export OPENAI_BASE_URL=https://llm.example.com/v1
|
||||
export OPENAI_API_KEY=$LLMGW_API_KEY
|
||||
|
||||
export ANTHROPIC_BASE_URL=https://llm.example.com
|
||||
export ANTHROPIC_API_KEY=$LLMGW_API_KEY
|
||||
```
|
||||
|
||||
## A first request
|
||||
|
||||
### curl
|
||||
|
||||
```
|
||||
curl https://llm.example.com/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "<model-name>",
|
||||
"messages": [{"role": "user", "content": "Say hello in Turkish."}]}'
|
||||
```
|
||||
|
||||
### Python
|
||||
|
||||
```
|
||||
pip install openai
|
||||
```
|
||||
|
||||
```python
|
||||
import os
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="https://llm.example.com/v1",
|
||||
api_key=os.environ["LLMGW_API_KEY"])
|
||||
|
||||
reply = client.chat.completions.create(
|
||||
model="<model-name>",
|
||||
messages=[{"role": "user", "content": "Say hello in Turkish."}],
|
||||
)
|
||||
print(reply.choices[0].message.content)
|
||||
```
|
||||
|
||||
### TypeScript
|
||||
|
||||
```
|
||||
npm install openai
|
||||
```
|
||||
|
||||
```typescript
|
||||
import OpenAI from "openai";
|
||||
|
||||
const client = new OpenAI({
|
||||
baseURL: "https://llm.example.com/v1",
|
||||
apiKey: process.env.LLMGW_API_KEY,
|
||||
});
|
||||
|
||||
const reply = await client.chat.completions.create({
|
||||
model: "<model-name>",
|
||||
messages: [{ role: "user", content: "Say hello in Turkish." }],
|
||||
});
|
||||
console.log(reply.choices[0].message.content);
|
||||
```
|
||||
|
||||
### Go
|
||||
|
||||
```
|
||||
go get github.com/openai/openai-go/v3
|
||||
```
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
|
||||
"github.com/openai/openai-go/v3"
|
||||
"github.com/openai/openai-go/v3/option"
|
||||
)
|
||||
|
||||
func main() {
|
||||
client := openai.NewClient(
|
||||
option.WithBaseURL("https://llm.example.com/v1/"),
|
||||
option.WithAPIKey(os.Getenv("LLMGW_API_KEY")),
|
||||
)
|
||||
reply, err := client.Chat.Completions.New(context.Background(), openai.ChatCompletionNewParams{
|
||||
Model: "<model-name>",
|
||||
Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("Say hello in Turkish.")},
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
fmt.Println(reply.Choices[0].Message.Content)
|
||||
}
|
||||
```
|
||||
|
||||
The import path carries the SDK's major version; use the one your module
|
||||
already has.
|
||||
|
||||
## Streaming
|
||||
|
||||
Set `stream` and read the chunks. llmgw may send a last chunk that carries
|
||||
only the token counts and no choices, so check for a choice before reading
|
||||
one — the OpenAI API does the same when usage is asked for.
|
||||
|
||||
```python
|
||||
stream = client.chat.completions.create(
|
||||
model="<model-name>",
|
||||
messages=[{"role": "user", "content": "Count to ten."}],
|
||||
stream=True,
|
||||
)
|
||||
for chunk in stream:
|
||||
if chunk.choices and chunk.choices[0].delta.content:
|
||||
print(chunk.choices[0].delta.content, end="", flush=True)
|
||||
print()
|
||||
```
|
||||
|
||||
```typescript
|
||||
const stream = await client.chat.completions.create({
|
||||
model: "<model-name>",
|
||||
messages: [{ role: "user", content: "Count to ten." }],
|
||||
stream: true,
|
||||
});
|
||||
for await (const chunk of stream) {
|
||||
process.stdout.write(chunk.choices[0]?.delta?.content ?? "");
|
||||
}
|
||||
```
|
||||
|
||||
```go
|
||||
stream := client.Chat.Completions.NewStreaming(ctx, openai.ChatCompletionNewParams{
|
||||
Model: "<model-name>",
|
||||
Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("Count to ten.")},
|
||||
})
|
||||
for stream.Next() {
|
||||
chunk := stream.Current()
|
||||
if len(chunk.Choices) > 0 {
|
||||
fmt.Print(chunk.Choices[0].Delta.Content)
|
||||
}
|
||||
}
|
||||
if err := stream.Err(); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
```
|
||||
|
||||
With curl, add `"stream": true` to the body and `-N` so curl prints the
|
||||
events as they come.
|
||||
|
||||
## Tools
|
||||
|
||||
Function calling is OpenAI's: describe the tools, run the ones the model
|
||||
asks for, send the results back. Whether a model calls tools well is the
|
||||
model's business, not the gateway's; ask your operator which of the served
|
||||
models are set up for it.
|
||||
|
||||
**Editors' agents need such a model.** Continue's Agent mode, Cline and
|
||||
Codex act only through tool calls: with a model whose calls come back as
|
||||
text — `<tool_call>{...}` or a JSON block in the answer, `tool_calls` empty —
|
||||
they chat but never read a file or run a command. Qwen2.5-Coder is one;
|
||||
Qwen3-Coder is not. The configuration the portal writes for Continue says
|
||||
`tool_use` only for a model whose catalog entry says tools, and gives
|
||||
Continue the model's context and longest answer, so it never sends a request
|
||||
the gateway refuses.
|
||||
|
||||
```python
|
||||
import json
|
||||
|
||||
tools = [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "The current weather in a city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}]
|
||||
|
||||
messages = [{"role": "user", "content": "Is it raining in Ankara?"}]
|
||||
reply = client.chat.completions.create(model="<model-name>", messages=messages, tools=tools)
|
||||
msg = reply.choices[0].message
|
||||
|
||||
if msg.tool_calls:
|
||||
messages.append(msg)
|
||||
for call in msg.tool_calls:
|
||||
args = json.loads(call.function.arguments)
|
||||
result = {"city": args["city"], "raining": False} # your code here
|
||||
messages.append({"role": "tool", "tool_call_id": call.id,
|
||||
"content": json.dumps(result)})
|
||||
reply = client.chat.completions.create(model="<model-name>", messages=messages, tools=tools)
|
||||
|
||||
print(reply.choices[0].message.content)
|
||||
```
|
||||
|
||||
## Embeddings
|
||||
|
||||
```python
|
||||
result = client.embeddings.create(
|
||||
model="<embedding-model>",
|
||||
input=["llmgw routes requests", "to GPUs it can reach"],
|
||||
)
|
||||
vectors = [d.embedding for d in result.data]
|
||||
```
|
||||
|
||||
```typescript
|
||||
const result = await client.embeddings.create({
|
||||
model: "<embedding-model>",
|
||||
input: ["llmgw routes requests", "to GPUs it can reach"],
|
||||
});
|
||||
const vectors = result.data.map((d) => d.embedding);
|
||||
```
|
||||
|
||||
An embedding model is a model of its own; a chat model does not answer
|
||||
`/v1/embeddings`.
|
||||
|
||||
## The Responses API
|
||||
|
||||
`/v1/responses` is served too, for clients that speak nothing else (Codex
|
||||
among them). llmgw translates it to chat completions, so it works with every
|
||||
model the gateway serves; what has no equivalent there — OpenAI's hosted
|
||||
tools, reasoning summaries — is ignored rather than refused.
|
||||
|
||||
```python
|
||||
response = client.responses.create(
|
||||
model="<model-name>",
|
||||
instructions="Answer in one sentence.",
|
||||
input="Why is the sky blue?",
|
||||
)
|
||||
print(response.output_text)
|
||||
```
|
||||
|
||||
Streamed, the events are the Responses API's own:
|
||||
|
||||
```python
|
||||
stream = client.responses.create(model="<model-name>", input="Count to ten.", stream=True)
|
||||
for event in stream:
|
||||
if event.type == "response.output_text.delta":
|
||||
print(event.delta, end="", flush=True)
|
||||
```
|
||||
|
||||
## Batches
|
||||
|
||||
For work that need not finish now: OpenAI's Files and Batch APIs. A batch runs
|
||||
at the lowest priority, so it uses what interactive traffic leaves free, and
|
||||
each line is counted against the key that submitted it.
|
||||
|
||||
```python
|
||||
import json
|
||||
|
||||
with open("requests.jsonl", "w") as f:
|
||||
for i, text in enumerate(["first document", "second document"]):
|
||||
f.write(json.dumps({
|
||||
"custom_id": f"doc-{i}",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": "<model-name>",
|
||||
"messages": [{"role": "user", "content": "Summarise: " + text}]},
|
||||
}) + "\n")
|
||||
|
||||
batch_file = client.files.create(file=open("requests.jsonl", "rb"), purpose="batch")
|
||||
batch = client.batches.create(
|
||||
input_file_id=batch_file.id,
|
||||
endpoint="/v1/chat/completions",
|
||||
completion_window="24h",
|
||||
)
|
||||
|
||||
# Later:
|
||||
batch = client.batches.retrieve(batch.id)
|
||||
if batch.status == "completed":
|
||||
for line in client.files.content(batch.output_file_id).text.splitlines():
|
||||
result = json.loads(line)
|
||||
print(result["custom_id"], result["response"]["body"]["choices"][0]["message"]["content"])
|
||||
```
|
||||
|
||||
## The Anthropic SDK
|
||||
|
||||
`/v1/messages` is Anthropic's Messages API, translated to chat completions and
|
||||
back, streamed or not, with tools and images. It serves the gateway's own
|
||||
models — `model` is a model this gateway serves, not a Claude model.
|
||||
|
||||
```
|
||||
pip install anthropic
|
||||
```
|
||||
|
||||
```python
|
||||
import os
|
||||
import anthropic
|
||||
|
||||
client = anthropic.Anthropic(base_url="https://llm.example.com",
|
||||
api_key=os.environ["LLMGW_API_KEY"])
|
||||
|
||||
message = client.messages.create(
|
||||
model="<model-name>",
|
||||
max_tokens=1024,
|
||||
messages=[{"role": "user", "content": "Say hello in Turkish."}],
|
||||
)
|
||||
print(message.content[0].text)
|
||||
|
||||
with client.messages.stream(
|
||||
model="<model-name>",
|
||||
max_tokens=1024,
|
||||
messages=[{"role": "user", "content": "Count to ten."}],
|
||||
) as stream:
|
||||
for text in stream.text_stream:
|
||||
print(text, end="", flush=True)
|
||||
```
|
||||
|
||||
```typescript
|
||||
import Anthropic from "@anthropic-ai/sdk";
|
||||
|
||||
const client = new Anthropic({
|
||||
baseURL: "https://llm.example.com",
|
||||
apiKey: process.env.LLMGW_API_KEY,
|
||||
});
|
||||
|
||||
const message = await client.messages.create({
|
||||
model: "<model-name>",
|
||||
max_tokens: 1024,
|
||||
messages: [{ role: "user", content: "Say hello in Turkish." }],
|
||||
});
|
||||
for (const block of message.content) {
|
||||
if (block.type === "text") console.log(block.text);
|
||||
}
|
||||
```
|
||||
|
||||
```go
|
||||
client := anthropic.NewClient(
|
||||
option.WithBaseURL("https://llm.example.com/"),
|
||||
option.WithAPIKey(os.Getenv("LLMGW_API_KEY")),
|
||||
)
|
||||
message, err := client.Messages.New(ctx, anthropic.MessageNewParams{
|
||||
Model: "<model-name>",
|
||||
MaxTokens: 1024,
|
||||
Messages: []anthropic.MessageParam{
|
||||
anthropic.NewUserMessage(anthropic.NewTextBlock("Say hello in Turkish.")),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
for _, block := range message.Content {
|
||||
if text, ok := block.AsAny().(anthropic.TextBlock); ok {
|
||||
fmt.Println(text.Text)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
(`github.com/anthropics/anthropic-sdk-go` and its `option` package.)
|
||||
|
||||
With curl:
|
||||
|
||||
```
|
||||
curl https://llm.example.com/v1/messages \
|
||||
-H "x-api-key: $LLMGW_API_KEY" \
|
||||
-H "anthropic-version: 2023-06-01" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "<model-name>", "max_tokens": 1024,
|
||||
"messages": [{"role": "user", "content": "Say hello in Turkish."}]}'
|
||||
```
|
||||
|
||||
## What llmgw adds
|
||||
|
||||
Each is a header, so a request stays exactly what the SDK would send anywhere
|
||||
else. Set one for every request on the client, or for one request:
|
||||
|
||||
| | Every request | One request |
|
||||
|---|---|---|
|
||||
| Python (both SDKs) | `OpenAI(..., default_headers={...})` | `create(..., extra_headers={...})` |
|
||||
| TypeScript (both SDKs) | `new OpenAI({ ..., defaultHeaders: {...} })` | `create({...}, { headers: {...} })` |
|
||||
| Go (both SDKs) | `openai.NewClient(..., option.WithHeader(k, v))` | `New(ctx, params, option.WithHeader(k, v))` |
|
||||
| curl | | `-H "k: v"` |
|
||||
|
||||
### Stored conversations: `x-llmgw-thread`
|
||||
|
||||
Send the same thread id with each request and only the new turn; llmgw puts
|
||||
the stored conversation in front of it. The conversation continues on
|
||||
whichever model serves the next turn, even a different one. It must be
|
||||
turned on for your project by an operator (`llmgw keys threads`); a request
|
||||
that names a thread on a gateway that keeps none is refused rather than
|
||||
answered without its history.
|
||||
|
||||
```python
|
||||
import uuid
|
||||
|
||||
thread = str(uuid.uuid4()) # anyone in your project who knows it reads it
|
||||
for turn in ["My name is Deniz.", "What is my name?"]:
|
||||
reply = client.chat.completions.create(
|
||||
model="<model-name>",
|
||||
messages=[{"role": "user", "content": turn}],
|
||||
extra_headers={"x-llmgw-thread": thread},
|
||||
)
|
||||
print(reply.choices[0].message.content)
|
||||
```
|
||||
|
||||
A thread belongs to the project, not to the key, so use an id nobody can
|
||||
guess: at most 128 letters, digits and `- _ . :`. When the conversation no
|
||||
longer fits the model, the oldest turns are left out and the reply carries
|
||||
`x-llmgw-thread-trimmed` with how many. `DELETE /v1/threads/{id}` forgets one:
|
||||
|
||||
```
|
||||
curl -X DELETE https://llm.example.com/v1/threads/$THREAD \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY"
|
||||
```
|
||||
|
||||
### Sessions: `x-session-id`
|
||||
|
||||
Requests with the same session id go to the same replica while it is
|
||||
healthy, so a conversation that resends its history finds it already in that
|
||||
replica's cache and answers sooner. Without the header, the body's `user`
|
||||
field is used. Nothing is stored; it is only a routing hint.
|
||||
|
||||
```python
|
||||
client = OpenAI(base_url="https://llm.example.com/v1", api_key=os.environ["LLMGW_API_KEY"],
|
||||
default_headers={"x-session-id": "chat-7f3a"})
|
||||
```
|
||||
|
||||
### Tags and metadata
|
||||
|
||||
Requests can be labelled so the usage they record can be told apart — by
|
||||
feature, by customer, by job. Send OpenAI's own `metadata`, string keys to
|
||||
string values, or the same as the `x-llmgw-tags` header for a client that
|
||||
cannot set `metadata`; the `user` field is recorded too. Which tags also
|
||||
become metric labels is the operator's choice (`router.tag_labels`).
|
||||
|
||||
```python
|
||||
reply = client.chat.completions.create(
|
||||
model="<model-name>",
|
||||
messages=[{"role": "user", "content": "Classify this ticket."}],
|
||||
metadata={"feature": "triage", "customer": "acme"},
|
||||
user="user-1234",
|
||||
)
|
||||
```
|
||||
|
||||
```
|
||||
curl https://llm.example.com/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "x-llmgw-tags: feature=triage,customer=acme" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "<model-name>", "messages": [{"role": "user", "content": "Classify this ticket."}]}'
|
||||
```
|
||||
|
||||
### Canaries: `x-llmgw-variant`
|
||||
|
||||
While an operator tries a model's newer weights on a share of its traffic, a
|
||||
session stays on one side. `x-llmgw-variant: canary` sends a request to the
|
||||
canary deliberately — to try it before it is promoted, or to compare the two.
|
||||
|
||||
```python
|
||||
reply = client.chat.completions.create(
|
||||
model="<model-name>",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
extra_headers={"x-llmgw-variant": "canary"},
|
||||
)
|
||||
```
|
||||
|
||||
### Memory notes: `/v1/memory`
|
||||
|
||||
Notes an agent keeps for itself across conversations: a key, some text, and
|
||||
optional tags, private to your project and to the agent named in
|
||||
`x-llmgw-agent` (`default` when it is not sent). An operator turns them on for
|
||||
a project (`llmgw keys memory`). They are plain REST, so any HTTP client
|
||||
will do:
|
||||
|
||||
```
|
||||
H="Authorization: Bearer $LLMGW_API_KEY"
|
||||
|
||||
curl -X PUT https://llm.example.com/v1/memory/favourite-colour -H "$H" \
|
||||
-H "x-llmgw-agent: helper" -H "Content-Type: application/json" \
|
||||
-d '{"text": "The user likes dark green.", "tags": ["preference"]}'
|
||||
|
||||
curl "https://llm.example.com/v1/memory?q=green" -H "$H" -H "x-llmgw-agent: helper"
|
||||
curl "https://llm.example.com/v1/memory?tag=preference" -H "$H" -H "x-llmgw-agent: helper"
|
||||
curl https://llm.example.com/v1/memory/favourite-colour -H "$H" -H "x-llmgw-agent: helper"
|
||||
curl -X DELETE https://llm.example.com/v1/memory/favourite-colour -H "$H" -H "x-llmgw-agent: helper"
|
||||
```
|
||||
|
||||
A key is up to 200 bytes, a note up to 16 KiB, with at most 16 tags of
|
||||
lowercase letters, digits, `-` and `_`, and 1000 notes per agent.
|
||||
|
||||
### MCP: `/mcp`
|
||||
|
||||
The same notes as tools an agent can be given — `remember`, `recall` and
|
||||
`forget` — from an MCP server at `https://llm.example.com/mcp` (the HTTP
|
||||
transport, answered as plain JSON). Add it to your agent's MCP servers with
|
||||
the key as a bearer token and, if you want notes kept apart, `x-llmgw-agent`.
|
||||
llmgw does not put these tools into requests by itself: an agent that is not
|
||||
given the server never sees its notes.
|
||||
|
||||
```
|
||||
curl https://llm.example.com/mcp \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc": "2.0", "id": 1, "method": "tools/list"}'
|
||||
```
|
||||
|
||||
### Headers llmgw sends back
|
||||
|
||||
| Header | Says |
|
||||
|---|---|
|
||||
| `x-request-id` | The request's id in the gateway's logs and usage records — quote it when asking an operator about a request. |
|
||||
| `x-llmgw-served-model` | The model that answered, when it was not the one asked for: a spent budget can move a request to a cheaper model. |
|
||||
| `x-llmgw-queue-ms` | How long the request waited for room before it was sent. |
|
||||
| `x-llmgw-thread-trimmed` | How many stored turns were left out to fit the model. |
|
||||
|
||||
## When something is refused
|
||||
|
||||
Errors come back in the shape of the API that was called — OpenAI's
|
||||
`{"error": {"message": …, "type": …}}` or Anthropic's — so each SDK raises
|
||||
its usual exception with the gateway's message in it. The ones worth knowing:
|
||||
|
||||
| Status | Usually means |
|
||||
|---|---|
|
||||
| 400 | Besides a malformed request: a prompt too long for the model, or a guardrail policy refused it (the code `content_policy_violation`). |
|
||||
| 401 | No key, a mistyped one, or one that was revoked or has expired. |
|
||||
| 403 | The key may not use that model. |
|
||||
| 429 | A rate limit or budget of the key, its project or its team was reached. `Retry-After` says when to try again, and is absent when waiting will not help — a spent monthly budget. |
|
||||
| 503 | No model of that name is served (the message lists the ones that are), the model is still loading, or every replica is busy and the request waited as long as it may. The SDKs retry by themselves; a name that is not served will not start to be. |
|
||||
|
||||
## Priority
|
||||
|
||||
With keys on, the key decides the request's priority band and routing
|
||||
profile, set by an operator; the `x-llmgw-band` and `x-llmgw-profile` headers
|
||||
are ignored, so a client cannot promote itself. On a gateway without keys
|
||||
they are read from the request.
|
||||
+339
@@ -0,0 +1,339 @@
|
||||
# The console
|
||||
|
||||
`sudo llmgw`, with nothing after it, on a machine llmgw runs on. On a machine
|
||||
with nothing installed it opens the setup screens; on a control plane, the
|
||||
console described here; on a node, the node's own console, which edits
|
||||
`agent.yaml`. Everything on these screens can also be done by command — each
|
||||
screen says which — so a script never needs them and a person never needs
|
||||
to remember the commands.
|
||||
|
||||
The keyboard drives it; there is no mouse. **?** opens this page from
|
||||
anywhere, and the **Manual** entry holds every other page, as `llmgw docs`
|
||||
does without the console.
|
||||
|
||||
## The layout
|
||||
|
||||
A bar across the top names the machine. Down the left, the entries, in four
|
||||
groups: **Watch**, **People**, **Help**, **Settings** and **Do**. The panel on
|
||||
the right shows the entry chosen. The bottom line always says which keys do
|
||||
what at that moment, and a question or a result opens as a dialog over the
|
||||
panel.
|
||||
|
||||
| Key | Does |
|
||||
|---|---|
|
||||
| ↑ ↓ | Choose an entry, or a row in the panel |
|
||||
| → or Enter | Go into the panel |
|
||||
| Enter on a row | What can be done to it: its details first, then the rest |
|
||||
| ← or Esc | Back out: from a page to its list, from the panel to the entries |
|
||||
| A | Add one, on a view that holds things: a quota, an account, a key |
|
||||
| PgUp PgDn | Scroll a long panel |
|
||||
| ? | This page |
|
||||
| Q | Leave; the cluster keeps running |
|
||||
| Ctrl-C | Leave at once |
|
||||
|
||||
What is watched is read again every three seconds.
|
||||
|
||||
## Forms and questions
|
||||
|
||||
A change that needs values opens a form, filled in with what is set now. Tab
|
||||
or ↓ moves to the next field, ← → changes a choice, and Enter on **Apply**
|
||||
sends it; a form of one field is sent by Enter on it. A mistake the form can
|
||||
see — a maximum below the minimum, a word where a number goes — is said in the
|
||||
form and nothing is sent. One the control plane refuses comes back as a dialog
|
||||
with its reason, word for word.
|
||||
|
||||
Anything that cannot be undone — removing, draining, deleting, revoking —
|
||||
asks first, and only **Y** goes ahead.
|
||||
|
||||
Serving a model and moving one to new weights are jobs on the control plane:
|
||||
started from here, they run there in the background whatever this screen
|
||||
does, and the console is free at once. **Jobs** follows them. Upgrading the
|
||||
cluster, a backup and adding a node change this machine or another over ssh;
|
||||
those close the screen and run in the terminal, so their output is there to
|
||||
read, and Enter brings the console back.
|
||||
|
||||
## Setting up
|
||||
|
||||
On a machine with nothing installed, the control plane is set up from one
|
||||
form (`sudo llmgw server init`): the administrator's email, other names nodes
|
||||
reach the machine by, and ticks for PostgreSQL, the model mirror, this machine
|
||||
as a node, the starter model and the guardrails' judge. **Offline bundle**
|
||||
and **Its sha256** are for a site with no internet: the path of a bundle made
|
||||
by `llmgw bundle create`, and the sha256 it printed. Given, the bundle goes
|
||||
into the mirror, every file checked, and the cluster is set up to fetch
|
||||
nothing from outside (install.md, "Offline sites"). A bundle with no sha256,
|
||||
or with the mirror unticked, is said in the form and nothing starts.
|
||||
|
||||
## Watch
|
||||
|
||||
**Overview** is `llmgw ctl status`: nodes, GPUs by class, fabrics, replicas,
|
||||
models, the store's members, and anything parked, unplaced or idle.
|
||||
|
||||
**Monitor** draws the cluster the way btop draws a machine: CPU, memory,
|
||||
GPU busy and GPU memory, requests and tokens a second, time to first token,
|
||||
requests in flight and queued, network and disk — Braille charts, filled
|
||||
and coloured green to red by height — over 5 minutes, an hour, a day or a
|
||||
week (**1**–**4**, or **R** for the next). Below them each node's CPU,
|
||||
memory and network now and as a sparkline, each card's busy, memory, heat
|
||||
and power, and each model's requests a second. The history comes from the
|
||||
Prometheus named by `monitoring.prometheus` in server.yaml; without one the
|
||||
charts draw what this console has seen since it opened, and say so. Network
|
||||
and disk come from each machine's node exporter, which `llmgw node add`
|
||||
installs. **N** sends a test notification to each webhook in server.yaml's
|
||||
`notify` section and says how each answered (`llmgw ctl notify test`). A node's and a model's own pages have their charts too, over the
|
||||
same range.
|
||||
|
||||
**Nodes** lists every node with its state, where it is (its rack, or row,
|
||||
room or site, and PDUs), load, GPUs, replicas, disk and agent version. A node
|
||||
in a held rack reads *ready, held*. On a row:
|
||||
|
||||
- **Details** — its address, location and any hold, system, load, engines and
|
||||
GPUs, and any warning.
|
||||
- **Cordon** places nothing new on it; **Uncordon** undoes it.
|
||||
- **Location…** sets its site, room, row, rack and PDUs (the `llmgw.io/`
|
||||
labels), which placement spreads replicas across.
|
||||
- **Hold rack …** places nothing new in the node's rack, and can drain it;
|
||||
**Release …** on a node in a held domain undoes it. See
|
||||
[facility-aware-placement.md](facility-aware-placement.md).
|
||||
- **Drain** moves its replicas off one at a time, each replacement up before
|
||||
the original stops.
|
||||
- **vllm build / sglang build** (nodes with GPUs) — which build of the engine
|
||||
the node installs: `auto` follows its hardware, `none` installs none.
|
||||
- **Upgrade its agent** to a release; its engines keep serving.
|
||||
- **Delete model files nothing uses** frees the node's disk.
|
||||
- **Remove from the cluster** — the control plane forgets it; the machine is
|
||||
not touched.
|
||||
|
||||
**GPUs** is every card: how busy, its memory, temperature, power, and the
|
||||
model holding it.
|
||||
|
||||
**Models** lists the catalog: ready against wanted, status, class, profiles,
|
||||
and **UPDATE** when the registry has new weights for it. On a row:
|
||||
|
||||
- **Details** — description, replicas, class, owner, profiles, source, where
|
||||
its weights came from, what it reserves against what it uses, and its
|
||||
replicas.
|
||||
- **Replicas (scale)** — the minimum kept running, the most it scales to,
|
||||
and for a model loaded on demand how long it may sit unused. A maximum of
|
||||
0 stops it.
|
||||
- **Class** — critical, standard or opportunistic: whose capacity gives way
|
||||
when there is not enough for everyone.
|
||||
- **Owner** — the organization, team or project whose capacity quota it
|
||||
counts against.
|
||||
- **Resources** — the memory and cores one replica reserves on its node.
|
||||
- **LoRA adapters** — `name=source[@scale]`, separated by commas.
|
||||
- **Restart its replicas**, one by one or in place.
|
||||
- **Update to …** — when the repository it came from has published different
|
||||
weights for its files: they are fetched into the mirror beside the ones in
|
||||
use, checked, swapped in, and the replicas replaced one at a time.
|
||||
**Go back to …** returns to the weights the update replaced.
|
||||
**Ask the registry for new weights now** checks without waiting for the
|
||||
next scheduled look (`mirror.update_check`, six hours by default).
|
||||
- **Try the new weights as a canary…** — instead of swapping them in, runs
|
||||
them beside the model as `MODEL-canary`, one replica, given the share of
|
||||
its requests you choose. With a canary running the menu has **Canary's
|
||||
share…**, **Promote the canary…** (the model moves to its weights, the
|
||||
canary is removed) and **Take the canary out…**. A canary that errs or
|
||||
answers more slowly than the model is taken out by itself. The MODEL
|
||||
column shows `(10% to canary)` and `(canary of …)`.
|
||||
- **Ask it a quality suite…** — one of the suites the control plane has (in
|
||||
`control.eval_dir`, or built in: `arithmetic`, `json-format`), asked
|
||||
through the router; the dialog counts the answers as they come in, then
|
||||
shows the score, each case and the model's last runs. Esc stops watching;
|
||||
the run carries on and is kept. **Its quality runs** lists them, as
|
||||
`llmgw ctl eval --history` does.
|
||||
- **Hold its updates to a quality suite…** — the suite and score its new
|
||||
weights must reach before an update or a promotion moves it to them
|
||||
(`llmgw ctl model gate`); **none** takes it away. With one set, **Update
|
||||
to …** and **Promote the canary…** say so.
|
||||
- **Why is it placed where it is** — `llmgw ctl explain-placement`.
|
||||
- **Try again**, for a model parked after failing every start.
|
||||
- **Remove** — its replicas stop and it leaves the catalog. A model from
|
||||
`server.yaml` is changed there instead; the console says so.
|
||||
|
||||
**A** on Models adds one from a definition file on this machine: the YAML a
|
||||
model takes in `server.yaml`.
|
||||
|
||||
**Replicas** lists every replica; on one, **Restart it**, replaced or in
|
||||
place.
|
||||
|
||||
**Usage** is the last 24 hours by team, project or organization (→ changes
|
||||
which): requests, tokens, errors and spend.
|
||||
|
||||
**Quotas** are capacity quotas: the GPUs, cores and memory each organization,
|
||||
team or project is guaranteed and may at most hold, and what it holds now.
|
||||
**A** adds one; on a row, change or remove it.
|
||||
|
||||
**Knowledge** lists the knowledge-base collections: each one's owner, the
|
||||
model that embeds it, how many documents and chunks it holds and who else it
|
||||
is granted to. **A** makes one. On a row: its **Documents**, **Grant…** a
|
||||
project, team, organization, person or key read or write access, **Revoke**
|
||||
a grant, **Answer from it by default…** for a part of the tree, **Try a
|
||||
search…** as an administrator, or **Delete…** it. People
|
||||
fill collections from their own machines with `llmgw-kb`
|
||||
([knowledge.md](knowledge.md)). Off on a SQLite store, and says why.
|
||||
|
||||
**Mirror** is what the model mirror holds, one row for each version of each
|
||||
model — a file, or the set of files a model of several is: shards, a vision
|
||||
projector, a config and tokenizer — with how far its download has got and
|
||||
which models use it. **A** downloads a URL into it; **I** imports a model
|
||||
from an Ollama on this machine into a mirror directory on this machine — the
|
||||
mirror beside the control plane — and shows the model definition to add for
|
||||
it (a mirror elsewhere: `llmgw mirror import-ollama` there); on a row, its files with
|
||||
their checksums and where each was fetched from, the source a model
|
||||
definition names for it, or remove it — with a warning when a model still
|
||||
uses it.
|
||||
|
||||
**Jobs** lists the models being served or updated since the control plane
|
||||
started: what each is doing now — the download's progress, a replica
|
||||
starting — and whether it ended done, failed, or waiting for room on busy
|
||||
cards. The entry shows how many are running. Enter on one shows its steps.
|
||||
`llmgw ctl jobs` is the same, and `llmgw ctl jobs ID` follows one.
|
||||
|
||||
**Logs** shows a node's agent log, newest at the bottom, errors and warnings
|
||||
coloured.
|
||||
|
||||
## People
|
||||
|
||||
The accounts, keys and organizations live in the store; these views read it
|
||||
and change it with the same commands as `llmgw users`, `llmgw keys` and
|
||||
`llmgw org`, so their checks and messages are the ones shown.
|
||||
|
||||
**Accounts** — who has signed up, their state, whether they see the cluster
|
||||
dashboard, and what they belong to. **A** creates an account, already
|
||||
approved. On a row: approve one waiting, suspend or reinstate, give or take
|
||||
the cluster dashboard, add to or take out of an organization, team or
|
||||
project, their limits, a new password (their sessions close), delete, and
|
||||
**Forget them…** for a request to be forgotten (`llmgw users forget`).
|
||||
|
||||
**API keys** — every key, its project, owner, state and last use. **A**
|
||||
issues one; the key is shown once, in the result, and never again. On a row:
|
||||
its limits, **Rotate…** — a new key with the same project, owner and limits,
|
||||
the old one working on for the overlap you give (`llmgw keys rotate`) — or
|
||||
revoke it.
|
||||
|
||||
**Accounts** rows also have **Staff role…**:
|
||||
- `viewer` looks at every page under /admin and changes nothing;
|
||||
- `operator` also runs the cluster from them (models through the mirror, a
|
||||
node's engines);
|
||||
- `billing` sees usage and quotas;
|
||||
- `admin` may do anything.
|
||||
|
||||
The STAFF column shows each person's role.
|
||||
|
||||
On **Usage**, **E** exports usage as CSV to a file on the control plane:
|
||||
by hour, day or month and any of organization, team, project, user, key and
|
||||
model.
|
||||
|
||||
**Audit** lists every change of the last week — through `ctl`, this console
|
||||
or the /admin pages, and people's own keys and passwords — with who made
|
||||
it, through what, and how it ended. Enter shows the fields it was given,
|
||||
with anything secret taken out. `llmgw ctl audit` is the same, with filters.
|
||||
|
||||
**Organizations** — the tenancy tree: organizations, their teams, the teams'
|
||||
projects, with members, keys and limits. **A** adds an organization. On a
|
||||
row: add a team or project inside it, rename, move, limits for the level and
|
||||
for each of its members, its capacity quota, and for a project how long
|
||||
conversations and agent notes are kept, and whether its requests are
|
||||
(**Keep requests…**: for how long, and what share), and whether they stay in
|
||||
the cluster (**Requests stay in the cluster…**: never an outside provider),
|
||||
and whether the same request is answered again from memory (**Cache
|
||||
answers…**: for how long, and whether only reproducible answers or any). **Longest key lifetime…** caps how
|
||||
long a key made there may live, and is the expiry of one made without any
|
||||
(`llmgw keys ttl`; `auth.max_key_ttl` in server.yaml caps every key).
|
||||
**Remove** takes away only something empty.
|
||||
|
||||
Limits are requests and tokens a minute, a budget per day or month, and the
|
||||
models allowed. Each level only narrows what is above it.
|
||||
|
||||
**Guardrails…** on a row lists the guardrail policies server.yaml defines.
|
||||
Each policy is marked in one of three ways:
|
||||
|
||||
- attached here (✓): Enter detaches it;
|
||||
- applying from above (·): from everywhere, or from the organization or
|
||||
team named;
|
||||
- neither: Enter attaches it.
|
||||
|
||||
The **GUARDRAILS** column names those attached to the row itself. Policies
|
||||
narrow like limits; see [guardrails.md](guardrails.md).
|
||||
|
||||
## Help
|
||||
|
||||
**Manual** lists every page there is: these, the installation and operation
|
||||
guides, the configuration reference, a page for every command, the README
|
||||
and the changes in each release. On a page:
|
||||
|
||||
| Key | Does |
|
||||
|---|---|
|
||||
| ↑ ↓, PgUp PgDn, Space | Scroll |
|
||||
| g, G | The top, the end |
|
||||
| T | Contents: the page's headings, Enter goes to one |
|
||||
| / | Find; the lines holding it are marked |
|
||||
| n, N | The next match, the one before |
|
||||
| [ ] | The topic before, the one after |
|
||||
| ← | Back to the list |
|
||||
|
||||
`llmgw docs` opens the same manual outside the console, on any machine,
|
||||
installed or not; `llmgw docs install` opens one page; piped, it prints the
|
||||
page's text.
|
||||
|
||||
## Settings
|
||||
|
||||
**server.yaml** (on a node, **agent.yaml**) is every key, as a tree of menus
|
||||
built from the configuration reference: the value set, or the one that
|
||||
applies when nothing is, its type, and what it is for. Enter changes a value
|
||||
or opens a block; on a list or map of blocks, Enter opens an entry and **D**
|
||||
removes one. **U** puts a key back to its default. **S** saves.
|
||||
|
||||
On a control plane, the cluster's keys are read from and saved to the
|
||||
cluster's settings in the store — the same on every control plane, in force
|
||||
on all of them within seconds (`llmgw ctl config` does the same from a
|
||||
shell) — and this machine's own keys (listeners, certificates, the store's
|
||||
address) to its server.yaml, which the server notices and applies by itself.
|
||||
Either is checked before anything is kept, and the file's comments survive.
|
||||
On a node, agent.yaml is read when the agent starts, so the console offers
|
||||
to restart it.
|
||||
|
||||
## Do
|
||||
|
||||
- **Add a node** — installs llmgw on a machine over ssh, checks it, gives it a
|
||||
certificate and starts its agent (`sudo llmgw node add`).
|
||||
- **Replicate the store** — moves the store onto two more machines under
|
||||
Patroni (offered only while it is not replicated).
|
||||
- **Serve a model** — Hugging Face's models: trending, most downloaded,
|
||||
most liked or newest (**S** changes which), or **/** to search. **F**
|
||||
switches between GGUF files, which llama.cpp runs on the CPU or a GPU, and
|
||||
original safetensors weights, which vLLM runs on GPUs. Enter on a
|
||||
repository plans it against the cluster as it is, before anything is
|
||||
fetched: the model's shape is read from one file's header, or its
|
||||
config.json, and every file is listed with the best way it runs here — on
|
||||
one card, on several of a node's cards, across several nodes' cards over a
|
||||
fabric (original weights only), or on a CPU node — what each card or node
|
||||
would hold, and whether that hardware is free now. The file that suits the
|
||||
cluster best is first, marked ★; one that fits nowhere says so. Busy is not
|
||||
unfit: which file is best follows the hardware the cluster has, not what it
|
||||
is running, so a file for a card that is in use is still offered for it —
|
||||
served, it waits for the card, or with a fallback runs elsewhere until then. Choosing a
|
||||
file offers every way it fits, best first, and what to do when that one is
|
||||
busy: the next best way, written as a lower-priority profile so placement
|
||||
falls back to it, or wait. A vision model offers its projector too, fetched
|
||||
beside it and loaded with it so the model reads images. What is chosen is
|
||||
fetched into the mirror, checked, added to the catalog and started, and a
|
||||
replica is asked to say hello (`llmgw ctl model serve`; the plan alone is
|
||||
`llmgw ctl model plan`).
|
||||
`mirror.registry` points it at another registry, for a site without the
|
||||
internet.
|
||||
- **Upgrade the cluster** to a release: this machine, then every node in turn
|
||||
(`sudo llmgw ctl upgrade`).
|
||||
- **Restart all replicas**, one at a time.
|
||||
- **Back up now** — the store and `/etc/llmgw` into one archive
|
||||
(`sudo llmgw server backup`).
|
||||
- **Wipe this machine** — removes llmgw's services, configuration, state and
|
||||
users from this machine, after listing them and asking for its name
|
||||
(`sudo llmgw wipe`; on a node too).
|
||||
- **Leave**.
|
||||
|
||||
## The node's console
|
||||
|
||||
On a node, `sudo llmgw` opens `agent.yaml` in the same settings view, and the
|
||||
manual. The cluster is watched and changed from a control plane.
|
||||
+119
-13
@@ -16,20 +16,38 @@ 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.
|
||||
Nothing on the serving path leaves the cluster — **unless the catalog has a
|
||||
model served outside it** (`external:` in a model definition): a request for
|
||||
that model, or one that falls back to it, sends its prompt to that provider.
|
||||
A guardrail policy with `no_store` refuses both for the scopes it covers.
|
||||
|
||||
**Guardrails** read every prompt and answer they are configured for, in the
|
||||
gateway's own memory, to judge them. What they record — in the usage row and
|
||||
the log — is the action, the policy, the category and how long it took: never
|
||||
the text, the value that matched, or the placeholder map of a redaction. See
|
||||
[guardrails.md](guardrails.md).
|
||||
|
||||
## 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.
|
||||
`Usage` and there is no field for text in it. What the client chose to say
|
||||
about the request is kept with it: OpenAI's `user` field and its `metadata`
|
||||
or `x-llmgw-tags` — short strings, set by the client, for attributing spend.
|
||||
A client that puts a person's name or address there puts it in the usage
|
||||
records; the field is the client's, and so is that choice. Tags are not sent
|
||||
on to the engine.
|
||||
|
||||
**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.
|
||||
|
||||
**The store's own secrets**, like the one engine keys are derived from, plain
|
||||
unless `secrets.master_key` seals them — with Vault's transit engine, so that
|
||||
neither the store nor its backups hold them.
|
||||
|
||||
**People and tenancy**: accounts, the organisation/team/project tree, and the
|
||||
limits attached to it. Passwords are hashed.
|
||||
|
||||
@@ -37,10 +55,32 @@ limits attached to it. Passwords are hashed.
|
||||
|
||||
**Conversations and agent notes** — only for a project that asked. See below.
|
||||
|
||||
**Kept requests** — only for a project that asked (`llmgw keys requests`):
|
||||
what each request sent the model, after the guardrails' redaction and with
|
||||
any stored conversation in it, and the answer the client got, for the hours
|
||||
it said, all of them or a share. Never under a `no_store` policy.
|
||||
|
||||
**Batch files** — when a client uses the Batch API: the file of requests it
|
||||
uploaded (prompts) and the files the batch wrote (answers), in the store,
|
||||
until deleted (`DELETE /v1/files/{id}`) or, for what a batch wrote, until
|
||||
`router.batch.retain_days` (30) pass. Only keys of the project that made
|
||||
them can read them, and a project under a `no_store` policy cannot make
|
||||
them.
|
||||
|
||||
**Reused answers** — only for a project that asked (`llmgw keys cache`):
|
||||
each answer that may be reused, as the client got it, and a hash of the
|
||||
request it answers, in the gateway process's memory — never in the store,
|
||||
never on disk — for the project's TTL or until pushed out by newer ones
|
||||
(`router.cache.max_bytes`), and gone with a restart. Never under a `no_store`
|
||||
policy. Forgetting a person does not reach into it; the answers age out by
|
||||
the TTL. The store keeps only the setting, and a usage record's mark that it
|
||||
was answered from the cache.
|
||||
|
||||
## What is not stored
|
||||
|
||||
Prompts. Completions. Anything the model was shown or produced, unless the
|
||||
project turned on conversations or notes.
|
||||
project turned on conversations, notes or kept requests, or sent a batch.
|
||||
(Reused answers are held in memory only, as above.)
|
||||
|
||||
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,
|
||||
@@ -86,22 +126,62 @@ 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:
|
||||
Around the serving path, a cluster set up the ordinary way does reach out.
|
||||
None of it carries a prompt, an answer, a key or a user — it is software and
|
||||
public model files coming in — but a review should know each one:
|
||||
|
||||
**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
|
||||
**Installing engines.** Each node's agent downloads the llama.cpp build for its
|
||||
hardware from GitHub, checked against the checksum GitHub publishes, and vLLM
|
||||
(SGLang when asked) from PyPI, with `uv` from GitHub. On every node, when it
|
||||
joins and when the engine version changes.
|
||||
|
||||
**Releases.** `llmgw node add`, `llmgw ctl upgrade` and `llmgw upgrade` fetch
|
||||
llmgw's own release binaries from the project's release server, checked against
|
||||
its signed checksums.
|
||||
|
||||
**Fetching weights.** When a model's `source` is a URL, the agent downloads it;
|
||||
the mirror fetches what it is asked to, from Hugging Face or another URL, and
|
||||
`server init` has it fetch the starter model and the guardrails' judge. Point
|
||||
sources at the mirror and the nodes fetch only from it — 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.
|
||||
**The registry.** With a mirror configured, the control plane asks Hugging Face
|
||||
every six hours whether the models served from it have newer files
|
||||
(`mirror.update_check`), and `/admin/mirror` searches it. `llmgw ctl plan`
|
||||
reads a model's architecture from huggingface.co to work out whether it fits
|
||||
your hardware.
|
||||
|
||||
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.
|
||||
There is no telemetry, no licence check, and no call home.
|
||||
|
||||
### A cluster with no way out
|
||||
|
||||
A site with no route to the internet installs from an **offline bundle**
|
||||
instead, and then nothing at all leaves the cluster (docs/install.md, "Offline
|
||||
sites"). The bundle is made on a machine that has internet — llmgw's release
|
||||
binaries with their signed checksums, the llama.cpp builds with the checksums
|
||||
GitHub publishes, the model files with the checksums Hugging Face publishes,
|
||||
each checked as it is downloaded — carried in, and imported into the mirror.
|
||||
With `offline.mirror` set in `server.yaml`, nodes install engines from the
|
||||
mirror and never from GitHub, releases come from the mirror, and the registry
|
||||
is not asked about newer weights.
|
||||
|
||||
Verification is not weakened by any of this. A release is held to the project's
|
||||
signing key wherever it is fetched from. An engine file is held to the checksum
|
||||
recorded when the bundle was made, which the control plane keeps and hands each
|
||||
node over the agent's authenticated connection: the mirror serves plain HTTP,
|
||||
and nothing a node runs is taken on the mirror's word. The bundle itself is
|
||||
checked against the sha256 its maker printed, carried to the site apart from
|
||||
the bundle.
|
||||
|
||||
What does not come in a bundle: vLLM and SGLang, which are Python environments
|
||||
from PyPI — an offline node serves with llama.cpp unless one is installed on it
|
||||
by hand — and the operating system's own packages (the OpenMP runtime the
|
||||
engine needs, PostgreSQL when chosen), which come from the machines' package
|
||||
sources: at an offline site, the distribution mirror it already has. Searching
|
||||
the registry from `/admin/mirror` and `llmgw ctl plan` need the internet and
|
||||
do not work offline; `mirror.registry` can point the first at an internal copy.
|
||||
|
||||
## What the software cannot do for you
|
||||
|
||||
@@ -131,3 +211,29 @@ touch, whatever the catalog says.
|
||||
| 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` |
|
||||
|
||||
## Forgetting a person
|
||||
|
||||
`llmgw users forget <email> --yes` (or **Forget them…** in the console) is for
|
||||
a request to be forgotten:
|
||||
|
||||
- their usage — every request record and every sum — is kept, so totals and
|
||||
budgets stay true, but under a random pseudonym that links to nothing;
|
||||
- their kept requests (`keys requests`), personal files and batches, and
|
||||
their own limits are deleted;
|
||||
- their email and id are taken out of the audit log, which keeps what was
|
||||
done;
|
||||
- then the account goes, as `users delete` removes it: sessions ended, keys
|
||||
revoked and left without an owner.
|
||||
|
||||
Conversations and agent notes belong to projects, not to the people in them,
|
||||
and are not touched: a project forgets them by its retention, and a thread by
|
||||
`DELETE /v1/threads/{id}`. Backups taken before still hold everything until
|
||||
they age out; that is what their retention is for.
|
||||
|
||||
## Keeping a project's requests in the cluster
|
||||
|
||||
`llmgw keys residency --project <path> --in-cluster` refuses the project's
|
||||
keys any model served by an outside provider, and never uses one for them as
|
||||
a fallback or a canary: their prompts do not leave the machines the cluster
|
||||
runs on. `--anywhere` lifts it.
|
||||
@@ -1,8 +1,9 @@
|
||||
# Quickstart: one machine, one model, ten minutes
|
||||
# Development: llmgw from source on one machine
|
||||
|
||||
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.**
|
||||
workstation, built from source, with no systemd and no `sudo` — for working on
|
||||
llmgw itself. **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.
|
||||
@@ -26,7 +27,7 @@ make test
|
||||
./bin/llmgw --version
|
||||
```
|
||||
|
||||
`make build` writes `./bin/llmgw`. Go 1.26 or later is the only requirement.
|
||||
`make build` writes `./bin/llmgw`. Go 1.26.6 or later is the only requirement (an older Go fetches it itself).
|
||||
|
||||
## 2. Check the path with a fake engine
|
||||
|
||||
@@ -1,13 +1,109 @@
|
||||
# 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.
|
||||
**Status: steps 0, 1 and 2 are built, with the manual hold of step 4;
|
||||
temperature conditions (steps 3 and 4) and shedding (step 5) are not**,
|
||||
until the thermal data of a real site is at hand — without it thresholds 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). [What is built](#what-is-built) says what works today and where it
|
||||
departs from the design below; the rest of the document is the design.
|
||||
|
||||
## What is built
|
||||
|
||||
**Where a node is.** Five labels: `llmgw.io/site`, `room`, `row`, `rack` and
|
||||
`pdu` (a comma separated list, for a dual-fed rack). They come from the node's
|
||||
`agent.yaml` under `labels:`, or from the control plane:
|
||||
|
||||
```
|
||||
llmgw ctl label gpu01 llmgw.io/rack=r12-07 llmgw.io/pdu=r12-07-a,r12-07-b
|
||||
llmgw ctl label gpu01 llmgw.io/rack- # back to agent.yaml's value, if any
|
||||
```
|
||||
|
||||
Labels set with `ctl label` are kept in the store, win over the agent's, and
|
||||
may be set before the node joins. A misspelt `llmgw.io/` key is refused. A
|
||||
domain is named `level:name` — `rack:r12-07`, `pdu:r12-07-a` — and names are
|
||||
taken as given, so two rooms that both have a rack `r1` must label them apart.
|
||||
A node without location labels is in no domain and places as it always did.
|
||||
|
||||
**Spreading.** A model's next replica goes first to a node sharing no rack and
|
||||
no PDU with the nodes the model already runs on; nodes sharing fewer come
|
||||
before nodes sharing more. It is an ordering, never a filter: a rack holding
|
||||
the only free cards still takes the replica. It ranks above preferring a node
|
||||
that already holds the weights — a second replica one breaker away from the
|
||||
first is worth less than the copy saved — and above bin packing.
|
||||
|
||||
**Keeping a replica together.** A profile's `keep_within: rack` (or `row`,
|
||||
`room`, `site`) keeps a multi-node replica's nodes in one such domain, inside
|
||||
the fabric or NVLink domain placement already requires; nodes without the
|
||||
label are passed over for it, with the reason.
|
||||
|
||||
**Holds.** `llmgw ctl facility hold rack:r12-07 --reason "PDU swap"` places
|
||||
nothing new in the domain, as a cordon of each of its nodes would, including
|
||||
nodes labelled into it later. `--drain` also drains its nodes, one replica at
|
||||
a time with the replacement up first; `llmgw ctl facility release` undoes
|
||||
both. Holds are kept in the store across restarts (the drain itself is node
|
||||
state, which is not: after a restart the hold still keeps work out, and
|
||||
replicas not yet moved stay where they are). Any domain can be held: a PDU
|
||||
holds every rack on it.
|
||||
|
||||
**Power budgets.** In `server.yaml`:
|
||||
|
||||
```yaml
|
||||
facility:
|
||||
margin: 0.1 # share of every cap kept free
|
||||
gpu_watts: 700 # a card that reports no power limit
|
||||
cpu_core_watts: 10 # a core reserved by a CPU replica (cpu_cores)
|
||||
node_watts: 350 # each node's own draw beyond its cards
|
||||
budgets:
|
||||
- domain: rack:r12-07
|
||||
cap_watts: 16000 # what one feed carries alone, derated
|
||||
other_watts: 800 # switches, storage
|
||||
- domain: pdu:r12-07-a
|
||||
cap_watts: 7400
|
||||
```
|
||||
|
||||
A replica's draw is the profile's `power_watts` (split over its nodes) when
|
||||
declared; otherwise each card at the power limit it reports through the agent's
|
||||
sensors (or `gpu_watts`), and each reserved core at `cpu_core_watts`. A
|
||||
domain's use is, node by node, `node_watts` plus the larger of what is placed
|
||||
there and what the node's sensors read in the last two minutes
|
||||
(`measured_fresh`), plus the domain's `other_watts`. A shard that would take
|
||||
any of its node's domains past `cap × (1 − margin)` is refused; a multi-node
|
||||
replica's shards are counted together, so two that each fit a rack alone do not
|
||||
both go into it. A draw nobody knows — no declared figure, no card limit, no
|
||||
`gpu_watts` — counts as nothing: missing data never blocks.
|
||||
|
||||
**What people see.** `llmgw ctl explain-placement` names the domain: *rack
|
||||
r12-07 is held: PDU swap*, *rack r12-07: 14.2 kW of a 14.4 kW limit (16.0 kW
|
||||
cap less 10%) in use, the shard needs ~2.8 kW*. The rejections of a placed
|
||||
replica are kept too, so a replica in an unexpected rack says why. `llmgw ctl
|
||||
nodes` has a LOCATION column, shows *ready, held*, and filters with `--domain
|
||||
rack:r12-07`. `llmgw ctl facility` lists every domain with its nodes, use,
|
||||
limit, reserved and measured draw, and holds. The console's Nodes view has a
|
||||
WHERE column; a node's page shows its location and hold, and its menu sets
|
||||
the location (**Location…**) and holds or releases its rack. The admin API:
|
||||
`GET /api/v1/facility`, `POST /api/v1/facility/holds`, `DELETE
|
||||
/api/v1/facility/holds/{domain}`, `POST /api/v1/nodes/{id}/labels`.
|
||||
|
||||
**Where it departs from the design below, and why.**
|
||||
|
||||
- Domains come from labels alone; there is no facility file declaring them.
|
||||
Budgets are in `server.yaml`, the one file the control plane already reads.
|
||||
A DCIM export would generate that block.
|
||||
- Feeds are not modelled. The cap given is what one feed carries alone,
|
||||
derated — the figure the design says matters — so the operator gives the
|
||||
one number instead of llmgw deriving it from two.
|
||||
- The manual hold of step 4 is built now, without conditions: it needs no
|
||||
sensor data and answers the commonest facility event, planned maintenance.
|
||||
Its command is `ctl facility hold`, with a `--drain` the design did not have.
|
||||
- `keep_within` is new: the request that a multi-node replica stay in one rack
|
||||
where the cabling requires it.
|
||||
- Spreading is an ordering of candidates, not the scoring of the design; power
|
||||
headroom does not yet order racks ("level load across racks"), and
|
||||
multi-node replicas are not spread across racks.
|
||||
- No `llmgw_domain_*` metrics or alerts yet; `ctl facility` and the API carry
|
||||
the same figures.
|
||||
|
||||
## Why
|
||||
|
||||
@@ -260,16 +356,17 @@ 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
|
||||
1. **Location and spreading (built).** 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
|
||||
2. **Power budgets from a static file (built, without headroom
|
||||
scoring).** 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.
|
||||
4. **Thermal enforcement.** `warm` scored down, `hot` refused; the manual hold
|
||||
and release are built.
|
||||
5. **Shedding and tokens per watt.** Opt-in shedding through step 2's
|
||||
preemption; measured power in profile preference.
|
||||
|
||||
|
||||
@@ -0,0 +1,334 @@
|
||||
# Guardrails
|
||||
|
||||
Written for the operator deciding what the gateway lets through: what a
|
||||
policy can check, what it can do, and how one is attached to a team.
|
||||
|
||||
`llmgw server init` writes a `baseline` policy that applies to every
|
||||
request, with Qwen3Guard-Gen 0.6B as its judge on CPUs; see
|
||||
[install.md](install.md#2-the-control-plane) for what it checks. The
|
||||
policies are part of the cluster's settings, kept in the store: `llmgw ctl
|
||||
config edit` (or the console's Settings) changes them, and the change is in
|
||||
force on every control plane within seconds, without a restart. A cluster
|
||||
whose settings were brought in from a server.yaml without a `guard:` section
|
||||
gets the `baseline` policy, without a judge. With no policies, nothing is
|
||||
checked and nothing costs anything. With one, every request is judged once it is authorised and
|
||||
within its limits, with its stored conversation spliced in, and before it is
|
||||
sent to a model; and the answer is judged before the client sees it.
|
||||
|
||||
## Turning them off and on
|
||||
|
||||
```
|
||||
llmgw ctl guard off --note "why" # no request is checked, whatever the policies say
|
||||
llmgw ctl guard on # every policy checks again, as it was
|
||||
```
|
||||
|
||||
That is `guard.enabled` in the cluster's settings, in force on every control
|
||||
plane within seconds without a restart. The policies and where they are
|
||||
attached are kept. While off, `llmgw ctl status`, `llmgw ctl guard list` and
|
||||
the guardrails page say so. To stop one policy rather than all, give it
|
||||
`mode: off`, or `mode: shadow` to have it record what it would do.
|
||||
|
||||
## A policy
|
||||
|
||||
```yaml
|
||||
guard:
|
||||
# Applied to every request.
|
||||
default: [baseline]
|
||||
policies:
|
||||
baseline:
|
||||
input:
|
||||
pii: {secret: block, card: redact}
|
||||
injection: flag
|
||||
output:
|
||||
pii: {secret: block}
|
||||
finance:
|
||||
mode: enforce # enforce (default), shadow, off
|
||||
fail: closed # open (default) or closed, for checks that can fail
|
||||
no_store: true # no conversations, notes, prefix routing or external models
|
||||
max_output_tokens: 2048
|
||||
system_prompt: "You answer questions about Acme's products only."
|
||||
drop_client_system: true
|
||||
stream: sentence # sentence (default) or buffer_all
|
||||
require_capabilities: true # no tools/images to models whose catalog entry lacks them
|
||||
validate_structured: true # answers must fit response_format's json_schema
|
||||
input:
|
||||
pii: {email: redact, phone: redact, iban: block, tckn: redact}
|
||||
injection: block
|
||||
max_chars: 40000
|
||||
deny:
|
||||
- {name: codename, regex: "(?i)project falcon", action: block}
|
||||
output:
|
||||
pii: {iban: redact, tckn: redact, secret: block}
|
||||
tools:
|
||||
allow: [search, fetch]
|
||||
deny: [shell]
|
||||
on_denied: strip # strip (default) or block
|
||||
max_calls_per_turn: 4
|
||||
validate_arguments: true
|
||||
```
|
||||
|
||||
**What is found.**
|
||||
|
||||
- **Personal data:** email, phone, iban, card, tckn, ip. Each is checked
|
||||
before it counts: IBANs by their mod-97 check, cards by Luhn and issuer
|
||||
prefix, TCKNs by their two check digits, phone numbers only when written
|
||||
to be dialled (`+90…`, `00…`, `0532…`).
|
||||
- **Secrets** (`secret`): OpenAI- and Anthropic-style keys, AWS access key
|
||||
ids, GitHub, Slack and Google tokens, JWTs, PEM private keys and llmgw's
|
||||
own keys.
|
||||
- **Prompt injection:** instructions to drop instructions (English and
|
||||
Turkish), chat-template control tokens in user text, invisible Unicode,
|
||||
long encoded blobs. These are found through disguises: fullwidth letters,
|
||||
Cyrillic look-alikes, zero-width characters, and base64, hex and URL
|
||||
encoding two levels deep.
|
||||
- **`deny` patterns:** matched against the text as written, and as a model
|
||||
reads it once those disguises are removed.
|
||||
|
||||
**What is done.**
|
||||
|
||||
- **`flag`** lets the request through and records it.
|
||||
- **`redact`** rewrites the value as a placeholder (`[EMAIL_1]`; the same
|
||||
value is the same placeholder throughout a request).
|
||||
- **`block`** stops it.
|
||||
|
||||
**Shapes.**
|
||||
|
||||
- `system_prompt` is put first.
|
||||
- `drop_client_system` removes the caller's own system and developer
|
||||
messages.
|
||||
- `max_output_tokens` is set where a request asks for none, and lowers one
|
||||
that asks for more.
|
||||
- A tool `deny`/`allow` list strips the tool from the request, or refuses
|
||||
it.
|
||||
- In the answer, a call to a tool that is not allowed, more calls than
|
||||
`max_calls_per_turn`, or arguments that do not fit the parameters the
|
||||
request declared (`validate_arguments`) stop it.
|
||||
- `require_capabilities` refuses tools offered to a model whose catalog
|
||||
entry does not say `capabilities.tools`, and images to one without
|
||||
`capabilities.vision`. It is opt-in because an entry that says nothing is
|
||||
read as "cannot".
|
||||
- `validate_structured` checks an answer against the `json_schema` the
|
||||
request asked for in `response_format`. Buffered, or streamed with
|
||||
`buffer_all`, an answer that does not fit is stopped; streamed a sentence
|
||||
at a time it has gone out already, and is flagged.
|
||||
|
||||
## A judge model
|
||||
|
||||
Patterns find what has a shape. For what does not — violence, crime,
|
||||
sexual content, self-harm, jailbreaks, whatever a moderation model was
|
||||
trained on — a policy can have a **judge model on the cluster** read the
|
||||
text.
|
||||
|
||||
The judge must be in the catalog under the name the policy gives, as a
|
||||
judge (`role: guard`), with a replica ready. When it is not, what the
|
||||
policy judges is refused with `guard_unavailable` (`fail: closed`) or served
|
||||
unjudged (`fail: open`), and `llmgw ctl status` and the console's Overview
|
||||
say which judge and why.
|
||||
|
||||
### Which judge
|
||||
|
||||
| judge | runs on | size | |
|
||||
|---|---|---|---|
|
||||
| `qwen3guard-0.6b` | CPUs | 0.8 GB | the default: about 2 s a check on two threads |
|
||||
| `qwen3guard-4b` | a GPU | 2.5 GB | for a cluster with a card to spare |
|
||||
| `qwen3guard-8b` | a GPU | 4.7 GB | |
|
||||
| `llama-guard-3-1b` | CPUs | 0.9 GB | the earlier default: refuses ordinary coding requests |
|
||||
| `llama-guard-3-8b` | a GPU | 4.6 GB | |
|
||||
|
||||
Measured on the reference cluster, as the judge's probability of "unsafe"
|
||||
for ordinary coding requests against harmful ones, Llama Guard 3 1B could
|
||||
not tell them apart: "continue" 0.51, "continue the code above" 0.89, "how
|
||||
do I kill all child processes in bash" 0.92, against 0.97 to 1.00 for
|
||||
harmful requests, at any quantisation. Qwen3Guard-Gen 0.6B gave the
|
||||
ordinary ones 0.15 at the most and the harmful ones 1.00. Its verdicts are
|
||||
safe, unsafe and **controversial**; controversial counts as safe unless the
|
||||
policy says `controversial: unsafe`.
|
||||
|
||||
A policy that names one of these is enough: the leader fetches its weights
|
||||
into the mirror and serves it as a judge, and serves it again if it is
|
||||
removed from the catalog. Another judge is served by hand with `--repo`.
|
||||
|
||||
```
|
||||
llmgw ctl guard judge status # each judge the policies ask, and what is wrong with it
|
||||
llmgw ctl guard judge enable # Qwen3Guard-Gen 0.6B on CPUs, or the judge the policies ask that is missing
|
||||
llmgw ctl guard judge enable --gpu # Qwen3Guard-Gen 4B on a GPU
|
||||
llmgw ctl guard judge use qwen3guard-4b qwen3guard-0.6b # the order every policy asks them in
|
||||
llmgw ctl guard judge enable --replace # when a model of that name is there, but not as a judge
|
||||
```
|
||||
|
||||
or **Do → Serve the guardrails' judge** in the console.
|
||||
|
||||
### A GPU and CPUs together
|
||||
|
||||
`judge.models` is a list asked in turn: the first judge that answers gives
|
||||
the verdict. The next is asked when one has no replica serving, fails, or
|
||||
does not answer within `latency_budget`; a judge with no replica at all is
|
||||
passed over at once. A judge on a GPU first and one on CPUs after it is
|
||||
the GPU's speed while it has its card, and the CPUs' when it does not.
|
||||
|
||||
```yaml
|
||||
judge:
|
||||
models: [qwen3guard-4b, qwen3guard-0.6b]
|
||||
input: block
|
||||
output: block
|
||||
latency_budget: 5s # each judge's wait
|
||||
```
|
||||
|
||||
A judge on a GPU is served as a `standard` model, on one card: it takes a
|
||||
card only from what may give one up, and waits while every card is busy,
|
||||
the judge on CPUs standing in. A judge on CPUs is `critical`, with two
|
||||
replicas, because every judged request waits on one. `llmgw ctl guard judge
|
||||
status` says, for each judge, which is asked when it gives no verdict.
|
||||
|
||||
### When it refuses ordinary requests
|
||||
|
||||
A refusal by a judge names the hazard it saw — the client is told `refused
|
||||
by the gateway's policy (judge: Violent)` — and the same is kept with the
|
||||
request: `llmgw ctl request <id>`, the server's log (`guard_rule`) and the
|
||||
guardrails page show it. To see what the judge says of any text, without
|
||||
anything being kept:
|
||||
|
||||
```
|
||||
llmgw ctl guard judge ask 'How do I sort a list in Python?'
|
||||
# qwen3guard-0.6b: safe, 1900 ms
|
||||
# it answered: "Safety: Safe\nCategories: None"
|
||||
```
|
||||
|
||||
A judge that says `unsafe` to ordinary text is usually Llama Guard 3 1B
|
||||
(`llmgw ctl guard judge use qwen3guard-0.6b`), the wrong weights (another
|
||||
repository's, without the judge's template: `judge enable --replace` serves
|
||||
the right ones), or a policy counting every category: give its judge
|
||||
`categories`. The baseline counts Violent, Non-violent Illegal Acts, Sexual
|
||||
Content or Sexual Acts, Suicide & Self-Harm and Jailbreak; PII is the
|
||||
patterns', and copyright, unethical acts and politically sensitive topics
|
||||
refuse too much ordinary work. While finding out, `judge.input: flag`
|
||||
records instead of refusing.
|
||||
|
||||
A judge's file served from the mirror page — Llama Guard, Qwen3Guard,
|
||||
ShieldGemma — is served as a judge too, under the name the policies ask.
|
||||
By hand, the catalog entry is:
|
||||
|
||||
```yaml
|
||||
# a catalog model (server.yaml models:, or ctl model add)
|
||||
- name: qwen3guard-0.6b
|
||||
source: http://mirror:8081/qwen3guard-0.6b/…/Qwen3Guard-Gen-0.6B.Q8_0.gguf
|
||||
role: guard
|
||||
guard:
|
||||
format: qwen3guard # qwen3guard, llamaguard3, shieldgemma or json
|
||||
scaling: {min: 2, max: 2}
|
||||
profiles: [{name: cpu, engine: llamacpp}]
|
||||
|
||||
# in a guard policy
|
||||
judge:
|
||||
models: [qwen3guard-0.6b] # or model: one judge
|
||||
input: block # flag or block; empty does not ask about input
|
||||
output: flag
|
||||
categories: [Violent, Jailbreak] # only these count; empty: any unsafe verdict
|
||||
controversial: safe # or unsafe
|
||||
latency_budget: 250ms # each judge's wait; after the last, fail: open serves (flagged), closed refuses
|
||||
```
|
||||
|
||||
A model with `role: guard` is the guardrails' alone:
|
||||
|
||||
- it is not in `/v1/models` or the chat page, and is refused (404) to
|
||||
every key, as a name and as a fallback;
|
||||
- it must always run — on demand is refused. One on CPUs is `critical`
|
||||
and should have two replicas on different nodes, because every judged
|
||||
request waits on it.
|
||||
|
||||
The gateway's calls to it go past admission, accounting and the guardrails
|
||||
themselves.
|
||||
|
||||
- **Formats:**
|
||||
- `qwen3guard` asks in Qwen3Guard's own words and reads "Safety: Safe |
|
||||
Unsafe | Controversial" and its "Categories:".
|
||||
- `llamaguard3` reads Meta's "safe" / "unsafe\nS1,S10"; give the catalog
|
||||
entry `categories` to have S1 recorded as its name.
|
||||
- `shieldgemma` asks Google's model whether the text breaks the listed
|
||||
principles (`categories`, each a sentence) and reads Yes or No.
|
||||
- `json` expects `{"safe": false, "categories": [...]}`.
|
||||
- **What is read:**
|
||||
- **Input:** the request's last user turn, after the patterns.
|
||||
- **Output:** the answer, together with that user turn. A streamed answer
|
||||
is read every 600 characters and at the end, so text between two
|
||||
readings has gone out already; `stream: buffer_all` holds it until it
|
||||
is read.
|
||||
- **A long text** is read by its start and its end (about 6000
|
||||
characters), so a message with a file pasted into it fits the judge's
|
||||
context rather than overflowing it and going unjudged.
|
||||
- **Cache:** the same text is not asked about twice within five minutes.
|
||||
|
||||
## Where a policy applies
|
||||
|
||||
`guard.default` applies everywhere. Others are attached to a part of the
|
||||
tenancy tree, and apply to everything below it:
|
||||
|
||||
```
|
||||
llmgw ctl guard attach project:Acme/finance/payroll finance
|
||||
llmgw ctl guard attach key:key-7f3a baseline
|
||||
llmgw ctl guard list
|
||||
llmgw ctl guard detach project:Acme/finance/payroll finance
|
||||
```
|
||||
|
||||
Policies narrow, as limits do. A policy attached lower can tighten what
|
||||
applies above it, never loosen it. Combined, they take:
|
||||
|
||||
- the strictest mode and failure;
|
||||
- for each kind, the strictest action;
|
||||
- every deny pattern and denied tool;
|
||||
- only the tools every allow list allows;
|
||||
- the smallest limits;
|
||||
- `no_store` if any of them asks;
|
||||
- every pinned prompt, widest first.
|
||||
|
||||
An attachment applies at once.
|
||||
|
||||
## What a client sees
|
||||
|
||||
**A refused request** is answered with 400:
|
||||
|
||||
```json
|
||||
{"error": {"type": "invalid_request_error", "code": "content_policy_violation",
|
||||
"message": "this request was refused by the gateway's policy (secret)"}}
|
||||
```
|
||||
|
||||
It names the category, never the text, and costs nothing.
|
||||
|
||||
**A stopped answer** ends with `finish_reason: "content_filter"`. A stream
|
||||
ends with that chunk and `[DONE]`, and the engine is cancelled. The
|
||||
Responses API reports it as `incomplete` with
|
||||
`incomplete_details.reason: "content_filter"`.
|
||||
|
||||
**Headers:** `x-llmgw-guard` (`flagged`, `redacted`, `blocked`) and
|
||||
`x-llmgw-guard-policy` say what was done.
|
||||
|
||||
**Streams.** A streamed answer is judged a sentence at a time. Text is held
|
||||
until a sentence, a line or 200 characters ends, released up to its last
|
||||
whitespace — so a value still arriving is never sent in two halves — and
|
||||
rewritten if needed. **What was released cannot be taken back**: a policy
|
||||
that must never let a word through uses `stream: buffer_all`, which holds
|
||||
the whole answer until it is judged.
|
||||
|
||||
## Shadow mode
|
||||
|
||||
`mode: shadow` judges and records — the usage row says `shadow-blocked`,
|
||||
`shadow-redacted` — and changes nothing. Run a new policy in shadow first,
|
||||
and read what it would have done before enforcing it.
|
||||
|
||||
## What is recorded
|
||||
|
||||
Each request's usage row has the action, policy, category and the time the
|
||||
checks took. The access log has the action, stage and policy. Neither has
|
||||
the text, a matched value, or the placeholder map: the map exists only in
|
||||
the gateway's memory while the request is served.
|
||||
|
||||
## In the console
|
||||
|
||||
`sudo llmgw` has a **Guardrails** view: each policy with its mode, what it
|
||||
does when it cannot judge, and what it checks; whether each judge model it
|
||||
leans on has replicas ready — a policy that fails closed refuses everything
|
||||
while its judge has none; where policies are attached, with A to attach one
|
||||
and Enter on a row to detach it; and what they did in the last day (W for an
|
||||
hour or a week), totals first and the newest requests after, each with the
|
||||
request id `llmgw ctl request` reads. The same is at `GET /api/v1/guard` and
|
||||
`GET /api/v1/guard/events?since=24h`, and on the portal's Guardrails page.
|
||||
+1477
-88
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,349 @@
|
||||
# Knowledge bases
|
||||
|
||||
A knowledge base is a collection of a company's documents that the cluster
|
||||
can answer from. People put their documents into it from their own machines
|
||||
with `llmgw-kb`. The tool reads the files on that machine, cuts them into
|
||||
passages ("chunks"), has the gateway embed each passage with the person's own
|
||||
key, and sends the text and the vectors. The files themselves never leave
|
||||
the machine.
|
||||
|
||||
## Why retrieval, and not training
|
||||
|
||||
Retrieval is the tool for *knowledge*: policies, procedures, contracts. The
|
||||
answer can name its source, a changed document is sent again in minutes, and
|
||||
a passage only ever reaches people allowed to read its collection.
|
||||
|
||||
Training a model on the documents (LoRA) changes how it *writes*, not what it
|
||||
reliably *knows*. It has no access control: whatever goes into the weights
|
||||
is available to everyone who may use them.
|
||||
|
||||
## What it needs
|
||||
|
||||
- **The PostgreSQL store with pgvector.** On a SQLite store knowledge bases
|
||||
are off, and `llmgw ctl kb list` and the console say why.
|
||||
- `llmgw server init` and the high-availability database setup install
|
||||
pgvector and enable it.
|
||||
- A cluster set up before this build needs it by hand, once:
|
||||
1. Install the package on every database machine:
|
||||
`sudo apt-get install postgresql-16-pgvector` (use the server's
|
||||
version number).
|
||||
2. On the leader, enable it in llmgw's database:
|
||||
`sudo -u postgres psql -d llmgw -c 'CREATE EXTENSION vector'`.
|
||||
3. Restart the control planes: `llmgw restart`.
|
||||
- **An embedding model in the catalog**, with `capabilities.embeddings`.
|
||||
Two good choices handle Turkish and run on CPUs: `BAAI/bge-m3` (1024
|
||||
dimensions) and `Qwen/Qwen3-Embedding-0.6B`.
|
||||
|
||||
Add it as any model is added (**Serve a model** in the console, or the
|
||||
catalog), with its entry saying that it embeds:
|
||||
|
||||
```yaml
|
||||
capabilities:
|
||||
embeddings: true
|
||||
```
|
||||
|
||||
A collection stays with the model it was made with: vectors from two
|
||||
models cannot be compared.
|
||||
|
||||
## For administrators
|
||||
|
||||
### Make a collection
|
||||
|
||||
A collection belongs to one part of the tenancy tree: an organization, a
|
||||
team, a project or a person. Whoever is in the owner, at any level below it,
|
||||
may read it and add to it.
|
||||
|
||||
```sh
|
||||
llmgw ctl kb create team:Acme/hr policies bge-m3
|
||||
llmgw ctl kb create user:USER-ID notes bge-m3 --org Acme
|
||||
```
|
||||
|
||||
The model's dimension is asked of the model itself, so the model has to be
|
||||
able to serve when the collection is made. `--dims 1024` gives it instead.
|
||||
|
||||
In the console: **Knowledge**, then **A**.
|
||||
|
||||
### Who may read and add
|
||||
|
||||
The owner, and everything below it, may read and add. A **grant** gives the
|
||||
same to another part of the same organization, to a person, or to one key:
|
||||
|
||||
```sh
|
||||
llmgw ctl kb grant Acme/policies project:Acme/sales/crm # read
|
||||
llmgw ctl kb grant Acme/policies team:Acme/legal --write # read and add
|
||||
llmgw ctl kb revoke Acme/policies project:Acme/sales/crm
|
||||
```
|
||||
|
||||
How access works:
|
||||
|
||||
- **Grants add up down the tree.** A grant to a team reaches its projects
|
||||
and their keys.
|
||||
- **Nothing crosses organizations.** A grant to another organization's
|
||||
team, project, key or person is refused.
|
||||
- **A person's own key** (a key with no project) reads nothing by itself.
|
||||
The person names the project they act in (`llmgw-kb login` asks for it,
|
||||
or `--project`) and must be a member of it. Joining a second team
|
||||
therefore does not quietly widen what every key of theirs can read.
|
||||
- **No access looks like no collection.** A collection a key may not read
|
||||
is answered exactly as one that does not exist, so even its name does not
|
||||
leak.
|
||||
|
||||
### Answer from collections by default
|
||||
|
||||
Clients that cannot add a header still get answers from the documents when
|
||||
their part of the tree has a default:
|
||||
|
||||
```sh
|
||||
llmgw ctl kb default team:Acme/hr Acme/policies Acme/forms
|
||||
llmgw ctl kb default key:KEY-ID Acme/policies
|
||||
llmgw ctl kb default team:Acme/hr --clear
|
||||
```
|
||||
|
||||
- **Which default applies.** The narrowest one: a key's over its project's,
|
||||
a project's over its team's, a team's over its organization's.
|
||||
- **A request can override it.** A request that names collections uses
|
||||
those instead, and `X-Llmgw-Knowledge: none` turns the default off for
|
||||
that request.
|
||||
- **A default never refuses a request.** A collection the caller may not
|
||||
read is left out for them. If the search fails (the embedding model is
|
||||
down, say), the request is answered without passages and the failure is
|
||||
logged.
|
||||
- **Changes take effect within 30 seconds.**
|
||||
|
||||
In the console: **Knowledge**, Enter on a collection, then **Answer from it
|
||||
by default…**.
|
||||
|
||||
### Look after it
|
||||
|
||||
```sh
|
||||
llmgw ctl kb list # collections, owners, documents, grants
|
||||
llmgw ctl kb docs Acme/policies # its documents, who added each and when
|
||||
llmgw ctl kb rm Acme/policies old/form.pdf # one document out
|
||||
llmgw ctl kb search Acme/policies how many days of annual leave
|
||||
llmgw ctl kb delete Acme/policies # the collection, its documents and grants
|
||||
```
|
||||
|
||||
Rules that apply:
|
||||
|
||||
- A project, team or organization that owns a collection is not deleted
|
||||
until its collections are. This is the same rule as for one with members.
|
||||
- A person's own collections go with their account (`users delete`,
|
||||
`users forget`).
|
||||
- Every document added or removed is in the audit log (`llmgw ctl audit`),
|
||||
by name. The text is never recorded.
|
||||
|
||||
## For everyone: `llmgw-kb`
|
||||
|
||||
### Get it
|
||||
|
||||
Download `llmgw-kb` for your system from the release page, next to `llmgw`:
|
||||
|
||||
| System | File |
|
||||
|---|---|
|
||||
| Windows (Intel/AMD) | `llmgw-kb-VERSION-windows-amd64.exe` |
|
||||
| Windows (ARM) | `llmgw-kb-VERSION-windows-arm64.exe` |
|
||||
| macOS (Apple silicon) | `llmgw-kb-VERSION-darwin-arm64` |
|
||||
| macOS (Intel) | `llmgw-kb-VERSION-darwin-amd64` |
|
||||
| Linux | `llmgw-kb-VERSION-linux-amd64` / `-arm64` |
|
||||
|
||||
Rename it to `llmgw-kb` (`llmgw-kb.exe` on Windows) and put it somewhere on
|
||||
your PATH. You do not need to install anything else.
|
||||
|
||||
On macOS, a downloaded program that is not signed is stopped the first time.
|
||||
Either right-click it and choose Open, or run
|
||||
`xattr -d com.apple.quarantine llmgw-kb`.
|
||||
|
||||
### Log in once
|
||||
|
||||
```sh
|
||||
llmgw-kb login
|
||||
```
|
||||
|
||||
It asks for three things, checks them, and saves them where only you can
|
||||
read them:
|
||||
|
||||
- **The gateway's address.**
|
||||
- **Your API key.** It is typed without being shown.
|
||||
- **A project.** Needed only for a personal key.
|
||||
|
||||
Instead of logging in you can set `LLMGW_URL`, `LLMGW_API_KEY` and
|
||||
`LLMGW_PROJECT`.
|
||||
|
||||
```sh
|
||||
llmgw-kb collections # what your key may read, and which you may add to
|
||||
```
|
||||
|
||||
### Send your documents
|
||||
|
||||
```sh
|
||||
llmgw-kb sync policies ~/Documents/HR
|
||||
llmgw-kb sync policies "C:\Users\ayse\Documents\İK" --include "*.pdf,*.docx"
|
||||
llmgw-kb sync policies ./handbook ./forms --exclude "drafts/**"
|
||||
```
|
||||
|
||||
Run it as often as you like. The first run sends everything. Later runs
|
||||
send only what changed, and do not even read files whose size and time have
|
||||
not changed.
|
||||
|
||||
| Option | What it does |
|
||||
|---|---|
|
||||
| `--dry-run` | Reads and cuts everything, sends nothing, and lists what would leave the machine and how much. |
|
||||
| `--delete` | Also removes documents whose files are gone. Only those this machine sent, from these folders; never anything someone else sent. |
|
||||
| `--retry` | Reads again files skipped before, for example after installing LibreOffice. |
|
||||
| `--include`, `--exclude` | Glob patterns. `*.pdf` matches a file name; `reports/**/*.xlsx` matches a path. |
|
||||
| `--hidden` | Reads hidden files and folders too (left out by default). |
|
||||
| `-v` | Lists unchanged files too. |
|
||||
|
||||
Documents are named by their path under the folder, with forward slashes,
|
||||
on every system. With several folders, each folder's own name comes first.
|
||||
|
||||
```sh
|
||||
llmgw-kb status policies # what this machine sent and skipped, and why
|
||||
llmgw-kb history # the latest runs
|
||||
llmgw-kb formats # what can be read here
|
||||
```
|
||||
|
||||
### What it reads
|
||||
|
||||
| Read here | Formats |
|
||||
|---|---|
|
||||
| Documents | PDF (text), Word (`.docx`), Excel (`.xlsx`), PowerPoint (`.pptx`, with speaker notes), OpenDocument (`.odt`, `.ods`, `.odp`), RTF, EPUB |
|
||||
| Web and mail | HTML, MHTML, e-mail (`.eml`, with its attachments) |
|
||||
| Text | Markdown, plain text, CSV, JSON, YAML, XML, subtitles, notebooks (`.ipynb`), source code |
|
||||
| Archives | ZIP, with every file inside it read in turn |
|
||||
|
||||
| Read with a helper on your machine | Helper |
|
||||
|---|---|
|
||||
| Word, Excel and PowerPoint 97–2003 (`.doc`, `.xls`, `.ppt`), Pages, Numbers, Keynote and others | [LibreOffice](https://www.libreoffice.org) |
|
||||
| Images and scanned PDFs | [Tesseract](https://tesseract-ocr.github.io) (with the Turkish language data), plus Poppler's `pdftoppm` for PDFs |
|
||||
| PDFs the built-in reader cannot parse | Poppler's `pdftotext` |
|
||||
|
||||
`llmgw-kb formats` says which helpers it found. Text in UTF-8, UTF-16, or
|
||||
the Windows Turkish and Western code pages is read correctly.
|
||||
|
||||
A file that cannot be read is skipped with the reason, and the run goes on.
|
||||
Reasons include a locked Office file, a format with no helper, and a
|
||||
scanned page with no OCR.
|
||||
|
||||
### What leaves your machine
|
||||
|
||||
For each new or changed document, two things leave the machine, both over
|
||||
the gateway's TLS connection:
|
||||
|
||||
- **The passages, to be embedded.** They go through `/v1/embeddings` like
|
||||
any request of yours. Your key's limits, budget and guardrails apply.
|
||||
- **The document itself.** Its name, its file's SHA-256 digest and size,
|
||||
and the passages with their vectors.
|
||||
|
||||
The file itself is never sent. `--dry-run` shows exactly what would be.
|
||||
|
||||
### The ledger
|
||||
|
||||
`llmgw-kb` keeps a small SQLite database of what it did with every
|
||||
document: the file it came from, its size, time and digest, and whether it
|
||||
was sent, skipped or failed, with the reason.
|
||||
|
||||
| System | Where it is |
|
||||
|---|---|
|
||||
| Windows | `%AppData%\llmgw-kb\ledger.db` |
|
||||
| macOS | `~/Library/Application Support/llmgw-kb/ledger.db` |
|
||||
| Linux | `~/.config/llmgw-kb/ledger.db` |
|
||||
|
||||
Deleting the ledger is safe: the next run hashes every file once, finds what
|
||||
the collection already has, and sends nothing that is unchanged. `--ledger`
|
||||
uses another file.
|
||||
|
||||
## Asking questions of a collection
|
||||
|
||||
A chat request names the collections to answer from in a header. Any client
|
||||
that can add a header can do it: the OpenAI and Anthropic SDKs, Continue, or
|
||||
curl.
|
||||
|
||||
```sh
|
||||
curl https://llm.example.com/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "X-Llmgw-Knowledge: politikalar" \
|
||||
-d '{"model": "qwen2.5-coder-14b-instruct",
|
||||
"messages": [{"role": "user", "content": "Kaç gün yıllık iznim var?"}]}'
|
||||
```
|
||||
|
||||
```python
|
||||
client.chat.completions.create(
|
||||
model="qwen2.5-coder-14b-instruct",
|
||||
messages=[{"role": "user", "content": "Kaç gün yıllık iznim var?"}],
|
||||
extra_headers={"X-Llmgw-Knowledge": "politikalar"},
|
||||
)
|
||||
```
|
||||
|
||||
What the gateway does with such a request:
|
||||
|
||||
1. **It finds the collections.** It looks up the ones named, by name or ID,
|
||||
among those the key may read. A collection it may not read is answered
|
||||
404, the same as one that does not exist.
|
||||
2. **It embeds the question.** The last user message is embedded by the
|
||||
collection's model, on the caller's own account: their limits, budget and
|
||||
guardrails apply.
|
||||
3. **It searches.** The collections are searched by meaning and by words
|
||||
(below).
|
||||
4. **It adds the passages.** The best passages are numbered and added to the
|
||||
system message, with an instruction to answer from them, cite them as
|
||||
[1], [2], and say so when they do not hold the answer. The system message
|
||||
is created if the request has none. Passages that would not fit the
|
||||
model's context are left out, the weakest first.
|
||||
5. **It judges and sends the request.** The request, passages included, then
|
||||
goes through the guardrails and to the model as usual. A document
|
||||
carrying instructions aimed at its reader is judged like anything else
|
||||
the model reads.
|
||||
|
||||
On the portal's chat page, **Documents** picks a collection the person may
|
||||
read in the chosen project, the default, or none. Under each answer, the
|
||||
sources it was given are listed by number.
|
||||
|
||||
How it is shaped:
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| `X-Llmgw-Knowledge` | Collections, comma separated; `none` for no passages, whatever the default. |
|
||||
| `X-Llmgw-Knowledge-K` | How many passages, 1 to 20; 5 by default. |
|
||||
| `X-Llmgw-Project` | For a personal key: the project it acts in. |
|
||||
| `x-llmgw-sources` (answer) | The passages given, as `document#chunk`, in their numbered order; empty when nothing matched. |
|
||||
|
||||
Behaviour to know about:
|
||||
|
||||
- **No passages matched.** The model is told that nothing in the documents
|
||||
answers the question, and asked to say so rather than answer from general
|
||||
knowledge.
|
||||
- **Other APIs.** Requests through `/v1/responses` and `/v1/messages`
|
||||
(Codex, Claude Code) work the same way, since they are answered as chat
|
||||
completions.
|
||||
- **Separate prefix caches.** On vLLM replicas, a request given passages
|
||||
carries a `cache_salt` derived from its organization. One organization's
|
||||
passages then never share cached blocks with another's, so how fast an
|
||||
answer starts says nothing about what another tenant asked.
|
||||
- **No answer cache.** Answers from knowledge bases are never served from
|
||||
the answer cache, because two keys of one project can read different
|
||||
collections.
|
||||
|
||||
### How a passage is chosen
|
||||
|
||||
Each passage scores the following, and the highest scores win:
|
||||
|
||||
- its similarity in meaning to the question (cosine, 0 to 1), plus
|
||||
- up to 0.3 for the share of the question's words it contains.
|
||||
|
||||
Words almost every passage has ("ne", "için", "the", "of") are not counted.
|
||||
The meaning does most of the work: "how much leave do fathers get" finds the
|
||||
paternity page, not the annual-leave page that shares more of its words.
|
||||
The word match helps with what meaning cannot carry, such as a form number,
|
||||
a product code or a name.
|
||||
|
||||
Metrics: `llmgw_kb_retrievals_total{result}` (found, nothing, refused) and
|
||||
`llmgw_kb_retrieval_seconds`.
|
||||
|
||||
## How it is kept apart
|
||||
|
||||
Passages live in PostgreSQL, partitioned by organization, with an index per
|
||||
collection. Every search and every write names the organization and the
|
||||
collections the caller may use. On top of that, row-level security on the
|
||||
table checks the same thing again, so a query that forgot its filter
|
||||
returns nothing rather than another tenant's passages.
|
||||
@@ -0,0 +1,228 @@
|
||||
# llmgw and Kubernetes
|
||||
|
||||
The control plane runs in Kubernetes, from the Helm chart in
|
||||
`deploy/helm/llmgw`. The GPU nodes do not: they stay ordinary machines,
|
||||
joined with `llmgw node add` as on any other cluster. This page says why,
|
||||
and how the two meet.
|
||||
|
||||
## What goes where
|
||||
|
||||
| Part | Where | Why |
|
||||
|---|---|---|
|
||||
| Control plane: the API on 8080, the agents' port 7443, the scheduler | A StatefulSet, from the chart | A Go binary with a store and three ports: what Kubernetes runs well |
|
||||
| Store | A file on a PersistentVolumeClaim (SQLite), or your PostgreSQL | Either, as on a host |
|
||||
| GPU and CPU nodes | Outside, as hosts | See below |
|
||||
| Model mirror | Outside, on a machine with a large disk | Every node pulls tens of gigabytes from it over plain HTTP; a pod with a volume that size is a host in all but name |
|
||||
|
||||
### Why the nodes stay outside
|
||||
|
||||
The agent on a node is not a service that happens to use a GPU. It installs
|
||||
the engines on the machine — the llama.cpp build for its driver, vLLM in a
|
||||
Python environment — checks them against their publisher's checksums and
|
||||
runs them; it chooses which GPUs each replica gets, by index and by NVLink
|
||||
topology; it starts engines as processes that outlive its own restarts, so a
|
||||
node upgrades while it serves; it writes the machine's firewall; and it
|
||||
replaces its own binary when the cluster upgrades. All of that is root on
|
||||
the host, which is exactly what a container is built to prevent.
|
||||
|
||||
Run as a DaemonSet it would need to be privileged, with the host's network,
|
||||
process namespace and filesystem, and the NVIDIA device plugin would have to
|
||||
be kept off its GPUs — or the two would hand the same card to two owners.
|
||||
What is left is the host agent with a container's indirection around it, and
|
||||
none of the scheduling, isolation or restart guarantees that make a pod
|
||||
worth having. So the chart does not ship one, and this is the
|
||||
recommendation:
|
||||
|
||||
- **Keep GPU machines out of Kubernetes**, and add them with `node add`.
|
||||
- **If they are already Kubernetes workers**, dedicate them: taint them so
|
||||
no other pod lands there, keep the NVIDIA device plugin off them (most
|
||||
installs select nodes by label — leave the label off), and run the llmgw
|
||||
agent on the host, where `node add` puts it. llmgw then owns the GPUs, and
|
||||
Kubernetes schedules nothing that would want them.
|
||||
|
||||
## The image
|
||||
|
||||
The project publishes binaries, not an image. Build one from the repository
|
||||
and push it where the cluster pulls from — at an offline site, the internal
|
||||
registry:
|
||||
|
||||
```
|
||||
docker build -f deploy/container/Dockerfile --build-arg VERSION=v0.11.1 \
|
||||
-t registry.internal/llmgw:v0.11.1 .
|
||||
docker push registry.internal/llmgw:v0.11.1
|
||||
```
|
||||
|
||||
It is the static binary on a distroless base: no shell, running as user
|
||||
65532, with a read-only root filesystem in the pod.
|
||||
|
||||
## The agent port's certificate
|
||||
|
||||
Agents and the control plane authenticate each other with certificates from
|
||||
llmgw's own authority, and the TLS ends in the pod — so the authority's key
|
||||
stays outside the cluster, on the machine you will add nodes from, and the
|
||||
pod is given only its own certificate.
|
||||
|
||||
Decide first the address agents will dial: the `-agents` Service is a
|
||||
LoadBalancer, so give it a fixed `agentService.loadBalancerIP` or a DNS name
|
||||
that points at it. That address has to be in the certificate:
|
||||
|
||||
```
|
||||
llmgw ctl cert init-ca --dir ./pki
|
||||
llmgw ctl cert issue-server llmgw-ctl --host 10.0.0.200 --host llmgw-agents.example.com \
|
||||
--dir ./pki --out ./pki
|
||||
kubectl -n llmgw create secret generic llmgw-agent-tls \
|
||||
--from-file=ca.pem=pki/ca.pem \
|
||||
--from-file=server.pem=pki/server.pem \
|
||||
--from-file=server-key.pem=pki/server-key.pem
|
||||
```
|
||||
|
||||
Keep `pki/ca-key.pem` as you would on a host: whoever holds it can make any
|
||||
node's identity.
|
||||
|
||||
## Installing
|
||||
|
||||
```yaml
|
||||
# values.yaml
|
||||
image:
|
||||
repository: registry.internal/llmgw
|
||||
tag: v0.11.1
|
||||
agentTLS:
|
||||
existingSecret: llmgw-agent-tls
|
||||
agentService:
|
||||
loadBalancerIP: 10.0.0.200
|
||||
ingress:
|
||||
enabled: true
|
||||
className: nginx
|
||||
hosts: [{host: llm.example.com, paths: [{path: /, pathType: Prefix}]}]
|
||||
tls: [{secretName: llm-example-com, hosts: [llm.example.com]}]
|
||||
config:
|
||||
mirror:
|
||||
api: http://models.internal:8082
|
||||
# The address the router reaches nodes' engines from is a Kubernetes
|
||||
# node's, once the pod's traffic leaves the cluster: the nodes' firewalls
|
||||
# must let it in.
|
||||
firewall_networks: ["10.0.1.0/24"]
|
||||
```
|
||||
|
||||
```
|
||||
helm install llmgw deploy/helm/llmgw -n llmgw -f values.yaml
|
||||
```
|
||||
|
||||
`config` is `server.yaml`: any key of [reference.md](reference.md), laid
|
||||
under what the chart sets itself — the listeners, where the certificates,
|
||||
the admin token and the store are mounted. A changed `config` restarts the
|
||||
pod.
|
||||
|
||||
**The store.** By default SQLite, on a 10 GiB claim. For PostgreSQL, put its
|
||||
address in a Secret and point the chart at it:
|
||||
|
||||
```
|
||||
kubectl -n llmgw create secret generic llmgw-store \
|
||||
--from-literal=dsn='postgres://llmgw:…@pg.internal:5432/llmgw?sslmode=verify-full'
|
||||
```
|
||||
|
||||
```yaml
|
||||
store:
|
||||
type: postgres
|
||||
postgres:
|
||||
existingSecret: llmgw-store
|
||||
```
|
||||
|
||||
**One control plane.** The chart runs one replica. Two against PostgreSQL
|
||||
would elect a leader and keep a standby, but a standby serves nothing and is
|
||||
never ready, so a StatefulSet rollout would wait on it for ever. Kubernetes
|
||||
replaces the one pod within seconds when it goes, and agents reconnect on
|
||||
their own; the replicas on the nodes keep serving meanwhile.
|
||||
|
||||
**Probes.** Alive is `/health`. Ready is `/leader`, which on PostgreSQL waits
|
||||
for the lease of the pod it replaced to run out — not `/ready`, which waits
|
||||
for a model replica, and would keep a cluster with no node yet from ever
|
||||
accepting one.
|
||||
|
||||
**Monitoring.** The metrics port, 9091, serves metrics and nothing else;
|
||||
`serviceMonitor.enabled: true` makes a ServiceMonitor for the Prometheus
|
||||
Operator, and `deploy/monitoring` holds the rules and dashboards.
|
||||
|
||||
## Using it
|
||||
|
||||
The admin API listens on the pod's loopback, as on a host — it can dump keys
|
||||
from the heap when profiling is on. Reach it through the pod:
|
||||
|
||||
```
|
||||
kubectl -n llmgw exec -it statefulset/llmgw -- llmgw users create you@example.com --password '…'
|
||||
kubectl -n llmgw exec -it statefulset/llmgw -- llmgw users role you@example.com admin
|
||||
kubectl -n llmgw exec -it statefulset/llmgw -- llmgw ctl status
|
||||
```
|
||||
|
||||
or from your machine, through a port-forward:
|
||||
|
||||
```
|
||||
kubectl -n llmgw port-forward statefulset/llmgw 9090:9090 &
|
||||
export LLMGW_ADMIN_TOKEN=$(kubectl -n llmgw get secret llmgw-admin -o jsonpath='{.data.admin-token}' | base64 -d)
|
||||
llmgw ctl status
|
||||
```
|
||||
|
||||
The token is made by the chart and kept across upgrades
|
||||
(`adminToken.existingSecret` brings your own).
|
||||
|
||||
## Adding nodes
|
||||
|
||||
From the machine holding `pki/`, with the port-forward above running and
|
||||
`LLMGW_ADMIN_TOKEN` set:
|
||||
|
||||
```
|
||||
sudo -E llmgw node add ubuntu@gpu01 \
|
||||
--server https://10.0.0.200:7443 --pki ./pki --admin 127.0.0.1:9090
|
||||
```
|
||||
|
||||
`--server` is the address in the certificate. The node is installed, checked
|
||||
and started over ssh exactly as on any cluster ([node-onboarding.md](node-onboarding.md)),
|
||||
and dials the control plane in the pod. Its firewall comes from
|
||||
`firewall_networks` in the server configuration `node add` reads (`--config`,
|
||||
`/etc/llmgw/server.yaml` by default): copy the chart's out of the cluster
|
||||
first, or give `--no-firewall` and keep the engines' ports open to the
|
||||
Kubernetes nodes by hand.
|
||||
|
||||
```
|
||||
kubectl -n llmgw get configmap llmgw -o jsonpath='{.data.server\.yaml}' > server.yaml
|
||||
sudo -E llmgw node add ubuntu@gpu01 --config ./server.yaml --server https://10.0.0.200:7443 \
|
||||
--pki ./pki --admin 127.0.0.1:9090
|
||||
```
|
||||
|
||||
## Offline
|
||||
|
||||
An offline site ([install.md](install.md), "Offline sites") works the same
|
||||
way with the control plane in a pod. Import the bundle on the mirror's
|
||||
machine with `llmgw bundle import`, record its engines' checksums in a
|
||||
file, and give that file to the chart as a ConfigMap:
|
||||
|
||||
```
|
||||
llmgw bundle trust llmgw-bundle.tar --sha256 3f2a… --file ./offline.sha256
|
||||
kubectl -n llmgw create configmap llmgw-offline --from-file=offline.sha256
|
||||
```
|
||||
|
||||
```yaml
|
||||
offlineChecksums:
|
||||
configMap: llmgw-offline
|
||||
config:
|
||||
offline:
|
||||
mirror: http://models.internal:8081
|
||||
```
|
||||
|
||||
A later bundle: `bundle trust` into the same file, and replace the ConfigMap
|
||||
(`kubectl create configmap … --dry-run=client -o yaml | kubectl apply -f -`).
|
||||
The control plane reads the file again when it changes; there is no restart.
|
||||
|
||||
## What has and has not been tried
|
||||
|
||||
The chart is checked with `helm lint` and `helm template` for SQLite,
|
||||
PostgreSQL, an Ingress, a ServiceMonitor and an offline cluster, and refuses
|
||||
what would not work — no image, no certificate, a PostgreSQL store without
|
||||
its address, more than one replica. The image has been run the way the
|
||||
chart runs it — read-only root, user 65532, the same files mounted at the
|
||||
same paths, SQLite on a volume — and answered `/health`, `/leader` and
|
||||
`/metrics`, with `llmgw ctl` and `llmgw users` working through `exec`.
|
||||
|
||||
It has **not** been installed on a Kubernetes cluster, and no node has yet
|
||||
joined a control plane running in one. Until it has, a control plane on a
|
||||
host is the well-trodden path.
|
||||
+77
-16
@@ -65,6 +65,11 @@ model already in it:
|
||||
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.
|
||||
- **Claude Code** — not on the page yet; it speaks Anthropic's Messages API,
|
||||
which the gateway translates at `/v1/messages`:
|
||||
`ANTHROPIC_BASE_URL=https://<gateway> ANTHROPIC_AUTH_TOKEN=<key>
|
||||
ANTHROPIC_MODEL=<model> ANTHROPIC_SMALL_FAST_MODEL=<model> claude`. It
|
||||
needs a model that calls tools well.
|
||||
|
||||
For Continue and Codex there is a one-line command,
|
||||
`curl -fsSL <gateway>/setup/<tool>?model=… | sh`. Its script is served without
|
||||
@@ -304,6 +309,28 @@ 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.
|
||||
|
||||
## The docs portal
|
||||
|
||||
The gateway also serves its manual at `/docs`: these pages, rendered, with
|
||||
[clients.md](clients.md) — how to use the gateway from the official OpenAI
|
||||
and Anthropic SDKs — first on its index. The landing page and the account
|
||||
page link to it, so a user's questions about streaming, tools or the Anthropic
|
||||
SDK have an answer that is not you.
|
||||
|
||||
It is public, like the landing page. The client guide is for the person who
|
||||
has no account yet, and the operator documents describe llmgw in general —
|
||||
no address, model, node or key of this installation — which is what the
|
||||
project's repository publishes anyway. A site that would rather not say even
|
||||
that turns sign-in on:
|
||||
|
||||
```yaml
|
||||
landing:
|
||||
docs_sign_in: true
|
||||
```
|
||||
|
||||
A visitor is then sent to sign in first. On a gateway with API keys and no
|
||||
accounts, the portal asks for a key instead, like the API.
|
||||
|
||||
## What the user sees
|
||||
|
||||
They open the link, press the button, and get the key once. The page shows them
|
||||
@@ -325,8 +352,15 @@ llmgw keys list --store … --project default
|
||||
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.
|
||||
**To rotate**, `llmgw keys rotate <key-id> --overlap 48h` issues a new key with
|
||||
the same project, owner and limits and lets the old one work for the overlap,
|
||||
then stop. Watch `LAST USED` on the old one stop moving before it runs out.
|
||||
|
||||
**To cap how long keys live**, `llmgw keys ttl project:Acme/ml/train 720h` (or
|
||||
`team:`, `org:`; `off` removes it), and `auth.max_key_ttl` in server.yaml for
|
||||
every key. A key is made within the shortest limit that applies, and one asked
|
||||
for without an expiry gets that limit. `keys list --all` shows each key's kind
|
||||
— service (a project's), member, personal — and when it expires.
|
||||
|
||||
```
|
||||
llmgw keys revoke --store … key-mpmip77twpt7qaxc
|
||||
@@ -345,16 +379,38 @@ is what someone gets from a borrowed laptop, and the password is what they do no
|
||||
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:
|
||||
Someone who has forgotten theirs follows **Forgotten your password?** on the
|
||||
sign-in page. With mail set up in the cluster's settings, a link to set a new one
|
||||
is mailed to their address — the page says the same whether or not the address
|
||||
has an account, so nobody can find out who does from it:
|
||||
|
||||
```
|
||||
llmgw ctl config set mail '{smtp: smtp.example.com:587, from: llmgw@example.com,
|
||||
username: llmgw, password_file: /etc/llmgw/secrets/smtp, base_url: https://llm.example.com}'
|
||||
```
|
||||
|
||||
`base_url` is where the link points, and it is never taken from the request:
|
||||
otherwise a stranger could ask for someone's link from a page of their own. A
|
||||
link works once, for an hour, and a newer one replaces it.
|
||||
|
||||
Without mail, the page says to ask an operator, who makes the link and hands it
|
||||
over however suits:
|
||||
|
||||
```
|
||||
llmgw users reset-link ayse@example.com --store … # --base-url without mail.base_url
|
||||
```
|
||||
|
||||
Or sets a temporary password to give them, and has them change it from their
|
||||
profile:
|
||||
|
||||
```
|
||||
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.
|
||||
Either way 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. People who sign in through single sign-on have no password
|
||||
here, and get no link.
|
||||
|
||||
## Trying a model in the browser
|
||||
|
||||
@@ -433,21 +489,26 @@ Three things are deliberately not there:
|
||||
| `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.
|
||||
rare enough that a shell is no hardship. Cordon and drain are CLI-only for the
|
||||
same reason; what the pages do change on the cluster — serving models through
|
||||
the mirror, a node's engines — needs the `operator` or `admin` staff role.
|
||||
|
||||
Every administrative action is logged with who did it, including the ones that
|
||||
were refused.
|
||||
Every administrative change is recorded in the audit log — who, when, through
|
||||
what, the fields given with secrets taken out, and how it ended, refusals
|
||||
included: `llmgw ctl audit`, or **Audit** in the console.
|
||||
|
||||
## 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.
|
||||
- **Single sign-on is optional.** With `auth.oidc` people sign in through
|
||||
your identity provider, and SCIM 2.0 (`auth.scim`) creates, updates and
|
||||
suspends accounts from it; see [install.md](install.md#single-sign-on).
|
||||
Without them accounts are local, with local passwords.
|
||||
- **Staff roles are coarse.** `viewer`, `operator`, `billing` and `admin`
|
||||
(`llmgw users role`) decide what someone may see and do under /admin; an
|
||||
organization's manager can view it but not change it. Every change is in
|
||||
the audit log (`llmgw ctl audit`).
|
||||
- **Rate limits are per gateway instance.** With two instances each counts its own,
|
||||
so the effective limit doubles. v1 is single-instance by design.
|
||||
+37
-19
@@ -8,10 +8,10 @@ explains every choice.
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| **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. |
|
||||
| **A control plane** | One machine, or three for high availability (the store replicated, either one able to take over). Linux, x86-64 or ARM64: Debian 12, Ubuntu 22.04+ or RHEL 10. No GPU needed; 2 GB RAM, 20 GB disk each — 8 GB RAM when it is also a node running the starter model and the two judges. |
|
||||
| **GPU nodes** (optional to start) | The same Linux, with the NVIDIA driver installed. vLLM needs driver 525 or newer and a Turing GPU or newer; a V100 gets an older vLLM. |
|
||||
| **root** | on the control plane. On the other two of a three-machine setup, `sudo` (its password is asked for if it needs one); on nodes, `sudo` without a password. |
|
||||
| **SSH** | from the control plane to the other machines. The setup screens take a key, the ssh agent or a password; `node add` uses ssh keys. A machine nothing may log into joins with a code instead (`llmgw join`). |
|
||||
| **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
|
||||
@@ -35,26 +35,41 @@ install it.
|
||||
## 2. Set up the control plane
|
||||
|
||||
```
|
||||
sudo llmgw server init
|
||||
sudo llmgw
|
||||
```
|
||||
|
||||
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`.
|
||||
With nothing installed, `llmgw` opens the setup screens. They ask whether
|
||||
this is **one machine or three** (high availability), and for three, the
|
||||
other two: address, user, and an ssh key, the ssh agent or a password. They
|
||||
check every machine before installing anything — system, disk, clock, ports,
|
||||
internet, whether the machines reach one another — show the plan, and
|
||||
install:
|
||||
|
||||
- the service, its certificates and its configuration;
|
||||
- your administrator account. **The password is printed once. Copy it now.**
|
||||
- the service, its certificates and its configuration, and a firewall for
|
||||
its internal ports;
|
||||
- with three machines, the store replicated across them (PostgreSQL under
|
||||
Patroni);
|
||||
- your administrator account. **The password is shown 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.
|
||||
- this machine as a first node, and a small model (Qwen2.5 0.5B) on it;
|
||||
- guardrails on every request, with a judge model (Qwen3Guard-Gen 0.6B, two
|
||||
replicas) on CPUs — a request it cannot check is refused, not let through — and API keys that expire within 90 days. What is hardened, and how
|
||||
to loosen it: [install.md](install.md#2-the-control-plane).
|
||||
|
||||
Run again later, `sudo llmgw` is the console: the cluster's state and
|
||||
everything you can do to it ([console.md](console.md)). Every step below is
|
||||
also there, under **Do**.
|
||||
|
||||
**In a script**, the same for one machine:
|
||||
|
||||
```
|
||||
sudo llmgw server init --admin you@example.com
|
||||
```
|
||||
|
||||
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.
|
||||
`--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,
|
||||
@@ -75,7 +90,7 @@ the account page.
|
||||
|
||||
## 3. Add the GPU nodes
|
||||
|
||||
On the control plane:
|
||||
In the console, **Do → Add a node**. Or on the control plane:
|
||||
|
||||
```
|
||||
sudo llmgw node add ubuntu@10.0.0.21
|
||||
@@ -197,6 +212,9 @@ through the proxy. There is an nginx configuration too:
|
||||
|
||||
## Next
|
||||
|
||||
- [install.md](install.md): the same steps with every option explained.
|
||||
- [install.md](install.md): the same steps with every option explained, and
|
||||
running the cluster after: monitoring, notifications, single sign-on,
|
||||
backups, upgrades.
|
||||
- [console.md](console.md): everything the console does.
|
||||
- [fabrics.md](fabrics.md): InfiniBand and NVLink.
|
||||
- [onboarding.md](onboarding.md): accounts, projects, budgets and keys.
|
||||
+1495
-9
File diff suppressed because it is too large.
Load diff
@@ -1,23 +1,30 @@
|
||||
module gitea.mynodes.xyz/ilker/llmgw
|
||||
|
||||
go 1.26.3
|
||||
go 1.26.6
|
||||
|
||||
require (
|
||||
connectrpc.com/connect v1.21.0
|
||||
filippo.io/age v1.3.2
|
||||
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/charmbracelet/x/ansi v0.11.6
|
||||
github.com/coreos/go-oidc/v3 v3.21.0
|
||||
github.com/go-jose/go-jose/v4 v4.1.4
|
||||
github.com/go-zeromq/zmq4 v0.17.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/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0
|
||||
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/vmihailenco/msgpack/v5 v5.4.1
|
||||
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
|
||||
@@ -26,17 +33,21 @@ require (
|
||||
go.yaml.in/yaml/v3 v3.0.5
|
||||
golang.org/x/crypto v0.57.0
|
||||
golang.org/x/net v0.59.0
|
||||
golang.org/x/oauth2 v0.37.0
|
||||
golang.org/x/sync v0.23.0
|
||||
golang.org/x/term v0.46.0
|
||||
golang.org/x/text v0.42.0
|
||||
google.golang.org/protobuf v1.36.12
|
||||
modernc.org/sqlite v1.59.0
|
||||
)
|
||||
|
||||
require (
|
||||
filippo.io/edwards25519 v1.2.0 // indirect
|
||||
filippo.io/hpke v0.4.0 // indirect
|
||||
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
|
||||
@@ -48,6 +59,7 @@ require (
|
||||
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/go-zeromq/goczmq/v4 v4.2.2 // 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
|
||||
@@ -70,13 +82,13 @@ require (
|
||||
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/vmihailenco/tagparser/v2 v2.0.0 // 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
|
||||
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
|
||||
|
||||
@@ -1,5 +1,13 @@
|
||||
c2sp.org/CCTV/age v0.0.0-20260829155415-4448f2097b2d h1:Blprhc2SbChNZtWcU+BLTM4YdoqYAS9V7cJgOwJKyAs=
|
||||
c2sp.org/CCTV/age v0.0.0-20260829155415-4448f2097b2d/go.mod h1:SrHC2C7r5GkDk8R+NFVzYy/sdj0Ypg9htaPXQq5Cqeo=
|
||||
connectrpc.com/connect v1.21.0 h1:LhqSJt7jHf5NJBo9Jq/t/9FjcYAideif0mg+qe2jCUs=
|
||||
connectrpc.com/connect v1.21.0/go.mod h1:A2ygJrukXwWy32vkCAAHNVguZrqZ+jeZ9rGRnGR4dN4=
|
||||
filippo.io/age v1.3.2 h1:r6RSZLFSMm6rzKepZ7ZAYkKCu14f3/Me8c7uKYh7C8c=
|
||||
filippo.io/age v1.3.2/go.mod h1:TH/Yr2sSRhCKbaH4XPxpUV0Us8Gv6txYUpiZQWz8Evk=
|
||||
filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo=
|
||||
filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc=
|
||||
filippo.io/hpke v0.4.0 h1:p575VVQ6ted4pL+it6M00V/f2qTZITO0zgmdKCkd5+A=
|
||||
filippo.io/hpke v0.4.0/go.mod h1:EmAN849/P3qdeK+PCMkDpDm83vRHM5cDipBJ8xbQLVY=
|
||||
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=
|
||||
@@ -28,6 +36,8 @@ github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfa
|
||||
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/coreos/go-oidc/v3 v3.21.0 h1:wZo4Q9Pum8dYEj0eMUPrqR+kvuGkeUplbLpNCkBqoWM=
|
||||
github.com/coreos/go-oidc/v3 v3.21.0/go.mod h1:DYCf24+ncYi+XkIH97GY1+dqoRlbaSI26KVTCI9SrY4=
|
||||
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=
|
||||
@@ -36,6 +46,8 @@ github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6
|
||||
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-jose/go-jose/v4 v4.1.4 h1:moDMcTHmvE6Groj34emNPLs/qtYXRVcd6S7NHbHz3kA=
|
||||
github.com/go-jose/go-jose/v4 v4.1.4/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08=
|
||||
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=
|
||||
@@ -43,6 +55,10 @@ 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/go-zeromq/goczmq/v4 v4.2.2 h1:HAJN+i+3NW55ijMJJhk7oWxHKXgAuSBkoFfvr8bYj4U=
|
||||
github.com/go-zeromq/goczmq/v4 v4.2.2/go.mod h1:Sm/lxrfxP/Oxqs0tnHD6WAhwkWrx+S+1MRrKzcxoaYE=
|
||||
github.com/go-zeromq/zmq4 v0.17.0 h1:r12/XdqPeRbuaF4C3QZJeWCt7a5vpJbslDH1rTXF+Kc=
|
||||
github.com/go-zeromq/zmq4 v0.17.0/go.mod h1:EQxjJD92qKnrsVMzAnx62giD6uJIPi1dMGZ781iCDtY=
|
||||
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=
|
||||
@@ -79,6 +95,8 @@ 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/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0 h1:7Q+xNAZFmnfYOMweHN3c/PDFUKKfY1pVJ26K++QvVfU=
|
||||
github.com/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0/go.mod h1:1fEHWurg7pvf5SG6XNE5Q8UZmOwex51Mkx3SLhrW5B4=
|
||||
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=
|
||||
@@ -124,6 +142,10 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV
|
||||
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/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8=
|
||||
github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok=
|
||||
github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=
|
||||
github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds=
|
||||
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=
|
||||
@@ -161,6 +183,8 @@ 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/oauth2 v0.37.0 h1:JUlcxA8oAtauLfiH8FX2/FkAWHAdi0QtGCGc+hofE98=
|
||||
golang.org/x/oauth2 v0.37.0/go.mod h1:IxwZNxUULJmpBFf9K/9NTMSIfZZuvuTy1gGxhigP/58=
|
||||
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.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
|
||||
@@ -398,3 +398,90 @@ func TestGuessingTheCurrentPasswordIsBounded(t *testing.T) {
|
||||
t.Errorf("guessing was unbounded: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSomebodyGuessingElsewhereCannotLockTheOwnerOut(t *testing.T) {
|
||||
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))
|
||||
|
||||
attacker, owner := WithClient(ctx, "203.0.113.9"), WithClient(ctx, "198.51.100.7")
|
||||
for range 5 {
|
||||
_, _, _ = svc.SignIn(attacker, "ayse@example.com", "nope")
|
||||
}
|
||||
if _, _, err := svc.SignIn(attacker, "ayse@example.com", "nope"); !errors.Is(err, ErrTooManyTries) {
|
||||
t.Errorf("the guessing client was not held back: %v", err)
|
||||
}
|
||||
if _, _, err := svc.SignIn(owner, "ayse@example.com", goodPassword); err != nil {
|
||||
t.Errorf("the owner, elsewhere, was locked out: %v", err)
|
||||
}
|
||||
// Spread over many clients, guessing is still bounded per address.
|
||||
held := false
|
||||
for i := range 40 {
|
||||
c := WithClient(ctx, "192.0.2."+strings.Repeat("1", 1+i%3)+string(rune('0'+i%10)))
|
||||
if _, _, err := svc.SignIn(c, "ayse@example.com", "nope"); errors.Is(err, ErrTooManyTries) {
|
||||
held = true
|
||||
}
|
||||
}
|
||||
if !held {
|
||||
t.Error("guesses from many clients were not bounded")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAHashThatCannotGetATurnSaysTheServerIsBusy(t *testing.T) {
|
||||
was := argonWait
|
||||
argonWait = 10 * time.Millisecond
|
||||
t.Cleanup(func() { argonWait = was })
|
||||
// Every turn taken.
|
||||
n := cap(argonSlots)
|
||||
for range n {
|
||||
argonSlots <- struct{}{}
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
for range n {
|
||||
<-argonSlots
|
||||
}
|
||||
})
|
||||
if _, err := HashPassword(goodPassword); !errors.Is(err, ErrBusy) {
|
||||
t.Errorf("a hash with no turn: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAForgottenPasswordIsResetOnceThroughALink(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")
|
||||
|
||||
// Waiting for approval: no link, said the same as for nobody.
|
||||
if _, _, err := svc.RequestReset(ctx, st, "ayse@example.com"); !errors.Is(err, ErrNoReset) {
|
||||
t.Fatalf("a pending account got a link: %v", err)
|
||||
}
|
||||
if _, _, err := svc.RequestReset(ctx, st, "nobody@example.com"); !errors.Is(err, ErrNoReset) {
|
||||
t.Fatalf("an unknown address: %v", err)
|
||||
}
|
||||
must(t, st.ApproveUser(ctx, u.ID, "proj", "ilker", time.Now()))
|
||||
session, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
must(t, err)
|
||||
|
||||
token, got, err := svc.RequestReset(ctx, st, "AYSE@example.com")
|
||||
if err != nil || token == "" || got.ID != u.ID {
|
||||
t.Fatalf("an active account: %q %+v %v", token, got, err)
|
||||
}
|
||||
if err := svc.ResetPassword(ctx, st, st, token, "short"); !errors.Is(err, ErrWeakPassword) {
|
||||
t.Fatalf("a weak password: %v", err)
|
||||
}
|
||||
next := goodPassword + "-new"
|
||||
must(t, svc.ResetPassword(ctx, st, st, token, next))
|
||||
if _, _, err := svc.Resolve(ctx, session); err == nil {
|
||||
t.Error("a session from before the reset still works")
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", next); err != nil {
|
||||
t.Errorf("the new password does not sign in: %v", err)
|
||||
}
|
||||
if err := svc.ResetPassword(ctx, st, st, token, next+"2"); !errors.Is(err, ErrNoReset) {
|
||||
t.Errorf("a link used twice: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"fmt"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"golang.org/x/crypto/argon2"
|
||||
)
|
||||
@@ -57,7 +58,10 @@ func HashPassword(password string) (string, error) {
|
||||
if _, err := rand.Read(salt); err != nil {
|
||||
return "", err
|
||||
}
|
||||
sum := idKey([]byte(password), salt, argonTime, argonMemory, argonThreads, argonKeyLen)
|
||||
sum, err := idKey([]byte(password), salt, argonTime, argonMemory, argonThreads, argonKeyLen)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
b64 := base64.RawStdEncoding.EncodeToString
|
||||
return fmt.Sprintf("$argon2id$v=%d$m=%d,t=%d,p=%d$%s$%s",
|
||||
@@ -94,7 +98,10 @@ func VerifyPassword(password, encoded string) (bool, error) {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
|
||||
got := idKey([]byte(password), salt, time, memory, threads, uint32(len(want)))
|
||||
got, err := idKey([]byte(password), salt, time, memory, threads, uint32(len(want)))
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return subtle.ConstantTimeCompare(got, want) == 1, nil
|
||||
}
|
||||
|
||||
@@ -119,8 +126,19 @@ func init() {
|
||||
// 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{}{}
|
||||
// ErrBusy is a hash that waited its turn too long: the server is answering
|
||||
// more sign-ins than it can hash, and says so rather than hold the request.
|
||||
var ErrBusy = errors.New("too many sign-ins at once; try again in a moment")
|
||||
|
||||
// argonWait bounds the wait for a turn; a variable so a test can shorten it.
|
||||
var argonWait = 10 * time.Second
|
||||
|
||||
func idKey(password, salt []byte, t, memory uint32, threads uint8, keyLen uint32) ([]byte, error) {
|
||||
select {
|
||||
case argonSlots <- struct{}{}:
|
||||
case <-time.After(argonWait):
|
||||
return nil, ErrBusy
|
||||
}
|
||||
defer func() { <-argonSlots }()
|
||||
return argon2.IDKey(password, salt, time, memory, threads, keyLen)
|
||||
return argon2.IDKey(password, salt, t, memory, threads, keyLen), nil
|
||||
}
|
||||
@@ -138,7 +138,15 @@ func (s *Service) SignUp(ctx context.Context, email, name, password string) erro
|
||||
// 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) {
|
||||
// Counted per address and the client trying it, so somebody guessing
|
||||
// from elsewhere cannot lock the owner out of their own account; and per
|
||||
// address alone, more loosely, so guesses spread over many clients are
|
||||
// still bounded.
|
||||
byClient := email
|
||||
if c := clientFrom(ctx); c != "" {
|
||||
byClient = email + "|" + c
|
||||
}
|
||||
if !s.allowAttempt(byClient, s.cfg.MaxAttempts) || !s.allowAttempt("all|"+email, 10*s.cfg.MaxAttempts) {
|
||||
return "", "", ErrTooManyTries
|
||||
}
|
||||
|
||||
@@ -154,11 +162,15 @@ func (s *Service) SignIn(ctx context.Context, email, password string) (token, cs
|
||||
}
|
||||
|
||||
ok, err := VerifyPassword(password, u.PasswordHash)
|
||||
if errors.Is(err, ErrBusy) {
|
||||
// Not an answer about the password: said as what it is.
|
||||
return "", "", err
|
||||
}
|
||||
if err != nil || !ok {
|
||||
return "", "", ErrBadCredentials
|
||||
}
|
||||
// The password was right, so this was not an attack.
|
||||
s.clearAttempts(email)
|
||||
s.clearAttempts(byClient)
|
||||
|
||||
// State is checked after the password, on purpose. Checking first would
|
||||
// tell anyone who guesses an address whether it has a pending account.
|
||||
@@ -224,7 +236,7 @@ func (s *Service) CheckCSRF(sess store.Session, submitted string) bool {
|
||||
// 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 {
|
||||
func (s *Service) allowAttempt(email string, limit int) bool {
|
||||
now := s.now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
@@ -251,7 +263,20 @@ func (s *Service) allowAttempt(email string) bool {
|
||||
return true
|
||||
}
|
||||
w.count++
|
||||
return w.count <= s.cfg.MaxAttempts
|
||||
return w.count <= limit
|
||||
}
|
||||
|
||||
type clientKey struct{}
|
||||
|
||||
// WithClient says which client a request comes from — its address — for the
|
||||
// sign-in throttle.
|
||||
func WithClient(ctx context.Context, client string) context.Context {
|
||||
return context.WithValue(ctx, clientKey{}, client)
|
||||
}
|
||||
|
||||
func clientFrom(ctx context.Context) string {
|
||||
c, _ := ctx.Value(clientKey{}).(string)
|
||||
return c
|
||||
}
|
||||
|
||||
// maxTrackedAddresses bounds the throttle's memory. Unbounded, it grew by one
|
||||
@@ -307,7 +332,7 @@ type PasswordChanger interface {
|
||||
func (s *Service) ChangePassword(ctx context.Context, pw PasswordChanger,
|
||||
u store.User, current, next string,
|
||||
) error {
|
||||
if !s.allowAttempt("pw:" + u.ID) {
|
||||
if !s.allowAttempt("pw:"+u.ID, s.cfg.MaxAttempts) {
|
||||
return ErrTooManyTries
|
||||
}
|
||||
ok, err := VerifyPassword(current, u.PasswordHash)
|
||||
@@ -348,3 +373,63 @@ func (s *Service) OpenSession(ctx context.Context, userID string) (token, csrf s
|
||||
}
|
||||
return token, csrf, nil
|
||||
}
|
||||
|
||||
// ResetLinkTTL is how long a password reset link works.
|
||||
const ResetLinkTTL = time.Hour
|
||||
|
||||
// ResetStore keeps password reset links.
|
||||
type ResetStore interface {
|
||||
UserByEmail(ctx context.Context, email string) (store.User, error)
|
||||
CreatePasswordReset(ctx context.Context, userID string, ttl time.Duration) (string, error)
|
||||
PasswordResetUser(ctx context.Context, token string) (string, error)
|
||||
UsePasswordReset(ctx context.Context, token string) (string, error)
|
||||
}
|
||||
|
||||
// ErrNoReset is an address that gets no reset link: none has an account, the
|
||||
// account signs in through single sign-on, or it is not active. Never told to
|
||||
// whoever asked, who learns the same either way.
|
||||
var ErrNoReset = errors.New("no reset link for that address")
|
||||
|
||||
// RequestReset mints a reset link token for the person an address belongs
|
||||
// to, when they sign in with a password and are active. A few a quarter of
|
||||
// an hour per address: each sends a mail.
|
||||
func (s *Service) RequestReset(ctx context.Context, rs ResetStore, email string) (string, store.User, error) {
|
||||
email = store.NormalizeEmail(email)
|
||||
if !validEmail(email) {
|
||||
return "", store.User{}, ErrNoReset
|
||||
}
|
||||
if !s.allowAttempt("reset:"+email, 3) {
|
||||
return "", store.User{}, ErrTooManyTries
|
||||
}
|
||||
u, err := rs.UserByEmail(ctx, email)
|
||||
if err != nil || u.PasswordHash == "" || u.State != store.UserActive {
|
||||
return "", store.User{}, ErrNoReset
|
||||
}
|
||||
token, err := rs.CreatePasswordReset(ctx, u.ID, ResetLinkTTL)
|
||||
if err != nil {
|
||||
return "", store.User{}, err
|
||||
}
|
||||
s.log.Info("a password reset link was made", "user", u.ID)
|
||||
return token, u, nil
|
||||
}
|
||||
|
||||
// ResetPassword sets a password from a reset link, which is spent, and signs
|
||||
// the person out everywhere: if they lost it, someone else may have it.
|
||||
func (s *Service) ResetPassword(ctx context.Context, rs ResetStore, pw PasswordChanger, token, next string) error {
|
||||
hash, err := HashPassword(next)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
userID, err := rs.UsePasswordReset(ctx, token)
|
||||
if err != nil {
|
||||
return ErrNoReset
|
||||
}
|
||||
if err := pw.SetPassword(ctx, userID, hash); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := pw.DeleteUserSessions(ctx, userID); err != nil {
|
||||
return err
|
||||
}
|
||||
s.log.Info("a password was set from a reset link", "user", userID)
|
||||
return nil
|
||||
}
|
||||
@@ -12,22 +12,22 @@ func TestSprayingAddressesCannotGrowTheThrottleOrEscapeIt(t *testing.T) {
|
||||
|
||||
// The target is being guessed at.
|
||||
for i := 0; i < s.cfg.MaxAttempts; i++ {
|
||||
s.allowAttempt("victim@example.com")
|
||||
s.allowAttempt("victim@example.com", s.cfg.MaxAttempts)
|
||||
}
|
||||
// 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))
|
||||
s.allowAttempt(fmt.Sprintf("x%d@example.com", i), s.cfg.MaxAttempts)
|
||||
}
|
||||
if n := len(s.attempts); n > maxTrackedAddresses {
|
||||
t.Errorf("the table grew to %d", n)
|
||||
}
|
||||
if s.allowAttempt("victim@example.com") {
|
||||
if s.allowAttempt("victim@example.com", s.cfg.MaxAttempts) {
|
||||
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") {
|
||||
if !s.allowAttempt("new@example.com", s.cfg.MaxAttempts) {
|
||||
t.Error("a new address was refused after the flood expired")
|
||||
}
|
||||
}
|
||||
@@ -6,9 +6,11 @@ import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
@@ -45,8 +47,13 @@ type Accountant struct {
|
||||
tokens *Limiter
|
||||
recorder *Recorder
|
||||
log *slog.Logger
|
||||
// notify is told when a budget runs out; nil tells no one.
|
||||
notify atomic.Pointer[notify.Notifier]
|
||||
}
|
||||
|
||||
// SetNotifier gives the accountant somewhere to say a budget is spent.
|
||||
func (a *Accountant) SetNotifier(n *notify.Notifier) { a.notify.Store(n) }
|
||||
|
||||
// New builds an accountant.
|
||||
func New(pricing *Pricing, recorder *Recorder, log *slog.Logger) *Accountant {
|
||||
if log == nil {
|
||||
@@ -182,6 +189,10 @@ func (a *Accountant) Check(p auth.Principal, model string, estimatedTokens int64
|
||||
if !price.Zero() {
|
||||
for _, l := range chain {
|
||||
if l.Budget != nil && a.budgets.Exhausted(l.Counter, l.Period, *l.Budget) {
|
||||
a.notify.Load().Notify(notify.BudgetSpent, l.Counter,
|
||||
fmt.Sprintf("the %s budget of %s is spent: requests to priced models are refused until it renews",
|
||||
or(string(l.Period), "total"), l.Counter),
|
||||
map[string]any{"counter": l.Counter, "period": string(l.Period), "budget": *l.Budget})
|
||||
return res, &LimitError{Err: ErrBudgetSpent}
|
||||
}
|
||||
}
|
||||
@@ -306,7 +317,13 @@ func (a *Accountant) Record(res Reservation, u store.Usage) {
|
||||
a.tokens.Charge(c, actual-res.EstimatedTokens)
|
||||
}
|
||||
|
||||
u.CostMicros = res.Price.Cost(u.PromptTokens, u.CompletionTokens)
|
||||
// Priced by the model that answered: a fallback served by another model
|
||||
// costs what that one costs, not what the one asked for would have.
|
||||
price := res.Price
|
||||
if u.ServedModel != "" && u.ServedModel != u.Model && a.pricing != nil {
|
||||
price = a.pricing.For(u.ServedModel)
|
||||
}
|
||||
u.CostMicros = price.Cost(u.PromptTokens, u.CompletionTokens)
|
||||
for _, b := range res.Budgets {
|
||||
a.budgets.Add(b.Counter, b.Period, u.CostMicros)
|
||||
}
|
||||
|
||||
@@ -5,11 +5,14 @@ import (
|
||||
"errors"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/notify/notifytest"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
@@ -658,3 +661,46 @@ func TestHeadroomIsTheTightestLimitAfterThisRequest(t *testing.T) {
|
||||
t.Errorf("a key with no limits advertised some: %+v", none)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFallbackIsChargedWhatTheModelThatAnsweredCosts(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := inTree("fb-k", teamBudget("fb", 1_000_000))
|
||||
res, err := a.Check(p, "paid", 400)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Asked for "paid", answered by "free": nothing is owed.
|
||||
a.Record(res, store.Usage{KeyID: "fb-k", Model: "paid", ServedModel: "free", PromptTokens: 200, CompletionTokens: 200})
|
||||
if spent, _, _ := a.Status(p); spent != 0 {
|
||||
t.Errorf("charged %d for a request the free model answered", spent)
|
||||
}
|
||||
res, _ = a.Check(p, "paid", 400)
|
||||
a.Record(res, store.Usage{KeyID: "fb-k", Model: "paid", ServedModel: "paid", PromptTokens: 200, CompletionTokens: 200})
|
||||
if spent, _, _ := a.Status(p); spent == 0 {
|
||||
t.Error("the paid model answering cost nothing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestASpentBudgetTellsSomeoneOnce(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
n, sent := notifytest.Catch(t)
|
||||
a.SetNotifier(n)
|
||||
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})
|
||||
for range 3 {
|
||||
if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Fatalf("not refused: %v", err)
|
||||
}
|
||||
}
|
||||
got := sent()
|
||||
if len(got) != 1 || got[0].Kind != notify.BudgetSpent || !strings.Contains(got[0].Message, "month") {
|
||||
t.Fatalf("sent %+v, want one budget_spent", got)
|
||||
}
|
||||
}
|
||||
+345
-39
@@ -14,13 +14,18 @@ import (
|
||||
"math/rand/v2"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"slices"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"golang.org/x/net/http2"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
@@ -31,16 +36,22 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
// Agent runs on a GPU node.
|
||||
type Agent struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
sup *supervisor.Supervisor
|
||||
reg *engines.Registry
|
||||
http connectHTTPClient
|
||||
cfg Config
|
||||
// answered says the current session heard from its control plane, which
|
||||
// starts the reconnect backoff again when it ends.
|
||||
answered atomic.Bool
|
||||
log *slog.Logger
|
||||
sup *supervisor.Supervisor
|
||||
reg *engines.Registry
|
||||
http connectHTTPClient
|
||||
|
||||
// engineOnce probes the configured engines once before registering.
|
||||
//
|
||||
@@ -77,7 +88,11 @@ type Agent struct {
|
||||
// serverList is every control plane to dial: the configured ones, and
|
||||
// any the leader has named since. Read through servers().
|
||||
serverList []string
|
||||
inv discovery.Inventory
|
||||
// fwMu keeps one firewall change at a time; fwRunner runs its scripts
|
||||
// (nil: this machine's shell).
|
||||
fwMu sync.Mutex
|
||||
fwRunner firewall.Runner
|
||||
inv discovery.Inventory
|
||||
// shardReplica remembers which replica a shard belongs to, so a status
|
||||
// report can name both. The control plane keys on the replica.
|
||||
shardReplica map[string]string
|
||||
@@ -115,6 +130,16 @@ type Agent struct {
|
||||
// — eight shards starting together produce one download — and all eight
|
||||
// are waiting, so all eight should say so.
|
||||
fetchWatchers map[string][]func(weights.Progress)
|
||||
|
||||
// kv mirrors the KV caches of the engines that publish events, for the
|
||||
// router's precise index; kvSubs are their listeners by shard, guarded
|
||||
// by mu.
|
||||
kv *kvevents.Mirror
|
||||
kvSubs map[string]*kvevents.Subscriber
|
||||
|
||||
// settingsCheck says whether the cluster's settings for this agent
|
||||
// would load beside its agent.yaml; see SetSettingsCheck.
|
||||
settingsCheck func([]byte) error
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
@@ -157,6 +182,8 @@ func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, er
|
||||
sensors: discovery.NewSensorSampler(),
|
||||
wear: discovery.NewHealthSampler(),
|
||||
serverList: loadControlPlanes(cfg.Servers, cfg.TLS.Enabled()),
|
||||
kv: kvevents.NewMirror(),
|
||||
kvSubs: map[string]*kvevents.Subscriber{},
|
||||
}
|
||||
a.prov = &provision.Installer{Root: cfg.EnginesDir, Log: log, OnStatus: a.provisionStatus}
|
||||
return a, nil
|
||||
@@ -201,6 +228,14 @@ func (a *Agent) Run(ctx context.Context) error {
|
||||
// never answers — from here they are the same thing, and trying
|
||||
// somewhere else is the answer to all three.
|
||||
next++
|
||||
// A session that was answered was a connection that worked, however
|
||||
// it ended: the backoff starts again rather than carrying on from
|
||||
// the last outage. Without this every failover the agent lived
|
||||
// through made the next one slower — the third waited half a minute
|
||||
// on a control plane that had been ready for twenty seconds.
|
||||
if a.answered.Swap(false) {
|
||||
attempt = 0
|
||||
}
|
||||
attempt++
|
||||
delay := a.reconnectDelay(attempt)
|
||||
// No wait when there is somewhere else to try: the delay exists to
|
||||
@@ -275,6 +310,12 @@ func (a *Agent) recoverOrphans() {
|
||||
a.mu.Lock()
|
||||
a.shardReplica[o.ShardID] = o.ReplicaID
|
||||
a.mu.Unlock()
|
||||
// The engine is still publishing to the port the last agent
|
||||
// listened on; listen there again.
|
||||
if o.Spec.KVEventsPort > 0 {
|
||||
spec := o.Spec
|
||||
a.listenKV(&spec, o.Spec.KVEventsPort)
|
||||
}
|
||||
a.log.Info("adopted a running engine",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID, "pid", o.PID)
|
||||
go a.settleAdopted(sh, o)
|
||||
@@ -310,7 +351,7 @@ func (a *Agent) settleAdopted(sh *supervisor.Shard, o supervisor.Orphan) {
|
||||
// The same two paths a freshly started shard takes: a rank that serves HTTP
|
||||
// is health-checked, a worker is judged by staying alive.
|
||||
if eng.ServesHTTP(o.Spec) {
|
||||
base := fmt.Sprintf("http://127.0.0.1:%d", o.Spec.HTTPPort)
|
||||
base := a.engineBase(o.Spec)
|
||||
err = a.sup.WaitReady(ctx, sh, base, o.Spec.Profile.ServedModelName, a.cfg.AdoptReadyTimeout)
|
||||
} else {
|
||||
err = a.sup.WaitAlive(ctx, sh)
|
||||
@@ -334,6 +375,12 @@ func (a *Agent) settleAdopted(sh *supervisor.Shard, o supervisor.Orphan) {
|
||||
_ = a.sup.Stop(ctx, o.ShardID)
|
||||
return
|
||||
}
|
||||
// What it computes on, for one started before the agent kept that.
|
||||
if a.sup.ComputeOf(o.ShardID) == nil {
|
||||
if c, ok := a.gpuCompute(ctx, o.Spec.GPUIndices, sh.Status().PID); ok {
|
||||
a.sup.SetCompute(o.ShardID, c)
|
||||
}
|
||||
}
|
||||
a.log.Info("adopted engine is serving",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID)
|
||||
}
|
||||
@@ -374,6 +421,17 @@ func (a *Agent) weights() *weights.Fetcher {
|
||||
}
|
||||
|
||||
// runningShards describes what this agent is supervising, for registration.
|
||||
// unready are the shards a registration reported as not yet ready.
|
||||
func unready(running []*llmgwv1.RunningShard) map[string]bool {
|
||||
out := map[string]bool{}
|
||||
for _, r := range running {
|
||||
if r.GetState() != agentsvc.StateToWire(supervisor.Ready) {
|
||||
out[r.GetShardId()] = true
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (a *Agent) runningShards() []*llmgwv1.RunningShard {
|
||||
var out []*llmgwv1.RunningShard
|
||||
for _, st := range a.sup.Statuses() {
|
||||
@@ -399,6 +457,7 @@ func (a *Agent) runningShards() []*llmgwv1.RunningShard {
|
||||
Pool: spec.Pool,
|
||||
Fabric: spec.Fabric,
|
||||
ProfileFp: spec.ProfileFP,
|
||||
Compute: computeToWire(a.sup.ComputeOf(st.ID)),
|
||||
LeaderAddr: a.cfg.MgmtAddr,
|
||||
LeaderPort: int32(spec.LeaderPort),
|
||||
})
|
||||
@@ -417,6 +476,17 @@ func int32Slice(v []int) []int32 {
|
||||
return out
|
||||
}
|
||||
|
||||
func int32Ints(v []int32) []int {
|
||||
out := make([]int, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = int(x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// tracerName names the spans the agent opens.
|
||||
const tracerName = "gitea.mynodes.xyz/ilker/llmgw/internal/agent"
|
||||
|
||||
// scan refreshes the node's hardware inventory.
|
||||
func (a *Agent) scan(ctx context.Context) error {
|
||||
// Carried from the agent's own configuration rather than duplicated into
|
||||
@@ -516,6 +586,9 @@ func (a *Agent) session(ctx context.Context, server string) error {
|
||||
a.sendLoop(ctx, stream, out)
|
||||
}()
|
||||
defer wg.Wait()
|
||||
go a.watchExits(ctx, out, unready(reg.RunningShards))
|
||||
go a.forwardKV(ctx, out)
|
||||
go a.reportCompute(ctx, out)
|
||||
|
||||
for {
|
||||
msg, err := stream.Receive()
|
||||
@@ -535,6 +608,7 @@ func (a *Agent) session(ctx context.Context, server string) error {
|
||||
// why it is not registering is one an operator cannot fix.
|
||||
return err
|
||||
}
|
||||
a.answered.Store(true)
|
||||
a.handle(ctx, msg, out)
|
||||
}
|
||||
}
|
||||
@@ -783,15 +857,23 @@ func (a *Agent) removeWeights(paths []string) {
|
||||
}
|
||||
|
||||
func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out chan<- *llmgwv1.AgentMessage) {
|
||||
// The control plane's span that sent this, when it is tracing: what is
|
||||
// done here for a start or a stop goes into the same trace.
|
||||
tctx := obs.Extract(ctx, msg.GetTrace())
|
||||
switch m := msg.Msg.(type) {
|
||||
case *llmgwv1.ControlMessage_Start:
|
||||
go a.startShard(ctx, m.Start, out)
|
||||
go a.startShard(tctx, m.Start, out)
|
||||
case *llmgwv1.ControlMessage_Stop:
|
||||
go a.stopShard(ctx, m.Stop.GetShardId(), out)
|
||||
go a.stopShard(tctx, m.Stop.GetShardId(), out)
|
||||
case *llmgwv1.ControlMessage_Drain:
|
||||
if err := a.sup.Drain(m.Drain.GetShardId()); err != nil {
|
||||
_, span := obs.Tracer(tracerName).Start(tctx, "drain shard",
|
||||
trace.WithSpanKind(trace.SpanKindConsumer),
|
||||
trace.WithAttributes(attribute.String("llmgw.shard", m.Drain.GetShardId())))
|
||||
err := a.sup.Drain(m.Drain.GetShardId())
|
||||
if err != nil {
|
||||
a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err)
|
||||
}
|
||||
obs.End(span, err)
|
||||
case *llmgwv1.ControlMessage_RemoveWeights:
|
||||
a.removeWeights(m.RemoveWeights.GetPaths())
|
||||
case *llmgwv1.ControlMessage_Provision:
|
||||
@@ -800,6 +882,18 @@ func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out cha
|
||||
go a.upgrade(ctx, m.Upgrade, out)
|
||||
case *llmgwv1.ControlMessage_ControlPlanes:
|
||||
a.setControlPlanes(m.ControlPlanes)
|
||||
case *llmgwv1.ControlMessage_Firewall:
|
||||
go a.applyFirewall(m.Firewall)
|
||||
case *llmgwv1.ControlMessage_AgentSettings:
|
||||
go a.applySettings(m.AgentSettings)
|
||||
case *llmgwv1.ControlMessage_RestartAgent:
|
||||
go func() {
|
||||
a.log.Info("restarting, as the control plane asks; the engines keep serving",
|
||||
"reason", m.RestartAgent.GetReason())
|
||||
// A moment for the log line to leave before the process does.
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
exitForUpgrade()
|
||||
}()
|
||||
case *llmgwv1.ControlMessage_Logs:
|
||||
reply := a.logsReply(m.Logs)
|
||||
// Off the receive loop: a full send queue must not stop commands being
|
||||
@@ -885,9 +979,64 @@ func (a *Agent) watchFetch(source string, fn func(weights.Progress)) (stop func(
|
||||
}
|
||||
}
|
||||
|
||||
// bindAddr is where engines listen: the management address the control
|
||||
// plane reaches them on, not every address the node has — a public one
|
||||
// included. Empty, on every address, when that is a name rather than an
|
||||
// address this machine holds.
|
||||
func (a *Agent) bindAddr() string {
|
||||
ip, err := netip.ParseAddr(a.cfg.MgmtAddr)
|
||||
if err != nil || ip.IsUnspecified() || !localAddr(ip) {
|
||||
return ""
|
||||
}
|
||||
return ip.String()
|
||||
}
|
||||
|
||||
// engineBase is where the agent asks an engine whether it is ready: where it
|
||||
// listens, loopback when that is every address.
|
||||
func (a *Agent) engineBase(s engines.ShardSpec) string {
|
||||
host := s.BindAddr
|
||||
if host == "" {
|
||||
host = "127.0.0.1"
|
||||
}
|
||||
return "http://" + net.JoinHostPort(host, strconv.Itoa(s.HTTPPort))
|
||||
}
|
||||
|
||||
// localAddr reports whether an address is one of this machine's.
|
||||
func localAddr(ip netip.Addr) bool {
|
||||
addrs, err := net.InterfaceAddrs()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
for _, a := range addrs {
|
||||
if p, err := netip.ParsePrefix(a.String()); err == nil && p.Addr().Unmap() == ip.Unmap() {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<- *llmgwv1.AgentMessage) {
|
||||
shardID, replicaID := m.GetShardId(), m.GetReplicaId()
|
||||
|
||||
// This node's whole part of the start, from the command arriving to the
|
||||
// engine answering, beneath the control plane's span that sent it. The
|
||||
// steps — weights, launch, readiness — are its children.
|
||||
ctx, span := obs.Tracer(tracerName).Start(ctx, "start shard",
|
||||
trace.WithSpanKind(trace.SpanKindConsumer),
|
||||
trace.WithAttributes(
|
||||
attribute.String("llmgw.node", a.cfg.NodeID),
|
||||
attribute.String("llmgw.replica", replicaID),
|
||||
attribute.String("llmgw.shard", shardID),
|
||||
attribute.Int("llmgw.rank", int(m.GetRank())),
|
||||
attribute.String("llmgw.model", m.GetProfile().GetServedModelName()),
|
||||
attribute.String("llmgw.profile", m.GetProfile().GetName()),
|
||||
attribute.String("llmgw.engine", m.GetProfile().GetEngine()),
|
||||
attribute.String("llmgw.gpu_class", m.GetProfile().GetGpuClass()),
|
||||
attribute.IntSlice("llmgw.gpus", int32Ints(m.GetGpuIndices())),
|
||||
))
|
||||
var failed error
|
||||
defer func() { obs.End(span, failed) }()
|
||||
|
||||
a.mu.Lock()
|
||||
a.shardReplica[shardID] = replicaID
|
||||
a.mu.Unlock()
|
||||
@@ -900,6 +1049,9 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
}
|
||||
}
|
||||
report := func(st supervisor.State, errMsg string, tail []string) {
|
||||
if st == supervisor.Failed {
|
||||
failed = errors.New(errMsg)
|
||||
}
|
||||
send(&llmgwv1.ShardStatus{
|
||||
State: agentsvc.StateToWire(st), Error: errMsg, LastOutput: tail,
|
||||
})
|
||||
@@ -907,6 +1059,7 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
|
||||
spec, timeout := agentsvc.ShardSpecFromStart(m)
|
||||
a.localizeSpec(&spec)
|
||||
spec.BindAddr = a.bindAddr()
|
||||
|
||||
eng, err := a.reg.Get(spec.Profile.Engine)
|
||||
if err != nil {
|
||||
@@ -920,6 +1073,12 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
|
||||
// Validate before launching. An engine that refuses a profile does so after
|
||||
// minutes of loading; catching it here costs nothing.
|
||||
// Checked again here, whatever sent the profile: this is the process
|
||||
// about to run as root.
|
||||
if err := engines.UnsafeArgs(spec.Profile.Engine, spec.Profile.Args); err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
if err := eng.Validate(spec.Profile, gpuInfo(inv, spec.GPUIndices)); err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
@@ -963,9 +1122,20 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
}
|
||||
spec.Profile.Adapters[i].Path = path
|
||||
}
|
||||
// Adapters that load on demand are loaded by the router, by a path that
|
||||
// has to mean the same file on every node: a link of the adapter's name
|
||||
// in the directory engines run in.
|
||||
if spec.Profile.RuntimeLoRA() {
|
||||
if err := engines.LinkRuntimeAdapters(a.sup.EngineDir(), spec.Profile.Adapters); err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
a.listenKV(&spec, 0)
|
||||
cmds, err := eng.Commands(spec)
|
||||
if err != nil {
|
||||
a.stopKV(shardID, replicaID)
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
@@ -978,6 +1148,7 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
report(sh.Status().State, "", nil)
|
||||
return
|
||||
}
|
||||
a.stopKV(shardID, replicaID)
|
||||
report(supervisor.Failed, err.Error(), sh.Logs(10))
|
||||
return
|
||||
}
|
||||
@@ -987,11 +1158,19 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
// staying alive, and the replica's real readiness is the leader answering,
|
||||
// which only happens once every worker has joined (plan §11, phase 4).
|
||||
if eng.ServesHTTP(spec) {
|
||||
base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort)
|
||||
base := a.engineBase(spec)
|
||||
err = a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout)
|
||||
} else {
|
||||
err = a.sup.WaitAlive(ctx, sh)
|
||||
}
|
||||
// Where it put the model, by its own account: said with the readiness,
|
||||
// and a GPU placement it computes on the CPU refused even where
|
||||
// nvidia-smi cannot be asked.
|
||||
compute, said, cerr := engineCompute(eng, spec.GPUIndices, sh.Logs(512))
|
||||
if err == nil && cerr != nil {
|
||||
err = cerr
|
||||
_ = a.sup.Stop(ctx, shardID)
|
||||
}
|
||||
if err == nil {
|
||||
err = a.onItsGPUs(ctx, eng, spec, sh.Status().PID)
|
||||
if err != nil {
|
||||
@@ -1004,37 +1183,92 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
report(supervisor.Failed, err.Error(), sh.Logs(20))
|
||||
return
|
||||
}
|
||||
report(supervisor.Ready, "", nil)
|
||||
|
||||
// Watch for the engine dying later. Nothing else would notice: the router
|
||||
// would simply see failing requests.
|
||||
go a.watch(ctx, sh, replicaID, shardID, out)
|
||||
// The engine's own account when it gave one, with its layers; else what
|
||||
// nvidia-smi sees the shard holding.
|
||||
if !said {
|
||||
compute, said = a.gpuCompute(ctx, spec.GPUIndices, sh.Status().PID)
|
||||
}
|
||||
if said {
|
||||
// Kept with the shard, so an agent restarted, or a control plane
|
||||
// taking over, still learns it.
|
||||
a.sup.SetCompute(shardID, compute)
|
||||
send(&llmgwv1.ShardStatus{State: agentsvc.StateToWire(supervisor.Ready), Compute: computeToWire(&compute)})
|
||||
} else {
|
||||
report(supervisor.Ready, "", nil)
|
||||
}
|
||||
// Its dying later is reported by the session's watchExits.
|
||||
}
|
||||
|
||||
// watch reports a shard that exits after it was ready.
|
||||
func (a *Agent) watch(ctx context.Context, sh *supervisor.Shard, replicaID, shardID string, out chan<- *llmgwv1.AgentMessage) {
|
||||
// watchExits reports, for as long as this session lasts, every shard that
|
||||
// was ready and has failed: the ones this agent started, and the ones it
|
||||
// adopted from the agent before it.
|
||||
//
|
||||
// Per session rather than per shard. A watcher tied to the stream a shard
|
||||
// was started on stopped with that stream, so after a reconnect nothing
|
||||
// reported the engine dying; and an adopted engine had no watcher at all —
|
||||
// on node06 one was killed and its replica stayed "ready" for the control
|
||||
// plane, routed to and failing, indefinitely.
|
||||
//
|
||||
// It also reports the shards the registration said were not ready yet once
|
||||
// they are. An adopted engine is health-checked after the agent registers,
|
||||
// and a shard started on the stream before is reported on that stream; so
|
||||
// when the registration came first nothing ever told the control plane, and
|
||||
// the replica stayed "starting" while it served — on node05, after a
|
||||
// control plane restart, the only replica of a model.
|
||||
func (a *Agent) watchExits(ctx context.Context, out chan<- *llmgwv1.AgentMessage, pending map[string]bool) {
|
||||
t := time.NewTicker(time.Second)
|
||||
defer t.Stop()
|
||||
|
||||
ready := map[string]bool{}
|
||||
for {
|
||||
for _, st := range a.sup.Statuses() {
|
||||
switch st.State {
|
||||
case supervisor.Ready:
|
||||
ready[st.ID] = true
|
||||
if !pending[st.ID] {
|
||||
continue
|
||||
}
|
||||
delete(pending, st.ID)
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: sh.Spec.ReplicaID, ShardId: st.ID,
|
||||
State: agentsvc.StateToWire(st.State),
|
||||
Compute: computeToWire(a.sup.ComputeOf(st.ID)),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
case supervisor.Failed:
|
||||
if !ready[st.ID] {
|
||||
continue // never ready here: its start reported it
|
||||
}
|
||||
delete(ready, st.ID)
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: sh.Spec.ReplicaID, ShardId: st.ID,
|
||||
State: agentsvc.StateToWire(st.State), Error: st.Err,
|
||||
ExitCode: int32(st.ExitCode), LastOutput: sh.Logs(20),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
default:
|
||||
delete(ready, st.ID)
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
st := sh.Status()
|
||||
if !st.State.Terminal() {
|
||||
continue
|
||||
}
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
|
||||
State: agentsvc.StateToWire(st.State), Error: st.Err,
|
||||
ExitCode: int32(st.ExitCode), LastOutput: sh.Logs(20),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1045,9 +1279,19 @@ func (a *Agent) stopShard(ctx context.Context, shardID string, out chan<- *llmgw
|
||||
delete(a.shardReplica, shardID)
|
||||
a.mu.Unlock()
|
||||
|
||||
if err := a.sup.Stop(ctx, shardID); err != nil {
|
||||
ctx, span := obs.Tracer(tracerName).Start(ctx, "stop shard",
|
||||
trace.WithSpanKind(trace.SpanKindConsumer),
|
||||
trace.WithAttributes(
|
||||
attribute.String("llmgw.node", a.cfg.NodeID),
|
||||
attribute.String("llmgw.replica", replicaID),
|
||||
attribute.String("llmgw.shard", shardID),
|
||||
))
|
||||
err := a.sup.Stop(ctx, shardID)
|
||||
if err != nil {
|
||||
a.log.Warn("stopping a shard failed", "shard", shardID, "err", err)
|
||||
}
|
||||
a.stopKV(shardID, replicaID)
|
||||
obs.End(span, err)
|
||||
a.sup.Remove(shardID)
|
||||
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
@@ -1138,18 +1382,20 @@ const (
|
||||
//
|
||||
// The interface NCCL bootstraps over is this node's, not the cluster's: the
|
||||
// control plane sends one name for every node, and a node whose management
|
||||
// network is on ens3 rather than eth0 would have NCCL bind nothing. And on
|
||||
// RoCE, which GID entry routes is a property of these adapters.
|
||||
// network is on ens3 rather than eth0 would have NCCL bind nothing. On RoCE,
|
||||
// which GID entry routes is a property of these adapters. And a MIG instance
|
||||
// is opened by its UUID, which only this node knows.
|
||||
func (a *Agent) localizeSpec(spec *engines.ShardSpec) {
|
||||
if a.cfg.MgmtInterface != "" {
|
||||
spec.MgmtInterface = a.cfg.MgmtInterface
|
||||
}
|
||||
a.mu.Lock()
|
||||
hcas, gpus := a.inv.HCAs, a.inv.GPUs
|
||||
a.mu.Unlock()
|
||||
spec.GPUUUIDs = migUUIDs(gpus, spec.GPUIndices)
|
||||
if len(spec.HCAs) == 0 {
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
hcas := a.inv.HCAs
|
||||
a.mu.Unlock()
|
||||
want := map[string]bool{}
|
||||
for _, h := range spec.HCAs {
|
||||
want[h] = true
|
||||
@@ -1168,6 +1414,27 @@ func (a *Agent) localizeSpec(spec *engines.ShardSpec) {
|
||||
}
|
||||
}
|
||||
|
||||
// migUUIDs names a shard's devices by UUID when one of them is a MIG
|
||||
// instance, and is nil otherwise: whole cards keep being given by index, as
|
||||
// they always were. A device this node does not have is left to fail where
|
||||
// it always has, in the engine, rather than guessed at here.
|
||||
func migUUIDs(gpus []discovery.GPU, indices []int) []string {
|
||||
var out []string
|
||||
mig := false
|
||||
for _, i := range indices {
|
||||
k := slices.IndexFunc(gpus, func(g discovery.GPU) bool { return g.Index == i })
|
||||
if k < 0 || gpus[k].UUID == "" {
|
||||
return nil
|
||||
}
|
||||
mig = mig || gpus[k].IsMIG()
|
||||
out = append(out, gpus[k].UUID)
|
||||
}
|
||||
if !mig {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ResolveManagement is the address and interface this node's agent would use,
|
||||
// found the way normalize finds them, for a tool that reports them.
|
||||
func ResolveManagement(c Config) (addr, iface string) {
|
||||
@@ -1184,3 +1451,42 @@ func ResolveManagement(c Config) (addr, iface string) {
|
||||
}
|
||||
return addr, iface
|
||||
}
|
||||
|
||||
// reportCompute tells the control plane, once a session, what each running
|
||||
// shard computes on as soon as that is known: an adopted engine's is found
|
||||
// after registration, which carried what was known then.
|
||||
func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMessage) {
|
||||
sent := map[string]bool{}
|
||||
t := time.NewTicker(5 * time.Second)
|
||||
defer t.Stop()
|
||||
for range 60 {
|
||||
for _, st := range a.sup.Statuses() {
|
||||
if st.State.Terminal() || sent[st.ID] {
|
||||
continue
|
||||
}
|
||||
c := a.sup.ComputeOf(st.ID)
|
||||
if c == nil {
|
||||
continue
|
||||
}
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
// No state: a report of compute alone, not a transition.
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: sh.Spec.ReplicaID, ShardId: st.ID, Compute: computeToWire(c),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
sent[st.ID] = true
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -37,11 +37,12 @@ import (
|
||||
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
|
||||
mu sync.Mutex
|
||||
calls []adoptCall
|
||||
adopted []reconciler.AdoptedShard
|
||||
unclaim []string // shard ids to report back as unclaimed
|
||||
fetches []fetchReport
|
||||
computes []reconciler.Compute
|
||||
}
|
||||
|
||||
type status struct {
|
||||
@@ -61,6 +62,12 @@ func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.
|
||||
|
||||
// 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) OnShardCompute(replicaID, shardID string, c reconciler.Compute) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.computes = append(r.computes, c)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) OnFetchProgress(replicaID, shardID, source string, done, total int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
|
||||
func TestEnginesListenOnTheManagementAddressWhenItIsThisMachines(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"127.0.0.1": "127.0.0.1", // one of this machine's
|
||||
"203.0.113.9": "", // not: every address rather than none
|
||||
"gpu07.lan": "", // a name
|
||||
"0.0.0.0": "",
|
||||
"": "",
|
||||
}
|
||||
for mgmt, want := range cases {
|
||||
a := &Agent{cfg: Config{MgmtAddr: mgmt}}
|
||||
if got := a.bindAddr(); got != want {
|
||||
t.Errorf("management address %q: listens on %q, expected %q", mgmt, got, want)
|
||||
}
|
||||
}
|
||||
a := &Agent{}
|
||||
if got := a.engineBase(engines.ShardSpec{HTTPPort: 8000}); got != "http://127.0.0.1:8000" {
|
||||
t.Errorf("on every address, asked at %s", got)
|
||||
}
|
||||
if got := a.engineBase(engines.ShardSpec{HTTPPort: 8000, BindAddr: "10.0.0.147"}); got != "http://10.0.0.147:8000" {
|
||||
t.Errorf("on one address, asked at %s", got)
|
||||
}
|
||||
}
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"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/obs"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
@@ -94,7 +95,19 @@ type Config struct {
|
||||
Engines engines.Config `koanf:"engines"`
|
||||
Superv supervisor.Config `koanf:"supervisor"`
|
||||
|
||||
// Tracing sends this node's part of each replica start — the weights
|
||||
// fetched, the engine launched, the wait for it to answer — to an
|
||||
// OpenTelemetry collector, inside the control plane's trace of that
|
||||
// start. Give it the collector server.yaml's tracing names. Off unless
|
||||
// an endpoint is set here or in OTEL_EXPORTER_OTLP_ENDPOINT.
|
||||
Tracing obs.Tracing `koanf:"tracing"`
|
||||
|
||||
Version string
|
||||
|
||||
// Settings is the cluster's settings for this agent it was started on,
|
||||
// as the control plane sent them: what a newly sent version is
|
||||
// compared with.
|
||||
Settings string `koanf:"-"`
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
|
||||
@@ -47,6 +47,11 @@ type GPU struct {
|
||||
Fabric *GPUFabric `json:"fabric,omitempty"`
|
||||
// Health carries anything that should keep the GPU out of placement.
|
||||
Health GPUHealth `json:"health"`
|
||||
// MIG is set when this is a MIG instance rather than a whole card: its
|
||||
// Index is then MIGIndex's, its UUID the instance's (MIG-…), and its
|
||||
// memory the instance's share. Utilisation stays unreported, because no
|
||||
// driver measures it per instance.
|
||||
MIG *MIGDevice `json:"mig,omitempty"`
|
||||
}
|
||||
|
||||
// GPUHealth records conditions that make a GPU unsafe to schedule onto.
|
||||
@@ -94,8 +99,30 @@ func (g GPU) Class() string {
|
||||
}
|
||||
|
||||
// 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)
|
||||
// A MIG instance is named from its card's memory and then its profile,
|
||||
// so two instances are one class only when both the card and the slice
|
||||
// are the same.
|
||||
mem := g.MemoryTotalMB
|
||||
if g.MIG != nil {
|
||||
mem = g.MIG.ParentMemoryTotalMB
|
||||
}
|
||||
gb := int(math.Round(float64(mem) / 1024))
|
||||
class := fmt.Sprintf("%s-%dg", model, gb)
|
||||
if g.MIG != nil {
|
||||
return migClass(class, g.MIG.Profile)
|
||||
}
|
||||
return class
|
||||
}
|
||||
|
||||
// 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, ":")
|
||||
}
|
||||
|
||||
// MarshalJSON adds the derived class, so a consumer of the JSON does not have
|
||||
|
||||
@@ -14,7 +14,6 @@ package discovery
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/NVIDIA/go-nvml/pkg/nvml"
|
||||
)
|
||||
@@ -83,11 +82,18 @@ func DiscoverGPUsNVML(_ context.Context) ([]GPU, error) {
|
||||
}
|
||||
|
||||
assessHealth(&g)
|
||||
gpus = append(gpus, g)
|
||||
// A card in MIG mode is its instances, not itself: CUDA will not
|
||||
// open the whole card while MIG is on.
|
||||
devs, _ := migFromNVML(g, nvmlCard{dev})
|
||||
gpus = append(gpus, devs...)
|
||||
}
|
||||
|
||||
// A second pass, because a link's remote end may be a GPU enumerated later.
|
||||
// A second pass, because a link's remote end may be a GPU enumerated
|
||||
// later. MIG instances have no NVLink of their own to use.
|
||||
for i := range gpus {
|
||||
if gpus[i].IsMIG() {
|
||||
continue
|
||||
}
|
||||
dev, ret := nvml.DeviceGetHandleByIndex(gpus[i].Index)
|
||||
if ret != nvml.SUCCESS {
|
||||
continue
|
||||
@@ -136,13 +142,73 @@ func busIDString(b []int8) string {
|
||||
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:]
|
||||
// nvmlCard adapts an NVML device to what MIG discovery reads of it.
|
||||
type nvmlCard struct{ dev nvml.Device }
|
||||
|
||||
func (c nvmlCard) MIGMode() (current, pending, ok bool) {
|
||||
cur, pend, ret := c.dev.GetMigMode()
|
||||
if ret != nvml.SUCCESS {
|
||||
// NOT_SUPPORTED on a card without MIG, which is most of them.
|
||||
return false, false, false
|
||||
}
|
||||
return strings.TrimPrefix(s, ":")
|
||||
return cur == nvml.DEVICE_MIG_ENABLE, pend == nvml.DEVICE_MIG_ENABLE, true
|
||||
}
|
||||
|
||||
func (c nvmlCard) MaxMIGDevices() int {
|
||||
n, ret := c.dev.GetMaxMigDeviceCount()
|
||||
if ret != nvml.SUCCESS {
|
||||
return 0
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (c nvmlCard) MIGDevice(i int) (nvmlMIGInstance, bool) {
|
||||
h, ret := c.dev.GetMigDeviceHandleByIndex(i)
|
||||
if ret != nvml.SUCCESS {
|
||||
// NOT_FOUND for an index with no instance: the count is the most a
|
||||
// card can hold, not how many it has.
|
||||
return nil, false
|
||||
}
|
||||
return nvmlMIG{h}, true
|
||||
}
|
||||
|
||||
// nvmlMIG adapts a MIG device handle.
|
||||
type nvmlMIG struct{ dev nvml.Device }
|
||||
|
||||
func (m nvmlMIG) UUID() string {
|
||||
u, _ := m.dev.GetUUID()
|
||||
return u
|
||||
}
|
||||
|
||||
func (m nvmlMIG) Name() string {
|
||||
n, _ := m.dev.GetName()
|
||||
return n
|
||||
}
|
||||
|
||||
func (m nvmlMIG) Memory() (totalMB, usedMB int, ok bool) {
|
||||
mem, ret := m.dev.GetMemoryInfo()
|
||||
if ret != nvml.SUCCESS {
|
||||
return 0, 0, false
|
||||
}
|
||||
return int(mem.Total / (1 << 20)), int(mem.Used / (1 << 20)), true
|
||||
}
|
||||
|
||||
func (m nvmlMIG) Instances() (gpuInstance, computeInstance int) {
|
||||
gi, ret := m.dev.GetGpuInstanceId()
|
||||
if ret != nvml.SUCCESS {
|
||||
gi = -1
|
||||
}
|
||||
ci, ret := m.dev.GetComputeInstanceId()
|
||||
if ret != nvml.SUCCESS {
|
||||
ci = -1
|
||||
}
|
||||
return gi, ci
|
||||
}
|
||||
|
||||
func (m nvmlMIG) Attributes() (slices, memoryMB int, ok bool) {
|
||||
a, ret := m.dev.GetAttributes()
|
||||
if ret != nvml.SUCCESS {
|
||||
return 0, 0, false
|
||||
}
|
||||
return int(a.GpuInstanceSliceCount), int(a.MemorySizeMB), true
|
||||
}
|
||||
@@ -60,7 +60,33 @@ func DiscoverGPUsSMI(ctx context.Context, binary string) ([]GPU, error) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("running %s: %w", binary, err)
|
||||
}
|
||||
return parseSMI(string(out))
|
||||
cards, err := parseSMI(string(out))
|
||||
if err != nil || len(cards) == 0 {
|
||||
return cards, err
|
||||
}
|
||||
// The query lists whole cards, MIG mode or not. A card in MIG mode is
|
||||
// replaced by its instances, which only -L names and only the full
|
||||
// report sizes; asked only when some card could be in MIG mode, since
|
||||
// the report is the slowest thing nvidia-smi does.
|
||||
if !anyMIGCapable(cards) {
|
||||
return cards, nil
|
||||
}
|
||||
listed, report := readMIG(ctx, binary)
|
||||
return expandMIG(cards, listed, report), nil
|
||||
}
|
||||
|
||||
// anyMIGCapable reports whether any card could be in MIG mode: MIG began
|
||||
// with compute capability 8.0 (the A100). The consumer cards of the same
|
||||
// generations have none, but telling them apart by name is guesswork, and
|
||||
// asking one costs only two more nvidia-smi runs at registration.
|
||||
func anyMIGCapable(cards []GPU) bool {
|
||||
for _, c := range cards {
|
||||
major, _, _ := strings.Cut(c.ComputeCapability, ".")
|
||||
if n, err := strconv.Atoi(major); err == nil && n >= 8 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func parseSMI(out string) ([]GPU, error) {
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/xml"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
@@ -34,9 +35,16 @@ func GPUProcesses(ctx context.Context, binary string) (map[int][]int, error) {
|
||||
return parseGPUProcesses(out)
|
||||
}
|
||||
|
||||
// GPUMemory is one card's memory as nvidia-smi reports it at a moment: what
|
||||
// GPUMemory is one device's memory as nvidia-smi reports it at a moment: what
|
||||
// is allocated on it, and which processes hold how much of that.
|
||||
//
|
||||
// A device is what discovery reports: a whole card, or each MIG instance of a
|
||||
// card in MIG mode, in the same order, so the readings line up with the
|
||||
// inventory one for one.
|
||||
type GPUMemory struct {
|
||||
// Index is the device's number as placement knows it: the card's
|
||||
// nvidia-smi index, or MIGIndex's for an instance.
|
||||
Index int
|
||||
UsedMB int
|
||||
Procs []GPUProcess
|
||||
}
|
||||
@@ -66,48 +74,62 @@ func GPUMemoryByProcess(ctx context.Context, binary string) ([]GPUMemory, error)
|
||||
return parseGPUMemory(out)
|
||||
}
|
||||
|
||||
// smiLog is the part of nvidia-smi's XML report this reads: each GPU, in
|
||||
// index order, its memory and the processes on it.
|
||||
type smiLog struct {
|
||||
GPUs []struct {
|
||||
Used string `xml:"fb_memory_usage>used"`
|
||||
Processes []struct {
|
||||
PID string `xml:"pid"`
|
||||
Type string `xml:"type"`
|
||||
UsedMB string `xml:"used_memory"`
|
||||
} `xml:"processes>process_info"`
|
||||
} `xml:"gpu"`
|
||||
}
|
||||
|
||||
// parseGPUProcesses reads the report. GPUs are listed in index order, which
|
||||
// is nvidia-smi's own numbering — the one placement assigns by.
|
||||
// parseGPUProcesses reads the report, by device as placement numbers them:
|
||||
// nvidia-smi's index for a whole card, MIGIndex's for a MIG instance.
|
||||
func parseGPUProcesses(report []byte) (map[int][]int, error) {
|
||||
mem, err := parseGPUMemory(report)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := map[int][]int{}
|
||||
for i, g := range mem {
|
||||
for _, g := range mem {
|
||||
for _, p := range g.Procs {
|
||||
out[i] = append(out[i], p.PID)
|
||||
out[g.Index] = append(out[g.Index], p.PID)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// parseGPUMemory reads the report into each GPU's memory and its holders.
|
||||
// parseGPUMemory reads the report into each device's memory and its holders.
|
||||
//
|
||||
// Cards are listed in index order, which is nvidia-smi's own numbering. A
|
||||
// card in MIG mode is read as its instances, each with its own memory and
|
||||
// the processes nvidia-smi attributes to its GPU and compute instance: a
|
||||
// process on one instance is not on its neighbours, and the card's total
|
||||
// would have charged every instance with all of them.
|
||||
func parseGPUMemory(report []byte) ([]GPUMemory, error) {
|
||||
var log smiLog
|
||||
var log migLog
|
||||
if err := xml.Unmarshal(report, &log); err != nil {
|
||||
return nil, fmt.Errorf("reading nvidia-smi's report: %w", err)
|
||||
}
|
||||
out := make([]GPUMemory, len(log.GPUs))
|
||||
out := make([]GPUMemory, 0, len(log.GPUs))
|
||||
for i, g := range log.GPUs {
|
||||
out[i].UsedMB = mib(g.Used)
|
||||
for _, p := range g.Processes {
|
||||
if pid, err := strconv.Atoi(strings.TrimSpace(p.PID)); err == nil {
|
||||
out[i].Procs = append(out[i].Procs, GPUProcess{PID: pid, UsedMB: mib(p.UsedMB)})
|
||||
procs := func(gi, ci int) []GPUProcess {
|
||||
var ps []GPUProcess
|
||||
for _, p := range g.Processes {
|
||||
pid, err := strconv.Atoi(strings.TrimSpace(p.PID))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if gi >= 0 && (atoiOr(p.GPUInstance, -1) != gi || atoiOr(p.ComputeInstance, -1) != ci) {
|
||||
continue
|
||||
}
|
||||
ps = append(ps, GPUProcess{PID: pid, UsedMB: mib(p.UsedMB)})
|
||||
}
|
||||
return ps
|
||||
}
|
||||
if !enabled(g.Current) || len(g.Devices) == 0 {
|
||||
out = append(out, GPUMemory{Index: i, UsedMB: mib(g.Used), Procs: procs(-1, -1)})
|
||||
continue
|
||||
}
|
||||
devs := slices.Clone(g.Devices)
|
||||
slices.SortFunc(devs, func(a, b migLogDevice) int { return atoiOr(a.Index, -1) - atoiOr(b.Index, -1) })
|
||||
for _, d := range devs {
|
||||
out = append(out, GPUMemory{
|
||||
Index: MIGIndex(i, atoiOr(d.Index, 0)),
|
||||
UsedMB: mib(d.Used),
|
||||
Procs: procs(atoiOr(d.GPUInstance, -1), atoiOr(d.ComputeInstance, -1)),
|
||||
})
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
|
||||
@@ -57,8 +57,8 @@ func TestGPUMemoryIsReadPerProcess(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := []GPUMemory{
|
||||
{UsedMB: 14514, Procs: []GPUProcess{{PID: 71315, UsedMB: 4932}, {PID: 4100, UsedMB: 9500}}},
|
||||
{},
|
||||
{Index: 0, UsedMB: 14514, Procs: []GPUProcess{{PID: 71315, UsedMB: 4932}, {PID: 4100, UsedMB: 9500}}},
|
||||
{Index: 1},
|
||||
}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("got %+v, want %+v", got, want)
|
||||
|
||||
@@ -0,0 +1,560 @@
|
||||
package discovery
|
||||
|
||||
// MIG — Multi-Instance GPU — carves one A100, H100 or later card into up to
|
||||
// seven instances, each with its own memory, its own share of the
|
||||
// multiprocessors and its own UUID. To CUDA an instance is a device of its
|
||||
// own, and to placement it is one too: a card in MIG mode is reported as its
|
||||
// instances instead of as the whole card, each a GPU of a class named from
|
||||
// the card's and the instance's profile ("a100-80g-mig-3g.40gb").
|
||||
//
|
||||
// llmgw does not create instances. Which split a card gets is an operator's
|
||||
// decision, made with nvidia-smi mig; discovery reads what is there.
|
||||
//
|
||||
// Two facts about MIG shape everything below. A process is given one MIG
|
||||
// device: CUDA enumerates a single instance however many are made visible.
|
||||
// And instances have no NVLink or peer-to-peer access between them. A replica
|
||||
// on MIG therefore runs on exactly one instance, never a tensor-parallel group
|
||||
// of them, and the catalog refuses a profile that asks for more.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/xml"
|
||||
"fmt"
|
||||
"math"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// MIGDevice is what a MIG instance carries beyond any GPU's fields: which
|
||||
// card it is carved from, and how.
|
||||
type MIGDevice struct {
|
||||
// ParentIndex is the card's nvidia-smi index, which is what its sensors,
|
||||
// topology and wear are reported by.
|
||||
ParentIndex int `json:"parent_index"`
|
||||
ParentUUID string `json:"parent_uuid,omitempty"`
|
||||
// ParentMemoryTotalMB is the whole card's memory: the class is named from
|
||||
// the card, so an instance of an 80 GB card is not mistaken for one of a
|
||||
// 40 GB card that happens to have the same profile name.
|
||||
ParentMemoryTotalMB int `json:"parent_memory_total_mb"`
|
||||
// Device is the instance's number on its card, as nvidia-smi -L lists it.
|
||||
Device int `json:"device"`
|
||||
// Profile is NVIDIA's name for the slice, like "3g.40gb": three of the
|
||||
// card's seven compute slices and forty gigabytes of its memory.
|
||||
Profile string `json:"profile"`
|
||||
// GPUInstance and ComputeInstance are the instance's ids in nvidia-smi
|
||||
// mig -lgi and -lci, which is how nvidia-smi attributes processes to it.
|
||||
GPUInstance int `json:"gpu_instance"`
|
||||
ComputeInstance int `json:"compute_instance"`
|
||||
}
|
||||
|
||||
// migIndexBase starts the numbers MIG instances are given.
|
||||
//
|
||||
// Placement, the agent protocol and every running shard name a node's devices
|
||||
// by an integer, which for a whole card is its nvidia-smi index. An instance
|
||||
// has none of its own, so it is given 1000 + 100 × card + device: card 0's
|
||||
// instance 2 is 1002, card 3's instance 0 is 1300. Readable, stable while the
|
||||
// instances are, and never one a whole card could have.
|
||||
const migIndexBase = 1000
|
||||
|
||||
// MIGIndex is the device number placement knows a MIG instance by.
|
||||
func MIGIndex(card, device int) int { return migIndexBase + 100*card + device }
|
||||
|
||||
// IsMIGIndex reports whether a device number names a MIG instance.
|
||||
func IsMIGIndex(i int) bool { return i >= migIndexBase }
|
||||
|
||||
// DeviceLabel is how a device number is shown to people: "3" for a whole
|
||||
// card, "0/mig2" for card 0's MIG instance 2.
|
||||
func DeviceLabel(i int) string {
|
||||
if !IsMIGIndex(i) {
|
||||
return strconv.Itoa(i)
|
||||
}
|
||||
i -= migIndexBase
|
||||
return fmt.Sprintf("%d/mig%d", i/100, i%100)
|
||||
}
|
||||
|
||||
// DeviceLabels is DeviceLabel for each of a list.
|
||||
func DeviceLabels(is []int) []string {
|
||||
out := make([]string, len(is))
|
||||
for k, i := range is {
|
||||
out[k] = DeviceLabel(i)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// IsMIG reports whether this GPU is a MIG instance rather than a whole card.
|
||||
func (g GPU) IsMIG() bool { return g.MIG != nil }
|
||||
|
||||
// CardIndex is the nvidia-smi index of the physical card this GPU is on: its
|
||||
// own index for a whole card, its parent's for a MIG instance.
|
||||
func (g GPU) CardIndex() int {
|
||||
if g.MIG != nil {
|
||||
return g.MIG.ParentIndex
|
||||
}
|
||||
return g.Index
|
||||
}
|
||||
|
||||
// migClassMarker joins a card's class to an instance's profile in a MIG
|
||||
// class: "a100-80g" + "-mig-" + "3g.40gb".
|
||||
const migClassMarker = "-mig-"
|
||||
|
||||
// IsMIGClass reports whether a GPU class names MIG instances.
|
||||
func IsMIGClass(class string) bool { return strings.Contains(class, migClassMarker) }
|
||||
|
||||
// migProfileClean keeps a profile name to what a class may hold: "1g.10gb+me"
|
||||
// becomes "1g.10gb-me".
|
||||
var migProfileClean = regexp.MustCompile(`[^a-z0-9.]+`)
|
||||
|
||||
// migClass is the class of a MIG instance: the card's class and the profile.
|
||||
func migClass(parentClass, profile string) string {
|
||||
p := migProfileClean.ReplaceAllString(strings.ToLower(strings.TrimSpace(profile)), "-")
|
||||
return parentClass + migClassMarker + strings.Trim(p, "-")
|
||||
}
|
||||
|
||||
// MIGProfileMemoryGB reads the memory out of a MIG class or profile:
|
||||
// 40 from "a100-80g-mig-3g.40gb" or "3g.40gb". Zero when there is none.
|
||||
func MIGProfileMemoryGB(s string) int {
|
||||
if i := strings.LastIndex(s, migClassMarker); i >= 0 {
|
||||
s = s[i+len(migClassMarker):]
|
||||
}
|
||||
m := migProfileMemory.FindStringSubmatch(strings.ToLower(s))
|
||||
if m == nil {
|
||||
return 0
|
||||
}
|
||||
gb, _ := strconv.Atoi(m[1])
|
||||
return gb
|
||||
}
|
||||
|
||||
var migProfileMemory = regexp.MustCompile(`^\d+g\.(\d+)gb`)
|
||||
|
||||
// --- nvidia-smi -L ---
|
||||
|
||||
// listedMIG is one MIG instance as nvidia-smi -L lists it.
|
||||
type listedMIG struct {
|
||||
Device int
|
||||
Profile string
|
||||
UUID string
|
||||
}
|
||||
|
||||
var (
|
||||
// "GPU 0: NVIDIA A100-SXM4-80GB (UUID: GPU-5d5ba0d6-...)"
|
||||
smiListGPU = regexp.MustCompile(`^GPU (\d+):`)
|
||||
// " MIG 3g.40gb Device 0: (UUID: MIG-c6d4f1ef-...)". Drivers before
|
||||
// R470 name the instance MIG-GPU-<card uuid>/<gi>/<ci>, which CUDA of that
|
||||
// era takes in the same place, so the UUID is read as whatever it is.
|
||||
smiListMIG = regexp.MustCompile(`^\s+MIG\s+(\S+)\s+Device\s+(\d+):\s*\(UUID:\s*([^)\s]+)\)`)
|
||||
)
|
||||
|
||||
// parseSMIList reads nvidia-smi -L into each card's MIG instances, by the
|
||||
// card's index, in device order. A card with none is absent.
|
||||
func parseSMIList(out string) map[int][]listedMIG {
|
||||
res := map[int][]listedMIG{}
|
||||
card := -1
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
if m := smiListGPU.FindStringSubmatch(line); m != nil {
|
||||
card, _ = strconv.Atoi(m[1])
|
||||
continue
|
||||
}
|
||||
m := smiListMIG.FindStringSubmatch(line)
|
||||
if m == nil || card < 0 {
|
||||
continue
|
||||
}
|
||||
dev, _ := strconv.Atoi(m[2])
|
||||
res[card] = append(res[card], listedMIG{Device: dev, Profile: m[1], UUID: m[3]})
|
||||
}
|
||||
for c := range res {
|
||||
slices.SortFunc(res[c], func(a, b listedMIG) int { return a.Device - b.Device })
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// --- nvidia-smi -q -x ---
|
||||
|
||||
// migLog is the part of nvidia-smi's XML report MIG is read from: each card's
|
||||
// mode, and each instance's ids and memory.
|
||||
type migLog struct {
|
||||
GPUs []migLogGPU `xml:"gpu"`
|
||||
}
|
||||
|
||||
type migLogGPU struct {
|
||||
ID string `xml:"id,attr"`
|
||||
Current string `xml:"mig_mode>current_mig"`
|
||||
Pending string `xml:"mig_mode>pending_mig"`
|
||||
Devices []migLogDevice `xml:"mig_devices>mig_device"`
|
||||
// Processes carry the GPU and compute instance they run in on a card
|
||||
// in MIG mode, and "N/A" for both otherwise.
|
||||
Processes []struct {
|
||||
GPUInstance string `xml:"gpu_instance_id"`
|
||||
ComputeInstance string `xml:"compute_instance_id"`
|
||||
PID string `xml:"pid"`
|
||||
UsedMB string `xml:"used_memory"`
|
||||
} `xml:"processes>process_info"`
|
||||
Used string `xml:"fb_memory_usage>used"`
|
||||
}
|
||||
|
||||
type migLogDevice struct {
|
||||
Index string `xml:"index"`
|
||||
GPUInstance string `xml:"gpu_instance_id"`
|
||||
ComputeInstance string `xml:"compute_instance_id"`
|
||||
Total string `xml:"fb_memory_usage>total"`
|
||||
Used string `xml:"fb_memory_usage>used"`
|
||||
}
|
||||
|
||||
// migState is one card's MIG mode and instances from the report.
|
||||
type migState struct {
|
||||
BusID string
|
||||
Enabled bool
|
||||
PendingEnabled bool
|
||||
Devices []migReported
|
||||
}
|
||||
|
||||
// migReported is one instance as the report describes it.
|
||||
type migReported struct {
|
||||
Device, GPUInstance, ComputeInstance int
|
||||
TotalMB, UsedMB int
|
||||
}
|
||||
|
||||
// enabled reads nvidia-smi's mode words: "Enabled", "Disabled", "N/A".
|
||||
func enabled(s string) bool { return strings.EqualFold(strings.TrimSpace(s), "Enabled") }
|
||||
|
||||
func atoiOr(s string, def int) int {
|
||||
n, err := strconv.Atoi(strings.TrimSpace(s))
|
||||
if err != nil {
|
||||
return def
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// parseMIGReport reads each card's MIG state from the XML report, in the
|
||||
// report's order, which is nvidia-smi's index order.
|
||||
func parseMIGReport(report []byte) ([]migState, error) {
|
||||
var log migLog
|
||||
if err := xml.Unmarshal(report, &log); err != nil {
|
||||
return nil, fmt.Errorf("reading nvidia-smi's report: %w", err)
|
||||
}
|
||||
out := make([]migState, len(log.GPUs))
|
||||
for i, g := range log.GPUs {
|
||||
out[i] = migState{BusID: g.ID, Enabled: enabled(g.Current), PendingEnabled: enabled(g.Pending)}
|
||||
for _, d := range g.Devices {
|
||||
out[i].Devices = append(out[i].Devices, migReported{
|
||||
Device: atoiOr(d.Index, -1), GPUInstance: atoiOr(d.GPUInstance, -1),
|
||||
ComputeInstance: atoiOr(d.ComputeInstance, -1),
|
||||
TotalMB: mib(d.Total), UsedMB: mib(d.Used),
|
||||
})
|
||||
}
|
||||
slices.SortFunc(out[i].Devices, func(a, b migReported) int { return a.Device - b.Device })
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// --- putting it together ---
|
||||
|
||||
// migProblemNoInstances keeps a card in MIG mode with nothing carved from it
|
||||
// out of placement: CUDA cannot use the whole card while MIG is on, and
|
||||
// there is no instance to use instead.
|
||||
const migProblemNoInstances = "MIG is enabled with no instances created: nothing can run on it " +
|
||||
"until instances are created (nvidia-smi mig -cgi … -C) or MIG is turned off"
|
||||
|
||||
// migProblemUnnamed is a card whose instances exist but could not be named:
|
||||
// an engine is given an instance by its UUID, and without one it cannot be.
|
||||
const migProblemUnnamed = "MIG is enabled and has instances, but nvidia-smi -L did not name them, " +
|
||||
"so none can be given to an engine"
|
||||
|
||||
// expandMIG replaces each card in MIG mode with its instances.
|
||||
//
|
||||
// listed is nvidia-smi -L's instances by card, which names them; report is
|
||||
// the XML report's state per card, which gives their memory and ids. Either
|
||||
// may be missing: -L alone still says which instances exist, and the report
|
||||
// alone still says a card is in MIG mode, so a card it says has instances
|
||||
// that -L did not name is kept as the whole card with the reason, rather
|
||||
// than placed as a card CUDA will not open.
|
||||
func expandMIG(cards []GPU, listed map[int][]listedMIG, report []migState) []GPU {
|
||||
if len(listed) == 0 && !slices.ContainsFunc(report, func(s migState) bool { return s.Enabled }) {
|
||||
return cards
|
||||
}
|
||||
byBus := map[string]migState{}
|
||||
for _, s := range report {
|
||||
byBus[normalizeBusID(s.BusID)] = s
|
||||
}
|
||||
out := make([]GPU, 0, len(cards))
|
||||
for pos, c := range cards {
|
||||
st, ok := byBus[normalizeBusID(c.PCIBusID)]
|
||||
if !ok && pos < len(report) {
|
||||
st = report[pos]
|
||||
}
|
||||
devs := listed[c.Index]
|
||||
if !st.Enabled && len(devs) == 0 {
|
||||
out = append(out, c)
|
||||
continue
|
||||
}
|
||||
if len(devs) == 0 {
|
||||
p := migProblemNoInstances
|
||||
if len(st.Devices) > 0 {
|
||||
p = migProblemUnnamed
|
||||
}
|
||||
c.Health.Problems = append(c.Health.Problems, p)
|
||||
out = append(out, c)
|
||||
continue
|
||||
}
|
||||
for _, d := range devs {
|
||||
var rep migReported
|
||||
found := false
|
||||
for _, r := range st.Devices {
|
||||
if r.Device == d.Device {
|
||||
rep, found = r, true
|
||||
}
|
||||
}
|
||||
mig := migInstance(c, d.Device, d.Profile, d.UUID)
|
||||
if found {
|
||||
mig.MIG.GPUInstance, mig.MIG.ComputeInstance = rep.GPUInstance, rep.ComputeInstance
|
||||
mig.MemoryTotalMB, mig.MemoryUsedMB = rep.TotalMB, rep.UsedMB
|
||||
}
|
||||
if mig.MemoryTotalMB == 0 {
|
||||
// The profile says it to the gigabyte, which is close enough
|
||||
// to place by when the report could not be read.
|
||||
mig.MemoryTotalMB = MIGProfileMemoryGB(d.Profile) * 1024
|
||||
}
|
||||
assessHealth(&mig)
|
||||
out = append(out, mig)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// migInstance builds a MIG instance of card: everything about the card that
|
||||
// is shared — driver, compute capability, NUMA node, bus, temperature and
|
||||
// ECC errors — and nothing that is not. Not NVLink, which an instance cannot
|
||||
// use; not utilisation, which no driver reports per instance; and not the
|
||||
// card's memory or problems, which the instance's own replace.
|
||||
func migInstance(card GPU, device int, profile, uuid string) GPU {
|
||||
g := GPU{
|
||||
Index: MIGIndex(card.Index, device), UUID: uuid,
|
||||
Name: strings.TrimSpace(card.Name + " MIG " + profile),
|
||||
ComputeCapability: card.ComputeCapability, NUMANode: card.NUMANode,
|
||||
PCIBusID: card.PCIBusID, DriverVersion: card.DriverVersion, CUDAVersion: card.CUDAVersion,
|
||||
NICPaths: card.NICPaths, Fabric: card.Fabric,
|
||||
Health: GPUHealth{ECCErrors: card.Health.ECCErrors, TemperatureC: card.Health.TemperatureC},
|
||||
MIG: &MIGDevice{
|
||||
ParentIndex: card.Index, ParentUUID: card.UUID, ParentMemoryTotalMB: card.MemoryTotalMB,
|
||||
Device: device, Profile: profile, GPUInstance: -1, ComputeInstance: -1,
|
||||
},
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
// readMIG asks nvidia-smi for what expandMIG needs. A failure of either
|
||||
// command leaves its half empty, which expandMIG is written to take.
|
||||
func readMIG(ctx context.Context, binary string) (map[int][]listedMIG, []migState) {
|
||||
var listed map[int][]listedMIG
|
||||
if out, err := exec.CommandContext(ctx, binary, "-L").Output(); err == nil {
|
||||
listed = parseSMIList(string(out))
|
||||
}
|
||||
var report []migState
|
||||
if out, err := exec.CommandContext(ctx, binary, "-q", "-x").Output(); err == nil {
|
||||
report, _ = parseMIGReport(out)
|
||||
}
|
||||
return listed, report
|
||||
}
|
||||
|
||||
// --- NVML ---
|
||||
|
||||
// nvmlMIGCard is the part of an NVML device MIG discovery reads, behind an
|
||||
// interface so the logic is tested without a driver: gpu_nvml.go adapts
|
||||
// nvml.Device to it.
|
||||
type nvmlMIGCard interface {
|
||||
// MIGMode is nvmlDeviceGetMigMode; ok is false on a card without MIG.
|
||||
MIGMode() (current, pending, ok bool)
|
||||
// MaxMIGDevices is nvmlDeviceGetMaxMigDeviceCount.
|
||||
MaxMIGDevices() int
|
||||
// MIGDevice is nvmlDeviceGetMigDeviceHandleByIndex; ok is false for an
|
||||
// index with no instance behind it, which is how NVML says so.
|
||||
MIGDevice(i int) (nvmlMIGInstance, bool)
|
||||
}
|
||||
|
||||
// nvmlMIGInstance is one MIG device handle.
|
||||
type nvmlMIGInstance interface {
|
||||
UUID() string
|
||||
// Name is "NVIDIA A100-SXM4-80GB MIG 3g.40gb" on the drivers that name
|
||||
// MIG devices, which is where the profile is read from.
|
||||
Name() string
|
||||
Memory() (totalMB, usedMB int, ok bool)
|
||||
Instances() (gpuInstance, computeInstance int)
|
||||
// Slices and MemoryMB are the instance's attributes, which name the
|
||||
// profile on a driver whose device name does not.
|
||||
Attributes() (slices, memoryMB int, ok bool)
|
||||
}
|
||||
|
||||
var nvmlMIGName = regexp.MustCompile(`MIG\s+(\S+)\s*$`)
|
||||
|
||||
// migFromNVML is what one card is to placement under NVML: its instances if
|
||||
// it is in MIG mode, itself if it is not. isMIG says which.
|
||||
func migFromNVML(card GPU, dev nvmlMIGCard) (out []GPU, isMIG bool) {
|
||||
current, _, ok := dev.MIGMode()
|
||||
if !ok || !current {
|
||||
return []GPU{card}, false
|
||||
}
|
||||
n := dev.MaxMIGDevices()
|
||||
for i := range n {
|
||||
h, ok := dev.MIGDevice(i)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
profile := ""
|
||||
if m := nvmlMIGName.FindStringSubmatch(h.Name()); m != nil {
|
||||
profile = m[1]
|
||||
} else if n, mem, ok := h.Attributes(); ok {
|
||||
// The convention NVIDIA's names follow: compute slices, then the
|
||||
// memory in whole gigabytes.
|
||||
profile = fmt.Sprintf("%dg.%dgb", n, int(math.Round(float64(mem)/1024)))
|
||||
}
|
||||
g := migInstance(card, i, profile, h.UUID())
|
||||
g.MIG.GPUInstance, g.MIG.ComputeInstance = h.Instances()
|
||||
if total, used, ok := h.Memory(); ok {
|
||||
g.MemoryTotalMB, g.MemoryUsedMB = total, used
|
||||
} else {
|
||||
g.MemoryTotalMB = MIGProfileMemoryGB(profile) * 1024
|
||||
}
|
||||
assessHealth(&g)
|
||||
out = append(out, g)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
card.Health.Problems = append(card.Health.Problems, migProblemNoInstances)
|
||||
return []GPU{card}, true
|
||||
}
|
||||
return out, true
|
||||
}
|
||||
|
||||
// --- what the doctor reports ---
|
||||
|
||||
// MIGCard is one card's MIG state, as `llmgw doctor` reports it.
|
||||
type MIGCard struct {
|
||||
Index int `json:"index"`
|
||||
Name string `json:"name,omitempty"`
|
||||
// Enabled is the mode in force; PendingEnabled the mode after the next
|
||||
// reset of the card, which differs from it until the card is reset.
|
||||
Enabled bool `json:"enabled"`
|
||||
PendingEnabled bool `json:"pending_enabled"`
|
||||
// Instances are the MIG devices an engine can be given, as nvidia-smi -L
|
||||
// lists them.
|
||||
Instances []MIGListed `json:"instances,omitempty"`
|
||||
// GPUInstances and ComputeInstances are nvidia-smi mig -lgi and -lci: a
|
||||
// GPU instance with no compute instance in it is not a device yet.
|
||||
GPUInstances []MIGInstance `json:"gpu_instances,omitempty"`
|
||||
ComputeInstances []MIGInstance `json:"compute_instances,omitempty"`
|
||||
}
|
||||
|
||||
// MIGListed is one instance an engine can be given.
|
||||
type MIGListed struct {
|
||||
Device int `json:"device"`
|
||||
Profile string `json:"profile"`
|
||||
UUID string `json:"uuid"`
|
||||
}
|
||||
|
||||
// MIGInstance is one row of nvidia-smi mig -lgi or -lci.
|
||||
type MIGInstance struct {
|
||||
GPU int `json:"gpu"`
|
||||
Profile string `json:"profile"`
|
||||
// ProfileID is NVIDIA's number for the profile.
|
||||
ProfileID int `json:"profile_id"`
|
||||
// GPUInstance is the GPU instance's id; ComputeInstance the compute
|
||||
// instance's within it, -1 in a -lgi row.
|
||||
GPUInstance int `json:"gpu_instance"`
|
||||
ComputeInstance int `json:"compute_instance"`
|
||||
Placement string `json:"placement,omitempty"`
|
||||
}
|
||||
|
||||
var (
|
||||
// -lgi: "| 0 MIG 3g.40gb 9 2 4:4 |"
|
||||
migLGIRow = regexp.MustCompile(`^\|\s*(\d+)\s+MIG\s+(\S+)\s+(\d+)\s+(\d+)\s+(\d+:\d+)\s*\|`)
|
||||
// -lci: "| 0 1 MIG 3g.40gb 2 0 0:3 |"
|
||||
migLCIRow = regexp.MustCompile(`^\|\s*(\d+)\s+(\d+)\s+MIG\s+(\S+)\s+(\d+)\s+(\d+)\s+(\d+:\d+)\s*\|`)
|
||||
)
|
||||
|
||||
// ParseMIGInstances reads nvidia-smi mig -lgi or -lci, whichever it is
|
||||
// given: the two tables differ by the GPU instance column -lci adds. "No GPU
|
||||
// instances found" and the like are no rows, not an error.
|
||||
func ParseMIGInstances(out string) []MIGInstance {
|
||||
var res []MIGInstance
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if m := migLCIRow.FindStringSubmatch(line); m != nil {
|
||||
res = append(res, MIGInstance{
|
||||
GPU: atoiOr(m[1], -1), GPUInstance: atoiOr(m[2], -1), Profile: m[3],
|
||||
ProfileID: atoiOr(m[4], -1), ComputeInstance: atoiOr(m[5], -1), Placement: m[6],
|
||||
})
|
||||
continue
|
||||
}
|
||||
if m := migLGIRow.FindStringSubmatch(line); m != nil {
|
||||
res = append(res, MIGInstance{
|
||||
GPU: atoiOr(m[1], -1), Profile: m[2], ProfileID: atoiOr(m[3], -1),
|
||||
GPUInstance: atoiOr(m[4], -1), ComputeInstance: -1, Placement: m[5],
|
||||
})
|
||||
}
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// DiscoverMIG reads every card's MIG state for the doctor: the mode from the
|
||||
// XML report, the devices from -L, the instances from nvidia-smi mig. Only
|
||||
// cards that support MIG are listed, and an error is returned only when the
|
||||
// report itself cannot be read.
|
||||
func DiscoverMIG(ctx context.Context, binary string) ([]MIGCard, error) {
|
||||
out, err := exec.CommandContext(ctx, binary, "-q", "-x").Output()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("running %s -q -x: %w", binary, err)
|
||||
}
|
||||
var list, lgi, lci string
|
||||
if b, err := exec.CommandContext(ctx, binary, "-L").Output(); err == nil {
|
||||
list = string(b)
|
||||
}
|
||||
// nvidia-smi mig exits non-zero when there is nothing to list, with the
|
||||
// table or the reason on its output either way.
|
||||
b, _ := exec.CommandContext(ctx, binary, "mig", "-lgi").CombinedOutput()
|
||||
lgi = string(b)
|
||||
b, _ = exec.CommandContext(ctx, binary, "mig", "-lci").CombinedOutput()
|
||||
lci = string(b)
|
||||
return buildMIGCards(out, list, lgi, lci)
|
||||
}
|
||||
|
||||
// buildMIGCards is DiscoverMIG from the commands' output.
|
||||
func buildMIGCards(report []byte, list, lgi, lci string) ([]MIGCard, error) {
|
||||
var log struct {
|
||||
GPUs []struct {
|
||||
Name string `xml:"product_name"`
|
||||
Current string `xml:"mig_mode>current_mig"`
|
||||
Pending string `xml:"mig_mode>pending_mig"`
|
||||
} `xml:"gpu"`
|
||||
}
|
||||
if err := xml.Unmarshal(report, &log); err != nil {
|
||||
return nil, fmt.Errorf("reading nvidia-smi's report: %w", err)
|
||||
}
|
||||
listed := parseSMIList(list)
|
||||
gis, cis := ParseMIGInstances(lgi), ParseMIGInstances(lci)
|
||||
var out []MIGCard
|
||||
for i, g := range log.GPUs {
|
||||
cur, pend := strings.TrimSpace(g.Current), strings.TrimSpace(g.Pending)
|
||||
// A card without MIG says N/A, or nothing on a driver that predates
|
||||
// it; neither is worth a line.
|
||||
if !enabled(cur) && !enabled(pend) && !strings.EqualFold(cur, "Disabled") {
|
||||
continue
|
||||
}
|
||||
c := MIGCard{Index: i, Name: strings.TrimSpace(g.Name), Enabled: enabled(cur), PendingEnabled: enabled(pend)}
|
||||
for _, l := range listed[i] {
|
||||
c.Instances = append(c.Instances, MIGListed(l))
|
||||
}
|
||||
for _, gi := range gis {
|
||||
if gi.GPU == i && gi.ComputeInstance < 0 {
|
||||
c.GPUInstances = append(c.GPUInstances, gi)
|
||||
}
|
||||
}
|
||||
for _, ci := range cis {
|
||||
if ci.GPU == i && ci.ComputeInstance >= 0 {
|
||||
c.ComputeInstances = append(c.ComputeInstances, ci)
|
||||
}
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,418 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// Two A100 80GB cards: card 0 in MIG mode, cut into two 3g.40gb instances,
|
||||
// card 1 whole. What nvidia-smi prints for that machine, by command. None of
|
||||
// it is from real MIG hardware, which this was written without; the shapes
|
||||
// follow NVIDIA's documentation and the output of drivers R470 to R550.
|
||||
const (
|
||||
migQuery = `0, GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77, NVIDIA A100-SXM4-80GB, 81920, 0, 8.0, 00000000:07:00.0, 550.54.15, 31, 0, [N/A]
|
||||
1, GPU-1f2e3d4c-5b6a-7980-a1b2-c3d4e5f60718, NVIDIA A100-SXM4-80GB, 81920, 0, 8.0, 00000000:0F:00.0, 550.54.15, 33, 0, 0
|
||||
`
|
||||
migList = `GPU 0: NVIDIA A100-SXM4-80GB (UUID: GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77)
|
||||
MIG 3g.40gb Device 0: (UUID: MIG-c6d4f1ef-42e4-5de3-91c7-45d71c87eb3f)
|
||||
MIG 3g.40gb Device 1: (UUID: MIG-0f7c3a2e-9d1b-5a8e-b3c4-1e2f3a4b5c6d)
|
||||
GPU 1: NVIDIA A100-SXM4-80GB (UUID: GPU-1f2e3d4c-5b6a-7980-a1b2-c3d4e5f60718)
|
||||
`
|
||||
migReport = `<?xml version="1.0" ?>
|
||||
<nvidia_smi_log>
|
||||
<driver_version>550.54.15</driver_version>
|
||||
<gpu id="00000000:07:00.0">
|
||||
<product_name>NVIDIA A100-SXM4-80GB</product_name>
|
||||
<mig_mode><current_mig>Enabled</current_mig><pending_mig>Enabled</pending_mig></mig_mode>
|
||||
<mig_devices>
|
||||
<mig_device>
|
||||
<index>1</index>
|
||||
<gpu_instance_id>1</gpu_instance_id>
|
||||
<compute_instance_id>0</compute_instance_id>
|
||||
<device_attributes><shared><multiprocessor_count>42</multiprocessor_count></shared></device_attributes>
|
||||
<fb_memory_usage><total>40192 MiB</total><reserved>0 MiB</reserved><used>21000 MiB</used><free>19192 MiB</free></fb_memory_usage>
|
||||
</mig_device>
|
||||
<mig_device>
|
||||
<index>0</index>
|
||||
<gpu_instance_id>2</gpu_instance_id>
|
||||
<compute_instance_id>0</compute_instance_id>
|
||||
<fb_memory_usage><total>40192 MiB</total><reserved>0 MiB</reserved><used>37 MiB</used><free>40155 MiB</free></fb_memory_usage>
|
||||
</mig_device>
|
||||
</mig_devices>
|
||||
<fb_memory_usage><total>81920 MiB</total><used>21037 MiB</used></fb_memory_usage>
|
||||
<processes>
|
||||
<process_info>
|
||||
<gpu_instance_id>1</gpu_instance_id>
|
||||
<compute_instance_id>0</compute_instance_id>
|
||||
<pid>4100</pid><type>C</type><process_name>vllm</process_name>
|
||||
<used_memory>20900 MiB</used_memory>
|
||||
</process_info>
|
||||
</processes>
|
||||
</gpu>
|
||||
<gpu id="00000000:0F:00.0">
|
||||
<product_name>NVIDIA A100-SXM4-80GB</product_name>
|
||||
<mig_mode><current_mig>Disabled</current_mig><pending_mig>Disabled</pending_mig></mig_mode>
|
||||
<mig_devices>None</mig_devices>
|
||||
<fb_memory_usage><total>81920 MiB</total><used>300 MiB</used></fb_memory_usage>
|
||||
<processes>
|
||||
<process_info>
|
||||
<gpu_instance_id>N/A</gpu_instance_id>
|
||||
<compute_instance_id>N/A</compute_instance_id>
|
||||
<pid>5200</pid><type>C</type><used_memory>250 MiB</used_memory>
|
||||
</process_info>
|
||||
</processes>
|
||||
</gpu>
|
||||
</nvidia_smi_log>`
|
||||
migLGI = `+-------------------------------------------------------+
|
||||
| GPU instances: |
|
||||
| GPU Name Profile Instance Placement |
|
||||
| ID ID Start:Size |
|
||||
|=======================================================|
|
||||
| 0 MIG 3g.40gb 9 1 4:4 |
|
||||
+-------------------------------------------------------+
|
||||
| 0 MIG 3g.40gb 9 2 0:4 |
|
||||
+-------------------------------------------------------+
|
||||
`
|
||||
migLCI = `+--------------------------------------------------------------------+
|
||||
| Compute instances: |
|
||||
| GPU GPU Name Profile Instance Placement |
|
||||
| Instance ID ID Start:Size |
|
||||
| ID |
|
||||
|====================================================================|
|
||||
| 0 1 MIG 3g.40gb 2 0 0:3 |
|
||||
+--------------------------------------------------------------------+
|
||||
| 0 2 MIG 3g.40gb 2 0 0:3 |
|
||||
+--------------------------------------------------------------------+
|
||||
`
|
||||
)
|
||||
|
||||
func migMachine(t *testing.T) []GPU {
|
||||
t.Helper()
|
||||
cards, err := parseSMI(migQuery)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
report, err := parseMIGReport([]byte(migReport))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return expandMIG(cards, parseSMIList(migList), report)
|
||||
}
|
||||
|
||||
func TestACardInMIGModeIsReportedAsItsInstancesInsteadOfItself(t *testing.T) {
|
||||
gpus := migMachine(t)
|
||||
var got []string
|
||||
for _, g := range gpus {
|
||||
got = append(got, DeviceLabel(g.Index)+" "+g.Class())
|
||||
}
|
||||
want := []string{"0/mig0 a100-80g-mig-3g.40gb", "0/mig1 a100-80g-mig-3g.40gb", "1 a100-80g"}
|
||||
if !slices.Equal(got, want) {
|
||||
t.Fatalf("devices = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
mig := gpus[1]
|
||||
if mig.Index != 1001 || mig.UUID != "MIG-0f7c3a2e-9d1b-5a8e-b3c4-1e2f3a4b5c6d" {
|
||||
t.Errorf("instance 1 is %d %s; an engine is given it by that UUID", mig.Index, mig.UUID)
|
||||
}
|
||||
if mig.MemoryTotalMB != 40192 || mig.MemoryUsedMB != 21000 {
|
||||
t.Errorf("instance 1 memory %d/%d, want its own 21000/40192 rather than the card's",
|
||||
mig.MemoryUsedMB, mig.MemoryTotalMB)
|
||||
}
|
||||
if m := mig.MIG; m == nil || m.ParentIndex != 0 || m.Device != 1 || m.Profile != "3g.40gb" ||
|
||||
m.GPUInstance != 1 || m.ComputeInstance != 0 || m.ParentUUID != "GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77" {
|
||||
t.Errorf("instance 1's parentage: %+v", mig.MIG)
|
||||
}
|
||||
if mig.ComputeCapability != "8.0" || mig.DriverVersion != "550.54.15" || mig.Health.TemperatureC != 31 {
|
||||
t.Errorf("what the card and its instances share was not carried: %+v", mig)
|
||||
}
|
||||
// Twenty gigabytes in use on an instance is a foreign tenant there, and
|
||||
// only there: instance 0 next to it is free.
|
||||
if mig.Health.OK() || !gpus[0].Health.OK() {
|
||||
t.Errorf("foreign memory judged per instance: 0 %v, 1 %v", gpus[0].Health.Problems, mig.Health.Problems)
|
||||
}
|
||||
// No driver reports utilisation per instance, and none is made up.
|
||||
if mig.UtilizationPercent != 0 {
|
||||
t.Errorf("an instance reported utilisation %d", mig.UtilizationPercent)
|
||||
}
|
||||
if gpus[2].IsMIG() || gpus[2].Index != 1 || gpus[2].CardIndex() != 1 || gpus[0].CardIndex() != 0 {
|
||||
t.Errorf("the whole card beside it changed: %+v", gpus[2])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMIGInstanceSurvivesTheTripToTheControlPlaneWithItsClass(t *testing.T) {
|
||||
g := migMachine(t)[0]
|
||||
b, err := json.Marshal(g)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(b), `"class":"a100-80g-mig-3g.40gb"`) {
|
||||
t.Errorf("the JSON carries no class for the instance: %s", b)
|
||||
}
|
||||
var back GPU
|
||||
if err := json.Unmarshal(b, &back); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if back.Class() != g.Class() || !reflect.DeepEqual(back.MIG, g.MIG) {
|
||||
t.Errorf("came back as %s %+v", back.Class(), back.MIG)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMIGClassesAreNamedFromTheCardAndTheProfile(t *testing.T) {
|
||||
for _, c := range []struct {
|
||||
name string
|
||||
cardMB int
|
||||
profile string
|
||||
want string
|
||||
}{
|
||||
{"NVIDIA A100-SXM4-80GB MIG 3g.40gb", 81920, "3g.40gb", "a100-80g-mig-3g.40gb"},
|
||||
// The same profile name on a 40 GB card is another class: its
|
||||
// instance is the same size, but the card is not, and neither is
|
||||
// the set of profiles that card can be cut to.
|
||||
{"NVIDIA A100-SXM4-40GB MIG 3g.20gb", 40960, "3g.20gb", "a100-40g-mig-3g.20gb"},
|
||||
{"NVIDIA H100 80GB HBM3 MIG 1g.10gb+me", 81559, "1g.10gb+me", "h100-80g-mig-1g.10gb-me"},
|
||||
{"NVIDIA H100 80GB HBM3 MIG 7g.80gb", 81559, "7g.80gb", "h100-80g-mig-7g.80gb"},
|
||||
} {
|
||||
g := GPU{Name: c.name, MemoryTotalMB: 10000, MIG: &MIGDevice{ParentMemoryTotalMB: c.cardMB, Profile: c.profile}}
|
||||
if got := g.Class(); got != c.want {
|
||||
t.Errorf("%s -> %s, want %s", c.name, got, c.want)
|
||||
}
|
||||
if !IsMIGClass(g.Class()) {
|
||||
t.Errorf("%s is not recognised as a MIG class", g.Class())
|
||||
}
|
||||
}
|
||||
if IsMIGClass("a100-80g") || IsMIGClass("6000-24g") {
|
||||
t.Error("a whole card's class was taken for a MIG one")
|
||||
}
|
||||
if MIGProfileMemoryGB("a100-80g-mig-3g.40gb") != 40 || MIGProfileMemoryGB("1g.10gb+me") != 10 || MIGProfileMemoryGB("a100-80g") != 0 {
|
||||
t.Error("the instance's memory is misread from its profile")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeviceNumbersAreShownAsCardAndInstance(t *testing.T) {
|
||||
for i, want := range map[int]string{0: "0", 7: "7", MIGIndex(0, 2): "0/mig2", MIGIndex(3, 0): "3/mig0", 1302: "3/mig2"} {
|
||||
if got := DeviceLabel(i); got != want {
|
||||
t.Errorf("DeviceLabel(%d) = %s, want %s", i, got, want)
|
||||
}
|
||||
}
|
||||
if IsMIGIndex(15) || !IsMIGIndex(MIGIndex(0, 0)) {
|
||||
t.Error("whole cards and instances are told apart wrongly")
|
||||
}
|
||||
}
|
||||
|
||||
func TestACardInMIGModeWithNoInstancesIsKeptOutOfPlacement(t *testing.T) {
|
||||
cards, _ := parseSMI(migQuery)
|
||||
report, _ := parseMIGReport([]byte(strings.Replace(migReport,
|
||||
migReport[strings.Index(migReport, "<mig_devices>"):strings.Index(migReport, "</mig_devices>")+len("</mig_devices>")],
|
||||
"<mig_devices>None</mig_devices>", 1)))
|
||||
gpus := expandMIG(cards, parseSMIList("GPU 0: NVIDIA A100-SXM4-80GB (UUID: GPU-x)\nGPU 1: NVIDIA A100-SXM4-80GB (UUID: GPU-y)\n"), report)
|
||||
if len(gpus) != 2 || gpus[0].IsMIG() {
|
||||
t.Fatalf("got %+v", gpus)
|
||||
}
|
||||
// CUDA will not open the whole card while MIG is on, and there is no
|
||||
// instance to open instead: placing anything here fails at start.
|
||||
if gpus[0].Health.OK() || !strings.Contains(gpus[0].Health.Problems[0], "no instances") {
|
||||
t.Errorf("card 0's problems: %v", gpus[0].Health.Problems)
|
||||
}
|
||||
if !gpus[1].Health.OK() {
|
||||
t.Errorf("card 1, not in MIG mode, was flagged: %v", gpus[1].Health.Problems)
|
||||
}
|
||||
|
||||
// Instances the report sees but -L could not name cannot be given to an
|
||||
// engine either; the card says why.
|
||||
report, _ = parseMIGReport([]byte(migReport))
|
||||
gpus = expandMIG(cards, nil, report)
|
||||
if len(gpus) != 2 || gpus[0].Health.OK() || !strings.Contains(gpus[0].Health.Problems[0], "did not name") {
|
||||
t.Errorf("with -L failing: %+v", gpus)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstancesAreFoundFromTheListAloneWhenTheReportCannotBeRead(t *testing.T) {
|
||||
cards, _ := parseSMI(migQuery)
|
||||
gpus := expandMIG(cards, parseSMIList(migList), nil)
|
||||
if len(gpus) != 3 || !gpus[0].IsMIG() || gpus[0].MemoryTotalMB != 40*1024 {
|
||||
t.Errorf("without the report: %+v", gpus)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnOlderDriversInstanceUUIDIsTakenAsItIs(t *testing.T) {
|
||||
// R450 and R460 name an instance by its card and ids, and CUDA of that
|
||||
// era takes exactly that in CUDA_VISIBLE_DEVICES.
|
||||
l := parseSMIList("GPU 0: A100-SXM4-40GB (UUID: GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77)\n" +
|
||||
" MIG 1g.5gb Device 0: (UUID: MIG-GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77/7/0)\n")
|
||||
if len(l[0]) != 1 || l[0][0].UUID != "MIG-GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77/7/0" || l[0][0].Profile != "1g.5gb" {
|
||||
t.Errorf("parsed %+v", l)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessesOnAMIGCardAreChargedToTheirOwnInstance(t *testing.T) {
|
||||
mem, err := parseGPUMemory([]byte(migReport))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := []GPUMemory{
|
||||
{Index: 1000, UsedMB: 37},
|
||||
{Index: 1001, UsedMB: 21000, Procs: []GPUProcess{{PID: 4100, UsedMB: 20900}}},
|
||||
{Index: 1, UsedMB: 300, Procs: []GPUProcess{{PID: 5200, UsedMB: 250}}},
|
||||
}
|
||||
if !reflect.DeepEqual(mem, want) {
|
||||
t.Errorf("got %+v\nwant %+v", mem, want)
|
||||
}
|
||||
// In the order discovery reports the devices, so a heartbeat's readings
|
||||
// line up with the inventory one for one.
|
||||
gpus := migMachine(t)
|
||||
for i := range gpus {
|
||||
if mem[i].Index != gpus[i].Index {
|
||||
t.Errorf("reading %d is for device %d, inventory has %d", i, mem[i].Index, gpus[i].Index)
|
||||
}
|
||||
}
|
||||
procs, _ := parseGPUProcesses([]byte(migReport))
|
||||
if !reflect.DeepEqual(procs, map[int][]int{1001: {4100}, 1: {5200}}) {
|
||||
t.Errorf("processes by device: %v", procs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMIGInstanceTakesItsCardsPlaceInTheTopologyButNoNVLink(t *testing.T) {
|
||||
gpus := migMachine(t)
|
||||
topo := Topology{
|
||||
NVLink: map[int]map[int]int{0: {1: 12}, 1: {0: 12}},
|
||||
NUMA: map[int]int{0: 0, 1: 1},
|
||||
NICPath: map[int]map[string]string{0: {"mlx5_0": "PIX"}},
|
||||
}
|
||||
topo.Apply(gpus)
|
||||
if len(gpus[0].NVLinkPeers) != 0 || len(gpus[1].NVLinkPeers) != 0 {
|
||||
t.Errorf("MIG instances were given NVLink peers: %v %v", gpus[0].NVLinkPeers, gpus[1].NVLinkPeers)
|
||||
}
|
||||
if gpus[1].NUMANode != 0 || gpus[1].NICPaths["mlx5_0"] != "PIX" {
|
||||
t.Errorf("instance 1 did not take its card's NUMA node and adapters: %+v", gpus[1])
|
||||
}
|
||||
if gpus[2].NUMANode != 1 || !slices.Equal(gpus[2].NVLinkPeers, []int{0}) {
|
||||
t.Errorf("the whole card's topology: %+v", gpus[2])
|
||||
}
|
||||
}
|
||||
|
||||
func TestMIGInstancesAreDiscoveredThroughNvidiaSMI(t *testing.T) {
|
||||
// The real code path against a stub binary that answers each command as
|
||||
// the machine above would.
|
||||
dir := t.TempDir()
|
||||
stub := filepath.Join(dir, "nvidia-smi")
|
||||
script := "#!/bin/sh\ncase \"$1\" in\n" +
|
||||
"-L) cat <<'EOF'\n" + migList + "EOF\n;;\n" +
|
||||
"-q) cat <<'EOF'\n" + migReport + "\nEOF\n;;\n" +
|
||||
"*) cat <<'EOF'\n" + migQuery + "EOF\n;;\nesac\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)
|
||||
}
|
||||
inv := Inventory{GPUs: gpus}
|
||||
if !slices.Equal(inv.Classes(), []string{"a100-80g", "a100-80g-mig-3g.40gb"}) || len(gpus) != 3 {
|
||||
t.Errorf("classes %v from %d devices", inv.Classes(), len(gpus))
|
||||
}
|
||||
}
|
||||
|
||||
// fakeMIGCard is NVML's view of one card, as nvmlDeviceGetMigMode,
|
||||
// nvmlDeviceGetMaxMigDeviceCount and nvmlDeviceGetMigDeviceHandleByIndex
|
||||
// would give it.
|
||||
type fakeMIGCard struct {
|
||||
mode, pending, supported bool
|
||||
max int
|
||||
devs map[int]fakeMIGInstance
|
||||
}
|
||||
|
||||
func (c fakeMIGCard) MIGMode() (bool, bool, bool) { return c.mode, c.pending, c.supported }
|
||||
func (c fakeMIGCard) MaxMIGDevices() int { return c.max }
|
||||
func (c fakeMIGCard) MIGDevice(i int) (nvmlMIGInstance, bool) {
|
||||
d, ok := c.devs[i]
|
||||
return d, ok
|
||||
}
|
||||
|
||||
type fakeMIGInstance struct {
|
||||
uuid, name string
|
||||
total, used int
|
||||
gi, ci int
|
||||
slices, attrMem int
|
||||
}
|
||||
|
||||
func (m fakeMIGInstance) UUID() string { return m.uuid }
|
||||
func (m fakeMIGInstance) Name() string { return m.name }
|
||||
func (m fakeMIGInstance) Memory() (int, int, bool) {
|
||||
return m.total, m.used, m.total > 0
|
||||
}
|
||||
func (m fakeMIGInstance) Instances() (int, int) { return m.gi, m.ci }
|
||||
func (m fakeMIGInstance) Attributes() (int, int, bool) {
|
||||
return m.slices, m.attrMem, m.slices > 0
|
||||
}
|
||||
|
||||
func TestNVMLReportsACardInMIGModeAsItsInstances(t *testing.T) {
|
||||
card := GPU{Index: 2, UUID: "GPU-card", Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559,
|
||||
ComputeCapability: "9.0", NUMANode: 1, Health: GPUHealth{TemperatureC: 40}}
|
||||
// Seven slots, three filled, one of them on a driver that does not put
|
||||
// the profile in the device's name.
|
||||
dev := fakeMIGCard{mode: true, pending: true, supported: true, max: 7, devs: map[int]fakeMIGInstance{
|
||||
0: {uuid: "MIG-aaa", name: "NVIDIA H100 80GB HBM3 MIG 3g.40gb", total: 40448, used: 10, gi: 2, ci: 0},
|
||||
1: {uuid: "MIG-bbb", name: "NVIDIA H100 80GB HBM3 MIG 3g.40gb", total: 40448, used: 10, gi: 1, ci: 0},
|
||||
4: {uuid: "MIG-ccc", name: "NVIDIA H100 80GB HBM3", total: 9984, gi: 9, ci: 0, slices: 1, attrMem: 9984},
|
||||
}}
|
||||
gpus, isMIG := migFromNVML(card, dev)
|
||||
if !isMIG || len(gpus) != 3 {
|
||||
t.Fatalf("got %d devices, MIG %v", len(gpus), isMIG)
|
||||
}
|
||||
var classes []string
|
||||
for _, g := range gpus {
|
||||
classes = append(classes, DeviceLabel(g.Index)+" "+g.UUID+" "+g.Class())
|
||||
}
|
||||
want := []string{"2/mig0 MIG-aaa h100-80g-mig-3g.40gb", "2/mig1 MIG-bbb h100-80g-mig-3g.40gb",
|
||||
"2/mig4 MIG-ccc h100-80g-mig-1g.10gb"}
|
||||
if !slices.Equal(classes, want) {
|
||||
t.Errorf("devices %q, want %q", classes, want)
|
||||
}
|
||||
if g := gpus[1]; g.MemoryTotalMB != 40448 || g.MIG.GPUInstance != 1 || g.NUMANode != 1 ||
|
||||
g.Health.TemperatureC != 40 || len(g.NVLinkPeers) != 0 {
|
||||
t.Errorf("instance 1: %+v", g)
|
||||
}
|
||||
|
||||
// MIG on and nothing carved: the card, kept out of placement.
|
||||
gpus, isMIG = migFromNVML(card, fakeMIGCard{mode: true, supported: true, max: 7})
|
||||
if !isMIG || len(gpus) != 1 || gpus[0].IsMIG() || gpus[0].Health.OK() {
|
||||
t.Errorf("an empty MIG card: %+v", gpus)
|
||||
}
|
||||
// MIG off, or a card without it: the card, as it always was.
|
||||
for _, d := range []fakeMIGCard{{supported: true, max: 7}, {}} {
|
||||
gpus, isMIG = migFromNVML(card, d)
|
||||
if isMIG || len(gpus) != 1 || !reflect.DeepEqual(gpus[0], card) {
|
||||
t.Errorf("a card not in MIG mode came back as %+v", gpus)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMIGInstanceTablesAreRead(t *testing.T) {
|
||||
gis := ParseMIGInstances(migLGI)
|
||||
if len(gis) != 2 || gis[0] != (MIGInstance{GPU: 0, Profile: "3g.40gb", ProfileID: 9, GPUInstance: 1,
|
||||
ComputeInstance: -1, Placement: "4:4"}) {
|
||||
t.Errorf("-lgi: %+v", gis)
|
||||
}
|
||||
cis := ParseMIGInstances(migLCI)
|
||||
if len(cis) != 2 || cis[1] != (MIGInstance{GPU: 0, GPUInstance: 2, Profile: "3g.40gb", ProfileID: 2,
|
||||
ComputeInstance: 0, Placement: "0:3"}) {
|
||||
t.Errorf("-lci: %+v", cis)
|
||||
}
|
||||
if got := ParseMIGInstances("No GPU instances found: Not Found\n"); len(got) != 0 {
|
||||
t.Errorf("nothing to list read as %+v", got)
|
||||
}
|
||||
|
||||
cards, err := buildMIGCards([]byte(migReport), migList, migLGI, migLCI)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cards) != 2 || !cards[0].Enabled || len(cards[0].Instances) != 2 || len(cards[0].GPUInstances) != 2 ||
|
||||
len(cards[0].ComputeInstances) != 2 || cards[1].Enabled {
|
||||
t.Errorf("cards: %+v", cards)
|
||||
}
|
||||
}
|
||||
@@ -166,19 +166,25 @@ func splitTabs(l string) []string {
|
||||
|
||||
// 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.
|
||||
//
|
||||
// The matrix is of whole cards. A MIG instance takes its card's NUMA node and
|
||||
// adapters, and none of its NVLink: instances cannot reach each other or any
|
||||
// other card over it, and peers listed for one would draw a tensor-parallel
|
||||
// group onto it.
|
||||
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] {
|
||||
card := g.CardIndex()
|
||||
if len(g.NVLinkPeers) == 0 && !g.IsMIG() {
|
||||
for peer := range t.NVLink[card] {
|
||||
g.NVLinkPeers = append(g.NVLinkPeers, peer)
|
||||
}
|
||||
sort.Ints(g.NVLinkPeers)
|
||||
}
|
||||
if n, ok := t.NUMA[g.Index]; ok {
|
||||
if n, ok := t.NUMA[card]; ok {
|
||||
g.NUMANode = n
|
||||
}
|
||||
if paths := t.NICPath[g.Index]; len(paths) > 0 {
|
||||
if paths := t.NICPath[card]; len(paths) > 0 {
|
||||
g.NICPaths = paths
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Compute is where an engine runs its model, as it said while loading it.
|
||||
type Compute struct {
|
||||
// Backend is CUDA, Vulkan, Metal, ROCm or CPU.
|
||||
Backend string
|
||||
// Devices are those it named: "CUDA0: Quadro RTX 6000 (7.5)".
|
||||
Devices []string
|
||||
// Offloaded of Total layers went to the devices; 0 and 0 when it did
|
||||
// not say.
|
||||
Offloaded, Total int
|
||||
// NoGPU is the engine saying it found no device it could use.
|
||||
NoGPU bool
|
||||
}
|
||||
|
||||
var (
|
||||
offloadedRE = regexp.MustCompile(`offloaded (\d+)/(\d+) layers to GPU`)
|
||||
usingRE = regexp.MustCompile(`using device ([A-Za-z]+\d+) \(([^)]*)\)`)
|
||||
ccRE = regexp.MustCompile(`Device (\d+): ([^,]+), compute capability ([0-9.]+)`)
|
||||
noDevicesRE = regexp.MustCompile(`found 0 (CUDA|Vulkan|ROCm) devices|no usable GPU found`)
|
||||
)
|
||||
|
||||
// ReadLlamaCompute reads llama.cpp's start-up output for where it put the
|
||||
// model, reporting false when the output says nothing of it.
|
||||
func ReadLlamaCompute(lines []string) (Compute, bool) {
|
||||
var c Compute
|
||||
known := false
|
||||
cc := map[string]string{} // device index → compute capability
|
||||
for _, l := range lines {
|
||||
if m := ccRE.FindStringSubmatch(l); m != nil {
|
||||
cc[m[1]] = m[3]
|
||||
}
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
for _, l := range lines {
|
||||
if m := offloadedRE.FindStringSubmatch(l); m != nil {
|
||||
c.Offloaded, _ = strconv.Atoi(m[1])
|
||||
c.Total, _ = strconv.Atoi(m[2])
|
||||
known = true
|
||||
}
|
||||
if m := usingRE.FindStringSubmatch(l); m != nil && !seen[m[1]] {
|
||||
seen[m[1]] = true
|
||||
d := m[1] + ": " + strings.TrimSpace(m[2])
|
||||
idx := strings.TrimLeft(m[1], "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
|
||||
if v, ok := cc[idx]; ok && strings.HasPrefix(m[1], "CUDA") {
|
||||
d += " (" + v + ")"
|
||||
}
|
||||
c.Devices = append(c.Devices, d)
|
||||
if c.Backend == "" {
|
||||
c.Backend = strings.TrimRight(m[1], "0123456789")
|
||||
}
|
||||
known = true
|
||||
}
|
||||
if noDevicesRE.MatchString(l) {
|
||||
c.NoGPU, known = true, true
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case c.OnTheCPU():
|
||||
c.Backend = "CPU"
|
||||
case c.Backend == "" && c.Offloaded > 0:
|
||||
c.Backend = "GPU"
|
||||
case c.Backend == "":
|
||||
c.Backend = "CPU"
|
||||
}
|
||||
return c, known
|
||||
}
|
||||
|
||||
// OnTheCPU is whether a model the engine was given GPUs for is computed on
|
||||
// the CPU, by its own account: no device found, or no layer offloaded.
|
||||
func (c Compute) OnTheCPU() bool {
|
||||
return c.NoGPU || (c.Total > 0 && c.Offloaded == 0)
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLlamaCppSaysWhereItPutTheModel(t *testing.T) {
|
||||
gpu := strings.Split(`ggml_cuda_init: found 1 CUDA devices:
|
||||
Device 0: Quadro RTX 6000, compute capability 7.5, VMM: yes
|
||||
load_backend: loaded CUDA backend from /opt/llmgw/engines/llamacpp/b10991/libggml-cuda.so
|
||||
llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free
|
||||
load_tensors: offloading 48 repeating layers to GPU
|
||||
load_tensors: offloaded 49/49 layers to GPU`, "\n")
|
||||
c, ok := ReadLlamaCompute(gpu)
|
||||
if !ok || c.Backend != "CUDA" || c.Offloaded != 49 || c.Total != 49 || c.OnTheCPU() ||
|
||||
len(c.Devices) != 1 || c.Devices[0] != "CUDA0: Quadro RTX 6000 (7.5)" {
|
||||
t.Errorf("on the GPU: %+v %v", c, ok)
|
||||
}
|
||||
|
||||
cpu := strings.Split(`warning: no usable GPU found, --gpu-layers option will be ignored
|
||||
warning: one possible reason is that llama.cpp was compiled without GPU support
|
||||
load_tensors: offloaded 0/49 layers to GPU`, "\n")
|
||||
c, ok = ReadLlamaCompute(cpu)
|
||||
if !ok || c.Backend != "CPU" || !c.OnTheCPU() {
|
||||
t.Errorf("fell back to the CPU: %+v %v", c, ok)
|
||||
}
|
||||
|
||||
if _, ok := ReadLlamaCompute([]string{"main: server is listening on http://127.0.0.1:20001"}); ok {
|
||||
t.Error("output that says nothing was read as saying something")
|
||||
}
|
||||
}
|
||||
@@ -44,6 +44,9 @@ type Profile struct {
|
||||
// Adapters are LoRA adapters to load beside the model, in order: an
|
||||
// adapter's index is the id requests apply it by.
|
||||
Adapters []Adapter
|
||||
// KVEvents has the engine publish its KV cache events to the agent,
|
||||
// for the router's precise index. vLLM only.
|
||||
KVEvents bool `json:"kv_events,omitempty"`
|
||||
}
|
||||
|
||||
// ShardSpec is one shard of one replica on this node.
|
||||
@@ -62,6 +65,12 @@ type ShardSpec struct {
|
||||
|
||||
// GPUIndices are the local GPUs this shard uses.
|
||||
GPUIndices []int
|
||||
// GPUUUIDs name the same devices by UUID, in the same order, and are set
|
||||
// by the agent only when one of them is a MIG instance. An instance has
|
||||
// no index CUDA knows: CUDA_VISIBLE_DEVICES=MIG-<uuid> is the only way to
|
||||
// give a process one, and the indices placement uses for it (1000 and
|
||||
// up) mean nothing to CUDA.
|
||||
GPUUUIDs []string `json:"gpu_uuids,omitempty"`
|
||||
// HCAs are the adapters of the chosen fabric, and only that fabric. A shard
|
||||
// must never be given HCAs from two fabrics.
|
||||
HCAs []string
|
||||
@@ -73,13 +82,22 @@ type ShardSpec struct {
|
||||
|
||||
HTTPPort int
|
||||
Profile Profile
|
||||
// BindAddr is the address the engine listens on: the node's management
|
||||
// address, which is what the router reaches it on, set by the agent.
|
||||
// Empty listens on every address.
|
||||
BindAddr string `json:"bind_addr,omitempty"`
|
||||
|
||||
// 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
|
||||
// the node's management address; without a key anybody who can reach it
|
||||
// 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"`
|
||||
// KVEventsPort is the loopback port the agent receives this shard's KV
|
||||
// cache events on, chosen by the agent; zero when it publishes none.
|
||||
// Kept with the spec so an agent that restarts listens where the engine
|
||||
// it adopts is still sending.
|
||||
KVEventsPort int `json:"kv_events_port,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.
|
||||
@@ -96,6 +114,18 @@ type ShardSpec struct {
|
||||
// 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
|
||||
|
||||
// Compute is where the engine said it put the model as it became ready,
|
||||
// kept with the shard's state so an agent restarted later still knows.
|
||||
Compute *Compute `json:",omitempty"`
|
||||
}
|
||||
|
||||
// Host is where the engine listens.
|
||||
func (s ShardSpec) Host() string {
|
||||
if s.BindAddr == "" {
|
||||
return "0.0.0.0"
|
||||
}
|
||||
return s.BindAddr
|
||||
}
|
||||
|
||||
// Leader reports whether this shard serves HTTP for the replica.
|
||||
@@ -156,6 +186,13 @@ type CPUFallback interface {
|
||||
Devices(ctx context.Context) ([]string, error)
|
||||
}
|
||||
|
||||
// HelpCommand is an engine whose program is not asked for its help by
|
||||
// running the command's path with --help, because the path is an interpreter
|
||||
// starting it through a launcher.
|
||||
type HelpCommand interface {
|
||||
HelpCommand(c Command) (path string, args []string)
|
||||
}
|
||||
|
||||
// GPUInfo is the subset of GPU facts validation needs.
|
||||
type GPUInfo struct {
|
||||
Index int
|
||||
@@ -175,6 +212,7 @@ func NewRegistry(cfg Config) *Registry {
|
||||
r.Add(&VLLM{cfg: cfg.VLLM})
|
||||
r.Add(&SGLang{cfg: cfg.SGLang})
|
||||
r.Add(&LlamaCpp{cfg: cfg.LlamaCpp})
|
||||
r.Add(&TRTLLM{cfg: cfg.TRTLLM})
|
||||
r.Add(&Sim{cfg: cfg.Sim})
|
||||
return r
|
||||
}
|
||||
@@ -220,7 +258,10 @@ type Config struct {
|
||||
VLLM VLLMConfig `koanf:"vllm"`
|
||||
SGLang SGLangConfig `koanf:"sglang"`
|
||||
LlamaCpp LlamaCppConfig `koanf:"llamacpp"`
|
||||
Sim SimConfig `koanf:"sim"`
|
||||
// TRTLLM is TensorRT-LLM, which the agent does not install: a node
|
||||
// offers it only when its interpreter is named here.
|
||||
TRTLLM TRTLLMConfig `koanf:"trtllm"`
|
||||
Sim SimConfig `koanf:"sim"`
|
||||
}
|
||||
|
||||
// --- shared helpers ---
|
||||
@@ -245,7 +286,7 @@ func NCCLEnv(s ShardSpec) map[string]string { return baseEnv(s) }
|
||||
// one instead of timing out.
|
||||
func baseEnv(s ShardSpec) map[string]string {
|
||||
env := map[string]string{
|
||||
"CUDA_VISIBLE_DEVICES": joinInts(s.GPUIndices),
|
||||
"CUDA_VISIBLE_DEVICES": visibleDevices(s),
|
||||
// The indices above are nvidia-smi's, which counts by PCI bus. CUDA
|
||||
// counts fastest first unless told otherwise, and on a node with two
|
||||
// kinds of GPU the shard would get a different card than placement
|
||||
@@ -317,6 +358,21 @@ func materializeEnv(generated map[string]string) []string {
|
||||
return out
|
||||
}
|
||||
|
||||
// visibleDevices is CUDA_VISIBLE_DEVICES for a shard: the indices placement
|
||||
// gave it, or their UUIDs when the agent named them, which it does for a MIG
|
||||
// instance.
|
||||
func visibleDevices(s ShardSpec) string {
|
||||
if len(s.GPUUUIDs) > 0 {
|
||||
return strings.Join(s.GPUUUIDs, ",")
|
||||
}
|
||||
return joinInts(s.GPUIndices)
|
||||
}
|
||||
|
||||
// OnMIG reports whether the shard was given a MIG instance.
|
||||
func (s ShardSpec) OnMIG() bool {
|
||||
return slices.ContainsFunc(s.GPUUUIDs, func(u string) bool { return strings.HasPrefix(u, "MIG-") })
|
||||
}
|
||||
|
||||
func joinInts(v []int) string {
|
||||
parts := make([]string, len(v))
|
||||
for i, n := range v {
|
||||
|
||||
@@ -323,8 +323,8 @@ func TestSGLangMultiNodeArgs(t *testing.T) {
|
||||
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 cmd.Args[0] != "-c" || !strings.Contains(cmd.Args[1], `run_module("sglang.launch_server"`) {
|
||||
t.Errorf("not launched as the 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)
|
||||
@@ -334,6 +334,58 @@ func TestSGLangMultiNodeArgs(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSGLangsKeyIsNotOnItsCommandLine(t *testing.T) {
|
||||
s := singleNodeSpec()
|
||||
s.APIKey = "k1"
|
||||
cmd := buildOne(t, &SGLang{}, s)
|
||||
if slices.Contains(cmd.Args, "k1") || slices.Contains(cmd.Args, "--api-key") {
|
||||
t.Errorf("the key is in the arguments: %v", cmd.Args)
|
||||
}
|
||||
if !slices.Contains(cmd.Env, sglangKeyEnv+"=k1") {
|
||||
t.Errorf("the key is not in the environment")
|
||||
}
|
||||
// A worker serves no HTTP, and is not given it.
|
||||
w := multiNodeSpec(1)
|
||||
w.APIKey = "k1"
|
||||
if cmd := buildOne(t, &SGLang{}, w); slices.Contains(cmd.Env, sglangKeyEnv+"=k1") {
|
||||
t.Error("a worker was given the key")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnginesListenOnTheAddressTheyAreGiven(t *testing.T) {
|
||||
s := singleNodeSpec()
|
||||
for _, e := range []Engine{&VLLM{}, &SGLang{}} {
|
||||
s.BindAddr = ""
|
||||
if got := argValue(buildOne(t, e, s).Args, "--host"); got != "0.0.0.0" {
|
||||
t.Errorf("%s: without an address listens on %q", e.Name(), got)
|
||||
}
|
||||
s.BindAddr = "10.0.0.147"
|
||||
if got := argValue(buildOne(t, e, s).Args, "--host"); got != "10.0.0.147" {
|
||||
t.Errorf("%s: listens on %q, expected the management address", e.Name(), got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSGLangIsStartedServingItsMetrics(t *testing.T) {
|
||||
// Without --enable-metrics SGLang has no /metrics, and every scrape of
|
||||
// the replica fails as though it were unreachable.
|
||||
if cmd := buildOne(t, &SGLang{}, singleNodeSpec()); !slices.Contains(cmd.Args, "--enable-metrics") {
|
||||
t.Errorf("not started with its metrics: %v", cmd.Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSGLangIsAskedForItsHelpThroughItsModule(t *testing.T) {
|
||||
g := &SGLang{cfg: SGLangConfig{Python: "/opt/sglang/bin/python"}}
|
||||
g.UseInstalled("/opt/llmgw/engines/sglang/0.5.21-cu13/bin/python")
|
||||
path, args := g.HelpCommand(Command{})
|
||||
if path != "/opt/llmgw/engines/sglang/0.5.21-cu13/bin/python" {
|
||||
t.Errorf("asked %s, not the interpreter it was installed into", path)
|
||||
}
|
||||
if !slices.Equal(args, []string{"-m", "sglang.launch_server", "--help"}) {
|
||||
t.Errorf("asked %v", args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSGLangSingleNodeOmitsMultiNodeFlags(t *testing.T) {
|
||||
cmd := buildOne(t, &SGLang{}, singleNodeSpec())
|
||||
for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} {
|
||||
@@ -425,14 +477,14 @@ func TestValidateRequiresModelPath(t *testing.T) {
|
||||
|
||||
func TestRegistry(t *testing.T) {
|
||||
r := NewRegistry(Config{})
|
||||
if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "sim", "vllm"}) {
|
||||
if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "sim", "trtllm", "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") {
|
||||
if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, trtllm, vllm") {
|
||||
t.Errorf("an unknown engine should list what is available: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAVLLMShardWithKVEventsPublishesToTheAgentOnLoopback(t *testing.T) {
|
||||
s := singleNodeSpec()
|
||||
s.Profile.KVEvents = true
|
||||
s.KVEventsPort = 5999
|
||||
s.Profile.Args = []string{"--max-model-len", "8192"}
|
||||
cmd := buildOne(t, &VLLM{}, s)
|
||||
|
||||
i := slices.Index(cmd.Args, "--kv-events-config")
|
||||
if i < 0 || !strings.Contains(cmd.Args[i+1], `"endpoint":"tcp://127.0.0.1:5999"`) {
|
||||
t.Fatalf("args %v should name the agent's loopback port", cmd.Args)
|
||||
}
|
||||
if j := slices.Index(cmd.Args, "--prefix-caching-hash-algo"); j < 0 || cmd.Args[j+1] != "sha256_cbor" {
|
||||
t.Fatalf("args %v should pin the hashing the router computes", cmd.Args)
|
||||
}
|
||||
if slices.Index(cmd.Args, "--max-model-len") < i {
|
||||
t.Fatal("profile arguments come last, so an operator can still override")
|
||||
}
|
||||
if !slices.Contains(cmd.Env, "PYTHONHASHSEED=1729") {
|
||||
t.Fatalf("env %v should fix the hash seed", cmd.Env)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoKVEventsWithoutAPortOrOffTheLeader(t *testing.T) {
|
||||
s := singleNodeSpec()
|
||||
s.Profile.KVEvents = true // but the agent could not listen
|
||||
if cmd := buildOne(t, &VLLM{}, s); slices.Contains(cmd.Args, "--kv-events-config") {
|
||||
t.Fatal("an engine told to publish where nobody listens only fills a queue")
|
||||
}
|
||||
w := multiNodeSpec(1)
|
||||
w.Profile.KVEvents, w.KVEventsPort = true, 5999
|
||||
if cmd := buildEngine(t, &VLLM{}, w); slices.Contains(cmd.Args, "--kv-events-config") {
|
||||
t.Fatal("a headless worker has no scheduler to publish from")
|
||||
}
|
||||
off := singleNodeSpec()
|
||||
off.KVEventsPort = 5999
|
||||
if cmd := buildOne(t, &VLLM{}, off); slices.Contains(cmd.Args, "--kv-events-config") {
|
||||
t.Fatal("kv_events is opt-in")
|
||||
}
|
||||
}
|
||||
@@ -93,13 +93,13 @@ func (l *LlamaCpp) Commands(s ShardSpec) ([]Command, error) {
|
||||
bin := l.bin()
|
||||
p := s.Profile
|
||||
|
||||
model, err := ggufIn(p.ModelPath)
|
||||
model, projector, err := ggufIn(p.ModelPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
args := []string{
|
||||
"--model", model,
|
||||
"--host", "0.0.0.0",
|
||||
"--host", s.Host(),
|
||||
"--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.
|
||||
@@ -120,6 +120,11 @@ func (l *LlamaCpp) Commands(s ShardSpec) ([]Command, error) {
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--alias", p.ServedModelName)
|
||||
}
|
||||
// A vision model's projector, fetched beside it: what turns an image in
|
||||
// a request into something the language model reads.
|
||||
if projector != "" && !hasFlag(p.Args, "--mmproj", "-mm") {
|
||||
args = append(args, "--mmproj", projector)
|
||||
}
|
||||
if !hasFlag(p.Args, "--parallel", "-np") {
|
||||
n := l.cfg.Parallel
|
||||
if n <= 0 {
|
||||
@@ -173,7 +178,11 @@ func (l *LlamaCpp) Commands(s ShardSpec) ([]Command, error) {
|
||||
// as the driver lists them, which on NVIDIA follows the same bus order
|
||||
// as nvidia-smi; were it not to, the shard would not hold memory on its
|
||||
// own GPUs, and starting it checks exactly that.
|
||||
if len(s.GPUIndices) > 0 {
|
||||
//
|
||||
// Not on a MIG instance, which Vulkan does not see at all: only CUDA
|
||||
// opens one, by the UUID CUDA_VISIBLE_DEVICES carries, and an index here
|
||||
// would name some other card.
|
||||
if len(s.GPUIndices) > 0 && !s.OnMIG() {
|
||||
env["GGML_VK_VISIBLE_DEVICES"] = joinInts(s.GPUIndices)
|
||||
}
|
||||
return []Command{{
|
||||
@@ -203,30 +212,46 @@ func hasFlag(args []string, names ...string) bool {
|
||||
//
|
||||
// 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) {
|
||||
//
|
||||
// A vision projector in the directory — an mmproj file — is not the model:
|
||||
// it is returned second, to be loaded beside it.
|
||||
func ggufIn(path string) (model, projector string, err 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
|
||||
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
|
||||
return "", "", err
|
||||
}
|
||||
var names []string
|
||||
var names, projectors []string
|
||||
for _, e := range entries {
|
||||
if !e.IsDir() && strings.HasSuffix(strings.ToLower(e.Name()), ".gguf") {
|
||||
lower := strings.ToLower(e.Name())
|
||||
if e.IsDir() || !strings.HasSuffix(lower, ".gguf") {
|
||||
continue
|
||||
}
|
||||
if strings.Contains(lower, "mmproj") {
|
||||
projectors = append(projectors, e.Name())
|
||||
} else {
|
||||
names = append(names, e.Name())
|
||||
}
|
||||
}
|
||||
if len(names) == 0 {
|
||||
return "", fmt.Errorf("%s holds no .gguf file", path)
|
||||
if len(projectors) > 0 {
|
||||
return "", "", fmt.Errorf("%s holds a vision projector and no model", path)
|
||||
}
|
||||
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
|
||||
if len(projectors) > 0 {
|
||||
sort.Strings(projectors)
|
||||
projector = filepath.Join(path, projectors[0])
|
||||
}
|
||||
return filepath.Join(path, names[0]), projector, nil
|
||||
}
|
||||
|
||||
// UseInstalled points the engine at a build the agent installed and checked.
|
||||
@@ -252,6 +277,9 @@ func (l *LlamaCpp) bin() string {
|
||||
return "llama-server"
|
||||
}
|
||||
|
||||
// Binary is the llama-server this engine runs.
|
||||
func (l *LlamaCpp) Binary() string { return l.bin() }
|
||||
|
||||
// FallsBackToCPU says that llama.cpp keeps running when its GPU backend does
|
||||
// not load: --n-gpu-layers asks for layers on a GPU it then has none of, and
|
||||
// the server comes up and answers on the CPU, at a tenth of the speed, with no
|
||||
|
||||
@@ -204,7 +204,7 @@ func TestASetDirectoryResolvesToItsFirstPart(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
got, err := ggufIn(dir)
|
||||
got, _, err := ggufIn(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -221,7 +221,7 @@ func TestAPlainFilePathIsUntouched(t *testing.T) {
|
||||
if err := os.WriteFile(p, []byte("x"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, _ := ggufIn(p); got != p {
|
||||
if got, _, _ := ggufIn(p); got != p {
|
||||
t.Errorf("expected the path unchanged, got %s", got)
|
||||
}
|
||||
}
|
||||
@@ -231,7 +231,7 @@ func TestADirectoryWithNoWeightsSaysSo(t *testing.T) {
|
||||
if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte("{}"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := ggufIn(dir); err == nil {
|
||||
if _, _, err := ggufIn(dir); err == nil {
|
||||
t.Error("a directory with no .gguf in it is not a model llama.cpp can load")
|
||||
}
|
||||
}
|
||||
@@ -240,7 +240,7 @@ 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)
|
||||
got, _, err := ggufIn(p)
|
||||
if err != nil || got != p {
|
||||
t.Errorf("expected it passed through, got %q %v", got, err)
|
||||
}
|
||||
@@ -361,3 +361,28 @@ func TestLlamaCppSpreadsAReplicaOverItsNodesGPUsByLayer(t *testing.T) {
|
||||
t.Error("a one-GPU replica was given a split mode")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAVisionProjectorIsLoadedBesideTheModelNotAsIt(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
// "mmproj" sorts before the model's name; it must still not be taken for it.
|
||||
for _, n := range []string{"gemma-3-4b-it-Q4_K_M.gguf", "mmproj-gemma-3-4b-it-f16.gguf"} {
|
||||
if err := os.WriteFile(filepath.Join(dir, n), []byte("x"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
cmds, err := (&LlamaCpp{}).Commands(ShardSpec{HTTPPort: 8080, Profile: Profile{ModelPath: dir}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
args := strings.Join(cmds[0].Args, " ")
|
||||
if !strings.Contains(args, "--model "+filepath.Join(dir, "gemma-3-4b-it-Q4_K_M.gguf")) ||
|
||||
!strings.Contains(args, "--mmproj "+filepath.Join(dir, "mmproj-gemma-3-4b-it-f16.gguf")) {
|
||||
t.Errorf("%s", args)
|
||||
}
|
||||
// One the profile names itself is left as it is.
|
||||
cmds, _ = (&LlamaCpp{}).Commands(ShardSpec{HTTPPort: 8080, Profile: Profile{ModelPath: dir,
|
||||
Args: []string{"--mmproj", "/elsewhere.gguf"}}})
|
||||
if strings.Count(strings.Join(cmds[0].Args, " "), "--mmproj") != 1 {
|
||||
t.Errorf("two projectors: %v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"maps"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// RuntimeLoRAEnv is the variable that lets vLLM load and unload adapters
|
||||
// while it runs, through POST /v1/load_lora_adapter and
|
||||
// /v1/unload_lora_adapter. Without it those routes do not exist.
|
||||
//
|
||||
// It doubles as the mark that a profile's adapters load on demand: the
|
||||
// catalog sets it for a model with max_loaded_adapters and refuses it set by
|
||||
// hand anywhere else, so it cannot disagree with what the router does.
|
||||
const RuntimeLoRAEnv = "VLLM_ALLOW_RUNTIME_LORA_UPDATING"
|
||||
|
||||
// RuntimeLoRADir is where an agent links each adapter it fetched for a
|
||||
// replica whose adapters load on demand, under the adapter's name, relative
|
||||
// to the directory engines run in.
|
||||
//
|
||||
// Relative on purpose. The router asks the engine to load an adapter by a
|
||||
// path, and does not know where any node keeps its cache — each node may
|
||||
// keep it somewhere else, and a replica spanning nodes reads the adapter on
|
||||
// every one of them. vLLM resolves a relative path that exists against its
|
||||
// working directory, which on every node is the one the agent starts
|
||||
// engines in, and the agent links the adapter there before the engine
|
||||
// starts. So one path, the same everywhere, names the right file on each.
|
||||
const RuntimeLoRADir = "llmgw-lora"
|
||||
|
||||
// WithRuntimeLoRA has a vLLM profile with adapters start holding none of
|
||||
// them, ready to load up to limit at a time on demand: the engine is allowed
|
||||
// runtime loading, and --max-loras — how many adapters one batch may mix —
|
||||
// is the cap, unless the operator's arguments already say. Anything else is
|
||||
// returned as it is.
|
||||
func (p Profile) WithRuntimeLoRA(limit int) Profile {
|
||||
if limit <= 0 || p.Engine != "vllm" || len(p.Adapters) == 0 {
|
||||
return p
|
||||
}
|
||||
env := maps.Clone(p.Env)
|
||||
if env == nil {
|
||||
env = map[string]string{}
|
||||
}
|
||||
env[RuntimeLoRAEnv] = "True"
|
||||
p.Env = env
|
||||
if !hasFlag(p.Args, "--max-loras") {
|
||||
p.Args = append([]string{"--max-loras", strconv.Itoa(limit)}, p.Args...)
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// RuntimeLoRA reports whether a profile's adapters load on demand rather
|
||||
// than at start.
|
||||
func (p Profile) RuntimeLoRA() bool {
|
||||
v, ok := p.Env[RuntimeLoRAEnv]
|
||||
return ok && v != "" && v != "0" && !strings.EqualFold(v, "false")
|
||||
}
|
||||
|
||||
// RuntimeLoRAPath is the lora_path the router gives vLLM to load an adapter
|
||||
// from: the link the agent made, or, for a Hugging Face repository id, the
|
||||
// id, which vLLM downloads itself and the agent leaves alone.
|
||||
func RuntimeLoRAPath(name, source string) string {
|
||||
if IsHubID(source) {
|
||||
return source
|
||||
}
|
||||
return RuntimeLoRADir + "/" + name
|
||||
}
|
||||
|
||||
// IsHubID reports whether a source is a bare Hugging Face repository id
|
||||
// ("org/repo") rather than a URL or a path: what the agent passes to an
|
||||
// engine untouched.
|
||||
func IsHubID(source string) bool {
|
||||
switch {
|
||||
case source == "",
|
||||
strings.HasPrefix(source, "/"), strings.HasPrefix(source, "./"), strings.HasPrefix(source, "../"),
|
||||
strings.HasPrefix(source, "~"),
|
||||
strings.HasPrefix(source, "http://"), strings.HasPrefix(source, "https://"),
|
||||
strings.HasPrefix(source, "hf:"):
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// LinkRuntimeAdapters makes RuntimeLoRADir under dir hold a link to each
|
||||
// fetched adapter, named as the adapter is. A link already there is
|
||||
// replaced, so a replica started after an adapter's source changed loads
|
||||
// the new one. Adapters still named by a repository id have nothing local
|
||||
// to link.
|
||||
func LinkRuntimeAdapters(dir string, ads []Adapter) error {
|
||||
if len(ads) == 0 {
|
||||
return nil
|
||||
}
|
||||
root := filepath.Join(dir, RuntimeLoRADir)
|
||||
if err := os.MkdirAll(root, 0o755); err != nil {
|
||||
return fmt.Errorf("making the adapters' directory: %w", err)
|
||||
}
|
||||
for _, a := range ads {
|
||||
if IsHubID(a.Path) {
|
||||
continue
|
||||
}
|
||||
if a.Name == "" || strings.ContainsAny(a.Name, `/\`) || a.Name == "." || a.Name == ".." {
|
||||
return fmt.Errorf("adapter %q: not a name a file can have", a.Name)
|
||||
}
|
||||
target, err := filepath.Abs(a.Path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("adapter %s: %w", a.Name, err)
|
||||
}
|
||||
link := filepath.Join(root, a.Name)
|
||||
// Made beside the old one and renamed over it, so an engine loading
|
||||
// the adapter at this moment sees one or the other and never neither.
|
||||
tmp := link + ".new"
|
||||
_ = os.Remove(tmp)
|
||||
if err := os.Symlink(target, tmp); err != nil {
|
||||
return fmt.Errorf("adapter %s: %w", a.Name, err)
|
||||
}
|
||||
if err := os.Rename(tmp, link); err != nil {
|
||||
_ = os.Remove(tmp)
|
||||
return fmt.Errorf("adapter %s: %w", a.Name, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// loraModules is vLLM's --lora-modules for adapters loaded at start: one
|
||||
// flag followed by every name=path, since vLLM keeps only the last of a
|
||||
// repeated flag.
|
||||
func loraModules(ads []Adapter) []string {
|
||||
out := []string{"--lora-modules"}
|
||||
for _, a := range ads {
|
||||
out = append(out, a.Name+"="+a.Path)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func adapterSpec() ShardSpec {
|
||||
s := singleNodeSpec()
|
||||
s.Profile.Engine = "vllm"
|
||||
s.Profile.Adapters = []Adapter{
|
||||
{Name: "sql", Path: "/cache/ab-sql"},
|
||||
{Name: "docs", Path: "/cache/cd-docs"},
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func TestVLLMLoadsEveryAdapterAtStartUnderOneFlag(t *testing.T) {
|
||||
cmd := buildOne(t, &VLLM{}, adapterSpec())
|
||||
if !slices.Contains(cmd.Args, "--enable-lora") {
|
||||
t.Errorf("LoRA not enabled: %v", cmd.Args)
|
||||
}
|
||||
// vLLM keeps only the last of a repeated --lora-modules, so both
|
||||
// adapters must follow one flag.
|
||||
i := slices.Index(cmd.Args, "--lora-modules")
|
||||
if i < 0 || !slices.Equal(cmd.Args[i+1:i+3], []string{"sql=/cache/ab-sql", "docs=/cache/cd-docs"}) {
|
||||
t.Errorf("adapters not registered at start: %v", cmd.Args)
|
||||
}
|
||||
if n := countFlag(cmd.Args, "--lora-modules"); n != 1 {
|
||||
t.Errorf("--lora-modules given %d times", n)
|
||||
}
|
||||
if _, set := envMap(cmd.Env)[RuntimeLoRAEnv]; set {
|
||||
t.Error("runtime loading was allowed for adapters loaded at start")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMWithAdaptersOnDemandStartsHoldingNone(t *testing.T) {
|
||||
s := adapterSpec()
|
||||
s.Profile = s.Profile.WithRuntimeLoRA(3)
|
||||
cmd := buildOne(t, &VLLM{}, s)
|
||||
if slices.Contains(cmd.Args, "--lora-modules") {
|
||||
t.Errorf("adapters registered at start: %v", cmd.Args)
|
||||
}
|
||||
if !slices.Contains(cmd.Args, "--enable-lora") || argValue(cmd.Args, "--max-loras") != "3" {
|
||||
t.Errorf("want LoRA enabled with the cap as --max-loras: %v", cmd.Args)
|
||||
}
|
||||
if envMap(cmd.Env)[RuntimeLoRAEnv] != "True" {
|
||||
t.Errorf("the load and unload routes were not turned on: %v", cmd.Env)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheOperatorsMaxLorasWinsOverTheCap(t *testing.T) {
|
||||
p := adapterSpec().Profile
|
||||
p.Args = []string{"--max-loras", "8"}
|
||||
got := p.WithRuntimeLoRA(3)
|
||||
if argValue(got.Args, "--max-loras") != "8" || countFlag(got.Args, "--max-loras") != 1 {
|
||||
t.Errorf("args %v", got.Args)
|
||||
}
|
||||
if p.Env != nil {
|
||||
t.Error("the profile's own env map was written to")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeLoRALeavesOtherProfilesAlone(t *testing.T) {
|
||||
llama := adapterSpec().Profile
|
||||
llama.Engine = "llamacpp"
|
||||
none := adapterSpec().Profile
|
||||
none.Adapters = nil
|
||||
for _, p := range []Profile{llama, none} {
|
||||
if got := p.WithRuntimeLoRA(2); got.RuntimeLoRA() || len(got.Args) != 0 {
|
||||
t.Errorf("%s with %d adapters: %+v", p.Engine, len(p.Adapters), got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheRouterLoadsAnAdapterByItsLinkOrItsRepository(t *testing.T) {
|
||||
for source, want := range map[string]string{
|
||||
"http://mirror:8081/sql-lora/v1/": "llmgw-lora/sql",
|
||||
"/srv/adapters/sql": "llmgw-lora/sql",
|
||||
"hf:org/repo/adapter.gguf": "llmgw-lora/sql",
|
||||
"someone/sql-lora": "someone/sql-lora",
|
||||
"https://huggingface.co/someone/sql": "llmgw-lora/sql",
|
||||
} {
|
||||
if got := RuntimeLoRAPath("sql", source); got != want {
|
||||
t.Errorf("%s: %s, want %s", source, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdaptersAreLinkedUnderTheirNamesAndRelinkedWhenTheyChange(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
old := filepath.Join(dir, "old-sql")
|
||||
cur := filepath.Join(dir, "new-sql")
|
||||
for _, d := range []string{old, cur} {
|
||||
if err := os.Mkdir(d, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ads := []Adapter{{Name: "sql", Path: old}, {Name: "hub", Path: "someone/hub-lora"}}
|
||||
if err := LinkRuntimeAdapters(dir, ads); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ads[0].Path = cur
|
||||
if err := LinkRuntimeAdapters(dir, ads); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := os.Readlink(filepath.Join(dir, RuntimeLoRADir, "sql"))
|
||||
if err != nil || got != cur {
|
||||
t.Errorf("link to %q (%v), want %q", got, err, cur)
|
||||
}
|
||||
if _, err := os.Lstat(filepath.Join(dir, RuntimeLoRADir, "hub")); !os.IsNotExist(err) {
|
||||
t.Errorf("a repository id was linked: %v", err)
|
||||
}
|
||||
if err := LinkRuntimeAdapters(dir, []Adapter{{Name: "../escape", Path: cur}}); err == nil {
|
||||
t.Error("an adapter named like a path was linked")
|
||||
}
|
||||
}
|
||||
|
||||
func countFlag(args []string, flag string) int {
|
||||
n := 0
|
||||
for _, a := range args {
|
||||
if a == flag {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package engines
|
||||
|
||||
import "testing"
|
||||
|
||||
const migUUID = "MIG-c6d4f1ef-42e4-5de3-91c7-45d71c87eb3f"
|
||||
|
||||
func TestAShardOnAMIGInstanceIsGivenItByUUID(t *testing.T) {
|
||||
// Placement numbers card 0's instance 1 as 1001, which CUDA has never
|
||||
// heard of: CUDA_VISIBLE_DEVICES=1001 is no device at all, and =1 would
|
||||
// be the wrong card.
|
||||
s := singleNodeSpec()
|
||||
s.GPUIndices, s.GPUUUIDs = []int{1001}, []string{migUUID}
|
||||
s.Profile.TP = 1
|
||||
|
||||
env := envMap(buildOne(t, &VLLM{}, s).Env)
|
||||
if env["CUDA_VISIBLE_DEVICES"] != migUUID {
|
||||
t.Errorf("vLLM: CUDA_VISIBLE_DEVICES=%q, want the instance's UUID", env["CUDA_VISIBLE_DEVICES"])
|
||||
}
|
||||
|
||||
l := llamaSpec([]int{1001})
|
||||
l.GPUUUIDs = s.GPUUUIDs
|
||||
cmds, err := (&LlamaCpp{}).Commands(l)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
env = envMap(cmds[0].Env)
|
||||
if env["CUDA_VISIBLE_DEVICES"] != migUUID {
|
||||
t.Errorf("llama.cpp: CUDA_VISIBLE_DEVICES=%q, want the instance's UUID", env["CUDA_VISIBLE_DEVICES"])
|
||||
}
|
||||
// Vulkan does not see MIG instances, and an index for it would name
|
||||
// some other card.
|
||||
if v, set := env["GGML_VK_VISIBLE_DEVICES"]; set {
|
||||
t.Errorf("GGML_VK_VISIBLE_DEVICES=%q on a MIG instance", v)
|
||||
}
|
||||
if v, _ := flagValue(cmds[0].Args, "--n-gpu-layers"); v != "999" {
|
||||
t.Errorf("the model was not offloaded to the instance: --n-gpu-layers %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAShardOnWholeCardsIsStillGivenThemByIndex(t *testing.T) {
|
||||
env := envMap(buildOne(t, &VLLM{}, singleNodeSpec()).Env)
|
||||
if env["CUDA_VISIBLE_DEVICES"] != "0,1,2,3" {
|
||||
t.Errorf("CUDA_VISIBLE_DEVICES=%q", env["CUDA_VISIBLE_DEVICES"])
|
||||
}
|
||||
cmds, _ := (&LlamaCpp{}).Commands(llamaSpec([]int{2}))
|
||||
if env := envMap(cmds[0].Env); env["GGML_VK_VISIBLE_DEVICES"] != "2" {
|
||||
t.Errorf("GGML_VK_VISIBLE_DEVICES=%q", env["GGML_VK_VISIBLE_DEVICES"])
|
||||
}
|
||||
}
|
||||
@@ -40,7 +40,15 @@ func (*SGLang) Name() string { return "sglang" }
|
||||
// HealthPath is SGLang's readiness endpoint.
|
||||
func (*SGLang) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is SGLang's Prometheus endpoint.
|
||||
// HelpCommand is how to ask SGLang for its help: through its module, in the
|
||||
// environment it was installed into, since the command itself is llmgw's
|
||||
// launcher and the interpreter's own --help describes Python.
|
||||
func (g *SGLang) HelpCommand(Command) (string, []string) {
|
||||
return g.python(), []string{"-m", "sglang.launch_server", "--help"}
|
||||
}
|
||||
|
||||
// MetricsPath is SGLang's Prometheus endpoint, served only with
|
||||
// --enable-metrics.
|
||||
func (*SGLang) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
@@ -79,11 +87,17 @@ func (g *SGLang) serverCommand(s ShardSpec) (Command, error) {
|
||||
python := g.python()
|
||||
p := s.Profile
|
||||
|
||||
args := []string{"-m", "sglang.launch_server",
|
||||
args := []string{"-c", sglangLauncher,
|
||||
"--model-path", p.ModelPath,
|
||||
"--host", "0.0.0.0",
|
||||
"--host", s.Host(),
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
"--tp", strconv.Itoa(p.TP),
|
||||
// The flag's whole name: SGLang's documentation writes --tp, which
|
||||
// argparse accepts only as an abbreviation of --tp-size and only
|
||||
// while no other flag begins with it.
|
||||
"--tp-size", strconv.Itoa(p.TP),
|
||||
// Without it SGLang serves no /metrics at all, and the router has
|
||||
// nothing to schedule this replica by.
|
||||
"--enable-metrics",
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--served-model-name", p.ServedModelName)
|
||||
@@ -96,11 +110,23 @@ func (g *SGLang) serverCommand(s ShardSpec) (Command, error) {
|
||||
)
|
||||
}
|
||||
args = append(args, p.Args...)
|
||||
// SGLang reads its key only from the command line, so on a node with
|
||||
// other users it is visible to them in the process list.
|
||||
env := baseEnv(s)
|
||||
if s.APIKey != "" && s.Leader() {
|
||||
args = append(args, "--api-key", s.APIKey)
|
||||
env[sglangKeyEnv] = s.APIKey
|
||||
}
|
||||
|
||||
return Command{Name: "sglang", Path: python, Args: args, Env: materializeEnv(baseEnv(s))}, nil
|
||||
return Command{Name: "sglang", Path: python, Args: args, Env: materializeEnv(env)}, nil
|
||||
}
|
||||
|
||||
// SGLang reads its key only from the command line, where every user on the
|
||||
// node can read it in the process list. So the key goes in the environment,
|
||||
// which only the engine's own user and root can read, and this launcher —
|
||||
// sglang.launch_server under another name — adds it to the arguments inside
|
||||
// the process: sys.argv is Python's copy, not what /proc shows.
|
||||
const (
|
||||
sglangKeyEnv = "LLMGW_ENGINE_KEY"
|
||||
sglangLauncher = `import os, runpy, sys
|
||||
key = os.environ.pop("` + sglangKeyEnv + `", "")
|
||||
sys.argv = ["sglang.launch_server"] + sys.argv[1:] + (["--api-key", key] if key else [])
|
||||
runpy.run_module("sglang.launch_server", run_name="__main__", alter_sys=True)
|
||||
`
|
||||
)
|
||||
@@ -2,6 +2,7 @@ package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
)
|
||||
@@ -63,7 +64,7 @@ func (s *Sim) Commands(spec ShardSpec) ([]Command, error) {
|
||||
p := spec.Profile
|
||||
|
||||
args := []string{"sim-engine",
|
||||
"--listen", "0.0.0.0:" + strconv.Itoa(spec.HTTPPort),
|
||||
"--listen", net.JoinHostPort(spec.Host(), strconv.Itoa(spec.HTTPPort)),
|
||||
"--model", p.ServedModelName,
|
||||
}
|
||||
if spec.NNodes > 1 {
|
||||
|
||||
@@ -0,0 +1,218 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// TRTLLMConfig locates a TensorRT-LLM the operator installed.
|
||||
//
|
||||
// llmgw does not install TensorRT-LLM. It is a pip package pinned to one
|
||||
// CUDA, one PyTorch and an MPI the machine must provide, and NVIDIA's own
|
||||
// advice is its container; an installer that guessed at all of that would
|
||||
// fail on most machines it met, a long way into a large download. So it is
|
||||
// brought by the operator — a virtual environment, or NVIDIA's container —
|
||||
// and named here.
|
||||
type TRTLLMConfig struct {
|
||||
// Python is the interpreter of the environment TensorRT-LLM is installed
|
||||
// in: <venv>/bin/python for a virtual environment, python3 in NVIDIA's
|
||||
// container. Naming it is what makes this node offer the engine; empty
|
||||
// offers nothing.
|
||||
//
|
||||
// The interpreter rather than trtllm-serve, because llmgw starts the
|
||||
// server through a few lines of its own that make it demand the
|
||||
// gateway's key, which trtllm-serve has no option for.
|
||||
Python string `koanf:"python"`
|
||||
}
|
||||
|
||||
// TRTLLM builds TensorRT-LLM launch commands: its OpenAI-compatible server,
|
||||
// trtllm-serve, on one node.
|
||||
//
|
||||
// Written against TensorRT-LLM's documented command line and server, and not
|
||||
// yet run on hardware by anyone working on llmgw: no datacenter GPU was at
|
||||
// hand. `llmgw ctl engine-check trtllm` asks an installed one whether it
|
||||
// takes the flags built here.
|
||||
type TRTLLM struct{ cfg TRTLLMConfig }
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*TRTLLM) Name() string { return "trtllm" }
|
||||
|
||||
// HealthPath is trtllm-serve's liveness endpoint: 200 once the executor is
|
||||
// up, 503 while it is not. The supervisor's /v1/models step is what says the
|
||||
// model is being served.
|
||||
func (*TRTLLM) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is trtllm-serve's Prometheus endpoint. Its /metrics is a
|
||||
// different thing — a JSON list of iteration statistics, drained by each
|
||||
// read — and a scrape there would both fail to parse and steal the numbers
|
||||
// from anyone else reading them.
|
||||
func (*TRTLLM) MetricsPath() string { return "/prometheus/metrics" }
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API; a TensorRT-LLM
|
||||
// replica is one node, and so one rank.
|
||||
func (*TRTLLM) ServesHTTP(s ShardSpec) bool { return s.Leader() }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
//
|
||||
// One node: trtllm-serve spans machines only when started under MPI or
|
||||
// Slurm from one of them (trtllm-llmapi-launch), and the agents each start
|
||||
// their own node's part, which is not that. Ampere or newer: the oldest
|
||||
// generation current TensorRT-LLM builds kernels for.
|
||||
func (*TRTLLM) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if p.Nodes > 1 {
|
||||
return fmt.Errorf("profile %q: trtllm runs a replica on one node; "+
|
||||
"use vllm or sglang for a %d-node replica", p.Name, p.Nodes)
|
||||
}
|
||||
if err := validateCommon(p, gpus); err != nil {
|
||||
return err
|
||||
}
|
||||
return requireComputeCapability(p, gpus, 8.0, "TensorRT-LLM")
|
||||
}
|
||||
|
||||
// Commands returns the processes for one shard, which is one process. Tensor
|
||||
// and pipeline parallel ranks on the node are trtllm-serve's own MPI workers,
|
||||
// started by it.
|
||||
func (t *TRTLLM) Commands(s ShardSpec) ([]Command, error) {
|
||||
if s.NNodes > 1 {
|
||||
return nil, fmt.Errorf("trtllm: a replica spanning %d nodes is not something llmgw starts; "+
|
||||
"TensorRT-LLM runs one node per replica here", s.NNodes)
|
||||
}
|
||||
p := s.Profile
|
||||
args := []string{"-c", trtllmLauncher,
|
||||
// The subcommand by name: trtllm-serve takes a bare model as "serve",
|
||||
// but saying it costs nothing and does not depend on that default.
|
||||
"serve", p.ModelPath,
|
||||
"--host", s.Host(),
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
"--tp_size", strconv.Itoa(p.TP),
|
||||
}
|
||||
if p.PP > 1 {
|
||||
args = append(args, "--pp_size", strconv.Itoa(p.PP))
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
// Without it the model is named after its directory, and the agent's
|
||||
// readiness check, which looks for this name in /v1/models, never
|
||||
// finds it.
|
||||
args = append(args, "--served_model_name", p.ServedModelName)
|
||||
}
|
||||
// What makes the queue and the cache measurable: the options file the
|
||||
// launcher writes turns on the Prometheus endpoint and the iteration
|
||||
// statistics its gauges are set from. A placeholder here, so the flag is
|
||||
// one engine-check sees and checks; the launcher replaces it.
|
||||
args = append(args, "--extra_llm_api_options", trtllmOptionsPlaceholder)
|
||||
args = append(args, p.Args...)
|
||||
|
||||
env := baseEnv(s)
|
||||
if s.APIKey != "" {
|
||||
env[trtllmKeyEnv] = s.APIKey
|
||||
}
|
||||
// TensorRT-LLM reports usage to NVIDIA unless told not to. A gateway on
|
||||
// somebody's own network does not phone home on their behalf, and a
|
||||
// model's profile cannot make it: these are set after the profile's env.
|
||||
env["TRTLLM_NO_USAGE_STATS"] = "1"
|
||||
env["DO_NOT_TRACK"] = "1"
|
||||
return []Command{{Name: "trtllm-serve", Path: t.python(), Args: args, Env: materializeEnv(env)}}, nil
|
||||
}
|
||||
|
||||
// python is the interpreter to start trtllm-serve with.
|
||||
func (t *TRTLLM) python() string {
|
||||
if t.cfg.Python != "" {
|
||||
return t.cfg.Python
|
||||
}
|
||||
return "python3"
|
||||
}
|
||||
|
||||
// Configured reports whether the configuration names TensorRT-LLM's
|
||||
// interpreter, which is the only way a node offers it.
|
||||
func (t *TRTLLM) Configured() bool { return t.cfg.Python != "" }
|
||||
|
||||
// Probe reports the installed TensorRT-LLM's version.
|
||||
//
|
||||
// From the package's metadata rather than by importing it: the import loads
|
||||
// CUDA, MPI and PyTorch and takes longer than a probe may, while what
|
||||
// placement needs is whether the package is there. Whether it runs is
|
||||
// engine-check's question, and the replica's.
|
||||
func (t *TRTLLM) Probe(ctx context.Context) (string, error) {
|
||||
return versionOfModule(ctx, t.python(), "-c",
|
||||
"import importlib.metadata as m, importlib.util as u, sys\n"+
|
||||
"if u.find_spec('tensorrt_llm') is None: sys.exit(1)\n"+
|
||||
"sys.stdout.write(m.version('tensorrt_llm'))")
|
||||
}
|
||||
|
||||
// HelpCommand is how to ask trtllm-serve for its help: through its module,
|
||||
// since the command itself is llmgw's launcher.
|
||||
func (t *TRTLLM) HelpCommand(Command) (string, []string) {
|
||||
return t.python(), []string{"-m", "tensorrt_llm.commands.serve", "serve", "--help"}
|
||||
}
|
||||
|
||||
const (
|
||||
trtllmKeyEnv = "LLMGW_ENGINE_KEY"
|
||||
trtllmOptionsPlaceholder = "@llmgw-options"
|
||||
)
|
||||
|
||||
// trtllmLauncher starts trtllm-serve inside the process, after two things
|
||||
// trtllm-serve has no flag for.
|
||||
//
|
||||
// The key. trtllm-serve checks none, and an engine listens on the node's
|
||||
// management address, where without one anybody who can reach the port talks
|
||||
// to the model past the gateway's authentication, quotas and accounting. So
|
||||
// every FastAPI application the process makes — trtllm-serve's server is one
|
||||
// — is given a middleware demanding it, before trtllm-serve imports anything.
|
||||
// /health is left open: it says only up or down, and a supervisor of
|
||||
// TensorRT-LLM's own may ask it without a key. The key comes in the
|
||||
// environment and is taken out of it, so the MPI workers and anything else
|
||||
// started from here never hold it.
|
||||
//
|
||||
// The options. The Prometheus endpoint exists only with return_perf_metrics,
|
||||
// and its queue gauges move only with enable_iter_perf_stats; both are
|
||||
// settings of TensorRT-LLM's options file, not flags. The file is written to
|
||||
// an unnamed temporary file and passed by descriptor: nothing is left on disk
|
||||
// however the server ends.
|
||||
//
|
||||
// trtllm-serve can start further HTTP front ends as copies of its own command
|
||||
// line (--num_serve_frontends); those would not run this, so the flag is
|
||||
// refused in a profile, and a copy started anyway would be given a program
|
||||
// path that cannot exist — /dev/null is no directory — and fail rather than
|
||||
// serve without the key. Click names itself by the path's last element, so
|
||||
// trtllm-serve's messages still say trtllm-serve.
|
||||
const trtllmLauncher = `import hmac, os, sys, tempfile
|
||||
key = os.environ.pop("` + trtllmKeyEnv + `", "")
|
||||
if key:
|
||||
import fastapi
|
||||
want = ("Bearer " + key).encode()
|
||||
denied = b'{"error":{"message":"a valid API key is required","type":"invalid_request_error","code":"invalid_api_key"}}'
|
||||
class RequireKey:
|
||||
def __init__(self, app):
|
||||
self.app = app
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] in ("http", "websocket") and scope.get("path") != "/health":
|
||||
got = b""
|
||||
for name, value in scope.get("headers") or ():
|
||||
if name == b"authorization":
|
||||
got = value
|
||||
if not hmac.compare_digest(got, want):
|
||||
if scope["type"] == "websocket":
|
||||
await send({"type": "websocket.close", "code": 1008})
|
||||
return
|
||||
await send({"type": "http.response.start", "status": 401,
|
||||
"headers": [(b"content-type", b"application/json")]})
|
||||
await send({"type": "http.response.body", "body": denied})
|
||||
return
|
||||
await self.app(scope, receive, send)
|
||||
plain_init = fastapi.FastAPI.__init__
|
||||
def init_with_key(self, *a, **kw):
|
||||
plain_init(self, *a, **kw)
|
||||
self.add_middleware(RequireKey)
|
||||
fastapi.FastAPI.__init__ = init_with_key
|
||||
args = sys.argv[1:]
|
||||
if "` + trtllmOptionsPlaceholder + `" in args:
|
||||
options = tempfile.TemporaryFile()
|
||||
options.write(b"return_perf_metrics: true\nenable_iter_perf_stats: true\n")
|
||||
options.flush()
|
||||
options.seek(0)
|
||||
args[args.index("` + trtllmOptionsPlaceholder + `")] = "/dev/fd/%d" % options.fileno()
|
||||
sys.argv = ["/dev/null/trtllm-serve"] + args
|
||||
from tensorrt_llm.commands.serve import main
|
||||
main()
|
||||
`
|
||||
@@ -0,0 +1,316 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func trtllmSpec() ShardSpec {
|
||||
s := singleNodeSpec()
|
||||
s.Profile.Engine = "trtllm"
|
||||
s.Profile.ServedModelName = "llama-3.1-70b"
|
||||
return s
|
||||
}
|
||||
|
||||
func TestTRTLLMRunsTrtllmServeWithTheFlagsItDocuments(t *testing.T) {
|
||||
s := trtllmSpec()
|
||||
s.BindAddr = "10.0.0.147"
|
||||
s.Profile.Args = []string{"--max_batch_size", "64"}
|
||||
cmd := buildOne(t, &TRTLLM{cfg: TRTLLMConfig{Python: "/opt/trtllm/bin/python"}}, s)
|
||||
|
||||
if cmd.Path != "/opt/trtllm/bin/python" {
|
||||
t.Errorf("Path=%q, expected the configured interpreter", cmd.Path)
|
||||
}
|
||||
if cmd.Args[0] != "-c" || !strings.Contains(cmd.Args[1], "from tensorrt_llm.commands.serve import main") {
|
||||
t.Fatalf("not started through the launcher: %v", cmd.Args[:2])
|
||||
}
|
||||
rest := cmd.Args[2:]
|
||||
if rest[0] != "serve" || rest[1] != "/models/llama" {
|
||||
t.Errorf("expected serve and the model first, got %v", rest[:2])
|
||||
}
|
||||
for flag, want := range map[string]string{
|
||||
"--host": "10.0.0.147", "--port": "20000", "--tp_size": "4",
|
||||
"--served_model_name": "llama-3.1-70b",
|
||||
"--extra_llm_api_options": trtllmOptionsPlaceholder,
|
||||
} {
|
||||
if got := argValue(rest, flag); got != want {
|
||||
t.Errorf("%s=%q, want %q (args %v)", flag, got, want, rest)
|
||||
}
|
||||
}
|
||||
// One pipeline stage is the default, and saying so would only be one
|
||||
// more flag for a version to rename.
|
||||
if slices.Contains(rest, "--pp_size") {
|
||||
t.Errorf("--pp_size given for a single stage: %v", rest)
|
||||
}
|
||||
// The profile's own, last, so a tuning flag is the operator's to set.
|
||||
if got := rest[len(rest)-2:]; !slices.Equal(got, []string{"--max_batch_size", "64"}) {
|
||||
t.Errorf("profile args are not last: %v", rest)
|
||||
}
|
||||
|
||||
s.Profile.TP, s.Profile.PP = 2, 2
|
||||
if got := argValue(buildOne(t, &TRTLLM{}, s).Args, "--pp_size"); got != "2" {
|
||||
t.Errorf("--pp_size=%q for two stages", got)
|
||||
}
|
||||
if got := buildOne(t, &TRTLLM{}, s).Path; got != "python3" {
|
||||
t.Errorf("unconfigured, the interpreter is %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTRTLLMsKeyIsNotOnItsCommandLine(t *testing.T) {
|
||||
s := trtllmSpec()
|
||||
s.APIKey = "k1"
|
||||
cmd := buildOne(t, &TRTLLM{}, s)
|
||||
for _, a := range cmd.Args {
|
||||
if strings.Contains(a, "k1") && a != cmd.Args[1] {
|
||||
t.Errorf("the key is in the arguments: %v", cmd.Args)
|
||||
}
|
||||
}
|
||||
if envMap(cmd.Env)[trtllmKeyEnv] != "k1" {
|
||||
t.Error("the key is not in the environment")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTRTLLMDoesNotReportUsageToNVIDIA(t *testing.T) {
|
||||
s := trtllmSpec()
|
||||
s.Profile.Env = map[string]string{"TRTLLM_NO_USAGE_STATS": "0", "TRTLLM_ENABLE_PDL": "1"}
|
||||
env := envMap(buildOne(t, &TRTLLM{}, s).Env)
|
||||
if env["TRTLLM_NO_USAGE_STATS"] != "1" || env["DO_NOT_TRACK"] != "1" {
|
||||
t.Errorf("usage reporting is not off: %q %q", env["TRTLLM_NO_USAGE_STATS"], env["DO_NOT_TRACK"])
|
||||
}
|
||||
// Tuning still reaches it.
|
||||
if env["TRTLLM_ENABLE_PDL"] != "1" {
|
||||
t.Error("a profile's TensorRT-LLM tuning was dropped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTRTLLMIsReadFromItsPrometheusEndpoint(t *testing.T) {
|
||||
// /metrics on trtllm-serve is JSON, and each read drains what it returns.
|
||||
if got := (&TRTLLM{}).MetricsPath(); got != "/prometheus/metrics" {
|
||||
t.Errorf("MetricsPath=%q", got)
|
||||
}
|
||||
if !(&TRTLLM{}).ServesHTTP(trtllmSpec()) {
|
||||
t.Error("the one rank of a replica serves its API")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTRTLLMRunsAReplicaOnOneNode(t *testing.T) {
|
||||
p := trtllmSpec().Profile
|
||||
p.Nodes, p.GPUsPerNode, p.TP = 2, 4, 8
|
||||
err := (&TRTLLM{}).Validate(p, nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "one node") {
|
||||
t.Errorf("a two-node profile was accepted: %v", err)
|
||||
}
|
||||
s := multiNodeSpec(0)
|
||||
s.Profile.Engine = "trtllm"
|
||||
if _, err := (&TRTLLM{}).Commands(s); err == nil {
|
||||
t.Error("a command was built for a shard of a two-node replica")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTRTLLMNeedsAmpereOrNewer(t *testing.T) {
|
||||
p := trtllmSpec().Profile
|
||||
p.TP = 2
|
||||
v100 := []GPUInfo{{Index: 0, Name: "Tesla V100", ComputeCapability: "7.0"}, {Index: 1, Name: "Tesla V100", ComputeCapability: "7.0"}}
|
||||
if err := (&TRTLLM{}).Validate(p, v100); err == nil {
|
||||
t.Error("a V100 was accepted")
|
||||
}
|
||||
h100 := []GPUInfo{{Index: 0, Name: "H100", ComputeCapability: "9.0"}, {Index: 1, Name: "H100", ComputeCapability: "9.0"}}
|
||||
if err := (&TRTLLM{}).Validate(p, h100); err != nil {
|
||||
t.Errorf("two H100s were refused: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTRTLLMProfilesMayNotOpenWhatTheAgentCloses(t *testing.T) {
|
||||
for _, args := range [][]string{
|
||||
{"--trust_remote_code"},
|
||||
{"--num_serve_frontends", "4"},
|
||||
{"--grpc"},
|
||||
{"--config", "/etc/shadow"},
|
||||
{"--extra_llm_api_options=/tmp/x.yaml"},
|
||||
{"--served_model_name", "other"},
|
||||
{"--middleware", "evil.mw"},
|
||||
{"--tp_size", "8"},
|
||||
} {
|
||||
if err := UnsafeArgs("trtllm", args); err == nil {
|
||||
t.Errorf("%v was allowed", args)
|
||||
}
|
||||
}
|
||||
if err := UnsafeArgs("trtllm", []string{"--max_batch_size", "64", "--set", "kv_cache_config.enable_block_reuse=false",
|
||||
"--free_gpu_memory_fraction", "0.85"}); err != nil {
|
||||
t.Errorf("tuning was refused: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTRTLLMIsOfferedOnlyWhenItsPackageIsThere(t *testing.T) {
|
||||
there := script(t, "python3", `echo "1.3.0rc29"`)
|
||||
r := NewRegistry(Config{})
|
||||
r.engines = map[string]Engine{"trtllm": &TRTLLM{cfg: TRTLLMConfig{Python: there}}}
|
||||
if got := r.ProbeConfigured(t.Context()); got["trtllm"] != "1.3.0rc29" {
|
||||
t.Errorf("expected the version, got %v", got)
|
||||
}
|
||||
|
||||
// An interpreter without the package, which every machine has.
|
||||
absent := script(t, "python3", "exit 1")
|
||||
r.engines = map[string]Engine{"trtllm": &TRTLLM{cfg: TRTLLMConfig{Python: absent}}}
|
||||
if got := r.ProbeConfigured(t.Context()); len(got) != 0 {
|
||||
t.Errorf("reported without its package: %v", got)
|
||||
}
|
||||
|
||||
// Not named, not offered: python3 on PATH is nobody's TensorRT-LLM.
|
||||
r.engines = map[string]Engine{"trtllm": &TRTLLM{}}
|
||||
if got := r.ProbeConfigured(t.Context()); len(got) != 0 {
|
||||
t.Errorf("offered without being configured: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// fakeTRTLLM is a Python path holding stand-ins for fastapi and
|
||||
// tensorrt_llm: enough of each for the launcher to do what it does in front
|
||||
// of the real server, and for the stand-in server to report what it saw.
|
||||
func fakeTRTLLM(t *testing.T) string {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
files := map[string]string{
|
||||
"fastapi/__init__.py": `
|
||||
class FastAPI:
|
||||
def __init__(self, *a, **kw):
|
||||
self.user_middleware = []
|
||||
def add_middleware(self, cls, **kw):
|
||||
self.user_middleware.append(cls)
|
||||
`,
|
||||
"tensorrt_llm/__init__.py": "",
|
||||
"tensorrt_llm/commands/__init__.py": "",
|
||||
"tensorrt_llm/commands/serve.py": `
|
||||
import asyncio, json, os, sys
|
||||
import fastapi
|
||||
|
||||
def main():
|
||||
args = sys.argv[1:]
|
||||
with open(args[args.index("--extra_llm_api_options") + 1]) as f:
|
||||
options = f.read()
|
||||
app = fastapi.FastAPI()
|
||||
|
||||
async def server(scope, receive, send):
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"ok"})
|
||||
|
||||
handler = server
|
||||
for mw in reversed(app.user_middleware):
|
||||
handler = mw(handler)
|
||||
|
||||
async def ask(path, auth=None):
|
||||
headers = [(b"authorization", auth.encode())] if auth else []
|
||||
sent = []
|
||||
async def send(msg):
|
||||
sent.append(msg)
|
||||
async def receive():
|
||||
return {"type": "http.request"}
|
||||
await handler({"type": "http", "path": path, "headers": headers}, receive, send)
|
||||
return sent[0]["status"]
|
||||
|
||||
async def run():
|
||||
return {
|
||||
"health": await ask("/health"),
|
||||
"anonymous": await ask("/v1/chat/completions"),
|
||||
"wrong": await ask("/v1/models", "Bearer nope"),
|
||||
"keyed": await ask("/prometheus/metrics", "Bearer k1"),
|
||||
}
|
||||
|
||||
out = asyncio.run(run())
|
||||
out.update(options=options, argv0=sys.argv[0], args=args,
|
||||
key_left=os.environ.get("LLMGW_ENGINE_KEY"))
|
||||
print(json.dumps(out))
|
||||
`,
|
||||
}
|
||||
for name, body := range files {
|
||||
p := filepath.Join(dir, name)
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(p, []byte(body), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
return dir
|
||||
}
|
||||
|
||||
type launched struct {
|
||||
Health, Anonymous, Wrong, Keyed int
|
||||
Options, Argv0 string
|
||||
Args []string
|
||||
KeyLeft *string `json:"key_left"`
|
||||
}
|
||||
|
||||
func launch(t *testing.T, s ShardSpec) launched {
|
||||
t.Helper()
|
||||
python, err := exec.LookPath("python3")
|
||||
if err != nil {
|
||||
t.Skip("no python3 here to run the launcher with")
|
||||
}
|
||||
cmd := buildOne(t, &TRTLLM{cfg: TRTLLMConfig{Python: python}}, s)
|
||||
c := exec.Command(cmd.Path, cmd.Args...) //nolint:gosec // the command under test
|
||||
c.Env = append(cmd.Env, "PYTHONPATH="+fakeTRTLLM(t))
|
||||
out, err := c.Output()
|
||||
if err != nil {
|
||||
var stderr string
|
||||
if ee, ok := err.(*exec.ExitError); ok {
|
||||
stderr = string(ee.Stderr)
|
||||
}
|
||||
t.Fatalf("the launcher failed: %v\n%s", err, stderr)
|
||||
}
|
||||
var got launched
|
||||
if err := json.Unmarshal(out, &got); err != nil {
|
||||
t.Fatalf("%v: %s", err, out)
|
||||
}
|
||||
return got
|
||||
}
|
||||
|
||||
func TestTheTRTLLMLauncherDemandsTheKey(t *testing.T) {
|
||||
s := trtllmSpec()
|
||||
s.APIKey = "k1"
|
||||
got := launch(t, s)
|
||||
|
||||
if got.Anonymous != 401 || got.Wrong != 401 {
|
||||
t.Errorf("served without the key: anonymous %d, wrong key %d", got.Anonymous, got.Wrong)
|
||||
}
|
||||
if got.Keyed != 200 {
|
||||
t.Errorf("refused the key: %d", got.Keyed)
|
||||
}
|
||||
// Up or down is all /health says, and it is asked without one.
|
||||
if got.Health != 200 {
|
||||
t.Errorf("/health answered %d", got.Health)
|
||||
}
|
||||
// Taken out of the environment, so nothing the server starts holds it.
|
||||
if got.KeyLeft != nil {
|
||||
t.Errorf("the key was left in the environment: %q", *got.KeyLeft)
|
||||
}
|
||||
// A copy of this command line, as trtllm-serve's extra front ends are,
|
||||
// names no program that exists and so cannot serve without the key.
|
||||
if _, err := os.Stat(got.Argv0); err == nil || filepath.Base(got.Argv0) != "trtllm-serve" {
|
||||
t.Errorf("argv[0] is %q", got.Argv0)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheTRTLLMLauncherTurnsOnTheMetricsTheRouterReads(t *testing.T) {
|
||||
got := launch(t, trtllmSpec())
|
||||
for _, want := range []string{"return_perf_metrics: true", "enable_iter_perf_stats: true"} {
|
||||
if !strings.Contains(got.Options, want) {
|
||||
t.Errorf("the options file lacks %q: %q", want, got.Options)
|
||||
}
|
||||
}
|
||||
path := argValue(got.Args, "--extra_llm_api_options")
|
||||
if !strings.HasPrefix(path, "/dev/fd/") {
|
||||
t.Errorf("the options were not passed by descriptor: %q", path)
|
||||
}
|
||||
if slices.Contains(got.Args, trtllmOptionsPlaceholder) {
|
||||
t.Error("the placeholder reached trtllm-serve")
|
||||
}
|
||||
// Without a key there is nothing to demand.
|
||||
if got.Anonymous != 200 {
|
||||
t.Errorf("with no key set, a request was refused: %d", got.Anonymous)
|
||||
}
|
||||
}
|
||||
@@ -1,30 +1,174 @@
|
||||
package engines
|
||||
|
||||
import "strings"
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// A model profile's env and args are the operator's escape hatch for tuning
|
||||
// an engine, and they reach a process the agent starts as root on every node
|
||||
// the model is placed on. Some of what an engine accepts is not tuning but a
|
||||
// way to run other code, or to read or write files, in that process: a
|
||||
// loader preload, an NCCL network plugin, a CUDA injection library, vLLM's
|
||||
// --trust-remote-code, llama.cpp's --log-file pointed at /etc/cron.d. Taken
|
||||
// from a model definition, which the admin API or a config file supplies,
|
||||
// they would turn "add a model" into "run this as root on every node".
|
||||
//
|
||||
// So env is an allowlist — the families of variables engines are tuned with
|
||||
// — less any whose name says it names a path, a plugin or a library, and args
|
||||
// are checked against what each engine takes that reads, writes or loads.
|
||||
// Both are refused where a model is defined and again where an engine starts.
|
||||
|
||||
// envFamilies are the prefixes of variables engines are tuned with.
|
||||
var envFamilies = []string{
|
||||
"CUDA_", "NCCL_", "NVIDIA_", "NVTE_", "GLOO_", "UCX_", "TORCH_", "PYTORCH_",
|
||||
"VLLM_", "SGLANG_", "GGML_", "LLAMA_", "OMP_", "MKL_", "OPENBLAS_", "KMP_",
|
||||
"HF_", "TRANSFORMERS_", "TOKENIZERS_", "TRITON_", "FLASHINFER_",
|
||||
"TRTLLM_", "TLLM_",
|
||||
}
|
||||
|
||||
// envExact are single variables that are tuning too.
|
||||
var envExact = map[string]bool{
|
||||
"PYTHONUNBUFFERED": true, "PYTHONHASHSEED": true, "TZ": true,
|
||||
}
|
||||
|
||||
// envNever are names in an allowed family that load code or move where the
|
||||
// engine reads and writes.
|
||||
var envNever = []string{
|
||||
"PLUGIN", "INJECTION", "PRELOAD", "LIBRARY", "_LIB",
|
||||
"_PATH", "_DIR", "_HOME", "_FILE", "_ROOT", "_CACHE", "_CONFIG",
|
||||
"_ADDRESS", "_ENDPOINT", "_TOKEN", "_URL",
|
||||
}
|
||||
|
||||
// UnsafeEnv reports whether a profile may not set an environment variable,
|
||||
// and why.
|
||||
//
|
||||
// A profile's env is the escape hatch for tuning no adapter can guess —
|
||||
// CUDA_*, NCCL_*, VLLM_*, HF_* — and it reaches an engine the agent starts as
|
||||
// root. Some variables are not tuning but a way to run other code in that
|
||||
// process: the dynamic loader's LD_PRELOAD, Python's PYTHONSTARTUP, a shell's
|
||||
// BASH_ENV. Set from a model definition, which an admin API or a config file
|
||||
// supplies, they would turn "add a model" into "run this as root on every
|
||||
// node". They are refused where a model is defined and again where an engine
|
||||
// starts.
|
||||
func UnsafeEnv(key string) (bool, string) {
|
||||
k := strings.ToUpper(key)
|
||||
switch {
|
||||
case strings.HasPrefix(k, "LD_"), strings.HasPrefix(k, "DYLD_"):
|
||||
return true, "it changes what the dynamic loader loads"
|
||||
case k == "PATH":
|
||||
return true, "it changes which binaries run"
|
||||
case strings.HasPrefix(k, "PYTHON") && k != "PYTHONUNBUFFERED" && k != "PYTHONHASHSEED":
|
||||
return true, "it changes what Python imports or runs at start"
|
||||
case k == "BASH_ENV" || k == "ENV" || k == "PERL5OPT" || k == "PERL5LIB" ||
|
||||
k == "NODE_OPTIONS" || k == "RUBYOPT" || k == "GCONV_PATH" || k == "OPENSSL_CONF":
|
||||
return true, "it makes an interpreter or library load code"
|
||||
k := strings.ToUpper(strings.TrimSpace(key))
|
||||
if k == "" {
|
||||
return true, "it has no name"
|
||||
}
|
||||
allowed := envExact[k]
|
||||
for _, f := range envFamilies {
|
||||
if strings.HasPrefix(k, f) {
|
||||
allowed = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !allowed {
|
||||
return true, "only the variables engines are tuned with are allowed (" +
|
||||
strings.Join(envFamilies, "*, ") + "*)"
|
||||
}
|
||||
for _, n := range envNever {
|
||||
if strings.Contains(k, n) && !envExact[k] {
|
||||
return true, "it names a path, a plugin, a library or an address, which is a way to make the engine " +
|
||||
"read, write or load something other than its model"
|
||||
}
|
||||
}
|
||||
// llama.cpp reads every flag from LLAMA_ARG_<FLAG> too.
|
||||
if rest, ok := strings.CutPrefix(k, "LLAMA_ARG_"); ok {
|
||||
flag := "--" + strings.ToLower(strings.ReplaceAll(rest, "_", "-"))
|
||||
if why := unsafeFlag("llamacpp", flag, ""); why != "" {
|
||||
return true, "it is " + flag + ", which " + why
|
||||
}
|
||||
}
|
||||
return false, ""
|
||||
}
|
||||
|
||||
// flagsNever are flags no profile may give, by engine; "" is every engine.
|
||||
// Each reads or writes files the agent did not choose, loads code, or moves
|
||||
// what the agent itself sets: where the engine listens and the key it
|
||||
// demands.
|
||||
var flagsNever = map[string]map[string]string{
|
||||
"": {
|
||||
"--host": "is set by the agent", "--port": "is set by the agent",
|
||||
"--api-key": "is set by the agent", "--api-key-file": "is set by the agent",
|
||||
"--model": "is the model's source, set by the agent", "-m": "is the model's source, set by the agent",
|
||||
"--ssl-keyfile": "reads a file", "--ssl-certfile": "reads a file",
|
||||
"--ssl-key-file": "reads a file", "--ssl-cert-file": "reads a file",
|
||||
"--trust-remote-code": "runs Python the model repository ships",
|
||||
},
|
||||
"llamacpp": {
|
||||
"--log-file": "writes a file", "--slot-save-path": "writes files",
|
||||
"--path": "serves a directory's files over HTTP", "--rpc": "sends the model's work to other machines",
|
||||
"--chat-template-file": "reads a file", "--system-prompt-file": "reads a file",
|
||||
"--grammar-file": "reads a file", "--json-schema-file": "reads a file",
|
||||
"--lora": "loads a file; adapters are set in the model's adapters",
|
||||
"--lora-scaled": "loads a file; adapters are set in the model's adapters",
|
||||
"--control-vector": "loads a file", "--control-vector-scaled": "loads a file",
|
||||
"--mmproj": "loads a file; the projector is fetched with the model",
|
||||
"--model-url": "downloads a model", "-mu": "downloads a model",
|
||||
"--hf-repo": "downloads a model", "-hf": "downloads a model", "--hf-file": "downloads a model",
|
||||
"--hf-token": "is a credential", "--alias": "is set by the agent",
|
||||
"--model-draft": "loads a file", "-md": "loads a file",
|
||||
"--prompt-cache": "writes a file", "--logdir": "writes files",
|
||||
},
|
||||
"vllm": {
|
||||
"--download-dir": "writes files", "--tokenizer": "reads another repository or path",
|
||||
"--chat-template": "reads a file", "--allowed-local-media-path": "serves local files to requests",
|
||||
"--lora-modules": "loads files; adapters are set in the model's adapters",
|
||||
"--tool-parser-plugin": "runs a Python file", "--middleware": "imports Python",
|
||||
"--worker-extension-cls": "imports Python", "--worker-cls": "imports Python",
|
||||
"--load-format": "changes how weights are loaded, some ways running code",
|
||||
"--otlp-traces-endpoint": "sends request traces elsewhere",
|
||||
"--served-model-name": "is set by the agent",
|
||||
"--enable-log-requests": "logs prompts",
|
||||
},
|
||||
"sglang": {
|
||||
"--tokenizer-path": "reads another repository or path", "--chat-template": "reads a file",
|
||||
"--file-storage-path": "writes files", "--log-requests": "logs prompts",
|
||||
"--model-path": "is the model's source, set by the agent", "--served-model-name": "is set by the agent",
|
||||
"--lora-paths": "loads files", "--download-dir": "writes files",
|
||||
},
|
||||
// trtllm-serve spells its flags with underscores, so the shared list's
|
||||
// hyphenated names do not reach them.
|
||||
"trtllm": {
|
||||
"--trust_remote_code": "runs Python the model repository ships",
|
||||
"--served_model_name": "is set by the agent",
|
||||
"--config": "reads a file; the agent writes the options file itself",
|
||||
"--extra_llm_api_options": "reads a file; the agent writes the options file itself",
|
||||
"--tokenizer": "reads another repository or path",
|
||||
"--custom_tokenizer": "imports Python", "--post_processor_hook": "imports Python",
|
||||
"--custom_module_dirs": "imports Python", "--middleware": "imports Python",
|
||||
"--chat_template": "reads a file",
|
||||
"--metadata_server_config_file": "reads a file",
|
||||
"--visual_gen_args": "reads a file",
|
||||
"--report_addr": "writes a file",
|
||||
"--otlp_traces_endpoint": "sends request traces elsewhere",
|
||||
"--disagg_cluster_uri": "joins the server to another cluster",
|
||||
"--server_role": "makes the server part of another cluster",
|
||||
"--allow_request_chat_template": "lets a request supply the chat template",
|
||||
"--grpc": "serves without the key the agent sets",
|
||||
"--num_serve_frontends": "starts front ends that do not demand the key the agent sets",
|
||||
"--telemetry": "sends usage to NVIDIA",
|
||||
"--tensor_parallel_size": "is set by the agent from tp",
|
||||
"--tp_size": "is set by the agent from tp",
|
||||
"--pipeline_parallel_size": "is set by the agent from pp",
|
||||
"--pp_size": "is set by the agent from pp",
|
||||
},
|
||||
}
|
||||
|
||||
// unsafeFlag is why an engine may not be given a flag, or "".
|
||||
func unsafeFlag(engine, flag, _ string) string {
|
||||
name, _, _ := strings.Cut(flag, "=")
|
||||
for _, set := range []string{"", engine} {
|
||||
if why, ok := flagsNever[set][name]; ok && why != "" {
|
||||
return why
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// UnsafeArgs refuses the first of a profile's arguments an engine may not be
|
||||
// given, saying why.
|
||||
func UnsafeArgs(engine string, args []string) error {
|
||||
for _, a := range args {
|
||||
if !strings.HasPrefix(a, "-") {
|
||||
continue // a value
|
||||
}
|
||||
if why := unsafeFlag(engine, a, ""); why != "" {
|
||||
name, _, _ := strings.Cut(a, "=")
|
||||
return fmt.Errorf("argument %s is not allowed: it %s", name, why)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package engines
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestTuningVariablesAreAllowedAndWaysToRunCodeAreNot(t *testing.T) {
|
||||
for _, k := range []string{
|
||||
"CUDA_VISIBLE_DEVICES", "CUDA_MODULE_LOADING", "NCCL_IB_HCA", "NCCL_DEBUG", "NCCL_NET",
|
||||
"VLLM_ATTENTION_BACKEND", "PYTORCH_CUDA_ALLOC_CONF", "OMP_NUM_THREADS", "HF_HUB_OFFLINE",
|
||||
"GGML_VK_VISIBLE_DEVICES", "TOKENIZERS_PARALLELISM", "PYTHONUNBUFFERED", "LLAMA_ARG_CTX_SIZE",
|
||||
} {
|
||||
if bad, why := UnsafeEnv(k); bad {
|
||||
t.Errorf("%s refused: %s", k, why)
|
||||
}
|
||||
}
|
||||
for _, k := range []string{
|
||||
"LD_PRELOAD", "PATH", "PYTHONPATH", "PYTHONSTARTUP", "BASH_ENV",
|
||||
// What the denylist missed: each loads a library or moves where the
|
||||
// engine reads and writes.
|
||||
"NCCL_NET_PLUGIN", "NCCL_TUNER_PLUGIN", "NCCL_PROFILER_PLUGIN", "CUDA_INJECTION64_PATH",
|
||||
"NVTX_INJECTION64_PATH", "TRITON_CACHE_DIR", "HF_HOME", "HF_TOKEN", "NCCL_DEBUG_FILE",
|
||||
"VLLM_PLUGINS", "UCX_MODULE_DIR", "TORCH_EXTENSIONS_DIR",
|
||||
// llama.cpp's flags as variables.
|
||||
"LLAMA_ARG_LOG_FILE", "LLAMA_ARG_HOST", "LLAMA_ARG_SLOT_SAVE_PATH",
|
||||
// Anything outside the families.
|
||||
"SOMETHING_ELSE", "",
|
||||
} {
|
||||
if bad, _ := UnsafeEnv(k); !bad {
|
||||
t.Errorf("%s allowed", k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestArgumentsThatReadWriteOrLoadAreRefused(t *testing.T) {
|
||||
ok := map[string][]string{
|
||||
"llamacpp": {"--ctx-size", "16384", "--embedding", "--parallel", "4", "--jinja", "--n-gpu-layers=999"},
|
||||
"vllm": {"--max-model-len", "8192", "--gpu-memory-utilization", "0.9", "--enable-auto-tool-choice"},
|
||||
"sglang": {"--mem-fraction-static", "0.8", "--tool-call-parser", "qwen25"},
|
||||
}
|
||||
for eng, args := range ok {
|
||||
if err := UnsafeArgs(eng, args); err != nil {
|
||||
t.Errorf("%s %v: %v", eng, args, err)
|
||||
}
|
||||
}
|
||||
for _, c := range []struct {
|
||||
eng string
|
||||
args []string
|
||||
}{
|
||||
{"vllm", []string{"--trust-remote-code"}},
|
||||
{"sglang", []string{"--trust-remote-code"}},
|
||||
{"vllm", []string{"--tool-parser-plugin", "/tmp/x.py"}},
|
||||
{"vllm", []string{"--middleware=evil.mod"}},
|
||||
{"vllm", []string{"--allowed-local-media-path", "/"}},
|
||||
{"llamacpp", []string{"--log-file", "/etc/cron.d/x"}},
|
||||
{"llamacpp", []string{"--slot-save-path=/root"}},
|
||||
{"llamacpp", []string{"--path", "/"}},
|
||||
{"llamacpp", []string{"--system-prompt-file", "/etc/shadow"}},
|
||||
{"llamacpp", []string{"--host", "0.0.0.0"}},
|
||||
{"vllm", []string{"--api-key", "x"}},
|
||||
} {
|
||||
if err := UnsafeArgs(c.eng, c.args); err == nil {
|
||||
t.Errorf("%s %v allowed", c.eng, c.args)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2,8 +2,11 @@ package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"maps"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
// VLLMConfig locates the vLLM installation and picks how it is launched.
|
||||
@@ -106,7 +109,7 @@ func (v *VLLM) serverCommand(s ShardSpec) (Command, error) {
|
||||
p := s.Profile
|
||||
|
||||
args := []string{"serve", p.ModelPath,
|
||||
"--host", "0.0.0.0",
|
||||
"--host", s.Host(),
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
"--tensor-parallel-size", strconv.Itoa(p.TP),
|
||||
}
|
||||
@@ -147,6 +150,28 @@ func (v *VLLM) serverCommand(s ShardSpec) (Command, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// Adapters loaded at start are registered under their names, and a
|
||||
// request names one as its model. Adapters loaded on demand start
|
||||
// unregistered: the router loads each when a request first wants it
|
||||
// here, by the link the agent made (RuntimeLoRADir). --max-loras is
|
||||
// vLLM's default of one at start — raise it in the profile's arguments
|
||||
// to mix more adapters in a batch — and the cap on demand, which
|
||||
// WithRuntimeLoRA put in the arguments.
|
||||
if len(p.Adapters) > 0 {
|
||||
args = append(args, "--enable-lora")
|
||||
if !p.RuntimeLoRA() {
|
||||
args = append(args, loraModules(p.Adapters)...)
|
||||
}
|
||||
}
|
||||
|
||||
// The KV cache events, for the router's precise index: published to the
|
||||
// agent on loopback, hashed the way the router hashes. Only the leader
|
||||
// has a scheduler to publish from.
|
||||
if p.KVEvents && s.KVEventsPort > 0 && s.Leader() {
|
||||
args = append(args, kvevents.VLLMArgs(fmt.Sprintf("tcp://127.0.0.1:%d", s.KVEventsPort))...)
|
||||
maps.Copy(env, kvevents.VLLMEnv())
|
||||
}
|
||||
|
||||
// Profile arguments last, so an operator can override anything above.
|
||||
args = append(args, p.Args...)
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -105,3 +106,49 @@ func TestALeaderStillAcceptsWithNoLeadershipConfigured(t *testing.T) {
|
||||
}
|
||||
t.Fatal("a service with no leadership configured refused an agent")
|
||||
}
|
||||
|
||||
// flipping is a leader that can stop leading partway through a test.
|
||||
type flipping struct{ leading atomic.Bool }
|
||||
|
||||
func (f *flipping) IsLeader() bool { return f.leading.Load() }
|
||||
|
||||
func TestAnAgentLeavesAControlPlaneThatStopsLeading(t *testing.T) {
|
||||
// Leadership lost while the process lives on: the store was slow and the
|
||||
// lease lapsed, or another control plane took it. Agents that stayed on
|
||||
// the old leader would leave the new one with no nodes, so the old one
|
||||
// has to let them go rather than wait for them to leave.
|
||||
was := &flipping{}
|
||||
was.leading.Store(true)
|
||||
wasInv, wasAddr := startPlain(t, was)
|
||||
nowInv, nowAddr := startPlain(t, standby(false))
|
||||
|
||||
ag := runAgent(t, Config{
|
||||
NodeID: "n1", MgmtAddr: "127.0.0.1",
|
||||
Servers: []string{wasAddr, nowAddr},
|
||||
})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
go ag.Run(ctx)
|
||||
|
||||
waitFor(t, 5*time.Second, "the agent never reached the first leader", func() bool {
|
||||
_, ok := wasInv.Get("n1")
|
||||
return ok
|
||||
})
|
||||
was.leading.Store(false)
|
||||
waitFor(t, 8*time.Second, "the agent stayed on a control plane that stopped leading", func() bool {
|
||||
_, ok := nowInv.Get("n1")
|
||||
return ok
|
||||
})
|
||||
}
|
||||
|
||||
func waitFor(t *testing.T, d time.Duration, msg string, ok func() bool) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(d)
|
||||
for time.Now().Before(deadline) {
|
||||
if ok() {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Fatal(msg)
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
|
||||
)
|
||||
|
||||
// applyFirewall keeps llmgw's firewall on this node as the cluster's
|
||||
// settings say: the engines' ports and the node exporter's for the networks
|
||||
// the node reaches its control planes on, and those named, or no table at
|
||||
// all. Off the receive loop, one at a time.
|
||||
func (a *Agent) applyFirewall(m *llmgwv1.Firewall) {
|
||||
a.fwMu.Lock()
|
||||
defer a.fwMu.Unlock()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cancel()
|
||||
r := a.fwRunner
|
||||
if r == nil {
|
||||
r = firewall.Local{}
|
||||
}
|
||||
if !m.GetEnabled() {
|
||||
had, err := firewall.Disable(ctx, r)
|
||||
switch {
|
||||
case err != nil:
|
||||
a.log.Warn("could not take llmgw's firewall off this node", "err", err)
|
||||
case had:
|
||||
a.log.Info("llmgw's firewall is off on this node, as the cluster's settings say")
|
||||
}
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
servers := append([]string(nil), a.serverList...)
|
||||
a.mu.Unlock()
|
||||
nets, err := firewall.Networks(ctx, r, controlPlaneIPs(ctx, servers))
|
||||
if err != nil {
|
||||
a.log.Warn("could not tell which networks reach the control planes; the firewall is left as it is", "err", err)
|
||||
return
|
||||
}
|
||||
for _, n := range m.GetNetworks() {
|
||||
if p, err := netip.ParsePrefix(n); err == nil {
|
||||
nets = append(nets, p)
|
||||
} else if ad, err := netip.ParseAddr(n); err == nil {
|
||||
nets = append(nets, netip.PrefixFrom(ad, ad.BitLen()))
|
||||
}
|
||||
}
|
||||
from, to := int(m.GetEnginePortsFrom()), int(m.GetEnginePortsTo())
|
||||
ports := []firewall.Port{firewall.NodeExporter}
|
||||
if from > 0 && to >= from {
|
||||
ports = append(ports, firewall.Engines(from, to))
|
||||
}
|
||||
changed, err := firewall.Install(ctx, r, firewall.Spec{Ports: ports, Allow: nets})
|
||||
switch {
|
||||
case err != nil:
|
||||
a.log.Warn("could not put llmgw's firewall on this node", "err", err)
|
||||
case changed:
|
||||
a.log.Info("llmgw's firewall is on this node", "networks", len(nets))
|
||||
}
|
||||
}
|
||||
|
||||
// controlPlaneIPs are the addresses of the control planes an agent dials.
|
||||
func controlPlaneIPs(ctx context.Context, servers []string) []string {
|
||||
var out []string
|
||||
for _, s := range servers {
|
||||
host := s
|
||||
if u, err := url.Parse(s); err == nil && u.Host != "" {
|
||||
host = u.Hostname()
|
||||
}
|
||||
host = strings.Trim(host, "[]")
|
||||
if a, err := netip.ParseAddr(host); err == nil {
|
||||
out = append(out, a.String())
|
||||
continue
|
||||
}
|
||||
ips, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, ip := range ips {
|
||||
out = append(out, ip.Unmap().String())
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
// fakeShell answers the firewall package's scripts and records what was
|
||||
// installed or removed.
|
||||
type fakeShell struct {
|
||||
present bool
|
||||
installed string
|
||||
deleted bool
|
||||
}
|
||||
|
||||
func (f *fakeShell) Run(_ context.Context, script string, stdin []byte) ([]byte, error) {
|
||||
switch {
|
||||
case strings.HasPrefix(script, "ip -o route get 10.0.0.140"):
|
||||
return []byte("10.0.0.140 dev eth0 src 10.0.0.141 uid 0\n2: eth0 inet 10.0.0.141/24 brd 10.0.0.255 scope global eth0\n"), nil
|
||||
case strings.HasPrefix(script, "sudo -n /usr/sbin/nft list table"):
|
||||
if f.present {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, errors.New("no table")
|
||||
case strings.Contains(script, "nft delete table") && !strings.Contains(script, "cat > "):
|
||||
f.deleted, f.present = true, false
|
||||
return nil, nil
|
||||
case strings.HasPrefix(script, "sudo cat "):
|
||||
return nil, nil
|
||||
default:
|
||||
f.installed, f.present = string(stdin), true
|
||||
return []byte("changed\n"), nil
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodeKeepsOrDropsItsFirewallAsTheClusterSays(t *testing.T) {
|
||||
sh := &fakeShell{}
|
||||
a := &Agent{log: slog.New(slog.DiscardHandler), fwRunner: sh, serverList: []string{"https://10.0.0.140:7443"}}
|
||||
|
||||
a.applyFirewall(&llmgwv1.Firewall{Enabled: true, Networks: []string{"192.168.5.0/24", "10.9.9.9"},
|
||||
EnginePortsFrom: 20000, EnginePortsTo: 20999})
|
||||
for _, want := range []string{"10.0.0.0/24", "192.168.5.0/24", "10.9.9.9/32", "20000-20999", "9100"} {
|
||||
if !strings.Contains(sh.installed, want) {
|
||||
t.Errorf("the rules lack %s:\n%s", want, sh.installed)
|
||||
}
|
||||
}
|
||||
|
||||
a.applyFirewall(&llmgwv1.Firewall{Enabled: false})
|
||||
if !sh.deleted {
|
||||
t.Error("turned off, the table stayed")
|
||||
}
|
||||
}
|
||||
@@ -4,10 +4,12 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
@@ -153,3 +155,70 @@ func joinIndices(is []int) string {
|
||||
}
|
||||
return strings.Join(s, ", ")
|
||||
}
|
||||
|
||||
// engineCompute is where an engine said it put its model, from its output
|
||||
// so far — llama.cpp's, which says — and false for one that said nothing. A
|
||||
// shard given GPUs that its engine says it computes on the CPU is an error,
|
||||
// found this way even where nvidia-smi cannot be asked.
|
||||
func engineCompute(eng engines.Engine, gpus []int, lines []string) (engines.Compute, bool, error) {
|
||||
if _, ok := eng.(engines.CPUFallback); !ok {
|
||||
return engines.Compute{}, false, nil
|
||||
}
|
||||
read, ok := engines.ReadLlamaCompute(lines)
|
||||
if !ok {
|
||||
return engines.Compute{}, false, nil
|
||||
}
|
||||
if len(gpus) > 0 && read.OnTheCPU() {
|
||||
why := "its GPU backend saw no device it could use"
|
||||
if read.Total > 0 && !read.NoGPU {
|
||||
why = fmt.Sprintf("it offloaded %d of %d layers", read.Offloaded, read.Total)
|
||||
}
|
||||
return read, true, fmt.Errorf("placed on GPU %v but computing on the CPU: %s, by llama.cpp's own account "+
|
||||
"(llmgw ctl engine-check on this node says why)", gpus, why)
|
||||
}
|
||||
return read, true, nil
|
||||
}
|
||||
|
||||
// computeToWire is a compute as it travels to the control plane.
|
||||
func computeToWire(c *engines.Compute) *llmgwv1.Compute {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
return &llmgwv1.Compute{Backend: c.Backend, Devices: c.Devices,
|
||||
LayersOffloaded: int32(c.Offloaded), LayersTotal: int32(c.Total)} //nolint:gosec // layer counts
|
||||
}
|
||||
|
||||
// gpuCompute is where a shard computes as nvidia-smi sees it: the CPU for a
|
||||
// shard given no GPU, else the GPUs whose memory its processes hold, and how
|
||||
// much. False when nvidia-smi cannot be asked. Engine-agnostic, and needing
|
||||
// no log line: llama.cpp says where its layers went only at -lv 4.
|
||||
func (a *Agent) gpuCompute(ctx context.Context, gpus []int, pid int) (engines.Compute, bool) {
|
||||
if len(gpus) == 0 {
|
||||
return engines.Compute{Backend: "CPU"}, true
|
||||
}
|
||||
mem, err := gpuMemory(ctx, a.cfg.Discovery.NvidiaSMI)
|
||||
if err != nil || mem == nil || pid <= 0 {
|
||||
return engines.Compute{}, false
|
||||
}
|
||||
group := processGroup(pid)
|
||||
c := engines.Compute{Backend: "CUDA"}
|
||||
for _, g := range mem {
|
||||
if !slices.Contains(gpus, g.Index) {
|
||||
continue
|
||||
}
|
||||
used := 0
|
||||
for _, p := range g.Procs {
|
||||
if holds([]int{p.PID}, pid, group) {
|
||||
used += p.UsedMB
|
||||
}
|
||||
}
|
||||
if used > 0 {
|
||||
c.Devices = append(c.Devices, fmt.Sprintf("GPU %d: %.1f GB", g.Index, float64(used)/1024))
|
||||
}
|
||||
}
|
||||
if len(c.Devices) == 0 {
|
||||
// Placed on GPUs and holding none of their memory.
|
||||
c.Backend = "CPU"
|
||||
}
|
||||
return c, true
|
||||
}
|
||||
@@ -122,3 +122,55 @@ func TestForeignGPUMemoryLeavesOutTheAgentsOwnShards(t *testing.T) {
|
||||
t.Errorf("no reading came back as %v", foreign)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppsOwnWordOnTheCPUIsEnoughToRefuse(t *testing.T) {
|
||||
eng := fallbackEngine{}
|
||||
// Without nvidia-smi to ask, the engine's own output says it.
|
||||
c, said, err := engineCompute(eng, []int{0}, []string{
|
||||
"warning: no usable GPU found, --gpu-layers option will be ignored",
|
||||
"load_tensors: offloaded 0/49 layers to GPU"})
|
||||
if err == nil || !said || !strings.Contains(err.Error(), "computing on the CPU") || c.Backend != "CPU" {
|
||||
t.Fatalf("a GPU shard llama.cpp put on the CPU: %v %+v", err, c)
|
||||
}
|
||||
c, said, err = engineCompute(eng, []int{0}, []string{
|
||||
"llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free",
|
||||
"load_tensors: offloaded 49/49 layers to GPU"})
|
||||
if err != nil || !said || c.Backend != "CUDA" || c.Offloaded != 49 || computeToWire(&c).GetLayersOffloaded() != 49 {
|
||||
t.Errorf("on its GPU: %v %+v", err, c)
|
||||
}
|
||||
// A CPU shard is not refused for being on the CPU, and says so.
|
||||
if c, _, err := engineCompute(eng, nil, []string{"load_tensors: offloaded 0/49 layers to GPU"}); err != nil || c.Backend != "CPU" {
|
||||
t.Errorf("a CPU shard: %v %+v", err, c)
|
||||
}
|
||||
// An engine that said nothing says nothing.
|
||||
if _, said, err := engineCompute(eng, []int{0}, []string{"listening"}); err != nil || said {
|
||||
t.Errorf("nothing said: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNvidiaSMISaysWhereAShardComputes(t *testing.T) {
|
||||
oldMem, oldGroup := gpuMemory, processGroup
|
||||
t.Cleanup(func() { gpuMemory, processGroup = oldMem, oldGroup })
|
||||
processGroup = func(pid int) int { return pid }
|
||||
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) {
|
||||
return []discovery.GPUMemory{
|
||||
{Index: 0, UsedMB: 20000, Procs: []discovery.GPUProcess{{PID: 700, UsedMB: 2048}, {PID: 900, UsedMB: 17000}}},
|
||||
{Index: 1, UsedMB: 0},
|
||||
}, nil
|
||||
}
|
||||
a := &Agent{log: slog.Default()}
|
||||
if c, ok := a.gpuCompute(context.Background(), []int{0}, 700); !ok || c.Backend != "CUDA" ||
|
||||
len(c.Devices) != 1 || c.Devices[0] != "GPU 0: 2.0 GB" {
|
||||
t.Errorf("on GPU 0: %+v %v", c, ok)
|
||||
}
|
||||
if c, ok := a.gpuCompute(context.Background(), []int{1}, 700); !ok || c.Backend != "CPU" {
|
||||
t.Errorf("placed on GPU 1, holding none of it: %+v", c)
|
||||
}
|
||||
if c, ok := a.gpuCompute(context.Background(), nil, 700); !ok || c.Backend != "CPU" {
|
||||
t.Errorf("a CPU shard: %+v", c)
|
||||
}
|
||||
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) { return nil, nil }
|
||||
if _, ok := a.gpuCompute(context.Background(), []int{0}, 700); ok {
|
||||
t.Error("no nvidia-smi read as an answer")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
// kvFlushInterval is how often KV cache changes go up the stream. Often
|
||||
// enough that a prefix cached by one request is known before the next few
|
||||
// arrive; seldom enough that a busy engine's churn is sent as its net effect.
|
||||
const kvFlushInterval = 250 * time.Millisecond
|
||||
|
||||
// listenKV starts receiving a shard's KV cache events, when its profile asks
|
||||
// for them and this rank publishes any, and records the port in the spec for
|
||||
// the engine's arguments. port is the one to reuse, or 0 for any.
|
||||
//
|
||||
// A failure is logged and the engine started without events: the router then
|
||||
// routes the replica by tier 1, which is what it did before, rather than the
|
||||
// replica not starting at all for want of an optimisation.
|
||||
func (a *Agent) listenKV(spec *engines.ShardSpec, port int) {
|
||||
if !spec.Profile.KVEvents || spec.Profile.Engine != "vllm" || !spec.Leader() {
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
if sub, ok := a.kvSubs[spec.ShardID]; ok {
|
||||
a.mu.Unlock()
|
||||
// A repeated start command: the shard already has its listener.
|
||||
spec.KVEventsPort = sub.Port()
|
||||
return
|
||||
}
|
||||
a.mu.Unlock()
|
||||
|
||||
sub, err := kvevents.Listen(spec.ReplicaID, port, a.kv, a.log)
|
||||
if err != nil {
|
||||
a.log.Warn("cannot receive kv cache events; the router will route this replica without them",
|
||||
"shard", spec.ShardID, "replica", spec.ReplicaID, "err", err)
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
a.kvSubs[spec.ShardID] = sub
|
||||
a.mu.Unlock()
|
||||
spec.KVEventsPort = sub.Port()
|
||||
}
|
||||
|
||||
// stopKV stops a shard's listener and tells the control plane its replica
|
||||
// holds nothing any more.
|
||||
func (a *Agent) stopKV(shardID, replicaID string) {
|
||||
a.mu.Lock()
|
||||
sub, ok := a.kvSubs[shardID]
|
||||
delete(a.kvSubs, shardID)
|
||||
a.mu.Unlock()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
sub.Close()
|
||||
a.kv.Remove(replicaID)
|
||||
}
|
||||
|
||||
// forwardKV sends KV cache changes up one session's stream until it ends,
|
||||
// beginning with a full copy of every replica: the control plane at the other
|
||||
// end may have just started, or just taken over, and know nothing.
|
||||
func (a *Agent) forwardKV(ctx context.Context, out chan<- *llmgwv1.AgentMessage) {
|
||||
s := a.kv.Open()
|
||||
defer s.Close()
|
||||
t := time.NewTicker(kvFlushInterval)
|
||||
defer t.Stop()
|
||||
for {
|
||||
for _, u := range s.Drain() {
|
||||
select {
|
||||
case out <- agentsvc.KVEventsMessage(u):
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-zeromq/zmq4"
|
||||
"github.com/vmihailenco/msgpack/v5"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
func kvAgent() *Agent {
|
||||
return &Agent{log: slog.New(slog.DiscardHandler), kv: kvevents.NewMirror(),
|
||||
kvSubs: map[string]*kvevents.Subscriber{}}
|
||||
}
|
||||
|
||||
// TestAnEnginesEventsReachTheControlPlaneStream runs the agent's half end to
|
||||
// end: a listener for the shard, a publisher standing in for vLLM connecting
|
||||
// to the port the engine would be given, and the stream to the control plane
|
||||
// carrying the result.
|
||||
func TestAnEnginesEventsReachTheControlPlaneStream(t *testing.T) {
|
||||
a := kvAgent()
|
||||
spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
|
||||
Profile: engines.Profile{Engine: "vllm", KVEvents: true}}
|
||||
a.listenKV(&spec, 0)
|
||||
if spec.KVEventsPort == 0 {
|
||||
t.Fatal("the shard should have been given a port to publish to")
|
||||
}
|
||||
defer a.stopKV(spec.ShardID, spec.ReplicaID)
|
||||
|
||||
pub := zmq4.NewPub(context.Background())
|
||||
defer func() { _ = pub.Close() }()
|
||||
if err := pub.Dial(fmt.Sprintf("tcp://127.0.0.1:%d", spec.KVEventsPort)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tokens := make([]int, 32)
|
||||
for i := range tokens {
|
||||
tokens[i] = 500 + i
|
||||
}
|
||||
hs := kvevents.BlockHashes(tokens, 16, 0)
|
||||
payload, _ := msgpack.Marshal([]any{1.0, []any{map[string]any{
|
||||
"type": "BlockStored", "block_hashes": []any{hs[0], hs[1]}, "parent_block_hash": nil,
|
||||
"token_ids": tokens, "block_size": 16, "lora_id": nil, "medium": "GPU",
|
||||
}}})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
out := make(chan *llmgwv1.AgentMessage, 16)
|
||||
go a.forwardKV(ctx, out)
|
||||
|
||||
deadline := time.After(5 * time.Second)
|
||||
var seq uint64
|
||||
for {
|
||||
// Sent until it arrives: a subscription reaches the publisher after
|
||||
// the connection, and what is published before it is dropped.
|
||||
var s [8]byte
|
||||
binary.BigEndian.PutUint64(s[:], seq)
|
||||
seq++
|
||||
_ = pub.Send(zmq4.NewMsgFrom([]byte("kv-events"), s[:], payload))
|
||||
select {
|
||||
case m := <-out:
|
||||
kv := m.GetKvEvents()
|
||||
if kv == nil || kv.GetReplicaId() != "r1" || kv.GetBlockSize() != 16 || len(kv.GetStored()) != 2 {
|
||||
t.Fatalf("message %v", m)
|
||||
}
|
||||
return
|
||||
case <-deadline:
|
||||
t.Fatal("nothing reached the stream")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAShardWithoutKVEventsGetsNoListener(t *testing.T) {
|
||||
a := kvAgent()
|
||||
spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
|
||||
Profile: engines.Profile{Engine: "vllm"}}
|
||||
a.listenKV(&spec, 0)
|
||||
if spec.KVEventsPort != 0 || len(a.kvSubs) != 0 {
|
||||
t.Fatal("kv_events is opt-in")
|
||||
}
|
||||
}
|
||||
|
||||
func TestARepeatedStartReusesTheShardsListener(t *testing.T) {
|
||||
a := kvAgent()
|
||||
spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
|
||||
Profile: engines.Profile{Engine: "vllm", KVEvents: true}}
|
||||
a.listenKV(&spec, 0)
|
||||
defer a.stopKV(spec.ShardID, spec.ReplicaID)
|
||||
again := spec
|
||||
again.KVEventsPort = 0
|
||||
a.listenKV(&again, 0)
|
||||
if again.KVEventsPort != spec.KVEventsPort || len(a.kvSubs) != 1 {
|
||||
t.Fatal("a start command repeated must not open a second listener")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
|
||||
func TestAShardPlacedOnAMIGInstanceIsLaunchedWithItsUUID(t *testing.T) {
|
||||
// The control plane sends device 1001: card 0's instance 1, by the number
|
||||
// placement gives it. Only this node knows the instance's UUID, which is
|
||||
// the one thing CUDA will open it by.
|
||||
a := &Agent{}
|
||||
a.inv.GPUs = []discovery.GPU{
|
||||
{Index: 1000, UUID: "MIG-aaa", MIG: &discovery.MIGDevice{ParentIndex: 0, Device: 0, Profile: "3g.40gb"}},
|
||||
{Index: 1001, UUID: "MIG-bbb", MIG: &discovery.MIGDevice{ParentIndex: 0, Device: 1, Profile: "3g.40gb"}},
|
||||
{Index: 1, UUID: "GPU-whole"},
|
||||
}
|
||||
|
||||
spec := engines.ShardSpec{GPUIndices: []int{1001}}
|
||||
a.localizeSpec(&spec)
|
||||
if got := engines.NCCLEnv(spec)["CUDA_VISIBLE_DEVICES"]; got != "MIG-bbb" {
|
||||
t.Errorf("CUDA_VISIBLE_DEVICES=%q, want MIG-bbb", got)
|
||||
}
|
||||
|
||||
// A whole card is given by index, as it always was.
|
||||
whole := engines.ShardSpec{GPUIndices: []int{1}}
|
||||
a.localizeSpec(&whole)
|
||||
if whole.GPUUUIDs != nil || engines.NCCLEnv(whole)["CUDA_VISIBLE_DEVICES"] != "1" {
|
||||
t.Errorf("a whole card was given as %v / %q", whole.GPUUUIDs, engines.NCCLEnv(whole)["CUDA_VISIBLE_DEVICES"])
|
||||
}
|
||||
|
||||
// A device this node does not have is not guessed at.
|
||||
gone := engines.ShardSpec{GPUIndices: []int{1002}}
|
||||
a.localizeSpec(&gone)
|
||||
if gone.GPUUUIDs != nil {
|
||||
t.Errorf("an unknown device was named %v", gone.GPUUUIDs)
|
||||
}
|
||||
}
|
||||
@@ -88,7 +88,13 @@ func (a *Agent) setProvision(p *llmgwv1.Provision) {
|
||||
// the installed build is used first — and quietly replace theirs.
|
||||
continue
|
||||
}
|
||||
wants = append(wants, provision.Want{Engine: e.GetEngine(), Version: e.GetVersion(), Variant: e.GetVariant()})
|
||||
w := provision.Want{Engine: e.GetEngine(), Version: e.GetVersion(), Variant: e.GetVariant()}
|
||||
if base := e.GetOfflineBase(); base != "" {
|
||||
// An offline cluster: from its mirror, checked against what the
|
||||
// control plane holds, never the publisher.
|
||||
w.Offline = &provision.Offline{Base: base, SHA256: maps.Clone(e.GetOfflineSha256())}
|
||||
}
|
||||
wants = append(wants, w)
|
||||
}
|
||||
a.mu.Lock()
|
||||
a.provWants = wants
|
||||
|
||||
@@ -20,7 +20,11 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
|
||||
)
|
||||
|
||||
// Want is one engine at one version.
|
||||
@@ -30,6 +34,24 @@ type Want struct {
|
||||
// Variant is a build chosen for this node ("cu129"); empty leaves it to
|
||||
// the node's hardware.
|
||||
Variant string
|
||||
// Offline, when set, is where the engine comes from instead of its
|
||||
// publisher: a cluster with no way to the internet installs from its
|
||||
// mirror's copy of an offline bundle.
|
||||
Offline *Offline
|
||||
}
|
||||
|
||||
// Offline is an engine version's files on a cluster's own mirror.
|
||||
//
|
||||
// The files come from the mirror and their checksums from the control plane,
|
||||
// which read them from the bundle they were imported from. The mirror serves
|
||||
// plain HTTP to every node; a checksum taken from beside the file there would
|
||||
// prove only that the file and the checksum came from the same place.
|
||||
type Offline struct {
|
||||
// Base is where the version's files are: <mirror>/_llmgw/engines/<engine>/<version>.
|
||||
Base string
|
||||
// SHA256 is each file's checksum, by name, as the publisher gave it
|
||||
// when the bundle was made.
|
||||
SHA256 map[string]string
|
||||
}
|
||||
|
||||
// The states an engine's installation can be in.
|
||||
@@ -152,6 +174,16 @@ func (in *Installer) one(ctx context.Context, w Want) Status {
|
||||
}
|
||||
return st
|
||||
}
|
||||
if w.Offline != nil && recipe.Repo == "" {
|
||||
// Said, not attempted: uv would spend minutes failing to reach PyPI,
|
||||
// and the failure would read like a network fault rather than the
|
||||
// design of an offline site.
|
||||
st.State = Unsupported
|
||||
st.Detail = "this cluster installs offline, and " + w.Engine + " comes from PyPI, which an offline bundle " +
|
||||
"does not carry: install it on the node by hand and name it in agent.yaml, or set engines." +
|
||||
w.Engine + ` to "" in server.yaml`
|
||||
return st
|
||||
}
|
||||
|
||||
// Installed before: whichever candidate is, best first, checked that it
|
||||
// still runs and trusted for nothing else.
|
||||
@@ -251,7 +283,17 @@ func (in *Installer) ready(st Status, b Build, bin, version string) Status {
|
||||
//
|
||||
// Into a scratch directory renamed into place at the end, so a failure halfway
|
||||
// leaves nothing that looks installed.
|
||||
func (in *Installer) install(ctx context.Context, r Recipe, w Want, b Build, dir string, st *Status) error {
|
||||
//
|
||||
// Traced as "install engine", a trace of its own: an install is not part of
|
||||
// any one start, and runs once per version rather than once per replica.
|
||||
func (in *Installer) install(ctx context.Context, r Recipe, w Want, b Build, dir string, st *Status) (err error) {
|
||||
ctx, span := obs.Tracer("gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision").Start(ctx, "install engine",
|
||||
trace.WithAttributes(
|
||||
attribute.String("llmgw.engine", w.Engine),
|
||||
attribute.String("llmgw.engine_version", w.Version),
|
||||
attribute.String("llmgw.engine_build", w.Variant),
|
||||
))
|
||||
defer func() { obs.End(span, err) }()
|
||||
if b.Python != "" {
|
||||
return in.installPython(ctx, r, w, b, dir, st)
|
||||
}
|
||||
@@ -259,10 +301,15 @@ func (in *Installer) install(ctx context.Context, r Recipe, w Want, b Build, dir
|
||||
return err
|
||||
}
|
||||
|
||||
in.progress(st, "asking GitHub for the checksums of %s %s", w.Engine, w.Version)
|
||||
sums, err := in.checksums(ctx, r.Repo, w.Version)
|
||||
if err != nil {
|
||||
return err
|
||||
var sums map[string]string
|
||||
if w.Offline != nil {
|
||||
sums = w.Offline.SHA256
|
||||
} else {
|
||||
in.progress(st, "asking GitHub for the checksums of %s %s", w.Engine, w.Version)
|
||||
var err error
|
||||
if sums, err = in.checksums(ctx, r.Repo, w.Version); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
tmp := dir + ".partial"
|
||||
@@ -275,14 +322,22 @@ func (in *Installer) install(ctx context.Context, r Recipe, w Want, b Build, dir
|
||||
got := map[string]string{}
|
||||
for _, asset := range b.Assets {
|
||||
want, ok := sums[asset]
|
||||
if !ok {
|
||||
switch {
|
||||
case !ok && w.Offline != nil:
|
||||
return fmt.Errorf("the offline bundle has no %s for %s %s: make one that has "+
|
||||
"(llmgw bundle create --engine %s=%s) and import it", asset, w.Engine, w.Version, w.Engine, w.Version)
|
||||
case !ok:
|
||||
// Refused rather than installed unchecked: this is a binary that
|
||||
// will run as root on every node.
|
||||
return fmt.Errorf("%s %s publishes no checksum for %s, so it cannot be verified",
|
||||
w.Engine, w.Version, asset)
|
||||
}
|
||||
url := strings.TrimSuffix(in.downloadBase(), "/") + "/" + r.Repo + "/releases/download/" + w.Version + "/" + asset
|
||||
if w.Offline != nil {
|
||||
url = strings.TrimSuffix(w.Offline.Base, "/") + "/" + asset
|
||||
}
|
||||
archive := filepath.Join(tmp, asset+".download")
|
||||
if err := in.download(ctx, r.Repo, w.Version, asset, archive, want, st); err != nil {
|
||||
if err := in.download(ctx, url, asset, archive, want, st); err != nil {
|
||||
return err
|
||||
}
|
||||
in.progress(st, "unpacking %s", asset)
|
||||
@@ -539,6 +594,26 @@ func (in *Installer) has(tool string) bool {
|
||||
return err == nil
|
||||
}
|
||||
|
||||
// Published is the sha256 an engine's publisher gives for each file of a
|
||||
// version, and where each is downloaded from: what an offline bundle is made
|
||||
// of, read the way a node online reads it.
|
||||
func (in *Installer) Published(ctx context.Context, engine, version string) (sums map[string]string, url func(asset string) string, err error) {
|
||||
r, ok := Recipes[engine]
|
||||
if !ok || r.Repo == "" {
|
||||
return nil, nil, fmt.Errorf("%s is not published as release files", engine)
|
||||
}
|
||||
if !validVersion.MatchString(version) {
|
||||
return nil, nil, fmt.Errorf("%q is not a version of %s", version, engine)
|
||||
}
|
||||
sums, err = in.checksums(ctx, r.Repo, version)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return sums, func(asset string) string {
|
||||
return strings.TrimSuffix(in.downloadBase(), "/") + "/" + r.Repo + "/releases/download/" + version + "/" + asset
|
||||
}, nil
|
||||
}
|
||||
|
||||
// checksums reads the sha256 GitHub publishes for each asset of a release.
|
||||
func (in *Installer) checksums(ctx context.Context, repo, tag string) (map[string]string, error) {
|
||||
url := strings.TrimSuffix(in.apiBase(), "/") + "/repos/" + repo + "/releases/tags/" + tag
|
||||
@@ -580,8 +655,7 @@ func (in *Installer) checksums(ctx context.Context, repo, tag string) (map[strin
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (in *Installer) download(ctx context.Context, repo, tag, asset, dst, want string, st *Status) error {
|
||||
url := strings.TrimSuffix(in.downloadBase(), "/") + "/" + repo + "/releases/download/" + tag + "/" + asset
|
||||
func (in *Installer) download(ctx context.Context, url, asset, dst, want string, st *Status) error {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
package provision
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
)
|
||||
|
||||
// offlineMirror is a cluster's mirror holding a bundle's engine files, and
|
||||
// nothing that answers for GitHub.
|
||||
type offlineMirror struct {
|
||||
files map[string][]byte
|
||||
mu sync.Mutex
|
||||
paths []string
|
||||
}
|
||||
|
||||
func (m *offlineMirror) serve(t *testing.T) *httptest.Server {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
m.mu.Lock()
|
||||
m.paths = append(m.paths, r.URL.Path)
|
||||
m.mu.Unlock()
|
||||
b, ok := m.files[r.URL.Path]
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write(b)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestAnOfflineNodeInstallsFromTheMirrorWithChecksumsFromTheControlPlane(t *testing.T) {
|
||||
rel := llamaRelease(t)
|
||||
asset := "llama-b1-bin-ubuntu-x64.tar.gz"
|
||||
m := &offlineMirror{files: map[string][]byte{"/_llmgw/engines/llamacpp/b1/" + asset: rel.assets[asset]}}
|
||||
srv := m.serve(t)
|
||||
sys := &fakeSystem{installed: map[string]bool{"libgomp1": true}}
|
||||
// GitHub is the mirror too, so a request for it would be seen and
|
||||
// answered 404: an offline install must not ask.
|
||||
in, _ := installer(t, srv, sys, ubuntu)
|
||||
|
||||
want := Want{Engine: "llamacpp", Version: "b1", Offline: &Offline{
|
||||
Base: srv.URL + "/_llmgw/engines/llamacpp/b1",
|
||||
SHA256: map[string]string{asset: strings.TrimPrefix(rel.digests[asset], "sha256:")},
|
||||
}}
|
||||
got := in.Converge(t.Context(), []Want{want})
|
||||
if got[0].State != Ready {
|
||||
t.Fatalf("not installed from the mirror: %+v", got[0])
|
||||
}
|
||||
for _, p := range m.paths {
|
||||
if strings.HasPrefix(p, "/repos/") || strings.Contains(p, "/releases/download/") {
|
||||
t.Errorf("an offline install asked for %s", p)
|
||||
}
|
||||
}
|
||||
|
||||
// A file on the mirror that is not what the bundle recorded is refused,
|
||||
// whatever the mirror says of it.
|
||||
m.files["/_llmgw/engines/llamacpp/b2/llama-b2-bin-ubuntu-x64.tar.gz"] = []byte("replaced on the mirror")
|
||||
got = in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b2", Offline: &Offline{
|
||||
Base: srv.URL + "/_llmgw/engines/llamacpp/b2",
|
||||
SHA256: map[string]string{"llama-b2-bin-ubuntu-x64.tar.gz": strings.Repeat("a", 64)},
|
||||
}}})
|
||||
if got[0].State != Failed || !strings.Contains(got[0].Detail, "checksum") {
|
||||
t.Errorf("a file not matching the bundle's checksum was not refused: %+v", got[0])
|
||||
}
|
||||
|
||||
// One the bundle has no checksum for is not fetched at all, and the
|
||||
// operator is told which bundle would have it.
|
||||
got = in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b3", Offline: &Offline{
|
||||
Base: srv.URL + "/_llmgw/engines/llamacpp/b3",
|
||||
}}})
|
||||
if got[0].State != Failed || !strings.Contains(got[0].Detail, "bundle create --engine llamacpp=b3") {
|
||||
t.Errorf("a version missing from the bundle did not say how to add it: %+v", got[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnOfflineNodeIsToldPythonEnginesComeFromPyPI(t *testing.T) {
|
||||
gpu := ubuntu
|
||||
gpu.NVIDIADriver, gpu.MinComputeCapability = "580.65", 8.9
|
||||
in := &Installer{Root: t.TempDir(), Host: gpu, Run: (&fakeSystem{}).run}
|
||||
got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0", Offline: &Offline{Base: "http://mirror"}}})
|
||||
if got[0].State != Unsupported || !strings.Contains(got[0].Detail, "PyPI") {
|
||||
t.Fatalf("vLLM offline: %+v", got[0])
|
||||
}
|
||||
|
||||
// A node it is not meant for still says nothing is wrong.
|
||||
in.Host = ubuntu
|
||||
got = in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0", Offline: &Offline{Base: "http://mirror"}}})
|
||||
if got[0].State != Skipped {
|
||||
t.Errorf("vLLM offline on a CPU node: %+v", got[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestABundleCarriesEveryBuildTheNodesMightChoose(t *testing.T) {
|
||||
variants := func(arch string) []string {
|
||||
bs, err := Builds("llamacpp", "b1", arch)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out []string
|
||||
for _, b := range bs {
|
||||
out = append(out, b.Variant)
|
||||
}
|
||||
return out
|
||||
}
|
||||
if got := variants("amd64"); !slices.Equal(got, []string{"cpu", "cuda-13.3", "cuda-12.8", "vulkan"}) {
|
||||
t.Errorf("amd64 builds: %v", got)
|
||||
}
|
||||
if got := variants("arm64"); !slices.Equal(got, []string{"cpu", "cuda-13.3", "vulkan"}) {
|
||||
t.Errorf("arm64 builds: %v", got)
|
||||
}
|
||||
|
||||
// Each one a node would choose is among them, file for file.
|
||||
bundled := map[string][]string{}
|
||||
bs, _ := Builds("llamacpp", "b1", "amd64")
|
||||
for _, b := range bs {
|
||||
bundled[b.Variant] = b.Assets
|
||||
}
|
||||
for _, h := range []discovery.Host{
|
||||
ubuntu,
|
||||
{Arch: "amd64", NVIDIADriver: "550.1", MinComputeCapability: 8.6},
|
||||
{Arch: "amd64", NVIDIADriver: "470.1", MinComputeCapability: 6.1},
|
||||
} {
|
||||
chosen, err := llamaCppCandidates("b1", "", h)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, b := range chosen {
|
||||
if !slices.Equal(bundled[b.Variant], b.Assets) {
|
||||
t.Errorf("%+v chooses %s %v, which the bundle carries as %v", h, b.Variant, b.Assets, bundled[b.Variant])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := Builds("vllm", "0.30.0", "amd64"); err == nil || !strings.Contains(err.Error(), "PyPI") {
|
||||
t.Errorf("a Python engine was offered for a bundle: %v", err)
|
||||
}
|
||||
if _, err := Builds("llamacpp", "../b1", "amd64"); err == nil {
|
||||
t.Error("a version that is a path was accepted")
|
||||
}
|
||||
}
|
||||
@@ -8,6 +8,8 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
@@ -33,6 +35,10 @@ const pythonInstallTimeout = time.Hour
|
||||
// lockName is the resolved requirements, kept in the environment.
|
||||
const lockName = "llmgw-requirements.txt"
|
||||
|
||||
// pylockName is the lock of a build with a second index: PEP 751's, with
|
||||
// every file's address.
|
||||
const pylockName = "pylock.llmgw.toml"
|
||||
|
||||
func (in *Installer) installPython(ctx context.Context, r Recipe, w Want, b Build, dir string, st *Status) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, pythonInstallTimeout)
|
||||
defer cancel()
|
||||
@@ -113,28 +119,16 @@ func (in *Installer) installPython(ctx context.Context, r Recipe, w Want, b Buil
|
||||
b.Wheel.Tag + "/" + asset
|
||||
reqs = append([]string{b.Wheel.Project + " @ " + url + "#sha256=" + sum}, reqs...)
|
||||
}
|
||||
// A second index, for the torch the engine was built on. Every package
|
||||
// is taken from whichever of the two has the version resolved — both are
|
||||
// the publishers' own — and each file is still pinned by its hash.
|
||||
var index []string
|
||||
if b.ExtraIndex != "" {
|
||||
index = []string{"--extra-index-url", b.ExtraIndex, "--index-strategy", "unsafe-best-match"}
|
||||
}
|
||||
|
||||
in.progress(st, "resolving %s against PyPI as of %s", strings.Join(reqs, " "), asOf[:10])
|
||||
lock := filepath.Join(tmp, lockName)
|
||||
args := append([]string{"pip", "compile", "--python", python, "--exclude-newer", asOf,
|
||||
"--generate-hashes", "--quiet", "--output-file", lock}, index...)
|
||||
args = append(args, "-")
|
||||
if out, err := in.runStdin(ctx, env, strings.Join(reqs, "\n")+"\n", uv.Binary, args...); err != nil {
|
||||
return fmt.Errorf("resolving %s: %v: %s", strings.Join(reqs, " "), err, tail(out))
|
||||
lock, err := in.resolve(ctx, env, uv.Binary, python, tmp, asOf, reqs, b)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
in.progress(st, "downloading and installing %s and its dependencies: several gigabytes, which takes minutes",
|
||||
strings.Join(reqs, " "))
|
||||
install := append([]string{"pip", "install", "--python", python,
|
||||
"--require-hashes", "--no-deps", "-r", lock}, index...)
|
||||
if out, err := in.run(ctx, env, uv.Binary, install...); err != nil {
|
||||
if out, err := in.run(ctx, env, uv.Binary, "pip", "install", "--python", python,
|
||||
"--require-hashes", "--no-deps", "-r", lock); err != nil {
|
||||
return fmt.Errorf("installing %s: %v: %s", strings.Join(reqs, " "), err, tail(out))
|
||||
}
|
||||
|
||||
@@ -168,6 +162,63 @@ func (in *Installer) installPython(ctx context.Context, r Recipe, w Want, b Buil
|
||||
return nil
|
||||
}
|
||||
|
||||
// resolve pins every package the requirements need, each file by its hash,
|
||||
// into a lock in the environment, and returns its path.
|
||||
//
|
||||
// With a second index — PyTorch's, for the torch a build was made against —
|
||||
// the packages it is for come from it and nowhere else, and every other
|
||||
// package from PyPI alone. Letting both indexes offer every package is how
|
||||
// torchtriton was taken in 2022: a name one index had, registered on the
|
||||
// other with a higher version. So the second index is explicit, named only
|
||||
// for its packages, and the lock is PEP 751's, which records each file's
|
||||
// address as well as its hash: installing it asks no index anything.
|
||||
func (in *Installer) resolve(ctx context.Context, env []string, uv, python, tmp, asOf string,
|
||||
reqs []string, b Build,
|
||||
) (string, error) {
|
||||
if b.ExtraIndex == "" {
|
||||
lock := filepath.Join(tmp, lockName)
|
||||
if out, err := in.runStdin(ctx, env, strings.Join(reqs, "\n")+"\n", uv, "pip", "compile",
|
||||
"--python", python, "--exclude-newer", asOf, "--generate-hashes", "--quiet",
|
||||
"--output-file", lock, "-"); err != nil {
|
||||
return "", fmt.Errorf("resolving %s: %v: %s", strings.Join(reqs, " "), err, tail(out))
|
||||
}
|
||||
return lock, nil
|
||||
}
|
||||
|
||||
project := tmp + ".resolve"
|
||||
if err := os.MkdirAll(project, 0o755); err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer os.RemoveAll(project) //nolint:errcheck // scratch
|
||||
if err := os.WriteFile(filepath.Join(project, "pyproject.toml"), []byte(resolveProject(reqs, b)), 0o644); err != nil { //nolint:gosec // not secret
|
||||
return "", err
|
||||
}
|
||||
lock := filepath.Join(tmp, pylockName)
|
||||
if out, err := in.run(ctx, env, uv, "pip", "compile", "--python", python, "--exclude-newer", asOf,
|
||||
"--format", "pylock.toml", "--quiet", "--output-file", lock,
|
||||
filepath.Join(project, "pyproject.toml")); err != nil {
|
||||
return "", fmt.Errorf("resolving %s: %v: %s", strings.Join(reqs, " "), err, tail(out))
|
||||
}
|
||||
return lock, nil
|
||||
}
|
||||
|
||||
// resolveProject is the pyproject.toml a build with a second index is
|
||||
// resolved from. Its packages are direct requirements too: an index is
|
||||
// given to a package only where the project names it.
|
||||
func resolveProject(reqs []string, b Build) string {
|
||||
var p strings.Builder
|
||||
p.WriteString("[project]\nname = \"llmgw-engine\"\nversion = \"0\"\ndependencies = [\n")
|
||||
for _, r := range append(slices.Clone(reqs), b.IndexPackages...) {
|
||||
fmt.Fprintf(&p, " %s,\n", strconv.Quote(r))
|
||||
}
|
||||
p.WriteString("]\n\n[[tool.uv.index]]\nname = \"extra\"\n")
|
||||
fmt.Fprintf(&p, "url = %s\nexplicit = true\n\n[tool.uv.sources]\n", strconv.Quote(b.ExtraIndex))
|
||||
for _, name := range b.IndexPackages {
|
||||
fmt.Fprintf(&p, "%s = { index = \"extra\" }\n", name)
|
||||
}
|
||||
return p.String()
|
||||
}
|
||||
|
||||
// published is when PyPI received a release's last file.
|
||||
func (in *Installer) published(ctx context.Context, project, version string) (time.Time, error) {
|
||||
url := strings.TrimSuffix(in.pypiBase(), "/") + "/pypi/" + project + "/" + version + "/json"
|
||||
|
||||
@@ -40,6 +40,14 @@ func (f *fakePython) run(ctx context.Context, env []string, name string, args ..
|
||||
return []byte("uv 0.12.20\n"), nil
|
||||
case strings.HasSuffix(name, "/uv") && args[0] == "venv":
|
||||
return nil, os.MkdirAll(filepath.Join(args[len(args)-1], "bin"), 0o755)
|
||||
case strings.HasSuffix(name, "/uv") && args[0] == "pip" && args[1] == "compile":
|
||||
// From a pyproject: kept in the record, as the request was.
|
||||
project, _ := os.ReadFile(args[len(args)-1])
|
||||
f.pmu.Lock()
|
||||
f.commands = append(f.commands, "pyproject: "+string(project))
|
||||
f.pmu.Unlock()
|
||||
out := args[slices.Index(args, "--output-file")+1]
|
||||
return nil, os.WriteFile(out, []byte("lock-version = \"1.0\"\n"), 0o644)
|
||||
case strings.HasSuffix(name, "/uv") && args[0] == "pip" && args[1] == "install":
|
||||
if f.failPip {
|
||||
return []byte("error: Hash mismatch for torch\n"), errExit
|
||||
@@ -220,8 +228,12 @@ func TestAnOlderDriverGetsTheSameVLLMBuiltForCUDA12(t *testing.T) {
|
||||
// Pinned to the digest GitHub publishes for the wheel.
|
||||
"vllm @ " + in.DownloadBase + "/vllm-project/vllm/releases/download/v0.30.0/" + wheel + "#sha256=" +
|
||||
strings.TrimPrefix(gh.digests[wheel], "sha256:"),
|
||||
"--extra-index-url https://download.pytorch.org/whl/cu129 --index-strategy unsafe-best-match -",
|
||||
"--require-hashes --no-deps -r",
|
||||
// The torch packages from PyTorch's index and only those; the lock
|
||||
// records every file's address.
|
||||
"--format pylock.toml",
|
||||
`url = "https://download.pytorch.org/whl/cu129"` + "\nexplicit = true",
|
||||
`torch = { index = "extra" }`, `torchvision = { index = "extra" }`,
|
||||
"--require-hashes --no-deps -r " + filepath.Join(in.Root, "vllm", "0.30.0-cu129.partial", pylockName),
|
||||
// Computing on the card, not only importing.
|
||||
"/bin/python -c import torch",
|
||||
} {
|
||||
@@ -230,8 +242,8 @@ func TestAnOlderDriverGetsTheSameVLLMBuiltForCUDA12(t *testing.T) {
|
||||
}
|
||||
}
|
||||
for _, c := range sys.commands {
|
||||
if strings.Contains(c, "pip install") && !strings.Contains(c, "download.pytorch.org/whl/cu129") {
|
||||
t.Errorf("installed without the index its torch comes from: %s", c)
|
||||
if strings.Contains(c, "unsafe-best-match") || strings.Contains(c, "--extra-index-url") {
|
||||
t.Errorf("every index offered every package: %s", c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,10 +50,12 @@ type Build struct {
|
||||
GPUBackend string
|
||||
// Wheel, when set, is the engine itself as a wheel from a GitHub
|
||||
// release rather than from PyPI, checked against the digest GitHub
|
||||
// publishes for it. ExtraIndex is a package index besides PyPI that the
|
||||
// rest resolves from too: the CUDA 12.9 torch the wheel was built on.
|
||||
Wheel *Wheel
|
||||
ExtraIndex string
|
||||
// publishes for it. ExtraIndex is a package index besides PyPI, for
|
||||
// IndexPackages only: the CUDA 12.9 torch the wheel was built on. Every
|
||||
// other package comes from PyPI alone.
|
||||
Wheel *Wheel
|
||||
ExtraIndex string
|
||||
IndexPackages []string
|
||||
// GPUCheck, on a node with an NVIDIA GPU, has the installed environment
|
||||
// compute something on one before the build counts as installed. A build
|
||||
// with no kernels for the card — a GPU generation its torch dropped —
|
||||
@@ -118,6 +120,49 @@ func Installable() []string {
|
||||
return out
|
||||
}
|
||||
|
||||
// Builds are every build of an engine version published for an
|
||||
// architecture: what an offline bundle carries, for nodes it has not seen.
|
||||
// Read from the same table the nodes choose from, asked as a machine with
|
||||
// no GPU and as one with any NVIDIA GPU, so a bundle and the nodes that
|
||||
// install from it cannot disagree about which files a build is. Archive
|
||||
// engines only: a Python engine is a PyPI environment, which no bundle
|
||||
// carries.
|
||||
func Builds(engine, version, arch string) ([]Build, error) {
|
||||
r, ok := Recipes[engine]
|
||||
switch {
|
||||
case !ok:
|
||||
return nil, fmt.Errorf("this agent does not know how to install %s", engine)
|
||||
case r.Repo == "":
|
||||
return nil, fmt.Errorf("%s is installed from PyPI, which an offline bundle does not carry", engine)
|
||||
case !validVersion.MatchString(version):
|
||||
return nil, fmt.Errorf("%q is not a version of %s", version, engine)
|
||||
}
|
||||
var out []Build
|
||||
seen := map[string]bool{}
|
||||
var errs []string
|
||||
for _, h := range []discovery.Host{
|
||||
{Arch: arch},
|
||||
// Any card and the newest driver: every GPU build is a candidate.
|
||||
{Arch: arch, NVIDIADriver: "9999"},
|
||||
} {
|
||||
bs, err := r.Candidates(version, "", h)
|
||||
if err != nil {
|
||||
errs = append(errs, err.Error())
|
||||
continue
|
||||
}
|
||||
for _, b := range bs {
|
||||
if !seen[b.Variant] {
|
||||
seen[b.Variant] = true
|
||||
out = append(out, b)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil, fmt.Errorf("no build of %s %s for %s: %s", engine, version, arch, strings.Join(errs, "; "))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// UVVersion is the uv that installs the Python engines.
|
||||
const UVVersion = "0.12.20"
|
||||
|
||||
@@ -234,6 +279,10 @@ var VLLMVariants = []string{"cu13", "cu129"}
|
||||
// torch is built for CUDA 13; the cu129 vLLM wheel needs this one.
|
||||
const pytorchCU129 = "https://download.pytorch.org/whl/cu129"
|
||||
|
||||
// pytorchPackages are what is taken from PyTorch's index: the builds for one
|
||||
// CUDA version, which PyPI does not have.
|
||||
var pytorchPackages = []string{"torch", "torchvision", "torchaudio"}
|
||||
|
||||
// vllmCUDA13Since is the first vLLM whose PyPI build is CUDA 13, read from
|
||||
// its releases: 0.19.1 publishes "+cu130" beside a CUDA 12 default, 0.20.0
|
||||
// "+cu129" beside a CUDA 13 one.
|
||||
@@ -332,6 +381,7 @@ func vllmCandidates(version, variant string, h discovery.Host) ([]Build, error)
|
||||
Prefix: "vllm-" + v + "+cu129-", Suffix: "_" + map[string]string{"amd64": "x86_64", "arm64": "aarch64"}[h.Arch] + ".whl",
|
||||
}
|
||||
cu129.ExtraIndex = pytorchCU129
|
||||
cu129.IndexPackages = pytorchPackages
|
||||
return pick("vLLM "+v, variant, []candidate{
|
||||
{pypi("cu13"), gpuNeeds{minDriverForCUDA13, minCC}, "vLLM " + v + " built for CUDA 13"},
|
||||
{cu129, gpuNeeds{minDriverForCUDA12, minCC}, "vLLM " + v + " built for CUDA 12.9"},
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
// The cluster's settings for this agent (BACKLOG E3): what agent.yaml can
|
||||
// say but the agent's id, control planes, address and certificate, kept by
|
||||
// the control planes and sent over the stream. Kept here beside agent.yaml,
|
||||
// which then holds only what the agent needs to start; at start the two are
|
||||
// read together, these over the file.
|
||||
|
||||
// SettingsFile keeps the settings the control plane last sent.
|
||||
var SettingsFile = "/var/lib/llmgw-agent/agent-settings.yaml"
|
||||
|
||||
// ReadSettings is what SettingsFile holds, nil when nothing.
|
||||
func ReadSettings() []byte {
|
||||
b, err := os.ReadFile(SettingsFile)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// SetSettingsCheck says how to tell whether settings would load beside this
|
||||
// agent's own file; settings that would not are not taken.
|
||||
func (a *Agent) SetSettingsCheck(f func([]byte) error) { a.settingsCheck = f }
|
||||
|
||||
// applySettings takes the settings the control plane sent: nothing when they
|
||||
// are what this agent runs on, and otherwise kept and started again on.
|
||||
func (a *Agent) applySettings(m *llmgwv1.AgentSettings) {
|
||||
next := m.GetYaml()
|
||||
if next == a.cfg.Settings {
|
||||
return
|
||||
}
|
||||
if a.settingsCheck != nil {
|
||||
if err := a.settingsCheck([]byte(next)); err != nil {
|
||||
// Kept running on what it has: settings that will not load
|
||||
// here must not take the node out of service.
|
||||
a.log.Error("the cluster's settings for this agent do not load beside its agent.yaml; running on what it has",
|
||||
"err", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(SettingsFile), 0o755); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
tmp := SettingsFile + ".new"
|
||||
if err := os.WriteFile(tmp, []byte(next), 0o600); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
if err := os.Rename(tmp, SettingsFile); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
a.log.Info("the cluster's settings for this agent changed; starting again on them, the engines keep serving")
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
exitForUpgrade()
|
||||
}
|
||||
|
||||
// WatchFile starts the agent again when its configuration file changes and
|
||||
// still loads: agent.yaml is read only at start, and a change to it that
|
||||
// waited for somebody to restart the agent was a change nobody saw take
|
||||
// effect.
|
||||
func (a *Agent) WatchFile(ctx context.Context, path string, loads func() error) {
|
||||
if path == "" {
|
||||
return
|
||||
}
|
||||
sum := func() [32]byte {
|
||||
b, _ := os.ReadFile(path) //nolint:gosec // the agent's own configuration
|
||||
return sha256.Sum256(b)
|
||||
}
|
||||
was := sum()
|
||||
t := time.NewTicker(watchEvery)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
now := sum()
|
||||
if now == was {
|
||||
continue
|
||||
}
|
||||
if err := loads(); err != nil {
|
||||
// Once per change: an editor saving halfway must not stop the
|
||||
// agent, and a broken file is reported, not restarted into.
|
||||
a.log.Error("agent.yaml changed and does not load; running on what it had", "path", path, "err", err)
|
||||
was = now
|
||||
continue
|
||||
}
|
||||
a.log.Info("agent.yaml changed; starting again on it, the engines keep serving", "path", path)
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
exitForUpgrade()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// watchEvery is how often agent.yaml is looked at.
|
||||
var watchEvery = 10 * time.Second
|
||||
@@ -0,0 +1,88 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
func settingsAgent(t *testing.T, running string) (*Agent, chan struct{}) {
|
||||
t.Helper()
|
||||
oldFile, oldExit := SettingsFile, exitForUpgrade
|
||||
SettingsFile = filepath.Join(t.TempDir(), "agent-settings.yaml")
|
||||
exited := make(chan struct{}, 1)
|
||||
exitForUpgrade = func() { exited <- struct{}{} }
|
||||
t.Cleanup(func() { SettingsFile, exitForUpgrade = oldFile, oldExit })
|
||||
cfg := Config{Server: "127.0.0.1:7443", NodeID: "n1", WeightsCache: t.TempDir(), EnginesDir: t.TempDir(), Settings: running}
|
||||
cfg.Superv.WorkDir = t.TempDir()
|
||||
a, err := New(cfg, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return a, exited
|
||||
}
|
||||
|
||||
// exitedSoon is whether the agent exited within a second: it waits a
|
||||
// moment for its log line to leave first.
|
||||
func exitedSoon(ch chan struct{}) bool {
|
||||
select {
|
||||
case <-ch:
|
||||
return true
|
||||
case <-time.After(time.Second):
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAgentOnItsSettingsStaysAndOneOffThemStartsAgain(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "heartbeat_interval: 5s\n")
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "heartbeat_interval: 5s\n"})
|
||||
if exitedSoon(exited) || ReadSettings() != nil {
|
||||
t.Error("settings it already runs on restarted it")
|
||||
}
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "heartbeat_interval: 2s\n"})
|
||||
if !exitedSoon(exited) || string(ReadSettings()) != "heartbeat_interval: 2s\n" {
|
||||
t.Errorf("new settings: kept %q", ReadSettings())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsThatWouldNotLoadAreNotTaken(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "")
|
||||
a.SetSettingsCheck(func([]byte) error { return errors.New("fabric ib-a lists no HCAs") })
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "fabrics: [{name: ib-a}]\n"})
|
||||
if exitedSoon(exited) || ReadSettings() != nil {
|
||||
t.Error("settings that do not load were taken")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAChangedAgentYamlThatLoadsStartsTheAgentAgain(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "")
|
||||
old := watchEvery
|
||||
watchEvery = 10 * time.Millisecond
|
||||
t.Cleanup(func() { watchEvery = old })
|
||||
p := filepath.Join(t.TempDir(), "agent.yaml")
|
||||
_ = os.WriteFile(p, []byte("node_id: n1\n"), 0o600)
|
||||
var loads atomic.Bool
|
||||
check := func() error {
|
||||
if loads.Load() {
|
||||
return nil
|
||||
}
|
||||
return errors.New("broken")
|
||||
}
|
||||
go a.WatchFile(t.Context(), p, check)
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
// Saved broken: reported, not restarted into.
|
||||
_ = os.WriteFile(p, []byte("node_id: [\n"), 0o600)
|
||||
if exitedSoon(exited) {
|
||||
t.Fatal("a broken agent.yaml restarted the agent")
|
||||
}
|
||||
loads.Store(true)
|
||||
_ = os.WriteFile(p, []byte("node_id: n1\nlabels: {a: b}\n"), 0o600)
|
||||
if !exitedSoon(exited) {
|
||||
t.Error("a changed agent.yaml did not restart the agent")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"os/user"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// Engines run as their own user, not as the agent's root: an engine reads a
|
||||
// model's files and answers requests, and a model's files are whatever a
|
||||
// registry published. Root in that process would be root on the node for
|
||||
// anything that can make an engine misbehave — a malformed file, a template,
|
||||
// a flag nobody thought of.
|
||||
//
|
||||
// The user is DefaultRunAs, made by `llmgw node add`, in the video and render
|
||||
// groups for the GPU devices. A node without it keeps running engines as
|
||||
// root, and says so, so an existing fleet is not broken by an upgrade.
|
||||
|
||||
// DefaultRunAs is the user engines run as.
|
||||
const DefaultRunAs = "llmgw-engine"
|
||||
|
||||
// runAs is how an engine process is started as another user: its
|
||||
// credential, a working directory it owns, and the environment that points
|
||||
// its caches there.
|
||||
type runAs struct {
|
||||
name string
|
||||
cred *syscall.Credential
|
||||
home string
|
||||
}
|
||||
|
||||
// lookupUser is user.Lookup, replaced by tests.
|
||||
var lookupUser = user.Lookup
|
||||
|
||||
// resolveRunAs finds the user engines run as. nil means as this process —
|
||||
// when it is not root, there is nobody to drop to; when the user does not
|
||||
// exist, why is returned.
|
||||
func resolveRunAs(name, workDir string) (*runAs, string) {
|
||||
if name == "" || name == "root" || os.Geteuid() != 0 {
|
||||
return nil, ""
|
||||
}
|
||||
u, err := lookupUser(name)
|
||||
if err != nil {
|
||||
return nil, fmt.Sprintf("there is no user %s, so engines run as root; `llmgw node add` makes it, "+
|
||||
"or: useradd --system --shell /usr/sbin/nologin %s && usermod -aG video,render %s", name, name, name)
|
||||
}
|
||||
r, err := credentialFor(u, workDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Sprintf("engines run as root: %v", err)
|
||||
}
|
||||
return r, ""
|
||||
}
|
||||
|
||||
func credentialFor(u *user.User, workDir string) (*runAs, error) {
|
||||
uid, err1 := strconv.ParseUint(u.Uid, 10, 32)
|
||||
gid, err2 := strconv.ParseUint(u.Gid, 10, 32)
|
||||
if err1 != nil || err2 != nil {
|
||||
return nil, fmt.Errorf("user %s has ids %s:%s", u.Username, u.Uid, u.Gid)
|
||||
}
|
||||
var groups []uint32
|
||||
if ids, err := u.GroupIds(); err == nil {
|
||||
for _, g := range ids {
|
||||
if n, err := strconv.ParseUint(g, 10, 32); err == nil {
|
||||
groups = append(groups, uint32(n))
|
||||
}
|
||||
}
|
||||
}
|
||||
// A home of its own under the work directory, which is root's: caches
|
||||
// (Triton, torch, Hugging Face) are written there and nowhere else.
|
||||
home := filepath.Join(workDir, "engine-home")
|
||||
if err := os.MkdirAll(home, 0o700); err != nil {
|
||||
return nil, fmt.Errorf("making %s: %w", home, err)
|
||||
}
|
||||
if err := os.Chown(home, int(uid), int(gid)); err != nil {
|
||||
return nil, fmt.Errorf("giving %s to %s: %w", home, u.Username, err)
|
||||
}
|
||||
// The work directory is 0750 root's; the engine has to pass through it
|
||||
// to its home.
|
||||
if err := os.Chmod(workDir, 0o751); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &runAs{name: u.Username, home: home,
|
||||
cred: &syscall.Credential{Uid: uint32(uid), Gid: uint32(gid), Groups: groups}}, nil
|
||||
}
|
||||
|
||||
// env is what is added to an engine's environment when it runs as r: the
|
||||
// last value of a key wins in exec.
|
||||
func (r *runAs) env() []string {
|
||||
return []string{"HOME=" + r.home, "XDG_CACHE_HOME=" + filepath.Join(r.home, ".cache"), "USER=" + r.name, "LOGNAME=" + r.name}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"os"
|
||||
"os/exec"
|
||||
"os/user"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
"syscall"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAnEngineIsGivenItsUsersCredentialHomeAndCaches(t *testing.T) {
|
||||
me, err := user.Current()
|
||||
if err != nil {
|
||||
t.Skip(err)
|
||||
}
|
||||
work := filepath.Join(t.TempDir(), "agent")
|
||||
if err := os.MkdirAll(work, 0o750); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
r, err := credentialFor(me, work)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
uid, _ := strconv.Atoi(me.Uid)
|
||||
if int(r.cred.Uid) != uid || r.home != filepath.Join(work, "engine-home") {
|
||||
t.Fatalf("%+v", r)
|
||||
}
|
||||
if st, err := os.Stat(r.home); err != nil || st.Mode().Perm() != 0o700 {
|
||||
t.Errorf("home: %v %v", st.Mode(), err)
|
||||
}
|
||||
s := &Supervisor{cfg: Config{WorkDir: work}, as: r}
|
||||
c := exec.Command("true")
|
||||
c.Env = []string{"HOME=/root", "PATH=/bin"}
|
||||
c.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
s.dropTo(c)
|
||||
if c.SysProcAttr.Credential == nil || c.Dir != r.home || !slices.Contains(c.Env, "HOME="+r.home) {
|
||||
t.Errorf("dir %s env %v cred %v", c.Dir, c.Env, c.SysProcAttr.Credential)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotRootOrNoSuchUserLeavesEnginesAsTheAgent(t *testing.T) {
|
||||
if os.Geteuid() == 0 {
|
||||
t.Skip("run as root")
|
||||
}
|
||||
if r, why := resolveRunAs(DefaultRunAs, t.TempDir()); r != nil || why != "" {
|
||||
t.Errorf("not root: %v %q", r, why)
|
||||
}
|
||||
}
|
||||
Loaded 100 of 770 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user