Compare commits
No files matched your search
@@ -0,0 +1,152 @@
|
||||
# Attaches binaries to a release when a version tag is pushed.
|
||||
#
|
||||
# Without this a release is a git tag, and installing llmgw on a node means
|
||||
# having a Go toolchain and the source. The onboarding script copies a binary
|
||||
# from wherever it is run, which works and quietly makes whoever runs it the
|
||||
# build server.
|
||||
name: Release
|
||||
|
||||
on:
|
||||
push:
|
||||
tags: ["v*"]
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
# The built-in token every job is given carries whatever the instance's
|
||||
# default mode allows, and under the restricted default that is read-only
|
||||
# on releases — so the build would succeed and the publish would be
|
||||
# refused. `contents` covers code and releases together, which is what
|
||||
# creating a release needs: `releases: write` alone was refused, because
|
||||
# creating one can create its tag.
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
# Tags are what the version string is built from, and a shallow clone
|
||||
# has none of them.
|
||||
fetch-depth: 0
|
||||
|
||||
- uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version-file: go.mod
|
||||
cache: true
|
||||
|
||||
- name: Build
|
||||
run: |
|
||||
set -euo pipefail
|
||||
version="${GITHUB_REF_NAME:-$(git describe --tags --always)}"
|
||||
mkdir -p dist
|
||||
|
||||
# CGO_ENABLED=0 is the whole point: the binary has to run on a node
|
||||
# whose libc is older than the builder's, and llmgw has no C
|
||||
# dependency — modernc.org/sqlite is pure Go precisely so this stays
|
||||
# true.
|
||||
for target in linux/amd64 linux/arm64 darwin/amd64 darwin/arm64; do
|
||||
os="${target%/*}"; arch="${target#*/}"
|
||||
name="llmgw-${version}-${os}-${arch}"
|
||||
echo "building $name"
|
||||
CGO_ENABLED=0 GOOS="$os" GOARCH="$arch" go build \
|
||||
-trimpath \
|
||||
-ldflags "-s -w -X gitea.mynodes.xyz/ilker/llmgw/internal/version.Version=${version}" \
|
||||
-o "dist/${name}" ./cmd/llmgw
|
||||
(cd dist && sha256sum "${name}" >> checksums.txt)
|
||||
done
|
||||
|
||||
# An NVML build for nodes with NVLink, which the portable binary
|
||||
# cannot serve.
|
||||
#
|
||||
# NVML is the only source of NVLink topology, its bindings need cgo,
|
||||
# and cgo is exactly what CGO_ENABLED=0 above exists to avoid. So the
|
||||
# portable binary silently falls back to taking the lowest free GPU
|
||||
# indices — right on an SXM box with NVSwitch, and wrong on a PCIe
|
||||
# box where NVLink bridges pair specific cards: a tensor group
|
||||
# straddling two bridges runs several times slower and nothing says
|
||||
# so. Shipping only the portable build meant nobody could have that
|
||||
# topology without building it themselves.
|
||||
#
|
||||
# Dynamically linked against the host's libc by necessity, so it is
|
||||
# published beside the portable one rather than instead of it.
|
||||
if command -v gcc >/dev/null; then
|
||||
nvml_name="llmgw-${version}-linux-amd64-nvml"
|
||||
CGO_ENABLED=1 GOOS=linux GOARCH=amd64 go build \
|
||||
-tags nvml -trimpath \
|
||||
-ldflags "-s -w -X gitea.mynodes.xyz/ilker/llmgw/internal/version.Version=${version}" \
|
||||
-o "dist/${nvml_name}" ./cmd/llmgw
|
||||
(cd dist && sha256sum "${nvml_name}" >> checksums.txt)
|
||||
else
|
||||
echo "no C toolchain; skipping the nvml build" >&2
|
||||
fi
|
||||
|
||||
# Proof the artifact runs at all before it is published. It cannot
|
||||
# test the cross-compiled ones, which is why the native build is
|
||||
# checked rather than assumed.
|
||||
./dist/llmgw-${version}-linux-amd64 --version
|
||||
cat dist/checksums.txt
|
||||
|
||||
- name: Attach to the release
|
||||
env:
|
||||
TOKEN: ${{ secrets.GITEA_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
# Gitea gives every job a GITEA_TOKEN, so this is not the usual
|
||||
# "you forgot to add a secret" — an empty one means the job was denied
|
||||
# a token, which is an instance setting rather than a repository one.
|
||||
if [ -z "${TOKEN:-}" ]; then
|
||||
echo "The job was given no GITEA_TOKEN, so nothing can be uploaded."
|
||||
echo "Check Settings -> Actions -> General; the binaries are in dist/."
|
||||
exit 1
|
||||
fi
|
||||
# workflow_dispatch can be fired from a branch, where GITHUB_REF_NAME
|
||||
# is the branch name. Uploading then would create a release called
|
||||
# "main". The binaries are still built and kept as a check that the
|
||||
# release build works; only the publishing is a tag's business.
|
||||
if [ "${GITHUB_REF_TYPE:-}" != "tag" ]; then
|
||||
echo "not a tag (ref=${GITHUB_REF_NAME}); built, not published"
|
||||
exit 0
|
||||
fi
|
||||
tag="${GITHUB_REF_NAME}"
|
||||
api="${GITHUB_SERVER_URL}/api/v1/repos/${GITHUB_REPOSITORY}"
|
||||
|
||||
# Gitea materialises a release row for every pushed tag, so this
|
||||
# nearly always finds one and the create below is for a tag pushed
|
||||
# some other way.
|
||||
#
|
||||
# The id is read with grep rather than sed because sed's .* is greedy:
|
||||
# "id" appears again inside the nested author object, so the old
|
||||
# pattern took the author's id when it matched at all — and it did not
|
||||
# match "id": 13 with a space, which is what actually happened. The
|
||||
# first "id" in a release object is the release's own.
|
||||
release_id() {
|
||||
curl -sS -H "Authorization: token $TOKEN" "$api/releases/tags/$tag" \
|
||||
| grep -o '"id"[[:space:]]*:[[:space:]]*[0-9]*' \
|
||||
| head -1 | grep -o '[0-9]*$'
|
||||
}
|
||||
|
||||
id=$(release_id || true)
|
||||
if [ -z "$id" ]; then
|
||||
echo "no release for $tag yet; creating one"
|
||||
out=$(curl -sS -w '\n%{http_code}' -X POST \
|
||||
-H "Authorization: token $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "{\"tag_name\":\"$tag\",\"name\":\"$tag\"}" \
|
||||
"$api/releases")
|
||||
echo "create returned $(printf '%s' "$out" | tail -1)"
|
||||
# Re-read rather than parsing the create response: a 409 means it
|
||||
# was there all along, which is a success for our purposes.
|
||||
id=$(release_id || true)
|
||||
[ -n "$id" ] || printf '%s\n' "$out" | head -5 >&2
|
||||
fi
|
||||
[ -n "$id" ] || { echo "could not find or create the release" >&2; exit 1; }
|
||||
echo "release id $id"
|
||||
|
||||
for f in dist/*; do
|
||||
echo "uploading $(basename "$f")"
|
||||
curl -sS -f -X POST -H "Authorization: token $TOKEN" \
|
||||
-F "attachment=@$f" \
|
||||
"$api/releases/$id/assets?name=$(basename "$f")" >/dev/null \
|
||||
|| { echo "upload of $(basename "$f") was refused" >&2; exit 1; }
|
||||
done
|
||||
echo "attached $(ls dist | wc -l) files to $tag"
|
||||
+286
-6
@@ -8,14 +8,294 @@ is only released once its **exit criterion** is met, not when its code lands.
|
||||
| v0.1.0 | 0 | Skeleton and sim-engine | released |
|
||||
| v0.2.0 | 1 | Static router: streaming proxy, filter/score/pick | released |
|
||||
| v0.3.0 | 1.5 | Admission control: priority bands, saturation detector | released |
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | in progress |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler | planned |
|
||||
| v0.6.0 | 4 | Multi-node replicas and fabric awareness | planned |
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | released as part of v0.5.0 |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler | released |
|
||||
| v0.6.0 | 4 | Multi-node replicas and fabric awareness | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
Until v1.0.0 the minor version is the breaking one: configuration formats and Go
|
||||
APIs change between minors without a deprecation period.
|
||||
|
||||
## Unreleased
|
||||
|
||||
### Memory notes for agents, and threads without API keys
|
||||
|
||||
Notes an agent writes for itself and reads back in later conversations, offered
|
||||
as REST (`/v1/memory`) and as MCP tools at `/mcp` (`remember`, `recall`,
|
||||
`forget`). A note is a key, text and tags, searched by words or tag, private to
|
||||
a project and to the agent id in `x-llmgw-agent`. llmgw does not inject the
|
||||
tools into chat requests: an agent that is not given the server never sees
|
||||
memory.
|
||||
|
||||
Off for every project until someone runs `llmgw keys memory --retain-hours N`,
|
||||
and the retention is required, as it is for threads.
|
||||
|
||||
On a gateway with no authentication, a `threads:` block (`store`,
|
||||
`retain_hours`) now keeps conversations and notes in one shared pool. Before,
|
||||
`x-llmgw-thread` was silently ignored there; now a gateway configured to keep
|
||||
nothing refuses it with an explanation instead. The shared pool is for one
|
||||
trusted machine: anyone who can reach the port and knows an id can read it.
|
||||
|
||||
### Conversations that can move between models
|
||||
|
||||
A thread a client names with `x-llmgw-thread`, carrying only its new turn.
|
||||
llmgw replays the stored history to whichever model is serving, which is what
|
||||
lets a conversation begun on a large model continue on a small one — including
|
||||
when a spent budget moved it there and the caller never chose.
|
||||
|
||||
Off for every project until someone runs `llmgw keys threads`. Until then llmgw
|
||||
stores no message content anywhere, as it never has, and that default is worth
|
||||
more than the feature. The retention is required rather than optional, because
|
||||
a conversation store with no expiry is a transcript of everything a team ever
|
||||
typed kept forever by omission, and it is measured from the last turn so a
|
||||
conversation in progress does not expire underneath it.
|
||||
|
||||
History that does not fit the serving model is trimmed oldest-first and the
|
||||
reply says how many turns went, in `x-llmgw-thread-trimmed`. System messages
|
||||
are never dropped — losing one silently changes what the model is for while it
|
||||
carries on answering — and neither is the turn the caller just sent.
|
||||
|
||||
Threads are scoped to a project, and reading, writing or deleting one that
|
||||
belongs to another answers exactly as one that never existed: two answers would
|
||||
let a caller map which ids are real in projects they cannot read.
|
||||
|
||||
Capturing the reply is opt-in and bounded, so a request without a thread
|
||||
accumulates nothing on the hottest path in the system, and one with a thread
|
||||
stops at a cap rather than growing to it.
|
||||
|
||||
`DELETE /v1/threads/{id}` forgets a conversation, as a client rather than as an
|
||||
operator: the person who said the thing should not have to ask someone to read
|
||||
it before it can go.
|
||||
|
||||
### Disk, and taking back only what llmgw put there
|
||||
|
||||
Disk was the one resource nothing reported. `statfs` was already being called
|
||||
on every model directory, but only to ask whether the filesystem was a network
|
||||
one — the free space in the same struct was discarded, so the only way to find
|
||||
out why a node was filling up was to ssh to it.
|
||||
|
||||
Nodes now report it, per filesystem rather than per configured directory: free
|
||||
space belongs to the disk, and two model directories on one disk are one pool
|
||||
of bytes. Shared filesystems are kept out of a node's own figure, because space
|
||||
on a filer belongs to everyone mounting it and twenty nodes would otherwise
|
||||
each report the same free terabyte. It shows in `ctl nodes`, the dashboard and
|
||||
Prometheus, and a node that reported nothing reads as unknown everywhere — an
|
||||
older agent must not look like a full disk.
|
||||
|
||||
A profile may declare `min_disk_gb`, and a node without the room is no longer a
|
||||
candidate, with a reason `explain-placement` prints. Skipped on a node that
|
||||
already holds the weights: it needs no room for them, and without that the node
|
||||
best placed to run a model is the one most likely to be rejected.
|
||||
|
||||
`ctl weights` lists model files llmgw downloaded for models no longer in the
|
||||
catalog, and `--collect` deletes them. The design is mostly about what it will
|
||||
not do. Only files llmgw fetched are candidates, and that is a calculation
|
||||
rather than a heuristic — the cache name is a hash of the source, so the wanted
|
||||
set is computable exactly. Anything an operator put on a node is left alone
|
||||
however unreferenced it looks. Each node checks every path again before
|
||||
unlinking it, refusing anything outside its cache directory or without the
|
||||
cache name, so a control plane that is wrong cannot talk a node into losing a
|
||||
model that took a week to copy.
|
||||
|
||||
### Fixes
|
||||
|
||||
- **A capability was lost at an indirection.** The controller is handed a
|
||||
forwarding dispatcher, and reaches past its interface by type assertion to
|
||||
find what lies beyond scheduling. Deleting weights was implemented on the
|
||||
agent service and not on the forwarder, so it was simply absent at runtime —
|
||||
with the whole suite green, because the dispatchers the tests use do not
|
||||
implement it either and its absence there is the expected case.
|
||||
- **`ctl tool-check` blamed the engine for what the model did.** It had one
|
||||
verdict for a failed attempt and told operators not to try another model.
|
||||
Measured: the engine renders the tool prompt correctly, the coder models do
|
||||
not use the `<tool_call>` wrapper it asks for, and `qwen2.5-3b-instruct` on
|
||||
the same engine returns a populated `tool_calls` array. A failed attempt is
|
||||
now two verdicts, told apart by whether the reply used the wrapper.
|
||||
|
||||
## v0.6.0 — 2026-09-18
|
||||
|
||||
Phase 4: a replica can be bigger than one machine, and the fabric it spans
|
||||
decides where it is allowed to go. Around that, everything a cluster needs
|
||||
before anyone other than its operator is let near it — keys, budgets, a way in
|
||||
through a browser, and enough visibility to answer "why is nothing scheduling".
|
||||
|
||||
### Multi-node replicas
|
||||
|
||||
A replica is now a set of shards with ranks, a leader and a bootstrap address,
|
||||
rather than one process. Two invariants are enforced rather than documented: a
|
||||
replica never spans two fabrics, and a tensor-parallel group never mixes GPU
|
||||
classes. Both are the kind of mistake that does not fail — it hangs, or it runs
|
||||
several times slower, and nothing says which.
|
||||
|
||||
A replica missing a shard is not ready. It was previously derived from the
|
||||
shards that had reported, so a two-node replica with one node up looked
|
||||
servable and took traffic into a model that could not answer.
|
||||
|
||||
Adoption after a restart happens one agent at a time, because that is how
|
||||
agents reconnect. Merging them as they arrive is what keeps a multi-node
|
||||
replica whole instead of being torn into fragments that are then killed as
|
||||
strays — and an adopted engine is health-checked before it is called ready,
|
||||
rather than being assumed alive because its process exists.
|
||||
|
||||
### Fabrics are distribution domains too
|
||||
|
||||
An isolated fabric is not only a compute boundary: weights present in one are
|
||||
not present in the other. Node weight inventories are reported and summarised
|
||||
per fabric, so "the model is on the cluster" is never mistaken for "the model
|
||||
is reachable from the nodes that would run it". A fabric name that matches
|
||||
nothing a node actually has is refused at registration rather than becoming an
|
||||
unschedulable replica hours later.
|
||||
|
||||
### Who may use it
|
||||
|
||||
API keys, tenants and per-model allowlists, on SQLite (`modernc.org/sqlite`, so
|
||||
still no cgo). Rate limits and budgets are enforced on the request path; usage
|
||||
is written off it, in batches, so accounting never adds latency to inference.
|
||||
Budgets survive a restart — counters are rebuilt from the store, because a
|
||||
limit anyone can clear by bouncing the process is not a limit.
|
||||
|
||||
A spent budget is now served a cheaper model where the catalog offers one,
|
||||
rather than refused. For a team trying models out, being dropped to something
|
||||
smaller is a far better answer than a wall, and the response says which model
|
||||
answered.
|
||||
|
||||
People can sign up, be approved, and issue and revoke their own keys. A key is
|
||||
shown once; only its hash is kept.
|
||||
|
||||
### The browser surface
|
||||
|
||||
A landing page that says how to connect, a read-only cluster dashboard, and an
|
||||
in-browser chat that never puts an API key in the browser. The dashboard shows
|
||||
who is using the cluster and what they are running — per person and per model,
|
||||
with requests, tokens, errors, spend and latency.
|
||||
|
||||
All of it is served from the client-facing port behind a session, deliberately
|
||||
not from the admin port, which also serves pprof — and pprof can dump the heap
|
||||
that holds every API key in plaintext.
|
||||
|
||||
### Models without a restart
|
||||
|
||||
The catalog is live. Models can be added and removed from the command line or
|
||||
the API, take effect immediately, and outlive a restart; configuration still
|
||||
wins over the store, and the whole catalog is validated before anything is
|
||||
written.
|
||||
|
||||
A model is only listed once it can actually serve. It was listed as soon as it
|
||||
was known, so the first thing a new user did was ask for a model that returned
|
||||
503.
|
||||
|
||||
### Loading on demand, and letting go
|
||||
|
||||
A model whose scaling allows it is started when someone asks for it, and
|
||||
stopped when nobody has for a while. An idle replica is evicted to make room
|
||||
for one that was asked for, with guards that exist because each is a way to
|
||||
lose work quietly: only on-demand models, only single-node replicas, never one
|
||||
that has never been used, and never one younger than its minimum lifetime.
|
||||
|
||||
Unknown is not idle. A replica with no recorded use is left alone rather than
|
||||
treated as idle since the epoch.
|
||||
|
||||
### Knowing whether it will run before it runs
|
||||
|
||||
- `ctl plan` works out whether a model fits the cluster and how it would be
|
||||
split, in the precision the chosen engine can actually serve — recommending
|
||||
a quantisation the engine cannot load would cause the failure the command
|
||||
exists to prevent.
|
||||
- `ctl tool-check` verifies llmgw's engine flags against the engines actually
|
||||
installed. Flags drift between releases, and the alternative is finding out
|
||||
during a model load.
|
||||
- Placement now knows a model's context length and memory footprint, and node
|
||||
RAM, so a CPU profile is not placed on a machine that cannot hold it.
|
||||
- Placement waits for the fleet to settle after a restart rather than deciding
|
||||
against a half-reported cluster.
|
||||
|
||||
### Weights
|
||||
|
||||
Nodes report the model files they can see, without llmgw owning the storage:
|
||||
a shared filer, an image exported over the fabric, rsync or a copy made by hand
|
||||
all look the same to placement, which only cares whether the files are there.
|
||||
Reporting nothing means "not asked to look", not "has nothing".
|
||||
|
||||
Where the simple arrangement is not enough, a node can fetch a model by URL —
|
||||
single-flight, resumable, written to a temporary name and renamed atomically,
|
||||
so a killed download never leaves something that looks like weights.
|
||||
|
||||
### Seeing the fleet
|
||||
|
||||
Prometheus metrics for nodes, GPUs and replicas, collected from live state at
|
||||
scrape time rather than from a background goroutine, and served on their own
|
||||
endpoint rather than through the one that also exposes pprof. The same view is
|
||||
in the dashboard, so nothing else needs installing to see it.
|
||||
|
||||
### Onboarding a node
|
||||
|
||||
Adding a node is one command. There are systemd units, a private CA and a
|
||||
from-nothing install guide that starts with a machine that has no CUDA and no
|
||||
binary.
|
||||
|
||||
### llama.cpp, and CPU profiles
|
||||
|
||||
A third engine adapter, and the reason for it: llama.cpp serves the same OpenAI
|
||||
API without a GPU, so placement, supervision, scraping, scheduling and admission
|
||||
can all be exercised on ordinary machines.
|
||||
|
||||
A profile may now omit `gpu_class`, which is how it says it runs on the CPU. Only
|
||||
engines that can actually serve without a card may omit it — for vLLM an omitted
|
||||
class stays an error, because otherwise a typo becomes a replica placed on a
|
||||
machine with no hardware and diagnosed minutes later during model loading. CPU
|
||||
replicas are spread across nodes rather than bin packed: unlike GPUs there is no
|
||||
allocation to stop two of them sharing a machine's cores.
|
||||
|
||||
### Metric mappings are per engine
|
||||
|
||||
`Running`, `Waiting` and `KVCacheUsage` now have a vocabulary per engine rather
|
||||
than one for the cluster. A mixed fleet has two names for the same three numbers,
|
||||
and a single global mapping left half of it unmeasurable.
|
||||
|
||||
A mapping that leaves a metric unnamed now means the engine does not publish it,
|
||||
and the field stays zero. Naming one the engine does not serve is still a contract
|
||||
violation. `Running` and `Waiting` can never be omitted.
|
||||
|
||||
### Mutual TLS on the agent stream
|
||||
|
||||
A private CA (`llmgw ctl cert init-ca | issue-server | issue-client | inspect`),
|
||||
certificates on both ends, and the control plane refusing any agent whose claimed
|
||||
node id does not match its certificate's common name. Public authorities cannot
|
||||
supply these: they stopped issuing the `clientAuth` usage.
|
||||
|
||||
Plaintext HTTP/2 remains the default. `agent_tls:` on the server and `tls:` on the
|
||||
agent turn it on.
|
||||
|
||||
### Fixes
|
||||
|
||||
- **An agent told to stop did not stop.** Cancelling the context is not enough —
|
||||
HTTP/2 watches it while the request body is being written, and a bidirectional
|
||||
stream spends its life blocked on the read side. Measured on a deployed node:
|
||||
`systemctl stop` took 90 seconds and ended in SIGKILL. Now 0.01 seconds.
|
||||
- **An agent could not explain why it was refused.** The session cancelled its own
|
||||
context and then asked that context whether it was shutting down, so every
|
||||
stream error read as a clean close: no message, and the backoff reset.
|
||||
- **A broken metric contract was invisible.** It was logged per scrape at debug
|
||||
level, so a cluster could look healthy while every routing decision was made
|
||||
blind. It is now a warning, once, with a matching line when it clears.
|
||||
- **Token estimates came from the routing prefix, not the prompt.** Every prompt
|
||||
over about 2 KB was undercounted, so tokens-per-minute limits let through far
|
||||
more than they were set to.
|
||||
- **Model-fitting double-counted overhead.** Whole-model overhead was charged per
|
||||
rank and a 10% reserve taken twice, so `plan` reported that a 70B model in
|
||||
bfloat16 does not fit two 80 GB cards. It does.
|
||||
- **`x-llmgw-served-model` was never sent.** It was set after the proxy had
|
||||
already written the response header, which is the one point at which headers
|
||||
stop being changeable. A client had no way to tell a fallback from the model
|
||||
it asked for.
|
||||
- **Fleet metrics published nothing.** They were registered after the registry
|
||||
had been built, so the collector existed and was never scraped.
|
||||
- **GPU discovery broke on older drivers.** A field added for utilisation was
|
||||
required rather than optional, so a driver that does not report it produced no
|
||||
inventory at all instead of an inventory without that number.
|
||||
- **Replica ids and ports were reused** while the processes holding them were
|
||||
still alive.
|
||||
|
||||
## v0.5.0 — 2026-09-16
|
||||
|
||||
Phases 2 and 3. The cluster runs itself: agents report their hardware, the
|
||||
@@ -50,10 +330,10 @@ not have. The failure path already ships an engine log tail with every status
|
||||
report, so the case that matters most is covered; reading a healthy shard's
|
||||
output on demand is not.
|
||||
|
||||
## Unreleased — phase 4 in progress
|
||||
### Phases 2 and 3, in detail
|
||||
|
||||
Agent groundwork. Everything here works without GPUs and is covered by tests;
|
||||
what is missing is the control plane connection, so there is no release yet.
|
||||
Written as the work landed, which is why it reads as a log rather than as
|
||||
notes. Everything here shipped in v0.5.0.
|
||||
|
||||
### Added
|
||||
|
||||
|
||||
@@ -6,6 +6,18 @@ heterogeneous GPUs.
|
||||
It serves LLMs running on your own GPU servers (vLLM / SGLang) behind one
|
||||
OpenAI-compatible endpoint. Design document: [llmgw-plan.md](llmgw-plan.md).
|
||||
|
||||
## Getting started
|
||||
|
||||
Starting from nothing — no packages installed, no drivers, no nodes — follow
|
||||
[docs/install.md](docs/install.md): it builds a control plane, adds a node and
|
||||
serves a model, and says plainly which parts llmgw does not install for you.
|
||||
|
||||
| Guide | For |
|
||||
|---|---|
|
||||
| [docs/install.md](docs/install.md) | Standing up a cluster from scratch |
|
||||
| [docs/node-onboarding.md](docs/node-onboarding.md) | Adding nodes to one that exists |
|
||||
| [docs/onboarding.md](docs/onboarding.md) | Giving people accounts and keys |
|
||||
|
||||
## Status
|
||||
|
||||
Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
@@ -17,13 +29,26 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
| v0.3.0 | 1.5 | Admission control (priority bands) | released |
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, engine supervision | released |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler, ctl | released |
|
||||
| v0.6.0 | 4 | Multi-node placement and fabric awareness | planned |
|
||||
| v0.6.0 | 4 | Multi-node placement and fabric awareness | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
As of v0.3.0 the router serves traffic and protects it: a batch workload no
|
||||
longer starves interactive requests. Still missing before v1.0.0: authentication,
|
||||
rate limits and budgets, and the control plane that replaces the static endpoint
|
||||
list.
|
||||
As of v0.6.0 a cluster runs unattended: agents report their hardware, the
|
||||
control plane places replicas — including ones larger than a single machine —
|
||||
models load when someone asks for them and stop when nobody does, and keys,
|
||||
budgets and rate limits decide who may use any of it.
|
||||
|
||||
Before v1.0.0: tracing, reference documentation, and filtering on `ctl nodes`
|
||||
and `ctl replicas` for fleets too large to read a whole table of. Two items
|
||||
that were on this list are not gaps: a model falls through to its fallbacks
|
||||
when there is no capacity as well as when a budget is spent, and routing
|
||||
profiles are chosen by the key rather than by a request header on purpose —
|
||||
a header would let any client promote itself into the top priority band.
|
||||
|
||||
One thing is worth saying plainly:
|
||||
multi-node replicas are proven on a simulated cluster and on real hardware
|
||||
without a fabric. The InfiniBand, NVML and NCCL paths are written against the
|
||||
formats and failure modes those tools actually have, but they have not yet run
|
||||
on a machine that has them.
|
||||
|
||||
## Building
|
||||
|
||||
|
||||
+425
-85
@@ -90,7 +90,7 @@ func (x ShardStatus_State) Number() protoreflect.EnumNumber {
|
||||
|
||||
// Deprecated: Use ShardStatus_State.Descriptor instead.
|
||||
func (ShardStatus_State) EnumDescriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7, 0}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9, 0}
|
||||
}
|
||||
|
||||
// AgentMessage is anything a node tells the control plane.
|
||||
@@ -201,6 +201,7 @@ type ControlMessage struct {
|
||||
// *ControlMessage_Stop
|
||||
// *ControlMessage_Drain
|
||||
// *ControlMessage_Ack
|
||||
// *ControlMessage_RemoveWeights
|
||||
Msg isControlMessage_Msg `protobuf_oneof:"msg"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
@@ -279,6 +280,15 @@ func (x *ControlMessage) GetAck() *Ack {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *ControlMessage) GetRemoveWeights() *RemoveWeights {
|
||||
if x != nil {
|
||||
if x, ok := x.Msg.(*ControlMessage_RemoveWeights); ok {
|
||||
return x.RemoveWeights
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type isControlMessage_Msg interface {
|
||||
isControlMessage_Msg()
|
||||
}
|
||||
@@ -299,6 +309,10 @@ type ControlMessage_Ack struct {
|
||||
Ack *Ack `protobuf:"bytes,4,opt,name=ack,proto3,oneof"`
|
||||
}
|
||||
|
||||
type ControlMessage_RemoveWeights struct {
|
||||
RemoveWeights *RemoveWeights `protobuf:"bytes,5,opt,name=remove_weights,json=removeWeights,proto3,oneof"`
|
||||
}
|
||||
|
||||
func (*ControlMessage_Start) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_Stop) isControlMessage_Msg() {}
|
||||
@@ -307,6 +321,8 @@ func (*ControlMessage_Drain) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_Ack) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_RemoveWeights) isControlMessage_Msg() {}
|
||||
|
||||
// Register is the first message on every stream. It carries the node's full
|
||||
// hardware inventory, so the control plane needs no separate discovery call and
|
||||
// a reconnecting agent automatically refreshes what the scheduler knows.
|
||||
@@ -336,6 +352,36 @@ type Register struct {
|
||||
// — so the new engines fail to allocate, with an error naming the wrong
|
||||
// process (plan §6.5).
|
||||
RunningShards []*RunningShard `protobuf:"bytes,9,rep,name=running_shards,json=runningShards,proto3" json:"running_shards,omitempty"`
|
||||
// memory_total_mb is the machine's RAM, and 0 means the agent did not say —
|
||||
// which is what every agent older than this field reports, and is treated as
|
||||
// no constraint rather than as a node with no memory.
|
||||
//
|
||||
// Placement needs it because a CPU profile has no card to be free, so every
|
||||
// ready node looks like a candidate: a 14B model whose working set is 24 GB
|
||||
// would be placed just as happily on a node with 8 GB.
|
||||
MemoryTotalMb int64 `protobuf:"varint,10,opt,name=memory_total_mb,json=memoryTotalMb,proto3" json:"memory_total_mb,omitempty"`
|
||||
// weights are the model files this node can see, from the directories its
|
||||
// configuration names.
|
||||
//
|
||||
// How they got there is the operator's business and deliberately outside
|
||||
// llmgw's: a shared filesystem, an image exported over the fabric, rsync, or
|
||||
// a copy made by hand all look identical here. What matters to placement is
|
||||
// only whether they are present, because a node without them either cannot
|
||||
// start the replica at all or must first move tens of gigabytes.
|
||||
//
|
||||
// Empty means the node was not asked to look, which is not the same as having
|
||||
// nothing — an agent with no model directories configured behaves exactly as
|
||||
// every agent did before this existed, and constrains nothing.
|
||||
Weights []*LocalWeights `protobuf:"bytes,11,rep,name=weights,proto3" json:"weights,omitempty"`
|
||||
// storage is the disk behind those weights, one entry per filesystem.
|
||||
//
|
||||
// Reported so an operator can see a node filling up before a fetch dies at
|
||||
// ninety percent, and so placement can refuse a node that cannot hold a
|
||||
// model rather than discovering it during the download.
|
||||
//
|
||||
// Empty means the node was not asked to look, exactly as for weights, and
|
||||
// constrains nothing.
|
||||
Storage []*Storage `protobuf:"bytes,12,rep,name=storage,proto3" json:"storage,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
@@ -433,6 +479,105 @@ func (x *Register) GetRunningShards() []*RunningShard {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Register) GetMemoryTotalMb() int64 {
|
||||
if x != nil {
|
||||
return x.MemoryTotalMb
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *Register) GetWeights() []*LocalWeights {
|
||||
if x != nil {
|
||||
return x.Weights
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Register) GetStorage() []*Storage {
|
||||
if x != nil {
|
||||
return x.Storage
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// Per filesystem rather than per configured directory: free space belongs to
|
||||
// the disk, and two model directories on one disk are one pool of bytes.
|
||||
type Storage struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
// path is one directory on this filesystem, so the operator recognises it.
|
||||
Path string `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"`
|
||||
// free_bytes is what a non-root process can write — not the root reserve,
|
||||
// which is space the agent cannot have.
|
||||
FreeBytes int64 `protobuf:"varint,2,opt,name=free_bytes,json=freeBytes,proto3" json:"free_bytes,omitempty"`
|
||||
TotalBytes int64 `protobuf:"varint,3,opt,name=total_bytes,json=totalBytes,proto3" json:"total_bytes,omitempty"`
|
||||
// shared marks a filesystem the node did not fill. Free space on one is
|
||||
// shared with every other node using it, so two nodes each reporting 500 GB
|
||||
// free on the same filer have 500 GB between them and not a terabyte.
|
||||
Shared bool `protobuf:"varint,4,opt,name=shared,proto3" json:"shared,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Storage) Reset() {
|
||||
*x = Storage{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *Storage) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*Storage) ProtoMessage() {}
|
||||
|
||||
func (x *Storage) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use Storage.ProtoReflect.Descriptor instead.
|
||||
func (*Storage) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3}
|
||||
}
|
||||
|
||||
func (x *Storage) GetPath() string {
|
||||
if x != nil {
|
||||
return x.Path
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *Storage) GetFreeBytes() int64 {
|
||||
if x != nil {
|
||||
return x.FreeBytes
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *Storage) GetTotalBytes() int64 {
|
||||
if x != nil {
|
||||
return x.TotalBytes
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *Storage) GetShared() bool {
|
||||
if x != nil {
|
||||
return x.Shared
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
//
|
||||
// It carries enough for the control plane to rebuild the replica it belongs to,
|
||||
@@ -463,7 +608,7 @@ type RunningShard struct {
|
||||
|
||||
func (x *RunningShard) Reset() {
|
||||
*x = RunningShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -475,7 +620,7 @@ func (x *RunningShard) String() string {
|
||||
func (*RunningShard) ProtoMessage() {}
|
||||
|
||||
func (x *RunningShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -488,7 +633,7 @@ func (x *RunningShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use RunningShard.ProtoReflect.Descriptor instead.
|
||||
func (*RunningShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4}
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetReplicaId() string {
|
||||
@@ -605,16 +750,20 @@ type GPU struct {
|
||||
NvlinkPeers []int32 `protobuf:"varint,11,rep,packed,name=nvlink_peers,json=nvlinkPeers,proto3" json:"nvlink_peers,omitempty"`
|
||||
// problems keeps a GPU out of placement: ECC faults, or memory already held
|
||||
// by another process.
|
||||
Problems []string `protobuf:"bytes,12,rep,name=problems,proto3" json:"problems,omitempty"`
|
||||
TemperatureC int32 `protobuf:"varint,13,opt,name=temperature_c,json=temperatureC,proto3" json:"temperature_c,omitempty"`
|
||||
EccErrors int32 `protobuf:"varint,14,opt,name=ecc_errors,json=eccErrors,proto3" json:"ecc_errors,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
Problems []string `protobuf:"bytes,12,rep,name=problems,proto3" json:"problems,omitempty"`
|
||||
TemperatureC int32 `protobuf:"varint,13,opt,name=temperature_c,json=temperatureC,proto3" json:"temperature_c,omitempty"`
|
||||
EccErrors int32 `protobuf:"varint,14,opt,name=ecc_errors,json=eccErrors,proto3" json:"ecc_errors,omitempty"`
|
||||
// utilization_percent is how busy the card was over the last sample period.
|
||||
// Zero when the driver is too old to report it, which is not the same as
|
||||
// idle: the metric is simply absent rather than wrong.
|
||||
UtilizationPercent int32 `protobuf:"varint,15,opt,name=utilization_percent,json=utilizationPercent,proto3" json:"utilization_percent,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *GPU) Reset() {
|
||||
*x = GPU{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -626,7 +775,7 @@ func (x *GPU) String() string {
|
||||
func (*GPU) ProtoMessage() {}
|
||||
|
||||
func (x *GPU) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -639,7 +788,7 @@ func (x *GPU) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use GPU.ProtoReflect.Descriptor instead.
|
||||
func (*GPU) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5}
|
||||
}
|
||||
|
||||
func (x *GPU) GetIndex() int32 {
|
||||
@@ -740,6 +889,13 @@ func (x *GPU) GetEccErrors() int32 {
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *GPU) GetUtilizationPercent() int32 {
|
||||
if x != nil {
|
||||
return x.UtilizationPercent
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
type FabricAttachment struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
// fabric is the operator-assigned name. It is not derived from the GID
|
||||
@@ -754,7 +910,7 @@ type FabricAttachment struct {
|
||||
|
||||
func (x *FabricAttachment) Reset() {
|
||||
*x = FabricAttachment{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -766,7 +922,7 @@ func (x *FabricAttachment) String() string {
|
||||
func (*FabricAttachment) ProtoMessage() {}
|
||||
|
||||
func (x *FabricAttachment) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -779,7 +935,7 @@ func (x *FabricAttachment) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use FabricAttachment.ProtoReflect.Descriptor instead.
|
||||
func (*FabricAttachment) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6}
|
||||
}
|
||||
|
||||
func (x *FabricAttachment) GetFabric() string {
|
||||
@@ -804,6 +960,72 @@ func (x *FabricAttachment) GetRateGbps() int32 {
|
||||
}
|
||||
|
||||
// Heartbeat says the node is alive and reports anything that changed cheaply.
|
||||
// LocalWeights is one model file or directory a node can see.
|
||||
type LocalWeights struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
// path is what an engine would be given, so it can be compared against a
|
||||
// model's source without interpretation.
|
||||
Path string `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"`
|
||||
SizeBytes int64 `protobuf:"varint,2,opt,name=size_bytes,json=sizeBytes,proto3" json:"size_bytes,omitempty"`
|
||||
// shared marks weights on a filesystem the node did not have to fill —
|
||||
// reported so a fabric with a shared filer is not mistaken for one where
|
||||
// every node happens to have fetched its own copy.
|
||||
Shared bool `protobuf:"varint,3,opt,name=shared,proto3" json:"shared,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *LocalWeights) Reset() {
|
||||
*x = LocalWeights{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *LocalWeights) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*LocalWeights) ProtoMessage() {}
|
||||
|
||||
func (x *LocalWeights) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use LocalWeights.ProtoReflect.Descriptor instead.
|
||||
func (*LocalWeights) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7}
|
||||
}
|
||||
|
||||
func (x *LocalWeights) GetPath() string {
|
||||
if x != nil {
|
||||
return x.Path
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *LocalWeights) GetSizeBytes() int64 {
|
||||
if x != nil {
|
||||
return x.SizeBytes
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *LocalWeights) GetShared() bool {
|
||||
if x != nil {
|
||||
return x.Shared
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type Heartbeat struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
NodeId string `protobuf:"bytes,1,opt,name=node_id,json=nodeId,proto3" json:"node_id,omitempty"`
|
||||
@@ -811,13 +1033,27 @@ type Heartbeat struct {
|
||||
// a card after registration.
|
||||
GpuMemoryUsedMb []int32 `protobuf:"varint,2,rep,packed,name=gpu_memory_used_mb,json=gpuMemoryUsedMb,proto3" json:"gpu_memory_used_mb,omitempty"`
|
||||
NewProblems []string `protobuf:"bytes,3,rep,name=new_problems,json=newProblems,proto3" json:"new_problems,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
// storage for the same reason the GPU memory is here: it changes while the
|
||||
// agent runs, and a figure sampled once at registration is a figure that is
|
||||
// wrong on exactly the node worth watching — the one busy filling up.
|
||||
//
|
||||
// Always sent, because it is a few numbers per filesystem.
|
||||
Storage []*Storage `protobuf:"bytes,4,rep,name=storage,proto3" json:"storage,omitempty"`
|
||||
// weights only when they changed, because the list is long enough that
|
||||
// sending it every few seconds would be paying for nothing most of the time.
|
||||
//
|
||||
// weights_changed is what says the list is meant: without it an empty list
|
||||
// would be ambiguous between "this node now holds nothing" and "nothing to
|
||||
// report", and the first would quietly empty a node's inventory.
|
||||
WeightsChanged bool `protobuf:"varint,5,opt,name=weights_changed,json=weightsChanged,proto3" json:"weights_changed,omitempty"`
|
||||
Weights []*LocalWeights `protobuf:"bytes,6,rep,name=weights,proto3" json:"weights,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Heartbeat) Reset() {
|
||||
*x = Heartbeat{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -829,7 +1065,7 @@ func (x *Heartbeat) String() string {
|
||||
func (*Heartbeat) ProtoMessage() {}
|
||||
|
||||
func (x *Heartbeat) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -842,7 +1078,7 @@ func (x *Heartbeat) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use Heartbeat.ProtoReflect.Descriptor instead.
|
||||
func (*Heartbeat) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8}
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetNodeId() string {
|
||||
@@ -866,6 +1102,27 @@ func (x *Heartbeat) GetNewProblems() []string {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetStorage() []*Storage {
|
||||
if x != nil {
|
||||
return x.Storage
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetWeightsChanged() bool {
|
||||
if x != nil {
|
||||
return x.WeightsChanged
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetWeights() []*LocalWeights {
|
||||
if x != nil {
|
||||
return x.Weights
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
type ShardStatus struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
@@ -884,7 +1141,7 @@ type ShardStatus struct {
|
||||
|
||||
func (x *ShardStatus) Reset() {
|
||||
*x = ShardStatus{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -896,7 +1153,7 @@ func (x *ShardStatus) String() string {
|
||||
func (*ShardStatus) ProtoMessage() {}
|
||||
|
||||
func (x *ShardStatus) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -909,7 +1166,7 @@ func (x *ShardStatus) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use ShardStatus.ProtoReflect.Descriptor instead.
|
||||
func (*ShardStatus) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9}
|
||||
}
|
||||
|
||||
func (x *ShardStatus) GetNodeId() string {
|
||||
@@ -990,7 +1247,7 @@ type StartShard struct {
|
||||
|
||||
func (x *StartShard) Reset() {
|
||||
*x = StartShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1002,7 +1259,7 @@ func (x *StartShard) String() string {
|
||||
func (*StartShard) ProtoMessage() {}
|
||||
|
||||
func (x *StartShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1015,7 +1272,7 @@ func (x *StartShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use StartShard.ProtoReflect.Descriptor instead.
|
||||
func (*StartShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10}
|
||||
}
|
||||
|
||||
func (x *StartShard) GetReplicaId() string {
|
||||
@@ -1135,7 +1392,7 @@ type EngineProfile struct {
|
||||
|
||||
func (x *EngineProfile) Reset() {
|
||||
*x = EngineProfile{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1147,7 +1404,7 @@ func (x *EngineProfile) String() string {
|
||||
func (*EngineProfile) ProtoMessage() {}
|
||||
|
||||
func (x *EngineProfile) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1160,7 +1417,7 @@ func (x *EngineProfile) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use EngineProfile.ProtoReflect.Descriptor instead.
|
||||
func (*EngineProfile) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11}
|
||||
}
|
||||
|
||||
func (x *EngineProfile) GetName() string {
|
||||
@@ -1240,6 +1497,57 @@ func (x *EngineProfile) GetServedModelName() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
//
|
||||
// The paths are computed by the control plane from the catalog, but the agent
|
||||
// checks them again before unlinking anything: it deletes only files inside its
|
||||
// own cache directory whose names this system generated. A control plane that
|
||||
// is wrong, compromised or simply newer than the agent must not be able to
|
||||
// talk it into deleting a model an operator spent a week copying there.
|
||||
type RemoveWeights struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Paths []string `protobuf:"bytes,1,rep,name=paths,proto3" json:"paths,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *RemoveWeights) Reset() {
|
||||
*x = RemoveWeights{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *RemoveWeights) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*RemoveWeights) ProtoMessage() {}
|
||||
|
||||
func (x *RemoveWeights) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use RemoveWeights.ProtoReflect.Descriptor instead.
|
||||
func (*RemoveWeights) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{12}
|
||||
}
|
||||
|
||||
func (x *RemoveWeights) GetPaths() []string {
|
||||
if x != nil {
|
||||
return x.Paths
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StopShard terminates a shard and frees its GPUs.
|
||||
type StopShard struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
@@ -1250,7 +1558,7 @@ type StopShard struct {
|
||||
|
||||
func (x *StopShard) Reset() {
|
||||
*x = StopShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[13]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1262,7 +1570,7 @@ func (x *StopShard) String() string {
|
||||
func (*StopShard) ProtoMessage() {}
|
||||
|
||||
func (x *StopShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[13]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1275,7 +1583,7 @@ func (x *StopShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use StopShard.ProtoReflect.Descriptor instead.
|
||||
func (*StopShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{13}
|
||||
}
|
||||
|
||||
func (x *StopShard) GetShardId() string {
|
||||
@@ -1295,7 +1603,7 @@ type DrainShard struct {
|
||||
|
||||
func (x *DrainShard) Reset() {
|
||||
*x = DrainShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[14]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1307,7 +1615,7 @@ func (x *DrainShard) String() string {
|
||||
func (*DrainShard) ProtoMessage() {}
|
||||
|
||||
func (x *DrainShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[14]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1320,7 +1628,7 @@ func (x *DrainShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use DrainShard.ProtoReflect.Descriptor instead.
|
||||
func (*DrainShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{14}
|
||||
}
|
||||
|
||||
func (x *DrainShard) GetShardId() string {
|
||||
@@ -1342,7 +1650,7 @@ type Ack struct {
|
||||
|
||||
func (x *Ack) Reset() {
|
||||
*x = Ack{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[15]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1354,7 +1662,7 @@ func (x *Ack) String() string {
|
||||
func (*Ack) ProtoMessage() {}
|
||||
|
||||
func (x *Ack) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[15]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1367,7 +1675,7 @@ func (x *Ack) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use Ack.ProtoReflect.Descriptor instead.
|
||||
func (*Ack) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{12}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{15}
|
||||
}
|
||||
|
||||
func (x *Ack) GetMessageId() string {
|
||||
@@ -1400,13 +1708,14 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\bregister\x18\x01 \x01(\v2\x12.llmgw.v1.RegisterH\x00R\bregister\x123\n" +
|
||||
"\theartbeat\x18\x02 \x01(\v2\x13.llmgw.v1.HeartbeatH\x00R\theartbeat\x12-\n" +
|
||||
"\x05shard\x18\x03 \x01(\v2\x15.llmgw.v1.ShardStatusH\x00R\x05shardB\x05\n" +
|
||||
"\x03msg\"\xc1\x01\n" +
|
||||
"\x03msg\"\x83\x02\n" +
|
||||
"\x0eControlMessage\x12,\n" +
|
||||
"\x05start\x18\x01 \x01(\v2\x14.llmgw.v1.StartShardH\x00R\x05start\x12)\n" +
|
||||
"\x04stop\x18\x02 \x01(\v2\x13.llmgw.v1.StopShardH\x00R\x04stop\x12,\n" +
|
||||
"\x05drain\x18\x03 \x01(\v2\x14.llmgw.v1.DrainShardH\x00R\x05drain\x12!\n" +
|
||||
"\x03ack\x18\x04 \x01(\v2\r.llmgw.v1.AckH\x00R\x03ackB\x05\n" +
|
||||
"\x03msg\"\xa0\x04\n" +
|
||||
"\x03ack\x18\x04 \x01(\v2\r.llmgw.v1.AckH\x00R\x03ack\x12@\n" +
|
||||
"\x0eremove_weights\x18\x05 \x01(\v2\x17.llmgw.v1.RemoveWeightsH\x00R\rremoveWeightsB\x05\n" +
|
||||
"\x03msg\"\xa7\x05\n" +
|
||||
"\bRegister\x12\x17\n" +
|
||||
"\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1a\n" +
|
||||
"\bhostname\x18\x02 \x01(\tR\bhostname\x12\x1b\n" +
|
||||
@@ -1416,13 +1725,24 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\afabrics\x18\x06 \x03(\v2\x1a.llmgw.v1.FabricAttachmentR\afabrics\x12#\n" +
|
||||
"\ragent_version\x18\a \x01(\tR\fagentVersion\x12O\n" +
|
||||
"\x0fengine_versions\x18\b \x03(\v2&.llmgw.v1.Register.EngineVersionsEntryR\x0eengineVersions\x12=\n" +
|
||||
"\x0erunning_shards\x18\t \x03(\v2\x16.llmgw.v1.RunningShardR\rrunningShards\x1a9\n" +
|
||||
"\x0erunning_shards\x18\t \x03(\v2\x16.llmgw.v1.RunningShardR\rrunningShards\x12&\n" +
|
||||
"\x0fmemory_total_mb\x18\n" +
|
||||
" \x01(\x03R\rmemoryTotalMb\x120\n" +
|
||||
"\aweights\x18\v \x03(\v2\x16.llmgw.v1.LocalWeightsR\aweights\x12+\n" +
|
||||
"\astorage\x18\f \x03(\v2\x11.llmgw.v1.StorageR\astorage\x1a9\n" +
|
||||
"\vLabelsEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\x1aA\n" +
|
||||
"\x13EngineVersionsEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\x97\x03\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"u\n" +
|
||||
"\aStorage\x12\x12\n" +
|
||||
"\x04path\x18\x01 \x01(\tR\x04path\x12\x1d\n" +
|
||||
"\n" +
|
||||
"free_bytes\x18\x02 \x01(\x03R\tfreeBytes\x12\x1f\n" +
|
||||
"\vtotal_bytes\x18\x03 \x01(\x03R\n" +
|
||||
"totalBytes\x12\x16\n" +
|
||||
"\x06shared\x18\x04 \x01(\bR\x06shared\"\x97\x03\n" +
|
||||
"\fRunningShard\x12\x1d\n" +
|
||||
"\n" +
|
||||
"replica_id\x18\x01 \x01(\tR\treplicaId\x12\x19\n" +
|
||||
@@ -1442,7 +1762,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\vleader_addr\x18\r \x01(\tR\n" +
|
||||
"leaderAddr\x12\x1f\n" +
|
||||
"\vleader_port\x18\x0e \x01(\x05R\n" +
|
||||
"leaderPort\"\xc8\x03\n" +
|
||||
"leaderPort\"\xf9\x03\n" +
|
||||
"\x03GPU\x12\x14\n" +
|
||||
"\x05index\x18\x01 \x01(\x05R\x05index\x12\x12\n" +
|
||||
"\x04uuid\x18\x02 \x01(\tR\x04uuid\x12\x12\n" +
|
||||
@@ -1460,15 +1780,24 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\bproblems\x18\f \x03(\tR\bproblems\x12#\n" +
|
||||
"\rtemperature_c\x18\r \x01(\x05R\ftemperatureC\x12\x1d\n" +
|
||||
"\n" +
|
||||
"ecc_errors\x18\x0e \x01(\x05R\teccErrors\"[\n" +
|
||||
"ecc_errors\x18\x0e \x01(\x05R\teccErrors\x12/\n" +
|
||||
"\x13utilization_percent\x18\x0f \x01(\x05R\x12utilizationPercent\"[\n" +
|
||||
"\x10FabricAttachment\x12\x16\n" +
|
||||
"\x06fabric\x18\x01 \x01(\tR\x06fabric\x12\x12\n" +
|
||||
"\x04hcas\x18\x02 \x03(\tR\x04hcas\x12\x1b\n" +
|
||||
"\trate_gbps\x18\x03 \x01(\x05R\brateGbps\"t\n" +
|
||||
"\trate_gbps\x18\x03 \x01(\x05R\brateGbps\"Y\n" +
|
||||
"\fLocalWeights\x12\x12\n" +
|
||||
"\x04path\x18\x01 \x01(\tR\x04path\x12\x1d\n" +
|
||||
"\n" +
|
||||
"size_bytes\x18\x02 \x01(\x03R\tsizeBytes\x12\x16\n" +
|
||||
"\x06shared\x18\x03 \x01(\bR\x06shared\"\xfc\x01\n" +
|
||||
"\tHeartbeat\x12\x17\n" +
|
||||
"\anode_id\x18\x01 \x01(\tR\x06nodeId\x12+\n" +
|
||||
"\x12gpu_memory_used_mb\x18\x02 \x03(\x05R\x0fgpuMemoryUsedMb\x12!\n" +
|
||||
"\fnew_problems\x18\x03 \x03(\tR\vnewProblems\"\x8c\x03\n" +
|
||||
"\fnew_problems\x18\x03 \x03(\tR\vnewProblems\x12+\n" +
|
||||
"\astorage\x18\x04 \x03(\v2\x11.llmgw.v1.StorageR\astorage\x12'\n" +
|
||||
"\x0fweights_changed\x18\x05 \x01(\bR\x0eweightsChanged\x120\n" +
|
||||
"\aweights\x18\x06 \x03(\v2\x16.llmgw.v1.LocalWeightsR\aweights\"\x8c\x03\n" +
|
||||
"\vShardStatus\x12\x17\n" +
|
||||
"\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1d\n" +
|
||||
"\n" +
|
||||
@@ -1525,7 +1854,9 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\x11served_model_name\x18\v \x01(\tR\x0fservedModelName\x1a6\n" +
|
||||
"\bEnvEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"&\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"%\n" +
|
||||
"\rRemoveWeights\x12\x14\n" +
|
||||
"\x05paths\x18\x01 \x03(\tR\x05paths\"&\n" +
|
||||
"\tStopShard\x12\x19\n" +
|
||||
"\bshard_id\x18\x01 \x01(\tR\ashardId\"'\n" +
|
||||
"\n" +
|
||||
@@ -1554,50 +1885,58 @@ func file_llmgw_v1_agent_proto_rawDescGZIP() []byte {
|
||||
}
|
||||
|
||||
var file_llmgw_v1_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
|
||||
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 16)
|
||||
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 19)
|
||||
var file_llmgw_v1_agent_proto_goTypes = []any{
|
||||
(ShardStatus_State)(0), // 0: llmgw.v1.ShardStatus.State
|
||||
(*AgentMessage)(nil), // 1: llmgw.v1.AgentMessage
|
||||
(*ControlMessage)(nil), // 2: llmgw.v1.ControlMessage
|
||||
(*Register)(nil), // 3: llmgw.v1.Register
|
||||
(*RunningShard)(nil), // 4: llmgw.v1.RunningShard
|
||||
(*GPU)(nil), // 5: llmgw.v1.GPU
|
||||
(*FabricAttachment)(nil), // 6: llmgw.v1.FabricAttachment
|
||||
(*Heartbeat)(nil), // 7: llmgw.v1.Heartbeat
|
||||
(*ShardStatus)(nil), // 8: llmgw.v1.ShardStatus
|
||||
(*StartShard)(nil), // 9: llmgw.v1.StartShard
|
||||
(*EngineProfile)(nil), // 10: llmgw.v1.EngineProfile
|
||||
(*StopShard)(nil), // 11: llmgw.v1.StopShard
|
||||
(*DrainShard)(nil), // 12: llmgw.v1.DrainShard
|
||||
(*Ack)(nil), // 13: llmgw.v1.Ack
|
||||
nil, // 14: llmgw.v1.Register.LabelsEntry
|
||||
nil, // 15: llmgw.v1.Register.EngineVersionsEntry
|
||||
nil, // 16: llmgw.v1.EngineProfile.EnvEntry
|
||||
(*Storage)(nil), // 4: llmgw.v1.Storage
|
||||
(*RunningShard)(nil), // 5: llmgw.v1.RunningShard
|
||||
(*GPU)(nil), // 6: llmgw.v1.GPU
|
||||
(*FabricAttachment)(nil), // 7: llmgw.v1.FabricAttachment
|
||||
(*LocalWeights)(nil), // 8: llmgw.v1.LocalWeights
|
||||
(*Heartbeat)(nil), // 9: llmgw.v1.Heartbeat
|
||||
(*ShardStatus)(nil), // 10: llmgw.v1.ShardStatus
|
||||
(*StartShard)(nil), // 11: llmgw.v1.StartShard
|
||||
(*EngineProfile)(nil), // 12: llmgw.v1.EngineProfile
|
||||
(*RemoveWeights)(nil), // 13: llmgw.v1.RemoveWeights
|
||||
(*StopShard)(nil), // 14: llmgw.v1.StopShard
|
||||
(*DrainShard)(nil), // 15: llmgw.v1.DrainShard
|
||||
(*Ack)(nil), // 16: llmgw.v1.Ack
|
||||
nil, // 17: llmgw.v1.Register.LabelsEntry
|
||||
nil, // 18: llmgw.v1.Register.EngineVersionsEntry
|
||||
nil, // 19: llmgw.v1.EngineProfile.EnvEntry
|
||||
}
|
||||
var file_llmgw_v1_agent_proto_depIdxs = []int32{
|
||||
3, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register
|
||||
7, // 1: llmgw.v1.AgentMessage.heartbeat:type_name -> llmgw.v1.Heartbeat
|
||||
8, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus
|
||||
9, // 3: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
|
||||
11, // 4: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
|
||||
12, // 5: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
|
||||
13, // 6: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
|
||||
14, // 7: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
|
||||
5, // 8: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
|
||||
6, // 9: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
15, // 10: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
|
||||
4, // 11: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
|
||||
0, // 12: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
0, // 13: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
10, // 14: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
|
||||
16, // 15: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
|
||||
1, // 16: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
|
||||
2, // 17: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
|
||||
17, // [17:18] is the sub-list for method output_type
|
||||
16, // [16:17] is the sub-list for method input_type
|
||||
16, // [16:16] is the sub-list for extension type_name
|
||||
16, // [16:16] is the sub-list for extension extendee
|
||||
0, // [0:16] is the sub-list for field type_name
|
||||
9, // 1: llmgw.v1.AgentMessage.heartbeat:type_name -> llmgw.v1.Heartbeat
|
||||
10, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus
|
||||
11, // 3: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
|
||||
14, // 4: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
|
||||
15, // 5: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
|
||||
16, // 6: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
|
||||
13, // 7: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights
|
||||
17, // 8: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
|
||||
6, // 9: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
|
||||
7, // 10: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
18, // 11: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
|
||||
5, // 12: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
|
||||
8, // 13: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights
|
||||
4, // 14: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage
|
||||
0, // 15: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
4, // 16: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage
|
||||
8, // 17: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights
|
||||
0, // 18: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
12, // 19: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
|
||||
19, // 20: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
|
||||
1, // 21: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
|
||||
2, // 22: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
|
||||
22, // [22:23] is the sub-list for method output_type
|
||||
21, // [21:22] is the sub-list for method input_type
|
||||
21, // [21:21] is the sub-list for extension type_name
|
||||
21, // [21:21] is the sub-list for extension extendee
|
||||
0, // [0:21] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_llmgw_v1_agent_proto_init() }
|
||||
@@ -1615,6 +1954,7 @@ func file_llmgw_v1_agent_proto_init() {
|
||||
(*ControlMessage_Stop)(nil),
|
||||
(*ControlMessage_Drain)(nil),
|
||||
(*ControlMessage_Ack)(nil),
|
||||
(*ControlMessage_RemoveWeights)(nil),
|
||||
}
|
||||
type x struct{}
|
||||
out := protoimpl.TypeBuilder{
|
||||
@@ -1622,7 +1962,7 @@ func file_llmgw_v1_agent_proto_init() {
|
||||
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
|
||||
RawDescriptor: unsafe.Slice(unsafe.StringData(file_llmgw_v1_agent_proto_rawDesc), len(file_llmgw_v1_agent_proto_rawDesc)),
|
||||
NumEnums: 1,
|
||||
NumMessages: 16,
|
||||
NumMessages: 19,
|
||||
NumExtensions: 0,
|
||||
NumServices: 1,
|
||||
},
|
||||
|
||||
@@ -33,6 +33,7 @@ message ControlMessage {
|
||||
StopShard stop = 2;
|
||||
DrainShard drain = 3;
|
||||
Ack ack = 4;
|
||||
RemoveWeights remove_weights = 5;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -67,6 +68,56 @@ message Register {
|
||||
// — so the new engines fail to allocate, with an error naming the wrong
|
||||
// process (plan §6.5).
|
||||
repeated RunningShard running_shards = 9;
|
||||
|
||||
// memory_total_mb is the machine's RAM, and 0 means the agent did not say —
|
||||
// which is what every agent older than this field reports, and is treated as
|
||||
// no constraint rather than as a node with no memory.
|
||||
//
|
||||
// Placement needs it because a CPU profile has no card to be free, so every
|
||||
// ready node looks like a candidate: a 14B model whose working set is 24 GB
|
||||
// would be placed just as happily on a node with 8 GB.
|
||||
int64 memory_total_mb = 10;
|
||||
|
||||
// weights are the model files this node can see, from the directories its
|
||||
// configuration names.
|
||||
//
|
||||
// How they got there is the operator's business and deliberately outside
|
||||
// llmgw's: a shared filesystem, an image exported over the fabric, rsync, or
|
||||
// a copy made by hand all look identical here. What matters to placement is
|
||||
// only whether they are present, because a node without them either cannot
|
||||
// start the replica at all or must first move tens of gigabytes.
|
||||
//
|
||||
// Empty means the node was not asked to look, which is not the same as having
|
||||
// nothing — an agent with no model directories configured behaves exactly as
|
||||
// every agent did before this existed, and constrains nothing.
|
||||
repeated LocalWeights weights = 11;
|
||||
|
||||
// storage is the disk behind those weights, one entry per filesystem.
|
||||
//
|
||||
// Reported so an operator can see a node filling up before a fetch dies at
|
||||
// ninety percent, and so placement can refuse a node that cannot hold a
|
||||
// model rather than discovering it during the download.
|
||||
//
|
||||
// Empty means the node was not asked to look, exactly as for weights, and
|
||||
// constrains nothing.
|
||||
repeated Storage storage = 12;
|
||||
}
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// Per filesystem rather than per configured directory: free space belongs to
|
||||
// the disk, and two model directories on one disk are one pool of bytes.
|
||||
message Storage {
|
||||
// path is one directory on this filesystem, so the operator recognises it.
|
||||
string path = 1;
|
||||
// free_bytes is what a non-root process can write — not the root reserve,
|
||||
// which is space the agent cannot have.
|
||||
int64 free_bytes = 2;
|
||||
int64 total_bytes = 3;
|
||||
// shared marks a filesystem the node did not fill. Free space on one is
|
||||
// shared with every other node using it, so two nodes each reporting 500 GB
|
||||
// free on the same filer have 500 GB between them and not a terabyte.
|
||||
bool shared = 4;
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
@@ -114,6 +165,10 @@ message GPU {
|
||||
repeated string problems = 12;
|
||||
int32 temperature_c = 13;
|
||||
int32 ecc_errors = 14;
|
||||
// utilization_percent is how busy the card was over the last sample period.
|
||||
// Zero when the driver is too old to report it, which is not the same as
|
||||
// idle: the metric is simply absent rather than wrong.
|
||||
int32 utilization_percent = 15;
|
||||
}
|
||||
|
||||
message FabricAttachment {
|
||||
@@ -126,12 +181,40 @@ message FabricAttachment {
|
||||
}
|
||||
|
||||
// Heartbeat says the node is alive and reports anything that changed cheaply.
|
||||
// LocalWeights is one model file or directory a node can see.
|
||||
message LocalWeights {
|
||||
// path is what an engine would be given, so it can be compared against a
|
||||
// model's source without interpretation.
|
||||
string path = 1;
|
||||
int64 size_bytes = 2;
|
||||
// shared marks weights on a filesystem the node did not have to fill —
|
||||
// reported so a fabric with a shared filer is not mistaken for one where
|
||||
// every node happens to have fetched its own copy.
|
||||
bool shared = 3;
|
||||
}
|
||||
|
||||
message Heartbeat {
|
||||
string node_id = 1;
|
||||
// gpu_memory_used_mb, indexed by GPU, catches a foreign process appearing on
|
||||
// a card after registration.
|
||||
repeated int32 gpu_memory_used_mb = 2;
|
||||
repeated string new_problems = 3;
|
||||
|
||||
// storage for the same reason the GPU memory is here: it changes while the
|
||||
// agent runs, and a figure sampled once at registration is a figure that is
|
||||
// wrong on exactly the node worth watching — the one busy filling up.
|
||||
//
|
||||
// Always sent, because it is a few numbers per filesystem.
|
||||
repeated Storage storage = 4;
|
||||
|
||||
// weights only when they changed, because the list is long enough that
|
||||
// sending it every few seconds would be paying for nothing most of the time.
|
||||
//
|
||||
// weights_changed is what says the list is meant: without it an empty list
|
||||
// would be ambiguous between "this node now holds nothing" and "nothing to
|
||||
// report", and the first would quietly empty a node's inventory.
|
||||
bool weights_changed = 5;
|
||||
repeated LocalWeights weights = 6;
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
@@ -200,6 +283,17 @@ message EngineProfile {
|
||||
string served_model_name = 11;
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
//
|
||||
// The paths are computed by the control plane from the catalog, but the agent
|
||||
// checks them again before unlinking anything: it deletes only files inside its
|
||||
// own cache directory whose names this system generated. A control plane that
|
||||
// is wrong, compromised or simply newer than the agent must not be able to
|
||||
// talk it into deleting a model an operator spent a week copying there.
|
||||
message RemoveWeights {
|
||||
repeated string paths = 1;
|
||||
}
|
||||
|
||||
// StopShard terminates a shard and frees its GPUs.
|
||||
message StopShard {
|
||||
string shard_id = 1;
|
||||
|
||||
Executable
+113
@@ -0,0 +1,113 @@
|
||||
#!/bin/sh
|
||||
# add-node.sh — build a node and put it in the fleet, in one command.
|
||||
#
|
||||
# provision-node.sh and onboard-node.sh stay separate because the two halves
|
||||
# need different access: one the hypervisor, one the certificate authority.
|
||||
# That is a real boundary and this script does not erase it — it drives both
|
||||
# where one person holds both, which is the common case and was the case here.
|
||||
#
|
||||
# The reason to have it is not keystrokes. The VM name, the hostname, the node
|
||||
# id and the certificate's common name must all be the same string, and the
|
||||
# address appears in three places; split across two invocations run minutes
|
||||
# apart, that agreement is maintained by whoever is typing. Here it is stated
|
||||
# once.
|
||||
set -eu
|
||||
|
||||
usage() {
|
||||
cat >&2 <<'USAGE'
|
||||
usage: add-node.sh --name NODE --ip ADDR/CIDR --server URL [options]
|
||||
|
||||
--name NODE node id, hostname and VM name; also the certificate's CN
|
||||
--ip ADDR/CIDR e.g. 10.0.0.145/24
|
||||
--server URL control plane, e.g. https://10.0.0.140:7443
|
||||
|
||||
--vmid N Proxmox VMID (default: the last octet of --ip)
|
||||
--cores N (default 4)
|
||||
--cpu TYPE QEMU CPU model (default host; see provision-node.sh)
|
||||
--memory MB (default 8192)
|
||||
--disk GB (default 20)
|
||||
--fabric NAME InfiniBand fabric; repeatable as NAME:hca1,hca2.
|
||||
Omit for a node with no fabric.
|
||||
--llamacpp-bin P path to llama-server on the node
|
||||
--vllm-bin P path to vllm
|
||||
--sglang-python P path to the python that has sglang
|
||||
--binary PATH llmgw binary to install (default: the one on PATH)
|
||||
--pki DIR certificate authority (default /etc/llmgw/pki)
|
||||
--ssh-user U (default ilker)
|
||||
--provision-only build the VM and stop
|
||||
--onboard-only skip provisioning; the machine already exists
|
||||
USAGE
|
||||
exit 2
|
||||
}
|
||||
|
||||
NAME= IP= SERVER= VMID= CORES=4 MEMORY=8192 DISK=20 CPU=
|
||||
BINARY= PKI=/etc/llmgw/pki SSH_USER=ilker
|
||||
PROVISION=1 ONBOARD=1
|
||||
set -- "$@"
|
||||
PASS_THROUGH=
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--name) NAME="$2"; shift 2;;
|
||||
--ip) IP="$2"; shift 2;;
|
||||
--server) SERVER="$2"; shift 2;;
|
||||
--vmid) VMID="$2"; shift 2;;
|
||||
--cores) CORES="$2"; shift 2;;
|
||||
--cpu) CPU="$2"; shift 2;;
|
||||
--memory) MEMORY="$2"; shift 2;;
|
||||
--disk) DISK="$2"; shift 2;;
|
||||
--binary) BINARY="$2"; shift 2;;
|
||||
--pki) PKI="$2"; shift 2;;
|
||||
--ssh-user) SSH_USER="$2"; shift 2;;
|
||||
--provision-only) ONBOARD=0; shift;;
|
||||
--onboard-only) PROVISION=0; shift;;
|
||||
# Everything the onboarding step understands and this one does not need to.
|
||||
--fabric|--llamacpp-bin|--vllm-bin|--sglang-python|--ib-root|--nvidia-smi|--mgmt-addr|--mgmt-if)
|
||||
PASS_THROUGH="$PASS_THROUGH $1 $2"; shift 2;;
|
||||
-h|--help) usage;;
|
||||
*) echo "unknown option: $1" >&2; usage;;
|
||||
esac
|
||||
done
|
||||
[ -n "$NAME" ] && [ -n "$IP" ] && [ -n "$SERVER" ] || usage
|
||||
|
||||
# The address without its prefix, which is what onboarding and ssh want.
|
||||
ADDR="${IP%%/*}"
|
||||
# The last octet as the VMID by default: it keeps `qm list` and `ctl nodes` in
|
||||
# the same order, and one fewer number to choose is one fewer to get wrong.
|
||||
[ -n "$VMID" ] || VMID="${ADDR##*.}"
|
||||
|
||||
here() { dirname "$0"; }
|
||||
PROVISION_SH="$(here)/provision-node.sh"
|
||||
ONBOARD_SH="$(here)/onboard-node.sh"
|
||||
[ -x "$PROVISION_SH" ] || PROVISION_SH="$(command -v llmgw-provision-node || true)"
|
||||
[ -x "$ONBOARD_SH" ] || ONBOARD_SH="$(command -v llmgw-onboard-node || true)"
|
||||
|
||||
if [ "$PROVISION" = 1 ]; then
|
||||
[ -x "$PROVISION_SH" ] || { echo "cannot find provision-node.sh" >&2; exit 1; }
|
||||
echo "== 1/2 building $NAME as VM $VMID"
|
||||
# shellcheck disable=SC2086
|
||||
"$PROVISION_SH" --vmid "$VMID" --name "$NAME" --ip "$IP" \
|
||||
--cores "$CORES" --memory "$MEMORY" --disk "$DISK" ${CPU:+--cpu "$CPU"}
|
||||
fi
|
||||
|
||||
if [ "$ONBOARD" = 0 ]; then
|
||||
echo "$NAME built. Onboard it from the CA host when you are ready."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
[ -x "$ONBOARD_SH" ] || { echo "cannot find onboard-node.sh" >&2; exit 1; }
|
||||
echo
|
||||
echo "== 2/2 onboarding $NAME"
|
||||
# Onboarding waits for ssh itself, which is what makes running the two halves
|
||||
# back to back safe: cloud-init is still regenerating host keys and restarting
|
||||
# sshd for the first minute of a new machine's life.
|
||||
# shellcheck disable=SC2086
|
||||
"$ONBOARD_SH" --node "$NAME" --addr "$ADDR" --server "$SERVER" \
|
||||
--pki "$PKI" --ssh-user "$SSH_USER" \
|
||||
${BINARY:+--binary "$BINARY"} $PASS_THROUGH
|
||||
|
||||
echo
|
||||
echo "Check it with: llmgw ctl nodes"
|
||||
echo
|
||||
echo "A node is not finished until the engines its profiles name are installed"
|
||||
echo "and the weights are on disk. Onboarding checks that a configured engine"
|
||||
echo "runs, but it installs neither — see docs/node-onboarding.md."
|
||||
Executable
+270
@@ -0,0 +1,270 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Give a machine an identity in an llmgw cluster.
|
||||
#
|
||||
# Run on the host that holds the certificate authority, because that is the one
|
||||
# thing that cannot be delegated: this script issues the credential by which the
|
||||
# control plane will believe the node is who it says it is.
|
||||
#
|
||||
# What it does, in the order it matters:
|
||||
# 1. issue a client certificate whose common name is the node id
|
||||
# 2. copy the CA root, the certificate and its key to the machine
|
||||
# 3. trust the root system-wide, so other tooling there stops complaining
|
||||
# 4. install the binary, the configuration and the unit
|
||||
# 5. start the agent
|
||||
#
|
||||
# The node id, the VM name, the hostname and the certificate common name are all
|
||||
# the same string. They do not have to be, and every time they are not, someone
|
||||
# spends an afternoon working out which of the four a message is talking about.
|
||||
set -euo pipefail
|
||||
|
||||
usage() {
|
||||
cat >&2 <<USAGE
|
||||
usage: $0 --node ID --addr HOST --server URL [options]
|
||||
|
||||
--node ID node id, hostname and certificate common name
|
||||
--addr HOST how to reach the machine over ssh
|
||||
--server URL the control plane's agent endpoint, e.g. https://10.0.0.140:7443
|
||||
--fabric NAME InfiniBand fabric this node is attached to; repeatable as
|
||||
--fabric NAME:hca1,hca2 . Omit for a node with no fabric,
|
||||
which can host single-node replicas and will never be given
|
||||
a multi-node one.
|
||||
--mgmt-addr ADDR address engines are reachable on (default: --addr)
|
||||
--mgmt-if NAME interface NCCL bootstraps over (default eth0)
|
||||
--llamacpp-bin P path to llama-server
|
||||
--vllm-bin P path to the vllm entry point
|
||||
--sglang-python P python interpreter of the SGLang virtualenv
|
||||
--nvidia-smi PATH where nvidia-smi lives, when it is not on PATH. Containerized
|
||||
driver installs routinely put it somewhere else.
|
||||
--ib-root DIR the sysfs InfiniBand directory (default /sys/class/infiniband)
|
||||
--pki DIR where the authority lives (default /etc/llmgw/pki)
|
||||
--binary PATH llmgw binary to install (default: the one running this).
|
||||
Releases carry static linux/amd64 and linux/arm64 builds;
|
||||
downloading one is better than copying whatever happens to
|
||||
be on the machine you are running this from.
|
||||
--ssh-user NAME (default ilker)
|
||||
USAGE
|
||||
exit 2
|
||||
}
|
||||
|
||||
NODE= ADDR= SERVER= MGMT= MGMT_IF=eth0
|
||||
NVIDIA_SMI= IB_ROOT=/sys/class/infiniband
|
||||
LLAMACPP_BIN= VLLM_BIN= SGLANG_PYTHON=
|
||||
SSH_WAIT=180
|
||||
PKI=/etc/llmgw/pki BINARY= SSH_USER=ilker
|
||||
FABRIC_NAME= FABRIC_HCAS="mlx5_0,mlx5_1"
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--node) NODE="$2"; shift 2;;
|
||||
--addr) ADDR="$2"; shift 2;;
|
||||
--server) SERVER="$2"; shift 2;;
|
||||
--fabric) FABRIC_NAME="${2%%:*}"; [ "${2#*:}" != "$2" ] && FABRIC_HCAS="${2#*:}"; shift 2;;
|
||||
--mgmt-addr) MGMT="$2"; shift 2;;
|
||||
--mgmt-if) MGMT_IF="$2"; shift 2;;
|
||||
--llamacpp-bin) LLAMACPP_BIN="$2"; shift 2;;
|
||||
--vllm-bin) VLLM_BIN="$2"; shift 2;;
|
||||
--sglang-python) SGLANG_PYTHON="$2"; shift 2;;
|
||||
--nvidia-smi) NVIDIA_SMI="$2"; shift 2;;
|
||||
--ib-root) IB_ROOT="$2"; shift 2;;
|
||||
--pki) PKI="$2"; shift 2;;
|
||||
--binary) BINARY="$2"; shift 2;;
|
||||
--ssh-user) SSH_USER="$2"; shift 2;;
|
||||
*) usage;;
|
||||
esac
|
||||
done
|
||||
[ -n "$NODE" ] && [ -n "$ADDR" ] && [ -n "$SERVER" ] || usage
|
||||
MGMT="${MGMT:-$ADDR}"
|
||||
BINARY="${BINARY:-$(command -v llmgw)}"
|
||||
[ -x "$BINARY" ] || { echo "no llmgw binary; pass --binary" >&2; exit 1; }
|
||||
[ -f "$PKI/ca-key.pem" ] || { echo "no authority in $PKI; run 'llmgw ctl cert init-ca' first" >&2; exit 1; }
|
||||
|
||||
# accept-new, not "no": a machine provisioned a minute ago has a host key
|
||||
# nobody has seen, and refusing it stopped onboarding dead with "Host key
|
||||
# verification failed" — an error about ssh, in the middle of a script whose
|
||||
# whole job is to talk to a new machine. accept-new learns an unknown key and
|
||||
# still refuses a changed one, which is the distinction that matters: the risk
|
||||
# being guarded against is a substituted machine, not a new one.
|
||||
SSHOPTS="-o BatchMode=yes -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new"
|
||||
SSH="ssh $SSHOPTS $SSH_USER@$ADDR"
|
||||
say() { printf ' %s\n' "$*"; }
|
||||
|
||||
echo "onboarding $NODE at $ADDR"
|
||||
|
||||
# Waiting rather than failing: cloud-init restarts sshd after it regenerates
|
||||
# the host keys, so a machine that answered a moment ago refuses connections
|
||||
# for a few seconds. Observed while onboarding node05 — the agent was already
|
||||
# registered and healthy while ssh was refusing.
|
||||
wait_for_ssh() {
|
||||
deadline=$(( $(date +%s) + SSH_WAIT ))
|
||||
until $SSH true 2>/dev/null; do
|
||||
[ "$(date +%s)" -lt "$deadline" ] || {
|
||||
echo "cannot reach $SSH_USER@$ADDR over ssh after ${SSH_WAIT}s" >&2
|
||||
exit 1
|
||||
}
|
||||
sleep 5
|
||||
done
|
||||
}
|
||||
wait_for_ssh
|
||||
|
||||
# 1. The certificate. Issued fresh every time: re-onboarding a machine should
|
||||
# not depend on finding an old key, and the old one keeps working until it
|
||||
# is revoked, which is a separate decision.
|
||||
say "issuing a certificate"
|
||||
TMP="$(mktemp -d)"
|
||||
trap 'rm -rf "$TMP"' EXIT
|
||||
"$BINARY" ctl cert issue-client "$NODE" --dir "$PKI" --out "$TMP" >/dev/null
|
||||
|
||||
# 2 and 3. Identity, and trusting the root so future tooling on that machine
|
||||
# does not treat our certificates as strangers.
|
||||
say "installing the certificate and trusting the root"
|
||||
scp -q $SSHOPTS "$PKI/ca.pem" "$TMP/$NODE.pem" "$TMP/$NODE-key.pem" "$SSH_USER@$ADDR:/tmp/"
|
||||
$SSH "set -e
|
||||
sudo install -d -m 0755 /etc/llmgw
|
||||
sudo install -m 0644 /tmp/ca.pem /etc/llmgw/ca.pem
|
||||
sudo install -m 0644 /tmp/$NODE.pem /etc/llmgw/agent.pem
|
||||
sudo install -m 0600 /tmp/$NODE-key.pem /etc/llmgw/agent-key.pem
|
||||
sudo chown root:root /etc/llmgw/*
|
||||
sudo install -m 0644 /tmp/ca.pem /usr/local/share/ca-certificates/llmgw-ca.crt
|
||||
sudo update-ca-certificates >/dev/null 2>&1 || true
|
||||
rm -f /tmp/ca.pem /tmp/$NODE.pem /tmp/$NODE-key.pem"
|
||||
|
||||
# 4. The binary, the configuration and the unit.
|
||||
say "installing the binary"
|
||||
scp -q $SSHOPTS "$BINARY" "$SSH_USER@$ADDR:/tmp/llmgw"
|
||||
$SSH 'sudo install -m 0755 /tmp/llmgw /usr/local/bin/llmgw && rm -f /tmp/llmgw'
|
||||
|
||||
say "writing the configuration"
|
||||
FABRICS="fabrics: []"
|
||||
if [ -n "$FABRIC_NAME" ]; then
|
||||
FABRICS="fabrics:
|
||||
- name: $FABRIC_NAME
|
||||
hcas: [$(echo "$FABRIC_HCAS" | sed 's/,/, /g')]"
|
||||
fi
|
||||
cat > "$TMP/agent.yaml" <<YAML
|
||||
# $NODE
|
||||
#
|
||||
# Written by onboard-node.sh. Fabric membership is configuration and there is no
|
||||
# way around it: OpenSM's default subnet prefix is the same on every fabric, so
|
||||
# a matching prefix proves nothing and a wrong guess produces a replica that
|
||||
# starts, passes health checks and then hangs on its first all-reduce.
|
||||
node_id: $NODE
|
||||
server: "$SERVER"
|
||||
|
||||
tls:
|
||||
ca: /etc/llmgw/ca.pem
|
||||
cert: /etc/llmgw/agent.pem
|
||||
key: /etc/llmgw/agent-key.pem
|
||||
|
||||
mgmt_addr: "$MGMT"
|
||||
mgmt_interface: $MGMT_IF
|
||||
|
||||
heartbeat_interval: 5s
|
||||
reconnect_base: 1s
|
||||
reconnect_max: 30s
|
||||
|
||||
$FABRICS
|
||||
|
||||
discovery:
|
||||
# Empty means look on PATH. Worth setting where the driver lives outside the
|
||||
# default paths, which containerized installs routinely do.
|
||||
nvidia_smi: "$NVIDIA_SMI"
|
||||
ib_root: $IB_ROOT
|
||||
prefer_nvml: false
|
||||
skip_gpus: false
|
||||
|
||||
# Where the engines live. Empty means "on PATH", which is rarely true: a node
|
||||
# that cannot find its engine registers healthy and then fails every shard
|
||||
# placed on it with "executable file not found".
|
||||
engines:
|
||||
llamacpp:
|
||||
bin: "$LLAMACPP_BIN"
|
||||
parallel: 4
|
||||
vllm:
|
||||
bin: "$VLLM_BIN"
|
||||
launch_mode: ray
|
||||
sglang:
|
||||
python: "$SGLANG_PYTHON"
|
||||
|
||||
supervisor:
|
||||
work_dir: /var/lib/llmgw-agent
|
||||
log_lines: 512
|
||||
stop_timeout: 60s
|
||||
health_interval: 2s
|
||||
YAML
|
||||
scp -q "$TMP/agent.yaml" "$SSH_USER@$ADDR:/tmp/agent.yaml"
|
||||
|
||||
cat > "$TMP/llmgw-agent.service" <<'UNIT'
|
||||
[Unit]
|
||||
Description=llmgw node agent
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
ExecStart=/usr/local/bin/llmgw agent --config /etc/llmgw/agent.yaml
|
||||
# Root, because the agent launches engines and needs device access. Engines
|
||||
# outlive a restart of this unit and are adopted again on the next start, so
|
||||
# KillMode must not take the whole process group with it.
|
||||
KillMode=process
|
||||
Restart=on-failure
|
||||
RestartSec=2s
|
||||
StateDirectory=llmgw-agent
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
UNIT
|
||||
scp -q "$TMP/llmgw-agent.service" "$SSH_USER@$ADDR:/tmp/"
|
||||
|
||||
# An engine path that is wrong, or right but unrunnable, is the failure this
|
||||
# node type is most likely to have, and it is invisible: the agent registers
|
||||
# healthy and every shard placed on it dies with "executable file not found"
|
||||
# or a missing shared library. Checked here because it is cheap here and
|
||||
# expensive later — llama-server on node05 was present, executable, and could
|
||||
# not start for want of libgomp1, and the node reported ready throughout.
|
||||
check_engine() {
|
||||
[ -n "$2" ] || return 0
|
||||
if $SSH "test -x '$2'" 2>/dev/null; then
|
||||
if $SSH "'$2' --version >/dev/null 2>&1 || '$2' --help >/dev/null 2>&1"; then
|
||||
say "$1: $2 runs"
|
||||
return 0
|
||||
fi
|
||||
ENGINE_PROBLEMS="$ENGINE_PROBLEMS
|
||||
$1 at $2 exists but will not run. Usually a missing shared library:
|
||||
ssh $SSH_USER@$ADDR ldd $2 | grep 'not found'"
|
||||
return 0
|
||||
fi
|
||||
ENGINE_PROBLEMS="$ENGINE_PROBLEMS
|
||||
$1 is configured at $2, which is not an executable on $NODE."
|
||||
}
|
||||
ENGINE_PROBLEMS=
|
||||
if [ -n "$LLAMACPP_BIN$VLLM_BIN$SGLANG_PYTHON" ]; then
|
||||
say "checking the engines"
|
||||
check_engine llamacpp "$LLAMACPP_BIN"
|
||||
check_engine vllm "$VLLM_BIN"
|
||||
check_engine sglang "$SGLANG_PYTHON"
|
||||
fi
|
||||
|
||||
say "starting the agent"
|
||||
$SSH 'set -e
|
||||
sudo install -m 0644 -o root -g root /tmp/agent.yaml /etc/llmgw/agent.yaml
|
||||
sudo install -m 0644 /tmp/llmgw-agent.service /etc/systemd/system/llmgw-agent.service
|
||||
rm -f /tmp/agent.yaml /tmp/llmgw-agent.service
|
||||
sudo systemctl daemon-reload
|
||||
sudo systemctl enable --now llmgw-agent >/dev/null
|
||||
sudo systemctl restart llmgw-agent'
|
||||
|
||||
echo
|
||||
if [ -z "$LLAMACPP_BIN$VLLM_BIN$SGLANG_PYTHON" ]; then
|
||||
say "no engine paths given; this node can only run the 'sim' engine until"
|
||||
say "one is installed and the configuration points at it"
|
||||
fi
|
||||
if [ -n "$ENGINE_PROBLEMS" ]; then
|
||||
echo "$NODE is onboarded but an engine will not run:$ENGINE_PROBLEMS"
|
||||
echo
|
||||
echo "The node will register as ready regardless. Fix this before a replica"
|
||||
echo "is placed on it, or the shard fails and the reason is three logs away."
|
||||
echo
|
||||
fi
|
||||
echo "$NODE onboarded. It should appear in 'llmgw ctl nodes' within a few seconds."
|
||||
echo "If it does not, the control plane's log says why — a certificate whose"
|
||||
echo "common name does not match the claimed node id is refused by design."
|
||||
Executable
+96
@@ -0,0 +1,96 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Create a VM on Proxmox for a new llmgw node.
|
||||
#
|
||||
# Run on the Proxmox host. It does the half of onboarding that is about
|
||||
# hardware; onboard-node.sh does the half that is about identity, and the two
|
||||
# are separate because they are done by different people with different access.
|
||||
#
|
||||
# Idempotent only in the sense that it refuses to touch an existing VMID. A node
|
||||
# that already exists is not something a provisioning script should have an
|
||||
# opinion about.
|
||||
set -euo pipefail
|
||||
|
||||
usage() {
|
||||
cat >&2 <<USAGE
|
||||
usage: $0 --vmid N --name NAME --ip ADDR [options]
|
||||
|
||||
--vmid N Proxmox VM id; must be free
|
||||
--name NAME VM name; use the node id, so the machine, the hostname and
|
||||
the certificate common name are the same string
|
||||
--ip ADDR address with prefix, e.g. 10.0.0.144/24
|
||||
--gw ADDR gateway (default 10.0.0.1)
|
||||
--cores N (default 4)
|
||||
--cpu TYPE QEMU CPU model (default host)
|
||||
--memory MB (default 8192)
|
||||
--disk GB (default 20)
|
||||
--bridge NAME (default vmbr0)
|
||||
--storage NAME (default local-lvm)
|
||||
--image PATH cloud image (default Ubuntu 24.04 in the ISO store)
|
||||
--user NAME cloud-init user (default ilker)
|
||||
--sshkeys PATH authorized_keys to install (default /root/.ssh/authorized_keys)
|
||||
USAGE
|
||||
exit 2
|
||||
}
|
||||
|
||||
VMID= NAME= IP= GW=10.0.0.1 CORES=4 MEMORY=8192 DISK=20
|
||||
# host, not Proxmox's kvm64 default. kvm64 hides AVX, AVX2, FMA and even
|
||||
# SSE4.2 from the guest, and an inference node without vector instructions is
|
||||
# not slightly slower, it is unusable: llama.cpp falls back to its scalar
|
||||
# backend and a 1.5B model dropped from 39 tokens a second to 5.8, with prompt
|
||||
# processing 19x worse. It cost an afternoon to find because nothing is broken
|
||||
# — every number is simply small.
|
||||
#
|
||||
# The trade is live migration: a VM pinned to the host CPU cannot move to a
|
||||
# machine with a different one. For a node whose job is arithmetic, that is
|
||||
# the right way round.
|
||||
CPU=host
|
||||
BRIDGE=vmbr0 STORAGE=local-lvm USER=ilker
|
||||
IMAGE=/var/lib/vz/template/iso/noble-server-cloudimg-amd64.img
|
||||
SSHKEYS=/root/.ssh/authorized_keys
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--vmid) VMID="$2"; shift 2;;
|
||||
--name) NAME="$2"; shift 2;;
|
||||
--ip) IP="$2"; shift 2;;
|
||||
--gw) GW="$2"; shift 2;;
|
||||
--cores) CORES="$2"; shift 2;;
|
||||
--cpu) CPU="$2"; shift 2;;
|
||||
--memory) MEMORY="$2"; shift 2;;
|
||||
--disk) DISK="$2"; shift 2;;
|
||||
--bridge) BRIDGE="$2"; shift 2;;
|
||||
--storage) STORAGE="$2"; shift 2;;
|
||||
--image) IMAGE="$2"; shift 2;;
|
||||
--user) USER="$2"; shift 2;;
|
||||
--sshkeys) SSHKEYS="$2"; shift 2;;
|
||||
*) usage;;
|
||||
esac
|
||||
done
|
||||
[ -n "$VMID" ] && [ -n "$NAME" ] && [ -n "$IP" ] || usage
|
||||
|
||||
if qm status "$VMID" >/dev/null 2>&1; then
|
||||
echo "VM $VMID already exists. Pick a free id; this script will not adopt one." >&2
|
||||
exit 1
|
||||
fi
|
||||
[ -f "$IMAGE" ] || { echo "no cloud image at $IMAGE" >&2; exit 1; }
|
||||
|
||||
echo "creating $NAME as VM $VMID at $IP"
|
||||
qm create "$VMID" \
|
||||
--name "$NAME" --cores "$CORES" --memory "$MEMORY" --cpu "$CPU" \
|
||||
--net0 "virtio,bridge=$BRIDGE" --ostype l26 --agent enabled=1
|
||||
|
||||
# Imported rather than installed: a cloud image plus cloud-init is the whole
|
||||
# reason a node takes a minute to build instead of an afternoon.
|
||||
qm importdisk "$VMID" "$IMAGE" "$STORAGE" >/dev/null
|
||||
qm set "$VMID" --scsihw virtio-scsi-pci --scsi0 "$STORAGE:vm-$VMID-disk-0,discard=on,ssd=1" >/dev/null
|
||||
qm disk resize "$VMID" scsi0 "${DISK}G" >/dev/null
|
||||
qm set "$VMID" --ide2 "$STORAGE:cloudinit" --boot order=scsi0 >/dev/null
|
||||
qm set "$VMID" \
|
||||
--ipconfig0 "ip=$IP,gw=$GW" \
|
||||
--nameserver "$GW 8.8.8.8" \
|
||||
--ciuser "$USER" \
|
||||
--sshkeys "$SSHKEYS" >/dev/null
|
||||
|
||||
qm start "$VMID"
|
||||
echo "started. cloud-init takes a minute; then run onboard-node.sh from the CA host."
|
||||
@@ -7,6 +7,21 @@
|
||||
node_id: gpu-a-01 # defaults to the hostname
|
||||
server: "http://llmgw-server.internal:7443"
|
||||
|
||||
# The agent's half of mutual TLS. The certificate's common name must be this
|
||||
# node's id, because the control plane compares the two and refuses the stream
|
||||
# when they disagree:
|
||||
#
|
||||
# llmgw ctl cert issue-client gpu-a-01 --dir /etc/llmgw/pki
|
||||
#
|
||||
# With this set, `server` above must be an https:// address — an agent holding
|
||||
# certificates and dialling http:// is refused at startup rather than connecting
|
||||
# in the clear without noticing.
|
||||
#
|
||||
# tls:
|
||||
# ca: /etc/llmgw/ca.pem
|
||||
# cert: /etc/llmgw/agent.pem
|
||||
# key: /etc/llmgw/agent-key.pem
|
||||
|
||||
# The address engines are reachable on, and the address NCCL bootstraps over.
|
||||
# Never a fabric address: NCCL rendezvouses over TCP before moving to RDMA.
|
||||
mgmt_addr: "10.0.0.11"
|
||||
@@ -42,6 +57,23 @@ discovery:
|
||||
prefer_nvml: true
|
||||
skip_gpus: false
|
||||
|
||||
# Where this node keeps model weights.
|
||||
#
|
||||
# Unset by default, and unset means the node is not asked to look — which is
|
||||
# not the same as having nothing, and constrains nothing. Storage differs at
|
||||
# every installation and llmgw deliberately does not own it: a Lustre mount, an
|
||||
# image exported over the fabric with SRP or NVMe-oF, local NVMe, or a copy
|
||||
# made by hand all look identical from here.
|
||||
#
|
||||
# What it buys is placement preferring a node that already holds the model over
|
||||
# one that must first move tens of gigabytes. That matters most where fabrics
|
||||
# are isolated from one another, because weights in one fabric are not weights
|
||||
# in another and the only path between them is the management network.
|
||||
#
|
||||
# model_dirs:
|
||||
# - /var/lib/llmgw-models
|
||||
# - /mnt/lustre/models
|
||||
|
||||
engines:
|
||||
vllm:
|
||||
bin: /opt/venvs/vllm/bin/vllm
|
||||
|
||||
+155
-6
@@ -7,12 +7,92 @@
|
||||
# __ for nesting: LLMGW_ROUTER__RETRIES=3
|
||||
|
||||
listen: ":8080" # client-facing OpenAI API
|
||||
admin_listen: ":9090" # metrics, health, pprof — do not expose this publicly
|
||||
admin_listen: "127.0.0.1:9090" # dashboard, management API, pprof — keep it local
|
||||
# Metrics only, on their own port, so Prometheus can reach them without pprof
|
||||
# being reachable too. Empty means off; metrics stay on admin_listen either way.
|
||||
# metrics_listen: "0.0.0.0:9091"
|
||||
# Where node agents connect. Separate from the client port so the fleet-facing
|
||||
# surface can be firewalled independently. Agents speak plaintext HTTP/2 today,
|
||||
# so this must stay on a trusted network until mTLS lands (plan §5.6).
|
||||
# surface can be firewalled independently.
|
||||
agent_listen: ":7443"
|
||||
|
||||
# Mutual TLS on the agent port (plan §5.6). With it, a node's identity is taken
|
||||
# from its certificate's common name and an agent claiming a different node id
|
||||
# is refused. Without it the port is plaintext HTTP/2 and anyone who can reach
|
||||
# it can claim to be any node, so it must stay on a trusted network.
|
||||
#
|
||||
# These certificates come from llmgw's own authority, not a public one: public
|
||||
# authorities stopped issuing the clientAuth usage, so no Let's Encrypt
|
||||
# certificate can authenticate an agent.
|
||||
#
|
||||
# llmgw ctl cert init-ca --dir /etc/llmgw/pki
|
||||
# llmgw ctl cert issue-server llmgw-ctl --host 10.0.0.140 --dir /etc/llmgw/pki
|
||||
#
|
||||
# Every address an agent might dial has to appear in --host, bare IP addresses
|
||||
# included: hostname verification reads subject alternative names and ignores
|
||||
# the common name entirely.
|
||||
#
|
||||
# This covers the fleet-facing port only. The client-facing API above is public
|
||||
# and belongs behind a publicly trusted certificate.
|
||||
#
|
||||
# agent_tls:
|
||||
# ca: /etc/llmgw/ca.pem
|
||||
# cert: /etc/llmgw/server.pem
|
||||
# key: /etc/llmgw/server-key.pem
|
||||
|
||||
# The public page at "/". Someone who reaches this address without a key needs a
|
||||
# route to getting one, and the gateway cannot guess who to send them to.
|
||||
#
|
||||
# landing:
|
||||
# contact: "platform@example.com" # or an https:// URL, or "#platform on Slack"
|
||||
# notice: "Internal use only."
|
||||
|
||||
# API keys on the client-facing port (plan §7.9). Off by default, because a
|
||||
# gateway on a trusted network with no keys issued still has to serve — but
|
||||
# nothing else protects that port, so turning it on is what makes it safe to
|
||||
# expose.
|
||||
#
|
||||
# llmgw keys init --store /var/lib/llmgw/llmgw.db
|
||||
# llmgw keys create --store /var/lib/llmgw/llmgw.db --project default "laptop"
|
||||
#
|
||||
# The key is printed once. Only its hash is stored, so there is nothing to
|
||||
# recover it from.
|
||||
#
|
||||
# auth:
|
||||
# enabled: true
|
||||
# store: /var/lib/llmgw/llmgw.db
|
||||
# # How long a resolved key is trusted before being read again. This is the
|
||||
# # price of keeping the database off the request path: it bounds how long a
|
||||
# # revoked key keeps working.
|
||||
# cache_ttl: 30s
|
||||
# # How long a rejected key is remembered as rejected, so a client looping
|
||||
# # with a bad key cannot become a database read per request.
|
||||
# negative_ttl: 5s
|
||||
#
|
||||
# # Serve the signup form. What it creates is a pending account that can do
|
||||
# # nothing until `llmgw users approve` places it in a project, so this is not
|
||||
# # a way in — it only moves key creation from you to the person who needs one.
|
||||
# # Off, there is no form and you issue keys with `llmgw keys create --claim`.
|
||||
# signup_open: true
|
||||
#
|
||||
# # What each model costs, in micro-units per thousand tokens. A unit is
|
||||
# # whatever you decide: self-hosted models have no per-token price, so it can
|
||||
# # be GPU-seconds, a notional rate, or 1000 everywhere — which turns budgets
|
||||
# # into plain token quotas.
|
||||
# #
|
||||
# # A model priced at zero is exempt from budget checks, which is what lets a
|
||||
# # key with nothing left still be served a cheap fallback.
|
||||
# default_price: { prompt_per_1k: 0, completion_per_1k: 0 }
|
||||
# pricing:
|
||||
# llama-3.3-70b: { prompt_per_1k: 1000, completion_per_1k: 3000 }
|
||||
# qwen2.5-32b: { prompt_per_1k: 400, completion_per_1k: 1200 }
|
||||
#
|
||||
# # Usage rows are written in batches off the request path. Beyond buffer_size
|
||||
# # they are dropped rather than blocking a request: losing some accounting is
|
||||
# # bad, turning a stalled database into an outage is worse.
|
||||
# flush_interval: 200ms
|
||||
# batch_size: 500
|
||||
# buffer_size: 8192
|
||||
|
||||
router:
|
||||
scrape_interval: 2s
|
||||
scrape_timeout: 1s
|
||||
@@ -33,6 +113,39 @@ router:
|
||||
failure_threshold: 3
|
||||
|
||||
default_profile: chat # chat | batch | long-context | debug
|
||||
|
||||
# Approximate prefix routing (plan §7.5 tier 1): send a request to the replica
|
||||
# that has probably already prefilled its prefix.
|
||||
#
|
||||
# It earns its keep when the working set does not fit in one replica's cache.
|
||||
# With a handful of prompts and a large cache, scattering is free — every
|
||||
# replica ends up holding everything — and this only adds constraint. Measured
|
||||
# with sim-engines whose caches were too small for the working set: cache hits
|
||||
# went from 22.5% to 81.7%.
|
||||
#
|
||||
# Watch llmgw_prefix_decisions_total to tune it. A high
|
||||
# load_balance_override count means the guard is doing its job and affinity is
|
||||
# not getting a chance; a high cache_miss_min_load means match_ratio or
|
||||
# min_match is too strict for your traffic.
|
||||
prefix:
|
||||
enabled: false
|
||||
# The shortest match worth acting on, in bytes. Below this the overlap is
|
||||
# coincidence — two English prompts share "The".
|
||||
min_match: 64
|
||||
# The share of the prompt that must match. A short shared header on a very
|
||||
# long prompt saves almost no prefill.
|
||||
match_ratio: 0.3
|
||||
# The imbalance guard, and both must trip before affinity is abandoned.
|
||||
# Absolute alone fires on an idle cluster where one replica has two requests
|
||||
# and another none; relative alone fires on a busy one where 200 against 180
|
||||
# is a rounding error. Without the guard the cache hit rate rises while p99
|
||||
# collapses.
|
||||
balance_abs: 8
|
||||
balance_rel: 1.5
|
||||
# An engine evicts on its own schedule, so an old record may describe a
|
||||
# cache that is gone.
|
||||
ttl: 10m
|
||||
max_nodes: 100000
|
||||
session_header: x-session-id # affinity key; falls back to the `user` field
|
||||
# Ask engines for usage even when the client did not, so token accounting
|
||||
# works regardless of client behaviour. The extra chunk is dropped on the way
|
||||
@@ -52,9 +165,9 @@ router:
|
||||
# traffic, because an engine queue is FIFO and knows nothing about tenants.
|
||||
admission:
|
||||
enabled: true
|
||||
# Until authentication lands, the band comes from a header and is trusted.
|
||||
# Do not expose this port to callers who should not be able to promote
|
||||
# themselves.
|
||||
# With auth on, the key decides the band and this header is ignored. With
|
||||
# auth off the header is trusted, so any client can promote itself into the
|
||||
# top band — do not expose the port in that configuration.
|
||||
band_header: x-llmgw-band
|
||||
default_band: gold
|
||||
# A band waits once pool saturation passes its ceiling, so the cluster
|
||||
@@ -96,15 +209,51 @@ control:
|
||||
port_range: [20000, 20999]
|
||||
mgmt_interface: eth0
|
||||
|
||||
# Each model needs a source — a weights path or a HuggingFace id — and at least
|
||||
# one profile if the control plane is to place it. Without profiles the model is
|
||||
# catalogued but never started, which is the router-only mode below.
|
||||
models:
|
||||
- name: llama-3.3-70b
|
||||
source: meta-llama/Llama-3.3-70B-Instruct
|
||||
aliases: [llama]
|
||||
fallbacks: [qwen2.5-32b]
|
||||
routing_profile: chat
|
||||
scaling: { min: 1, max: 4 }
|
||||
profiles:
|
||||
# Two nodes of four H100s each. A multi-node profile requires a fabric,
|
||||
# and a replica never spans two of them.
|
||||
- name: tp8
|
||||
engine: vllm
|
||||
gpu_class: h100-80g
|
||||
tp: 8
|
||||
nodes: 2
|
||||
gpus_per_node: 4
|
||||
priority: 100
|
||||
weight: 8
|
||||
startup_timeout: 30m
|
||||
# Where a whole node is free, one node is cheaper to schedule and faster
|
||||
# to start than two.
|
||||
- name: tp4-a100
|
||||
engine: vllm
|
||||
gpu_class: a100-40g
|
||||
tp: 4
|
||||
priority: 50
|
||||
weight: 3
|
||||
startup_timeout: 20m
|
||||
|
||||
- name: qwen2.5-32b
|
||||
source: Qwen/Qwen2.5-32B-Instruct
|
||||
aliases: [qwen]
|
||||
routing_profile: batch
|
||||
scaling: { min: 1, max: 2 }
|
||||
profiles:
|
||||
- name: tp2
|
||||
engine: sglang
|
||||
gpu_class: h100-80g
|
||||
tp: 2
|
||||
priority: 100
|
||||
weight: 2
|
||||
startup_timeout: 15m
|
||||
|
||||
# A static replica list, used only when no model declares a profile. That is the
|
||||
# router-only mode for testing against sim-engines; with the control plane
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
# Scraping llmgw.
|
||||
#
|
||||
# Point this at metrics_listen, not admin_listen. They are different ports for
|
||||
# a reason: the admin listener also carries pprof, which can dump the process
|
||||
# heap, and API keys are in that heap in plaintext.
|
||||
#
|
||||
# In the control plane's configuration:
|
||||
#
|
||||
# admin_listen: "127.0.0.1:9090" # dashboard, management API, pprof
|
||||
# metrics_listen: "0.0.0.0:9091" # metrics only, safe to reach
|
||||
#
|
||||
# One target for the whole fleet. Node and GPU figures are exported by the
|
||||
# control plane, which already holds them — agents report hardware at
|
||||
# registration and state on every heartbeat — so there is nothing to scrape on
|
||||
# the nodes and no Prometheus configuration per machine.
|
||||
global:
|
||||
scrape_interval: 15s
|
||||
|
||||
scrape_configs:
|
||||
- job_name: llmgw
|
||||
static_configs:
|
||||
- targets: ["llmgw-ctl.internal:9091"]
|
||||
labels:
|
||||
cluster: primary
|
||||
|
||||
# Worth alerting on, in rough order of how often it matters.
|
||||
#
|
||||
# A node that has stopped reporting:
|
||||
# llmgw_node_last_seen_seconds > 60
|
||||
#
|
||||
# A model that should be serving and is not. Covers a replica that will never
|
||||
# start as well as one that has died, which look the same from a distance:
|
||||
# llmgw_model_replicas_ready == 0
|
||||
#
|
||||
# A card kept out of placement by a fault, which is otherwise invisible until
|
||||
# a model will not fit:
|
||||
# llmgw_gpu_healthy == 0
|
||||
#
|
||||
# A fleet holding models nobody is asking for. Memory used says a model is
|
||||
# loaded; utilisation says whether it is doing anything:
|
||||
# avg_over_time(llmgw_gpu_utilization_percent[1h]) < 5
|
||||
#
|
||||
# Requests queueing because the cluster is full, which is the pair of figures
|
||||
# that needed two systems to correlate before these were on one endpoint:
|
||||
# llmgw_admission_queue_depth > 0
|
||||
@@ -0,0 +1,34 @@
|
||||
# systemd units
|
||||
|
||||
Two roles, two units. Copy the one you need to `/etc/systemd/system/`, put the
|
||||
matching configuration in `/etc/llmgw/`, and `systemctl enable --now` it.
|
||||
|
||||
## The control plane — `llmgw-server.service`
|
||||
|
||||
Runs unprivileged. It needs to read its own TLS key, so that file must belong to
|
||||
the `llmgw` user and nothing else:
|
||||
|
||||
```
|
||||
useradd --system --home-dir /var/lib/llmgw --shell /usr/sbin/nologin llmgw
|
||||
chown llmgw:llmgw /etc/llmgw/server-key.pem
|
||||
chmod 0400 /etc/llmgw/server-key.pem
|
||||
```
|
||||
|
||||
`StateDirectory=llmgw` creates `/var/lib/llmgw` with the right owner, which is
|
||||
where the store belongs. `ProtectSystem=strict` makes the rest of the filesystem
|
||||
read-only, so anything else it writes to has to be named explicitly.
|
||||
|
||||
## The agent — `llmgw-agent.service`
|
||||
|
||||
Runs as root, because it launches engines and needs device access.
|
||||
|
||||
`KillMode=process` is the line that matters and the reason this file exists
|
||||
rather than being generated. Without it, systemd kills the whole control group
|
||||
on stop — including the engines, which hold a loaded model and are serving
|
||||
traffic. Engines are meant to outlive an agent restart: the router talks to them
|
||||
directly, and the agent adopts them again when it comes back. Taking them down
|
||||
with the agent turns a thirty-second upgrade into a full model reload on every
|
||||
node.
|
||||
|
||||
systemd will log `Found left-over process ... in control group while starting
|
||||
unit. Ignoring.` on the next start. That is the mechanism working, not a fault.
|
||||
@@ -0,0 +1,17 @@
|
||||
[Unit]
|
||||
Description=llmgw node agent
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
ExecStart=/usr/local/bin/llmgw agent --config /etc/llmgw/agent.yaml
|
||||
# Root, because the agent launches engines and needs device access. Engines
|
||||
# outlive a restart of this unit and are adopted again on the next start, so
|
||||
# KillMode must not take the whole process group with it.
|
||||
KillMode=process
|
||||
Restart=on-failure
|
||||
RestartSec=2s
|
||||
StateDirectory=llmgw-agent
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -0,0 +1,19 @@
|
||||
[Unit]
|
||||
Description=llmgw control plane
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
ExecStart=/usr/local/bin/llmgw server --config /etc/llmgw/server.yaml
|
||||
User=llmgw
|
||||
Group=llmgw
|
||||
Restart=on-failure
|
||||
RestartSec=2s
|
||||
NoNewPrivileges=true
|
||||
ProtectSystem=strict
|
||||
ProtectHome=true
|
||||
PrivateTmp=true
|
||||
StateDirectory=llmgw
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
+753
@@ -0,0 +1,753 @@
|
||||
# Installing llmgw from nothing
|
||||
|
||||
Written for whoever is standing up the cluster: you have servers, you have root
|
||||
on them, and nothing is installed yet. By the end there is a control plane, a
|
||||
node in the fleet, and a model answering requests.
|
||||
|
||||
It assumes Debian or Ubuntu and adjusts nothing for other distributions beyond
|
||||
the package names. Everything here has been done on a real cluster except where
|
||||
this page says otherwise, and where it says otherwise it says so plainly.
|
||||
|
||||
## What llmgw does not install
|
||||
|
||||
Read this first, because it is the difference between an afternoon and a week.
|
||||
|
||||
llmgw runs engines and routes requests. It does **not** install GPU drivers,
|
||||
CUDA, inference engines, or model weights. A node missing any of them still
|
||||
registers, still reports `ready`, and still fails every replica placed on it —
|
||||
so the useful order is to make a node able to serve *before* onboarding it, not
|
||||
after.
|
||||
|
||||
| You provide | llmgw provides |
|
||||
| --- | --- |
|
||||
| NVIDIA driver and `nvidia-smi` | GPU discovery, classes, health |
|
||||
| InfiniBand stack (`rdma-core`/OFED) | Fabric discovery and the invariant that a replica never spans one |
|
||||
| vLLM, SGLang or llama.cpp | Launching them, the right flags, supervision |
|
||||
| Model weights on disk | Placement, routing, accounting |
|
||||
|
||||
The reason is boring and unavoidable: drivers and engines are platform decisions
|
||||
with kernel modules and version matrices attached, and distributing weights is
|
||||
its own problem — hundreds of gigabytes that want a shared filesystem, not a
|
||||
copy from somebody's laptop. It is a known gap, not an oversight.
|
||||
|
||||
## 0. Before you begin
|
||||
|
||||
Four machines is a reasonable first cluster: one control plane and three nodes.
|
||||
One of each also works.
|
||||
|
||||
**The control plane** needs no GPU. It is a single process holding a SQLite
|
||||
file; 2 GB of RAM and a few gigabytes of disk is generous. It must be reachable
|
||||
from every node on the agent port, and from your users on the client port.
|
||||
|
||||
**Nodes** need whatever the model needs, plus about 8 GB for the operating
|
||||
system and the agent. They dial the control plane; the control plane never dials
|
||||
them, so nodes need no inbound access except your own SSH.
|
||||
|
||||
**Ports**
|
||||
|
||||
| Port | Who connects | Expose to |
|
||||
| --- | --- | --- |
|
||||
| 8080 | clients, the web UI | your users, usually via a reverse proxy |
|
||||
| 7443 | node agents | the node network only |
|
||||
| 9090 | `llmgw ctl`, metrics, pprof | loopback only — it serves pprof, which can dump the process heap, and API keys live in that heap |
|
||||
|
||||
Decide now which address nodes will use to reach the control plane. It goes into
|
||||
a certificate and changing it later means reissuing one.
|
||||
|
||||
## 1. The binary
|
||||
|
||||
One static binary, no runtime dependencies, the same on every machine.
|
||||
|
||||
```
|
||||
VERSION=v0.6.0 # check the releases page for the current tag
|
||||
curl -fsSLO https://gitea.mynodes.xyz/ilker/llmgw/releases/download/$VERSION/llmgw-$VERSION-linux-amd64
|
||||
curl -fsSL https://gitea.mynodes.xyz/ilker/llmgw/releases/download/$VERSION/checksums.txt | sha256sum -c --ignore-missing
|
||||
sudo install -m0755 llmgw-$VERSION-linux-amd64 /usr/local/bin/llmgw
|
||||
llmgw --version
|
||||
```
|
||||
|
||||
`arm64` and `darwin` builds are published alongside. They are built with
|
||||
`CGO_ENABLED=0` deliberately, so a binary built today runs on a node whose libc
|
||||
is older than the builder's.
|
||||
|
||||
To build instead of download, you need Go 1.26 or later and nothing else:
|
||||
|
||||
```
|
||||
git clone https://gitea.mynodes.xyz/ilker/llmgw && cd llmgw
|
||||
CGO_ENABLED=0 go build -o llmgw ./cmd/llmgw
|
||||
```
|
||||
|
||||
## 2. The certificate authority
|
||||
|
||||
Agents authenticate to the control plane with mutual TLS, and a node's identity
|
||||
is taken from its certificate's common name. A public authority cannot supply
|
||||
these: Let's Encrypt and its peers stopped issuing certificates carrying the
|
||||
`clientAuth` usage, so nothing they hand out can authenticate an agent.
|
||||
|
||||
Run this wherever you intend to keep the authority. It talks to nothing and
|
||||
works on files, so it does not have to be a cluster machine — and there is an
|
||||
argument that it should not be.
|
||||
|
||||
```
|
||||
sudo llmgw ctl cert init-ca --dir /etc/llmgw/pki
|
||||
```
|
||||
|
||||
That writes `ca.pem` and `ca-key.pem`. **Whoever holds `ca-key.pem` can mint any
|
||||
node's identity**, so it stays here; `ca.pem` is copied to every machine and is
|
||||
not a secret.
|
||||
|
||||
Now the control plane's own certificate. Every address an agent might dial must
|
||||
appear in `--host`, bare IP addresses included — hostname verification reads
|
||||
subject alternative names and ignores the common name entirely, so a certificate
|
||||
naming the host only in its subject is rejected by every client:
|
||||
|
||||
```
|
||||
sudo llmgw ctl cert issue-server llmgw-ctl \
|
||||
--host 10.0.0.140 --host llmgw-ctl.internal \
|
||||
--dir /etc/llmgw/pki --out /etc/llmgw
|
||||
```
|
||||
|
||||
This covers the fleet-facing port only. The client-facing API is public and
|
||||
belongs behind an ordinary publicly trusted certificate, usually terminated by a
|
||||
reverse proxy in front of port 8080.
|
||||
|
||||
## 3. The control plane
|
||||
|
||||
### The user and directories
|
||||
|
||||
```
|
||||
sudo useradd --system --home /var/lib/llmgw --shell /usr/sbin/nologin llmgw
|
||||
sudo install -d -o llmgw -g llmgw -m 0750 /var/lib/llmgw
|
||||
sudo install -d -m 0755 /etc/llmgw
|
||||
sudo chown root:llmgw /etc/llmgw/ca.pem /etc/llmgw/server.pem /etc/llmgw/server-key.pem
|
||||
sudo chmod 0640 /etc/llmgw/server-key.pem
|
||||
```
|
||||
|
||||
### The configuration
|
||||
|
||||
`/etc/llmgw/server.yaml`. This is the smallest file that is not a toy: mutual
|
||||
TLS on, authentication on, one model, and limits that apply to anyone nobody has
|
||||
given limits to.
|
||||
|
||||
```yaml
|
||||
listen: ":8080"
|
||||
admin_listen: "127.0.0.1:9090"
|
||||
agent_listen: ":7443"
|
||||
|
||||
agent_tls:
|
||||
ca: /etc/llmgw/ca.pem
|
||||
cert: /etc/llmgw/server.pem
|
||||
key: /etc/llmgw/server-key.pem
|
||||
|
||||
auth:
|
||||
enabled: true
|
||||
store: /var/lib/llmgw/llmgw.db
|
||||
signup_open: true
|
||||
# Applied to any key no scope gives a limit to. Without these, a project
|
||||
# created without limits is metered by nothing at all, and the first sign of
|
||||
# that is the bill or the outage.
|
||||
default_rpm: 60
|
||||
default_tpm: 200000
|
||||
|
||||
landing:
|
||||
contact: "platform@example.com"
|
||||
|
||||
# Start with none. Models are added while the cluster runs — see step 7 —
|
||||
# and anything listed here is declarative: it is in the catalog on every start
|
||||
# and cannot be removed without editing this file.
|
||||
models: []
|
||||
```
|
||||
|
||||
`admin_listen` is on loopback on purpose. It serves pprof, which can dump the
|
||||
process heap, and API keys sit in that heap in plaintext. To see the dashboard
|
||||
from elsewhere, promote an account and use `/admin` on the client port — do not
|
||||
move this listener.
|
||||
|
||||
The store file is created on first start; there is no separate migration step.
|
||||
|
||||
### The service
|
||||
|
||||
```
|
||||
sudo cp deploy/systemd/llmgw-server.service /etc/systemd/system/
|
||||
sudo systemctl daemon-reload
|
||||
sudo systemctl enable --now llmgw-server
|
||||
systemctl status llmgw-server
|
||||
```
|
||||
|
||||
### Your account
|
||||
|
||||
Accounts are created and approved from a shell on this machine, deliberately:
|
||||
having a shell here is what makes someone an administrator, and there are no
|
||||
roles to get subtly wrong.
|
||||
|
||||
```
|
||||
sudo -u llmgw llmgw users create you@example.com --password 'change me now'
|
||||
# prints a user id, usr-xxxxxxxx
|
||||
sudo -u llmgw llmgw users promote usr-xxxxxxxx
|
||||
```
|
||||
|
||||
`promote` is what grants `/admin`. Sign in at `http://<control-plane>:8080/login`,
|
||||
change the password, and issue yourself a key from the profile page.
|
||||
|
||||
### Check it
|
||||
|
||||
```
|
||||
curl -s -o /dev/null -w '%{http_code}\n' http://127.0.0.1:8080/health # 200, with an empty body
|
||||
sudo -u llmgw llmgw ctl status
|
||||
```
|
||||
|
||||
`ctl status` will report no nodes and no models, which is correct — nothing has
|
||||
been added yet. It is worth running now so its output is familiar before you
|
||||
need it.
|
||||
|
||||
## 4. Preparing a node
|
||||
|
||||
**Do this before onboarding.** A node missing its driver or engine registers
|
||||
healthy and fails on first placement, and the reason is three logs away.
|
||||
|
||||
### The driver
|
||||
|
||||
```
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y nvidia-driver-550 # or whatever your cards need
|
||||
sudo reboot
|
||||
nvidia-smi # must print your GPUs
|
||||
```
|
||||
|
||||
`nvidia-smi` must be on `PATH` for the user the agent runs as, which is root.
|
||||
Containerised driver installs often put it elsewhere; if so, note the path and
|
||||
set `discovery.nvidia_smi` in the agent configuration rather than moving it.
|
||||
|
||||
A node with no GPU at all is fine and is not an error. It reports zero cards and
|
||||
can host CPU profiles, which is a good way to try llmgw before committing
|
||||
hardware to it.
|
||||
|
||||
### NVLink, if your GPUs have it
|
||||
|
||||
The published binary cannot see NVLink topology. NVML is the only source of it,
|
||||
its bindings need cgo, and the portable binary is built without cgo on purpose
|
||||
so it runs on a node whose libc is older than the builder's.
|
||||
|
||||
Without that topology, a tensor-parallel group takes the lowest free GPU
|
||||
indices. On an SXM board with NVSwitch that is correct, because every card is
|
||||
equidistant. On a PCIe box where NVLink bridges pair specific cards, a group
|
||||
straddling two bridges falls back to PCIe and runs several times slower — and
|
||||
nothing reports it, because from the outside it is simply a replica that is
|
||||
slow.
|
||||
|
||||
For those nodes, install the `-nvml` build published beside the portable one:
|
||||
|
||||
```
|
||||
sudo install -m0755 llmgw-$VERSION-linux-amd64-nvml /usr/local/bin/llmgw
|
||||
```
|
||||
|
||||
It is dynamically linked, so it needs a libc no older than the builder's and the
|
||||
NVIDIA driver present. Check it took effect with:
|
||||
|
||||
```
|
||||
llmgw ctl nodes --json | grep -c nvlink_peers # 0 means the topology is unknown
|
||||
```
|
||||
|
||||
`llmgw_gpu_nvlink_peers` in the metrics says the same thing across a fleet: a
|
||||
flat zero everywhere means the topology is unknown rather than absent.
|
||||
|
||||
### InfiniBand, if you have it
|
||||
|
||||
```
|
||||
sudo apt-get install -y rdma-core ibverbs-utils infiniband-diags
|
||||
ibstat # must list your HCAs with State: Active
|
||||
ls /sys/class/infiniband
|
||||
```
|
||||
|
||||
The agent checks what you declare against what the node has. An adapter that is
|
||||
not there, or that has no ACTIVE port — physically linked but with no subnet
|
||||
manager, which survives a reboot and looks like a software problem — takes that
|
||||
fabric out of service on this node, loudly, in the log. Without that check the
|
||||
replica starts, passes its health checks, and hangs on its first all-reduce.
|
||||
|
||||
Fabric membership is configuration, not discovery, and there is no way around
|
||||
that: OpenSM's default subnet prefix is the same on every fabric, so a matching
|
||||
prefix proves nothing and a wrong guess produces a replica that starts, passes
|
||||
health checks, and then hangs on its first all-reduce. You name the fabric when
|
||||
you onboard the node, and two nodes are only in the same fabric if you said so.
|
||||
|
||||
### An engine
|
||||
|
||||
Pick what you will actually serve. For a first cluster, llama.cpp is the least
|
||||
trouble and runs without a GPU.
|
||||
|
||||
```
|
||||
# llama.cpp, from a release build
|
||||
sudo apt-get install -y libgomp1 # it will not start without this
|
||||
sudo mkdir -p /opt/llama.cpp
|
||||
# unpack a llama.cpp release into /opt/llama.cpp
|
||||
/opt/llama.cpp/llama-server --version
|
||||
```
|
||||
|
||||
That `libgomp1` line is not padding. A `llama-server` binary missing it exists,
|
||||
is executable, and cannot start — and the node reports `ready` throughout.
|
||||
|
||||
For vLLM, a virtualenv on the node:
|
||||
|
||||
```
|
||||
sudo apt-get install -y python3-venv
|
||||
sudo python3 -m venv /opt/venvs/vllm
|
||||
sudo /opt/venvs/vllm/bin/pip install vllm
|
||||
/opt/venvs/vllm/bin/vllm --version
|
||||
```
|
||||
|
||||
Check the engine still takes the flags llmgw builds for it, which drift between
|
||||
releases:
|
||||
|
||||
```
|
||||
llmgw ctl engine-check vllm --bin /opt/venvs/vllm/bin/vllm
|
||||
```
|
||||
|
||||
### Weights
|
||||
|
||||
Put them where the node can read them. `scp`, a shared filesystem, or the
|
||||
engine's own downloader are all valid, and which you want depends on how many
|
||||
nodes you have. If you would rather not choose, skip to the next section and
|
||||
let llmgw fetch them.
|
||||
|
||||
```
|
||||
sudo install -d -o root -g root /var/lib/llmgw-models
|
||||
# then place your .gguf or model directory there
|
||||
```
|
||||
|
||||
For a first model, something small: a 3B at `q4_k_m` is about 2 GB, runs on
|
||||
CPU, and answers fast enough to prove the whole path works.
|
||||
|
||||
### Letting llmgw fetch them
|
||||
|
||||
The simple arrangement, and the default: every node keeps its own copy. Point a
|
||||
model's `source` at a file to download and the agent fetches it once, before the
|
||||
engine is launched:
|
||||
|
||||
```yaml
|
||||
source: hf:Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/qwen2.5-coder-3b-instruct-q4_k_m.gguf
|
||||
# or an ordinary URL
|
||||
source: https://models.internal/qwen2.5-coder-3b-q4_k_m.gguf
|
||||
```
|
||||
|
||||
Eight shards of one replica starting together produce one download, not eight.
|
||||
An interrupted fetch resumes rather than starting over, and a partial file is
|
||||
never visible under the model's name — so a half-downloaded model cannot be
|
||||
picked up by the next scan and handed to an engine.
|
||||
|
||||
Three kinds of source, and only one of them is fetched:
|
||||
|
||||
| `source` | What happens |
|
||||
| --- | --- |
|
||||
| `/var/lib/llmgw-models/x.gguf` | Checked, never copied. However it got there. |
|
||||
| `hf:org/repo/file.gguf`, `https://…` | Fetched once into the cache, then reused |
|
||||
| `Qwen/Qwen2.5-7B-Instruct` | Left alone — vLLM and SGLang download repositories themselves |
|
||||
|
||||
Which means the exotic arrangements need no configuration at all. Weights on a
|
||||
Lustre mount, or an image exported over the fabric with SRP or NVMe-oF, are
|
||||
simply *present*, and present means nothing is fetched.
|
||||
|
||||
### Telling llmgw where they are
|
||||
|
||||
Optional, and unset by default. Naming the directories in the agent's
|
||||
`model_dirs` lets a node report which models it can see, and placement then
|
||||
prefers a node that already holds one over a node that must first move tens of
|
||||
gigabytes:
|
||||
|
||||
```yaml
|
||||
model_dirs:
|
||||
- /var/lib/llmgw-models
|
||||
- /mnt/lustre/models
|
||||
```
|
||||
|
||||
llmgw does not manage that storage and does not care how the files arrived. A
|
||||
Lustre or NFS mount, an image exported over the fabric with SRP or NVMe-oF,
|
||||
local NVMe, or an `scp` all look identical to it. A node with nothing configured
|
||||
reports nothing and is treated exactly as every node was before this existed.
|
||||
|
||||
**It matters most when your fabrics are not connected to each other.** Weights
|
||||
in `ib-a` are not weights in `ib-b`: the only path between them is the
|
||||
management network, which is the slow one. `llmgw ctl status` shows what each
|
||||
fabric holds, so a model seeded into one fabric does not quietly leave the rest
|
||||
of the cluster unable to serve it.
|
||||
|
||||
How much disk a node needs follows from the models rather than a rule of thumb.
|
||||
`llmgw ctl model plan <repo>` prints the sizes; as a starting point, a 70B at
|
||||
bf16 is 141 GB and the same model at q4 is 43 GB.
|
||||
|
||||
## 5. Adding the node
|
||||
|
||||
Two halves, because they need different access: building the machine needs the
|
||||
hypervisor, and issuing its identity needs the certificate authority. Where one
|
||||
person holds both — the common case — one command does it:
|
||||
|
||||
```
|
||||
llmgw-add-node \
|
||||
--name llmgw-node01 --ip 10.0.0.141/24 \
|
||||
--server https://10.0.0.140:7443 \
|
||||
--cores 16 --memory 131072 --disk 200 \
|
||||
--fabric ib-a:mlx5_0,mlx5_1 \
|
||||
--llamacpp-bin /opt/llama.cpp/llama-server
|
||||
```
|
||||
|
||||
**The VM name, the hostname, the node id and the certificate's common name are
|
||||
all the same string.** They do not have to be, and every time they are not,
|
||||
somebody spends an afternoon working out which of the four a message refers to.
|
||||
|
||||
For a machine that already exists — bare metal, or a VM you built yourself — skip
|
||||
provisioning and issue the identity only:
|
||||
|
||||
```
|
||||
llmgw-onboard-node \
|
||||
--node llmgw-node01 --addr 10.0.0.141 \
|
||||
--server https://10.0.0.140:7443 \
|
||||
--fabric ib-a:mlx5_0,mlx5_1 \
|
||||
--llamacpp-bin /opt/llama.cpp/llama-server
|
||||
```
|
||||
|
||||
Onboarding needs SSH from the CA host to the node. It waits for SSH rather than
|
||||
failing on it, and learns the host key of a machine it has not seen — a new
|
||||
machine spends its first minute regenerating host keys and restarting sshd.
|
||||
|
||||
It also checks that a configured engine actually runs, and says so loudly when
|
||||
it does not. That check exists because the failure is otherwise invisible.
|
||||
|
||||
Within a few seconds:
|
||||
|
||||
```
|
||||
llmgw ctl nodes
|
||||
```
|
||||
|
||||
```
|
||||
NODE STATE GPUS MEM CLASSES FABRICS REPLICAS LAST SEEN
|
||||
llmgw-node01 ready 8/8 free 503G h100-80g ib-a 0 2s
|
||||
```
|
||||
|
||||
If it does not appear, the control plane's log says why. A certificate whose
|
||||
common name does not match the claimed node id is refused by design — that check
|
||||
is what makes a node's identity something it proves rather than something it
|
||||
picks.
|
||||
|
||||
Repeat for each node. Nodes in the same fabric must be given the same `--fabric`
|
||||
name; nodes with no InfiniBand take no `--fabric` at all and are each their own
|
||||
island, which is fine for single-node models.
|
||||
|
||||
## 6. Setting the node's CPU model, if it is a VM
|
||||
|
||||
Proxmox defaults a VM to `kvm64`, which hides AVX, AVX2, FMA and even SSE4.2
|
||||
from the guest. Inference then falls back to scalar code and is not slightly
|
||||
slower but unusable. Measured on one 8-core VM, same weights, same threads, the
|
||||
only change being `--cpu host`:
|
||||
|
||||
| Qwen2.5-Coder | kvm64 | host |
|
||||
| --- | ---: | ---: |
|
||||
| 1.5B q4_K_M | 5.80 | **39.09** tok/s |
|
||||
| 3B q4_K_M | 2.90 | **20.51** |
|
||||
| 7B q4_K_M | 1.33 | **9.88** |
|
||||
|
||||
`llmgw-provision-node` sets `--cpu host` for you. A VM you built yourself
|
||||
probably did not, and nothing about it looks broken — every number is simply
|
||||
small. Check with `grep -qw avx2 /proc/cpuinfo` on the node.
|
||||
|
||||
## 7. Adding a model
|
||||
|
||||
Write the definition. This is the same shape a model takes in `server.yaml`, so
|
||||
an entry moves between the two by being copied:
|
||||
|
||||
```yaml
|
||||
# qwen-coder-3b.yaml
|
||||
name: qwen2.5-coder-3b-instruct
|
||||
aliases: [coder]
|
||||
source: /var/lib/llmgw-models/qwen2.5-coder-3b-instruct-q4_k_m.gguf
|
||||
routing_profile: chat
|
||||
# What a client may use in one request. Worth stating: llama.cpp divides
|
||||
# --ctx-size across its parallel slots, so a pool of 65536 over two slots is
|
||||
# 32768 each, and nothing else will tell you that.
|
||||
context_length: 32768
|
||||
max_output: 4096
|
||||
scaling: { min: 1, max: 1 }
|
||||
profiles:
|
||||
- name: cpu
|
||||
engine: llamacpp
|
||||
weight: 1
|
||||
priority: 100
|
||||
startup_timeout: 5m
|
||||
min_memory_gb: 8
|
||||
args: ["--ctx-size", "65536"]
|
||||
```
|
||||
|
||||
```
|
||||
llmgw ctl model add -f qwen-coder-3b.yaml
|
||||
llmgw ctl models
|
||||
```
|
||||
|
||||
```
|
||||
MODEL SERVING REPLICAS STATUS ALIASES
|
||||
qwen2.5-coder-3b-instruct yes 1/1 ready coder
|
||||
```
|
||||
|
||||
No restart, and it survives one: models added this way are stored, while
|
||||
anything in `server.yaml` is declarative and returns on every start.
|
||||
|
||||
### Models that load when someone asks
|
||||
|
||||
A cluster with more models than GPUs can hold every model in its catalog and
|
||||
only the wanted ones in its cards:
|
||||
|
||||
```yaml
|
||||
scaling: { min: 0, max: 1, idle_timeout: 15m }
|
||||
```
|
||||
|
||||
`min: 0` with a `max` above it means load on demand. The first request for the
|
||||
model starts it and waits briefly; if it is not serving by then the client gets
|
||||
a 503 with `Retry-After` and an honest "the model is loading", rather than being
|
||||
told nothing serves it — a client told to retry succeeds where one told the
|
||||
model does not exist gives up. Once nobody has asked for it in `idle_timeout`,
|
||||
the replicas are stopped and the cards go back to the pool.
|
||||
|
||||
`min: 0, max: 0` is different and equally useful: configured, and deliberately
|
||||
not running. A model with no `scaling` stanza at all gets that, because making
|
||||
every unscaled model loadable on demand would let any client request start a
|
||||
70B nobody intended to run.
|
||||
|
||||
Two things to weigh. A cold start is visible to whoever triggers it, so this
|
||||
suits the long tail people are experimenting with rather than the two or three
|
||||
models used daily — keep those at `min: 1`. And stopping a replica throws away
|
||||
the KV cache prefix routing spent its effort filling, so an idle timeout wants
|
||||
to be minutes rather than seconds.
|
||||
|
||||
When the cluster is full, a model somebody asks for will stop an idle one to
|
||||
make room. Only idle replicas of other on-demand models: a model at `min: 1` is
|
||||
a promise that it stays up, and breaking that to serve a different request would
|
||||
make every guarantee in the catalog conditional on load. Only single-node
|
||||
replicas, because evicting a multi-node one frees several machines and costs
|
||||
tens of minutes to rebuild — a decision worth making deliberately. One at a
|
||||
time. And never a replica that has just started, which is what stops two models
|
||||
on a cluster with room for one from taking turns evicting each other.
|
||||
|
||||
For a GPU model, name the class and the width instead:
|
||||
|
||||
```yaml
|
||||
profiles:
|
||||
- name: tp2
|
||||
engine: vllm
|
||||
gpu_class: h100-80g # as `llmgw ctl nodes` reports it
|
||||
tp: 2
|
||||
priority: 100
|
||||
```
|
||||
|
||||
If it does not place, ask:
|
||||
|
||||
```
|
||||
llmgw ctl explain-placement qwen2.5-coder-3b-instruct
|
||||
```
|
||||
|
||||
which names every candidate and why each was rejected, rather than leaving you
|
||||
to infer it.
|
||||
|
||||
## 8. The first request
|
||||
|
||||
```
|
||||
curl http://<control-plane>:8080/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "coder", "messages": [{"role":"user","content":"hello"}]}'
|
||||
```
|
||||
|
||||
Or open `/chat` in a browser and use it without a key at all — the page carries
|
||||
your session, and the request still goes through admission, rate limits, budget
|
||||
and accounting exactly as an API call does.
|
||||
|
||||
## Watching it
|
||||
|
||||
The control plane exports both halves on one endpoint: what the cluster is
|
||||
serving — requests, latency, time to first token, tokens, admission, per
|
||||
endpoint — and what it is made of, per node and per GPU, including memory,
|
||||
temperature, ECC faults and utilisation.
|
||||
|
||||
**The dashboard is the answer for most clusters.** `/admin` shows capacity per
|
||||
GPU class, free whole nodes per fabric, which replicas are holding hardware
|
||||
nothing is asking for, who is using the cluster and what they are running, and
|
||||
the same node and model tables `ctl` prints — with no monitoring stack to stand
|
||||
up beside it. One binary, no dependencies, which is the point.
|
||||
|
||||
Prometheus is worth adding when you want history, alerts, or one view across
|
||||
several clusters. The dashboard shows what is true now; it keeps nothing.
|
||||
|
||||
Metrics then get their own listener, because the admin one also carries pprof
|
||||
and pprof can dump the process heap, where API keys sit in plaintext:
|
||||
|
||||
```yaml
|
||||
admin_listen: "127.0.0.1:9090" # dashboard, management API, pprof — keep local
|
||||
metrics_listen: "0.0.0.0:9091" # metrics only, safe to reach
|
||||
```
|
||||
|
||||
One scrape target for the whole fleet. The node and GPU figures come from the
|
||||
control plane, which already holds them, so there is nothing to install on the
|
||||
nodes and no Prometheus configuration per machine.
|
||||
|
||||
A starting configuration, with the alerts worth having, is in
|
||||
[deploy/monitoring/prometheus.yml](../deploy/monitoring/prometheus.yml).
|
||||
|
||||
Two figures repay knowing about. `llmgw_gpu_utilization_percent` is absent
|
||||
rather than zero when the driver does not report it — a flat zero across a
|
||||
fleet means the measurement is missing, not that the fleet is idle. And
|
||||
`llmgw_gpu_nvlink_peers` at zero everywhere means the topology is unknown,
|
||||
which is what a binary built without the nvml tag reports.
|
||||
|
||||
## Conversations across models
|
||||
|
||||
Off everywhere until a project asks. llmgw otherwise stores no message content
|
||||
at all — usage rows carry token counts, costs and latencies, never text, and
|
||||
even its own chat page keeps the conversation in the browser. A gateway that
|
||||
keeps nothing anyone typed is a much easier thing to have in front of a team,
|
||||
so keeping it is a decision someone makes rather than a default they discover.
|
||||
|
||||
```
|
||||
llmgw keys threads --project default --retain-hours 24
|
||||
```
|
||||
|
||||
The retention is required. A conversation store with no expiry is a transcript
|
||||
of everything a team ever typed, kept forever because nobody chose a number.
|
||||
The window is measured from the last turn, so a conversation still in progress
|
||||
does not expire underneath it.
|
||||
|
||||
A client then sends a thread id and only its new turn:
|
||||
|
||||
```
|
||||
curl $BASE/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "x-llmgw-thread: my-conversation" \
|
||||
-d '{"model": "<model-name>", "messages": [{"role":"user","content":"and then?"}]}'
|
||||
```
|
||||
|
||||
llmgw puts the stored history in front of that turn and sends it to whichever
|
||||
model is serving — **including a different one from the one the conversation
|
||||
started on**. That is the point: a thread begun on a large model can continue
|
||||
on a small one, and a caller dropped to a cheaper model by a spent budget keeps
|
||||
their conversation.
|
||||
|
||||
When the history does not fit the model now serving, the oldest turns are
|
||||
dropped until it does, and the reply carries `x-llmgw-thread-trimmed` with how
|
||||
many. System messages are never dropped, and neither is the turn the caller
|
||||
just sent.
|
||||
|
||||
`DELETE /v1/threads/{id}` forgets one. It is a client endpoint on purpose: the
|
||||
person who said the thing should be able to unsay it without asking an
|
||||
operator to read it first. A thread belonging to another project answers the
|
||||
same as one that never existed.
|
||||
|
||||
Turning it off again stops new storage, and existing conversations go at the
|
||||
next hourly prune:
|
||||
|
||||
```
|
||||
llmgw keys threads --project default --off
|
||||
```
|
||||
|
||||
## Memory notes for agents
|
||||
|
||||
Also off until a project asks, for the same reason. Where conversations are what
|
||||
was said, notes are what an agent chose to keep, and they persist across
|
||||
conversations:
|
||||
|
||||
```
|
||||
llmgw keys memory --project default --retain-hours 720
|
||||
```
|
||||
|
||||
The retention is required and is measured from the last write to a note. An agent
|
||||
then saves, searches and deletes notes through `/v1/memory` or, as tools, through
|
||||
the MCP server at `/mcp` (`remember`, `recall`, `forget`). Both need the API key
|
||||
of the project, and notes are private to the project and to the agent id sent in
|
||||
`x-llmgw-agent`. A note in another project or agent answers the same as one that
|
||||
does not exist. llmgw does not add these tools to requests itself.
|
||||
|
||||
Turning it off stops new storage, and existing notes go at the next hourly
|
||||
prune:
|
||||
|
||||
```
|
||||
llmgw keys memory --project default --off
|
||||
```
|
||||
|
||||
Without authentication, `threads.store` and `threads.retain_hours` in
|
||||
`server.yaml` keep both conversations and notes in one shared pool. That is for
|
||||
one trusted machine only; see [quickstart.md](quickstart.md).
|
||||
|
||||
## Reclaiming disk
|
||||
|
||||
`llmgw ctl nodes` has a DISK column: free space against the size it is free
|
||||
out of, for the filesystems holding model weights. A dash means the node was
|
||||
not asked — an agent with no `model_dirs`, or one older than this — and not
|
||||
that it has no space. Shared filesystems are excluded from a node's figure,
|
||||
because space on a filer belongs to everyone mounting it.
|
||||
|
||||
Weights llmgw downloaded for models that are no longer in the catalog can be
|
||||
listed and removed:
|
||||
|
||||
```
|
||||
llmgw ctl weights # what could be freed, per node
|
||||
llmgw ctl weights --collect # delete it
|
||||
```
|
||||
|
||||
**Only files llmgw fetched itself are ever candidates.** They are recognisable
|
||||
because llmgw names them after a hash of their source, so the set still wanted
|
||||
is computed from the catalog exactly rather than guessed. Anything you put on
|
||||
a node — by scp, by a shared mount, by hand — is left alone however
|
||||
unreferenced it looks, and no llmgw command will remove it. The nodes check
|
||||
every path again before deleting it, so a control plane that is wrong cannot
|
||||
talk one into losing a model that took a week to copy.
|
||||
|
||||
Which also means llmgw will not reclaim the storage most likely to be filling
|
||||
your disk. Weights you placed yourself are yours to remove.
|
||||
|
||||
A profile may declare `min_disk_gb`, and a node without that much free is then
|
||||
not a candidate for it — checked only on nodes that do not already hold the
|
||||
weights, since those need no room for them. Without the field, disk is not
|
||||
considered, which is how every profile written before it keeps behaving.
|
||||
|
||||
## When something does not work
|
||||
|
||||
**The node is `ready` but every replica fails.** Almost always a missing engine
|
||||
or weights. `llmgw ctl replicas` shows the shard's error and the engine's last
|
||||
output, which is the actual diagnosis.
|
||||
|
||||
**A model will not place.** `llmgw ctl explain-placement <model>` names each
|
||||
node and the reason: wrong GPU class, not enough free cards, too little memory,
|
||||
no fabric, a label it does not carry.
|
||||
|
||||
**Everything is slow.** On a VM, check AVX2 (step 6). On CPU, check the
|
||||
quantisation — the difference between `q8_0` and `q4_k_m` is roughly a factor of
|
||||
two. `llmgw ctl status` shows how full the cluster is.
|
||||
|
||||
**Tool calls come back as text.** The model made the call and the engine did not
|
||||
parse it. `llmgw ctl tool-check --endpoint http://<node>:<port>` distinguishes
|
||||
that from a model ignoring its tools, which has a different fix.
|
||||
|
||||
**Requests fail at exactly sixty seconds.** `timeouts.response_header` is a
|
||||
liveness check for streaming and a total generation budget for non-streaming
|
||||
requests. On a slow model, set `timeouts.generation` higher.
|
||||
|
||||
## What to read next
|
||||
|
||||
- `docs/node-onboarding.md` — adding nodes in more detail, and what onboarding
|
||||
deliberately does not do
|
||||
- `docs/onboarding.md` — giving people accounts and keys
|
||||
- `llmgw-plan.md` — why the system is shaped the way it is
|
||||
|
||||
## Known gaps, so you find them here rather than at 2am
|
||||
|
||||
- **A fetch is silent while it runs.** llmgw will fetch weights for you, but
|
||||
it reports no progress, so a first start on a large model looks like a hang.
|
||||
`ls -l` in the model directory is the only progress bar there is.
|
||||
- **Every node fetches its own copy.** There is no seeding between nodes, so
|
||||
bringing up eight nodes on a new model is eight downloads from the same
|
||||
origin. On a fabric where the nodes can see each other this is wasteful, and
|
||||
it is the obvious next thing to build.
|
||||
- **Engines are yours to install.** Onboarding verifies one runs; it installs
|
||||
nothing.
|
||||
- **One control plane.** No HA yet; it is a single process with a SQLite file.
|
||||
Back up `/var/lib/llmgw/llmgw.db`.
|
||||
- **Multi-node replicas have not been run on real InfiniBand.** The scheduling,
|
||||
the fabric invariant and the rendezvous are exercised against simulated
|
||||
hardware only. Single-node placement on real GPUs is the well-trodden path;
|
||||
if you are the first to run a TP group across real cards, expect to find
|
||||
something.
|
||||
@@ -0,0 +1,230 @@
|
||||
# Adding a node
|
||||
|
||||
Written for whoever runs the cluster.
|
||||
|
||||
## The short version
|
||||
|
||||
Where one person holds both the hypervisor and the certificate authority —
|
||||
which is the common case — one command does the whole thing in about
|
||||
45 seconds:
|
||||
|
||||
```
|
||||
llmgw-add-node \
|
||||
--name llmgw-node05 --ip 10.0.0.145/24 \
|
||||
--server https://10.0.0.140:7443 \
|
||||
--cores 8 --memory 32768 --disk 60 \
|
||||
--llamacpp-bin /opt/llama.cpp/llama-server
|
||||
```
|
||||
|
||||
The VMID defaults to the last octet of the address, so `qm list` and
|
||||
`llmgw ctl nodes` stay in the same order. `--provision-only` and
|
||||
`--onboard-only` split it back into the two halves below when the two kinds
|
||||
of access belong to different people.
|
||||
|
||||
The rest of this page is what that command does, and what it still does not.
|
||||
|
||||
## The two halves
|
||||
|
||||
Two scripts, because the two halves of onboarding are done by different people
|
||||
with different access: one needs the hypervisor, the other needs the
|
||||
certificate authority.
|
||||
|
||||
## 1. The machine
|
||||
|
||||
On the Proxmox host:
|
||||
|
||||
```
|
||||
llmgw-provision-node --vmid 144 --name llmgw-node04 --ip 10.0.0.144/24
|
||||
```
|
||||
|
||||
Name it after the node id. **The VM name, the hostname, the node id and the
|
||||
certificate's common name are all the same string.** They do not have to be, and
|
||||
every time they are not, somebody spends an afternoon working out which of the
|
||||
four a message is talking about.
|
||||
|
||||
It imports a cloud image and lets cloud-init do the rest, which is why a node
|
||||
takes a minute rather than an afternoon. It refuses to touch an existing VMID:
|
||||
adopting a machine is not a decision a provisioning script should make.
|
||||
|
||||
Releases carry static `linux/amd64` and `linux/arm64` binaries. Download one and
|
||||
pass it with `--binary` rather than copying whatever happens to be on the
|
||||
machine you are running onboarding from — otherwise whoever runs the script
|
||||
quietly becomes the build server, and the fleet ends up running whatever was in
|
||||
their working tree.
|
||||
|
||||
## 2. The identity
|
||||
|
||||
On the host that holds the certificate authority — this is the part that cannot
|
||||
be delegated, because it issues the credential by which the control plane
|
||||
believes the node is who it says it is:
|
||||
|
||||
```
|
||||
llmgw-onboard-node \
|
||||
--node llmgw-node04 \
|
||||
--addr 10.0.0.144 \
|
||||
--server https://10.0.0.140:7443 \
|
||||
--fabric ib-b
|
||||
```
|
||||
|
||||
It issues a client certificate whose common name is the node id, installs it
|
||||
along with the CA root, trusts that root system-wide so other tooling on the
|
||||
machine stops treating llmgw's certificates as strangers, installs the binary
|
||||
and the unit, and starts the agent. The node appears in `llmgw ctl nodes` within
|
||||
a few seconds.
|
||||
|
||||
If it does not, the control plane's log says why. A certificate whose common
|
||||
name does not match the claimed node id is refused by design — that check is
|
||||
what makes a node's identity something it proves rather than something it picks.
|
||||
|
||||
**It needs ssh from the CA host to the node.** Provisioning copies that host's
|
||||
`authorized_keys` into the new machine, so a node built by the script above
|
||||
already satisfies it; an older machine may need the key added.
|
||||
|
||||
It waits for ssh rather than failing on it, and learns the host key of a machine
|
||||
it has not seen (`accept-new`, which still refuses a *changed* key — the risk
|
||||
being guarded against is a substituted machine, not a new one). Both are needed
|
||||
to run the halves back to back: for the first minute of its life a new machine
|
||||
is still regenerating host keys and restarting sshd, so it refuses connections
|
||||
it accepted a moment earlier.
|
||||
|
||||
It also checks that a configured engine actually runs, which is cheap here and
|
||||
expensive later — see below.
|
||||
|
||||
## The CPU model matters, and its default is wrong
|
||||
|
||||
Nodes are provisioned with `--cpu host`. Proxmox's default, `kvm64`, hides
|
||||
AVX, AVX2, FMA and even SSE4.2 from the guest, and llama.cpp then loads its
|
||||
scalar backend instead of the one built for the host's instruction set.
|
||||
Measured on node05, an 8-core VM on a Xeon E5-2680 v4, same weights, same
|
||||
threads, the only change being the CPU model:
|
||||
|
||||
| Qwen2.5-Coder | kvm64 tok/s | host tok/s |
|
||||
| -------------- | ----------: | ---------: |
|
||||
| 1.5B q4_K_M | 5.80 | 39.09 |
|
||||
| 3B q4_K_M | 2.90 | 20.51 |
|
||||
| 7B q4_K_M | 1.33 | 9.88 |
|
||||
|
||||
Prompt processing is worse still: 7.6 tokens a second against 143.8, a factor
|
||||
of nineteen.
|
||||
|
||||
It is worth calling out because nothing looks broken. The node registers, the
|
||||
engine starts, replies are correct — every number is simply small, and the
|
||||
obvious conclusion is that CPU inference is hopeless rather than that the VM
|
||||
cannot multiply efficiently.
|
||||
|
||||
The trade is live migration: a VM pinned to the host CPU cannot move to a
|
||||
machine with a different one. For a node whose job is arithmetic that is the
|
||||
right way round, and `--cpu` overrides it where it is not.
|
||||
|
||||
## What onboarding does not do
|
||||
|
||||
It gives a machine an identity and an agent. It does **not** install inference
|
||||
engines or model weights, and a node missing either looks healthy in
|
||||
`llmgw ctl nodes` right up until a replica is placed on it and the shard dies
|
||||
with "executable file not found".
|
||||
|
||||
Onboarding does check that an engine you point it at can be executed, and says
|
||||
so loudly when it cannot. That check exists because the failure is invisible
|
||||
otherwise: `llama-server` on node05 was present, executable, and could not start
|
||||
for want of `libgomp1`, and the node reported `ready` in `llmgw ctl nodes` the
|
||||
whole time. Being told at onboarding beats finding out when a replica is placed
|
||||
and the reason is three logs away.
|
||||
|
||||
So the node is not finished until:
|
||||
|
||||
- the engines its catalog profiles name are installed (`vllm`, `sglang`,
|
||||
`llama-server`), at the paths the agent configuration gives
|
||||
- the weights are on disk, at the `source` each model declares
|
||||
|
||||
Both are deliberately outside this script. Engines are a platform decision with
|
||||
versions and drivers attached, and distributing weights is its own problem —
|
||||
hundreds of gigabytes that want a shared filesystem or a cache, not an scp from
|
||||
a laptop. llmgw will fetch weights for a model whose `source` is a URL, one
|
||||
copy per node, which covers the simple case; anything larger stays the
|
||||
operator's arrangement, and llmgw is content to find the files already there.
|
||||
|
||||
The `sim` engine needs neither, which is what makes it useful for exercising
|
||||
everything else.
|
||||
|
||||
## Checking an engine after an upgrade
|
||||
|
||||
Engines rename and remove flags between versions, and when that happens the
|
||||
failure arrives minutes into a model load, in the engine's own error message, on
|
||||
a node nobody was watching.
|
||||
|
||||
```
|
||||
llmgw ctl engine-check vllm
|
||||
llmgw ctl engine-check llamacpp --help-output "$(llama-server --help 2>&1)"
|
||||
```
|
||||
|
||||
It builds the widest command llmgw would ever produce — a multi-node one, since
|
||||
a flag that only appears there is exactly the one nobody notices has gone — and
|
||||
checks every flag against the engine's own `--help`. No GPU required: a CLI
|
||||
prints its help on any machine.
|
||||
|
||||
Some programs do not give up their whole help for `--help`. vLLM prints a
|
||||
summary of section names and expects `--help=all`; checking against the summary
|
||||
reports every flag as missing, confidently and wrongly. `engine-check` knows
|
||||
that one. If you add an engine, check that what you get back is the full list.
|
||||
|
||||
vLLM also refuses to do anything at all on a machine with no accelerator,
|
||||
including print its help, unless `VLLM_TARGET_DEVICE` says what to assume.
|
||||
`engine-check` sets it.
|
||||
|
||||
Give it the **whole** help. Truncating llama.cpp's 725 lines to 400 reports five
|
||||
flags as missing that the running cluster uses; the command warns when the text
|
||||
it was given does not document `--help` itself, which is the cheapest sign it
|
||||
was cut short.
|
||||
|
||||
## Fabrics
|
||||
|
||||
`--fabric` is how a node says which InfiniBand fabric it is attached to, and
|
||||
there is no way around stating it. OpenSM's default subnet prefix is `fe80::` on
|
||||
every fabric, so a matching prefix proves nothing and a wrong guess produces a
|
||||
replica that starts, passes health checks, and then hangs on its first
|
||||
all-reduce.
|
||||
|
||||
Omit it for a node with no fabric. Such a node is its own island: it can host
|
||||
single-node replicas and will never be given a multi-node one.
|
||||
|
||||
Two adapters per fabric is the default; `--fabric ib-b:mlx5_2,mlx5_3` names them
|
||||
explicitly. An adapter must never appear under two fabrics — the agent refuses
|
||||
that configuration, because the consequence is the hang above.
|
||||
|
||||
## Where the driver is
|
||||
|
||||
```
|
||||
--nvidia-smi /opt/nvidia/bin/nvidia-smi
|
||||
--ib-root /sys/class/infiniband
|
||||
```
|
||||
|
||||
Both default to the usual places. Containerized driver installs routinely put
|
||||
`nvidia-smi` somewhere else, and an agent that cannot find it registers with no
|
||||
GPUs — which looks exactly like a node that has none.
|
||||
|
||||
## Removing a node
|
||||
|
||||
```
|
||||
llmgw ctl drain llmgw-node04 # move its replicas off, one at a time
|
||||
llmgw ctl nodes # wait until it holds none
|
||||
```
|
||||
|
||||
Then stop the agent. Draining brings each replacement up before stopping the
|
||||
original, which is the difference between a rolling move and an outage.
|
||||
|
||||
The node's certificate stays valid until it expires. There is no revocation list
|
||||
— if a machine is decommissioned rather than retired, reissue the authority or
|
||||
accept that a certificate naming a node that no longer exists can still connect
|
||||
and claim to be it.
|
||||
|
||||
## Testing without hardware
|
||||
|
||||
Every part of this works on machines with no GPUs and no InfiniBand. Point
|
||||
`--nvidia-smi` at a script that prints the CSV discovery asks for, point
|
||||
`--ib-root` at a sysfs-shaped directory, and use `engine: sim` in the catalog.
|
||||
|
||||
The simulated engine rendezvouses for real: rank 0 binds the dist-init port and
|
||||
does not serve until every worker has joined, and workers serve nothing. That
|
||||
exercises the multi-node launch sequence, the fabric constraint, and failure and
|
||||
rebuild — which is most of what goes wrong — without exercising NCCL, which is
|
||||
the rest.
|
||||
@@ -0,0 +1,319 @@
|
||||
# Onboarding a user
|
||||
|
||||
Written for whoever operates an llmgw deployment: the person with a shell on the
|
||||
control plane machine.
|
||||
|
||||
There are two ways a person gets a key. Pick one and put it in the config.
|
||||
|
||||
| | Self-service | Operator-issued |
|
||||
|---|---|---|
|
||||
| Config | `auth.signup_open: true` | `auth.signup_open: false` |
|
||||
| They do | sign up, wait, then make their own keys | wait for a link |
|
||||
| You do | approve once | issue every key |
|
||||
| Good for | more than a handful of people | a few, or machines |
|
||||
|
||||
Self-service does not mean open. A signup produces an account that can do
|
||||
nothing at all; **approval is what makes it work**, and approval needs a shell
|
||||
here rather than a session cookie, because it is the step that decides who gets
|
||||
to spend the cluster.
|
||||
|
||||
Everything below runs as the user the server runs as, so the files it creates
|
||||
stay readable by it:
|
||||
|
||||
```
|
||||
sudo -u llmgw llmgw <command> --store /var/lib/llmgw/llmgw.db
|
||||
```
|
||||
|
||||
## Self-service
|
||||
|
||||
Turn on the form:
|
||||
|
||||
```yaml
|
||||
auth:
|
||||
enabled: true
|
||||
store: /var/lib/llmgw/llmgw.db
|
||||
signup_open: true
|
||||
```
|
||||
|
||||
The landing page then links to it, and the rest is:
|
||||
|
||||
```
|
||||
llmgw users list --pending --store …
|
||||
ID EMAIL NAME STATE PROJECT SIGNED UP
|
||||
usr-k4mq… ayse@example.com Ayşe pending - 4m12s ago
|
||||
|
||||
llmgw users approve usr-k4mq… --project default --by ilker --store …
|
||||
```
|
||||
|
||||
That is the whole operator loop. From then on they sign in at `/login`, and
|
||||
`/profile` is where they create keys, see when each was last used, and revoke the
|
||||
ones they are done with. A key is shown once on the page that creates it and
|
||||
never again.
|
||||
|
||||
**This gateway sends no email.** There is no "your account was approved" message
|
||||
and no password reset, so tell people what to expect: they sign up, they tell
|
||||
you, you approve, they sign in. If that loop is too manual for your size, it is
|
||||
the thing to automate — not the approval step.
|
||||
|
||||
To close an account:
|
||||
|
||||
```
|
||||
llmgw users suspend usr-k4mq… --store …
|
||||
```
|
||||
|
||||
which closes their sessions immediately. Their **API keys keep working** — those
|
||||
are separate credentials, and revoking them is a separate step:
|
||||
|
||||
```
|
||||
llmgw keys list --project default --store …
|
||||
llmgw keys revoke key-… --store …
|
||||
```
|
||||
|
||||
That separation is deliberate: a suspended person's running jobs should not all
|
||||
fail at once without someone deciding they should.
|
||||
|
||||
## Operator-issued
|
||||
|
||||
With `signup_open: false` there is no form, and you issue keys yourself. This is
|
||||
also the path for things that are not people — CI, a notebook server, a script.
|
||||
|
||||
### Decide where the key belongs
|
||||
|
||||
Keys hang off a project, projects off a team, teams off an organization. Limits
|
||||
can be set at any of those levels and compose downward by taking the **most
|
||||
restrictive** value — a key cannot raise what its team set, and a team cannot
|
||||
raise what its organization set. That is the only rule under which a budget
|
||||
means anything.
|
||||
|
||||
If you have no structure yet, one of each is fine:
|
||||
|
||||
```
|
||||
llmgw keys init --store … --org mynodes --team default --project default
|
||||
```
|
||||
|
||||
It is idempotent: running it again renames rather than duplicating, so it can
|
||||
live in whatever provisions the machine.
|
||||
|
||||
Split it when you want separate limits or separate reporting, not before.
|
||||
|
||||
### Issue the key as a one-time link
|
||||
|
||||
```
|
||||
llmgw keys create --store … \
|
||||
--project default \
|
||||
--band gold \
|
||||
--base-url https://llmgw.example \
|
||||
--claim \
|
||||
"ayse-laptop"
|
||||
|
||||
https://llmgw.example/claim/3mKq…
|
||||
```
|
||||
|
||||
Send the link, not the key.
|
||||
|
||||
A key pasted into a chat window stays in that chat window forever, searchable by
|
||||
everyone who joins the channel later. A link works once and then expires, so the
|
||||
worst case is someone opening it before the intended person — which you find out
|
||||
about, because they will tell you the link is dead.
|
||||
|
||||
The name is for you, not the user: it appears in `keys list` and is how you will
|
||||
know which key to revoke.
|
||||
|
||||
**Why the link is safe to store.** The key is sealed with the link's own token,
|
||||
and the token is never written down — only its hash is. A database leak yields
|
||||
two useless hashes and a blob nothing in the database can open.
|
||||
|
||||
Without `--claim` the key is printed to stdout instead, for when you are
|
||||
provisioning something that is not a person:
|
||||
|
||||
```
|
||||
llmgw keys create --store … --project ci "github-runner" > /run/secrets/llmgw-key
|
||||
```
|
||||
|
||||
## Limits
|
||||
|
||||
Limits are optional. A key with none is unlimited, which is the right default on
|
||||
a private cluster and the wrong one the moment the gateway faces anything else.
|
||||
|
||||
```
|
||||
# Everyone in this project
|
||||
llmgw keys limits --store … --scope project --id default \
|
||||
--rpm 60 --tpm 200000
|
||||
|
||||
# One key, tighter
|
||||
llmgw keys limits --store … --scope key --id key-mpmip77twpt7qaxc \
|
||||
--rpm 10
|
||||
|
||||
# Restrict which models
|
||||
llmgw keys limits --store … --scope project --id default \
|
||||
--models qwen-0.5b,llama-3.3-70b
|
||||
```
|
||||
|
||||
An **absent** model list permits everything. An **empty** one permits nothing —
|
||||
`--models ""` is a way to suspend a project without revoking its keys.
|
||||
|
||||
Allowlists compose by intersection: a key cannot add a model its team withheld.
|
||||
A model's canonical name and any of its aliases name the same thing, so either
|
||||
form works.
|
||||
|
||||
Budgets are counted in micro-units, and what a unit means is yours to decide.
|
||||
Self-hosted models have no per-token price, so it can be GPU-seconds, a notional
|
||||
rate, or 1 000 000 everywhere — which turns a budget into a plain token quota.
|
||||
Models are priced in the server's `auth.pricing` section; a model priced at zero
|
||||
is exempt from budget checks, which is what lets a key with nothing left still be
|
||||
served a cheap fallback.
|
||||
|
||||
```
|
||||
llmgw keys limits --store … --scope project --id default \
|
||||
--budget 50000000 --period month # 50 units a month
|
||||
```
|
||||
|
||||
## The landing page
|
||||
|
||||
Someone who reaches the gateway without a key sees the landing page. Tell it who
|
||||
to send them to, or it can only shrug:
|
||||
|
||||
```yaml
|
||||
landing:
|
||||
contact: "platform@example.com" # or an https:// URL, or "#platform on Slack"
|
||||
notice: "Internal use only."
|
||||
```
|
||||
|
||||
An address becomes a mailto link, an `http(s)` URL stays a link, anything else is
|
||||
shown as text. Other schemes are never linked — configuration is not always
|
||||
written by whoever runs the gateway, and a `javascript:` href on a page served to
|
||||
every visitor is a short path to a stored XSS.
|
||||
|
||||
## What the user sees
|
||||
|
||||
They open the link, press the button, and get the key once. The page shows them
|
||||
the base URL and a working `curl`, and points at two endpoints:
|
||||
|
||||
- `GET /v1/models` — what this key may use
|
||||
- `GET /me` — its limits and what it has spent
|
||||
|
||||
`/me` is there so nobody has to ask you "why was I rate limited" or "how much
|
||||
have I used". It is outside `/v1` because it is not part of the OpenAI API.
|
||||
|
||||
## Rotation and revocation
|
||||
|
||||
```
|
||||
llmgw keys list --store … --project default
|
||||
```
|
||||
|
||||
`LAST USED` is written in batches rather than per request, so it lags by up to a
|
||||
minute. It is the column that tells you whether a key is still in use, which is
|
||||
what makes rotation safe:
|
||||
|
||||
**To rotate**, issue the new key, let both work while the user switches, watch
|
||||
`LAST USED` on the old one stop moving, then revoke it.
|
||||
|
||||
```
|
||||
llmgw keys revoke --store … key-mpmip77twpt7qaxc
|
||||
```
|
||||
|
||||
Revocation is a soft delete: the key stops working, the row stays, and its usage
|
||||
history keeps its subject. A running server keeps accepting a revoked key until
|
||||
its cache expires — `auth.cache_ttl`, 30 seconds by default. That is the price of
|
||||
keeping the database off the request path, and it is the number to lower if you
|
||||
need revocation to bite faster.
|
||||
|
||||
## Passwords
|
||||
|
||||
People change their own from `/profile`. It asks for the current one — a session
|
||||
is what someone gets from a borrowed laptop, and the password is what they do not
|
||||
have — and it signs every other browser out, which is the point if they are
|
||||
changing it because something went wrong.
|
||||
|
||||
There is no self-service reset, because this gateway sends no email and a reset
|
||||
link with nowhere to send it is not a reset. When someone loses theirs:
|
||||
|
||||
```
|
||||
llmgw users reset-password usr-k4mq… --password "a temporary passphrase" --store …
|
||||
```
|
||||
|
||||
Tell them out of band and have them change it from their profile. Their sessions
|
||||
are closed — if they lost the password, the reason might be that someone else has
|
||||
it. Their API keys keep working; those are separate credentials.
|
||||
|
||||
## Trying a model in the browser
|
||||
|
||||
`/chat` is a playground for anyone signed in. It exists so somebody can see
|
||||
whether a model works without setting up a client, and — more often — so they
|
||||
can tell whether a problem is the gateway or their code.
|
||||
|
||||
**The browser never holds a key.** The page carries a session; the server turns
|
||||
that into the same identity an API key would produce, resolved from the person's
|
||||
project, and the request goes down the ordinary path: admission, model
|
||||
allowlist, rate limits, budget, accounting. Anything else would make the page a
|
||||
hole in every limit the rest of the gateway enforces.
|
||||
|
||||
Usage from the page is attributed to `session:<user-id>` rather than a key, so
|
||||
it is billed to the right project and visibly did not come from a credential
|
||||
somebody holds. One turn is capped at 512 tokens: a playground should not be the
|
||||
thing that ties up a replica for ten minutes, and anyone who wants a longer
|
||||
answer has an API key for that.
|
||||
|
||||
Nothing is stored. The conversation lives in the page and is gone when it is
|
||||
closed.
|
||||
|
||||
## Seeing the cluster
|
||||
|
||||
The dashboard — nodes, replicas, models, why something did not get placed, usage
|
||||
per key — is served twice, and the difference matters.
|
||||
|
||||
On the **admin listener** (`127.0.0.1:9090` by default) it needs no
|
||||
authentication: having a shell on the machine is the credential.
|
||||
|
||||
```
|
||||
ssh -L 9090:127.0.0.1:9090 <control-plane> # then http://127.0.0.1:9090/
|
||||
```
|
||||
|
||||
That listener also serves pprof, which can dump the process heap — where API
|
||||
keys sit in plaintext. **Never expose that port.** It is bound to loopback for
|
||||
that reason.
|
||||
|
||||
On the **client-facing port** the same page is at `/admin`, behind a session and
|
||||
an admin flag, so it can be reached from a browser without the port carrying
|
||||
pprof along with it:
|
||||
|
||||
```
|
||||
llmgw users promote usr-k4mq… --store …
|
||||
```
|
||||
|
||||
Someone without the flag gets a 404 rather than a 403, because 403 confirms the
|
||||
page exists.
|
||||
|
||||
## Managing accounts from the browser
|
||||
|
||||
`/admin/users` lists who is waiting and who is active. Approving (into a project)
|
||||
and suspending are there, because that is the work an operator does constantly
|
||||
and it should not need a terminal.
|
||||
|
||||
Three things are deliberately not there:
|
||||
|
||||
| | Why |
|
||||
|---|---|
|
||||
| `llmgw users promote` | An admin who can make admins turns a stolen session into permanent access |
|
||||
| `llmgw users reset-password` | Setting someone's password is taking over their account |
|
||||
| `llmgw keys create` | Issuing a credential on another person's behalf |
|
||||
|
||||
Each of those turns a stolen admin cookie into more than it was, and each is
|
||||
rare enough that a shell is no hardship. Cordon, drain and everything else that
|
||||
touches the cluster are also CLI-only for the same reason: the dashboard shows
|
||||
the cluster, it does not steer it.
|
||||
|
||||
Every administrative action is logged with who did it, including the ones that
|
||||
were refused.
|
||||
|
||||
## What this does not do
|
||||
|
||||
- **No email.** No approval notice, no self-service password reset, no address
|
||||
verification. The loop is manual on purpose at this size; if it stops being
|
||||
the right size, that is the thing to automate.
|
||||
- **No SSO.** Accounts are local, with local passwords.
|
||||
- **No roles.** There is one flag, `is_admin`, and it grants looking at the
|
||||
dashboard. Everything else an administrator does is the CLI, and having a
|
||||
shell here is what makes someone one.
|
||||
- **Rate limits are per gateway instance.** With two instances each counts its own,
|
||||
so the effective limit doubles. v1 is single-instance by design.
|
||||
@@ -0,0 +1,375 @@
|
||||
# Quickstart: one machine, one model, ten minutes
|
||||
|
||||
Written for someone who wants to see llmgw answer a request on a single
|
||||
workstation before reading anything else. No systemd, no `sudo`. For a real
|
||||
cluster, see [install.md](install.md).
|
||||
|
||||
What you end up with: an OpenAI-compatible endpoint on `localhost:8080` in front
|
||||
of a real model served by llama.cpp.
|
||||
|
||||
## Two ways to run
|
||||
|
||||
Steps 1 to 4 use the simple one, **router-only**: you start the engines yourself
|
||||
and tell the router where they are with `--endpoint`. No certificates, no node
|
||||
agent.
|
||||
|
||||
Step 5 uses the **control plane and node agent**, the same architecture as a real
|
||||
cluster on one machine: the agent starts and supervises the engine, and the
|
||||
control plane places the model.
|
||||
|
||||
## 1. Build
|
||||
|
||||
```
|
||||
git clone https://gitea.mynodes.xyz/ilker/llmgw && cd llmgw
|
||||
make build
|
||||
make test
|
||||
./bin/llmgw --version
|
||||
```
|
||||
|
||||
`make build` writes `./bin/llmgw`. Go 1.26 or later is the only requirement.
|
||||
|
||||
## 2. Check the path with a fake engine
|
||||
|
||||
Before involving a GPU, prove the router works. `sim-engine` is an
|
||||
OpenAI-compatible fake that streams placeholder words.
|
||||
|
||||
```
|
||||
./bin/llmgw sim-engine --listen :8001 --model llama &
|
||||
./bin/llmgw server --endpoint e1=http://127.0.0.1:8001=llama &
|
||||
|
||||
curl -N localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-H 'x-session-id: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}'
|
||||
```
|
||||
|
||||
Chunks of `alpha beta gamma ...` ending in `data: [DONE]` mean the router is
|
||||
fine. Stop both before continuing (see [Stopping](#stopping)).
|
||||
|
||||
## 3. Serve a real model
|
||||
|
||||
Two processes, two ports. Clients only ever talk to the router; the router
|
||||
knows where the engine is because you tell it with `--endpoint`:
|
||||
|
||||
```
|
||||
curl :8080 -> llmgw router :8080 -> llama serve :8082 -> model
|
||||
(client) (the gateway) (the engine)
|
||||
```
|
||||
|
||||
Requests go to 8080. Port 8082 is the engine's own, and you only touch it to
|
||||
check that the engine is up.
|
||||
|
||||
llmgw does not install engines or weights. You need:
|
||||
|
||||
- a `llama-server` build from llama.cpp (CUDA build if you want the GPU), and
|
||||
- a `.gguf` file.
|
||||
|
||||
Check the card first with `nvidia-smi`. An 8 GB card holds a 1B model at bf16
|
||||
(2.4 GB) or a 9B model at `q4_k_m` (5.4 GB).
|
||||
|
||||
Start the engine on its own port. The model name given to `--alias` is the name
|
||||
clients will use:
|
||||
|
||||
```
|
||||
llama serve \
|
||||
-m /path/to/Llama-3.2-1B-Instruct-BF16.gguf \
|
||||
--host 0.0.0.0 --port 8082 -ngl 99 --alias llama
|
||||
```
|
||||
|
||||
`-ngl 99` offloads every layer to the GPU. Confirm it is up and knows the model:
|
||||
|
||||
```
|
||||
curl -s localhost:8082/v1/models
|
||||
```
|
||||
|
||||
Then point the router at it. The format is `id=URL=model`:
|
||||
|
||||
```
|
||||
./bin/llmgw server --endpoint e1=http://127.0.0.1:8082=llama
|
||||
```
|
||||
|
||||
## 4. Ask it something
|
||||
|
||||
Streaming:
|
||||
|
||||
```
|
||||
curl -N localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-H 'x-session-id: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}'
|
||||
```
|
||||
|
||||
One reply at once:
|
||||
|
||||
```
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"max_tokens":100}'
|
||||
```
|
||||
|
||||
Also useful: `curl -s localhost:8080/v1/models` and
|
||||
`curl -s localhost:8080/health`. `llmgw ctl` does not work in this mode: it
|
||||
talks to the control plane's admin API, which router-only mode does not have,
|
||||
so `ctl status` answers `404`.
|
||||
|
||||
On the machine this was written on, a 1B bf16 model on an RTX 4070 laptop GPU
|
||||
answered at roughly 85 tokens per second.
|
||||
|
||||
## Several replicas of one model
|
||||
|
||||
Give the router more than one `--endpoint` with the same model name and it
|
||||
spreads requests across them. `x-session-id` keeps a conversation on the same
|
||||
replica so its cache is reused:
|
||||
|
||||
```
|
||||
./bin/llmgw server \
|
||||
--endpoint e1=http://127.0.0.1:8082=llama \
|
||||
--endpoint e2=http://127.0.0.1:8083=llama
|
||||
```
|
||||
|
||||
## Stopping
|
||||
|
||||
```
|
||||
pkill -f "llmgw server"
|
||||
pkill -f "llmgw sim-engine"
|
||||
```
|
||||
|
||||
Stop the engine the way you started it (Ctrl-C in its terminal).
|
||||
|
||||
## What goes wrong
|
||||
|
||||
**`listening on :8080: bind: address already in use`.** An earlier router is
|
||||
still running. `ss -ltnp | grep 8080` names the process; stop it first. This is
|
||||
the usual result of moving from step 2 to step 3 without stopping the fake
|
||||
engine's router.
|
||||
|
||||
**`/v1/models` on the engine returns an empty list.** The engine was started
|
||||
without `-m`, so it is a router with nothing to load. Restart it with a model.
|
||||
|
||||
**`error: nothing to serve`.** `llmgw server` with no `--endpoint` and no model
|
||||
carrying `profiles` has nothing to do. In router-only mode, always pass at least
|
||||
one `--endpoint`.
|
||||
|
||||
**Requests fail with a model-not-found error.** The name in the request must
|
||||
match the name after the last `=` in `--endpoint`, and it must be a name the
|
||||
engine actually serves (its `--alias`).
|
||||
|
||||
**Answers are poor.** A 1B model is a plumbing test, not a chat assistant. Swap
|
||||
in a larger `.gguf`; nothing else changes.
|
||||
|
||||
## 5. Let llmgw run the engine (control plane and node agent)
|
||||
|
||||
Steps 2 and 3 start the engine by hand. Here a **node agent** starts and
|
||||
supervises llama.cpp for you, and a control plane decides where the model runs.
|
||||
Still one machine, still no `sudo`. Stop everything from steps 2 and 3 first:
|
||||
the GPU has to be free and port 8080 unused.
|
||||
|
||||
```
|
||||
curl :8080 -> llmgw server (control plane + router) <- llmgw agent (this node)
|
||||
|
|
||||
starts llama-server :20000
|
||||
```
|
||||
|
||||
The agent dials the control plane over mutual TLS, so this needs a small private
|
||||
certificate authority. Everything goes in a scratch directory, `devcluster/`,
|
||||
which is in `.gitignore`.
|
||||
|
||||
### Certificates
|
||||
|
||||
```
|
||||
mkdir -p devcluster/pki
|
||||
./bin/llmgw ctl cert init-ca --dir devcluster/pki
|
||||
./bin/llmgw ctl cert issue-server llmgw-ctl --host 127.0.0.1 --host localhost \
|
||||
--dir devcluster/pki --out devcluster
|
||||
./bin/llmgw ctl cert issue-client llmgw-node01 --dir devcluster/pki --out devcluster
|
||||
```
|
||||
|
||||
The name given to `issue-client` becomes the certificate's common name and must
|
||||
equal `node_id` in the agent config. The control plane refuses an agent whose
|
||||
claimed id differs from its certificate.
|
||||
|
||||
### The engine wrapper
|
||||
|
||||
The agent starts the engine as `<bin> --model ... --port ...`. A `llama` binary
|
||||
that wants `llama serve ...` refuses every flag, so put a wrapper in front of it:
|
||||
|
||||
```
|
||||
cat > devcluster/llama-server <<'EOF'
|
||||
#!/bin/sh
|
||||
exec /path/to/llama serve "$@"
|
||||
EOF
|
||||
chmod +x devcluster/llama-server
|
||||
|
||||
./bin/llmgw ctl engine-check llamacpp --bin devcluster/llama-server
|
||||
```
|
||||
|
||||
`engine-check` should end with `every flag llmgw builds is accepted`. If your
|
||||
build already ships a `llama-server` binary, point `bin` at that and skip the
|
||||
wrapper.
|
||||
|
||||
### Control plane config: `devcluster/server.yaml`
|
||||
|
||||
```yaml
|
||||
listen: "127.0.0.1:8080"
|
||||
admin_listen: "127.0.0.1:9090"
|
||||
agent_listen: "127.0.0.1:7443"
|
||||
|
||||
agent_tls:
|
||||
ca: devcluster/pki/ca.pem
|
||||
cert: devcluster/server.pem
|
||||
key: devcluster/server-key.pem
|
||||
|
||||
models:
|
||||
- name: llama
|
||||
source: /path/to/Llama-3.2-1B-Instruct-BF16.gguf
|
||||
routing_profile: chat
|
||||
context_length: 8192
|
||||
scaling: { min: 1, max: 1 }
|
||||
profiles:
|
||||
- name: gpu
|
||||
engine: llamacpp
|
||||
gpu_class: 4070-8g
|
||||
weight: 1
|
||||
priority: 100
|
||||
startup_timeout: 5m
|
||||
```
|
||||
|
||||
Three things in there are easy to get wrong:
|
||||
|
||||
- **The model must have at least one profile.** That is what switches the control
|
||||
plane on. With `models: []` the server starts as a plain router and exits with
|
||||
`nothing to serve`.
|
||||
- **`gpu_class` decides CPU or GPU.** Without it the replica is placed in the CPU
|
||||
pool and runs on the CPU. Get the class name from `./bin/llmgw agent inventory`
|
||||
(the `CLASS` column).
|
||||
- **There is no `auth:` block.** Leave it out for a local trial; with auth on,
|
||||
requests need an API key.
|
||||
|
||||
### Agent config: `devcluster/agent.yaml`
|
||||
|
||||
```yaml
|
||||
node_id: llmgw-node01
|
||||
server: "https://127.0.0.1:7443"
|
||||
mgmt_addr: "127.0.0.1"
|
||||
|
||||
tls:
|
||||
ca: devcluster/pki/ca.pem
|
||||
cert: devcluster/llmgw-node01.pem
|
||||
key: devcluster/llmgw-node01-key.pem
|
||||
|
||||
engines:
|
||||
llamacpp:
|
||||
bin: /absolute/path/to/devcluster/llama-server
|
||||
|
||||
supervisor:
|
||||
work_dir: devcluster/agent-work
|
||||
```
|
||||
|
||||
### Run it
|
||||
|
||||
Two terminals:
|
||||
|
||||
```
|
||||
./bin/llmgw server --config devcluster/server.yaml
|
||||
./bin/llmgw agent --config devcluster/agent.yaml
|
||||
```
|
||||
|
||||
The model is not placed at once. The control plane waits about 35 seconds for the
|
||||
fleet to settle after it starts, then places the replica, and the agent then
|
||||
loads the weights. Watch it with:
|
||||
|
||||
```
|
||||
./bin/llmgw ctl nodes
|
||||
./bin/llmgw ctl models
|
||||
./bin/llmgw ctl replicas
|
||||
./bin/llmgw ctl explain-placement llama
|
||||
```
|
||||
|
||||
`STATUS ready` in `ctl models` means it is serving. `nvidia-smi` should now show
|
||||
the model in GPU memory. The same curl from step 4 works against port 8080.
|
||||
|
||||
`ctl` works here, unlike in router-only mode, because there is a control plane
|
||||
for it to talk to.
|
||||
|
||||
Stop it with Ctrl-C in each terminal. The agent leaves the engine running when
|
||||
the agent itself stops, so if you restart the agent and the GPU stays full,
|
||||
`pgrep -af "llama serve"` will show the engine still up.
|
||||
|
||||
## 6. Memory: conversations and notes
|
||||
|
||||
Off unless you ask. llmgw otherwise stores nothing anyone typed. On a single
|
||||
machine with no API keys, one block in the server config turns on two things:
|
||||
|
||||
```yaml
|
||||
threads:
|
||||
store: devcluster/data/threads.db
|
||||
retain_hours: 24
|
||||
```
|
||||
|
||||
Restart the server after adding it. `retain_hours` is required: everything kept
|
||||
expires that long after it was last written.
|
||||
|
||||
Everything kept this way sits in one shared pool with no owner, and anyone who
|
||||
can reach the port and knows an id can read it. That is fine on one machine and
|
||||
wrong on a shared network. With auth on, threads and notes are per project
|
||||
instead and the block is ignored; see [install.md](install.md).
|
||||
|
||||
### Conversations (threads)
|
||||
|
||||
Send the same `x-llmgw-thread` id and only the new turn. llmgw puts the stored
|
||||
history in front of it:
|
||||
|
||||
```
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' -H 'x-llmgw-thread: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"My name is Deniz."}]}'
|
||||
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' -H 'x-llmgw-thread: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"What is my name?"}]}'
|
||||
|
||||
curl -X DELETE localhost:8080/v1/threads/demo # forget it
|
||||
```
|
||||
|
||||
Without the block, a request carrying `x-llmgw-thread` is refused with
|
||||
`this gateway keeps no conversations` rather than answered without its history.
|
||||
|
||||
### Notes an agent keeps for itself
|
||||
|
||||
Notes outlive any one conversation. Each is a key, some text and optional tags,
|
||||
kept per agent; name the agent with `x-llmgw-agent` (default `default`), so two
|
||||
agents using the same key do not read each other's notes.
|
||||
|
||||
```
|
||||
curl -X PUT localhost:8080/v1/memory/favorite-color \
|
||||
-H 'content-type: application/json' \
|
||||
-d '{"text":"The user likes dark green.","tags":["preference"]}'
|
||||
|
||||
curl 'localhost:8080/v1/memory?q=green' # search by words
|
||||
curl 'localhost:8080/v1/memory?tag=preference' # or by tag
|
||||
curl localhost:8080/v1/memory/favorite-color
|
||||
curl -X DELETE localhost:8080/v1/memory/favorite-color
|
||||
```
|
||||
|
||||
The same notes are offered to an agent as tools over MCP at
|
||||
`http://localhost:8080/mcp`: `remember`, `recall` and `forget`. Add that URL to
|
||||
your agent's MCP server list (HTTP transport; how you set request headers such
|
||||
as `x-llmgw-agent` depends on the client). llmgw does not put these tools into
|
||||
chat requests on its own, so an agent that is not given the server never sees
|
||||
memory.
|
||||
|
||||
Limits: a key is up to 200 bytes, a note's text up to 16 KiB, at most 16 tags,
|
||||
1000 notes per agent. Tags are lowercase letters, digits, `-` and `_`.
|
||||
|
||||
## Where this stops
|
||||
|
||||
Router-only mode (steps 2 to 4) has no authentication, no placement and no model
|
||||
management: you start and restart engines yourself. Step 5 adds placement, model
|
||||
management and a node agent, and can add authentication as `install.md`
|
||||
describes. For a real cluster, with several nodes, systemd and a persistent
|
||||
certificate authority, follow [install.md](install.md).
|
||||
|
||||
One caveat found while writing this page: `install.md` recommends starting with
|
||||
an empty `models: []` and adding models later with `ctl model add`. The control
|
||||
plane only turns on when the config declares a model with `profiles`, so that
|
||||
first start does not produce a control plane.
|
||||
@@ -14,23 +14,33 @@ require (
|
||||
github.com/prometheus/client_model v0.6.3
|
||||
github.com/prometheus/common v0.71.0
|
||||
github.com/spf13/cobra v1.10.2
|
||||
golang.org/x/net v0.59.0
|
||||
golang.org/x/sync v0.23.0
|
||||
google.golang.org/protobuf v1.36.12
|
||||
modernc.org/sqlite v1.59.0
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/beorn7/perks v1.0.1 // indirect
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/fsnotify/fsnotify v1.9.0 // indirect
|
||||
github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
github.com/knadh/koanf/maps v0.1.2 // indirect
|
||||
github.com/mattn/go-isatty v0.0.24 // indirect
|
||||
github.com/mitchellh/copystructure v1.2.0 // indirect
|
||||
github.com/mitchellh/reflectwalk v1.0.2 // indirect
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
|
||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||
github.com/prometheus/procfs v0.21.1 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/spf13/pflag v1.0.9 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
golang.org/x/net v0.59.0 // indirect
|
||||
golang.org/x/crypto v0.57.0 // indirect
|
||||
golang.org/x/sys v0.48.0 // indirect
|
||||
golang.org/x/text v0.42.0 // indirect
|
||||
modernc.org/libc v1.75.7 // indirect
|
||||
modernc.org/mathutil v1.7.1 // indirect
|
||||
modernc.org/memory v1.12.1 // indirect
|
||||
)
|
||||
@@ -7,12 +7,20 @@ github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6r
|
||||
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k=
|
||||
github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
|
||||
github.com/go-viper/mapstructure/v2 v2.4.0 h1:EBsztssimR/CONLSZZ04E8qAkxNYq4Qp9LvH92wZUgs=
|
||||
github.com/go-viper/mapstructure/v2 v2.4.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3 h1:LMLX+LgTNWpfvCBdFebv6EsYotImrt/Ppc5cXIriCSo=
|
||||
github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3/go.mod h1:jl5iWTm0/hd5PjEYEOuwAJ57L/CibdZfrqZ5XA5GrCk=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
|
||||
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
||||
github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk=
|
||||
@@ -29,12 +37,16 @@ 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/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
|
||||
github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
|
||||
github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw=
|
||||
github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s=
|
||||
github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ=
|
||||
github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU=
|
||||
github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE=
|
||||
github.com/prometheus/client_model v0.6.3 h1:O0jaTVAYNxTHYInEPFJt5I3+sN8zqBtVMPTB1qyxiEo=
|
||||
@@ -43,6 +55,8 @@ github.com/prometheus/common v0.71.0 h1:9KDAKb7Mj3HEVKyFCK6Dc/HIwlBzZIN2l7/lrHl3
|
||||
github.com/prometheus/common v0.71.0/go.mod h1:CLJ5H8TEsGX8bl31BdMkfhIZ+QmZ9tBPPotUxUbfcmk=
|
||||
github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI=
|
||||
github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
|
||||
github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
|
||||
github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
|
||||
@@ -57,6 +71,10 @@ go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ=
|
||||
go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
|
||||
go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw=
|
||||
go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg=
|
||||
golang.org/x/crypto v0.57.0 h1:3ZVCjf8Ggz7zneR/EHRVx68Ctf+2pmIMP2UFhh9cC6M=
|
||||
golang.org/x/crypto v0.57.0/go.mod h1:Fdz0i5U6CoizGwLda9DttjSk6qlZo25zYNtR+ycvuZA=
|
||||
golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c=
|
||||
golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o=
|
||||
golang.org/x/net v0.59.0 h1:5zfYln+w5XCxwrnMMJPufRgNoXEaGxl0wo5GqPXyues=
|
||||
golang.org/x/net v0.59.0/go.mod h1:2DA/G1UfVbCpQPeWTmMPGY7Cs2PkBkwu743bVX5PIVg=
|
||||
golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk=
|
||||
@@ -65,6 +83,36 @@ golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
|
||||
golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og=
|
||||
golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI=
|
||||
golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E=
|
||||
golang.org/x/tools v0.49.0 h1:3NI7VXzL9+1WZD52Dx2ttoPwD5DWrFGpl9mFZDlmisI=
|
||||
golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo=
|
||||
google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc=
|
||||
google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
modernc.org/cc/v4 v4.29.2 h1:h6+9ciCnPKutf4I03CvheAvDLX7+IHlqR6Iy6J+cgd8=
|
||||
modernc.org/cc/v4 v4.29.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
|
||||
modernc.org/ccgo/v4 v4.35.0 h1:F+TUsmw09QxLzmi3aeYYGxjAXarmZaKgj3mKQHNaA8w=
|
||||
modernc.org/ccgo/v4 v4.35.0/go.mod h1:qrVGs9S3Sr2Ztcg9ve+kTAYMp5a3YvWjo+SoN06kJ5I=
|
||||
modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
|
||||
modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
|
||||
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
|
||||
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
|
||||
modernc.org/gc/v3 v3.1.5 h1:21ldfPfRYE31Tb7B3mwAK8gy1AxP4+dKjrOQPfqakoc=
|
||||
modernc.org/gc/v3 v3.1.5/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
|
||||
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
|
||||
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
|
||||
modernc.org/libc v1.75.7 h1:o3DTP9/0p9pKmY2WCKQaySW6wIiZhNM7wc2lUoyhfew=
|
||||
modernc.org/libc v1.75.7/go.mod h1:bO5o2ztHxBb2rjz0PgdHN0sSMw57CgxGFLZ3Qd/QpVQ=
|
||||
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||
modernc.org/memory v1.12.1 h1:nFMiWrpStgZczNl6XI9GnIk/rWhYIyHGUaR04pGbp9g=
|
||||
modernc.org/memory v1.12.1/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
|
||||
modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg=
|
||||
modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
|
||||
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
|
||||
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
|
||||
modernc.org/sqlite v1.59.0 h1:X1es1GpqBlS/5T+vbM4HLUdaa8OtQx468DF2vrx+38A=
|
||||
modernc.org/sqlite v1.59.0/go.mod h1:+paeT2A3iPRHkQDwG7oA6Tk0zQd5woMEI8q7orfry8k=
|
||||
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
||||
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
|
||||
@@ -0,0 +1,396 @@
|
||||
package account
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func newStore(t *testing.T) *store.SQLite {
|
||||
t.Helper()
|
||||
s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "a.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = s.Close() })
|
||||
|
||||
ctx := context.Background()
|
||||
for _, err := range []error{
|
||||
s.PutOrg(ctx, store.Org{ID: "org", Name: "Org"}),
|
||||
s.PutTeam(ctx, store.Team{ID: "team", OrgID: "org", Name: "Team"}),
|
||||
s.PutProject(ctx, store.Project{ID: "proj", TeamID: "team", Name: "Proj"}),
|
||||
} {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func newService(t *testing.T) (*Service, *store.SQLite) {
|
||||
t.Helper()
|
||||
s := newStore(t)
|
||||
return New(s, DefaultConfig(), slog.New(slog.DiscardHandler)), s
|
||||
}
|
||||
|
||||
const goodPassword = "correct horse battery staple"
|
||||
|
||||
func TestSignupLeavesAnAccountWaiting(t *testing.T) {
|
||||
// A public signup form that produced a working account would be a way in.
|
||||
svc, st := newService(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := svc.SignUp(ctx, "Ayse@Example.COM", "Ayşe", goodPassword); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u, err := st.UserByEmail(ctx, "ayse@example.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if u.State != store.UserPending {
|
||||
t.Errorf("a new account is %q", u.State)
|
||||
}
|
||||
if u.Email != "ayse@example.com" {
|
||||
t.Errorf("the address was not normalized: %q", u.Email)
|
||||
}
|
||||
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); !errors.Is(err, ErrNotApproved) {
|
||||
t.Errorf("a pending account signed in: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSignupDoesNotRevealExistingAccounts(t *testing.T) {
|
||||
// A form that says "this address already has an account" is an account
|
||||
// enumerator anyone can run, and the person signing up learns nothing they
|
||||
// need: either way they wait for approval.
|
||||
svc, _ := newService(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := svc.SignUp(ctx, "ayse@example.com", "A", goodPassword); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := svc.SignUp(ctx, "ayse@example.com", "Someone Else", goodPassword); err != nil {
|
||||
t.Errorf("the second signup reported %v, which confirms the address is taken", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalIsWhatMakesAnAccountWork(t *testing.T) {
|
||||
svc, st := newService(t)
|
||||
ctx := context.Background()
|
||||
must(t, svc.SignUp(ctx, "ayse@example.com", "A", goodPassword))
|
||||
|
||||
u, _ := st.UserByEmail(ctx, "ayse@example.com")
|
||||
must(t, st.ApproveUser(ctx, u.ID, "proj", "ilker", time.Now()))
|
||||
|
||||
token, csrf, err := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if token == "" || csrf == "" {
|
||||
t.Fatal("sign-in produced no session")
|
||||
}
|
||||
_, got, err := svc.Resolve(ctx, token)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.ProjectID != "proj" {
|
||||
t.Errorf("the approved account landed in project %q", got.ProjectID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalNeedsAProject(t *testing.T) {
|
||||
// An approved account with nowhere to put its keys passes every check and
|
||||
// then fails at the first key it tries to make.
|
||||
_, st := newStoreAndUser(t)
|
||||
if err := st.ApproveUser(context.Background(), "usr-1", "", "ilker", time.Now()); err == nil {
|
||||
t.Error("an account was approved into no project")
|
||||
}
|
||||
}
|
||||
|
||||
func newStoreAndUser(t *testing.T) (*Service, *store.SQLite) {
|
||||
t.Helper()
|
||||
svc, st := newService(t)
|
||||
hash, _ := HashPassword(goodPassword)
|
||||
must(t, st.CreateUser(context.Background(), store.User{
|
||||
ID: "usr-1", Email: "ayse@example.com", Name: "A",
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: time.Now(),
|
||||
}))
|
||||
return svc, st
|
||||
}
|
||||
|
||||
func TestSuspensionTakesEffectOnOpenSessions(t *testing.T) {
|
||||
// Leaving existing sessions open would make suspension advisory.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
token, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := svc.Resolve(ctx, token); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
must(t, st.SetUserState(ctx, "usr-1", store.UserSuspended))
|
||||
if _, _, err := svc.Resolve(ctx, token); !errors.Is(err, store.ErrNoSession) {
|
||||
t.Errorf("a suspended account kept its session: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSignOutEndsTheSession(t *testing.T) {
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
token, _, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
must(t, svc.SignOut(ctx, token))
|
||||
if _, _, err := svc.Resolve(ctx, token); !errors.Is(err, store.ErrNoSession) {
|
||||
t.Errorf("the session survived sign-out: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWrongPasswordAndAMissingAccountAreIndistinguishable(t *testing.T) {
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
_, _, wrong := svc.SignIn(ctx, "ayse@example.com", "wrong password here")
|
||||
_, _, missing := svc.SignIn(ctx, "nobody@example.com", "wrong password here")
|
||||
if !errors.Is(wrong, ErrBadCredentials) || !errors.Is(missing, ErrBadCredentials) {
|
||||
t.Fatalf("wrong=%v missing=%v", wrong, missing)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGuessingIsBounded(t *testing.T) {
|
||||
// Without this a public login form is an open brute-force target.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.MaxAttempts = 3
|
||||
svc = New(st, cfg, slog.New(slog.DiscardHandler))
|
||||
|
||||
for i := range 3 {
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", "nope"); !errors.Is(err, ErrBadCredentials) {
|
||||
t.Fatalf("attempt %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", "nope"); !errors.Is(err, ErrTooManyTries) {
|
||||
t.Errorf("guessing was not bounded: %v", err)
|
||||
}
|
||||
// And a correct password is refused too while the window is open, or the
|
||||
// limit would be trivially bypassed by interleaving.
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); !errors.Is(err, ErrTooManyTries) {
|
||||
t.Errorf("the window did not apply to a correct password: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASuccessfulSignInClearsTheWindow(t *testing.T) {
|
||||
// Otherwise a person who mistypes their password a few times is locked out
|
||||
// for the rest of the window even after getting it right.
|
||||
st := newStore(t)
|
||||
hash, _ := HashPassword(goodPassword)
|
||||
ctx := context.Background()
|
||||
must(t, st.CreateUser(ctx, store.User{
|
||||
ID: "usr-1", Email: "a@example.com", Name: "A",
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: time.Now(),
|
||||
}))
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.MaxAttempts = 3
|
||||
svc := New(st, cfg, slog.New(slog.DiscardHandler))
|
||||
|
||||
for range 2 {
|
||||
_, _, _ = svc.SignIn(ctx, "a@example.com", "nope")
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "a@example.com", goodPassword); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for range 3 {
|
||||
if _, _, err := svc.SignIn(ctx, "a@example.com", "nope"); errors.Is(err, ErrTooManyTries) {
|
||||
t.Fatal("the window was not cleared by a successful sign-in")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCSRFTokensAreCheckedAgainstTheSession(t *testing.T) {
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
token, csrf, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
sess, _, _ := svc.Resolve(ctx, token)
|
||||
|
||||
if !svc.CheckCSRF(sess, csrf) {
|
||||
t.Error("the session's own token was rejected")
|
||||
}
|
||||
for _, bad := range []string{"", "wrong", csrf + "x", csrf[:len(csrf)-1]} {
|
||||
if svc.CheckCSRF(sess, bad) {
|
||||
t.Errorf("%q was accepted as a csrf token", bad)
|
||||
}
|
||||
}
|
||||
// A token from another session must not work in this one.
|
||||
other, otherCSRF, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
_ = other
|
||||
if svc.CheckCSRF(sess, otherCSRF) {
|
||||
t.Error("another session's token was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasswordsAreHashedNotStored(t *testing.T) {
|
||||
svc, st := newService(t)
|
||||
ctx := context.Background()
|
||||
must(t, svc.SignUp(ctx, "ayse@example.com", "A", goodPassword))
|
||||
|
||||
u, _ := st.UserByEmail(ctx, "ayse@example.com")
|
||||
if strings.Contains(u.PasswordHash, goodPassword) {
|
||||
t.Fatal("the password is in the database")
|
||||
}
|
||||
if !strings.HasPrefix(u.PasswordHash, "$argon2id$") {
|
||||
t.Errorf("hash format: %q", u.PasswordHash)
|
||||
}
|
||||
ok, err := VerifyPassword(goodPassword, u.PasswordHash)
|
||||
if err != nil || !ok {
|
||||
t.Errorf("the hash does not verify: %v %v", ok, err)
|
||||
}
|
||||
if ok, _ := VerifyPassword("something else entirely", u.PasswordHash); ok {
|
||||
t.Error("a wrong password verified")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTwoHashesOfOnePasswordDiffer(t *testing.T) {
|
||||
// A salt that did not vary would make the table a rainbow-table lookup.
|
||||
a, _ := HashPassword(goodPassword)
|
||||
b, _ := HashPassword(goodPassword)
|
||||
if a == b {
|
||||
t.Fatal("hashing is not salted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestShortPasswordsAreRefused(t *testing.T) {
|
||||
if _, err := HashPassword("short"); !errors.Is(err, ErrWeakPassword) {
|
||||
t.Errorf("a five-character password was accepted: %v", err)
|
||||
}
|
||||
svc, _ := newService(t)
|
||||
if err := svc.SignUp(context.Background(), "a@example.com", "A", "short"); !errors.Is(err, ErrWeakPassword) {
|
||||
t.Errorf("signup accepted it: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMalformedHashesAreRejectedRatherThanMatched(t *testing.T) {
|
||||
for _, bad := range []string{"", "plaintext", "$argon2id$", "$bcrypt$v=19$m=1,t=1,p=1$aa$bb"} {
|
||||
if ok, err := VerifyPassword(goodPassword, bad); ok || err == nil {
|
||||
t.Errorf("VerifyPassword(%q) = %v, %v", bad, ok, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestObviousNonAddressesAreRefused(t *testing.T) {
|
||||
svc, _ := newService(t)
|
||||
ctx := context.Background()
|
||||
for _, bad := range []string{"", "ayse", "ayse@", "@example.com", "a b@example.com", "a@b", "a@@b.com"} {
|
||||
if err := svc.SignUp(ctx, bad, "A", goodPassword); !errors.Is(err, ErrInvalidEmail) {
|
||||
t.Errorf("SignUp(%q) = %v", bad, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func must(t *testing.T, err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChangingAPasswordNeedsTheOldOne(t *testing.T) {
|
||||
// A session is what an attacker gets from a borrowed laptop or a stolen
|
||||
// cookie. The password is what they do not have, and without this check the
|
||||
// first thing they would do is take the account.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
u, _ := st.UserByID(ctx, "usr-1")
|
||||
|
||||
if err := svc.ChangePassword(ctx, st, u, "not the password", "a brand new passphrase"); !errors.Is(err, ErrBadCredentials) {
|
||||
t.Errorf("a wrong current password was accepted: %v", err)
|
||||
}
|
||||
// And the old one still works.
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); err != nil {
|
||||
t.Errorf("the password changed anyway: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChangingAPasswordSignsOutEverywhere(t *testing.T) {
|
||||
// Leaving other sessions open would mean a password change does nothing
|
||||
// about the reason someone is changing it.
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
laptop, _, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
phone, _, _ := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
|
||||
u, _ := st.UserByID(ctx, "usr-1")
|
||||
must(t, svc.ChangePassword(ctx, st, u, goodPassword, "a brand new passphrase"))
|
||||
|
||||
for name, token := range map[string]string{"laptop": laptop, "phone": phone} {
|
||||
if _, _, err := svc.Resolve(ctx, token); !errors.Is(err, store.ErrNoSession) {
|
||||
t.Errorf("the %s session survived the change: %v", name, err)
|
||||
}
|
||||
}
|
||||
// The new password works and the old one does not.
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", "a brand new passphrase"); err != nil {
|
||||
t.Errorf("the new password does not work: %v", err)
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword); !errors.Is(err, ErrBadCredentials) {
|
||||
t.Errorf("the old password still works: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANewPasswordMustActuallyBeNew(t *testing.T) {
|
||||
svc, st := newStoreAndUser(t)
|
||||
ctx := context.Background()
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
u, _ := st.UserByID(ctx, "usr-1")
|
||||
|
||||
if err := svc.ChangePassword(ctx, st, u, goodPassword, goodPassword); !errors.Is(err, ErrSamePassword) {
|
||||
t.Errorf("changing a password to itself: %v", err)
|
||||
}
|
||||
if err := svc.ChangePassword(ctx, st, u, goodPassword, "short"); !errors.Is(err, ErrWeakPassword) {
|
||||
t.Errorf("a short new password was accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGuessingTheCurrentPasswordIsBounded(t *testing.T) {
|
||||
// Someone with a stolen session should not get unlimited attempts at the
|
||||
// password that would let them keep it.
|
||||
st := newStore(t)
|
||||
ctx := context.Background()
|
||||
hash, _ := HashPassword(goodPassword)
|
||||
must(t, st.CreateUser(ctx, store.User{
|
||||
ID: "usr-1", Email: "a@example.com", Name: "A",
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: time.Now(),
|
||||
}))
|
||||
must(t, st.ApproveUser(ctx, "usr-1", "proj", "ilker", time.Now()))
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.MaxAttempts = 3
|
||||
svc := New(st, cfg, slog.New(slog.DiscardHandler))
|
||||
u, _ := st.UserByID(ctx, "usr-1")
|
||||
|
||||
for range 3 {
|
||||
_ = svc.ChangePassword(ctx, st, u, "wrong", "a brand new passphrase")
|
||||
}
|
||||
if err := svc.ChangePassword(ctx, st, u, "wrong", "a brand new passphrase"); !errors.Is(err, ErrTooManyTries) {
|
||||
t.Errorf("guessing was unbounded: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
// Package account is the human half of authentication: people sign up, an
|
||||
// operator approves them, and they issue their own API keys from a profile page
|
||||
// (plan §7.9).
|
||||
//
|
||||
// This is deliberately a different mechanism from the one in internal/auth.
|
||||
// That one authenticates machines holding 256-bit random keys; this one
|
||||
// authenticates people holding passwords they chose. The two have opposite
|
||||
// requirements — a fast hash is correct there and wrong here — and keeping them
|
||||
// apart is what stops one set of assumptions leaking into the other.
|
||||
package account
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/crypto/argon2"
|
||||
)
|
||||
|
||||
// Argon2 parameters. OWASP's current floor for argon2id, which costs roughly
|
||||
// 50ms per attempt on ordinary hardware — slow enough to make an offline
|
||||
// dictionary attack expensive, fast enough that a login does not feel broken.
|
||||
//
|
||||
// They are stored with each hash, so raising them later leaves existing
|
||||
// passwords working until their owners next sign in.
|
||||
const (
|
||||
argonTime = 3
|
||||
argonMemory = 64 * 1024 // KiB
|
||||
argonThreads = 4
|
||||
argonKeyLen = 32
|
||||
saltLen = 16
|
||||
)
|
||||
|
||||
// MinPasswordLen is the only password rule.
|
||||
//
|
||||
// Length is the one requirement that reliably buys entropy. Composition rules —
|
||||
// a digit, a symbol, a capital — push people towards Password1! and produce
|
||||
// less entropy than they cost, so there are none.
|
||||
const MinPasswordLen = 12
|
||||
|
||||
// ErrWeakPassword says a password is too short to accept.
|
||||
var ErrWeakPassword = fmt.Errorf("a password must be at least %d characters", MinPasswordLen)
|
||||
|
||||
// ErrBadHash says a stored hash could not be parsed.
|
||||
var ErrBadHash = errors.New("unreadable password hash")
|
||||
|
||||
// HashPassword derives a storable hash.
|
||||
func HashPassword(password string) (string, error) {
|
||||
if len([]rune(password)) < MinPasswordLen {
|
||||
return "", ErrWeakPassword
|
||||
}
|
||||
salt := make([]byte, saltLen)
|
||||
if _, err := rand.Read(salt); err != nil {
|
||||
return "", err
|
||||
}
|
||||
sum := argon2.IDKey([]byte(password), salt, argonTime, argonMemory, argonThreads, argonKeyLen)
|
||||
|
||||
b64 := base64.RawStdEncoding.EncodeToString
|
||||
return fmt.Sprintf("$argon2id$v=%d$m=%d,t=%d,p=%d$%s$%s",
|
||||
argon2.Version, argonMemory, argonTime, argonThreads, b64(salt), b64(sum)), nil
|
||||
}
|
||||
|
||||
// VerifyPassword reports whether a password matches a stored hash.
|
||||
//
|
||||
// The comparison is constant time. The parameters come from the hash rather
|
||||
// than from the constants above, so a hash written under older settings still
|
||||
// verifies.
|
||||
func VerifyPassword(password, encoded string) (bool, error) {
|
||||
parts := strings.Split(encoded, "$")
|
||||
if len(parts) != 6 || parts[1] != "argon2id" {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
|
||||
var version int
|
||||
if _, err := fmt.Sscanf(parts[2], "v=%d", &version); err != nil || version != argon2.Version {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
var memory, time uint32
|
||||
var threads uint8
|
||||
if _, err := fmt.Sscanf(parts[3], "m=%d,t=%d,p=%d", &memory, &time, &threads); err != nil {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
|
||||
salt, err := base64.RawStdEncoding.DecodeString(parts[4])
|
||||
if err != nil {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
want, err := base64.RawStdEncoding.DecodeString(parts[5])
|
||||
if err != nil {
|
||||
return false, ErrBadHash
|
||||
}
|
||||
|
||||
got := argon2.IDKey([]byte(password), salt, time, memory, threads, uint32(len(want)))
|
||||
return subtle.ConstantTimeCompare(got, want) == 1, nil
|
||||
}
|
||||
|
||||
// dummyHash is verified against when no account matches, so a login attempt for
|
||||
// an address that does not exist costs the same time as one that does.
|
||||
//
|
||||
// Without it, a stopwatch enumerates every account on the system.
|
||||
var dummyHash string
|
||||
|
||||
func init() {
|
||||
h, err := HashPassword(strings.Repeat("x", MinPasswordLen))
|
||||
if err != nil {
|
||||
panic("account: cannot build the dummy hash: " + err.Error())
|
||||
}
|
||||
dummyHash = h
|
||||
}
|
||||
@@ -0,0 +1,330 @@
|
||||
package account
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/base32"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// What a sign-in or sign-up can fail with.
|
||||
var (
|
||||
ErrBadCredentials = errors.New("that email and password do not match an account")
|
||||
ErrNotApproved = errors.New("this account is waiting for approval")
|
||||
ErrSuspended = errors.New("this account has been suspended")
|
||||
ErrTooManyTries = errors.New("too many attempts; try again shortly")
|
||||
ErrInvalidEmail = errors.New("that does not look like an email address")
|
||||
)
|
||||
|
||||
// Store is what the service needs from persistence.
|
||||
type Store interface {
|
||||
CreateUser(ctx context.Context, u store.User) error
|
||||
UserByEmail(ctx context.Context, email string) (store.User, error)
|
||||
UserByID(ctx context.Context, id string) (store.User, error)
|
||||
TouchLogin(ctx context.Context, id string, at time.Time) error
|
||||
|
||||
CreateSession(ctx context.Context, userID, token, csrf string, ttl time.Duration) error
|
||||
SessionByToken(ctx context.Context, token string) (store.Session, store.User, error)
|
||||
DeleteSession(ctx context.Context, token string) error
|
||||
TouchSession(ctx context.Context, token string, at time.Time) error
|
||||
}
|
||||
|
||||
// Config tunes the service.
|
||||
type Config struct {
|
||||
// SignupOpen decides whether anyone may create an account. Off, the form is
|
||||
// not served at all; accounts are made with `llmgw users create`.
|
||||
SignupOpen bool
|
||||
// SessionTTL is how long a sign-in lasts.
|
||||
SessionTTL time.Duration
|
||||
// MaxAttempts and AttemptWindow bound password guessing per address.
|
||||
MaxAttempts int
|
||||
AttemptWindow time.Duration
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible values.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
SignupOpen: true,
|
||||
SessionTTL: store.SessionTTL,
|
||||
MaxAttempts: 10,
|
||||
AttemptWindow: 15 * time.Minute,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() {
|
||||
d := DefaultConfig()
|
||||
if c.SessionTTL <= 0 {
|
||||
c.SessionTTL = d.SessionTTL
|
||||
}
|
||||
if c.MaxAttempts <= 0 {
|
||||
c.MaxAttempts = d.MaxAttempts
|
||||
}
|
||||
if c.AttemptWindow <= 0 {
|
||||
c.AttemptWindow = d.AttemptWindow
|
||||
}
|
||||
}
|
||||
|
||||
// Service is signup, sign-in and sessions.
|
||||
type Service struct {
|
||||
store Store
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
attempts map[string]*attemptWindow
|
||||
}
|
||||
|
||||
type attemptWindow struct {
|
||||
count int
|
||||
since time.Time
|
||||
}
|
||||
|
||||
// New builds the service.
|
||||
func New(s Store, cfg Config, log *slog.Logger) *Service {
|
||||
cfg.normalize()
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Service{store: s, cfg: cfg, log: log, now: time.Now, attempts: map[string]*attemptWindow{}}
|
||||
}
|
||||
|
||||
// SignupOpen reports whether the form should be served.
|
||||
func (s *Service) SignupOpen() bool { return s.cfg.SignupOpen }
|
||||
|
||||
// SignUp creates a pending account.
|
||||
//
|
||||
// It returns nil whether or not the address was already taken. A signup form
|
||||
// that says "this address already has an account" is an account enumerator that
|
||||
// anyone can run, and the person signing up learns nothing they need: either
|
||||
// way, what happens next is that they wait for approval.
|
||||
func (s *Service) SignUp(ctx context.Context, email, name, password string) error {
|
||||
email = store.NormalizeEmail(email)
|
||||
if !validEmail(email) {
|
||||
return ErrInvalidEmail
|
||||
}
|
||||
hash, err := HashPassword(password)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if name = strings.TrimSpace(name); name == "" {
|
||||
name = email
|
||||
}
|
||||
|
||||
err = s.store.CreateUser(ctx, store.User{
|
||||
ID: randomID("usr"), Email: email, Name: name,
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: s.now(),
|
||||
})
|
||||
if errors.Is(err, store.ErrEmailTaken) {
|
||||
// Logged so an operator can see it; not told to the caller.
|
||||
s.log.Info("signup for an address that already has an account")
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.log.Info("a new account is waiting for approval", "email", email)
|
||||
return nil
|
||||
}
|
||||
|
||||
// SignIn checks a password and opens a session, returning the session token and
|
||||
// the CSRF token that must accompany the forms in it.
|
||||
func (s *Service) SignIn(ctx context.Context, email, password string) (token, csrf string, err error) {
|
||||
email = store.NormalizeEmail(email)
|
||||
if !s.allowAttempt(email) {
|
||||
return "", "", ErrTooManyTries
|
||||
}
|
||||
|
||||
u, err := s.store.UserByEmail(ctx, email)
|
||||
if errors.Is(err, store.ErrNotFound) {
|
||||
// Verify anyway, against a hash that cannot match. Returning early here
|
||||
// would let a stopwatch enumerate every account on the system.
|
||||
_, _ = VerifyPassword(password, dummyHash)
|
||||
return "", "", ErrBadCredentials
|
||||
}
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
|
||||
ok, err := VerifyPassword(password, u.PasswordHash)
|
||||
if err != nil || !ok {
|
||||
return "", "", ErrBadCredentials
|
||||
}
|
||||
// The password was right, so this was not an attack.
|
||||
s.clearAttempts(email)
|
||||
|
||||
// State is checked after the password, on purpose. Checking first would
|
||||
// tell anyone who guesses an address whether it has a pending account.
|
||||
switch u.State {
|
||||
case store.UserPending:
|
||||
return "", "", ErrNotApproved
|
||||
case store.UserSuspended:
|
||||
return "", "", ErrSuspended
|
||||
}
|
||||
|
||||
if token, err = store.NewSessionToken(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if csrf, err = store.NewSessionToken(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if err := s.store.CreateSession(ctx, u.ID, token, csrf, s.cfg.SessionTTL); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if err := s.store.TouchLogin(ctx, u.ID, s.now()); err != nil {
|
||||
s.log.Warn("could not record a sign-in", "err", err)
|
||||
}
|
||||
return token, csrf, nil
|
||||
}
|
||||
|
||||
// Resolve turns a session cookie into the person holding it.
|
||||
func (s *Service) Resolve(ctx context.Context, token string) (store.Session, store.User, error) {
|
||||
if token == "" {
|
||||
return store.Session{}, store.User{}, store.ErrNoSession
|
||||
}
|
||||
sess, u, err := s.store.SessionByToken(ctx, token)
|
||||
if err != nil {
|
||||
return store.Session{}, store.User{}, err
|
||||
}
|
||||
if err := s.store.TouchSession(ctx, token, s.now()); err != nil {
|
||||
s.log.Debug("could not touch a session", "err", err)
|
||||
}
|
||||
return sess, u, nil
|
||||
}
|
||||
|
||||
// SignOut closes one session.
|
||||
func (s *Service) SignOut(ctx context.Context, token string) error {
|
||||
if token == "" {
|
||||
return nil
|
||||
}
|
||||
return s.store.DeleteSession(ctx, token)
|
||||
}
|
||||
|
||||
// CheckCSRF compares a submitted token against the session's.
|
||||
//
|
||||
// A cookie authenticates the browser, not the page that asked. Without this, any
|
||||
// site the person visits can post a form to ours and have the browser attach the
|
||||
// session automatically.
|
||||
func (s *Service) CheckCSRF(sess store.Session, submitted string) bool {
|
||||
if submitted == "" || len(sess.CSRFHash) == 0 {
|
||||
return false
|
||||
}
|
||||
return subtle.ConstantTimeCompare(store.HashKey(submitted), sess.CSRFHash) == 1
|
||||
}
|
||||
|
||||
// allowAttempt bounds password guessing per address.
|
||||
//
|
||||
// Per address rather than per IP: an attacker has many addresses to come from
|
||||
// and only one to guess against, and limiting by IP would lock out everyone
|
||||
// behind one office NAT the moment somebody fat-fingers their password.
|
||||
func (s *Service) allowAttempt(email string) bool {
|
||||
now := s.now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
w, ok := s.attempts[email]
|
||||
if !ok || now.Sub(w.since) > s.cfg.AttemptWindow {
|
||||
s.attempts[email] = &attemptWindow{count: 1, since: now}
|
||||
return true
|
||||
}
|
||||
w.count++
|
||||
return w.count <= s.cfg.MaxAttempts
|
||||
}
|
||||
|
||||
func (s *Service) clearAttempts(email string) {
|
||||
s.mu.Lock()
|
||||
delete(s.attempts, email)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// validEmail is deliberately crude. A real check is delivery, which this does
|
||||
// not do; this only rejects what is obviously not an address so the database
|
||||
// does not fill with them.
|
||||
func validEmail(s string) bool {
|
||||
at := strings.IndexByte(s, '@')
|
||||
if at <= 0 || at != strings.LastIndexByte(s, '@') || at == len(s)-1 {
|
||||
return false
|
||||
}
|
||||
if strings.ContainsAny(s, " \t\r\n<>\"") {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(s[at+1:], ".")
|
||||
}
|
||||
|
||||
func randomID(prefix string) string {
|
||||
var buf [10]byte
|
||||
_, _ = rand.Read(buf[:])
|
||||
return prefix + "-" + strings.ToLower(
|
||||
base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:]))
|
||||
}
|
||||
|
||||
// ErrSamePassword says the new password is the old one.
|
||||
var ErrSamePassword = errors.New("that is the password you already have")
|
||||
|
||||
// PasswordChanger is the extra storage a password change needs.
|
||||
type PasswordChanger interface {
|
||||
SetPassword(ctx context.Context, userID, hash string) error
|
||||
DeleteUserSessions(ctx context.Context, userID string) error
|
||||
}
|
||||
|
||||
// ChangePassword replaces a person's password and signs them out everywhere.
|
||||
//
|
||||
// The current password is required even though the caller already holds a
|
||||
// session. A session is what an attacker gets from a borrowed laptop or a stolen
|
||||
// cookie; the password is what they do not have, and without this check the
|
||||
// first thing they would do is take the account.
|
||||
//
|
||||
// Every session is closed, this browser's included. Leaving the others open
|
||||
// would mean a password change does nothing about the reason someone is
|
||||
// changing it. The caller re-issues a session for the browser in front of them.
|
||||
func (s *Service) ChangePassword(ctx context.Context, pw PasswordChanger,
|
||||
u store.User, current, next string,
|
||||
) error {
|
||||
if !s.allowAttempt("pw:" + u.ID) {
|
||||
return ErrTooManyTries
|
||||
}
|
||||
ok, err := VerifyPassword(current, u.PasswordHash)
|
||||
if err != nil || !ok {
|
||||
return ErrBadCredentials
|
||||
}
|
||||
s.clearAttempts("pw:" + u.ID)
|
||||
|
||||
if same, _ := VerifyPassword(next, u.PasswordHash); same {
|
||||
return ErrSamePassword
|
||||
}
|
||||
hash, err := HashPassword(next)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := pw.SetPassword(ctx, u.ID, hash); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := pw.DeleteUserSessions(ctx, u.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
s.log.Info("a password was changed", "user", u.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// OpenSession starts a session for a person who has just proved who they are by
|
||||
// some other means — immediately after changing their password, which closed
|
||||
// the one they were using.
|
||||
func (s *Service) OpenSession(ctx context.Context, userID string) (token, csrf string, err error) {
|
||||
if token, err = store.NewSessionToken(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if csrf, err = store.NewSessionToken(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if err := s.store.CreateSession(ctx, userID, token, csrf, s.cfg.SessionTTL); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return token, csrf, nil
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// Refusals a key can hit. They are distinguished so the client is told which
|
||||
// limit it hit, and so retrying is only suggested where retrying helps.
|
||||
var (
|
||||
ErrRateLimited = errors.New("request rate limit exceeded")
|
||||
ErrTokenLimited = errors.New("token rate limit exceeded")
|
||||
ErrBudgetSpent = errors.New("budget exhausted")
|
||||
)
|
||||
|
||||
// LimitError carries which limit was hit and, where it makes sense, when to
|
||||
// come back.
|
||||
type LimitError struct {
|
||||
Err error
|
||||
// RetryAfter is zero when waiting does not help, which is the case for a
|
||||
// spent budget: it will not refill until the period rolls over, and telling
|
||||
// a client to retry in a second would be a lie.
|
||||
RetryAfter time.Duration
|
||||
}
|
||||
|
||||
func (e *LimitError) Error() string { return e.Err.Error() }
|
||||
func (e *LimitError) Unwrap() error { return e.Err }
|
||||
|
||||
// Accountant enforces limits and records what was used.
|
||||
//
|
||||
// One object because the three halves are not independent: a budget check needs
|
||||
// the price, the price needs the model, and the record closes the loop by
|
||||
// feeding the counters that the next check reads.
|
||||
type Accountant struct {
|
||||
pricing *Pricing
|
||||
budgets *Budgets
|
||||
requests *Limiter
|
||||
tokens *Limiter
|
||||
recorder *Recorder
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// New builds an accountant.
|
||||
func New(pricing *Pricing, recorder *Recorder, log *slog.Logger) *Accountant {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Accountant{
|
||||
pricing: pricing,
|
||||
budgets: NewBudgets(),
|
||||
requests: NewLimiter(),
|
||||
tokens: NewLimiter(),
|
||||
recorder: recorder,
|
||||
log: log,
|
||||
}
|
||||
}
|
||||
|
||||
// RestoreBudgets seeds the in-memory counters from the store.
|
||||
//
|
||||
// Called at startup. A restart must not reset a budget, which it would if the
|
||||
// counters only ever lived in memory.
|
||||
func (a *Accountant) RestoreBudgets(ctx context.Context, s store.Store, period store.Period) error {
|
||||
start := PeriodStart(period, time.Now())
|
||||
if start.IsZero() {
|
||||
return nil
|
||||
}
|
||||
spend, err := s.SpendSince(ctx, start)
|
||||
if err != nil {
|
||||
return fmt.Errorf("restoring budgets: %w", err)
|
||||
}
|
||||
a.budgets.Restore(period, spend)
|
||||
a.log.Info("restored budgets from the store", "keys", len(spend), "since", start)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Reservation is what Check returns: the estimate it charged, so Record can
|
||||
// settle the difference against the real figure.
|
||||
type Reservation struct {
|
||||
KeyID string
|
||||
EstimatedTokens int64
|
||||
Price Price
|
||||
// Period is the key's budget period, carried from the check so the record
|
||||
// counts against the same window the check read. Deriving it later would
|
||||
// mean guessing, and a monthly budget credited to a daily counter is a
|
||||
// budget that never binds.
|
||||
Period store.Period
|
||||
}
|
||||
|
||||
// Check decides whether a request may proceed.
|
||||
//
|
||||
// It runs before scheduling and touches no database: every counter it reads is
|
||||
// in memory. The token limit is checked against an estimate, because the real
|
||||
// figure only exists once the request has run; Record settles the difference.
|
||||
func (a *Accountant) Check(p auth.Principal, model string, estimatedTokens int64) (Reservation, error) {
|
||||
price := a.pricing.For(model)
|
||||
res := Reservation{
|
||||
KeyID: p.KeyID, EstimatedTokens: estimatedTokens,
|
||||
Price: price, Period: p.Limits.BudgetPeriod,
|
||||
}
|
||||
if p.KeyID == "" {
|
||||
return res, nil
|
||||
}
|
||||
lim := p.Limits
|
||||
|
||||
// Budget first. It is the limit whose refusal cannot be retried away, so
|
||||
// telling the client about it first is more useful than reporting a rate
|
||||
// limit it would hit again in a second.
|
||||
//
|
||||
// A model that costs nothing is exempt, which is what lets a key with no
|
||||
// budget left still be served a cheap fallback (plan §7.7).
|
||||
if lim.Budget != nil && !price.Zero() &&
|
||||
a.budgets.Exhausted(p.KeyID, lim.BudgetPeriod, *lim.Budget) {
|
||||
return res, &LimitError{Err: ErrBudgetSpent}
|
||||
}
|
||||
|
||||
if lim.RPM != nil && !a.requests.Allow(p.KeyID, *lim.RPM, 1) {
|
||||
return res, &LimitError{Err: ErrRateLimited, RetryAfter: time.Second}
|
||||
}
|
||||
if lim.TPM != nil && !a.tokens.Allow(p.KeyID, *lim.TPM, estimatedTokens) {
|
||||
// The request was not charged against the request-rate limit in vain:
|
||||
// it was, and that is correct — it did arrive.
|
||||
return res, &LimitError{Err: ErrTokenLimited, RetryAfter: time.Second}
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// Record settles a finished request: it reconciles the token estimate, adds to
|
||||
// the budget, and queues the usage row.
|
||||
func (a *Accountant) Record(res Reservation, u store.Usage) {
|
||||
if res.KeyID == "" {
|
||||
return
|
||||
}
|
||||
actual := u.PromptTokens + u.CompletionTokens
|
||||
|
||||
// The estimate was charged up front; charge the difference now. A caller
|
||||
// whose requests are consistently larger than estimated is still bounded,
|
||||
// one window later than with a perfect estimate.
|
||||
a.tokens.Charge(res.KeyID, actual-res.EstimatedTokens)
|
||||
|
||||
u.CostMicros = res.Price.Cost(u.PromptTokens, u.CompletionTokens)
|
||||
a.budgets.Add(res.KeyID, res.Period, u.CostMicros)
|
||||
|
||||
if a.recorder != nil {
|
||||
a.recorder.Record(u)
|
||||
}
|
||||
}
|
||||
|
||||
// Status returns what a key has used, for headers and `ctl`.
|
||||
func (a *Accountant) Status(p auth.Principal) (spent int64, rpm, tpm float64) {
|
||||
if p.KeyID == "" {
|
||||
return 0, 0, 0
|
||||
}
|
||||
return a.budgets.Spent(p.KeyID, p.Limits.BudgetPeriod),
|
||||
a.requests.Rate(p.KeyID), a.tokens.Rate(p.KeyID)
|
||||
}
|
||||
@@ -0,0 +1,290 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func intp(v int) *int { return &v }
|
||||
func i64p(v int64) *int64 { return &v }
|
||||
|
||||
func principal(id string, tune func(*store.Effective)) auth.Principal {
|
||||
e := store.Effective{KeyID: id, ProjectID: "proj"}
|
||||
if tune != nil {
|
||||
tune(&e)
|
||||
}
|
||||
return auth.Principal{KeyID: id, ProjectID: "proj", Limits: e}
|
||||
}
|
||||
|
||||
func pricing(t *testing.T) *Pricing {
|
||||
t.Helper()
|
||||
return NewPricing(map[string]Price{
|
||||
"paid": {PromptPer1k: 1000, CompletionPer1k: 3000},
|
||||
"free": {},
|
||||
}, Price{})
|
||||
}
|
||||
|
||||
func TestCostRoundsUpSoSmallRequestsAreNotFree(t *testing.T) {
|
||||
// Rounding down would make a stream of tiny requests cost nothing, which is
|
||||
// exactly the traffic shape a quota exists to bound.
|
||||
p := Price{PromptPer1k: 1000, CompletionPer1k: 3000}
|
||||
if got := p.Cost(1, 0); got != 1 {
|
||||
t.Errorf("one prompt token cost %d", got)
|
||||
}
|
||||
if got := p.Cost(1000, 1000); got != 4000 {
|
||||
t.Errorf("a thousand of each cost %d, want 4000", got)
|
||||
}
|
||||
if got := p.Cost(0, 0); got != 0 {
|
||||
t.Errorf("an empty request cost %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestZeroPricedModelsStayAvailableWhenTheBudgetIsSpent(t *testing.T) {
|
||||
// This is what makes budget fallback work: a key with nothing left can
|
||||
// still be served something that costs nothing.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) {
|
||||
e.Budget = i64p(100)
|
||||
e.BudgetPeriod = store.PeriodMonth
|
||||
})
|
||||
|
||||
res, err := a.Check(p, "paid", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 0})
|
||||
|
||||
if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("a spent budget still served a paid model: %v", err)
|
||||
}
|
||||
if _, err := a.Check(p, "free", 10); err != nil {
|
||||
t.Errorf("a spent budget refused a free model: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetRefusalDoesNotSuggestRetrying(t *testing.T) {
|
||||
// A budget does not refill until the period rolls over, so a Retry-After of
|
||||
// a second would be a lie that turns one client into a retry loop.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) {
|
||||
e.Budget = i64p(1)
|
||||
e.BudgetPeriod = store.PeriodDay
|
||||
})
|
||||
res, _ := a.Check(p, "paid", 10)
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 5000})
|
||||
|
||||
_, err := a.Check(p, "paid", 10)
|
||||
var le *LimitError
|
||||
if !errors.As(err, &le) {
|
||||
t.Fatalf("expected a limit error, got %v", err)
|
||||
}
|
||||
if le.RetryAfter != 0 {
|
||||
t.Errorf("a spent budget suggested retrying in %v", le.RetryAfter)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitsBind(t *testing.T) {
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) { e.RPM = intp(3) })
|
||||
|
||||
for i := range 3 {
|
||||
if _, err := a.Check(p, "free", 1); err != nil {
|
||||
t.Fatalf("request %d was refused: %v", i, err)
|
||||
}
|
||||
}
|
||||
if _, err := a.Check(p, "free", 1); !errors.Is(err, ErrRateLimited) {
|
||||
t.Errorf("the fourth request was allowed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenLimitsUseTheEstimateAndThenSettle(t *testing.T) {
|
||||
// The real token count only exists once the request has run, so the check
|
||||
// has to work on an estimate. A caller that consistently underestimates
|
||||
// must still end up bounded.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", func(e *store.Effective) { e.TPM = intp(1000) })
|
||||
|
||||
// Claims 10 tokens, actually uses 900.
|
||||
res, err := a.Check(p, "free", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 400, CompletionTokens: 500})
|
||||
|
||||
// The settled figure is what the next check sees, so a second request of
|
||||
// the same shape cannot fit.
|
||||
if _, err := a.Check(p, "free", 200); !errors.Is(err, ErrTokenLimited) {
|
||||
t.Errorf("the underestimate was never settled: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoLimitIsNotALimitOfZero(t *testing.T) {
|
||||
// The distinction the store keeps with a nil pointer. Getting it wrong here
|
||||
// would silently forbid every request for every key with no limits set.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
p := principal("k1", nil)
|
||||
|
||||
for range 50 {
|
||||
if _, err := a.Check(p, "paid", 10_000); err != nil {
|
||||
t.Fatalf("a key with no limits was refused: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnUnauthenticatedRequestIsNotAccounted(t *testing.T) {
|
||||
// With auth off there is no key to charge, and Check must not refuse.
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
res, err := a.Check(auth.Principal{}, "paid", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Record(res, store.Usage{PromptTokens: 1000})
|
||||
if spent, _, _ := a.Status(auth.Principal{}); spent != 0 {
|
||||
t.Errorf("an anonymous request was charged %d", spent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetSurvivesARestart(t *testing.T) {
|
||||
// The counters live in memory; the store is what makes them survive.
|
||||
ctx := context.Background()
|
||||
path := filepath.Join(t.TempDir(), "acct.db")
|
||||
s, err := store.OpenSQLite(ctx, path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
must(t, s.PutOrg(ctx, store.Org{ID: "o", Name: "o"}))
|
||||
must(t, s.PutTeam(ctx, store.Team{ID: "t", OrgID: "o", Name: "t"}))
|
||||
must(t, s.PutProject(ctx, store.Project{ID: "proj", TeamID: "t", Name: "p"}))
|
||||
_, k, _ := store.NewKey("k1", "n", "proj")
|
||||
k.CreatedAt = time.Now()
|
||||
must(t, s.CreateKey(ctx, k))
|
||||
|
||||
must(t, s.WriteUsage(ctx, []store.Usage{
|
||||
{At: time.Now(), KeyID: "k1", ProjectID: "proj", Model: "paid",
|
||||
ServedModel: "paid", CostMicros: 5000, Status: 200},
|
||||
}))
|
||||
|
||||
a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
|
||||
must(t, a.RestoreBudgets(ctx, s, store.PeriodMonth))
|
||||
|
||||
p := principal("k1", func(e *store.Effective) {
|
||||
e.Budget = i64p(5000)
|
||||
e.BudgetPeriod = store.PeriodMonth
|
||||
})
|
||||
if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
|
||||
t.Errorf("the budget was reset by the restart: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetsResetAtThePeriodBoundary(t *testing.T) {
|
||||
b := NewBudgets()
|
||||
now := time.Date(2026, 1, 31, 23, 59, 0, 0, time.UTC)
|
||||
b.now = func() time.Time { return now }
|
||||
|
||||
b.Add("k1", store.PeriodMonth, 900)
|
||||
if got := b.Spent("k1", store.PeriodMonth); got != 900 {
|
||||
t.Fatalf("spent %d", got)
|
||||
}
|
||||
|
||||
now = now.Add(2 * time.Minute) // into February
|
||||
if got := b.Spent("k1", store.PeriodMonth); got != 0 {
|
||||
t.Errorf("the counter carried %d across the month boundary", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDailyAndMonthlyBoundariesDiffer(t *testing.T) {
|
||||
b := NewBudgets()
|
||||
now := time.Date(2026, 3, 15, 12, 0, 0, 0, time.UTC)
|
||||
b.now = func() time.Time { return now }
|
||||
|
||||
b.Add("daily", store.PeriodDay, 100)
|
||||
b.Add("monthly", store.PeriodMonth, 100)
|
||||
|
||||
now = now.Add(24 * time.Hour)
|
||||
if got := b.Spent("daily", store.PeriodDay); got != 0 {
|
||||
t.Errorf("the daily counter carried %d across a day", got)
|
||||
}
|
||||
if got := b.Spent("monthly", store.PeriodMonth); got != 100 {
|
||||
t.Errorf("the monthly counter lost %d within a month", 100-got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheWindowDoesNotAllowTwiceTheLimit(t *testing.T) {
|
||||
// The reason for a sliding window. With a fixed one, a caller sends its
|
||||
// whole allowance at the end of a window and again at the start of the
|
||||
// next, so the observed rate is double the limit exactly when the system is
|
||||
// least able to take it.
|
||||
l := NewLimiter()
|
||||
base := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC)
|
||||
now := base.Add(59 * time.Second)
|
||||
l.now = func() time.Time { return now }
|
||||
|
||||
const limit = 60
|
||||
sent := 0
|
||||
for range limit {
|
||||
if l.Allow("k", limit, 1) {
|
||||
sent++
|
||||
}
|
||||
}
|
||||
if sent != limit {
|
||||
t.Fatalf("the first burst only managed %d", sent)
|
||||
}
|
||||
|
||||
// One second later, in the next fixed window.
|
||||
now = base.Add(61 * time.Second)
|
||||
second := 0
|
||||
for range limit {
|
||||
if l.Allow("k", limit, 1) {
|
||||
second++
|
||||
}
|
||||
}
|
||||
// Some allowance returns as the old window slides out of view, but nowhere
|
||||
// near the whole of it.
|
||||
if second > limit/2 {
|
||||
t.Errorf("a second burst of %d passed immediately after %d; the window is not sliding",
|
||||
second, sent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterIsSafeUnderConcurrency(t *testing.T) {
|
||||
// Checking and charging have to be one operation, or two concurrent
|
||||
// requests both see room and both take it.
|
||||
l := NewLimiter()
|
||||
const limit = 100
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
allowed := 0
|
||||
for range 500 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if l.Allow("k", limit, 1) {
|
||||
mu.Lock()
|
||||
allowed++
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if allowed != limit {
|
||||
t.Errorf("%d requests passed a limit of %d", allowed, limit)
|
||||
}
|
||||
}
|
||||
|
||||
func must(t *testing.T, err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// PeriodStart is when the current budget period began.
|
||||
//
|
||||
// Periods are anchored to UTC rather than to a per-key start date, so every
|
||||
// key's day and month boundary is the same instant. A key whose period started
|
||||
// when it was created would make "spend this month" mean something different
|
||||
// for every key, and make the restore query unanswerable in one pass.
|
||||
func PeriodStart(p store.Period, now time.Time) time.Time {
|
||||
u := now.UTC()
|
||||
switch p {
|
||||
case store.PeriodDay:
|
||||
return time.Date(u.Year(), u.Month(), u.Day(), 0, 0, 0, 0, time.UTC)
|
||||
case store.PeriodMonth:
|
||||
return time.Date(u.Year(), u.Month(), 1, 0, 0, 0, 0, time.UTC)
|
||||
default:
|
||||
return time.Time{}
|
||||
}
|
||||
}
|
||||
|
||||
// Budgets tracks spend per key in memory, so a budget check costs no query.
|
||||
//
|
||||
// The database catches up through the batch writer. That means the in-memory
|
||||
// figure leads the stored one, which is the right direction: a budget that only
|
||||
// counted what had already been written would let a burst overspend by however
|
||||
// much the writer was behind.
|
||||
type Budgets struct {
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
spend map[string]counter
|
||||
}
|
||||
|
||||
type counter struct {
|
||||
micros int64
|
||||
since time.Time
|
||||
}
|
||||
|
||||
// NewBudgets builds an empty tracker.
|
||||
func NewBudgets() *Budgets {
|
||||
return &Budgets{now: time.Now, spend: map[string]counter{}}
|
||||
}
|
||||
|
||||
// Restore seeds the counters from what the store already recorded.
|
||||
//
|
||||
// A restart must not reset a budget, which it would if the counter lived only
|
||||
// in memory. The caller passes the spend since each period's start.
|
||||
func (b *Budgets) Restore(period store.Period, spend map[string]int64) {
|
||||
start := PeriodStart(period, b.now())
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
for key, micros := range spend {
|
||||
b.spend[key] = counter{micros: micros, since: start}
|
||||
}
|
||||
}
|
||||
|
||||
// Spent returns a key's spend in the current period, resetting it if the period
|
||||
// has rolled over.
|
||||
func (b *Budgets) Spent(keyID string, period store.Period) int64 {
|
||||
start := PeriodStart(period, b.now())
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
return b.currentLocked(keyID, start).micros
|
||||
}
|
||||
|
||||
// Add records spend.
|
||||
func (b *Budgets) Add(keyID string, period store.Period, micros int64) {
|
||||
if micros == 0 {
|
||||
return
|
||||
}
|
||||
start := PeriodStart(period, b.now())
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
c := b.currentLocked(keyID, start)
|
||||
c.micros += micros
|
||||
b.spend[keyID] = c
|
||||
}
|
||||
|
||||
// Exhausted reports whether a key has spent its budget.
|
||||
func (b *Budgets) Exhausted(keyID string, period store.Period, budget int64) bool {
|
||||
return b.Spent(keyID, period) >= budget
|
||||
}
|
||||
|
||||
// currentLocked returns the counter, rolled over if the period changed.
|
||||
func (b *Budgets) currentLocked(keyID string, start time.Time) counter {
|
||||
c, ok := b.spend[keyID]
|
||||
if !ok || c.since.Before(start) {
|
||||
// Either never seen, or the period rolled over while the process was
|
||||
// running. Both start from zero.
|
||||
c = counter{since: start}
|
||||
b.spend[keyID] = c
|
||||
}
|
||||
return c
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Package accounting turns finished requests into spend, enforces what a key
|
||||
// is allowed, and gets both into the store without touching the request path
|
||||
// (plan §7.9.3, §7.9.4).
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Price is what one model costs, in micro-units per thousand tokens.
|
||||
//
|
||||
// Micro-units rather than a float because spend is accumulated by repeated
|
||||
// addition, and repeated addition of floats drifts. What a unit means is the
|
||||
// operator's: self-hosted models have no per-token price, so it may be
|
||||
// GPU-seconds, a notional rate, or 1000 everywhere — which turns a budget into
|
||||
// a token quota.
|
||||
type Price struct {
|
||||
PromptPer1k int64 `koanf:"prompt_per_1k"`
|
||||
CompletionPer1k int64 `koanf:"completion_per_1k"`
|
||||
}
|
||||
|
||||
// Zero reports whether this model spends nothing.
|
||||
//
|
||||
// It matters beyond arithmetic: a zero-cost model stays available to a key
|
||||
// whose budget is spent, which is what makes the plan's budget fallback work
|
||||
// (§7.7).
|
||||
func (p Price) Zero() bool { return p.PromptPer1k == 0 && p.CompletionPer1k == 0 }
|
||||
|
||||
// Cost computes what a request spent, rounding up.
|
||||
//
|
||||
// Up rather than to nearest: rounding down would make a stream of tiny requests
|
||||
// free, which is exactly the shape of traffic a quota exists to bound.
|
||||
func (p Price) Cost(promptTokens, completionTokens int64) int64 {
|
||||
return divUp(promptTokens*p.PromptPer1k, 1000) + divUp(completionTokens*p.CompletionPer1k, 1000)
|
||||
}
|
||||
|
||||
func divUp(n, d int64) int64 {
|
||||
if n <= 0 {
|
||||
return 0
|
||||
}
|
||||
return (n + d - 1) / d
|
||||
}
|
||||
|
||||
// Pricing resolves a model name to its price.
|
||||
type Pricing struct {
|
||||
mu sync.RWMutex
|
||||
byModel map[string]Price
|
||||
fallback Price
|
||||
}
|
||||
|
||||
// NewPricing builds a price table. Names are matched case-insensitively,
|
||||
// because a model name reaches us from whatever the client typed.
|
||||
func NewPricing(byModel map[string]Price, fallback Price) *Pricing {
|
||||
table := make(map[string]Price, len(byModel))
|
||||
for name, p := range byModel {
|
||||
table[strings.ToLower(name)] = p
|
||||
}
|
||||
return &Pricing{byModel: table, fallback: fallback}
|
||||
}
|
||||
|
||||
// For returns a model's price, or the fallback.
|
||||
func (p *Pricing) For(model string) Price {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
if v, ok := p.byModel[strings.ToLower(model)]; ok {
|
||||
return v
|
||||
}
|
||||
return p.fallback
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Limiter enforces per-key rate limits.
|
||||
//
|
||||
// A sliding window rather than a fixed one. A fixed window lets a caller send
|
||||
// its whole minute's allowance in the last second of one window and again in
|
||||
// the first second of the next, so the observed rate is twice the configured
|
||||
// limit exactly when the system is least able to take it.
|
||||
//
|
||||
// The window is approximated from two buckets, which is the usual compromise:
|
||||
// exact enough that the doubling cannot happen, and bounded memory per key
|
||||
// rather than a timestamp per request.
|
||||
type Limiter struct {
|
||||
now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
keys map[string]*window
|
||||
}
|
||||
|
||||
type window struct {
|
||||
// current and previous are counts in this minute and the one before.
|
||||
current, previous int64
|
||||
start time.Time
|
||||
}
|
||||
|
||||
// NewLimiter builds an empty limiter.
|
||||
func NewLimiter() *Limiter {
|
||||
return &Limiter{now: time.Now, keys: map[string]*window{}}
|
||||
}
|
||||
|
||||
const windowLen = time.Minute
|
||||
|
||||
// Allow charges n against a key's limit and reports whether it fits.
|
||||
//
|
||||
// Charging and checking are one operation on purpose: doing them separately
|
||||
// lets two concurrent requests both see room and both take it.
|
||||
func (l *Limiter) Allow(key string, limit int, n int64) bool {
|
||||
if limit <= 0 {
|
||||
// No limit configured. Zero would be a limit of zero, which is why
|
||||
// callers pass a nil pointer rather than 0 for "unset".
|
||||
return true
|
||||
}
|
||||
now := l.now()
|
||||
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
w := l.windowLocked(key, now)
|
||||
if w.rate(now)+float64(n) > float64(limit) {
|
||||
return false
|
||||
}
|
||||
w.current += n
|
||||
return true
|
||||
}
|
||||
|
||||
// Charge records usage that was already allowed, for reconciliation.
|
||||
//
|
||||
// Token limits are checked against an estimate before the request runs, because
|
||||
// the real figure only exists once it has. The difference is charged here, so a
|
||||
// caller whose requests are consistently larger than estimated is still bounded
|
||||
// — one window later than it would be with a perfect estimate.
|
||||
func (l *Limiter) Charge(key string, n int64) {
|
||||
if n == 0 {
|
||||
return
|
||||
}
|
||||
now := l.now()
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
w := l.windowLocked(key, now)
|
||||
w.current += n
|
||||
if w.current < 0 {
|
||||
w.current = 0
|
||||
}
|
||||
}
|
||||
|
||||
// Rate returns a key's current observed rate.
|
||||
func (l *Limiter) Rate(key string) float64 {
|
||||
now := l.now()
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.windowLocked(key, now).rate(now)
|
||||
}
|
||||
|
||||
func (l *Limiter) windowLocked(key string, now time.Time) *window {
|
||||
w, ok := l.keys[key]
|
||||
if !ok {
|
||||
w = &window{start: now.Truncate(windowLen)}
|
||||
l.keys[key] = w
|
||||
}
|
||||
w.roll(now)
|
||||
return w
|
||||
}
|
||||
|
||||
// roll advances the window to now.
|
||||
func (w *window) roll(now time.Time) {
|
||||
start := now.Truncate(windowLen)
|
||||
switch {
|
||||
case start.Equal(w.start):
|
||||
return
|
||||
case start.Sub(w.start) == windowLen:
|
||||
w.previous, w.current = w.current, 0
|
||||
default:
|
||||
// More than one window has passed with no traffic.
|
||||
w.previous, w.current = 0, 0
|
||||
}
|
||||
w.start = start
|
||||
}
|
||||
|
||||
// rate weights the previous window by how much of it is still in view.
|
||||
func (w *window) rate(now time.Time) float64 {
|
||||
elapsed := now.Sub(w.start)
|
||||
if elapsed >= windowLen {
|
||||
return float64(w.current)
|
||||
}
|
||||
remaining := 1 - float64(elapsed)/float64(windowLen)
|
||||
return float64(w.previous)*remaining + float64(w.current)
|
||||
}
|
||||
|
||||
// Forget drops a key's window. Used when a key is revoked; keeping it would
|
||||
// leak memory for every key that ever called.
|
||||
func (l *Limiter) Forget(key string) {
|
||||
l.mu.Lock()
|
||||
delete(l.keys, key)
|
||||
l.mu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// RecorderConfig tunes the batch writer.
|
||||
type RecorderConfig struct {
|
||||
// FlushInterval and BatchSize decide when a batch is written. Whichever
|
||||
// comes first.
|
||||
FlushInterval time.Duration
|
||||
BatchSize int
|
||||
// Capacity bounds the buffer. Beyond it, records are dropped rather than
|
||||
// blocking a request — see Record.
|
||||
Capacity int
|
||||
}
|
||||
|
||||
// DefaultRecorderConfig returns sensible values.
|
||||
func DefaultRecorderConfig() RecorderConfig {
|
||||
return RecorderConfig{FlushInterval: 200 * time.Millisecond, BatchSize: 500, Capacity: 8192}
|
||||
}
|
||||
|
||||
func (c *RecorderConfig) normalize() {
|
||||
d := DefaultRecorderConfig()
|
||||
if c.FlushInterval <= 0 {
|
||||
c.FlushInterval = d.FlushInterval
|
||||
}
|
||||
if c.BatchSize <= 0 {
|
||||
c.BatchSize = d.BatchSize
|
||||
}
|
||||
if c.Capacity <= 0 {
|
||||
c.Capacity = d.Capacity
|
||||
}
|
||||
}
|
||||
|
||||
// Toucher is the part of the authenticator the recorder flushes alongside
|
||||
// usage, so last-used timestamps ride the same transaction schedule.
|
||||
type Toucher interface {
|
||||
FlushTouches(ctx context.Context) error
|
||||
}
|
||||
|
||||
// Recorder batches usage records into the store.
|
||||
//
|
||||
// The buffered channel is the whole design. A request must never wait on a
|
||||
// database write, so Record is a non-blocking send and the writer drains it on
|
||||
// its own schedule.
|
||||
type Recorder struct {
|
||||
cfg RecorderConfig
|
||||
store store.Store
|
||||
touch Toucher
|
||||
log *slog.Logger
|
||||
|
||||
ch chan store.Usage
|
||||
done chan struct{}
|
||||
wg sync.WaitGroup
|
||||
|
||||
mu sync.Mutex
|
||||
dropped int64
|
||||
written int64
|
||||
}
|
||||
|
||||
// NewRecorder builds a recorder. Start begins writing.
|
||||
func NewRecorder(s store.Store, touch Toucher, cfg RecorderConfig, log *slog.Logger) *Recorder {
|
||||
cfg.normalize()
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Recorder{
|
||||
cfg: cfg, store: s, touch: touch, log: log,
|
||||
ch: make(chan store.Usage, cfg.Capacity),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Record queues one request's accounting.
|
||||
//
|
||||
// Never blocks. If the buffer is full — a database stalled long enough to fill
|
||||
// eight thousand records — the record is dropped and counted. Losing some
|
||||
// accounting is bad; making every request wait on a stalled database is worse,
|
||||
// because it converts a reporting problem into an outage.
|
||||
func (r *Recorder) Record(u store.Usage) {
|
||||
select {
|
||||
case r.ch <- u:
|
||||
default:
|
||||
r.mu.Lock()
|
||||
r.dropped++
|
||||
n := r.dropped
|
||||
r.mu.Unlock()
|
||||
if n == 1 || n%1000 == 0 {
|
||||
r.log.Warn("dropping usage records: the store is not keeping up",
|
||||
"dropped_total", n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Start runs the writer until ctx ends.
|
||||
func (r *Recorder) Start(ctx context.Context) {
|
||||
r.wg.Add(1)
|
||||
go func() {
|
||||
defer r.wg.Done()
|
||||
defer close(r.done)
|
||||
r.run(ctx)
|
||||
}()
|
||||
}
|
||||
|
||||
// Stop waits for the writer to drain and finish.
|
||||
func (r *Recorder) Stop() { r.wg.Wait() }
|
||||
|
||||
func (r *Recorder) run(ctx context.Context) {
|
||||
t := time.NewTicker(r.cfg.FlushInterval)
|
||||
defer t.Stop()
|
||||
|
||||
batch := make([]store.Usage, 0, r.cfg.BatchSize)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// Drain what is already queued before leaving. These are completed
|
||||
// requests that have already been served; throwing them away would
|
||||
// mean a clean shutdown silently loses the last moments of billing.
|
||||
r.drain(&batch)
|
||||
r.flush(context.WithoutCancel(ctx), batch)
|
||||
return
|
||||
|
||||
case u := <-r.ch:
|
||||
batch = append(batch, u)
|
||||
if len(batch) >= r.cfg.BatchSize {
|
||||
r.flush(ctx, batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
|
||||
case <-t.C:
|
||||
r.drain(&batch)
|
||||
if len(batch) > 0 {
|
||||
r.flush(ctx, batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
if err := r.touchFlush(ctx); err != nil {
|
||||
r.log.Warn("failed to write last-used timestamps", "err", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// drain moves everything queued into the batch without blocking.
|
||||
func (r *Recorder) drain(batch *[]store.Usage) {
|
||||
for len(*batch) < r.cfg.BatchSize {
|
||||
select {
|
||||
case u := <-r.ch:
|
||||
*batch = append(*batch, u)
|
||||
default:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) flush(ctx context.Context, batch []store.Usage) {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
if err := r.store.WriteUsage(ctx, batch); err != nil {
|
||||
// The records are gone either way; what matters is that this is loud.
|
||||
// Silently losing billing data is the kind of thing discovered at the
|
||||
// end of a month.
|
||||
r.log.Error("failed to write usage records", "count", len(batch), "err", err)
|
||||
return
|
||||
}
|
||||
r.mu.Lock()
|
||||
r.written += int64(len(batch))
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *Recorder) touchFlush(ctx context.Context) error {
|
||||
if r.touch == nil {
|
||||
return nil
|
||||
}
|
||||
return r.touch.FlushTouches(ctx)
|
||||
}
|
||||
|
||||
// Stats reports what the recorder has done, for metrics and tests.
|
||||
func (r *Recorder) Stats() (written, dropped int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.written, r.dropped
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
package accounting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// fakeStore records what was written and can be made slow or broken.
|
||||
type fakeStore struct {
|
||||
store.Store
|
||||
|
||||
mu sync.Mutex
|
||||
batches [][]store.Usage
|
||||
|
||||
block chan struct{}
|
||||
fail atomic.Bool
|
||||
touched atomic.Int64
|
||||
}
|
||||
|
||||
func (f *fakeStore) WriteUsage(_ context.Context, rows []store.Usage) error {
|
||||
if f.block != nil {
|
||||
<-f.block
|
||||
}
|
||||
if f.fail.Load() {
|
||||
return errors.New("store is down")
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.batches = append(f.batches, append([]store.Usage(nil), rows...))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeStore) rows() []store.Usage {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
var out []store.Usage
|
||||
for _, b := range f.batches {
|
||||
out = append(out, b...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (f *fakeStore) batchCount() int {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return len(f.batches)
|
||||
}
|
||||
|
||||
type fakeToucher struct{ n *atomic.Int64 }
|
||||
|
||||
func (f fakeToucher) FlushTouches(context.Context) error {
|
||||
f.n.Add(1)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestRecordsAreBatchedNotWrittenOneByOne(t *testing.T) {
|
||||
// The whole point of the buffer: a hundred requests must not be a hundred
|
||||
// transactions.
|
||||
f := &fakeStore{}
|
||||
r := NewRecorder(f, nil, RecorderConfig{FlushInterval: 20 * time.Millisecond, BatchSize: 500},
|
||||
slog.New(slog.DiscardHandler))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
|
||||
for i := range 100 {
|
||||
r.Record(store.Usage{KeyID: "k", PromptTokens: int64(i)})
|
||||
}
|
||||
waitFor(t, time.Second, func() bool { return len(f.rows()) == 100 })
|
||||
|
||||
if n := f.batchCount(); n > 5 {
|
||||
t.Errorf("100 records took %d transactions", n)
|
||||
}
|
||||
cancel()
|
||||
r.Stop()
|
||||
}
|
||||
|
||||
func TestRecordNeverBlocksARequest(t *testing.T) {
|
||||
// A stalled database must not become an outage. Losing accounting is bad;
|
||||
// making every request wait on the store is worse.
|
||||
f := &fakeStore{block: make(chan struct{})}
|
||||
r := NewRecorder(f, nil, RecorderConfig{
|
||||
FlushInterval: time.Millisecond, BatchSize: 1, Capacity: 8,
|
||||
}, slog.New(slog.DiscardHandler))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
defer func() { close(f.block); cancel(); r.Stop() }()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
for range 1000 {
|
||||
r.Record(store.Usage{KeyID: "k"})
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Record blocked while the store was stalled")
|
||||
}
|
||||
|
||||
_, dropped := r.Stats()
|
||||
if dropped == 0 {
|
||||
t.Error("the overflow was not counted, so nobody would ever know")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAShutdownDoesNotLoseWhatIsAlreadyQueued(t *testing.T) {
|
||||
// Those are completed requests that were already served. Throwing them away
|
||||
// means a clean shutdown silently loses the last moments of billing.
|
||||
f := &fakeStore{}
|
||||
r := NewRecorder(f, nil, RecorderConfig{FlushInterval: time.Hour, BatchSize: 10_000},
|
||||
slog.New(slog.DiscardHandler))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
|
||||
for range 50 {
|
||||
r.Record(store.Usage{KeyID: "k"})
|
||||
}
|
||||
cancel()
|
||||
r.Stop()
|
||||
|
||||
if got := len(f.rows()); got != 50 {
|
||||
t.Errorf("a clean shutdown wrote %d of 50 queued records", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFailedWriteIsLoudRatherThanSilent(t *testing.T) {
|
||||
// Silently losing billing data is the kind of thing discovered at the end
|
||||
// of a month.
|
||||
f := &fakeStore{}
|
||||
f.fail.Store(true)
|
||||
logged := &countingHandler{}
|
||||
r := NewRecorder(f, nil, RecorderConfig{FlushInterval: 10 * time.Millisecond, BatchSize: 1},
|
||||
slog.New(logged))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
r.Record(store.Usage{KeyID: "k"})
|
||||
|
||||
waitFor(t, time.Second, func() bool { return logged.count(slog.LevelError) > 0 })
|
||||
cancel()
|
||||
r.Stop()
|
||||
|
||||
if written, _ := r.Stats(); written != 0 {
|
||||
t.Errorf("a failed write was counted as %d written", written)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLastUsedRidesTheSameSchedule(t *testing.T) {
|
||||
var n atomic.Int64
|
||||
f := &fakeStore{}
|
||||
r := NewRecorder(f, fakeToucher{n: &n}, RecorderConfig{FlushInterval: 10 * time.Millisecond},
|
||||
slog.New(slog.DiscardHandler))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r.Start(ctx)
|
||||
waitFor(t, time.Second, func() bool { return n.Load() > 0 })
|
||||
cancel()
|
||||
r.Stop()
|
||||
}
|
||||
|
||||
type countingHandler struct {
|
||||
mu sync.Mutex
|
||||
levels []slog.Level
|
||||
}
|
||||
|
||||
func (c *countingHandler) Enabled(context.Context, slog.Level) bool { return true }
|
||||
func (c *countingHandler) WithAttrs([]slog.Attr) slog.Handler { return c }
|
||||
func (c *countingHandler) WithGroup(string) slog.Handler { return c }
|
||||
func (c *countingHandler) Handle(_ context.Context, r slog.Record) error {
|
||||
c.mu.Lock()
|
||||
c.levels = append(c.levels, r.Level)
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *countingHandler) count(l slog.Level) int {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
n := 0
|
||||
for _, got := range c.levels {
|
||||
if got == l {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func waitFor(t *testing.T, limit time.Duration, pred func() bool) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(limit)
|
||||
for time.Now().Before(deadline) {
|
||||
if pred() {
|
||||
return
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("condition never became true")
|
||||
}
|
||||
+360
-19
@@ -3,7 +3,9 @@ package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"crypto/tls"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
@@ -11,6 +13,9 @@ import (
|
||||
"math/rand/v2"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -21,8 +26,10 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
// Config configures the agent.
|
||||
@@ -45,6 +52,36 @@ type Config struct {
|
||||
ReconnectBase time.Duration `koanf:"reconnect_base"`
|
||||
ReconnectMax time.Duration `koanf:"reconnect_max"`
|
||||
|
||||
// AdoptReadyTimeout bounds how long an adopted engine is given to answer a
|
||||
// health check before the agent stops waiting on it.
|
||||
//
|
||||
// Shorter than a startup timeout on purpose: this engine is already
|
||||
// running, so it is not loading a model, it is either serving or it is not.
|
||||
AdoptReadyTimeout time.Duration `koanf:"adopt_ready_timeout"`
|
||||
|
||||
// TLS authenticates both ends of the control-plane stream. Leaving it unset
|
||||
// keeps the plaintext behaviour, which only belongs on a trusted network.
|
||||
TLS tlsconf.Config `koanf:"tls"`
|
||||
|
||||
// ModelDirs are the directories this node keeps model weights in.
|
||||
//
|
||||
// Empty by default, and empty means the node is not asked to look — which
|
||||
// is not the same as having nothing. Storage differs at every installation,
|
||||
// so llmgw observes it rather than owning it: a Lustre mount, an image
|
||||
// exported over the fabric, local NVMe and a copy made by hand all look
|
||||
// identical from here, and placement only needs to know whether the file
|
||||
// is present.
|
||||
ModelDirs []string `koanf:"model_dirs"`
|
||||
|
||||
// WeightsCache is where a remote source is downloaded to. Defaults to the
|
||||
// first entry of ModelDirs, because a node that says where its models live
|
||||
// has already answered this question.
|
||||
//
|
||||
// Fetching is the simple arrangement: every node keeps its own copy. The
|
||||
// exotic ones need no code — weights on a shared filer or an image exported
|
||||
// over the fabric are simply present, and present means nothing is fetched.
|
||||
WeightsCache string `koanf:"weights_cache"`
|
||||
|
||||
Discovery discovery.Options `koanf:"discovery"`
|
||||
Engines engines.Config `koanf:"engines"`
|
||||
Superv supervisor.Config `koanf:"supervisor"`
|
||||
@@ -59,6 +96,7 @@ func DefaultConfig() Config {
|
||||
HeartbeatInterval: 5 * time.Second,
|
||||
ReconnectBase: time.Second,
|
||||
ReconnectMax: 30 * time.Second,
|
||||
AdoptReadyTimeout: 60 * time.Second,
|
||||
Superv: supervisor.DefaultConfig(),
|
||||
// NVML gives NVLink topology, so prefer it when this build has it.
|
||||
Discovery: discovery.Options{PreferNVML: true},
|
||||
@@ -73,6 +111,12 @@ func (c *Config) normalize() error {
|
||||
if c.Server == "" {
|
||||
return errors.New("server is required")
|
||||
}
|
||||
if err := c.TLS.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := c.normalizeServerURL(); err != nil {
|
||||
return err
|
||||
}
|
||||
if c.MgmtInterface == "" {
|
||||
c.MgmtInterface = d.MgmtInterface
|
||||
}
|
||||
@@ -82,12 +126,40 @@ func (c *Config) normalize() error {
|
||||
if c.ReconnectBase <= 0 {
|
||||
c.ReconnectBase = d.ReconnectBase
|
||||
}
|
||||
if c.AdoptReadyTimeout <= 0 {
|
||||
c.AdoptReadyTimeout = d.AdoptReadyTimeout
|
||||
}
|
||||
if c.ReconnectMax < c.ReconnectBase {
|
||||
c.ReconnectMax = max(c.ReconnectBase, d.ReconnectMax)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// normalizeServerURL makes the configured address a URL connect can use, and
|
||||
// refuses the combination that would quietly undo the certificates: a scheme
|
||||
// that says plaintext while TLS is configured. Silently connecting anyway would
|
||||
// mean a cluster believing it has mutual authentication when it has none.
|
||||
func (c *Config) normalizeServerURL() error {
|
||||
secure := c.TLS.Enabled()
|
||||
switch {
|
||||
case strings.HasPrefix(c.Server, "https://"):
|
||||
if !secure {
|
||||
return errors.New("server is an https:// address but no tls certificates are configured")
|
||||
}
|
||||
case strings.HasPrefix(c.Server, "http://"):
|
||||
if secure {
|
||||
return errors.New("tls is configured but server is an http:// address; use https://")
|
||||
}
|
||||
case strings.Contains(c.Server, "://"):
|
||||
return fmt.Errorf("server %q must be an http:// or https:// address", c.Server)
|
||||
case secure:
|
||||
c.Server = "https://" + c.Server
|
||||
default:
|
||||
c.Server = "http://" + c.Server
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Agent runs on a GPU node.
|
||||
type Agent struct {
|
||||
cfg Config
|
||||
@@ -101,6 +173,12 @@ type Agent struct {
|
||||
// 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
|
||||
fetcher *weights.Fetcher
|
||||
|
||||
// lastWeights is the fingerprint of the file list most recently sent, so a
|
||||
// heartbeat carries the list only when it changed. Touched only by the
|
||||
// heartbeat goroutine, which is the only thing that reads or writes it.
|
||||
lastWeights string
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
@@ -225,11 +303,115 @@ func (a *Agent) recoverOrphans() {
|
||||
a.mu.Unlock()
|
||||
a.log.Info("adopted a running engine",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID, "pid", o.PID)
|
||||
_ = sh
|
||||
go a.settleAdopted(sh, o)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// settleAdopted brings an adopted engine to Ready, or lets it fail.
|
||||
//
|
||||
// Adoption stopped at "the process is alive", which left the shard in Starting
|
||||
// with nothing that would ever move it. The control plane then waited out the
|
||||
// profile's startup timeout and killed a perfectly healthy engine so it could
|
||||
// start an identical one — so an agent restart, which exists precisely so that
|
||||
// engines are not interrupted, cost every model on the node a full reload.
|
||||
//
|
||||
// Measured on node05: the engine answered /health the whole time, and the
|
||||
// replica sat in Starting for three minutes until the timeout reaped it.
|
||||
//
|
||||
// Run in the background because recovery happens before the agent connects, and
|
||||
// a slow health check must not delay registration — the control plane needs to
|
||||
// hear about these shards promptly whatever state they are in.
|
||||
func (a *Agent) settleAdopted(sh *supervisor.Shard, o supervisor.Orphan) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), a.cfg.AdoptReadyTimeout)
|
||||
defer cancel()
|
||||
|
||||
eng, err := a.reg.Get(o.Spec.Profile.Engine)
|
||||
if err != nil {
|
||||
a.log.Warn("cannot judge an adopted engine; leaving it to the control plane",
|
||||
"shard", o.ShardID, "err", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 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)
|
||||
err = a.sup.WaitReady(ctx, sh, base, o.Spec.Profile.ServedModelName, a.cfg.AdoptReadyTimeout)
|
||||
} else {
|
||||
err = a.sup.WaitAlive(ctx, sh)
|
||||
}
|
||||
if err != nil {
|
||||
// Not killed here. An adopted engine that will not answer is the
|
||||
// control plane's decision to replace, and it has the timeout for it;
|
||||
// killing it locally would race that.
|
||||
a.log.Warn("an adopted engine did not become ready",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID, "err", err)
|
||||
return
|
||||
}
|
||||
a.log.Info("adopted engine is serving",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID)
|
||||
}
|
||||
|
||||
// usableFabrics reports the fabrics this node can actually carry traffic on.
|
||||
//
|
||||
// Loud about what it drops. A fabric silently missing from a registration looks
|
||||
// identical to one never configured, and the operator who typed the adapter
|
||||
// name is the only person who can fix it — so the reason is logged here, on the
|
||||
// machine where both halves are visible, rather than inferred later from a
|
||||
// replica that will not converge.
|
||||
func (a *Agent) usableFabrics(inv discovery.Inventory) []inventory.FabricAttachment {
|
||||
declared := map[string][]string{}
|
||||
for _, f := range a.cfg.Fabrics {
|
||||
declared[f.Fabric] = f.HCAs
|
||||
}
|
||||
|
||||
var out []inventory.FabricAttachment
|
||||
for _, c := range discovery.CheckFabrics(declared, inv.HCAs) {
|
||||
for _, p := range c.Problems {
|
||||
a.log.Error("a declared fabric adapter cannot carry traffic",
|
||||
"fabric", c.Fabric, "problem", p)
|
||||
}
|
||||
if !c.OK() {
|
||||
a.log.Error("fabric not advertised; no declared adapter on it can carry traffic",
|
||||
"fabric", c.Fabric,
|
||||
"consequence", "no multi-node replica will be placed here")
|
||||
continue
|
||||
}
|
||||
out = append(out, inventory.FabricAttachment{
|
||||
Fabric: c.Fabric, HCAs: c.Usable, RateGbps: c.RateGbps,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// weights builds the fetcher lazily, so an agent that never starts a shard
|
||||
// never needs a cache directory.
|
||||
func (a *Agent) weights() *weights.Fetcher {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if a.fetcher != nil {
|
||||
return a.fetcher
|
||||
}
|
||||
cache := a.cfg.WeightsCache
|
||||
if cache == "" && len(a.cfg.ModelDirs) > 0 {
|
||||
cache = a.cfg.ModelDirs[0]
|
||||
}
|
||||
a.fetcher = &weights.Fetcher{
|
||||
CacheDir: cache,
|
||||
Log: a.log,
|
||||
OnProgress: func(p weights.Progress) {
|
||||
// Logged rather than silent: a replica that is downloading a
|
||||
// hundred gigabytes and one that is wedged look identical from the
|
||||
// control plane, and the difference matters to whoever is waiting.
|
||||
a.log.Info("fetching weights",
|
||||
"source", p.Source, "percent", p.Percent(),
|
||||
"bytes", p.BytesDone, "total", p.BytesTotal)
|
||||
},
|
||||
}
|
||||
return a.fetcher
|
||||
}
|
||||
|
||||
// runningShards describes what this agent is supervising, for registration.
|
||||
func (a *Agent) runningShards() []*llmgwv1.RunningShard {
|
||||
var out []*llmgwv1.RunningShard
|
||||
@@ -275,7 +457,11 @@ func int32Slice(v []int) []int32 {
|
||||
|
||||
// scan refreshes the node's hardware inventory.
|
||||
func (a *Agent) scan(ctx context.Context) error {
|
||||
inv, err := discovery.Discover(ctx, a.cfg.Discovery, a.log)
|
||||
// Carried from the agent's own configuration rather than duplicated into
|
||||
// the discovery block, so an operator names the directories once.
|
||||
opts := a.cfg.Discovery
|
||||
opts.ModelDirs = a.cfg.ModelDirs
|
||||
inv, err := discovery.Discover(ctx, opts, a.log)
|
||||
a.mu.Lock()
|
||||
a.inv = inv
|
||||
a.mu.Unlock()
|
||||
@@ -286,6 +472,10 @@ func (a *Agent) scan(ctx context.Context) error {
|
||||
func (a *Agent) session(ctx context.Context) error {
|
||||
client := llmgwv1connect.NewAgentServiceClient(a.http, a.cfg.Server)
|
||||
|
||||
// parent is kept because the derived context is cancelled from inside this
|
||||
// function; asking the derived one whether we are shutting down would
|
||||
// always say yes, and every stream error would read as a clean close.
|
||||
parent := ctx
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
@@ -296,9 +486,14 @@ func (a *Agent) session(ctx context.Context) error {
|
||||
inv := a.inv
|
||||
a.mu.Unlock()
|
||||
|
||||
// Declared fabric membership against what the node actually has. A fabric
|
||||
// whose adapters are absent or have no active port is not advertised at
|
||||
// all, so the control plane will not place a multi-node replica on it —
|
||||
// rather than placing one that starts, passes health checks and hangs on
|
||||
// its first all-reduce.
|
||||
reg := agentsvc.RegisterFromInventory(
|
||||
a.cfg.NodeID, a.cfg.Hostname, a.cfg.MgmtAddr, a.cfg.Version,
|
||||
a.cfg.Labels, inv, a.cfg.Fabrics)
|
||||
a.cfg.Labels, inv, a.usableFabrics(inv))
|
||||
// Every registration reports what is already running, not just the first
|
||||
// one: a reconnect may be a control plane restart, and it has to recover
|
||||
// the same way (plan §6.5).
|
||||
@@ -310,6 +505,22 @@ func (a *Agent) session(ctx context.Context) error {
|
||||
return fmt.Errorf("registering: %w", err)
|
||||
}
|
||||
|
||||
// Closing the response is what actually unblocks Receive below. Cancelling
|
||||
// the context is not enough: HTTP/2 watches it while the request body is
|
||||
// being written, and a bidirectional stream spends its life blocked on the
|
||||
// read side instead. Without this an agent told to stop keeps reading until
|
||||
// something kills it — which for systemd means the full stop timeout on
|
||||
// every restart.
|
||||
watchDone := make(chan struct{})
|
||||
defer close(watchDone)
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
_ = stream.CloseResponse()
|
||||
case <-watchDone:
|
||||
}
|
||||
}()
|
||||
|
||||
// Sends are serialized through this channel: connect streams are not safe
|
||||
// for concurrent writes, and heartbeats race with status reports.
|
||||
out := make(chan *llmgwv1.AgentMessage, 32)
|
||||
@@ -325,9 +536,18 @@ func (a *Agent) session(ctx context.Context) error {
|
||||
msg, err := stream.Receive()
|
||||
if err != nil {
|
||||
cancel()
|
||||
if errors.Is(err, io.EOF) || ctx.Err() != nil {
|
||||
switch {
|
||||
case parent.Err() != nil:
|
||||
// The agent is stopping; the stream ending is the consequence.
|
||||
return nil
|
||||
case errors.Is(err, io.EOF):
|
||||
// The control plane closed the stream, which a restart does.
|
||||
// Reconnect immediately rather than backing off.
|
||||
return nil
|
||||
}
|
||||
// Everything else is worth both reporting and backing off from. A
|
||||
// refused certificate arrives here, and an agent that cannot say
|
||||
// why it is not registering is one an operator cannot fix.
|
||||
return err
|
||||
}
|
||||
a.handle(ctx, msg, out)
|
||||
@@ -367,9 +587,103 @@ func (a *Agent) heartbeat() *llmgwv1.AgentMessage {
|
||||
for _, g := range inv.GPUs {
|
||||
hb.GpuMemoryUsedMb = append(hb.GpuMemoryUsedMb, int32(g.MemoryUsedMB))
|
||||
}
|
||||
|
||||
// Rescanned here rather than reused from registration. Disk is not a fact
|
||||
// about a machine, it is a fact about a moment: a node that had 200 GB when
|
||||
// its agent started can be full an hour later, and it is precisely the node
|
||||
// busy filling up whose figure matters.
|
||||
for _, st := range discovery.ScanStorage(a.cfg.ModelDirs) {
|
||||
hb.Storage = append(hb.Storage, &llmgwv1.Storage{
|
||||
Path: st.Path, FreeBytes: st.FreeBytes,
|
||||
TotalBytes: st.TotalBytes, Shared: st.Shared,
|
||||
})
|
||||
}
|
||||
|
||||
// The file list is longer and changes rarely, so it goes only when it has.
|
||||
// The flag is what makes an empty list mean "this node now holds nothing"
|
||||
// rather than "nothing to report" — without it, the first heartbeat after
|
||||
// a model was deleted would be indistinguishable from silence.
|
||||
weights := discovery.ScanWeights(a.cfg.ModelDirs)
|
||||
if fp := weightsFingerprint(weights); fp != a.lastWeights {
|
||||
a.lastWeights = fp
|
||||
hb.WeightsChanged = true
|
||||
for _, w := range weights {
|
||||
hb.Weights = append(hb.Weights, &llmgwv1.LocalWeights{
|
||||
Path: w.Path, SizeBytes: w.SizeBytes, Shared: w.Shared,
|
||||
})
|
||||
}
|
||||
}
|
||||
return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Heartbeat{Heartbeat: hb}}
|
||||
}
|
||||
|
||||
// weightsFingerprint is enough to notice a model appearing, disappearing or
|
||||
// finishing its download.
|
||||
//
|
||||
// Paths and sizes, because a fetch in progress has the same name as the file
|
||||
// it will become — under a different name until it is renamed, but a shared
|
||||
// filesystem can also grow a file in place.
|
||||
func weightsFingerprint(ws []discovery.LocalWeights) string {
|
||||
h := sha256.New()
|
||||
for _, w := range ws {
|
||||
fmt.Fprintf(h, "%s:%d\n", w.Path, w.SizeBytes)
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
// removeWeights deletes cached model files the control plane says are no
|
||||
// longer needed.
|
||||
//
|
||||
// Every path is checked here as well as there, and that duplication is the
|
||||
// point. The control plane computes these from the catalog, but it is the node
|
||||
// that holds the files, and a node must not be talkable into deleting a model
|
||||
// an operator spent a week copying onto it — not by a control plane that is
|
||||
// wrong, not by one that is newer than this agent, not by one that has been
|
||||
// tampered with. So this deletes only what it can prove llmgw itself wrote:
|
||||
// inside its own cache directory, under a name this system generates.
|
||||
//
|
||||
// Failures are logged and skipped rather than reported upward. A file that is
|
||||
// already gone is the expected outcome of a retry, and one that will not
|
||||
// delete is a disk to look at, not a reason to fail a command whose purpose
|
||||
// was to free space.
|
||||
func (a *Agent) removeWeights(paths []string) {
|
||||
cache := a.cfg.WeightsCache
|
||||
if cache == "" && len(a.cfg.ModelDirs) > 0 {
|
||||
cache = a.cfg.ModelDirs[0]
|
||||
}
|
||||
if cache == "" {
|
||||
a.log.Warn("asked to remove weights with no cache directory configured")
|
||||
return
|
||||
}
|
||||
cacheAbs, err := filepath.Abs(cache)
|
||||
if err != nil {
|
||||
a.log.Warn("cannot resolve the cache directory", "dir", cache, "err", err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, p := range paths {
|
||||
abs, err := filepath.Abs(p)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
// Directly inside the cache directory, not merely under it: a nested
|
||||
// path would let ".." or a symlinked subdirectory reach elsewhere, and
|
||||
// nothing this package writes is nested anyway.
|
||||
if filepath.Dir(abs) != cacheAbs {
|
||||
a.log.Warn("refused to remove a file outside the weights cache", "path", p)
|
||||
continue
|
||||
}
|
||||
if !weights.Owned(abs) {
|
||||
a.log.Warn("refused to remove a file llmgw did not fetch", "path", p)
|
||||
continue
|
||||
}
|
||||
if err := os.Remove(abs); err != nil && !os.IsNotExist(err) {
|
||||
a.log.Warn("could not remove cached weights", "path", p, "err", err)
|
||||
continue
|
||||
}
|
||||
a.log.Info("removed cached weights no model asks for", "path", p)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out chan<- *llmgwv1.AgentMessage) {
|
||||
switch m := msg.Msg.(type) {
|
||||
case *llmgwv1.ControlMessage_Start:
|
||||
@@ -380,6 +694,8 @@ func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out cha
|
||||
if err := a.sup.Drain(m.Drain.GetShardId()); err != nil {
|
||||
a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err)
|
||||
}
|
||||
case *llmgwv1.ControlMessage_RemoveWeights:
|
||||
a.removeWeights(m.RemoveWeights.GetPaths())
|
||||
case *llmgwv1.ControlMessage_Ack:
|
||||
if !m.Ack.GetOk() {
|
||||
a.log.Warn("control plane rejected a message",
|
||||
@@ -429,6 +745,21 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
return
|
||||
}
|
||||
|
||||
// The weights, before anything is launched. A source that is already a
|
||||
// path is checked and used; one that names a remote file is fetched once
|
||||
// and shared by every shard that wants it.
|
||||
//
|
||||
// Checked here rather than left to the engine because the engine's answer
|
||||
// arrives minutes into a load, in its own words, on a node nobody is
|
||||
// watching — and because eight shards starting together would otherwise
|
||||
// each fetch the same file.
|
||||
if path, err := a.weights().Ensure(ctx, spec.Profile.ModelPath); err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
} else {
|
||||
spec.Profile.ModelPath = path
|
||||
}
|
||||
|
||||
cmds, err := eng.Commands(spec)
|
||||
if err != nil {
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
@@ -536,23 +867,33 @@ func gpuInfo(inv discovery.Inventory, indices []int) []engines.GPUInfo {
|
||||
return out
|
||||
}
|
||||
|
||||
// HTTP2Client returns the client the agent should use to reach the control
|
||||
// plane over plaintext HTTP/2.
|
||||
// Client returns the client the agent reaches the control plane with.
|
||||
//
|
||||
// connect's bidirectional streams require HTTP/2, and an agent talking to a
|
||||
// control plane over a trusted network has no TLS to negotiate it with. Go's
|
||||
// default transport will not speak HTTP/2 without TLS, so it has to be asked
|
||||
// explicitly.
|
||||
// Both branches speak HTTP/2, which connect's bidirectional streams require and
|
||||
// which Go will not do by default in either case: without TLS it has to be
|
||||
// asked explicitly, and with TLS it has to be negotiated through ALPN.
|
||||
//
|
||||
// mTLS is plan §5.6's target and replaces this; until then the stream must not
|
||||
// cross an untrusted network.
|
||||
func HTTP2Client(timeout time.Duration) *http.Client {
|
||||
return &http.Client{
|
||||
Transport: &http2.Transport{
|
||||
AllowHTTP: true,
|
||||
DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
|
||||
return (&net.Dialer{Timeout: timeout}).DialContext(ctx, network, addr)
|
||||
// The plaintext branch exists for local runs against sim-engines. Anything
|
||||
// using it must stay on a trusted network, because nothing there stops a node
|
||||
// from claiming to be another one.
|
||||
func Client(cfg tlsconf.Config, timeout time.Duration) (*http.Client, error) {
|
||||
if !cfg.Enabled() {
|
||||
return &http.Client{
|
||||
Transport: &http2.Transport{
|
||||
AllowHTTP: true,
|
||||
DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
|
||||
return (&net.Dialer{Timeout: timeout}).DialContext(ctx, network, addr)
|
||||
},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ServerName is left to the transport, which takes it from the address
|
||||
// being dialled. Hardcoding it here would break the moment the control
|
||||
// plane moved.
|
||||
tc, err := cfg.ClientConfig("")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &http.Client{Transport: &http2.Transport{TLSClientConfig: tc}}, nil
|
||||
}
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
@@ -34,6 +35,7 @@ type statusRecorder struct {
|
||||
ch chan status
|
||||
|
||||
mu sync.Mutex
|
||||
calls []adoptCall
|
||||
adopted []reconciler.AdoptedShard
|
||||
unclaim []string // shard ids to report back as unclaimed
|
||||
}
|
||||
@@ -53,9 +55,10 @@ func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.
|
||||
}
|
||||
}
|
||||
|
||||
func (r *statusRecorder) Adopt(shards []reconciler.AdoptedShard) []reconciler.AdoptedShard {
|
||||
func (r *statusRecorder) Adopt(nodeID string, shards []reconciler.AdoptedShard) []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.calls = append(r.calls, adoptCall{nodeID: nodeID, shards: shards})
|
||||
r.adopted = append(r.adopted, shards...)
|
||||
|
||||
var unclaimed []reconciler.AdoptedShard
|
||||
@@ -67,6 +70,19 @@ func (r *statusRecorder) Adopt(shards []reconciler.AdoptedShard) []reconciler.Ad
|
||||
return unclaimed
|
||||
}
|
||||
|
||||
// adoptCall records that the service reached the reconciler at all, which an
|
||||
// empty shard list used to prevent.
|
||||
type adoptCall struct {
|
||||
nodeID string
|
||||
shards []reconciler.AdoptedShard
|
||||
}
|
||||
|
||||
func (r *statusRecorder) adoptCalls() []adoptCall {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return slices.Clone(r.calls)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) adoptedShards() []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
@@ -208,7 +224,7 @@ func startWith(t *testing.T, log *slog.Logger, cfgFns ...func(*Config)) *harness
|
||||
|
||||
inv := inventory.NewStore()
|
||||
rec := newRecorder()
|
||||
svc := agentsvc.New(inv, rec, log)
|
||||
svc := agentsvc.New(inv, rec, false, log)
|
||||
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
@@ -285,8 +301,17 @@ func (h *harness) waitRegistered(t *testing.T) {
|
||||
}
|
||||
|
||||
// shardCommand builds a single-node vLLM shard on a free port.
|
||||
//
|
||||
// The model path is a real file. It used to be "/models/x", which stopped
|
||||
// working once the agent began checking that weights exist before launching an
|
||||
// engine — the check doing exactly what it is for, on a test that was relying
|
||||
// on nobody looking.
|
||||
func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand {
|
||||
t.Helper()
|
||||
modelPath := filepath.Join(t.TempDir(), "model.gguf")
|
||||
if err := os.WriteFile(modelPath, []byte("weights"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return reconciler.ShardCommand{
|
||||
ReplicaID: "rep-1", ShardID: shardID, NodeID: "n1",
|
||||
Spec: engines.ShardSpec{
|
||||
@@ -294,7 +319,7 @@ func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand {
|
||||
GPUIndices: []int{0}, HTTPPort: freePort(t),
|
||||
Profile: engines.Profile{
|
||||
Name: "p", Engine: "vllm", TP: 1,
|
||||
ModelPath: "/models/x", ServedModelName: "test-model",
|
||||
ModelPath: modelPath, ServedModelName: "test-model",
|
||||
},
|
||||
},
|
||||
StartupTimeout: 15 * time.Second,
|
||||
@@ -312,9 +337,15 @@ func freePort(t *testing.T) int {
|
||||
}
|
||||
|
||||
func TestRegistrationCarriesInventory(t *testing.T) {
|
||||
// The node must actually have the adapters it declares. It used to declare
|
||||
// mlx5_0 and mlx5_1 on a node whose sysfs held neither, which passed
|
||||
// because nothing compared the two — and that is the gap a real cluster
|
||||
// would have found as a replica hanging on its first all-reduce.
|
||||
ibRoot := ibTreeWithActivePorts(t, "mlx5_0", "mlx5_1")
|
||||
h := start(t, func(c *Config) {
|
||||
c.Hostname = "gpu-a-01"
|
||||
c.Labels = map[string]string{"rack": "r12"}
|
||||
c.Discovery.IBRoot = ibRoot
|
||||
c.Fabrics = []inventory.FabricAttachment{
|
||||
{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200},
|
||||
}
|
||||
@@ -598,3 +629,110 @@ func TestUnclaimedShardIsStopped(t *testing.T) {
|
||||
func processAliveForTest(pid int) bool {
|
||||
return syscall.Kill(pid, 0) == nil
|
||||
}
|
||||
|
||||
func TestAShardWithMissingWeightsFailsBeforeLaunching(t *testing.T) {
|
||||
// The commonest setup mistake on a fresh node. Left to the engine it
|
||||
// surfaces minutes into a load, in the engine's own words, on a machine
|
||||
// nobody is watching — so it is caught here, where the message can say
|
||||
// what is actually wrong.
|
||||
a := &Agent{cfg: Config{WeightsCache: t.TempDir()}, log: slog.New(slog.DiscardHandler)}
|
||||
_, err := a.weights().Ensure(t.Context(), filepath.Join(t.TempDir(), "not-there.gguf"))
|
||||
if err == nil {
|
||||
t.Fatal("a missing weights file was accepted")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not-there.gguf") {
|
||||
t.Errorf("the error does not name the file: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheCacheDefaultsToTheFirstModelDirectory(t *testing.T) {
|
||||
// A node that has said where its models live has already answered this;
|
||||
// asking twice is a second thing to get out of step.
|
||||
dir := t.TempDir()
|
||||
a := &Agent{cfg: Config{ModelDirs: []string{dir, "/elsewhere"}}, log: slog.New(slog.DiscardHandler)}
|
||||
if got := a.weights().CacheDir; got != dir {
|
||||
t.Errorf("cache = %q, want the first model directory %q", got, dir)
|
||||
}
|
||||
}
|
||||
|
||||
// ibTreeWithActivePorts builds a sysfs InfiniBand tree with the named adapters,
|
||||
// each with one active port, shaped like the real thing.
|
||||
func ibTreeWithActivePorts(t *testing.T, names ...string) string {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
for _, name := range names {
|
||||
dir := filepath.Join(root, name, "ports", "1")
|
||||
if err := os.MkdirAll(filepath.Join(dir, "gids"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for file, content := range map[string]string{
|
||||
"state": "4: ACTIVE",
|
||||
"phys_state": "5: LinkUp",
|
||||
"rate": "200 Gb/sec (4X HDR)",
|
||||
"link_layer": "InfiniBand",
|
||||
"lid": "0x1a",
|
||||
"gids/0": "fec0:0000:0000:000a:0002:c903:00a0:1234",
|
||||
} {
|
||||
if err := os.WriteFile(filepath.Join(dir, file), []byte(content+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return root
|
||||
}
|
||||
|
||||
func TestAFabricWhoseAdaptersAreAbsentIsNotAdvertised(t *testing.T) {
|
||||
// The failure this prevents: a typo, or a card that enumerated differently
|
||||
// after a reboot, produced NCCL_IB_HCA naming an adapter that is not there.
|
||||
// The replica starts, passes its health checks, and hangs on its first
|
||||
// all-reduce — the exact case the fabric invariant exists to prevent,
|
||||
// arriving through a door nobody was watching.
|
||||
h := start(t, func(c *Config) {
|
||||
c.Discovery.IBRoot = ibTreeWithActivePorts(t, "mlx5_0")
|
||||
c.Fabrics = []inventory.FabricAttachment{
|
||||
{Fabric: "ib-typo", HCAs: []string{"mlx5_9"}, RateGbps: 200},
|
||||
}
|
||||
})
|
||||
|
||||
n, ok := h.inv.Get("n1")
|
||||
if !ok {
|
||||
t.Fatal("node missing from inventory")
|
||||
}
|
||||
if n.HasFabric("ib-typo") {
|
||||
t.Errorf("a fabric with no usable adapter was advertised: %+v", n.Fabrics)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFabricRateIsTheSlowestActiveLink(t *testing.T) {
|
||||
// A collective runs at the speed of its slowest link, so reporting the
|
||||
// fastest would have the scheduler believe a fabric is quicker than it is.
|
||||
root := t.TempDir()
|
||||
for name, rate := range map[string]string{"mlx5_0": "200", "mlx5_1": "100"} {
|
||||
dir := filepath.Join(root, name, "ports", "1")
|
||||
if err := os.MkdirAll(filepath.Join(dir, "gids"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for file, content := range map[string]string{
|
||||
"state": "4: ACTIVE", "phys_state": "5: LinkUp",
|
||||
"rate": rate + " Gb/sec (4X HDR)", "link_layer": "InfiniBand",
|
||||
"lid": "0x1a", "gids/0": "fec0::1",
|
||||
} {
|
||||
if err := os.WriteFile(filepath.Join(dir, file), []byte(content+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
h := start(t, func(c *Config) {
|
||||
c.Discovery.IBRoot = root
|
||||
c.Fabrics = []inventory.FabricAttachment{
|
||||
{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}},
|
||||
}
|
||||
})
|
||||
|
||||
n, _ := h.inv.Get("n1")
|
||||
for _, f := range n.Fabrics {
|
||||
if f.Fabric == "ib-a" && f.RateGbps != 100 {
|
||||
t.Errorf("fabric rate = %d, want the slower link's 100", f.RateGbps)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -26,6 +26,8 @@ type Options struct {
|
||||
PreferNVML bool `koanf:"prefer_nvml"`
|
||||
// SkipGPUs omits GPU discovery, for nodes that have none.
|
||||
SkipGPUs bool `koanf:"skip_gpus"`
|
||||
// ModelDirs are scanned for weights. Empty means do not look.
|
||||
ModelDirs []string `koanf:"-"`
|
||||
}
|
||||
|
||||
// Discover builds the node's inventory.
|
||||
@@ -39,6 +41,19 @@ func Discover(ctx context.Context, opts Options, log *slog.Logger) (Inventory, e
|
||||
log = slog.Default()
|
||||
}
|
||||
inv := Inventory{Source: "none"}
|
||||
// First, because it has nothing to do with GPUs and every path below can
|
||||
// return early. Read after them, a node with no nvidia-smi reported no
|
||||
// memory at all: GPU discovery fails, Discover returns, and the fact that
|
||||
// has nothing to do with cards goes with it. That is exactly the node whose
|
||||
// memory matters most, since a CPU profile is what it will be asked to run.
|
||||
inv.MemoryTotalMB = MemoryTotalMB()
|
||||
// Also before the GPU paths, and for the same reason: which weights a node
|
||||
// holds has nothing to do with its cards, and a node with none is exactly
|
||||
// where a CPU profile will look for them.
|
||||
inv.Weights = ScanWeights(opts.ModelDirs)
|
||||
// And the disk they sit on. A node can hold a model today and be unable to
|
||||
// take the next one, which the file list alone never shows.
|
||||
inv.Storage = ScanStorage(opts.ModelDirs)
|
||||
|
||||
if !opts.SkipGPUs {
|
||||
switch {
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"testing"
|
||||
)
|
||||
|
||||
@@ -99,3 +100,20 @@ func TestNVMLSupportedMatchesBuild(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMemoryIsReportedEvenWhenGPUDiscoveryFails(t *testing.T) {
|
||||
// A node with no nvidia-smi is exactly the node whose memory matters, since
|
||||
// a CPU profile is what it will be asked to run. Reading memory after GPU
|
||||
// discovery meant that node reported none: the GPU error returns early and
|
||||
// takes a fact that has nothing to do with cards with it. Seen on node05,
|
||||
// which has 32 GB and was reporting "-".
|
||||
inv, err := Discover(t.Context(), Options{NvidiaSMI: "/nonexistent/nvidia-smi"}, nil)
|
||||
if err == nil {
|
||||
t.Fatal("GPU discovery was expected to fail with no nvidia-smi")
|
||||
}
|
||||
// On Linux this is a real number; elsewhere there is no /proc/meminfo and 0
|
||||
// is the honest answer. Either way it must not be lost to the GPU error.
|
||||
if runtime.GOOS == "linux" && inv.MemoryTotalMB <= 0 {
|
||||
t.Errorf("memory was not reported alongside a GPU failure: %d", inv.MemoryTotalMB)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// FabricCheck is one declared fabric measured against what the node actually
|
||||
// has.
|
||||
type FabricCheck struct {
|
||||
Fabric string
|
||||
// Usable are the declared adapters that exist and have an active port.
|
||||
Usable []string
|
||||
// RateGbps is the slowest active port among them, because a collective
|
||||
// runs at the speed of its slowest link.
|
||||
RateGbps int
|
||||
// Problems are the declared adapters that cannot carry traffic, and why.
|
||||
Problems []string
|
||||
}
|
||||
|
||||
// OK reports whether this fabric can carry a collective at all.
|
||||
func (c FabricCheck) OK() bool { return len(c.Usable) > 0 }
|
||||
|
||||
// CheckFabrics measures declared fabric membership against discovered adapters.
|
||||
//
|
||||
// Fabric membership has to be configuration — OpenSM's default subnet prefix is
|
||||
// the same on every fabric, so a matching prefix proves nothing — but taking it
|
||||
// entirely on trust was a mistake of a different kind. A typo, or a card that
|
||||
// enumerated differently after a reboot, produced NCCL_IB_HCA naming an adapter
|
||||
// that does not exist; and an adapter whose cable is out, or whose subnet
|
||||
// manager is down, is present in sysfs and carries nothing.
|
||||
//
|
||||
// Either way the replica starts, passes its health checks, and hangs on its
|
||||
// first all-reduce — which is the failure the whole fabric invariant exists to
|
||||
// prevent, arriving through a door nobody was watching. Declared and real are
|
||||
// now compared where both are visible, which is only here.
|
||||
func CheckFabrics(declared map[string][]string, discovered []HCA) []FabricCheck {
|
||||
byName := make(map[string]HCA, len(discovered))
|
||||
for _, h := range discovered {
|
||||
byName[h.Name] = h
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(declared))
|
||||
for f := range declared {
|
||||
names = append(names, f)
|
||||
}
|
||||
sort.Strings(names)
|
||||
|
||||
out := make([]FabricCheck, 0, len(names))
|
||||
for _, fabric := range names {
|
||||
c := FabricCheck{Fabric: fabric}
|
||||
for _, want := range declared[fabric] {
|
||||
hca, found := byName[want]
|
||||
if !found {
|
||||
c.Problems = append(c.Problems, fmt.Sprintf(
|
||||
"%s is not on this node (found: %s)", want, listNames(discovered)))
|
||||
continue
|
||||
}
|
||||
rate, active := activeRate(hca)
|
||||
if !active {
|
||||
c.Problems = append(c.Problems, fmt.Sprintf(
|
||||
"%s has no active port (%s)", want, portStates(hca)))
|
||||
continue
|
||||
}
|
||||
c.Usable = append(c.Usable, want)
|
||||
if c.RateGbps == 0 || rate < c.RateGbps {
|
||||
c.RateGbps = rate
|
||||
}
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// activeRate returns the fastest active port's rate on an adapter.
|
||||
//
|
||||
// Only ACTIVE ports carry traffic. A port that is physically linked but has no
|
||||
// subnet manager sits in INIT, looks healthy to anything counting adapters, and
|
||||
// passes no data — which is precisely the case worth catching, because it
|
||||
// survives a reboot and looks like a software problem.
|
||||
func activeRate(h HCA) (int, bool) {
|
||||
best, ok := 0, false
|
||||
for _, p := range h.Ports {
|
||||
if !strings.EqualFold(p.State, "ACTIVE") {
|
||||
continue
|
||||
}
|
||||
ok = true
|
||||
if p.RateGbps > best {
|
||||
best = p.RateGbps
|
||||
}
|
||||
}
|
||||
return best, ok
|
||||
}
|
||||
|
||||
func portStates(h HCA) string {
|
||||
if len(h.Ports) == 0 {
|
||||
return "no ports"
|
||||
}
|
||||
var parts []string
|
||||
for _, p := range h.Ports {
|
||||
parts = append(parts, fmt.Sprintf("port %d %s/%s", p.Num, p.State, p.PhysState))
|
||||
}
|
||||
return strings.Join(parts, ", ")
|
||||
}
|
||||
|
||||
func listNames(hcas []HCA) string {
|
||||
if len(hcas) == 0 {
|
||||
return "none"
|
||||
}
|
||||
var names []string
|
||||
for _, h := range hcas {
|
||||
names = append(names, h.Name)
|
||||
}
|
||||
return strings.Join(names, ", ")
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func hca(name string, ports ...Port) HCA { return HCA{Name: name, Ports: ports} }
|
||||
|
||||
func active(rate int) Port {
|
||||
return Port{Num: 1, State: "ACTIVE", PhysState: "LinkUp", RateGbps: rate, LinkLayer: LinkInfiniBand}
|
||||
}
|
||||
|
||||
func TestADeclaredAdapterThatIsNotThereIsReported(t *testing.T) {
|
||||
// A typo, or a card that enumerated differently after a reboot. Taken on
|
||||
// trust it becomes NCCL_IB_HCA naming an adapter that does not exist, and
|
||||
// the replica hangs on its first all-reduce.
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-a": {"mlx5_0", "mlx5_9"}},
|
||||
[]HCA{hca("mlx5_0", active(200))},
|
||||
)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("got %d checks", len(got))
|
||||
}
|
||||
c := got[0]
|
||||
if len(c.Usable) != 1 || c.Usable[0] != "mlx5_0" {
|
||||
t.Errorf("usable = %v, want only the adapter that exists", c.Usable)
|
||||
}
|
||||
if len(c.Problems) != 1 || !strings.Contains(c.Problems[0], "mlx5_9") {
|
||||
t.Errorf("problems = %v", c.Problems)
|
||||
}
|
||||
// The message names what is actually present, because the operator is
|
||||
// looking at a typo and cannot see the node.
|
||||
if !strings.Contains(c.Problems[0], "mlx5_0") {
|
||||
t.Errorf("the problem does not say what was found: %q", c.Problems[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAdapterWithNoActivePortCannotCarryTraffic(t *testing.T) {
|
||||
// A port that is physically linked but has no subnet manager sits in INIT.
|
||||
// It looks healthy to anything counting adapters and passes no data, which
|
||||
// is the case worth catching: it survives a reboot and looks like software.
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-a": {"mlx5_0"}},
|
||||
[]HCA{hca("mlx5_0", Port{Num: 1, State: "INIT", PhysState: "LinkUp", RateGbps: 200})},
|
||||
)
|
||||
c := got[0]
|
||||
if c.OK() {
|
||||
t.Error("a fabric with no active port was reported usable")
|
||||
}
|
||||
if len(c.Problems) != 1 || !strings.Contains(c.Problems[0], "INIT") {
|
||||
t.Errorf("the problem does not say what state the port is in: %v", c.Problems)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheRateIsTheSlowestActiveLink(t *testing.T) {
|
||||
// A collective runs at the speed of its slowest link, so reporting the
|
||||
// fastest would have the scheduler believe a fabric is quicker than it is.
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-a": {"fast", "slow"}},
|
||||
[]HCA{hca("fast", active(400)), hca("slow", active(100))},
|
||||
)
|
||||
if got[0].RateGbps != 100 {
|
||||
t.Errorf("rate = %d, want 100", got[0].RateGbps)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAdapterWithSeveralPortsUsesItsBestActiveOne(t *testing.T) {
|
||||
// A dual-port card with one cable in is a perfectly good adapter, and
|
||||
// judging it by the dead port would take a working fabric out of service.
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-a": {"mlx5_0"}},
|
||||
[]HCA{hca("mlx5_0",
|
||||
Port{Num: 1, State: "DOWN", PhysState: "Disabled"},
|
||||
active(200),
|
||||
)},
|
||||
)
|
||||
if !got[0].OK() {
|
||||
t.Fatalf("a card with one live port was rejected: %+v", got[0])
|
||||
}
|
||||
if got[0].RateGbps != 200 {
|
||||
t.Errorf("rate = %d, want the active port's 200", got[0].RateGbps)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFabricWithNothingUsableIsNotOK(t *testing.T) {
|
||||
got := CheckFabrics(
|
||||
map[string][]string{"ib-dead": {"mlx5_0"}},
|
||||
[]HCA{hca("mlx5_0", Port{Num: 1, State: "DOWN", PhysState: "Polling"})},
|
||||
)
|
||||
if got[0].OK() {
|
||||
t.Error("a fabric whose only adapter is down was reported usable")
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodeWithNoFabricsDeclaredProducesNoChecks(t *testing.T) {
|
||||
// Every node without InfiniBand, which must stay entirely unaffected.
|
||||
if got := CheckFabrics(nil, []HCA{hca("mlx5_0", active(200))}); len(got) != 0 {
|
||||
t.Errorf("got %d checks for a node that declared none", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestChecksComeBackInAStableOrder(t *testing.T) {
|
||||
// Map iteration would otherwise reorder log lines between restarts, which
|
||||
// makes a diff of two registrations unreadable.
|
||||
declared := map[string][]string{"ib-c": {"m"}, "ib-a": {"m"}, "ib-b": {"m"}}
|
||||
got := CheckFabrics(declared, nil)
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("got %d", len(got))
|
||||
}
|
||||
if got[0].Fabric != "ib-a" || got[1].Fabric != "ib-b" || got[2].Fabric != "ib-c" {
|
||||
t.Errorf("order = %s, %s, %s", got[0].Fabric, got[1].Fabric, got[2].Fabric)
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,14 @@ type GPU struct {
|
||||
// MemoryUsedMB is what is already allocated. Non-zero on an idle node means
|
||||
// something other than our engines is using the card (plan §5.1).
|
||||
MemoryUsedMB int `json:"memory_used_mb"`
|
||||
// UtilizationPercent is how busy the card was over the last sample period.
|
||||
//
|
||||
// The number people actually watch, and the one nothing collected: memory
|
||||
// used says a model is loaded, and this says whether it is doing anything.
|
||||
// A fleet of cards at ninety percent occupancy and two percent utilisation
|
||||
// is a fleet holding models nobody is asking for, which is invisible
|
||||
// without it.
|
||||
UtilizationPercent int `json:"utilization_percent,omitempty"`
|
||||
// ComputeCapability gates quantization choices: an FP8 profile must not
|
||||
// land on a card below 8.9 (plan §6.2).
|
||||
ComputeCapability string `json:"compute_capability"`
|
||||
@@ -100,6 +108,15 @@ type Inventory struct {
|
||||
// Source names the backend that produced the GPU list, so a consumer can
|
||||
// tell "no GPUs" from "discovery did not work".
|
||||
Source string `json:"source"`
|
||||
// MemoryTotalMB is the machine's RAM. Zero means unknown, which is what an
|
||||
// agent older than this field reports, and is treated as no constraint.
|
||||
MemoryTotalMB int `json:"memory_total_mb,omitempty"`
|
||||
// Weights are the model files this node can see. Empty means it was not
|
||||
// asked to look.
|
||||
Weights []LocalWeights `json:"weights,omitempty"`
|
||||
// Storage is the disk behind those weights, one entry per filesystem.
|
||||
// Empty means the same thing it means for Weights: not asked.
|
||||
Storage []Storage `json:"storage,omitempty"`
|
||||
}
|
||||
|
||||
// Classes returns the distinct GPU classes present, sorted.
|
||||
|
||||
@@ -21,8 +21,21 @@ var smiQueryFields = []string{
|
||||
"driver_version",
|
||||
"temperature.gpu",
|
||||
"ecc.errors.uncorrected.volatile.total",
|
||||
// Last, and optional. An nvidia-smi too old to know this field returns a
|
||||
// shorter row, and every column before it still lines up — which is why it
|
||||
// goes on the end rather than anywhere else, and why the parser requires
|
||||
// only the mandatory prefix.
|
||||
"utilization.gpu",
|
||||
}
|
||||
|
||||
// mandatorySMIFields is how many columns a row must have to be usable.
|
||||
//
|
||||
// Everything before utilization.gpu. A driver that does not know the last field
|
||||
// gives a short row, and rejecting those would make an llmgw upgrade silently
|
||||
// stop discovering GPUs on an older node — a far worse outcome than not knowing
|
||||
// how busy a card is.
|
||||
const mandatorySMIFields = 10
|
||||
|
||||
// DiscoverGPUsSMI enumerates GPUs by shelling out to nvidia-smi.
|
||||
//
|
||||
// This is the cgo-free path (plan §3.2): the NVML bindings need cgo, which
|
||||
@@ -55,7 +68,7 @@ func parseSMI(out string) ([]GPU, error) {
|
||||
|
||||
gpus := make([]GPU, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
if len(row) < len(smiQueryFields) {
|
||||
if len(row) < mandatorySMIFields {
|
||||
continue
|
||||
}
|
||||
idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
|
||||
@@ -76,6 +89,9 @@ func parseSMI(out string) ([]GPU, error) {
|
||||
g.MemoryUsedMB, _ = parseCSVInt(row[4])
|
||||
g.Health.TemperatureC, _ = parseCSVInt(row[8])
|
||||
g.Health.ECCErrors, _ = parseCSVInt(row[9])
|
||||
if len(row) > 10 {
|
||||
g.UtilizationPercent, _ = parseCSVInt(row[10])
|
||||
}
|
||||
|
||||
assessHealth(&g)
|
||||
gpus = append(gpus, g)
|
||||
|
||||
@@ -142,3 +142,38 @@ func TestDiscoverGPUsSMIMissingBinary(t *testing.T) {
|
||||
t.Error("expected an error when nvidia-smi is missing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnOlderDriverWithoutUtilisationStillDiscovers(t *testing.T) {
|
||||
// An nvidia-smi too old for utilization.gpu returns a shorter row. Rejecting
|
||||
// those would make an llmgw upgrade silently stop discovering GPUs on an
|
||||
// older node, which is far worse than not knowing how busy a card is — and
|
||||
// is what the first version of this change did.
|
||||
old := "0, GPU-aaa, NVIDIA H100 80GB HBM3, 81559, 1024, 9.0, 0000:01:00.0, 550.54, 41, 0"
|
||||
gpus, err := parseSMI(old)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("a row without utilisation produced %d GPUs", len(gpus))
|
||||
}
|
||||
if gpus[0].MemoryTotalMB != 81559 || gpus[0].Health.TemperatureC != 41 {
|
||||
t.Errorf("columns shifted: %+v", gpus[0])
|
||||
}
|
||||
if gpus[0].UtilizationPercent != 0 {
|
||||
t.Errorf("utilisation = %d, want 0 when unreported", gpus[0].UtilizationPercent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUtilisationIsReadWhenPresent(t *testing.T) {
|
||||
row := "0, GPU-aaa, NVIDIA H100 80GB HBM3, 81559, 1024, 9.0, 0000:01:00.0, 550.54, 41, 0, 87"
|
||||
gpus, err := parseSMI(row)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("got %d GPUs", len(gpus))
|
||||
}
|
||||
if gpus[0].UtilizationPercent != 87 {
|
||||
t.Errorf("utilisation = %d, want 87", gpus[0].UtilizationPercent)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// MemoryTotalMB reports the machine's total RAM.
|
||||
//
|
||||
// It exists because placement had no idea how much a node could hold. A CPU
|
||||
// profile has no card to be free, so every ready node looked like a candidate:
|
||||
// a 14B model whose working set is 24 GB would have been placed just as
|
||||
// happily on a node with 8 GB, which reports itself perfectly healthy right up
|
||||
// until the OOM killer arrives. Labels were the stopgap; this is the fact they
|
||||
// were standing in for.
|
||||
//
|
||||
// Read from /proc/meminfo, which is Linux-only — every node is. On anything
|
||||
// else it returns 0, which callers treat as "unknown" and therefore no
|
||||
// constraint, exactly as an older agent that does not report it at all.
|
||||
func MemoryTotalMB() int { return memTotalMBFrom("/proc/meminfo") }
|
||||
|
||||
func memTotalMBFrom(path string) int {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
sc := bufio.NewScanner(f)
|
||||
for sc.Scan() {
|
||||
rest, ok := strings.CutPrefix(sc.Text(), "MemTotal:")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fields := strings.Fields(rest)
|
||||
if len(fields) == 0 {
|
||||
return 0
|
||||
}
|
||||
kb, err := strconv.Atoi(fields[0])
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
// MemTotal is what the kernel can hand out, already short of the
|
||||
// installed total by the firmware's reservations. That is the right
|
||||
// number for deciding what fits.
|
||||
return kb / 1024
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMemTotalIsReadFromMeminfo(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "meminfo")
|
||||
// The real shape, including the fields before and after, so the parser is
|
||||
// not tested against something tidier than it will see.
|
||||
if err := os.WriteFile(p, []byte(
|
||||
"MemTotal: 32876244 kB\nMemFree: 1234 kB\nMemAvailable: 5678 kB\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := memTotalMBFrom(p), 32876244/1024; got != want {
|
||||
t.Errorf("memTotalMBFrom = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreadableMeminfoReportsUnknownRatherThanZeroMemory(t *testing.T) {
|
||||
// Zero means "did not say" everywhere downstream, and placement treats it
|
||||
// as no constraint. A node with genuinely no memory is not a case.
|
||||
if got := memTotalMBFrom(filepath.Join(t.TempDir(), "absent")); got != 0 {
|
||||
t.Errorf("got %d for a missing file", got)
|
||||
}
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "meminfo")
|
||||
for _, junk := range []string{"", "MemTotal:\n", "MemTotal: not-a-number kB\n", "Other: 1 kB\n"} {
|
||||
if err := os.WriteFile(p, []byte(junk), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := memTotalMBFrom(p); got != 0 {
|
||||
t.Errorf("got %d for %q", got, junk)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// One entry per filesystem rather than per configured directory, because free
|
||||
// space belongs to the disk and not to the path. Two model directories on the
|
||||
// same disk are one pool of bytes, and reporting them separately would let the
|
||||
// control plane add 24 GB to 24 GB and place a 40 GB model on a node with 24.
|
||||
type Storage struct {
|
||||
// Path is one directory on this filesystem — the first one configured that
|
||||
// landed here, so the operator recognises it.
|
||||
Path string `json:"path"`
|
||||
// FreeBytes is what a non-root process can still write. Deliberately not
|
||||
// the root reserve: the agent does not run as root, so the reserve is space
|
||||
// it cannot have.
|
||||
FreeBytes int64 `json:"free_bytes"`
|
||||
// TotalBytes is the filesystem's size, so a reading of "8 GB free" can be
|
||||
// told apart on a 10 TB filer and on a 10 GB root partition.
|
||||
TotalBytes int64 `json:"total_bytes"`
|
||||
// Shared marks a filesystem the node did not have to fill. Free space on
|
||||
// one is shared with every other node using it, so two nodes reporting
|
||||
// 500 GB free on the same filer have 500 GB between them, not a terabyte.
|
||||
Shared bool `json:"shared,omitempty"`
|
||||
}
|
||||
|
||||
// ScanStorage reports the filesystems behind the given model directories.
|
||||
//
|
||||
// Empty in, empty out, and a directory that cannot be read is skipped rather
|
||||
// than reported as full: a mount that has not come back yet looks exactly like
|
||||
// a missing directory, and claiming zero free bytes would take the node out of
|
||||
// placement for a reason that is not true. Absent is not zero here, as
|
||||
// everywhere else in this system.
|
||||
func ScanStorage(dirs []string) []Storage {
|
||||
var out []Storage
|
||||
seen := map[uint64]int{} // filesystem id -> index in out
|
||||
|
||||
for _, dir := range dirs {
|
||||
if dir == "" {
|
||||
continue
|
||||
}
|
||||
dev, ok := filesystemID(dir)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if i, dup := seen[dev]; dup {
|
||||
// Same disk, already counted. Keep the first path: it is the one
|
||||
// the operator wrote first and so the one they will recognise.
|
||||
_ = i
|
||||
continue
|
||||
}
|
||||
|
||||
var st syscall.Statfs_t
|
||||
if err := syscall.Statfs(dir, &st); err != nil {
|
||||
continue
|
||||
}
|
||||
// Bsize is int32 on some platforms and int64 on others, which is why
|
||||
// every field here is widened explicitly rather than assumed.
|
||||
block := int64(st.Bsize)
|
||||
seen[dev] = len(out)
|
||||
out = append(out, Storage{
|
||||
Path: dir,
|
||||
FreeBytes: int64(st.Bavail) * block, //nolint:gosec,unconvert // widened on purpose
|
||||
TotalBytes: int64(st.Blocks) * block, //nolint:gosec,unconvert // widened on purpose
|
||||
Shared: isNetworkFS(dir),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// filesystemID identifies the disk a path is on.
|
||||
//
|
||||
// By device number rather than by the statfs filesystem id, because the field
|
||||
// holding that id is spelled differently on Linux and on macOS and this has to
|
||||
// compile on both — the agent runs on one and is developed on the other.
|
||||
func filesystemID(path string) (uint64, bool) {
|
||||
fi, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
st, ok := fi.Sys().(*syscall.Stat_t)
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
return uint64(st.Dev), true //nolint:gosec,unconvert // int32 on macOS, uint64 on Linux
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTwoDirectoriesOnOneDiskAreOnePoolOfBytes(t *testing.T) {
|
||||
// The reason this matters: free space belongs to the disk, not to the path.
|
||||
// Counting two directories separately lets the control plane add 24 GB to
|
||||
// 24 GB and place a 40 GB model on a node that has 24.
|
||||
base := t.TempDir()
|
||||
a := filepath.Join(base, "models-a")
|
||||
b := filepath.Join(base, "models-b")
|
||||
for _, d := range []string{a, b} {
|
||||
if err := os.Mkdir(d, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
got := ScanStorage([]string{a, b})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("two directories on one filesystem reported as %d pools: %+v", len(got), got)
|
||||
}
|
||||
if got[0].Path != a {
|
||||
t.Errorf("kept %q rather than the first directory configured", got[0].Path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStorageReportsSomethingPlausible(t *testing.T) {
|
||||
got := ScanStorage([]string{t.TempDir()})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("a real directory reported %d filesystems", len(got))
|
||||
}
|
||||
s := got[0]
|
||||
if s.TotalBytes <= 0 {
|
||||
t.Error("a filesystem with no size at all")
|
||||
}
|
||||
if s.FreeBytes < 0 || s.FreeBytes > s.TotalBytes {
|
||||
t.Errorf("free %d of total %d", s.FreeBytes, s.TotalBytes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingDirectoryIsNotAFullDisk(t *testing.T) {
|
||||
// A mount that has not come back yet looks exactly like a missing
|
||||
// directory. Reporting it as zero free bytes would take the node out of
|
||||
// placement for a reason that is not true; reporting nothing says only
|
||||
// that nothing is known, which is the honest answer.
|
||||
got := ScanStorage([]string{filepath.Join(t.TempDir(), "never-mounted")})
|
||||
if len(got) != 0 {
|
||||
t.Errorf("a missing directory was reported as a filesystem: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoDirectoriesConstrainNothing(t *testing.T) {
|
||||
if got := ScanStorage(nil); len(got) != 0 {
|
||||
t.Errorf("scanning nothing found %+v", got)
|
||||
}
|
||||
if got := ScanStorage([]string{""}); len(got) != 0 {
|
||||
t.Errorf("an empty path found %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"io/fs"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// LocalWeights is one model file or directory a node can see.
|
||||
type LocalWeights struct {
|
||||
Path string `json:"path"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
// Shared marks weights on a filesystem this node did not have to fill.
|
||||
//
|
||||
// Worth distinguishing: twenty nodes each holding their own copy and twenty
|
||||
// nodes seeing one filer are the same list of paths and completely
|
||||
// different situations — one has twenty copies of a 141 GB model and the
|
||||
// other has one.
|
||||
Shared bool `json:"shared,omitempty"`
|
||||
}
|
||||
|
||||
// ScanWeights lists the model files under the given directories.
|
||||
//
|
||||
// Storage is the part of this system that will differ at every installation —
|
||||
// a Lustre mount, an image exported over the fabric, local NVMe, an rsync from
|
||||
// a laptop — so llmgw does not try to own it. It observes the result instead:
|
||||
// whatever put the file there, the question placement needs answered is only
|
||||
// whether it is present.
|
||||
//
|
||||
// Scanning nothing is the default and constrains nothing. A node with no
|
||||
// directories configured behaves exactly as every node did before this existed,
|
||||
// which is what makes the whole feature optional.
|
||||
func ScanWeights(dirs []string) []LocalWeights {
|
||||
var out []LocalWeights
|
||||
seen := map[string]struct{}{}
|
||||
|
||||
for _, dir := range dirs {
|
||||
if dir == "" {
|
||||
continue
|
||||
}
|
||||
shared := isNetworkFS(dir)
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
// A directory that is not there is not an error to report upward:
|
||||
// a mount that has not come back yet looks exactly like this, and
|
||||
// the node is simply one without those weights until it does.
|
||||
continue
|
||||
}
|
||||
for _, e := range entries {
|
||||
path := filepath.Join(dir, e.Name())
|
||||
if _, dup := seen[path]; dup {
|
||||
continue
|
||||
}
|
||||
size, ok := weightSize(path, e)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
seen[path] = struct{}{}
|
||||
out = append(out, LocalWeights{Path: path, SizeBytes: size, Shared: shared})
|
||||
}
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path })
|
||||
return out
|
||||
}
|
||||
|
||||
// weightSize measures a candidate and says whether it is one at all.
|
||||
//
|
||||
// A model is either a single file — a .gguf, usually — or a directory of
|
||||
// shards and a config, which is what every HuggingFace checkout looks like.
|
||||
// Anything else in the directory is somebody's notes and is not reported.
|
||||
func weightSize(path string, e fs.DirEntry) (int64, bool) {
|
||||
if e.IsDir() {
|
||||
// A directory counts only if it actually holds weights. An empty one
|
||||
// is a half-finished download, and reporting it as present would have
|
||||
// placement choose a node whose engine then fails to load.
|
||||
var total int64
|
||||
var found bool
|
||||
_ = filepath.WalkDir(path, func(p string, d fs.DirEntry, err error) error {
|
||||
if err != nil || d.IsDir() {
|
||||
return nil //nolint:nilerr // an unreadable subtree is not fatal
|
||||
}
|
||||
if isWeightFile(p) {
|
||||
found = true
|
||||
}
|
||||
if info, err := d.Info(); err == nil {
|
||||
total += info.Size()
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return total, found
|
||||
}
|
||||
if !isWeightFile(path) {
|
||||
return 0, false
|
||||
}
|
||||
info, err := e.Info()
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return info.Size(), true
|
||||
}
|
||||
|
||||
func isWeightFile(path string) bool {
|
||||
switch strings.ToLower(filepath.Ext(path)) {
|
||||
case ".gguf", ".safetensors", ".bin", ".pt", ".pth":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isNetworkFS reports whether a path is on a filesystem the node did not fill.
|
||||
//
|
||||
// By magic number rather than by parsing mounts, because the mount table says
|
||||
// what was asked for and this says what the kernel actually has. The list is
|
||||
// the network filesystems that turn up in front of GPUs: NFS, Lustre, GPFS,
|
||||
// CephFS, GlusterFS and SMB. Anything unrecognised is treated as local, which
|
||||
// is the conservative direction — a shared filer mistaken for local storage
|
||||
// costs an unnecessary copy, while local storage mistaken for shared would have
|
||||
// the cluster believe a model is everywhere when it is on one node.
|
||||
func isNetworkFS(path string) bool {
|
||||
var st syscall.Statfs_t
|
||||
if err := syscall.Statfs(path, &st); err != nil {
|
||||
return false
|
||||
}
|
||||
switch uint32(st.Type) { //nolint:unconvert // Type is int64 on some platforms
|
||||
case 0x6969, // NFS
|
||||
0x0BD00BD0, // Lustre
|
||||
0x47504653, // GPFS
|
||||
0x00C36400, // CephFS
|
||||
0x65735546, // FUSE, which is how GlusterFS and many object mounts appear
|
||||
0xFF534D42: // SMB/CIFS
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func write(t *testing.T, path string, size int) {
|
||||
t.Helper()
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, make([]byte, size), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanFindsFilesAndDirectoriesOfWeights(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
write(t, filepath.Join(dir, "qwen-3b-q4.gguf"), 1024)
|
||||
// A HuggingFace checkout: a directory of shards rather than one file.
|
||||
write(t, filepath.Join(dir, "Llama-3.3-70B", "model-00001.safetensors"), 2048)
|
||||
write(t, filepath.Join(dir, "Llama-3.3-70B", "config.json"), 64)
|
||||
|
||||
got := ScanWeights([]string{dir})
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("found %d entries: %+v", len(got), got)
|
||||
}
|
||||
byName := map[string]int64{}
|
||||
for _, w := range got {
|
||||
byName[filepath.Base(w.Path)] = w.SizeBytes
|
||||
}
|
||||
if byName["qwen-3b-q4.gguf"] != 1024 {
|
||||
t.Errorf("gguf size = %d", byName["qwen-3b-q4.gguf"])
|
||||
}
|
||||
// The directory's size is everything in it, config included: what matters
|
||||
// is how much would have to be moved to put it somewhere else.
|
||||
if byName["Llama-3.3-70B"] != 2048+64 {
|
||||
t.Errorf("checkout size = %d, want %d", byName["Llama-3.3-70B"], 2048+64)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEmptyDirectoryIsNotWeights(t *testing.T) {
|
||||
// A half-finished download looks exactly like this, and reporting it as
|
||||
// present would have placement choose a node whose engine then fails to
|
||||
// load — the failure this whole feature exists to avoid.
|
||||
dir := t.TempDir()
|
||||
if err := os.MkdirAll(filepath.Join(dir, "downloading"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
write(t, filepath.Join(dir, "downloading", "README.md"), 10)
|
||||
|
||||
if got := ScanWeights([]string{dir}); len(got) != 0 {
|
||||
t.Errorf("a directory with no weight files was reported: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotesAndScriptsAreNotReportedAsModels(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
write(t, filepath.Join(dir, "notes.txt"), 10)
|
||||
write(t, filepath.Join(dir, "fetch.sh"), 10)
|
||||
write(t, filepath.Join(dir, "real.gguf"), 10)
|
||||
|
||||
got := ScanWeights([]string{dir})
|
||||
if len(got) != 1 || filepath.Base(got[0].Path) != "real.gguf" {
|
||||
t.Errorf("scan returned %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingDirectoryIsNotAnError(t *testing.T) {
|
||||
// A mount that has not come back yet looks exactly like a directory that is
|
||||
// not there, and the node is simply one without those weights until it is.
|
||||
got := ScanWeights([]string{filepath.Join(t.TempDir(), "never-existed")})
|
||||
if len(got) != 0 {
|
||||
t.Errorf("got %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanningNothingIsTheDefaultAndConstrainsNothing(t *testing.T) {
|
||||
// What makes the feature optional: an agent with no directories configured
|
||||
// behaves exactly as every agent did before this existed.
|
||||
if got := ScanWeights(nil); len(got) != 0 {
|
||||
t.Errorf("got %+v", got)
|
||||
}
|
||||
if got := ScanWeights([]string{""}); len(got) != 0 {
|
||||
t.Errorf("an empty directory name produced %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheSamePathIsReportedOnce(t *testing.T) {
|
||||
// Two configured directories can overlap, or be the same mount reached two
|
||||
// ways. Counting a model twice would overstate what a fabric holds.
|
||||
dir := t.TempDir()
|
||||
write(t, filepath.Join(dir, "m.gguf"), 10)
|
||||
if got := ScanWeights([]string{dir, dir}); len(got) != 1 {
|
||||
t.Errorf("the same directory twice produced %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalStorageIsNeverMistakenForShared(t *testing.T) {
|
||||
// The conservative direction. A shared filer read as local costs an
|
||||
// unnecessary copy; local read as shared would have the cluster believe a
|
||||
// model is everywhere when it is on one node.
|
||||
dir := t.TempDir()
|
||||
write(t, filepath.Join(dir, "m.gguf"), 10)
|
||||
got := ScanWeights([]string{dir})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("got %+v", got)
|
||||
}
|
||||
if got[0].Shared {
|
||||
t.Error("a temporary directory on local disk was reported as shared")
|
||||
}
|
||||
}
|
||||
@@ -141,6 +141,8 @@ func NewRegistry(cfg Config) *Registry {
|
||||
r := &Registry{engines: map[string]Engine{}}
|
||||
r.Add(&VLLM{cfg: cfg.VLLM})
|
||||
r.Add(&SGLang{cfg: cfg.SGLang})
|
||||
r.Add(&LlamaCpp{cfg: cfg.LlamaCpp})
|
||||
r.Add(&Sim{cfg: cfg.Sim})
|
||||
return r
|
||||
}
|
||||
|
||||
@@ -168,8 +170,10 @@ func (r *Registry) Names() []string {
|
||||
|
||||
// Config configures the engine adapters.
|
||||
type Config struct {
|
||||
VLLM VLLMConfig `koanf:"vllm"`
|
||||
SGLang SGLangConfig `koanf:"sglang"`
|
||||
VLLM VLLMConfig `koanf:"vllm"`
|
||||
SGLang SGLangConfig `koanf:"sglang"`
|
||||
LlamaCpp LlamaCppConfig `koanf:"llamacpp"`
|
||||
Sim SimConfig `koanf:"sim"`
|
||||
}
|
||||
|
||||
// --- shared helpers ---
|
||||
|
||||
@@ -397,14 +397,14 @@ func TestValidateRequiresModelPath(t *testing.T) {
|
||||
|
||||
func TestRegistry(t *testing.T) {
|
||||
r := NewRegistry(Config{})
|
||||
if got := r.Names(); len(got) != 2 || got[0] != "sglang" || got[1] != "vllm" {
|
||||
if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "sim", "vllm"}) {
|
||||
t.Errorf("registry contents: %v", got)
|
||||
}
|
||||
if _, err := r.Get("vllm"); err != nil {
|
||||
t.Errorf("vllm missing: %v", err)
|
||||
}
|
||||
err := func() error { _, err := r.Get("nope"); return err }()
|
||||
if err == nil || !strings.Contains(err.Error(), "sglang, vllm") {
|
||||
if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, vllm") {
|
||||
t.Errorf("an unknown engine should list what is available: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// LlamaCppConfig locates the llama.cpp installation.
|
||||
type LlamaCppConfig struct {
|
||||
// Bin is the llama-server binary. It is the server, not the CLI: the
|
||||
// gateway talks OpenAI over HTTP and needs the one that serves it.
|
||||
Bin string `koanf:"bin"`
|
||||
// Parallel is how many requests a server handles at once, when a profile
|
||||
// does not say. llama.cpp defaults to one, which would make every second
|
||||
// request wait behind the first no matter what the router decided.
|
||||
Parallel int `koanf:"parallel"`
|
||||
}
|
||||
|
||||
// LlamaCpp builds llama.cpp launch commands.
|
||||
//
|
||||
// It is the engine that makes a cluster without GPUs useful: llama.cpp runs a
|
||||
// small model on CPU at a speed that is fine for testing the whole path, and it
|
||||
// serves the same OpenAI API as the others, so nothing downstream can tell the
|
||||
// difference.
|
||||
type LlamaCpp struct{ cfg LlamaCppConfig }
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*LlamaCpp) Name() string { return "llamacpp" }
|
||||
|
||||
// HealthPath is llama.cpp's readiness endpoint. It returns 503 while the model
|
||||
// is still loading, which is exactly what the supervisor needs.
|
||||
func (*LlamaCpp) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is llama.cpp's Prometheus endpoint, which only exists when the
|
||||
// server was started with --metrics. Commands always passes it.
|
||||
func (*LlamaCpp) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API. Every llama.cpp shard
|
||||
// does, because a llama.cpp replica is always exactly one shard.
|
||||
func (*LlamaCpp) ServesHTTP(ShardSpec) bool { return true }
|
||||
|
||||
// Validate checks the profile against the GPUs it would run on.
|
||||
//
|
||||
// The two refusals here are about not pretending: llama.cpp has no equivalent
|
||||
// of a tensor-parallel group spanning machines, and its multi-GPU support is a
|
||||
// different mechanism (--split-mode, --tensor-split) that this adapter does not
|
||||
// drive. Accepting either would produce a replica that starts and then serves
|
||||
// from one machine while the catalog believes it has several.
|
||||
func (*LlamaCpp) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if err := validateCommon(p, gpus); err != nil {
|
||||
return err
|
||||
}
|
||||
if p.Nodes > 1 {
|
||||
return fmt.Errorf("profile %q: llamacpp does not span nodes; use vllm or sglang for a %d-node replica",
|
||||
p.Name, p.Nodes)
|
||||
}
|
||||
if p.TP > 1 {
|
||||
return fmt.Errorf("profile %q: llamacpp does not do tensor parallelism; "+
|
||||
"it splits a model across local GPUs with --tensor-split, which this adapter does not drive",
|
||||
p.Name)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Commands returns the processes for one shard, which is always one process.
|
||||
func (l *LlamaCpp) Commands(s ShardSpec) ([]Command, error) {
|
||||
bin := l.cfg.Bin
|
||||
if bin == "" {
|
||||
bin = "llama-server"
|
||||
}
|
||||
p := s.Profile
|
||||
|
||||
args := []string{
|
||||
"--model", p.ModelPath,
|
||||
"--host", "0.0.0.0",
|
||||
"--port", strconv.Itoa(s.HTTPPort),
|
||||
// Without this there is no /metrics, and an endpoint the router cannot
|
||||
// measure is one it has to route to blind.
|
||||
"--metrics",
|
||||
}
|
||||
// The model's own chat template rather than a generic one, which is what
|
||||
// renders tool definitions in the form the model was trained to read.
|
||||
//
|
||||
// Pinned rather than relied on: build 10991 has it on by default, older
|
||||
// builds do not, and an engine upgrade should not be able to change how
|
||||
// prompts are built underneath a running fleet. Passing it explicitly was
|
||||
// a no-op on the build here — it was added while chasing tool calls that
|
||||
// came back as text, and did not fix that, which turned out to be the
|
||||
// model rather than the template.
|
||||
if !hasFlag(p.Args, "--jinja", "--no-jinja") {
|
||||
args = append(args, "--jinja")
|
||||
}
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--alias", p.ServedModelName)
|
||||
}
|
||||
if !hasFlag(p.Args, "--parallel", "-np") {
|
||||
n := l.cfg.Parallel
|
||||
if n <= 0 {
|
||||
n = 4
|
||||
}
|
||||
args = append(args, "--parallel", strconv.Itoa(n))
|
||||
}
|
||||
// With GPUs assigned, offload everything; llama.cpp keeps the model on the
|
||||
// CPU otherwise, which is not what a GPU placement asked for. Without them
|
||||
// the flag is omitted entirely rather than set to zero, so a build with no
|
||||
// GPU support does not see an argument it rejects.
|
||||
if len(s.GPUIndices) > 0 && !hasFlag(p.Args, "--n-gpu-layers", "-ngl") {
|
||||
args = append(args, "--n-gpu-layers", "999")
|
||||
}
|
||||
args = append(args, p.Args...)
|
||||
|
||||
return []Command{{
|
||||
Name: "llama-server", Path: bin, Args: args, Env: materializeEnv(baseEnv(s)),
|
||||
}}, nil
|
||||
}
|
||||
|
||||
// hasFlag reports whether the profile already sets one of these flags, in
|
||||
// either "--flag value" or "--flag=value" form.
|
||||
func hasFlag(args []string, names ...string) bool {
|
||||
for _, a := range args {
|
||||
for _, n := range names {
|
||||
if a == n || strings.HasPrefix(a, n+"=") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func llamaSpec(gpus []int) ShardSpec {
|
||||
return ShardSpec{
|
||||
ReplicaID: "rep-1", ShardID: "sh-1", Rank: 0, NNodes: 1,
|
||||
GPUIndices: gpus, HTTPPort: 8000, MgmtInterface: "eth0",
|
||||
Profile: Profile{
|
||||
Name: "cpu", Engine: "llamacpp", TP: 1,
|
||||
ModelPath: "/models/qwen.gguf", ServedModelName: "qwen-0.5b",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppCommand(t *testing.T) {
|
||||
cmds, err := (&LlamaCpp{}).Commands(llamaSpec(nil))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("a llama.cpp shard is one process, got %d", len(cmds))
|
||||
}
|
||||
c := cmds[0]
|
||||
if c.Path != "llama-server" {
|
||||
t.Errorf("path %q", c.Path)
|
||||
}
|
||||
|
||||
// --metrics is not optional: llama-server serves no /metrics without it, and
|
||||
// an endpoint the router cannot measure is one it routes to blind.
|
||||
if !slices.Contains(c.Args, "--metrics") {
|
||||
t.Error("the server was started without --metrics")
|
||||
}
|
||||
if v, _ := flagValue(c.Args, "--model"); v != "/models/qwen.gguf" {
|
||||
t.Errorf("--model %q", v)
|
||||
}
|
||||
if v, _ := flagValue(c.Args, "--port"); v != "8000" {
|
||||
t.Errorf("--port %q", v)
|
||||
}
|
||||
if v, _ := flagValue(c.Args, "--alias"); v != "qwen-0.5b" {
|
||||
t.Errorf("--alias %q", v)
|
||||
}
|
||||
// llama.cpp serves one request at a time unless told otherwise, which would
|
||||
// make every second request wait no matter what the router decided.
|
||||
if v, ok := flagValue(c.Args, "--parallel"); !ok || v == "1" {
|
||||
t.Errorf("--parallel %q", v)
|
||||
}
|
||||
if slices.Contains(c.Args, "--n-gpu-layers") {
|
||||
t.Error("a CPU placement was given GPU layers")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppOffloadsWhenGivenGPUs(t *testing.T) {
|
||||
cmds, _ := (&LlamaCpp{}).Commands(llamaSpec([]int{2}))
|
||||
if v, ok := flagValue(cmds[0].Args, "--n-gpu-layers"); !ok || v != "999" {
|
||||
t.Errorf("a GPU placement should offload the model, got %q", v)
|
||||
}
|
||||
// The shard must still be pinned to its own card.
|
||||
if !slices.Contains(cmds[0].Env, "CUDA_VISIBLE_DEVICES=2") {
|
||||
t.Error("the shard was not pinned to its GPU")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileArgumentsWinOverDefaults(t *testing.T) {
|
||||
// The escape hatch has to actually work, or an operator tuning slot count
|
||||
// ends up with two --parallel flags and whichever llama.cpp happens to
|
||||
// prefer.
|
||||
s := llamaSpec(nil)
|
||||
s.Profile.Args = []string{"--parallel", "16", "--ctx-size", "4096"}
|
||||
|
||||
cmds, _ := (&LlamaCpp{}).Commands(s)
|
||||
n := 0
|
||||
for _, a := range cmds[0].Args {
|
||||
if a == "--parallel" {
|
||||
n++
|
||||
}
|
||||
}
|
||||
if n != 1 {
|
||||
t.Errorf("--parallel appears %d times", n)
|
||||
}
|
||||
if v, _ := flagValue(cmds[0].Args, "--parallel"); v != "16" {
|
||||
t.Errorf("the profile's value lost: %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShortFormArgumentsAreRecognized(t *testing.T) {
|
||||
s := llamaSpec([]int{0})
|
||||
s.Profile.Args = []string{"-np", "8", "-ngl", "20"}
|
||||
|
||||
cmds, _ := (&LlamaCpp{}).Commands(s)
|
||||
for _, flag := range []string{"--parallel", "--n-gpu-layers"} {
|
||||
if slices.Contains(cmds[0].Args, flag) {
|
||||
t.Errorf("%s was added although the profile set its short form", flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppRefusesWhatItCannotDo(t *testing.T) {
|
||||
// Accepting either of these would produce a replica that starts and then
|
||||
// serves from one machine while the catalog believes it has several.
|
||||
multi := Profile{Name: "multi", Engine: "llamacpp", TP: 1, Nodes: 2, GPUsPerNode: 1, ModelPath: "/m"}
|
||||
if err := (&LlamaCpp{}).Validate(multi, nil); err == nil {
|
||||
t.Error("a multi-node llama.cpp profile was accepted")
|
||||
}
|
||||
|
||||
tp := Profile{Name: "tp", Engine: "llamacpp", TP: 2, GPUsPerNode: 2, ModelPath: "/m"}
|
||||
if err := (&LlamaCpp{}).Validate(tp, h100(2)); err == nil {
|
||||
t.Error("a tensor-parallel llama.cpp profile was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCPUProfileValidates(t *testing.T) {
|
||||
p := Profile{Name: "cpu", Engine: "llamacpp", TP: 1, ModelPath: "/models/qwen.gguf"}
|
||||
if err := (&LlamaCpp{}).Validate(p, nil); err != nil {
|
||||
t.Errorf("a CPU profile on a node with no GPUs was rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppAlwaysServesHTTP(t *testing.T) {
|
||||
// Unlike the others there is no worker rank to get this wrong for.
|
||||
if !(&LlamaCpp{}).ServesHTTP(llamaSpec(nil)) {
|
||||
t.Error("a llama.cpp shard must answer the API")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBinaryIsConfigurable(t *testing.T) {
|
||||
l := &LlamaCpp{cfg: LlamaCppConfig{Bin: "/opt/llama.cpp/bin/llama-server", Parallel: 2}}
|
||||
cmds, _ := l.Commands(llamaSpec(nil))
|
||||
if cmds[0].Path != "/opt/llama.cpp/bin/llama-server" {
|
||||
t.Errorf("path %q", cmds[0].Path)
|
||||
}
|
||||
if v, _ := flagValue(cmds[0].Args, "--parallel"); v != "2" {
|
||||
t.Errorf("--parallel %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnknownEngineNamesLlamaCpp(t *testing.T) {
|
||||
_, err := NewRegistry(Config{}).Get("llama.cpp")
|
||||
if err == nil || !strings.Contains(err.Error(), "llamacpp") {
|
||||
t.Errorf("the suggestion should name the real engine: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// flagValue reads a "--flag value" pair out of a command line.
|
||||
func flagValue(args []string, flag string) (string, bool) {
|
||||
for i, a := range args {
|
||||
if a == flag && i+1 < len(args) {
|
||||
return args[i+1], true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func TestLlamaCppUsesTheModelsChatTemplate(t *testing.T) {
|
||||
// Pinned so an engine upgrade cannot change how prompts are built under a
|
||||
// running fleet: build 10991 defaults to jinja, older builds do not.
|
||||
//
|
||||
// Not a fix for tool calling, though it was added while chasing that. On
|
||||
// this build the flag is a no-op and the behaviour is identical with and
|
||||
// without it.
|
||||
l := &LlamaCpp{}
|
||||
cmds, err := l.Commands(ShardSpec{HTTPPort: 8080, Profile: Profile{ModelPath: "/m.gguf"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !slices.Contains(cmds[0].Args, "--jinja") {
|
||||
t.Errorf("--jinja is missing; tool calls will come back as text: %v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppLetsAProfileTurnTheTemplateOff(t *testing.T) {
|
||||
// A model whose bundled template is broken has to be serviceable without
|
||||
// editing the adapter.
|
||||
l := &LlamaCpp{}
|
||||
cmds, err := l.Commands(ShardSpec{
|
||||
HTTPPort: 8080,
|
||||
Profile: Profile{ModelPath: "/m.gguf", Args: []string{"--no-jinja"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if slices.Contains(cmds[0].Args, "--jinja") {
|
||||
t.Errorf("the profile said --no-jinja and got both: %v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// SimConfig locates the simulated engine.
|
||||
type SimConfig struct {
|
||||
// Bin is the llmgw binary itself; the engine is one of its subcommands.
|
||||
// Empty means whatever started this agent, which is right almost always.
|
||||
Bin string `koanf:"bin"`
|
||||
}
|
||||
|
||||
// Sim launches llmgw's own simulated engine.
|
||||
//
|
||||
// It is in the production registry on purpose. A cluster can be brought up and
|
||||
// exercised end to end — placement, fabrics, multi-node sequencing, failure and
|
||||
// rebuild — on machines that have no GPUs, which is the only way most of that
|
||||
// path gets tested at all. The multi-node launch sequence went four phases
|
||||
// without ever running, because the only things exercising it were stub
|
||||
// processes that exercised the command line and nothing else.
|
||||
//
|
||||
// Nothing about it pretends to be real: the model name it serves is whatever it
|
||||
// is told, the tokens are invented, and it says so in its own logs.
|
||||
type Sim struct{ cfg SimConfig }
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*Sim) Name() string { return "sim" }
|
||||
|
||||
// HealthPath is the simulated engine's readiness endpoint.
|
||||
func (*Sim) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is its Prometheus endpoint, in vLLM's vocabulary.
|
||||
func (*Sim) MetricsPath() string { return "/metrics" }
|
||||
|
||||
// ServesHTTP reports whether this rank answers the API.
|
||||
//
|
||||
// Only the leader, exactly as vLLM and SGLang behave — which is the behaviour
|
||||
// this engine exists to let us test.
|
||||
func (*Sim) ServesHTTP(s ShardSpec) bool { return s.Leader() }
|
||||
|
||||
// Validate accepts anything the arithmetic allows.
|
||||
//
|
||||
// No compute-capability checks and no quantization rules: there is no hardware
|
||||
// to be wrong about, and inventing refusals would make this a worse stand-in
|
||||
// for the engines it stands in for.
|
||||
func (*Sim) Validate(p Profile, gpus []GPUInfo) error {
|
||||
return validateCommon(p, gpus)
|
||||
}
|
||||
|
||||
// Commands returns the single process that makes up one shard.
|
||||
func (s *Sim) Commands(spec ShardSpec) ([]Command, error) {
|
||||
bin := s.cfg.Bin
|
||||
if bin == "" {
|
||||
exe, err := os.Executable()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("locating the llmgw binary for the sim engine: %w", err)
|
||||
}
|
||||
bin = exe
|
||||
}
|
||||
p := spec.Profile
|
||||
|
||||
args := []string{"sim-engine",
|
||||
"--listen", "0.0.0.0:" + strconv.Itoa(spec.HTTPPort),
|
||||
"--model", p.ServedModelName,
|
||||
}
|
||||
if spec.NNodes > 1 {
|
||||
args = append(args,
|
||||
"--nnodes", strconv.Itoa(spec.NNodes),
|
||||
"--node-rank", strconv.Itoa(spec.Rank),
|
||||
"--dist-init-addr", spec.LeaderAddr+":"+strconv.Itoa(spec.LeaderPort),
|
||||
)
|
||||
}
|
||||
args = append(args, p.Args...)
|
||||
|
||||
return []Command{{
|
||||
Name: "sim-engine", Path: bin, Args: args,
|
||||
Env: materializeEnv(baseEnv(spec)),
|
||||
// Workers dial the leader, so the leader has to be listening first.
|
||||
// Both ranks express the same wait because both know the same address.
|
||||
WaitForLeaderPort: spec.NNodes > 1 && !spec.Leader(),
|
||||
}}, nil
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func simSpec(rank, nnodes int) ShardSpec {
|
||||
return ShardSpec{
|
||||
ReplicaID: "rep-1", ShardID: "sh-1", Rank: rank, NNodes: nnodes,
|
||||
LeaderAddr: "10.0.0.5", LeaderPort: 20001,
|
||||
GPUIndices: []int{0, 1}, HTTPPort: 20000, MgmtInterface: "eth0",
|
||||
HCAs: []string{"mlx5_0", "mlx5_1"},
|
||||
Profile: Profile{
|
||||
Name: "tp4", Engine: "sim", TP: 4, GPUsPerNode: 2, Nodes: nnodes,
|
||||
ModelPath: "/models/x", ServedModelName: "big-model",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyTheLeaderServes(t *testing.T) {
|
||||
// The behaviour that made replicas fail while they were coming up perfectly
|
||||
// well: health checking a worker's port, which never opens.
|
||||
if !(&Sim{}).ServesHTTP(simSpec(0, 2)) {
|
||||
t.Error("the leader does not serve")
|
||||
}
|
||||
if (&Sim{}).ServesHTTP(simSpec(1, 2)) {
|
||||
t.Error("a worker claims to serve")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMultiNodeShardCarriesItsRendezvous(t *testing.T) {
|
||||
cmds, err := (&Sim{cfg: SimConfig{Bin: "/usr/local/bin/llmgw"}}).Commands(simSpec(1, 2))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(cmds) != 1 {
|
||||
t.Fatalf("%d commands", len(cmds))
|
||||
}
|
||||
c := cmds[0]
|
||||
|
||||
if c.Args[0] != "sim-engine" {
|
||||
t.Errorf("the engine is a subcommand of llmgw: %v", c.Args)
|
||||
}
|
||||
for flag, want := range map[string]string{
|
||||
"--nnodes": "2", "--node-rank": "1", "--dist-init-addr": "10.0.0.5:20001",
|
||||
} {
|
||||
if got, _ := flagValue(c.Args, flag); got != want {
|
||||
t.Errorf("%s = %q, want %q", flag, got, want)
|
||||
}
|
||||
}
|
||||
// A worker must wait for the leader's port, or the launch is a race.
|
||||
if !c.WaitForLeaderPort {
|
||||
t.Error("a worker does not wait for the leader")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheLeaderDoesNotWaitForItself(t *testing.T) {
|
||||
cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(0, 2))
|
||||
if cmds[0].WaitForLeaderPort {
|
||||
t.Error("the leader waits for a port it is about to open")
|
||||
}
|
||||
}
|
||||
|
||||
func TestASingleNodeShardCarriesNoRendezvous(t *testing.T) {
|
||||
cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(0, 1))
|
||||
for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} {
|
||||
if slices.Contains(cmds[0].Args, flag) {
|
||||
t.Errorf("a single-node shard was given %s", flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheShardIsStillPinnedToItsFabric(t *testing.T) {
|
||||
// The simulated engine does not use NCCL, but the environment is generated
|
||||
// by the same code that generates it for the ones that do — so a cluster
|
||||
// exercised with this engine is exercising that too.
|
||||
cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(1, 2))
|
||||
env := cmds[0].Env
|
||||
for _, want := range []string{
|
||||
"NCCL_IB_HCA=mlx5_0,mlx5_1", "NCCL_SOCKET_IFNAME=eth0",
|
||||
"CUDA_VISIBLE_DEVICES=0,1", "NCCL_IB_DISABLE=0",
|
||||
} {
|
||||
if !slices.Contains(env, want) {
|
||||
t.Errorf("the environment is missing %s", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,317 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/http2/h2c"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
// pki issues a cluster's certificates into a temporary directory and returns
|
||||
// the configurations the control plane and one agent would load.
|
||||
func pki(t *testing.T, nodeID string) (server, client tlsconf.Config) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
|
||||
ca, err := tlsconf.NewCA("test CA")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := ca.Write(dir); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
srv, err := ca.IssueServer("control-plane", []string{"127.0.0.1"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := srv.Write(dir, "server"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cli, err := ca.IssueClient(nodeID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := cli.Write(dir, "agent"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
caPath := filepath.Join(dir, "ca.pem")
|
||||
return tlsconf.Config{CA: caPath, Cert: filepath.Join(dir, "server.pem"), Key: filepath.Join(dir, "server-key.pem")},
|
||||
tlsconf.Config{CA: caPath, Cert: filepath.Join(dir, "agent.pem"), Key: filepath.Join(dir, "agent-key.pem")}
|
||||
}
|
||||
|
||||
// startTLS runs a real service behind a real TLS listener, with identity
|
||||
// required — the deployed configuration rather than an approximation of it.
|
||||
func startTLS(t *testing.T, serverTLS tlsconf.Config) (*agentsvc.Service, *inventory.Store, string) {
|
||||
t.Helper()
|
||||
|
||||
inv := inventory.NewStore()
|
||||
svc := agentsvc.New(inv, newRecorder(), true, slog.New(slog.DiscardHandler))
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
|
||||
tc, err := serverTLS.ServerConfig()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
srv := httptest.NewUnstartedServer(mux)
|
||||
srv.Listener = raw
|
||||
srv.TLS = tc
|
||||
srv.StartTLS()
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
return svc, inv, "https://" + raw.Addr().String()
|
||||
}
|
||||
|
||||
func runAgent(t *testing.T, cfg Config) *Agent {
|
||||
t.Helper()
|
||||
|
||||
cfg.MgmtAddr = "127.0.0.1"
|
||||
cfg.HeartbeatInterval = 50 * time.Millisecond
|
||||
cfg.ReconnectBase = 20 * time.Millisecond
|
||||
cfg.ReconnectMax = 100 * time.Millisecond
|
||||
cfg.Discovery.SkipGPUs = true
|
||||
cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
|
||||
cfg.Superv = supervisor.Config{
|
||||
WorkDir: t.TempDir(), LogLines: 64,
|
||||
StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond,
|
||||
}
|
||||
|
||||
client, err := Client(cfg.TLS, 5*time.Second)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ag, err := New(cfg, client, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() { _ = ag.Run(ctx) }()
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
ag.Supervisor().StopAll(context.Background())
|
||||
})
|
||||
return ag
|
||||
}
|
||||
|
||||
func registered(svc *agentsvc.Service, inv *inventory.Store, nodeID string) func() bool {
|
||||
return func() bool {
|
||||
_, ok := inv.Get(nodeID)
|
||||
return ok && len(svc.ConnectedNodes()) == 1
|
||||
}
|
||||
}
|
||||
|
||||
func eventually(t *testing.T, limit time.Duration, pred func() bool) bool {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(limit)
|
||||
for time.Now().Before(deadline) {
|
||||
if pred() {
|
||||
return true
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func TestAgentRegistersOverMutualTLS(t *testing.T) {
|
||||
serverTLS, clientTLS := pki(t, "llmgw-node01")
|
||||
svc, inv, url := startTLS(t, serverTLS)
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "llmgw-node01"
|
||||
cfg.Server = url
|
||||
cfg.TLS = clientTLS
|
||||
runAgent(t, cfg)
|
||||
|
||||
if !eventually(t, 10*time.Second, registered(svc, inv, "llmgw-node01")) {
|
||||
t.Fatal("the agent never registered over mutual TLS")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodeCannotClaimAnotherNodesIdentity(t *testing.T) {
|
||||
// The point of the whole exercise. node01 holds a valid certificate — it is
|
||||
// a legitimate member of the cluster — and still cannot register as node02,
|
||||
// because the id it claims is checked against the one it can prove.
|
||||
serverTLS, clientTLS := pki(t, "llmgw-node01")
|
||||
svc, inv, url := startTLS(t, serverTLS)
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "llmgw-node02"
|
||||
cfg.Server = url
|
||||
cfg.TLS = clientTLS
|
||||
runAgent(t, cfg)
|
||||
|
||||
// The agent reconnects on failure, so this has to stay false throughout
|
||||
// rather than merely start false.
|
||||
if eventually(t, 2*time.Second, func() bool {
|
||||
_, ok := inv.Get("llmgw-node02")
|
||||
return ok || len(svc.ConnectedNodes()) > 0
|
||||
}) {
|
||||
t.Fatal("an agent registered under a node id its certificate does not authorize")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnauthenticatedAgentIsRefused(t *testing.T) {
|
||||
// The defensive half of the check: identity is required, but the connection
|
||||
// carried none. In the shipped wiring TLS and this flag are tied together,
|
||||
// so reaching here means someone untied them — and the answer must still be
|
||||
// refusal rather than a silently trusted node id.
|
||||
inv := inventory.NewStore()
|
||||
svc := agentsvc.New(inv, newRecorder(), true, slog.New(slog.DiscardHandler))
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
|
||||
srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "llmgw-node01"
|
||||
cfg.Server = srv.URL
|
||||
runAgent(t, cfg)
|
||||
|
||||
if eventually(t, 2*time.Second, registered(svc, inv, "llmgw-node01")) {
|
||||
t.Fatal("an agent with no proven identity reached the control plane")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlaintextAgentIsRejectedWhenTLSIsConfigured(t *testing.T) {
|
||||
// A misconfiguration worth catching at startup rather than at runtime: an
|
||||
// agent holding certificates but dialling http:// would connect in the
|
||||
// clear and never learn that it had.
|
||||
_, clientTLS := pki(t, "llmgw-node01")
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "llmgw-node01"
|
||||
cfg.Server = "http://127.0.0.1:9443"
|
||||
cfg.TLS = clientTLS
|
||||
|
||||
if _, err := New(cfg, nil, slog.New(slog.DiscardHandler)); err == nil {
|
||||
t.Fatal("an agent with certificates was allowed to dial a plaintext address")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBareAddressGainsTheRightScheme(t *testing.T) {
|
||||
_, clientTLS := pki(t, "llmgw-node01")
|
||||
|
||||
secure := Config{NodeID: "n", Server: "10.0.0.140:9443", TLS: clientTLS}
|
||||
if err := secure.normalize(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if secure.Server != "https://10.0.0.140:9443" {
|
||||
t.Errorf("a configured agent dials %q", secure.Server)
|
||||
}
|
||||
|
||||
plain := Config{NodeID: "n", Server: "10.0.0.140:9443"}
|
||||
if err := plain.normalize(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plain.Server != "http://10.0.0.140:9443" {
|
||||
t.Errorf("an unconfigured agent dials %q", plain.Server)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefusedAgentLearnsWhy(t *testing.T) {
|
||||
// An agent that cannot register must be able to say why. Reporting a
|
||||
// refusal as a clean disconnect leaves an operator with a silent reconnect
|
||||
// loop and nothing to act on — and resets the backoff, so the loop is a
|
||||
// fast one.
|
||||
serverTLS, clientTLS := pki(t, "llmgw-node01")
|
||||
_, _, url := startTLS(t, serverTLS)
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "llmgw-node02"
|
||||
cfg.Server = url
|
||||
cfg.TLS = clientTLS
|
||||
cfg.Discovery.SkipGPUs = true
|
||||
cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
|
||||
cfg.Superv = supervisor.Config{WorkDir: t.TempDir(), LogLines: 8}
|
||||
|
||||
client, err := Client(cfg.TLS, 5*time.Second)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ag, err := New(cfg, client, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
err = ag.session(context.Background())
|
||||
if err == nil {
|
||||
t.Fatal("a refused session reported success")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "llmgw-node01") || !strings.Contains(err.Error(), "llmgw-node02") {
|
||||
t.Errorf("the error names neither certificate nor claim: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnAgentWithNothingRunningStillTellsTheControlPlane covers the gap that
|
||||
// made a reconciler fix dead code in production.
|
||||
//
|
||||
// "I am running nothing" is information: it is how the control plane learns
|
||||
// that an engine died while its agent survived. The service used to return
|
||||
// early on an empty list, so that message never reached the reconciler — while
|
||||
// the reconciler's own tests, which call it directly, kept passing.
|
||||
//
|
||||
// This test goes through the real stream for that reason.
|
||||
func TestAnAgentWithNothingRunningStillTellsTheControlPlane(t *testing.T) {
|
||||
serverTLS, clientTLS := pki(t, "llmgw-node01")
|
||||
|
||||
inv := inventory.NewStore()
|
||||
rec := newRecorder()
|
||||
svc := agentsvc.New(inv, rec, true, slog.New(slog.DiscardHandler))
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
|
||||
tc, err := serverTLS.ServerConfig()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
srv := httptest.NewUnstartedServer(mux)
|
||||
srv.Listener = raw
|
||||
srv.TLS = tc
|
||||
srv.StartTLS()
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "llmgw-node01"
|
||||
cfg.Server = "https://" + raw.Addr().String()
|
||||
cfg.TLS = clientTLS
|
||||
runAgent(t, cfg)
|
||||
|
||||
// The agent has no shards at all, which is the whole point.
|
||||
if !eventually(t, 10*time.Second, func() bool {
|
||||
return len(rec.adoptCalls()) > 0
|
||||
}) {
|
||||
t.Fatal("an agent running nothing never reached the reconciler")
|
||||
}
|
||||
for _, call := range rec.adoptCalls() {
|
||||
if call.nodeID != "llmgw-node01" {
|
||||
t.Errorf("adopt was told node %q", call.nodeID)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
)
|
||||
|
||||
// agentWithCache is an agent that only knows where its weights live, which is
|
||||
// all removeWeights needs.
|
||||
func agentWithCache(t *testing.T, dir string) *Agent {
|
||||
t.Helper()
|
||||
return &Agent{
|
||||
cfg: Config{ModelDirs: []string{dir}},
|
||||
log: slog.New(slog.DiscardHandler),
|
||||
}
|
||||
}
|
||||
|
||||
func write(t *testing.T, path string) string {
|
||||
t.Helper()
|
||||
if err := os.WriteFile(path, []byte("weights"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
func TestTheAgentRefusesToDeleteWhatItDidNotFetch(t *testing.T) {
|
||||
// The control plane computes these paths, but the node holds the files. A
|
||||
// control plane that is wrong, newer than this agent, or tampered with must
|
||||
// not be able to talk it into deleting a model an operator spent a week
|
||||
// copying here. So the agent checks again, on its own.
|
||||
dir := t.TempDir()
|
||||
a := agentWithCache(t, dir)
|
||||
|
||||
operators := write(t, filepath.Join(dir, "qwen2.5-coder-14b-instruct-q8_0.gguf"))
|
||||
fetched := write(t, weights.CachePathFor(dir, "https://models.internal/gone.gguf"))
|
||||
|
||||
a.removeWeights([]string{operators, fetched})
|
||||
|
||||
if _, err := os.Stat(operators); err != nil {
|
||||
t.Errorf("a file llmgw never fetched was deleted: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(fetched); !os.IsNotExist(err) {
|
||||
t.Errorf("the fetched file was not removed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheAgentRefusesPathsOutsideItsCache(t *testing.T) {
|
||||
// Even with a name that looks like one of ours. The cache directory is the
|
||||
// boundary, and traversal out of it is the attack this closes.
|
||||
dir := t.TempDir()
|
||||
elsewhere := t.TempDir()
|
||||
a := agentWithCache(t, dir)
|
||||
|
||||
// A correctly named cache file, in the wrong directory.
|
||||
outside := write(t, weights.CachePathFor(elsewhere, "https://models.internal/x.gguf"))
|
||||
// And the same thing reached by traversal from inside the cache.
|
||||
traversal := filepath.Join(dir, "..", filepath.Base(elsewhere),
|
||||
filepath.Base(outside))
|
||||
|
||||
a.removeWeights([]string{outside, traversal})
|
||||
|
||||
if _, err := os.Stat(outside); err != nil {
|
||||
t.Errorf("a file outside the cache was deleted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemovingAFileThatIsAlreadyGoneIsNotAFailure(t *testing.T) {
|
||||
// The expected outcome of a retry. It must not stop the rest of the batch.
|
||||
dir := t.TempDir()
|
||||
a := agentWithCache(t, dir)
|
||||
|
||||
missing := weights.CachePathFor(dir, "https://models.internal/never.gguf")
|
||||
present := write(t, weights.CachePathFor(dir, "https://models.internal/here.gguf"))
|
||||
|
||||
a.removeWeights([]string{missing, present})
|
||||
|
||||
if _, err := os.Stat(present); !os.IsNotExist(err) {
|
||||
t.Errorf("a missing file stopped the batch: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithNoCacheConfiguredNothingIsDeleted(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
a := &Agent{cfg: Config{}, log: slog.New(slog.DiscardHandler)}
|
||||
kept := write(t, weights.CachePathFor(dir, "https://models.internal/x.gguf"))
|
||||
|
||||
a.removeWeights([]string{kept})
|
||||
|
||||
if _, err := os.Stat(kept); err != nil {
|
||||
t.Errorf("an agent with no cache directory deleted something: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
)
|
||||
|
||||
// TestAgentStopsWhenContextEnds pins down something cancellation does not do by
|
||||
// itself. A bidirectional stream spends its life blocked reading, and HTTP/2
|
||||
// stops watching the context once the request headers are away, so an agent
|
||||
// that only cancels keeps reading until something kills it. Under systemd that
|
||||
// is the full stop timeout on every restart, followed by SIGKILL.
|
||||
func TestAgentStopsWhenContextEnds(t *testing.T) {
|
||||
serverTLS, clientTLS := pki(t, "llmgw-node01")
|
||||
svc, inv, url := startTLS(t, serverTLS)
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "llmgw-node01"
|
||||
cfg.Server = url
|
||||
cfg.TLS = clientTLS
|
||||
cfg.HeartbeatInterval = 50 * time.Millisecond
|
||||
cfg.Discovery.SkipGPUs = true
|
||||
cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
|
||||
cfg.Superv = supervisor.Config{
|
||||
WorkDir: t.TempDir(), LogLines: 8,
|
||||
StopTimeout: time.Second, HealthInterval: 20 * time.Millisecond,
|
||||
}
|
||||
|
||||
client, err := Client(cfg.TLS, 5*time.Second)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ag, err := New(cfg, client, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan struct{})
|
||||
go func() { _ = ag.Run(ctx); close(done) }()
|
||||
|
||||
// Cancelling before the stream is established would prove nothing: the
|
||||
// blocked read is the case that matters.
|
||||
if !eventually(t, 10*time.Second, registered(svc, inv, "llmgw-node01")) {
|
||||
t.Fatal("the agent never registered")
|
||||
}
|
||||
|
||||
cancel()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("the agent did not stop within 5s of being told to")
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,9 @@
|
||||
package supervisor
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"os"
|
||||
"strconv"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
@@ -22,13 +24,42 @@ func killGroup(pid int, sig syscall.Signal) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// processAlive reports whether a pid exists.
|
||||
// processAlive reports whether a pid belongs to a process that is still doing
|
||||
// something.
|
||||
//
|
||||
// A zombie is not. It has exited and is only waiting to be reaped, but
|
||||
// kill(pid, 0) still succeeds on it — so a naive check reports a dead engine as
|
||||
// running forever. That happens for real whenever an engine's parent dies
|
||||
// before it does and nothing reaps the orphan: in a container, PID 1 is usually
|
||||
// the workload rather than an init, and orphans simply accumulate.
|
||||
func processAlive(pid int) bool {
|
||||
if pid <= 0 {
|
||||
return false
|
||||
}
|
||||
// Signal 0 performs the permission and existence checks without delivering
|
||||
// anything.
|
||||
err := syscall.Kill(pid, 0)
|
||||
return err == nil || err == syscall.EPERM
|
||||
if err := syscall.Kill(pid, 0); err != nil {
|
||||
// EPERM means it exists but belongs to someone else, which for our
|
||||
// purposes is still "there".
|
||||
return err == syscall.EPERM
|
||||
}
|
||||
return !isZombie(pid)
|
||||
}
|
||||
|
||||
// isZombie reports whether a pid has exited but not been reaped.
|
||||
//
|
||||
// Only Linux exposes this; elsewhere there is no procfs and the answer is
|
||||
// false, which degrades to the old behaviour rather than to something worse.
|
||||
func isZombie(pid int) bool {
|
||||
b, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat")
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
// The comm field is parenthesised and may itself contain spaces or
|
||||
// parentheses, so the state character is found from the last ')'.
|
||||
i := bytes.LastIndexByte(b, ')')
|
||||
if i < 0 || i+2 >= len(b) {
|
||||
return false
|
||||
}
|
||||
return b[i+2] == 'Z'
|
||||
}
|
||||
@@ -557,15 +557,27 @@ func TestSidecarRunsBeforeEngine(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if b, err := os.ReadFile(order); err == nil && strings.Count(string(b), "\n") == 2 {
|
||||
if got := strings.Fields(string(b)); got[0] != "sidecar" || got[1] != "engine" {
|
||||
t.Fatalf("start order was %v", got)
|
||||
}
|
||||
break
|
||||
// Generous, because this waits on two processes starting under the race
|
||||
// detector while every other package is competing for the same cores. Three
|
||||
// seconds was enough alone and not in a full run.
|
||||
var started []string
|
||||
deadline := time.Now().Add(20 * time.Second)
|
||||
for time.Now().Before(deadline) && len(started) < 2 {
|
||||
if b, err := os.ReadFile(order); err == nil {
|
||||
started = strings.Fields(string(b))
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
if len(started) < 2 {
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
// Asserted rather than broken out of: the original fell through on timeout,
|
||||
// so a slow machine turned the ordering check into a no-op and the failure
|
||||
// surfaced later as something else.
|
||||
if len(started) != 2 {
|
||||
t.Fatalf("only %d of 2 processes started: %v", len(started), started)
|
||||
}
|
||||
if started[0] != "sidecar" || started[1] != "engine" {
|
||||
t.Fatalf("start order was %v", started)
|
||||
}
|
||||
|
||||
// The shard's pid is the engine's, not the sidecar's: the engine's exit is
|
||||
@@ -574,7 +586,7 @@ func TestSidecarRunsBeforeEngine(t *testing.T) {
|
||||
t.Fatal("no engine pid recorded")
|
||||
}
|
||||
// Both processes appear in one log tail, distinguishable by prefix.
|
||||
deadline = time.Now().Add(3 * time.Second)
|
||||
deadline = time.Now().Add(20 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
logs := strings.Join(sh.Logs(20), "\n")
|
||||
if strings.Contains(logs, "[sidecar] hello-from-sidecar") &&
|
||||
@@ -727,3 +739,50 @@ func TestWaitAliveFailsOnEarlyExit(t *testing.T) {
|
||||
t.Errorf("waited %v for a process that had already exited", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAdoptedShardCanBecomeReady(t *testing.T) {
|
||||
// Adoption stopped at "the process is alive", leaving the shard in Starting
|
||||
// with nothing that would ever move it. The control plane then waited out
|
||||
// the startup timeout and killed a perfectly healthy engine to start an
|
||||
// identical one — so an agent restart, which exists precisely so engines
|
||||
// are not interrupted, cost every model on the node a full reload.
|
||||
//
|
||||
// Seen on node05: the engine answered /health throughout and the replica
|
||||
// sat in Starting for three minutes until the timeout reaped it.
|
||||
s, _ := newSupervisor(t)
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20201}
|
||||
if _, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
health := &scriptedHealth{}
|
||||
s2, err := New(Config{
|
||||
WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second,
|
||||
HealthInterval: 10 * time.Millisecond,
|
||||
}, health, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
orphans, err := s2.FindOrphans()
|
||||
if err != nil || len(orphans) != 1 {
|
||||
t.Fatalf("orphan not found: %+v (%v)", orphans, err)
|
||||
}
|
||||
adopted, err := s2.Adopt(orphans[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := adopted.Status().State; got != Starting {
|
||||
t.Fatalf("a freshly adopted shard is %s, expected Starting", got)
|
||||
}
|
||||
|
||||
// The engine answers, as node05's did. The shard must be able to reach
|
||||
// Ready from here; before this it had no path to it at all.
|
||||
health.ready.Store(true)
|
||||
if err := s2.WaitReady(t.Context(), adopted, "http://127.0.0.1:20201", "m", 2*time.Second); err != nil {
|
||||
t.Fatalf("an adopted, healthy engine never became ready: %v", err)
|
||||
}
|
||||
if got := adopted.Status().State; got != Ready {
|
||||
t.Errorf("state %s after a passing health check", got)
|
||||
}
|
||||
_ = s2.Stop(t.Context(), "a")
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
// Package weights makes a model's files present on a node.
|
||||
//
|
||||
// The simple arrangement, and the default: every node has its own copy. A
|
||||
// source that is already a local path is used as it stands; one that names a
|
||||
// remote file is fetched into a cache directory once and reused. Nothing here
|
||||
// manages storage, decides a retention policy, or knows what a filesystem is —
|
||||
// it answers one question, which is where the engine should be pointed.
|
||||
//
|
||||
// The exotic arrangements are all handled by not acting. Weights on a shared
|
||||
// filer, an image exported over the fabric, or anything else an operator has
|
||||
// mounted are simply present, and present means there is nothing to do. That is
|
||||
// why the interesting case needs no code: the escape hatch from a distribution
|
||||
// mechanism is a distribution mechanism that checks first.
|
||||
package weights
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Progress reports how a fetch is going, for anything that wants to say so.
|
||||
type Progress struct {
|
||||
Source string
|
||||
BytesDone int64
|
||||
BytesTotal int64
|
||||
}
|
||||
|
||||
// Percent is how far along, or -1 when the total is unknown.
|
||||
func (p Progress) Percent() int {
|
||||
if p.BytesTotal <= 0 {
|
||||
return -1
|
||||
}
|
||||
return int(p.BytesDone * 100 / p.BytesTotal)
|
||||
}
|
||||
|
||||
// Fetcher makes sources local.
|
||||
type Fetcher struct {
|
||||
CacheDir string
|
||||
Client *http.Client
|
||||
Log *slog.Logger
|
||||
// OnProgress is called occasionally during a fetch. A missing model is the
|
||||
// difference between a replica that is slow to start and one that is stuck,
|
||||
// and without this the two are indistinguishable from outside.
|
||||
OnProgress func(Progress)
|
||||
|
||||
// inflight makes concurrent requests for one source share a single fetch.
|
||||
//
|
||||
// Eight shards of a tensor-parallel replica start at the same instant and
|
||||
// want the same file. Without this they would each open a connection and
|
||||
// each write the same bytes, turning one download into eight and, on a
|
||||
// shared origin, into a self-inflicted denial of service.
|
||||
mu sync.Mutex
|
||||
inflight map[string]*fetchCall
|
||||
}
|
||||
|
||||
type fetchCall struct {
|
||||
done chan struct{}
|
||||
path string
|
||||
err error
|
||||
}
|
||||
|
||||
// ErrNotLocal means a source names a file that is not here and cannot be
|
||||
// fetched — a bare path with nothing to fetch it from.
|
||||
var ErrNotLocal = errors.New("the weights are not on this node and the source says nothing about where to get them")
|
||||
|
||||
// Ensure returns a local path for a source, fetching it if it is not here.
|
||||
//
|
||||
// A plain path is never fetched, only checked. That is deliberate: a path is
|
||||
// what an operator writes when the file is already there, whether it arrived on
|
||||
// local disk, a Lustre mount or an image exported over the fabric, and a
|
||||
// downloader that second-guessed that would copy a hundred gigabytes that were
|
||||
// already present.
|
||||
func (f *Fetcher) Ensure(ctx context.Context, source string) (string, error) {
|
||||
if source == "" {
|
||||
return "", fmt.Errorf("no source given")
|
||||
}
|
||||
if !isRemote(source) {
|
||||
if !isPath(source) {
|
||||
// A bare model id — "Qwen/Qwen2.5-7B-Instruct" — is not a file and
|
||||
// was never meant to be one. vLLM and SGLang resolve those
|
||||
// themselves, and doing it for them would be both a duplicate
|
||||
// downloader and a change in documented behaviour.
|
||||
return source, nil
|
||||
}
|
||||
if _, err := os.Stat(source); err == nil {
|
||||
return source, nil
|
||||
}
|
||||
return "", fmt.Errorf("%w: %s", ErrNotLocal, source)
|
||||
}
|
||||
|
||||
target, err := f.cachePath(source)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := os.Stat(target); err == nil {
|
||||
return target, nil
|
||||
}
|
||||
|
||||
f.mu.Lock()
|
||||
if f.inflight == nil {
|
||||
f.inflight = map[string]*fetchCall{}
|
||||
}
|
||||
if call, running := f.inflight[source]; running {
|
||||
f.mu.Unlock()
|
||||
select {
|
||||
case <-call.done:
|
||||
return call.path, call.err
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
}
|
||||
}
|
||||
call := &fetchCall{done: make(chan struct{})}
|
||||
f.inflight[source] = call
|
||||
f.mu.Unlock()
|
||||
|
||||
call.path, call.err = f.fetch(ctx, source, target)
|
||||
close(call.done)
|
||||
|
||||
f.mu.Lock()
|
||||
delete(f.inflight, source)
|
||||
f.mu.Unlock()
|
||||
return call.path, call.err
|
||||
}
|
||||
|
||||
// fetch downloads a source to a target, resuming and verifying.
|
||||
func (f *Fetcher) fetch(ctx context.Context, source, target string) (string, error) {
|
||||
url, err := resolveURL(source)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Written beside the target and renamed at the end. A partial file under
|
||||
// the real name would be found by the next scan, reported as present, and
|
||||
// handed to an engine that fails on it minutes later — the exact failure
|
||||
// this whole area exists to prevent.
|
||||
part := target + ".part"
|
||||
existing := int64(0)
|
||||
if st, err := os.Stat(part); err == nil {
|
||||
existing = st.Size()
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if existing > 0 {
|
||||
// Resume. A forty gigabyte fetch that fails at ninety percent must not
|
||||
// start again from nothing.
|
||||
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", existing))
|
||||
}
|
||||
|
||||
client := f.Client
|
||||
if client == nil {
|
||||
client = &http.Client{}
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("fetching %s: %w", source, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
switch resp.StatusCode {
|
||||
case http.StatusOK:
|
||||
existing = 0 // the server ignored the range; start over
|
||||
case http.StatusPartialContent:
|
||||
default:
|
||||
return "", fmt.Errorf("fetching %s: %s returned %s", source, url, resp.Status)
|
||||
}
|
||||
|
||||
flags := os.O_CREATE | os.O_WRONLY
|
||||
if existing > 0 {
|
||||
flags |= os.O_APPEND
|
||||
} else {
|
||||
flags |= os.O_TRUNC
|
||||
}
|
||||
out, err := os.OpenFile(part, flags, 0o644)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
total := existing + resp.ContentLength
|
||||
f.report(Progress{Source: source, BytesDone: existing, BytesTotal: total})
|
||||
written, err := f.copyWithProgress(ctx, out, resp.Body, source, existing, total)
|
||||
closeErr := out.Close()
|
||||
if err != nil {
|
||||
// The part file is kept on purpose, so the next attempt resumes.
|
||||
return "", fmt.Errorf("fetching %s after %d bytes: %w", source, written, err)
|
||||
}
|
||||
if closeErr != nil {
|
||||
return "", closeErr
|
||||
}
|
||||
|
||||
if err := os.Rename(part, target); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if f.Log != nil {
|
||||
f.Log.Info("weights fetched", "source", source, "path", target, "bytes", written)
|
||||
}
|
||||
return target, nil
|
||||
}
|
||||
|
||||
func (f *Fetcher) copyWithProgress(ctx context.Context, dst io.Writer, src io.Reader,
|
||||
source string, done, total int64,
|
||||
) (int64, error) {
|
||||
buf := make([]byte, 4<<20)
|
||||
last := time.Now()
|
||||
for {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return done, err
|
||||
}
|
||||
n, err := src.Read(buf)
|
||||
if n > 0 {
|
||||
if _, werr := dst.Write(buf[:n]); werr != nil {
|
||||
return done, werr
|
||||
}
|
||||
done += int64(n)
|
||||
// Occasionally, not per read: a four-megabyte buffer over a
|
||||
// hundred gigabytes is twenty-five thousand callbacks, and nobody
|
||||
// needs to be told that often.
|
||||
if time.Since(last) > 2*time.Second {
|
||||
f.report(Progress{Source: source, BytesDone: done, BytesTotal: total})
|
||||
last = time.Now()
|
||||
}
|
||||
}
|
||||
if errors.Is(err, io.EOF) {
|
||||
f.report(Progress{Source: source, BytesDone: done, BytesTotal: total})
|
||||
return done, nil
|
||||
}
|
||||
if err != nil {
|
||||
return done, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (f *Fetcher) report(p Progress) {
|
||||
if f.OnProgress != nil {
|
||||
f.OnProgress(p)
|
||||
}
|
||||
}
|
||||
|
||||
// cachePath is where a remote source lands.
|
||||
//
|
||||
// Named by a hash of the source with the original filename kept on the end, so
|
||||
// two models with the same filename from different places cannot collide while
|
||||
// a human looking at the directory can still tell what is in it.
|
||||
func (f *Fetcher) cachePath(source string) (string, error) {
|
||||
if f.CacheDir == "" {
|
||||
return "", fmt.Errorf("no cache directory configured, so %s cannot be fetched", source)
|
||||
}
|
||||
return CachePathFor(f.CacheDir, source), nil
|
||||
}
|
||||
|
||||
// isPath reports whether a source names a file on this machine.
|
||||
//
|
||||
// By its leading character, because that is the only honest signal available:
|
||||
// "Qwen/Qwen2.5-7B-Instruct" and "models/qwen.gguf" are indistinguishable as
|
||||
// strings, and an operator writing a local file writes an absolute path. The
|
||||
// cost of being wrong in this direction is that a relative path is passed to
|
||||
// the engine untouched, which is what happened before this code existed.
|
||||
func isPath(source string) bool {
|
||||
return strings.HasPrefix(source, "/") ||
|
||||
strings.HasPrefix(source, "./") ||
|
||||
strings.HasPrefix(source, "../")
|
||||
}
|
||||
|
||||
// isRemote reports whether a source has to be fetched at all.
|
||||
func isRemote(source string) bool {
|
||||
return strings.HasPrefix(source, "hf:") ||
|
||||
strings.HasPrefix(source, "http://") ||
|
||||
strings.HasPrefix(source, "https://")
|
||||
}
|
||||
|
||||
// resolveURL turns a source into something to GET.
|
||||
//
|
||||
// hf: names one file in a repository, which is what a GGUF is. A whole
|
||||
// repository of shards is deliberately not handled here: vLLM and SGLang
|
||||
// download those themselves and do it well, and reimplementing a snapshot
|
||||
// downloader to compete with them would be work spent to arrive back where we
|
||||
// started.
|
||||
func resolveURL(source string) (string, error) {
|
||||
if rest, ok := strings.CutPrefix(source, "hf:"); ok {
|
||||
parts := strings.Split(rest, "/")
|
||||
if len(parts) < 3 {
|
||||
return "", fmt.Errorf(
|
||||
"%s names a repository, not a file; give the file as hf:org/repo/model.gguf, "+
|
||||
"or use the repository id directly and let the engine download it", source)
|
||||
}
|
||||
repo := strings.Join(parts[:2], "/")
|
||||
file := strings.Join(parts[2:], "/")
|
||||
return "https://huggingface.co/" + repo + "/resolve/main/" + file, nil
|
||||
}
|
||||
if _, err := url.Parse(source); err != nil {
|
||||
return "", fmt.Errorf("%s is not a usable source: %w", source, err)
|
||||
}
|
||||
return source, nil
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
package weights
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func serveBytes(t *testing.T, body []byte, hits *atomic.Int64) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if hits != nil {
|
||||
hits.Add(1)
|
||||
}
|
||||
http.ServeContent(w, r, "model.gguf", time.Time{}, strings.NewReader(string(body)))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestALocalPathIsUsedNeverCopied(t *testing.T) {
|
||||
// A path is what an operator writes when the file is already there —
|
||||
// local disk, a Lustre mount, an image exported over the fabric. A
|
||||
// downloader that second-guessed that would copy a hundred gigabytes that
|
||||
// were already present.
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "model.gguf")
|
||||
if err := os.WriteFile(path, []byte("weights"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
|
||||
got, err := f.Ensure(t.Context(), path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != path {
|
||||
t.Errorf("returned %q, want the original path", got)
|
||||
}
|
||||
if entries, _ := os.ReadDir(f.CacheDir); len(entries) != 0 {
|
||||
t.Errorf("the cache was written to for a local file: %v", entries)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingLocalPathSaysWhatIsWrong(t *testing.T) {
|
||||
// The commonest setup mistake, and "no such file" from an engine twenty
|
||||
// minutes into a load is a poor way to learn it.
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
_, err := f.Ensure(t.Context(), filepath.Join(t.TempDir(), "absent.gguf"))
|
||||
if !errors.Is(err, ErrNotLocal) {
|
||||
t.Fatalf("err = %v, want ErrNotLocal", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "absent.gguf") {
|
||||
t.Errorf("the error does not name the file: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestARemoteFileIsFetchedOnceAndReused(t *testing.T) {
|
||||
var hits atomic.Int64
|
||||
srv := serveBytes(t, []byte("the weights"), &hits)
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
|
||||
first, err := f.Ensure(t.Context(), srv.URL+"/model.gguf")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body, err := os.ReadFile(first)
|
||||
if err != nil || string(body) != "the weights" {
|
||||
t.Fatalf("fetched content = %q, err %v", body, err)
|
||||
}
|
||||
|
||||
second, err := f.Ensure(t.Context(), srv.URL+"/model.gguf")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if second != first {
|
||||
t.Errorf("the second call returned a different path: %q vs %q", second, first)
|
||||
}
|
||||
if hits.Load() != 1 {
|
||||
t.Errorf("the file was fetched %d times; the cache did not hold", hits.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentRequestsProduceOneDownload(t *testing.T) {
|
||||
// Eight shards of a tensor-parallel replica start at the same instant and
|
||||
// want the same file. Eight downloads is a self-inflicted denial of service
|
||||
// on whatever is serving them.
|
||||
var hits atomic.Int64
|
||||
srv := serveBytes(t, []byte("shared weights"), &hits)
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
paths := make([]string, 8)
|
||||
errs := make([]error, 8)
|
||||
for i := range 8 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
paths[i], errs[i] = f.Ensure(t.Context(), srv.URL+"/model.gguf")
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
for i, err := range errs {
|
||||
if err != nil {
|
||||
t.Fatalf("caller %d failed: %v", i, err)
|
||||
}
|
||||
if paths[i] != paths[0] {
|
||||
t.Errorf("caller %d got %q, caller 0 got %q", i, paths[i], paths[0])
|
||||
}
|
||||
}
|
||||
if n := hits.Load(); n != 1 {
|
||||
t.Errorf("%d downloads for eight concurrent callers", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPartialDownloadIsNeverVisibleAsTheModel(t *testing.T) {
|
||||
// A partial file under the real name would be found by the next scan,
|
||||
// reported as present, and handed to an engine that fails on it minutes
|
||||
// later — the exact failure all of this exists to prevent.
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Length", "1000")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("only the first bytes"))
|
||||
// and then the connection dies without the rest
|
||||
if fl, ok := w.(http.Flusher); ok {
|
||||
fl.Flush()
|
||||
}
|
||||
panic(http.ErrAbortHandler)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cache := t.TempDir()
|
||||
f := &Fetcher{CacheDir: cache}
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/model.gguf"); err == nil {
|
||||
t.Fatal("a truncated download was reported as success")
|
||||
}
|
||||
entries, _ := os.ReadDir(cache)
|
||||
for _, e := range entries {
|
||||
if !strings.HasSuffix(e.Name(), ".part") {
|
||||
t.Errorf("a non-part file was left behind: %q", e.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnInterruptedDownloadResumes(t *testing.T) {
|
||||
// A forty gigabyte fetch that fails at ninety percent must not start again
|
||||
// from nothing.
|
||||
full := []byte(strings.Repeat("x", 4096))
|
||||
var ranges []string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
ranges = append(ranges, r.Header.Get("Range"))
|
||||
http.ServeContent(w, r, "model.gguf", time.Time{}, strings.NewReader(string(full)))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cache := t.TempDir()
|
||||
f := &Fetcher{CacheDir: cache}
|
||||
target, _ := f.cachePath(srv.URL + "/model.gguf")
|
||||
// Half a file already on disk, as an interrupted attempt would leave.
|
||||
if err := os.WriteFile(target+".part", full[:2048], 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got, err := f.Ensure(t.Context(), srv.URL+"/model.gguf")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body, _ := os.ReadFile(got)
|
||||
if len(body) != len(full) {
|
||||
t.Errorf("resumed file is %d bytes, want %d", len(body), len(full))
|
||||
}
|
||||
if len(ranges) == 0 || !strings.Contains(ranges[0], "2048") {
|
||||
t.Errorf("no range request was made: %v", ranges)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnHFSourceNamingARepositoryExplainsItself(t *testing.T) {
|
||||
// A whole repository of shards is deliberately not handled here: vLLM and
|
||||
// SGLang download those themselves and do it well. Saying so beats a
|
||||
// confusing 404.
|
||||
_, err := resolveURL("hf:Qwen/Qwen2.5-Coder-32B-Instruct")
|
||||
if err == nil {
|
||||
t.Fatal("a bare repository was accepted as a file")
|
||||
}
|
||||
for _, want := range []string{"model.gguf", "let the engine"} {
|
||||
if !strings.Contains(err.Error(), want) {
|
||||
t.Errorf("the error does not suggest what to do (%q): %v", want, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnHFFileResolvesToItsDownloadURL(t *testing.T) {
|
||||
got, err := resolveURL("hf:Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/qwen2.5-coder-3b-instruct-q4_k_m.gguf")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "https://huggingface.co/Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/resolve/main/qwen2.5-coder-3b-instruct-q4_k_m.gguf"
|
||||
if got != want {
|
||||
t.Errorf("resolved to\n %s\nwant\n %s", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTwoSourcesWithTheSameFilenameDoNotCollide(t *testing.T) {
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
a, _ := f.cachePath("https://one.example/model.gguf")
|
||||
b, _ := f.cachePath("https://two.example/model.gguf")
|
||||
if a == b {
|
||||
t.Fatalf("both sources cache to %q", a)
|
||||
}
|
||||
// And a person looking at the directory can still tell what is in it.
|
||||
for _, p := range []string{a, b} {
|
||||
if !strings.HasSuffix(p, "model.gguf") {
|
||||
t.Errorf("cache path %q does not keep the filename", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProgressIsReportedWithAPercentage(t *testing.T) {
|
||||
// A missing model is the difference between a replica that is slow to start
|
||||
// and one that is stuck, and without this they look identical.
|
||||
srv := serveBytes(t, []byte(strings.Repeat("y", 8192)), nil)
|
||||
var seen []Progress
|
||||
f := &Fetcher{
|
||||
CacheDir: t.TempDir(),
|
||||
OnProgress: func(p Progress) { seen = append(seen, p) },
|
||||
}
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/model.gguf"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(seen) == 0 {
|
||||
t.Fatal("no progress was reported")
|
||||
}
|
||||
last := seen[len(seen)-1]
|
||||
if last.Percent() != 100 {
|
||||
t.Errorf("final progress is %d%%, want 100 (%+v)", last.Percent(), last)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPercentIsUnknownRatherThanZeroWithoutATotal(t *testing.T) {
|
||||
// A server that sends no Content-Length is common, and reporting 0% for
|
||||
// a download that is nearly finished is worse than admitting ignorance.
|
||||
if got := (Progress{BytesDone: 500}).Percent(); got != -1 {
|
||||
t.Errorf("Percent with no total = %d, want -1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoCacheDirectoryIsAClearRefusal(t *testing.T) {
|
||||
f := &Fetcher{}
|
||||
_, err := f.Ensure(t.Context(), "https://example.invalid/model.gguf")
|
||||
if err == nil || !strings.Contains(err.Error(), "cache directory") {
|
||||
t.Errorf("err = %v, want a complaint about the cache directory", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestABareModelIdIsLeftToTheEngine(t *testing.T) {
|
||||
// "Qwen/Qwen2.5-7B-Instruct" is a documented source that vLLM and SGLang
|
||||
// resolve themselves. Treating it as a missing file broke every existing
|
||||
// vLLM profile, which is how this was found.
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
for _, id := range []string{
|
||||
"Qwen/Qwen2.5-7B-Instruct",
|
||||
"meta-llama/Llama-3.3-70B-Instruct",
|
||||
} {
|
||||
got, err := f.Ensure(t.Context(), id)
|
||||
if err != nil {
|
||||
t.Errorf("%s was refused: %v", id, err)
|
||||
}
|
||||
if got != id {
|
||||
t.Errorf("%s was rewritten to %q", id, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAbsolutePathThatIsMissingIsStillAnError(t *testing.T) {
|
||||
// The other half: an operator writing a path means a file, and a missing
|
||||
// one should be said plainly rather than handed to an engine.
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
if _, err := f.Ensure(t.Context(), "/models/definitely-not-here.gguf"); !errors.Is(err, ErrNotLocal) {
|
||||
t.Errorf("err = %v, want ErrNotLocal", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package weights
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// cacheNamePattern matches a filename this package created.
|
||||
//
|
||||
// The eight-byte hash of the source, hex encoded, then a dash, then the
|
||||
// original filename — exactly what cachePath writes. It is the only mark of
|
||||
// ownership there is, and it has to be exact: everything else in a model
|
||||
// directory was put there by the operator, and deleting one of those files
|
||||
// could cost a week of someone's bandwidth.
|
||||
var cacheNamePattern = regexp.MustCompile(`^[0-9a-f]{16}-`)
|
||||
|
||||
// CachePathFor is where a source would be cached in the given directory.
|
||||
//
|
||||
// Exported so the control plane can work out which files in a node's reported
|
||||
// list are ones llmgw fetched for a model it still wants, without asking the
|
||||
// node anything. The agent and the control plane must agree on this name
|
||||
// exactly, which is why there is one function and not two.
|
||||
func CachePathFor(dir, source string) string {
|
||||
sum := sha256.Sum256([]byte(source))
|
||||
name := filepath.Base(source)
|
||||
if name == "" || name == "." || name == "/" {
|
||||
name = "weights"
|
||||
}
|
||||
return filepath.Join(dir, hex.EncodeToString(sum[:8])+"-"+name)
|
||||
}
|
||||
|
||||
// Owned reports whether a path is one this package created.
|
||||
//
|
||||
// The dividing line of the whole feature. A file without the cache name was
|
||||
// put there by the operator — by scp, by a shared mount, by hand — and llmgw
|
||||
// has no business deleting it however unreferenced it looks. A file with the
|
||||
// cache name is one llmgw downloaded, and llmgw can take it back.
|
||||
func Owned(path string) bool {
|
||||
return cacheNamePattern.MatchString(filepath.Base(path))
|
||||
}
|
||||
|
||||
// Orphans lists cached files that no wanted source accounts for.
|
||||
//
|
||||
// `have` is what a node reports on disk and `wanted` is every source in the
|
||||
// catalog; both are passed in so this is a calculation rather than a scan, and
|
||||
// so it can be run on the control plane against a node's reported inventory.
|
||||
//
|
||||
// Only files this package created are ever returned. A source that is not
|
||||
// fetched — a plain path, or a model id an engine resolves itself — produces
|
||||
// no cache name and so protects nothing and deletes nothing, which is correct:
|
||||
// those files are not in the cache to begin with.
|
||||
func Orphans(have []string, wanted []string) []string {
|
||||
keep := map[string]struct{}{}
|
||||
for _, src := range wanted {
|
||||
if !isRemote(src) {
|
||||
// Nothing was fetched for it, so nothing in the cache belongs to
|
||||
// it. Computing a name here would reserve a file that does not
|
||||
// exist and could never be an orphan anyway.
|
||||
continue
|
||||
}
|
||||
// By base name rather than full path: the control plane does not know
|
||||
// each node's cache directory, and the name is what identifies the
|
||||
// file wherever the node keeps it.
|
||||
keep[filepath.Base(CachePathFor("", src))] = struct{}{}
|
||||
}
|
||||
|
||||
var out []string
|
||||
seen := map[string]struct{}{}
|
||||
for _, path := range have {
|
||||
base := filepath.Base(path)
|
||||
if !Owned(base) {
|
||||
continue
|
||||
}
|
||||
if _, wanted := keep[base]; wanted {
|
||||
continue
|
||||
}
|
||||
if _, dup := seen[path]; dup {
|
||||
continue
|
||||
}
|
||||
seen[path] = struct{}{}
|
||||
out = append(out, path)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// PartialName reports whether a path is the temporary file a fetch writes to.
|
||||
//
|
||||
// Those are not orphans in the same sense — a fetch may be running right now —
|
||||
// so they are named separately rather than swept up with the rest. A `.part`
|
||||
// left behind by a killed agent is real garbage, but only its own node can
|
||||
// tell whether it is stale, and nothing here can.
|
||||
func PartialName(path string) bool {
|
||||
return strings.HasSuffix(path, ".part")
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package weights
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestOnlyFilesLLMGWFetchedAreEverOrphans(t *testing.T) {
|
||||
// The invariant the whole feature rests on. Everything without the cache
|
||||
// name was put there by the operator — scp, a shared mount, a copy by hand
|
||||
// — and deleting one could cost a week of someone's bandwidth. It does not
|
||||
// matter how unreferenced it looks.
|
||||
const dir = "/var/lib/llmgw-models"
|
||||
have := []string{
|
||||
filepath.Join(dir, "qwen2.5-coder-14b-instruct-q8_0-00001-of-00002.gguf"),
|
||||
filepath.Join(dir, "some-model.safetensors"),
|
||||
filepath.Join(dir, "notes.txt"),
|
||||
CachePathFor(dir, "https://models.internal/gone.gguf"),
|
||||
}
|
||||
|
||||
got := Orphans(have, nil)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected only the fetched file, got %v", got)
|
||||
}
|
||||
if !Owned(got[0]) {
|
||||
t.Errorf("returned a file llmgw did not fetch: %s", got[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFetchedFileStillInTheCatalogIsKept(t *testing.T) {
|
||||
const dir = "/var/lib/llmgw-models"
|
||||
const wanted = "hf:Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/q4.gguf"
|
||||
const gone = "https://models.internal/old.gguf"
|
||||
|
||||
have := []string{CachePathFor(dir, wanted), CachePathFor(dir, gone)}
|
||||
got := Orphans(have, []string{wanted})
|
||||
|
||||
if slices.Contains(got, CachePathFor(dir, wanted)) {
|
||||
t.Error("a file the catalog still points at was called an orphan")
|
||||
}
|
||||
if !slices.Contains(got, CachePathFor(dir, gone)) {
|
||||
t.Errorf("a file nothing points at was not found: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheCacheNameMatchesWhatAFetchActuallyWrites(t *testing.T) {
|
||||
// The control plane proposes deletions by recomputing this name from the
|
||||
// catalog. If it ever drifted from what a fetch writes, the proposal would
|
||||
// name a file that is in use — so the two must be one function, and this
|
||||
// is what says so.
|
||||
f := &Fetcher{CacheDir: "/cache"}
|
||||
const source = "https://models.internal/x.gguf"
|
||||
|
||||
got, err := f.cachePath(source)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if want := CachePathFor("/cache", source); got != want {
|
||||
t.Errorf("a fetch writes %q but the catalog would look for %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneOfTwoLookalikeFilesIsNotMistakenForTheOther(t *testing.T) {
|
||||
// That the two cache to different names is covered where the naming lives.
|
||||
// What matters here is the consequence: deleting the unreferenced one must
|
||||
// not take the referenced one with it, and both end in "model.gguf".
|
||||
a := CachePathFor("/cache", "https://one.example/model.gguf")
|
||||
b := CachePathFor("/cache", "https://two.example/model.gguf")
|
||||
|
||||
got := Orphans([]string{a, b}, []string{"https://one.example/model.gguf"})
|
||||
if !slices.Equal(got, []string{b}) {
|
||||
t.Errorf("orphans = %v, want only %s", got, b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASourceThatIsNeverFetchedProtectsNothingAndDeletesNothing(t *testing.T) {
|
||||
// A plain path, or a model id an engine resolves itself, puts nothing in
|
||||
// the cache. Computing a cache name for it would reserve a file that does
|
||||
// not exist, and it must not stop a real orphan being found either.
|
||||
const dir = "/models"
|
||||
orphan := CachePathFor(dir, "https://models.internal/gone.gguf")
|
||||
|
||||
got := Orphans([]string{orphan, filepath.Join(dir, "local.gguf")},
|
||||
[]string{"/models/local.gguf", "Qwen/Qwen2.5-7B-Instruct"})
|
||||
|
||||
if !slices.Equal(got, []string{orphan}) {
|
||||
t.Errorf("orphans = %v, want only %s", got, orphan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEmptyCatalogDoesNotMakeOperatorFilesDisposable(t *testing.T) {
|
||||
// The dangerous case: every model removed from the catalog at once. Files
|
||||
// llmgw fetched become orphans, which is right. Files it did not must stay
|
||||
// untouched, which is the whole point.
|
||||
const dir = "/models"
|
||||
have := []string{
|
||||
filepath.Join(dir, "hand-placed.gguf"),
|
||||
CachePathFor(dir, "https://models.internal/a.gguf"),
|
||||
}
|
||||
got := Orphans(have, nil)
|
||||
if len(got) != 1 || got[0] != CachePathFor(dir, "https://models.internal/a.gguf") {
|
||||
t.Errorf("orphans = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPartialsAreNotSweptUpWithTheRest(t *testing.T) {
|
||||
// A fetch may be running right now, and only the node itself can tell a
|
||||
// live download from one a killed agent left behind.
|
||||
if !PartialName("/models/abc-model.gguf.part") {
|
||||
t.Error("a .part file was not recognised")
|
||||
}
|
||||
if PartialName("/models/abc-model.gguf") {
|
||||
t.Error("a finished file was called partial")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
// Package auth authenticates client requests against the key store (plan §7.9).
|
||||
//
|
||||
// The rule that shapes everything here: authentication runs on every request,
|
||||
// and the request path does not touch the database. A key is looked up once,
|
||||
// resolved against its tenancy chain once, and then answered from memory until
|
||||
// the entry expires or is invalidated.
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// Principal is who a request is, once authenticated.
|
||||
type Principal struct {
|
||||
KeyID string
|
||||
KeyPrefix string
|
||||
ProjectID string
|
||||
TeamID string
|
||||
OrgID string
|
||||
|
||||
// Band and RoutingProfile come from the key, and are authoritative: with
|
||||
// authentication on, a client can no longer promote itself with a header.
|
||||
Band string
|
||||
RoutingProfile string
|
||||
|
||||
Limits store.Effective
|
||||
}
|
||||
|
||||
// Allows reports whether this principal may use a model.
|
||||
func (p Principal) Allows(model string) bool { return p.Limits.Allows(model) }
|
||||
|
||||
type contextKey struct{}
|
||||
|
||||
// WithPrincipal attaches an authenticated identity to a context.
|
||||
func WithPrincipal(ctx context.Context, p Principal) context.Context {
|
||||
return context.WithValue(ctx, contextKey{}, p)
|
||||
}
|
||||
|
||||
// FromContext returns the authenticated identity, if there is one.
|
||||
//
|
||||
// Absent means authentication is off, which is a deployment choice rather than
|
||||
// a failure: a gateway on a trusted network with no keys configured still
|
||||
// serves. Callers must therefore treat the missing case as "unrestricted", and
|
||||
// the middleware is what guarantees it cannot mean "unchecked".
|
||||
func FromContext(ctx context.Context) (Principal, bool) {
|
||||
p, ok := ctx.Value(contextKey{}).(Principal)
|
||||
return p, ok
|
||||
}
|
||||
|
||||
// Config tunes the authenticator.
|
||||
type Config struct {
|
||||
// TTL is how long a resolved key is trusted before being read again. It
|
||||
// bounds how stale a revocation can be, so it trades a database read per
|
||||
// key per TTL against how long a revoked key keeps working.
|
||||
TTL time.Duration
|
||||
// NegativeTTL is how long a rejected key is remembered as rejected, so a
|
||||
// client looping with a bad key cannot turn into a database read per
|
||||
// request.
|
||||
NegativeTTL time.Duration
|
||||
// TouchInterval is how often a key's last-used time is written, at most.
|
||||
TouchInterval time.Duration
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible values.
|
||||
func DefaultConfig() Config {
|
||||
return Config{TTL: 30 * time.Second, NegativeTTL: 5 * time.Second, TouchInterval: time.Minute}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() {
|
||||
d := DefaultConfig()
|
||||
if c.TTL <= 0 {
|
||||
c.TTL = d.TTL
|
||||
}
|
||||
if c.NegativeTTL <= 0 {
|
||||
c.NegativeTTL = d.NegativeTTL
|
||||
}
|
||||
if c.TouchInterval <= 0 {
|
||||
c.TouchInterval = d.TouchInterval
|
||||
}
|
||||
}
|
||||
|
||||
// Errors an authenticator returns. They are distinguished because the status
|
||||
// codes differ: a missing credential is 401, a forbidden model is 403.
|
||||
var (
|
||||
ErrMissing = errors.New("no api key")
|
||||
ErrInvalid = errors.New("invalid api key")
|
||||
ErrExpired = errors.New("api key expired")
|
||||
ErrRevoked = errors.New("api key revoked")
|
||||
)
|
||||
|
||||
// Authenticator resolves keys and caches the result.
|
||||
type Authenticator struct {
|
||||
store store.Store
|
||||
cfg Config
|
||||
now func() time.Time
|
||||
|
||||
mu sync.RWMutex
|
||||
cache map[string]entry
|
||||
|
||||
// touch batches last-used timestamps; nothing reads the field in real time
|
||||
// and writing it per request would put a database write back on the path.
|
||||
touchMu sync.Mutex
|
||||
touch map[string]time.Time
|
||||
}
|
||||
|
||||
type entry struct {
|
||||
principal Principal
|
||||
err error
|
||||
expires time.Time
|
||||
}
|
||||
|
||||
// New builds an authenticator over a store.
|
||||
func New(s store.Store, cfg Config) *Authenticator {
|
||||
cfg.normalize()
|
||||
return &Authenticator{
|
||||
store: s, cfg: cfg, now: time.Now,
|
||||
cache: map[string]entry{},
|
||||
touch: map[string]time.Time{},
|
||||
}
|
||||
}
|
||||
|
||||
// Authenticate resolves a presented credential.
|
||||
func (a *Authenticator) Authenticate(ctx context.Context, presented string) (Principal, error) {
|
||||
if presented == "" {
|
||||
return Principal{}, ErrMissing
|
||||
}
|
||||
// Shape first: a bearer token from another system is rejected without a
|
||||
// lookup, and never reaches the store.
|
||||
if err := store.ParseKey(presented); err != nil {
|
||||
return Principal{}, ErrInvalid
|
||||
}
|
||||
|
||||
hash := store.HashKey(presented)
|
||||
id := string(hash)
|
||||
|
||||
if e, ok := a.lookup(id); ok {
|
||||
if e.err == nil {
|
||||
a.markUsed(e.principal.KeyID)
|
||||
}
|
||||
return e.principal, e.err
|
||||
}
|
||||
|
||||
p, err := a.resolve(ctx, hash)
|
||||
if err != nil && !isRejection(err) {
|
||||
// A database problem is not an authentication answer, so it must not be
|
||||
// cached as one. Failing the request is right — the alternative is
|
||||
// serving traffic whose authorization is unknown.
|
||||
return Principal{}, err
|
||||
}
|
||||
|
||||
ttl := a.cfg.TTL
|
||||
if err != nil {
|
||||
ttl = a.cfg.NegativeTTL
|
||||
}
|
||||
a.store_(id, entry{principal: p, err: err, expires: a.now().Add(ttl)})
|
||||
if err == nil {
|
||||
a.markUsed(p.KeyID)
|
||||
}
|
||||
return p, err
|
||||
}
|
||||
|
||||
func (a *Authenticator) resolve(ctx context.Context, hash []byte) (Principal, error) {
|
||||
k, err := a.store.KeyByHash(ctx, hash)
|
||||
if errors.Is(err, store.ErrNotFound) {
|
||||
return Principal{}, ErrInvalid
|
||||
}
|
||||
if err != nil {
|
||||
return Principal{}, err
|
||||
}
|
||||
switch {
|
||||
case !k.RevokedAt.IsZero():
|
||||
return Principal{}, ErrRevoked
|
||||
case !k.ExpiresAt.IsZero() && !a.now().Before(k.ExpiresAt):
|
||||
return Principal{}, ErrExpired
|
||||
}
|
||||
|
||||
eff, err := a.store.Resolve(ctx, k.ID)
|
||||
if err != nil {
|
||||
return Principal{}, err
|
||||
}
|
||||
return Principal{
|
||||
KeyID: k.ID, KeyPrefix: k.Prefix, ProjectID: eff.ProjectID,
|
||||
TeamID: eff.TeamID, OrgID: eff.OrgID,
|
||||
Band: eff.Band, RoutingProfile: eff.RoutingProfile,
|
||||
Limits: eff,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (a *Authenticator) lookup(id string) (entry, bool) {
|
||||
a.mu.RLock()
|
||||
e, ok := a.cache[id]
|
||||
a.mu.RUnlock()
|
||||
if !ok || a.now().After(e.expires) {
|
||||
return entry{}, false
|
||||
}
|
||||
return e, true
|
||||
}
|
||||
|
||||
func (a *Authenticator) store_(id string, e entry) {
|
||||
a.mu.Lock()
|
||||
a.cache[id] = e
|
||||
a.mu.Unlock()
|
||||
}
|
||||
|
||||
// Invalidate drops every cached decision.
|
||||
//
|
||||
// Coarse on purpose: a limit change anywhere in the tree can change any key
|
||||
// below it, and working out which keys those are costs more than re-resolving
|
||||
// the few that are actually in use.
|
||||
func (a *Authenticator) Invalidate() {
|
||||
a.mu.Lock()
|
||||
a.cache = map[string]entry{}
|
||||
a.mu.Unlock()
|
||||
}
|
||||
|
||||
func (a *Authenticator) markUsed(keyID string) {
|
||||
if keyID == "" {
|
||||
return
|
||||
}
|
||||
now := a.now()
|
||||
a.touchMu.Lock()
|
||||
if last, ok := a.touch[keyID]; !ok || now.Sub(last) >= a.cfg.TouchInterval {
|
||||
a.touch[keyID] = now
|
||||
}
|
||||
a.touchMu.Unlock()
|
||||
}
|
||||
|
||||
// FlushTouches writes the pending last-used timestamps. The batch writer calls
|
||||
// it on its own schedule.
|
||||
func (a *Authenticator) FlushTouches(ctx context.Context) error {
|
||||
a.touchMu.Lock()
|
||||
pending := a.touch
|
||||
a.touch = map[string]time.Time{}
|
||||
a.touchMu.Unlock()
|
||||
|
||||
if len(pending) == 0 {
|
||||
return nil
|
||||
}
|
||||
return a.store.TouchKeys(ctx, pending)
|
||||
}
|
||||
|
||||
func isRejection(err error) bool {
|
||||
return errors.Is(err, ErrInvalid) || errors.Is(err, ErrExpired) ||
|
||||
errors.Is(err, ErrRevoked) || errors.Is(err, ErrMissing)
|
||||
}
|
||||
|
||||
// BearerToken extracts the credential from a request.
|
||||
//
|
||||
// Both spellings are accepted because both are in the wild: the OpenAI SDKs
|
||||
// send Authorization, and some tools only manage to set an api-key header.
|
||||
func BearerToken(r *http.Request) string {
|
||||
if h := r.Header.Get("Authorization"); h != "" {
|
||||
if token, ok := strings.CutPrefix(h, "Bearer "); ok {
|
||||
return strings.TrimSpace(token)
|
||||
}
|
||||
// A bare value is common enough in hand-written clients to accept.
|
||||
return strings.TrimSpace(h)
|
||||
}
|
||||
return strings.TrimSpace(r.Header.Get("api-key"))
|
||||
}
|
||||
@@ -0,0 +1,440 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// fixture builds a store with one org/team/project and one key.
|
||||
func fixture(t *testing.T) (*store.SQLite, string, string) {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
s, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "t.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = s.Close() })
|
||||
|
||||
must(t, s.PutOrg(ctx, store.Org{ID: "org1", Name: "Acme"}))
|
||||
must(t, s.PutTeam(ctx, store.Team{ID: "team1", OrgID: "org1", Name: "R"}))
|
||||
must(t, s.PutProject(ctx, store.Project{ID: "proj1", TeamID: "team1", Name: "P"}))
|
||||
|
||||
secret, k, err := store.NewKey("key1", "laptop", "proj1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
k.CreatedAt = time.Now()
|
||||
k.Band = "gold"
|
||||
must(t, s.CreateKey(ctx, k))
|
||||
return s, secret, k.ID
|
||||
}
|
||||
|
||||
func must(t *testing.T, err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAValidKeyResolvesItsWholeChain(t *testing.T) {
|
||||
s, secret, keyID := fixture(t)
|
||||
a := New(s, DefaultConfig())
|
||||
|
||||
p, err := a.Authenticate(context.Background(), secret)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if p.KeyID != keyID || p.ProjectID != "proj1" || p.TeamID != "team1" || p.OrgID != "org1" {
|
||||
t.Errorf("principal %+v", p)
|
||||
}
|
||||
if p.Band != "gold" {
|
||||
t.Errorf("band %q", p.Band)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRejections(t *testing.T) {
|
||||
s, secret, keyID := fixture(t)
|
||||
ctx := context.Background()
|
||||
a := New(s, DefaultConfig())
|
||||
|
||||
if _, err := a.Authenticate(ctx, ""); !errors.Is(err, ErrMissing) {
|
||||
t.Errorf("no credential: %v", err)
|
||||
}
|
||||
// A bearer token from some other system must be rejected on shape alone.
|
||||
if _, err := a.Authenticate(ctx, "ghp_0123456789abcdefghij"); !errors.Is(err, ErrInvalid) {
|
||||
t.Errorf("foreign token: %v", err)
|
||||
}
|
||||
if _, err := a.Authenticate(ctx, alter(secret)); !errors.Is(err, ErrInvalid) {
|
||||
t.Errorf("wrong key: %v", err)
|
||||
}
|
||||
|
||||
must(t, s.RevokeKey(ctx, keyID, time.Now()))
|
||||
a.Invalidate()
|
||||
if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrRevoked) {
|
||||
t.Errorf("revoked key: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// countingStore counts how often the database is actually read.
|
||||
type countingStore struct {
|
||||
store.Store
|
||||
lookups atomic.Int64
|
||||
fail atomic.Bool
|
||||
}
|
||||
|
||||
func (c *countingStore) KeyByHash(ctx context.Context, hash []byte) (store.Key, error) {
|
||||
c.lookups.Add(1)
|
||||
if c.fail.Load() {
|
||||
return store.Key{}, errors.New("database is down")
|
||||
}
|
||||
return c.Store.KeyByHash(ctx, hash)
|
||||
}
|
||||
|
||||
// alter changes a key by one character, guaranteeing the result differs.
|
||||
//
|
||||
// Replacing the last character with a fixed one does not: keys are base62, so
|
||||
// one in sixty-two already ends in that character and the "wrong" key is the
|
||||
// right one. That test passed locally for a day and failed in CI, which is
|
||||
// exactly how long a one-in-sixty-two flake takes to look like a real bug.
|
||||
func alter(secret string) string {
|
||||
last := secret[len(secret)-1]
|
||||
replacement := byte('a')
|
||||
if last == replacement {
|
||||
replacement = 'b'
|
||||
}
|
||||
return secret[:len(secret)-1] + string(replacement)
|
||||
}
|
||||
|
||||
func TestTheRequestPathDoesNotReadTheDatabase(t *testing.T) {
|
||||
// The rule the whole package is shaped around. Without the cache this is a
|
||||
// database read per request, on every request.
|
||||
s, secret, _ := fixture(t)
|
||||
c := &countingStore{Store: s}
|
||||
a := New(c, DefaultConfig())
|
||||
|
||||
for range 100 {
|
||||
if _, err := a.Authenticate(context.Background(), secret); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if n := c.lookups.Load(); n != 1 {
|
||||
t.Errorf("100 requests caused %d database reads", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestARejectedKeyIsNotADatabaseReadPerRequest(t *testing.T) {
|
||||
// A client looping with a bad key would otherwise be a free way to load the
|
||||
// database from outside.
|
||||
s, _, _ := fixture(t)
|
||||
c := &countingStore{Store: s}
|
||||
a := New(c, DefaultConfig())
|
||||
|
||||
bad := store.KeyPrefix + "0000000000000000000000000000000000000000000"
|
||||
for range 50 {
|
||||
if _, err := a.Authenticate(context.Background(), bad); !errors.Is(err, ErrInvalid) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if n := c.lookups.Load(); n != 1 {
|
||||
t.Errorf("50 bad requests caused %d database reads", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCachedDecisionsExpire(t *testing.T) {
|
||||
// The cache bounds how long a revocation can take to take effect, so it has
|
||||
// to actually expire.
|
||||
s, secret, keyID := fixture(t)
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
a := New(s, Config{TTL: time.Second, NegativeTTL: time.Second})
|
||||
a.now = func() time.Time { return now }
|
||||
|
||||
if _, err := a.Authenticate(ctx, secret); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
must(t, s.RevokeKey(ctx, keyID, now))
|
||||
|
||||
// Still inside the TTL: the stale answer is the point of the cache.
|
||||
if _, err := a.Authenticate(ctx, secret); err != nil {
|
||||
t.Errorf("the cache expired early: %v", err)
|
||||
}
|
||||
|
||||
now = now.Add(2 * time.Second)
|
||||
if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrRevoked) {
|
||||
t.Errorf("a revoked key still works after the TTL: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAStoreFailureIsNotAnAuthenticationAnswer(t *testing.T) {
|
||||
// Two things must not happen: the failure must not be cached as a
|
||||
// rejection, and it must not read to the caller as a bad key.
|
||||
s, secret, _ := fixture(t)
|
||||
c := &countingStore{Store: s}
|
||||
a := New(c, DefaultConfig())
|
||||
|
||||
c.fail.Store(true)
|
||||
_, err := a.Authenticate(context.Background(), secret)
|
||||
if err == nil {
|
||||
t.Fatal("a broken database authenticated a request")
|
||||
}
|
||||
if isRejection(err) {
|
||||
t.Errorf("a database failure was reported as a credential problem: %v", err)
|
||||
}
|
||||
|
||||
c.fail.Store(false)
|
||||
if _, err := a.Authenticate(context.Background(), secret); err != nil {
|
||||
t.Errorf("the failure was cached: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpiryIsNotCached(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s, _, _ := fixture(t)
|
||||
|
||||
secret, k, _ := store.NewKey("key2", "temp", "proj1")
|
||||
k.CreatedAt = time.Now()
|
||||
k.ExpiresAt = time.Now().Add(time.Hour)
|
||||
must(t, s.CreateKey(ctx, k))
|
||||
|
||||
now := time.Now()
|
||||
a := New(s, Config{TTL: time.Minute})
|
||||
a.now = func() time.Time { return now }
|
||||
|
||||
if _, err := a.Authenticate(ctx, secret); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Past the key's expiry but inside the cache TTL. The cached decision is
|
||||
// allowed to be stale by at most the TTL, which is the documented trade.
|
||||
now = now.Add(2 * time.Hour)
|
||||
if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrExpired) {
|
||||
t.Errorf("an expired key authenticated: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentAuthentication(t *testing.T) {
|
||||
s, secret, _ := fixture(t)
|
||||
a := New(s, DefaultConfig())
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, 64)
|
||||
for range 64 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := a.Authenticate(context.Background(), secret); err != nil {
|
||||
errs <- err
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMiddlewareGuardsEveryRouteButTheExemptOnes(t *testing.T) {
|
||||
s, secret, _ := fixture(t)
|
||||
a := New(s, DefaultConfig())
|
||||
|
||||
var reached atomic.Int64
|
||||
inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
reached.Add(1)
|
||||
if _, ok := FromContext(r.Context()); !ok && r.URL.Path != "/health" {
|
||||
t.Error("an authenticated request carried no principal")
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
srv := httptest.NewServer(Middleware(a, nil, "/health")(inner))
|
||||
defer srv.Close()
|
||||
|
||||
get := func(path, token string) int {
|
||||
req, _ := http.NewRequest(http.MethodGet, srv.URL+path, nil) //nolint:noctx // test
|
||||
if token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
return resp.StatusCode
|
||||
}
|
||||
|
||||
if code := get("/health", ""); code != http.StatusOK {
|
||||
t.Errorf("health needs no key, got %d", code)
|
||||
}
|
||||
if code := get("/v1/models", ""); code != http.StatusUnauthorized {
|
||||
t.Errorf("an unauthenticated request got %d", code)
|
||||
}
|
||||
if code := get("/v1/models", "sk-llmgw-bogus"); code != http.StatusUnauthorized {
|
||||
t.Errorf("a bad key got %d", code)
|
||||
}
|
||||
if code := get("/v1/models", secret); code != http.StatusOK {
|
||||
t.Errorf("a good key got %d", code)
|
||||
}
|
||||
if n := reached.Load(); n != 2 {
|
||||
t.Errorf("the handler was reached %d times; only health and the valid key should pass", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestARevokedKeyLooksExactlyLikeAnUnknownOne(t *testing.T) {
|
||||
// Telling a caller that a key exists but was revoked confirms the key was
|
||||
// real, which is information worth withholding.
|
||||
s, secret, keyID := fixture(t)
|
||||
must(t, s.RevokeKey(context.Background(), keyID, time.Now()))
|
||||
|
||||
a := New(s, DefaultConfig())
|
||||
srv := httptest.NewServer(Middleware(a, nil)(http.HandlerFunc(
|
||||
func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })))
|
||||
defer srv.Close()
|
||||
|
||||
body := func(token string) string {
|
||||
req, _ := http.NewRequest(http.MethodGet, srv.URL, nil) //nolint:noctx // test
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var e struct {
|
||||
Error struct{ Message string } `json:"error"`
|
||||
}
|
||||
_ = json.NewDecoder(resp.Body).Decode(&e)
|
||||
return e.Error.Message
|
||||
}
|
||||
|
||||
unknown := store.KeyPrefix + "0000000000000000000000000000000000000000000"
|
||||
if body(secret) != body(unknown) {
|
||||
t.Errorf("revoked says %q, unknown says %q", body(secret), body(unknown))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAStoreFailureDoesNotReadAsABadKey(t *testing.T) {
|
||||
s, secret, _ := fixture(t)
|
||||
c := &countingStore{Store: s}
|
||||
c.fail.Store(true)
|
||||
a := New(c, DefaultConfig())
|
||||
|
||||
srv := httptest.NewServer(Middleware(a, nil)(http.HandlerFunc(
|
||||
func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })))
|
||||
defer srv.Close()
|
||||
|
||||
req, _ := http.NewRequest(http.MethodGet, srv.URL, nil) //nolint:noctx // test
|
||||
req.Header.Set("Authorization", "Bearer "+secret)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
// 503 rather than 401: sending someone hunting for a key problem that does
|
||||
// not exist wastes the one thing an incident is short of.
|
||||
if resp.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Errorf("status %d", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBearerTokenAcceptsWhatClientsActuallySend(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"Bearer sk-llmgw-abc": "sk-llmgw-abc",
|
||||
"sk-llmgw-abc": "sk-llmgw-abc",
|
||||
"Bearer sk-llmgw-abc": "sk-llmgw-abc",
|
||||
"": "",
|
||||
}
|
||||
for header, want := range cases {
|
||||
r := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
if header != "" {
|
||||
r.Header.Set("Authorization", header)
|
||||
}
|
||||
if got := BearerToken(r); got != want {
|
||||
t.Errorf("Authorization %q -> %q, want %q", header, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// Some tools only manage to set this one.
|
||||
r := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
r.Header.Set("api-key", "sk-llmgw-xyz")
|
||||
if got := BearerToken(r); got != "sk-llmgw-xyz" {
|
||||
t.Errorf("api-key header -> %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLastUsedIsBatchedNotWrittenPerRequest(t *testing.T) {
|
||||
s, secret, keyID := fixture(t)
|
||||
ctx := context.Background()
|
||||
a := New(s, Config{TTL: time.Minute, TouchInterval: time.Hour})
|
||||
|
||||
for range 20 {
|
||||
if _, err := a.Authenticate(ctx, secret); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
// Nothing written yet: the field is not read in real time.
|
||||
k, _ := s.KeyByHash(ctx, store.HashKey(secret))
|
||||
if !k.LastUsedAt.IsZero() {
|
||||
t.Error("last-used was written on the request path")
|
||||
}
|
||||
|
||||
must(t, a.FlushTouches(ctx))
|
||||
k, _ = s.KeyByHash(ctx, store.HashKey(secret))
|
||||
if k.LastUsedAt.IsZero() {
|
||||
t.Error("last-used was never written")
|
||||
}
|
||||
if k.ID != keyID {
|
||||
t.Errorf("wrong key touched: %s", k.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExemptingTheLandingPageDoesNotExemptEverything(t *testing.T) {
|
||||
// A near-miss worth a test. "/" ends in a slash, so a prefix rule would
|
||||
// match every path there is: authentication still installed, still logging,
|
||||
// and letting the entire API through.
|
||||
cases := []struct {
|
||||
path string
|
||||
want bool
|
||||
}{
|
||||
{"/", true},
|
||||
{"/_static/base.css", true},
|
||||
{"/health", true},
|
||||
{"/v1/models", false},
|
||||
{"/v1/chat/completions", false},
|
||||
{"/healthz", false}, // not a prefix rule
|
||||
{"/_staticx", false}, // prefix rules end at the slash
|
||||
}
|
||||
exempt := []string{"/health", "/", "/_static/"}
|
||||
for _, c := range cases {
|
||||
if got := isExempt(c.path, exempt); got != c.want {
|
||||
t.Errorf("isExempt(%q) = %v, want %v", c.path, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAlterAlwaysChangesTheKey(t *testing.T) {
|
||||
// The helper the test above depends on. If it can return its input, the
|
||||
// test it serves silently stops testing anything.
|
||||
for range 500 {
|
||||
secret, _, err := store.NewKey("id", "n", "p")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if alter(secret) == secret {
|
||||
t.Fatalf("alter returned its input for %q", secret)
|
||||
}
|
||||
}
|
||||
// Including the case that made the original flaky: a key ending in the
|
||||
// replacement character.
|
||||
if got := alter("sk-llmgw-aaaa"); got == "sk-llmgw-aaaa" {
|
||||
t.Error("alter is a no-op when the key already ends in the replacement")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Middleware authenticates every request before it reaches the router.
|
||||
//
|
||||
// It is a wrapper rather than a check inside the handler so there is no path
|
||||
// into the API that forgets to call it. Adding a route cannot accidentally add
|
||||
// an unauthenticated one.
|
||||
func Middleware(a *Authenticator, log *slog.Logger, exempt ...string) func(http.Handler) http.Handler {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if isExempt(r.URL.Path, exempt) {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
|
||||
p, err := a.Authenticate(r.Context(), BearerToken(r))
|
||||
if err != nil {
|
||||
writeAuthError(w, err)
|
||||
// The key prefix, never the key: a rejected credential in a log
|
||||
// is still a credential.
|
||||
log.Debug("rejected a request", "path", r.URL.Path,
|
||||
"reason", err, "remote", r.RemoteAddr)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r.WithContext(WithPrincipal(r.Context(), p)))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// isExempt reports whether a path is served without a key.
|
||||
//
|
||||
// An entry ending in "/" is a prefix rule, with one exception that matters
|
||||
// enormously: "/" itself is exact. Treating it as a prefix would match every
|
||||
// path there is, so exempting the landing page would silently exempt the whole
|
||||
// API — authentication still installed, still logging, and letting everything
|
||||
// through.
|
||||
//
|
||||
// Metrics and pprof are not in this list because they live on the admin
|
||||
// listener, a separate port precisely so it can be firewalled rather than
|
||||
// authenticated.
|
||||
func isExempt(path string, exempt []string) bool {
|
||||
for _, e := range exempt {
|
||||
if path == e {
|
||||
return true
|
||||
}
|
||||
if e != "/" && strings.HasSuffix(e, "/") && strings.HasPrefix(path, e) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func writeAuthError(w http.ResponseWriter, err error) {
|
||||
status, kind, msg := http.StatusUnauthorized, "invalid_request_error", "invalid api key"
|
||||
switch {
|
||||
case errors.Is(err, ErrMissing):
|
||||
msg = "no api key was provided"
|
||||
case errors.Is(err, ErrExpired):
|
||||
msg = "this api key has expired"
|
||||
case errors.Is(err, ErrRevoked):
|
||||
// Deliberately the same message as an unknown key. Telling a caller
|
||||
// that a key exists but was revoked confirms the key was real.
|
||||
msg = "invalid api key"
|
||||
case errors.Is(err, ErrInvalid):
|
||||
default:
|
||||
// A store failure is ours, not the caller's, and must not read as a
|
||||
// rejected credential — that would send someone hunting for a key
|
||||
// problem that does not exist.
|
||||
status, kind, msg = http.StatusServiceUnavailable, "server_error",
|
||||
"authentication is temporarily unavailable"
|
||||
}
|
||||
|
||||
if status == http.StatusUnauthorized {
|
||||
// What the OpenAI SDKs look for to stop retrying.
|
||||
w.Header().Set("WWW-Authenticate", `Bearer realm="llmgw"`)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{
|
||||
"message": msg, "type": kind, "code": nil,
|
||||
}})
|
||||
}
|
||||
|
||||
// WriteForbiddenModel reports a model this key may not use.
|
||||
//
|
||||
// Separate from authentication because it happens later: the model is only
|
||||
// known once the body has been parsed.
|
||||
func WriteForbiddenModel(w http.ResponseWriter, model string) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{
|
||||
"message": "this api key is not permitted to use model " + model,
|
||||
"type": "invalid_request_error", "code": nil,
|
||||
}})
|
||||
}
|
||||
@@ -0,0 +1,437 @@
|
||||
// Package bench measures what a replica can actually do, so the numbers the
|
||||
// router depends on stop being guesses (plan §6.2.1).
|
||||
//
|
||||
// Three of them matter: the weight a replica gets when the scheduler compares
|
||||
// it to others, the queue depth admission treats as healthy, and the KV usage
|
||||
// it treats as full. All three are currently hand-written in configuration, and
|
||||
// a hand-written capacity number is a hand-written opinion about hardware
|
||||
// nobody measured.
|
||||
//
|
||||
// The measurement is a load ramp. Concurrency climbs until the SLO breaks, and
|
||||
// the last level that still met it is the capacity — not peak throughput, which
|
||||
// a GPU will happily deliver long after latency has become unusable.
|
||||
package bench
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/rand/v2"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SLO is the latency target capacity is measured against.
|
||||
type SLO struct {
|
||||
TTFTp99 time.Duration
|
||||
TPOTp99 time.Duration
|
||||
}
|
||||
|
||||
// Met reports whether a step's latencies satisfy the target.
|
||||
func (s SLO) Met(ttft, tpot time.Duration) bool {
|
||||
if s.TTFTp99 > 0 && ttft > s.TTFTp99 {
|
||||
return false
|
||||
}
|
||||
if s.TPOTp99 > 0 && tpot > s.TPOTp99 {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// Config describes a run.
|
||||
type Config struct {
|
||||
// Target is the base URL: the gateway, or one replica directly. Pointing at
|
||||
// a replica measures that replica; pointing at the gateway measures the
|
||||
// whole path including scheduling, which is a different question.
|
||||
Target string
|
||||
// MetricsURL is an engine's /metrics, sampled during the run so queue depth
|
||||
// and KV usage can be reported alongside latency. Only meaningful when
|
||||
// Target is a single replica.
|
||||
MetricsURL string
|
||||
// Engine names which metric vocabulary MetricsURL speaks. Getting it wrong
|
||||
// is silent: nothing matches, every sample reads zero, and the suggested
|
||||
// thresholds are built on nothing.
|
||||
Engine string
|
||||
Model string
|
||||
APIKey string
|
||||
|
||||
// Concurrency is the ramp. Each level runs for Duration.
|
||||
Concurrency []int
|
||||
Duration time.Duration
|
||||
|
||||
PromptTokens int
|
||||
MaxTokens int
|
||||
|
||||
// SharedPrefix sends every request the same prompt.
|
||||
//
|
||||
// Off by default, and that default is the whole point: with a shared prefix
|
||||
// the engine's cache answers most of the prefill and the measurement
|
||||
// describes the cache rather than the hardware. Turning it on measures the
|
||||
// cached case deliberately, which is a fair question to ask separately.
|
||||
SharedPrefix bool
|
||||
|
||||
SLO SLO
|
||||
}
|
||||
|
||||
// DefaultConfig returns a ramp that covers most single-replica hardware.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Concurrency: []int{1, 2, 4, 8, 16, 32, 64, 128},
|
||||
Duration: 20 * time.Second,
|
||||
PromptTokens: 512,
|
||||
MaxTokens: 128,
|
||||
SLO: SLO{TTFTp99: 2 * time.Second, TPOTp99: 50 * time.Millisecond},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) normalize() error {
|
||||
d := DefaultConfig()
|
||||
if c.Target == "" {
|
||||
return errors.New("bench: a target is required")
|
||||
}
|
||||
if c.Model == "" {
|
||||
return errors.New("bench: a model is required")
|
||||
}
|
||||
c.Target = strings.TrimRight(c.Target, "/")
|
||||
if len(c.Concurrency) == 0 {
|
||||
c.Concurrency = d.Concurrency
|
||||
}
|
||||
if c.Duration <= 0 {
|
||||
c.Duration = d.Duration
|
||||
}
|
||||
if c.PromptTokens <= 0 {
|
||||
c.PromptTokens = d.PromptTokens
|
||||
}
|
||||
if c.MaxTokens <= 0 {
|
||||
c.MaxTokens = d.MaxTokens
|
||||
}
|
||||
sort.Ints(c.Concurrency)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Percentiles summarizes a latency sample.
|
||||
type Percentiles struct {
|
||||
P50, P90, P99, Max time.Duration
|
||||
}
|
||||
|
||||
// Step is one concurrency level's result.
|
||||
type Step struct {
|
||||
Concurrency int
|
||||
Completed int
|
||||
Failed int
|
||||
// Requests and output tokens per second, over the step's wall clock.
|
||||
RequestsPerSec float64
|
||||
TokensPerSec float64
|
||||
TTFT Percentiles
|
||||
// TPOT is time per output token after the first, which is what a reader
|
||||
// experiences as the speed of the stream.
|
||||
TPOT Percentiles
|
||||
Waiting float64
|
||||
Running float64
|
||||
KVUsage float64
|
||||
MeetsSLO bool
|
||||
// FirstError is kept so a step that failed says why rather than just
|
||||
// reporting zeroes.
|
||||
FirstError string
|
||||
}
|
||||
|
||||
// Result is the whole ramp plus what it concluded.
|
||||
type Result struct {
|
||||
Model string
|
||||
Target string
|
||||
Steps []Step
|
||||
SLO SLO
|
||||
RanAt time.Time
|
||||
Elapsed time.Duration
|
||||
|
||||
// MaxConcurrencyAtSLO is the capacity: the highest level that still met the
|
||||
// target. Zero means even one request did not.
|
||||
MaxConcurrencyAtSLO int
|
||||
// ThroughputTokPerSec and the latencies are taken at that level.
|
||||
ThroughputTokPerSec float64
|
||||
TTFTp99, TPOTp99 time.Duration
|
||||
// SuggestedTargetWaiting and SuggestedKVCeiling are what admission should
|
||||
// treat as healthy, observed rather than guessed.
|
||||
SuggestedTargetWaiting int
|
||||
SuggestedKVCeiling float64
|
||||
}
|
||||
|
||||
// Run executes the ramp.
|
||||
func Run(ctx context.Context, cfg Config, progress func(Step)) (Result, error) {
|
||||
if err := cfg.normalize(); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
res := Result{Model: cfg.Model, Target: cfg.Target, SLO: cfg.SLO, RanAt: time.Now()}
|
||||
start := time.Now()
|
||||
|
||||
client := &http.Client{Timeout: 10 * time.Minute}
|
||||
for _, c := range cfg.Concurrency {
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
step := runStep(ctx, client, cfg, c)
|
||||
res.Steps = append(res.Steps, step)
|
||||
if progress != nil {
|
||||
progress(step)
|
||||
}
|
||||
// Stop climbing once the target is missed. Everything above is slower
|
||||
// by construction, and the cost of finding out is real GPU time.
|
||||
if !step.MeetsSLO {
|
||||
break
|
||||
}
|
||||
}
|
||||
res.Elapsed = time.Since(start)
|
||||
res.conclude()
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (r *Result) conclude() {
|
||||
for _, s := range r.Steps {
|
||||
if !s.MeetsSLO || s.Completed == 0 {
|
||||
continue
|
||||
}
|
||||
if s.Concurrency <= r.MaxConcurrencyAtSLO {
|
||||
continue
|
||||
}
|
||||
r.MaxConcurrencyAtSLO = s.Concurrency
|
||||
r.ThroughputTokPerSec = s.TokensPerSec
|
||||
r.TTFTp99, r.TPOTp99 = s.TTFT.P99, s.TPOT.P99
|
||||
// The queue depth and cache pressure seen while still meeting the SLO
|
||||
// are exactly what admission should treat as healthy.
|
||||
r.SuggestedTargetWaiting = int(s.Waiting + 0.5)
|
||||
r.SuggestedKVCeiling = roundTo(s.KVUsage, 0.05)
|
||||
}
|
||||
}
|
||||
|
||||
func runStep(ctx context.Context, client *http.Client, cfg Config, concurrency int) Step {
|
||||
step := Step{Concurrency: concurrency}
|
||||
ctx, cancel := context.WithTimeout(ctx, cfg.Duration)
|
||||
defer cancel()
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
ttfts []time.Duration
|
||||
tpots []time.Duration
|
||||
tokens int64
|
||||
failed int
|
||||
firstErr string
|
||||
)
|
||||
|
||||
sampler := newSampler(cfg.MetricsURL, cfg.Engine)
|
||||
sampler.start(ctx)
|
||||
|
||||
start := time.Now()
|
||||
var wg sync.WaitGroup
|
||||
for range concurrency {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for ctx.Err() == nil {
|
||||
r, err := one(ctx, client, cfg)
|
||||
mu.Lock()
|
||||
switch {
|
||||
case err != nil:
|
||||
// A cancelled context is the step ending, not a failure.
|
||||
if ctx.Err() == nil {
|
||||
failed++
|
||||
if firstErr == "" {
|
||||
firstErr = err.Error()
|
||||
}
|
||||
}
|
||||
default:
|
||||
ttfts = append(ttfts, r.ttft)
|
||||
if r.tpot > 0 {
|
||||
tpots = append(tpots, r.tpot)
|
||||
}
|
||||
tokens += r.tokens
|
||||
}
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
elapsed := time.Since(start)
|
||||
sampler.stop()
|
||||
|
||||
step.Completed = len(ttfts)
|
||||
step.Failed = failed
|
||||
step.FirstError = firstErr
|
||||
step.TTFT = percentiles(ttfts)
|
||||
step.TPOT = percentiles(tpots)
|
||||
if elapsed > 0 {
|
||||
step.RequestsPerSec = float64(step.Completed) / elapsed.Seconds()
|
||||
step.TokensPerSec = float64(tokens) / elapsed.Seconds()
|
||||
}
|
||||
step.Waiting, step.Running, step.KVUsage = sampler.averages()
|
||||
|
||||
// A step that could not complete anything has not met anything.
|
||||
step.MeetsSLO = step.Completed > 0 && step.Failed == 0 &&
|
||||
cfg.SLO.Met(step.TTFT.P99, step.TPOT.P99)
|
||||
return step
|
||||
}
|
||||
|
||||
type outcome struct {
|
||||
ttft time.Duration
|
||||
tpot time.Duration
|
||||
tokens int64
|
||||
}
|
||||
|
||||
// one sends a single streaming completion and times it.
|
||||
func one(ctx context.Context, client *http.Client, cfg Config) (outcome, error) {
|
||||
body, err := json.Marshal(map[string]any{
|
||||
"model": cfg.Model,
|
||||
"stream": true,
|
||||
"max_tokens": cfg.MaxTokens,
|
||||
"messages": []map[string]string{
|
||||
{"role": "user", "content": prompt(cfg)},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return outcome{}, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
|
||||
cfg.Target+"/v1/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return outcome{}, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if cfg.APIKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+cfg.APIKey)
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return outcome{}, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
return outcome{}, fmt.Errorf("%s: %s", resp.Status, strings.TrimSpace(string(b)))
|
||||
}
|
||||
|
||||
var out outcome
|
||||
var firstAt, lastAt time.Time
|
||||
sc := bufio.NewScanner(resp.Body)
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 1<<20)
|
||||
for sc.Scan() {
|
||||
line := sc.Text()
|
||||
if !strings.HasPrefix(line, "data: ") || strings.HasSuffix(line, "[DONE]") {
|
||||
continue
|
||||
}
|
||||
// Only chunks that carry content count: a role-only first chunk is not
|
||||
// a token, and counting it would flatter the time to first token.
|
||||
if !hasContent(line[6:]) {
|
||||
continue
|
||||
}
|
||||
now := time.Now()
|
||||
if out.tokens == 0 {
|
||||
firstAt = now
|
||||
out.ttft = now.Sub(start)
|
||||
}
|
||||
lastAt = now
|
||||
out.tokens++
|
||||
}
|
||||
if err := sc.Err(); err != nil {
|
||||
return outcome{}, err
|
||||
}
|
||||
if out.tokens == 0 {
|
||||
return outcome{}, errors.New("the response carried no content")
|
||||
}
|
||||
if out.tokens > 1 {
|
||||
out.tpot = lastAt.Sub(firstAt) / time.Duration(out.tokens-1)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func hasContent(payload string) bool {
|
||||
var chunk struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"delta"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
||||
return false
|
||||
}
|
||||
for _, c := range chunk.Choices {
|
||||
if c.Delta.Content != "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// prompt builds a request body of roughly the configured size.
|
||||
//
|
||||
// Unique per request unless a shared prefix was asked for: repeating one prompt
|
||||
// lets the engine's cache answer most of the prefill, and the result then
|
||||
// describes the cache rather than the hardware.
|
||||
//
|
||||
// "Shared" has to mean byte-identical. Dropping only the random header while
|
||||
// still picking words at random would leave every prompt different after the
|
||||
// first few tokens, so the cache would miss and the flag would quietly measure
|
||||
// the thing it exists to avoid measuring.
|
||||
func prompt(cfg Config) string {
|
||||
var b strings.Builder
|
||||
if cfg.SharedPrefix {
|
||||
for i := 0; b.Len() < cfg.PromptTokens*4; i++ {
|
||||
b.WriteString(words[i%len(words)])
|
||||
b.WriteByte(' ')
|
||||
}
|
||||
} else {
|
||||
fmt.Fprintf(&b, "%016x %016x ", rand.Uint64(), rand.Uint64()) //nolint:gosec // not security
|
||||
// Roughly four characters to a token, the same crude ratio the router
|
||||
// uses; exactness is not what this number is for.
|
||||
for b.Len() < cfg.PromptTokens*4 {
|
||||
b.WriteString(words[rand.IntN(len(words))]) //nolint:gosec // not security
|
||||
b.WriteByte(' ')
|
||||
}
|
||||
}
|
||||
b.WriteString("\n\nSummarize the text above.")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
var words = strings.Fields(`the quick brown fox jumps over lazy dog while distant
|
||||
thunder rolls across an empty valley and somewhere a door closes softly in the
|
||||
dark house where nobody has lived for years although the garden still produces
|
||||
apples every autumn which fall and rot beneath the trees unnoticed by anyone
|
||||
except the birds that come each morning to pick at them before the frost`)
|
||||
|
||||
func percentiles(v []time.Duration) Percentiles {
|
||||
if len(v) == 0 {
|
||||
return Percentiles{}
|
||||
}
|
||||
s := make([]time.Duration, len(v))
|
||||
copy(s, v)
|
||||
sort.Slice(s, func(i, j int) bool { return s[i] < s[j] })
|
||||
at := func(p float64) time.Duration {
|
||||
i := int(float64(len(s)-1) * p)
|
||||
return s[i]
|
||||
}
|
||||
return Percentiles{P50: at(0.50), P90: at(0.90), P99: at(0.99), Max: s[len(s)-1]}
|
||||
}
|
||||
|
||||
// roundTo rounds to a multiple of step, through integers.
|
||||
//
|
||||
// The obvious float arithmetic returns 0.7000000000000001 for what should be
|
||||
// 0.70, which is harmless on a terminal and ugly in the JSON someone pastes
|
||||
// into a configuration file.
|
||||
func roundTo(v, step float64) float64 {
|
||||
if step <= 0 {
|
||||
return v
|
||||
}
|
||||
const scale = 10000
|
||||
units := int64(v/step + 0.5)
|
||||
return float64(int64(float64(units)*step*scale+0.5)) / scale
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
package bench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCapacityIsTheLastLevelThatMetTheSLONotThePeak(t *testing.T) {
|
||||
// A GPU keeps delivering more tokens per second long after latency has
|
||||
// become unusable. Taking peak throughput as capacity would tell the router
|
||||
// a replica can serve a load at which nobody would want the answers.
|
||||
r := Result{SLO: SLO{TTFTp99: time.Second}}
|
||||
r.Steps = []Step{
|
||||
{Concurrency: 1, Completed: 10, TokensPerSec: 100, MeetsSLO: true,
|
||||
TTFT: Percentiles{P99: 100 * time.Millisecond}},
|
||||
{Concurrency: 8, Completed: 80, TokensPerSec: 700, MeetsSLO: true,
|
||||
TTFT: Percentiles{P99: 900 * time.Millisecond}, Waiting: 3.4, KVUsage: 0.71},
|
||||
// Faster in tokens, far too slow in latency.
|
||||
{Concurrency: 32, Completed: 200, TokensPerSec: 1500, MeetsSLO: false,
|
||||
TTFT: Percentiles{P99: 9 * time.Second}},
|
||||
}
|
||||
r.conclude()
|
||||
|
||||
if r.MaxConcurrencyAtSLO != 8 {
|
||||
t.Errorf("capacity = %d, want 8", r.MaxConcurrencyAtSLO)
|
||||
}
|
||||
if r.ThroughputTokPerSec != 700 {
|
||||
t.Errorf("throughput = %v; it must come from the level that met the SLO", r.ThroughputTokPerSec)
|
||||
}
|
||||
if r.SuggestedTargetWaiting != 3 {
|
||||
t.Errorf("target_waiting = %d, want 3", r.SuggestedTargetWaiting)
|
||||
}
|
||||
if r.SuggestedKVCeiling != 0.70 {
|
||||
t.Errorf("kv_ceiling = %v, want 0.70", r.SuggestedKVCeiling)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAStepWithFailuresHasNotMetAnything(t *testing.T) {
|
||||
// Latency percentiles over the requests that survived say nothing useful
|
||||
// when the rest were refused.
|
||||
r := Result{SLO: SLO{TTFTp99: time.Second}}
|
||||
r.Steps = []Step{
|
||||
{Concurrency: 4, Completed: 10, Failed: 90, TokensPerSec: 50, MeetsSLO: false,
|
||||
TTFT: Percentiles{P99: 50 * time.Millisecond}},
|
||||
}
|
||||
r.conclude()
|
||||
if r.MaxConcurrencyAtSLO != 0 {
|
||||
t.Errorf("a step that failed 90 of 100 requests was counted as capacity")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSLOIsMissedWhenEitherHalfIs(t *testing.T) {
|
||||
s := SLO{TTFTp99: time.Second, TPOTp99: 50 * time.Millisecond}
|
||||
cases := []struct {
|
||||
ttft, tpot time.Duration
|
||||
want bool
|
||||
}{
|
||||
{500 * time.Millisecond, 20 * time.Millisecond, true},
|
||||
{2 * time.Second, 20 * time.Millisecond, false},
|
||||
{500 * time.Millisecond, 80 * time.Millisecond, false},
|
||||
{2 * time.Second, 80 * time.Millisecond, false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := s.Met(c.ttft, c.tpot); got != c.want {
|
||||
t.Errorf("Met(%s, %s) = %v", c.ttft, c.tpot, got)
|
||||
}
|
||||
}
|
||||
|
||||
// A target of zero is no target, not a target of zero.
|
||||
none := SLO{}
|
||||
if !none.Met(time.Hour, time.Hour) {
|
||||
t.Error("an unset SLO refused everything")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPercentiles(t *testing.T) {
|
||||
var v []time.Duration
|
||||
for i := 1; i <= 100; i++ {
|
||||
v = append(v, time.Duration(i)*time.Millisecond)
|
||||
}
|
||||
p := percentiles(v)
|
||||
if p.P50 != 50*time.Millisecond || p.P90 != 90*time.Millisecond {
|
||||
t.Errorf("p50=%s p90=%s", p.P50, p.P90)
|
||||
}
|
||||
if p.Max != 100*time.Millisecond {
|
||||
t.Errorf("max=%s", p.Max)
|
||||
}
|
||||
if got := percentiles(nil); got != (Percentiles{}) {
|
||||
t.Errorf("empty sample gave %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptsAreUniqueUnlessAskedNotToBe(t *testing.T) {
|
||||
// The default matters more than it looks. With one repeated prompt the
|
||||
// engine's cache answers most of the prefill, and the measurement describes
|
||||
// the cache rather than the hardware.
|
||||
cfg := Config{PromptTokens: 64}
|
||||
seen := map[string]bool{}
|
||||
for range 100 {
|
||||
seen[prompt(cfg)] = true
|
||||
}
|
||||
if len(seen) < 100 {
|
||||
t.Errorf("%d distinct prompts out of 100; the cache would answer the rest", len(seen))
|
||||
}
|
||||
|
||||
cfg.SharedPrefix = true
|
||||
first := prompt(cfg)
|
||||
for range 20 {
|
||||
if got := prompt(cfg); !strings.HasPrefix(got, first[:200]) {
|
||||
t.Error("--shared-prefix did not produce a shared prefix")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptSizeFollowsTheRequest(t *testing.T) {
|
||||
small := prompt(Config{PromptTokens: 32})
|
||||
large := prompt(Config{PromptTokens: 512})
|
||||
if len(large) < len(small)*4 {
|
||||
t.Errorf("a 512-token prompt is %d chars against %d for 32", len(large), len(small))
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfigIsValidatedBeforeAnyLoadIsGenerated(t *testing.T) {
|
||||
// The failure has to come before the traffic, not halfway through a ramp.
|
||||
if _, err := Run(context.Background(), Config{Model: "m"}, nil); err == nil {
|
||||
t.Error("a run with no target started")
|
||||
}
|
||||
if _, err := Run(context.Background(), Config{Target: "http://x"}, nil); err == nil {
|
||||
t.Error("a run with no model started")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheRampStopsOnceTheSLOBreaks(t *testing.T) {
|
||||
// Everything above a missed level is slower by construction, and finding
|
||||
// that out costs real GPU time.
|
||||
r := Result{}
|
||||
r.Steps = []Step{
|
||||
{Concurrency: 1, Completed: 1, MeetsSLO: true},
|
||||
{Concurrency: 2, Completed: 1, MeetsSLO: false},
|
||||
}
|
||||
r.conclude()
|
||||
if r.MaxConcurrencyAtSLO != 1 {
|
||||
t.Errorf("capacity = %d", r.MaxConcurrencyAtSLO)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundTo(t *testing.T) {
|
||||
cases := map[float64]float64{0.71: 0.70, 0.73: 0.75, 0.88: 0.90, 0.0: 0.0}
|
||||
for in, want := range cases {
|
||||
if got := roundTo(in, 0.05); got != want {
|
||||
t.Errorf("roundTo(%v) = %v, want %v", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package bench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape"
|
||||
)
|
||||
|
||||
// sampler reads an engine's metrics during a step.
|
||||
//
|
||||
// It is what turns a latency measurement into something admission can use: the
|
||||
// queue depth and cache pressure observed while the SLO was still being met are
|
||||
// the values target_waiting and kv_ceiling should have, rather than the round
|
||||
// numbers somebody typed.
|
||||
//
|
||||
// Optional, because a run pointed at the gateway has no single engine to read.
|
||||
type sampler struct {
|
||||
url string
|
||||
engine string
|
||||
client *http.Client
|
||||
cancel context.CancelFunc
|
||||
|
||||
mu sync.Mutex
|
||||
wg sync.WaitGroup
|
||||
waiting float64
|
||||
running float64
|
||||
kv float64
|
||||
n int
|
||||
}
|
||||
|
||||
func newSampler(metricsURL, engine string) *sampler {
|
||||
return &sampler{
|
||||
url: metricsURL,
|
||||
engine: engine,
|
||||
client: &http.Client{Timeout: 2 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *sampler) start(ctx context.Context) {
|
||||
if s.url == "" {
|
||||
return
|
||||
}
|
||||
ctx, s.cancel = context.WithCancel(ctx)
|
||||
s.wg.Add(1)
|
||||
go func() {
|
||||
defer s.wg.Done()
|
||||
t := time.NewTicker(time.Second)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
s.sample(ctx)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *sampler) sample(ctx context.Context) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.url, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
resp, err := s.client.Do(req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return
|
||||
}
|
||||
|
||||
// The same mapping the router reads with, so a bench result cannot disagree
|
||||
// with what the scheduler sees. The engine has to be named: reading a
|
||||
// llama.cpp endpoint with vLLM's names finds nothing, and the run would
|
||||
// report a queue depth of zero and suggest thresholds built on it.
|
||||
stats, err := scrape.Parse(resp.Body, scrape.MappingFor(s.engine))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.waiting += float64(stats.Waiting)
|
||||
s.running += float64(stats.Running)
|
||||
s.kv += stats.KVCacheUsage
|
||||
s.n++
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *sampler) stop() {
|
||||
if s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
s.wg.Wait()
|
||||
}
|
||||
|
||||
func (s *sampler) averages() (waiting, running, kv float64) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.n == 0 {
|
||||
return 0, 0, 0
|
||||
}
|
||||
n := float64(s.n)
|
||||
return s.waiting / n, s.running / n, s.kv / n
|
||||
}
|
||||
@@ -64,7 +64,12 @@ func newAgentCmd() *cobra.Command {
|
||||
}
|
||||
cfg.Version = version.Version
|
||||
|
||||
a, err := agent.New(cfg, agent.HTTP2Client(30*time.Second), log)
|
||||
httpClient, err := agent.Client(cfg.TLS, 30*time.Second)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
a, err := agent.New(cfg, httpClient, log)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/bench"
|
||||
)
|
||||
|
||||
func newBenchCmd(asJSON *bool) *cobra.Command {
|
||||
cfg := bench.DefaultConfig()
|
||||
var (
|
||||
steps []int
|
||||
target string
|
||||
metrics string
|
||||
engine string
|
||||
key string
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "bench <model>",
|
||||
Short: "Measure what a replica can do, and what the router should be told",
|
||||
Long: "Drives real traffic. Point it at a replica to measure that hardware, or at\n" +
|
||||
"the gateway to measure the whole path — they answer different questions,\n" +
|
||||
"and only the first produces a capacity profile.\n\n" +
|
||||
"Concurrency climbs until the SLO breaks. The last level that still met it\n" +
|
||||
"is the capacity: peak throughput arrives long after latency has become\n" +
|
||||
"unusable, so throughput alone would overstate what the hardware can serve.\n\n" +
|
||||
"This costs real GPU time and competes with real traffic. Run it on a\n" +
|
||||
"replica you have cordoned, or during a window where that is acceptable.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
cfg.Model = args[0]
|
||||
cfg.Target = target
|
||||
cfg.MetricsURL = metrics
|
||||
cfg.Engine = engine
|
||||
cfg.APIKey = key
|
||||
if len(steps) > 0 {
|
||||
cfg.Concurrency = steps
|
||||
}
|
||||
|
||||
if metrics == "" {
|
||||
cmd.PrintErrln("No --metrics endpoint: queue depth and KV usage will not be " +
|
||||
"measured, so target_waiting and kv_ceiling cannot be suggested.")
|
||||
}
|
||||
cmd.PrintErrf("benchmarking %s against %s\n\n", cfg.Model, cfg.Target)
|
||||
|
||||
// Fixed columns rather than a tabwriter: rows are printed as each
|
||||
// level finishes, and a tabwriter cannot align a column it has not
|
||||
// seen the rest of yet — it lays out whatever has been flushed, so
|
||||
// every row after the first came out ragged.
|
||||
const row = "%-5s %-7s %-7s %-9s %-9s %-9s %-6s %-5s %s\n"
|
||||
fmt.Fprintf(cmd.ErrOrStderr(), row,
|
||||
"CONC", "REQ/S", "TOK/S", "TTFT p50", "TTFT p99", "TPOT p99", "WAIT", "KV", "SLO")
|
||||
|
||||
progress := func(s bench.Step) {
|
||||
verdict := "ok"
|
||||
if !s.MeetsSLO {
|
||||
verdict = "MISSED"
|
||||
}
|
||||
if s.Failed > 0 {
|
||||
verdict = fmt.Sprintf("%d failed", s.Failed)
|
||||
}
|
||||
fmt.Fprintf(cmd.ErrOrStderr(), row,
|
||||
fmt.Sprint(s.Concurrency),
|
||||
fmt.Sprintf("%.1f", s.RequestsPerSec),
|
||||
fmt.Sprintf("%.0f", s.TokensPerSec),
|
||||
ms(s.TTFT.P50), ms(s.TTFT.P99), ms(s.TPOT.P99),
|
||||
fmt.Sprintf("%.1f", s.Waiting),
|
||||
fmt.Sprintf("%.0f%%", s.KVUsage*100),
|
||||
verdict)
|
||||
}
|
||||
|
||||
res, err := bench.Run(cmd.Context(), cfg, progress)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(res)
|
||||
}
|
||||
printConclusion(cmd, res, cfg)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
f := cmd.Flags()
|
||||
f.StringVar(&target, "target", "http://127.0.0.1:8080", "base URL: a replica, or the gateway")
|
||||
f.StringVar(&metrics, "metrics", "", "the replica's /metrics, to observe queue depth and KV usage")
|
||||
f.StringVar(&engine, "engine", "", "which metric names --metrics uses: vllm | sglang | llamacpp")
|
||||
f.StringVar(&key, "key", "", "API key, when the gateway requires one")
|
||||
f.IntSliceVar(&steps, "concurrency", nil, "levels to try (default 1,2,4,…,128)")
|
||||
f.DurationVar(&cfg.Duration, "duration", cfg.Duration, "how long to hold each level")
|
||||
f.IntVar(&cfg.PromptTokens, "prompt-tokens", cfg.PromptTokens, "approximate prompt size")
|
||||
f.IntVar(&cfg.MaxTokens, "max-tokens", cfg.MaxTokens, "tokens to generate per request")
|
||||
f.BoolVar(&cfg.SharedPrefix, "shared-prefix", false,
|
||||
"send every request the same prompt, measuring the cached case instead of the hardware")
|
||||
f.DurationVar(&cfg.SLO.TTFTp99, "ttft-p99", cfg.SLO.TTFTp99, "time-to-first-token target")
|
||||
f.DurationVar(&cfg.SLO.TPOTp99, "tpot-p99", cfg.SLO.TPOTp99, "per-output-token target")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func printConclusion(cmd *cobra.Command, res bench.Result, cfg bench.Config) {
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintln(out)
|
||||
|
||||
if res.MaxConcurrencyAtSLO == 0 {
|
||||
fmt.Fprintf(out, "No level met the SLO (TTFT p99 ≤ %s, TPOT p99 ≤ %s).\n",
|
||||
cfg.SLO.TTFTp99, cfg.SLO.TPOTp99)
|
||||
if len(res.Steps) > 0 && res.Steps[0].FirstError != "" {
|
||||
fmt.Fprintf(out, "First error: %s\n", res.Steps[0].FirstError)
|
||||
}
|
||||
fmt.Fprintln(out, "Either the hardware cannot serve this model at this SLO, or the")
|
||||
fmt.Fprintln(out, "target is stricter than the model can meet at any concurrency.")
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Fprintf(out, "Capacity: %d concurrent requests at %.0f tokens/s\n",
|
||||
res.MaxConcurrencyAtSLO, res.ThroughputTokPerSec)
|
||||
fmt.Fprintf(out, " TTFT p99 %s · TPOT p99 %s · measured over %s\n\n",
|
||||
ms(res.TTFTp99), ms(res.TPOTp99), res.Elapsed.Round(time.Second))
|
||||
|
||||
fmt.Fprintln(out, "Put this in the model's profile:")
|
||||
fmt.Fprintf(out, " weight: %d\n", res.MaxConcurrencyAtSLO)
|
||||
if res.SuggestedTargetWaiting > 0 || res.SuggestedKVCeiling > 0 {
|
||||
fmt.Fprintln(out, "\nAnd in router.scoring, from what the engine reported while still")
|
||||
fmt.Fprintln(out, "meeting the target:")
|
||||
if res.SuggestedTargetWaiting > 0 {
|
||||
fmt.Fprintf(out, " target_waiting: %d\n", res.SuggestedTargetWaiting)
|
||||
}
|
||||
if res.SuggestedKVCeiling > 0 {
|
||||
fmt.Fprintf(out, " kv_ceiling: %.2f\n", res.SuggestedKVCeiling)
|
||||
}
|
||||
}
|
||||
|
||||
if last := res.Steps[len(res.Steps)-1]; !last.MeetsSLO {
|
||||
fmt.Fprintf(out, "\n%d concurrent was the first level to miss it (TTFT p99 %s).\n",
|
||||
last.Concurrency, ms(last.TTFT.P99))
|
||||
} else {
|
||||
fmt.Fprintf(out, "\nEvery level tried met the SLO, so %d is a floor rather than a\n",
|
||||
res.MaxConcurrencyAtSLO)
|
||||
fmt.Fprintln(out, "ceiling — raise --concurrency to find where it actually breaks.")
|
||||
}
|
||||
if cfg.SharedPrefix {
|
||||
fmt.Fprintln(out, "\nRun with --shared-prefix: this measures the engine's prefix cache,")
|
||||
fmt.Fprintln(out, "not its hardware. Do not use these numbers as capacity.")
|
||||
}
|
||||
}
|
||||
|
||||
func ms(d time.Duration) string {
|
||||
switch {
|
||||
case d == 0:
|
||||
return "-"
|
||||
case d < time.Millisecond:
|
||||
return d.Round(10 * time.Microsecond).String()
|
||||
case d < time.Second:
|
||||
return d.Round(time.Millisecond).String()
|
||||
default:
|
||||
return d.Round(10 * time.Millisecond).String()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"crypto/x509"
|
||||
"encoding/pem"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
// defaultPKIDir is where the authority lives on the machine that issues
|
||||
// certificates. Only that machine ever holds ca-key.pem; everything else gets
|
||||
// ca.pem and its own leaf.
|
||||
const defaultPKIDir = "/etc/llmgw/pki"
|
||||
|
||||
func newCertCmd() *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "cert",
|
||||
Short: "Issue the certificates the control plane and agents authenticate with",
|
||||
Long: "llmgw's internal connections use mutual TLS, and a public authority cannot\n" +
|
||||
"supply them: Let's Encrypt and its peers stopped issuing certificates\n" +
|
||||
"carrying the clientAuth usage, so nothing they hand out can authenticate an\n" +
|
||||
"agent. These commands run a small private authority instead.\n\n" +
|
||||
"They work on files and talk to nothing, so they run wherever you keep the\n" +
|
||||
"authority — not necessarily on a cluster machine.",
|
||||
}
|
||||
cmd.AddCommand(newInitCACmd(), newIssueServerCmd(), newIssueClientCmd(), newCertInspectCmd())
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newInitCACmd() *cobra.Command {
|
||||
var (
|
||||
dir string
|
||||
cn string
|
||||
force bool
|
||||
)
|
||||
cmd := &cobra.Command{
|
||||
Use: "init-ca",
|
||||
Short: "Create the cluster's certificate authority",
|
||||
Long: "Writes ca.pem and ca-key.pem. Copy ca.pem to every machine; leave\n" +
|
||||
"ca-key.pem here, because whoever holds it can mint any node's identity.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
keyPath := filepath.Join(dir, "ca-key.pem")
|
||||
if _, err := os.Stat(keyPath); err == nil && !force {
|
||||
// Overwriting the key orphans every certificate ever issued from
|
||||
// it, across every machine, with no warning until connections
|
||||
// start failing.
|
||||
return fmt.Errorf("%s already exists; re-running would invalidate every "+
|
||||
"certificate issued from it (pass --force if that is what you want)", keyPath)
|
||||
}
|
||||
|
||||
ca, err := tlsconf.NewCA(cn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := ca.Write(dir); err != nil {
|
||||
return err
|
||||
}
|
||||
cmd.Printf("authority %q written to %s\n", cn, dir)
|
||||
cmd.Printf(" ca.pem trust this everywhere (valid until %s)\n",
|
||||
ca.Cert.NotAfter.Format(time.DateOnly))
|
||||
cmd.Printf(" ca-key.pem keep it here, mode 0600\n")
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&dir, "dir", defaultPKIDir, "directory holding the authority")
|
||||
cmd.Flags().StringVar(&cn, "cn", "llmgw cluster CA", "common name for the authority")
|
||||
cmd.Flags().BoolVar(&force, "force", false, "replace an existing authority")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newIssueServerCmd() *cobra.Command {
|
||||
var (
|
||||
dir string
|
||||
out string
|
||||
hosts []string
|
||||
name string
|
||||
)
|
||||
cmd := &cobra.Command{
|
||||
Use: "issue-server <common-name>",
|
||||
Short: "Issue the control plane's certificate",
|
||||
Long: "Every address an agent might dial must appear in --host, including bare\n" +
|
||||
"IP addresses. Hostname verification reads subject alternative names and\n" +
|
||||
"ignores the common name entirely, so a certificate that names the host\n" +
|
||||
"only in its subject is rejected by every client.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
ca, err := loadCA(dir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
leaf, err := ca.IssueServer(args[0], hosts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if name == "" {
|
||||
name = "server"
|
||||
}
|
||||
if err := leaf.Write(out, name); err != nil {
|
||||
return err
|
||||
}
|
||||
return describeIssued(cmd, out, name, hosts)
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&dir, "dir", defaultPKIDir, "directory holding the authority")
|
||||
cmd.Flags().StringVar(&out, "out", ".", "directory to write the certificate into")
|
||||
cmd.Flags().StringVar(&name, "name", "server", "base filename for the certificate")
|
||||
cmd.Flags().StringArrayVar(&hosts, "host", nil, "address agents dial; repeatable (IP or DNS name)")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newIssueClientCmd() *cobra.Command {
|
||||
var dir, out string
|
||||
cmd := &cobra.Command{
|
||||
Use: "issue-client <node-id>",
|
||||
Short: "Issue one agent's certificate",
|
||||
Long: "The node id becomes the certificate's common name, and the control plane\n" +
|
||||
"refuses any agent whose claimed id does not match it. That is what makes a\n" +
|
||||
"node's identity something it proves rather than something it picks.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
ca, err := loadCA(dir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
leaf, err := ca.IssueClient(args[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := leaf.Write(out, args[0]); err != nil {
|
||||
return err
|
||||
}
|
||||
return describeIssued(cmd, out, args[0], nil)
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&dir, "dir", defaultPKIDir, "directory holding the authority")
|
||||
cmd.Flags().StringVar(&out, "out", ".", "directory to write the certificate into")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newCertInspectCmd() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "inspect <file.pem>",
|
||||
Short: "Print what a certificate actually says",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
cert, err := readCert(args[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cmd.Printf("subject %s\n", cert.Subject)
|
||||
cmd.Printf("issuer %s\n", cert.Issuer)
|
||||
cmd.Printf("valid %s .. %s", cert.NotBefore.Format(time.DateOnly),
|
||||
cert.NotAfter.Format(time.DateOnly))
|
||||
if d := time.Until(cert.NotAfter); d < 0 {
|
||||
cmd.Printf(" (EXPIRED)")
|
||||
} else {
|
||||
cmd.Printf(" (%d days left)", int(d.Hours()/24))
|
||||
}
|
||||
cmd.Println()
|
||||
|
||||
if cert.IsCA {
|
||||
cmd.Println("role certificate authority")
|
||||
}
|
||||
if u := usageNames(cert.ExtKeyUsage); u != "" {
|
||||
cmd.Printf("usage %s\n", u)
|
||||
}
|
||||
if sans := sanNames(cert); sans != "" {
|
||||
cmd.Printf("hosts %s\n", sans)
|
||||
} else if !cert.IsCA && !hasUsage(cert, x509.ExtKeyUsageClientAuth) {
|
||||
cmd.Println("hosts none — no client will accept this as a server certificate")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func loadCA(dir string) (*tlsconf.CA, error) {
|
||||
ca, err := tlsconf.LoadCA(filepath.Join(dir, "ca.pem"), filepath.Join(dir, "ca-key.pem"))
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("no authority in %s — run 'llmgw ctl cert init-ca --dir %s' first", dir, dir)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return ca, nil
|
||||
}
|
||||
|
||||
func describeIssued(cmd *cobra.Command, dir, name string, hosts []string) error {
|
||||
cert, err := readCert(filepath.Join(dir, name+".pem"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cmd.Printf("issued %s\n", cert.Subject.CommonName)
|
||||
cmd.Printf(" %s\n", filepath.Join(dir, name+".pem"))
|
||||
cmd.Printf(" %s\n", filepath.Join(dir, name+"-key.pem"))
|
||||
if len(hosts) > 0 {
|
||||
cmd.Printf(" hosts: %s\n", strings.Join(hosts, ", "))
|
||||
}
|
||||
cmd.Printf(" expires %s\n", cert.NotAfter.Format(time.DateOnly))
|
||||
return nil
|
||||
}
|
||||
|
||||
func readCert(path string) (*x509.Certificate, error) {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
block, _ := pem.Decode(b)
|
||||
if block == nil {
|
||||
return nil, fmt.Errorf("%s is not PEM", path)
|
||||
}
|
||||
return x509.ParseCertificate(block.Bytes)
|
||||
}
|
||||
|
||||
func hasUsage(cert *x509.Certificate, want x509.ExtKeyUsage) bool {
|
||||
for _, u := range cert.ExtKeyUsage {
|
||||
if u == want {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func usageNames(usages []x509.ExtKeyUsage) string {
|
||||
var out []string
|
||||
for _, u := range usages {
|
||||
switch u {
|
||||
case x509.ExtKeyUsageServerAuth:
|
||||
out = append(out, "serverAuth (identifies the control plane)")
|
||||
case x509.ExtKeyUsageClientAuth:
|
||||
out = append(out, "clientAuth (identifies an agent)")
|
||||
default:
|
||||
out = append(out, fmt.Sprintf("%v", u))
|
||||
}
|
||||
}
|
||||
return strings.Join(out, ", ")
|
||||
}
|
||||
|
||||
func sanNames(cert *x509.Certificate) string {
|
||||
names := append([]string{}, cert.DNSNames...)
|
||||
for _, ip := range cert.IPAddresses {
|
||||
names = append(names, ip.String())
|
||||
}
|
||||
return strings.Join(names, ", ")
|
||||
}
|
||||
+103
-12
@@ -37,6 +37,16 @@ func (c *ctlClient) get(ctx context.Context, path string, out any) error {
|
||||
return c.do(ctx, http.MethodGet, path, nil, out)
|
||||
}
|
||||
|
||||
// post sends an already-encoded body, so a caller that has marshalled a typed
|
||||
// value keeps its field names rather than having them re-derived here.
|
||||
func (c *ctlClient) post(ctx context.Context, path string, body []byte, out any) error {
|
||||
return c.do(ctx, http.MethodPost, path, json.RawMessage(body), out)
|
||||
}
|
||||
|
||||
func (c *ctlClient) delete(ctx context.Context, path string, out any) error {
|
||||
return c.do(ctx, http.MethodDelete, path, nil, out)
|
||||
}
|
||||
|
||||
func (c *ctlClient) do(ctx context.Context, method, path string, body, out any) error {
|
||||
var rdr io.Reader
|
||||
if body != nil {
|
||||
@@ -54,6 +64,9 @@ func (c *ctlClient) do(ctx context.Context, method, path string, body, out any)
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
// Who made the change, kept with a stored model so a catalog that grew
|
||||
// over months can say where each entry came from.
|
||||
req.Header.Set("X-Llmgw-Actor", actor())
|
||||
|
||||
resp, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
@@ -106,6 +119,15 @@ func newCtlCmd() *cobra.Command {
|
||||
"Allow placement on a node again"),
|
||||
newNodeStateCmd(client, "drain", "draining",
|
||||
"Move work off a node, one replica at a time"),
|
||||
// cert works on files rather than the admin API, so --server does not
|
||||
// apply to it; it lives here because this is where operators look.
|
||||
newBenchCmd(&asJSON),
|
||||
newStatusCmd(client, &asJSON),
|
||||
newModelCmd(client),
|
||||
newEngineCheckCmd(),
|
||||
newToolCheckCmd(),
|
||||
newWeightsCmd(client, &asJSON),
|
||||
newCertCmd(),
|
||||
)
|
||||
return cmd
|
||||
}
|
||||
@@ -128,10 +150,11 @@ func newNodesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "NODE\tSTATE\tGPUS\tCLASSES\tFABRICS\tREPLICAS\tLAST SEEN")
|
||||
fmt.Fprintln(w, "NODE\tSTATE\tGPUS\tMEM\tDISK\tCLASSES\tFABRICS\tREPLICAS\tLAST SEEN")
|
||||
for _, n := range nodes {
|
||||
fmt.Fprintf(w, "%s\t%s\t%d/%d free\t%s\t%s\t%d\t%s\n",
|
||||
n.ID, n.State, n.GPUFree, n.GPUTotal,
|
||||
fmt.Fprintf(w, "%s\t%s\t%d/%d free\t%s\t%s\t%s\t%s\t%d\t%s\n",
|
||||
n.ID, n.State, n.GPUFree, n.GPUTotal, memColumn(n.MemoryGB),
|
||||
diskColumn(n.DiskFreeGB, n.DiskTotalGB),
|
||||
join(n.Classes), join(n.Fabrics), n.Replicas, n.LastSeenAgo)
|
||||
}
|
||||
_ = w.Flush()
|
||||
@@ -151,6 +174,72 @@ func newNodesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
}
|
||||
}
|
||||
|
||||
// modelStatus says what is happening to a model, in the words that decide what
|
||||
// to do about it.
|
||||
//
|
||||
// The previous rendering distinguished "a replica is coming up" from "nothing
|
||||
// could be placed" by a parenthetical — "0/1 ready (1 total)" against "0/1
|
||||
// ready" — and those need opposite responses: one means wait, the other means
|
||||
// look at explain-placement. On four models the difference is findable; on
|
||||
// forty it is invisible, and a model that will never start looks exactly like
|
||||
// one that is thirty seconds away.
|
||||
//
|
||||
// A model appearing here before it serves is deliberate: this is the operator's
|
||||
// view, and a model coming up is precisely what someone wants to see. Clients
|
||||
// are told separately and only when it is true — /v1/models lists what is
|
||||
// actually served, never what is merely catalogued.
|
||||
func modelStatus(m adminapi.ModelView) string {
|
||||
switch {
|
||||
case m.Min == 0 && m.Total == 0:
|
||||
return "stopped"
|
||||
case m.Ready >= m.Min && m.Min > 0:
|
||||
return "ready"
|
||||
case m.Total > m.Ready:
|
||||
return "starting"
|
||||
default:
|
||||
return "not placed"
|
||||
}
|
||||
}
|
||||
|
||||
// idleColumn shows how long a replica has gone unused.
|
||||
//
|
||||
// "-" for a replica nothing has reported use of, which is emphatically not the
|
||||
// same as idle forever: it is what every replica looks like before its first
|
||||
// request, and what they all look like after a router restart. Showing that as
|
||||
// a very large number would make an untouched fleet read as an abandoned one.
|
||||
func idleColumn(sec int64) string {
|
||||
if sec < 0 {
|
||||
return "-"
|
||||
}
|
||||
return (time.Duration(sec) * time.Second).Round(time.Second).String()
|
||||
}
|
||||
|
||||
// memColumn shows a node's memory, distinguishing "did not say" from a number.
|
||||
// An older agent reports nothing, and that is not the same as a small machine.
|
||||
func memColumn(gb int) string {
|
||||
if gb <= 0 {
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%dG", gb)
|
||||
}
|
||||
|
||||
// diskColumn renders free space against the size it is free out of.
|
||||
//
|
||||
// Both, because they answer different questions and the first is useless
|
||||
// without the second: 8 GB free is comfortable on a 10 TB filer and nearly
|
||||
// full on a 10 GB root partition, and it is the second case that stops a
|
||||
// model from being placed.
|
||||
//
|
||||
// A dash for a node that reported nothing, which an older agent and a node
|
||||
// with no model directories both do. It has to read as "not known" rather
|
||||
// than as "no space", or every unconfigured node looks like it is full.
|
||||
func diskColumn(freeGB, totalGB int) string {
|
||||
if totalGB <= 0 {
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%dG/%dG", freeGB, totalGB)
|
||||
}
|
||||
|
||||
func newModelsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "models",
|
||||
@@ -169,14 +258,15 @@ func newModelsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "MODEL\tREPLICAS\tALIASES\tPROFILES\tFALLBACKS")
|
||||
fmt.Fprintln(w, "MODEL\tSERVING\tREPLICAS\tSTATUS\tALIASES\tPROFILES\tFALLBACKS")
|
||||
for _, m := range models {
|
||||
replicas := fmt.Sprintf("%d/%d ready", m.Ready, m.Min)
|
||||
if m.Total > m.Ready {
|
||||
replicas += fmt.Sprintf(" (%d total)", m.Total)
|
||||
serving := "no"
|
||||
if m.Ready > 0 {
|
||||
serving = "yes"
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n",
|
||||
m.Name, replicas, join(m.Aliases), join(m.Profiles), join(m.Fallbacks))
|
||||
fmt.Fprintf(w, "%s\t%s\t%d/%d\t%s\t%s\t%s\t%s\n",
|
||||
m.Name, serving, m.Ready, m.Min, modelStatus(m),
|
||||
join(m.Aliases), join(m.Profiles), join(m.Fallbacks))
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
@@ -206,15 +296,16 @@ func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "REPLICA\tMODEL\tSTATE\tPOOL\tNODES\tAGE\tENDPOINT")
|
||||
fmt.Fprintln(w, "REPLICA\tMODEL\tSTATE\tPOOL\tNODES\tAGE\tIDLE\tENDPOINT")
|
||||
for _, r := range reps {
|
||||
var nodes []string
|
||||
for _, s := range r.Shards {
|
||||
nodes = append(nodes, fmt.Sprintf("%s:%v", s.Node, s.GPUs))
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
r.ID, r.Model, r.State, r.Pool, join(nodes),
|
||||
(time.Duration(r.AgeSec) * time.Second).String(), r.Endpoint)
|
||||
(time.Duration(r.AgeSec) * time.Second).String(),
|
||||
idleColumn(r.IdleSec), r.Endpoint)
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
// newWeightsCmd reports and reclaims model files nothing needs.
|
||||
//
|
||||
// The dangerous half of storage management, so its shape is defensive: listing
|
||||
// is the default and deleting takes a flag, only files llmgw fetched are ever
|
||||
// candidates, and the node checks every path again before unlinking it.
|
||||
//
|
||||
// What it will never touch is the more important half of the contract. Weights
|
||||
// an operator put on a node — by scp, by a shared mount, by hand — are none of
|
||||
// this system's business however unreferenced they look, and a model that took
|
||||
// a week to copy must not be removable by a command someone ran to tidy up.
|
||||
func newWeightsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "weights",
|
||||
Short: "Show cached model files no model in the catalog needs",
|
||||
Long: "Lists weights llmgw downloaded for models that are no longer in the\n" +
|
||||
"catalog, per node, with what deleting them would free.\n\n" +
|
||||
"Only files llmgw fetched are considered. Anything an operator put on a\n" +
|
||||
"node is left alone, whatever the catalog says — llmgw observes storage\n" +
|
||||
"rather than owning it, and only takes back what it put there itself.\n\n" +
|
||||
"Nothing is deleted without --collect.",
|
||||
Args: cobra.NoArgs,
|
||||
}
|
||||
var collect bool
|
||||
var node string
|
||||
|
||||
cmd.Flags().BoolVar(&collect, "collect", false, "delete them, rather than only listing")
|
||||
cmd.Flags().StringVar(&node, "node", "", "limit to one node")
|
||||
|
||||
cmd.RunE = func(cmd *cobra.Command, _ []string) error {
|
||||
if collect {
|
||||
return collectOrphans(cmd, client(), node, *asJSON)
|
||||
}
|
||||
return listOrphanWeights(cmd, client(), node, *asJSON)
|
||||
}
|
||||
return cmd
|
||||
}
|
||||
|
||||
func listOrphanWeights(cmd *cobra.Command, c *ctlClient, node string, asJSON bool) error {
|
||||
var found []adminapi.OrphanView
|
||||
if err := c.get(cmd.Context(), "/api/v1/weights/orphans", &found); err != nil {
|
||||
return err
|
||||
}
|
||||
if node != "" {
|
||||
found = onlyNode(found, node)
|
||||
}
|
||||
if asJSON {
|
||||
return emitJSON(found)
|
||||
}
|
||||
if len(found) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(),
|
||||
"Nothing to reclaim: every file llmgw fetched is still wanted by a model.")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "NODE\tFILES\tWOULD FREE")
|
||||
var total int64
|
||||
for _, o := range found {
|
||||
fmt.Fprintf(w, "%s\t%d\t%s\n", o.Node, len(o.Paths), fileSize(o.Bytes))
|
||||
total += o.Bytes
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
for _, o := range found {
|
||||
for _, p := range o.Paths {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n %s: %s", o.Node, p)
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n\n%s in total. Delete with --collect.\n", fileSize(total))
|
||||
return nil
|
||||
}
|
||||
|
||||
func collectOrphans(cmd *cobra.Command, c *ctlClient, node string, asJSON bool) error {
|
||||
body, err := json.Marshal(map[string]string{"node": node})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var results []struct {
|
||||
Node string `json:"node"`
|
||||
Files int `json:"files"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
if err := c.post(cmd.Context(), "/api/v1/weights/collect", body, &results); err != nil {
|
||||
return err
|
||||
}
|
||||
if asJSON {
|
||||
return emitJSON(results)
|
||||
}
|
||||
if len(results) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), "Nothing to reclaim.")
|
||||
return nil
|
||||
}
|
||||
var freed int64
|
||||
for _, r := range results {
|
||||
if r.Error != "" {
|
||||
// Reported rather than fatal: the other nodes' space is still worth
|
||||
// freeing, and this one will be offered again next time.
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: %s\n", r.Node, r.Error)
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: asked to remove %d files, %s\n",
|
||||
r.Node, r.Files, fileSize(r.Bytes))
|
||||
freed += r.Bytes
|
||||
}
|
||||
// "asked", not "removed": each node decides, and one that refuses a path it
|
||||
// cannot prove it fetched is the feature working rather than failing.
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n%s requested. The nodes verify each path before deleting it.\n",
|
||||
fileSize(freed))
|
||||
return nil
|
||||
}
|
||||
|
||||
func onlyNode(in []adminapi.OrphanView, node string) []adminapi.OrphanView {
|
||||
out := make([]adminapi.OrphanView, 0, len(in))
|
||||
for _, o := range in {
|
||||
if o.Node == node {
|
||||
out = append(out, o)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// fileSize renders a size the way someone reads one when deciding whether a
|
||||
// file is worth deleting.
|
||||
//
|
||||
// Separate from the per-token cache size renderer in modelplan, which never
|
||||
// leaves KiB and MiB. Model files are tens of gigabytes, and "8601.6 MiB" is
|
||||
// not a number anyone weighs a decision against.
|
||||
func fileSize(b int64) string {
|
||||
switch {
|
||||
case b >= 1<<30:
|
||||
return fmt.Sprintf("%.1f GB", float64(b)/(1<<30))
|
||||
case b >= 1<<20:
|
||||
return fmt.Sprintf("%.0f MB", float64(b)/(1<<20))
|
||||
case b > 0:
|
||||
return fmt.Sprintf("%d B", b)
|
||||
default:
|
||||
// The node reported no size. Saying "0 B" would suggest deleting it
|
||||
// gains nothing, which is a different claim from not knowing.
|
||||
return "size unknown"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
)
|
||||
|
||||
func TestTheForwardingDispatcherCarriesEveryCapability(t *testing.T) {
|
||||
// dispatcherRef exists to break a construction cycle, and it satisfies the
|
||||
// dispatcher interface by forwarding each method by hand. That makes it a
|
||||
// place where capabilities go missing silently: the controller reaches
|
||||
// past the interface by type assertion, and an assertion against this type
|
||||
// fails for anything nobody remembered to forward.
|
||||
//
|
||||
// It happened. `ctl weights --collect` returned "this control plane cannot
|
||||
// ask nodes to delete files" against a real cluster, while every unit test
|
||||
// passed — because the test dispatchers do not implement the capability
|
||||
// either, so its absence looked exactly like the expected case.
|
||||
var ref any = &dispatcherRef{}
|
||||
if _, ok := ref.(control.WeightRemover); !ok {
|
||||
t.Error("dispatcherRef drops WeightRemover, so the control plane cannot reclaim disk")
|
||||
}
|
||||
|
||||
// And the thing it forwards to must have it, or the forwarding has nothing
|
||||
// to reach.
|
||||
var svc any = &agentsvc.Service{}
|
||||
if _, ok := svc.(control.WeightRemover); !ok {
|
||||
t.Error("the agent service does not implement WeightRemover")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
|
||||
func newEngineCheckCmd() *cobra.Command {
|
||||
var (
|
||||
help string
|
||||
bin string
|
||||
nodes int
|
||||
tp int
|
||||
perNode int
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "engine-check <engine>",
|
||||
Short: "Check that an engine still accepts the flags llmgw builds for it",
|
||||
Long: "llmgw's launch arguments are read from each engine's documentation, and\n" +
|
||||
"engines rename and remove flags between versions. When that happens the\n" +
|
||||
"failure arrives minutes into a model load, in an engine's own error\n" +
|
||||
"message, on a node nobody was watching.\n\n" +
|
||||
"This asks the engine itself. It runs the engine's --help and checks every\n" +
|
||||
"flag llmgw would pass against it, which needs no GPU — a CLI prints its\n" +
|
||||
"help on any machine. Run it after upgrading an engine.\n\n" +
|
||||
"With --help-output it reads a saved --help instead of running anything,\n" +
|
||||
"so the check works from a laptop against help captured on a node.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
name := args[0]
|
||||
reg := engines.NewRegistry(engines.Config{})
|
||||
eng, err := reg.Get(name)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// The binary lives where the agent's configuration says, which is
|
||||
// rarely on PATH: a check that cannot find the engine reports
|
||||
// nothing rather than reporting a problem.
|
||||
cfg := engines.Config{}
|
||||
if bin != "" {
|
||||
switch name {
|
||||
case "llamacpp":
|
||||
cfg.LlamaCpp.Bin = bin
|
||||
case "vllm":
|
||||
cfg.VLLM.Bin = bin
|
||||
case "sglang":
|
||||
cfg.SGLang.Python = bin
|
||||
case "sim":
|
||||
cfg.Sim.Bin = bin
|
||||
}
|
||||
reg = engines.NewRegistry(cfg)
|
||||
if eng, err = reg.Get(name); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
spec := representativeShard(name, nodes, tp, perNode)
|
||||
cmds, err := eng.Commands(spec)
|
||||
if err != nil {
|
||||
return fmt.Errorf("building a command for %s: %w", name, err)
|
||||
}
|
||||
|
||||
out, errOut := cmd.OutOrStdout(), cmd.ErrOrStderr()
|
||||
var checked, skipped int
|
||||
var unknown []string
|
||||
|
||||
// Per command, and that distinction is the whole point. A shard can
|
||||
// be several processes — vLLM under Ray is `ray start` and then
|
||||
// `vllm serve` — and they are different programs with different
|
||||
// vocabularies. Checking the sidecar's flags against the engine's
|
||||
// help reports --head and --num-gpus as missing, which is true of
|
||||
// vllm and irrelevant, because ray is what receives them.
|
||||
for i, c := range cmds {
|
||||
fmt.Fprintf(errOut, " %s %s\n", c.Path, strings.Join(c.Args, " "))
|
||||
|
||||
text := ""
|
||||
switch {
|
||||
case help != "" && i == len(cmds)-1:
|
||||
// A supplied help belongs to the engine, which is the last
|
||||
// command; the sidecars have their own and we do not have it.
|
||||
text = help
|
||||
case help != "":
|
||||
skipped++
|
||||
fmt.Fprintf(errOut, " skipped: --help-output describes the engine, not %s\n", c.Name)
|
||||
continue
|
||||
default:
|
||||
var err error
|
||||
if text, err = runHelp(cmd.Context(), c); err != nil {
|
||||
skipped++
|
||||
fmt.Fprintf(errOut, " skipped: %v\n", err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// A truncated help produces confident nonsense: every flag below
|
||||
// the cut looks missing. Real help documents --help itself,
|
||||
// which is the cheapest signal that we have all of it.
|
||||
if !mentions(text, "--help") {
|
||||
fmt.Fprintf(errOut, " warning: this help does not document --help, "+
|
||||
"which usually means it was truncated; flags below the cut will look missing\n")
|
||||
}
|
||||
|
||||
seen := map[string]bool{}
|
||||
for _, a := range c.Args {
|
||||
flag, _, _ := strings.Cut(a, "=")
|
||||
if !strings.HasPrefix(flag, "--") || seen[flag] {
|
||||
continue
|
||||
}
|
||||
seen[flag] = true
|
||||
checked++
|
||||
if !mentions(text, flag) {
|
||||
unknown = append(unknown, c.Name+": "+flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
sort.Strings(unknown)
|
||||
|
||||
fmt.Fprintf(out, "%s: %d flags checked", name, checked)
|
||||
if skipped > 0 {
|
||||
fmt.Fprintf(out, ", %d commands skipped", skipped)
|
||||
}
|
||||
fmt.Fprintln(out)
|
||||
|
||||
if len(unknown) == 0 {
|
||||
if checked == 0 {
|
||||
return fmt.Errorf("nothing was checked; install the engine here or pass --help-output")
|
||||
}
|
||||
fmt.Fprintln(out, "every flag llmgw builds is accepted")
|
||||
return nil
|
||||
}
|
||||
fmt.Fprintf(out, "\n%d flags not mentioned by the program they are passed to:\n", len(unknown))
|
||||
for _, f := range unknown {
|
||||
fmt.Fprintf(out, " %s\n", f)
|
||||
}
|
||||
return fmt.Errorf("%s: %d flags are not accepted", name, len(unknown))
|
||||
},
|
||||
}
|
||||
|
||||
f := cmd.Flags()
|
||||
f.StringVar(&help, "help-output", "", "a saved --help to check against, instead of running the engine")
|
||||
f.StringVar(&bin, "bin", "", "where the engine lives, when it is not on PATH")
|
||||
f.IntVar(&nodes, "nodes", 2, "shape the command for a replica of this many nodes")
|
||||
f.IntVar(&tp, "tp", 8, "tensor parallel degree")
|
||||
f.IntVar(&perNode, "gpus-per-node", 4, "GPUs each node contributes")
|
||||
return cmd
|
||||
}
|
||||
|
||||
// representativeShard is the widest command llmgw would ever build, because a
|
||||
// flag that only appears in the multi-node case is exactly the one nobody
|
||||
// notices has been renamed.
|
||||
func representativeShard(engine string, nodes, tp, perNode int) engines.ShardSpec {
|
||||
if nodes < 1 {
|
||||
nodes = 1
|
||||
}
|
||||
return engines.ShardSpec{
|
||||
ReplicaID: "check", ShardID: "check-r0", Rank: 0, NNodes: nodes,
|
||||
LeaderAddr: "10.0.0.1", LeaderPort: 20001,
|
||||
GPUIndices: seq(perNode), HCAs: []string{"mlx5_0"},
|
||||
MgmtInterface: "eth0", HTTPPort: 20000,
|
||||
Profile: engines.Profile{
|
||||
Name: "check", Engine: engine, TP: tp, Nodes: nodes, GPUsPerNode: perNode,
|
||||
ModelPath: "/models/x", ServedModelName: "x",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func seq(n int) []int {
|
||||
out := make([]int, n)
|
||||
for i := range out {
|
||||
out[i] = i
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// helpArgs is how each program is asked for its *complete* help.
|
||||
//
|
||||
// Not every program gives it all up for --help. vLLM prints a summary listing
|
||||
// section names and tells you to ask again per section; checking against that
|
||||
// reports every flag as missing, confidently and wrongly. This is exactly what
|
||||
// happened the first time this ran against a real vLLM.
|
||||
var helpArgs = map[string][]string{
|
||||
"vllm": {"--help=all"},
|
||||
}
|
||||
|
||||
// runHelp asks one program to describe itself.
|
||||
func runHelp(ctx context.Context, c engines.Command) (string, error) {
|
||||
// Everything up to the first flag, plus the help request. Passing the real
|
||||
// arguments would start the engine.
|
||||
var args []string
|
||||
for _, a := range c.Args {
|
||||
if strings.HasPrefix(a, "-") {
|
||||
break
|
||||
}
|
||||
args = append(args, a)
|
||||
}
|
||||
ask, ok := helpArgs[filepath.Base(c.Path)]
|
||||
if !ok {
|
||||
ask = []string{"--help"}
|
||||
}
|
||||
args = append(args, ask...)
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
ec := exec.CommandContext(ctx, c.Path, args...)
|
||||
// vLLM refuses to do anything at all on a machine with no accelerator,
|
||||
// including print its help, unless it is told what to assume. Asking an
|
||||
// engine about its flags should not require the hardware it runs on.
|
||||
ec.Env = append(ec.Environ(), "VLLM_TARGET_DEVICE=cpu")
|
||||
|
||||
raw, err := ec.CombinedOutput()
|
||||
if err != nil && len(raw) == 0 {
|
||||
return "", fmt.Errorf("could not run %s: %w (install it here, or pass --help-output)",
|
||||
c.Path, err)
|
||||
}
|
||||
return string(raw), nil
|
||||
}
|
||||
|
||||
// mentions reports whether the help text documents a flag.
|
||||
//
|
||||
// Matched on a word boundary, so --port does not satisfy --port-range and a
|
||||
// missing flag cannot hide behind a longer one that happens to contain it.
|
||||
func mentions(help, flag string) bool {
|
||||
for i := 0; ; {
|
||||
j := strings.Index(help[i:], flag)
|
||||
if j < 0 {
|
||||
return false
|
||||
}
|
||||
end := i + j + len(flag)
|
||||
if end == len(help) || !isFlagChar(help[end]) {
|
||||
return true
|
||||
}
|
||||
i = end
|
||||
}
|
||||
}
|
||||
|
||||
func isFlagChar(b byte) bool {
|
||||
return b == '-' || b == '_' || (b >= 'a' && b <= 'z') ||
|
||||
(b >= 'A' && b <= 'Z') || (b >= '0' && b <= '9')
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAFlagIsNotSatisfiedByALongerOne(t *testing.T) {
|
||||
// The failure this guards against is silent: --port would be reported as
|
||||
// present because --port-range is documented, and a renamed flag would sail
|
||||
// through the check that exists to catch renamed flags.
|
||||
help := " --port-range N a range\n --model-name S a name\n"
|
||||
|
||||
if mentions(help, "--port") {
|
||||
t.Error("--port was satisfied by --port-range")
|
||||
}
|
||||
if mentions(help, "--model") {
|
||||
t.Error("--model was satisfied by --model-name")
|
||||
}
|
||||
if !mentions(help, "--port-range") {
|
||||
t.Error("--port-range was not found")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFlagsAreFoundInTheShapesHelpTextUses(t *testing.T) {
|
||||
help := `
|
||||
-m, --model FNAME model path
|
||||
--host HOST ip address
|
||||
--parallel, -np N slots
|
||||
--metrics enable prometheus
|
||||
--alias=NAME served name
|
||||
`
|
||||
for _, flag := range []string{"--model", "--host", "--parallel", "--metrics", "--alias"} {
|
||||
if !mentions(help, flag) {
|
||||
t.Errorf("%s was not found in a help text that documents it", flag)
|
||||
}
|
||||
}
|
||||
for _, flag := range []string{"--nonsense", "--mode", "--para"} {
|
||||
if mentions(help, flag) {
|
||||
t.Errorf("%s was found in a help text that does not document it", flag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheRepresentativeShardIsTheWidestOne(t *testing.T) {
|
||||
// A flag that only appears in the multi-node case is exactly the one nobody
|
||||
// notices has been renamed, so the command built for checking has to be the
|
||||
// widest llmgw would ever produce.
|
||||
spec := representativeShard("vllm", 2, 8, 4)
|
||||
if spec.NNodes != 2 || spec.Profile.Nodes != 2 {
|
||||
t.Errorf("not a multi-node shard: %+v", spec)
|
||||
}
|
||||
if len(spec.GPUIndices) != 4 {
|
||||
t.Errorf("%d GPUs, want 4", len(spec.GPUIndices))
|
||||
}
|
||||
if spec.LeaderAddr == "" || spec.LeaderPort == 0 {
|
||||
t.Error("no rendezvous point, so the flags that carry it will not appear")
|
||||
}
|
||||
}
|
||||
|
||||
func TestATruncatedHelpIsSuspected(t *testing.T) {
|
||||
// Found by truncating llama.cpp's 725-line help to 400 lines and being told,
|
||||
// confidently, that five flags the running cluster uses do not exist.
|
||||
full := " --help show this\n --model F path\n"
|
||||
if !mentions(full, "--help") {
|
||||
t.Fatal("the guard's own signal is missing from a complete help")
|
||||
}
|
||||
truncated := " --model F path\n"
|
||||
if mentions(truncated, "--help") {
|
||||
t.Error("a truncated help was taken for a complete one")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,327 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/base32"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func newTenancyCmd() *cobra.Command {
|
||||
var dbPath string
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "keys",
|
||||
Short: "Manage tenants and API keys",
|
||||
Long: "Works directly on the store file rather than through the admin API, so keys\n" +
|
||||
"can be issued before the server is running and without exposing an endpoint\n" +
|
||||
"that mints credentials.",
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store")
|
||||
|
||||
open := func(ctx context.Context) (*store.SQLite, error) { return store.OpenSQLite(ctx, dbPath) }
|
||||
|
||||
cmd.AddCommand(
|
||||
newKeyInitCmd(open),
|
||||
newKeyCreateCmd(open),
|
||||
newKeyListCmd(open),
|
||||
newKeyRevokeCmd(open),
|
||||
newLimitsCmd(open),
|
||||
newThreadsCmd(open),
|
||||
newMemoryCmd(open),
|
||||
)
|
||||
return cmd
|
||||
}
|
||||
|
||||
type opener func(context.Context) (*store.SQLite, error)
|
||||
|
||||
func newKeyInitCmd(open opener) *cobra.Command {
|
||||
var org, team, project string
|
||||
cmd := &cobra.Command{
|
||||
Use: "init",
|
||||
Short: "Create the tenancy tree a key hangs off",
|
||||
Long: "Idempotent: running it again renames rather than duplicating, so it can be\n" +
|
||||
"kept in whatever provisions the machine.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
ctx := cmd.Context()
|
||||
if err := s.PutOrg(ctx, store.Org{ID: slug(org), Name: org}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.PutTeam(ctx, store.Team{ID: slug(team), OrgID: slug(org), Name: team}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.PutProject(ctx, store.Project{
|
||||
ID: slug(project), TeamID: slug(team), Name: project,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s / %s / %s\n", slug(org), slug(team), slug(project))
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&org, "org", "default", "organization name")
|
||||
cmd.Flags().StringVar(&team, "team", "default", "team name")
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project name")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newKeyCreateCmd(open opener) *cobra.Command {
|
||||
var project, band, profile, baseURL string
|
||||
var expires, claimTTL time.Duration
|
||||
var claim bool
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "create <name>",
|
||||
Short: "Issue an API key",
|
||||
Long: "The key is printed once and never again: only its hash is stored, so there\n" +
|
||||
"is nothing to recover it from. That is the point — a leaked database hands\n" +
|
||||
"over no working credential.\n\n" +
|
||||
"With --claim it prints a one-time link instead. Send the link rather than\n" +
|
||||
"the key: a key pasted into a chat window stays in that chat window forever,\n" +
|
||||
"while a link works once and then expires. The key is sealed with the link's\n" +
|
||||
"own token, which is not stored, so a pending claim is not a usable\n" +
|
||||
"credential sitting in the database either.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
secret, k, err := store.NewKey(randomID("key"), args[0], project)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
k.CreatedAt = time.Now()
|
||||
k.Band, k.RoutingProfile = band, profile
|
||||
if expires > 0 {
|
||||
k.ExpiresAt = k.CreatedAt.Add(expires)
|
||||
}
|
||||
if err := s.CreateKey(cmd.Context(), k); err != nil {
|
||||
return fmt.Errorf("creating the key (does project %q exist? run `llmgw keys init`): %w",
|
||||
project, err)
|
||||
}
|
||||
|
||||
// The one output that matters goes to stdout and everything else to
|
||||
// stderr, so `keys create ... > out.txt` captures it and nothing
|
||||
// else. Cobra's Print writes to stderr, which would have made the
|
||||
// useful half the half that cannot be piped.
|
||||
out := cmd.OutOrStdout()
|
||||
|
||||
if claim {
|
||||
if baseURL == "" {
|
||||
return errors.New("--claim needs --base-url, the address users reach this gateway at")
|
||||
}
|
||||
token, err := store.NewClaimToken()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.CreateClaim(cmd.Context(), k.ID, token, secret, claimTTL); err != nil {
|
||||
return err
|
||||
}
|
||||
link := strings.TrimRight(baseURL, "/") + "/claim/" + token
|
||||
fmt.Fprintln(out, link)
|
||||
cmd.PrintErrf("\nid %s\nproject %s\nexpires %s\n",
|
||||
k.ID, k.ProjectID, time.Now().Add(orDefault(claimTTL, store.ClaimTTL)).Format(time.RFC3339))
|
||||
cmd.PrintErrln("\nSend the link. It works once, and the key is shown only to whoever opens it.")
|
||||
return nil
|
||||
}
|
||||
|
||||
fmt.Fprintln(out, secret)
|
||||
cmd.PrintErrf("\nid %s\nproject %s\n", k.ID, k.ProjectID)
|
||||
if !k.ExpiresAt.IsZero() {
|
||||
cmd.PrintErrf("expires %s\n", k.ExpiresAt.Format(time.RFC3339))
|
||||
}
|
||||
cmd.PrintErrln("\nThis is the only time the key is shown. Store it now.")
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project the key belongs to")
|
||||
cmd.Flags().StringVar(&band, "band", "", "admission band; empty inherits")
|
||||
cmd.Flags().StringVar(&profile, "profile", "", "routing profile; empty inherits")
|
||||
cmd.Flags().DurationVar(&expires, "expires-in", 0, "lifetime; 0 never expires")
|
||||
cmd.Flags().BoolVar(&claim, "claim", false, "print a one-time link instead of the key")
|
||||
cmd.Flags().StringVar(&baseURL, "base-url", "", "public address of the gateway, for --claim")
|
||||
cmd.Flags().DurationVar(&claimTTL, "claim-ttl", store.ClaimTTL, "how long the link stays usable")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func orDefault(d, fallback time.Duration) time.Duration {
|
||||
if d <= 0 {
|
||||
return fallback
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func newKeyListCmd(open opener) *cobra.Command {
|
||||
var project string
|
||||
cmd := &cobra.Command{
|
||||
Use: "list",
|
||||
Short: "List a project's keys",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
keys, err := s.Keys(cmd.Context(), project)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(keys) == 0 {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "no keys in project %q\n", project)
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
|
||||
fmt.Fprintln(w, "ID\tNAME\tKEY\tBAND\tSTATE\tLAST USED")
|
||||
now := time.Now()
|
||||
for _, k := range keys {
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
k.ID, k.Name, k.Prefix, dash(k.Band), keyState(k, now), ago(k.LastUsedAt, now))
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project to list")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newKeyRevokeCmd(open opener) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "revoke <key-id>",
|
||||
Short: "Stop a key working, without deleting it",
|
||||
Long: "A soft delete. Usage rows reference the key, and an audit trail that loses\n" +
|
||||
"its subject is not an audit trail.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
if err := s.RevokeKey(cmd.Context(), args[0], time.Now()); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "revoked %s\n", args[0])
|
||||
cmd.PrintErrln("A running server keeps accepting it until its cache expires (auth.cache_ttl).")
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newLimitsCmd(open opener) *cobra.Command {
|
||||
var scope, id, period string
|
||||
var rpm, tpm int
|
||||
var budget int64
|
||||
var models []string
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "limits",
|
||||
Short: "Set limits at one level of the tenancy tree",
|
||||
Long: "Limits compose downward by taking the most restrictive value, so a key\n" +
|
||||
"cannot raise what its organization set. Model allowlists intersect: a key\n" +
|
||||
"cannot add a model its team withheld.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
l := store.Limits{Scope: store.Scope(scope), ScopeID: id}
|
||||
if cmd.Flags().Changed("rpm") {
|
||||
l.RPM = &rpm
|
||||
}
|
||||
if cmd.Flags().Changed("tpm") {
|
||||
l.TPM = &tpm
|
||||
}
|
||||
if cmd.Flags().Changed("budget") {
|
||||
l.Budget = &budget
|
||||
l.BudgetPeriod = store.Period(period)
|
||||
}
|
||||
if cmd.Flags().Changed("models") {
|
||||
// An explicitly empty list is an allowlist that permits
|
||||
// nothing, which is different from setting no list at all.
|
||||
if models == nil {
|
||||
models = []string{}
|
||||
}
|
||||
l.AllowedModels = models
|
||||
}
|
||||
if err := s.PutLimits(cmd.Context(), l); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "limits set on %s %s\n", scope, id)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&scope, "scope", "project", "org | team | project | key")
|
||||
cmd.Flags().StringVar(&id, "id", "default", "id at that scope")
|
||||
cmd.Flags().IntVar(&rpm, "rpm", 0, "requests per minute")
|
||||
cmd.Flags().IntVar(&tpm, "tpm", 0, "tokens per minute")
|
||||
cmd.Flags().Int64Var(&budget, "budget", 0, "budget per period, in micro-units")
|
||||
cmd.Flags().StringVar(&period, "period", "month", "day | month")
|
||||
cmd.Flags().StringSliceVar(&models, "models", nil, "allowed models; empty allows none")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func keyState(k store.Key, now time.Time) string {
|
||||
switch {
|
||||
case !k.RevokedAt.IsZero():
|
||||
return "revoked"
|
||||
case !k.ExpiresAt.IsZero() && !now.Before(k.ExpiresAt):
|
||||
return "expired"
|
||||
default:
|
||||
return "active"
|
||||
}
|
||||
}
|
||||
|
||||
func ago(t, now time.Time) string {
|
||||
if t.IsZero() {
|
||||
return "never"
|
||||
}
|
||||
return now.Sub(t).Round(time.Second).String() + " ago"
|
||||
}
|
||||
|
||||
// slug turns a human name into an id.
|
||||
func slug(name string) string {
|
||||
var b strings.Builder
|
||||
for _, r := range strings.ToLower(strings.TrimSpace(name)) {
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
|
||||
b.WriteRune(r)
|
||||
case r == ' ' || r == '-' || r == '_':
|
||||
b.WriteByte('-')
|
||||
}
|
||||
}
|
||||
if b.Len() == 0 {
|
||||
return "default"
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func randomID(prefix string) string {
|
||||
var buf [10]byte
|
||||
_, _ = rand.Read(buf[:])
|
||||
return prefix + "-" + strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:]))
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// newMemoryCmd turns note storage on for a project, or off.
|
||||
//
|
||||
// A command rather than a configuration key for the same reason threads are:
|
||||
// keeping what an agent decided to remember is a decision about data, not about
|
||||
// deployment, and it should look like an act rather than a default nobody read.
|
||||
func newMemoryCmd(open opener) *cobra.Command {
|
||||
var (
|
||||
project string
|
||||
hours int
|
||||
off bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "memory",
|
||||
Short: "Keep notes an agent writes for itself, across conversations",
|
||||
Long: "With this on, an agent can save and search notes through /v1/memory or the\n" +
|
||||
"MCP tools at /mcp (remember, recall, forget), and find them again in later\n" +
|
||||
"conversations.\n\n" +
|
||||
"Off for every project until this is run. The retention is required: notes\n" +
|
||||
"expire that many hours after they were last written, so a store nobody\n" +
|
||||
"chose a number for does not keep everything forever.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if !off && hours <= 0 {
|
||||
return fmt.Errorf("--retain-hours is required: say how long notes are kept, or pass --off")
|
||||
}
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
keep := hours
|
||||
if off {
|
||||
keep = 0
|
||||
}
|
||||
if err := s.SetMemoryRetention(cmd.Context(), project, keep); err != nil {
|
||||
return err
|
||||
}
|
||||
if keep == 0 {
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s keeps no notes. Existing ones are removed at the next prune.\n", project)
|
||||
return nil
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s keeps notes for %d hours after they were last written.\n", project, keep)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project id")
|
||||
cmd.Flags().IntVar(&hours, "retain-hours", 0, "how long a note is kept after it was last written")
|
||||
cmd.Flags().BoolVar(&off, "off", false, "stop keeping notes for this project")
|
||||
return cmd
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os/user"
|
||||
|
||||
"github.com/knadh/koanf/parsers/yaml"
|
||||
"github.com/knadh/koanf/providers/file"
|
||||
"github.com/knadh/koanf/v2"
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
)
|
||||
|
||||
// newModelCmd adds and removes models without a restart.
|
||||
//
|
||||
// Adding a model used to mean editing a file on the control plane and
|
||||
// restarting it: a settling window, adoption churn, and on one occasion two
|
||||
// minutes of outage. For a fixed production set that is a reasonable trade; for
|
||||
// a platform whose promise is that people can try things, it is the difference
|
||||
// between trying one and not bothering.
|
||||
func newModelCmd(client func() *ctlClient) *cobra.Command {
|
||||
c := &cobra.Command{
|
||||
Use: "model",
|
||||
Short: "Add or remove models while the cluster is running",
|
||||
}
|
||||
c.AddCommand(newModelAddCmd(client), newModelRemoveCmd(client), newModelPlanCmd(client))
|
||||
return c
|
||||
}
|
||||
|
||||
func newModelAddCmd(client func() *ctlClient) *cobra.Command {
|
||||
var path string
|
||||
c := &cobra.Command{
|
||||
Use: "add -f <file.yaml>",
|
||||
Short: "Add a model from a definition file",
|
||||
Long: "The definition is the same YAML a model takes in server.yaml, so an\n" +
|
||||
"entry can be moved between the two by copying it.\n\n" +
|
||||
"The catalog is rebuilt and validated before anything is written, so a\n" +
|
||||
"definition that collides with an existing name, or names a fallback that\n" +
|
||||
"does not exist, is refused rather than stored and then ignored on every\n" +
|
||||
"start. Once accepted it takes effect on the reconciler's next pass — no\n" +
|
||||
"restart — and survives one.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if path == "" {
|
||||
return fmt.Errorf("--file is required: the model definition, as YAML")
|
||||
}
|
||||
// Read by koanf, the same parser and the same tags that read
|
||||
// server.yaml — so an entry can be moved between the two by copying
|
||||
// it, and a field that works in one cannot silently mean nothing in
|
||||
// the other.
|
||||
k := koanf.New(".")
|
||||
if err := k.Load(file.Provider(path), yaml.Parser()); err != nil {
|
||||
return fmt.Errorf("%s: %w", path, err)
|
||||
}
|
||||
var m catalog.Model
|
||||
if err := k.UnmarshalWithConf("", &m, koanf.UnmarshalConf{Tag: "koanf"}); err != nil {
|
||||
return fmt.Errorf("%s: %w", path, err)
|
||||
}
|
||||
if m.Name == "" {
|
||||
return fmt.Errorf("%s: the definition has no name", path)
|
||||
}
|
||||
body, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var res map[string]string
|
||||
if err := client().post(cmd.Context(), "/api/v1/models", body, &res); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s added. It will be placed on the next reconcile pass; watch it with:\n"+
|
||||
" llmgw ctl models\n", m.Name)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
c.Flags().StringVarP(&path, "file", "f", "", "YAML file holding the model definition")
|
||||
return c
|
||||
}
|
||||
|
||||
func newModelRemoveCmd(client func() *ctlClient) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "remove <name>",
|
||||
Short: "Remove a model that was added at runtime",
|
||||
Long: "Only models added at runtime. One declared in server.yaml would come\n" +
|
||||
"back on the next restart, so removing it here would leave the cluster\n" +
|
||||
"disagreeing with its own configuration until then.\n\n" +
|
||||
"Its replicas are stopped on the next pass, after the endpoint has been\n" +
|
||||
"withdrawn, so traffic stops before the engine does.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
var res map[string]string
|
||||
if err := client().delete(cmd.Context(), "/api/v1/models/"+args[0], &res); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s removed.\n", args[0])
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// actor names who made a change, for the audit trail on a stored model.
|
||||
func actor() string {
|
||||
if u, err := user.Current(); err == nil && u.Username != "" {
|
||||
return u.Username
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
|
||||
)
|
||||
|
||||
// newModelPlanCmd answers whether a model fits this cluster, and how.
|
||||
//
|
||||
// The arithmetic it replaces is the last thing standing between having GPUs and
|
||||
// serving a model on them. Deciding tp, gpu_class and a context that will
|
||||
// actually fit means knowing the model's layer count, its KV head count, what
|
||||
// its weights weigh at a chosen quantisation, and how much of a card the
|
||||
// runtime keeps — and getting it wrong is not an error message, it is an
|
||||
// out-of-memory twenty minutes into a load.
|
||||
//
|
||||
// It reads the model from its own repository and the hardware from the running
|
||||
// fleet, so neither half is a number somebody retyped.
|
||||
func newModelPlanCmd(client func() *ctlClient) *cobra.Command {
|
||||
var (
|
||||
context int
|
||||
params int64
|
||||
precs []string
|
||||
)
|
||||
c := &cobra.Command{
|
||||
Use: "plan <huggingface-repo>",
|
||||
Short: "Work out whether a model fits this cluster, and how it should be split",
|
||||
Long: "Reads the model's architecture from its repository and the hardware from\n" +
|
||||
"the running fleet, then prints every way it could run — fewest GPUs\n" +
|
||||
"first — and a catalog entry for the best one.\n\n" +
|
||||
"Every candidate is listed rather than one answer, because the choice is\n" +
|
||||
"not the planner's to make: with cards to spare you want the\n" +
|
||||
"full-precision fit, and with one card free you want to know a quantised\n" +
|
||||
"one exists.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
repo := args[0]
|
||||
m, err := modelplan.FetchModel(cmd.Context(), repo)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if params > 0 {
|
||||
m.Params = params
|
||||
}
|
||||
if m.Params == 0 {
|
||||
return fmt.Errorf("%s does not publish a parameter count; pass --params", repo)
|
||||
}
|
||||
|
||||
var chosen []modelplan.Precision
|
||||
for _, name := range precs {
|
||||
p, err := modelplan.PrecisionByName(name)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
chosen = append(chosen, p)
|
||||
}
|
||||
|
||||
var nodes []adminapi.NodeView
|
||||
if err := client().get(cmd.Context(), "/api/v1/nodes", &nodes); err != nil {
|
||||
return err
|
||||
}
|
||||
hw := hardwareFrom(nodes)
|
||||
if len(hw) == 0 {
|
||||
return fmt.Errorf("no node is ready, so there is nothing to plan against")
|
||||
}
|
||||
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintf(out, "%s: %.1fB parameters, %d layers, %d KV heads, trained to %d tokens\n",
|
||||
m.Name, float64(m.Params)/1e9, m.Layers, m.KVHeads, m.MaxContext)
|
||||
|
||||
want := context
|
||||
if want <= 0 {
|
||||
want = m.MaxContext
|
||||
}
|
||||
fmt.Fprintf(out, "Planned for a %d-token pool. %s a token of context.\n\n",
|
||||
want, humanBytes(m.KVBytesPerToken(modelplan.BF16)))
|
||||
|
||||
fits := modelplan.Fits(m, hw, want, chosen)
|
||||
if len(fits) == 0 {
|
||||
fmt.Fprintln(out, "It does not fit this cluster at any offered precision.")
|
||||
fmt.Fprintln(out, "\nWhat would change that:")
|
||||
fmt.Fprintln(out, " - a smaller context pool, with --context")
|
||||
fmt.Fprintln(out, " - a heavier quantisation, with --precision q4_k_m")
|
||||
fmt.Fprintln(out, " - more cards free; llmgw ctl status shows what is in use")
|
||||
return errDoesNotFit
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, " GPUS\tPRECISION\tPER CARD\tFREE\tWHERE")
|
||||
for _, f := range fits {
|
||||
where := f.Class
|
||||
if f.CPUOnly {
|
||||
where = "cpu"
|
||||
} else if f.Nodes > 1 {
|
||||
where = fmt.Sprintf("%s, %d nodes on %s", f.Class, f.Nodes, f.Fabric)
|
||||
}
|
||||
fmt.Fprintf(w, " %d\t%s\t%.1f GB\t%.1f GB\t%s\n",
|
||||
f.TP, f.Precision.Name, f.PerGPUGB, f.Headroom, where)
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
fmt.Fprintf(out, "\nA catalog entry for the first of those:\n\n%s\n",
|
||||
catalogEntry(m, fits[0], repo, want))
|
||||
fmt.Fprintln(out, "Save it and add it with: llmgw ctl model add -f <file>")
|
||||
return nil
|
||||
},
|
||||
}
|
||||
c.Flags().IntVar(&context, "context", 0, "token pool to plan for (default: the model's trained context)")
|
||||
c.Flags().Int64Var(¶ms, "params", 0, "parameter count, when the repository does not publish one")
|
||||
c.Flags().StringSliceVar(&precs, "precision", nil, "precisions to consider (default: bf16, fp8, q8_0, q4_k_m)")
|
||||
return c
|
||||
}
|
||||
|
||||
// errDoesNotFit gives a script something to test without parsing prose.
|
||||
var errDoesNotFit = fmt.Errorf("the model does not fit this cluster")
|
||||
|
||||
// hardwareFrom turns the fleet into pools a replica could actually occupy.
|
||||
//
|
||||
// Grouped by class and fabric together, because a replica may never span two
|
||||
// fabrics: eight free cards split across two fabrics are not eight cards for a
|
||||
// model that needs eight, and treating them as one pool would plan something
|
||||
// that hangs on its first all-reduce.
|
||||
func hardwareFrom(nodes []adminapi.NodeView) []modelplan.Hardware {
|
||||
type key struct{ class, fabric string }
|
||||
pools := map[key]*modelplan.Hardware{}
|
||||
|
||||
for _, n := range nodes {
|
||||
if n.State != "ready" {
|
||||
continue
|
||||
}
|
||||
fabric := ""
|
||||
if len(n.Fabrics) > 0 {
|
||||
fabric = n.Fabrics[0]
|
||||
}
|
||||
class := ""
|
||||
if len(n.Classes) > 0 {
|
||||
class = n.Classes[0]
|
||||
}
|
||||
k := key{class, fabric}
|
||||
if pools[k] == nil {
|
||||
pools[k] = &modelplan.Hardware{Class: class, Fabric: fabric}
|
||||
}
|
||||
p := pools[k]
|
||||
p.Nodes++
|
||||
// The widest a single node could contribute, not the total: a tensor
|
||||
// group takes the same number of cards from every node it spans.
|
||||
if n.GPUFree > p.GPUsPerNode {
|
||||
p.GPUsPerNode = n.GPUFree
|
||||
}
|
||||
if n.MemoryGB > p.HostMemGB {
|
||||
p.HostMemGB = n.MemoryGB
|
||||
}
|
||||
// Card memory from the class label, which carries it by construction:
|
||||
// "h100-80g" is eighty gigabytes and says so.
|
||||
if class != "" && p.GPUMemGB == 0 {
|
||||
p.GPUMemGB = memFromClass(class)
|
||||
}
|
||||
}
|
||||
|
||||
out := make([]modelplan.Hardware, 0, len(pools))
|
||||
for _, p := range pools {
|
||||
if p.Class != "" && (p.GPUMemGB == 0 || p.GPUsPerNode == 0) {
|
||||
continue // a GPU pool with no free cards is not a pool
|
||||
}
|
||||
out = append(out, *p)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// memFromClass reads the card size out of a class label such as "h100-80g".
|
||||
func memFromClass(class string) int {
|
||||
i := strings.LastIndex(class, "-")
|
||||
if i < 0 || !strings.HasSuffix(class, "g") {
|
||||
return 0
|
||||
}
|
||||
var gb int
|
||||
if _, err := fmt.Sscanf(class[i+1:], "%dg", &gb); err != nil {
|
||||
return 0
|
||||
}
|
||||
return gb
|
||||
}
|
||||
|
||||
// catalogEntry renders a fit as the YAML that would serve it.
|
||||
func catalogEntry(m modelplan.Model, f modelplan.Fit, source string, context int) string {
|
||||
var b strings.Builder
|
||||
name := strings.ToLower(m.Name)
|
||||
fmt.Fprintf(&b, "name: %s\n", name)
|
||||
if f.Precision.GGUF {
|
||||
// A GGUF quantisation lives in a file, not a repository, and llama.cpp
|
||||
// is given a path. Printing the repository id here would produce an
|
||||
// entry that looks right and cannot load — which is the failure this
|
||||
// whole command exists to prevent.
|
||||
fmt.Fprintf(&b, "# %s is a GGUF quantisation: download the %s file and give its path.\n",
|
||||
f.Precision.Name, f.Precision.Name)
|
||||
fmt.Fprintf(&b, "source: /var/lib/llmgw-models/%s-%s.gguf\n", name, f.Precision.Name)
|
||||
} else {
|
||||
fmt.Fprintf(&b, "source: %s\n", source)
|
||||
}
|
||||
fmt.Fprintf(&b, "routing_profile: chat\n")
|
||||
fmt.Fprintf(&b, "context_length: %d\n", context)
|
||||
fmt.Fprintf(&b, "scaling: { min: 1, max: 1 }\n")
|
||||
fmt.Fprintf(&b, "profiles:\n")
|
||||
|
||||
switch {
|
||||
case f.CPUOnly:
|
||||
fmt.Fprintf(&b, " - name: cpu\n engine: %s\n", f.Precision.Engine)
|
||||
fmt.Fprintf(&b, " min_memory_gb: %d\n", int(f.PerGPUGB)+2)
|
||||
default:
|
||||
fmt.Fprintf(&b, " - name: tp%d\n engine: %s\n", f.TP, f.Precision.Engine)
|
||||
fmt.Fprintf(&b, " gpu_class: %s\n tp: %d\n", f.Class, f.TP)
|
||||
if f.Nodes > 1 {
|
||||
fmt.Fprintf(&b, " nodes: %d\n gpus_per_node: %d\n", f.Nodes, f.GPUsPerNode)
|
||||
}
|
||||
if q := f.Precision.Quantization; q != "" {
|
||||
fmt.Fprintf(&b, " args: [\"--quantization\", \"%s\"]\n", q)
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(&b, " priority: 100\n weight: 1\n")
|
||||
// Measured rather than guessed is the rule everywhere else in this system,
|
||||
// and a planner should say so rather than pretend its estimate is one.
|
||||
fmt.Fprintf(&b, " startup_timeout: 10m\n")
|
||||
b.WriteString("\n# weight should come from `llmgw ctl bench` once it is running;\n")
|
||||
b.WriteString("# 1 is a placeholder, not a measurement.\n")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// humanBytes renders a per-token cache size, which is always small.
|
||||
func humanBytes(n int64) string {
|
||||
switch {
|
||||
case n >= 1<<20:
|
||||
return fmt.Sprintf("%.1f MiB", float64(n)/(1<<20))
|
||||
default:
|
||||
return fmt.Sprintf("%.0f KiB", float64(n)/1024)
|
||||
}
|
||||
}
|
||||
@@ -40,6 +40,8 @@ func Root() *cobra.Command {
|
||||
newServerCmd(),
|
||||
newAgentCmd(),
|
||||
newCtlCmd(),
|
||||
newTenancyCmd(),
|
||||
newUsersCmd(),
|
||||
newSimEngineCmd(),
|
||||
)
|
||||
return root
|
||||
|
||||
+162
-15
@@ -2,6 +2,7 @@ package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
@@ -23,16 +24,21 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fleetmetrics"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/web"
|
||||
)
|
||||
|
||||
func newServerCmd() *cobra.Command {
|
||||
var (
|
||||
configPath string
|
||||
listen string
|
||||
adminListen string
|
||||
endpoints []string
|
||||
configPath string
|
||||
listen string
|
||||
adminListen string
|
||||
metricsListen string
|
||||
endpoints []string
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
@@ -54,6 +60,9 @@ func newServerCmd() *cobra.Command {
|
||||
if adminListen != "" {
|
||||
cfg.AdminListen = adminListen
|
||||
}
|
||||
if metricsListen != "" {
|
||||
cfg.MetricsListen = metricsListen
|
||||
}
|
||||
for _, spec := range endpoints {
|
||||
ep, err := parseEndpointFlag(spec)
|
||||
if err != nil {
|
||||
@@ -92,8 +101,14 @@ func newServerCmd() *cobra.Command {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
live := catalog.NewLive(cat)
|
||||
// The router derives its own view, and it is empty until told.
|
||||
// Told here rather than at construction so there is exactly one
|
||||
// path by which it learns the catalog, used at startup and at
|
||||
// every change after.
|
||||
rt.SetCatalog(cat)
|
||||
|
||||
ctl, svc, err := startControlPlane(ctx, cfg, cat, rt, log)
|
||||
ctl, svc, err := startControlPlane(ctx, cfg, live, rt, cfg.AgentTLS.Enabled(), log)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -101,8 +116,23 @@ func newServerCmd() *cobra.Command {
|
||||
admin = ctl
|
||||
agentPath, agentHandler = svc.Handler()
|
||||
|
||||
// Before AdminHandler, which is where the registry is built —
|
||||
// registering after it silently publishes nothing, which is
|
||||
// exactly what the first version of this did.
|
||||
//
|
||||
// On the same /metrics as the traffic figures: correlating
|
||||
// "the cluster was full" with "requests were queueing" should
|
||||
// not need two systems.
|
||||
if err := rt.RegisterMetrics(fleetmetrics.New(ctl)); err != nil {
|
||||
return err
|
||||
}
|
||||
// So a request for a model that is configured but not running
|
||||
// starts it, rather than being told nothing serves it.
|
||||
rt.SetActivator(ctl.Reconciler())
|
||||
|
||||
log.Info("control plane running",
|
||||
"models", len(cfg.Models), "agent_listen", cfg.AgentListen)
|
||||
"models", len(cfg.Models), "agent_listen", cfg.AgentListen,
|
||||
"agent_mtls", cfg.AgentTLS.Enabled())
|
||||
}
|
||||
|
||||
adminMux, err := rt.AdminHandler()
|
||||
@@ -110,7 +140,42 @@ func newServerCmd() *cobra.Command {
|
||||
return err
|
||||
}
|
||||
if admin != nil {
|
||||
adminapi.New(admin).Register(adminMux)
|
||||
api := adminapi.New(admin)
|
||||
// Catalog changes need somewhere to persist, so this is wired
|
||||
// only when there is a store. Without one, the endpoints say so
|
||||
// rather than accepting a model that vanishes on restart.
|
||||
if st := rt.Store(); st != nil {
|
||||
svc := modelsvc.New(cfg.Models, st, admin, log)
|
||||
// Rebuild the catalog from configuration plus what was
|
||||
// stored, and put that in force — otherwise models added
|
||||
// before the last restart would be silently absent until
|
||||
// someone added them again.
|
||||
if built, err := svc.Load(ctx); err != nil {
|
||||
log.Error("could not load stored models; continuing with configuration only", "err", err)
|
||||
} else {
|
||||
admin.ReplaceCatalog(ctx, built)
|
||||
}
|
||||
api.SetModels(svc)
|
||||
}
|
||||
api.Register(adminMux)
|
||||
|
||||
// The dashboard lives here and only here. This listener is
|
||||
// bound to loopback by default and also serves pprof, which can
|
||||
// dump the process heap — where API keys sit in plaintext.
|
||||
// Reaching the page from elsewhere is an SSH tunnel's job:
|
||||
// ssh -L 9090:127.0.0.1:9090 <host>
|
||||
var usage web.UsageSource
|
||||
if st := rt.Store(); st != nil {
|
||||
usage = st
|
||||
}
|
||||
adminMux.HandleFunc("GET /", web.Dashboard(
|
||||
api, usage, version.Version, 5*time.Second, true, log))
|
||||
adminMux.Handle("GET /_static/", web.Static())
|
||||
|
||||
// The same page on the client-facing port at /admin, behind a
|
||||
// session and an admin flag. Handing the router the data rather
|
||||
// than the port keeps pprof where it is.
|
||||
rt.SetClusterView(api, usage)
|
||||
}
|
||||
|
||||
apiSrv := &http.Server{
|
||||
@@ -134,9 +199,35 @@ func newServerCmd() *cobra.Command {
|
||||
return fmt.Errorf("listening on %s: %w", cfg.AdminListen, err)
|
||||
}
|
||||
|
||||
// A metrics-only listener, when one is configured. Started after
|
||||
// AdminHandler, which is where the registry it serves is built.
|
||||
var metricsSrv *http.Server
|
||||
var metricsLn net.Listener
|
||||
if cfg.MetricsListen != "" {
|
||||
h := rt.MetricsHandler()
|
||||
if h == nil {
|
||||
_ = apiLn.Close()
|
||||
_ = adminLn.Close()
|
||||
return fmt.Errorf("metrics handler is not ready; this is a wiring fault, not a configuration one")
|
||||
}
|
||||
metricsLn, err = net.Listen("tcp", cfg.MetricsListen)
|
||||
if err != nil {
|
||||
_ = apiLn.Close()
|
||||
_ = adminLn.Close()
|
||||
return fmt.Errorf("listening on %s: %w", cfg.MetricsListen, err)
|
||||
}
|
||||
metricsSrv = &http.Server{Handler: h, ReadHeaderTimeout: 10 * time.Second}
|
||||
go func() {
|
||||
if err := metricsSrv.Serve(metricsLn); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
log.Error("metrics listener stopped", "err", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
log.Info("router listening",
|
||||
"api", apiLn.Addr().String(),
|
||||
"admin", adminLn.Addr().String(),
|
||||
"metrics", metricsAddr(metricsLn),
|
||||
"endpoints", len(cfg.Endpoints),
|
||||
"default_profile", cfg.Router.DefaultProfile,
|
||||
)
|
||||
@@ -147,17 +238,34 @@ func newServerCmd() *cobra.Command {
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(agentPath, agentHandler)
|
||||
agentSrv = &http.Server{
|
||||
// Agents speak plaintext HTTP/2, which Go will not do
|
||||
// without being asked. mTLS replaces this (plan §5.6);
|
||||
// until then the agent port must stay on a trusted network.
|
||||
Handler: h2c.NewHandler(mux, &http2.Server{}),
|
||||
Handler: mux,
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
if !cfg.AgentTLS.Enabled() {
|
||||
// Without TLS, agents speak plaintext HTTP/2, which Go will
|
||||
// not do unless asked. Nothing then stops a node from
|
||||
// claiming to be another one, so this belongs on a trusted
|
||||
// network only.
|
||||
agentSrv.Handler = h2c.NewHandler(mux, &http2.Server{})
|
||||
}
|
||||
if agentLn, err = net.Listen("tcp", cfg.AgentListen); err != nil {
|
||||
_ = apiLn.Close()
|
||||
_ = adminLn.Close()
|
||||
return fmt.Errorf("listening on %s: %w", cfg.AgentListen, err)
|
||||
}
|
||||
if cfg.AgentTLS.Enabled() {
|
||||
tc, terr := cfg.AgentTLS.ServerConfig()
|
||||
if terr != nil {
|
||||
_ = apiLn.Close()
|
||||
_ = adminLn.Close()
|
||||
_ = agentLn.Close()
|
||||
return terr
|
||||
}
|
||||
// Serve on an already-wrapped listener rather than
|
||||
// ServeTLS: agentSrv.TLSConfig stays nil, which is exactly
|
||||
// the case where net/http sets up HTTP/2 for ALPN itself.
|
||||
agentLn = tls.NewListener(agentLn, tc)
|
||||
}
|
||||
}
|
||||
|
||||
g, gctx := errgroup.WithContext(ctx)
|
||||
@@ -174,6 +282,9 @@ func newServerCmd() *cobra.Command {
|
||||
if agentSrv != nil {
|
||||
_ = agentSrv.Shutdown(sctx)
|
||||
}
|
||||
if metricsSrv != nil {
|
||||
_ = metricsSrv.Shutdown(sctx)
|
||||
}
|
||||
_ = adminSrv.Shutdown(sctx)
|
||||
return apiSrv.Shutdown(sctx)
|
||||
})
|
||||
@@ -184,7 +295,8 @@ func newServerCmd() *cobra.Command {
|
||||
f := cmd.Flags()
|
||||
f.StringVar(&configPath, "config", "", "configuration file (optional)")
|
||||
f.StringVar(&listen, "listen", "", "override the client-facing address")
|
||||
f.StringVar(&adminListen, "admin-listen", "", "override the metrics/pprof address")
|
||||
f.StringVar(&adminListen, "admin-listen", "", "override the dashboard/pprof address")
|
||||
f.StringVar(&metricsListen, "metrics-listen", "", "serve metrics only on this address, safe to expose")
|
||||
f.StringArrayVar(&endpoints, "endpoint", nil, "static endpoint: id=URL=model[,model] (repeatable)")
|
||||
|
||||
return cmd
|
||||
@@ -195,8 +307,8 @@ func newServerCmd() *cobra.Command {
|
||||
// The service is both the AgentService handler and the reconciler's Dispatcher:
|
||||
// it holds the open streams, so it is the only thing that can deliver a command
|
||||
// to a node.
|
||||
func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Catalog,
|
||||
sink control.EndpointSink, log *slog.Logger,
|
||||
func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Live,
|
||||
sink control.EndpointSink, requireIdentity bool, log *slog.Logger,
|
||||
) (*control.Controller, *agentsvc.Service, error) {
|
||||
ccfg := control.DefaultConfig()
|
||||
ccfg.Reconciler = reconciler.Config{
|
||||
@@ -217,13 +329,28 @@ func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Cata
|
||||
// the service once it exists.
|
||||
fwd := &dispatcherRef{}
|
||||
ctl := control.New(ccfg, cat, fwd, sink, log)
|
||||
svc := agentsvc.New(ctl.Inventory(), ctl.Reconciler(), log)
|
||||
// The router is both where endpoints go and where activity comes back
|
||||
// from. It only reports usage when it is the thing serving requests, which
|
||||
// is why this is a type assertion rather than a parameter.
|
||||
if u, ok := sink.(control.UsageSource); ok {
|
||||
ctl.SetUsageSource(u)
|
||||
}
|
||||
svc := agentsvc.New(ctl.Inventory(), ctl.Reconciler(), requireIdentity, log)
|
||||
fwd.set(svc)
|
||||
|
||||
ctl.Start(ctx)
|
||||
return ctl, svc, nil
|
||||
}
|
||||
|
||||
// metricsAddr renders the metrics listener for the startup line, saying so
|
||||
// plainly when there is none rather than printing an empty string.
|
||||
func metricsAddr(ln net.Listener) string {
|
||||
if ln == nil {
|
||||
return "off (admin listener only)"
|
||||
}
|
||||
return ln.Addr().String()
|
||||
}
|
||||
|
||||
// dispatcherRef lets the controller be built before the service exists.
|
||||
type dispatcherRef struct {
|
||||
mu sync.RWMutex
|
||||
@@ -245,6 +372,26 @@ func (r *dispatcherRef) get() (reconciler.Dispatcher, error) {
|
||||
return r.d, nil
|
||||
}
|
||||
|
||||
// RemoveWeights forwards a capability beyond scheduling.
|
||||
//
|
||||
// It has to be spelled out here rather than inherited, because this type
|
||||
// satisfies reconciler.Dispatcher by forwarding each method by hand — so
|
||||
// anything the controller reaches by type assertion stops at this boundary
|
||||
// unless it is forwarded too. A capability added to the service and not to
|
||||
// this type is simply absent at runtime, with nothing at compile time to say
|
||||
// so.
|
||||
func (r *dispatcherRef) RemoveWeights(ctx context.Context, nodeID string, paths []string) error {
|
||||
d, err := r.get()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rm, ok := d.(control.WeightRemover)
|
||||
if !ok {
|
||||
return errors.New("the agent service cannot remove weights")
|
||||
}
|
||||
return rm.RemoveWeights(ctx, nodeID, paths)
|
||||
}
|
||||
|
||||
func (r *dispatcherRef) StartShard(ctx context.Context, cmd reconciler.ShardCommand) error {
|
||||
d, err := r.get()
|
||||
if err != nil {
|
||||
|
||||
@@ -17,15 +17,39 @@ import (
|
||||
func newSimEngineCmd() *cobra.Command {
|
||||
cfg := sim.Default()
|
||||
var listen string
|
||||
var rv sim.Rendezvous
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "sim-engine",
|
||||
Short: "OpenAI-compatible fake inference engine (for testing)",
|
||||
Long: "Simulates the engine behaviours that matter to the router:\n" +
|
||||
"continuous batching, KV block accounting, prefix caching, and vLLM-shaped\n" +
|
||||
"/metrics. Lets the router be developed and tested without GPUs.",
|
||||
"/metrics. Lets the router be developed and tested without GPUs.\n\n" +
|
||||
"With --nnodes above 1 it also simulates a multi-node replica: rank 0\n" +
|
||||
"listens at --dist-init-addr and does not serve until every other rank has\n" +
|
||||
"joined, and only rank 0 answers the API. That is the part of the\n" +
|
||||
"multi-node path stub processes cannot exercise — the sequencing, the\n" +
|
||||
"waiting, and one dead shard taking the replica with it.",
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
log := logger()
|
||||
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
// Before anything binds the API port: a worker must never look
|
||||
// ready, and a leader must not serve until the replica exists.
|
||||
if err := rv.Join(ctx, log); err != nil {
|
||||
return err
|
||||
}
|
||||
if rv.Multi() && !rv.Leader() {
|
||||
// A worker contributes GPUs and answers nothing. Health
|
||||
// checking its port is what used to fail replicas that were
|
||||
// coming up perfectly well.
|
||||
log.Info("worker rank running; this rank serves no API", "rank", rv.Rank)
|
||||
<-ctx.Done()
|
||||
return nil
|
||||
}
|
||||
|
||||
eng := sim.New(cfg)
|
||||
defer eng.Close()
|
||||
|
||||
@@ -47,9 +71,6 @@ func newSimEngineCmd() *cobra.Command {
|
||||
"prefix_cache_blocks", cfg.PrefixCacheBlocks,
|
||||
)
|
||||
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() { errCh <- srv.Serve(ln) }()
|
||||
|
||||
@@ -70,6 +91,11 @@ func newSimEngineCmd() *cobra.Command {
|
||||
|
||||
f := cmd.Flags()
|
||||
f.StringVar(&listen, "listen", ":8000", "address to listen on")
|
||||
f.IntVar(&rv.NNodes, "nnodes", 1, "ranks in this replica; above 1 simulates a multi-node launch")
|
||||
f.IntVar(&rv.Rank, "node-rank", 0, "this rank; 0 is the leader and the only one that serves")
|
||||
f.StringVar(&rv.Addr, "dist-init-addr", "", "the leader's rendezvous address")
|
||||
f.DurationVar(&rv.Timeout, "rendezvous-timeout", 5*time.Minute,
|
||||
"how long to wait for the replica to assemble before failing")
|
||||
f.StringVar(&cfg.Model, "model", cfg.Model, "model name to serve")
|
||||
f.IntVar(&cfg.MaxNumSeqs, "max-num-seqs", cfg.MaxNumSeqs, "concurrent sequences (continuous batch capacity)")
|
||||
f.IntVar(&cfg.BlockSize, "block-size", cfg.BlockSize, "tokens per KV cache block")
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
// newStatusCmd prints the whole fleet in one screen.
|
||||
//
|
||||
// `ctl nodes` is the right tool for four nodes and the wrong one for two
|
||||
// hundred. The questions that matter at that size are aggregates — how full am
|
||||
// I, what is broken, what is holding cards nobody is using — and reading them
|
||||
// off a long table is a chore, which means it gets skipped, which means nobody
|
||||
// notices until something fails.
|
||||
func newStatusCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "status",
|
||||
Short: "The whole fleet in one answer: capacity, health, and what is idle",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var s adminapi.Summary
|
||||
if err := client().get(cmd.Context(), "/api/v1/summary", &s); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(s)
|
||||
}
|
||||
printSummary(cmd, s)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func printSummary(cmd *cobra.Command, s adminapi.Summary) {
|
||||
out := cmd.OutOrStdout()
|
||||
|
||||
fmt.Fprintf(out, "NODES %d ready", s.Nodes.Ready)
|
||||
for _, x := range []struct {
|
||||
n int
|
||||
label string
|
||||
}{
|
||||
{s.Nodes.Cordoned, "cordoned"},
|
||||
{s.Nodes.Draining, "draining"},
|
||||
{s.Nodes.Unreachable, "unreachable"},
|
||||
} {
|
||||
if x.n > 0 {
|
||||
fmt.Fprintf(out, ", %d %s", x.n, x.label)
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(out, " (of %d)\n", s.Nodes.Total)
|
||||
|
||||
if len(s.GPUs) > 0 {
|
||||
fmt.Fprintln(out, "\nGPUS")
|
||||
w := tab()
|
||||
fmt.Fprintln(w, " CLASS\tIN USE\tFREE\tUNHEALTHY\tUSAGE")
|
||||
for _, g := range s.GPUs {
|
||||
fmt.Fprintf(w, " %s\t%d/%d\t%d\t%s\t%s\n",
|
||||
g.Class, g.Used, g.Total, g.Free(), unhealthyCell(g.Unhealthy), bar(g.Used, g.Total))
|
||||
}
|
||||
_ = w.Flush()
|
||||
}
|
||||
|
||||
if len(s.Fabrics) > 0 {
|
||||
fmt.Fprintln(out, "\nFABRICS")
|
||||
w := tab()
|
||||
// The free-block column is the one that decides whether a multi-node
|
||||
// replica can land: free cards on the wrong fabric are not capacity for
|
||||
// a model that spans nodes, because a replica may never cross one.
|
||||
fmt.Fprintln(w, " FABRIC\tNODES\tFULLY FREE\tWEIGHTS")
|
||||
for _, f := range s.Fabrics {
|
||||
fmt.Fprintf(w, " %s\t%d\t%d\t%s\n",
|
||||
f.Fabric, f.Nodes, f.LargestFreeBlock, weightsCell(f))
|
||||
}
|
||||
_ = w.Flush()
|
||||
}
|
||||
|
||||
fmt.Fprintf(out, "\nREPLICAS %d ready", s.Replicas.Ready)
|
||||
for _, x := range []struct {
|
||||
n int
|
||||
label string
|
||||
}{
|
||||
{s.Replicas.Starting, "starting"},
|
||||
{s.Replicas.Draining, "draining"},
|
||||
{s.Replicas.Failed, "failed"},
|
||||
} {
|
||||
if x.n > 0 {
|
||||
fmt.Fprintf(out, ", %d %s", x.n, x.label)
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(out, "\nMODELS %d of %d have the replicas they asked for\n",
|
||||
s.Models.Satisfied, s.Models.Total)
|
||||
|
||||
if len(s.Unplaceable) > 0 {
|
||||
fmt.Fprintln(out, "\nNOT PLACED")
|
||||
for _, m := range s.Unplaceable {
|
||||
fmt.Fprintf(out, " %s\n", m)
|
||||
}
|
||||
fmt.Fprintln(out, "\n Why, per model: llmgw ctl explain-placement <model>")
|
||||
}
|
||||
|
||||
if len(s.IdleReplicas) > 0 {
|
||||
fmt.Fprintln(out, "\nIDLE")
|
||||
w := tab()
|
||||
fmt.Fprintln(w, " REPLICA\tMODEL\tGPUS\tUNUSED FOR")
|
||||
for _, r := range s.IdleReplicas {
|
||||
fmt.Fprintf(w, " %s\t%s\t%d\t%s\n", r.ID, r.Model, r.GPUs,
|
||||
(time.Duration(r.IdleSec) * time.Second).Round(time.Minute))
|
||||
}
|
||||
_ = w.Flush()
|
||||
fmt.Fprintln(out, "\n These hold hardware nothing is asking for.")
|
||||
}
|
||||
}
|
||||
|
||||
// weightsCell says what a fabric holds, and stays quiet when nobody asked.
|
||||
//
|
||||
// Fabrics are not connected to one another, so weights in one are not weights
|
||||
// in another: a model seeded into a single fabric leaves the rest of the
|
||||
// cluster unable to serve it without first moving tens of gigabytes over the
|
||||
// management network. This is the column that shows that before someone
|
||||
// discovers it during a scale-up.
|
||||
func weightsCell(f adminapi.FabricUse) string {
|
||||
if f.NodesWithWeights == 0 {
|
||||
// Nobody on this fabric was asked to look, which is the default and is
|
||||
// not the same as an empty fabric.
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%d model(s) on %d/%d nodes", f.Models, f.NodesWithWeights, f.Nodes)
|
||||
}
|
||||
|
||||
// bar draws usage so a column of classes can be read at a glance rather than
|
||||
// compared digit by digit.
|
||||
func bar(used, total int) string {
|
||||
const width = 20
|
||||
if total <= 0 {
|
||||
return ""
|
||||
}
|
||||
filled := used * width / total
|
||||
if filled > width {
|
||||
filled = width
|
||||
}
|
||||
b := make([]rune, width)
|
||||
for i := range b {
|
||||
if i < filled {
|
||||
b[i] = '#'
|
||||
} else {
|
||||
b[i] = '.'
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("%s %3d%%", string(b), used*100/total)
|
||||
}
|
||||
|
||||
// unhealthyCell keeps a zero quiet. A column of noughts trains the eye to skip
|
||||
// the column, which is the opposite of what it is for.
|
||||
func unhealthyCell(n int) string {
|
||||
if n == 0 {
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%d", n)
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
|
||||
)
|
||||
|
||||
func TestBarStaysWithinItsWidth(t *testing.T) {
|
||||
// It is drawn from counts that can disagree briefly — a GPU freed between
|
||||
// two reads — and a bar that overruns its column wrecks the table it is in.
|
||||
for _, tc := range [][2]int{{0, 0}, {0, 8}, {8, 8}, {9, 8}, {100, 8}} {
|
||||
got := bar(tc[0], tc[1])
|
||||
if tc[1] == 0 {
|
||||
if got != "" {
|
||||
t.Errorf("bar(%d,%d) = %q, want empty for an empty class", tc[0], tc[1], got)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if n := strings.Count(got, "#") + strings.Count(got, "."); n != 20 {
|
||||
t.Errorf("bar(%d,%d) drew %d cells: %q", tc[0], tc[1], n, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdleColumnDistinguishesUnknownFromIdle(t *testing.T) {
|
||||
// Every replica looks unreported before its first request and after a
|
||||
// router restart. Showing that as a very large duration would make an
|
||||
// untouched fleet read as an abandoned one.
|
||||
if got := idleColumn(-1); got != "-" {
|
||||
t.Errorf("idleColumn(-1) = %q, want a dash", got)
|
||||
}
|
||||
if got := idleColumn(90); got != "1m30s" {
|
||||
t.Errorf("idleColumn(90) = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelStatusSeparatesWaitingFromBroken(t *testing.T) {
|
||||
// These used to differ by a parenthetical — "0/1 ready (1 total)" against
|
||||
// "0/1 ready" — and they need opposite responses: one means wait, the other
|
||||
// means look at explain-placement. On forty models that difference is
|
||||
// invisible.
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
view adminapi.ModelView
|
||||
want string
|
||||
}{
|
||||
{"serving what it asked for", adminapi.ModelView{Min: 2, Ready: 2, Total: 2}, "ready"},
|
||||
{"a replica is coming up", adminapi.ModelView{Min: 1, Ready: 0, Total: 1}, "starting"},
|
||||
{"nothing could be placed", adminapi.ModelView{Min: 1, Ready: 0, Total: 0}, "not placed"},
|
||||
{"deliberately not running", adminapi.ModelView{Min: 0, Ready: 0, Total: 0}, "stopped"},
|
||||
{"short but rebuilding", adminapi.ModelView{Min: 2, Ready: 1, Total: 2}, "starting"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := modelStatus(tc.view); got != tc.want {
|
||||
t.Errorf("modelStatus(%+v) = %q, want %q", tc.view, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHardwarePoolsNeverMergeAcrossFabrics(t *testing.T) {
|
||||
// A replica may never span two fabrics, so eight free cards split across
|
||||
// two of them are not eight cards for a model that needs eight. Merging
|
||||
// them would plan something that starts, passes health checks and hangs on
|
||||
// its first all-reduce.
|
||||
nodes := []adminapi.NodeView{
|
||||
{ID: "a", State: "ready", Classes: []string{"h100-80g"}, Fabrics: []string{"ib-a"}, GPUFree: 4, GPUTotal: 4},
|
||||
{ID: "b", State: "ready", Classes: []string{"h100-80g"}, Fabrics: []string{"ib-a"}, GPUFree: 4, GPUTotal: 4},
|
||||
{ID: "c", State: "ready", Classes: []string{"h100-80g"}, Fabrics: []string{"ib-b"}, GPUFree: 4, GPUTotal: 4},
|
||||
}
|
||||
hw := hardwareFrom(nodes)
|
||||
if len(hw) != 2 {
|
||||
t.Fatalf("expected one pool per fabric, got %d: %+v", len(hw), hw)
|
||||
}
|
||||
for _, p := range hw {
|
||||
switch p.Fabric {
|
||||
case "ib-a":
|
||||
if p.Nodes != 2 {
|
||||
t.Errorf("ib-a has %d nodes, want 2", p.Nodes)
|
||||
}
|
||||
case "ib-b":
|
||||
if p.Nodes != 1 {
|
||||
t.Errorf("ib-b has %d nodes, want 1", p.Nodes)
|
||||
}
|
||||
default:
|
||||
t.Errorf("unexpected fabric %q", p.Fabric)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreadyNodesAreNotCountedAsCapacity(t *testing.T) {
|
||||
// Planning against a node that is draining or unreachable produces a fit
|
||||
// that cannot be placed, which is worse than saying nothing fits.
|
||||
nodes := []adminapi.NodeView{
|
||||
{ID: "a", State: "ready", Classes: []string{"h100-80g"}, GPUFree: 8, GPUTotal: 8},
|
||||
{ID: "b", State: "draining", Classes: []string{"h100-80g"}, GPUFree: 8, GPUTotal: 8},
|
||||
{ID: "c", State: "unreachable", Classes: []string{"h100-80g"}, GPUFree: 8, GPUTotal: 8},
|
||||
}
|
||||
hw := hardwareFrom(nodes)
|
||||
if len(hw) != 1 || hw[0].Nodes != 1 {
|
||||
t.Fatalf("expected only the ready node, got %+v", hw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCardMemoryComesFromTheClassLabel(t *testing.T) {
|
||||
// The label carries it by construction — "h100-80g" is eighty gigabytes and
|
||||
// says so — which is why the planner needs no table of card models.
|
||||
for label, want := range map[string]int{
|
||||
"h100-80g": 80, "a100-40g": 40, "l40s-48g": 48, "rtx4090-24g": 24,
|
||||
"nonsense": 0, "": 0, "h100-": 0,
|
||||
} {
|
||||
if got := memFromClass(label); got != want {
|
||||
t.Errorf("memFromClass(%q) = %d, want %d", label, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestACPUOnlyNodeBecomesAMemoryPool(t *testing.T) {
|
||||
// How this cluster began: a node with no cards is not an empty pool, it is
|
||||
// a pool measured in RAM.
|
||||
hw := hardwareFrom([]adminapi.NodeView{
|
||||
{ID: "cpu1", State: "ready", MemoryGB: 31},
|
||||
})
|
||||
if len(hw) != 1 {
|
||||
t.Fatalf("the CPU node produced %d pools", len(hw))
|
||||
}
|
||||
if hw[0].Class != "" || hw[0].HostMemGB != 31 {
|
||||
t.Errorf("pool = %+v, want a classless pool with 31 GB", hw[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAGPUNodeWithNoFreeCardsIsNotAPool(t *testing.T) {
|
||||
// Fully allocated hardware is not capacity, and offering it would plan a
|
||||
// replica that cannot be placed.
|
||||
hw := hardwareFrom([]adminapi.NodeView{
|
||||
{ID: "full", State: "ready", Classes: []string{"h100-80g"}, GPUFree: 0, GPUTotal: 8},
|
||||
})
|
||||
for _, p := range hw {
|
||||
if p.Class != "" {
|
||||
t.Errorf("a node with no free cards became a pool: %+v", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheEntryNamesAnEngineThatCanServeThePrecision(t *testing.T) {
|
||||
// The defect this prevents, found by reading the tool's own output: it
|
||||
// chose q8_0 — a GGUF quantisation — and emitted a vLLM profile with no
|
||||
// quantisation at all. Following that loads full-precision weights, which
|
||||
// the same plan had just shown do not fit on one card. The tool would have
|
||||
// caused the out-of-memory it exists to prevent.
|
||||
m := modelplan.Model{Name: "Qwen2.5-Coder-32B", Params: 32_800_000_000, Layers: 64, KVHeads: 8, HeadDim: 128}
|
||||
|
||||
gguf := modelplan.Fit{Precision: modelplan.Q8GGUF, Class: "h100-80g", TP: 1, Nodes: 1}
|
||||
entry := catalogEntry(m, gguf, "Qwen/Qwen2.5-Coder-32B-Instruct", 32768)
|
||||
if !strings.Contains(entry, "engine: llamacpp") {
|
||||
t.Errorf("a GGUF precision produced a non-llamacpp profile:\n%s", entry)
|
||||
}
|
||||
if strings.Contains(entry, "source: Qwen/") {
|
||||
t.Errorf("a GGUF entry points at a repository rather than a file:\n%s", entry)
|
||||
}
|
||||
if !strings.Contains(entry, ".gguf") {
|
||||
t.Errorf("a GGUF entry has no file path:\n%s", entry)
|
||||
}
|
||||
|
||||
fp8 := modelplan.Fit{Precision: modelplan.FP8, Class: "h100-80g", TP: 2, Nodes: 1}
|
||||
entry = catalogEntry(m, fp8, "Qwen/Qwen2.5-Coder-32B-Instruct", 32768)
|
||||
if !strings.Contains(entry, "engine: vllm") {
|
||||
t.Errorf("fp8 did not produce a vLLM profile:\n%s", entry)
|
||||
}
|
||||
// And the quantisation has to reach the engine, or vLLM loads bf16.
|
||||
if !strings.Contains(entry, "--quantization") || !strings.Contains(entry, "fp8") {
|
||||
t.Errorf("the fp8 quantisation is not passed to the engine:\n%s", entry)
|
||||
}
|
||||
if !strings.Contains(entry, "source: Qwen/") {
|
||||
t.Errorf("a vLLM entry should load from the repository:\n%s", entry)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// newThreadsCmd turns conversation storage on for a project, or off.
|
||||
//
|
||||
// A command rather than a configuration key because it is a decision about
|
||||
// data, not about deployment: someone has to choose to keep what a team types,
|
||||
// and choosing it should look like an act rather than like a default nobody
|
||||
// read.
|
||||
//
|
||||
// Until it is run for a project, llmgw stores no message content for that
|
||||
// project — usage rows carry token counts and costs, never text — and that is
|
||||
// how every project starts.
|
||||
func newThreadsCmd(open opener) *cobra.Command {
|
||||
var (
|
||||
project string
|
||||
hours int
|
||||
off bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "threads",
|
||||
Short: "Keep conversations for a project, so a thread can move between models",
|
||||
Long: "With this on, a client sends x-llmgw-thread and only its new turn, and\n" +
|
||||
"llmgw replays the conversation to whichever model serves — including a\n" +
|
||||
"different one from the one it started on.\n\n" +
|
||||
"Off for every project until this is run. llmgw otherwise stores no\n" +
|
||||
"message content anywhere, which is a much easier thing to have deployed\n" +
|
||||
"in front of a team, so turning it on is deliberate.\n\n" +
|
||||
"The retention is required rather than optional: a conversation store with\n" +
|
||||
"no expiry is a transcript of everything a team ever typed, kept forever\n" +
|
||||
"because nobody chose a number.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if !off && hours <= 0 {
|
||||
return fmt.Errorf("--retain-hours is required: say how long conversations are kept, or pass --off")
|
||||
}
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
keep := hours
|
||||
if off {
|
||||
keep = 0
|
||||
}
|
||||
if err := s.SetThreadRetention(cmd.Context(), project, keep); err != nil {
|
||||
return err
|
||||
}
|
||||
if keep == 0 {
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s keeps no conversations. Existing ones are removed at the next prune.\n", project)
|
||||
return nil
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s keeps conversations for %d hours, measured from the last turn.\n", project, keep)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project id")
|
||||
cmd.Flags().IntVar(&hours, "retain-hours", 0, "how long a conversation is kept after its last turn")
|
||||
cmd.Flags().BoolVar(&off, "off", false, "stop keeping conversations for this project")
|
||||
return cmd
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// newToolCheckCmd asks a running engine whether tool calls come back as tool
|
||||
// calls.
|
||||
//
|
||||
// A model's capabilities are declared in the catalog, which makes them claims.
|
||||
// This is the one that can be checked, and the one worth checking: the gateway
|
||||
// forwards a tools array to any model, and an engine that does not return
|
||||
// structured tool_calls looks exactly like a model that ignores its tools.
|
||||
// Every OpenAI-compatible client then behaves as though the model is broken.
|
||||
//
|
||||
// Found the hard way. qwen2.5-coder produced exactly the right call —
|
||||
// {"name": "list_files", "arguments": {"path": "."}} — as plain reply text,
|
||||
// with tool_calls null, and diagnosing that took an afternoon of eliminating
|
||||
// the gateway, the model size, the chat template and tool_choice in turn.
|
||||
//
|
||||
// The answer was neither the engine nor the model alone, which is why this
|
||||
// command now distinguishes them by what the reply looks like rather than
|
||||
// asserting one. The engine rendered the tool prompt correctly, including the
|
||||
// <tool_call> wrapper its template asks the model to use; the model wrapped
|
||||
// the call in a markdown fence instead, and the parser — which looks for that
|
||||
// wrapper — left it in the content. A differently trained model fixes that. A
|
||||
// bigger one from the same family does not: qwen2.5-coder-7b invented
|
||||
// <function-call> and failed the same way.
|
||||
func newToolCheckCmd() *cobra.Command {
|
||||
var (
|
||||
endpoint string
|
||||
model string
|
||||
timeout time.Duration
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "tool-check",
|
||||
Short: "Check that an engine returns tool calls as tool_calls",
|
||||
Long: "The catalog declares whether a model supports tools, which makes it a\n" +
|
||||
"claim. This checks it against a running engine.\n\n" +
|
||||
"A model that cannot use tools, and an engine that does not parse the\n" +
|
||||
"calls it makes, are indistinguishable from outside: both answer in prose\n" +
|
||||
"while tool_calls stays null. The difference matters, because one is fixed\n" +
|
||||
"by choosing another model and the other by changing engine build.\n\n" +
|
||||
"Point it at an engine's own address, not the gateway, so nothing in\n" +
|
||||
"between is part of the answer.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if endpoint == "" {
|
||||
return fmt.Errorf("--endpoint is required, e.g. http://10.0.0.145:20002")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(cmd.Context(), timeout)
|
||||
defer cancel()
|
||||
|
||||
res, err := probeToolCall(ctx, endpoint, model)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
report(cmd, res)
|
||||
if !res.Structured {
|
||||
// A non-zero exit so this is usable from a script that gates a
|
||||
// rollout on it.
|
||||
return errToolCallsUnsupported
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&endpoint, "endpoint", "", "engine base URL, e.g. http://10.0.0.145:20002")
|
||||
cmd.Flags().StringVar(&model, "model", "", "model name to send, when the engine needs one")
|
||||
cmd.Flags().DurationVar(&timeout, "timeout", 5*time.Minute, "how long to wait; CPU inference is slow")
|
||||
return cmd
|
||||
}
|
||||
|
||||
var errToolCallsUnsupported = fmt.Errorf("this engine does not return structured tool calls")
|
||||
|
||||
// ToolProbe is what one check found.
|
||||
type ToolProbe struct {
|
||||
// Structured is the answer: did tool_calls come back populated.
|
||||
Structured bool
|
||||
// Attempted means the model tried to call the tool, in whatever form. It
|
||||
// separates "the model ignored its tools" from "the engine did not parse
|
||||
// the call the model made" — the same symptom, different fix.
|
||||
Attempted bool
|
||||
// Wrapped means the attempt used the <tool_call> tags that the Qwen and
|
||||
// Hermes templates instruct the model to use.
|
||||
//
|
||||
// This is what separates the two halves of a failed attempt. Wrapped and
|
||||
// unparsed is the engine's parser failing on output that was correct.
|
||||
// Unwrapped is the model answering in a shape of its own — the engine had
|
||||
// no way to recognise it, and the fix is a model trained for the format,
|
||||
// not a newer engine.
|
||||
Wrapped bool
|
||||
Content string
|
||||
Raw string
|
||||
}
|
||||
|
||||
func probeToolCall(ctx context.Context, endpoint, model string) (ToolProbe, error) {
|
||||
// A question with one obvious tool answer and no other sensible reply, so a
|
||||
// model that can use tools has no excuse not to.
|
||||
body := map[string]any{
|
||||
"messages": []map[string]string{
|
||||
{"role": "user", "content": "What is the weather in Istanbul? Use the tool."},
|
||||
},
|
||||
"tools": []map[string]any{{
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather in a city",
|
||||
"parameters": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{"city": map[string]any{"type": "string"}},
|
||||
"required": []string{"city"},
|
||||
},
|
||||
},
|
||||
}},
|
||||
"tool_choice": "auto",
|
||||
"max_tokens": 200,
|
||||
}
|
||||
if model != "" {
|
||||
body["model"] = model
|
||||
}
|
||||
buf, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return ToolProbe{}, err
|
||||
}
|
||||
|
||||
url := endpoint + "/v1/chat/completions"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(buf))
|
||||
if err != nil {
|
||||
return ToolProbe{}, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return ToolProbe{}, fmt.Errorf("asking %s: %w", url, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if err != nil {
|
||||
return ToolProbe{}, err
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return ToolProbe{}, fmt.Errorf("%s returned %d: %s", url, resp.StatusCode, trim(string(raw), 300))
|
||||
}
|
||||
return readToolProbe(raw)
|
||||
}
|
||||
|
||||
func readToolProbe(raw []byte) (ToolProbe, error) {
|
||||
var parsed struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
ToolCalls []json.RawMessage `json:"tool_calls"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &parsed); err != nil {
|
||||
return ToolProbe{}, fmt.Errorf("the engine's reply is not a chat completion: %w", err)
|
||||
}
|
||||
if len(parsed.Choices) == 0 {
|
||||
return ToolProbe{}, fmt.Errorf("the engine returned no choices")
|
||||
}
|
||||
msg := parsed.Choices[0].Message
|
||||
p := ToolProbe{
|
||||
Structured: len(msg.ToolCalls) > 0,
|
||||
Content: msg.Content,
|
||||
Raw: trim(string(raw), 600),
|
||||
}
|
||||
// The model tried, and the engine handed the attempt back as text. The
|
||||
// tell is the tool's name in a reply that was supposed to be a call.
|
||||
p.Attempted = p.Structured || bytes.Contains([]byte(msg.Content), []byte("get_weather"))
|
||||
p.Wrapped = bytes.Contains([]byte(msg.Content), []byte("<tool_call>"))
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func report(cmd *cobra.Command, p ToolProbe) {
|
||||
out := cmd.OutOrStdout()
|
||||
switch {
|
||||
case p.Structured:
|
||||
fmt.Fprintln(out, "tool calls work: the engine returned a populated tool_calls array.")
|
||||
case p.Attempted && p.Wrapped:
|
||||
fmt.Fprintln(out, "the model called the tool correctly and the engine did not parse it.")
|
||||
fmt.Fprintln(out, "\n The call is in the reply text instead of tool_calls:")
|
||||
fmt.Fprintf(out, " %s\n", trim(p.Content, 200))
|
||||
fmt.Fprintln(out, "\n It used the <tool_call> wrapper its template asks for, so the model")
|
||||
fmt.Fprintln(out, " did its part. This is the engine's parser. Check its build; another")
|
||||
fmt.Fprintln(out, " model will fail the same way.")
|
||||
case p.Attempted:
|
||||
fmt.Fprintln(out, "the model called the tool in a shape the engine does not recognise.")
|
||||
fmt.Fprintln(out, "\n The call is in the reply text instead of tool_calls:")
|
||||
fmt.Fprintf(out, " %s\n", trim(p.Content, 200))
|
||||
fmt.Fprintln(out, "\n The engine rendered the tools correctly — the model knew the tool")
|
||||
fmt.Fprintln(out, " existed and what to pass it — but did not use the <tool_call> wrapper")
|
||||
fmt.Fprintln(out, " its template asked for, so there was nothing for the parser to find.")
|
||||
fmt.Fprintln(out, "\n Fixed by a model trained for this format, not by a bigger one: the")
|
||||
fmt.Fprintln(out, " instruct models generally manage it where the coder variants do not.")
|
||||
default:
|
||||
fmt.Fprintln(out, "the model did not try to call the tool.")
|
||||
fmt.Fprintf(out, "\n It replied: %s\n", trim(p.Content, 200))
|
||||
fmt.Fprintln(out, "\n This is the model, not the engine. A model trained for tool use")
|
||||
fmt.Fprintln(out, " would call one here; try another before suspecting the engine.")
|
||||
}
|
||||
}
|
||||
|
||||
func trim(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "…"
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAStructuredToolCallIsRecognised(t *testing.T) {
|
||||
p, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"tool_calls":[{"id":"c1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Istanbul\"}"}}]}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !p.Structured {
|
||||
t.Error("a populated tool_calls array was not recognised")
|
||||
}
|
||||
if !p.Attempted {
|
||||
t.Error("a structured call must also count as an attempt")
|
||||
}
|
||||
}
|
||||
|
||||
func TestACallReturnedAsTextIsAnAttempt(t *testing.T) {
|
||||
// The exact reply that cost an afternoon: the model made the right call and
|
||||
// the engine handed it back as prose.
|
||||
p, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"content":"{\"name\": \"get_weather\", \"arguments\": {\"city\": \".\"}}"}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if p.Structured {
|
||||
t.Fatal("text was reported as a structured call")
|
||||
}
|
||||
if !p.Attempted {
|
||||
t.Error("the model plainly tried; that must be distinguished from ignoring the tools")
|
||||
}
|
||||
if p.Wrapped {
|
||||
t.Error("bare JSON was read as having used the <tool_call> wrapper")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWrappedCallBlamesTheEngineAndABareOneDoesNot(t *testing.T) {
|
||||
// Both look identical from outside — tool_calls null, the call sitting in
|
||||
// the content — and they have opposite fixes. The whole point of the
|
||||
// command is telling them apart, and the earlier version did not: it
|
||||
// blamed the engine for both, which sends someone off rebuilding llama.cpp
|
||||
// when the answer is a different model.
|
||||
wrapped, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"content":"<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Istanbul\"}}\n</tool_call>"}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !wrapped.Attempted || !wrapped.Wrapped {
|
||||
t.Errorf("a correctly wrapped call was not recognised as one: %+v", wrapped)
|
||||
}
|
||||
|
||||
// What qwen2.5-coder-3b actually returns, measured: the right call in a
|
||||
// markdown fence, which the parser has no way to find.
|
||||
fenced, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"content":"` + "```" + `json\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Istanbul\"}}\n` + "```" + `"}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !fenced.Attempted {
|
||||
t.Error("a fenced call is still an attempt")
|
||||
}
|
||||
if fenced.Wrapped {
|
||||
t.Error("a markdown fence was mistaken for the <tool_call> wrapper")
|
||||
}
|
||||
|
||||
// And what qwen2.5-coder-7b returns: a tag it invented, which is not the
|
||||
// one the template asked for and is not parsed either.
|
||||
invented, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"content":"<function-call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Istanbul\"}}\n</function-call>"}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !invented.Attempted || invented.Wrapped {
|
||||
t.Errorf("an invented tag was read as the template's wrapper: %+v", invented)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAModelThatIgnoresItsToolsIsBlamedOnTheModel(t *testing.T) {
|
||||
// The other diagnosis, and it has a different fix: try another model rather
|
||||
// than another engine build.
|
||||
p, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"content":"The weather in Istanbul is usually mild in autumn."}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if p.Structured || p.Attempted {
|
||||
t.Errorf("a plain answer was read as a tool call: %+v", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAReplyThatIsNotACompletionIsAnError(t *testing.T) {
|
||||
for _, bad := range []string{`{}`, `not json`, `{"choices":[]}`} {
|
||||
if _, err := readToolProbe([]byte(bad)); err == nil {
|
||||
t.Errorf("no error for %q", bad)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrimKeepsOutputReadable(t *testing.T) {
|
||||
if got := trim(strings.Repeat("x", 100), 10); len(got) > 14 {
|
||||
t.Errorf("trim returned %d characters", len(got))
|
||||
}
|
||||
if got := trim("short", 10); got != "short" {
|
||||
t.Errorf("trim altered a short string: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/account"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func newUsersCmd() *cobra.Command {
|
||||
var dbPath string
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "users",
|
||||
Short: "Approve and manage the people who sign up",
|
||||
Long: "Signing up creates an account that cannot do anything. Approval is what\n" +
|
||||
"makes it work, and it happens here rather than in a browser: approving is\n" +
|
||||
"the step that decides who gets to spend the cluster, and it needs a shell\n" +
|
||||
"on this machine rather than a session cookie.",
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store")
|
||||
open := func(cmd *cobra.Command) (*store.SQLite, error) {
|
||||
return store.OpenSQLite(cmd.Context(), dbPath)
|
||||
}
|
||||
|
||||
cmd.AddCommand(newUserListCmd(open), newUserApproveCmd(open),
|
||||
newUserStateCmd(open, "suspend", store.UserSuspended,
|
||||
"Stop an account working, keeping it and its history"),
|
||||
newUserStateCmd(open, "reinstate", store.UserActive,
|
||||
"Let a suspended account work again"),
|
||||
newUserCreateCmd(open), newUserResetPasswordCmd(open),
|
||||
newUserAdminCmd(open, "promote", true, "Let an account see the cluster dashboard at /admin"),
|
||||
newUserAdminCmd(open, "demote", false, "Take away access to the cluster dashboard"))
|
||||
return cmd
|
||||
}
|
||||
|
||||
type userOpener func(*cobra.Command) (*store.SQLite, error)
|
||||
|
||||
func newUserListCmd(open userOpener) *cobra.Command {
|
||||
var pending bool
|
||||
cmd := &cobra.Command{
|
||||
Use: "list",
|
||||
Short: "List accounts",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
s, err := open(cmd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
state := store.UserState("")
|
||||
if pending {
|
||||
state = store.UserPending
|
||||
}
|
||||
users, err := s.Users(cmd.Context(), state)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(users) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), "no accounts")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
|
||||
fmt.Fprintln(w, "ID\tEMAIL\tNAME\tSTATE\tADMIN\tPROJECT\tSIGNED UP\tLAST LOGIN")
|
||||
now := time.Now()
|
||||
for _, u := range users {
|
||||
admin := "-"
|
||||
if u.IsAdmin {
|
||||
admin = "yes"
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
u.ID, u.Email, u.Name, u.State, admin, dash(u.ProjectID),
|
||||
ago(u.CreatedAt, now), ago(u.LastLoginAt, now))
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
}
|
||||
cmd.Flags().BoolVar(&pending, "pending", false, "only accounts waiting for approval")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newUserApproveCmd(open userOpener) *cobra.Command {
|
||||
var project, by string
|
||||
cmd := &cobra.Command{
|
||||
Use: "approve <user-id>",
|
||||
Short: "Let an account sign in and issue its own keys",
|
||||
Long: "The project is required: it decides which limits this person's keys\n" +
|
||||
"inherit. An account approved into no project would pass every check and\n" +
|
||||
"then fail at the first key it tried to make.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
s, err := open(cmd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
if err := s.ApproveUser(cmd.Context(), args[0], project, by, time.Now()); err != nil {
|
||||
return fmt.Errorf("approving %s: %w", args[0], err)
|
||||
}
|
||||
u, err := s.UserByID(cmd.Context(), args[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s approved into project %s\n", u.Email, project)
|
||||
cmd.PrintErrln("They can sign in and create their own keys from /profile.")
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project the person's keys belong to")
|
||||
cmd.Flags().StringVar(&by, "by", "", "who approved, for the record")
|
||||
_ = cmd.MarkFlagRequired("project")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newUserStateCmd(open userOpener, verb string, state store.UserState, short string) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: verb + " <user-id>",
|
||||
Short: short,
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
s, err := open(cmd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
if err := s.SetUserState(cmd.Context(), args[0], state); err != nil {
|
||||
return err
|
||||
}
|
||||
if state == store.UserSuspended {
|
||||
// Otherwise the person stays signed in and suspension is
|
||||
// advisory until their session happens to expire.
|
||||
if err := s.DeleteUserSessions(cmd.Context(), args[0]); err != nil {
|
||||
return err
|
||||
}
|
||||
cmd.PrintErrln("Their sessions were closed. Their API keys still work — revoke those separately.")
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s is now %s\n", args[0], state)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newUserAdminCmd(open userOpener, verb string, admin bool, short string) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: verb + " <user-id>",
|
||||
Short: short,
|
||||
Long: "The dashboard is read-only, so this grants looking and nothing else.\n" +
|
||||
"Everything that changes the cluster or issues a credential stays here on\n" +
|
||||
"the CLI, where it needs a shell rather than a session — which keeps the\n" +
|
||||
"worst case of a stolen admin cookie to \"saw the cluster\".",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
s, err := open(cmd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
if err := s.SetAdmin(cmd.Context(), args[0], admin); err != nil {
|
||||
return err
|
||||
}
|
||||
u, err := s.UserByID(cmd.Context(), args[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if admin {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s can now see /admin\n", u.Email)
|
||||
} else {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s can no longer see /admin\n", u.Email)
|
||||
cmd.PrintErrln("Their session stays open; the page is read-only. " +
|
||||
"Suspend the account to close it.")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newUserResetPasswordCmd(open userOpener) *cobra.Command {
|
||||
var password string
|
||||
cmd := &cobra.Command{
|
||||
Use: "reset-password <user-id>",
|
||||
Short: "Set a new password for someone who lost theirs",
|
||||
Long: "There is no self-service reset, because this gateway sends no email and\n" +
|
||||
"a reset link with nowhere to send it is not a reset. Set one here, tell\n" +
|
||||
"them out of band, and have them change it from their profile.\n\n" +
|
||||
"Every session they have open is closed: if they lost the password, the\n" +
|
||||
"reason might be that somebody else has it.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if password == "" {
|
||||
return fmt.Errorf("--password is required")
|
||||
}
|
||||
s, err := open(cmd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
hash, err := account.HashPassword(password)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.SetPassword(cmd.Context(), args[0], hash); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.DeleteUserSessions(cmd.Context(), args[0]); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "password set for %s\n", args[0])
|
||||
cmd.PrintErrln("Their sessions were closed. Their API keys still work.")
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&password, "password", "", "the new password")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newUserCreateCmd(open userOpener) *cobra.Command {
|
||||
var project, password string
|
||||
cmd := &cobra.Command{
|
||||
Use: "create <email>",
|
||||
Short: "Create an account directly, already approved",
|
||||
Long: "For when signup is closed, or for the first account on a new deployment.\n" +
|
||||
"The password is set here and should be changed by its owner.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if password == "" {
|
||||
return fmt.Errorf("--password is required")
|
||||
}
|
||||
s, err := open(cmd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
hash, err := account.HashPassword(password)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
id := randomID("usr")
|
||||
u := store.User{
|
||||
ID: id, Email: args[0], Name: args[0],
|
||||
PasswordHash: hash, State: store.UserPending, CreatedAt: time.Now(),
|
||||
}
|
||||
if err := s.CreateUser(cmd.Context(), u); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.ApproveUser(cmd.Context(), id, project, "cli", time.Now()); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s\n", id)
|
||||
cmd.PrintErrf("%s created and approved into project %s\n", args[0], project)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project the person's keys belong to")
|
||||
cmd.Flags().StringVar(&password, "password", "", "initial password")
|
||||
return cmd
|
||||
}
|
||||
+184
-13
@@ -17,20 +17,54 @@ import (
|
||||
"github.com/knadh/koanf/v2"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
// Server is the whole server configuration.
|
||||
type Server struct {
|
||||
// Listen is the client-facing OpenAI API address.
|
||||
Listen string `koanf:"listen"`
|
||||
// AdminListen serves metrics, health and pprof. Kept separate so the
|
||||
// client-facing port can be exposed without exposing diagnostics.
|
||||
// AdminListen serves the dashboard, the management API, health and pprof.
|
||||
//
|
||||
// Loopback by default, and it should stay there: pprof can dump the process
|
||||
// heap, and API keys are in that heap in plaintext. Every comment beside
|
||||
// this setting has always said not to expose it, while the default bound
|
||||
// every interface — a default that contradicts its own documentation is a
|
||||
// footgun, so the default now matches the advice.
|
||||
AdminListen string `koanf:"admin_listen"`
|
||||
|
||||
// MetricsListen serves Prometheus metrics and nothing else.
|
||||
//
|
||||
// Separate from AdminListen because the reason that one is loopback is
|
||||
// pprof, not metrics — and metrics are the one thing on it that something
|
||||
// else genuinely needs to reach. Sharing a listener meant a fleet could
|
||||
// either be monitored or keep its heap private, which is not a choice
|
||||
// anybody should have to make.
|
||||
//
|
||||
// Empty means off, and metrics remain available on AdminListen either way,
|
||||
// so nothing that scrapes it today has to change.
|
||||
MetricsListen string `koanf:"metrics_listen"`
|
||||
|
||||
// AgentListen is where node agents connect. Separate from the client port
|
||||
// so the fleet-facing surface can be firewalled independently.
|
||||
AgentListen string `koanf:"agent_listen"`
|
||||
|
||||
// Auth turns on API keys on the client-facing port.
|
||||
Auth Auth `koanf:"auth"`
|
||||
|
||||
// Threads keeps conversations and agent memory without API keys, for a
|
||||
// gateway on one trusted machine. With auth on they live in the auth store,
|
||||
// per project, and this block is ignored.
|
||||
Threads Threads `koanf:"threads"`
|
||||
|
||||
// Landing configures the page served at "/".
|
||||
Landing Landing `koanf:"landing"`
|
||||
|
||||
// AgentTLS authenticates the fleet-facing port. It deliberately does not
|
||||
// cover the client-facing one: that surface is fronted by a publicly
|
||||
// trusted certificate, while this one is mutual and private.
|
||||
AgentTLS tlsconf.Config `koanf:"agent_tls"`
|
||||
|
||||
Router Router `koanf:"router"`
|
||||
Control Control `koanf:"control"`
|
||||
Models []catalog.Model `koanf:"models"`
|
||||
@@ -41,6 +75,92 @@ type Server struct {
|
||||
Endpoints []Endpoint `koanf:"endpoints"`
|
||||
}
|
||||
|
||||
// Landing configures the public page at "/".
|
||||
//
|
||||
// It exists for one thing the gateway cannot know: who to ask for a key. A page
|
||||
// that explains the API and then leaves the reader with no route to using it is
|
||||
// where onboarding stops.
|
||||
type Landing struct {
|
||||
// Contact is where someone without a key should go — an address, a URL, a
|
||||
// channel name. Rendered as a link when it is one, as text otherwise.
|
||||
Contact string `koanf:"contact"`
|
||||
// Notice is an optional line of your own, for whatever the reader needs to
|
||||
// know that this page cannot guess.
|
||||
Notice string `koanf:"notice"`
|
||||
}
|
||||
|
||||
// Threads configures conversation and memory storage for a gateway without
|
||||
// authentication.
|
||||
//
|
||||
// Everything stored is in one pool with no owner, readable by anyone who can
|
||||
// reach the client port and knows an id. That is acceptable on one machine and
|
||||
// wrong on a shared network, which is why it takes a setting rather than being
|
||||
// the default.
|
||||
type Threads struct {
|
||||
// Store is the SQLite file. Empty means nothing is kept.
|
||||
Store string `koanf:"store"`
|
||||
// RetainHours is how long idle conversations and notes are kept. Required
|
||||
// with Store: a store with no expiry keeps everything anyone typed, by
|
||||
// omission.
|
||||
RetainHours int `koanf:"retain_hours"`
|
||||
}
|
||||
|
||||
// Auth configures client authentication (plan §7.9).
|
||||
//
|
||||
// Off by default, because a gateway on a trusted network with no keys issued
|
||||
// still has to serve. Turning it on is what makes the client-facing port safe
|
||||
// to expose.
|
||||
type Auth struct {
|
||||
Enabled bool `koanf:"enabled"`
|
||||
// Store is the SQLite file holding keys, tenancy and usage.
|
||||
Store string `koanf:"store"`
|
||||
// CacheTTL bounds how long a revoked key keeps working, traded against a
|
||||
// database read per key per TTL.
|
||||
CacheTTL time.Duration `koanf:"cache_ttl"`
|
||||
// NegativeTTL is how long a rejected key is remembered as rejected, so a
|
||||
// client looping with a bad key cannot become a database read per request.
|
||||
NegativeTTL time.Duration `koanf:"negative_ttl"`
|
||||
|
||||
// DefaultRPM and DefaultTPM apply to a key that no scope gives a limit to.
|
||||
//
|
||||
// Unlimited used to be the default, and a project created without limits is
|
||||
// the easy path — the "default" project on this cluster still is one. The
|
||||
// first sign of that is the bill or the outage. Zero keeps the old
|
||||
// behaviour for anyone who wants it, but it should be a choice rather than
|
||||
// an omission. A fallback, not a ceiling: an explicit limit at any scope
|
||||
// wins, above or below.
|
||||
DefaultRPM int `koanf:"default_rpm"`
|
||||
DefaultTPM int `koanf:"default_tpm"`
|
||||
|
||||
// SignupOpen serves the signup form. Accounts it creates are pending until
|
||||
// an operator approves them, so this is not a way in — but turning it off
|
||||
// means accounts are made with `llmgw users create` instead.
|
||||
SignupOpen bool `koanf:"signup_open"`
|
||||
|
||||
// Pricing is what each model costs, for budgets. A model absent from the
|
||||
// table costs whatever DefaultPrice says, which is nothing unless set.
|
||||
Pricing map[string]Price `koanf:"pricing"`
|
||||
DefaultPrice Price `koanf:"default_price"`
|
||||
|
||||
// FlushInterval and BatchSize decide when usage records are written.
|
||||
FlushInterval time.Duration `koanf:"flush_interval"`
|
||||
BatchSize int `koanf:"batch_size"`
|
||||
// BufferSize bounds the queue. Beyond it records are dropped rather than
|
||||
// blocking a request: losing accounting is bad, turning a stalled database
|
||||
// into an outage is worse.
|
||||
BufferSize int `koanf:"buffer_size"`
|
||||
}
|
||||
|
||||
// Price is a model's cost in micro-units per thousand tokens.
|
||||
//
|
||||
// What a unit means is the operator's choice — self-hosted models have no
|
||||
// per-token price, so it may be GPU-seconds, a notional rate, or 1000
|
||||
// everywhere, which turns a budget into a token quota (plan §7.9.4).
|
||||
type Price struct {
|
||||
PromptPer1k int64 `koanf:"prompt_per_1k"`
|
||||
CompletionPer1k int64 `koanf:"completion_per_1k"`
|
||||
}
|
||||
|
||||
// Control configures the control plane.
|
||||
type Control struct {
|
||||
// Interval is how often desired and actual state are compared.
|
||||
@@ -61,6 +181,15 @@ type Control struct {
|
||||
|
||||
// Router holds routing behaviour.
|
||||
type Router struct {
|
||||
// ActivationWait is how long a request waits for an on-demand model to
|
||||
// load before the client is told to come back. Zero takes the default.
|
||||
//
|
||||
// The one real decision in on-demand loading. Refusing outright makes the
|
||||
// feature useless, since the first request for every model would fail;
|
||||
// waiting without a bound turns a request into a job, and a client that
|
||||
// budgeted thirty seconds ends up holding a connection for ten minutes.
|
||||
ActivationWait time.Duration `koanf:"activation_wait"`
|
||||
|
||||
ScrapeInterval time.Duration `koanf:"scrape_interval"`
|
||||
ScrapeTimeout time.Duration `koanf:"scrape_timeout"`
|
||||
// StaleAfter is how long an endpoint may go unscraped before being scored
|
||||
@@ -78,6 +207,11 @@ type Router struct {
|
||||
|
||||
// DefaultProfile names the routing profile used when a model picks none.
|
||||
DefaultProfile string `koanf:"default_profile"`
|
||||
// Prefix turns on the approximate prefix tree (plan §7.5 tier 1). Off, the
|
||||
// router uses session affinity alone, which is tier 0 and already most of
|
||||
// the win for multi-turn traffic.
|
||||
Prefix Prefix `koanf:"prefix"`
|
||||
|
||||
// SessionHeader carries the session id for affinity (plan §7.5 tier 0).
|
||||
SessionHeader string `koanf:"session_header"`
|
||||
// InjectUsage asks engines for usage even when the client did not, so token
|
||||
@@ -110,6 +244,24 @@ type Admission struct {
|
||||
Interval time.Duration `koanf:"interval"`
|
||||
}
|
||||
|
||||
// Prefix configures approximate prefix routing.
|
||||
type Prefix struct {
|
||||
Enabled bool `koanf:"enabled"`
|
||||
// MaxNodes bounds the tree, MinMatch is the shortest match worth acting on
|
||||
// in bytes, and TTL is how long an entry is believed — an engine evicts on
|
||||
// its own schedule, so an old record may describe a cache that is gone.
|
||||
MaxNodes int `koanf:"max_nodes"`
|
||||
MinMatch int `koanf:"min_match"`
|
||||
TTL time.Duration `koanf:"ttl"`
|
||||
|
||||
// MatchRatio is the share of the prompt that must match to be worth acting
|
||||
// on. BalanceAbs and BalanceRel are the imbalance guard, and both must trip
|
||||
// before affinity is abandoned.
|
||||
MatchRatio float64 `koanf:"match_ratio"`
|
||||
BalanceAbs float64 `koanf:"balance_abs"`
|
||||
BalanceRel float64 `koanf:"balance_rel"`
|
||||
}
|
||||
|
||||
// Band is one priority class.
|
||||
type Band struct {
|
||||
Name string `koanf:"name"`
|
||||
@@ -119,11 +271,18 @@ type Band struct {
|
||||
|
||||
// Timeouts bounds each phase of an upstream request.
|
||||
type Timeouts struct {
|
||||
Dial time.Duration `koanf:"dial"`
|
||||
Dial time.Duration `koanf:"dial"`
|
||||
// ResponseHeader bounds the wait for headers on a streaming request. It is
|
||||
// a liveness check there, because headers arrive before generation starts.
|
||||
ResponseHeader time.Duration `koanf:"response_header"`
|
||||
FirstToken time.Duration `koanf:"first_token"`
|
||||
Idle time.Duration `koanf:"idle"`
|
||||
Total time.Duration `koanf:"total"`
|
||||
// Generation bounds a non-streaming request, where the engine sends no
|
||||
// headers until the completion is done. Setting response_header alone used
|
||||
// to cap both, so 60s read as a liveness guard and behaved as a limit of
|
||||
// about a hundred tokens on a slow model.
|
||||
Generation time.Duration `koanf:"generation"`
|
||||
FirstToken time.Duration `koanf:"first_token"`
|
||||
Idle time.Duration `koanf:"idle"`
|
||||
Total time.Duration `koanf:"total"`
|
||||
}
|
||||
|
||||
// Scoring holds the thresholds the scorers normalize against. These want to be
|
||||
@@ -137,19 +296,22 @@ type Scoring struct {
|
||||
|
||||
// Endpoint is one engine replica. Static in phase 1.
|
||||
type Endpoint struct {
|
||||
ID string `koanf:"id"`
|
||||
URL string `koanf:"url"`
|
||||
Models []string `koanf:"models"`
|
||||
Pool string `koanf:"pool"`
|
||||
Weight float64 `koanf:"weight"`
|
||||
Priority int `koanf:"priority"`
|
||||
ID string `koanf:"id"`
|
||||
URL string `koanf:"url"`
|
||||
Models []string `koanf:"models"`
|
||||
// Engine names the inference engine, which decides what its metrics are
|
||||
// called. Empty means vLLM's names.
|
||||
Engine string `koanf:"engine"`
|
||||
Pool string `koanf:"pool"`
|
||||
Weight float64 `koanf:"weight"`
|
||||
Priority int `koanf:"priority"`
|
||||
}
|
||||
|
||||
// Default returns the configuration used when nothing is specified.
|
||||
func Default() Server {
|
||||
return Server{
|
||||
Listen: ":8080",
|
||||
AdminListen: ":9090",
|
||||
AdminListen: "127.0.0.1:9090",
|
||||
AgentListen: ":7443",
|
||||
Control: Control{
|
||||
Interval: 5 * time.Second,
|
||||
@@ -244,6 +406,15 @@ func (s Server) Validate() error {
|
||||
if s.Listen == "" {
|
||||
return fmt.Errorf("listen is empty")
|
||||
}
|
||||
if err := s.AgentTLS.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if s.Auth.Enabled && s.Auth.Store == "" {
|
||||
return fmt.Errorf("auth.store is required when auth is enabled")
|
||||
}
|
||||
if s.Threads.Store != "" && s.Threads.RetainHours <= 0 {
|
||||
return fmt.Errorf("threads.retain_hours must be positive when threads.store is set")
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{}, len(s.Endpoints))
|
||||
for i, e := range s.Endpoints {
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestTheDocumentedAgentExampleLoads parses the file shipped in deploy/examples.
|
||||
//
|
||||
// It exists because the example was wrong for the entire life of the project
|
||||
// and nothing noticed. The fabric struct carried no koanf tags, so the `name:`
|
||||
// key the example documents bound to nothing, and an agent configured exactly
|
||||
// as instructed refused to start with "fabrics[0]: name is required" — naming
|
||||
// the key the operator had just written.
|
||||
//
|
||||
// Every fabric until then had been configured in a Go test, where the tags do
|
||||
// not matter. A published example that has never been parsed is a guess.
|
||||
func TestTheDocumentedAgentExampleLoads(t *testing.T) {
|
||||
path := filepath.Join("..", "..", "deploy", "examples", "agent.yaml")
|
||||
if _, err := os.Stat(path); err != nil {
|
||||
t.Skipf("no example to check: %v", err)
|
||||
}
|
||||
|
||||
cfg, err := LoadAgent(path)
|
||||
if err != nil {
|
||||
t.Fatalf("the documented example does not load: %v", err)
|
||||
}
|
||||
if len(cfg.Fabrics) == 0 {
|
||||
t.Fatal("the example documents a fabric and none was parsed")
|
||||
}
|
||||
f := cfg.Fabrics[0]
|
||||
if f.Fabric == "" {
|
||||
t.Error("the fabric name did not bind")
|
||||
}
|
||||
if len(f.HCAs) == 0 {
|
||||
t.Error("the fabric's adapters did not bind")
|
||||
}
|
||||
if f.RateGbps == 0 {
|
||||
t.Error("the fabric's rate did not bind")
|
||||
}
|
||||
if cfg.NodeID == "" || cfg.Server == "" {
|
||||
t.Errorf("node_id=%q server=%q", cfg.NodeID, cfg.Server)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheDocumentedServerExampleLoads(t *testing.T) {
|
||||
path := filepath.Join("..", "..", "deploy", "examples", "server.yaml")
|
||||
if _, err := os.Stat(path); err != nil {
|
||||
t.Skipf("no example to check: %v", err)
|
||||
}
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatalf("the documented example does not load: %v", err)
|
||||
}
|
||||
if err := cfg.Validate(); err != nil {
|
||||
t.Fatalf("the documented example does not validate: %v", err)
|
||||
}
|
||||
if cfg.Listen == "" || len(cfg.Models) == 0 {
|
||||
t.Errorf("listen=%q models=%d", cfg.Listen, len(cfg.Models))
|
||||
}
|
||||
}
|
||||
@@ -7,19 +7,23 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// API serves the management endpoints.
|
||||
type API struct {
|
||||
ctl *control.Controller
|
||||
models Models
|
||||
ctl *control.Controller
|
||||
}
|
||||
|
||||
// New builds the API over a controller.
|
||||
@@ -27,9 +31,19 @@ func New(ctl *control.Controller) *API { return &API{ctl: ctl} }
|
||||
|
||||
// Register mounts the endpoints on a mux under /api/v1.
|
||||
func (a *API) Register(mux *http.ServeMux) {
|
||||
mux.HandleFunc("GET /api/v1/summary", a.summary)
|
||||
// Adding a model is a write, and the only one this API has. It lives here
|
||||
// rather than on the client-facing listener for the same reason the rest of
|
||||
// it does: reaching this port already means holding the control plane.
|
||||
mux.HandleFunc("POST /api/v1/models", a.addModel)
|
||||
mux.HandleFunc("DELETE /api/v1/models/{name}", a.removeModel)
|
||||
mux.HandleFunc("GET /api/v1/nodes", a.listNodes)
|
||||
mux.HandleFunc("GET /api/v1/models", a.listModels)
|
||||
mux.HandleFunc("GET /api/v1/replicas", a.listReplicas)
|
||||
// Listing is a read and deleting is a separate verb, so a mistyped
|
||||
// command cannot remove anything.
|
||||
mux.HandleFunc("GET /api/v1/weights/orphans", a.listOrphans)
|
||||
mux.HandleFunc("POST /api/v1/weights/collect", a.collectWeights)
|
||||
mux.HandleFunc("GET /api/v1/explain-placement", a.explainPlacement)
|
||||
mux.HandleFunc("POST /api/v1/nodes/{id}/state", a.setNodeState)
|
||||
}
|
||||
@@ -44,10 +58,29 @@ type NodeView struct {
|
||||
State string `json:"state"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
|
||||
GPUTotal int `json:"gpu_total"`
|
||||
GPUFree int `json:"gpu_free"`
|
||||
Classes []string `json:"gpu_classes,omitempty"`
|
||||
Fabrics []string `json:"fabrics,omitempty"`
|
||||
GPUTotal int `json:"gpu_total"`
|
||||
GPUFree int `json:"gpu_free"`
|
||||
// MemoryGB is the machine's RAM, 0 when the agent did not report it. It is
|
||||
// what decides whether a CPU profile can be placed here, so it belongs
|
||||
// where someone looks when a model will not schedule.
|
||||
MemoryGB int `json:"memory_gb,omitempty"`
|
||||
// Models are the weights this node can see. Absent means it was not asked
|
||||
// to look, which is not the same as holding none.
|
||||
Models []string `json:"models_on_disk,omitempty"`
|
||||
// DiskFreeGB and DiskTotalGB are the filesystems holding those weights,
|
||||
// summed. Zero means the agent did not report any, which an older agent
|
||||
// and a node with no model directories both do — so it reads as unknown
|
||||
// rather than as a full disk.
|
||||
//
|
||||
// Summed because an operator asking "can this node take another model"
|
||||
// wants one number; the per-filesystem detail is in the node's own record
|
||||
// for anyone who needs it. Shared filesystems are left out of the sum: the
|
||||
// space on a filer is not this node's to spend, and adding it would let a
|
||||
// fleet of twenty nodes each believe it has the same free terabyte.
|
||||
DiskFreeGB int `json:"disk_free_gb,omitempty"`
|
||||
DiskTotalGB int `json:"disk_total_gb,omitempty"`
|
||||
Classes []string `json:"gpu_classes,omitempty"`
|
||||
Fabrics []string `json:"fabrics,omitempty"`
|
||||
|
||||
// Unhealthy lists GPUs kept out of placement and why, which is usually the
|
||||
// answer when a node looks idle but nothing will schedule onto it.
|
||||
@@ -73,16 +106,21 @@ type ModelView struct {
|
||||
|
||||
// ReplicaView is a replica as `ctl replicas` shows it.
|
||||
type ReplicaView struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Profile string `json:"profile"`
|
||||
State string `json:"state"`
|
||||
Pool string `json:"pool,omitempty"`
|
||||
Fabric string `json:"fabric,omitempty"`
|
||||
Endpoint string `json:"endpoint,omitempty"`
|
||||
AgeSec int64 `json:"age_seconds"`
|
||||
Err string `json:"error,omitempty"`
|
||||
Shards []ShardView `json:"shards"`
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Profile string `json:"profile"`
|
||||
State string `json:"state"`
|
||||
Pool string `json:"pool,omitempty"`
|
||||
Fabric string `json:"fabric,omitempty"`
|
||||
Endpoint string `json:"endpoint,omitempty"`
|
||||
AgeSec int64 `json:"age_seconds"`
|
||||
// IdleSec is how long since a request was routed here. Negative means the
|
||||
// router has not reported any use of this replica, which is not the same
|
||||
// as idle — it is what a replica looks like before anything has asked for
|
||||
// it, and after a router restart.
|
||||
IdleSec int64 `json:"idle_seconds"`
|
||||
Err string `json:"error,omitempty"`
|
||||
Shards []ShardView `json:"shards"`
|
||||
}
|
||||
|
||||
// ShardView is one shard of a replica.
|
||||
@@ -117,6 +155,15 @@ type RejectionView struct {
|
||||
// --- handlers ---
|
||||
|
||||
func (a *API) listNodes(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, a.Nodes())
|
||||
}
|
||||
|
||||
// Nodes builds the node view.
|
||||
//
|
||||
// Exported because the dashboard renders the same data server-side. Going back
|
||||
// out through HTTP to fetch what is already in this process would add a way for
|
||||
// the two to disagree.
|
||||
func (a *API) Nodes() []NodeView {
|
||||
replicasOn := map[string]int{}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
for _, n := range rep.NodeIDs() {
|
||||
@@ -131,9 +178,17 @@ func (a *API) listNodes(w http.ResponseWriter, _ *http.Request) {
|
||||
ID: n.ID, Hostname: n.Hostname, MgmtAddr: n.MgmtAddr,
|
||||
State: n.State.String(), Labels: n.Labels,
|
||||
GPUTotal: len(n.GPUs), GPUFree: len(inv.FreeGPUs(n.ID)),
|
||||
Fabrics: n.FabricNames(), AgentVersion: n.AgentVersion,
|
||||
MemoryGB: n.MemoryTotalMB / 1024,
|
||||
Fabrics: n.FabricNames(), AgentVersion: n.AgentVersion,
|
||||
Replicas: replicasOn[n.ID],
|
||||
}
|
||||
if free, total := n.LocalDisk(); total > 0 {
|
||||
v.DiskFreeGB = int(free / (1 << 30))
|
||||
v.DiskTotalGB = int(total / (1 << 30))
|
||||
}
|
||||
for _, w := range n.Weights {
|
||||
v.Models = append(v.Models, w.Path)
|
||||
}
|
||||
seen := map[string]struct{}{}
|
||||
for _, g := range n.GPUs {
|
||||
if _, dup := seen[g.Class()]; !dup {
|
||||
@@ -150,10 +205,15 @@ func (a *API) listNodes(w http.ResponseWriter, _ *http.Request) {
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
return out
|
||||
}
|
||||
|
||||
func (a *API) listModels(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, a.Models())
|
||||
}
|
||||
|
||||
// Models builds the catalog view.
|
||||
func (a *API) Models() []ModelView {
|
||||
ready, total := map[string]int{}, map[string]int{}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
total[rep.Model]++
|
||||
@@ -175,12 +235,15 @@ func (a *API) listModels(w http.ResponseWriter, _ *http.Request) {
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
return out
|
||||
}
|
||||
|
||||
func (a *API) listReplicas(w http.ResponseWriter, r *http.Request) {
|
||||
model := r.URL.Query().Get("model")
|
||||
writeJSON(w, http.StatusOK, a.Replicas(r.URL.Query().Get("model")))
|
||||
}
|
||||
|
||||
// Replicas builds the replica view, optionally filtered to one model.
|
||||
func (a *API) Replicas(model string) []ReplicaView {
|
||||
out := make([]ReplicaView, 0)
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
if model != "" && rep.Model != model {
|
||||
@@ -190,7 +253,13 @@ func (a *API) listReplicas(w http.ResponseWriter, r *http.Request) {
|
||||
ID: rep.ID, Model: rep.Model, Profile: rep.Profile,
|
||||
State: rep.State.String(), Pool: rep.Pool, Fabric: rep.Fabric,
|
||||
Endpoint: rep.Endpoint, Err: rep.Err,
|
||||
AgeSec: int64(time.Since(rep.CreatedAt).Seconds()),
|
||||
AgeSec: int64(time.Since(rep.CreatedAt).Seconds()),
|
||||
IdleSec: -1,
|
||||
}
|
||||
if rep.LastUsedAt.IsZero() {
|
||||
v.IdleSec = -1
|
||||
} else {
|
||||
v.IdleSec = int64(time.Since(rep.LastUsedAt).Seconds())
|
||||
}
|
||||
for _, s := range rep.Shards {
|
||||
v.Shards = append(v.Shards, ShardView{
|
||||
@@ -200,7 +269,57 @@ func (a *API) listReplicas(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
return out
|
||||
}
|
||||
|
||||
// Models is the service that owns catalog changes, when there is one.
|
||||
//
|
||||
// Optional because a control plane with no store cannot persist a model, and
|
||||
// pretending otherwise would accept a model that vanishes on restart.
|
||||
type Models interface {
|
||||
Add(ctx context.Context, m catalog.Model, by string) error
|
||||
Remove(ctx context.Context, name string) error
|
||||
}
|
||||
|
||||
// SetModels supplies the catalog-changing service.
|
||||
func (a *API) SetModels(m Models) { a.models = m }
|
||||
|
||||
func (a *API) addModel(w http.ResponseWriter, r *http.Request) {
|
||||
if a.models == nil {
|
||||
writeError(w, http.StatusNotImplemented,
|
||||
"this control plane has no store, so a model added here would not survive a restart")
|
||||
return
|
||||
}
|
||||
var m catalog.Model
|
||||
if err := json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&m); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "the model definition is not valid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := a.models.Add(r.Context(), m, r.Header.Get("X-Llmgw-Actor")); err != nil {
|
||||
// Every failure here is the definition being wrong, which is the
|
||||
// caller's to fix — a collision, a missing fallback, a profile that
|
||||
// cannot be placed. None of them are server faults.
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]string{"added": m.Name})
|
||||
}
|
||||
|
||||
func (a *API) removeModel(w http.ResponseWriter, r *http.Request) {
|
||||
if a.models == nil {
|
||||
writeError(w, http.StatusNotImplemented, "this control plane has no store")
|
||||
return
|
||||
}
|
||||
name := r.PathValue("name")
|
||||
if err := a.models.Remove(r.Context(), name); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]string{"removed": name})
|
||||
}
|
||||
|
||||
func (a *API) summary(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, a.Summary())
|
||||
}
|
||||
|
||||
func (a *API) explainPlacement(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -214,6 +333,15 @@ func (a *API) explainPlacement(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusOK, a.Explain(model))
|
||||
}
|
||||
|
||||
// Explain builds the placement explanation for one model.
|
||||
//
|
||||
// It does not check that the model exists; the handler above does that so it
|
||||
// can answer 404, while the dashboard only ever asks about models it just read
|
||||
// out of the catalog.
|
||||
func (a *API) Explain(model string) ExplainView {
|
||||
v := ExplainView{Model: model}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
if rep.Model == model && rep.State.Counts() {
|
||||
@@ -228,7 +356,7 @@ func (a *API) explainPlacement(w http.ResponseWriter, r *http.Request) {
|
||||
})
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, v)
|
||||
return v
|
||||
}
|
||||
|
||||
// setNodeState implements cordon, uncordon and drain.
|
||||
|
||||
@@ -3,6 +3,7 @@ package adminapi
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -35,7 +36,12 @@ func setup(t *testing.T, models []catalog.Model, nodes ...inventory.Node) (*API,
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctl := control.New(control.DefaultConfig(), cat, nullDispatcher{}, nullSink{}, slog.New(slog.DiscardHandler))
|
||||
cfg := control.DefaultConfig()
|
||||
// No startup settling window: these tests build the whole fleet before
|
||||
// anything runs, which is the state settling waits for. Leaving it on would
|
||||
// make every test here assert that placement had not happened yet.
|
||||
cfg.Reconciler.SettleFor = -1
|
||||
ctl := control.New(cfg, catalog.NewLive(cat), nullDispatcher{}, nullSink{}, slog.New(slog.DiscardHandler))
|
||||
for _, n := range nodes {
|
||||
ctl.Inventory().Upsert(n)
|
||||
}
|
||||
@@ -241,3 +247,109 @@ func TestEmptyClusterReturnsArraysNotNull(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type fakeModels struct {
|
||||
added []catalog.Model
|
||||
actors []string
|
||||
removed []string
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeModels) Add(_ context.Context, m catalog.Model, by string) error {
|
||||
if f.err != nil {
|
||||
return f.err
|
||||
}
|
||||
f.added = append(f.added, m)
|
||||
f.actors = append(f.actors, by)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeModels) Remove(_ context.Context, name string) error {
|
||||
if f.err != nil {
|
||||
return f.err
|
||||
}
|
||||
f.removed = append(f.removed, name)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAddingAModelOverTheAPI(t *testing.T) {
|
||||
a, _ := setup(t, testModels())
|
||||
fm := &fakeModels{}
|
||||
a.SetModels(fm)
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
|
||||
body := `{"name":"new","source":"/models/new","profiles":[{"name":"p","engine":"llamacpp"}]}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/models", strings.NewReader(body))
|
||||
req.Header.Set("X-Llmgw-Actor", "ilker")
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if len(fm.added) != 1 || fm.added[0].Name != "new" {
|
||||
t.Fatalf("added = %+v", fm.added)
|
||||
}
|
||||
// Who made the change, kept so a catalog that grew over months can say
|
||||
// where each entry came from.
|
||||
if len(fm.actors) != 1 || fm.actors[0] != "ilker" {
|
||||
t.Errorf("actor = %v", fm.actors)
|
||||
}
|
||||
}
|
||||
|
||||
func TestABadModelDefinitionIsTheCallersFault(t *testing.T) {
|
||||
// Every failure here is the definition being wrong — a collision, a missing
|
||||
// fallback, a profile that cannot be placed. None are server faults, and
|
||||
// returning 500 for them would send someone looking at the wrong logs.
|
||||
a, _ := setup(t, testModels())
|
||||
a.SetModels(&fakeModels{err: errors.New("model \"x\" collides")})
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/models", strings.NewReader(`{"name":"x"}`))
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Errorf("status %d, expected 400", rec.Code)
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "collides") {
|
||||
t.Errorf("the reason was lost: %s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithoutAStoreTheAPIRefusesRatherThanForgets(t *testing.T) {
|
||||
// Accepting a model that vanishes on the next restart is worse than saying
|
||||
// it cannot be kept.
|
||||
a, _ := setup(t, testModels()) // SetModels never called
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/models", strings.NewReader(`{"name":"x"}`))
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusNotImplemented {
|
||||
t.Errorf("status %d, expected 501", rec.Code)
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "restart") {
|
||||
t.Errorf("the reason does not explain the consequence: %s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemovingAModelOverTheAPI(t *testing.T) {
|
||||
a, _ := setup(t, testModels())
|
||||
fm := &fakeModels{}
|
||||
a.SetModels(fm)
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
|
||||
req := httptest.NewRequest(http.MethodDelete, "/api/v1/models/gone", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if len(fm.removed) != 1 || fm.removed[0] != "gone" {
|
||||
t.Errorf("removed = %v", fm.removed)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Summary is the whole fleet in one answer.
|
||||
//
|
||||
// A per-node listing is the right tool for four nodes and the wrong one for two
|
||||
// hundred: the questions that matter at that size — how full am I, what is
|
||||
// broken, what is holding cards it is not using — are aggregates, and reading
|
||||
// them off a long table is a chore that gets skipped. Everything here is
|
||||
// derived from the same state the tables show, so it can never disagree with
|
||||
// them; it is a different question asked of the same facts.
|
||||
type Summary struct {
|
||||
Nodes NodeTotals `json:"nodes"`
|
||||
GPUs []ClassUse `json:"gpus"`
|
||||
// Fabrics matter because a multi-node replica may never span two of them,
|
||||
// so free capacity in the wrong fabric is not capacity for the model that
|
||||
// needs it.
|
||||
Fabrics []FabricUse `json:"fabrics,omitempty"`
|
||||
Replicas ReplicaTotal `json:"replicas"`
|
||||
Models ModelTotals `json:"models"`
|
||||
// Unplaceable names models the scheduler currently cannot satisfy. This is
|
||||
// the line that turns "something is wrong" into a place to look.
|
||||
Unplaceable []string `json:"unplaceable,omitempty"`
|
||||
// IdleReplicas are serving replicas nothing has asked for in a while, which
|
||||
// on a fleet with more models than GPUs is the first place to find room.
|
||||
IdleReplicas []IdleReplica `json:"idle_replicas,omitempty"`
|
||||
}
|
||||
|
||||
// NodeTotals counts nodes by the states that need different responses.
|
||||
type NodeTotals struct {
|
||||
Total int `json:"total"`
|
||||
Ready int `json:"ready"`
|
||||
Cordoned int `json:"cordoned"`
|
||||
Draining int `json:"draining"`
|
||||
Unreachable int `json:"unreachable"`
|
||||
}
|
||||
|
||||
// ClassUse is how much of one kind of card is in use.
|
||||
type ClassUse struct {
|
||||
Class string `json:"class"`
|
||||
Total int `json:"total"`
|
||||
Used int `json:"used"`
|
||||
// Unhealthy GPUs are neither used nor available, and counting them as free
|
||||
// is how a cluster looks like it has room it does not have.
|
||||
Unhealthy int `json:"unhealthy"`
|
||||
}
|
||||
|
||||
// Free is what can still be allocated.
|
||||
func (c ClassUse) Free() int { return c.Total - c.Used - c.Unhealthy }
|
||||
|
||||
// FabricUse is capacity within one fabric.
|
||||
type FabricUse struct {
|
||||
Fabric string `json:"fabric"`
|
||||
Nodes int `json:"nodes"`
|
||||
// ModelsPresent is how many nodes here hold weights at all, and Models how
|
||||
// many distinct models they hold between them.
|
||||
//
|
||||
// Per fabric because that is the unit at which the answer is actionable.
|
||||
// Fabrics are not connected to one another, so weights in one are not
|
||||
// weights in another: a model seeded into a single fabric leaves the rest
|
||||
// of the cluster unable to serve it without first moving tens of gigabytes
|
||||
// over the management network, and nothing said so.
|
||||
NodesWithWeights int `json:"nodes_with_weights"`
|
||||
Models int `json:"models_present"`
|
||||
// LargestFreeBlock is the most nodes on this fabric that are entirely free,
|
||||
// which is the number that decides whether a multi-node replica can land.
|
||||
LargestFreeBlock int `json:"largest_free_block"`
|
||||
}
|
||||
|
||||
// ReplicaTotal counts replicas by state.
|
||||
type ReplicaTotal struct {
|
||||
Ready int `json:"ready"`
|
||||
Starting int `json:"starting"`
|
||||
Draining int `json:"draining"`
|
||||
Failed int `json:"failed"`
|
||||
}
|
||||
|
||||
// ModelTotals counts models by whether they have what they asked for.
|
||||
type ModelTotals struct {
|
||||
Total int `json:"total"`
|
||||
Satisfied int `json:"satisfied"`
|
||||
Short int `json:"short"`
|
||||
}
|
||||
|
||||
// IdleReplica is a replica holding hardware nothing is asking for.
|
||||
type IdleReplica struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
IdleSec int64 `json:"idle_seconds"`
|
||||
GPUs int `json:"gpus"`
|
||||
}
|
||||
|
||||
// idleThreshold is how long a replica must go unasked-for to be worth naming.
|
||||
//
|
||||
// Minutes rather than seconds, because a replica that flaps in and out of this
|
||||
// list is noise, and because unloading one throws away its KV cache — the thing
|
||||
// prefix routing spent its effort filling.
|
||||
const idleThreshold = 10 * time.Minute
|
||||
|
||||
// Summary builds the fleet summary.
|
||||
func (a *API) Summary() Summary {
|
||||
var s Summary
|
||||
|
||||
nodes := a.Nodes()
|
||||
byClass := map[string]*ClassUse{}
|
||||
s.Nodes.Total = len(nodes)
|
||||
for _, n := range nodes {
|
||||
switch n.State {
|
||||
case "ready":
|
||||
s.Nodes.Ready++
|
||||
case "cordoned":
|
||||
s.Nodes.Cordoned++
|
||||
case "draining":
|
||||
s.Nodes.Draining++
|
||||
default:
|
||||
s.Nodes.Unreachable++
|
||||
}
|
||||
for _, c := range n.Classes {
|
||||
if byClass[c] == nil {
|
||||
byClass[c] = &ClassUse{Class: c}
|
||||
}
|
||||
}
|
||||
// A node reports one class in practice; attribute its cards to it.
|
||||
if len(n.Classes) == 1 {
|
||||
u := byClass[n.Classes[0]]
|
||||
u.Total += n.GPUTotal
|
||||
u.Unhealthy += len(n.Unhealthy)
|
||||
u.Used += n.GPUTotal - n.GPUFree - len(n.Unhealthy)
|
||||
}
|
||||
}
|
||||
for _, u := range byClass {
|
||||
s.GPUs = append(s.GPUs, *u)
|
||||
}
|
||||
sort.Slice(s.GPUs, func(i, j int) bool { return s.GPUs[i].Class < s.GPUs[j].Class })
|
||||
|
||||
byFabric := map[string]*FabricUse{}
|
||||
for _, n := range nodes {
|
||||
for _, f := range n.Fabrics {
|
||||
if byFabric[f] == nil {
|
||||
byFabric[f] = &FabricUse{Fabric: f}
|
||||
}
|
||||
byFabric[f].Nodes++
|
||||
if len(n.Models) > 0 {
|
||||
byFabric[f].NodesWithWeights++
|
||||
}
|
||||
// Entirely free, which is what a multi-node replica needs: a node
|
||||
// with one card in use cannot contribute a full-width shard.
|
||||
if n.State == "ready" && n.GPUFree == n.GPUTotal && n.GPUTotal > 0 {
|
||||
byFabric[f].LargestFreeBlock++
|
||||
}
|
||||
}
|
||||
}
|
||||
// Distinct models per fabric, counted across its nodes: one copy on a
|
||||
// shared filer and twenty local copies are the same answer to "can this
|
||||
// fabric serve it".
|
||||
perFabricModels := map[string]map[string]struct{}{}
|
||||
for _, n := range nodes {
|
||||
for _, f := range n.Fabrics {
|
||||
if perFabricModels[f] == nil {
|
||||
perFabricModels[f] = map[string]struct{}{}
|
||||
}
|
||||
for _, m := range n.Models {
|
||||
perFabricModels[f][m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
for name, f := range byFabric {
|
||||
f.Models = len(perFabricModels[name])
|
||||
s.Fabrics = append(s.Fabrics, *f)
|
||||
}
|
||||
sort.Slice(s.Fabrics, func(i, j int) bool { return s.Fabrics[i].Fabric < s.Fabrics[j].Fabric })
|
||||
|
||||
for _, r := range a.Replicas("") {
|
||||
switch r.State {
|
||||
case "ready":
|
||||
s.Replicas.Ready++
|
||||
case "draining":
|
||||
s.Replicas.Draining++
|
||||
case "failed":
|
||||
s.Replicas.Failed++
|
||||
default:
|
||||
s.Replicas.Starting++
|
||||
}
|
||||
if r.State == "ready" && r.IdleSec >= int64(idleThreshold.Seconds()) {
|
||||
gpus := 0
|
||||
for _, sh := range r.Shards {
|
||||
gpus += len(sh.GPUs)
|
||||
}
|
||||
s.IdleReplicas = append(s.IdleReplicas, IdleReplica{
|
||||
ID: r.ID, Model: r.Model, IdleSec: r.IdleSec, GPUs: gpus,
|
||||
})
|
||||
}
|
||||
}
|
||||
sort.Slice(s.IdleReplicas, func(i, j int) bool {
|
||||
return s.IdleReplicas[i].IdleSec > s.IdleReplicas[j].IdleSec
|
||||
})
|
||||
|
||||
for _, m := range a.Models() {
|
||||
s.Models.Total++
|
||||
if m.Ready >= m.Min {
|
||||
s.Models.Satisfied++
|
||||
continue
|
||||
}
|
||||
s.Models.Short++
|
||||
// Short of replicas and unable to place more is the actionable case;
|
||||
// short and still starting is just a cluster doing its job.
|
||||
if m.Total <= m.Ready {
|
||||
s.Unplaceable = append(s.Unplaceable, m.Name)
|
||||
}
|
||||
}
|
||||
sort.Strings(s.Unplaceable)
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
func TestSummaryCountsCapacityByClass(t *testing.T) {
|
||||
// The question a per-node table answers badly at two hundred nodes: how
|
||||
// full am I.
|
||||
a, _ := setup(t, testModels(), gpuNode("n1", 8), gpuNode("n2", 8))
|
||||
s := a.Summary()
|
||||
|
||||
if s.Nodes.Total != 2 || s.Nodes.Ready != 2 {
|
||||
t.Errorf("nodes = %+v", s.Nodes)
|
||||
}
|
||||
if len(s.GPUs) != 1 {
|
||||
t.Fatalf("expected one GPU class, got %+v", s.GPUs)
|
||||
}
|
||||
if g := s.GPUs[0]; g.Total != 16 || g.Free() != 16 {
|
||||
t.Errorf("class totals wrong: %+v", g)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnhealthyGPUsAreNotCountedAsFree(t *testing.T) {
|
||||
// Counting them free is how a cluster looks like it has room it does not,
|
||||
// which is worse than showing less capacity than exists.
|
||||
n := gpuNode("n1", 4)
|
||||
n.GPUs[0].Health.Problems = []string{"3 uncorrected ECC errors"}
|
||||
a, _ := setup(t, testModels(), n)
|
||||
|
||||
g := a.Summary().GPUs[0]
|
||||
if g.Unhealthy != 1 {
|
||||
t.Errorf("unhealthy = %d", g.Unhealthy)
|
||||
}
|
||||
if g.Free() != 3 {
|
||||
t.Errorf("free = %d; a faulty card was counted as available", g.Free())
|
||||
}
|
||||
}
|
||||
|
||||
func TestFabricFreeBlocksCountWholeNodesOnly(t *testing.T) {
|
||||
// A multi-node replica may never span fabrics, and a node with one card in
|
||||
// use cannot contribute a full-width shard. Free cards scattered across
|
||||
// partly-used nodes are not capacity for the model that needs them.
|
||||
free := gpuNode("free", 8, "ib-a")
|
||||
partial := gpuNode("partial", 8, "ib-a")
|
||||
a, ctl := setup(t, testModels(), free, partial)
|
||||
// Take one GPU on the second node.
|
||||
ctl.Inventory().Allocate("partial", inventory.Allocation{
|
||||
ShardID: "s", ReplicaID: "r", GPUIndices: []int{0},
|
||||
})
|
||||
|
||||
s := a.Summary()
|
||||
if len(s.Fabrics) != 1 {
|
||||
t.Fatalf("fabrics = %+v", s.Fabrics)
|
||||
}
|
||||
if got := s.Fabrics[0].LargestFreeBlock; got != 1 {
|
||||
t.Errorf("fully free nodes = %d, want 1: a partly used node is not one", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodeWithNoGPUsDoesNotInventAClass(t *testing.T) {
|
||||
// The CPU-only node this fleet acquired. It has no class, and a summary
|
||||
// that invents one for it would report capacity nobody has.
|
||||
a, _ := setup(t, testModels(), inventory.Node{ID: "cpu1", State: inventory.Ready})
|
||||
s := a.Summary()
|
||||
for _, g := range s.GPUs {
|
||||
if g.Class == "" {
|
||||
t.Errorf("an empty class was reported: %+v", g)
|
||||
}
|
||||
}
|
||||
if s.Nodes.Total != 1 {
|
||||
t.Errorf("the node was not counted: %+v", s.Nodes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdleReplicasNeedRealIdleness(t *testing.T) {
|
||||
// A replica the router has said nothing about carries -1, and must not be
|
||||
// read as idle since the epoch — otherwise a router restart nominates the
|
||||
// whole fleet for eviction at once.
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
idle int64
|
||||
want bool
|
||||
}{
|
||||
{"never reported", -1, false},
|
||||
{"just used", 5, false},
|
||||
{"idle for an hour", int64(time.Hour.Seconds()), true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := tc.idle >= int64(idleThreshold.Seconds())
|
||||
if got != tc.want {
|
||||
t.Errorf("idle=%d counted as idle=%v, want %v", tc.idle, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummaryNamesModelsThatCannotBePlaced(t *testing.T) {
|
||||
// The line that turns "something is wrong" into a place to look.
|
||||
a, _ := setup(t, testModels()) // no nodes at all
|
||||
s := a.Summary()
|
||||
if len(s.Unplaceable) == 0 {
|
||||
t.Fatal("a model with no hardware to run on was not named")
|
||||
}
|
||||
if s.Models.Short == 0 {
|
||||
t.Error("no model was counted short")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWeightsAreCountedPerFabric(t *testing.T) {
|
||||
// Fabrics are not connected, so weights in one are not weights in another.
|
||||
// A model seeded into a single fabric leaves the rest of the cluster unable
|
||||
// to serve it without first moving tens of gigabytes over the management
|
||||
// network, and nothing said so.
|
||||
a1 := gpuNode("a1", 4, "ib-a")
|
||||
a1.Weights = []discovery.LocalWeights{{Path: "/models/llama.gguf"}}
|
||||
a2 := gpuNode("a2", 4, "ib-a")
|
||||
a2.Weights = []discovery.LocalWeights{{Path: "/models/llama.gguf"}}
|
||||
b1 := gpuNode("b1", 4, "ib-b") // the same fabric capacity, none of the weights
|
||||
|
||||
api, _ := setup(t, testModels(), a1, a2, b1)
|
||||
s := api.Summary()
|
||||
|
||||
byName := map[string]FabricUse{}
|
||||
for _, f := range s.Fabrics {
|
||||
byName[f.Fabric] = f
|
||||
}
|
||||
if got := byName["ib-a"]; got.Models != 1 || got.NodesWithWeights != 2 {
|
||||
t.Errorf("ib-a = %+v, want one model on two nodes", got)
|
||||
}
|
||||
if got := byName["ib-b"]; got.Models != 0 || got.NodesWithWeights != 0 {
|
||||
t.Errorf("ib-b = %+v; it holds nothing and should say so", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneSharedCopyCountsAsMuchAsTwenty(t *testing.T) {
|
||||
// A filer mounted on every node and a local copy on every node are the same
|
||||
// answer to "can this fabric serve it", which is the question being asked.
|
||||
var nodes []inventory.Node
|
||||
for _, id := range []string{"n1", "n2", "n3"} {
|
||||
n := gpuNode(id, 4, "ib-a")
|
||||
n.Weights = []discovery.LocalWeights{{Path: "/mnt/shared/llama.gguf", Shared: true}}
|
||||
nodes = append(nodes, n)
|
||||
}
|
||||
api, _ := setup(t, testModels(), nodes...)
|
||||
|
||||
f := api.Summary().Fabrics[0]
|
||||
if f.Models != 1 {
|
||||
t.Errorf("counted %d models; three nodes seeing one file is one model", f.Models)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"sort"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
)
|
||||
|
||||
// OrphanView is cached weights on one node that no catalog model accounts for.
|
||||
type OrphanView struct {
|
||||
Node string `json:"node"`
|
||||
Paths []string `json:"paths"`
|
||||
// Bytes is what deleting them would free, as far as the node reported. A
|
||||
// file whose size was not reported counts as zero rather than being left
|
||||
// out, so the list and the number always describe the same set.
|
||||
Bytes int64 `json:"bytes"`
|
||||
}
|
||||
|
||||
// Orphans finds cached weights nothing in the catalog points at.
|
||||
//
|
||||
// Computed here from what nodes already report, so listing costs no round trip
|
||||
// and a node that is offline simply contributes nothing — its files are not
|
||||
// deletable while it is gone anyway.
|
||||
//
|
||||
// Only files llmgw fetched are ever considered. Everything else in a model
|
||||
// directory was put there by the operator and is none of this system's
|
||||
// business, however unreferenced it looks.
|
||||
func (a *API) Orphans() []OrphanView {
|
||||
var wanted []string
|
||||
for _, m := range a.ctl.Catalog().Models() {
|
||||
wanted = append(wanted, m.Source)
|
||||
}
|
||||
|
||||
out := make([]OrphanView, 0)
|
||||
for _, n := range a.ctl.Inventory().Nodes() {
|
||||
size := map[string]int64{}
|
||||
have := make([]string, 0, len(n.Weights))
|
||||
for _, w := range n.Weights {
|
||||
have = append(have, w.Path)
|
||||
size[w.Path] = w.SizeBytes
|
||||
}
|
||||
|
||||
paths := weights.Orphans(have, wanted)
|
||||
if len(paths) == 0 {
|
||||
continue
|
||||
}
|
||||
v := OrphanView{Node: n.ID, Paths: paths}
|
||||
for _, p := range paths {
|
||||
v.Bytes += size[p]
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Node < out[j].Node })
|
||||
return out
|
||||
}
|
||||
|
||||
func (a *API) listOrphans(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, a.Orphans())
|
||||
}
|
||||
|
||||
// collectWeights deletes them.
|
||||
//
|
||||
// A proposal, not an instruction: each node checks every path again before
|
||||
// unlinking it, and refuses anything it cannot prove it fetched. The two
|
||||
// checks are deliberate duplication — the control plane knows the catalog and
|
||||
// the node knows its disk, and only the node can be sure.
|
||||
func (a *API) collectWeights(w http.ResponseWriter, r *http.Request) {
|
||||
rm := a.ctl.WeightRemover()
|
||||
if rm == nil {
|
||||
writeError(w, http.StatusNotImplemented,
|
||||
"this control plane cannot ask nodes to delete files")
|
||||
return
|
||||
}
|
||||
|
||||
var body struct {
|
||||
// Node limits it to one node. Empty means every node, which is the
|
||||
// common case after removing a model from the catalog.
|
||||
Node string `json:"node,omitempty"`
|
||||
}
|
||||
// An empty body is the "all nodes" case and must not be an error.
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
|
||||
type result struct {
|
||||
Node string `json:"node"`
|
||||
Files int `json:"files"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
out := make([]result, 0)
|
||||
for _, o := range a.Orphans() {
|
||||
if body.Node != "" && o.Node != body.Node {
|
||||
continue
|
||||
}
|
||||
res := result{Node: o.Node, Files: len(o.Paths), Bytes: o.Bytes}
|
||||
// A node that has gone since the list was built is reported rather than
|
||||
// failing the whole sweep: the other nodes' space is still worth
|
||||
// freeing, and this one will be offered again next time.
|
||||
if err := rm.RemoveWeights(r.Context(), o.Node, o.Paths); err != nil {
|
||||
res.Error = err.Error()
|
||||
res.Files, res.Bytes = 0, 0
|
||||
}
|
||||
out = append(out, res)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
@@ -23,20 +23,35 @@ func NodeFromRegister(r *llmgwv1.Register) inventory.Node {
|
||||
State: inventory.Ready,
|
||||
AgentVersion: r.GetAgentVersion(),
|
||||
EngineVersions: r.GetEngineVersions(),
|
||||
MemoryTotalMB: int(r.GetMemoryTotalMb()),
|
||||
}
|
||||
for _, w := range r.GetWeights() {
|
||||
n.Weights = append(n.Weights, discovery.LocalWeights{
|
||||
Path: w.GetPath(), SizeBytes: w.GetSizeBytes(), Shared: w.GetShared(),
|
||||
})
|
||||
}
|
||||
for _, s := range r.GetStorage() {
|
||||
n.Storage = append(n.Storage, discovery.Storage{
|
||||
Path: s.GetPath(),
|
||||
FreeBytes: s.GetFreeBytes(),
|
||||
TotalBytes: s.GetTotalBytes(),
|
||||
Shared: s.GetShared(),
|
||||
})
|
||||
}
|
||||
for _, g := range r.GetGpus() {
|
||||
n.GPUs = append(n.GPUs, discovery.GPU{
|
||||
Index: int(g.GetIndex()),
|
||||
UUID: g.GetUuid(),
|
||||
Name: g.GetName(),
|
||||
MemoryTotalMB: int(g.GetMemoryTotalMb()),
|
||||
MemoryUsedMB: int(g.GetMemoryUsedMb()),
|
||||
ComputeCapability: g.GetComputeCapability(),
|
||||
NUMANode: int(g.GetNumaNode()),
|
||||
PCIBusID: g.GetPciBusId(),
|
||||
DriverVersion: g.GetDriverVersion(),
|
||||
CUDAVersion: g.GetCudaVersion(),
|
||||
NVLinkPeers: intSlice(g.GetNvlinkPeers()),
|
||||
Index: int(g.GetIndex()),
|
||||
UUID: g.GetUuid(),
|
||||
Name: g.GetName(),
|
||||
MemoryTotalMB: int(g.GetMemoryTotalMb()),
|
||||
MemoryUsedMB: int(g.GetMemoryUsedMb()),
|
||||
ComputeCapability: g.GetComputeCapability(),
|
||||
NUMANode: int(g.GetNumaNode()),
|
||||
PCIBusID: g.GetPciBusId(),
|
||||
DriverVersion: g.GetDriverVersion(),
|
||||
UtilizationPercent: int(g.GetUtilizationPercent()),
|
||||
CUDAVersion: g.GetCudaVersion(),
|
||||
NVLinkPeers: intSlice(g.GetNvlinkPeers()),
|
||||
Health: discovery.GPUHealth{
|
||||
ECCErrors: int(g.GetEccErrors()),
|
||||
TemperatureC: int(g.GetTemperatureC()),
|
||||
@@ -61,23 +76,38 @@ func RegisterFromInventory(nodeID, hostname, mgmtAddr, agentVersion string,
|
||||
r := &llmgwv1.Register{
|
||||
NodeId: nodeID, Hostname: hostname, MgmtAddr: mgmtAddr,
|
||||
Labels: labels, AgentVersion: agentVersion,
|
||||
MemoryTotalMb: int64(inv.MemoryTotalMB),
|
||||
}
|
||||
for _, w := range inv.Weights {
|
||||
r.Weights = append(r.Weights, &llmgwv1.LocalWeights{
|
||||
Path: w.Path, SizeBytes: w.SizeBytes, Shared: w.Shared,
|
||||
})
|
||||
}
|
||||
for _, s := range inv.Storage {
|
||||
r.Storage = append(r.Storage, &llmgwv1.Storage{
|
||||
Path: s.Path,
|
||||
FreeBytes: s.FreeBytes,
|
||||
TotalBytes: s.TotalBytes,
|
||||
Shared: s.Shared,
|
||||
})
|
||||
}
|
||||
for _, g := range inv.GPUs {
|
||||
r.Gpus = append(r.Gpus, &llmgwv1.GPU{
|
||||
Index: int32(g.Index),
|
||||
Uuid: g.UUID,
|
||||
Name: g.Name,
|
||||
MemoryTotalMb: int32(g.MemoryTotalMB),
|
||||
MemoryUsedMb: int32(g.MemoryUsedMB),
|
||||
ComputeCapability: g.ComputeCapability,
|
||||
NumaNode: int32(g.NUMANode),
|
||||
PciBusId: g.PCIBusID,
|
||||
DriverVersion: g.DriverVersion,
|
||||
CudaVersion: g.CUDAVersion,
|
||||
NvlinkPeers: int32Slice(g.NVLinkPeers),
|
||||
Problems: g.Health.Problems,
|
||||
TemperatureC: int32(g.Health.TemperatureC),
|
||||
EccErrors: int32(g.Health.ECCErrors),
|
||||
Index: int32(g.Index),
|
||||
Uuid: g.UUID,
|
||||
Name: g.Name,
|
||||
MemoryTotalMb: int32(g.MemoryTotalMB),
|
||||
MemoryUsedMb: int32(g.MemoryUsedMB),
|
||||
ComputeCapability: g.ComputeCapability,
|
||||
NumaNode: int32(g.NUMANode),
|
||||
PciBusId: g.PCIBusID,
|
||||
DriverVersion: g.DriverVersion,
|
||||
UtilizationPercent: int32(g.UtilizationPercent),
|
||||
CudaVersion: g.CUDAVersion,
|
||||
NvlinkPeers: int32Slice(g.NVLinkPeers),
|
||||
Problems: g.Health.Problems,
|
||||
TemperatureC: int32(g.Health.TemperatureC),
|
||||
EccErrors: int32(g.Health.ECCErrors),
|
||||
})
|
||||
}
|
||||
for _, f := range fabrics {
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
func nodeWith(s *Service, n inventory.Node) { s.inv.Upsert(n) }
|
||||
|
||||
func newTestService() *Service {
|
||||
return &Service{inv: inventory.NewStore(), conns: map[string]*conn{}}
|
||||
}
|
||||
|
||||
func TestDiskIsRefreshedWhileTheAgentRuns(t *testing.T) {
|
||||
// It was sampled once, at registration, and never again. Found on the live
|
||||
// cluster: an agent running for two days still reported the disk it had
|
||||
// when it started, and a file deleted in between was still being offered
|
||||
// for deletion. Disk is a fact about a moment, and the node whose figure
|
||||
// matters most is the one busy filling up.
|
||||
s := newTestService()
|
||||
nodeWith(s, inventory.Node{
|
||||
ID: "n1",
|
||||
State: inventory.Ready,
|
||||
Storage: []discovery.Storage{
|
||||
{Path: "/models", FreeBytes: 200 << 30, TotalBytes: 400 << 30},
|
||||
},
|
||||
})
|
||||
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{
|
||||
NodeId: "n1",
|
||||
Storage: []*llmgwv1.Storage{
|
||||
{Path: "/models", FreeBytes: 5 << 30, TotalBytes: 400 << 30},
|
||||
},
|
||||
})
|
||||
|
||||
n, _ := s.inv.Get("n1")
|
||||
free, total := n.LocalDisk()
|
||||
if free != 5<<30 || total != 400<<30 {
|
||||
t.Errorf("disk is %d of %d, want the figure the heartbeat carried", free, total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEmptyWeightListOnlyCountsWhenTheAgentMeansIt(t *testing.T) {
|
||||
// Most heartbeats carry no file list because nothing changed. Treating
|
||||
// that as "this node holds nothing" would empty a node's inventory five
|
||||
// seconds after it registered, and placement would stop preferring the
|
||||
// node that actually has the model.
|
||||
s := newTestService()
|
||||
nodeWith(s, inventory.Node{
|
||||
ID: "n1",
|
||||
State: inventory.Ready,
|
||||
Weights: []discovery.LocalWeights{{Path: "/models/a.gguf", SizeBytes: 100}},
|
||||
})
|
||||
|
||||
// Silence about weights: the list must survive.
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{
|
||||
NodeId: "n1",
|
||||
Storage: []*llmgwv1.Storage{{Path: "/models", FreeBytes: 1, TotalBytes: 2}},
|
||||
})
|
||||
if n, _ := s.inv.Get("n1"); len(n.Weights) != 1 {
|
||||
t.Fatalf("a heartbeat with no weight list emptied the inventory: %+v", n.Weights)
|
||||
}
|
||||
|
||||
// And when the agent says it means an empty list, it is taken.
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", WeightsChanged: true})
|
||||
if n, _ := s.inv.Get("n1"); len(n.Weights) != 0 {
|
||||
t.Errorf("a deliberate empty list was ignored: %+v", n.Weights)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANewlyFetchedModelAppearsWithoutARestart(t *testing.T) {
|
||||
// The other direction of the same bug: a model fetched after registration
|
||||
// was invisible until the agent restarted, so placement could not prefer
|
||||
// the node that had just spent ten minutes downloading it.
|
||||
s := newTestService()
|
||||
nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready})
|
||||
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{
|
||||
NodeId: "n1",
|
||||
WeightsChanged: true,
|
||||
Weights: []*llmgwv1.LocalWeights{
|
||||
{Path: "/models/new.gguf", SizeBytes: 4 << 30},
|
||||
},
|
||||
})
|
||||
|
||||
n, _ := s.inv.Get("n1")
|
||||
if !n.HasWeights("/models/new.gguf") {
|
||||
t.Errorf("a model fetched after registration did not appear: %+v", n.Weights)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAHeartbeatAboutNothingChangesNothing(t *testing.T) {
|
||||
s := newTestService()
|
||||
before := inventory.Node{
|
||||
ID: "n1",
|
||||
State: inventory.Ready,
|
||||
Weights: []discovery.LocalWeights{{Path: "/models/a.gguf", SizeBytes: 100}},
|
||||
Storage: []discovery.Storage{{Path: "/models", FreeBytes: 7, TotalBytes: 9}},
|
||||
}
|
||||
nodeWith(s, before)
|
||||
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1"})
|
||||
|
||||
n, _ := s.inv.Get("n1")
|
||||
if len(n.Weights) != 1 {
|
||||
t.Errorf("weights changed: %+v", n.Weights)
|
||||
}
|
||||
if free, _ := n.LocalDisk(); free != 7 {
|
||||
t.Errorf("disk changed: %d", free)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/x509"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
// peerKey carries the verified client identity through to the stream handler.
|
||||
//
|
||||
// connect's BidiStream exposes the peer's address but not its certificate, so
|
||||
// the identity is lifted out of the TLS state by middleware and put on the
|
||||
// context before the connect handler ever sees the request.
|
||||
type peerKey struct{}
|
||||
|
||||
// PeerIdentity is who the other end proved itself to be.
|
||||
type PeerIdentity struct {
|
||||
// CommonName is the certificate's subject CN, which is the node id a node
|
||||
// is allowed to claim.
|
||||
CommonName string
|
||||
// Verified is false when the connection carried no client certificate,
|
||||
// which only happens when TLS is switched off.
|
||||
Verified bool
|
||||
}
|
||||
|
||||
// WithPeerIdentity records a verified identity on the context.
|
||||
func WithPeerIdentity(ctx context.Context, id PeerIdentity) context.Context {
|
||||
return context.WithValue(ctx, peerKey{}, id)
|
||||
}
|
||||
|
||||
// peerFrom reads the identity a request was authenticated with.
|
||||
func peerFrom(ctx context.Context) PeerIdentity {
|
||||
id, _ := ctx.Value(peerKey{}).(PeerIdentity)
|
||||
return id
|
||||
}
|
||||
|
||||
// AuthMiddleware lifts the client certificate's common name onto the request
|
||||
// context.
|
||||
//
|
||||
// The certificate chain has already been verified by the TLS handshake — this
|
||||
// only makes the result reachable from the handler. Without a certificate the
|
||||
// identity is marked unverified and the handler decides what to do, which keeps
|
||||
// the insecure development path working without a second code path through the
|
||||
// service itself.
|
||||
func AuthMiddleware(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var id PeerIdentity
|
||||
if r.TLS != nil && len(r.TLS.PeerCertificates) > 0 {
|
||||
id = identityOf(r.TLS.PeerCertificates[0])
|
||||
}
|
||||
next.ServeHTTP(w, r.WithContext(WithPeerIdentity(r.Context(), id)))
|
||||
})
|
||||
}
|
||||
|
||||
func identityOf(cert *x509.Certificate) PeerIdentity {
|
||||
return PeerIdentity{CommonName: cert.Subject.CommonName, Verified: true}
|
||||
}
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
@@ -23,7 +24,7 @@ import (
|
||||
// agents report they are already running. The reconciler implements it.
|
||||
type StatusSink interface {
|
||||
OnShardStatus(replicaID, shardID string, state supervisor.State, errMsg string)
|
||||
Adopt(shards []reconciler.AdoptedShard) (unclaimed []reconciler.AdoptedShard)
|
||||
Adopt(nodeID string, shards []reconciler.AdoptedShard) (unclaimed []reconciler.AdoptedShard)
|
||||
}
|
||||
|
||||
// Service is the control plane's AgentService implementation, and at the same
|
||||
@@ -34,6 +35,10 @@ type Service struct {
|
||||
status StatusSink
|
||||
log *slog.Logger
|
||||
|
||||
// requireIdentity refuses agents that present no certificate. It is off
|
||||
// only for the insecure development path.
|
||||
requireIdentity bool
|
||||
|
||||
mu sync.RWMutex
|
||||
conns map[string]*conn
|
||||
}
|
||||
@@ -52,16 +57,28 @@ type conn struct {
|
||||
func (c *conn) close() { c.once.Do(func() { close(c.closed) }) }
|
||||
|
||||
// New builds the service.
|
||||
func New(inv *inventory.Store, status StatusSink, log *slog.Logger) *Service {
|
||||
//
|
||||
// requireIdentity refuses any agent that has not proved who it is. It should be
|
||||
// true whenever TLS is configured, and is false only for local runs without
|
||||
// certificates.
|
||||
func New(inv *inventory.Store, status StatusSink, requireIdentity bool, log *slog.Logger) *Service {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Service{inv: inv, status: status, log: log, conns: map[string]*conn{}}
|
||||
return &Service{
|
||||
inv: inv, status: status, log: log,
|
||||
requireIdentity: requireIdentity,
|
||||
conns: map[string]*conn{},
|
||||
}
|
||||
}
|
||||
|
||||
// Handler returns the HTTP path and handler to mount.
|
||||
//
|
||||
// The handler is wrapped so the verified client identity reaches the stream:
|
||||
// connect exposes the peer's address but not its certificate.
|
||||
func (s *Service) Handler() (string, http.Handler) {
|
||||
return llmgwv1connect.NewAgentServiceHandler(s)
|
||||
path, h := llmgwv1connect.NewAgentServiceHandler(s)
|
||||
return path, AuthMiddleware(h)
|
||||
}
|
||||
|
||||
// Connect serves one agent's stream for the lifetime of that agent.
|
||||
@@ -85,6 +102,24 @@ func (s *Service) Connect(
|
||||
return connect.NewError(connect.CodeInvalidArgument, errors.New("node_id is required"))
|
||||
}
|
||||
|
||||
// The certificate decides who this is, not the message.
|
||||
//
|
||||
// Without this check an agent could register as any node in the fleet and
|
||||
// receive that node's shards — which is the same exposure the stream had
|
||||
// before TLS, just with encryption on top. mTLS is only worth having if the
|
||||
// identity it proves is the one actually used.
|
||||
if peer := peerFrom(ctx); peer.Verified {
|
||||
if peer.CommonName != nodeID {
|
||||
s.log.Warn("rejecting an agent claiming a node id its certificate does not authorize",
|
||||
"certificate_cn", peer.CommonName, "claimed", nodeID)
|
||||
return connect.NewError(connect.CodePermissionDenied,
|
||||
fmt.Errorf("certificate authorizes node %q, not %q", peer.CommonName, nodeID))
|
||||
}
|
||||
} else if s.requireIdentity {
|
||||
return connect.NewError(connect.CodeUnauthenticated,
|
||||
errors.New("no verified client certificate"))
|
||||
}
|
||||
|
||||
node := NodeFromRegister(reg)
|
||||
s.inv.Upsert(node)
|
||||
s.log.Info("agent registered",
|
||||
@@ -133,11 +168,14 @@ func (s *Service) Connect(
|
||||
|
||||
// adopt hands the agent's running shards to the reconciler and stops whatever
|
||||
// nothing in the catalog claims.
|
||||
//
|
||||
// An empty list is not skipped. "I am running nothing" is information — it is
|
||||
// how the control plane learns that an engine died while its agent survived —
|
||||
// and returning early on it made the reconciler's half of that exchange
|
||||
// unreachable in production while its unit tests, which call the reconciler
|
||||
// directly, kept passing.
|
||||
func (s *Service) adopt(ctx context.Context, nodeID string, shards []*llmgwv1.RunningShard, c *conn) {
|
||||
if len(shards) == 0 {
|
||||
return
|
||||
}
|
||||
unclaimed := s.status.Adopt(AdoptedFromRegister(nodeID, shards))
|
||||
unclaimed := s.status.Adopt(nodeID, AdoptedFromRegister(nodeID, shards))
|
||||
for _, u := range unclaimed {
|
||||
// The model or profile left the catalog while this was running. Nobody
|
||||
// will ever route to it, and it is holding GPUs.
|
||||
@@ -202,7 +240,8 @@ func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) {
|
||||
func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
|
||||
used := hb.GetGpuMemoryUsedMb()
|
||||
problems := hb.GetNewProblems()
|
||||
if len(used) == 0 && len(problems) == 0 {
|
||||
storage := hb.GetStorage()
|
||||
if len(used) == 0 && len(problems) == 0 && len(storage) == 0 && !hb.GetWeightsChanged() {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -222,6 +261,33 @@ func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
|
||||
// conservative direction.
|
||||
n.GPUs[0].Health.Problems = append(n.GPUs[0].Health.Problems, problems...)
|
||||
}
|
||||
|
||||
// Disk is a fact about a moment, not about a machine. Sampled only at
|
||||
// registration it is wrong on exactly the node worth watching — the one
|
||||
// busy filling up — and it stays wrong until that agent restarts.
|
||||
if len(storage) > 0 {
|
||||
n.Storage = n.Storage[:0]
|
||||
for _, st := range storage {
|
||||
n.Storage = append(n.Storage, discovery.Storage{
|
||||
Path: st.GetPath(),
|
||||
FreeBytes: st.GetFreeBytes(),
|
||||
TotalBytes: st.GetTotalBytes(),
|
||||
Shared: st.GetShared(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Only when the agent says the list is meant. An empty list without the
|
||||
// flag is silence, and treating it as "this node holds nothing" would
|
||||
// empty a node's inventory on every heartbeat.
|
||||
if hb.GetWeightsChanged() {
|
||||
n.Weights = n.Weights[:0]
|
||||
for _, w := range hb.GetWeights() {
|
||||
n.Weights = append(n.Weights, discovery.LocalWeights{
|
||||
Path: w.GetPath(), SizeBytes: w.GetSizeBytes(), Shared: w.GetShared(),
|
||||
})
|
||||
}
|
||||
}
|
||||
s.inv.Upsert(n)
|
||||
}
|
||||
|
||||
@@ -303,6 +369,22 @@ func (s *Service) DrainShard(_ context.Context, nodeID, shardID string) error {
|
||||
})
|
||||
}
|
||||
|
||||
// RemoveWeights asks a node to delete cached files no model needs.
|
||||
//
|
||||
// The node checks every path again before unlinking it, so this is a proposal
|
||||
// rather than an instruction — which is the right shape for a message that
|
||||
// deletes data the control plane cannot see.
|
||||
func (s *Service) RemoveWeights(_ context.Context, nodeID string, paths []string) error {
|
||||
if len(paths) == 0 {
|
||||
return nil
|
||||
}
|
||||
return s.send(nodeID, &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_RemoveWeights{
|
||||
RemoveWeights: &llmgwv1.RemoveWeights{Paths: paths},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) send(nodeID string, msg *llmgwv1.ControlMessage) error {
|
||||
s.mu.RLock()
|
||||
c, ok := s.conns[nodeID]
|
||||
|
||||
@@ -14,10 +14,26 @@ import (
|
||||
|
||||
// Scaling is how many replicas a model should have.
|
||||
type Scaling struct {
|
||||
Min int `koanf:"min"`
|
||||
Max int `koanf:"max"`
|
||||
Min int `json:"min,omitempty" koanf:"min"`
|
||||
Max int `json:"max,omitempty" koanf:"max"`
|
||||
// IdleTimeout is how long an on-demand model may go unasked-for before its
|
||||
// replicas are stopped. Zero takes the cluster default.
|
||||
//
|
||||
// Minutes rather than seconds, always. Stopping a replica throws away the
|
||||
// KV cache prefix routing spent its effort filling, and a model that flaps
|
||||
// in and out loses that entirely — so the saving has to be worth more than
|
||||
// the cache, which it only is over a real idle period.
|
||||
IdleTimeout time.Duration `json:"idle_timeout,omitempty" koanf:"idle_timeout"`
|
||||
}
|
||||
|
||||
// OnDemand reports whether this model is loaded when someone asks for it.
|
||||
//
|
||||
// min 0 with a max above it, which needs no new field and reads correctly: keep
|
||||
// none running, allow up to max. min 0 and max 0 is a model that is configured
|
||||
// and deliberately not running, which is a different and equally useful thing —
|
||||
// the difference between "load it when wanted" and "do not load it".
|
||||
func (s Scaling) OnDemand() bool { return s.Min == 0 && s.Max > 0 }
|
||||
|
||||
// SLO is the latency target a model is expected to meet. `ctl bench` measures
|
||||
// capacity against it (plan §6.2.1).
|
||||
type SLO struct {
|
||||
@@ -27,42 +43,127 @@ type SLO struct {
|
||||
|
||||
// Profile is how a model runs on one GPU class.
|
||||
type Profile struct {
|
||||
Name string `koanf:"name"`
|
||||
Engine string `koanf:"engine"`
|
||||
GPUClass string `koanf:"gpu_class"`
|
||||
Name string `json:"name,omitempty" koanf:"name"`
|
||||
Engine string `json:"engine,omitempty" koanf:"engine"`
|
||||
GPUClass string `json:"gpu_class,omitempty" koanf:"gpu_class"`
|
||||
|
||||
TP int `koanf:"tp"`
|
||||
PP int `koanf:"pp"`
|
||||
TP int `json:"tp,omitempty" koanf:"tp"`
|
||||
PP int `json:"pp,omitempty" koanf:"pp"`
|
||||
// GPUsPerNode defaults to TP for a single-node profile.
|
||||
GPUsPerNode int `koanf:"gpus_per_node"`
|
||||
GPUsPerNode int `json:"gpus_per_node,omitempty" koanf:"gpus_per_node"`
|
||||
// Nodes above 1 makes this a multi-node profile, which requires a fabric.
|
||||
Nodes int `koanf:"nodes"`
|
||||
Nodes int `json:"nodes,omitempty" koanf:"nodes"`
|
||||
|
||||
Args []string `koanf:"args"`
|
||||
Env map[string]string `koanf:"env"`
|
||||
Args []string `json:"args,omitempty" koanf:"args"`
|
||||
Env map[string]string `json:"env,omitempty" koanf:"env"`
|
||||
|
||||
// MinMemoryGB is the host RAM this profile needs to run.
|
||||
//
|
||||
// The number that matters is the model's working set: the weights, plus the
|
||||
// KV cache for the whole context pool, plus the engine's own overhead. For
|
||||
// a GPU profile that lives in the cards and this is small; for a CPU
|
||||
// profile it is the whole thing, and getting it wrong means an OOM on a
|
||||
// node that reported itself healthy.
|
||||
//
|
||||
// Zero means unconstrained, which is every profile that existed before
|
||||
// this. A node that does not report its memory is not excluded either —
|
||||
// an older agent must not have its nodes quietly emptied by an upgrade of
|
||||
// the control plane.
|
||||
MinMemoryGB int `json:"min_memory_gb,omitempty" koanf:"min_memory_gb"`
|
||||
|
||||
// MinDiskGB is the free space a node needs before it can hold these
|
||||
// weights.
|
||||
//
|
||||
// Only checked when the node does not already have them: a node holding
|
||||
// the file needs no room for it, and requiring the space twice would keep
|
||||
// a model off the one node best placed to run it.
|
||||
//
|
||||
// Declared rather than derived, because the size of a model that has never
|
||||
// been fetched anywhere is not knowable from here — and a wrong guess
|
||||
// fails late, halfway through a download, which is the failure this exists
|
||||
// to prevent. Zero means unconstrained, and a node that reports no storage
|
||||
// is not excluded: unknown is not full.
|
||||
MinDiskGB int `json:"min_disk_gb,omitempty" koanf:"min_disk_gb"`
|
||||
|
||||
// NodeSelector restricts this profile to nodes carrying every one of these
|
||||
// labels, with these values.
|
||||
//
|
||||
// A CPU profile otherwise fits everywhere, because there is no card to be
|
||||
// free — which is right for a 0.5B model and wrong for a 14B one. Nothing
|
||||
// in the inventory says how much memory a node has, so the operator says
|
||||
// which nodes are meant for a model instead of the scheduler inferring it.
|
||||
// Agents have reported labels since the beginning; this is what reads them.
|
||||
NodeSelector map[string]string `json:"node_selector,omitempty" koanf:"node_selector"`
|
||||
|
||||
// Weight is the router's capacity weight. It should come from ctl bench
|
||||
// rather than a guess (plan §6.2.1).
|
||||
Weight float64 `koanf:"weight"`
|
||||
Weight float64 `json:"weight,omitempty" koanf:"weight"`
|
||||
// Priority orders profiles during placement; highest that fits wins.
|
||||
Priority int `koanf:"priority"`
|
||||
Priority int `json:"priority,omitempty" koanf:"priority"`
|
||||
// StartupTimeout bounds model loading. A 70B model across two nodes takes
|
||||
// tens of minutes, so this is per profile rather than global.
|
||||
StartupTimeout time.Duration `koanf:"startup_timeout"`
|
||||
StartupTimeout time.Duration `json:"startup_timeout,omitempty" koanf:"startup_timeout"`
|
||||
}
|
||||
|
||||
// CPUOnly reports whether this profile needs no GPUs.
|
||||
//
|
||||
// Naming no GPU class is how a profile says so, and only engines that can
|
||||
// actually serve without a card are allowed to leave it out — otherwise an
|
||||
// omitted class would quietly become a replica placed on a machine with no
|
||||
// hardware to run it.
|
||||
func (p Profile) CPUOnly() bool { return p.GPUClass == "" && cpuCapable(p.Engine) }
|
||||
|
||||
// cpuCapable reports whether an engine can serve a model without a GPU.
|
||||
//
|
||||
// llama.cpp is the one that can, and it is what makes a cluster with no cards
|
||||
// testable end to end.
|
||||
func cpuCapable(engine string) bool { return engine == "llamacpp" }
|
||||
|
||||
// Model is a catalog entry as clients see it.
|
||||
type Model struct {
|
||||
Name string `koanf:"name"`
|
||||
Aliases []string `koanf:"aliases"`
|
||||
Name string `json:"name,omitempty" koanf:"name"`
|
||||
Aliases []string `json:"aliases,omitempty" koanf:"aliases"`
|
||||
// Source is the weights path or a HuggingFace id.
|
||||
Source string `koanf:"source"`
|
||||
Profiles []Profile `koanf:"profiles"`
|
||||
Scaling Scaling `koanf:"scaling"`
|
||||
Fallbacks []string `koanf:"fallbacks"`
|
||||
Source string `json:"source,omitempty" koanf:"source"`
|
||||
Profiles []Profile `json:"profiles,omitempty" koanf:"profiles"`
|
||||
Scaling Scaling `json:"scaling,omitempty" koanf:"scaling"`
|
||||
Fallbacks []string `json:"fallbacks,omitempty" koanf:"fallbacks"`
|
||||
// RoutingProfile selects how the router scores endpoints for this model.
|
||||
RoutingProfile string `koanf:"routing_profile"`
|
||||
SLO SLO `koanf:"slo"`
|
||||
RoutingProfile string `json:"routing_profile,omitempty" koanf:"routing_profile"`
|
||||
SLO SLO `json:"slo,omitempty" koanf:"slo"`
|
||||
|
||||
// ContextLength is how many tokens one request may use, prompt and
|
||||
// completion together, and MaxOutput bounds the completion alone. Zero
|
||||
// means unknown, which is what every model was until now.
|
||||
//
|
||||
// The gateway knew neither, so an over-long prompt was not refused: it was
|
||||
// admitted, given a replica, and failed at the engine, and a client sizing
|
||||
// its prompts from /v1/models had nothing to read. Worse, the number is not
|
||||
// the one an operator writes in the engine's arguments — llama.cpp divides
|
||||
// --ctx-size across its parallel slots, so a pool of 8192 across four slots
|
||||
// is 2048 a request, and nothing said so.
|
||||
ContextLength int `json:"context_length,omitempty" koanf:"context_length"`
|
||||
MaxOutput int `json:"max_output,omitempty" koanf:"max_output"`
|
||||
|
||||
// Capabilities is what this model can do beyond plain completion.
|
||||
Capabilities Capabilities `json:"capabilities,omitempty" koanf:"capabilities"`
|
||||
}
|
||||
|
||||
// Capabilities describes what a client may ask a model for.
|
||||
//
|
||||
// Declared by the operator rather than discovered, because there is no way to
|
||||
// ask an engine what its model can do. That makes these claims, and a wrong
|
||||
// claim is worse than none — which is why `ctl engine-check` verifies tool
|
||||
// calling against a running engine rather than trusting this.
|
||||
//
|
||||
// The reason it exists: the gateway forwards a tools array to any model, and a
|
||||
// model that cannot use them answers in prose. Every OpenAI-compatible client
|
||||
// then behaves as though the model is broken, and the only way to find out was
|
||||
// to try it. That cost an afternoon.
|
||||
type Capabilities struct {
|
||||
Tools bool `json:"tools,omitempty" koanf:"tools"`
|
||||
Vision bool `json:"vision,omitempty" koanf:"vision"`
|
||||
Embeddings bool `json:"embeddings,omitempty" koanf:"embeddings"`
|
||||
}
|
||||
|
||||
// Catalog is the set of models the cluster should serve.
|
||||
@@ -122,8 +223,14 @@ func (m Model) normalize() Model {
|
||||
if m.Scaling.Min < 0 {
|
||||
m.Scaling.Min = 0
|
||||
}
|
||||
if m.Scaling.Max == 0 {
|
||||
m.Scaling.Max = max(m.Scaling.Min, 1)
|
||||
// Max 0 used to be an omission to fill in, because a model with no maximum
|
||||
// could not run. It is now also a statement: with min 0 it means configured
|
||||
// and deliberately not running, which is a different and useful thing from
|
||||
// "load it when wanted". Filling it in would silently turn every switched-
|
||||
// off model into an on-demand one — including the ones already written that
|
||||
// way in a running cluster's configuration.
|
||||
if m.Scaling.Max == 0 && m.Scaling.Min > 0 {
|
||||
m.Scaling.Max = m.Scaling.Min
|
||||
}
|
||||
for i := range m.Profiles {
|
||||
p := &m.Profiles[i]
|
||||
@@ -136,7 +243,11 @@ func (m Model) normalize() Model {
|
||||
if p.Nodes <= 0 {
|
||||
p.Nodes = 1
|
||||
}
|
||||
if p.GPUsPerNode <= 0 {
|
||||
switch {
|
||||
case p.CPUOnly():
|
||||
// Nothing to divide up: this profile asks for no cards at all.
|
||||
p.GPUsPerNode = 0
|
||||
case p.GPUsPerNode <= 0:
|
||||
// The common case: a single-node profile uses TP GPUs on one node.
|
||||
p.GPUsPerNode = p.TP / p.Nodes
|
||||
if p.GPUsPerNode <= 0 {
|
||||
@@ -177,8 +288,16 @@ func (m Model) validate() error {
|
||||
}
|
||||
seen[p.Name] = struct{}{}
|
||||
|
||||
if p.GPUClass == "" {
|
||||
return fmt.Errorf("model %q profile %q: gpu_class is required", m.Name, p.Name)
|
||||
if p.GPUClass == "" && !cpuCapable(p.Engine) {
|
||||
return fmt.Errorf("model %q profile %q: gpu_class is required for engine %q",
|
||||
m.Name, p.Name, p.Engine)
|
||||
}
|
||||
if p.CPUOnly() {
|
||||
if p.Nodes > 1 {
|
||||
return fmt.Errorf("model %q profile %q: a CPU profile cannot span %d nodes",
|
||||
m.Name, p.Name, p.Nodes)
|
||||
}
|
||||
continue
|
||||
}
|
||||
// The arithmetic has to add up, or the engine is handed a shard spec it
|
||||
// cannot satisfy and fails minutes later during model loading.
|
||||
@@ -231,12 +350,18 @@ func (m Model) PlacementProfiles() []placement.Profile {
|
||||
out := make([]placement.Profile, 0, len(m.Profiles))
|
||||
for _, p := range m.Profiles {
|
||||
out = append(out, placement.Profile{
|
||||
Name: p.Name,
|
||||
Engine: p.Engine,
|
||||
GPUClass: p.GPUClass,
|
||||
GPUsPerNode: p.GPUsPerNode,
|
||||
Nodes: p.Nodes,
|
||||
Priority: p.Priority,
|
||||
Name: p.Name,
|
||||
Engine: p.Engine,
|
||||
GPUClass: p.GPUClass,
|
||||
GPUsPerNode: p.GPUsPerNode,
|
||||
Nodes: p.Nodes,
|
||||
Priority: p.Priority,
|
||||
NodeSelector: p.NodeSelector,
|
||||
MinMemoryGB: p.MinMemoryGB,
|
||||
MinDiskGB: p.MinDiskGB,
|
||||
// The model's source, so placement can prefer nodes that already
|
||||
// hold the weights rather than one that must first move them.
|
||||
Source: m.Source,
|
||||
})
|
||||
}
|
||||
return out
|
||||
@@ -287,3 +412,24 @@ func (c *Catalog) String() string {
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// Simulated reports whether this model produces synthetic text rather than
|
||||
// inference output.
|
||||
//
|
||||
// Derived from the profiles rather than configured, because an operator who
|
||||
// added a sim profile and forgot the flag would be shipping placeholder text to
|
||||
// users under a real model's name. A model counts as simulated only when every
|
||||
// profile is simulated: one real profile means a request can land on a real
|
||||
// engine, and a model that sometimes answers for real must not be labelled as
|
||||
// though it never does.
|
||||
func (m Model) Simulated() bool {
|
||||
if len(m.Profiles) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, p := range m.Profiles {
|
||||
if p.Engine != "sim" {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -124,8 +124,51 @@ func TestDefaults(t *testing.T) {
|
||||
if p.Weight <= 0 {
|
||||
t.Errorf("weight=%v", p.Weight)
|
||||
}
|
||||
if m.Scaling.Max < 1 {
|
||||
t.Errorf("scaling max=%d", m.Scaling.Max)
|
||||
// Max is no longer filled in for a model that names no minimum. With min 0
|
||||
// it now means "configured and deliberately not running", and defaulting it
|
||||
// to 1 would make every model without a scaling stanza loadable on demand —
|
||||
// so any client request could start a 70B nobody intended to run.
|
||||
if m.Scaling.Min == 0 && m.Scaling.Max != 0 {
|
||||
t.Errorf("scaling max=%d; an unscaled model should stay off, not become on-demand", m.Scaling.Max)
|
||||
}
|
||||
if m.Scaling.OnDemand() {
|
||||
t.Error("a model with no scaling stanza was made on-demand")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaxDefaultsToTheMinimumWhenOneIsGiven(t *testing.T) {
|
||||
// A model that says it wants two replicas and nothing about a maximum wants
|
||||
// two, not one and not unbounded.
|
||||
built, err := New([]Model{{
|
||||
Name: "m", Source: "/m", Scaling: Scaling{Min: 2},
|
||||
Profiles: []Profile{{Name: "p", Engine: "llamacpp"}},
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m, _ := built.Resolve("m")
|
||||
if m.Scaling.Max != 2 {
|
||||
t.Errorf("max=%d, want 2", m.Scaling.Max)
|
||||
}
|
||||
if m.Scaling.OnDemand() {
|
||||
t.Error("a model with min 2 was reported on-demand")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnDemandNeedsBothZeroMinimumAndARealMaximum(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
s Scaling
|
||||
want bool
|
||||
}{
|
||||
{"off", Scaling{Min: 0, Max: 0}, false},
|
||||
{"on demand", Scaling{Min: 0, Max: 1}, true},
|
||||
{"always on", Scaling{Min: 1, Max: 1}, false},
|
||||
{"always on, room to grow", Scaling{Min: 1, Max: 4}, false},
|
||||
} {
|
||||
if got := tc.s.OnDemand(); got != tc.want {
|
||||
t.Errorf("%s: OnDemand() = %v, want %v", tc.name, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -233,3 +276,29 @@ func TestNamesAndModelsOrder(t *testing.T) {
|
||||
t.Errorf("Names order=%v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimulatedNeedsEveryProfileToBeSimulated(t *testing.T) {
|
||||
// The question the flag answers is "can this request come back as real
|
||||
// inference?". One real profile means yes, however many sim profiles sit
|
||||
// beside it, so a mixed model must not be labelled simulated.
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
engines []string
|
||||
wantSim bool
|
||||
}{
|
||||
{"all simulated", []string{"sim", "sim"}, true},
|
||||
{"one real profile", []string{"sim", "llamacpp"}, false},
|
||||
{"no simulation", []string{"llamacpp"}, false},
|
||||
{"no profiles at all", nil, false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var m Model
|
||||
for _, e := range tc.engines {
|
||||
m.Profiles = append(m.Profiles, Profile{Engine: e})
|
||||
}
|
||||
if got := m.Simulated(); got != tc.wantSim {
|
||||
t.Errorf("Simulated() = %v, want %v", got, tc.wantSim)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Loaded 100 of 244 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user