Compare commits
No files matched your search
+75
-4
@@ -30,17 +30,31 @@ What is left, most needed first. The detail is in the item each names.
|
||||
**Soon: the customer's experience, and HA.**
|
||||
4. ~~**C24**~~: done 2026-10-21. The store's, catalog's and mirror's errors, percent signs, footer times and stale
|
||||
forms are in the page's language. Left: the /docs documents, and a language kept per person (cookie only today).
|
||||
5. **J4**: the model mirror on more than one machine, the last single point of an HA install. (M)
|
||||
5. **J4**: the model mirror on more than one machine. Code done 2026-10-29: mirrors kept in step (`mirror serve
|
||||
--peer`, removals carried), `mirror.peers` and failover in the control plane, nodes given the file on a mirror that
|
||||
answers, `mirror add-peer` and the setup's second mirror. Tried 2026-10-29 on two VMs:
|
||||
- `add-peer` set the second up over ssh;
|
||||
- a file came across in ~30 s, checksum matching;
|
||||
- a version added on the second while the first was off reached the first when it came back;
|
||||
- one removed on the first while the second was off was removed there when it came back, not copied back.
|
||||
|
||||
Left: the same with a control plane and a node, a model served while the first mirror is off. That part (failover
|
||||
in the control plane, the node given the other mirror's file) is so far tried only in tests.
|
||||
6. **F5**: large imports limited by the embedding engine's slots. (M)
|
||||
|
||||
**From Kong** (section L): L1 done; L9, then L2 and L3, then L4/L5. L6 and L7 are llmgw-80's.
|
||||
**From Kong** (section L): L1–L5 and L9 done (v0.15.0); left of L4 is routing by cost between outside
|
||||
providers, of L5 cache-write and audio prices. L6 and L7 are llmgw-80's; L8 (A2A) when asked.
|
||||
|
||||
**Models we can trust** (section M, 2026-10-25): M1 done (pickles refused; the explorer lists only what this
|
||||
cluster's hardware holds, with CPU and GPU tok/s). Next M2 (every download pinned and checked against Hugging
|
||||
Face), then M5 (engines without the internet), M3, M4, M6.
|
||||
|
||||
**Governance.**
|
||||
7. **C1**: a tenant-admin role, and forms hidden from viewers.
|
||||
8. **Audit before and after**: what a change changed, and changes made with the CLI recorded too.
|
||||
|
||||
**When there is room.**
|
||||
9. **I1**: answers to similar questions from the cache. (M)
|
||||
9. ~~**I1**~~: done 2026-10-25 (in memory, not pgvector); left: the 40-rewordings run on the PoC.
|
||||
10. **H1**: tried on the PoC.
|
||||
11. **C4**: Parquet, and a scheduled push for usage export.
|
||||
|
||||
@@ -59,7 +73,7 @@ of several control planes. 15. **J6**: rate limits per process (a note).
|
||||
**Not code, the operator's to decide or do:**
|
||||
- Where Alertmanager sends (its receiver is "none"): with a judge down, nobody hears.
|
||||
- The fabric test for K.
|
||||
- The PoC's control planes run a dev build of main (v0.14.9-11-g4eb49d0); its nodes run v0.14.8.
|
||||
- The PoC runs a dev build of main (v0.15.0-7-ga583359) on its control planes and nodes, the mirror v0.15.0.
|
||||
|
||||
---
|
||||
|
||||
@@ -1784,6 +1798,63 @@ A customer that already runs Kong puts llmgw behind it: `ai-proxy` with
|
||||
route. A page in docs/ with the configuration, tried against Kong 3.9 in a
|
||||
container, and what Kong's token accounting sees.
|
||||
|
||||
## M. Models from Hugging Face we can trust (written 2026-10-25)
|
||||
|
||||
The user asked how we know a model fetched from Hugging Face is safe. Two
|
||||
questions: whether its files can run code on our machines, which we can
|
||||
close, and whether the model itself was trained to misbehave, which no scan
|
||||
proves — only who published it, an eval before it serves, and the
|
||||
guardrails narrow that. What existed: `--trust-remote-code` refused, no
|
||||
repository Python fetched, ComfyUI nodes pinned to a commit, engines
|
||||
hash-pinned and run as `llmgw-engine`, mirror versions pinned to a commit
|
||||
and never overwritten. What did not: everything below.
|
||||
|
||||
### M1. Weights in a pickle format refused (S) — done 2026-10-25
|
||||
`.bin`, `.pt`, `.pth`, `.ckpt`, `.pkl` run code as they are loaded, which is
|
||||
how the malicious models found on Hugging Face work. The node refuses such a
|
||||
source, or a mirror set holding one, before fetching anything; a model being
|
||||
added or given such a source is refused; vLLM and SGLang downloading a model
|
||||
by its id are given `--load-format safetensors`. The explorer hides a repository with neither GGUF nor safetensors
|
||||
(under "runs here", ticked by default) and does not offer a repository's
|
||||
pickle files. `allow_pickle: true` on a model lets it through, on purpose. On the
|
||||
PoC, `triposr` (`hf:stabilityai/TripoSR/model.ckpt`) is the one model that
|
||||
needs it.
|
||||
|
||||
### M2. Every download pinned and checked against Hugging Face (M)
|
||||
- `hf:org/repo/file` follows `main` and is checked by nothing; a bare
|
||||
`org/repo` given to vLLM is downloaded by vLLM, of whatever `main` is.
|
||||
Both: resolve to the commit when the model is added, keep it in the
|
||||
catalog's origin, and fetch that commit only (`--revision` for vLLM).
|
||||
- The sha256 Hugging Face publishes for each LFS file compared on every
|
||||
path, and a file that differs removed: `mirror add`, the manager, the
|
||||
node's `hf:` and http sources, mirror set files with no sha recorded.
|
||||
- Done when: a file changed on Hugging Face after a model was added is
|
||||
refused by the node with both checksums named.
|
||||
|
||||
### M3. Hugging Face's own scan (S)
|
||||
Hugging Face scans every file (pickle imports, ClamAV, partners' scanners)
|
||||
and marks a file unsafe. Ask for it with the repository listing; refuse an
|
||||
unsafe file, show the result on the mirror page and in the catalog's origin.
|
||||
The API field has to be confirmed first.
|
||||
|
||||
### M4. Publishers we trust, and the rest asked about (S)
|
||||
An allow-list of organisations in the cluster's settings (Qwen, meta-llama,
|
||||
mistralai, google, microsoft, …). A model from anyone else is added only by
|
||||
an admin who says so, and who said it is kept in the origin. The PoC already
|
||||
serves four community models whose training nobody here checked.
|
||||
|
||||
### M5. Engines without the internet (M)
|
||||
The firewall is inbound only, so an engine that was made to run code can
|
||||
send what it reads anywhere. With weights from the mirror the engines need
|
||||
no egress: an nft output rule for the engine user (allow the mirror, the
|
||||
cluster, loopback), systemd hardening of the agent's engines
|
||||
(`NoNewPrivileges`, `ProtectSystem`, private /tmp). Agent itself still root.
|
||||
|
||||
### M6. An eval before a new model serves (S)
|
||||
`eval` and `canary` exist; make a new model from an untrusted publisher
|
||||
(M4) serve only after its eval suite passes, so a model trained to answer
|
||||
wrongly on a trigger has at least been asked the suite.
|
||||
|
||||
## D. Facts worth knowing before starting
|
||||
|
||||
- Admin API: `127.0.0.1:9090`, behind the token in
|
||||
|
||||
+223
-1
@@ -19,12 +19,234 @@ is only released once its **exit criterion** is met, not when its code lands.
|
||||
| v0.12.0 | — | The cluster's settings in the store, in force when changed; the firewall switch; wipe | released |
|
||||
| v0.13.0 | — | The whole system stopped and started from one command; Qwen3Guard judges, on a GPU and CPUs together | released |
|
||||
| v0.14.0 | — | Knowledge bases: documents from people's own machines, kept apart per tenant, answered from | released (v0.14.1: adapters, the tree over the API, the firewall on every machine; v0.14.2: pictures from ComfyUI; v0.14.3: 3D models as jobs; v0.14.4: any ComfyUI workflow as a job, presets, licences; v0.14.5: picture and 3D models that add themselves; v0.14.6: setup that explains every part and lets it be chosen, the gateway on all three HA machines; v0.14.7: the judge sized for its machine; v0.14.8: sign-in with Active Directory, several keys for one outside provider; v0.14.9: sign-in set up from the console, a node's console that shows the machine) |
|
||||
| v0.15.0 | — | The company's MCP tools and VS Code through the gateway; guardrails that ask outside services and put redacted values back; answers to similar questions | released |
|
||||
| v0.15.0 | — | The company's MCP tools and VS Code through the gateway; guardrails that ask outside services and put redacted values back; answers to similar questions | released (v0.15.1: Strata, a 125B model of experts on one card; a second mirror kept in step; pickled weights refused; a served model's context changed; v0.15.2: a served model moved to more cards) |
|
||||
| v1.0.0 | 5 | What phase 5 has left: precise KV routing proven on vLLM, the block cost picker, defragmentation across NVLink domains, prefill/decode disaggregation | planned |
|
||||
|
||||
Until v1.0.0 the minor version is the breaking one: configuration formats and Go
|
||||
APIs change between minors without a deprecation period.
|
||||
|
||||
## v0.15.2 — 2026-10-29
|
||||
|
||||
A served model moved to more cards, or fewer, while it is served.
|
||||
|
||||
**Upgrading.**
|
||||
- No migrations, and nothing changes until the command is used. The control
|
||||
planes carry it; the nodes need nothing new.
|
||||
|
||||
### Added
|
||||
|
||||
- **A served model moved to more cards, or fewer.** `llmgw ctl model gpus
|
||||
MODEL N [--profile P] [--in-place]`, `POST /api/v1/models/{name}/gpus`, and
|
||||
GPUs… in the console's model menu. How many cards of one node each replica
|
||||
uses: tensor parallel for vLLM, SGLang and TensorRT-LLM, checked against
|
||||
the model's KV heads, and the layers split for llama.cpp. Refused for a
|
||||
count no node has, or whose cards would not hold the model and its
|
||||
context, with the figures; its replicas are then started again on it, one
|
||||
by one, beside the old — or, with `--in-place` when there is no room
|
||||
beside them, each stopped first.
|
||||
|
||||
## v0.15.1 — 2026-10-29
|
||||
|
||||
A 125B model on one card, a second mirror, and weights that run code
|
||||
refused. Strata runs as an engine: a mixture-of-experts model with its
|
||||
experts in RAM, served from one 24 GB card. A second model mirror keeps in
|
||||
step with the first, so an HA install has no single machine left. Pickled
|
||||
weights are refused unless a model says otherwise. A served model's context
|
||||
can be changed in place. The explorer says how fast a model would answer
|
||||
here before it is fetched, and a card in the wrong PCIe slot is said so.
|
||||
|
||||
**Upgrading.**
|
||||
- No migrations. Upgrade the control planes before the nodes, for the
|
||||
pickle refusal below.
|
||||
- A model whose weights are in a pickle format (`.bin`, `.pt`, `.pth`,
|
||||
`.ckpt`, `.pkl`) stops being placed: its node refuses the source. Give it
|
||||
its safetensors or GGUF, or add it again with `allow_pickle: true`
|
||||
(`llmgw ctl model add -f`) if its files were checked. Look at the
|
||||
models' sources before upgrading the nodes: one ending in
|
||||
`.ckpt` or `.bin`, or a mirror version holding `pytorch_model.bin`. Upgrade the control planes first: a node on
|
||||
this version refuses those weights until its control plane can say
|
||||
`allow_pickle`.
|
||||
- Strata is opt-in: nothing installs it until `engines` or an
|
||||
`engine_overrides` entry names a version (`0.2.7`).
|
||||
|
||||
### Added
|
||||
|
||||
- **Strata as an engine: a 125B model of experts on one card.** Strata
|
||||
(github.com/Niko1221/Strata) keeps a model's experts in RAM and the ones
|
||||
used most in VRAM. On an RTX 6000 24 GB with 251 GB of RAM it serves
|
||||
Qwen3.8-Flash-Next 125B (IQ2_XS) at 70–87 tokens a second to one client;
|
||||
more clients share about that, its slots buying fairness rather than
|
||||
throughput on one card. It runs a few families only — Qwen3.8-Flash-Next
|
||||
and its variants, Qwen3.6-35B-A3B, DeepSeek V4.1-Flash — each named in a
|
||||
profile's args with `--family` and `--size`
|
||||
(`--family qwen --size IQ2_XS --context 32768`). The source is the
|
||||
directory of the model's GGUF files, or the repository, which Strata then
|
||||
downloads itself, pinned. The agent installs it on an NVIDIA node of
|
||||
x86-64: its source pinned by digest, in a Python environment, and its
|
||||
release's engine checked against the digest GitHub publishes — the CUDA 13
|
||||
build for Turing and newer on driver 580, the CUDA 12.9 build for Pascal
|
||||
to Ada, a V100 included, on driver 575. A replica prepares its model the
|
||||
first time (minutes; its draft layer is a 5 GB download) and loads in
|
||||
about a minute and a half after. One card of one node per replica, for
|
||||
now; not in an offline bundle, its packages coming from PyPI.
|
||||
|
||||
- **A served model's context, changed.** `llmgw ctl model context MODEL
|
||||
TOKENS [--slots N]`, `POST /api/v1/models/{name}/context`, and Context
|
||||
length… in the console's model menu. The catalog, the engine's arguments and
|
||||
the memory change together, checked against what the model was trained for
|
||||
and what its cards hold; its replicas are then started again one by one. The
|
||||
models API shows `context_length` and `max_output`.
|
||||
|
||||
- **A second model mirror, kept in step.** In a three-machine install the setup screens put one on a database machine
|
||||
(`--second-mirror 1|2`). An install with one mirror gets one with `sudo llmgw mirror add-peer user@host`.
|
||||
- **In step:** each mirror reads the other every thirty seconds. It copies what it lacks, checksum checked, and
|
||||
removes what the other removed, even if it was down at the time (`llmgw mirror serve --peer`).
|
||||
- **Control planes:** they name both (`mirror.peers`) and use whichever answers.
|
||||
- **Nodes:** while one mirror is down, a node is given a model's file on the other.
|
||||
|
||||
This takes away the last single point of an HA install: with the first machine off, models are still added and
|
||||
fetched.
|
||||
|
||||
### Fixed
|
||||
|
||||
- **An engine that computes in a process of its own group was not its
|
||||
replica's.** A GPU replica read as running on the CPU, and the memory it
|
||||
held on its card as somebody else's, which placement then kept free of
|
||||
other work. A process the replica started, at any remove, is now the
|
||||
replica's. And what a replica computes on, once sent, was never sent
|
||||
again: a wrong reading kept from before an agent's restart stayed wrong.
|
||||
It is sent again when it changes, and read again when a GPU replica was
|
||||
kept as the CPU.
|
||||
|
||||
- **VS Code's sign-in lost its place through the sign-in page**, so the
|
||||
prompt to confirm came every time; a link to the same gateway by plain
|
||||
http could take an https setting down to http; and a new build of the
|
||||
extension was not installed over an old one (the command now says
|
||||
`--force`).
|
||||
- **The install guide's recordings were not checked for secrets**, only
|
||||
the page; every frame is now.
|
||||
- **MCP server fields could not be emptied**, an address could carry a
|
||||
password, the cluster's own server could not be changed from the
|
||||
console, and a grant to a key that does not exist was said done.
|
||||
- **A guardrail service two policies named was asked as the first said**,
|
||||
not the stricter; a `null` answer read as safe; and with `restore` the
|
||||
request log kept the answer with the redacted values back in it.
|
||||
- **What an engine downloads by a Hugging Face id was not held to the
|
||||
pickle refusal** — an adapter by its id, a TensorRT-LLM model's; a model
|
||||
with either now needs `allow_pickle`.
|
||||
- **`server init` run again on a machine already set up was refused its
|
||||
memory**, as on an 8 GB install fixing its certificate with `--host`, and
|
||||
judges that would not be installed (no mirror) were counted.
|
||||
- **Removing an MCP server left its sealed credential in the store**, and
|
||||
replacing a credential under another header kept the old one sent. A
|
||||
server's secrets now go when it no longer reads them.
|
||||
- **A cloned repository could point the VS Code extension at another
|
||||
gateway.** `llmgw.gatewayUrl` is a machine setting now, never taken from a
|
||||
workspace.
|
||||
- **A node's memory bandwidth could be measured in its cache** on a busy
|
||||
node, and was 0 after a reconnect until the next heartbeat. The buffer is
|
||||
at least eight times the processor's last cache.
|
||||
- **A project kept in the cluster could have its questions embedded outside
|
||||
it.** The similar-question cache now asks the embedding model as the
|
||||
question's project would, external fallbacks refused, and skips a model
|
||||
that is itself outside. Its comparing no longer holds the cache's lock,
|
||||
which every project's lookups waited on.
|
||||
- **Restored values were left as placeholders in streamed tool calls**, in
|
||||
completions and in reasoning, and a value with a quote in a whole answer's
|
||||
tool-call arguments broke their JSON.
|
||||
- **An outside guardrail service read what `redact` had taken out**, and a
|
||||
streamed answer was sent to it twice. The docs now say a service stops a
|
||||
streamed answer only under `stream: buffer_all`.
|
||||
- **A replica found slow once could stay unused.** Its time to first token
|
||||
is forgotten after five minutes without an answer, so it is tried again.
|
||||
- **Weights that run code got past the refusal in four ways.** A mirror
|
||||
version already on a node was used without being looked at again; a
|
||||
`?query` or `#fragment` on an `hf:` source hid the file's name; `.pt2` was
|
||||
not on the list; and a profile could write `--load_format`,
|
||||
`--trust_remote_code` or an abbreviation such as `--trust-remote`, which
|
||||
vLLM and argparse read as the flag, or give SGLang `--load-format`. All
|
||||
refused now. TripoSR's preset says `allow_pickle`, without which it could
|
||||
not be added, and a shared ComfyUI environment takes its first member's.
|
||||
- **A temperature written as a string was refused.** Since the GenAI traces
|
||||
read it, `"temperature": "0.7"` got a 400 before reaching the engine; it
|
||||
is read loosely again, for the trace only.
|
||||
- **A reused answer could carry another caller's personal data.** With
|
||||
`restore: true` and the answer cache on, two callers' requests that
|
||||
redacted to the same text shared one answer — with the first caller's
|
||||
address put back in it. Such requests are now never answered from the
|
||||
cache nor kept in it.
|
||||
- **Granting a server's tools could take others away; revoking could close
|
||||
a level to every tool.** A level with no tool list sees what the levels
|
||||
above grant; giving it its first list made it see only that list. And a
|
||||
level left with an empty list was closed to every tool, which the screens
|
||||
did not show. Now a first list starts from the nearest list above, taking
|
||||
a server away leaves the rest or goes back to no list, the screens say
|
||||
when a level above keeps granted tools out of reach and which levels are
|
||||
closed, and the console reads and writes the level at once rather than
|
||||
from what it loaded.
|
||||
- **Similar questions about different pictures got each other's answer.**
|
||||
A question with a picture, a sound or a file in it is reused only when it
|
||||
is the same request.
|
||||
|
||||
### Added
|
||||
|
||||
- **A card in the wrong PCIe slot is said so.** Each agent reads how every
|
||||
card is linked — the fastest generation card and slot allow together, and
|
||||
the lanes in use — against the card's own most, from the kernel (the
|
||||
driver's figure is what card and slot allow together); a card on fewer lanes or
|
||||
an older generation is a node warning, with what it carries and what it
|
||||
could, on the dashboard, the console and in `llmgw doctor`.
|
||||
- **And on a web page, `/admin/mcp`.** For administrators: the same list,
|
||||
adding a server with its credential sealed, testing it and listing its
|
||||
tools, granting them in the tree beside that level's other grants and
|
||||
limits, revoking, changing, replacing the credential, turning it off,
|
||||
removing; and the line that points an agent at this gateway's `/mcp`.
|
||||
In English and Turkish; every change in the audit log, the secret not.
|
||||
- **The company's MCP servers in the console.** A screen of its own: each
|
||||
server with its state, who has its tools, its rate, the names of the
|
||||
credentials it is sent and its address; add one with a bearer token or
|
||||
a header kept sealed, test it and list its tools, grant them to a team,
|
||||
project, organization, person or key beside that level's other grants,
|
||||
revoke, change, replace the credential, turn off, remove.
|
||||
- **And on the cards, as this cluster would lay it out.** The explorer's
|
||||
**GPU tok/s** and a repository's files give one reply's speed for each
|
||||
layout the planner offers here — one card, tensor parallel across a
|
||||
node's cards, llama.cpp's layer split, a pipeline across nodes on a
|
||||
fabric — the fastest of them, with the layout in the tooltip. From each
|
||||
card's memory bandwidth by its model name, NVLink or PCIe between a
|
||||
node's cards, and InfiniBand or RoCE between nodes; a cluster of mixed
|
||||
classes is judged class by class. A model of experts reads 1.6 times its
|
||||
active share, as Qwen3-Coder-30B-A3B did on cuda01's RTX 6000 (144
|
||||
tok/s; qwen3-vl-8b Q8: 62).
|
||||
- **How fast a model would answer on a CPU, before it is fetched.** Each
|
||||
agent measures how fast its node's memory reads when it starts (`Host`
|
||||
gets `mem_bandwidth_gbs`, the node's view `mem_bandwidth_gbs`); the
|
||||
explorer's **CPU tok/s** column and a repository's GGUF files say about
|
||||
how many tokens a second one reply gets on the fastest node that holds
|
||||
the model — the bandwidth over the bytes a token reads, its active part
|
||||
for a model of experts. Calibrated on the PoC against llama.cpp's own
|
||||
timings: within 15% from 1.5B to 14B.
|
||||
- **The explorer lists what this cluster can run.** Judged by the nodes
|
||||
as they are now — every card's memory and compute capability, every
|
||||
node's RAM — through the same planner as serving a model, rather than a
|
||||
size picked from a list. A model too big for any node, in FP8 without a
|
||||
card of 8.9, in MLX, or in neither GGUF nor safetensors, is left out and
|
||||
counted; the page says what it judged by. With cards, safetensors repositories are
|
||||
listed beside GGUF. On a repository's page a file too big for the
|
||||
cluster is shown with what it needs and no Fetch button.
|
||||
- **Weights that run code are refused.** Python pickles — `.bin`, `.pt`,
|
||||
`.pth`, `.ckpt` — run whatever they hold when loaded, which is how the
|
||||
malicious models on Hugging Face work. The node refuses such a source,
|
||||
or a mirror version with one in it, before fetching it; adding such a
|
||||
model is refused; vLLM and SGLang downloading a model by its id read its
|
||||
safetensors only; the mirror page marks the file. `allow_pickle: true`
|
||||
on a model lets it through on purpose. The mirror's explorer no longer
|
||||
lists a repository with nothing a node loads (no GGUF, no safetensors)
|
||||
while "runs here" is ticked, and a repository's pickle files are not
|
||||
offered, only counted. BACKLOG M lists what is next:
|
||||
checksums against Hugging Face everywhere, its own scan, trusted
|
||||
publishers, engines without the internet.
|
||||
|
||||
## v0.15.0 — 2026-10-25
|
||||
|
||||
The company's tools and editors through the gateway, and guardrails that
|
||||
|
||||
@@ -41,7 +41,7 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
| v0.12.0 | — | The cluster's settings in the store, in force when changed; the firewall switch; wipe | released |
|
||||
| v0.13.0 | — | The whole system stopped and started from one command; Qwen3Guard judges, on a GPU and CPUs together | released |
|
||||
| v0.14.0 | — | Knowledge bases: documents from people's own machines, kept apart per tenant, answered from | released (v0.14.1: adapters, the tree over the API, the firewall on every machine; v0.14.2: pictures from ComfyUI; v0.14.3: 3D models as jobs; v0.14.4: any ComfyUI workflow as a job, presets, licences; v0.14.5: picture and 3D models that add themselves; v0.14.6: setup that explains every part and lets it be chosen, the gateway on all three HA machines; v0.14.7: the judge sized for its machine; v0.14.8: sign-in with Active Directory, several keys for one outside provider; v0.14.9: sign-in set up from the console, a node's console that shows the machine) |
|
||||
| v0.15.0 | — | The company's MCP tools and VS Code through the gateway; guardrails that ask outside services and put redacted values back; answers to similar questions | released |
|
||||
| v0.15.0 | — | The company's MCP tools and VS Code through the gateway; guardrails that ask outside services and put redacted values back; answers to similar questions | released (v0.15.1: Strata, a 125B model of experts on one card; a second mirror kept in step; pickled weights refused; a served model's context changed; v0.15.2: a served model moved to more cards) |
|
||||
| v1.0.0 | 5 | What phase 5 has left: precise KV routing proven on vLLM, the block cost picker, defragmentation across NVLink domains, prefill/decode disaggregation | planned |
|
||||
|
||||
As of v0.6.0 a cluster runs unattended: agents report their hardware, the
|
||||
|
||||
@@ -1206,9 +1206,12 @@ type Host struct {
|
||||
CpuFlags []string `protobuf:"bytes,5,rep,name=cpu_flags,json=cpuFlags,proto3" json:"cpu_flags,omitempty"`
|
||||
// nvidia_driver is the driver version, empty when there is none. A card
|
||||
// without a driver is not a GPU anything can use.
|
||||
NvidiaDriver string `protobuf:"bytes,6,opt,name=nvidia_driver,json=nvidiaDriver,proto3" json:"nvidia_driver,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
NvidiaDriver string `protobuf:"bytes,6,opt,name=nvidia_driver,json=nvidiaDriver,proto3" json:"nvidia_driver,omitempty"`
|
||||
// mem_bandwidth_gbs is how fast the memory reads, in GB/s, measured when
|
||||
// the agent started: what decides a model's speed on the CPU. 0 unknown.
|
||||
MemBandwidthGbs float64 `protobuf:"fixed64,7,opt,name=mem_bandwidth_gbs,json=memBandwidthGbs,proto3" json:"mem_bandwidth_gbs,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Host) Reset() {
|
||||
@@ -1283,6 +1286,13 @@ func (x *Host) GetNvidiaDriver() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *Host) GetMemBandwidthGbs() float64 {
|
||||
if x != nil {
|
||||
return x.MemBandwidthGbs
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// Per filesystem rather than per configured directory: free space belongs to
|
||||
@@ -1579,7 +1589,14 @@ type GPU struct {
|
||||
// is then 1000 + 100 × card + device, uuid the instance's (MIG-…), and
|
||||
// memory the instance's share; utilization_percent stays zero, because no
|
||||
// driver measures it per instance.
|
||||
Mig *MIGDevice `protobuf:"bytes,22,opt,name=mig,proto3" json:"mig,omitempty"`
|
||||
Mig *MIGDevice `protobuf:"bytes,22,opt,name=mig,proto3" json:"mig,omitempty"`
|
||||
// How the card is linked to its host: the fastest generation card and
|
||||
// slot allow together, the card's own fastest, and the lanes in use and
|
||||
// the card's most. Zero when the driver did not say.
|
||||
PcieGen int32 `protobuf:"varint,23,opt,name=pcie_gen,json=pcieGen,proto3" json:"pcie_gen,omitempty"`
|
||||
PcieGpuGen int32 `protobuf:"varint,24,opt,name=pcie_gpu_gen,json=pcieGpuGen,proto3" json:"pcie_gpu_gen,omitempty"`
|
||||
PcieWidth int32 `protobuf:"varint,25,opt,name=pcie_width,json=pcieWidth,proto3" json:"pcie_width,omitempty"`
|
||||
PcieWidthMax int32 `protobuf:"varint,26,opt,name=pcie_width_max,json=pcieWidthMax,proto3" json:"pcie_width_max,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
@@ -1768,6 +1785,34 @@ func (x *GPU) GetMig() *MIGDevice {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *GPU) GetPcieGen() int32 {
|
||||
if x != nil {
|
||||
return x.PcieGen
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *GPU) GetPcieGpuGen() int32 {
|
||||
if x != nil {
|
||||
return x.PcieGpuGen
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *GPU) GetPcieWidth() int32 {
|
||||
if x != nil {
|
||||
return x.PcieWidth
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *GPU) GetPcieWidthMax() int32 {
|
||||
if x != nil {
|
||||
return x.PcieWidthMax
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// MIGDevice is which card a MIG instance is carved from, and how.
|
||||
type MIGDevice struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
@@ -2943,8 +2988,12 @@ type Heartbeat struct {
|
||||
// is not handed out. Empty from an agent that does not measure it, or from
|
||||
// one that could not this time.
|
||||
GpuForeignMemoryMb []int32 `protobuf:"varint,19,rep,packed,name=gpu_foreign_memory_mb,json=gpuForeignMemoryMb,proto3" json:"gpu_foreign_memory_mb,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
// mem_bandwidth_gbs is how fast the node's memory reads, in GB/s, once
|
||||
// the agent has measured it — in the background after it starts, so a
|
||||
// registration usually has none yet. 0 means not measured.
|
||||
MemBandwidthGbs float64 `protobuf:"fixed64,20,opt,name=mem_bandwidth_gbs,json=memBandwidthGbs,proto3" json:"mem_bandwidth_gbs,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Heartbeat) Reset() {
|
||||
@@ -3110,6 +3159,13 @@ func (x *Heartbeat) GetGpuForeignMemoryMb() []int32 {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetMemBandwidthGbs() float64 {
|
||||
if x != nil {
|
||||
return x.MemBandwidthGbs
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// Provision is the software a node should have: which engines, at which
|
||||
// versions. Nothing more.
|
||||
//
|
||||
@@ -3790,7 +3846,10 @@ type EngineProfile struct {
|
||||
Adapters []*LoraAdapter `protobuf:"bytes,12,rep,name=adapters,proto3" json:"adapters,omitempty"`
|
||||
// kv_events has a vLLM engine publish its KV cache events to the agent, for
|
||||
// the router's precise index.
|
||||
KvEvents bool `protobuf:"varint,13,opt,name=kv_events,json=kvEvents,proto3" json:"kv_events,omitempty"`
|
||||
KvEvents bool `protobuf:"varint,13,opt,name=kv_events,json=kvEvents,proto3" json:"kv_events,omitempty"`
|
||||
// allow_pickle lets the node fetch weights in a pickle format, which run
|
||||
// code when loaded: the catalog's allow_pickle, set on purpose.
|
||||
AllowPickle bool `protobuf:"varint,14,opt,name=allow_pickle,json=allowPickle,proto3" json:"allow_pickle,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
@@ -3916,6 +3975,13 @@ func (x *EngineProfile) GetKvEvents() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (x *EngineProfile) GetAllowPickle() bool {
|
||||
if x != nil {
|
||||
return x.AllowPickle
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
//
|
||||
// The paths are computed by the control plane from the catalog, but the agent
|
||||
@@ -5180,7 +5246,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\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\"\xb1\x01\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xdd\x01\n" +
|
||||
"\x04Host\x12\x13\n" +
|
||||
"\x05os_id\x18\x01 \x01(\tR\x04osId\x12\x1d\n" +
|
||||
"\n" +
|
||||
@@ -5189,7 +5255,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"osCodename\x12\x12\n" +
|
||||
"\x04arch\x18\x04 \x01(\tR\x04arch\x12\x1b\n" +
|
||||
"\tcpu_flags\x18\x05 \x03(\tR\bcpuFlags\x12#\n" +
|
||||
"\rnvidia_driver\x18\x06 \x01(\tR\fnvidiaDriver\"u\n" +
|
||||
"\rnvidia_driver\x18\x06 \x01(\tR\fnvidiaDriver\x12*\n" +
|
||||
"\x11mem_bandwidth_gbs\x18\a \x01(\x01R\x0fmemBandwidthGbs\"u\n" +
|
||||
"\aStorage\x12\x12\n" +
|
||||
"\x04path\x18\x01 \x01(\tR\x04path\x12\x1d\n" +
|
||||
"\n" +
|
||||
@@ -5219,7 +5286,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"leaderPort\x12\x1d\n" +
|
||||
"\n" +
|
||||
"profile_fp\x18\x0f \x01(\tR\tprofileFp\x12+\n" +
|
||||
"\acompute\x18\x10 \x01(\v2\x11.llmgw.v1.ComputeR\acompute\"\xca\x06\n" +
|
||||
"\acompute\x18\x10 \x01(\v2\x11.llmgw.v1.ComputeR\acompute\"\xcc\a\n" +
|
||||
"\x03GPU\x12\x14\n" +
|
||||
"\x05index\x18\x01 \x01(\x05R\x05index\x12\x12\n" +
|
||||
"\x04uuid\x18\x02 \x01(\tR\x04uuid\x12\x12\n" +
|
||||
@@ -5245,7 +5312,13 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\fcluster_uuid\x18\x13 \x01(\tR\vclusterUuid\x12\x1b\n" +
|
||||
"\tclique_id\x18\x14 \x01(\tR\bcliqueId\x12)\n" +
|
||||
"\x10fabric_bandwidth\x18\x15 \x01(\tR\x0ffabricBandwidth\x12%\n" +
|
||||
"\x03mig\x18\x16 \x01(\v2\x13.llmgw.v1.MIGDeviceR\x03mig\x1a;\n" +
|
||||
"\x03mig\x18\x16 \x01(\v2\x13.llmgw.v1.MIGDeviceR\x03mig\x12\x19\n" +
|
||||
"\bpcie_gen\x18\x17 \x01(\x05R\apcieGen\x12 \n" +
|
||||
"\fpcie_gpu_gen\x18\x18 \x01(\x05R\n" +
|
||||
"pcieGpuGen\x12\x1d\n" +
|
||||
"\n" +
|
||||
"pcie_width\x18\x19 \x01(\x05R\tpcieWidth\x12$\n" +
|
||||
"\x0epcie_width_max\x18\x1a \x01(\x05R\fpcieWidthMax\x1a;\n" +
|
||||
"\rNicPathsEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\x84\x02\n" +
|
||||
@@ -5361,7 +5434,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\x06status\x18\x03 \x01(\tR\x06status\x12!\n" +
|
||||
"\fcluster_uuid\x18\x04 \x01(\tR\vclusterUuid\x12\x1b\n" +
|
||||
"\tclique_id\x18\x05 \x01(\tR\bcliqueId\x12\x1c\n" +
|
||||
"\tbandwidth\x18\x06 \x01(\tR\tbandwidth\"\xa9\a\n" +
|
||||
"\tbandwidth\x18\x06 \x01(\tR\tbandwidth\"\xd5\a\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" +
|
||||
@@ -5383,7 +5456,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"gpu_fabric\x18\x10 \x03(\v2\x18.llmgw.v1.GPUFabricStateR\tgpuFabric\x12\"\n" +
|
||||
"\x04imex\x18\x11 \x01(\v2\x0e.llmgw.v1.IMEXR\x04imex\x12\x1b\n" +
|
||||
"\tcpu_count\x18\x12 \x01(\x05R\bcpuCount\x121\n" +
|
||||
"\x15gpu_foreign_memory_mb\x18\x13 \x03(\x05R\x12gpuForeignMemoryMb\x1aA\n" +
|
||||
"\x15gpu_foreign_memory_mb\x18\x13 \x03(\x05R\x12gpuForeignMemoryMb\x12*\n" +
|
||||
"\x11mem_bandwidth_gbs\x18\x14 \x01(\x01R\x0fmemBandwidthGbs\x1aA\n" +
|
||||
"\x13EngineVersionsEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\";\n" +
|
||||
@@ -5462,7 +5536,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"profile_fp\x18\x0f \x01(\tR\tprofileFp\x12\x17\n" +
|
||||
"\aapi_key\x18\x10 \x01(\tR\x06apiKey\x12\x1f\n" +
|
||||
"\vcpu_threads\x18\x11 \x01(\x05R\n" +
|
||||
"cpuThreads\"\xcd\x03\n" +
|
||||
"cpuThreads\"\xf0\x03\n" +
|
||||
"\rEngineProfile\x12\x12\n" +
|
||||
"\x04name\x18\x01 \x01(\tR\x04name\x12\x16\n" +
|
||||
"\x06engine\x18\x02 \x01(\tR\x06engine\x12\x1b\n" +
|
||||
@@ -5478,7 +5552,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
" \x01(\tR\tmodelPath\x12*\n" +
|
||||
"\x11served_model_name\x18\v \x01(\tR\x0fservedModelName\x121\n" +
|
||||
"\badapters\x18\f \x03(\v2\x15.llmgw.v1.LoraAdapterR\badapters\x12\x1b\n" +
|
||||
"\tkv_events\x18\r \x01(\bR\bkvEvents\x1a6\n" +
|
||||
"\tkv_events\x18\r \x01(\bR\bkvEvents\x12!\n" +
|
||||
"\fallow_pickle\x18\x0e \x01(\bR\vallowPickle\x1a6\n" +
|
||||
"\bEnvEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"%\n" +
|
||||
|
||||
@@ -238,6 +238,9 @@ message Host {
|
||||
// nvidia_driver is the driver version, empty when there is none. A card
|
||||
// without a driver is not a GPU anything can use.
|
||||
string nvidia_driver = 6;
|
||||
// mem_bandwidth_gbs is how fast the memory reads, in GB/s, measured when
|
||||
// the agent started: what decides a model's speed on the CPU. 0 unknown.
|
||||
double mem_bandwidth_gbs = 7;
|
||||
}
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
@@ -334,6 +337,13 @@ message GPU {
|
||||
// memory the instance's share; utilization_percent stays zero, because no
|
||||
// driver measures it per instance.
|
||||
MIGDevice mig = 22;
|
||||
// How the card is linked to its host: the fastest generation card and
|
||||
// slot allow together, the card's own fastest, and the lanes in use and
|
||||
// the card's most. Zero when the driver did not say.
|
||||
int32 pcie_gen = 23;
|
||||
int32 pcie_gpu_gen = 24;
|
||||
int32 pcie_width = 25;
|
||||
int32 pcie_width_max = 26;
|
||||
}
|
||||
|
||||
// MIGDevice is which card a MIG instance is carved from, and how.
|
||||
@@ -588,6 +598,11 @@ message Heartbeat {
|
||||
// is not handed out. Empty from an agent that does not measure it, or from
|
||||
// one that could not this time.
|
||||
repeated int32 gpu_foreign_memory_mb = 19;
|
||||
|
||||
// mem_bandwidth_gbs is how fast the node's memory reads, in GB/s, once
|
||||
// the agent has measured it — in the background after it starts, so a
|
||||
// registration usually has none yet. 0 means not measured.
|
||||
double mem_bandwidth_gbs = 20;
|
||||
}
|
||||
|
||||
// Provision is the software a node should have: which engines, at which
|
||||
@@ -749,6 +764,9 @@ message EngineProfile {
|
||||
// kv_events has a vLLM engine publish its KV cache events to the agent, for
|
||||
// the router's precise index.
|
||||
bool kv_events = 13;
|
||||
// allow_pickle lets the node fetch weights in a pickle format, which run
|
||||
// code when loaded: the catalog's allow_pickle, set on purpose.
|
||||
bool allow_pickle = 14;
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
|
||||
@@ -2719,6 +2719,57 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/models/{name}/context": {
|
||||
"post": {
|
||||
"operationId": "postModelsNameContext",
|
||||
"summary": "Change a served model's context: the catalog, its engines' arguments and memory, then its replicas restarted",
|
||||
"tags": [
|
||||
"models"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "name",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "The model's name in the catalog.",
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/modelserve.ContextRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/modelserve.ContextChange"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/models/{name}/eval": {
|
||||
"post": {
|
||||
"operationId": "postModelsNameEval",
|
||||
@@ -2913,6 +2964,58 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/models/{name}/gpus": {
|
||||
"post": {
|
||||
"operationId": "postModelsNameGpus",
|
||||
"summary": "Spread a served model over more cards or fewer: its profile, checked against its KV heads, its cards' memory and the nodes, then its replicas restarted",
|
||||
"description": "Tensor parallel for vLLM, SGLang and TensorRT-LLM, a layer split for llama.cpp, on one node. With no room beside the running replicas, in_place stops each first.",
|
||||
"tags": [
|
||||
"models"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "name",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "The model's name in the catalog.",
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/modelserve.GPUsRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/modelserve.GPUsChange"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/models/{name}/headroom": {
|
||||
"get": {
|
||||
"operationId": "getModelsNameHeadroom",
|
||||
@@ -7550,6 +7653,13 @@
|
||||
"rate_per_minute": {
|
||||
"type": "integer",
|
||||
"description": "RatePerMinute bounds each caller's calls to its tools; a negative number takes the bound away."
|
||||
},
|
||||
"clear": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "Clear empties what an empty field above would leave as it is: \"description\", \"timeout\" (back to the default), \"rate\"."
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -7742,6 +7852,13 @@
|
||||
"name": {
|
||||
"type": "string"
|
||||
},
|
||||
"context_length": {
|
||||
"type": "integer",
|
||||
"description": "ContextLength is what one request may use, MaxOutput what it may answer; 0 when the catalog does not say."
|
||||
},
|
||||
"max_output": {
|
||||
"type": "integer"
|
||||
},
|
||||
"canary": {
|
||||
"$ref": "#/components/schemas/catalog.Canary",
|
||||
"description": "Canary is a share of its requests sent to a candidate; CanaryOf is set on the candidate, naming the model it is a canary of."
|
||||
@@ -8011,6 +8128,10 @@
|
||||
"type": "integer",
|
||||
"description": "CPUCount is the machine's CPUs, 0 when the agent did not report them."
|
||||
},
|
||||
"mem_bandwidth_gbs": {
|
||||
"type": "number",
|
||||
"description": "MemBandwidthGBs is how fast its memory reads, in GB/s, measured when its agent started; 0 when not reported. A model's speed on the CPU."
|
||||
},
|
||||
"engines": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
@@ -9930,6 +10051,10 @@
|
||||
"type": "string",
|
||||
"description": "Source is the weights path or a HuggingFace id."
|
||||
},
|
||||
"allow_pickle": {
|
||||
"type": "boolean",
|
||||
"description": "AllowPickle lets the nodes fetch this model's weights, or its adapters', in a pickle format (.bin, .pt, .pth, .ckpt), which run code as they are loaded. Off, such a source is refused when the model is added and by the node before it fetches anything. For a model that has no safetensors or GGUF form and whose files were checked."
|
||||
},
|
||||
"profiles": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
@@ -10309,6 +10434,10 @@
|
||||
"type": "string",
|
||||
"description": "HFRepo is the repository its node loads the weights by name from."
|
||||
},
|
||||
"allow_pickle": {
|
||||
"type": "boolean",
|
||||
"description": "AllowPickle is the member model's allow_pickle."
|
||||
},
|
||||
"workflow": {
|
||||
"type": "string",
|
||||
"description": "Workflow and Checkpoint are what a request for it is sent."
|
||||
@@ -10986,6 +11115,137 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"modelserve.ContextChange": {
|
||||
"type": "object",
|
||||
"description": "ContextChange is what was changed, and the replicas started again on it.",
|
||||
"properties": {
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"before": {
|
||||
"type": "integer"
|
||||
},
|
||||
"after": {
|
||||
"type": "integer"
|
||||
},
|
||||
"max_output": {
|
||||
"type": "integer"
|
||||
},
|
||||
"trained": {
|
||||
"type": "integer"
|
||||
},
|
||||
"profiles": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/modelserve.ContextProfile"
|
||||
}
|
||||
},
|
||||
"restarted": {
|
||||
"type": "integer"
|
||||
},
|
||||
"note": {
|
||||
"type": "string",
|
||||
"description": "Note is what could not be checked, said once."
|
||||
}
|
||||
}
|
||||
},
|
||||
"modelserve.ContextProfile": {
|
||||
"type": "object",
|
||||
"description": "ContextProfile is what one profile needs before and after.",
|
||||
"properties": {
|
||||
"profile": {
|
||||
"type": "string"
|
||||
},
|
||||
"engine": {
|
||||
"type": "string"
|
||||
},
|
||||
"before": {
|
||||
"type": "string"
|
||||
},
|
||||
"after": {
|
||||
"type": "string"
|
||||
},
|
||||
"args": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"modelserve.ContextRequest": {
|
||||
"type": "object",
|
||||
"description": "ContextRequest is the context to change to: tokens one request may use, and for llama.cpp how many at once (its pool is the two multiplied).",
|
||||
"properties": {
|
||||
"tokens": {
|
||||
"type": "integer"
|
||||
},
|
||||
"slots": {
|
||||
"type": "integer"
|
||||
}
|
||||
}
|
||||
},
|
||||
"modelserve.GPUsChange": {
|
||||
"type": "object",
|
||||
"description": "GPUsChange is what was changed, and the replicas started again on it.",
|
||||
"properties": {
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"profile": {
|
||||
"type": "string"
|
||||
},
|
||||
"engine": {
|
||||
"type": "string"
|
||||
},
|
||||
"gpu_class": {
|
||||
"type": "string"
|
||||
},
|
||||
"before": {
|
||||
"type": "integer"
|
||||
},
|
||||
"after": {
|
||||
"type": "integer"
|
||||
},
|
||||
"args": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"per_card": {
|
||||
"type": "string",
|
||||
"description": "PerCard is what each card holds after, when it could be worked out."
|
||||
},
|
||||
"restarted": {
|
||||
"type": "integer"
|
||||
},
|
||||
"in_place": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"note": {
|
||||
"type": "string",
|
||||
"description": "Note is what could not be checked, said once."
|
||||
}
|
||||
}
|
||||
},
|
||||
"modelserve.GPUsRequest": {
|
||||
"type": "object",
|
||||
"description": "GPUsRequest is how many cards a replica is to use.",
|
||||
"properties": {
|
||||
"gpus": {
|
||||
"type": "integer"
|
||||
},
|
||||
"profile": {
|
||||
"type": "string",
|
||||
"description": "Profile is the one to change, when the model has more than one profile on cards."
|
||||
},
|
||||
"in_place": {
|
||||
"type": "boolean",
|
||||
"description": "InPlace stops each replica before its replacement starts: for when no node has the cards free beside it. It does not answer while the new one loads."
|
||||
}
|
||||
}
|
||||
},
|
||||
"modelserve.Job": {
|
||||
"type": "object",
|
||||
"description": "Job is one model being served, and where it stands.",
|
||||
|
||||
@@ -173,6 +173,23 @@ search…** as an administrator, or **Delete…** it. People
|
||||
fill collections from their own machines with `llmgw-kb`
|
||||
([knowledge.md](knowledge.md)). Off on a SQLite store, and says why.
|
||||
|
||||
**MCP servers** lists the company's own MCP servers that agents reach through
|
||||
the gateway's `/mcp`, their tools named `server.tool`: whether each is on,
|
||||
who has been granted its tools, its rate per caller, the names of the
|
||||
credentials it is sent (never their values) and its address. **A** adds one:
|
||||
its name, which is its tools' prefix, its address, and a bearer token or
|
||||
another header, kept sealed in the store. On a row: **Test, and list its
|
||||
tools**, with its credential, as a call would; **Grant its tools to…** a
|
||||
team, project, organization, person or key — all of them, `jira.*`, or some
|
||||
by name — beside whatever else that level is granted, and saying when a
|
||||
level above keeps them out of reach; **Revoke** a grant;
|
||||
**Change…** its address, timeout or rate; **Replace its credential…**;
|
||||
**Turn off** (its tools are listed nowhere and every call refused) or on;
|
||||
**Remove…**. A server's tools are seen nowhere until granted. The screen
|
||||
also says how a person points an agent at the gateway, with a key of their
|
||||
own from `/profile/keys`. `llmgw ctl mcp` does the same from a shell, and
|
||||
the web page **MCP servers** (`/admin/mcp`, administrators only) in a browser.
|
||||
|
||||
**Mirror** is what the model mirror holds, one row for each version of each
|
||||
model — a file, or the set of files a model of several is: shards, a vision
|
||||
projector, a config and tokenizer — with how far its download has got and
|
||||
|
||||
+12
-4
@@ -301,10 +301,18 @@ policies:
|
||||
string; empty sends `{"input": …}`). `blocked` and `reason` are dotted paths
|
||||
into the JSON answer, numbers for list items; true, a number above zero, a
|
||||
non-empty list or a string other than "false" is unsafe. On input the
|
||||
service reads the latest user turn; on output the whole answer, once, at its
|
||||
end — a streamed answer released a sentence at a time is then stopped or
|
||||
flagged as the judge's is. The rule recorded is `service:<name>:<reason>`.
|
||||
The service is told the text only, never who sent it.
|
||||
service reads the latest user turn as the model will, with what `redact`
|
||||
took out still out; on output the whole answer, once, at its end, as the
|
||||
client reads it — with `restore`, the values put back. Being at the end, a
|
||||
service can stop a streamed answer only where it is held back
|
||||
(`stream: buffer_all`): released a sentence at a time, the client has read
|
||||
it by then, and what the service finds ends the stream with
|
||||
`content_filter`, flags it and is recorded. The rule recorded is
|
||||
`service:<name>:<reason>`. The service is told the text only, never who
|
||||
sent it. Under `fail: open`, the default, a service that cannot be asked
|
||||
lets the text through — and a caller can make it fail, sending what it
|
||||
refuses to read; give `fail: closed` where the service must decide. Of an
|
||||
answer asked for with `n` above 1, the service reads the first.
|
||||
|
||||
## Where a policy applies
|
||||
|
||||
|
||||
+135
-8
@@ -876,6 +876,27 @@ Nothing to do here if you built the mirror in step 3 — the agent fetches what
|
||||
model's `source` names, and that source is your mirror. This section is for the
|
||||
arrangements where the files are already on the machine.
|
||||
|
||||
**Weights in a pickle format are refused.** `.bin`, `.pt`, `.pth`, `.ckpt` and
|
||||
`.pkl` are Python pickles, and loading one runs whatever it holds, as the
|
||||
engine's user, on the node — which is how the malicious models found on Hugging
|
||||
Face work. A model whose source is one, or a mirror version holding one, is
|
||||
refused when it is added and by the node before it fetches anything; vLLM and
|
||||
SGLang given a model by its Hugging Face id read only its safetensors. Take the
|
||||
repository's safetensors or GGUF. A model that has no other form, and whose
|
||||
files were checked, can say so:
|
||||
|
||||
```yaml
|
||||
- name: triposr
|
||||
source: http://mirror:8081/triposr/<rev>/model.ckpt
|
||||
allow_pickle: true
|
||||
```
|
||||
|
||||
A local path is judged by its name as any source is; a local directory is
|
||||
not looked into. What an engine downloads itself by a Hugging Face id, past
|
||||
the node — an adapter given by its id, a TensorRT-LLM model's — is refused
|
||||
when the model is added, unless it says `allow_pickle`: mirror its
|
||||
safetensors and give that instead.
|
||||
|
||||
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.
|
||||
@@ -1034,6 +1055,27 @@ A node with no location is placed as before, as a rack of its own.
|
||||
See [facility-aware-placement.md](facility-aware-placement.md) for power caps
|
||||
and holds.
|
||||
|
||||
### A card in the wrong slot
|
||||
|
||||
An x16 card in a slot wired x8 or x4, behind a riser, or in a slot of an
|
||||
older PCIe generation than the card's, links slower than it was built to and
|
||||
works otherwise. A model on the card still decodes at the card's own memory
|
||||
speed, but it loads slower, a replica split over several cards by PCIe is
|
||||
slower, and so is anything that spills to the host. Each card's agent reads
|
||||
how its link stands — the fastest generation card and slot allow together,
|
||||
and the lanes in use — against what the card could do, and the node says so
|
||||
among its warnings (the dashboard, the console's Nodes, `llmgw doctor`):
|
||||
|
||||
```
|
||||
GPU 1 (NVIDIA H100 80GB HBM3) linked at x8 of its x16 lanes: 31.5 GB/s where it
|
||||
could have 63.0. Check the slot (wired fewer lanes, or an older generation), a
|
||||
riser, or the BIOS's PCIe setting
|
||||
```
|
||||
|
||||
The generation a card reports now drops while it idles, to save power; what is
|
||||
compared is the fastest it may reach, which does not. `lspci -vv -s BUS-ID`
|
||||
shows the same: `LnkCap` is the card's, `LnkSta` the link's.
|
||||
|
||||
## 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
|
||||
@@ -1587,6 +1629,55 @@ Not yet: models declared in server.yaml cannot name a peer's model as a
|
||||
fallback (the catalog in the file is checked before the peers are read), and
|
||||
a peer behind a private CA needs that CA in the system's trust store.
|
||||
|
||||
### Changing a served model's context
|
||||
|
||||
How much one request may use is set when a model is served. To change it
|
||||
afterwards:
|
||||
|
||||
```
|
||||
llmgw ctl model context qwen-7b 32768 # tokens a request
|
||||
llmgw ctl model context qwen-7b 32768 --slots 4 # llama.cpp: four at once
|
||||
```
|
||||
|
||||
The catalog, the engine's arguments (`--ctx-size` and `--parallel` for
|
||||
llama.cpp, `--max-model-len` for vLLM, `--context-length` for SGLang) and
|
||||
the memory its KV cache takes change together. A context longer than the
|
||||
model was trained for is refused, and so is one its cards cannot hold — with
|
||||
the figures, so you know how far to go. Its replicas are then started again on
|
||||
it, one by one, a new one ready before the old one stops. The console has the
|
||||
same under a model's menu, Context length…. A model declared in the cluster's
|
||||
settings is changed there (`llmgw ctl config edit`).
|
||||
|
||||
### Moving a served model to more cards
|
||||
|
||||
How many cards of one node each replica uses can be changed while the model
|
||||
is served:
|
||||
|
||||
```
|
||||
llmgw ctl model gpus qwen3-coder-30b 2 # two cards a replica
|
||||
llmgw ctl model gpus llama-70b 4 --profile h100 # the model has several profiles on cards
|
||||
llmgw ctl model gpus qwen3-coder-30b 2 --in-place # no room beside the running replica
|
||||
```
|
||||
|
||||
vLLM, SGLang and TensorRT-LLM spread the model by tensor parallel (`tp`),
|
||||
which needs the count to divide the model's KV heads or be divided by them:
|
||||
a model with 8 splits 2, 4, 8 or 16 ways, not 3 or 12. llama.cpp has no
|
||||
tensor parallel; it splits the layers over the cards (`gpus_per_node`), and
|
||||
a `--tensor-split` written for the old cards is dropped. ComfyUI and Strata
|
||||
run on one card. A count no node has is refused, and so is one whose cards
|
||||
would not hold the weights and the context — with the figures; fewer cards
|
||||
with a long context may need `llmgw ctl model context` first.
|
||||
|
||||
Its replicas are then started again on it, each new one ready beside the old
|
||||
before the old stops. That needs the cards free beside it: on a node of two
|
||||
cards with the model on one, two do not fit until it is gone, and the command
|
||||
says so. `--in-place` stops each replica first and starts its replacement on
|
||||
the cards it held and the free ones, and the model does not answer while the
|
||||
new one loads — a model with several replicas keeps answering on the others.
|
||||
A replica spread over several nodes is laid out by hand, with a fabric
|
||||
(`nodes`, `gpus_per_node` in the profile). The console has the same under a
|
||||
model's menu, GPUs….
|
||||
|
||||
### Trying new weights on some of the traffic first
|
||||
|
||||
When a model's repository publishes new weights, `llmgw ctl model update`
|
||||
@@ -2137,8 +2228,10 @@ llmgw keys cache --project Acme/ml/train --off
|
||||
- **Policies still apply:** the limits are checked and the guardrails judge
|
||||
the request before the cache is looked in, and the output policy judges a
|
||||
reused answer as it would a new one. A request continuing a stored
|
||||
conversation, or under a `no_store` policy, is never answered from the
|
||||
cache nor kept in it.
|
||||
conversation, under a `no_store` policy, or whose redacted values are put
|
||||
back in its answer (`restore: true`) is never answered from the cache nor
|
||||
kept in it: redacted, two callers' requests read the same, and the answer
|
||||
has the first caller's values in it.
|
||||
- **What a hit costs:** nothing. It is recorded as a request — status 200,
|
||||
the model, the key, the rate limit's count — with no tokens and no cost,
|
||||
marked as a cache hit (`llmgw ctl request <id>` says so), so it spends no
|
||||
@@ -2178,8 +2271,8 @@ router:
|
||||
`router.cache.embedding_model`, against the questions already answered
|
||||
in the same conversation — the same system prompt, the same earlier turns,
|
||||
tools, `response_format` and sampling settings, the same model and
|
||||
project. A like question in another conversation, or after a tool's
|
||||
result, is not taken for this one. The exact cache is asked first, and an
|
||||
project. A like question in another conversation, after a tool's result,
|
||||
or with a picture, a sound or a file in it, is not taken for this one. The exact cache is asked first, and an
|
||||
exact hit embeds nothing.
|
||||
- **How alike:** the cosine similarity of the two embeddings, at least the
|
||||
project's `--similar`. 0.95 is strict; below 0.9 different questions start
|
||||
@@ -2306,6 +2399,7 @@ triposr`. What that adds is this entry, and its workflow is built in:
|
||||
```yaml
|
||||
- name: triposr
|
||||
source: hf:stabilityai/TripoSR/model.ckpt
|
||||
allow_pickle: true # published as this checkpoint only
|
||||
comfyui:
|
||||
workflow: triposr # built in; a model of your own names a JSON file on every control plane
|
||||
output: meshes
|
||||
@@ -2477,8 +2571,31 @@ lost: a standby leads within seconds, Patroni promotes the standby database
|
||||
if it was the primary, and the floating address follows whichever leads. The
|
||||
mirror's management API answers the three control planes' addresses, not
|
||||
only its own machine's, so a standby that comes to lead can still fetch and
|
||||
add models. The mirror itself stays on the first machine: while it is down,
|
||||
models the nodes already have keep serving, and new ones wait for it.
|
||||
add models.
|
||||
|
||||
**A second mirror.** The mirror is on the first machine, and without a
|
||||
second one, models the nodes already have keep serving while that machine
|
||||
is down, and new ones wait for it. With a second, the cluster carries on.
|
||||
- **Setting it up:** the screens offer one on either database machine
|
||||
("A second mirror"); by command, `--second-mirror 1` (or `2`). For an
|
||||
install that has one mirror, run this on the mirror's machine:
|
||||
|
||||
```
|
||||
sudo llmgw mirror add-peer ubuntu@10.0.0.7 --control-plane 10.0.0.5 --control-plane 10.0.0.6 --control-plane 10.0.0.7
|
||||
```
|
||||
- **Kept in step:** each mirror reads the other every thirty seconds and
|
||||
copies what it lacks, keeping a copy only if its checksum matches.
|
||||
- A model removed from one is removed from the other, even if that one
|
||||
was down at the time and comes back later.
|
||||
- A model added again after it was removed stands on both.
|
||||
- Neither mirror leads.
|
||||
- **Control planes:** they name both (`mirror.peers` in the cluster's
|
||||
settings) and send the mirror's management calls to the first that
|
||||
answers.
|
||||
- **Nodes:** while the mirror a model's address names is down, a node is
|
||||
given the same file on the other.
|
||||
- **Disk:** the second machine needs room for every model, as the first
|
||||
does.
|
||||
|
||||
## Two control planes
|
||||
|
||||
@@ -2586,10 +2703,20 @@ llmgw keys limits --scope team --id Acme/support --tools 'jira.*'
|
||||
- **Who may call what:** tools are named `server.tool` and granted as
|
||||
patterns (`jira.*`, `db.query`) at any level of the tenancy tree, as
|
||||
models are.
|
||||
- A level that names any narrows what is above it.
|
||||
- A level that names any narrows what is above it: a tool is seen where
|
||||
every level above that names tools names it too. A team under an
|
||||
organization that names `hr.*` sees no `jira.*` whatever the team is
|
||||
given, until the organization names it as well.
|
||||
- Unlike models, **none granted anywhere is none**: a tool reaches into a
|
||||
company system, so it is granted, never assumed.
|
||||
- The console's limits form has the same field.
|
||||
- `keys limits --tools` writes a level's list whole, every server's
|
||||
together; an empty list (`--tools ''`) closes the level to every tool.
|
||||
The console's **MCP servers** screen and the `/admin/mcp` page grant
|
||||
and take back one server's tools: a level given its first list starts
|
||||
from the nearest list above it, so it keeps what it saw; one left
|
||||
naming none goes back to naming none, or stays closed when a level
|
||||
above names tools, so taking one server away opens it to no other. They
|
||||
say when a level above keeps the granted tools out of reach.
|
||||
- **The server is told who calls:** `X-Llmgw-Key`, and `X-Llmgw-User` and
|
||||
`X-Llmgw-Project` where there are some. It is called with llmgw's
|
||||
credentials, never the caller's key.
|
||||
|
||||
+29
-1
@@ -146,7 +146,35 @@ agent is stopped.
|
||||
## 4. Add a model
|
||||
|
||||
Open `/admin/mirror` and choose **Explore models**. It lists what is popular on
|
||||
Hugging Face, showing only what your nodes can run.
|
||||
Hugging Face, showing only what your nodes can run: judged by the cluster as it
|
||||
is now — the largest card and its compute capability, the most memory on one
|
||||
node — so a model too big for them, in a form those cards lack (FP8 needs 8.9),
|
||||
or in neither GGUF nor safetensors is left out, and the page says how many.
|
||||
Without GPUs it lists GGUF only. Untick **only what these nodes can run** to
|
||||
see everything. On a model's page a file too big for the cluster is shown
|
||||
without a **Fetch** button, with what it would need.
|
||||
|
||||
The **CPU tok/s** column is about how fast one reply would come on the CPU of
|
||||
the fastest node whose memory holds the model at Q4_K_M; a model's page gives
|
||||
it for each GGUF file by its own size. Every node's agent measures how fast its
|
||||
memory reads when it starts (it is in the agent's log as `memory_gbs`), and a
|
||||
token on a CPU reads every weight once: the estimate is that speed over the
|
||||
model's size — for a model of experts named like `30B-A3B`, its active part.
|
||||
On the test cluster it came within 15% of what llama.cpp measured. Several
|
||||
replies at once share the reads and each is slower.
|
||||
|
||||
**GPU tok/s** is the same on the cards, for the layout this cluster would
|
||||
give the model — the fewest cards that hold it, of each class, the fastest of
|
||||
those; its tooltip says which ("2× h100-80g, NVLink, vLLM", "2 nodes × 1×
|
||||
a100-80g over ib0 (InfiniBand 200 Gb/s), vLLM"). A card's memory bandwidth is
|
||||
known from its model name (an H100 SXM and an H100 PCIe differ by half); a card
|
||||
not in the table gets no figure. Tensor parallel across a node's cards divides
|
||||
what each reads, and adds two all-reduces a layer, short over NVLink and
|
||||
longer over PCIe; llama.cpp across cards splits by layer, so the cards take
|
||||
turns and a reply is no faster than on one; a pipeline across nodes takes turns
|
||||
too, and adds a hop over the fabric for each node. One card's figures were
|
||||
measured on the test cluster's RTX 6000; tensor parallel and the fabric are
|
||||
from what NCCL is known to take, not measured here.
|
||||
|
||||
1. Click a model to read its card.
|
||||
2. Pick one of its files, usually a `q4_k_m` for a first try, and click
|
||||
|
||||
@@ -180,6 +180,7 @@ rather than repeated.
|
||||
| `parallel` | `int` | [LlamaCppConfig](#llamacppconfig) |
|
||||
| `threads` | `int` | [LlamaCppConfig](#llamacppconfig) |
|
||||
| `api` | `string` | [Mirror](#mirror) |
|
||||
| `peers` | `[]string` | [Mirror](#mirror) |
|
||||
| `tls` | `tlsconf.Config` | [Mirror](#mirror) |
|
||||
| `registry` | `string` | [Mirror](#mirror) |
|
||||
| `timeout` | `time.Duration` | [Mirror](#mirror) |
|
||||
@@ -188,6 +189,7 @@ rather than repeated.
|
||||
| `aliases` | `[]string` | [Model](#model) |
|
||||
| `description` | `string` | [Model](#model) |
|
||||
| `source` | `string` | [Model](#model) |
|
||||
| `allow_pickle` | `bool` | [Model](#model) |
|
||||
| `profiles` | `[]Profile` | [Model](#model) |
|
||||
| `scaling` | `Scaling` | [Model](#model) |
|
||||
| `class` | `string` | [Model](#model) |
|
||||
@@ -394,6 +396,9 @@ rather than repeated.
|
||||
| `timeout` | `time.Duration` | [Service](#service) |
|
||||
| `stopped` | `bool` | [ServingSettings](#servingsettings) |
|
||||
| `bin` | `string` | [SimConfig](#simconfig) |
|
||||
| `dir` | `string` | [StrataConfig](#strataconfig) |
|
||||
| `python` | `string` | [StrataConfig](#strataconfig) |
|
||||
| `engine` | `string` | [StrataConfig](#strataconfig) |
|
||||
| `python` | `string` | [TRTLLMConfig](#trtllmconfig) |
|
||||
| `store` | `string` | [Threads](#threads) |
|
||||
| `retain_hours` | `int` | [Threads](#threads) |
|
||||
@@ -474,6 +479,7 @@ rather than repeated.
|
||||
| `trtllm` | `TRTLLMConfig` | [engines.Config](#enginesconfig) |
|
||||
| `sim` | `SimConfig` | [engines.Config](#enginesconfig) |
|
||||
| `comfyui` | `ComfyUIConfig` | [engines.Config](#enginesconfig) |
|
||||
| `strata` | `StrataConfig` | [engines.Config](#enginesconfig) |
|
||||
| `enabled` | `*bool` | [guard.Config](#guardconfig) |
|
||||
| `default` | `[]string` | [guard.Config](#guardconfig) |
|
||||
| `policies` | `map[string]Policy` | [guard.Config](#guardconfig) |
|
||||
@@ -1721,6 +1727,14 @@ The management port, not the one nodes fetch bytes from. Those are two
|
||||
listeners on the mirror precisely so this one can be locked down without
|
||||
re-certifying every agent's fetcher.
|
||||
|
||||
### `peers`
|
||||
|
||||
`[]string`
|
||||
|
||||
Peers are the management addresses of mirrors kept in step with this
|
||||
one (llmgw mirror serve --peer): asked when it cannot be reached, and
|
||||
a node given the same file on one of them while it is down.
|
||||
|
||||
### `tls`
|
||||
|
||||
`tlsconf.Config`
|
||||
@@ -1790,6 +1804,16 @@ choice by guesswork.
|
||||
|
||||
Source is the weights path or a HuggingFace id.
|
||||
|
||||
### `allow_pickle`
|
||||
|
||||
`bool`
|
||||
|
||||
AllowPickle lets the nodes fetch this model's weights, or its
|
||||
adapters', in a pickle format (.bin, .pt, .pth, .ckpt), which run code
|
||||
as they are loaded. Off, such a source is refused when the model is
|
||||
added and by the node before it fetches anything. For a model that has
|
||||
no safetensors or GGUF form and whose files were checked.
|
||||
|
||||
### `profiles`
|
||||
|
||||
`[]Profile`
|
||||
@@ -3395,6 +3419,27 @@ entries ask. In force on every control plane within seconds.
|
||||
Bin is the llmgw binary itself; the engine is one of its subcommands.
|
||||
Empty means whatever started this agent, which is right almost always.
|
||||
|
||||
## StrataConfig
|
||||
|
||||
### `dir`
|
||||
|
||||
`string`
|
||||
|
||||
Dir is Strata's source directory, where setup.py and serve/ are.
|
||||
|
||||
### `python`
|
||||
|
||||
`string`
|
||||
|
||||
Python is the interpreter of its environment; Dir's .venv when unset.
|
||||
|
||||
### `engine`
|
||||
|
||||
`string`
|
||||
|
||||
Engine is the directory of its prebuilt engine, where the strata
|
||||
binary and the CUDA libraries it ships are; Dir's engine/ when unset.
|
||||
|
||||
## TRTLLMConfig
|
||||
|
||||
### `python`
|
||||
@@ -3967,6 +4012,15 @@ ComfyUI makes pictures; offered by a node whose ComfyUI is named here.
|
||||
|
||||
Takes a [ComfyUIConfig](#comfyuiconfig) block.
|
||||
|
||||
### `strata`
|
||||
|
||||
`StrataConfig`
|
||||
|
||||
Strata serves a few mixture-of-experts models from RAM and one card;
|
||||
offered by a node whose Strata is named here, or that installed it.
|
||||
|
||||
Takes a [StrataConfig](#strataconfig) block.
|
||||
|
||||
## guard.Config
|
||||
|
||||
### `enabled`
|
||||
|
||||
+59
-8
@@ -12,6 +12,7 @@ import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"maps"
|
||||
"math"
|
||||
"math/rand/v2"
|
||||
"net"
|
||||
"net/http"
|
||||
@@ -392,8 +393,11 @@ func (a *Agent) settleAdopted(sh *supervisor.Shard, o supervisor.Orphan) {
|
||||
_ = a.sup.Stop(ctx, o.ShardID)
|
||||
return
|
||||
}
|
||||
// What it computes on, for one started before the agent kept that.
|
||||
if a.sup.ComputeOf(o.ShardID) == nil {
|
||||
// What it computes on, for one started before the agent kept that — or
|
||||
// kept as the CPU for a GPU placement by an agent that did not follow
|
||||
// an engine into a process group of its own, which the check above has
|
||||
// just found on its GPUs.
|
||||
if c := a.sup.ComputeOf(o.ShardID); c == nil || (len(o.Spec.GPUIndices) > 0 && c.Backend == "CPU") {
|
||||
if c, ok := a.gpuCompute(ctx, o.Spec.GPUIndices, sh.Status().PID); ok {
|
||||
a.sup.SetCompute(o.ShardID, c)
|
||||
}
|
||||
@@ -513,6 +517,7 @@ func (a *Agent) scan(ctx context.Context) error {
|
||||
inv, err := discovery.Discover(ctx, opts, a.log)
|
||||
host := discovery.DiscoverHost()
|
||||
host.MinComputeCapability = discovery.MinComputeCapability(inv.GPUs)
|
||||
host.MemBandwidthGBs = a.memBandwidth()
|
||||
a.mu.Lock()
|
||||
a.inv = inv
|
||||
a.prov.Host = host
|
||||
@@ -569,6 +574,9 @@ func (a *Agent) session(ctx context.Context, server string) error {
|
||||
reg.Host = &llmgwv1.Host{
|
||||
OsId: h.OSID, OsVersion: h.OSVersion, OsCodename: h.OSCodename,
|
||||
Arch: h.Arch, CpuFlags: h.CPUFlags, NvidiaDriver: h.NVIDIADriver,
|
||||
// What is measured by now, not what the startup scan found before
|
||||
// the measurement ended: a reconnect would otherwise say 0.
|
||||
MemBandwidthGbs: a.memBandwidth(),
|
||||
}
|
||||
|
||||
if err := stream.Send(&llmgwv1.AgentMessage{
|
||||
@@ -666,6 +674,7 @@ func (a *Agent) heartbeat() *llmgwv1.AgentMessage {
|
||||
// without anybody restarting anything.
|
||||
hb.MemoryTotalMb = int64(discovery.MemoryTotalMB())
|
||||
hb.CpuCount = int32(runtime.NumCPU())
|
||||
hb.MemBandwidthGbs = a.memBandwidth()
|
||||
|
||||
a.reprobeEngines()
|
||||
a.mu.Lock()
|
||||
@@ -1124,7 +1133,7 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
},
|
||||
})
|
||||
})
|
||||
if path, err := a.weights().Ensure(ctx, spec.Profile.ModelPath); err != nil {
|
||||
if path, err := a.weights().Ensure(ctx, spec.Profile.ModelPath, spec.Profile.AllowPickle); err != nil {
|
||||
stop()
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
@@ -1145,7 +1154,7 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
send(&llmgwv1.ShardStatus{State: agentsvc.StateToWire(supervisor.Starting),
|
||||
Fetch: &llmgwv1.FetchProgress{Source: p.Source, BytesDone: p.BytesDone, BytesTotal: p.BytesTotal}})
|
||||
})
|
||||
path, err := a.weights().Ensure(ctx, m["source"])
|
||||
path, err := a.weights().Ensure(ctx, m["source"], m["allow_pickle"] == "true")
|
||||
stop()
|
||||
if err != nil {
|
||||
report(supervisor.Failed, fmt.Sprintf("model %s: %v", m["name"], err), nil)
|
||||
@@ -1161,7 +1170,7 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
// loads them at start, and one missing is a replica that answers for the
|
||||
// base model under the adapter's name.
|
||||
for i, ad := range spec.Profile.Adapters {
|
||||
path, err := a.weights().Ensure(ctx, ad.Path)
|
||||
path, err := a.weights().Ensure(ctx, ad.Path, spec.Profile.AllowPickle)
|
||||
if err != nil {
|
||||
report(supervisor.Failed, fmt.Sprintf("adapter %s: %v", ad.Name, err), nil)
|
||||
return
|
||||
@@ -1505,18 +1514,25 @@ func ResolveManagement(c Config) (addr, iface string) {
|
||||
// shard computes on as soon as that is known: an adopted engine's is found
|
||||
// after registration, which carried what was known then.
|
||||
func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMessage) {
|
||||
sent := map[string]bool{}
|
||||
// What was sent for each shard, so a compute that changes is sent again:
|
||||
// an adopted shard's kept one is sent at once, and adopting it may then
|
||||
// find it was kept wrong.
|
||||
sent := map[string]string{}
|
||||
t := time.NewTicker(5 * time.Second)
|
||||
defer t.Stop()
|
||||
for range 60 {
|
||||
for _, st := range a.sup.Statuses() {
|
||||
if st.State.Terminal() || sent[st.ID] {
|
||||
if st.State.Terminal() {
|
||||
continue
|
||||
}
|
||||
c := a.sup.ComputeOf(st.ID)
|
||||
if c == nil {
|
||||
continue
|
||||
}
|
||||
said := fmt.Sprint(*c)
|
||||
if sent[st.ID] == said {
|
||||
continue
|
||||
}
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
@@ -1527,7 +1543,7 @@ func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMess
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
sent[st.ID] = true
|
||||
sent[st.ID] = said
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
@@ -1539,3 +1555,38 @@ func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMess
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// memBandwidth is the node's memory read rate in GB/s, 0 until measured.
|
||||
// Measured in the background, the first time it is asked for: it reads half
|
||||
// a gigabyte, which must not hold up an agent reaching its control plane,
|
||||
// and the heartbeat carries it when it is known. Not at once but after a
|
||||
// minute or so, chosen at random, and again every hour, keeping the best:
|
||||
// VMs on one host that all restart together measure each other, and read a
|
||||
// third less than any of them gets alone. Once a process rather than an
|
||||
// agent, so tests running many agents in one read it once.
|
||||
func (a *Agent) memBandwidth() float64 {
|
||||
bwOnce.Do(func() {
|
||||
go func() {
|
||||
time.Sleep(bandwidthDelay())
|
||||
for {
|
||||
if bw := measureBandwidth(); bw > math.Float64frombits(bwBits.Load()) {
|
||||
bwBits.Store(math.Float64bits(bw))
|
||||
a.log.Info("this node's memory", "read_gbs", math.Round(bw*10)/10)
|
||||
}
|
||||
time.Sleep(time.Hour)
|
||||
}
|
||||
}()
|
||||
})
|
||||
return math.Float64frombits(bwBits.Load())
|
||||
}
|
||||
|
||||
// bandwidthDelay is how long an agent waits before measuring, so agents
|
||||
// started together measure apart.
|
||||
var bandwidthDelay = func() time.Duration { return 10*time.Second + rand.N(110*time.Second) }
|
||||
|
||||
var (
|
||||
bwOnce sync.Once
|
||||
bwBits atomic.Uint64
|
||||
// measureBandwidth is replaced in tests, which have no use for it.
|
||||
measureBandwidth = discovery.MeasureMemoryBandwidth
|
||||
)
|
||||
@@ -671,7 +671,7 @@ func TestAShardWithMissingWeightsFailsBeforeLaunching(t *testing.T) {
|
||||
// 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"))
|
||||
_, err := a.weights().Ensure(t.Context(), filepath.Join(t.TempDir(), "not-there.gguf"), false)
|
||||
if err == nil {
|
||||
t.Fatal("a missing weights file was accepted")
|
||||
}
|
||||
|
||||
@@ -52,6 +52,8 @@ type GPU struct {
|
||||
// memory the instance's share. Utilisation stays unreported, because no
|
||||
// driver measures it per instance.
|
||||
MIG *MIGDevice `json:"mig,omitempty"`
|
||||
// PCIe is how the card is linked to its host, zero when not read.
|
||||
PCIe PCIeLink `json:"pcie,omitempty"`
|
||||
}
|
||||
|
||||
// GPUHealth records conditions that make a GPU unsafe to schedule onto.
|
||||
|
||||
@@ -67,10 +67,23 @@ func DiscoverGPUsNVML(_ context.Context) ([]GPU, error) {
|
||||
if major, minor, ret := dev.GetCudaComputeCapability(); ret == nvml.SUCCESS {
|
||||
g.ComputeCapability = fmt.Sprintf("%d.%d", major, minor)
|
||||
}
|
||||
if v, ret := dev.GetMaxPcieLinkGeneration(); ret == nvml.SUCCESS {
|
||||
g.PCIe.Gen = v
|
||||
}
|
||||
if v, ret := dev.GetGpuMaxPcieLinkGeneration(); ret == nvml.SUCCESS {
|
||||
g.PCIe.GPUGen = v
|
||||
}
|
||||
if v, ret := dev.GetCurrPcieLinkWidth(); ret == nvml.SUCCESS {
|
||||
g.PCIe.Width = v
|
||||
}
|
||||
if v, ret := dev.GetMaxPcieLinkWidth(); ret == nvml.SUCCESS {
|
||||
g.PCIe.WidthMax = v
|
||||
}
|
||||
if pci, ret := dev.GetPciInfo(); ret == nvml.SUCCESS {
|
||||
g.PCIBusID = busIDString(pci.BusId[:])
|
||||
byBusID[normalizeBusID(g.PCIBusID)] = i
|
||||
}
|
||||
g.PCIe = withSysfs(g.PCIe, "/sys/bus/pci/devices", g.PCIBusID)
|
||||
if t, ret := dev.GetTemperature(nvml.TEMPERATURE_GPU); ret == nvml.SUCCESS {
|
||||
g.Health.TemperatureC = int(t)
|
||||
}
|
||||
|
||||
@@ -64,6 +64,7 @@ func DiscoverGPUsSMI(ctx context.Context, binary string) ([]GPU, error) {
|
||||
if err != nil || len(cards) == 0 {
|
||||
return cards, err
|
||||
}
|
||||
readPCIe(ctx, binary, cards)
|
||||
// The query lists whole cards, MIG mode or not. A card in MIG mode is
|
||||
// replaced by its instances, which only -L names and only the full
|
||||
// report sizes; asked only when some card could be in MIG mode, since
|
||||
|
||||
@@ -31,6 +31,10 @@ type Host struct {
|
||||
// MinComputeCapability is the oldest GPU's, 0 when there are none or it
|
||||
// is not known. Not read here but from the GPU inventory, by the agent.
|
||||
MinComputeCapability float64 `json:"min_compute_capability,omitempty"`
|
||||
// MemBandwidthGBs is how fast the memory reads, in GB/s, 0 when not
|
||||
// measured: what decides a model's speed on the CPU. Not read here but
|
||||
// measured by the agent when it starts (MeasureMemoryBandwidth).
|
||||
MemBandwidthGBs float64 `json:"mem_bandwidth_gbs,omitempty"`
|
||||
}
|
||||
|
||||
// Family is the packaging family: "deb", "rpm", or "" when neither is known.
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// How fast this machine reads its memory, which on a CPU is what decides how
|
||||
// fast a model answers: every token of a reply reads all the weights it uses
|
||||
// once, so a node that reads 40 GB/s writes a 4 GB model's tokens about ten
|
||||
// times a second, whatever its core count. Measured once when the agent
|
||||
// starts, in a fraction of a second, by every CPU summing a share of a buffer
|
||||
// too large for any cache — a VM sharing its host's memory with others gets
|
||||
// its share, which is the number that matters for it.
|
||||
|
||||
// bandwidthBuffer is what is read at most; the least is eight times the
|
||||
// processor's last cache, and never under minBuffer. Smaller, it would be
|
||||
// the cache that was measured — some hold a gigabyte — and the figure kept
|
||||
// would be too high.
|
||||
const (
|
||||
bandwidthBuffer = 2 << 30
|
||||
minBuffer = 256 << 20
|
||||
)
|
||||
|
||||
// MeasureMemoryBandwidth is the memory's read rate in GB/s, the best of a
|
||||
// few passes; 0 when the machine has too little free memory to measure.
|
||||
func MeasureMemoryBandwidth() float64 {
|
||||
size := int64(bandwidthBuffer)
|
||||
// Never more than a quarter of what is free: a node already holding
|
||||
// models must not be pushed into swap by a measurement. Without room
|
||||
// for a buffer past the cache, none is measured: the hourly retry
|
||||
// measures it when there is.
|
||||
if avail := memAvailable("/proc/meminfo"); avail > 0 {
|
||||
size = min(size, avail/4)
|
||||
}
|
||||
if size < max(minBuffer, 8*lastCacheBytes("/sys/devices/system/cpu")) {
|
||||
return 0
|
||||
}
|
||||
buf := make([]uint64, size/8)
|
||||
for i := range buf {
|
||||
buf[i] = uint64(i)
|
||||
}
|
||||
workers := runtime.NumCPU()
|
||||
chunk := (len(buf) + workers - 1) / workers
|
||||
best := 0.0
|
||||
for pass := 0; pass < 4; pass++ {
|
||||
var wg sync.WaitGroup
|
||||
start := time.Now()
|
||||
for w := 0; w < workers; w++ {
|
||||
lo, hi := w*chunk, min((w+1)*chunk, len(buf))
|
||||
if lo >= hi {
|
||||
continue
|
||||
}
|
||||
wg.Add(1)
|
||||
go func(part []uint64) {
|
||||
defer wg.Done()
|
||||
sink(sum(part))
|
||||
}(buf[lo:hi])
|
||||
}
|
||||
wg.Wait()
|
||||
if s := time.Since(start).Seconds(); s > 0 {
|
||||
best = max(best, float64(size)/s/1e9)
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
// sum reads every word, eight at a time so the loop is not what limits it.
|
||||
func sum(p []uint64) uint64 {
|
||||
var a, b, c, d, e, f, g, h uint64
|
||||
i := 0
|
||||
for ; i+8 <= len(p); i += 8 {
|
||||
a += p[i]
|
||||
b += p[i+1]
|
||||
c += p[i+2]
|
||||
d += p[i+3]
|
||||
e += p[i+4]
|
||||
f += p[i+5]
|
||||
g += p[i+6]
|
||||
h += p[i+7]
|
||||
}
|
||||
for ; i < len(p); i++ {
|
||||
a += p[i]
|
||||
}
|
||||
return a + b + c + d + e + f + g + h
|
||||
}
|
||||
|
||||
// sinkValue keeps the sums alive, so the reads are not optimised away.
|
||||
var sinkValue struct {
|
||||
sync.Mutex
|
||||
v uint64
|
||||
}
|
||||
|
||||
func sink(v uint64) {
|
||||
sinkValue.Lock()
|
||||
sinkValue.v += v
|
||||
sinkValue.Unlock()
|
||||
}
|
||||
|
||||
// memAvailable is MemAvailable in bytes, 0 when unknown.
|
||||
func memAvailable(path string) int64 {
|
||||
f, err := os.Open(path) //nolint:gosec // a fixed system path
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
defer f.Close()
|
||||
sc := bufio.NewScanner(f)
|
||||
for sc.Scan() {
|
||||
fields := strings.Fields(sc.Text())
|
||||
if len(fields) >= 2 && fields[0] == "MemAvailable:" {
|
||||
kb, _ := strconv.ParseInt(fields[1], 10, 64)
|
||||
return kb << 10
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// lastCacheBytes is the largest cache any CPU reports, 0 when unknown.
|
||||
func lastCacheBytes(root string) int64 {
|
||||
paths, _ := filepath.Glob(filepath.Join(root, "cpu*", "cache", "index*", "size"))
|
||||
var most int64
|
||||
for _, p := range paths {
|
||||
b, err := os.ReadFile(p) //nolint:gosec // under /sys
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
v := strings.TrimSpace(string(b))
|
||||
mult := int64(1)
|
||||
switch {
|
||||
case strings.HasSuffix(v, "K"):
|
||||
mult, v = 1<<10, strings.TrimSuffix(v, "K")
|
||||
case strings.HasSuffix(v, "M"):
|
||||
mult, v = 1<<20, strings.TrimSuffix(v, "M")
|
||||
}
|
||||
if n, err := strconv.ParseInt(v, 10, 64); err == nil {
|
||||
most = max(most, n*mult)
|
||||
}
|
||||
}
|
||||
return most
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMemoryBandwidthIsMeasured(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("reads half a gigabyte")
|
||||
}
|
||||
bw := MeasureMemoryBandwidth()
|
||||
// Any machine running this reads more than a gigabyte a second and less
|
||||
// than a few terabytes.
|
||||
if bw < 1 || bw > 5000 {
|
||||
t.Errorf("measured %.1f GB/s", bw)
|
||||
}
|
||||
t.Logf("%.1f GB/s", bw)
|
||||
}
|
||||
|
||||
func TestTheLastCacheIsReadFromSys(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
for path, size := range map[string]string{"cpu0/cache/index0/size": "48K", "cpu0/cache/index3/size": "96M"} {
|
||||
p := filepath.Join(root, path)
|
||||
_ = os.MkdirAll(filepath.Dir(p), 0o755)
|
||||
_ = os.WriteFile(p, []byte(size+"\n"), 0o644)
|
||||
}
|
||||
if got := lastCacheBytes(root); got != 96<<20 {
|
||||
t.Errorf("%d", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// A card in the wrong slot: an x16 card in a slot wired x8 or x4, behind a
|
||||
// riser, or in a slot of an older PCIe generation than the card's, links
|
||||
// slower than it was built to, and nothing says so. A model on the card
|
||||
// still decodes at the card's own memory speed, but it loads slower, a
|
||||
// tensor-parallel replica over PCIe is slower, and so is anything that
|
||||
// spills to the host. The driver says how the link stands and what the
|
||||
// card could do; the difference is reported where the node's warnings are.
|
||||
|
||||
// PCIeLink is how a card is linked to its host.
|
||||
type PCIeLink struct {
|
||||
// Gen is the fastest generation the card and the slot together allow —
|
||||
// not the current one, which drops while the card idles, to save power.
|
||||
// GPUGen is the card's own fastest, 0 when the driver does not say.
|
||||
Gen int `json:"gen,omitempty"`
|
||||
GPUGen int `json:"gpu_gen,omitempty"`
|
||||
// Width is the lanes in use now, WidthMax the card's most.
|
||||
Width int `json:"width,omitempty"`
|
||||
WidthMax int `json:"width_max,omitempty"`
|
||||
}
|
||||
|
||||
// laneGBs is one lane's usable rate in GB/s, by generation.
|
||||
var laneGBs = map[int]float64{1: 0.25, 2: 0.5, 3: 0.985, 4: 1.969, 5: 3.938, 6: 7.877}
|
||||
|
||||
// GBs is what a link of gen and width carries, in GB/s each way.
|
||||
func pcieGBs(gen, width int) float64 { return laneGBs[gen] * float64(width) }
|
||||
|
||||
// Problem says how a card's link falls short of what the card could do,
|
||||
// empty when it does not or the driver did not say.
|
||||
func (p PCIeLink) Problem() string {
|
||||
gen := p.Gen
|
||||
if p.GPUGen > 0 && gen == 0 {
|
||||
gen = p.GPUGen
|
||||
}
|
||||
slowGen := p.GPUGen > 0 && p.Gen > 0 && p.Gen < p.GPUGen
|
||||
narrow := p.WidthMax > 0 && p.Width > 0 && p.Width < p.WidthMax
|
||||
if !slowGen && !narrow {
|
||||
return ""
|
||||
}
|
||||
could := pcieGBs(max(p.GPUGen, gen), p.WidthMax)
|
||||
has := pcieGBs(gen, p.Width)
|
||||
var why []string
|
||||
if narrow {
|
||||
why = append(why, fmt.Sprintf("x%d of its x%d lanes", p.Width, p.WidthMax))
|
||||
}
|
||||
if slowGen {
|
||||
why = append(why, fmt.Sprintf("PCIe Gen%d of its Gen%d", p.Gen, p.GPUGen))
|
||||
}
|
||||
return fmt.Sprintf("linked at %s: %.1f GB/s where it could have %.1f. "+
|
||||
"Check the slot (wired fewer lanes, or an older generation), a riser, or the BIOS's PCIe setting",
|
||||
strings.Join(why, " and "), has, could)
|
||||
}
|
||||
|
||||
// pcieFields are asked apart from the rest, and the card's own fastest
|
||||
// generation apart again: a driver that does not know a field refuses the
|
||||
// whole query, and these must not cost a node its cards.
|
||||
var pcieFields = [][]string{
|
||||
{"index", "pcie.link.gen.max", "pcie.link.width.current", "pcie.link.width.max", "pcie.link.gen.gpumax"},
|
||||
{"index", "pcie.link.gen.max", "pcie.link.width.current", "pcie.link.width.max"},
|
||||
}
|
||||
|
||||
// readPCIe fills each card's link from nvidia-smi, leaving it empty when
|
||||
// the driver says nothing of it.
|
||||
func readPCIe(ctx context.Context, binary string, cards []GPU) {
|
||||
for _, fields := range pcieFields {
|
||||
out, err := exec.CommandContext(ctx, binary, "--query-gpu="+strings.Join(fields, ","),
|
||||
"--format=csv,noheader,nounits").Output()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
links := parsePCIe(string(out))
|
||||
for i := range cards {
|
||||
if l, ok := links[cards[i].Index]; ok && cards[i].MIG == nil {
|
||||
cards[i].PCIe = l
|
||||
}
|
||||
}
|
||||
break
|
||||
}
|
||||
for i := range cards {
|
||||
if cards[i].MIG == nil {
|
||||
cards[i].PCIe = withSysfs(cards[i].PCIe, "/sys/bus/pci/devices", cards[i].PCIBusID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// withSysfs takes the card's own lane count, and its own fastest
|
||||
// generation, from the kernel: the driver's width.max is what card and slot
|
||||
// allow together, so a x16 card in a slot wired x8 reads x8 of x8 there,
|
||||
// while the device's max_link_width is the card's, 16.
|
||||
func withSysfs(l PCIeLink, root, busID string) PCIeLink {
|
||||
dev := sysfsBusID(busID)
|
||||
if dev == "" {
|
||||
return l
|
||||
}
|
||||
read := func(name string) string {
|
||||
b, err := os.ReadFile(filepath.Join(root, dev, name)) //nolint:gosec // under /sys
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
if n, err := strconv.Atoi(read("max_link_width")); err == nil && n > 0 {
|
||||
l.WidthMax = n
|
||||
}
|
||||
if n, err := strconv.Atoi(read("current_link_width")); err == nil && n > 0 {
|
||||
l.Width = n
|
||||
}
|
||||
if g := genOfSpeed(read("max_link_speed")); g > 0 && l.GPUGen == 0 {
|
||||
l.GPUGen = g
|
||||
}
|
||||
return l
|
||||
}
|
||||
|
||||
// sysfsBusID is nvidia-smi's 00000000:3B:00.0 as the kernel names it,
|
||||
// 0000:3b:00.0.
|
||||
func sysfsBusID(id string) string {
|
||||
id = strings.ToLower(strings.TrimSpace(id))
|
||||
dom, rest, ok := strings.Cut(id, ":")
|
||||
if !ok || len(dom) < 4 {
|
||||
return ""
|
||||
}
|
||||
return dom[len(dom)-4:] + ":" + rest
|
||||
}
|
||||
|
||||
// genOfSpeed is the PCIe generation of a link speed as sysfs writes it.
|
||||
func genOfSpeed(s string) int {
|
||||
switch {
|
||||
case strings.HasPrefix(s, "2.5 GT/s"):
|
||||
return 1
|
||||
case strings.HasPrefix(s, "5.0 GT/s"), strings.HasPrefix(s, "5 GT/s"):
|
||||
return 2
|
||||
case strings.HasPrefix(s, "8.0 GT/s"), strings.HasPrefix(s, "8 GT/s"):
|
||||
return 3
|
||||
case strings.HasPrefix(s, "16.0 GT/s"), strings.HasPrefix(s, "16 GT/s"):
|
||||
return 4
|
||||
case strings.HasPrefix(s, "32.0 GT/s"), strings.HasPrefix(s, "32 GT/s"):
|
||||
return 5
|
||||
case strings.HasPrefix(s, "64.0 GT/s"), strings.HasPrefix(s, "64 GT/s"):
|
||||
return 6
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func parsePCIe(out string) map[int]PCIeLink {
|
||||
r := csv.NewReader(strings.NewReader(out))
|
||||
r.TrimLeadingSpace = true
|
||||
r.FieldsPerRecord = -1
|
||||
rows, _ := r.ReadAll()
|
||||
num := func(s string) int {
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(s))
|
||||
return n
|
||||
}
|
||||
links := map[int]PCIeLink{}
|
||||
for _, row := range rows {
|
||||
if len(row) < 4 {
|
||||
continue
|
||||
}
|
||||
idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
l := PCIeLink{Gen: num(row[1]), Width: num(row[2]), WidthMax: num(row[3])}
|
||||
if len(row) > 4 {
|
||||
l.GPUGen = num(row[4])
|
||||
}
|
||||
links[idx] = l
|
||||
}
|
||||
return links
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestACardInTheWrongSlotIsSaidSo(t *testing.T) {
|
||||
for _, c := range []struct {
|
||||
link PCIeLink
|
||||
want string // empty: nothing to say
|
||||
}{
|
||||
// The PoC's Quadro RTX 6000, idle: current Gen1, but its fastest is 3,
|
||||
// as the card's is, on all 16 lanes. Nothing wrong.
|
||||
{PCIeLink{Gen: 3, GPUGen: 3, Width: 16, WidthMax: 16}, ""},
|
||||
// An x16 card in a slot wired x8.
|
||||
{PCIeLink{Gen: 4, GPUGen: 4, Width: 8, WidthMax: 16}, "x8 of its x16 lanes: 15.8 GB/s where it could have 31.5"},
|
||||
// A Gen5 card in a Gen4 slot.
|
||||
{PCIeLink{Gen: 4, GPUGen: 5, Width: 16, WidthMax: 16}, "PCIe Gen4 of its Gen5: 31.5 GB/s where it could have 63.0"},
|
||||
// Both, behind a riser.
|
||||
{PCIeLink{Gen: 3, GPUGen: 4, Width: 4, WidthMax: 16}, "x4 of its x16 lanes and PCIe Gen3 of its Gen4"},
|
||||
// A driver that says nothing.
|
||||
{PCIeLink{}, ""},
|
||||
} {
|
||||
got := c.link.Problem()
|
||||
if (c.want == "") != (got == "") || !strings.Contains(got, c.want) {
|
||||
t.Errorf("%+v: %q, want %q", c.link, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheLinkIsReadFromNvidiaSMI(t *testing.T) {
|
||||
links := parsePCIe("0, 3, 16, 16, 3\n1, 4, 8, 16, [N/A]\n")
|
||||
if links[0] != (PCIeLink{Gen: 3, GPUGen: 3, Width: 16, WidthMax: 16}) || links[1] != (PCIeLink{Gen: 4, Width: 8, WidthMax: 16}) {
|
||||
t.Errorf("%+v", links)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheCardsOwnLanesComeFromTheKernel(t *testing.T) {
|
||||
// The driver says x8 of x8 for a x16 card in a slot wired x8; the
|
||||
// device's own max_link_width says 16.
|
||||
root := t.TempDir()
|
||||
dev := filepath.Join(root, "0000:3b:00.0")
|
||||
_ = os.MkdirAll(dev, 0o755)
|
||||
for f, v := range map[string]string{"max_link_width": "16", "current_link_width": "8", "max_link_speed": "16.0 GT/s PCIe"} {
|
||||
_ = os.WriteFile(filepath.Join(dev, f), []byte(v+"\n"), 0o644)
|
||||
}
|
||||
l := withSysfs(PCIeLink{Gen: 4, Width: 8, WidthMax: 8}, root, "00000000:3B:00.0")
|
||||
if l.WidthMax != 16 || l.Width != 8 || l.GPUGen != 4 || !strings.Contains(l.Problem(), "x8 of its x16") {
|
||||
t.Errorf("%+v: %q", l, l.Problem())
|
||||
}
|
||||
}
|
||||
@@ -47,6 +47,9 @@ type Profile struct {
|
||||
// KVEvents has the engine publish its KV cache events to the agent,
|
||||
// for the router's precise index. vLLM only.
|
||||
KVEvents bool `json:"kv_events,omitempty"`
|
||||
// AllowPickle lets the node fetch weights in a pickle format, which run
|
||||
// code when loaded; the catalog's allow_pickle.
|
||||
AllowPickle bool `json:"allow_pickle,omitempty"`
|
||||
}
|
||||
|
||||
// ShardSpec is one shard of one replica on this node.
|
||||
@@ -215,6 +218,7 @@ func NewRegistry(cfg Config) *Registry {
|
||||
r.Add(&TRTLLM{cfg: cfg.TRTLLM})
|
||||
r.Add(&Sim{cfg: cfg.Sim})
|
||||
r.Add(&ComfyUI{cfg: cfg.ComfyUI})
|
||||
r.Add(&Strata{cfg: cfg.Strata})
|
||||
return r
|
||||
}
|
||||
|
||||
@@ -265,6 +269,9 @@ type Config struct {
|
||||
Sim SimConfig `koanf:"sim"`
|
||||
// ComfyUI makes pictures; offered by a node whose ComfyUI is named here.
|
||||
ComfyUI ComfyUIConfig `koanf:"comfyui"`
|
||||
// Strata serves a few mixture-of-experts models from RAM and one card;
|
||||
// offered by a node whose Strata is named here, or that installed it.
|
||||
Strata StrataConfig `koanf:"strata"`
|
||||
}
|
||||
|
||||
// --- shared helpers ---
|
||||
|
||||
@@ -477,14 +477,14 @@ func TestValidateRequiresModelPath(t *testing.T) {
|
||||
|
||||
func TestRegistry(t *testing.T) {
|
||||
r := NewRegistry(Config{})
|
||||
if got := r.Names(); !slices.Equal(got, []string{"comfyui", "llamacpp", "sglang", "sim", "trtllm", "vllm"}) {
|
||||
if got := r.Names(); !slices.Equal(got, []string{"comfyui", "llamacpp", "sglang", "sim", "strata", "trtllm", "vllm"}) {
|
||||
t.Errorf("registry contents: %v", got)
|
||||
}
|
||||
if _, err := r.Get("vllm"); err != nil {
|
||||
t.Errorf("vllm missing: %v", err)
|
||||
}
|
||||
err := func() error { _, err := r.Get("nope"); return err }()
|
||||
if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, trtllm, vllm") {
|
||||
if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, strata, trtllm, vllm") {
|
||||
t.Errorf("an unknown engine should list what is available: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAModelDownloadedByTheEngineIsReadAsSafetensorsOnly(t *testing.T) {
|
||||
for path, want := range map[string]bool{
|
||||
"Qwen/Qwen2.5-7B-Instruct": true,
|
||||
"/var/lib/llmgw/weights/abc/": false, "./local": false,
|
||||
"http://m/q/v/": false, "q.gguf": false,
|
||||
} {
|
||||
got := slices.Contains(safetensorsOnly(Profile{ModelPath: path}), "safetensors")
|
||||
if got != want {
|
||||
t.Errorf("%q: safetensors only = %v, want %v", path, got, want)
|
||||
}
|
||||
}
|
||||
if len(safetensorsOnly(Profile{ModelPath: "Qwen/Qwen2.5-7B-Instruct", AllowPickle: true})) != 0 {
|
||||
t.Error("allow_pickle still forced safetensors")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAForbiddenFlagIsKnownHoweverItIsSpelled(t *testing.T) {
|
||||
for engine, args := range map[string][]string{
|
||||
"vllm": {"--trust_remote_code"},
|
||||
"sglang": {"--load-format", "pt"},
|
||||
} {
|
||||
if UnsafeArgs(engine, args) == nil {
|
||||
t.Errorf("%s %v was let through", engine, args)
|
||||
}
|
||||
}
|
||||
for _, a := range [][]string{{"--load_format", "pt"}, {"--trust-remote"}, {"--load-form=pt"}} {
|
||||
if UnsafeArgs("vllm", a) == nil {
|
||||
t.Errorf("vllm %v was let through", a)
|
||||
}
|
||||
}
|
||||
// What is not one of them, nor the start of one, still passes.
|
||||
for _, a := range [][]string{{"--max-model-len", "8192"}, {"--tokenizer-mode", "mistral"}, {"--gpu-memory-utilization=0.9"}} {
|
||||
if err := UnsafeArgs("vllm", a); err != nil {
|
||||
t.Errorf("vllm %v was refused: %v", a, err)
|
||||
}
|
||||
}
|
||||
// Short flags are matched as written; llama.cpp takes no abbreviations,
|
||||
// so its --chat-template is not taken for --chat-template-file.
|
||||
for _, a := range [][]string{{"-m", "x"}, {"-hf", "o/r"}, {"-mu", "http://x"}} {
|
||||
if UnsafeArgs("llamacpp", a) == nil {
|
||||
t.Errorf("llamacpp %v was let through", a)
|
||||
}
|
||||
}
|
||||
if err := UnsafeArgs("llamacpp", []string{"--chat-template", "chatml"}); err != nil {
|
||||
t.Errorf("llama.cpp's own --chat-template was refused: %v", err)
|
||||
}
|
||||
// ComfyUI's argparse reads --li as --listen.
|
||||
if UnsafeArgs("comfyui", []string{"--li", "0.0.0.0"}) == nil {
|
||||
t.Error("comfyui --li was let through")
|
||||
}
|
||||
}
|
||||
@@ -102,6 +102,7 @@ func (g *SGLang) serverCommand(s ShardSpec) (Command, error) {
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--served-model-name", p.ServedModelName)
|
||||
}
|
||||
args = append(args, safetensorsOnly(p)...)
|
||||
if s.NNodes > 1 {
|
||||
args = append(args,
|
||||
"--nnodes", strconv.Itoa(s.NNodes),
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// StrataConfig locates a Strata installation.
|
||||
type StrataConfig struct {
|
||||
// Dir is Strata's source directory, where setup.py and serve/ are.
|
||||
Dir string `koanf:"dir"`
|
||||
// Python is the interpreter of its environment; Dir's .venv when unset.
|
||||
Python string `koanf:"python"`
|
||||
// Engine is the directory of its prebuilt engine, where the strata
|
||||
// binary and the CUDA libraries it ships are; Dir's engine/ when unset.
|
||||
Engine string `koanf:"engine"`
|
||||
}
|
||||
|
||||
// Strata runs Strata (github.com/Niko1221/Strata): a mixture-of-experts
|
||||
// engine that keeps a model's experts in RAM and the ones used most in VRAM,
|
||||
// so a 125B model serves from one 24 GB card. It runs a few model families
|
||||
// only, each prepared from its GGUF files by Strata's own setup, so the agent
|
||||
// starts llmgw strata-serve, which prepares the model and then becomes
|
||||
// Strata's server.
|
||||
type Strata struct {
|
||||
cfg StrataConfig
|
||||
installed atomic.Pointer[string]
|
||||
}
|
||||
|
||||
// Name identifies the engine in the catalog.
|
||||
func (*Strata) Name() string { return "strata" }
|
||||
|
||||
// HealthPath answers 200 only once the model is loaded.
|
||||
func (*Strata) HealthPath() string { return "/health" }
|
||||
|
||||
// MetricsPath is Strata's Prometheus text, in vLLM's names; its /metrics is
|
||||
// JSON for its own monitor unless asked for this format.
|
||||
func (*Strata) MetricsPath() string { return "/metrics?format=prometheus" }
|
||||
|
||||
// ServesHTTP is true: a Strata replica is one process on one node.
|
||||
func (*Strata) ServesHTTP(ShardSpec) bool { return true }
|
||||
|
||||
// UseInstalled points the engine at the environment the agent installed:
|
||||
// its interpreter, with Strata's code beside its bin in strata/ and the
|
||||
// engine in strata-engine/.
|
||||
func (s *Strata) UseInstalled(python string) { s.installed.Store(&python) }
|
||||
|
||||
// Configured reports whether the configuration names a Strata.
|
||||
func (s *Strata) Configured() bool { return s.cfg.Dir != "" }
|
||||
|
||||
// StrataSourceDir and StrataEngineDir are where an installed environment
|
||||
// keeps Strata's code and its prebuilt engine, beside bin/.
|
||||
const (
|
||||
StrataSourceDir = "strata"
|
||||
StrataEngineDir = "strata-engine"
|
||||
)
|
||||
|
||||
func (s *Strata) env() string {
|
||||
if p := s.installed.Load(); p != nil {
|
||||
return filepath.Dir(filepath.Dir(*p))
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (s *Strata) dir() string {
|
||||
if e := s.env(); e != "" {
|
||||
return filepath.Join(e, StrataSourceDir)
|
||||
}
|
||||
return s.cfg.Dir
|
||||
}
|
||||
|
||||
func (s *Strata) engineDir() string {
|
||||
if e := s.env(); e != "" {
|
||||
return filepath.Join(e, StrataEngineDir)
|
||||
}
|
||||
if s.cfg.Engine != "" {
|
||||
return s.cfg.Engine
|
||||
}
|
||||
return filepath.Join(s.cfg.Dir, "engine")
|
||||
}
|
||||
|
||||
func (s *Strata) python() string {
|
||||
if p := s.installed.Load(); p != nil {
|
||||
return *p
|
||||
}
|
||||
if s.cfg.Python != "" {
|
||||
return s.cfg.Python
|
||||
}
|
||||
return filepath.Join(s.cfg.Dir, ".venv", "bin", "python")
|
||||
}
|
||||
|
||||
// Validate checks the profile: one card, on one node. Strata splits a model
|
||||
// over a node's cards by layer, but that is not driven from here yet.
|
||||
func (*Strata) Validate(p Profile, gpus []GPUInfo) error {
|
||||
if p.Nodes > 1 || p.TP > 1 || p.PP > 1 || p.GPUsPerNode > 1 {
|
||||
return fmt.Errorf("profile %q: strata runs on one GPU of one node", p.Name)
|
||||
}
|
||||
if err := requireComputeCapability(p, gpus, 7.0, "strata"); err != nil {
|
||||
return err
|
||||
}
|
||||
return validateCommon(p, gpus)
|
||||
}
|
||||
|
||||
// Commands is llmgw strata-serve, which prepares the model with Strata's
|
||||
// setup and then execs Strata's server.
|
||||
func (s *Strata) Commands(sh ShardSpec) ([]Command, error) {
|
||||
dir := s.dir()
|
||||
if dir == "" {
|
||||
return nil, fmt.Errorf("strata: no Strata directory is configured (engines.strata.dir)")
|
||||
}
|
||||
bin, err := os.Executable()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p := sh.Profile
|
||||
args := []string{"strata-serve",
|
||||
"--listen", net.JoinHostPort(sh.Host(), strconv.Itoa(sh.HTTPPort)),
|
||||
"--model", p.ServedModelName,
|
||||
"--weights", p.ModelPath,
|
||||
// Relative to where engines run — their user's home — so they are
|
||||
// theirs to write. The data is shared by the node's replicas: a
|
||||
// prepared model and its draft layer are gigabytes, the same for
|
||||
// every replica of it.
|
||||
"--state", filepath.Join("strata", sh.ReplicaID),
|
||||
"--data", filepath.Join("strata", "data"),
|
||||
"--python", s.python(),
|
||||
"--strata", dir,
|
||||
"--engine", s.engineDir(),
|
||||
"--",
|
||||
}
|
||||
args = append(args, p.Args...)
|
||||
env := baseEnv(sh)
|
||||
if sh.APIKey != "" {
|
||||
// Strata's server reads its key from here, and strata-serve passes
|
||||
// it on in the environment: not on a command line every user on the
|
||||
// node can read.
|
||||
env["STRATA_API_KEY"] = sh.APIKey
|
||||
}
|
||||
return []Command{{Name: "strata", Path: bin, Args: args, Env: materializeEnv(env)}}, nil
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestStrataStartsStrataServeWithTheKeyInTheEnvironment(t *testing.T) {
|
||||
s := &Strata{}
|
||||
s.UseInstalled("/opt/llmgw/engines/strata/0.2.7-cuda13/bin/python")
|
||||
spec := ShardSpec{
|
||||
ReplicaID: "r1", HTTPPort: 18090, BindAddr: "10.0.0.147", GPUIndices: []int{0},
|
||||
APIKey: "secret-key",
|
||||
Profile: Profile{Name: "p", Engine: "strata", TP: 1, ModelPath: "/var/lib/llmgw/models/q38",
|
||||
ServedModelName: "qwen3.8-flash-next", Args: []string{"--family", "qwen", "--size", "IQ2_XS"}},
|
||||
}
|
||||
cmds, err := s.Commands(spec)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c := cmds[0]
|
||||
args := strings.Join(c.Args, " ")
|
||||
env := "/opt/llmgw/engines/strata/0.2.7-cuda13"
|
||||
for _, want := range []string{
|
||||
"strata-serve", "--listen 10.0.0.147:18090", "--model qwen3.8-flash-next",
|
||||
"--weights /var/lib/llmgw/models/q38", "--state " + filepath.Join("strata", "r1"),
|
||||
"--python " + env + "/bin/python", "--strata " + env + "/strata", "--engine " + env + "/strata-engine",
|
||||
"-- --family qwen --size IQ2_XS",
|
||||
} {
|
||||
if !strings.Contains(args, want) {
|
||||
t.Errorf("args lack %q: %s", want, args)
|
||||
}
|
||||
}
|
||||
if strings.Contains(args, "secret-key") {
|
||||
t.Error("the key is on the command line, which every user on the node can read")
|
||||
}
|
||||
if !slices.Contains(c.Env, "STRATA_API_KEY=secret-key") {
|
||||
t.Error("the key is not in the environment, where Strata's server reads it")
|
||||
}
|
||||
if !slices.Contains(c.Env, "CUDA_VISIBLE_DEVICES=0") {
|
||||
t.Error("the shard is not pinned to its GPU")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrataRunsOnOneCard(t *testing.T) {
|
||||
s := &Strata{}
|
||||
p := Profile{Name: "p", Engine: "strata", TP: 1, GPUsPerNode: 2, ModelPath: "/m"}
|
||||
if err := s.Validate(p, nil); err == nil {
|
||||
t.Error("a two-card profile was accepted")
|
||||
}
|
||||
p.GPUsPerNode = 1
|
||||
if err := s.Validate(p, []GPUInfo{{Name: "Quadro RTX 6000", ComputeCapability: "7.5"}}); err != nil {
|
||||
t.Errorf("one Turing card was refused: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrataSetupFlagsThatMoveFilesAreRefused(t *testing.T) {
|
||||
for _, bad := range []string{"--gguf-dir", "--gguf", "--data-dir=/etc", "--model", "--mod", "--feed", "--host", "--gpus"} {
|
||||
if err := UnsafeArgs("strata", []string{bad, "x"}); err == nil {
|
||||
t.Errorf("%s was accepted", bad)
|
||||
}
|
||||
}
|
||||
if err := UnsafeArgs("strata", []string{"--family", "qwen", "--size", "IQ2_XS", "--context", "32768",
|
||||
"--parallel", "2", "--vram-reserve-mib", "2048", "--vision", "no"}); err != nil {
|
||||
t.Errorf("an ordinary profile was refused: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
@@ -101,6 +102,13 @@ var flagsNever = map[string]map[string]string{
|
||||
"--hf-token": "is a credential", "--alias": "is set by the agent",
|
||||
"--model-draft": "loads a file", "-md": "loads a file",
|
||||
"--prompt-cache": "writes a file", "--logdir": "writes files",
|
||||
// The short spellings of the same.
|
||||
"-mm": "loads a file; the projector is fetched with the model", "--mmproj-url": "downloads a file",
|
||||
"-mmu": "downloads a file", "-hfr": "downloads a model", "-hff": "downloads a model",
|
||||
"--hf-repo-draft": "downloads a model", "-hfd": "downloads a model", "-hfrd": "downloads a model",
|
||||
"--hf-repo-v": "downloads a model", "-hfv": "downloads a model", "--hf-file-v": "downloads a model",
|
||||
"-hffv": "downloads a model", "-hft": "is a credential", "-jf": "reads a file", "-a": "is set by the agent",
|
||||
"--model-vocoder": "loads a file", "-mv": "loads a file",
|
||||
},
|
||||
"vllm": {
|
||||
"--download-dir": "writes files", "--tokenizer": "reads another repository or path",
|
||||
@@ -118,6 +126,7 @@ var flagsNever = map[string]map[string]string{
|
||||
"--file-storage-path": "writes files", "--log-requests": "logs prompts",
|
||||
"--model-path": "is the model's source, set by the agent", "--served-model-name": "is set by the agent",
|
||||
"--lora-paths": "loads files", "--download-dir": "writes files",
|
||||
"--load-format": "changes how weights are loaded, some ways running code",
|
||||
},
|
||||
// ComfyUI's main.py, behind the gate that sets where it listens and
|
||||
// where its files are.
|
||||
@@ -132,6 +141,22 @@ var flagsNever = map[string]map[string]string{
|
||||
"--tls-keyfile": "reads a file", "--tls-certfile": "reads a file",
|
||||
"--comfy-api-base": "sends work to another service",
|
||||
},
|
||||
// Strata's setup.py, which a profile's arguments go to: llmgw
|
||||
// strata-serve sets where the model is, where its files go and where
|
||||
// the server listens, and the engine is the one the agent installed.
|
||||
// Its own --model, the size, is spelt --size in a profile, since
|
||||
// --model is the source everywhere else.
|
||||
"strata": {
|
||||
"--gguf-dir": "is the model's source, set by the agent", "--models-dir": "is set by the agent",
|
||||
"--data-dir": "is set by the agent", "--move-data": "moves files", "--set-default-data-dir": "writes settings",
|
||||
"--feed": "downloads an engine", "--prebuilt": "loads an engine", "--build": "compiles an engine",
|
||||
"--update": "downloads an engine", "--rollback-engine": "changes the engine",
|
||||
"--dsv41-package": "loads a file", "--inspect": "reads another repository or path",
|
||||
"--source": "downloads from elsewhere", "--setup": "is not for a server",
|
||||
"--browser": "opens a browser", "--gpu": "is set by the agent", "--gpus": "is set by the agent",
|
||||
"--gpu-roles": "is set by the agent", "--backend": "is set by the agent",
|
||||
"--experimental-speed-projection": "changes what the model refuses",
|
||||
},
|
||||
// trtllm-serve spells its flags with underscores, so the shared list's
|
||||
// hyphenated names do not reach them.
|
||||
"trtllm": {
|
||||
@@ -163,14 +188,43 @@ var flagsNever = map[string]map[string]string{
|
||||
// unsafeFlag is why an engine may not be given a flag, or "".
|
||||
func unsafeFlag(engine, flag, _ string) string {
|
||||
name, _, _ := strings.Cut(flag, "=")
|
||||
if !strings.HasPrefix(name, "-") {
|
||||
return ""
|
||||
}
|
||||
// Written as the engine would read it: vLLM takes --load_format for
|
||||
// --load-format, TensorRT-LLM the other way round. Short flags (-m,
|
||||
// -hf) are matched as written.
|
||||
norm := func(f string) string {
|
||||
if strings.HasPrefix(f, "--") {
|
||||
return "--" + strings.ReplaceAll(f[2:], "_", "-")
|
||||
}
|
||||
return f
|
||||
}
|
||||
given := norm(name)
|
||||
// argparse also takes any unambiguous beginning of a long flag for the
|
||||
// flag — --trust-remote is --trust-remote-code, --li is ComfyUI's
|
||||
// --listen — so for the engines that parse with it, a beginning of a
|
||||
// refused flag is refused. llama.cpp and TensorRT-LLM take no such
|
||||
// thing, and --chat-template there is a flag of its own.
|
||||
abbreviates := argparseEngines[engine] && strings.HasPrefix(given, "--") && len(given) >= 4
|
||||
for _, set := range []string{"", engine} {
|
||||
if why, ok := flagsNever[set][name]; ok && why != "" {
|
||||
return why
|
||||
for never, why := range flagsNever[set] {
|
||||
if why == "" {
|
||||
continue
|
||||
}
|
||||
n := norm(never)
|
||||
if n == given || (abbreviates && strings.HasPrefix(n, "--") && strings.HasPrefix(n, given)) {
|
||||
return why
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// argparseEngines read their flags with Python's argparse, which takes a
|
||||
// long flag's beginning for the flag.
|
||||
var argparseEngines = map[string]bool{"vllm": true, "sglang": true, "comfyui": true, "strata": true}
|
||||
|
||||
// UnsafeArgs refuses the first of a profile's arguments an engine may not be
|
||||
// given, saying why.
|
||||
func UnsafeArgs(engine string, args []string) error {
|
||||
@@ -185,3 +239,29 @@ func UnsafeArgs(engine string, args []string) error {
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// safetensorsOnly has vLLM or SGLang read only safetensors from a model it
|
||||
// downloads itself, by its Hugging Face id: left to choose, either takes
|
||||
// the repository's pickle files (pytorch_model.bin) when it has no
|
||||
// safetensors, and loading those runs whatever they hold. A path was
|
||||
// fetched by the agent, which checked it, or put there by the operator.
|
||||
func safetensorsOnly(p Profile) []string {
|
||||
if p.AllowPickle || !hubID(p.ModelPath) {
|
||||
return nil
|
||||
}
|
||||
return []string{"--load-format", "safetensors"}
|
||||
}
|
||||
|
||||
// hubID reports whether a model path is a Hugging Face id, org/name, rather
|
||||
// than a file or directory.
|
||||
func hubID(s string) bool {
|
||||
if s == "" || strings.HasPrefix(s, "/") || strings.HasPrefix(s, ".") || strings.Contains(s, "://") {
|
||||
return false
|
||||
}
|
||||
// org/name:Q4_K_M is a GGUF in a repository, which vLLM reads as one.
|
||||
if strings.Count(s, "/") != 1 || strings.Contains(s, ":") {
|
||||
return false
|
||||
}
|
||||
_, err := os.Stat(s)
|
||||
return err != nil
|
||||
}
|
||||
@@ -119,6 +119,7 @@ func (v *VLLM) serverCommand(s ShardSpec) (Command, error) {
|
||||
if p.ServedModelName != "" {
|
||||
args = append(args, "--served-model-name", p.ServedModelName)
|
||||
}
|
||||
args = append(args, safetensorsOnly(p)...)
|
||||
|
||||
env := baseEnv(s)
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ var (
|
||||
gpuProcesses = discovery.GPUProcesses
|
||||
gpuMemory = discovery.GPUMemoryByProcess
|
||||
processGroup = pgidOf
|
||||
parentOf = ppidOf
|
||||
// gpuCheckWait is how long a shard has, once it answers, to show up on
|
||||
// its GPUs: the weights are on them by then, but nvidia-smi can lag.
|
||||
gpuCheckWait = 15 * time.Second
|
||||
@@ -116,19 +117,44 @@ func (a *Agent) foreignGPUMemory(ctx context.Context, shards []int) (foreign, us
|
||||
return foreign, used
|
||||
}
|
||||
|
||||
// holds reports whether any of pids is the shard's process or in its group:
|
||||
// an engine that computes in a child process is still the shard.
|
||||
// holds reports whether any of pids is the shard's process, in its group,
|
||||
// or a process it started: an engine that computes in a child process is
|
||||
// still the shard, and Strata's server starts its engine in a group of the
|
||||
// engine's own.
|
||||
func holds(pids []int, pid, group int) bool {
|
||||
for _, p := range pids {
|
||||
if p == pid || (group > 0 && processGroup(p) == group) {
|
||||
if p == pid || (group > 0 && processGroup(p) == group) || descends(p, pid) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// descends reports whether process p was started, at any remove, by pid.
|
||||
func descends(p, pid int) bool {
|
||||
// Bounded: a chain this deep is no engine's, and a /proc read racing an
|
||||
// exit must not loop.
|
||||
for range 16 {
|
||||
p = parentOf(p)
|
||||
switch {
|
||||
case p == pid:
|
||||
return true
|
||||
case p <= 1:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// pgidOf reads a process's group from /proc; zero when it cannot.
|
||||
func pgidOf(pid int) int {
|
||||
func pgidOf(pid int) int { return statField(pid, 2) }
|
||||
|
||||
// ppidOf reads a process's parent from /proc; zero when it cannot.
|
||||
func ppidOf(pid int) int { return statField(pid, 1) }
|
||||
|
||||
// statField is one of the numbers after the command name in /proc's stat
|
||||
// line: 1 is the parent, 2 the group.
|
||||
func statField(pid, field int) int {
|
||||
b, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat")
|
||||
if err != nil {
|
||||
return 0
|
||||
@@ -141,10 +167,10 @@ func pgidOf(pid int) int {
|
||||
return 0
|
||||
}
|
||||
f := strings.Fields(s[i+1:])
|
||||
if len(f) < 3 {
|
||||
if len(f) <= field {
|
||||
return 0
|
||||
}
|
||||
n, _ := strconv.Atoi(f[2])
|
||||
n, _ := strconv.Atoi(f[field])
|
||||
return n
|
||||
}
|
||||
|
||||
|
||||
@@ -123,6 +123,30 @@ func TestForeignGPUMemoryLeavesOutTheAgentsOwnShards(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEngineStartedInAGroupOfItsOwnIsStillTheShards(t *testing.T) {
|
||||
// cuda01: Strata's server (259709) starts its engine (261072) in a
|
||||
// group of the engine's own, and the engine holds the card. It read as
|
||||
// a shard on the CPU, and as 22 GB of somebody else's memory.
|
||||
oldMem, oldGroup, oldParent := gpuMemory, processGroup, parentOf
|
||||
t.Cleanup(func() { gpuMemory, processGroup, parentOf = oldMem, oldGroup, oldParent })
|
||||
processGroup = func(pid int) int { return pid }
|
||||
parentOf = func(pid int) int {
|
||||
return map[int]int{261072: 259709, 259709: 258734, 258734: 1, 999: 1}[pid]
|
||||
}
|
||||
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) {
|
||||
return []discovery.GPUMemory{{Index: 0, UsedMB: 23000,
|
||||
Procs: []discovery.GPUProcess{{PID: 261072, UsedMB: 22226}, {PID: 999, UsedMB: 700}}}}, nil
|
||||
}
|
||||
a := &Agent{log: slog.Default()}
|
||||
if c, ok := a.gpuCompute(context.Background(), []int{0}, 259709); !ok || c.Backend != "CUDA" ||
|
||||
len(c.Devices) != 1 || c.Devices[0] != "GPU 0: 21.7 GB" {
|
||||
t.Errorf("the engine's child on GPU 0: %+v %v", c, ok)
|
||||
}
|
||||
if foreign, _ := a.foreignGPUMemory(context.Background(), []int{259709}); !slices.Equal(foreign, []int32{774}) {
|
||||
t.Errorf("foreign %v; want only what the other process and nobody holds, [774]", foreign)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppsOwnWordOnTheCPUIsEnoughToRefuse(t *testing.T) {
|
||||
eng := fallbackEngine{}
|
||||
// Without nvidia-smi to ask, the engine's own output says it.
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
// Half a gigabyte read per test binary, and contending with the tests
|
||||
// for the CPUs: these never look at the figure.
|
||||
measureBandwidth = func() float64 { return 0 }
|
||||
bandwidthDelay = func() time.Duration { return 0 }
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
@@ -174,7 +174,7 @@ func (in *Installer) one(ctx context.Context, w Want) Status {
|
||||
}
|
||||
return st
|
||||
}
|
||||
if w.Offline != nil && recipe.Repo == "" {
|
||||
if w.Offline != nil && (recipe.Repo == "" || builds[0].Python != "") {
|
||||
// Said, not attempted: uv would spend minutes failing to reach PyPI,
|
||||
// and the failure would read like a network fault rather than the
|
||||
// design of an offline site.
|
||||
|
||||
@@ -109,6 +109,11 @@ func (in *Installer) installPython(ctx context.Context, r Recipe, w Want, b Buil
|
||||
}
|
||||
reqs = append(reqs, more...)
|
||||
}
|
||||
if len(b.Assets) > 0 {
|
||||
if err := in.unpackAssets(ctx, r, w, b, tmp, st); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if b.Wheel != nil {
|
||||
// The engine from its GitHub release, pinned to the digest GitHub
|
||||
// publishes for the file: the same trust as an archive, and uv
|
||||
@@ -372,6 +377,48 @@ func (in *Installer) unpackSource(ctx context.Context, w Want, b Build, env stri
|
||||
return requirementsIn(string(f), b.IndexPackages)
|
||||
}
|
||||
|
||||
// unpackAssets fetches a Python engine's release binaries, each checked
|
||||
// against the digest GitHub publishes for it, and unpacks them into the
|
||||
// environment's AssetDir.
|
||||
func (in *Installer) unpackAssets(ctx context.Context, r Recipe, w Want, b Build, env string, st *Status) error {
|
||||
tag := b.AssetTag
|
||||
if tag == "" {
|
||||
tag = w.Version
|
||||
}
|
||||
in.progress(st, "asking GitHub for the checksums of %s %s", r.Repo, tag)
|
||||
sums, err := in.checksums(ctx, r.Repo, tag)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
dst := filepath.Join(env, b.AssetDir)
|
||||
if err := os.MkdirAll(dst, 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, asset := range b.Assets {
|
||||
want, ok := sums[asset]
|
||||
if !ok {
|
||||
// Refused rather than installed unchecked: this is a binary that
|
||||
// will run on every node.
|
||||
return fmt.Errorf("%s %s publishes no checksum for %s, so it cannot be verified", r.Repo, tag, asset)
|
||||
}
|
||||
url := strings.TrimSuffix(in.downloadBase(), "/") + "/" + r.Repo + "/releases/download/" + tag + "/" + asset
|
||||
archive := filepath.Join(env+".source", asset)
|
||||
if err := os.MkdirAll(filepath.Dir(archive), 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := in.download(ctx, url, asset, archive, want, st); err != nil {
|
||||
return err
|
||||
}
|
||||
in.progress(st, "unpacking %s", asset)
|
||||
err := untarGz(archive, dst)
|
||||
_ = os.RemoveAll(filepath.Dir(archive))
|
||||
if err != nil {
|
||||
return fmt.Errorf("unpacking %s: %w", asset, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// requirementsIn reads a requirements file's lines. Options (-r, -e, --index-url
|
||||
// and the like) are refused: they would fetch from somewhere other than PyPI,
|
||||
// or install what the lock never saw.
|
||||
|
||||
@@ -76,6 +76,11 @@ type Build struct {
|
||||
// was published, which PyPI cannot be asked for.
|
||||
Source *Source
|
||||
AsOf time.Time
|
||||
// AssetTag and AssetDir, for a Python engine that also ships a binary
|
||||
// in its releases, are the release tag Assets are taken from, checked
|
||||
// against the digests GitHub publishes, and the directory in the
|
||||
// environment they are unpacked into.
|
||||
AssetTag, AssetDir string
|
||||
// FreeGB is the disk space the installation needs before it starts: a
|
||||
// Python engine is several gigabytes, and one that fills the disk halfway
|
||||
// leaves the node worse off than one that never began.
|
||||
@@ -125,6 +130,7 @@ var Recipes = map[string]Recipe{
|
||||
"vllm": {PyPI: "vllm", Candidates: vllmCandidates},
|
||||
"sglang": {PyPI: "sglang", Candidates: sglangCandidates},
|
||||
"comfyui": {Candidates: comfyuiCandidates},
|
||||
"strata": {Repo: StrataRepo, Candidates: strataCandidates},
|
||||
"uv": {Repo: "astral-sh/uv", Candidates: uvCandidates},
|
||||
}
|
||||
|
||||
@@ -172,6 +178,9 @@ func Builds(engine, version, arch string) ([]Build, error) {
|
||||
continue
|
||||
}
|
||||
for _, b := range bs {
|
||||
if b.Python != "" {
|
||||
return nil, fmt.Errorf("%s is a Python environment from PyPI, which an offline bundle does not carry", engine)
|
||||
}
|
||||
if !seen[b.Variant] {
|
||||
seen[b.Variant] = true
|
||||
out = append(out, b)
|
||||
@@ -672,3 +681,80 @@ print("comfyui version", "` + version + `", "on torch", torch.__version__)`},
|
||||
{cu126, gpuNeeds{minDriverForCUDA12, 5.0}, "ComfyUI " + version + " on CUDA 12.6's torch"},
|
||||
}, h)
|
||||
}
|
||||
|
||||
// --- Strata ---
|
||||
|
||||
// StrataRepo is where Strata is published: its source, and in each release
|
||||
// the prebuilt engine for each CUDA.
|
||||
const StrataRepo = "Niko1221/Strata"
|
||||
|
||||
// strataReleases are the Strata versions this agent installs: the source
|
||||
// archive's digest, taken when it was added here and confirmed from a second
|
||||
// machine, when the version was published, and the engine its release
|
||||
// carries, whose files GitHub publishes digests for.
|
||||
var strataReleases = map[string]struct {
|
||||
sha256 string
|
||||
published string
|
||||
engine string
|
||||
}{
|
||||
"0.2.7": {"41ae861478f5e5d9506864f97dfb5dd7cfe6b83193300e2b0d4f04c0380e3af9", "2026-10-27T14:54:42Z", "0.1.52"},
|
||||
}
|
||||
|
||||
// strataCandidates are the Strata builds for a machine: its source, in a
|
||||
// Python environment, with the prebuilt engine for the node's driver and
|
||||
// cards — built for CUDA 13, Turing to Blackwell; or for CUDA 12.9, Pascal
|
||||
// to Ada, which is what a V100, or a driver older than 580, gets.
|
||||
func strataCandidates(version, variant string, h discovery.Host) ([]Build, error) {
|
||||
if _, err := gpuArch("Strata", h); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if h.Arch != "amd64" {
|
||||
return nil, NotHere{"Strata publishes its Linux CUDA engine for x86-64 only"}
|
||||
}
|
||||
rel, ok := strataReleases[version]
|
||||
if !ok {
|
||||
known := make([]string, 0, len(strataReleases))
|
||||
for v := range strataReleases {
|
||||
known = append(known, v)
|
||||
}
|
||||
slices.Sort(known)
|
||||
return nil, fmt.Errorf("this agent knows the source digest of Strata %s, not of %s; "+
|
||||
"install that one by hand and name it in agent.yaml (engines.strata.dir)",
|
||||
strings.Join(known, ", "), version)
|
||||
}
|
||||
published, err := time.Parse(time.RFC3339, rel.published)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
build := func(cuda string) Build {
|
||||
return Build{
|
||||
Variant: cuda,
|
||||
Python: pythonForEngines,
|
||||
Source: &Source{
|
||||
URL: "/" + StrataRepo + "/archive/refs/tags/v" + version + ".tar.gz",
|
||||
SHA256: rel.sha256, Top: "Strata-" + version, Dir: "strata",
|
||||
Requirements: "requirements.txt",
|
||||
},
|
||||
AsOf: published.Add(24 * time.Hour),
|
||||
Assets: []string{"strata-engine-" + rel.engine + "-linux-x86_64-cuda-" + cuda + ".tar.gz"},
|
||||
AssetTag: "v" + version,
|
||||
AssetDir: "strata-engine",
|
||||
// Strata is started as a script, so what the agent is given is
|
||||
// the environment's interpreter; the proof is importing what its
|
||||
// server needs, and finding its engine beside it.
|
||||
Binary: "python",
|
||||
VersionArgs: []string{"-c", `import os, sys, numpy, jinja2, regex
|
||||
e = os.path.dirname(os.path.dirname(os.path.abspath(sys.executable)))
|
||||
for f in ("strata/setup.py", "strata/serve/server.py", "strata-engine/strata"):
|
||||
assert os.path.isfile(os.path.join(e, f)), f + " is missing"
|
||||
print("strata version", "` + version + `", "engine", "` + rel.engine + `")`},
|
||||
FreeGB: 5,
|
||||
}
|
||||
}
|
||||
return pick("Strata "+version, variant, []candidate{
|
||||
{build("cuda13"), gpuNeeds{minDriverForCUDA13, 7.5}, "Strata " + version + " built for CUDA 13"},
|
||||
// 12.9 rather than 12: its machine code is for 6.0 to 8.9 and its
|
||||
// PTX was compiled by 12.9, so the driver must be 575.
|
||||
{build("cuda12"), gpuNeeds{575, 6.0}, "Strata " + version + " built for CUDA 12.9"},
|
||||
}, h)
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package provision
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
)
|
||||
|
||||
// strataRelease is Strata's source archive as GitHub makes one, and its
|
||||
// release's engine for CUDA 13, with the recipe pointed at the archive.
|
||||
func strataRelease(t *testing.T, rel *release) {
|
||||
t.Helper()
|
||||
archive := tarGz(t, map[string]string{
|
||||
"Strata-0.2.7/setup.py": "print('setup')\n",
|
||||
"Strata-0.2.7/serve/server.py": "print('server')\n",
|
||||
"Strata-0.2.7/requirements.txt": "numpy==2.5.3; python_version >= \"3.12\"\njinja2==3.1.6\n",
|
||||
}, nil)
|
||||
rel.assets["v0.2.7.tar.gz"] = archive
|
||||
sum := sha256.Sum256(archive)
|
||||
engine := tarGz(t, map[string]string{"strata": "\x7fELF", "libcudart.so.13": "lib"}, nil)
|
||||
name := "strata-engine-0.1.52-linux-x86_64-cuda-cuda13.tar.gz"
|
||||
rel.assets[name] = engine
|
||||
esum := sha256.Sum256(engine)
|
||||
rel.digests[name] = "sha256:" + hex.EncodeToString(esum[:])
|
||||
|
||||
old := strataReleases
|
||||
strataReleases = map[string]struct {
|
||||
sha256 string
|
||||
published string
|
||||
engine string
|
||||
}{"0.2.7": {hex.EncodeToString(sum[:]), "2026-10-27T14:54:42Z", "0.1.52"}}
|
||||
t.Cleanup(func() { strataReleases = old })
|
||||
}
|
||||
|
||||
func TestStrataIsInstalledWithItsEngineBesideItsCode(t *testing.T) {
|
||||
sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
|
||||
in, rel := pythonInstaller(t, turingNode, sys)
|
||||
strataRelease(t, rel)
|
||||
|
||||
got := in.Converge(t.Context(), []Want{{Engine: "strata", Version: "0.2.7"}})
|
||||
if got[0].State != Ready {
|
||||
t.Fatalf("not ready: %+v", got[0])
|
||||
}
|
||||
dir := filepath.Join(in.Root, "strata", "0.2.7-cuda13")
|
||||
if got[0].Binary != filepath.Join(dir, "bin", "python") {
|
||||
t.Errorf("binary %s, want the environment's interpreter", got[0].Binary)
|
||||
}
|
||||
for _, f := range []string{"strata/setup.py", "strata/serve/server.py", "strata-engine/strata", "strata-engine/libcudart.so.13"} {
|
||||
if _, err := os.Stat(filepath.Join(dir, f)); err != nil {
|
||||
t.Errorf("%s is not in the environment: %v", f, err)
|
||||
}
|
||||
}
|
||||
all := strings.Join(sys.commands, "\n")
|
||||
for _, want := range []string{"--exclude-newer 2026-10-28T14:54:42Z", "jinja2==3.1.6"} {
|
||||
if !strings.Contains(all, want) {
|
||||
t.Errorf("the installation never used %q:\n%s", want, all)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrataRefusesAnEngineWhoseDigestDoesNotMatch(t *testing.T) {
|
||||
sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
|
||||
in, rel := pythonInstaller(t, turingNode, sys)
|
||||
strataRelease(t, rel)
|
||||
rel.digests["strata-engine-0.1.52-linux-x86_64-cuda-cuda13.tar.gz"] = "sha256:" + strings.Repeat("0", 64)
|
||||
|
||||
got := in.Converge(t.Context(), []Want{{Engine: "strata", Version: "0.2.7", Variant: "cuda13"}})
|
||||
if got[0].State != Failed || !strings.Contains(got[0].Detail, "does not match the checksum") {
|
||||
t.Fatalf("an engine that does not match its digest was not refused: %+v", got[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrataPicksTheEngineForTheCardsAndDriver(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
host discovery.Host
|
||||
want []string
|
||||
err string
|
||||
}{
|
||||
{"turing on 580", turingNode, []string{"cuda13", "cuda12"}, ""},
|
||||
{"volta", discovery.Host{Arch: "amd64", NVIDIADriver: "580.82.07", MinComputeCapability: 7.0}, []string{"cuda12"}, ""},
|
||||
{"driver 575", discovery.Host{Arch: "amd64", NVIDIADriver: "575.51.03", MinComputeCapability: 8.6}, []string{"cuda12"}, ""},
|
||||
{"driver 570", discovery.Host{Arch: "amd64", NVIDIADriver: "570.86.10", MinComputeCapability: 8.6}, nil, "575"},
|
||||
{"no gpu", discovery.Host{Arch: "amd64"}, nil, "no"},
|
||||
{"arm", discovery.Host{Arch: "arm64", NVIDIADriver: "580.82.07", MinComputeCapability: 9.0}, nil, "x86-64"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
bs, err := strataCandidates("0.2.7", "", tc.host)
|
||||
if tc.err != "" {
|
||||
if err == nil || !strings.Contains(err.Error(), tc.err) {
|
||||
t.Fatalf("got %v, %v; want an error saying %q", bs, err, tc.err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got []string
|
||||
for _, b := range bs {
|
||||
got = append(got, b.Variant)
|
||||
}
|
||||
if strings.Join(got, ",") != strings.Join(tc.want, ",") {
|
||||
t.Errorf("builds %v, want %v", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrataIsNotInAnOfflineBundle(t *testing.T) {
|
||||
if _, err := Builds("strata", "0.2.7", "amd64"); err == nil {
|
||||
t.Fatal("an offline bundle would carry Strata, which needs PyPI")
|
||||
}
|
||||
}
|
||||
@@ -90,14 +90,20 @@ var ErrNotLocal = errors.New("the weights are not on this node and the source sa
|
||||
// already cached, downloaded, or downloaded by another shard this one waited
|
||||
// for — since that is most of the difference between a replica that starts
|
||||
// in a minute and one that takes an hour.
|
||||
func (f *Fetcher) Ensure(ctx context.Context, source string) (path string, err error) {
|
||||
//
|
||||
// A source in a pickle format, or a set with such a file in it, is refused
|
||||
// before anything is fetched unless allowPickle; see Pickled.
|
||||
func (f *Fetcher) Ensure(ctx context.Context, source string, allowPickle bool) (path string, err error) {
|
||||
ctx, span := obs.Tracer(tracerName).Start(ctx, "fetch weights",
|
||||
trace.WithAttributes(attribute.String("llmgw.weights.source", redact(source))))
|
||||
defer func() {
|
||||
span.SetAttributes(attribute.String("llmgw.weights.path", path))
|
||||
obs.End(span, err)
|
||||
}()
|
||||
return f.ensure(ctx, source)
|
||||
if !allowPickle && !IsSet(source) && Pickled(source) {
|
||||
return "", refusePickle(source)
|
||||
}
|
||||
return f.ensure(ctx, source, allowPickle)
|
||||
}
|
||||
|
||||
// tracerName names the spans this package opens.
|
||||
@@ -119,7 +125,7 @@ func redact(source string) string {
|
||||
return u.String()
|
||||
}
|
||||
|
||||
func (f *Fetcher) ensure(ctx context.Context, source string) (string, error) {
|
||||
func (f *Fetcher) ensure(ctx context.Context, source string, allowPickle bool) (string, error) {
|
||||
if source == "" {
|
||||
return "", fmt.Errorf("no source given")
|
||||
}
|
||||
@@ -127,7 +133,7 @@ func (f *Fetcher) ensure(ctx context.Context, source string) (string, error) {
|
||||
// file, and a set has to be handled before any of it.
|
||||
if IsSet(source) {
|
||||
outcome(ctx, "set")
|
||||
return f.ensureSet(ctx, source)
|
||||
return f.ensureSet(ctx, source, allowPickle)
|
||||
}
|
||||
if !isRemote(source) {
|
||||
if !isPath(source) {
|
||||
|
||||
@@ -37,7 +37,7 @@ func TestALocalPathIsUsedNeverCopied(t *testing.T) {
|
||||
}
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
|
||||
got, err := f.Ensure(t.Context(), path)
|
||||
got, err := f.Ensure(t.Context(), path, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -53,7 +53,7 @@ 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"))
|
||||
_, err := f.Ensure(t.Context(), filepath.Join(t.TempDir(), "absent.gguf"), false)
|
||||
if !errors.Is(err, ErrNotLocal) {
|
||||
t.Fatalf("err = %v, want ErrNotLocal", err)
|
||||
}
|
||||
@@ -67,7 +67,7 @@ func TestARemoteFileIsFetchedOnceAndReused(t *testing.T) {
|
||||
srv := serveBytes(t, []byte("the weights"), &hits)
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
|
||||
first, err := f.Ensure(t.Context(), srv.URL+"/model.gguf")
|
||||
first, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -76,7 +76,7 @@ func TestARemoteFileIsFetchedOnceAndReused(t *testing.T) {
|
||||
t.Fatalf("fetched content = %q, err %v", body, err)
|
||||
}
|
||||
|
||||
second, err := f.Ensure(t.Context(), srv.URL+"/model.gguf")
|
||||
second, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -103,7 +103,7 @@ func TestConcurrentRequestsProduceOneDownload(t *testing.T) {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
paths[i], errs[i] = f.Ensure(t.Context(), srv.URL+"/model.gguf")
|
||||
paths[i], errs[i] = f.Ensure(t.Context(), srv.URL+"/model.gguf", false)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
@@ -139,7 +139,7 @@ func TestAPartialDownloadIsNeverVisibleAsTheModel(t *testing.T) {
|
||||
|
||||
cache := t.TempDir()
|
||||
f := &Fetcher{CacheDir: cache}
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/model.gguf"); err == nil {
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false); err == nil {
|
||||
t.Fatal("a truncated download was reported as success")
|
||||
}
|
||||
entries, _ := os.ReadDir(cache)
|
||||
@@ -169,7 +169,7 @@ func TestAnInterruptedDownloadResumes(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got, err := f.Ensure(t.Context(), srv.URL+"/model.gguf")
|
||||
got, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -232,7 +232,7 @@ func TestProgressIsReportedWithAPercentage(t *testing.T) {
|
||||
CacheDir: t.TempDir(),
|
||||
OnProgress: func(p Progress) { seen = append(seen, p) },
|
||||
}
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/model.gguf"); err != nil {
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(seen) == 0 {
|
||||
@@ -254,7 +254,7 @@ func TestPercentIsUnknownRatherThanZeroWithoutATotal(t *testing.T) {
|
||||
|
||||
func TestNoCacheDirectoryIsAClearRefusal(t *testing.T) {
|
||||
f := &Fetcher{}
|
||||
_, err := f.Ensure(t.Context(), "https://example.invalid/model.gguf")
|
||||
_, err := f.Ensure(t.Context(), "https://example.invalid/model.gguf", false)
|
||||
if err == nil || !strings.Contains(err.Error(), "cache directory") {
|
||||
t.Errorf("err = %v, want a complaint about the cache directory", err)
|
||||
}
|
||||
@@ -269,7 +269,7 @@ func TestABareModelIdIsLeftToTheEngine(t *testing.T) {
|
||||
"Qwen/Qwen2.5-7B-Instruct",
|
||||
"meta-llama/Llama-3.3-70B-Instruct",
|
||||
} {
|
||||
got, err := f.Ensure(t.Context(), id)
|
||||
got, err := f.Ensure(t.Context(), id, false)
|
||||
if err != nil {
|
||||
t.Errorf("%s was refused: %v", id, err)
|
||||
}
|
||||
@@ -283,7 +283,7 @@ 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) {
|
||||
if _, err := f.Ensure(t.Context(), "/models/definitely-not-here.gguf", false); !errors.Is(err, ErrNotLocal) {
|
||||
t.Errorf("err = %v, want ErrNotLocal", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package weights
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"path"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Weights in Python's pickle format are a program as much as a file: loading
|
||||
// one runs whatever it says, as the engine's user, on the node. That is how
|
||||
// the malicious models found on Hugging Face work, and nothing in the file's
|
||||
// name or size says so. safetensors and GGUF are only tensors and are read,
|
||||
// never run, so those are what a node fetches unless the catalog says
|
||||
// otherwise for a model, with allow_pickle, on purpose.
|
||||
|
||||
// pickleExts are the extensions of PyTorch's pickle formats, and of pickle
|
||||
// itself. A .bin beside safetensors is the same model in pickle form.
|
||||
var pickleExts = []string{".bin", ".pt", ".pt2", ".pth", ".ckpt", ".pkl", ".pickle"}
|
||||
|
||||
// Pickled reports whether a file name or URL is in a pickle format.
|
||||
func Pickled(name string) bool {
|
||||
// A query or a fragment is not the file's name, on a URL or an hf:
|
||||
// source alike: model.bin?download=true is still model.bin.
|
||||
if i := strings.IndexAny(name, "?#"); i >= 0 {
|
||||
name = name[:i]
|
||||
}
|
||||
ext := strings.ToLower(path.Ext(name))
|
||||
for _, e := range pickleExts {
|
||||
if ext == e {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ErrPickle is a source refused for being in a pickle format.
|
||||
var ErrPickle = errors.New("weights in a pickle format run code when they are loaded")
|
||||
|
||||
func refusePickle(name string) error {
|
||||
return fmt.Errorf("%w: %s is refused; use the model's safetensors or GGUF, or set allow_pickle: true "+
|
||||
"on the model if it was checked and has no other form", ErrPickle, redact(name))
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package weights
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPickleFormatsAreKnownByTheirNames(t *testing.T) {
|
||||
for name, want := range map[string]bool{
|
||||
"pytorch_model.bin": true, "model.ckpt": true, "x.PT": true, "w.pth": true, "a.pkl": true,
|
||||
"https://h/r/resolve/main/model.ckpt?download=true": true,
|
||||
"hf:o/r/model.bin?download=true": true, "hf:o/r/model.ckpt#x": true, "m.pt2": true,
|
||||
"model.safetensors": false, "m-Q4_K_M.gguf": false, "Qwen/Qwen2.5-7B-Instruct": false,
|
||||
"config.json": false, "": false,
|
||||
} {
|
||||
if got := Pickled(name); got != want {
|
||||
t.Errorf("Pickled(%q) = %v, want %v", name, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPickleFileIsRefusedBeforeItIsFetched(t *testing.T) {
|
||||
var hits atomic.Int64
|
||||
srv := serveBytes(t, []byte("not really a pickle"), &hits)
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
|
||||
_, err := f.Ensure(t.Context(), srv.URL+"/model.ckpt", false)
|
||||
if !errors.Is(err, ErrPickle) {
|
||||
t.Fatalf("a .ckpt was not refused: %v", err)
|
||||
}
|
||||
if hits.Load() != 0 {
|
||||
t.Errorf("the refused file was fetched %d times", hits.Load())
|
||||
}
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/model.ckpt", true); err != nil {
|
||||
t.Errorf("allowed, the .ckpt was still refused: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASetWithAPickleFileIsRefusedWhole(t *testing.T) {
|
||||
f, srv, hits := mirrorOf(t, map[string]string{
|
||||
"pytorch_model.bin": "pickle",
|
||||
"config.json": "{}",
|
||||
})
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); !errors.Is(err, ErrPickle) {
|
||||
t.Fatalf("a set holding pytorch_model.bin was not refused: %v", err)
|
||||
}
|
||||
if *hits != 0 {
|
||||
t.Errorf("%d files of the refused set were fetched", *hits)
|
||||
}
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", true); err != nil {
|
||||
t.Errorf("allowed, the set was still refused: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASetAlreadyHereIsLookedAtAgain(t *testing.T) {
|
||||
// Fetched once for a model that allowed pickles, or before they were
|
||||
// refused: another asking for it without leave is refused all the same.
|
||||
f, srv, _ := mirrorOf(t, map[string]string{"pytorch_model.bin": "pickle", "config.json": "{}"})
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); !errors.Is(err, ErrPickle) {
|
||||
t.Errorf("a set already on the node was not looked at: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -46,7 +46,7 @@ type manifest struct {
|
||||
// The whole set or nothing. A directory with three of four shards in it loads
|
||||
// as a broken model rather than failing to load, so a partial set is left
|
||||
// unfinished under its temporary name and the next attempt starts again.
|
||||
func (f *Fetcher) ensureSet(ctx context.Context, source string) (string, error) {
|
||||
func (f *Fetcher) ensureSet(ctx context.Context, source string, allowPickle bool) (string, error) {
|
||||
dir, err := f.cachePath(strings.TrimSuffix(source, "/"))
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -56,6 +56,16 @@ func (f *Fetcher) ensureSet(ctx context.Context, source string) (string, error)
|
||||
// exactly the failure the marker exists to prevent.
|
||||
done := filepath.Join(dir, ".complete")
|
||||
if _, err := os.Stat(done); err == nil {
|
||||
// Already here, maybe fetched for a model that allowed pickles, or
|
||||
// before they were refused: looked at again, not trusted.
|
||||
if !allowPickle {
|
||||
entries, _ := os.ReadDir(dir)
|
||||
for _, e := range entries {
|
||||
if Pickled(e.Name()) {
|
||||
return "", refusePickle(source + e.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
return dir, nil
|
||||
}
|
||||
|
||||
@@ -66,6 +76,15 @@ func (f *Fetcher) ensureSet(ctx context.Context, source string) (string, error)
|
||||
if len(m.Files) == 0 {
|
||||
return "", fmt.Errorf("%s holds no files", source)
|
||||
}
|
||||
// Refused whole rather than skipped: an engine given the directory
|
||||
// loads the pickle when there is no safetensors beside it.
|
||||
if !allowPickle {
|
||||
for _, file := range m.Files {
|
||||
if Pickled(file.Name) {
|
||||
return "", refusePickle(source + file.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return "", err
|
||||
|
||||
@@ -53,7 +53,7 @@ func TestASetIsFetchedIntoOneDirectory(t *testing.T) {
|
||||
"config.json": "{}",
|
||||
})
|
||||
|
||||
got, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
|
||||
got, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -74,11 +74,11 @@ func TestASetIsFetchedIntoOneDirectory(t *testing.T) {
|
||||
func TestAFinishedSetIsNotFetchedAgain(t *testing.T) {
|
||||
f, srv, hits := mirrorOf(t, map[string]string{"a.gguf": "aaaa", "b.gguf": "bb"})
|
||||
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err != nil {
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
after := *hits
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err != nil {
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if *hits != after {
|
||||
@@ -100,7 +100,7 @@ func TestAHalfFetchedSetIsNotTreatedAsFinished(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
|
||||
got, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -117,7 +117,7 @@ func TestAFileAlreadyThereAtTheRightSizeIsNotRefetched(t *testing.T) {
|
||||
_ = os.MkdirAll(dir, 0o755)
|
||||
_ = os.WriteFile(filepath.Join(dir, "a.gguf"), []byte("aaaa"), 0o644)
|
||||
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err != nil {
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Resuming a set means not paying for the shards that already landed.
|
||||
@@ -133,7 +133,7 @@ func TestAFilenameThatEscapesTheDirectoryIsRefused(t *testing.T) {
|
||||
defer srv.Close()
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
|
||||
_, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
|
||||
_, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
|
||||
// The mirror is a place an operator copies files into by hand, so what it
|
||||
// reports is not automatically safe to join onto a path.
|
||||
if err == nil || !strings.Contains(err.Error(), "offered a file named") {
|
||||
@@ -148,7 +148,7 @@ func TestADirectoryListingIsNotAManifest(t *testing.T) {
|
||||
defer srv.Close()
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
|
||||
_, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
|
||||
_, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
|
||||
// "invalid character '<'" sends somebody looking at the wrong thing.
|
||||
if err == nil || !strings.Contains(err.Error(), "has to be a mirror version") {
|
||||
t.Errorf("the reason should name the cause, got %v", err)
|
||||
@@ -182,12 +182,12 @@ func TestASetFileThatDoesNotMatchTheMirrorsChecksumIsRefused(t *testing.T) {
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err == nil || !strings.Contains(err.Error(), "sha256") {
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err == nil || !strings.Contains(err.Error(), "sha256") {
|
||||
t.Fatalf("a damaged shard was accepted: %v", err)
|
||||
}
|
||||
served = good
|
||||
// The size differs, so it is fetched again; this time it matches.
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err != nil {
|
||||
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
|
||||
t.Fatalf("the shard as mirrored: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -198,7 +198,7 @@ func TestASetStillDownloadingOnTheMirrorSaysSo(t *testing.T) {
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
_, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
|
||||
_, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
|
||||
if err == nil || !strings.Contains(err.Error(), "not complete on the mirror yet") ||
|
||||
!strings.Contains(err.Error(), "model-00002-of-00002") {
|
||||
t.Fatalf("got %v", err)
|
||||
|
||||
@@ -13,7 +13,7 @@ func TestAFetchSaysInItsTraceWhetherItDownloadedAndHowMuch(t *testing.T) {
|
||||
f := &Fetcher{CacheDir: t.TempDir()}
|
||||
src := strings.Replace(srv.URL, "http://", "http://user:secret@", 1) + "/m/model.gguf?token=hunter2"
|
||||
|
||||
if _, err := f.Ensure(t.Context(), src); err != nil {
|
||||
if _, err := f.Ensure(t.Context(), src, false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fetch := rec.One(t, [16]byte{}, "fetch weights")
|
||||
@@ -30,7 +30,7 @@ func TestAFetchSaysInItsTraceWhetherItDownloadedAndHowMuch(t *testing.T) {
|
||||
}
|
||||
|
||||
// The second start on this node finds it here.
|
||||
if _, err := f.Ensure(t.Context(), src); err != nil {
|
||||
if _, err := f.Ensure(t.Context(), src, false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
again := rec.Named("fetch weights")
|
||||
|
||||
@@ -51,6 +51,7 @@ func newMirrorCmd() *cobra.Command {
|
||||
newMirrorRemoveCmd(&dir),
|
||||
newMirrorServeCmd(&dir),
|
||||
newMirrorInitCmd(&dir),
|
||||
newMirrorAddPeerCmd(&dir),
|
||||
)
|
||||
return c
|
||||
}
|
||||
@@ -255,6 +256,7 @@ func newMirrorServeCmd(dir *string) *cobra.Command {
|
||||
var tlsc tlsconf.Config
|
||||
var allow, clients []string
|
||||
var allowPrivate bool
|
||||
var peers []string
|
||||
|
||||
c := &cobra.Command{
|
||||
Use: "serve",
|
||||
@@ -317,6 +319,21 @@ func newMirrorServeCmd(dir *string) *cobra.Command {
|
||||
ReadTimeout: time.Minute,
|
||||
})
|
||||
|
||||
// Another mirror kept in step with this one, each copying what
|
||||
// the other has: the cluster's models while one machine is down.
|
||||
if len(peers) > 0 {
|
||||
p := &mirror.Peers{M: mgr, Log: log}
|
||||
for _, addr := range peers {
|
||||
c, err := mirror.NewClient(addr, tlsc, 30*time.Second)
|
||||
if err != nil {
|
||||
return fmt.Errorf("--peer %s: %w", addr, err)
|
||||
}
|
||||
p.Clients = append(p.Clients, c)
|
||||
}
|
||||
go p.Run(cmd.Context())
|
||||
fmt.Fprintf(out, "kept in step with %s, every %s\n", strings.Join(peers, ", "), mirror.PeerInterval)
|
||||
}
|
||||
|
||||
errs := make(chan error, 2)
|
||||
go func() { errs <- fmt.Errorf("serving weights: %w", files.ListenAndServe()) }()
|
||||
fmt.Fprintf(out, "serving %s on %s\n", *dir, listen)
|
||||
@@ -343,6 +360,9 @@ func newMirrorServeCmd(dir *string) *cobra.Command {
|
||||
},
|
||||
}
|
||||
c.Flags().StringVar(&listen, "listen", ":8081", "address to serve the weights on")
|
||||
c.Flags().StringArrayVar(&peers, "peer", nil,
|
||||
"another mirror's management API to keep in step with, each copying what the other has; repeatable. "+
|
||||
"Give each the other's address in --api-allow (or its certificate in --api-client)")
|
||||
c.Flags().StringVar(&apiListen, "api-listen", ":8082", "address to serve the management API on")
|
||||
c.Flags().StringVar(&base, "base", "http://localhost:8081",
|
||||
"the address nodes use to reach this mirror, as it should appear in catalog sources")
|
||||
|
||||
@@ -0,0 +1,209 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
|
||||
)
|
||||
|
||||
// A second mirror, so models can still be added and fetched while the first
|
||||
// machine is down. Set up from the machine of the first, which it is kept in
|
||||
// step with: the first is told of it, it of the first, and the control
|
||||
// planes of both.
|
||||
|
||||
func newMirrorAddPeerCmd(dir *string) *cobra.Command {
|
||||
var controlPlanes []string
|
||||
var base string
|
||||
cmd := &cobra.Command{
|
||||
Use: "add-peer [user@]host",
|
||||
Short: "Set up a second mirror on another machine, kept in step with this one",
|
||||
Long: `Run on the machine of the cluster's mirror, as root. It logs in to HOST with
|
||||
ssh, installs llmgw there if it is not, and sets up a mirror that keeps in
|
||||
step with this one: each copies what the other holds, checksum checked, and
|
||||
a model removed from one is removed from the other. This mirror is told of
|
||||
it, and the cluster's settings name it (mirror.peers), so while either
|
||||
machine is down the control planes use the other and nodes fetch from it.
|
||||
|
||||
HOST needs room for every model this mirror holds: it copies them all.`,
|
||||
Example: " sudo llmgw mirror add-peer ubuntu@10.0.0.12 --control-plane 10.0.0.10 --control-plane 10.0.0.12",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if os.Geteuid() != 0 {
|
||||
return errors.New("add-peer changes this mirror's service and the cluster's settings: run it as root")
|
||||
}
|
||||
user, host, ok := strings.Cut(args[0], "@")
|
||||
if !ok {
|
||||
user, host = os.Getenv("SUDO_USER"), args[0]
|
||||
}
|
||||
return addMirrorPeer(cmd.Context(), cmd.OutOrStdout(), &onboard.SSH{Addr: host, User: user}, host,
|
||||
*dir, base, controlPlanes, "/etc/llmgw/server.yaml")
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringArrayVar(&controlPlanes, "control-plane", nil,
|
||||
"address of a control plane that manages the mirrors; repeatable (default: this machine)")
|
||||
cmd.Flags().StringVar(&base, "base", "", "the address nodes fetch from on HOST (default: http://HOST:8081)")
|
||||
return cmd
|
||||
}
|
||||
|
||||
// addMirrorPeer sets up a mirror on a machine reached by remote, kept in
|
||||
// step with the one on this machine.
|
||||
func addMirrorPeer(ctx context.Context, out io.Writer, remote onboard.Remote, host, dir, base string,
|
||||
controlPlanes []string, configPath string,
|
||||
) error {
|
||||
ip, err := ipv4Of(host)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
here := primaryAddr()
|
||||
if len(controlPlanes) == 0 {
|
||||
controlPlanes = []string{here}
|
||||
}
|
||||
if base == "" {
|
||||
base = "http://" + net.JoinHostPort(ip, "8081")
|
||||
}
|
||||
thisAPI := "http://" + net.JoinHostPort(here, "8082")
|
||||
peerAPI := "http://" + net.JoinHostPort(ip, "8082")
|
||||
|
||||
// llmgw there, the same as here, when it has none.
|
||||
if _, err := remote.Run(ctx, "command -v llmgw >/dev/null", nil); err != nil {
|
||||
exe, err := os.Executable()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
bin, err := os.ReadFile(exe) //nolint:gosec // this program
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(out, " .. llmgw to %s\n", host)
|
||||
if _, err := remote.Run(ctx, "sudo install -m 755 /dev/stdin /usr/local/bin/llmgw", bin); err != nil {
|
||||
return fmt.Errorf("installing llmgw on %s: %w", host, err)
|
||||
}
|
||||
}
|
||||
|
||||
// The mirror there: its management API for the control planes and this
|
||||
// mirror, which reads it; kept in step with this one.
|
||||
allow := append([]string{}, controlPlanes...)
|
||||
if !slices.Contains(allow, here) {
|
||||
allow = append(allow, here)
|
||||
}
|
||||
script := "sudo llmgw mirror init --dir " + shellQuote(dir) + " --base " + shellQuote(base) + " --peer " + shellQuote(thisAPI)
|
||||
for _, a := range allow {
|
||||
script += " --control-plane " + shellQuote(a)
|
||||
}
|
||||
fmt.Fprintf(out, " .. the mirror on %s\n", host)
|
||||
if b, err := remote.Run(ctx, script, nil); err != nil {
|
||||
return fmt.Errorf("the mirror on %s: %w\n%s", host, err, b)
|
||||
}
|
||||
|
||||
// This mirror told of it, and to let it read this one's API.
|
||||
changed, err := addToMirrorUnit(mirrorUnitPath, peerAPI, ip)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if changed {
|
||||
in := &serverInit{out: out}
|
||||
if err := in.systemctl(ctx, "daemon-reload"); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := in.systemctl(ctx, "restart", "llmgw-mirror"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(out, " ok this mirror keeps in step with %s\n", peerAPI)
|
||||
|
||||
// The control planes told of both.
|
||||
if err := addMirrorPeerSetting(configPath, peerAPI); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(out, " ok mirror.peers names %s; the control planes use it while this one is down\n", peerAPI)
|
||||
return nil
|
||||
}
|
||||
|
||||
// addToMirrorUnit adds a peer and the address allowed to read this mirror's
|
||||
// API to its systemd unit, and says whether it changed.
|
||||
func addToMirrorUnit(path, peer, allow string) (bool, error) {
|
||||
b, err := os.ReadFile(path) //nolint:gosec // the mirror's unit
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("this machine has no mirror to keep in step (%s): run add-peer on the mirror's machine", path)
|
||||
}
|
||||
unit := string(b)
|
||||
re := regexp.MustCompile(`(?m)^ExecStart=.*$`)
|
||||
line := re.FindString(unit)
|
||||
if line == "" {
|
||||
return false, fmt.Errorf("%s has no ExecStart", path)
|
||||
}
|
||||
next := line
|
||||
if !strings.Contains(next, "--peer "+peer) {
|
||||
next += " --peer " + peer
|
||||
}
|
||||
if !strings.Contains(next, "--api-allow "+allow+" ") && !strings.HasSuffix(next, "--api-allow "+allow) &&
|
||||
strings.Contains(next, "--api-listen :") {
|
||||
next += " --api-allow " + allow
|
||||
}
|
||||
if next == line {
|
||||
return false, nil
|
||||
}
|
||||
return true, os.WriteFile(path, []byte(strings.Replace(unit, line, next, 1)), 0o644)
|
||||
}
|
||||
|
||||
// addMirrorPeerSetting names a peer in mirror.peers: in the cluster's
|
||||
// settings when the store holds them, else in server.yaml.
|
||||
func addMirrorPeerSetting(path, peer string) error {
|
||||
add := func(doc *settings.Doc) error {
|
||||
var peers []string
|
||||
if n := doc.Node([]string{"mirror", "peers"}); n != nil {
|
||||
_ = n.Decode(&peers)
|
||||
}
|
||||
if slices.Contains(peers, peer) {
|
||||
return nil
|
||||
}
|
||||
list := "[" + strings.Join(append(peers, peer), ", ") + "]"
|
||||
return doc.SetYAML([]string{"mirror", "peers"}, list)
|
||||
}
|
||||
changed, err := changeStoredSettings(path, "a second mirror (llmgw mirror add-peer)", add)
|
||||
if err != nil || changed {
|
||||
return err
|
||||
}
|
||||
doc, err := settings.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := add(doc); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = doc.Save(nil)
|
||||
return err
|
||||
}
|
||||
|
||||
// ensureSecondMirror sets up the mirror kept in step with this machine's on
|
||||
// the store machine chosen, reached as the store's setup reached it.
|
||||
func (in *serverInit) ensureSecondMirror(ctx context.Context) error {
|
||||
if !in.standbys {
|
||||
return errors.New("it needs --standbys: without them this mirror's API answers this machine alone, which the second cannot read")
|
||||
}
|
||||
i := in.secondMirror - 1
|
||||
if i < 0 || i >= len(in.patroniNodes) {
|
||||
return fmt.Errorf("there is no database machine %d to put it on", in.secondMirror)
|
||||
}
|
||||
if in.noMirror || in.mirrorURL != "" {
|
||||
return errors.New("it keeps in step with a mirror on this machine, and there is none")
|
||||
}
|
||||
n := in.patroniNodes[i]
|
||||
fmt.Fprintf(in.out, " .. a second mirror on %s\n", n.Name)
|
||||
if err := addMirrorPeer(ctx, in.out, n.Remote, n.IP, in.mirrorRoot(), "", in.mirrorAllow(), in.configPath()); err != nil {
|
||||
return err
|
||||
}
|
||||
in.done("a second mirror on %s (%s), kept in step with this one", n.Name, n.IP)
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAPeerIsAddedToTheMirrorsUnitOnce(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "llmgw-mirror.service")
|
||||
unit := mirrorUnitFor("/srv/models", ":8082", "http://10.0.0.10:8081", []string{"10.0.0.10"}, nil, "test")
|
||||
if err := os.WriteFile(path, []byte(unit), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
changed, err := addToMirrorUnit(path, "http://10.0.0.12:8082", "10.0.0.12")
|
||||
if err != nil || !changed {
|
||||
t.Fatalf("changed %v %v", changed, err)
|
||||
}
|
||||
b, _ := os.ReadFile(path)
|
||||
if !strings.Contains(string(b), "--peer http://10.0.0.12:8082") || !strings.Contains(string(b), "--api-allow 10.0.0.12") {
|
||||
t.Errorf("unit:\n%s", b)
|
||||
}
|
||||
if changed, _ := addToMirrorUnit(path, "http://10.0.0.12:8082", "10.0.0.12"); changed {
|
||||
t.Error("added twice")
|
||||
}
|
||||
// A mirror whose API is on loopback alone is not opened by the back door.
|
||||
lo := mirrorUnitFor("/srv/models", "127.0.0.1:8082", "http://10.0.0.10:8081", nil, nil, "test")
|
||||
_ = os.WriteFile(path, []byte(lo), 0o644)
|
||||
_, _ = addToMirrorUnit(path, "http://10.0.0.12:8082", "10.0.0.12")
|
||||
if b, _ := os.ReadFile(path); strings.Contains(string(b), "--api-allow") {
|
||||
t.Errorf("a loopback mirror was opened:\n%s", b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASecondMirrorIsNamedInTheSettings(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "server.yaml")
|
||||
if err := os.WriteFile(path, []byte("listen: \":8080\"\nmirror:\n api: http://10.0.0.10:8082\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for range 2 {
|
||||
if err := addMirrorPeerSetting(path, "http://10.0.0.12:8082"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
b, _ := os.ReadFile(path)
|
||||
if strings.Count(string(b), "10.0.0.12:8082") != 1 || !strings.Contains(string(b), "peers") {
|
||||
t.Errorf("server.yaml:\n%s", b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheSetupOffersASecondMirrorOnAStoreMachine(t *testing.T) {
|
||||
in := newDefaultServerInit()
|
||||
f := newInstallForm(&in, "ctl", []string{"10.0.0.5"}, true)
|
||||
f.secondMirror.Choice = 1
|
||||
if !f.secondMirror.Show() {
|
||||
t.Fatal("not offered with the mirror on this machine")
|
||||
}
|
||||
f.mirrorWhere.Choice = 1
|
||||
if f.secondMirror.Show() {
|
||||
t.Error("offered with the mirror elsewhere")
|
||||
}
|
||||
f.mirrorWhere.Choice = 0
|
||||
in.secondMirror = 1
|
||||
in.patroni = []string{"ubuntu@10.0.0.11", "ubuntu@10.0.0.12"}
|
||||
if got := strings.Join(in.initFlags(), " "); !strings.Contains(got, "--second-mirror 1") {
|
||||
t.Errorf("flags %s", got)
|
||||
}
|
||||
// Without standbys this mirror answers itself alone: the second could not read it.
|
||||
in.standbys = false
|
||||
if err := in.ensureSecondMirror(t.Context()); err == nil || !strings.Contains(err.Error(), "standbys") {
|
||||
t.Errorf("without standbys: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -29,6 +29,7 @@ import (
|
||||
func newMirrorInitCmd(dir *string) *cobra.Command {
|
||||
var (
|
||||
controlPlanes []string
|
||||
peers []string
|
||||
base string
|
||||
noFirewall bool
|
||||
)
|
||||
@@ -61,7 +62,7 @@ Safe to run again: it changes only what differs.`,
|
||||
if _, err := parseAllow(controlPlanes); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := mirrorInit(cmd.Context(), cmd.OutOrStdout(), *dir, base, controlPlanes); err != nil {
|
||||
if err := mirrorInit(cmd.Context(), cmd.OutOrStdout(), *dir, base, controlPlanes, peers); err != nil {
|
||||
return err
|
||||
}
|
||||
if !noFirewall {
|
||||
@@ -80,6 +81,9 @@ Safe to run again: it changes only what differs.`,
|
||||
}
|
||||
cmd.Flags().StringArrayVar(&controlPlanes, "control-plane", nil,
|
||||
"address of the control plane that manages this mirror; repeatable")
|
||||
cmd.Flags().StringArrayVar(&peers, "peer", nil,
|
||||
"another mirror's management API to keep in step with, http://host:8082; repeatable. "+
|
||||
"Its machine must be among --control-plane, to read this one")
|
||||
cmd.Flags().BoolVar(&noFirewall, "no-firewall", false,
|
||||
"leave the firewall alone, instead of keeping the mirror's ports to the cluster's networks")
|
||||
cmd.Flags().StringVar(&base, "base", "",
|
||||
@@ -87,7 +91,7 @@ Safe to run again: it changes only what differs.`,
|
||||
return cmd
|
||||
}
|
||||
|
||||
func mirrorInit(ctx context.Context, out io.Writer, dir, base string, controlPlanes []string) error {
|
||||
func mirrorInit(ctx context.Context, out io.Writer, dir, base string, controlPlanes, peers []string) error {
|
||||
in := &serverInit{user: "llmgw", stateDir: "/var/lib/llmgw", out: out}
|
||||
if err := in.ensureUser(ctx); err != nil {
|
||||
return err
|
||||
@@ -105,7 +109,7 @@ func mirrorInit(ctx context.Context, out io.Writer, dir, base string, controlPla
|
||||
}
|
||||
base = "http://" + net.JoinHostPort(addr, "8081")
|
||||
}
|
||||
unit := mirrorUnitFor(dir, ":8082", base, controlPlanes,
|
||||
unit := mirrorUnitFor(dir, ":8082", base, controlPlanes, peers,
|
||||
"'llmgw mirror init'. Bytes on :8081 for every node; the\n# management API on :8082, for the control plane only.")
|
||||
changed, err := in.writeUnit(mirrorUnitPath, unit)
|
||||
if err != nil {
|
||||
|
||||
+111
-1
@@ -18,6 +18,7 @@ import (
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/sourcecheck"
|
||||
)
|
||||
|
||||
@@ -33,7 +34,7 @@ func newModelCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
Use: "model",
|
||||
Short: "Add or remove models while the cluster is running",
|
||||
}
|
||||
c.AddCommand(newModelAddCmd(client), newModelPresetsCmd(client, asJSON), newModelRemoveCmd(client), newModelPlanCmd(client),
|
||||
c.AddCommand(newModelContextCmd(client), newModelGPUsCmd(client), newModelAddCmd(client), newModelPresetsCmd(client, asJSON), newModelRemoveCmd(client), newModelPlanCmd(client),
|
||||
newResourcesCmd(client), newClassCmd(client), newOwnerCmd(client),
|
||||
newTargetCmd(client, asJSON), newAdaptersCmd(client),
|
||||
newModelServeCmd(client), newModelUpdatesCmd(client, asJSON), newModelUpdateCmd(client),
|
||||
@@ -386,3 +387,112 @@ func newAdaptersCmd(client func() *ctlClient) *cobra.Command {
|
||||
c.Flags().StringArrayVar(&remove, "remove", nil, "an adapter's name; repeatable")
|
||||
return c
|
||||
}
|
||||
|
||||
func newModelContextCmd(client func() *ctlClient) *cobra.Command {
|
||||
var slots int
|
||||
c := &cobra.Command{
|
||||
Use: "context <model> <tokens>",
|
||||
Short: "Change how much one request to a served model may use",
|
||||
Long: "The context lives in three places: the catalog (what a request may use, and\n" +
|
||||
"what it may answer), the engine's arguments (llama.cpp's --ctx-size and\n" +
|
||||
"--parallel, vLLM's --max-model-len), and the memory the KV cache takes. This\n" +
|
||||
"changes all three together. It refuses more than the model was trained for,\n" +
|
||||
"and a context its cards would not hold, saying what each would need. Its\n" +
|
||||
"replicas are then started again on it, each new one before the old stops.\n\n" +
|
||||
"--slots is how many requests a llama.cpp replica answers at once; its pool is\n" +
|
||||
"slots × tokens. Left out, it keeps what it has.",
|
||||
Example: " llmgw ctl model context qwen3-coder-30b 65536\n" +
|
||||
" llmgw ctl model context qwen2.5-7b 16384 --slots 4",
|
||||
Args: cobra.ExactArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
tokens, err := strconv.Atoi(args[1])
|
||||
if err != nil {
|
||||
return fmt.Errorf("the context is a whole number of tokens, not %q", args[1])
|
||||
}
|
||||
body, _ := json.Marshal(modelserve.ContextRequest{Tokens: tokens, Slots: slots})
|
||||
var v modelserve.ContextChange
|
||||
if err := client().patient(time.Minute).post(cmd.Context(), "/api/v1/models/"+args[0]+"/context", body, &v); err != nil {
|
||||
return err
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintf(out, "%s: %d → %d tokens a request, answers up to %d", v.Model, v.Before, v.After, v.MaxOutput)
|
||||
if v.Trained > 0 {
|
||||
fmt.Fprintf(out, " (trained for %d)", v.Trained)
|
||||
}
|
||||
fmt.Fprintln(out)
|
||||
for _, p := range v.Profiles {
|
||||
line := " " + p.Profile + " (" + p.Engine + "): " + strings.Join(p.Args, " ")
|
||||
switch {
|
||||
case p.Before != "" && p.After != "":
|
||||
line += "; " + p.Before + " → " + p.After
|
||||
case p.After != "":
|
||||
line += "; " + p.After
|
||||
}
|
||||
fmt.Fprintln(out, line)
|
||||
}
|
||||
if v.Note != "" {
|
||||
fmt.Fprintln(out, " note: "+v.Note)
|
||||
}
|
||||
if v.Restarted > 0 {
|
||||
fmt.Fprintf(out, "%d replica(s) started again on it, each new one before the old stops: llmgw ctl replicas --model %s\n", v.Restarted, v.Model)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
c.Flags().IntVar(&slots, "slots", 0, "llama.cpp: requests a replica answers at once (default: as it is)")
|
||||
return c
|
||||
}
|
||||
|
||||
func newModelGPUsCmd(client func() *ctlClient) *cobra.Command {
|
||||
var profile string
|
||||
var inPlace bool
|
||||
c := &cobra.Command{
|
||||
Use: "gpus <model> <cards>",
|
||||
Short: "Spread a served model over more cards, or fewer",
|
||||
Long: "How many cards of one node each replica uses. vLLM, SGLang and TensorRT-LLM\n" +
|
||||
"spread the model by tensor parallel, which needs the count to divide the\n" +
|
||||
"model's KV heads or be divided by them; llama.cpp splits its layers over the\n" +
|
||||
"cards. It refuses a count a node does not have, and one whose cards would not\n" +
|
||||
"hold the model and its context, saying what each would need. Its replicas are\n" +
|
||||
"then started again on it, each new one beside the old before it stops.\n\n" +
|
||||
"With no node free enough beside the running replicas — a node of two cards,\n" +
|
||||
"the model on one — --in-place stops each replica first: the model does not\n" +
|
||||
"answer while the new one loads.\n\n" +
|
||||
"--profile names the profile for a model with more than one on cards.",
|
||||
Example: " llmgw ctl model gpus qwen3-coder-30b 2\n" +
|
||||
" llmgw ctl model gpus llama-70b 4 --profile h100 --in-place",
|
||||
Args: cobra.ExactArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
n, err := strconv.Atoi(args[1])
|
||||
if err != nil {
|
||||
return fmt.Errorf("the cards are a whole number, not %q", args[1])
|
||||
}
|
||||
body, _ := json.Marshal(modelserve.GPUsRequest{GPUs: n, Profile: profile, InPlace: inPlace})
|
||||
var v modelserve.GPUsChange
|
||||
if err := client().patient(time.Minute).post(cmd.Context(), "/api/v1/models/"+args[0]+"/gpus", body, &v); err != nil {
|
||||
return err
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintf(out, "%s: %d → %d %s card(s) a replica (%s, %s)\n", v.Model, v.Before, v.After, v.Class, v.Profile, v.Engine)
|
||||
if v.PerCard != "" {
|
||||
fmt.Fprintln(out, " "+v.PerCard)
|
||||
}
|
||||
if len(v.Args) > 0 {
|
||||
fmt.Fprintln(out, " args: "+strings.Join(v.Args, " "))
|
||||
}
|
||||
if v.Note != "" {
|
||||
fmt.Fprintln(out, " note: "+v.Note)
|
||||
}
|
||||
switch {
|
||||
case v.Restarted > 0 && v.InPlace:
|
||||
fmt.Fprintf(out, "%d replica(s) stopped and started again on it, one by one: llmgw ctl replicas --model %s\n", v.Restarted, v.Model)
|
||||
case v.Restarted > 0:
|
||||
fmt.Fprintf(out, "%d replica(s) started again on it, each new one before the old stops: llmgw ctl replicas --model %s\n", v.Restarted, v.Model)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
c.Flags().StringVar(&profile, "profile", "", "the profile to change, for a model with several on cards")
|
||||
c.Flags().BoolVar(&inPlace, "in-place", false, "stop each replica before its replacement starts, when there is no room beside it")
|
||||
return c
|
||||
}
|
||||
@@ -186,7 +186,7 @@ func (na *nodeAdd) prepare() error {
|
||||
}
|
||||
// Where the mirror serves model files, for the node to be asked to reach.
|
||||
if cfg, err := loadEffectiveConfig(na.serverCfg); err == nil && cfg.Mirror.API != "" {
|
||||
if mc, err := mirror.NewClient(cfg.Mirror.API, cfg.Mirror.TLS, 5*time.Second); err == nil {
|
||||
if mc, err := mirror.NewFailoverClient(append([]string{cfg.Mirror.API}, cfg.Mirror.Peers...), cfg.Mirror.TLS, 5*time.Second); err == nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
if list, err := mc.List(ctx); err == nil {
|
||||
o.MirrorBase = list.Base
|
||||
|
||||
@@ -55,6 +55,7 @@ func Root() *cobra.Command {
|
||||
newBundleCmd(),
|
||||
newSimEngineCmd(),
|
||||
newComfyUIGateCmd(),
|
||||
newStrataServeCmd(),
|
||||
newDoctorCmd(),
|
||||
newNodeCmd(),
|
||||
newSelfUpgradeCmd(),
|
||||
|
||||
+14
-3
@@ -567,10 +567,21 @@ func newServerCmd() *cobra.Command {
|
||||
// rather than turning into a broken admin page hours later when
|
||||
// somebody finally opens it.
|
||||
if cfg.Mirror.API != "" {
|
||||
mc, err := mirror.NewClient(cfg.Mirror.API, cfg.Mirror.TLS, cfg.Mirror.Timeout)
|
||||
mirrors := append([]string{cfg.Mirror.API}, cfg.Mirror.Peers...)
|
||||
mc, err := mirror.NewFailoverClient(mirrors, cfg.Mirror.TLS, cfg.Mirror.Timeout)
|
||||
if err != nil {
|
||||
return fmt.Errorf("the model mirror: %w", err)
|
||||
}
|
||||
// Mirrors kept in step: a node is given a model's file on one
|
||||
// that answers while the one the catalog names is down.
|
||||
if len(cfg.Mirror.Peers) > 0 {
|
||||
src, err := mirror.NewSources(mirrors, cfg.Mirror.TLS, log)
|
||||
if err != nil {
|
||||
return fmt.Errorf("the model mirrors: %w", err)
|
||||
}
|
||||
go src.Run(ctx)
|
||||
agentsvc.SetSourceRewrite(src.Rewrite)
|
||||
}
|
||||
rt.SetMirror(mc)
|
||||
if fwPub != nil {
|
||||
fwPub.setMirror(mc)
|
||||
@@ -578,7 +589,7 @@ func newServerCmd() *cobra.Command {
|
||||
// Adapters trained on the nodes are kept there too; sending one
|
||||
// takes longer than the mirror's usual answers.
|
||||
if adminAPI != nil {
|
||||
if up, err := mirror.NewClient(cfg.Mirror.API, cfg.Mirror.TLS, 10*time.Minute); err == nil {
|
||||
if up, err := mirror.NewFailoverClient(mirrors, cfg.Mirror.TLS, 10*time.Minute); err == nil {
|
||||
adminAPI.SetTrainMirror(up)
|
||||
}
|
||||
}
|
||||
@@ -596,7 +607,7 @@ func newServerCmd() *cobra.Command {
|
||||
// Updates: the catalog in force, the service that
|
||||
// changes it, and what replaces replicas.
|
||||
Current: func() []catalog.Model { return admin.Catalog().Models() },
|
||||
Models: catalogSvc, Restart: admin.Reconciler().Restart,
|
||||
Models: catalogSvc, Restart: admin.Reconciler().Restart, RestartInPlace: admin.Reconciler().RestartInPlace,
|
||||
Pool: adminAPI.PoolNodes, Notify: notifier,
|
||||
// Canaries: how the router saw each side serve.
|
||||
CanaryStats: rt.CanaryStats, ResetCanary: rt.ResetCanary,
|
||||
|
||||
+28
-10
@@ -57,10 +57,12 @@ WantedBy=multi-user.target
|
||||
// it again finishes an install that stopped halfway and changes nothing on
|
||||
// one that did not.
|
||||
type serverInit struct {
|
||||
etcDir string // /etc/llmgw
|
||||
stateDir string // /var/lib/llmgw
|
||||
pkiDir string // where the authority lives
|
||||
unitPath string // "" leaves systemd alone
|
||||
// serviceActive says whether llmgw-server runs; nil asks systemctl.
|
||||
serviceActive func() bool
|
||||
etcDir string // /etc/llmgw
|
||||
stateDir string // /var/lib/llmgw
|
||||
pkiDir string // where the authority lives
|
||||
unitPath string // "" leaves systemd alone
|
||||
|
||||
// user owns the state and reads the key; "" means no user is created and
|
||||
// nothing is chowned, which is what a test wants.
|
||||
@@ -84,6 +86,9 @@ type serverInit struct {
|
||||
// nodes labelled for them; judgeCPUs their cores, once added.
|
||||
judgeMachines []judgeMachine
|
||||
judgeCPUs []int
|
||||
// secondMirror is which machine of the store, 1 or 2, keeps a second
|
||||
// mirror in step with this machine's; 0 is none.
|
||||
secondMirror int
|
||||
// judgeNote is what timing the judge found, said before the next
|
||||
// steps when it is not fine.
|
||||
judgeNote string
|
||||
@@ -241,6 +246,8 @@ func newServerInitCmd() *cobra.Command {
|
||||
f.StringVar(&judgeCats, "judge-categories", strings.Join(defaultJudgeCategories, ","),
|
||||
"what the judge refuses, comma-separated, of: "+strings.Join(judgeCategories, ", "))
|
||||
f.IntVar(&in.judge.Replicas, "judge-replicas", 2, "how many copies of the judge run, each on CPUs")
|
||||
f.IntVar(&in.secondMirror, "second-mirror", 0,
|
||||
"with --patroni, a second mirror on the store's machine 1 or 2, kept in step with this one's (0: none)")
|
||||
f.Var(judgeMachineFlag{&in.judgeMachines}, "judge-node",
|
||||
"user@host of a machine for the guardrails' judges alone: added as a node, and the judges run only there; repeatable")
|
||||
f.DurationVar(&in.judge.Budget, "judge-budget", 5*time.Second, "how long a request waits for the judge")
|
||||
@@ -355,6 +362,12 @@ func (in *serverInit) run(ctx context.Context) error {
|
||||
return fmt.Errorf("standby control planes: %w", err)
|
||||
}
|
||||
}
|
||||
if hasSystemd && in.secondMirror > 0 {
|
||||
if err := in.ensureSecondMirror(ctx); err != nil {
|
||||
// Not fatal: the first mirror serves; said, and how to try again.
|
||||
fmt.Fprintf(in.out, " !! second mirror: %v; later: sudo llmgw mirror add-peer <user>@<machine>\n", err)
|
||||
}
|
||||
}
|
||||
if hasSystemd && in.judgeWanted() && in.configJudgesWith(guardModelName) {
|
||||
in.checkJudge(ctx)
|
||||
}
|
||||
@@ -365,12 +378,7 @@ func (in *serverInit) run(ctx context.Context) error {
|
||||
// judgeWanted is whether this install fetches the judge model, and so
|
||||
// whether the baseline policy asks it.
|
||||
func (in *serverInit) judgeWanted() bool {
|
||||
if in.noGuardModel || (in.noMirror && in.mirrorURL == "") || in.unitPath == "" {
|
||||
return false
|
||||
}
|
||||
// Offline, only a judge the bundle carries: a policy asking a judge
|
||||
// that never comes refuses every request.
|
||||
if in.bundle != "" && !in.bundleCarries(guardModelName) {
|
||||
if !in.judgeComes() || in.unitPath == "" {
|
||||
return false
|
||||
}
|
||||
_, err := exec.LookPath("systemctl")
|
||||
@@ -1000,3 +1008,13 @@ func virtualInterface(name string) bool {
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// judgeComes is whether this setup has a judge to install: one not left
|
||||
// out, a mirror to fetch it from, and, offline, a bundle that carries it —
|
||||
// a policy asking a judge that never comes refuses every request.
|
||||
func (in *serverInit) judgeComes() bool {
|
||||
if in.noGuardModel || (in.noMirror && in.mirrorURL == "") {
|
||||
return false
|
||||
}
|
||||
return in.bundle == "" || in.bundleCarries(guardModelName)
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"math"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path"
|
||||
"slices"
|
||||
"strconv"
|
||||
@@ -176,7 +177,9 @@ func (in *serverInit) servingMemoryMB(judgesElsewhere bool) (mb int, what string
|
||||
gb += in.starterModelSpec().minMemoryGB()
|
||||
parts = append(parts, "the starter model")
|
||||
}
|
||||
if !in.noGuardModel && !judgesElsewhere {
|
||||
// Only judges that will be installed: none without a mirror to fetch
|
||||
// them from, or from a bundle that does not carry one.
|
||||
if !judgesElsewhere && in.judgeComes() {
|
||||
n := in.judgeFor().Replicas
|
||||
gb += n * judgeMemoryGB
|
||||
parts = append(parts, fmt.Sprintf("%d judges", n))
|
||||
@@ -193,6 +196,12 @@ func (in *serverInit) checkMemory() error {
|
||||
if in.unitPath == "" {
|
||||
return nil // a development run, served by nothing
|
||||
}
|
||||
// Run again on a machine already serving — to change --host, say, as
|
||||
// the docs advise — it installs no models it does not already serve. A
|
||||
// first run that stopped half way is checked again: it still would.
|
||||
if in.alreadyServing() {
|
||||
return nil
|
||||
}
|
||||
meminfo, err := os.ReadFile("/proc/meminfo")
|
||||
if err != nil {
|
||||
return nil
|
||||
@@ -387,6 +396,9 @@ func (in *serverInit) initFlags() []string {
|
||||
for _, m := range in.judgeMachines {
|
||||
add("--judge-node", m.target)
|
||||
}
|
||||
if in.secondMirror > 0 {
|
||||
add("--second-mirror", fmt.Sprint(in.secondMirror))
|
||||
}
|
||||
switch {
|
||||
case in.mirrorURL != "":
|
||||
add("--mirror", in.mirrorURL)
|
||||
@@ -471,3 +483,15 @@ func shellWord(s string) string {
|
||||
}
|
||||
return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
|
||||
}
|
||||
|
||||
// alreadyServing reports whether this machine's gateway is set up and
|
||||
// running: a server.yaml, and its unit active.
|
||||
func (in *serverInit) alreadyServing() bool {
|
||||
if _, err := os.Stat(in.configPath()); err != nil {
|
||||
return false
|
||||
}
|
||||
if in.serviceActive != nil {
|
||||
return in.serviceActive()
|
||||
}
|
||||
return exec.Command("systemctl", "is-active", "--quiet", "llmgw-server").Run() == nil
|
||||
}
|
||||
@@ -123,10 +123,10 @@ func (in *serverInit) writeUnit(path, body string) (bool, error) {
|
||||
func (in *serverInit) mirrorUnit() string {
|
||||
base := "http://" + net.JoinHostPort(in.reachableAddr(), "8081")
|
||||
if in.standbys {
|
||||
return mirrorUnitFor(in.mirrorRoot(), ":8082", base, in.mirrorAllow(),
|
||||
return mirrorUnitFor(in.mirrorRoot(), ":8082", base, in.mirrorAllow(), nil,
|
||||
"'llmgw server init'. Bytes on :8081 for every node; the\n# management API for the control planes only, whichever leads.")
|
||||
}
|
||||
return mirrorUnitFor(in.mirrorRoot(), mirrorAPI, base, nil,
|
||||
return mirrorUnitFor(in.mirrorRoot(), mirrorAPI, base, nil, nil,
|
||||
"'llmgw server init'. Bytes on :8081 for every node; the\n# management API on loopback only, for the control plane beside it.")
|
||||
}
|
||||
|
||||
@@ -154,11 +154,14 @@ func (in *serverInit) mirrorAllow() []string {
|
||||
|
||||
// mirrorUnitFor is a mirror's systemd unit. allow is who, besides the
|
||||
// machine itself, may use the management API.
|
||||
func mirrorUnitFor(dir, apiListen, base string, allow []string, about string) string {
|
||||
func mirrorUnitFor(dir, apiListen, base string, allow, peers []string, about string) string {
|
||||
exec := "/usr/local/bin/llmgw mirror serve --dir " + dir + " --listen :8081 --api-listen " + apiListen + " --base " + base
|
||||
for _, a := range allow {
|
||||
exec += " --api-allow " + a
|
||||
}
|
||||
for _, p := range peers {
|
||||
exec += " --peer " + p
|
||||
}
|
||||
return `[Unit]
|
||||
Description=llmgw model mirror
|
||||
After=network-online.target
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/preflight"
|
||||
@@ -24,4 +25,28 @@ func TestTheMemoryAMachineNeedsCountsWhatItServes(t *testing.T) {
|
||||
if mb, _ := (&serverInit{noModel: true, noGuardModel: true}).servingMemoryMB(false); mb != 0 {
|
||||
t.Errorf("nothing served: %d", mb)
|
||||
}
|
||||
// Without a mirror no judge is fetched, so none is counted.
|
||||
if mb, what := (&serverInit{noMirror: true}).servingMemoryMB(false); mb != 0 {
|
||||
t.Errorf("no mirror, nothing to fetch: %d %q", mb, what)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetupRunAgainOnAMachineAlreadySetUpIsNotRefusedItsMemory(t *testing.T) {
|
||||
// The docs fix a certificate with server init --host on a machine set up
|
||||
// before the check: it installs nothing it does not already serve.
|
||||
dir := t.TempDir()
|
||||
cfg := dir + "/server.yaml"
|
||||
if err := os.WriteFile(cfg, []byte("listen: :8080\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
in := &serverInit{unitPath: dir + "/llmgw-server.service", etcDir: dir, serviceActive: func() bool { return true }}
|
||||
if err := in.checkMemory(); err != nil {
|
||||
t.Errorf("refused: %v", err)
|
||||
}
|
||||
// A first run that stopped half way left server.yaml and serves nothing:
|
||||
// it still would install the models, and is checked.
|
||||
in.serviceActive = func() bool { return false }
|
||||
if in.alreadyServing() {
|
||||
t.Error("a machine whose gateway is not running was taken for one serving")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"os/signal"
|
||||
"syscall"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/strata"
|
||||
)
|
||||
|
||||
func newStrataServeCmd() *cobra.Command {
|
||||
var cfg strata.Config
|
||||
cmd := &cobra.Command{
|
||||
Use: "strata-serve [flags] [-- setup.py flags]",
|
||||
Short: "Run Strata as an engine: prepare the model, then serve it (started by the agent)",
|
||||
Hidden: true,
|
||||
Long: "Started by the node agent for a replica whose engine is strata. Prepares the\n" +
|
||||
"model from the fetched GGUF files with Strata's setup, which skips what it has\n" +
|
||||
"done before, then becomes Strata's server on the replica's address. The key is\n" +
|
||||
"read from " + strata.KeyEnv + ".",
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
cfg.Args = args
|
||||
cfg.Log = logger()
|
||||
return strata.Run(ctx, cfg)
|
||||
},
|
||||
}
|
||||
f := cmd.Flags()
|
||||
f.StringVar(&cfg.Listen, "listen", "", "the replica's address")
|
||||
f.StringVar(&cfg.Model, "model", "", "the model's name, for /v1/models")
|
||||
f.StringVar(&cfg.Weights, "weights", "", "the fetched GGUF files' directory")
|
||||
f.StringVar(&cfg.State, "state", "", "the replica's own directory")
|
||||
f.StringVar(&cfg.Data, "data", "", "the node's directory for prepared models, shared by its replicas")
|
||||
f.StringVar(&cfg.Python, "python", "python3", "the interpreter of Strata's environment")
|
||||
f.StringVar(&cfg.Dir, "strata", "", "Strata's directory, where setup.py is")
|
||||
f.StringVar(&cfg.Engine, "engine", "", "Strata's engine directory, where the strata binary is")
|
||||
for _, n := range []string{"listen", "weights", "state", "data", "strata", "engine"} {
|
||||
_ = cmd.MarkFlagRequired(n)
|
||||
}
|
||||
return cmd
|
||||
}
|
||||
@@ -332,6 +332,55 @@ func (m consoleModel) modelMenu(md adminapi.ModelView) *menuDialog {
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
{"Context length…", func(m *consoleModel) tea.Cmd {
|
||||
m.openForm(&formDialog{title: "Context of " + md.Name,
|
||||
note: "What one request may use. The engine's arguments and the memory change with it, checked against " +
|
||||
"what the cards hold; its replicas are then started again on it, one by one.",
|
||||
fields: []formField{
|
||||
textF("tokens", strconv.Itoa(md.ContextLength), "a request's prompt and answer together"),
|
||||
textF("slots", "", "llama.cpp: requests at once; empty keeps it"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
tokens, err := strconv.Atoi(v[0])
|
||||
if err != nil || tokens < 256 {
|
||||
return nil, fmt.Errorf("the context is a whole number of tokens, at least 256")
|
||||
}
|
||||
slots := 0
|
||||
if v[1] != "" {
|
||||
if slots, err = strconv.Atoi(v[1]); err != nil || slots < 1 {
|
||||
return nil, fmt.Errorf("slots is a whole number from 1")
|
||||
}
|
||||
}
|
||||
return m.call("Context of "+md.Name, http.MethodPost, path+"/context",
|
||||
modelserve.ContextRequest{Tokens: tokens, Slots: slots}, contextSays), nil
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
{"GPUs…", func(m *consoleModel) tea.Cmd {
|
||||
profiles := append([]string{"(its only one on cards)"}, md.Profiles...)
|
||||
m.openForm(&formDialog{title: "Cards of " + md.Name,
|
||||
note: "How many cards of one node each replica uses: tensor parallel for vLLM, SGLang and TensorRT-LLM, " +
|
||||
"the layers split for llama.cpp. Checked against the model's heads, its cards' memory and the nodes; " +
|
||||
"its replicas are then started again on it, one by one.",
|
||||
fields: []formField{
|
||||
choiceF("profile", profiles, profiles[0]),
|
||||
textF("cards", "", "a replica's, on one node"),
|
||||
choiceF("in place", []string{"no", "yes"}, "no"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
n, err := strconv.Atoi(v[1])
|
||||
if err != nil || n < 1 {
|
||||
return nil, fmt.Errorf("the cards are a whole number from 1")
|
||||
}
|
||||
profile := v[0]
|
||||
if profile == profiles[0] {
|
||||
profile = ""
|
||||
}
|
||||
return m.call("Cards of "+md.Name, http.MethodPost, path+"/gpus",
|
||||
modelserve.GPUsRequest{GPUs: n, Profile: profile, InPlace: v[2] == "yes"}, gpusSays), nil
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
{"Resources…", func(m *consoleModel) tea.Cmd {
|
||||
profiles := md.Profiles
|
||||
if len(profiles) == 0 {
|
||||
@@ -898,3 +947,46 @@ func locationRequest(n adminapi.NodeView, vals []string) adminapi.NodeLabelsRequ
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
// contextSays is a context change in a sentence or two.
|
||||
func contextSays(raw json.RawMessage) string {
|
||||
var v modelserve.ContextChange
|
||||
if json.Unmarshal(raw, &v) != nil {
|
||||
return "Done."
|
||||
}
|
||||
s := fmt.Sprintf("%s: %d → %d tokens a request, answers up to %d.", v.Model, v.Before, v.After, v.MaxOutput)
|
||||
for _, p := range v.Profiles {
|
||||
if p.After != "" {
|
||||
s += "\n" + p.Profile + ": " + p.After
|
||||
}
|
||||
}
|
||||
if v.Note != "" {
|
||||
s += "\n" + v.Note
|
||||
}
|
||||
if v.Restarted > 0 {
|
||||
s += fmt.Sprintf("\n%d replica(s) started again on it, one by one.", v.Restarted)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// gpusSays is a change of cards in a sentence or two.
|
||||
func gpusSays(raw json.RawMessage) string {
|
||||
var v modelserve.GPUsChange
|
||||
if json.Unmarshal(raw, &v) != nil {
|
||||
return "Done."
|
||||
}
|
||||
s := fmt.Sprintf("%s: %d → %d %s card(s) a replica.", v.Model, v.Before, v.After, v.Class)
|
||||
if v.PerCard != "" {
|
||||
s += "\n" + v.PerCard
|
||||
}
|
||||
if v.Note != "" {
|
||||
s += "\n" + v.Note
|
||||
}
|
||||
switch {
|
||||
case v.Restarted > 0 && v.InPlace:
|
||||
s += fmt.Sprintf("\n%d replica(s) stopped and started again on it, one by one.", v.Restarted)
|
||||
case v.Restarted > 0:
|
||||
s += fmt.Sprintf("\n%d replica(s) started again on it, one by one.", v.Restarted)
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -62,6 +62,7 @@ type installForm struct {
|
||||
admin, hosts, language, signup *setupField
|
||||
storeKind *setupField
|
||||
mirrorWhere, mirrorDirF, mirrorURL *setupField
|
||||
secondMirror *setupField
|
||||
starterPick, starterSource, starterCtx *setupField
|
||||
judgeFail, judgeCopies, judgeBudget *setupField
|
||||
judgeCats []*setupField
|
||||
@@ -181,18 +182,26 @@ func newInstallForm(in *serverInit, host string, found []string, ha bool) *insta
|
||||
f.mirrorDirF.Show = func() bool { return f.mirrorWhere.Choice == 0 }
|
||||
f.mirrorURL = textSetupField("Its management API", in.mirrorURL, "http://10.0.0.20:8082", "")
|
||||
f.mirrorURL.Show = func() bool { return f.mirrorWhere.Choice == 1 }
|
||||
f.secondMirror = choiceField("A second mirror", []string{"none", "on database machine 1", "on database machine 2"},
|
||||
in.secondMirror, "Kept in step with this one, each copying what the other has: while either machine is down, "+
|
||||
"models are still added and fetched. It needs room for every model, as this one does.")
|
||||
f.secondMirror.Show = func() bool { return ha && f.mirrorWhere.Choice == 0 }
|
||||
f.mirror = &installPart{
|
||||
title: "Model mirror", on: !in.noMirror || in.mirrorURL != "",
|
||||
about: "Downloads each model once, from Hugging Face or wherever it comes from, and serves the files to " +
|
||||
"every node on the network, so ten nodes do not download the same 40 GB ten times. " +
|
||||
"The starter model, the judge and an offline bundle go into it.",
|
||||
fields: []*setupField{f.mirrorWhere, f.mirrorDirF, f.mirrorURL},
|
||||
fields: []*setupField{f.mirrorWhere, f.mirrorDirF, f.mirrorURL, f.secondMirror},
|
||||
off: "each model's files must already be on a mirror somewhere; no starter model, no judge",
|
||||
summary: func() string {
|
||||
if f.mirrorWhere.Choice == 1 {
|
||||
return "at " + orText(f.mirrorURL.Value(), "(an address)")
|
||||
}
|
||||
return "on this machine, in " + orText(f.mirrorDirF.Value(), mirrorDir)
|
||||
s := "on this machine, in " + orText(f.mirrorDirF.Value(), mirrorDir)
|
||||
if ha && f.secondMirror.Choice > 0 {
|
||||
s += "; a second " + f.secondMirror.Options[f.secondMirror.Choice]
|
||||
}
|
||||
return s
|
||||
},
|
||||
}
|
||||
|
||||
@@ -636,6 +645,10 @@ func (f *installForm) apply(in *serverInit) error {
|
||||
in.postgres = !f.ha && f.storeKind.Choice == 1
|
||||
in.noMirror = !f.mirror.on || mirrorURL != ""
|
||||
in.mirrorURL, in.mirrorPath = mirrorURL, mirrorPath
|
||||
in.secondMirror = 0
|
||||
if f.ha && f.mirror.on && f.mirrorWhere.Choice == 0 {
|
||||
in.secondMirror = f.secondMirror.Choice
|
||||
}
|
||||
in.noNode = !f.node.on
|
||||
in.noModel = !f.starter.on
|
||||
in.starter, in.starterCtx = starter, ctx
|
||||
|
||||
@@ -46,6 +46,7 @@ const (
|
||||
cQuotas
|
||||
cGuard
|
||||
cKB
|
||||
cMCP
|
||||
cFederation
|
||||
cMirror
|
||||
cJobs
|
||||
@@ -70,7 +71,7 @@ const (
|
||||
cCount
|
||||
)
|
||||
|
||||
var consoleLabels = [cCount]string{"Overview", "Monitor", "Nodes", "GPUs", "Models", "Replicas", "Usage", "Quotas", "Guardrails", "Knowledge", "Federation", "Mirror", "Jobs", "Logs", "Accounts", "API keys", "Organizations", "Sign-in", "Audit", "Manual",
|
||||
var consoleLabels = [cCount]string{"Overview", "Monitor", "Nodes", "GPUs", "Models", "Replicas", "Usage", "Quotas", "Guardrails", "Knowledge", "MCP servers", "Federation", "Mirror", "Jobs", "Logs", "Accounts", "API keys", "Organizations", "Sign-in", "Audit", "Manual",
|
||||
"server.yaml", "Add a node", "Add a control plane", "Replicate the store", "Serve a model", "Serve the guardrails' judge", "Upgrade the cluster",
|
||||
"Restart all replicas", "Back up now", "Wipe this machine", "Leave"}
|
||||
|
||||
@@ -184,6 +185,7 @@ type consoleModel struct {
|
||||
guard guardState
|
||||
signIn signInState
|
||||
kb kbState
|
||||
mcp mcpState
|
||||
|
||||
st *storeHolder
|
||||
people peopleData
|
||||
@@ -342,6 +344,9 @@ func (m consoleModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
if m.menu == cKB {
|
||||
cmds = append(cmds, m.fetchKB())
|
||||
}
|
||||
if m.menu == cMCP {
|
||||
cmds = append(cmds, m.fetchMCP())
|
||||
}
|
||||
if m.menu == cMirror {
|
||||
cmds = append(cmds, m.fetchMirror())
|
||||
}
|
||||
@@ -404,6 +409,8 @@ func (m consoleModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
m.quotas, m.quotasErr = msg.rows, msg.err
|
||||
case consoleKBMsg:
|
||||
m.kb = kbState{view: msg.view, err: msg.err, loaded: true}
|
||||
case consoleMCPMsg:
|
||||
m.mcp = mcpState(msg)
|
||||
case consoleSignInMsg:
|
||||
m.signIn = signInState{view: msg.view, err: msg.err, loaded: true, ldap: msg.ldap, oidc: msg.oidc}
|
||||
case consoleGuardMsg:
|
||||
@@ -568,6 +575,8 @@ func (m consoleModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.menuD = m.signInMenu(m.signInRowKey())
|
||||
case cKB:
|
||||
m.menuD = m.kbMenu(m.kb.view.Collections[m.row])
|
||||
case cMCP:
|
||||
m.menuD = m.mcpMenu(m.mcp.servers[m.row])
|
||||
case cMirror:
|
||||
m.menuD = m.mirrorMenu(m.mirrorGroups()[m.row])
|
||||
case cJobs:
|
||||
@@ -633,6 +642,10 @@ func (m consoleModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
if m.kb.view.Available {
|
||||
m.kbCreateForm()
|
||||
}
|
||||
case cMCP:
|
||||
if m.mcp.loaded && m.mcp.err == nil {
|
||||
m.mcpAddForm()
|
||||
}
|
||||
case cModels:
|
||||
m.addModelForm()
|
||||
case cMirror:
|
||||
@@ -705,6 +718,8 @@ func (m consoleModel) onEnter() tea.Cmd {
|
||||
return m.fetchSignIn()
|
||||
case cKB:
|
||||
return m.fetchKB()
|
||||
case cMCP:
|
||||
return m.fetchMCP()
|
||||
case cMirror:
|
||||
return m.fetchMirror()
|
||||
case cMonitor:
|
||||
@@ -1085,6 +1100,8 @@ func (m consoleModel) rowCount() int {
|
||||
return len(signInRows)
|
||||
case cKB:
|
||||
return len(m.kb.view.Collections)
|
||||
case cMCP:
|
||||
return len(m.mcp.servers)
|
||||
case cMirror:
|
||||
return len(m.mirrorGroups())
|
||||
case cJobs:
|
||||
@@ -1326,6 +1343,8 @@ func (m consoleModel) panel() (title, body, keys string) {
|
||||
return m.signInPanel(pw)
|
||||
case cKB:
|
||||
return m.kbPanel(pw)
|
||||
case cMCP:
|
||||
return m.mcpPanel(pw)
|
||||
case cAccounts, cKeys, cOrgs:
|
||||
return m.peoplePanel(pw)
|
||||
case cManual:
|
||||
|
||||
@@ -0,0 +1,534 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// The company's own MCP servers, as the console shows them: each with its
|
||||
// address, the names of the credentials it is sent (never their values),
|
||||
// its rate, and who has been granted its tools; Enter for what can be done
|
||||
// to one, A to add one. Clients reach them all through the gateway's /mcp
|
||||
// with an llmgw key, as server.tool.
|
||||
|
||||
type mcpState struct {
|
||||
servers []adminapi.MCPServerView
|
||||
err error
|
||||
loaded bool
|
||||
// grants are the tool grants in the tenancy tree, nil when the store
|
||||
// could not be read from here; where says what each scope's ID is
|
||||
// called, and ids the other way, for the grant form.
|
||||
grants []mcpGrant
|
||||
where map[string]string
|
||||
ids map[string]store.LimitOwner
|
||||
// keys are the IDs of the keys there are, for a grant to one.
|
||||
keys map[string]bool
|
||||
}
|
||||
|
||||
// mcpGrant is one level of the tree that grants tools, and which.
|
||||
type mcpGrant struct {
|
||||
owner store.LimitOwner
|
||||
label string
|
||||
tools []string
|
||||
}
|
||||
|
||||
type consoleMCPMsg mcpState
|
||||
|
||||
func (m consoleModel) fetchMCP() tea.Cmd {
|
||||
c, h := m.client, m.st
|
||||
return func() tea.Msg {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var msg consoleMCPMsg
|
||||
msg.err = c.get(ctx, "/api/v1/mcp/servers", &msg.servers)
|
||||
msg.loaded = true
|
||||
if h == nil {
|
||||
return msg
|
||||
}
|
||||
s, err := h.get(ctx)
|
||||
if err != nil {
|
||||
return msg
|
||||
}
|
||||
limits, err := s.AllLimits(ctx)
|
||||
if err != nil {
|
||||
return msg
|
||||
}
|
||||
msg.where, msg.ids = map[string]string{}, map[string]store.LimitOwner{}
|
||||
if tree, err := s.Tree(ctx); err == nil {
|
||||
for _, o := range tree {
|
||||
msg.name(store.ScopeOrg, o.ID, o.Name)
|
||||
for _, t := range o.Teams {
|
||||
msg.name(store.ScopeTeam, t.ID, o.Name+"/"+t.Name)
|
||||
for _, p := range t.Projects {
|
||||
msg.name(store.ScopeProject, p.ID, o.Name+"/"+t.Name+"/"+p.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if users, err := s.Users(ctx, ""); err == nil {
|
||||
for _, u := range users {
|
||||
msg.name(store.ScopeUser, u.ID, u.Email)
|
||||
}
|
||||
}
|
||||
msg.keys = map[string]bool{}
|
||||
if keys, err := s.AllKeys(ctx); err == nil {
|
||||
for _, k := range keys {
|
||||
msg.keys[k.ID] = true
|
||||
msg.where[store.LimitOwner{Scope: store.ScopeKey, ID: k.ID}.CounterKey()] = k.Name
|
||||
}
|
||||
}
|
||||
msg.grants = []mcpGrant{}
|
||||
for _, l := range limits {
|
||||
if l.AllowedTools == nil {
|
||||
continue
|
||||
}
|
||||
o := store.LimitOwner{Scope: l.Scope, ID: l.ScopeID}
|
||||
label := msg.where[o.CounterKey()]
|
||||
if label == "" {
|
||||
label = l.ScopeID
|
||||
}
|
||||
msg.grants = append(msg.grants, mcpGrant{owner: o, label: string(l.Scope) + " " + label, tools: l.AllowedTools})
|
||||
}
|
||||
slices.SortFunc(msg.grants, func(a, b mcpGrant) int { return strings.Compare(a.label, b.label) })
|
||||
return msg
|
||||
}
|
||||
}
|
||||
|
||||
func (s *consoleMCPMsg) name(scope store.Scope, id, label string) {
|
||||
o := store.LimitOwner{Scope: scope, ID: id}
|
||||
s.where[o.CounterKey()] = label
|
||||
s.ids[string(scope)+":"+strings.ToLower(label)] = o
|
||||
if m := store.MemberScope(scope); m != "" {
|
||||
s.where[store.LimitOwner{Scope: m, ID: id}.CounterKey()] = label
|
||||
}
|
||||
}
|
||||
|
||||
// closed are the levels whose list names no tool at all.
|
||||
func (k mcpState) closed() []mcpGrant {
|
||||
var out []mcpGrant
|
||||
for _, g := range k.grants {
|
||||
if len(g.tools) == 0 {
|
||||
out = append(out, g)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// mcpGrantsOf is who may use a server's tools, and which of them.
|
||||
func (k mcpState) grantsOf(server string) []mcpGrant {
|
||||
var out []mcpGrant
|
||||
for _, g := range k.grants {
|
||||
var mine []string
|
||||
for _, t := range g.tools {
|
||||
if strings.HasPrefix(t, server+".") || t == "*" {
|
||||
mine = append(mine, t)
|
||||
}
|
||||
}
|
||||
if len(mine) > 0 {
|
||||
out = append(out, mcpGrant{owner: g.owner, label: g.label, tools: mine})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (m consoleModel) mcpPanel(pw int) (string, string, string) {
|
||||
keys := "→ or Enter open · ↑↓ choose · Q leave"
|
||||
if m.inPanel {
|
||||
keys = "↑↓ server · Enter actions · A add a server · ← back · Q leave"
|
||||
}
|
||||
k := m.mcp
|
||||
var b strings.Builder
|
||||
b.WriteString("\n " + sDim.Render("The company's MCP servers, offered to agents through the gateway's /mcp as server.tool.") + "\n\n")
|
||||
switch {
|
||||
case k.err != nil:
|
||||
b.WriteString(" " + sBad.Render(k.err.Error()))
|
||||
return "MCP servers", b.String(), keys
|
||||
case !k.loaded:
|
||||
b.WriteString(" " + sDim.Render("asking the control plane…"))
|
||||
return "MCP servers", b.String(), keys
|
||||
case len(k.servers) == 0:
|
||||
b.WriteString(" " + sDim.Render("None yet. A to add one: its address, and the credential it wants."))
|
||||
b.WriteString("\n\n" + mcpClientHint())
|
||||
return "MCP servers", b.String(), keys
|
||||
}
|
||||
var rows [][]string
|
||||
for _, s := range k.servers {
|
||||
state := sOK.Render("on")
|
||||
if !s.Enabled {
|
||||
state = sDim.Render("off")
|
||||
}
|
||||
granted := "—"
|
||||
if k.grants != nil {
|
||||
var who []string
|
||||
for _, g := range k.grantsOf(s.Name) {
|
||||
who = append(who, g.label)
|
||||
}
|
||||
granted = sBad.Render("nobody yet")
|
||||
if len(who) > 0 {
|
||||
granted = strings.Join(who, ", ")
|
||||
}
|
||||
}
|
||||
rate := "—"
|
||||
if s.RatePerMinute > 0 {
|
||||
rate = fmt.Sprintf("%d/min", s.RatePerMinute)
|
||||
}
|
||||
rows = append(rows, []string{s.Name, state, granted, rate, dash(strings.Join(mcpCredentials(s), ", ")), s.URL})
|
||||
}
|
||||
b.WriteString(table([]string{"SERVER", "STATE", "GRANTED TO", "RATE", "SENT", "ADDRESS"}, rows, m.row, pw, m.inPanel))
|
||||
b.WriteString("\n\n " + sDim.Render("A server's tools are seen only where they are granted; nothing is granted by default.") + "\n")
|
||||
if closed := k.closed(); len(closed) > 0 {
|
||||
var names []string
|
||||
for _, g := range closed {
|
||||
names = append(names, g.label)
|
||||
}
|
||||
b.WriteString(" " + sBad.Render("Closed to every tool: ") + strings.Join(names, ", ") +
|
||||
sDim.Render(" (a server's menu opens them)") + "\n")
|
||||
}
|
||||
b.WriteString("\n")
|
||||
b.WriteString(mcpClientHint())
|
||||
return "MCP servers", b.String(), keys
|
||||
}
|
||||
|
||||
// mcpClientHint is how a person's agent is pointed at the gateway.
|
||||
func mcpClientHint() string {
|
||||
return " " + sDim.Render("In a person's agent: the gateway's https://…/mcp, with their key from /profile/keys") + "\n" +
|
||||
" " + sDim.Render(" claude mcp add --transport http company https://HOST/mcp --header \"Authorization: Bearer sk-llmgw-…\"")
|
||||
}
|
||||
|
||||
// mcpCredentials are the names of the headers a server is sent: which are
|
||||
// kept as secrets and which as written. Never their values.
|
||||
func mcpCredentials(s adminapi.MCPServerView) []string {
|
||||
var out []string
|
||||
for h := range s.HeaderRefs {
|
||||
out = append(out, h+" (secret)")
|
||||
}
|
||||
for h := range s.Headers {
|
||||
out = append(out, h)
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// mcpPath is a server's admin API address.
|
||||
func mcpPath(name string) string { return "/api/v1/mcp/servers/" + url.PathEscape(name) }
|
||||
|
||||
// mcpAuth turns the form's credential fields into what the server is sent:
|
||||
// a bearer token, another header, or nothing.
|
||||
func mcpAuth(kind, header, secret string) (map[string]string, error) {
|
||||
switch kind {
|
||||
case "none":
|
||||
return nil, nil
|
||||
case "bearer":
|
||||
if secret == "" {
|
||||
return nil, nil
|
||||
}
|
||||
return map[string]string{"Authorization": "Bearer " + strings.TrimPrefix(secret, "Bearer ")}, nil
|
||||
default:
|
||||
if header == "" {
|
||||
return nil, fmt.Errorf("which header carries the credential?")
|
||||
}
|
||||
if secret == "" {
|
||||
return nil, nil
|
||||
}
|
||||
return map[string]string{header: secret}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func atoiOr(s string, def int) (int, error) {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return def, nil
|
||||
}
|
||||
return strconv.Atoi(strings.TrimSpace(s))
|
||||
}
|
||||
|
||||
// mcpAddForm registers a server.
|
||||
func (m *consoleModel) mcpAddForm() {
|
||||
m.openForm(&formDialog{title: "Add an MCP server",
|
||||
note: "Reached by the control plane over HTTP (streamable HTTP or SSE). Its credential is kept in\n" +
|
||||
"the store, sealed, and never reaches a client. Its tools are seen nowhere until granted.",
|
||||
fields: []formField{
|
||||
textF("name", "", "lowercase, the tools' prefix: jira → jira.search"),
|
||||
textF("address", "", "https://jira-mcp.internal/mcp"),
|
||||
textF("about", "", "what it is for, in a line"),
|
||||
choiceF("credential", []string{"bearer", "header", "none"}, "bearer"),
|
||||
textF("header", "", "for credential: header, its name (X-Api-Key)"),
|
||||
pwF("secret", "the token or key; kept sealed in the store"),
|
||||
textF("timeout", "", "seconds a call may take; empty is the default"),
|
||||
textF("rate", "", "calls a minute per caller; empty is no bound"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
name, addr := strings.TrimSpace(v[0]), strings.TrimSpace(v[1])
|
||||
if name == "" || addr == "" {
|
||||
return nil, fmt.Errorf("a name and an address")
|
||||
}
|
||||
if err := adminapi.CheckMCPAddress(addr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
secrets, err := mcpAuth(v[3], strings.TrimSpace(v[4]), v[5])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
timeout, err := atoiOr(v[6], 0)
|
||||
if err != nil || timeout < 0 {
|
||||
return nil, fmt.Errorf("timeout is whole seconds")
|
||||
}
|
||||
rate, err := atoiOr(v[7], 0)
|
||||
if err != nil || rate < 0 {
|
||||
return nil, fmt.Errorf("rate is calls a minute")
|
||||
}
|
||||
on := true
|
||||
req := adminapi.MCPServerRequest{URL: addr, Description: v[2], Secrets: secrets, Enabled: &on,
|
||||
TimeoutSeconds: timeout, RatePerMinute: rate}
|
||||
return m.call("Add "+name, http.MethodPut, mcpPath(name), req, func(json.RawMessage) string {
|
||||
return "Added " + name + ". Test it from its menu, then grant its tools: nobody sees them yet."
|
||||
}), nil
|
||||
}})
|
||||
}
|
||||
|
||||
// mcpMenu is what can be done to one server.
|
||||
func (m consoleModel) mcpMenu(s adminapi.MCPServerView) *menuDialog {
|
||||
items := []menuItem{
|
||||
{"Test, and list its tools", func(m *consoleModel) tea.Cmd { return m.mcpTest(s) }},
|
||||
{"Grant its tools to…", func(m *consoleModel) tea.Cmd { m.mcpGrantForm(s); return nil }},
|
||||
}
|
||||
for _, g := range m.mcp.grantsOf(s.Name) {
|
||||
if !slices.ContainsFunc(g.tools, func(t string) bool { return t != "*" }) {
|
||||
continue // every tool, granted with the Limits form; revoked there too
|
||||
}
|
||||
items = append(items, menuItem{"Revoke from " + g.label + " (" + strings.Join(g.tools, ", ") + ")…",
|
||||
func(m *consoleModel) tea.Cmd {
|
||||
m.confirmD = &confirmDialog{title: "Revoke " + s.Name + " from " + g.label,
|
||||
text: "Taken from it and from everything beneath it. Its other grants stay; left naming\n" +
|
||||
"no tool, it is closed rather than opened to what the levels above allow.",
|
||||
yes: func(m *consoleModel) tea.Cmd { return m.mcpSetTools(s.Name, g.owner, nil) }}
|
||||
return nil
|
||||
}})
|
||||
}
|
||||
for _, g := range m.mcp.closed() {
|
||||
items = append(items, menuItem{"Open " + g.label + ", closed to every tool…", func(m *consoleModel) tea.Cmd {
|
||||
m.confirmD = &confirmDialog{title: "Open " + g.label,
|
||||
text: "It then names no tools of its own and sees what the levels above it allow.",
|
||||
yes: func(m *consoleModel) tea.Cmd { return m.mcpOpen(g.owner) }}
|
||||
return nil
|
||||
}})
|
||||
}
|
||||
items = append(items,
|
||||
menuItem{"Change…", func(m *consoleModel) tea.Cmd { m.mcpEditForm(s); return nil }},
|
||||
menuItem{"Replace its credential…", func(m *consoleModel) tea.Cmd { m.mcpCredentialForm(s); return nil }},
|
||||
)
|
||||
if s.Enabled {
|
||||
items = append(items, menuItem{"Turn off", func(m *consoleModel) tea.Cmd {
|
||||
off := false
|
||||
return m.call("Turn off "+s.Name, http.MethodPut, mcpPath(s.Name), adminapi.MCPServerRequest{URL: s.URL, Enabled: &off},
|
||||
func(json.RawMessage) string {
|
||||
return s.Name + " is off: its tools are listed nowhere and every call is refused."
|
||||
})
|
||||
}})
|
||||
} else {
|
||||
items = append(items, menuItem{"Turn on", func(m *consoleModel) tea.Cmd {
|
||||
on := true
|
||||
return m.call("Turn on "+s.Name, http.MethodPut, mcpPath(s.Name), adminapi.MCPServerRequest{URL: s.URL, Enabled: &on}, nil)
|
||||
}})
|
||||
}
|
||||
items = append(items, menuItem{"Remove…", func(m *consoleModel) tea.Cmd {
|
||||
m.confirmD = &confirmDialog{title: "Remove " + s.Name,
|
||||
text: "Its tools go from every agent at once, and its sealed credential with it.\nThe grants stay in the tree until revoked.",
|
||||
yes: func(m *consoleModel) tea.Cmd {
|
||||
return m.call("Remove "+s.Name, http.MethodDelete, mcpPath(s.Name), nil, nil)
|
||||
}}
|
||||
return nil
|
||||
}})
|
||||
return &menuDialog{title: s.Name + " · " + s.URL, items: items}
|
||||
}
|
||||
|
||||
// mcpTest asks the server for its tools, with its credentials, as a call
|
||||
// through the gateway would.
|
||||
func (m *consoleModel) mcpTest(s adminapi.MCPServerView) tea.Cmd {
|
||||
return m.call("Test "+s.Name, http.MethodPost, mcpPath(s.Name)+"/test", nil, func(raw json.RawMessage) string {
|
||||
var t adminapi.MCPTestView
|
||||
_ = json.Unmarshal(raw, &t)
|
||||
if !t.OK {
|
||||
return "It did not answer: " + t.Error
|
||||
}
|
||||
if len(t.Tools) == 0 {
|
||||
return "It answers, and offers no tools."
|
||||
}
|
||||
return fmt.Sprintf("It answers, with %d tools:\n\n %s", len(t.Tools), strings.Join(t.Tools, "\n "))
|
||||
})
|
||||
}
|
||||
|
||||
func (m *consoleModel) mcpEditForm(s adminapi.MCPServerView) {
|
||||
rate := ""
|
||||
if s.RatePerMinute > 0 {
|
||||
rate = strconv.Itoa(s.RatePerMinute)
|
||||
}
|
||||
timeout := ""
|
||||
if s.TimeoutSeconds > 0 {
|
||||
timeout = strconv.Itoa(s.TimeoutSeconds)
|
||||
}
|
||||
m.openForm(&formDialog{title: "Change " + s.Name,
|
||||
note: "Its credential stays as it is; Replace its credential for that.",
|
||||
fields: []formField{
|
||||
textF("address", s.URL, ""),
|
||||
textF("about", s.Description, ""),
|
||||
textF("timeout", timeout, "seconds a call may take"),
|
||||
textF("rate", rate, "calls a minute per caller; 0 is no bound"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
if err := adminapi.CheckMCPAddress(strings.TrimSpace(v[0])); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
timeout, err := atoiOr(v[2], 0)
|
||||
if err != nil || timeout < 0 {
|
||||
return nil, fmt.Errorf("timeout is whole seconds")
|
||||
}
|
||||
rate, err := atoiOr(v[3], 0)
|
||||
if err != nil || rate < 0 {
|
||||
return nil, fmt.Errorf("rate is calls a minute")
|
||||
}
|
||||
req := adminapi.MCPServerRequest{URL: strings.TrimSpace(v[0]), Description: v[1], TimeoutSeconds: timeout,
|
||||
RatePerMinute: rate}
|
||||
// An emptied field empties it, rather than leaving it as it was.
|
||||
if strings.TrimSpace(v[1]) == "" {
|
||||
req.Clear = append(req.Clear, "description")
|
||||
}
|
||||
if timeout == 0 {
|
||||
req.Clear = append(req.Clear, "timeout")
|
||||
}
|
||||
if rate == 0 {
|
||||
req.Clear = append(req.Clear, "rate")
|
||||
}
|
||||
return m.call("Change "+s.Name, http.MethodPut, mcpPath(s.Name), req, nil), nil
|
||||
}})
|
||||
}
|
||||
|
||||
func (m *consoleModel) mcpCredentialForm(s adminapi.MCPServerView) {
|
||||
m.openForm(&formDialog{title: "Replace the credential " + s.Name + " is sent",
|
||||
note: "Kept in the store, sealed. The old value is overwritten, not shown.",
|
||||
fields: []formField{
|
||||
choiceF("credential", []string{"bearer", "header"}, "bearer"),
|
||||
textF("header", "", "for credential: header, its name (X-Api-Key)"),
|
||||
pwF("secret", "the token or key"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
if v[2] == "" {
|
||||
return nil, fmt.Errorf("the new secret?")
|
||||
}
|
||||
secrets, err := mcpAuth(v[0], strings.TrimSpace(v[1]), v[2])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return m.call("Credential of "+s.Name, http.MethodPut, mcpPath(s.Name),
|
||||
adminapi.MCPServerRequest{URL: s.URL, Secrets: secrets, HeaderRefs: map[string]string{}}, nil), nil
|
||||
}})
|
||||
}
|
||||
|
||||
// mcpGrantForm grants a server's tools to part of the tree, a person or a
|
||||
// key, beside whatever else that level is granted.
|
||||
func (m *consoleModel) mcpGrantForm(s adminapi.MCPServerView) {
|
||||
if m.mcp.ids == nil {
|
||||
m.dialog, m.dialogTitle = "error", "Grant "+s.Name
|
||||
m.dialogText = "The store cannot be read from this machine, so its grants cannot be changed here.\n" +
|
||||
"On the control plane: llmgw keys limits --scope team --id ORG/TEAM --tools '" + s.Name + ".*'"
|
||||
return
|
||||
}
|
||||
m.openForm(&formDialog{title: "Grant " + s.Name + "'s tools",
|
||||
note: "Seen there, and below, where every level above that names tools names these too.\n" +
|
||||
"A level given its first grant keeps what it saw from above; other grants stay.",
|
||||
fields: []formField{
|
||||
choiceF("to", []string{"team", "project", "org", "user", "key"}, "team"),
|
||||
textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, an e-mail, or a key's ID"),
|
||||
textF("tools", s.Name+".*", "patterns: "+s.Name+".* for all, or "+s.Name+".search,"+s.Name+".get"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
which := strings.TrimSpace(v[1])
|
||||
if which == "" {
|
||||
return nil, fmt.Errorf("which %s?", v[0])
|
||||
}
|
||||
owner, ok := m.mcp.ids[v[0]+":"+strings.ToLower(which)]
|
||||
if v[0] == "key" {
|
||||
owner, ok = store.LimitOwner{Scope: store.ScopeKey, ID: which}, m.mcp.keys[which]
|
||||
}
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no %s %q", v[0], which)
|
||||
}
|
||||
var tools []string
|
||||
for _, t := range strings.Split(v[2], ",") {
|
||||
if t = strings.TrimSpace(t); t != "" {
|
||||
if !strings.HasPrefix(t, s.Name+".") {
|
||||
return nil, fmt.Errorf("%q is not one of %s's tools: they start %s.", t, s.Name, s.Name)
|
||||
}
|
||||
tools = append(tools, t)
|
||||
}
|
||||
}
|
||||
if len(tools) == 0 {
|
||||
return nil, fmt.Errorf("which tools? %s.* for all", s.Name)
|
||||
}
|
||||
return m.mcpSetTools(s.Name, owner, tools), nil
|
||||
}})
|
||||
}
|
||||
|
||||
// mcpSetTools grants a level these of a server's tools, or with none takes
|
||||
// them away, in the store (store.GrantTools, store.RevokeTools): read and
|
||||
// written at once, so another administrator's change since this screen
|
||||
// was loaded is kept.
|
||||
func (m *consoleModel) mcpSetTools(server string, owner store.LimitOwner, tools []string) tea.Cmd {
|
||||
h, where := m.st, m.mcp.where
|
||||
title := "Grant " + server
|
||||
if len(tools) == 0 {
|
||||
title = "Revoke " + server
|
||||
}
|
||||
return func() tea.Msg {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
s, err := h.get(ctx)
|
||||
if err != nil {
|
||||
return actionDoneMsg{title: title, err: err}
|
||||
}
|
||||
if len(tools) == 0 {
|
||||
if err := s.RevokeTools(ctx, owner.Scope, owner.ID, server); err != nil {
|
||||
return actionDoneMsg{title: title, err: err}
|
||||
}
|
||||
return actionDoneMsg{title: title, text: server + "'s tools taken back there; its other grants stay."}
|
||||
}
|
||||
g, err := s.GrantTools(ctx, owner.Scope, owner.ID, server, tools)
|
||||
if err != nil {
|
||||
return actionDoneMsg{title: title, err: err}
|
||||
}
|
||||
text := "Granted " + strings.Join(tools, ", ") + "."
|
||||
for _, b := range g.BlockedBy {
|
||||
text += "\n\nBut " + string(b.Scope) + " " + or(where[b.CounterKey()], b.ID) + " names tools and not these:\n" +
|
||||
server + "'s stay out of reach there until it grants them too."
|
||||
}
|
||||
return actionDoneMsg{title: title, text: text}
|
||||
}
|
||||
}
|
||||
|
||||
// mcpOpen has a closed level name no tools of its own (store.OpenTools).
|
||||
func (m *consoleModel) mcpOpen(owner store.LimitOwner) tea.Cmd {
|
||||
h := m.st
|
||||
return func() tea.Msg {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
s, err := h.get(ctx)
|
||||
if err == nil {
|
||||
err = s.OpenTools(ctx, owner.Scope, owner.ID)
|
||||
}
|
||||
if err != nil {
|
||||
return actionDoneMsg{title: "Open", err: err}
|
||||
}
|
||||
return actionDoneMsg{title: "Open", text: "Opened: it names no tools of its own now."}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func TestAnMCPServerIsAddedAndTestedFromTheConsole(t *testing.T) {
|
||||
m, s := consoleWith(t, cMCP)
|
||||
m.mcp = mcpState{loaded: true, grants: []mcpGrant{}, servers: []adminapi.MCPServerView{{
|
||||
Name: "jira", URL: "https://jira-mcp.internal/mcp", Enabled: true, RatePerMinute: 30,
|
||||
HeaderRefs: map[string]string{"Authorization": "store:mcp/jira/Authorization"},
|
||||
}}}
|
||||
var tm tea.Model = m
|
||||
view := tm.View()
|
||||
for _, want := range []string{"MCP servers", "jira", "Authorization (secret)",
|
||||
"30/min", "nobody yet", "claude mcp add"} {
|
||||
if !strings.Contains(view, want) {
|
||||
t.Errorf("missing %q in:\n%s", want, view)
|
||||
}
|
||||
}
|
||||
if strings.Contains(view, "store:mcp/jira") {
|
||||
t.Error("where the secret is kept is shown")
|
||||
}
|
||||
|
||||
// A: a server with a bearer token, which is sent once and kept sealed.
|
||||
tm = press(tm, "a")
|
||||
f := tm.(consoleModel).formD
|
||||
if f == nil {
|
||||
t.Fatal("no form")
|
||||
}
|
||||
f.fields[0].text.SetValue("hr")
|
||||
f.fields[1].text.SetValue("https://hr-mcp.internal/mcp")
|
||||
f.fields[5].text.SetValue("t0ken")
|
||||
f.focus = len(f.fields)
|
||||
_ = pressRun(tm, "enter")
|
||||
got := s.last()
|
||||
var req adminapi.MCPServerRequest
|
||||
if !strings.HasPrefix(got, "PUT /api/v1/mcp/servers/hr ") ||
|
||||
json.Unmarshal([]byte(strings.SplitN(got, " ", 3)[2]), &req) != nil ||
|
||||
req.URL != "https://hr-mcp.internal/mcp" || req.Secrets["Authorization"] != "Bearer t0ken" ||
|
||||
req.Enabled == nil || !*req.Enabled {
|
||||
t.Fatalf("sent %q", got)
|
||||
}
|
||||
|
||||
// Enter on jira: test it, then turn it off.
|
||||
tm = press(m, "enter")
|
||||
md := tm.(consoleModel).menuD
|
||||
if md == nil || md.items[0].label != "Test, and list its tools" || !strings.Contains(md.title, "https://jira-mcp.internal/mcp") {
|
||||
t.Fatalf("menu %+v", md)
|
||||
}
|
||||
_ = pressRun(tm, "enter")
|
||||
if got := s.last(); got != "POST /api/v1/mcp/servers/jira/test " {
|
||||
t.Fatalf("sent %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnMCPServersToolsAreGrantedBesideOtherGrants(t *testing.T) {
|
||||
// The servers, from a control plane; the grants, from the store.
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode([]adminapi.MCPServerView{
|
||||
{Name: "hr", URL: "https://hr/mcp", Enabled: true}, {Name: "jira", URL: "https://jira/mcp", Enabled: true}})
|
||||
}))
|
||||
defer srv.Close()
|
||||
dir := t.TempDir()
|
||||
cfgPath := dir + "/server.yaml"
|
||||
if err := os.WriteFile(cfgPath, []byte("auth:\n store: "+dir+"/llmgw.db\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := config.Load(cfgPath); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m := newConsoleModel(srv.URL, cfgPath)
|
||||
defer m.st.close()
|
||||
m.loaded, m.inPanel, m.menu = true, true, cMCP
|
||||
ctx := context.Background()
|
||||
st, err := m.st.get(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
org, _ := st.CreateOrg(ctx, "Acme")
|
||||
team, _ := st.CreateTeam(ctx, org, "ml")
|
||||
if err := st.PutLimits(ctx, store.Limits{Scope: store.ScopeTeam, ScopeID: team, AllowedTools: []string{"jira.*"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
toolsOf := func() []string {
|
||||
all, _ := st.AllLimits(ctx)
|
||||
return all[store.LimitOwner{Scope: store.ScopeTeam, ID: team}.CounterKey()].AllowedTools
|
||||
}
|
||||
refresh := func(tm tea.Model) tea.Model {
|
||||
tm, _ = tm.Update(tm.(consoleModel).fetchMCP()())
|
||||
return tm
|
||||
}
|
||||
|
||||
var tm tea.Model = refresh(m)
|
||||
if v := tm.View(); !strings.Contains(v, "team Acme/ml") || !strings.Contains(v, "nobody yet") {
|
||||
t.Fatalf("jira's grant, and hr's none, should show:\n%s", v)
|
||||
}
|
||||
// hr is the first row: grant its tools to the team.
|
||||
tm = press(tm, "enter", "down", "enter")
|
||||
f := tm.(consoleModel).formD
|
||||
if f == nil || f.fields[2].value() != "hr.*" {
|
||||
t.Fatalf("no grant form, or not for hr: %+v", f)
|
||||
}
|
||||
f.fields[1].text.SetValue("Acme/ml")
|
||||
f.focus = len(f.fields)
|
||||
tm = pressRun(tm, "enter")
|
||||
if got := toolsOf(); !slices.Equal(got, []string{"jira.*", "hr.*"}) {
|
||||
t.Fatalf("the team's grants are %v; jira's should have stayed beside hr's", got)
|
||||
}
|
||||
|
||||
// Revoke hr from the team: jira's stays.
|
||||
tm = refresh(press(tm, "x"))
|
||||
tm = press(tm, "enter", "down", "down", "enter")
|
||||
if c := tm.(consoleModel).confirmD; c == nil || !strings.Contains(c.title, "Revoke hr from team Acme/ml") {
|
||||
t.Fatalf("no revoke: %+v", c)
|
||||
}
|
||||
_ = pressRun(tm, "y")
|
||||
if got := toolsOf(); !slices.Equal(got, []string{"jira.*"}) {
|
||||
t.Fatalf("after the revoke the team's grants are %v", got)
|
||||
}
|
||||
}
|
||||
@@ -32,7 +32,7 @@ func (m consoleModel) mirrorClient() (*mirror.Client, error) {
|
||||
if cfg.Mirror.API == "" {
|
||||
return nil, errors.New("no mirror: mirror.api in server.yaml is empty")
|
||||
}
|
||||
return mirror.NewClient(cfg.Mirror.API, cfg.Mirror.TLS, 10*time.Second)
|
||||
return mirror.NewFailoverClient(append([]string{cfg.Mirror.API}, cfg.Mirror.Peers...), cfg.Mirror.TLS, 10*time.Second)
|
||||
}
|
||||
|
||||
func (m consoleModel) fetchMirror() tea.Cmd {
|
||||
|
||||
@@ -345,6 +345,10 @@ type Mirror struct {
|
||||
// listeners on the mirror precisely so this one can be locked down without
|
||||
// re-certifying every agent's fetcher.
|
||||
API string `koanf:"api"`
|
||||
// Peers are the management addresses of mirrors kept in step with this
|
||||
// one (llmgw mirror serve --peer): asked when it cannot be reached, and
|
||||
// a node given the same file on one of them while it is down.
|
||||
Peers []string `koanf:"peers"`
|
||||
// TLS authenticates to the mirror's management API, with the same CA the
|
||||
// agents use. Empty leaves the connection in plain HTTP, which is only
|
||||
// acceptable on a network where nothing untrusted can reach that port —
|
||||
|
||||
@@ -205,6 +205,8 @@ func (a *API) Register(mux Mux) {
|
||||
mux.HandleFunc("GET /api/v1/updates", a.listUpdates)
|
||||
mux.HandleFunc("POST /api/v1/updates/check", a.checkUpdates)
|
||||
mux.HandleFunc("POST /api/v1/models/{name}/update", a.updateModel)
|
||||
mux.HandleFunc("POST /api/v1/models/{name}/context", a.setModelContext)
|
||||
mux.HandleFunc("POST /api/v1/models/{name}/gpus", a.setModelGPUs)
|
||||
mux.HandleFunc("POST /api/v1/models/{name}/rollback", a.rollbackModel)
|
||||
a.registerCanary(mux)
|
||||
a.registerEval(mux)
|
||||
@@ -248,6 +250,9 @@ type NodeView struct {
|
||||
MemoryGB int `json:"memory_gb,omitempty"`
|
||||
// CPUCount is the machine's CPUs, 0 when the agent did not report them.
|
||||
CPUCount int `json:"cpu_count,omitempty"`
|
||||
// MemBandwidthGBs is how fast its memory reads, in GB/s, measured when
|
||||
// its agent started; 0 when not reported. A model's speed on the CPU.
|
||||
MemBandwidthGBs float64 `json:"mem_bandwidth_gbs,omitempty"`
|
||||
// Engines are the inference engines installed here, as the agent found
|
||||
// them. Absent means the agent is older than the probe and said nothing,
|
||||
// which is not the same as having none — placement treats it that way too,
|
||||
@@ -442,6 +447,10 @@ func loadView(l *inventory.Load, now time.Time) *LoadView {
|
||||
// ModelView is a model as `ctl models` shows it.
|
||||
type ModelView struct {
|
||||
Name string `json:"name"`
|
||||
// ContextLength is what one request may use, MaxOutput what it may
|
||||
// answer; 0 when the catalog does not say.
|
||||
ContextLength int `json:"context_length,omitempty"`
|
||||
MaxOutput int `json:"max_output,omitempty"`
|
||||
// Canary is a share of its requests sent to a candidate; CanaryOf is
|
||||
// set on the candidate, naming the model it is a canary of.
|
||||
Canary *catalog.Canary `json:"canary,omitempty"`
|
||||
@@ -703,6 +712,15 @@ func (a *API) Nodes() []NodeView {
|
||||
v.EngineBuilds[e] = b
|
||||
}
|
||||
}
|
||||
v.MemBandwidthGBs = n.Host.MemBandwidthGBs
|
||||
for _, g := range n.GPUs {
|
||||
if g.MIG != nil {
|
||||
continue // said of the card it is carved from
|
||||
}
|
||||
if p := g.PCIe.Problem(); p != "" {
|
||||
v.Warnings = append(v.Warnings, fmt.Sprintf("GPU %d (%s) %s", g.Index, g.Name, p))
|
||||
}
|
||||
}
|
||||
if h := n.Host; h.OSID != "" {
|
||||
v.OS = strings.TrimSpace(h.OSID + " " + h.OSVersion + " " + h.Arch)
|
||||
if h.Arch == "amd64" && !h.HasCPUFlag("avx2") {
|
||||
@@ -825,6 +843,7 @@ func (a *API) Models() []ModelView {
|
||||
h := health[m.Name]
|
||||
v := ModelView{
|
||||
Name: m.Name, Description: m.Description, Aliases: m.Aliases, Source: m.Source,
|
||||
ContextLength: m.ContextLength, MaxOutput: m.MaxOutput,
|
||||
Min: m.Scaling.Min, Max: m.Scaling.Max,
|
||||
Ready: ready[m.Name], Total: total[m.Name],
|
||||
Fallbacks: m.Fallbacks,
|
||||
@@ -1736,6 +1755,46 @@ func (a *API) job(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, job)
|
||||
}
|
||||
|
||||
// setModelContext changes a served model's context, its engines' arguments
|
||||
// and memory with it, and starts its replicas again on it.
|
||||
func (a *API) setModelContext(w http.ResponseWriter, r *http.Request) {
|
||||
if a.serve == nil {
|
||||
writeError(w, http.StatusNotImplemented, "this control plane cannot change a model's context")
|
||||
return
|
||||
}
|
||||
var req modelserve.ContextRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
change, err := a.serve.SetContext(r.Context(), r.PathValue("name"), req, actorOf(r))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, change)
|
||||
}
|
||||
|
||||
// setModelGPUs spreads a served model over more cards, or fewer, and
|
||||
// starts its replicas again on it.
|
||||
func (a *API) setModelGPUs(w http.ResponseWriter, r *http.Request) {
|
||||
if a.serve == nil {
|
||||
writeError(w, http.StatusNotImplemented, "this control plane cannot change a model's cards")
|
||||
return
|
||||
}
|
||||
var req modelserve.GPUsRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
change, err := a.serve.SetGPUs(r.Context(), r.PathValue("name"), req, actorOf(r))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, change)
|
||||
}
|
||||
|
||||
func (a *API) listUpdates(w http.ResponseWriter, _ *http.Request) {
|
||||
out := make([]modelserve.UpdateInfo, 0)
|
||||
if a.serve != nil {
|
||||
@@ -1829,9 +1888,12 @@ func (a *API) setOrigin(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// PoolNodes is the fleet's ready nodes as a plan sees them: their cards,
|
||||
// which of them are free, the fabric they share, and their memory.
|
||||
func (a *API) PoolNodes() []modelplan.PoolNode {
|
||||
func (a *API) PoolNodes() []modelplan.PoolNode { return PoolNodesOf(a.Nodes()) }
|
||||
|
||||
// PoolNodesOf is nodes as a plan sees them, the ready ones only.
|
||||
func PoolNodesOf(nodes []NodeView) []modelplan.PoolNode {
|
||||
var out []modelplan.PoolNode
|
||||
for _, n := range a.Nodes() {
|
||||
for _, n := range nodes {
|
||||
if n.State != "ready" {
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -735,3 +735,17 @@ func TestRemovingANodeRevokesItsCertificate(t *testing.T) {
|
||||
t.Error("the node was removed though its certificate was not revoked")
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodesCardInTheWrongSlotIsAWarning(t *testing.T) {
|
||||
n := gpuNode("g1", 2)
|
||||
n.GPUs[0].PCIe = discovery.PCIeLink{Gen: 5, GPUGen: 5, Width: 16, WidthMax: 16}
|
||||
n.GPUs[1].PCIe = discovery.PCIeLink{Gen: 5, GPUGen: 5, Width: 8, WidthMax: 16}
|
||||
api, _ := setup(t, testModels(), n)
|
||||
var warned []string
|
||||
for _, v := range api.Nodes() {
|
||||
warned = append(warned, v.Warnings...)
|
||||
}
|
||||
if len(warned) != 1 || !strings.Contains(warned[0], "GPU 1 (NVIDIA H100 80GB HBM3) linked at x8 of its x16 lanes") {
|
||||
t.Errorf("warnings %q", warned)
|
||||
}
|
||||
}
|
||||
@@ -57,6 +57,9 @@ type MCPServerRequest struct {
|
||||
// RatePerMinute bounds each caller's calls to its tools; a negative
|
||||
// number takes the bound away.
|
||||
RatePerMinute int `json:"rate_per_minute,omitempty"`
|
||||
// Clear empties what an empty field above would leave as it is:
|
||||
// "description", "timeout" (back to the default), "rate".
|
||||
Clear []string `json:"clear,omitempty"`
|
||||
}
|
||||
|
||||
// MCPTestView is what the server said when asked: its tools, or why not.
|
||||
@@ -150,11 +153,18 @@ func (a *API) registerMCP(mux Mux) {
|
||||
// mcpServerFrom is the server a request makes, over the one of that name if
|
||||
// there is one: its secrets kept in the store, and named by reference.
|
||||
func (a *API) mcpServerFrom(ctx context.Context, name string, req MCPServerRequest, by string) (store.MCPServer, error) {
|
||||
return MCPServerFrom(ctx, a.mcp, name, req, by)
|
||||
}
|
||||
|
||||
// MCPServerFrom is the server a request makes, over the one of that name in
|
||||
// st if there is one: its secrets kept in the store, sealed, and named by
|
||||
// reference. Not stored; the caller puts it. For the API and the admin page.
|
||||
func MCPServerFrom(ctx context.Context, st MCPStore, name string, req MCPServerRequest, by string) (store.MCPServer, error) {
|
||||
if !store.ValidMCPServerName(name) {
|
||||
return store.MCPServer{}, errors.New("a server's name is lowercase letters, digits, - and _, starting with a letter: " +
|
||||
"its tools are NAME.tool")
|
||||
}
|
||||
m, err := a.mcp.MCPServer(ctx, name)
|
||||
m, err := st.MCPServer(ctx, name)
|
||||
switch {
|
||||
case errors.Is(err, store.ErrNotFound):
|
||||
m = store.MCPServer{Name: name, Enabled: true, CreatedBy: by}
|
||||
@@ -164,8 +174,8 @@ func (a *API) mcpServerFrom(ctx context.Context, name string, req MCPServerReque
|
||||
if req.URL != "" {
|
||||
m.URL = req.URL
|
||||
}
|
||||
if u, err := url.Parse(m.URL); m.URL != AdminMCPAddress && (err != nil || (u.Scheme != "https" && u.Scheme != "http") || u.Host == "") {
|
||||
return m, errors.New("the server's url is an http:// or https:// address, or " + AdminMCPAddress + " for this cluster's own")
|
||||
if err := CheckMCPAddress(m.URL); err != nil {
|
||||
return m, err
|
||||
}
|
||||
if req.Description != "" {
|
||||
m.Description = req.Description
|
||||
@@ -184,7 +194,7 @@ func (a *API) mcpServerFrom(ctx context.Context, name string, req MCPServerReque
|
||||
return m, errors.New("a secret header needs a name and a value")
|
||||
}
|
||||
key := "mcp/" + name + "/" + header
|
||||
if err := a.mcp.SetClusterSecret(ctx, key, value); err != nil {
|
||||
if err := st.SetClusterSecret(ctx, key, value); err != nil {
|
||||
return m, err
|
||||
}
|
||||
m.HeaderRefs[header] = "store:" + key
|
||||
@@ -202,12 +212,49 @@ func (a *API) mcpServerFrom(ctx context.Context, name string, req MCPServerReque
|
||||
case req.RatePerMinute < 0:
|
||||
m.RatePerMinute = 0
|
||||
}
|
||||
for _, f := range req.Clear {
|
||||
switch f {
|
||||
case "description":
|
||||
m.Description = ""
|
||||
case "timeout":
|
||||
m.Timeout = 0
|
||||
case "rate":
|
||||
m.RatePerMinute = 0
|
||||
default:
|
||||
return m, errors.New("clear names description, timeout or rate, not " + f)
|
||||
}
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// CheckMCPAddress refuses an address a server cannot be reached at: an
|
||||
// http(s) one, or this cluster's own. Not one with a user and password in
|
||||
// it: a credential is a header, kept sealed, and an address is shown and
|
||||
// recorded in the audit log as written.
|
||||
func CheckMCPAddress(addr string) error {
|
||||
if addr == AdminMCPAddress {
|
||||
return nil
|
||||
}
|
||||
u, err := url.Parse(addr)
|
||||
if err != nil || (u.Scheme != "https" && u.Scheme != "http") || u.Host == "" {
|
||||
return errors.New("the server's url is an http:// or https:// address, or " + AdminMCPAddress + " for this cluster's own")
|
||||
}
|
||||
if u.User != nil {
|
||||
return errors.New("the server's url has a user and password in it: give the credential as a header, which is kept sealed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// testMCPServer asks a server for its tools, with its credentials, as /mcp
|
||||
// would.
|
||||
func (a *API) testMCPServer(ctx context.Context, m store.MCPServer) MCPTestView {
|
||||
return TestMCPServer(ctx, m, http.HandlerFunc(a.MCP))
|
||||
}
|
||||
|
||||
// TestMCPServer asks a server for its tools, with its credentials, as /mcp
|
||||
// would; local answers for this cluster's own (AdminMCPAddress), and
|
||||
// without it that one is not asked.
|
||||
func TestMCPServer(ctx context.Context, m store.MCPServer, local http.Handler) MCPTestView {
|
||||
ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
|
||||
defer cancel()
|
||||
c := &mcpclient.Client{URL: m.URL, Header: func(ctx context.Context) (http.Header, error) {
|
||||
@@ -225,7 +272,10 @@ func (a *API) testMCPServer(ctx context.Context, m store.MCPServer) MCPTestView
|
||||
return h, nil
|
||||
}}
|
||||
if m.URL == AdminMCPAddress {
|
||||
c.URL, c.HTTP = InProcessURL, &http.Client{Transport: InProcess(http.HandlerFunc(a.MCP))}
|
||||
if local == nil {
|
||||
return MCPTestView{Error: "this cluster's own tools are answered only by a control plane"}
|
||||
}
|
||||
c.URL, c.HTTP = InProcessURL, &http.Client{Transport: InProcess(local)}
|
||||
}
|
||||
tools, err := c.Tools(ctx)
|
||||
if err != nil {
|
||||
|
||||
@@ -213,6 +213,11 @@ var Routes = []Route{
|
||||
{Method: "POST", Path: "/api/v1/models/{name}/update", Summary: "Move a model to its repository's newer weights",
|
||||
Description: "With a quality suite, or the model's standing gate, the new weights are asked it first as a hidden canary.",
|
||||
Request: modelserve.Change{}, OptionalBody: true, Response: modelserve.Job{}, Status: http.StatusAccepted},
|
||||
{Method: "POST", Path: "/api/v1/models/{name}/context", Summary: "Change a served model's context: the catalog, its engines' arguments and memory, then its replicas restarted",
|
||||
Request: modelserve.ContextRequest{}, Response: modelserve.ContextChange{}},
|
||||
{Method: "POST", Path: "/api/v1/models/{name}/gpus", Summary: "Spread a served model over more cards or fewer: its profile, checked against its KV heads, its cards' memory and the nodes, then its replicas restarted",
|
||||
Description: "Tensor parallel for vLLM, SGLang and TensorRT-LLM, a layer split for llama.cpp, on one node. With no room beside the running replicas, in_place stops each first.",
|
||||
Request: modelserve.GPUsRequest{}, Response: modelserve.GPUsChange{}},
|
||||
{Method: "POST", Path: "/api/v1/models/{name}/rollback", Summary: "Go back to the weights the last update replaced",
|
||||
Response: Result{}},
|
||||
{Method: "POST", Path: "/api/v1/models/{name}/origin", Summary: "Record the repository a model's weights came from",
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"fmt"
|
||||
"maps"
|
||||
"runtime"
|
||||
"slices"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
@@ -34,6 +36,7 @@ func NodeFromRegister(r *llmgwv1.Register) inventory.Node {
|
||||
n.Host = discovery.Host{
|
||||
OSID: h.GetOsId(), OSVersion: h.GetOsVersion(), OSCodename: h.GetOsCodename(),
|
||||
Arch: h.GetArch(), CPUFlags: h.GetCpuFlags(), NVIDIADriver: h.GetNvidiaDriver(),
|
||||
MemBandwidthGBs: h.GetMemBandwidthGbs(),
|
||||
}
|
||||
}
|
||||
for _, w := range r.GetWeights() {
|
||||
@@ -72,6 +75,8 @@ func NodeFromRegister(r *llmgwv1.Register) inventory.Node {
|
||||
Problems: g.GetProblems(),
|
||||
},
|
||||
MIG: migFrom(g.GetMig()),
|
||||
PCIe: discovery.PCIeLink{Gen: int(g.GetPcieGen()), GPUGen: int(g.GetPcieGpuGen()),
|
||||
Width: int(g.GetPcieWidth()), WidthMax: int(g.GetPcieWidthMax())},
|
||||
})
|
||||
}
|
||||
n.IMEX = discovery.IMEX{State: r.GetImex().GetState(), Channels: int(r.GetImex().GetChannels())}
|
||||
@@ -126,6 +131,10 @@ func RegisterFromInventory(nodeID, hostname, mgmtAddr, agentVersion string,
|
||||
Problems: g.Health.Problems,
|
||||
TemperatureC: int32(g.Health.TemperatureC),
|
||||
EccErrors: int32(g.Health.ECCErrors),
|
||||
PcieGen: int32(g.PCIe.Gen),
|
||||
PcieGpuGen: int32(g.PCIe.GPUGen),
|
||||
PcieWidth: int32(g.PCIe.Width),
|
||||
PcieWidthMax: int32(g.PCIe.WidthMax),
|
||||
})
|
||||
if f := g.Fabric; f != nil {
|
||||
last := r.Gpus[len(r.Gpus)-1]
|
||||
@@ -182,8 +191,31 @@ func AdoptedFromRegister(nodeID string, shards []*llmgwv1.RunningShard) []reconc
|
||||
}
|
||||
|
||||
// StartShardFromCommand converts a reconciler command into a wire message.
|
||||
// sourceRewrite is where a model is fetched while its mirror is down: set
|
||||
// when mirrors are kept in step, else nil and the catalog's address as it is.
|
||||
var sourceRewrite atomic.Pointer[func(string) string]
|
||||
|
||||
// SetSourceRewrite has every shard started from now on fetch its model, and
|
||||
// adapters, from the address f gives: one on a mirror that answers.
|
||||
func SetSourceRewrite(f func(string) string) { sourceRewrite.Store(&f) }
|
||||
|
||||
func rewriteSource(src string) string {
|
||||
if f := sourceRewrite.Load(); f != nil && src != "" {
|
||||
return (*f)(src)
|
||||
}
|
||||
return src
|
||||
}
|
||||
|
||||
func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
|
||||
p := cmd.Spec.Profile
|
||||
p.ModelPath = rewriteSource(p.ModelPath)
|
||||
if len(p.Adapters) > 0 {
|
||||
ads := slices.Clone(p.Adapters)
|
||||
for i := range ads {
|
||||
ads[i].Path = rewriteSource(ads[i].Path)
|
||||
}
|
||||
p.Adapters = ads
|
||||
}
|
||||
return &llmgwv1.StartShard{
|
||||
ReplicaId: cmd.ReplicaID,
|
||||
ShardId: cmd.ShardID,
|
||||
@@ -201,8 +233,9 @@ func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
|
||||
GpusPerNode: int32(p.GPUsPerNode), Nodes: int32(p.Nodes),
|
||||
Args: p.Args, Env: p.Env,
|
||||
ModelPath: p.ModelPath, ServedModelName: p.ServedModelName,
|
||||
Adapters: adaptersToWire(p.Adapters),
|
||||
KvEvents: p.KVEvents,
|
||||
Adapters: adaptersToWire(p.Adapters),
|
||||
KvEvents: p.KVEvents,
|
||||
AllowPickle: p.AllowPickle,
|
||||
},
|
||||
StartupTimeoutMs: cmd.StartupTimeout.Milliseconds(),
|
||||
Pool: cmd.Pool,
|
||||
@@ -238,8 +271,9 @@ func ShardSpecFromStart(m *llmgwv1.StartShard) (engines.ShardSpec, time.Duration
|
||||
GPUsPerNode: int(p.GetGpusPerNode()), Nodes: int(p.GetNodes()),
|
||||
Args: p.GetArgs(), Env: p.GetEnv(),
|
||||
ModelPath: p.GetModelPath(), ServedModelName: p.GetServedModelName(),
|
||||
Adapters: adaptersFromWire(p.GetAdapters()),
|
||||
KVEvents: p.GetKvEvents(),
|
||||
Adapters: adaptersFromWire(p.GetAdapters()),
|
||||
KVEvents: p.GetKvEvents(),
|
||||
AllowPickle: p.GetAllowPickle(),
|
||||
},
|
||||
}
|
||||
return spec, time.Duration(m.GetStartupTimeoutMs()) * time.Millisecond
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
@@ -18,3 +19,20 @@ func TestAdaptersTravelToTheAgentInOrder(t *testing.T) {
|
||||
t.Errorf("adapters %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// While its mirror is down, a shard is sent its model and adapters on the
|
||||
// mirror the rewrite names, and the command it was made from is left as it is.
|
||||
func TestAShardIsSentItsModelOnAMirrorThatAnswers(t *testing.T) {
|
||||
SetSourceRewrite(func(s string) string { return strings.Replace(s, "http://mirror-a", "http://mirror-b", 1) })
|
||||
t.Cleanup(func() { SetSourceRewrite(func(s string) string { return s }) })
|
||||
ads := []engines.Adapter{{Name: "sql", Path: "http://mirror-a/sql.gguf"}}
|
||||
cmd := reconciler.ShardCommand{ReplicaID: "r", ShardID: "s", Spec: engines.ShardSpec{
|
||||
Profile: engines.Profile{Name: "cpu", Engine: "llamacpp", ModelPath: "http://mirror-a/qwen/v1/", Adapters: ads}}}
|
||||
spec, _ := ShardSpecFromStart(StartShardFromCommand(cmd))
|
||||
if spec.Profile.ModelPath != "http://mirror-b/qwen/v1/" || spec.Profile.Adapters[0].Path != "http://mirror-b/sql.gguf" {
|
||||
t.Errorf("sent %+v", spec.Profile)
|
||||
}
|
||||
if ads[0].Path != "http://mirror-a/sql.gguf" {
|
||||
t.Error("the catalog's adapter was changed in place")
|
||||
}
|
||||
}
|
||||
@@ -325,3 +325,18 @@ func TestAProblemOnANodeWithoutGPUsDoesNotPanic(t *testing.T) {
|
||||
t.Fatal("the node is gone")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAHeartbeatCarriesTheMemorysMeasuredSpeed(t *testing.T) {
|
||||
// Measured in the background after the agent starts, so it arrives in a
|
||||
// heartbeat, not in the registration.
|
||||
s := newTestService()
|
||||
nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready, MemoryTotalMB: 7 * 1024})
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", MemBandwidthGbs: 41.9})
|
||||
if n, _ := s.inv.Get("n1"); n.Host.MemBandwidthGBs != 41.9 {
|
||||
t.Errorf("bandwidth %v, want 41.9", n.Host.MemBandwidthGBs)
|
||||
}
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1"})
|
||||
if n, _ := s.inv.Get("n1"); n.Host.MemBandwidthGBs != 41.9 {
|
||||
t.Errorf("a heartbeat without it erased it: %v", n.Host.MemBandwidthGBs)
|
||||
}
|
||||
}
|
||||
@@ -547,10 +547,11 @@ func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
|
||||
storage := hb.GetStorage()
|
||||
memMB := int(hb.GetMemoryTotalMb())
|
||||
cpus := int(hb.GetCpuCount())
|
||||
bw := hb.GetMemBandwidthGbs()
|
||||
if len(used) == 0 && len(foreign) == 0 && len(problems) == 0 && len(storage) == 0 &&
|
||||
!hb.GetWeightsChanged() && memMB == 0 && cpus == 0 && !hb.GetEnginesChanged() &&
|
||||
!hb.GetProvisionChanged() && hb.GetLoad() == nil && len(hb.GetIbPorts()) == 0 &&
|
||||
!hb.GetFabricsChanged() && len(hb.GetGpuFabric()) == 0 && hb.GetImex() == nil {
|
||||
!hb.GetFabricsChanged() && len(hb.GetGpuFabric()) == 0 && hb.GetImex() == nil && bw == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -622,6 +623,9 @@ func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
|
||||
"from_mb", n.MemoryTotalMB, "to_mb", memMB)
|
||||
n.MemoryTotalMB = memMB
|
||||
}
|
||||
if bw > 0 {
|
||||
n.Host.MemBandwidthGBs = bw
|
||||
}
|
||||
if cpus > 0 && cpus != n.CPUCount {
|
||||
if n.CPUCount > 0 {
|
||||
s.log.Info("node CPUs changed", "node", nodeID, "from", n.CPUCount, "to", cpus)
|
||||
|
||||
@@ -323,9 +323,15 @@ type Model struct {
|
||||
// choice by guesswork.
|
||||
Description string `json:"description,omitempty" koanf:"description"`
|
||||
// Source is the weights path or a HuggingFace id.
|
||||
Source string `json:"source,omitempty" koanf:"source"`
|
||||
Profiles []Profile `json:"profiles,omitempty" koanf:"profiles"`
|
||||
Scaling Scaling `json:"scaling,omitempty" koanf:"scaling"`
|
||||
Source string `json:"source,omitempty" koanf:"source"`
|
||||
// AllowPickle lets the nodes fetch this model's weights, or its
|
||||
// adapters', in a pickle format (.bin, .pt, .pth, .ckpt), which run code
|
||||
// as they are loaded. Off, such a source is refused when the model is
|
||||
// added and by the node before it fetches anything. For a model that has
|
||||
// no safetensors or GGUF form and whose files were checked.
|
||||
AllowPickle bool `json:"allow_pickle,omitempty" koanf:"allow_pickle"`
|
||||
Profiles []Profile `json:"profiles,omitempty" koanf:"profiles"`
|
||||
Scaling Scaling `json:"scaling,omitempty" koanf:"scaling"`
|
||||
// Class is how much this model's capacity is worth against another's,
|
||||
// when there is not enough for both: critical, standard or opportunistic.
|
||||
// Empty is standard, or opportunistic for a model loaded on demand. See
|
||||
@@ -1185,6 +1191,22 @@ func (m Model) validate() error {
|
||||
}
|
||||
continue
|
||||
}
|
||||
// Strata prepares one of the few families it runs from their GGUF
|
||||
// files, on one card, and has to be told which.
|
||||
if p.Engine == "strata" {
|
||||
if p.Nodes > 1 || p.TP > 1 || p.PP > 1 || p.GPUsPerNode > 1 {
|
||||
return fmt.Errorf("model %q profile %q: strata runs a replica on one GPU of one node", m.Name, p.Name)
|
||||
}
|
||||
if !slices.ContainsFunc(p.Args, func(a string) bool { return a == "--family" || strings.HasPrefix(a, "--family=") }) {
|
||||
return fmt.Errorf("model %q profile %q: strata needs the model's family and size in args "+
|
||||
"(--family qwen --size IQ2_XS)", m.Name, p.Name)
|
||||
}
|
||||
if strings.HasSuffix(strings.ToLower(strings.TrimSuffix(m.SourceOf(p), "/")), ".gguf") {
|
||||
return fmt.Errorf("model %q profile %q: strata prepares a model from all of its GGUF files; "+
|
||||
"give the directory they are in as the source", m.Name, p.Name)
|
||||
}
|
||||
continue
|
||||
}
|
||||
// TensorRT-LLM spans machines only under MPI started from one of them,
|
||||
// which the agents, each starting its own node, do not do; and it
|
||||
// loads a Hugging Face checkpoint, not llama.cpp's format.
|
||||
@@ -1336,10 +1358,16 @@ func (m Model) EngineProfile(name string) (engines.Profile, bool) {
|
||||
Name string `json:"name"`
|
||||
Source string `json:"source"`
|
||||
HFRepo string `json:"hf_repo,omitempty"`
|
||||
// A string: the node reads the members as strings.
|
||||
AllowPickle string `json:"allow_pickle,omitempty"`
|
||||
}
|
||||
var ms []member
|
||||
for _, sm := range m.Shared.Members {
|
||||
ms = append(ms, member{Name: sm.Name, Source: sm.Source, HFRepo: sm.HFRepo})
|
||||
mb := member{Name: sm.Name, Source: sm.Source, HFRepo: sm.HFRepo}
|
||||
if sm.AllowPickle {
|
||||
mb.AllowPickle = "true"
|
||||
}
|
||||
ms = append(ms, mb)
|
||||
}
|
||||
b, _ := json.Marshal(ms)
|
||||
env = maps.Clone(env)
|
||||
@@ -1369,6 +1397,7 @@ func (m Model) EngineProfile(name string) (engines.Profile, bool) {
|
||||
ServedModelName: m.Name,
|
||||
Adapters: m.engineAdapters(),
|
||||
KVEvents: p.KVEvents,
|
||||
AllowPickle: m.AllowPickle,
|
||||
}.WithRuntimeLoRA(m.MaxLoadedAdapters), true
|
||||
}
|
||||
return engines.Profile{}, false
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
)
|
||||
|
||||
// PickleSources are the sources of a model's weights and adapters that are
|
||||
// in a pickle format, which run code as they are loaded (weights.Pickled).
|
||||
func (m Model) PickleSources() []string {
|
||||
var out []string
|
||||
add := func(s string) {
|
||||
if s != "" && weights.Pickled(s) && !slices.Contains(out, s) {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
add(m.Source)
|
||||
for _, p := range m.Profiles {
|
||||
add(p.Source)
|
||||
}
|
||||
for _, a := range m.Adapters {
|
||||
add(a.Source)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// CheckPickle refuses a model whose weights are in a pickle format, unless
|
||||
// it says allow_pickle. Asked when a model is added or its sources change,
|
||||
// not when the catalog is built: a model already running is refused by its
|
||||
// node, with the same words, rather than dropped from the catalog.
|
||||
// Sources in except were there before and are not asked about again.
|
||||
func (m Model) CheckPickle(except ...string) error {
|
||||
if m.AllowPickle {
|
||||
return nil
|
||||
}
|
||||
for _, s := range m.PickleSources() {
|
||||
if slices.Contains(except, s) {
|
||||
continue
|
||||
}
|
||||
return fmt.Errorf("model %q: %s is in a pickle format, which runs code when it is loaded; "+
|
||||
"use the model's safetensors or GGUF, or set allow_pickle: true if it was checked and has no other form",
|
||||
m.Name, s)
|
||||
}
|
||||
// What an engine downloads itself, by a Hugging Face id, no node looks
|
||||
// at: vLLM reads a model by id as safetensors only, but not an adapter,
|
||||
// and TensorRT-LLM neither.
|
||||
for _, s := range m.UnlookedSources() {
|
||||
if slices.Contains(except, s) {
|
||||
continue
|
||||
}
|
||||
return fmt.Errorf("model %q: %s is a Hugging Face id its engine downloads itself, and may load in a pickle "+
|
||||
"format, which runs code; mirror its safetensors and give that, or set allow_pickle: true if it was checked",
|
||||
m.Name, s)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UnlookedSources are those an engine fetches by a Hugging Face id, past
|
||||
// the node's check: adapters, and a TensorRT-LLM model's.
|
||||
func (m Model) UnlookedSources() []string {
|
||||
var out []string
|
||||
for _, a := range m.Adapters {
|
||||
if hubID(a.Source) {
|
||||
out = append(out, a.Source)
|
||||
}
|
||||
}
|
||||
for _, p := range m.Profiles {
|
||||
if p.Engine == "trtllm" {
|
||||
if src := m.SourceOf(p); hubID(src) && !slices.Contains(out, src) {
|
||||
out = append(out, src)
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// hubID reports whether a source is a Hugging Face id, as the agent reads
|
||||
// one (engines.IsHubID) — gpt2 too — and not a file named so, nor a GGUF
|
||||
// in a repository (org/name:Q4_K_M).
|
||||
func hubID(s string) bool {
|
||||
if !engines.IsHubID(s) || strings.Contains(s, ":") {
|
||||
return false
|
||||
}
|
||||
ext := strings.ToLower(path.Ext(s))
|
||||
return ext != ".gguf" && ext != ".safetensors" && !weights.Pickled(s)
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAModelInAPickleFormatIsRefusedUnlessItSaysSo(t *testing.T) {
|
||||
m := Model{Name: "triposr", Source: "hf:stabilityai/TripoSR/model.ckpt",
|
||||
Adapters: []Adapter{{Name: "a", Source: "http://m/a/adapter_model.bin"}}}
|
||||
err := m.CheckPickle()
|
||||
if err == nil || !strings.Contains(err.Error(), "model.ckpt") || !strings.Contains(err.Error(), "allow_pickle") {
|
||||
t.Fatalf("not refused, or not said why: %v", err)
|
||||
}
|
||||
if got := m.PickleSources(); len(got) != 2 {
|
||||
t.Errorf("pickle sources %v, want the model's and the adapter's", got)
|
||||
}
|
||||
// A source that was there before is not asked about again: changing
|
||||
// the scaling of a model that runs does not need allow_pickle.
|
||||
if err := m.CheckPickle(m.PickleSources()...); err != nil {
|
||||
t.Errorf("sources already there were refused: %v", err)
|
||||
}
|
||||
m.AllowPickle = true
|
||||
if err := m.CheckPickle(); err != nil {
|
||||
t.Errorf("allow_pickle was not honoured: %v", err)
|
||||
}
|
||||
if err := (Model{Name: "q", Source: "http://m/q/v/q-Q4_K_M.gguf"}).CheckPickle(); err != nil {
|
||||
t.Errorf("a GGUF was refused: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheTripoSRPresetSaysItsCheckpointIsAllowed(t *testing.T) {
|
||||
p, ok := PresetNamed("triposr")
|
||||
if !ok {
|
||||
t.Fatal("no triposr preset")
|
||||
}
|
||||
if err := p.Model.CheckPickle(); err != nil {
|
||||
t.Errorf("the preset llmgw ships is refused by its own check: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWhatAnEngineDownloadsByIDIsAskedAbout(t *testing.T) {
|
||||
m := Model{Name: "m", Source: "http://mirror/m/v/", Profiles: []Profile{{Name: "gpu", Engine: "vllm"}},
|
||||
Adapters: []Adapter{{Name: "a", Source: "someone/lora-x"}}}
|
||||
if err := m.CheckPickle(); err == nil || !strings.Contains(err.Error(), "someone/lora-x") {
|
||||
t.Errorf("an adapter by its id: %v", err)
|
||||
}
|
||||
m.Adapters = []Adapter{{Name: "a", Source: "http://mirror/a/v/adapter_model.safetensors"}}
|
||||
if err := m.CheckPickle(); err != nil {
|
||||
t.Errorf("a mirrored adapter: %v", err)
|
||||
}
|
||||
trt := Model{Name: "t", Source: "meta-llama/Llama-3.1-8B-Instruct", Profiles: []Profile{{Name: "gpu", Engine: "trtllm"}}}
|
||||
if err := trt.CheckPickle(); err == nil {
|
||||
t.Error("TensorRT-LLM downloading by id was not asked about")
|
||||
}
|
||||
vllm := Model{Name: "v", Source: "meta-llama/Llama-3.1-8B-Instruct", Profiles: []Profile{{Name: "gpu", Engine: "vllm"}}}
|
||||
if err := vllm.CheckPickle(); err != nil {
|
||||
t.Errorf("vLLM reads a model by id as safetensors only: %v", err)
|
||||
}
|
||||
for s, want := range map[string]bool{"Qwen/Qwen2.5-7B": true, "meta-llama/Llama-3.1-8B-Instruct": true,
|
||||
"org/model.gguf": false, "o/r:Q4_K_M": false, "/models/x": false, "hf:o/r/f.gguf": false,
|
||||
"gpt2": true, "~me/x": false} {
|
||||
if hubID(s) != want {
|
||||
t.Errorf("hubID(%q) = %v", s, !want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEveryPresetPassesItsOwnPickleCheck(t *testing.T) {
|
||||
// What llmgw ships must be addable: TripoSR was not, once.
|
||||
for _, p := range Presets() {
|
||||
if err := p.Model.CheckPickle(); err != nil {
|
||||
t.Errorf("%s: %v", p.Model.Name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -52,6 +52,9 @@ func Presets() []Preset {
|
||||
Name: "triposr",
|
||||
Description: "TripoSR: a 3D model (GLB) from a picture of one object on a plain background, as a job",
|
||||
Source: "hf:stabilityai/TripoSR/model.ckpt",
|
||||
// TripoSR is published as this checkpoint only; it is
|
||||
// Stability AI's, pinned, and the one pickle llmgw ships.
|
||||
AllowPickle: true,
|
||||
ComfyUI: &ComfyUI{
|
||||
Workflow: "triposr",
|
||||
Output: OutputMeshes,
|
||||
|
||||
@@ -34,6 +34,8 @@ type SharedMember struct {
|
||||
Source string `json:"source"`
|
||||
// HFRepo is the repository its node loads the weights by name from.
|
||||
HFRepo string `json:"hf_repo,omitempty"`
|
||||
// AllowPickle is the member model's allow_pickle.
|
||||
AllowPickle bool `json:"allow_pickle,omitempty"`
|
||||
// Workflow and Checkpoint are what a request for it is sent.
|
||||
Workflow string `json:"workflow"`
|
||||
Checkpoint string `json:"checkpoint"`
|
||||
@@ -80,7 +82,7 @@ func envModel(env Environment, members []Model) Model {
|
||||
hf = m.ComfyUI.HuggingFaceRepo
|
||||
}
|
||||
shared.Members = append(shared.Members, SharedMember{Name: m.Name, Aliases: m.Aliases,
|
||||
Source: m.Source, HFRepo: hf, Workflow: workflow, Checkpoint: ckpt})
|
||||
Source: m.Source, HFRepo: hf, AllowPickle: m.AllowPickle, Workflow: workflow, Checkpoint: ckpt})
|
||||
idle = max(idle, m.Scaling.IdleTimeout)
|
||||
minN, maxN = max(minN, m.Scaling.Min), max(maxN, m.Scaling.Max)
|
||||
if ClassRank(m.EffectiveClass()) > ClassRank(class) || class == "" {
|
||||
@@ -97,6 +99,7 @@ func envModel(env Environment, members []Model) Model {
|
||||
Role: RoleEnvironment,
|
||||
Description: "the shared ComfyUI replica of environment " + env.Name,
|
||||
Source: first.Source,
|
||||
AllowPickle: first.AllowPickle,
|
||||
ComfyUI: &ComfyUI{Environment: env.Name, Nodes: env.Nodes, Output: OutputEnvironment},
|
||||
Profiles: profs,
|
||||
Scaling: Scaling{Min: minN, Max: maxN, IdleTimeout: idle},
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAStrataProfileNamesItsFamilyAndRunsOnOneCard(t *testing.T) {
|
||||
ok := Model{Name: "q38", Source: "http://mirror:8081/q38/ab12cd34/", Profiles: []Profile{{
|
||||
Name: "p", Engine: "strata", GPUClass: "6000-24g",
|
||||
Args: []string{"--family", "qwen", "--size", "IQ2_XS", "--context", "32768"},
|
||||
}}}
|
||||
if _, err := New([]Model{ok}); err != nil {
|
||||
t.Fatalf("a one-card Strata profile was refused: %v", err)
|
||||
}
|
||||
for want, bad := range map[string]Model{
|
||||
"one GPU of one node": {Name: "x", Source: "http://mirror:8081/x/1/", Profiles: []Profile{{
|
||||
Name: "p", Engine: "strata", GPUClass: "6000-24g", GPUsPerNode: 2, TP: 2,
|
||||
Args: []string{"--family", "qwen"}}}},
|
||||
"family and size": {Name: "x", Source: "http://mirror:8081/x/1/", Profiles: []Profile{{
|
||||
Name: "p", Engine: "strata", GPUClass: "6000-24g", Args: []string{"--size", "IQ2_XS"}}}},
|
||||
"all of its GGUF files": {Name: "x", Source: "http://mirror:8081/x/1/x-00001-of-00002.gguf", Profiles: []Profile{{
|
||||
Name: "p", Engine: "strata", GPUClass: "6000-24g", Args: []string{"--family", "qwen"}}}},
|
||||
"--gguf-dir": {Name: "x", Source: "http://mirror:8081/x/1/", Profiles: []Profile{{
|
||||
Name: "p", Engine: "strata", GPUClass: "6000-24g", Args: []string{"--family", "qwen", "--gguf-dir", "/etc"}}}},
|
||||
} {
|
||||
if _, err := New([]Model{bad}); err == nil || !strings.Contains(err.Error(), want) {
|
||||
t.Errorf("expected %q, got %v", want, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,233 @@
|
||||
package modelserve
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
|
||||
)
|
||||
|
||||
// A model's context changed after it was served: what one request may use,
|
||||
// in the three places it lives — the catalog, the engine's arguments, and
|
||||
// the memory its KV cache takes — together, checked against what the cards
|
||||
// hold, then its replicas started again on it one by one.
|
||||
|
||||
// ContextRequest is the context to change to: tokens one request may use,
|
||||
// and for llama.cpp how many at once (its pool is the two multiplied).
|
||||
type ContextRequest struct {
|
||||
Tokens int `json:"tokens"`
|
||||
Slots int `json:"slots,omitempty"`
|
||||
}
|
||||
|
||||
// ContextProfile is what one profile needs before and after.
|
||||
type ContextProfile struct {
|
||||
Profile string `json:"profile"`
|
||||
Engine string `json:"engine"`
|
||||
Before string `json:"before,omitempty"`
|
||||
After string `json:"after,omitempty"`
|
||||
Args []string `json:"args"`
|
||||
}
|
||||
|
||||
// ContextChange is what was changed, and the replicas started again on it.
|
||||
type ContextChange struct {
|
||||
Model string `json:"model"`
|
||||
Before int `json:"before"`
|
||||
After int `json:"after"`
|
||||
MaxOutput int `json:"max_output"`
|
||||
Trained int `json:"trained,omitempty"`
|
||||
Profiles []ContextProfile `json:"profiles"`
|
||||
Restarted int `json:"restarted"`
|
||||
// Note is what could not be checked, said once.
|
||||
Note string `json:"note,omitempty"`
|
||||
}
|
||||
|
||||
// shape is what sizes a context: a token's KV cache, the context the model
|
||||
// was trained for, and its weights; and its KV heads, which decide how many
|
||||
// cards tensor parallel can spread it over (0 unknown).
|
||||
type shape struct {
|
||||
kvPerToken int64
|
||||
trained int
|
||||
weights int64
|
||||
kvHeads int
|
||||
}
|
||||
|
||||
// SetContext changes a served model's context.
|
||||
func (s *Server) SetContext(ctx context.Context, name string, req ContextRequest, by string) (ContextChange, error) {
|
||||
if s.deps.Current == nil || s.deps.Models == nil {
|
||||
return ContextChange{}, errors.New("this control plane cannot change the catalog")
|
||||
}
|
||||
var m catalog.Model
|
||||
found := false
|
||||
for _, c := range s.deps.Current() {
|
||||
if c.Name == name {
|
||||
m, found = c, true
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case !found:
|
||||
return ContextChange{}, fmt.Errorf("no model %q in the catalog", name)
|
||||
case s.deps.Models.Declared(name):
|
||||
return ContextChange{}, fmt.Errorf("model %q is declared in the cluster's settings: change its context there (llmgw ctl config edit)", name)
|
||||
case req.Tokens < 256:
|
||||
return ContextChange{}, errors.New("a context of at least 256 tokens")
|
||||
case req.Slots < 0 || req.Slots > 64:
|
||||
return ContextChange{}, errors.New("slots are 1 to 64")
|
||||
}
|
||||
sh, note := s.shapeOf(ctx, m)
|
||||
if sh.trained > 0 && req.Tokens > sh.trained {
|
||||
return ContextChange{}, fmt.Errorf("%s was trained for %d tokens; a longer context answers worse, and llama.cpp or vLLM refuses it",
|
||||
name, sh.trained)
|
||||
}
|
||||
out := ContextChange{Model: name, Before: m.ContextLength, After: req.Tokens, Trained: sh.trained, Note: note}
|
||||
profiles := slices.Clone(m.Profiles)
|
||||
for i, p := range profiles {
|
||||
next, change, err := resizeProfile(p, req, sh)
|
||||
if err != nil {
|
||||
return ContextChange{}, err
|
||||
}
|
||||
profiles[i] = next
|
||||
out.Profiles = append(out.Profiles, change)
|
||||
}
|
||||
out.MaxOutput = m.MaxOutput
|
||||
if m.MaxOutput == 0 || m.MaxOutput == maxOutput(m.ContextLength) {
|
||||
out.MaxOutput = maxOutput(req.Tokens)
|
||||
}
|
||||
out.MaxOutput = min(out.MaxOutput, req.Tokens)
|
||||
if _, err := s.deps.Models.SetContext(ctx, name, req.Tokens, out.MaxOutput, profiles, by); err != nil {
|
||||
return ContextChange{}, err
|
||||
}
|
||||
// Started again on it, a new replica before the old one stops.
|
||||
var ids []string
|
||||
if s.deps.Fleet != nil {
|
||||
for _, r := range s.deps.Fleet.Replicas() {
|
||||
if r.Model == name {
|
||||
ids = append(ids, r.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(ids) > 0 && s.deps.Restart != nil {
|
||||
if err := s.deps.Restart(ids...); err != nil {
|
||||
return out, fmt.Errorf("the context is changed, and its replicas could not be started again on it: %w", err)
|
||||
}
|
||||
out.Restarted = len(ids)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// shapeOf reads what sizes a model's context: a GGUF file's header from the
|
||||
// mirror, else the original weights' config from their repository.
|
||||
func (s *Server) shapeOf(ctx context.Context, m catalog.Model) (shape, string) {
|
||||
var sh shape
|
||||
if m.Origin != nil {
|
||||
for _, f := range m.Origin.Files {
|
||||
sh.weights += f.Bytes
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case strings.HasSuffix(strings.ToLower(m.Source), ".gguf"):
|
||||
h, err := gguf.FromURL(ctx, nil, m.Source)
|
||||
if err != nil {
|
||||
return sh, "its header could not be read (" + err.Error() + "), so the memory it needs is not checked"
|
||||
}
|
||||
sh.kvPerToken, sh.trained = h.KVBytesPerToken(), int(h.ContextLength) //nolint:gosec // a header's context
|
||||
sh.kvHeads = int(h.HeadCountKV) //nolint:gosec // a header's head count
|
||||
case m.Origin != nil && m.Origin.Repo != "":
|
||||
mm, err := modelplan.FetchModel(ctx, m.Origin.Repo)
|
||||
if err != nil {
|
||||
return sh, "its config could not be read (" + err.Error() + "), so the memory it needs is not checked"
|
||||
}
|
||||
sh.kvPerToken, sh.trained = mm.KVBytesPerToken(modelplan.BF16), mm.MaxContext
|
||||
sh.kvHeads = mm.KVHeads
|
||||
default:
|
||||
return sh, "where it came from is not known, so the memory it needs is not checked"
|
||||
}
|
||||
if sh.weights == 0 {
|
||||
return sh, "its files' sizes are not known, so the memory it needs is not checked"
|
||||
}
|
||||
return sh, ""
|
||||
}
|
||||
|
||||
// resizeProfile is a profile with the context asked for: its engine's
|
||||
// arguments, and its memory where it is on CPUs, refused where its cards
|
||||
// would not hold it.
|
||||
func resizeProfile(p catalog.Profile, req ContextRequest, sh shape) (catalog.Profile, ContextProfile, error) {
|
||||
change := ContextProfile{Profile: p.Name, Engine: p.Engine}
|
||||
var pool int
|
||||
switch p.Engine {
|
||||
case "llamacpp":
|
||||
slots := req.Slots
|
||||
if slots == 0 {
|
||||
slots = max(argInt(p.Args, "--parallel"), 1)
|
||||
}
|
||||
pool = req.Tokens * slots
|
||||
p.Args = setArg(setArg(p.Args, "--ctx-size", strconv.Itoa(pool)), "--parallel", strconv.Itoa(slots))
|
||||
case "vllm":
|
||||
pool = req.Tokens
|
||||
p.Args = setArg(p.Args, "--max-model-len", strconv.Itoa(pool))
|
||||
case "sglang":
|
||||
pool = req.Tokens
|
||||
p.Args = setArg(p.Args, "--context-length", strconv.Itoa(pool))
|
||||
default:
|
||||
return p, change, fmt.Errorf("a %s profile (%s) does not take a context this way", p.Engine, p.Name)
|
||||
}
|
||||
change.Args = p.Args
|
||||
if sh.kvPerToken == 0 || sh.weights == 0 {
|
||||
return p, change, nil
|
||||
}
|
||||
if p.CPUOnly() {
|
||||
need := sh.weights + int64(pool)*sh.kvPerToken + int64(1<<30) + sh.weights/20
|
||||
gb := int(math.Ceil(modelplan.GB(need)))
|
||||
change.Before, change.After = fmt.Sprintf("%d GB of RAM", p.MinMemoryGB), fmt.Sprintf("%d GB of RAM", gb)
|
||||
p.MinMemoryGB = gb
|
||||
return p, change, nil
|
||||
}
|
||||
cards := max(p.GPUsPerNode, p.TP, 1) * max(p.Nodes, p.PP, 1)
|
||||
need := modelplan.CardNeed(p.Engine, sh.weights, int64(pool)*sh.kvPerToken, cards)
|
||||
change.After = fmt.Sprintf("%.1f GB a card of %d", need, cards)
|
||||
if budget := modelplan.CardBudget(p.GPUClass); budget > 0 {
|
||||
change.After += fmt.Sprintf(", of %.0f usable", budget)
|
||||
if need > budget {
|
||||
return p, change, fmt.Errorf("with %d tokens the %s profile needs %.1f GB on each of its %d %s cards, more than the %.0f GB they hold: "+
|
||||
"a shorter context, or the model on more cards", req.Tokens, p.Name, need, cards, p.GPUClass, budget)
|
||||
}
|
||||
}
|
||||
return p, change, nil
|
||||
}
|
||||
|
||||
// setArg sets a flag's value in arguments, or adds it.
|
||||
func setArg(args []string, flag, value string) []string {
|
||||
out := slices.Clone(args)
|
||||
for i := 0; i < len(out); i++ {
|
||||
if out[i] == flag && i+1 < len(out) {
|
||||
out[i+1] = value
|
||||
return out
|
||||
}
|
||||
if v, ok := strings.CutPrefix(out[i], flag+"="); ok && v != "" {
|
||||
out[i] = flag + "=" + value
|
||||
return out
|
||||
}
|
||||
}
|
||||
return append(out, flag, value)
|
||||
}
|
||||
|
||||
// argInt is a flag's whole-number value, 0 when absent.
|
||||
func argInt(args []string, flag string) int {
|
||||
for i, a := range args {
|
||||
if a == flag && i+1 < len(args) {
|
||||
n, _ := strconv.Atoi(args[i+1])
|
||||
return n
|
||||
}
|
||||
if v, ok := strings.CutPrefix(a, flag+"="); ok {
|
||||
n, _ := strconv.Atoi(v)
|
||||
return n
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package modelserve
|
||||
|
||||
import (
|
||||
"context"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
)
|
||||
|
||||
func TestAContextChangeResizesTheEngineAndTheMemoryAndRestartsTheReplicas(t *testing.T) {
|
||||
s, _, models, _, fleet := updSetup(t)
|
||||
m := models.models["qwen"]
|
||||
m.ContextLength, m.MaxOutput = 4096, maxOutput(4096)
|
||||
m.Origin = &catalog.Origin{Repo: "o/r", Files: []catalog.OriginFile{{Path: "q4.gguf", Bytes: 2 << 30}}}
|
||||
m.Profiles[0].Args = []string{"--ctx-size", "8192", "--parallel", "2"}
|
||||
models.models["qwen"] = m
|
||||
|
||||
got, err := s.SetContext(context.Background(), "qwen", ContextRequest{Tokens: 32768}, "admin")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.Before != 4096 || got.After != 32768 || got.Restarted != 2 || got.Note != "" {
|
||||
t.Fatalf("%+v", got)
|
||||
}
|
||||
after := models.models["qwen"]
|
||||
if after.ContextLength != 32768 || after.MaxOutput != maxOutput(32768) {
|
||||
t.Errorf("the catalog says %d tokens, %d out", after.ContextLength, after.MaxOutput)
|
||||
}
|
||||
// Two slots kept: the pool is both.
|
||||
if a := after.Profiles[0].Args; !slices.Equal(a, []string{"--ctx-size", "65536", "--parallel", "2"}) {
|
||||
t.Errorf("args %v", a)
|
||||
}
|
||||
// 2 GiB of weights, 64k tokens at 28 KiB each, an engine's share.
|
||||
if gb := after.Profiles[0].MinMemoryGB; gb < 4 || gb > 6 {
|
||||
t.Errorf("memory %d GB", gb)
|
||||
}
|
||||
for _, r := range fleet.Replicas() {
|
||||
if !strings.Contains(r.ID, "-new") {
|
||||
t.Errorf("%s was not started again", r.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAContextTheCardsCannotHoldIsRefusedWithTheFigures(t *testing.T) {
|
||||
sh := shape{kvPerToken: 128 << 10, weights: 8 << 30}
|
||||
p := catalog.Profile{Name: "gpu", Engine: "vllm", GPUClass: "t4-16g", GPUsPerNode: 1}
|
||||
if _, _, err := resizeProfile(p, ContextRequest{Tokens: 131072}, sh); err == nil || !strings.Contains(err.Error(), "GB") {
|
||||
t.Fatalf("not refused: %v", err)
|
||||
}
|
||||
next, change, err := resizeProfile(p, ContextRequest{Tokens: 4096}, sh)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !slices.Equal(next.Args, []string{"--max-model-len", "4096"}) || !strings.Contains(change.After, "a card of 1") {
|
||||
t.Errorf("%v %+v", next.Args, change)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAContextIsRefusedForADeclaredModelOrPastWhatItWasTrainedFor(t *testing.T) {
|
||||
s, _, models, _, _ := updSetup(t)
|
||||
models.declared["qwen"] = true
|
||||
if _, err := s.SetContext(context.Background(), "qwen", ContextRequest{Tokens: 8192}, "admin"); err == nil {
|
||||
t.Error("a declared model's context was changed")
|
||||
}
|
||||
models.declared["qwen"] = false
|
||||
if _, err := s.SetContext(context.Background(), "qwen", ContextRequest{Tokens: 100}, "admin"); err == nil {
|
||||
t.Error("100 tokens taken")
|
||||
}
|
||||
if _, err := s.SetContext(context.Background(), "nope", ContextRequest{Tokens: 8192}, "admin"); err == nil {
|
||||
t.Error("no such model, and no error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineArgumentsAreSetInEitherSpelling(t *testing.T) {
|
||||
if a := setArg([]string{"--ctx-size=4096", "-ngl", "99"}, "--ctx-size", "8192"); !slices.Equal(a, []string{"--ctx-size=8192", "-ngl", "99"}) {
|
||||
t.Error(a)
|
||||
}
|
||||
if a := setArg([]string{"-ngl", "99"}, "--parallel", "4"); !slices.Equal(a, []string{"-ngl", "99", "--parallel", "4"}) {
|
||||
t.Error(a)
|
||||
}
|
||||
if argInt([]string{"--parallel=3"}, "--parallel") != 3 || argInt(nil, "--parallel") != 0 {
|
||||
t.Error("argInt")
|
||||
}
|
||||
p := catalog.Profile{Name: "x", Engine: "comfyui"}
|
||||
if _, _, err := resizeProfile(p, ContextRequest{Tokens: 4096}, shape{}); err == nil {
|
||||
t.Error("a comfyui profile took a context")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
package modelserve
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
|
||||
)
|
||||
|
||||
// A served model moved to more cards, or fewer: how many each replica
|
||||
// spreads it over, on one node. The profile changes — tensor parallel for
|
||||
// vLLM, SGLang and TensorRT-LLM, llama.cpp's layer split for llama.cpp —
|
||||
// checked against the model's KV heads, what each card would hold, and the
|
||||
// nodes there are; then its replicas start again on it.
|
||||
|
||||
// GPUsRequest is how many cards a replica is to use.
|
||||
type GPUsRequest struct {
|
||||
GPUs int `json:"gpus"`
|
||||
// Profile is the one to change, when the model has more than one
|
||||
// profile on cards.
|
||||
Profile string `json:"profile,omitempty"`
|
||||
// InPlace stops each replica before its replacement starts: for when no
|
||||
// node has the cards free beside it. It does not answer while the new
|
||||
// one loads.
|
||||
InPlace bool `json:"in_place,omitempty"`
|
||||
}
|
||||
|
||||
// GPUsChange is what was changed, and the replicas started again on it.
|
||||
type GPUsChange struct {
|
||||
Model string `json:"model"`
|
||||
Profile string `json:"profile"`
|
||||
Engine string `json:"engine"`
|
||||
Class string `json:"gpu_class"`
|
||||
Before int `json:"before"`
|
||||
After int `json:"after"`
|
||||
Args []string `json:"args"`
|
||||
// PerCard is what each card holds after, when it could be worked out.
|
||||
PerCard string `json:"per_card,omitempty"`
|
||||
Restarted int `json:"restarted"`
|
||||
InPlace bool `json:"in_place,omitempty"`
|
||||
// Note is what could not be checked, said once.
|
||||
Note string `json:"note,omitempty"`
|
||||
}
|
||||
|
||||
// tpFlags are each engine's own spelling of a tensor-parallel width, which
|
||||
// the agent sets from the profile: a profile's own would contradict it.
|
||||
var tpFlags = map[string][]string{
|
||||
"vllm": {"--tensor-parallel-size", "-tp"},
|
||||
"sglang": {"--tp-size", "--tensor-parallel-size", "--tp"},
|
||||
"trtllm": {"--tp_size"},
|
||||
// llama.cpp's proportions were for the cards it had.
|
||||
"llamacpp": {"--tensor-split", "-ts"},
|
||||
}
|
||||
|
||||
// SetGPUs changes how many cards a served model's replicas use.
|
||||
func (s *Server) SetGPUs(ctx context.Context, name string, req GPUsRequest, by string) (GPUsChange, error) {
|
||||
if s.deps.Current == nil || s.deps.Models == nil {
|
||||
return GPUsChange{}, errors.New("this control plane cannot change the catalog")
|
||||
}
|
||||
var m catalog.Model
|
||||
found := false
|
||||
for _, c := range s.deps.Current() {
|
||||
if c.Name == name {
|
||||
m, found = c, true
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case !found:
|
||||
return GPUsChange{}, fmt.Errorf("no model %q in the catalog", name)
|
||||
case s.deps.Models.Declared(name):
|
||||
return GPUsChange{}, fmt.Errorf("model %q is declared in the cluster's settings: change its profile there (llmgw ctl config edit)", name)
|
||||
case req.GPUs < 1 || req.GPUs > 64:
|
||||
return GPUsChange{}, errors.New("a replica uses 1 to 64 cards")
|
||||
}
|
||||
i, err := gpuProfile(m, req.Profile)
|
||||
if err != nil {
|
||||
return GPUsChange{}, err
|
||||
}
|
||||
p := m.Profiles[i]
|
||||
before := max(p.GPUsPerNode, p.TP, 1) * max(p.Nodes, p.PP, 1)
|
||||
out := GPUsChange{Model: name, Profile: p.Name, Engine: p.Engine, Class: p.GPUClass, Before: before, After: req.GPUs}
|
||||
if before == req.GPUs {
|
||||
return out, fmt.Errorf("%s's %s profile already uses %d card(s)", name, p.Name, before)
|
||||
}
|
||||
|
||||
sh, note := s.shapeOf(ctx, m)
|
||||
next, err := respread(p, req.GPUs, sh)
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
out.Args = next.Args
|
||||
switch kv := int64(contextOf(m, p)) * sh.kvPerToken; {
|
||||
case sh.weights == 0 || sh.kvPerToken == 0 || contextOf(m, p) == 0:
|
||||
if note == "" {
|
||||
note = "its context is not known, so what each card holds is not checked"
|
||||
}
|
||||
default:
|
||||
need := modelplan.CardNeed(p.Engine, sh.weights, kv, req.GPUs)
|
||||
out.PerCard = fmt.Sprintf("%.1f GB a card of %d", need, req.GPUs)
|
||||
if budget := modelplan.CardBudget(p.GPUClass); budget > 0 {
|
||||
out.PerCard += fmt.Sprintf(", of %.0f usable", budget)
|
||||
if need > budget {
|
||||
return out, fmt.Errorf("on %d %s cards each would hold %.1f GB, more than the %.0f GB they hold: "+
|
||||
"more cards, or a shorter context (llmgw ctl model context)", req.GPUs, p.GPUClass, need, budget)
|
||||
}
|
||||
}
|
||||
}
|
||||
out.Note = note
|
||||
|
||||
ids := s.replicasOf(name)
|
||||
inPlace, err := s.roomFor(p.GPUClass, req.GPUs, name, len(ids) > 0, req.InPlace)
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
profiles := slices.Clone(m.Profiles)
|
||||
profiles[i] = next
|
||||
if _, err := s.deps.Models.SetProfiles(ctx, name, profiles, by); err != nil {
|
||||
return out, err
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
restart := s.deps.Restart
|
||||
if inPlace {
|
||||
restart, out.InPlace = s.deps.RestartInPlace, true
|
||||
}
|
||||
if restart == nil {
|
||||
return out, errors.New("the profile is changed, and this control plane cannot start its replicas again: llmgw ctl restart")
|
||||
}
|
||||
if err := restart(ids...); err != nil {
|
||||
return out, fmt.Errorf("the profile is changed, and its replicas could not be started again on it: %w", err)
|
||||
}
|
||||
out.Restarted = len(ids)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// gpuProfile is the model's profile on cards to change: the one named, or
|
||||
// its only one.
|
||||
func gpuProfile(m catalog.Model, named string) (int, error) {
|
||||
var onCards []int
|
||||
for i, p := range m.Profiles {
|
||||
if named != "" && p.Name == named {
|
||||
if p.CPUOnly() {
|
||||
return 0, fmt.Errorf("%s's %s profile runs on CPUs; give it a gpu_class first", m.Name, p.Name)
|
||||
}
|
||||
return i, nil
|
||||
}
|
||||
if !p.CPUOnly() {
|
||||
onCards = append(onCards, i)
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case named != "":
|
||||
return 0, fmt.Errorf("%s has no profile %q", m.Name, named)
|
||||
case len(onCards) == 0:
|
||||
return 0, fmt.Errorf("%s runs on CPUs only; a profile on cards needs a gpu_class", m.Name)
|
||||
case len(onCards) > 1:
|
||||
var names []string
|
||||
for _, i := range onCards {
|
||||
names = append(names, m.Profiles[i].Name)
|
||||
}
|
||||
return 0, fmt.Errorf("%s has several profiles on cards (%s): name one", m.Name, strings.Join(names, ", "))
|
||||
}
|
||||
return onCards[0], nil
|
||||
}
|
||||
|
||||
// respread is a profile spread over n cards of one node.
|
||||
func respread(p catalog.Profile, n int, sh shape) (catalog.Profile, error) {
|
||||
for _, f := range tpFlags[p.Engine] {
|
||||
p.Args = dropArg(p.Args, f)
|
||||
}
|
||||
switch p.Engine {
|
||||
case "vllm", "sglang", "trtllm":
|
||||
if !modelplan.SplitsHeads(sh.kvHeads, n) {
|
||||
return p, fmt.Errorf("the model has %d KV heads, which tensor parallel cannot spread over %d cards: "+
|
||||
"a number that divides them, or that they divide", sh.kvHeads, n)
|
||||
}
|
||||
p.TP, p.GPUsPerNode = n, n
|
||||
case "llamacpp":
|
||||
// llama.cpp has no tensor parallel: it splits the layers over a
|
||||
// node's cards, in proportion to their free memory.
|
||||
p.TP, p.GPUsPerNode = 1, n
|
||||
default:
|
||||
return p, fmt.Errorf("a %s replica runs on one card", p.Engine)
|
||||
}
|
||||
// One node: a replica spanning nodes is laid out by hand, with a fabric.
|
||||
p.Nodes, p.PP = 1, 0
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// contextOf is the tokens a profile's KV cache holds: llama.cpp's pool, or
|
||||
// the engine's limit, or the model's context.
|
||||
func contextOf(m catalog.Model, p catalog.Profile) int {
|
||||
switch p.Engine {
|
||||
case "llamacpp":
|
||||
if n := argInt(p.Args, "--ctx-size"); n > 0 {
|
||||
return n
|
||||
}
|
||||
case "vllm":
|
||||
if n := argInt(p.Args, "--max-model-len"); n > 0 {
|
||||
return n
|
||||
}
|
||||
case "sglang":
|
||||
if n := argInt(p.Args, "--context-length"); n > 0 {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return m.ContextLength
|
||||
}
|
||||
|
||||
func (s *Server) replicasOf(model string) []string {
|
||||
var ids []string
|
||||
if s.deps.Fleet != nil {
|
||||
for _, r := range s.deps.Fleet.Replicas() {
|
||||
if r.Model == model {
|
||||
ids = append(ids, r.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// roomFor says whether the new replicas fit: a node with n cards of the
|
||||
// class at all, and free beside the running ones — or, asked to start them
|
||||
// in place, once the model's own are stopped. True is in place.
|
||||
func (s *Server) roomFor(class string, n int, model string, running, inPlace bool) (bool, error) {
|
||||
if s.deps.Pool == nil {
|
||||
return inPlace, nil
|
||||
}
|
||||
own := map[string]int{}
|
||||
if s.deps.Fleet != nil {
|
||||
for _, r := range s.deps.Fleet.Replicas() {
|
||||
if r.Model == model {
|
||||
for _, sh := range r.Shards {
|
||||
own[sh.NodeID] += len(sh.GPUIndices)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
most, freeMost, besideOK, inPlaceOK := 0, 0, false, false
|
||||
for _, node := range s.deps.Pool() {
|
||||
for _, k := range modelplan.Kinds(node) {
|
||||
if k.Class != class {
|
||||
continue
|
||||
}
|
||||
most, freeMost = max(most, k.GPUs), max(freeMost, k.FreeGPUs)
|
||||
besideOK = besideOK || k.FreeGPUs >= n
|
||||
inPlaceOK = inPlaceOK || k.FreeGPUs+own[node.ID] >= n
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case most == 0:
|
||||
return false, fmt.Errorf("no ready node has %s cards", class)
|
||||
case most < n:
|
||||
return false, fmt.Errorf("no node has %d %s cards; the most one has is %d", n, class, most)
|
||||
case !running || besideOK:
|
||||
return false, nil
|
||||
case !inPlaceOK:
|
||||
return false, fmt.Errorf("no node has %d %s cards free, even with this model's own stopped (%d at most free now): "+
|
||||
"other models hold them", n, class, freeMost)
|
||||
case !inPlace:
|
||||
return false, fmt.Errorf("no node has %d %s cards free beside the running replica (%d at most); "+
|
||||
"--in-place stops each replica first, and the model does not answer while the new one loads", n, class, freeMost)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// dropArg is arguments without a flag and its value.
|
||||
func dropArg(args []string, flag string) []string {
|
||||
out := make([]string, 0, len(args))
|
||||
for i := 0; i < len(args); i++ {
|
||||
switch {
|
||||
case args[i] == flag:
|
||||
i++ // its value
|
||||
case strings.HasPrefix(args[i], flag+"="):
|
||||
default:
|
||||
out = append(out, args[i])
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package modelserve
|
||||
|
||||
import (
|
||||
"context"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
|
||||
)
|
||||
|
||||
// gpuSetup is the update setup with qwen on one L40S card of llama.cpp, and
|
||||
// a pool of one node of four, of which free are free.
|
||||
func gpuSetup(t *testing.T, free int) (*Server, *updModels, *restartFleet, *[]string) {
|
||||
t.Helper()
|
||||
s, _, models, _, fleet := updSetup(t)
|
||||
m := models.models["qwen"]
|
||||
m.ContextLength = 8192
|
||||
m.Origin = &catalog.Origin{Repo: "o/r", Files: []catalog.OriginFile{{Path: "q4.gguf", Bytes: 20 << 30}}}
|
||||
m.Profiles = []catalog.Profile{
|
||||
{Name: "cpu", Engine: "llamacpp", MinMemoryGB: 8},
|
||||
{Name: "gpu", Engine: "llamacpp", GPUClass: "l40s-48g", TP: 1, GPUsPerNode: 1,
|
||||
Args: []string{"--ctx-size", "8192", "--tensor-split", "1"}},
|
||||
}
|
||||
models.models["qwen"] = m
|
||||
s.deps.Pool = func() []modelplan.PoolNode {
|
||||
return []modelplan.PoolNode{{ID: "g1", Class: "l40s-48g", GPUMemGB: 48, GPUs: 4, FreeGPUs: free}}
|
||||
}
|
||||
var inPlace []string
|
||||
s.deps.RestartInPlace = func(ids ...string) error { inPlace = append(inPlace, ids...); return nil }
|
||||
return s, models, fleet, &inPlace
|
||||
}
|
||||
|
||||
func TestAModelIsSpreadOverMoreCardsAndItsReplicasStartedAgainBesideTheOld(t *testing.T) {
|
||||
s, models, fleet, inPlace := gpuSetup(t, 3)
|
||||
got, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2}, "admin")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.Profile != "gpu" || got.Before != 1 || got.After != 2 || got.Restarted != 2 || got.InPlace {
|
||||
t.Fatalf("%+v", got)
|
||||
}
|
||||
p := models.models["qwen"].Profiles[1]
|
||||
// llama.cpp splits layers over a node's cards: gpus_per_node, not tp,
|
||||
// and the old card's proportions are gone.
|
||||
if p.GPUsPerNode != 2 || p.TP != 1 || p.Nodes != 1 || !slices.Equal(p.Args, []string{"--ctx-size", "8192"}) {
|
||||
t.Errorf("profile %+v", p)
|
||||
}
|
||||
if !strings.Contains(got.PerCard, "a card of 2") {
|
||||
t.Errorf("per card %q", got.PerCard)
|
||||
}
|
||||
if len(*inPlace) != 0 {
|
||||
t.Errorf("stopped first though there was room: %v", *inPlace)
|
||||
}
|
||||
for _, r := range fleet.Replicas() {
|
||||
if !strings.Contains(r.ID, "-new") {
|
||||
t.Errorf("%s was not started again", r.ID)
|
||||
}
|
||||
}
|
||||
if models.models["qwen"].Profiles[0].GPUsPerNode != 0 {
|
||||
t.Error("the CPU profile was changed too")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithNoRoomBesideIt_InPlaceIsAskedFor(t *testing.T) {
|
||||
// One card free and the model's replicas on the others: two do not fit
|
||||
// beside them.
|
||||
s, _, fleet, inPlace := gpuSetup(t, 1)
|
||||
fleet.reps[0].Shards = []reconciler.Shard{{NodeID: "g1", GPUIndices: []int{1}}}
|
||||
fleet.reps[1].Shards = []reconciler.Shard{{NodeID: "g1", GPUIndices: []int{2}}}
|
||||
_, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2}, "admin")
|
||||
if err == nil || !strings.Contains(err.Error(), "--in-place") {
|
||||
t.Fatalf("not asked for in place: %v", err)
|
||||
}
|
||||
got, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2, InPlace: true}, "admin")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !got.InPlace || len(*inPlace) != 2 {
|
||||
t.Errorf("%+v, stopped first: %v", got, *inPlace)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMoreCardsThanANodeHasOrOthersHoldAreRefused(t *testing.T) {
|
||||
s, models, _, _ := gpuSetup(t, 4)
|
||||
if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 8}, "admin"); err == nil ||
|
||||
!strings.Contains(err.Error(), "the most one has is 4") {
|
||||
t.Errorf("eight cards on a node of four: %v", err)
|
||||
}
|
||||
s, _, _, _ = gpuSetup(t, 0)
|
||||
if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2, InPlace: true}, "admin"); err == nil ||
|
||||
!strings.Contains(err.Error(), "other models hold them") {
|
||||
t.Errorf("cards other models hold: %v", err)
|
||||
}
|
||||
if models.models["qwen"].Profiles[1].GPUsPerNode != 1 {
|
||||
t.Error("a refused change was kept")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFewerCardsThanTheModelNeedsAreRefusedWithTheFigures(t *testing.T) {
|
||||
s, models, _, _ := gpuSetup(t, 3)
|
||||
m := models.models["qwen"]
|
||||
m.Profiles[1].GPUsPerNode, m.Origin.Files[0].Bytes = 4, 120<<30
|
||||
models.models["qwen"] = m
|
||||
_, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2}, "admin")
|
||||
if err == nil || !strings.Contains(err.Error(), "GB") {
|
||||
t.Errorf("120 GB on two 48 GB cards: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTensorParallelFollowsTheKVHeads(t *testing.T) {
|
||||
p := catalog.Profile{Name: "gpu", Engine: "vllm", GPUClass: "h100-80g", TP: 2, GPUsPerNode: 2,
|
||||
Args: []string{"--tensor-parallel-size", "2", "--max-model-len", "32768"}}
|
||||
if _, err := respread(p, 3, shape{kvHeads: 8}); err == nil || !strings.Contains(err.Error(), "8 KV heads") {
|
||||
t.Errorf("three ways for eight heads: %v", err)
|
||||
}
|
||||
next, err := respread(p, 4, shape{kvHeads: 8})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The agent sets the width from tp; the profile's own spelling went.
|
||||
if next.TP != 4 || next.GPUsPerNode != 4 || !slices.Equal(next.Args, []string{"--max-model-len", "32768"}) {
|
||||
t.Errorf("%+v", next)
|
||||
}
|
||||
for _, eng := range []string{"comfyui", "strata"} {
|
||||
if _, err := respread(catalog.Profile{Name: "x", Engine: eng}, 2, shape{}); err == nil {
|
||||
t.Errorf("%s spread over two cards", eng)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAModelOnCPUsOrWithSeveralCardProfilesSaysWhy(t *testing.T) {
|
||||
s, models, _, _ := gpuSetup(t, 3)
|
||||
m := models.models["qwen"]
|
||||
m.Profiles = append(m.Profiles, catalog.Profile{Name: "gpu2", Engine: "vllm", GPUClass: "l40s-48g", TP: 1})
|
||||
models.models["qwen"] = m
|
||||
if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2}, "admin"); err == nil ||
|
||||
!strings.Contains(err.Error(), "name one") {
|
||||
t.Errorf("two profiles on cards and none named: %v", err)
|
||||
}
|
||||
if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 2, Profile: "cpu"}, "admin"); err == nil ||
|
||||
!strings.Contains(err.Error(), "gpu_class") {
|
||||
t.Errorf("a CPU profile: %v", err)
|
||||
}
|
||||
if _, err := s.SetGPUs(context.Background(), "qwen", GPUsRequest{GPUs: 1, Profile: "gpu"}, "admin"); err == nil ||
|
||||
!strings.Contains(err.Error(), "already") {
|
||||
t.Errorf("the same count: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -82,6 +82,9 @@ type Deps struct {
|
||||
Current func() []catalog.Model
|
||||
Models Models
|
||||
Restart func(ids ...string) error
|
||||
// RestartInPlace stops each replica before its replacement starts, for
|
||||
// when there is no room beside it.
|
||||
RestartInPlace func(ids ...string) error
|
||||
|
||||
// Pool is the fleet's ready nodes as a plan sees them: what decides
|
||||
// which file suits the hardware, and how it is spread.
|
||||
|
||||
@@ -34,6 +34,8 @@ type Models interface {
|
||||
Declared(name string) bool
|
||||
SetSource(ctx context.Context, name, source string, origin *catalog.Origin, by string) (catalog.Model, error)
|
||||
SetOrigin(ctx context.Context, name string, origin *catalog.Origin) error
|
||||
SetContext(ctx context.Context, name string, contextLength, maxOutput int, profiles []catalog.Profile, by string) (catalog.Model, error)
|
||||
SetProfiles(ctx context.Context, name string, profiles []catalog.Profile, by string) (catalog.Model, error)
|
||||
// For a canary: the candidate added beside the model and taken out
|
||||
// again, and the share of requests it is given.
|
||||
Add(ctx context.Context, m catalog.Model, by string) error
|
||||
|
||||
@@ -104,6 +104,23 @@ func (c *updModels) SetCanary(_ context.Context, n string, cn *catalog.Canary, _
|
||||
c.models[n] = m
|
||||
return m, nil
|
||||
}
|
||||
func (c *updModels) SetContext(_ context.Context, n string, ctxLen, maxOut int, ps []catalog.Profile, _ string) (catalog.Model, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
m := c.models[n]
|
||||
m.ContextLength, m.MaxOutput, m.Profiles = ctxLen, maxOut, ps
|
||||
c.models[n] = m
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (c *updModels) SetProfiles(_ context.Context, n string, ps []catalog.Profile, _ string) (catalog.Model, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
m := c.models[n]
|
||||
m.Profiles = ps
|
||||
c.models[n] = m
|
||||
return m, nil
|
||||
}
|
||||
func (c *updModels) all() []catalog.Model {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
@@ -154,6 +154,9 @@ func (s *Service) add(ctx context.Context, m catalog.Model, by string, onlyNew b
|
||||
next = append(next, e)
|
||||
}
|
||||
next = append(next, m)
|
||||
if err := m.CheckPickle(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
built, err := catalog.New(next)
|
||||
if err != nil {
|
||||
@@ -266,6 +269,17 @@ func (s *Service) Resources(ctx context.Context, name, profile string, memoryGB,
|
||||
// SetProfiles replaces a model's profiles, as a model added with them
|
||||
// would have: for a model whose engines are to run otherwise, its args among
|
||||
// it. Running replicas keep what they were started with until restarted.
|
||||
// SetContext changes how much one request may use and what it may answer,
|
||||
// with the profiles that carry it: their engines' arguments and memory.
|
||||
func (s *Service) SetContext(ctx context.Context, name string, contextLength, maxOutput int,
|
||||
profiles []catalog.Profile, by string,
|
||||
) (catalog.Model, error) {
|
||||
return s.update(ctx, name, by, "change its context", func(m *catalog.Model) error {
|
||||
m.ContextLength, m.MaxOutput, m.Profiles = contextLength, maxOutput, profiles
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) SetProfiles(ctx context.Context, name string, profiles []catalog.Profile, by string) (catalog.Model, error) {
|
||||
if len(profiles) == 0 {
|
||||
return catalog.Model{}, errors.New("a model needs a profile")
|
||||
@@ -406,9 +420,13 @@ func (s *Service) update(ctx context.Context, name, by, what string, change func
|
||||
found := false
|
||||
for _, e := range stored {
|
||||
if e.Name == name {
|
||||
before := append(e.PickleSources(), e.UnlookedSources()...)
|
||||
if err := change(&e); err != nil {
|
||||
return catalog.Model{}, err
|
||||
}
|
||||
if err := e.CheckPickle(before...); err != nil {
|
||||
return catalog.Model{}, err
|
||||
}
|
||||
found = true
|
||||
}
|
||||
next = append(next, e)
|
||||
|
||||
@@ -230,6 +230,7 @@ var Reference = []Block{
|
||||
}},
|
||||
{Package: "config", Name: "Mirror", Title: "Mirror", Doc: "", Fields: []Field{
|
||||
{Key: "api", Type: "string", Doc: "API is the mirror's management address, scheme included, for example\nhttps://models.internal:8082. Empty means no mirror is configured.\n\nThe management port, not the one nodes fetch bytes from. Those are two\nlisteners on the mirror precisely so this one can be locked down without\nre-certifying every agent's fetcher.", Nested: ""},
|
||||
{Key: "peers", Type: "[]string", Doc: "Peers are the management addresses of mirrors kept in step with this\none (llmgw mirror serve --peer): asked when it cannot be reached, and\na node given the same file on one of them while it is down.", Nested: ""},
|
||||
{Key: "tls", Type: "tlsconf.Config", Doc: "TLS authenticates to the mirror's management API, with the same CA the\nagents use. Empty leaves the connection in plain HTTP, which is only\nacceptable on a network where nothing untrusted can reach that port —\nthe mirror says as much in its own log when it starts that way.", Nested: "tlsconf.Config"},
|
||||
{Key: "registry", Type: "string", Doc: "Registry is the model registry the mirror page browses, empty for the\npublic Hugging Face. Set it to an internal mirror of the registry where\nthe control plane cannot reach the internet — or leave the whole mirror\nblock out, which turns the explorer off rather than giving it a search\nbox that can only time out.", Nested: ""},
|
||||
{Key: "timeout", Type: "time.Duration", Doc: "Timeout bounds a management call. It does not bound a download: the\nmirror answers a fetch as soon as it has started one, and the bytes keep\narriving long after this deadline has passed.", Nested: ""},
|
||||
@@ -240,6 +241,7 @@ var Reference = []Block{
|
||||
{Key: "aliases", Type: "[]string", Doc: "", Nested: ""},
|
||||
{Key: "description", Type: "string", Doc: "Description is what this model is for, in a line.\n\nWritten by whoever added it, because nobody else can say it. A registry\noffers a pipeline tag and a wall of README, and neither answers the\nquestion a person picking from a list is asking: which of these nine\nnames should I send this request to.\n\nIt reaches clients through /v1/models, which is the only view most of\nthem ever see. A gateway serving nine models by name alone makes that\nchoice by guesswork.", Nested: ""},
|
||||
{Key: "source", Type: "string", Doc: "Source is the weights path or a HuggingFace id.", Nested: ""},
|
||||
{Key: "allow_pickle", Type: "bool", Doc: "AllowPickle lets the nodes fetch this model's weights, or its\nadapters', in a pickle format (.bin, .pt, .pth, .ckpt), which run code\nas they are loaded. Off, such a source is refused when the model is\nadded and by the node before it fetches anything. For a model that has\nno safetensors or GGUF form and whose files were checked.", Nested: ""},
|
||||
{Key: "profiles", Type: "[]Profile", Doc: "", Nested: ""},
|
||||
{Key: "scaling", Type: "Scaling", Doc: "", Nested: "catalog.Scaling"},
|
||||
{Key: "class", Type: "string", Doc: "Class is how much this model's capacity is worth against another's,\nwhen there is not enough for both: critical, standard or opportunistic.\nEmpty is standard, or opportunistic for a model loaded on demand. See\nEffectiveClass.", Nested: ""},
|
||||
@@ -503,6 +505,11 @@ var Reference = []Block{
|
||||
{Package: "engines", Name: "SimConfig", Title: "SimConfig", Doc: "", Fields: []Field{
|
||||
{Key: "bin", Type: "string", Doc: "Bin is the llmgw binary itself; the engine is one of its subcommands.\nEmpty means whatever started this agent, which is right almost always.", Nested: ""},
|
||||
}},
|
||||
{Package: "engines", Name: "StrataConfig", Title: "StrataConfig", Doc: "", Fields: []Field{
|
||||
{Key: "dir", Type: "string", Doc: "Dir is Strata's source directory, where setup.py and serve/ are.", Nested: ""},
|
||||
{Key: "python", Type: "string", Doc: "Python is the interpreter of its environment; Dir's .venv when unset.", Nested: ""},
|
||||
{Key: "engine", Type: "string", Doc: "Engine is the directory of its prebuilt engine, where the strata\nbinary and the CUDA libraries it ships are; Dir's engine/ when unset.", Nested: ""},
|
||||
}},
|
||||
{Package: "engines", Name: "TRTLLMConfig", Title: "TRTLLMConfig", Doc: "", Fields: []Field{
|
||||
{Key: "python", Type: "string", Doc: "Python is the interpreter of the environment TensorRT-LLM is installed\nin: <venv>/bin/python for a virtual environment, python3 in NVIDIA's\ncontainer. Naming it is what makes this node offer the engine; empty\noffers nothing.\n\nThe interpreter rather than trtllm-serve, because llmgw starts the\nserver through a few lines of its own that make it demand the\ngateway's key, which trtllm-serve has no option for.", Nested: ""},
|
||||
}},
|
||||
@@ -612,6 +619,7 @@ var Reference = []Block{
|
||||
{Key: "trtllm", Type: "TRTLLMConfig", Doc: "TRTLLM is TensorRT-LLM, which the agent does not install: a node\noffers it only when its interpreter is named here.", Nested: "engines.TRTLLMConfig"},
|
||||
{Key: "sim", Type: "SimConfig", Doc: "", Nested: "engines.SimConfig"},
|
||||
{Key: "comfyui", Type: "ComfyUIConfig", Doc: "ComfyUI makes pictures; offered by a node whose ComfyUI is named here.", Nested: "engines.ComfyUIConfig"},
|
||||
{Key: "strata", Type: "StrataConfig", Doc: "Strata serves a few mixture-of-experts models from RAM and one card;\noffered by a node whose Strata is named here, or that installed it.", Nested: "engines.StrataConfig"},
|
||||
}},
|
||||
{Package: "guard", Name: "Config", Title: "guard.Config", Doc: "", Fields: []Field{
|
||||
{Key: "enabled", Type: "*bool", Doc: "Enabled turns every policy on or off at once: false checks nothing,\nwhatever is defined or attached. Unset is on.", Nested: ""},
|
||||
|
||||
@@ -23,6 +23,7 @@ var Sources = []string{
|
||||
"internal/agent/engines/trtllm.go",
|
||||
"internal/agent/engines/sim.go",
|
||||
"internal/agent/engines/comfyui.go",
|
||||
"internal/agent/engines/strata.go",
|
||||
"internal/agent/supervisor/supervisor.go",
|
||||
"internal/tlsconf/tlsconf.go",
|
||||
"internal/control/catalog/catalog.go",
|
||||
|
||||
@@ -385,6 +385,12 @@ func Evaluate(f Facts) []Finding {
|
||||
add(Warn, "gpu", fmt.Sprintf("GPU %s is kept out of placement: %s",
|
||||
discovery.DeviceLabel(g.Index), strings.Join(g.Health.Problems, "; ")), "")
|
||||
}
|
||||
if g.MIG == nil {
|
||||
if p := g.PCIe.Problem(); p != "" {
|
||||
add(Warn, "gpu", fmt.Sprintf("GPU %d (%s) %s", g.Index, g.Name, p),
|
||||
"lspci -vv -s "+g.PCIBusID+" shows LnkCap (the card) against LnkSta (the link)")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
evaluateMIG(f, add)
|
||||
|
||||
@@ -41,6 +41,9 @@ type Header struct {
|
||||
// them; otherwise a head is EmbeddingLength / HeadCount.
|
||||
KeyLength uint64
|
||||
ValueLength uint64
|
||||
// ContextLength is the context the model was trained for
|
||||
// (<arch>.context_length); 0 when the file does not say.
|
||||
ContextLength uint64
|
||||
// Embedding is whether the model pools its tokens into one vector
|
||||
// (<arch>.pooling_type set and not none): an embedding model, which
|
||||
// llama.cpp serves with --embedding and does not chat.
|
||||
@@ -271,6 +274,8 @@ func Read(r io.Reader) (Header, error) {
|
||||
h.ValueLength = num
|
||||
case "pooling_type":
|
||||
h.Embedding = num > 0
|
||||
case "context_length":
|
||||
h.ContextLength = num
|
||||
}
|
||||
}
|
||||
// The tokenizer's arrays — megabytes of vocabulary — come after the
|
||||
|
||||
@@ -418,7 +418,15 @@ func (g *Gate) Input(ctx context.Context, s Subject, body []byte) (Verdict, erro
|
||||
// Outside services next, then the judge: what a pattern already stopped
|
||||
// need not wait for either.
|
||||
if user != "" {
|
||||
sf, serr := g.askServices(ctx, p, true, user)
|
||||
// As the model will read it: what redact took out stays out of the
|
||||
// outside service's sight too.
|
||||
told := user
|
||||
for _, rw := range rewrites {
|
||||
if rw.seg.text != "" && strings.Contains(told, rw.seg.text) {
|
||||
told = strings.ReplaceAll(told, rw.seg.text, rw.text)
|
||||
}
|
||||
}
|
||||
sf, serr := g.askServices(ctx, p, true, told)
|
||||
if serr != nil && !v.Shadow {
|
||||
v.Took = g.now().Sub(start)
|
||||
return v, serr
|
||||
|
||||
@@ -54,6 +54,10 @@ type Subject struct {
|
||||
Stream bool
|
||||
// Tools and Vision are what the model's catalog entry says it can do.
|
||||
Tools, Vision bool
|
||||
// Restored are the values the request's redaction took out, by
|
||||
// placeholder, which restore puts back in the answer: outside services
|
||||
// reading the answer are not told them.
|
||||
Restored map[string]string
|
||||
}
|
||||
|
||||
// Verdict is a judgement on one stage of one request.
|
||||
|
||||
+43
-18
@@ -68,27 +68,32 @@ func (g *Gate) Output(ctx context.Context, s Subject) OutputFilter {
|
||||
return nil
|
||||
}
|
||||
memo := g.tools.get(s.RequestID)
|
||||
return &outFilter{ctx: ctx, g: g, p: p, names: names, shadow: p.Mode == "shadow",
|
||||
return &outFilter{ctx: ctx, g: g, p: p, names: names, shadow: p.Mode == "shadow", restored: s.Restored,
|
||||
schemas: memo.schemas, answer: memo.answer, user: memo.user, choices: map[int]*choiceState{}}
|
||||
}
|
||||
|
||||
// outFilter judges one answer.
|
||||
type outFilter struct {
|
||||
ctx context.Context
|
||||
g *Gate
|
||||
p Policy
|
||||
names string
|
||||
shadow bool
|
||||
schemas map[string]json.RawMessage
|
||||
answer json.RawMessage
|
||||
full strings.Builder // the whole answer, when it must fit a schema
|
||||
user string // the request's last user turn, for the judge
|
||||
sofar strings.Builder // the answer so far, for the judge
|
||||
served strings.Builder // the answer so far, for outside services, asked once at its end
|
||||
judged int // how much of it the judge has read
|
||||
v Verdict
|
||||
took time.Duration
|
||||
stopped bool
|
||||
// servicesAsked says the outside services have read the answer.
|
||||
servicesAsked bool
|
||||
// restored are the request's redacted values, by placeholder, put back
|
||||
// in the answer before this filter reads it.
|
||||
restored map[string]string
|
||||
ctx context.Context
|
||||
g *Gate
|
||||
p Policy
|
||||
names string
|
||||
shadow bool
|
||||
schemas map[string]json.RawMessage
|
||||
answer json.RawMessage
|
||||
full strings.Builder // the whole answer, when it must fit a schema
|
||||
user string // the request's last user turn, for the judge
|
||||
sofar strings.Builder // the answer so far, for the judge
|
||||
served strings.Builder // the answer so far, for outside services, asked once at its end
|
||||
judged int // how much of it the judge has read
|
||||
v Verdict
|
||||
took time.Duration
|
||||
stopped bool
|
||||
|
||||
held []string // events not yet released
|
||||
sent bool // whether a chunk this filter made has gone out
|
||||
@@ -103,6 +108,8 @@ type choiceState struct {
|
||||
// replaced: it reaches the client only in a chunk this filter makes.
|
||||
carried bool
|
||||
calls map[int]*call
|
||||
// done says the choice has finished: its text is all here.
|
||||
done bool
|
||||
}
|
||||
|
||||
type call struct {
|
||||
@@ -212,6 +219,13 @@ func (f *outFilter) askServices(text string) finding {
|
||||
if !slices.ContainsFunc(f.p.Services, Service.onOutput) || text == "" {
|
||||
return finding{}
|
||||
}
|
||||
// As the model wrote it: with restore, the values put back for the
|
||||
// client are taken out again, as redact took them out of the request.
|
||||
for ph, v := range f.restored {
|
||||
if v != "" {
|
||||
text = strings.ReplaceAll(text, v, ph)
|
||||
}
|
||||
}
|
||||
start := time.Now()
|
||||
fs, err := f.g.askServices(f.ctx, f.p, false, text)
|
||||
f.took += time.Since(start)
|
||||
@@ -379,6 +393,7 @@ func (f *outFilter) OnEvent(data string) ([]string, bool) {
|
||||
}
|
||||
if ch.FinishReason != nil && *ch.FinishReason != "" {
|
||||
finished = true
|
||||
st.done = true
|
||||
}
|
||||
}
|
||||
if f.shadow {
|
||||
@@ -445,8 +460,14 @@ func (f *outFilter) flush(final bool) (bool, map[int]string, bool) {
|
||||
}
|
||||
if i == 0 {
|
||||
f.served.WriteString(head)
|
||||
if final && f.askServices(f.served.String()).action == Block {
|
||||
blocked = true
|
||||
// Once, when the first choice is whole — not when another
|
||||
// choice finishes first — and not again at the end of the
|
||||
// stream after its finish.
|
||||
if final && st.done && !f.servicesAsked {
|
||||
f.servicesAsked = true
|
||||
if f.askServices(f.served.String()).action == Block {
|
||||
blocked = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if head != "" {
|
||||
@@ -566,6 +587,10 @@ func withoutContent(data string) string {
|
||||
}
|
||||
|
||||
func (f *outFilter) OnEnd() []string {
|
||||
// The stream is over: whatever a choice has is all it will have.
|
||||
for _, st := range f.choices {
|
||||
st.done = true
|
||||
}
|
||||
switch {
|
||||
case f.stopped:
|
||||
return nil
|
||||
|
||||
@@ -220,10 +220,21 @@ func Merge(names []string, ps []Policy) Policy {
|
||||
out.RequireCapabilities = out.RequireCapabilities || p.RequireCapabilities
|
||||
out.ValidateStructured = out.ValidateStructured || p.ValidateStructured
|
||||
out.Judge = mergeJudge(out.Judge, p.Judge)
|
||||
// Every service any policy asks, each once.
|
||||
// Every service any policy asks, each once; one named by two is
|
||||
// asked the stricter way of the two: both directions either asks,
|
||||
// blocking where either blocks.
|
||||
for _, sv := range p.Services {
|
||||
if !slices.ContainsFunc(out.Services, func(e Service) bool { return e.Name == sv.Name }) {
|
||||
i := slices.IndexFunc(out.Services, func(e Service) bool { return e.Name == sv.Name })
|
||||
if i < 0 {
|
||||
out.Services = append(out.Services, sv)
|
||||
continue
|
||||
}
|
||||
e := &out.Services[i]
|
||||
if e.action() == Flag && sv.action() == Block {
|
||||
e.Action = "block"
|
||||
}
|
||||
if e.onInput() != sv.onInput() || e.onOutput() != sv.onOutput() {
|
||||
e.On = "both"
|
||||
}
|
||||
}
|
||||
out.MaxOutputTokens = minSet(out.MaxOutputTokens, p.MaxOutputTokens)
|
||||
|
||||
@@ -144,7 +144,8 @@ func (s Service) ask(ctx context.Context, text string) (unsafe bool, reason stri
|
||||
return false, "", fmt.Errorf("answer is not JSON: %w", err)
|
||||
}
|
||||
v, ok := at(doc, s.Blocked)
|
||||
if !ok {
|
||||
if !ok || v == nil {
|
||||
// null is no answer: a service that cannot say is not saying safe.
|
||||
return false, "", fmt.Errorf("the answer has no %s", s.Blocked)
|
||||
}
|
||||
if s.Reason != "" {
|
||||
|
||||
@@ -98,3 +98,89 @@ func TestAServiceIsCheckedWhenThePolicyIsRead(t *testing.T) {
|
||||
t.Errorf("at: %v %v", v, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// recorder is a service that finds nothing unsafe and keeps what it read.
|
||||
func recorder(t *testing.T) (*httptest.Server, *[]string) {
|
||||
var seen []string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var in struct{ Input string }
|
||||
_ = json.NewDecoder(r.Body).Decode(&in)
|
||||
seen = append(seen, in.Input)
|
||||
_, _ = w.Write([]byte(`{"flagged": false}`))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv, &seen
|
||||
}
|
||||
|
||||
func TestAnOutsideServiceReadsWhatTheModelReadsNotWhatWasRedacted(t *testing.T) {
|
||||
srv, seen := recorder(t)
|
||||
g := gate(t, lib(map[string]Policy{"p": {
|
||||
Input: Rules{PII: map[string]string{"email": "redact"}},
|
||||
Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged"}},
|
||||
}}, "p"))
|
||||
if _, err := g.Input(context.Background(), Subject{}, chat("user", "mail ayse@example.com")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(*seen) != 1 || strings.Contains((*seen)[0], "ayse@") || !strings.Contains((*seen)[0], "[EMAIL_1]") {
|
||||
t.Errorf("the service read %q", *seen)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAStreamedAnswerIsReadByAServiceOnce(t *testing.T) {
|
||||
srv, seen := recorder(t)
|
||||
g := gate(t, lib(map[string]Policy{"p": {
|
||||
Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged", On: "output"}},
|
||||
}}, "p"))
|
||||
got, _, _ := stream(g.Output(context.Background(), Subject{Stream: true}), "Hello ", "there.")
|
||||
if got != "Hello there." {
|
||||
t.Errorf("client got %q", got)
|
||||
}
|
||||
if len(*seen) != 1 || (*seen)[0] != "Hello there." {
|
||||
t.Errorf("the service was asked %d times: %q", len(*seen), *seen)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAServiceNamedByTwoPoliciesIsAskedTheStricterWay(t *testing.T) {
|
||||
a := Policy{Services: []Service{{Name: "cls", URL: "https://c", Blocked: "x", On: "input", Action: "flag"}}}
|
||||
b := Policy{Services: []Service{{Name: "cls", URL: "https://c", Blocked: "x", On: "output", Action: "block"}}}
|
||||
got := Merge([]string{"a", "b"}, []Policy{a, b}).Services
|
||||
if len(got) != 1 || got[0].action() != Block || !got[0].onInput() || !got[0].onOutput() {
|
||||
t.Errorf("%+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnOutputServiceReadsTheAnswerWithoutTheRestoredValues(t *testing.T) {
|
||||
srv, seen := recorder(t)
|
||||
g := gate(t, lib(map[string]Policy{"p": {
|
||||
Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged", On: "output"}},
|
||||
}}, "p"))
|
||||
f := g.Output(context.Background(), Subject{Restored: map[string]string{"[EMAIL_1]": "ayse@example.com"}})
|
||||
f.Buffered([]byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":"Sent to ayse@example.com."},"finish_reason":"stop"}]}`))
|
||||
if len(*seen) != 1 || strings.Contains((*seen)[0], "ayse@") || !strings.Contains((*seen)[0], "[EMAIL_1]") {
|
||||
t.Errorf("the service read %q", *seen)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOfTwoChoicesTheServiceReadsTheFirstWhole(t *testing.T) {
|
||||
// The second finishes first; the first is still arriving then.
|
||||
srv, seen := recorder(t)
|
||||
g := gate(t, lib(map[string]Policy{"p": {
|
||||
Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged", On: "output"}},
|
||||
}}, "p"))
|
||||
f := g.Output(context.Background(), Subject{Stream: true})
|
||||
ev := func(idx int, text string, finish bool) string {
|
||||
c := map[string]any{"index": idx, "delta": map[string]any{"content": text}, "finish_reason": nil}
|
||||
if finish {
|
||||
c["finish_reason"] = "stop"
|
||||
}
|
||||
b, _ := json.Marshal(map[string]any{"choices": []any{c}})
|
||||
return string(b)
|
||||
}
|
||||
for _, e := range []string{ev(0, "The first ", false), ev(1, "Second.", true), ev(0, "answer.", true)} {
|
||||
f.OnEvent(e)
|
||||
}
|
||||
f.OnEnd()
|
||||
if len(*seen) != 1 || (*seen)[0] != "The first answer." {
|
||||
t.Errorf("the service read %q", *seen)
|
||||
}
|
||||
}
|
||||
+67
-3
@@ -18,6 +18,7 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -69,15 +70,78 @@ type Model struct {
|
||||
// for one that has none.
|
||||
GGUF *GGUFInfo `json:"gguf"`
|
||||
// Safetensors is the registry's summary of original weights.
|
||||
Safetensors *struct {
|
||||
Total int64 `json:"total"`
|
||||
} `json:"safetensors,omitempty"`
|
||||
Safetensors *SafetensorsInfo `json:"safetensors,omitempty"`
|
||||
// Siblings are the repository's files, when asked for with WithFiles.
|
||||
Siblings []struct {
|
||||
RFilename string `json:"rfilename"`
|
||||
} `json:"siblings,omitempty"`
|
||||
}
|
||||
|
||||
// SafetensorsInfo is the registry's summary of a repository's safetensors.
|
||||
type SafetensorsInfo struct {
|
||||
Total int64 `json:"total"`
|
||||
// Parameters counts them by dtype: {"BF16": 8e9}. A packed quantisation
|
||||
// (AWQ, GPTQ, MLX) is counted as its I32 or U32 at the count of weights
|
||||
// it unpacks to, not of the integers that hold them.
|
||||
Parameters map[string]int64 `json:"parameters,omitempty"`
|
||||
}
|
||||
|
||||
// dtypeBytes is what one parameter of a safetensors dtype takes.
|
||||
var dtypeBytes = map[string]float64{
|
||||
"F64": 8, "I64": 8, "U64": 8, "F32": 4, "BF16": 2, "F16": 2, "I16": 2, "U16": 2,
|
||||
"F8_E4M3": 1, "F8_E5M2": 1, "F8_E8M0": 1, "I8": 1, "U8": 1, "BOOL": 1,
|
||||
}
|
||||
|
||||
// WeightBytes is about how much the safetensors weights take: each dtype's
|
||||
// count at its size, and a packed quantisation at 4 bits a weight, 8 for
|
||||
// one tagged 8-bit — I32 or U32 always, U8 for MXFP4 and NVFP4, which the
|
||||
// registry counts at the weights they unpack to. 0 when the registry has
|
||||
// no summary.
|
||||
func (m Model) WeightBytes() int64 {
|
||||
if m.Safetensors == nil {
|
||||
return 0
|
||||
}
|
||||
if len(m.Safetensors.Parameters) == 0 {
|
||||
return m.Safetensors.Total * 2
|
||||
}
|
||||
packed := 0.5
|
||||
if m.HasTag("8-bit") {
|
||||
packed = 1
|
||||
}
|
||||
fp4 := m.HasTag("mxfp4") || m.HasTag("nvfp4") || m.HasTag("fp4")
|
||||
var b float64
|
||||
for dt, n := range m.Safetensors.Parameters {
|
||||
switch per, ok := dtypeBytes[dt]; {
|
||||
case dt == "U8" && fp4:
|
||||
b += float64(n) * 0.5
|
||||
case ok:
|
||||
b += float64(n) * per
|
||||
case dt == "I32" || dt == "U32":
|
||||
b += float64(n) * packed
|
||||
default:
|
||||
b += float64(n) * 2
|
||||
}
|
||||
}
|
||||
return int64(b)
|
||||
}
|
||||
|
||||
// FP8 reports whether the safetensors weights are in an 8-bit float, which
|
||||
// vLLM runs only on cards of compute capability 8.9 or newer.
|
||||
func (m Model) FP8() bool {
|
||||
if m.Safetensors == nil {
|
||||
return false
|
||||
}
|
||||
for dt := range m.Safetensors.Parameters {
|
||||
if strings.HasPrefix(dt, "F8_") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// HasTag reports whether the registry tags the repository so.
|
||||
func (m Model) HasTag(tag string) bool { return slices.Contains(m.Tags, tag) }
|
||||
|
||||
// GGUFInfo is the registry's summary of a repository's GGUF weights.
|
||||
type GGUFInfo struct {
|
||||
// Total is the parameter count.
|
||||
|
||||
@@ -52,6 +52,9 @@ type ModelsResponse struct {
|
||||
// when it is not.
|
||||
Base string `json:"base"`
|
||||
Models []Status `json:"models"`
|
||||
// Removed are the versions removed here, model/version to when (ms),
|
||||
// for a peer mirror to remove too.
|
||||
Removed map[string]int64 `json:"removed,omitempty"`
|
||||
}
|
||||
|
||||
func (a *API) list(w http.ResponseWriter, _ *http.Request) {
|
||||
@@ -60,7 +63,7 @@ func (a *API) list(w http.ResponseWriter, _ *http.Request) {
|
||||
fail(w, http.StatusInternalServerError, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, ModelsResponse{Base: a.base, Models: all})
|
||||
writeJSON(w, http.StatusOK, ModelsResponse{Base: a.base, Models: all, Removed: a.m.Tombstones()})
|
||||
}
|
||||
|
||||
// FetchRequest asks the mirror to download a model.
|
||||
|
||||
+55
-17
@@ -4,6 +4,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -24,6 +25,29 @@ import (
|
||||
type Client struct {
|
||||
base string
|
||||
hc *http.Client
|
||||
// others are the mirrors kept in step with this one, asked in turn when
|
||||
// it cannot be reached: the same models, on another machine.
|
||||
others []string
|
||||
}
|
||||
|
||||
// NewFailoverClient is NewClient over several mirrors kept in step: the
|
||||
// first answers, and the next is asked when it cannot be reached.
|
||||
func NewFailoverClient(addrs []string, tc tlsconf.Config, timeout time.Duration) (*Client, error) {
|
||||
if len(addrs) == 0 {
|
||||
return nil, fmt.Errorf("no mirror address")
|
||||
}
|
||||
c, err := NewClient(addrs[0], tc, timeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, a := range addrs[1:] {
|
||||
u, err := url.Parse(a)
|
||||
if err != nil || u.Host == "" {
|
||||
return nil, fmt.Errorf("the mirror address %q is not a URL with a host", a)
|
||||
}
|
||||
c.others = append(c.others, strings.TrimSuffix(a, "/"))
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// NewClient dials a mirror's management API, with mutual TLS when configured.
|
||||
@@ -149,28 +173,42 @@ func (c *Client) Remove(ctx context.Context, model, version string) ([]string, e
|
||||
}
|
||||
|
||||
func (c *Client) do(ctx context.Context, method, path string, body, out any) error {
|
||||
var rdr io.Reader
|
||||
var b []byte
|
||||
if body != nil {
|
||||
b, err := json.Marshal(body)
|
||||
var err error
|
||||
if b, err = json.Marshal(body); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// The mirror first, then those kept in step with it, while each cannot
|
||||
// be reached at all: one that answers, even with a refusal, is the one.
|
||||
var resp *http.Response
|
||||
var errs []error
|
||||
for _, base := range append([]string{c.base}, c.others...) {
|
||||
var rdr io.Reader
|
||||
if body != nil {
|
||||
rdr = bytes.NewReader(b)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, base+path, rdr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rdr = bytes.NewReader(b)
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
if resp, err = c.hc.Do(req); err == nil {
|
||||
break
|
||||
}
|
||||
// Wrapped with the address because the failure an operator sees on
|
||||
// a page is otherwise "connection refused" with no clue which of
|
||||
// the machines refused it.
|
||||
errs = append(errs, fmt.Errorf("reaching the mirror at %s: %w", base, err))
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.base+path, rdr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
|
||||
resp, err := c.hc.Do(req)
|
||||
if err != nil {
|
||||
// Wrapped with the address because the failure an operator sees on a
|
||||
// page is otherwise "connection refused" with no clue which of the two
|
||||
// machines refused it.
|
||||
return fmt.Errorf("reaching the mirror at %s: %w", c.base, err)
|
||||
if resp == nil {
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
|
||||
@@ -43,6 +43,9 @@ type Status struct {
|
||||
|
||||
StartedAt time.Time `json:"started_at,omitzero"`
|
||||
FinishedAt time.Time `json:"finished_at,omitzero"`
|
||||
// AddedAt is when the file came to this mirror, as its modification
|
||||
// time says: set against a removal elsewhere, the later stands.
|
||||
AddedAt time.Time `json:"added_at,omitzero"`
|
||||
}
|
||||
|
||||
// Percent reports progress, and false when the total is unknown.
|
||||
@@ -106,6 +109,16 @@ var ErrAlreadyFetching = errors.New("a download for that version is already runn
|
||||
// the API call that asked for it, which is the point of doing it here rather
|
||||
// than making the caller hold a connection open for an hour.
|
||||
func (m *Manager) Start(source string, e Entry) (Status, error) {
|
||||
return m.start(source, e, "", m.client)
|
||||
}
|
||||
|
||||
// startCopy fetches a file a peer mirror holds, from a private address,
|
||||
// and keeps it only if its checksum is the peer's.
|
||||
func (m *Manager) startCopy(source string, e Entry, sha string) (Status, error) {
|
||||
return m.start(source, e, sha, fetchClient(true))
|
||||
}
|
||||
|
||||
func (m *Manager) start(source string, e Entry, expect string, client *http.Client) (Status, error) {
|
||||
if source == "" {
|
||||
return Status{}, fmt.Errorf("a source URL is required")
|
||||
}
|
||||
@@ -132,16 +145,16 @@ func (m *Manager) Start(source string, e Entry) (Status, error) {
|
||||
out := *st
|
||||
m.mu.Unlock()
|
||||
|
||||
go m.run(source, e, key)
|
||||
go m.run(source, e, key, expect, client)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (m *Manager) run(source string, e Entry, key string) {
|
||||
func (m *Manager) run(source string, e Entry, key, expect string, client *http.Client) {
|
||||
// Its own context, deliberately unbounded by any request. A download is
|
||||
// minutes to hours and the thing that asked for it is long gone.
|
||||
ctx := context.Background()
|
||||
|
||||
added, err := addWith(ctx, m.client, m.root, source, e, func(done, total int64) {
|
||||
added, err := addWith(ctx, client, m.root, source, e, func(done, total int64) {
|
||||
m.mu.Lock()
|
||||
if st, ok := m.running[key]; ok {
|
||||
st.BytesDone, st.BytesTotal = done, total
|
||||
@@ -149,6 +162,16 @@ func (m *Manager) run(source string, e Entry, key string) {
|
||||
m.mu.Unlock()
|
||||
})
|
||||
|
||||
if err == nil && expect != "" && added.SHA256 != expect {
|
||||
// Not what the peer holds: a file changed or cut short on the way.
|
||||
_ = Remove(m.root, e)
|
||||
err = fmt.Errorf("the copy's sha256 %s is not the peer's %s", added.SHA256, expect)
|
||||
}
|
||||
if err == nil {
|
||||
// Added again after it was removed: the new one stands.
|
||||
m.clearTombstone(e.Model, e.Version)
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
st, ok := m.running[key]
|
||||
@@ -284,6 +307,9 @@ func (m *Manager) All() ([]Status, error) {
|
||||
Model: e.Model, Version: e.Version, File: e.File,
|
||||
State: Present, BytesDone: e.Bytes, BytesTotal: e.Bytes,
|
||||
}
|
||||
if fi, err := os.Stat(filepath.Join(m.root, e.Model, e.Version, e.File)); err == nil {
|
||||
st.AddedAt = fi.ModTime()
|
||||
}
|
||||
// The checksum and the source are not recoverable from a directory
|
||||
// listing, and the checksum is the one field worth having: it is what
|
||||
// an operator compares against what upstream published. Dropping it in
|
||||
@@ -364,6 +390,9 @@ func (m *Manager) RemoveVersion(model, version string) ([]string, error) {
|
||||
return nil, fmt.Errorf("%s/%s: %w", model, version, os.ErrNotExist)
|
||||
}
|
||||
m.Forget(model, version)
|
||||
// Kept, so a peer mirror that still has it removes it too, and does
|
||||
// not copy it back here.
|
||||
m.recordTombstone(model, version, time.Now())
|
||||
m.log.Info("mirror removed a version", "model", model, "version", version, "files", len(removed))
|
||||
return removed, nil
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package mirror
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Mirrors kept in step, so the cluster has its models while one machine is
|
||||
// down: each reads its peers every while, copies the versions a peer has
|
||||
// that it lacks, checksum checked, and removes those a peer removed after
|
||||
// they came here. Neither leads; what was done last stands.
|
||||
|
||||
// PeerInterval is how often each peer is read.
|
||||
const PeerInterval = 30 * time.Second
|
||||
|
||||
// Peers keeps one mirror in step with others.
|
||||
type Peers struct {
|
||||
M *Manager
|
||||
Clients []*Client
|
||||
Log *slog.Logger
|
||||
}
|
||||
|
||||
// Run reads the peers until ctx ends.
|
||||
func (p *Peers) Run(ctx context.Context) {
|
||||
t := time.NewTicker(PeerInterval)
|
||||
defer t.Stop()
|
||||
for {
|
||||
for _, c := range p.Clients {
|
||||
p.SyncOnce(ctx, c)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SyncOnce brings this mirror in step with one peer.
|
||||
func (p *Peers) SyncOnce(ctx context.Context, c *Client) {
|
||||
ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
|
||||
defer cancel()
|
||||
theirs, err := c.List(ctx)
|
||||
if err != nil {
|
||||
p.Log.Warn("could not read a peer mirror; trying again shortly", "peer", c.Addr(), "err", err)
|
||||
return
|
||||
}
|
||||
mine, err := p.M.All()
|
||||
if err != nil {
|
||||
p.Log.Warn("could not read this mirror", "err", err)
|
||||
return
|
||||
}
|
||||
here := map[string]Status{}
|
||||
for _, st := range mine {
|
||||
here[st.Model+"/"+st.Version+"/"+st.File] = st
|
||||
}
|
||||
removed := p.M.Tombstones()
|
||||
|
||||
// What the peer removed after it came here goes here too.
|
||||
for key, at := range theirs.Removed {
|
||||
model, version, ok := strings.Cut(key, "/")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
newest := time.Time{}
|
||||
has := false
|
||||
for _, st := range mine {
|
||||
if st.Model == model && st.Version == version {
|
||||
has = true
|
||||
if st.AddedAt.After(newest) {
|
||||
newest = st.AddedAt
|
||||
}
|
||||
}
|
||||
}
|
||||
if has && newest.UnixMilli() < at {
|
||||
if _, err := p.M.RemoveVersion(model, version); err == nil {
|
||||
p.Log.Info("removed a version a peer mirror removed", "peer", c.Addr(), "model", model, "version", version)
|
||||
}
|
||||
} else if !has && at > removed[key] {
|
||||
p.M.recordTombstone(model, version, time.UnixMilli(at))
|
||||
}
|
||||
}
|
||||
|
||||
// Versions whole there and missing here: every file present, each with
|
||||
// its checksum, so a copy can be checked.
|
||||
type version struct {
|
||||
files []Status
|
||||
whole bool
|
||||
}
|
||||
versions := map[string]*version{}
|
||||
for _, st := range theirs.Models {
|
||||
key := st.Model + "/" + st.Version
|
||||
v := versions[key]
|
||||
if v == nil {
|
||||
v = &version{whole: true}
|
||||
versions[key] = v
|
||||
}
|
||||
v.files = append(v.files, st)
|
||||
v.whole = v.whole && st.State == Present && st.SHA256 != ""
|
||||
}
|
||||
for key, v := range versions {
|
||||
if !v.whole {
|
||||
continue
|
||||
}
|
||||
for _, st := range v.files {
|
||||
if _, ok := here[key+"/"+st.File]; ok {
|
||||
continue
|
||||
}
|
||||
// Removed here after the peer had it: not copied back. The
|
||||
// peer removes its own when it reads this one.
|
||||
if at, gone := removed[key]; gone && st.AddedAt.UnixMilli() < at {
|
||||
continue
|
||||
}
|
||||
source := strings.TrimRight(theirs.Base, "/") + "/" + st.Model + "/" + st.Version + "/" + st.File
|
||||
e := Entry{Model: st.Model, Version: st.Version, File: st.File}
|
||||
if _, err := p.M.startCopy(source, e, st.SHA256); err == nil {
|
||||
p.Log.Info("copying a version from a peer mirror", "peer", c.Addr(), "model", st.Model,
|
||||
"version", st.Version, "file", st.File)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Loaded 100 of 157 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user