Compare commits
70
Commits
No files matched your search
@@ -39,6 +39,14 @@ jobs:
|
||||
- name: Build
|
||||
run: go build ./...
|
||||
|
||||
# llmgw-kb runs on people's own machines, Windows included, which no
|
||||
# test here runs on: building it for each catches a Unix-only call.
|
||||
- name: Build llmgw-kb for every system
|
||||
run: |
|
||||
for target in windows/amd64 windows/arm64 darwin/amd64 darwin/arm64 linux/amd64 linux/arm64; do
|
||||
CGO_ENABLED=0 GOOS="${target%/*}" GOARCH="${target#*/}" go build -o /dev/null ./cmd/llmgw-kb
|
||||
done
|
||||
|
||||
# The Terraform provider is a module of its own, which ./... above does
|
||||
# not reach. Its tests check its client against docs/admin-api.openapi.json,
|
||||
# so a change to the admin API that breaks the provider fails here. The
|
||||
|
||||
@@ -55,6 +55,21 @@ jobs:
|
||||
(cd dist && sha256sum "${name}" >> checksums.txt)
|
||||
done
|
||||
|
||||
# llmgw-kb, which people run on their own machines to put their
|
||||
# documents into a knowledge base: Windows too, which llmgw itself
|
||||
# does not build for.
|
||||
for target in windows/amd64 windows/arm64 darwin/amd64 darwin/arm64 linux/amd64 linux/arm64; do
|
||||
os="${target%/*}"; arch="${target#*/}"
|
||||
ext=""; [ "$os" = windows ] && ext=".exe"
|
||||
name="llmgw-kb-${version}-${os}-${arch}${ext}"
|
||||
echo "building $name"
|
||||
CGO_ENABLED=0 GOOS="$os" GOARCH="$arch" go build \
|
||||
-trimpath \
|
||||
-ldflags "-s -w -X gitea.mynodes.xyz/ilker/llmgw/internal/version.Version=${version}" \
|
||||
-o "dist/${name}" ./cmd/llmgw-kb
|
||||
(cd dist && sha256sum "${name}" >> checksums.txt)
|
||||
done
|
||||
|
||||
# The reverse proxy configurations, so a release page carries
|
||||
# everything install.md asks for and the repository is not needed.
|
||||
for f in deploy/proxy/Caddyfile deploy/proxy/nginx.conf; do
|
||||
@@ -91,6 +106,7 @@ jobs:
|
||||
# test the cross-compiled ones, which is why the native build is
|
||||
# checked rather than assumed.
|
||||
./dist/llmgw-${version}-linux-amd64 --version
|
||||
./dist/llmgw-kb-${version}-linux-amd64 --version
|
||||
cat dist/checksums.txt
|
||||
|
||||
# Signed, so a machine that trusts the key built into it installs
|
||||
|
||||
+291
-8
@@ -569,9 +569,12 @@ portal is not done — the page shows; the console and ctl change.
|
||||
client tests against a fake admin API and the OpenAPI document;
|
||||
acceptance tests with Terraform and OpenTofu against the fake; one
|
||||
apply, empty plan, in-place change and destroy against a real `llmgw
|
||||
server` run by hand. Left: organizations, teams, projects, people and
|
||||
keys are not in the admin API (they are `llmgw org`/`users`/`keys`
|
||||
against the store), so not in the provider; a model's definition beyond
|
||||
server` run by hand. Organizations, teams, projects, keys and
|
||||
memberships are in the admin API since 2026-10-08 (`/api/v1/orgs`,
|
||||
`teams`, `projects`, `keys`, `users/{id}/memberships`) and in the
|
||||
provider (`llmgw_org`, `llmgw_team`, `llmgw_project`, `llmgw_key`,
|
||||
`llmgw_membership`); accounts are still made by sign-up, invitation or
|
||||
OIDC/SCIM. Left: a model's definition beyond
|
||||
scaling, class, owner and adapters is changed by replacing it (the API
|
||||
has no edit, and `resources` is not used yet); `canary`, `slo`,
|
||||
`scaling.schedule` and `prewarm`, `guard` role and `origin` are not in
|
||||
@@ -793,10 +796,19 @@ leader 30s later with /health answering throughout; rollback did the same.
|
||||
facility state from the store (it kept what it read at start).
|
||||
- Retires `appendGuardIfMissing` (server init): the import does it.
|
||||
|
||||
### E2. Generated files from settings (L) — firewall on/off done 2026-10-06
|
||||
### E2. Generated files from settings (L) — firewall on/off done 2026-10-06, every machine 2026-10-09
|
||||
`firewall.enabled` in the cluster's settings reaches nodes over the agent
|
||||
stream and control planes through llmgw-firewall-sync.path. Left: the mirror
|
||||
machine and store members, the allowed set from membership, the other files.
|
||||
stream and control planes through llmgw-firewall-sync.path. Since
|
||||
2026-10-09 the mirror machine too (the leader PUTs /v1/firewall, the mirror
|
||||
writes `.firewall.json`, the same path unit applies it) and store members
|
||||
that run no llmgw (the leader keeps it in `cluster_state`, a root timer
|
||||
reads the member's own copy with psql); `firewall_networks` reach all of
|
||||
them, and the mirror lets in the network of every node and control plane
|
||||
the leader knows. Tried on the PoC: off and on, and a network added, on
|
||||
control planes, mirror, store member and nodes. Left: the allowed set
|
||||
replaced by membership rather than added to (a network taken out of the
|
||||
settings stays let in), the other files (the mirror's unit flags,
|
||||
keepalived, etcd and Patroni configuration).
|
||||
|
||||
Every machine runs the agent (root) with roles — node, control plane, mirror,
|
||||
store member — and the agent renders and applies what was written once by
|
||||
@@ -807,17 +819,288 @@ Patroni configuration. The firewall's allowed set is the cluster's membership
|
||||
removed changes every machine's rules. Retires `firewall.Admit` and
|
||||
`firewall_networks` in server.yaml (imported into the store).
|
||||
|
||||
### E3. agent.yaml down to bootstrap (M)
|
||||
### E3. agent.yaml down to bootstrap (M) — done 2026-10-08 (branch `settings-e`)
|
||||
agent.yaml keeps the control planes' addresses, its certificate and its id.
|
||||
Labels, fabrics, engines, heartbeat and supervisor settings move to the
|
||||
store per node and reach the agent over its stream; a change to the file is
|
||||
noticed and the agent restarts itself.
|
||||
Built as a section of the cluster's settings rather than a store of its own,
|
||||
so history, audit, `ctl config`, the console and the portal page cover it:
|
||||
`agents.default` and `agents.nodes.<id>` (config.AgentSettings; bootstrap
|
||||
keys node_id/server(s)/hostname/mgmt_addr/tls refused; each node's merge
|
||||
loaded as an agent would, so a bad fabric is refused on save). A live key:
|
||||
the control plane sends ControlMessage.AgentSettings on register and, a
|
||||
second apart per node, on change; the agent compares with what it started on,
|
||||
checks it loads beside its agent.yaml (else logs and keeps running), keeps it
|
||||
in /var/lib/llmgw-agent/agent-settings.yaml and restarts (engines adopted).
|
||||
At start: agent.yaml < kept settings < LLMGW_ env; kept settings that no
|
||||
longer load are set aside, never a reason not to start. agent.yaml is
|
||||
watched every 10 s and restarted on when it changes and loads. Left: file
|
||||
keys are not imported into the settings (they keep working under them);
|
||||
`node add` still writes a full agent.yaml; labels have two homes now
|
||||
(`ctl label`, which wins, and agents.*.labels).
|
||||
|
||||
### E4. Web settings page (S)
|
||||
### E4. Web settings page (S) — done 2026-10-08 (branch `settings-e`)
|
||||
The same document, edited on /admin/settings, for whoever has no terminal.
|
||||
Built: administrators only (staff roles get 404: webhook secrets), checked by
|
||||
the same function as PUT /api/v1/settings, made from a base version or
|
||||
refused on a newer one (the edit kept in the editor), history with view and
|
||||
roll back as a new version, audit with the size not the body, English and
|
||||
Turkish, a Settings link in the menu.
|
||||
|
||||
---
|
||||
|
||||
## F. Company documents: knowledge bases (RAG), then adapters (2026-10-08)
|
||||
|
||||
A company wants a model to answer from its own documents. Retrieval is the
|
||||
tool for *knowledge* (sources shown, re-indexed in minutes, access per
|
||||
tenant); LoRA training is for *behaviour* (vocabulary, format) and has no
|
||||
access control — whatever goes into weights reaches everyone who may use the
|
||||
adapter. RAG first. Decisions taken: access per collection (no per-document
|
||||
ACL yet); a personal key or session names a project or team and gets that
|
||||
chain only, never the union of its person's memberships; PostgreSQL store
|
||||
only, off on SQLite with a sentence saying why.
|
||||
|
||||
### F1. Store: collections, grants, chunks (M) — done 2026-10-08 (branch `rag`)
|
||||
Built (`internal/store/knowledge.go`, migration 0035 PostgreSQL only):
|
||||
- `kb_collections` (owner scope/id, org, embed_model, dims), `kb_grants`
|
||||
(read|write, any scope, never across organizations; a key needs a
|
||||
project; a person needs a membership in the organization),
|
||||
`kb_documents` (digest: the same again is no change, a new one replaces
|
||||
in one transaction).
|
||||
- `kb_chunks` is created when the store opens, not in a migration (a
|
||||
migration needing an extension would stop the control plane; one that
|
||||
skipped it would never run again): partitioned by organization (hashed
|
||||
partition name), `tsv` generated (`simple`), GIN on the parent, one
|
||||
partial HNSW index per collection on its organization's partition
|
||||
(`embedding::vector(dims)`, `WHERE collection = '<id>'`). EXPLAIN on
|
||||
20k rows: partition pruned, that index used, policy applied.
|
||||
Later (scale, 2026-10-08): a partial GIN per collection too
|
||||
(`kb_tsv_<id>`, backfilled at open); `vec_type` per collection, `halfvec`
|
||||
on pgvector ≥ 0.7 (up to 4000 dims); `SET LOCAL hnsw.ef_search` to
|
||||
candidates(k) clamped 40..1000 (EXPLAIN: 40 rows without it, all with
|
||||
it); the text arm takes candidates by `ts_rank` in a subquery and scores
|
||||
only those.
|
||||
- RLS FORCEd on `kb_chunks` against `llmgw.kb_org`/`llmgw.kb_collections`
|
||||
(set with set_config per transaction), one role — the table owner — not
|
||||
a separate one. Chunks are deleted by the foreign-key cascade from their
|
||||
document, which RLS does not block.
|
||||
- `CollectionsFor(Effective, write)`: owner or grant anywhere on the chain;
|
||||
no organization (personal key without a project) reads nothing.
|
||||
- `SearchChunks`: per collection nearest-neighbour (literal id and cast so
|
||||
the partial index is used) plus a full-text arm with any word
|
||||
(`a | b`, lowercased by the database, not Go), fused by reciprocal rank.
|
||||
- Deleting a project/team/org that owns collections is refused like one
|
||||
with members (changed from "deleted with them": every other delete in the
|
||||
tree refuses); DeleteUser/ForgetUser delete the person's collections;
|
||||
deleting an organization drops its partition; grants to a deleted scope go.
|
||||
- `server init` installs `postgresql-N-pgvector` / `pgvector_N` and enables
|
||||
it; pgha installs it on every member at setup and enables it on the
|
||||
leader on every run. Neither stops on failure.
|
||||
Left: a cluster set up before (the PoC) gets no package from a rerun of
|
||||
pgha on a running cluster — install by hand or add an upgrade step;
|
||||
`server restore`/`CopyTo` do not carry kb tables (SQLite → PostgreSQL only,
|
||||
and SQLite has none); tests need `LLMGW_TEST_POSTGRES` pointing at a
|
||||
pgvector database as a non-superuser (a superuser bypasses RLS).
|
||||
|
||||
### F2. Ingest on the user's machine (L) — done 2026-10-08 (branch `rag`)
|
||||
Uploading whole files to the control plane was dropped: a PDF is 20–100×
|
||||
its text, and the work belongs where the files are. Instead a small client,
|
||||
`llmgw-kb`, run by the person on their own machine with their own key:
|
||||
- **Its own binary** (`cmd/llmgw-kb`), for Windows, macOS and Linux on amd64
|
||||
and arm64. `llmgw` itself does not build on Windows (agent syscalls), and
|
||||
employees' laptops should not carry the server, agent and setup screens. A
|
||||
test keeps server packages out of its imports.
|
||||
- `llmgw-kb sync <collection> <dir|file…> [--include/--exclude globs]
|
||||
[--dry-run] [--delete]`: scan → local catalog (path, size, mtime, sha256;
|
||||
cache under os.UserCacheDir, digests not recomputed for unchanged files)
|
||||
→ the collection's manifest from the gateway (name → digest only) → diff
|
||||
→ for new/changed files only: extract, chunk, embed through the gateway's
|
||||
`/v1/embeddings` with the person's key (its accounting, limits, guard and
|
||||
band apply), and send name, digest, chunk texts and vectors; `--delete`
|
||||
removes documents whose files are gone. `--dry-run` lists what would leave
|
||||
the machine and how much. A file that cannot be read (locked, unsupported,
|
||||
no text) is skipped with the reason; the run goes on.
|
||||
- **Formats, as many as possible.** Built in, pure Go: PDF (text layer),
|
||||
DOCX/XLSX/PPTX, ODT/ODS/ODP, RTF, EPUB, HTML/XHTML, EML (text and HTML
|
||||
parts), Jupyter notebooks, subtitles (SRT/VTT), CSV/TSV, JSON/YAML/XML,
|
||||
Markdown/RST/AsciiDoc/Org, plain text and source code; ZIP archives
|
||||
recursed into. With a helper found on the machine: legacy DOC/XLS/PPT, RTF
|
||||
and anything else LibreOffice opens (`soffice --headless --convert-to`),
|
||||
images and scanned PDFs through Tesseract (`tur+eng`). Text in UTF-8,
|
||||
UTF-16 (BOM) or Windows-1254/1252 is converted to UTF-8.
|
||||
- **Across systems:** document names relative to the synced root, `/`
|
||||
separated, Unicode NFC (macOS hands out NFD), a warning for two names that
|
||||
differ only in case; CRLF → LF; key from `LLMGW_API_KEY` or `llmgw-kb
|
||||
login` (os.UserConfigDir, 0600 where that means something); the system's
|
||||
CA store, `--ca` for another.
|
||||
- **Gateway endpoints** (client listener, person's key):
|
||||
`GET /v1/kb/collections` (what this key may read/write),
|
||||
`GET /v1/kb/collections/{id}/manifest`, `PUT
|
||||
/v1/kb/collections/{id}/documents` (one document per request: name,
|
||||
digest, bytes, chunks with text and embedding; dims checked, chunk count
|
||||
and size bounded), `DELETE …/documents?name=`. Write needs a write grant
|
||||
or ownership on the key's chain (`CollectionsFor(…, write)`); otherwise
|
||||
404.
|
||||
- **Administration** stays with operators: `llmgw ctl kb
|
||||
create|list|delete|grant|revoke|docs`, admin API `/api/v1/kb`, audit,
|
||||
console screen; `docs/knowledge.md`.
|
||||
- Trade-off accepted: vectors come from the client, so a writer could store
|
||||
vectors that do not match their text — only in collections they may
|
||||
already write.
|
||||
Built: `internal/kb/extract` (formats above; helpers found on PATH and in
|
||||
each system's install places), `internal/kb/chunk` (≈1800 characters under
|
||||
the heading path, a paragraph of overlap, sections kept apart once a chunk
|
||||
is a third full), `internal/kb/kbsync` (ledger in SQLite at
|
||||
os.UserConfigDir/llmgw-kb/ledger.db: documents keyed by gateway, collection
|
||||
and name with file, size, mtime, digest, state, reason; runs), `cmd/llmgw-kb`
|
||||
(login, collections, sync, status, history, formats), `internal/router/kb.go`
|
||||
(the four gateway endpoints, `X-Llmgw-Project` for personal keys, audit),
|
||||
`internal/control/adminapi/kb.go` + `llmgw ctl kb` + console Knowledge
|
||||
screen, `docs/knowledge.md`; CI builds the client for all six targets and
|
||||
releases attach it. Tested: every format from generated fixtures and a
|
||||
Turkish PDF (OCR path with Tesseract here), sync against a fake gateway and
|
||||
end to end against the router on pgvector.
|
||||
Tried on the PoC 2026-10-08 (control planes 140 and 150 on a dev build of
|
||||
this branch; their previous binary is /usr/local/bin/llmgw.pre-rag): Ubuntu's
|
||||
pgvector 0.6.0 installed on 149/150 and enabled; bge-m3's weights were never
|
||||
in the mirror (fetched from CompendiumLabs/bge-m3-gguf@1864f47f) and it is
|
||||
held at 1 replica; tree Demo/ik/belgeler and Demo/satis/crm with a key each;
|
||||
collection Demo/politikalar filled from a Mac through an SSH tunnel with ten
|
||||
Turkish files in nine formats (9 sent in 5 s, the .doc skipped for want of
|
||||
LibreOffice; a rerun sent nothing in 0 s; a changed file alone was sent).
|
||||
Nine Turkish questions each found their page first once search scored by
|
||||
similarity plus word coverage (rank fusion had got two wrong). The other
|
||||
team's key saw no collection and got 404 by id; with a read grant it read
|
||||
and could not write; revoked, nothing again. On the live store the llmgw
|
||||
role (not superuser, no BYPASSRLS) sees 0 chunks with no scope set and 0
|
||||
with the organization alone.
|
||||
Left: the gateway's collection list shows the owner as team:ID, not its
|
||||
path; no LibreOffice on the dev machine, so the legacy
|
||||
Office path is untested beyond its error; the client is unsigned (Gatekeeper
|
||||
and SmartScreen warn); Outlook `.msg` is not read; one document over ~60 MiB
|
||||
of JSON (a few thousand pages) is refused rather than split.
|
||||
PostgreSQL tests of store, router and adminapi share LLMGW_TEST_POSTGRES and
|
||||
the store's TRUNCATE each test; run them with `go test -p 1` or one package
|
||||
at a time, or a router test loses its rows mid-run.
|
||||
|
||||
### F3. Retrieval on the request path (M) — done 2026-10-08 (branch `rag`)
|
||||
- `llmgw.knowledge: [collections]` in the body (stripped before the
|
||||
engine), or default collections per project/key.
|
||||
- Allowed set from the caller's chain (§7.9.2), cached like limits; grants
|
||||
add up down the tree. Personal keys and sessions need `x-llmgw-project`
|
||||
(or the playground's picker) naming a scope the person is a member of.
|
||||
- A collection not allowed answers 404, exactly as a missing one.
|
||||
- Hybrid search with the filter inside the query (per collection, its own
|
||||
partial index; pgvector 0.6 is enough —
|
||||
scored by similarity plus the share of the question's words held (built
|
||||
in F1/F2 as SearchChunks), `/v1/rerank` when a reranker is in the catalog. Top
|
||||
chunks into a system message with document and position;
|
||||
`x-llmgw-sources` lists them.
|
||||
- Leaks closed: answer-cache key includes the allowed set; per-organization
|
||||
`cache_salt` on RAG requests; injected passages under `no_store` and the
|
||||
guard's prompt-injection check; returned doc ids in the audit.
|
||||
- Done: router tests that tenant B never receives tenant A's chunks
|
||||
(including via cache), that a personal key without a project is refused,
|
||||
metrics (retrieval latency, chunks returned), docs.
|
||||
|
||||
Built (`internal/router/rag.go`): `X-Llmgw-Knowledge` (header, as the
|
||||
other llmgw extensions are, not a body field), `X-Llmgw-Knowledge-K`;
|
||||
collections from the key's chain (kbEffective, shared with the kb
|
||||
endpoints), 404 for one not readable; the question embedded through the
|
||||
ordinary path on the caller's principal; SearchChunks; passages trimmed to
|
||||
the model's context; added to the one leading system message (or a new
|
||||
one); `x-llmgw-sources`; before guardInput; never cached; metrics. Default collections built (kb_defaults on org/team/project/key, narrowest
|
||||
wins, 30 s cache, `ctl kb default`, console action; never refuse); the
|
||||
portal chat's Documents picker (collection, default or none; sources listed
|
||||
under the answer; a header the page sends itself is dropped). `cache_salt` per organization on requests given passages, sent to vLLM
|
||||
replicas only (chosen replica's engine). Checked on the PoC 2026-10-08:
|
||||
vLLM 0.30.0 (cuda01) with TinyLlama-1.1B AWQ run by hand beside the 35B in
|
||||
the card's last 2.3 GB — a 668-token prompt reused 656 tokens under the same
|
||||
salt and 0 under another's, with identical text (llmgw cannot place it
|
||||
there: placement takes whole GPUs). Live on the PoC: a team default answers
|
||||
a request without the header, `none` turns it off, an explicit header with
|
||||
k=2 gives two sources and no more, llmgw-kb lists owners by path. Left: returned passages
|
||||
are not audited per request (the header and the kept request record them).
|
||||
|
||||
### F4. Adapter training (L, after F1–F3) — done 2026-10-08 (branch `adapters-f4`)
|
||||
A run-to-completion job placed by the control plane like a replica and
|
||||
supervised by the agent like an engine, in a Python environment the agent
|
||||
installs (TRL + PEFT, QLoRA). Data: Q/A pairs generated from a collection by
|
||||
a cluster model through `/v1/batches`, reviewed by people. Output to the
|
||||
mirror, into the model's `adapters`, through canary and a quality suite.
|
||||
Runs when the scaling schedule frees a GPU; 8–14B dense bases on one card.
|
||||
Never train on a collection some users of the adapter may not read.
|
||||
Tried 2026-10-08 on the PoC, by hand: Qwen2.5-0.5B-Instruct (fp16) with a
|
||||
LoRA (r=16, every linear layer, 8.8M trainable), trained beside the 35B in
|
||||
cuda01's last 2.3 GB (peak 2158 MiB on the card, 64 s for 138 steps, PEFT in a
|
||||
venv over the vLLM engine's torch). Data: 90 Turkish Q/A pairs a cluster
|
||||
model wrote from the nine Demo/politikalar documents, through the gateway.
|
||||
On nine questions it had not seen, phrased otherwise: base 0/9, 6 epochs
|
||||
4/9, 4 epochs at a lower rate 2/9. Wrong answers mixed facts across
|
||||
documents (the meal card's amount, the office day). The same nine through
|
||||
retrieval: 9/9 with sources. So F4 is for format, vocabulary and tone, and
|
||||
facts come from F1–F3; a bigger base on a free card would do better but not
|
||||
change that. The training itself fits a small slot on a shared card, which
|
||||
the job design can rely on for small bases.
|
||||
|
||||
Built: `ctl kb dataset` (POST /api/v1/kb/collections/{id}/dataset, NDJSON
|
||||
with each pair's collection); train_jobs/train_artifacts (migrations 0037,
|
||||
0038 owner); internal/control/train (queue, one job per node, artifacts in
|
||||
1 MiB pieces, a node that goes fails its job); internal/agent/train (the
|
||||
newest vLLM's python3 + `uv pip install --no-deps --target` peft 0.21.2 and
|
||||
accelerate 1.15.0; llama.cpp's conversion code for the installed release,
|
||||
untarred from its source; the card with the most free memory, if enough;
|
||||
runs as llmgw-engine); `ctl train start|list|show --follow|cancel|rm|
|
||||
download|publish`; the mirror's PUT /v1/models/{model}/{version}/{file};
|
||||
publish uploads the GGUF (llama.cpp) or PEFT files (vLLM) and sets the
|
||||
adapter after checkAdapters. Owners: catalog.Adapter.Owner, enforced in the
|
||||
router (404, hidden from /v1/models); a job's owner defaults to its single
|
||||
collection's owner, must be able to read every collection the pairs came
|
||||
from, and pairs from a knowledge base are never public. Tried on the PoC:
|
||||
dataset of 90 pairs in 89 s; a job trained beside the 35B (peak 2214 MiB,
|
||||
138 steps in about a minute) and sent back PEFT and GGUF.
|
||||
Left: a job dies with its agent's stream (no resume); no console screen for
|
||||
training yet; quality gating of a published adapter is by hand (ctl eval).
|
||||
|
||||
---
|
||||
|
||||
### F5. Large imports (M) — measured 2026-10-09, open
|
||||
The metallurgy corpus (2,314 files, 108,779 chunks, bge-m3 on the one
|
||||
GPU) went in at about 25 chunks/s, the same over a phone hotspot and the
|
||||
LAN. The engine's slots were the limit: llama.cpp `--parallel 4` embeds
|
||||
four texts per pass, while llmgw-kb sends 32 a request from 6 workers. The
|
||||
usage rows show 4.6 s a request on average, 1.8 s of it in the gateway's
|
||||
queue and 0.04 s in the guard. The GPU sat at 18–50% and the store machine
|
||||
at a third of its CPUs.
|
||||
- An embedding model's profile with many slots by default (`--parallel`
|
||||
32–64, `--ctx-size` to match, `--batch-size`/`--ubatch-size` at least the
|
||||
longest chunk), or `ctl model serve` choosing that for a model with
|
||||
`capabilities.embeddings`. With the default 512, one 516-token chunk failed.
|
||||
- Embedding on the gateway's side of a document PUT, so the client sends
|
||||
text only: today an embeddings answer is ~21 KB of JSON a chunk, and
|
||||
the text and vector go up again (~11 KB).
|
||||
- The store next: HNSW inserts took half of one of pg01's two cores; at a
|
||||
million chunks the vectors alone outgrow its 4 GB. halfvec (pgvector ≥
|
||||
0.7) and a larger store machine.
|
||||
- Done: a batch of embeddings checked text by text against the context
|
||||
(branch `embeddings-batch`); it was refused as one over-long prompt.
|
||||
- Done: `query_prefix` on an embedding model, before questions only.
|
||||
- Qwen3-Embedding-4B on the RTX 6000 (2026-10-09): llama.cpp Q8 1,575
|
||||
tok/s as first set up (32k cache, 16 slots), 2,700 with a 16k cache,
|
||||
FlashAttention and ubatch 1–2k; f16 no better. vLLM fp16 (`--runner
|
||||
pooling`) 6,200 direct, 4,700–4,900 through the gateway in the import,
|
||||
and its vectors match llama.cpp's (cosine 0.9995). The 103,579 chunks
|
||||
took 3 h. Wanted: an embedding model's GPU profile on vLLM and CPU
|
||||
profile on llama.cpp, which needs a source per profile (the catalog has
|
||||
one per model; the import swapped the entry by hand), and an admission
|
||||
queue limit that suits bulk imports (30 s turned some away).
|
||||
- Retrieval on 40 questions written from the corpus, document in the top
|
||||
5: bge-m3 31/40 (Turkish questions 13/21), Qwen3-Embedding-4B 37/40
|
||||
(18/21), with its query prefix 36/40 and first place 31/40 instead of
|
||||
28 (MRR 0.738 → 0.789 → 0.838). Next: translating the question with the
|
||||
chat model, and a reranker.
|
||||
|
||||
## D. Facts worth knowing before starting
|
||||
|
||||
- Admin API: `127.0.0.1:9090`, no auth, `ctlClient` in `internal/cli/ctl.go`.
|
||||
|
||||
+314
@@ -17,11 +17,325 @@ is only released once its **exit criterion** is met, not when its code lands.
|
||||
| v0.10.0 | — | InfiniBand fabrics and NVLink, doctor, node add, engines installed by the nodes | released |
|
||||
| v0.11.0 | — | High availability from setup, the setup and console screens, usage sums and per-tenant alerts | released |
|
||||
| v0.12.0 | — | The cluster's settings in the store, in force when changed; the firewall switch; wipe | released |
|
||||
| v0.13.0 | — | The whole system stopped and started from one command; Qwen3Guard judges, on a GPU and CPUs together | released |
|
||||
| v0.14.0 | — | Knowledge bases: documents from people's own machines, kept apart per tenant, answered from | released (v0.14.1: adapters, the tree over the API, the firewall on every machine) |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
Until v1.0.0 the minor version is the breaking one: configuration formats and Go
|
||||
APIs change between minors without a deprecation period.
|
||||
|
||||
## Unreleased
|
||||
|
||||
## v0.14.1 — 2026-10-09
|
||||
|
||||
Adapters trained on the nodes from a knowledge base and kept private to
|
||||
their owners; the cluster's and the agents' settings edited in one place;
|
||||
the tenancy tree and keys over the admin API and in Terraform; knowledge
|
||||
bases that hold up with many documents; the firewall setting on every
|
||||
machine.
|
||||
|
||||
**Upgrading.**
|
||||
- **Nodes** need this release to follow `agents.*` in the cluster's
|
||||
settings; until a node is upgraded its agent.yaml is what it runs on.
|
||||
- **A separate mirror machine and store members that run nothing else**
|
||||
follow the firewall setting once what reads it is installed: run `sudo
|
||||
llmgw mirror init` on the mirror again with the flags it was set up with,
|
||||
and `llmgw server init --patroni` again for the store's members. A control
|
||||
plane set up before v0.12.0 needs `sudo llmgw firewall sync --install`.
|
||||
- **The store's PostgreSQL memory** applies after each member is restarted,
|
||||
one at a time, the leader last; until then it runs as before.
|
||||
- The store gains tables and columns; nothing to do. Existing knowledge-base
|
||||
collections get their full-text index when a control plane starts, and
|
||||
keep `vector` storage; only new collections use `halfvec`, and only on
|
||||
pgvector 0.7 or later.
|
||||
- Cached knowledge-base passages in vLLM are made again once, under the new
|
||||
`cache_salt`.
|
||||
|
||||
### Added
|
||||
|
||||
- **The cluster's settings on the portal** (BACKLOG E4): `/admin/settings`
|
||||
for administrators. It edits the document `llmgw ctl config edit` opens:
|
||||
checked as a control plane would load it, refused when made from a
|
||||
version since replaced, with the history to view and roll back. Changes
|
||||
are audited by size, never by contents.
|
||||
- **The nodes' agent settings in the cluster's settings** (BACKLOG E3):
|
||||
`agents.default` and `agents.nodes.<id>` hold everything agent.yaml can
|
||||
say but the agent's id, control planes, address and certificate. They are
|
||||
sent to the agents as they change; an agent starts again on new settings,
|
||||
or on an edited agent.yaml, with its engines serving throughout. Settings
|
||||
that would not load are refused on save, and set aside on a node rather
|
||||
than keep its agent from starting. Nodes need this release to follow
|
||||
them.
|
||||
- **LoRA adapters trained on the nodes** (BACKLOG F4): `llmgw ctl kb dataset`
|
||||
writes question/answer pairs from a collection with a model of the
|
||||
cluster. `llmgw ctl train start` trains an adapter on a node with its
|
||||
installed vLLM's torch and a small overlay, on a card with enough memory
|
||||
free (beside an engine, if it leaves that much), converting to GGUF for
|
||||
llama.cpp. `llmgw ctl train publish` keeps it on the mirror (which now
|
||||
accepts uploads) and makes it the model's adapter.
|
||||
- **Adapters with owners.** An adapter's `owner` (org, team, project, user
|
||||
or key) is the only part of the tree that sees it; to anyone else it is a
|
||||
model that does not exist. An adapter trained on a knowledge base belongs
|
||||
to a part of the tree that may read it, and is never made public.
|
||||
- The knowledge bases' vLLM `cache_salt` follows the collections searched,
|
||||
not the organization, so two teams of one organization no longer share
|
||||
cached passages.
|
||||
- **Organizations, teams, projects, API keys and memberships in the admin
|
||||
API**: `/api/v1/orgs` (the tree), `POST /api/v1/teams` and `projects`,
|
||||
`PATCH` to rename or move, `DELETE` for what is empty; `/api/v1/keys` to
|
||||
issue (the key is in the answer, once), list and revoke; `/api/v1/users`
|
||||
and `/api/v1/users/{id}/memberships`. They were only `llmgw org`, `keys`
|
||||
and `users` on a control plane before. The Terraform provider manages them
|
||||
as `llmgw_org`, `llmgw_team`, `llmgw_project`, `llmgw_key` (the key in a
|
||||
sensitive `secret`) and `llmgw_membership`.
|
||||
- **The firewall setting reaches every machine.** A separate mirror and the
|
||||
store's members follow `firewall.enabled` and `firewall_networks` as nodes
|
||||
and control planes did: the leading control plane tells the mirror over
|
||||
its API, and keeps the state in the store for members that run no llmgw,
|
||||
which a timer on each reads. The mirror also lets in the network of every
|
||||
node the leader knows. `mirror init` and `server init --patroni` install
|
||||
what follows it; run them again on a cluster set up before.
|
||||
|
||||
### Changed
|
||||
|
||||
- **The store's PostgreSQL is sized to its machine's memory.** Before, it ran
|
||||
on PostgreSQL's defaults (128MB cache) under Patroni, or on 96MB on one
|
||||
machine. Now each member gets a share of its own RAM, a smaller one where
|
||||
a control plane or agent shares the machine, and the share is read again
|
||||
every time `server init --patroni` runs. The new cache is applied by
|
||||
restarting one member at a time, the leader last.
|
||||
- **Knowledge-base searches hold up with many documents.** New collections
|
||||
store their index as `halfvec` where pgvector is 0.7 or later (half the
|
||||
memory, and embeddings up to 4000 dimensions instead of 2000); older
|
||||
collections, and clusters on pgvector 0.6, keep `vector`. Each collection
|
||||
has its own full-text index, as it has its own vector index, made for
|
||||
existing collections when the control plane starts. The text arm ranks by
|
||||
the cheap measure first and scores only those candidates.
|
||||
|
||||
### Fixed
|
||||
|
||||
- A knowledge-base search asking for more than 10 passages
|
||||
(`X-Llmgw-Knowledge-K`) got at most 40 candidates per collection from the
|
||||
vector index, whatever it asked for: pgvector's `hnsw.ef_search` is now
|
||||
raised to the number wanted, up to 1000.
|
||||
- `llmgw server backup` and `restore` against a replicated store over TLS
|
||||
failed looking for `~/.postgresql/root.crt`: pg_dump was given
|
||||
`sslmode=verify-full` without the cluster's authority.
|
||||
|
||||
## v0.14.0 — 2026-10-08
|
||||
|
||||
Knowledge bases: a company's documents, put in by people from their own
|
||||
machines, kept apart per tenant, and answered from.
|
||||
|
||||
**Upgrading.**
|
||||
- **pgvector on a cluster set up before this release.** Knowledge bases
|
||||
need pgvector, which `server init` and the high-availability database
|
||||
setup now install, but a rerun on a running cluster does not:
|
||||
1. Install `postgresql-<version>-pgvector` on every database machine.
|
||||
2. As postgres, run `CREATE EXTENSION vector` in llmgw's database.
|
||||
3. Restart the control planes.
|
||||
Until then everything else works, and `llmgw ctl kb list` says why
|
||||
knowledge bases are off. On a SQLite store they stay off.
|
||||
- **An embedding model.** A collection needs a catalog model with
|
||||
`capabilities.embeddings`, for example bge-m3, whose weights are in the
|
||||
mirror.
|
||||
- **The client.** People download `llmgw-kb` for their system from this
|
||||
release.
|
||||
|
||||
### Added
|
||||
|
||||
- **Collections and who may read them.** A collection is owned by an
|
||||
organization, team, project or person. Grants add up down the tenancy
|
||||
tree and never cross an organization. A collection a key may not read is
|
||||
answered as one that does not exist. Chunks are kept in PostgreSQL with
|
||||
pgvector, partitioned by organization, with an HNSW index per collection,
|
||||
and row-level security on the table as a second line behind every
|
||||
query's own filter. A project, team or organization that owns collections
|
||||
is not deleted; a person's collections go with their account.
|
||||
- **`llmgw-kb`**, its own download for Windows, macOS and Linux, run by
|
||||
people with their own key. It reads documents where they are and sends
|
||||
text and vectors, never the files, for what changed only; a SQLite ledger
|
||||
on the machine remembers what it did.
|
||||
- It reads PDF, Word, Excel, PowerPoint, OpenDocument, RTF, EPUB, HTML,
|
||||
MHTML, mail with attachments, notebooks, subtitles, ZIP archives, and
|
||||
text and source code in any common encoding.
|
||||
- Legacy Office files are read with LibreOffice, and scanned pages with
|
||||
Tesseract, when those are installed.
|
||||
- `--dry-run` says what would leave the machine; `--delete` removes only
|
||||
what that machine sent.
|
||||
- **Answers from documents.** A chat request names collections in
|
||||
`X-Llmgw-Knowledge`, or its part of the tree has a default (`llmgw ctl kb
|
||||
default`). The best passages, chosen by meaning plus a little for the
|
||||
question's words, are added to the system message, numbered, and
|
||||
`x-llmgw-sources` says which were used. This happens before the
|
||||
guardrails judge the request, and the answer is never served from the
|
||||
answer cache. On vLLM replicas the organization's `cache_salt` keeps
|
||||
tenants' prefix caches apart. It works through the Responses and Messages
|
||||
APIs too.
|
||||
- **Administration**: `llmgw ctl kb` (create, list, grant, revoke, default,
|
||||
docs, rm, search, delete), `/api/v1/kb`, the console's Knowledge screen,
|
||||
and a Documents picker on the portal's chat page.
|
||||
- Metrics `llmgw_kb_retrievals_total` and `llmgw_kb_retrieval_seconds`;
|
||||
[docs/knowledge.md](docs/knowledge.md).
|
||||
|
||||
## v0.13.0 — 2026-10-08
|
||||
|
||||
The whole system stopped, started and restarted from one command, a
|
||||
guardrails judge that no longer refuses ordinary coding requests, and the
|
||||
rest of the field report from v0.10.6.
|
||||
|
||||
**Upgrading.** A cluster keeps the judge its policies name: one installed
|
||||
before this release goes on asking Llama Guard 3 1B until `llmgw ctl guard
|
||||
judge use qwen3guard-0.6b`, which the leader then fetches and serves by
|
||||
itself, and which replaces the policies' Llama Guard category codes.
|
||||
|
||||
### Added
|
||||
|
||||
- **A judge on a GPU and one on CPUs together**: a policy's `judge.models`
|
||||
is a list asked in turn, the first that answers giving the verdict; the
|
||||
next is asked when one has no replica serving, fails, or misses its
|
||||
`latency_budget`. `llmgw ctl guard judge use qwen3guard-4b
|
||||
qwen3guard-0.6b` sets the order for every policy with a judge, and
|
||||
`guard judge status` says which judge stands in for which. A judge on a
|
||||
GPU is served as a `standard` model, so it takes a card only from what
|
||||
may give one up.
|
||||
- **A judge a policy names is fetched if it is missing**: the leader
|
||||
fetches the weights of a judge llmgw knows (`qwen3guard-0.6b`, `-4b`,
|
||||
`-8b`, `llama-guard-3-1b`, `-8b`) into the mirror and serves it, when a
|
||||
policy names it and the catalog does not have it — and again if it is
|
||||
removed. `guard judge enable --gpu` serves Qwen3Guard-Gen 4B on the class
|
||||
with the most free cards.
|
||||
- **The whole system stops, starts and restarts from one command**: `llmgw
|
||||
stop` takes every model down and answers every request for one with 503
|
||||
(`system_stopped`), keeping the control planes, agents and mirror running
|
||||
and no model's definition changed; `llmgw start` brings every model back.
|
||||
It is the cluster's setting `serving.stopped`, live on every control plane.
|
||||
`llmgw restart` restarts every node's agent and every control plane, the
|
||||
leader last, with the engines serving throughout. `llmgw status` says
|
||||
whether the system serves. One machine's services are `llmgw service
|
||||
status|start|stop|restart`.
|
||||
- **A replica says what it computes on**: `ctl replicas` (ON), the console
|
||||
and the admin API show `CUDA` with the GPUs its processes hold memory on
|
||||
and how much, as nvidia-smi sees them, or `CPU` — for every engine, kept by
|
||||
the agent across its own and the control plane's restarts. llama.cpp's
|
||||
layer count is added when it says it, which it does only at `-lv 4`.
|
||||
- **`ctl engine-check llamacpp` computes a token on the GPU**: a small GGUF
|
||||
(`--model`, or the smallest in the weights cache) is loaded with every layer
|
||||
on the GPU at `-lv 4` and asked for one token; a build whose CUDA backend
|
||||
does not load fails, by llama.cpp's own account of its layers. With no
|
||||
`--bin` it checks the build the agent installed.
|
||||
- **Engine versions and compute capability on nodes**: `ctl nodes` shows
|
||||
`llamacpp b10991` and `6000-24g (7.5)`; the admin API gives
|
||||
`engine_versions` and `compute_capabilities`. A profile's `engine_version`
|
||||
restricts it to nodes whose engine version starts with it, and
|
||||
explain-placement says which version a node has, or which engines.
|
||||
- **`ctl restart --max-surge N` and `--max-unavailable N`**: that many of a
|
||||
model's replicas restart at once, beside themselves or stopped first.
|
||||
- **Multi-node readiness apart from the node's**: doctor gives both, with
|
||||
fabric findings in a section of their own; `node add` says the multi-node
|
||||
one; `ctl nodes` shows a declared fabric's adapters that are down.
|
||||
- **A forgotten password is set through a link**: `/forgot` mails one when
|
||||
the cluster's `mail` settings say how (SMTP, sender, the public
|
||||
`base_url` it points at), saying the same whether or not the address has an
|
||||
account; `/reset` sets the password once and closes every session. Without
|
||||
mail, `llmgw users reset-link` makes the link for an operator to hand over.
|
||||
|
||||
### Fixed
|
||||
|
||||
- **A replica restarted in place keeps its room**: the replica was stopped
|
||||
and the pass went on to other models before placing its replacement, so a
|
||||
model waiting for room took the freed cards. On the reference cluster a
|
||||
GPU judge waiting for the only card took it from the model being
|
||||
restarted there. The replacement is now placed in the same step.
|
||||
- **The judge no longer refuses ordinary coding requests**: Llama Guard 3
|
||||
1B, the judge `server init` installed, could not tell "continue", "devam
|
||||
et" or "how do I kill all child processes in bash" (0.51 to 0.92 unsafe)
|
||||
from harmful requests (0.97 to 1.00), at any quantisation. The judge is
|
||||
now Qwen3Guard-Gen 0.6B (format `qwen3guard`), which gave the same
|
||||
ordinary requests 0.15 at the most and the harmful ones 1.00, on CPUs,
|
||||
in 0.8 GB. Its third verdict, controversial, counts as safe unless the
|
||||
policy says `controversial: unsafe`. `guard judge use qwen3guard-0.6b`
|
||||
moves an installed cluster to it, and replaces Llama Guard's category
|
||||
codes, which Qwen3Guard never names, with the baseline's.
|
||||
- **A long message is judged rather than refused**: a judge read the whole
|
||||
last user message, and one with a file pasted into it overflowed the
|
||||
judge's context, so it gave no verdict and `fail: closed` refused the
|
||||
request. It now reads the start and the end, about 6000 characters.
|
||||
- **A replica no longer stays "starting" while it serves after a restart**:
|
||||
an agent that registered before its adopted engines had answered their
|
||||
health check never said they had, so the control plane kept them starting
|
||||
and routed nothing to them — on the reference cluster, after `llmgw
|
||||
restart`, the only replica of a model. The agent now reports each shard
|
||||
the registration called not ready once it is.
|
||||
- **A failover no longer replaces healthy replicas**: a standby that took
|
||||
over placed at once, before the agents had reconnected and said what runs,
|
||||
so it started a second copy of replicas it did not know of and then drained
|
||||
the healthy ones as surplus — on the reference cluster the guardrails'
|
||||
judge and a model's replica, within twenty seconds of a leader restarting.
|
||||
Placement now waits after a takeover as it does at start.
|
||||
|
||||
### Changed
|
||||
|
||||
- **A GPU placement llama.cpp says it cannot compute on is refused** when it
|
||||
warns that it found no usable GPU, even where nvidia-smi cannot be asked;
|
||||
where it can, holding no memory on its GPUs refuses it as before.
|
||||
|
||||
## v0.12.2 — 2026-10-06
|
||||
|
||||
The guardrails turned off and on as a whole, and a judge that refuses
|
||||
ordinary requests explained: the hazard it saw is told to the client and
|
||||
kept with the request, and `llmgw ctl guard judge ask` shows what it says of
|
||||
any text. Nothing to do to upgrade; the store gains a column.
|
||||
|
||||
### Added
|
||||
|
||||
- **The guardrails turn off and on as a whole**: `llmgw ctl guard off` /
|
||||
`on`, which is `guard.enabled` in the cluster's settings — in force within
|
||||
seconds, the policies and their attachments kept. `ctl status`, `ctl guard
|
||||
list` and the guardrails page say when they are off.
|
||||
- **`llmgw ctl guard judge ask TEXT`** has the judge read a text and shows
|
||||
its answer as given and as read — for finding out why ordinary requests
|
||||
are refused. Nothing of the text is kept, not even in the audit.
|
||||
|
||||
### Changed
|
||||
|
||||
- **A judge's refusal names the hazard**: the client is told `(judge:
|
||||
S6,Specialized Advice)` rather than `(judge)`, and the verdict is kept with
|
||||
the request (`guard_rule`): `ctl request`, the server's log and the
|
||||
guardrails page show it. Llama Guard 3's codes come with their names,
|
||||
S9 Indiscriminate Weapons, without a categories list.
|
||||
|
||||
## v0.12.1 — 2026-10-06
|
||||
|
||||
The guardrails' judge, served and watched without server init: a judge
|
||||
missing, served as the wrong kind of model, or not answering is said on
|
||||
`ctl status` and the console, and `llmgw ctl guard judge enable` puts it
|
||||
right. Nothing to do to upgrade.
|
||||
|
||||
### Added
|
||||
|
||||
- **`llmgw ctl guard judge status` and `enable`**: each judge model the
|
||||
policies ask, whether it is in the catalog, as a judge, with replicas
|
||||
ready and answering — and serving the missing one, Llama Guard 3 1B on
|
||||
CPUs under the name the policies give it, with `--replace` for a model of
|
||||
that name that is not a judge. No more `server init` to put a judge back.
|
||||
Also **Do → Serve the guardrails' judge** in the console.
|
||||
- **A missing or broken judge is said where it is looked for**: `llmgw ctl
|
||||
status` and the console's Overview list each judge with a problem — not in
|
||||
the catalog, not a judge, no replica ready, or failing to answer — and
|
||||
what that does to requests (refused or unjudged).
|
||||
|
||||
### Fixed
|
||||
|
||||
- **A Llama Guard file served from the mirror page is the judge**: it was
|
||||
added as a model clients use, under whatever name the form had, so the
|
||||
policies never found it and refused everything they judged
|
||||
(`guard_unavailable`). It is now added as a judge, under the name the
|
||||
policies ask.
|
||||
- **`ctl` says so when the API answers with a page**: a command a control
|
||||
plane does not know — the leader on an older release, through a standby —
|
||||
failed with "invalid character '<'"; it now says which and why.
|
||||
|
||||
## v0.12.0 — 2026-10-06
|
||||
|
||||
A minor version, because where the configuration lives changes: everything in
|
||||
|
||||
@@ -6,6 +6,7 @@ LDFLAGS := -X '$(MODULE)/internal/version.Version=$(VERSION)'
|
||||
|
||||
build:
|
||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw ./cmd/llmgw
|
||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw-kb ./cmd/llmgw-kb
|
||||
|
||||
test:
|
||||
go test ./...
|
||||
|
||||
@@ -3,8 +3,9 @@
|
||||
A single-binary LLM gateway for clusters with multiple InfiniBand fabrics,
|
||||
NVLink domains and heterogeneous GPUs.
|
||||
|
||||
It serves LLMs running on your own GPU servers (vLLM / SGLang) behind one
|
||||
OpenAI-compatible endpoint. Design document: [llmgw-plan.md](llmgw-plan.md).
|
||||
It serves LLMs running on your own GPU and CPU servers (vLLM, SGLang,
|
||||
llama.cpp, TensorRT-LLM) behind one OpenAI- and Anthropic-compatible
|
||||
endpoint, and answers from your own documents. Design document: [llmgw-plan.md](llmgw-plan.md).
|
||||
|
||||
## Getting started
|
||||
|
||||
@@ -37,6 +38,9 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
| v0.9.0 | — | Managing people and the tenancy tree | released |
|
||||
| v0.10.0 | — | InfiniBand fabrics and NVLink, doctor, node add, engines installed by the nodes | released |
|
||||
| v0.11.0 | — | High availability from setup, the setup and console screens, usage sums and per-tenant alerts | released |
|
||||
| v0.12.0 | — | The cluster's settings in the store, in force when changed; the firewall switch; wipe | released |
|
||||
| v0.13.0 | — | The whole system stopped and started from one command; Qwen3Guard judges, on a GPU and CPUs together | released |
|
||||
| v0.14.0 | — | Knowledge bases: documents from people's own machines, kept apart per tenant, answered from | released (v0.14.1: adapters, the tree over the API, the firewall on every machine) |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
As of v0.6.0 a cluster runs unattended: agents report their hardware, the
|
||||
@@ -44,38 +48,67 @@ control plane places replicas — including ones larger than a single machine
|
||||
models load when someone asks for them and stop when nobody does, and keys,
|
||||
budgets and rate limits decide who may use any of it.
|
||||
|
||||
Since v0.11.0, not yet released (see [CHANGELOG.md](CHANGELOG.md),
|
||||
"Unreleased"):
|
||||
Since v0.11.0 (each release's detail is in [CHANGELOG.md](CHANGELOG.md)):
|
||||
|
||||
- **Security hardening** — signed releases only, a token on the admin API,
|
||||
engines not run as root, TLS for the agent port, etcd and Patroni, a
|
||||
firewall for llmgw's internal ports, encrypted backups, bounded
|
||||
connections, sign-in throttling, a Content-Security-Policy.
|
||||
- **Guardrails** — policies that check requests and answers (PII, secrets,
|
||||
prompt injection, deny patterns, a judge model such as Llama Guard), in
|
||||
observe or enforce mode, attached to any part of the tenancy tree:
|
||||
- **Adapters, and the tenancy tree over the API** (v0.14.1). `llmgw ctl
|
||||
kb dataset` and `llmgw ctl train` make a LoRA adapter from a collection
|
||||
on a node's spare GPU memory, kept on the mirror and seen only by its
|
||||
owner. Organizations, teams, projects, keys and memberships are in the
|
||||
admin API and the Terraform provider. The agents' settings are edited
|
||||
with the cluster's, on the portal too, and the firewall setting reaches
|
||||
every machine.
|
||||
- **Knowledge bases** (v0.14.0). People put a company's documents in from
|
||||
their own machines with `llmgw-kb`, a separate download for Windows,
|
||||
macOS and Linux that reads PDF, Office, OpenDocument, mail and many more
|
||||
formats and sends text and vectors, never the files. Collections are kept
|
||||
apart per tenant, down to row-level security in PostgreSQL. A chat
|
||||
request, or a part of the tree by default, is answered from them with
|
||||
its sources: [docs/knowledge.md](docs/knowledge.md).
|
||||
- **One command for the whole system** (v0.13.0): `llmgw stop`, `start` and
|
||||
`restart`. Qwen3Guard is the guardrails' judge, on a GPU and on CPUs
|
||||
together.
|
||||
- **Settings that apply when changed** (v0.12.0). The cluster's
|
||||
configuration lives in the store, the same on every control plane and in
|
||||
force within seconds (`llmgw ctl config`, the console's Settings). The
|
||||
firewall can be switched on and off for the whole cluster.
|
||||
- **Security hardening** (v0.11.2): signed releases only, a token on the
|
||||
admin API, engines not run as root, TLS for the agent port, etcd and
|
||||
Patroni, a firewall for llmgw's internal ports, encrypted backups,
|
||||
bounded connections, sign-in throttling, a Content-Security-Policy.
|
||||
- **Guardrails** (v0.11.2). Policies check requests and answers for PII,
|
||||
secrets, prompt injection, deny patterns, and with a judge model; they
|
||||
observe or enforce, and attach to any part of the tenancy tree:
|
||||
[docs/guardrails.md](docs/guardrails.md).
|
||||
- **Administration** — an audit log of every change (`ctl audit`), staff
|
||||
roles (viewer, operator, billing, admin), single sign-on through OpenID
|
||||
Connect and provisioning through SCIM 2.0, key lifetime limits and
|
||||
rotation with overlap (`keys ttl`, `keys rotate`), usage export as CSV
|
||||
(`ctl usage export`).
|
||||
- **Notifications** — webhooks, signed or Slack-shaped, for a model parked,
|
||||
a node lost, a budget spent and new weights published (`notify.webhooks`,
|
||||
`ctl notify test`), and Alertmanager receiver templates.
|
||||
- **Traffic and models** — request tags, canaries of new weights with
|
||||
automatic rollback, quality suites gating updates, a response cache,
|
||||
kept and replayed requests, batches, Anthropic's Messages API, more of
|
||||
OpenAI's (structured output, images, rerank, audio, moderations),
|
||||
scaling by the clock, SLO-aware admission, TensorRT-LLM, MIG, models
|
||||
imported from Ollama, routing by what each vLLM replica's cache holds,
|
||||
LoRA adapters loaded on demand, other llmgw clusters as providers, and
|
||||
placement that knows racks, PDUs and power caps.
|
||||
- **Operations** — `server restore` and off-machine backup copies, secrets
|
||||
in Vault, forgetting a person and keeping a project's requests in the
|
||||
cluster, offline installs from a bundle, a Helm chart, tracing from the
|
||||
control plane to the engines, the admin API in OpenAPI with a Terraform
|
||||
provider, and the portal in English and Turkish with its docs at /docs.
|
||||
- **Administration** (v0.11.2):
|
||||
- an audit log of every change (`ctl audit`);
|
||||
- staff roles (viewer, operator, billing, admin);
|
||||
- single sign-on through OpenID Connect and provisioning through SCIM
|
||||
2.0;
|
||||
- key lifetime limits, and rotation with overlap (`keys ttl`, `keys
|
||||
rotate`);
|
||||
- usage export as CSV (`ctl usage export`).
|
||||
- **Notifications** (v0.11.2): webhooks, signed or Slack-shaped, for a
|
||||
model parked, a node lost, a budget spent and new weights published, and
|
||||
Alertmanager receiver templates.
|
||||
- **Traffic and models** (v0.11.2):
|
||||
- request tags, canaries of new weights with automatic rollback, and
|
||||
quality suites gating updates;
|
||||
- a response cache, kept and replayed requests, and batches;
|
||||
- Anthropic's Messages API, and more of OpenAI's (structured output,
|
||||
images, rerank, audio, moderations);
|
||||
- scaling by the clock, and SLO-aware admission;
|
||||
- TensorRT-LLM, MIG, and models imported from Ollama;
|
||||
- routing by what each vLLM replica's cache holds, and LoRA adapters
|
||||
loaded on demand;
|
||||
- other llmgw clusters as providers, and placement that knows racks,
|
||||
PDUs and power caps.
|
||||
- **Operations** (v0.11.2):
|
||||
- `server restore`, off-machine backup copies, and secrets in Vault;
|
||||
- forgetting a person, and keeping a project's requests in the cluster;
|
||||
- offline installs from a bundle, and a Helm chart;
|
||||
- tracing from the control plane to the engines;
|
||||
- the admin API in OpenAPI, with a Terraform provider;
|
||||
- the portal in English and Turkish, with its docs at /docs.
|
||||
|
||||
Every one of these is in the console as well as on the command line:
|
||||
[docs/console.md](docs/console.md).
|
||||
@@ -124,6 +157,9 @@ Every page is also in the binary: `llmgw docs`, or **Manual** in the console.
|
||||
official OpenAI and Anthropic SDKs, and llmgw's extensions
|
||||
- [docs/guardrails.md](docs/guardrails.md) — checking requests and answers:
|
||||
policies, detectors, a judge model, and where each applies
|
||||
- [docs/knowledge.md](docs/knowledge.md) — knowledge bases: collections,
|
||||
who may read them, `llmgw-kb` on people's machines, and answers from
|
||||
documents
|
||||
- [docs/choosing-models.md](docs/choosing-models.md) — which models to run,
|
||||
how many of each, and where the weights should live
|
||||
|
||||
@@ -152,7 +188,7 @@ Every page is also in the binary: `llmgw docs`, or **Manual** in the console.
|
||||
## Building
|
||||
|
||||
```sh
|
||||
make build # bin/llmgw
|
||||
make build # bin/llmgw and bin/llmgw-kb
|
||||
make test # unit tests
|
||||
make race # with the race detector
|
||||
```
|
||||
@@ -169,6 +205,7 @@ llmgw ctl cordon gpu-a-01 # stop new placement, leave what runs
|
||||
sudo llmgw doctor # on a node: what its GPUs, NVLink and fabric need
|
||||
llmgw ctl audit --since 24h # who changed what
|
||||
llmgw ctl usage export --from 2026-09-01 --to 2026-10-01 --by project,model
|
||||
llmgw ctl kb list # knowledge-base collections, grants, defaults
|
||||
```
|
||||
|
||||
Bare `llmgw` opens the console, which does all of this from one screen.
|
||||
|
||||
+911
-235
File diff suppressed because it is too large.
Load diff
@@ -28,6 +28,8 @@ message AgentMessage {
|
||||
FabricTestReply fabric_test = 6;
|
||||
UpgradeReply upgrade = 7;
|
||||
KVEvents kv_events = 8;
|
||||
TrainStatus train_status = 9;
|
||||
TrainArtifact train_artifact = 10;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,6 +74,10 @@ message ControlMessage {
|
||||
UpgradeAgent upgrade = 10;
|
||||
ControlPlanes control_planes = 11;
|
||||
Firewall firewall = 13;
|
||||
RestartAgent restart_agent = 14;
|
||||
AgentSettings agent_settings = 15;
|
||||
TrainStart train_start = 16;
|
||||
TrainCancel train_cancel = 17;
|
||||
}
|
||||
// trace is the W3C trace context (traceparent, and tracestate when there
|
||||
// is one) of the control plane's span that sent this command, so what the
|
||||
@@ -81,6 +87,12 @@ message ControlMessage {
|
||||
map<string, string> trace = 12;
|
||||
}
|
||||
|
||||
// RestartAgent has the agent start again, as llmgw restart asks of every
|
||||
// node: its engines keep serving and are taken over again.
|
||||
message RestartAgent {
|
||||
string reason = 1;
|
||||
}
|
||||
|
||||
// Firewall is whether llmgw's firewall is kept on the node — its nftables
|
||||
// table, which lets the cluster's networks alone reach the engines' ports and
|
||||
// the node exporter — sent when an agent registers and whenever the
|
||||
@@ -96,6 +108,17 @@ message Firewall {
|
||||
int32 engine_ports_to = 4;
|
||||
}
|
||||
|
||||
// AgentSettings is what the cluster's settings say of this node's agent
|
||||
// (agents.default, with agents.nodes.<id> over it), sent when it registers
|
||||
// and whenever they change. The agent keeps it beside agent.yaml, which
|
||||
// holds only what it needs to start; over that file, it is what the agent
|
||||
// runs on. A change from what it runs on has the agent start again on it,
|
||||
// its engines serving throughout.
|
||||
message AgentSettings {
|
||||
// yaml is agent.yaml's keys, as YAML; empty for none.
|
||||
string yaml = 1;
|
||||
}
|
||||
|
||||
// ControlPlanes is every control plane an agent may dial, sent by the leader
|
||||
// when an agent registers and whenever the list changes. The agent keeps it
|
||||
// beside its configuration, so a second control plane added after the node
|
||||
@@ -266,6 +289,11 @@ message RunningShard {
|
||||
// restart could not be compared with the catalog — which is exactly when the
|
||||
// catalog is most likely to have been edited.
|
||||
string profile_fp = 15;
|
||||
|
||||
// compute is where its engine said it put the model when it became ready,
|
||||
// kept by the agent across its own restarts: a control plane that takes
|
||||
// over, or restarts, learns it here.
|
||||
Compute compute = 16;
|
||||
}
|
||||
|
||||
message GPU {
|
||||
@@ -605,6 +633,18 @@ message ProvisionStatus {
|
||||
string detail = 5;
|
||||
}
|
||||
|
||||
// Compute is where an engine runs its model, as it said when it loaded it.
|
||||
message Compute {
|
||||
// backend is "CUDA", "Vulkan", "Metal" or "CPU".
|
||||
string backend = 1;
|
||||
// devices are the devices it named: "CUDA0: Quadro RTX 6000 (7.5)".
|
||||
repeated string devices = 2;
|
||||
// layers_offloaded of layers_total went to the devices; both 0 when the
|
||||
// engine did not say.
|
||||
int32 layers_offloaded = 3;
|
||||
int32 layers_total = 4;
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
message ShardStatus {
|
||||
string node_id = 1;
|
||||
@@ -626,6 +666,10 @@ message ShardStatus {
|
||||
// going to look.
|
||||
FetchProgress fetch = 8;
|
||||
|
||||
// compute is what the engine says it computes on, read from its own output
|
||||
// as it became ready: sent with STATE_READY when the engine says.
|
||||
Compute compute = 9;
|
||||
|
||||
enum State {
|
||||
STATE_UNSPECIFIED = 0;
|
||||
STATE_PENDING = 1;
|
||||
@@ -844,3 +888,55 @@ message LoraAdapter {
|
||||
string name = 1;
|
||||
string source = 2;
|
||||
}
|
||||
|
||||
// TrainStart has a node train a LoRA adapter (BACKLOG F4): on a base model's
|
||||
// weights from its registry, on question/answer pairs, with a Python trainer
|
||||
// the agent carries, in an environment it makes over an installed engine's.
|
||||
// One at a time per node.
|
||||
message TrainStart {
|
||||
string job_id = 1;
|
||||
// base_repo and base_revision name the base model's weights on its
|
||||
// registry (Hugging Face): safetensors, not the GGUF a catalog may serve.
|
||||
string base_repo = 2;
|
||||
string base_revision = 3;
|
||||
// data is the training pairs, NDJSON of {"q", "a"}.
|
||||
bytes data = 4;
|
||||
// params is the trainer's settings as JSON: rank, alpha, epochs, lr,
|
||||
// batch, max_len.
|
||||
string params = 5;
|
||||
// gpu_memory_mib is what the job needs of a card; the agent starts it
|
||||
// only when one has that much free, which lets it run beside an engine.
|
||||
int32 gpu_memory_mib = 6;
|
||||
// gguf asks for the adapter in llama.cpp's format too.
|
||||
bool gguf = 7;
|
||||
}
|
||||
|
||||
// TrainCancel stops a training job.
|
||||
message TrainCancel {
|
||||
string job_id = 1;
|
||||
}
|
||||
|
||||
// TrainStatus is how a training job is going, sent as it changes.
|
||||
message TrainStatus {
|
||||
string job_id = 1;
|
||||
// state is preparing, running, done or failed.
|
||||
string state = 2;
|
||||
int32 step = 3;
|
||||
int32 steps = 4;
|
||||
double loss = 5;
|
||||
// message is what is being done, or why it failed.
|
||||
string message = 6;
|
||||
// peak_mib is the most the job held of the card.
|
||||
int32 peak_mib = 7;
|
||||
}
|
||||
|
||||
// TrainArtifact is a piece of what a training job made, sent before its
|
||||
// final status: a gzipped tar of the PEFT adapter (kind peft), or the GGUF
|
||||
// adapter (kind gguf).
|
||||
message TrainArtifact {
|
||||
string job_id = 1;
|
||||
string kind = 2;
|
||||
int32 seq = 3;
|
||||
bytes data = 4;
|
||||
bool last = 5;
|
||||
}
|
||||
@@ -0,0 +1,467 @@
|
||||
// Command llmgw-kb puts a person's documents into an llmgw knowledge base
|
||||
// from their own machine (BACKLOG F2): it reads the files there, cuts them
|
||||
// into chunks, has the gateway embed them with the person's key, and sends
|
||||
// the text and vectors — only for what changed since the last run, which a
|
||||
// ledger on the machine remembers.
|
||||
//
|
||||
// Its own binary, for Windows, macOS and Linux: it carries none of the
|
||||
// server, agent or setup code that llmgw does.
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"golang.org/x/term"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/extract"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||
)
|
||||
|
||||
func main() {
|
||||
if err := root().Execute(); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "error:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
// config is what `login` keeps: the gateway, the key, and for a personal key
|
||||
// the project it acts in.
|
||||
type config struct {
|
||||
Gateway string `json:"gateway"`
|
||||
Key string `json:"key"`
|
||||
Project string `json:"project,omitempty"`
|
||||
CA string `json:"ca,omitempty"`
|
||||
}
|
||||
|
||||
func configPath() (string, error) {
|
||||
dir, err := os.UserConfigDir()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Join(dir, "llmgw-kb", "config.json"), nil
|
||||
}
|
||||
|
||||
// loadConfig reads the saved configuration, with the environment over it:
|
||||
// LLMGW_URL, LLMGW_API_KEY, LLMGW_PROJECT.
|
||||
func loadConfig() (config, error) {
|
||||
var c config
|
||||
if p, err := configPath(); err == nil {
|
||||
if b, err := os.ReadFile(p); err == nil {
|
||||
if err := json.Unmarshal(b, &c); err != nil {
|
||||
return c, fmt.Errorf("%s: %w", p, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
for env, dst := range map[string]*string{"LLMGW_URL": &c.Gateway, "LLMGW_API_KEY": &c.Key, "LLMGW_PROJECT": &c.Project} {
|
||||
if v := strings.TrimSpace(os.Getenv(env)); v != "" {
|
||||
*dst = v
|
||||
}
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func client(flags *globalFlags) (*kbsync.Client, error) {
|
||||
c, err := loadConfig()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if flags.gateway != "" {
|
||||
c.Gateway = flags.gateway
|
||||
}
|
||||
if flags.project != "" {
|
||||
c.Project = flags.project
|
||||
}
|
||||
if flags.ca != "" {
|
||||
c.CA = flags.ca
|
||||
}
|
||||
if c.Gateway == "" || c.Key == "" {
|
||||
return nil, errors.New("no gateway or key: run `llmgw-kb login`, or set LLMGW_URL and LLMGW_API_KEY")
|
||||
}
|
||||
h, err := kbsync.NewHTTP(c.CA)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &kbsync.Client{Base: c.Gateway, Key: c.Key, Project: c.Project, HTTP: h}, nil
|
||||
}
|
||||
|
||||
type globalFlags struct {
|
||||
gateway, project, ca, ledger string
|
||||
}
|
||||
|
||||
func (g *globalFlags) openLedger(ctx context.Context) (*kbsync.Ledger, error) {
|
||||
p := g.ledger
|
||||
if p == "" {
|
||||
var err error
|
||||
if p, err = kbsync.DefaultLedgerPath(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return kbsync.OpenLedger(ctx, p)
|
||||
}
|
||||
|
||||
func root() *cobra.Command {
|
||||
g := &globalFlags{}
|
||||
cmd := &cobra.Command{
|
||||
Use: "llmgw-kb",
|
||||
Short: "Put your documents into an llmgw knowledge base, from this machine",
|
||||
Long: `llmgw-kb reads documents on this machine, cuts them into chunks, has the
|
||||
gateway embed them with your key, and sends the text and vectors to a
|
||||
knowledge base. Files are never uploaded. Run it as often as you like: a
|
||||
ledger on this machine remembers what was sent, and only what changed is
|
||||
read and sent again.`,
|
||||
Version: version.Version,
|
||||
SilenceUsage: true,
|
||||
SilenceErrors: true,
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&g.gateway, "gateway", "", "the gateway's address (default: from login, or LLMGW_URL)")
|
||||
cmd.PersistentFlags().StringVar(&g.project, "project", "", "the project a personal key acts in (default: from login, or LLMGW_PROJECT)")
|
||||
cmd.PersistentFlags().StringVar(&g.ca, "ca", "", "a PEM file of certificates to trust besides the system's")
|
||||
cmd.PersistentFlags().StringVar(&g.ledger, "ledger", "", "the ledger file (default: in your configuration directory)")
|
||||
cmd.AddCommand(loginCmd(g), collectionsCmd(g), syncCmd(g), statusCmd(g), historyCmd(g), formatsCmd())
|
||||
return cmd
|
||||
}
|
||||
|
||||
func loginCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "login",
|
||||
Short: "Save the gateway's address and your key",
|
||||
Long: `Asks for the gateway's address and your API key, checks them, and saves them
|
||||
in your configuration directory, readable by you only. The key is typed
|
||||
without being shown, so it does not end up in your shell's history.`,
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
c, _ := loadConfig()
|
||||
in := bufio.NewReader(os.Stdin)
|
||||
ask := func(q, def string) string {
|
||||
if def != "" {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s [%s]: ", q, def)
|
||||
} else {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: ", q)
|
||||
}
|
||||
s, _ := in.ReadString('\n')
|
||||
if s = strings.TrimSpace(s); s == "" {
|
||||
return def
|
||||
}
|
||||
return s
|
||||
}
|
||||
c.Gateway = strings.TrimRight(ask("Gateway address (https://…)", or(g.gateway, c.Gateway)), "/")
|
||||
fmt.Fprint(cmd.OutOrStdout(), "API key: ")
|
||||
var key string
|
||||
if term.IsTerminal(int(os.Stdin.Fd())) {
|
||||
b, err := term.ReadPassword(int(os.Stdin.Fd()))
|
||||
fmt.Fprintln(cmd.OutOrStdout())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key = strings.TrimSpace(string(b))
|
||||
} else {
|
||||
key, _ = in.ReadString('\n')
|
||||
key = strings.TrimSpace(key)
|
||||
}
|
||||
if key != "" {
|
||||
c.Key = key
|
||||
}
|
||||
c.Project = ask("Project, for a personal key (empty for a project's key)", or(g.project, c.Project))
|
||||
if g.ca != "" {
|
||||
c.CA = g.ca
|
||||
}
|
||||
h, err := kbsync.NewHTTP(c.CA)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cl := &kbsync.Client{Base: c.Gateway, Key: c.Key, Project: c.Project, HTTP: h, Wait: time.Second}
|
||||
ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
|
||||
defer cancel()
|
||||
cs, err := cl.Collections(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("the gateway refused: %w", err)
|
||||
}
|
||||
p, err := configPath()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0o700); err != nil {
|
||||
return err
|
||||
}
|
||||
b, _ := json.MarshalIndent(c, "", " ")
|
||||
if err := os.WriteFile(p, b, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "Saved in %s. This key reaches %d collection(s); `llmgw-kb collections` lists them.\n", p, len(cs))
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func or(a, b string) string {
|
||||
if a != "" {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func collectionsCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "collections",
|
||||
Short: "List the collections your key may read, and which it may add to",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cs, err := c.Collections(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(cs) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), "No collections. An administrator creates them and grants access (llmgw ctl kb).")
|
||||
return nil
|
||||
}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "NAME\tACCESS\tOWNER\tEMBEDDED BY")
|
||||
for _, x := range cs {
|
||||
access := "read"
|
||||
if x.Writable {
|
||||
access = "read, write"
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%s\t%s (%d)\n", x.Name, access, or(x.OwnerName, x.Owner), x.EmbedModel, x.Dims)
|
||||
}
|
||||
return tw.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func syncCmd(g *globalFlags) *cobra.Command {
|
||||
var o kbsync.Options
|
||||
var verbose bool
|
||||
var ocr string
|
||||
var maxMiB int64
|
||||
cmd := &cobra.Command{
|
||||
Use: "sync COLLECTION PATH...",
|
||||
Short: "Send a collection what changed in these folders and files",
|
||||
Long: `Reads the files under each PATH that the collection does not already have as
|
||||
they are, and sends their text and vectors. Run it again whenever you like;
|
||||
unchanged files are not even read.
|
||||
|
||||
Documents are named by their path under PATH, with forward slashes; with
|
||||
several PATHs, each folder's own name comes first.
|
||||
|
||||
--dry-run reads and cuts everything, sends nothing, and says what would
|
||||
leave this machine
|
||||
--delete removes documents this machine sent whose files are gone
|
||||
(only under these PATHs; nothing anybody else sent)
|
||||
--retry reads again files skipped before (after installing LibreOffice
|
||||
or Tesseract, say)
|
||||
|
||||
Without --include, every file of a known format is read (llmgw-kb formats);
|
||||
hidden files, version control folders and Office lock files are not.`,
|
||||
Args: cobra.MinimumNArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), os.Interrupt)
|
||||
defer stop()
|
||||
l, err := g.openLedger(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
o.Collection, o.Roots = args[0], args[1:]
|
||||
o.Extract = extract.Options{Helpers: extract.FindHelpers(), OCRLanguages: ocr, MaxBytes: maxMiB << 20}
|
||||
out := cmd.OutOrStdout()
|
||||
c.Waiting = func(why string, d time.Duration) {
|
||||
fmt.Fprintf(out, " .. waiting %s: %s\n", d, why)
|
||||
}
|
||||
start := time.Now()
|
||||
rep, err := kbsync.Sync(ctx, c, l, o, func(e kbsync.Event) {
|
||||
switch e.Kind {
|
||||
case "unchanged":
|
||||
if !verbose {
|
||||
return
|
||||
}
|
||||
fmt.Fprintf(out, " = %s\n", e.Name)
|
||||
case "sent", "would-send":
|
||||
fmt.Fprintf(out, " %-8s %s (%d chunks, %s)\n", label(e.Kind), e.Name, e.Chunks, size(e.Bytes))
|
||||
default:
|
||||
line := fmt.Sprintf(" %-8s %s", label(e.Kind), e.Name)
|
||||
if e.Detail != "" {
|
||||
line += ": " + e.Detail
|
||||
}
|
||||
fmt.Fprintln(out, line)
|
||||
}
|
||||
})
|
||||
if rep.Collection.Name != "" {
|
||||
verb := "sent"
|
||||
if o.DryRun {
|
||||
verb = "would send"
|
||||
}
|
||||
fmt.Fprintf(out, "\n%s: %d files; %s %d (%d chunks, %s), %d unchanged, %d skipped, %d failed, %d deleted, in %s\n",
|
||||
rep.Collection.Name, rep.Files, verb, rep.Sent, rep.Chunks, size(rep.Bytes), rep.Unchanged,
|
||||
rep.Skipped, rep.Failed, rep.Deleted, time.Since(start).Round(time.Second))
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rep.Failed > 0 {
|
||||
return fmt.Errorf("%d document(s) failed; run again to retry them", rep.Failed)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
f := cmd.Flags()
|
||||
f.StringSliceVar(&o.Include, "include", nil, "read only files matching these globs (\"*.pdf,docs/**/*.md\")")
|
||||
f.StringSliceVar(&o.Exclude, "exclude", nil, "never read files matching these globs")
|
||||
f.BoolVar(&o.Hidden, "hidden", false, "read hidden files and folders too")
|
||||
f.BoolVar(&o.DryRun, "dry-run", false, "send nothing; say what would be sent")
|
||||
f.BoolVar(&o.Delete, "delete", false, "remove documents this machine sent whose files are gone")
|
||||
f.BoolVar(&o.Retry, "retry", false, "read again files skipped before")
|
||||
f.IntVar(&o.Workers, "workers", 4, "files read and sent at once")
|
||||
f.StringVar(&ocr, "ocr-languages", "tur+eng", "Tesseract's languages for images and scanned pages")
|
||||
f.Int64Var(&maxMiB, "max-size", 256, "skip files larger than this many MiB")
|
||||
f.BoolVarP(&verbose, "verbose", "v", false, "list unchanged files too")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func label(kind string) string {
|
||||
switch kind {
|
||||
case "would-send":
|
||||
return "send"
|
||||
case "would-delete":
|
||||
return "delete"
|
||||
}
|
||||
return kind
|
||||
}
|
||||
|
||||
func size(b int64) string {
|
||||
switch {
|
||||
case b >= 1<<20:
|
||||
return fmt.Sprintf("%.1f MiB", float64(b)/(1<<20))
|
||||
case b >= 1<<10:
|
||||
return fmt.Sprintf("%.0f KiB", float64(b)/(1<<10))
|
||||
}
|
||||
return fmt.Sprintf("%d B", b)
|
||||
}
|
||||
|
||||
func statusCmd(g *globalFlags) *cobra.Command {
|
||||
var all bool
|
||||
return &cobra.Command{
|
||||
Use: "status COLLECTION",
|
||||
Short: "Show what this machine has sent to a collection, and what it skipped",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
m, err := c.Manifest(cmd.Context(), args[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
l, err := g.openLedger(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
es, err := l.Entries(cmd.Context(), strings.TrimRight(c.Base, "/"), m.Collection.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
counts := map[kbsync.State]int{}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "STATE\tDOCUMENT\tCHUNKS\tWHEN\tWHY")
|
||||
for _, e := range es {
|
||||
counts[e.State]++
|
||||
if e.State == kbsync.Sent && !all {
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%d\t%s\t%s\n", e.State, e.Name, e.Chunks, e.At.Format("2006-01-02 15:04"), e.Reason)
|
||||
}
|
||||
_ = tw.Flush()
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n%s holds %d documents; from this machine %d sent, %d skipped, %d failed.",
|
||||
m.Collection.Name, len(m.Documents), counts[kbsync.Sent], counts[kbsync.Skipped], counts[kbsync.Failed])
|
||||
if !all {
|
||||
fmt.Fprint(cmd.OutOrStdout(), " (--all lists the sent ones too)")
|
||||
}
|
||||
fmt.Fprintln(cmd.OutOrStdout())
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func historyCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "history",
|
||||
Short: "List the latest syncs from this machine",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
l, err := g.openLedger(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
runs, err := l.Runs(cmd.Context(), 20)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "STARTED\tTOOK\tCOLLECTION\tSENT\tUNCHANGED\tSKIPPED\tFAILED\tDELETED")
|
||||
for _, r := range runs {
|
||||
took := "unfinished"
|
||||
if !r.Finished.IsZero() {
|
||||
took = r.Finished.Sub(r.Started).Round(time.Second).String()
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%s\t%d\t%d\t%d\t%d\t%d\n", r.Started.Format("2006-01-02 15:04"), took,
|
||||
r.Collection, r.Sent, r.Unchanged, r.Skipped, r.Failed, r.Delete)
|
||||
}
|
||||
return tw.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func formatsCmd() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "formats",
|
||||
Short: "List the formats read, and the helpers found on this machine",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
builtIn, helper := extract.Supported()
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintf(out, "Read here:\n%s\n\n", wrap(builtIn))
|
||||
fmt.Fprintf(out, "Read with a helper program:\n%s\n\n", wrap(helper))
|
||||
fmt.Fprintln(out, "Helpers:")
|
||||
for _, l := range extract.FindHelpers().Describe() {
|
||||
fmt.Fprintln(out, " "+l)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func wrap(words []string) string {
|
||||
var b strings.Builder
|
||||
line := " "
|
||||
for _, w := range words {
|
||||
if len(line)+len(w) > 78 {
|
||||
b.WriteString(strings.TrimRight(line, " ") + "\n")
|
||||
line = " "
|
||||
}
|
||||
line += w + " "
|
||||
}
|
||||
b.WriteString(strings.TrimRight(line, " "))
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The client goes onto employees' laptops: none of the server, agent, store
|
||||
// or setup code may come with it, and it must build where llmgw does not.
|
||||
func TestTheClientCarriesNoServerCode(t *testing.T) {
|
||||
out, err := exec.Command("go", "list", "-deps", ".").Output()
|
||||
if err != nil {
|
||||
t.Skipf("go list: %v", err)
|
||||
}
|
||||
for _, dep := range strings.Fields(string(out)) {
|
||||
for _, banned := range []string{"/internal/store", "/internal/router", "/internal/agent", "/internal/control",
|
||||
"/internal/cli", "/internal/guard", "github.com/jackc/pgx"} {
|
||||
if strings.Contains(dep, banned) {
|
||||
t.Errorf("llmgw-kb depends on %s", dep)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -39,7 +39,7 @@ func (d Doc) Source() string {
|
||||
// than with either reader, because the console's manual and the gateway's
|
||||
// docs portal both list them, and two lists drift.
|
||||
var docOrder = []string{"console", "quick-install", "install", "node-onboarding", "onboarding",
|
||||
"clients", "guardrails", "choosing-models", "fabrics", "facility-aware-placement", "data",
|
||||
"clients", "guardrails", "knowledge", "choosing-models", "fabrics", "facility-aware-placement", "data",
|
||||
"reference", "development"}
|
||||
|
||||
// Manual is every document, in reading order: those docOrder names, then any
|
||||
|
||||
+16
-7
@@ -8,8 +8,9 @@ description is.
|
||||
anything they do can be done with `curl`, a generated client or the
|
||||
[Terraform provider](../terraform-provider-llmgw/README.md). It adds and
|
||||
removes models, scales them, cordons and upgrades nodes, sets capacity
|
||||
quotas, attaches guardrail policies, and reads the fleet, usage and the audit
|
||||
log.
|
||||
quotas, attaches guardrail policies, adds organizations, teams and projects,
|
||||
issues and revokes API keys, puts people in the tree, and reads the fleet,
|
||||
usage and the audit log.
|
||||
|
||||
## Reaching it
|
||||
|
||||
@@ -99,6 +100,14 @@ curl -s "${H[@]}" -X POST $API/models/qwen-small/scale -d '{"min": 0, "max": 2,
|
||||
# A team's capacity quota.
|
||||
curl -s "${H[@]}" -X PUT $API/quotas -d '{"scope": "team:Acme/ml", "limit": {"gpus": 8}}'
|
||||
|
||||
# An organization, a team in it and a project in that; then a key for the
|
||||
# project. The answer holds the key once: only its hash is kept.
|
||||
curl -s "${H[@]}" -X POST $API/orgs -d '{"name": "Acme"}'
|
||||
curl -s "${H[@]}" -X POST $API/teams -d '{"name": "ml", "in": "Acme"}'
|
||||
curl -s "${H[@]}" -X POST $API/projects -d '{"name": "search", "in": "Acme/ml"}'
|
||||
curl -s "${H[@]}" -X POST $API/keys -d '{"name": "ci", "project": "Acme/ml/search", "expires_in": "720h"}'
|
||||
curl -s "${H[@]}" -X POST $API/users/ayse@example.com/memberships -d '{"target": "team:Acme/ml"}'
|
||||
|
||||
# Remove the model; a DELETE carries the content type too.
|
||||
curl -s "${H[@]}" -X DELETE $API/models/qwen-small
|
||||
```
|
||||
@@ -149,11 +158,11 @@ go generate ./internal/docgen
|
||||
|
||||
## What it does not cover
|
||||
|
||||
- **Organizations, teams, projects, people and API keys.** These live in the
|
||||
store and are managed with `llmgw org`, `llmgw users` and `llmgw keys` on the
|
||||
control plane, or in the web console's `/admin` pages on the client port
|
||||
behind a session. The admin API reads the tree only where a quota or a
|
||||
policy names a part of it.
|
||||
- **Making people's accounts.** People sign up, are invited, or come from
|
||||
the identity provider (OIDC, SCIM); the API lists them and puts them in
|
||||
the tree, and `llmgw users` approves, disables and deletes them. Key
|
||||
lifetime limits and spending limits are `llmgw keys ttl` and `llmgw keys
|
||||
limits`.
|
||||
- **The client-facing API** (`/v1/chat/completions` and the rest), which is
|
||||
OpenAI's and Anthropic's, on `listen`, with API keys.
|
||||
- **The rest of the admin listener**: `/metrics`, `/health`, the dashboard
|
||||
|
||||
+2595
-2
File diff suppressed because it is too large.
Load diff
@@ -164,6 +164,15 @@ which): requests, tokens, errors and spend.
|
||||
team or project is guaranteed and may at most hold, and what it holds now.
|
||||
**A** adds one; on a row, change or remove it.
|
||||
|
||||
**Knowledge** lists the knowledge-base collections: each one's owner, the
|
||||
model that embeds it, how many documents and chunks it holds and who else it
|
||||
is granted to. **A** makes one. On a row: its **Documents**, **Grant…** a
|
||||
project, team, organization, person or key read or write access, **Revoke**
|
||||
a grant, **Answer from it by default…** for a part of the tree, **Try a
|
||||
search…** as an administrator, or **Delete…** it. People
|
||||
fill collections from their own machines with `llmgw-kb`
|
||||
([knowledge.md](knowledge.md)). Off on a SQLite store, and says why.
|
||||
|
||||
**Mirror** is what the model mirror holds, one row for each version of each
|
||||
model — a file, or the set of files a model of several is: shards, a vision
|
||||
projector, a config and tokenizer — with how far its download has got and
|
||||
|
||||
+124
-16
@@ -4,7 +4,7 @@ Written for the operator deciding what the gateway lets through: what a
|
||||
policy can check, what it can do, and how one is attached to a team.
|
||||
|
||||
`llmgw server init` writes a `baseline` policy that applies to every
|
||||
request, with Llama Guard 3 1B as its judge on a CPU; see
|
||||
request, with Qwen3Guard-Gen 0.6B as its judge on CPUs; see
|
||||
[install.md](install.md#2-the-control-plane) for what it checks. The
|
||||
policies are part of the cluster's settings, kept in the store: `llmgw ctl
|
||||
config edit` (or the console's Settings) changes them, and the change is in
|
||||
@@ -15,6 +15,19 @@ checked and nothing costs anything. With one, every request is judged once it is
|
||||
within its limits, with its stored conversation spliced in, and before it is
|
||||
sent to a model; and the answer is judged before the client sees it.
|
||||
|
||||
## Turning them off and on
|
||||
|
||||
```
|
||||
llmgw ctl guard off --note "why" # no request is checked, whatever the policies say
|
||||
llmgw ctl guard on # every policy checks again, as it was
|
||||
```
|
||||
|
||||
That is `guard.enabled` in the cluster's settings, in force on every control
|
||||
plane within seconds without a restart. The policies and where they are
|
||||
attached are kept. While off, `llmgw ctl status`, `llmgw ctl guard list` and
|
||||
the guardrails page say so. To stop one policy rather than all, give it
|
||||
`mode: off`, or `mode: shadow` to have it record what it would do.
|
||||
|
||||
## A policy
|
||||
|
||||
```yaml
|
||||
@@ -101,44 +114,136 @@ guard:
|
||||
|
||||
## A judge model
|
||||
|
||||
Patterns find what has a shape. For what does not — violence, self-harm,
|
||||
abuse, whatever a moderation model was trained on — a policy can have a
|
||||
**judge model on the cluster** read the text:
|
||||
Patterns find what has a shape. For what does not — violence, crime,
|
||||
sexual content, self-harm, jailbreaks, whatever a moderation model was
|
||||
trained on — a policy can have a **judge model on the cluster** read the
|
||||
text.
|
||||
|
||||
The judge must be in the catalog under the name the policy gives, as a
|
||||
judge (`role: guard`), with a replica ready. When it is not, what the
|
||||
policy judges is refused with `guard_unavailable` (`fail: closed`) or served
|
||||
unjudged (`fail: open`), and `llmgw ctl status` and the console's Overview
|
||||
say which judge and why.
|
||||
|
||||
### Which judge
|
||||
|
||||
| judge | runs on | size | |
|
||||
|---|---|---|---|
|
||||
| `qwen3guard-0.6b` | CPUs | 0.8 GB | the default: about 2 s a check on two threads |
|
||||
| `qwen3guard-4b` | a GPU | 2.5 GB | for a cluster with a card to spare |
|
||||
| `qwen3guard-8b` | a GPU | 4.7 GB | |
|
||||
| `llama-guard-3-1b` | CPUs | 0.9 GB | the earlier default: refuses ordinary coding requests |
|
||||
| `llama-guard-3-8b` | a GPU | 4.6 GB | |
|
||||
|
||||
Measured on the reference cluster, as the judge's probability of "unsafe"
|
||||
for ordinary coding requests against harmful ones, Llama Guard 3 1B could
|
||||
not tell them apart: "continue" 0.51, "continue the code above" 0.89, "how
|
||||
do I kill all child processes in bash" 0.92, against 0.97 to 1.00 for
|
||||
harmful requests, at any quantisation. Qwen3Guard-Gen 0.6B gave the
|
||||
ordinary ones 0.15 at the most and the harmful ones 1.00. Its verdicts are
|
||||
safe, unsafe and **controversial**; controversial counts as safe unless the
|
||||
policy says `controversial: unsafe`.
|
||||
|
||||
A policy that names one of these is enough: the leader fetches its weights
|
||||
into the mirror and serves it as a judge, and serves it again if it is
|
||||
removed from the catalog. Another judge is served by hand with `--repo`.
|
||||
|
||||
```
|
||||
llmgw ctl guard judge status # each judge the policies ask, and what is wrong with it
|
||||
llmgw ctl guard judge enable # Qwen3Guard-Gen 0.6B on CPUs, or the judge the policies ask that is missing
|
||||
llmgw ctl guard judge enable --gpu # Qwen3Guard-Gen 4B on a GPU
|
||||
llmgw ctl guard judge use qwen3guard-4b qwen3guard-0.6b # the order every policy asks them in
|
||||
llmgw ctl guard judge enable --replace # when a model of that name is there, but not as a judge
|
||||
```
|
||||
|
||||
or **Do → Serve the guardrails' judge** in the console.
|
||||
|
||||
### A GPU and CPUs together
|
||||
|
||||
`judge.models` is a list asked in turn: the first judge that answers gives
|
||||
the verdict. The next is asked when one has no replica serving, fails, or
|
||||
does not answer within `latency_budget`; a judge with no replica at all is
|
||||
passed over at once. A judge on a GPU first and one on CPUs after it is
|
||||
the GPU's speed while it has its card, and the CPUs' when it does not.
|
||||
|
||||
```yaml
|
||||
judge:
|
||||
models: [qwen3guard-4b, qwen3guard-0.6b]
|
||||
input: block
|
||||
output: block
|
||||
latency_budget: 5s # each judge's wait
|
||||
```
|
||||
|
||||
A judge on a GPU is served as a `standard` model, on one card: it takes a
|
||||
card only from what may give one up, and waits while every card is busy,
|
||||
the judge on CPUs standing in. A judge on CPUs is `critical`, with two
|
||||
replicas, because every judged request waits on one. `llmgw ctl guard judge
|
||||
status` says, for each judge, which is asked when it gives no verdict.
|
||||
|
||||
### When it refuses ordinary requests
|
||||
|
||||
A refusal by a judge names the hazard it saw — the client is told `refused
|
||||
by the gateway's policy (judge: Violent)` — and the same is kept with the
|
||||
request: `llmgw ctl request <id>`, the server's log (`guard_rule`) and the
|
||||
guardrails page show it. To see what the judge says of any text, without
|
||||
anything being kept:
|
||||
|
||||
```
|
||||
llmgw ctl guard judge ask 'How do I sort a list in Python?'
|
||||
# qwen3guard-0.6b: safe, 1900 ms
|
||||
# it answered: "Safety: Safe\nCategories: None"
|
||||
```
|
||||
|
||||
A judge that says `unsafe` to ordinary text is usually Llama Guard 3 1B
|
||||
(`llmgw ctl guard judge use qwen3guard-0.6b`), the wrong weights (another
|
||||
repository's, without the judge's template: `judge enable --replace` serves
|
||||
the right ones), or a policy counting every category: give its judge
|
||||
`categories`. The baseline counts Violent, Non-violent Illegal Acts, Sexual
|
||||
Content or Sexual Acts, Suicide & Self-Harm and Jailbreak; PII is the
|
||||
patterns', and copyright, unethical acts and politically sensitive topics
|
||||
refuse too much ordinary work. While finding out, `judge.input: flag`
|
||||
records instead of refusing.
|
||||
|
||||
A judge's file served from the mirror page — Llama Guard, Qwen3Guard,
|
||||
ShieldGemma — is served as a judge too, under the name the policies ask.
|
||||
By hand, the catalog entry is:
|
||||
|
||||
```yaml
|
||||
# a catalog model (server.yaml models:, or ctl model add)
|
||||
- name: llama-guard-3-1b
|
||||
source: http://mirror:8081/llama-guard-3-1b/…/Llama-Guard-3-1B-Q4_K_M.gguf
|
||||
- name: qwen3guard-0.6b
|
||||
source: http://mirror:8081/qwen3guard-0.6b/…/Qwen3Guard-Gen-0.6B.Q8_0.gguf
|
||||
role: guard
|
||||
guard:
|
||||
format: llamaguard3 # llamaguard3, shieldgemma or json
|
||||
categories: [Violent Crimes, Non-Violent Crimes, Sex-Related Crimes, …] # names for S1, S2, …
|
||||
format: qwen3guard # qwen3guard, llamaguard3, shieldgemma or json
|
||||
scaling: {min: 2, max: 2}
|
||||
profiles: [{name: cpu, engine: llamacpp}]
|
||||
|
||||
# in a guard policy
|
||||
judge:
|
||||
model: llama-guard-3-1b
|
||||
models: [qwen3guard-0.6b] # or model: one judge
|
||||
input: block # flag or block; empty does not ask about input
|
||||
output: flag
|
||||
categories: [S1, S10] # only these count; empty: any unsafe verdict
|
||||
latency_budget: 250ms # then the policy's fail: open serves (flagged), closed refuses
|
||||
categories: [Violent, Jailbreak] # only these count; empty: any unsafe verdict
|
||||
controversial: safe # or unsafe
|
||||
latency_budget: 250ms # each judge's wait; after the last, fail: open serves (flagged), closed refuses
|
||||
```
|
||||
|
||||
A model with `role: guard` is the guardrails' alone:
|
||||
|
||||
- it is not in `/v1/models` or the chat page, and is refused (404) to
|
||||
every key, as a name and as a fallback;
|
||||
- it is `critical` by default and must always run — on demand is refused.
|
||||
Give it two replicas on different nodes, because every judged request
|
||||
waits on it.
|
||||
- it must always run — on demand is refused. One on CPUs is `critical`
|
||||
and should have two replicas on different nodes, because every judged
|
||||
request waits on it.
|
||||
|
||||
The gateway's calls to it go past admission, accounting and the guardrails
|
||||
themselves.
|
||||
|
||||
- **Formats:**
|
||||
- `llamaguard3` reads Meta's "safe" / "unsafe
|
||||
S1,S10".
|
||||
- `qwen3guard` asks in Qwen3Guard's own words and reads "Safety: Safe |
|
||||
Unsafe | Controversial" and its "Categories:".
|
||||
- `llamaguard3` reads Meta's "safe" / "unsafe\nS1,S10"; give the catalog
|
||||
entry `categories` to have S1 recorded as its name.
|
||||
- `shieldgemma` asks Google's model whether the text breaks the listed
|
||||
principles (`categories`, each a sentence) and reads Yes or No.
|
||||
- `json` expects `{"safe": false, "categories": [...]}`.
|
||||
@@ -148,6 +253,9 @@ S1,S10".
|
||||
is read every 600 characters and at the end, so text between two
|
||||
readings has gone out already; `stream: buffer_all` holds it until it
|
||||
is read.
|
||||
- **A long text** is read by its start and its end (about 6000
|
||||
characters), so a message with a file pasted into it fits the judge's
|
||||
context rather than overflowing it and going unjudged.
|
||||
- **Cache:** the same text is not asked about twice within five minutes.
|
||||
|
||||
## Where a policy applies
|
||||
|
||||
+158
-14
@@ -129,8 +129,9 @@ That is the whole step. In order, it:
|
||||
`nvidia-smi` finds any, on CPU otherwise — and waits for it to install its
|
||||
engine
|
||||
9. closes llmgw's internal ports to everything but the cluster's networks
|
||||
10. fetches Llama Guard 3 1B, the guardrails' judge, and runs two replicas
|
||||
on the CPUs of nodes with room — this one, or the first that join
|
||||
10. fetches Qwen3Guard-Gen 0.6B, the guardrails' judge, and runs two
|
||||
replicas on the CPUs of nodes with room — this one, or the first that
|
||||
join ([guardrails.md](guardrails.md#a-judge-model) has one on a GPU too)
|
||||
11. fetches Qwen2.5 0.5B into the mirror and serves it on this node, so the
|
||||
first thing anyone opens `/chat` to find is a model that answers. A
|
||||
machine without room for it is told so within two minutes, and why.
|
||||
@@ -200,6 +201,16 @@ this by running the same command again, on the control plane that holds the
|
||||
authority: etcd is rebuilt on TLS with Patroni in maintenance mode, so
|
||||
PostgreSQL serves throughout and no failover happens.
|
||||
|
||||
Each member's PostgreSQL is sized to the machine's memory, read each time
|
||||
the command runs: a quarter of it as cache (`shared_buffers`), three quarters
|
||||
counted on for the planner, and room to sort and build indexes, all capped. A
|
||||
member that also runs a control plane or an agent gets less: about 15% as
|
||||
cache, so the gateway is not pushed into swap. A knowledge base's vector
|
||||
index is searched well only from memory, so give a store machine more RAM and
|
||||
run the command again. The new cache takes a PostgreSQL restart, which Patroni
|
||||
does one member at a time, the leader last. The leader's restart costs a few
|
||||
seconds without writes, and there is no failover.
|
||||
|
||||
Every step looks before it acts. Running it again finishes an install that
|
||||
stopped halfway and changes nothing on one that did not; it never overwrites
|
||||
`server.yaml`, the authority, or an existing account's password.
|
||||
@@ -257,8 +268,14 @@ llmgw ctl config history # who changed what, when, and why
|
||||
llmgw ctl config rollback 4 # back to version 4, as a new version
|
||||
```
|
||||
|
||||
or on the console's **Settings** screen. A change is checked before it is
|
||||
kept — one that would not load is refused — and recorded in the audit log.
|
||||
or on the console's **Settings** screen, or in a browser on the portal's
|
||||
**Settings** page (`/admin/settings`, for administrators only: the document
|
||||
can hold webhook secrets). The page edits the same document, lists its
|
||||
history, and puts an earlier version back. A change is checked before it is
|
||||
kept, so one that would not load is refused. A change made from a version
|
||||
someone else has replaced since is refused too, rather than overwriting
|
||||
theirs. Every change is recorded in the audit log, with the document's size
|
||||
but not its contents.
|
||||
Guardrails change live; anything else has every control plane start again on
|
||||
the new settings, the standbys first and the leader half a minute later, so
|
||||
one of them is always serving. Engines keep running through it.
|
||||
@@ -268,6 +285,33 @@ cluster keys — with a `baseline` guard policy, if the file had none — and
|
||||
from then on those keys in any server.yaml are not read (the server says so
|
||||
when it starts).
|
||||
|
||||
**The nodes' agents** are set the same way, under `agents`. `agents.default`
|
||||
is for every node, and `agents.nodes.<node-id>` is for one node, laid over
|
||||
the default:
|
||||
|
||||
```
|
||||
llmgw ctl config set agents.default.heartbeat_interval 5s
|
||||
llmgw ctl config set agents.nodes.cuda01.labels.gpu rtx6000
|
||||
llmgw ctl config edit # fabrics, engines, supervisor, discovery…
|
||||
```
|
||||
|
||||
Each node's `agent.yaml` then needs only what its agent must know before it
|
||||
can be told anything: `node_id`, `server` or `servers`, `hostname`,
|
||||
`mgmt_addr` and `tls`. Those keys are refused under `agents`.
|
||||
|
||||
- **Precedence.** Anything else in the file is still read, with the
|
||||
cluster's settings over it.
|
||||
- **Applying a change.** The control plane sends each node its settings, a
|
||||
second apart. An agent whose settings changed keeps them beside its
|
||||
`agent.yaml` (in `/var/lib/llmgw-agent/agent-settings.yaml`) and starts
|
||||
again on them, its engines serving throughout.
|
||||
- **Bad settings.** Settings that would not load are refused when they are
|
||||
saved. If they somehow do not load beside one node's own file, that agent
|
||||
keeps running on what it has and logs why.
|
||||
- **Editing `agent.yaml` itself.** The agent notices the change within ten
|
||||
seconds and starts again on it. If the edited file does not load, the
|
||||
agent reports it and keeps running on what it had.
|
||||
|
||||
**This machine's own keys** stay in `/etc/llmgw/server.yaml`: `listen`,
|
||||
`admin_listen`, `metrics_listen`, `agent_listen`, `agent_tls`, `auth.store`,
|
||||
`secrets`, `tracing`, `monitoring`, `mirror.api` and `mirror.tls`, the backup
|
||||
@@ -1226,6 +1270,51 @@ prices and budgets, rate limits, usage records, and key allowlists.
|
||||
whose allowlist leaves the external model out cannot reach it, as a fallback or
|
||||
by name. The cluster page and `ctl models` mark it `external`.
|
||||
|
||||
### Pictures from ComfyUI
|
||||
|
||||
A ComfyUI server is an external model with `api: comfyui`. Clients ask it for
|
||||
pictures at OpenAI's `/v1/images/generations`; llmgw fills a ComfyUI workflow
|
||||
with the request, queues it, waits, and sends the pictures back as `b64_json`:
|
||||
|
||||
```yaml
|
||||
- name: flux-schnell
|
||||
source: external
|
||||
external:
|
||||
url: http://10.0.0.147:8188 # ComfyUI, started with --listen
|
||||
model: flux1-schnell-fp8.safetensors # the checkpoint, as ComfyUI lists it
|
||||
api: comfyui
|
||||
workflow: flux-schnell # or flux-dev, or a JSON file
|
||||
```
|
||||
|
||||
`flux-schnell` (the default) and `flux-dev` are built in. They suit the
|
||||
single-file FLUX checkpoints that ComfyUI loads with `CheckpointLoaderSimple`.
|
||||
For any other model, export the workflow from ComfyUI with **Save (API)** and
|
||||
name the file. In the file, put placeholders where the request's values go:
|
||||
|
||||
- `{{prompt}}` and `{{negative_prompt}}` are the prompts.
|
||||
- `{{width}}` and `{{height}}` are the size.
|
||||
- `{{n}}` is the batch size.
|
||||
- `{{seed}}` is the seed.
|
||||
- `{{steps|4}}` and `{{guidance|3.5}}` take the value after the `|` when the
|
||||
request leaves them out.
|
||||
- `{{checkpoint}}` is the `model` above.
|
||||
|
||||
A request is OpenAI's: `prompt`, `n` (at most 4), `size` (each side a multiple
|
||||
of 16, from 256 to 2048) and `response_format: b64_json`. Beyond those, it may
|
||||
send `seed`, `steps`, `negative_prompt` and `guidance`. The seed used comes
|
||||
back in `x-llmgw-seed`. `url` is not offered, because pictures are kept
|
||||
nowhere. ComfyUI also writes the workflow into each PNG; llmgw strips it
|
||||
before the picture is sent.
|
||||
|
||||
Requests for pictures pass the same checks as any other request: keys,
|
||||
allowlists, limits, admission, and the guardrails on the prompt. A client that
|
||||
leaves takes its workflow off the ComfyUI queue. Usage is recorded, but for
|
||||
now it is charged as the prompt's tokens; a per-picture price is not yet
|
||||
available.
|
||||
|
||||
Start ComfyUI with `--disable-smart-memory` on a card that llmgw also places
|
||||
models on. Without it, ComfyUI keeps the model on the card after each picture.
|
||||
|
||||
### Federating clusters
|
||||
|
||||
Another llmgw cluster — another site, another team's GPUs — can be a provider
|
||||
@@ -1702,9 +1791,11 @@ the file again.
|
||||
|
||||
On a machine that also runs the control plane, PostgreSQL's defaults assume it
|
||||
owns the box. A gateway swapped out to give the database cache is a worse
|
||||
outcome than a smaller cache, so cap it: `shared_buffers = 96MB` and
|
||||
`max_connections = 40` are enough for a store whose working set is measured in
|
||||
megabytes.
|
||||
outcome than a smaller cache, so size it to the machine. `server init` does this
|
||||
for a PostgreSQL it installs: `shared_buffers` about 15% of memory (at least
|
||||
96MB), `effective_cache_size` half, `work_mem` and `maintenance_work_mem`
|
||||
scaled to it, and `max_connections = 40`. A store with a large knowledge base
|
||||
wants a machine of its own, where its vector index fits in memory.
|
||||
|
||||
## Tagging requests
|
||||
|
||||
@@ -2150,13 +2241,29 @@ llmgw ctl config set firewall.enabled false # every machine's table taken out
|
||||
llmgw ctl config set firewall.enabled true # and put back as it was
|
||||
```
|
||||
|
||||
It is in force within seconds: each node's agent takes its table out or puts
|
||||
it back, and each control plane's `llmgw-firewall-sync.path` (installed by
|
||||
`server init` and `server add`; `sudo llmgw firewall sync --install` on one
|
||||
set up before) does the same there. The rules files stay, so turning it back
|
||||
on restores them. Nodes added while it is off get none. A separate mirror
|
||||
machine and the store's members keep theirs for now. `firewall_networks`
|
||||
changes reach the nodes the same way.
|
||||
It is in force on every machine within a minute, each in its own way:
|
||||
|
||||
- **nodes**: the agent takes its table out or puts it back;
|
||||
- **control planes**: `llmgw-firewall-sync.path`, installed by `server
|
||||
init` and `server add` (`sudo llmgw firewall sync --install` on one set up
|
||||
before);
|
||||
- **a separate mirror machine**: the leading control plane tells the mirror,
|
||||
which writes it to `.firewall.json` in its directory, and the same path
|
||||
unit applies it. `mirror init` installs it; on a mirror set up before, run
|
||||
`sudo llmgw mirror init` again with the same flags;
|
||||
- **store members that are nothing else**: they run no llmgw, so
|
||||
`llmgw-firewall-follow.timer` reads what the leading control plane keeps in
|
||||
the store, from the member's own copy, once a minute. `server init
|
||||
--patroni` installs it; run it again on a cluster set up before.
|
||||
|
||||
The rules files stay, so turning it back on restores them. Nodes added while
|
||||
it is off get none. Networks added to `firewall_networks` reach every machine
|
||||
the same way. One taken out of it stays let in until it is taken out of
|
||||
`/etc/llmgw/firewall.nft` by hand: the rules are only ever added to.
|
||||
|
||||
The leading control plane also tells the mirror machine where every node and
|
||||
control plane is, so a node on a network the mirror had not let in can fetch
|
||||
its weights without a `firewall_networks` entry.
|
||||
|
||||
Only those ports are touched. The table's policy is to accept, so whatever
|
||||
else the machine serves — and whatever firewall it already has — is left as
|
||||
@@ -2294,6 +2401,43 @@ A cluster set up online becomes offline the same way: `bundle import` on the
|
||||
mirror and control plane, then the `offline:` block in `server.yaml` and a
|
||||
restart.
|
||||
|
||||
## Stopping and starting the system
|
||||
|
||||
From a control plane:
|
||||
|
||||
```
|
||||
llmgw status # does it serve; its control planes, nodes, models, store
|
||||
llmgw stop --note "maintenance" # every model taken down, every request answered 503
|
||||
llmgw start # every model back as its catalog entry asks
|
||||
sudo llmgw restart # every agent and control plane started again, engines serving throughout
|
||||
```
|
||||
|
||||
`llmgw stop` frees every GPU and answers requests for a model with 503
|
||||
(`system_stopped`, with Retry-After) until `llmgw start`. The control planes,
|
||||
agents and mirror keep running, so the console, the admin pages and `llmgw
|
||||
start` keep working, and no model's definition changes. It is the cluster's
|
||||
setting `serving.stopped`, so a restart does not undo it, and it is in the
|
||||
audit with its note. `--wait` waits until every replica is down, or every
|
||||
model has its replicas again.
|
||||
|
||||
`llmgw restart` restarts every node's agent, then the other control planes
|
||||
over ssh (`--ssh-user`, the one sudo was run by), then this one — the leader
|
||||
last, so it hands over to a standby already restarted. Engines keep serving
|
||||
and are taken over again; nothing is placed anew.
|
||||
|
||||
One machine's own services are `llmgw service`:
|
||||
|
||||
```
|
||||
sudo llmgw service status
|
||||
sudo llmgw service stop # agent, control plane, mirror — whichever this machine has
|
||||
sudo llmgw service start
|
||||
sudo llmgw service restart agent # or name one: server, agent, mirror
|
||||
```
|
||||
|
||||
Stopping a node's agent leaves its engines serving — the agent takes them
|
||||
over again when it starts — unless `llmgw service stop --engines` stops them
|
||||
too, and the control plane then places their replicas elsewhere.
|
||||
|
||||
## Removing llmgw from a machine
|
||||
|
||||
```
|
||||
|
||||
@@ -0,0 +1,408 @@
|
||||
# Knowledge bases
|
||||
|
||||
A knowledge base is a collection of a company's documents that the cluster
|
||||
can answer from. People put their documents into it from their own machines
|
||||
with `llmgw-kb`. The tool reads the files on that machine, cuts them into
|
||||
passages ("chunks"), has the gateway embed each passage with the person's own
|
||||
key, and sends the text and the vectors. The files themselves never leave
|
||||
the machine.
|
||||
|
||||
## Why retrieval, and not training
|
||||
|
||||
Retrieval is the tool for *knowledge*: policies, procedures, contracts. The
|
||||
answer can name its source, a changed document is sent again in minutes, and
|
||||
a passage only ever reaches people allowed to read its collection.
|
||||
|
||||
Training a model on the documents (LoRA) changes how it *writes*, not what it
|
||||
reliably *knows*. It has no access control: whatever goes into the weights
|
||||
is available to everyone who may use them.
|
||||
|
||||
## What it needs
|
||||
|
||||
- **The PostgreSQL store with pgvector.** On a SQLite store knowledge bases
|
||||
are off, and `llmgw ctl kb list` and the console say why.
|
||||
- `llmgw server init` and the high-availability database setup install
|
||||
pgvector and enable it.
|
||||
- A cluster set up before this build needs it by hand, once:
|
||||
1. Install the package on every database machine:
|
||||
`sudo apt-get install postgresql-16-pgvector` (use the server's
|
||||
version number).
|
||||
2. On the leader, enable it in llmgw's database:
|
||||
`sudo -u postgres psql -d llmgw -c 'CREATE EXTENSION vector'`.
|
||||
3. Restart the control planes: `llmgw restart`.
|
||||
- **An embedding model in the catalog**, with `capabilities.embeddings`.
|
||||
Two good choices handle Turkish and run on CPUs: `BAAI/bge-m3` (1024
|
||||
dimensions) and `Qwen/Qwen3-Embedding-0.6B`.
|
||||
|
||||
Add it as any model is added (**Serve a model** in the console, or the
|
||||
catalog), with its entry saying that it embeds:
|
||||
|
||||
```yaml
|
||||
capabilities:
|
||||
embeddings: true
|
||||
```
|
||||
|
||||
A collection stays with the model it was made with: vectors from two
|
||||
models cannot be compared.
|
||||
|
||||
Some embedding models are trained with an instruction before the
|
||||
question and none before the documents. Give it as the entry's
|
||||
`query_prefix`; llmgw puts it before questions only. For
|
||||
Qwen3-Embedding:
|
||||
|
||||
```yaml
|
||||
query_prefix: "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: "
|
||||
```
|
||||
|
||||
With pgvector 0.7 or later a collection's index is kept at half
|
||||
precision (`halfvec`), which halves its memory and takes models of up to
|
||||
4000 dimensions; with an older pgvector the limit is 2000. A collection
|
||||
keeps the kind it was made with.
|
||||
|
||||
## For administrators
|
||||
|
||||
### Make a collection
|
||||
|
||||
A collection belongs to one part of the tenancy tree: an organization, a
|
||||
team, a project or a person. Whoever is in the owner, at any level below it,
|
||||
may read it and add to it.
|
||||
|
||||
```sh
|
||||
llmgw ctl kb create team:Acme/hr policies bge-m3
|
||||
llmgw ctl kb create user:USER-ID notes bge-m3 --org Acme
|
||||
```
|
||||
|
||||
The model's dimension is asked of the model itself, so the model has to be
|
||||
able to serve when the collection is made. `--dims 1024` gives it instead.
|
||||
|
||||
In the console: **Knowledge**, then **A**.
|
||||
|
||||
### Who may read and add
|
||||
|
||||
The owner, and everything below it, may read and add. A **grant** gives the
|
||||
same to another part of the same organization, to a person, or to one key:
|
||||
|
||||
```sh
|
||||
llmgw ctl kb grant Acme/policies project:Acme/sales/crm # read
|
||||
llmgw ctl kb grant Acme/policies team:Acme/legal --write # read and add
|
||||
llmgw ctl kb revoke Acme/policies project:Acme/sales/crm
|
||||
```
|
||||
|
||||
How access works:
|
||||
|
||||
- **Grants add up down the tree.** A grant to a team reaches its projects
|
||||
and their keys.
|
||||
- **Nothing crosses organizations.** A grant to another organization's
|
||||
team, project, key or person is refused.
|
||||
- **A person's own key** (a key with no project) reads nothing by itself.
|
||||
The person names the project they act in (`llmgw-kb login` asks for it,
|
||||
or `--project`) and must be a member of it. Joining a second team
|
||||
therefore does not quietly widen what every key of theirs can read.
|
||||
- **No access looks like no collection.** A collection a key may not read
|
||||
is answered exactly as one that does not exist, so even its name does not
|
||||
leak.
|
||||
|
||||
### Answer from collections by default
|
||||
|
||||
Clients that cannot add a header still get answers from the documents when
|
||||
their part of the tree has a default:
|
||||
|
||||
```sh
|
||||
llmgw ctl kb default team:Acme/hr Acme/policies Acme/forms
|
||||
llmgw ctl kb default key:KEY-ID Acme/policies
|
||||
llmgw ctl kb default team:Acme/hr --clear
|
||||
```
|
||||
|
||||
- **Which default applies.** The narrowest one: a key's over its project's,
|
||||
a project's over its team's, a team's over its organization's.
|
||||
- **A request can override it.** A request that names collections uses
|
||||
those instead, and `X-Llmgw-Knowledge: none` turns the default off for
|
||||
that request.
|
||||
- **A default never refuses a request.** A collection the caller may not
|
||||
read is left out for them. If the search fails (the embedding model is
|
||||
down, say), the request is answered without passages and the failure is
|
||||
logged.
|
||||
- **Changes take effect within 30 seconds.**
|
||||
|
||||
In the console: **Knowledge**, Enter on a collection, then **Answer from it
|
||||
by default…**.
|
||||
|
||||
### Look after it
|
||||
|
||||
```sh
|
||||
llmgw ctl kb list # collections, owners, documents, grants
|
||||
llmgw ctl kb docs Acme/policies # its documents, who added each and when
|
||||
llmgw ctl kb rm Acme/policies old/form.pdf # one document out
|
||||
llmgw ctl kb search Acme/policies how many days of annual leave
|
||||
llmgw ctl kb delete Acme/policies # the collection, its documents and grants
|
||||
```
|
||||
|
||||
Rules that apply:
|
||||
|
||||
- A project, team or organization that owns a collection is not deleted
|
||||
until its collections are. This is the same rule as for one with members.
|
||||
- A person's own collections go with their account (`users delete`,
|
||||
`users forget`).
|
||||
- Every document added or removed is in the audit log (`llmgw ctl audit`),
|
||||
by name. The text is never recorded.
|
||||
|
||||
## For everyone: `llmgw-kb`
|
||||
|
||||
### Get it
|
||||
|
||||
Download `llmgw-kb` for your system from the release page, next to `llmgw`:
|
||||
|
||||
| System | File |
|
||||
|---|---|
|
||||
| Windows (Intel/AMD) | `llmgw-kb-VERSION-windows-amd64.exe` |
|
||||
| Windows (ARM) | `llmgw-kb-VERSION-windows-arm64.exe` |
|
||||
| macOS (Apple silicon) | `llmgw-kb-VERSION-darwin-arm64` |
|
||||
| macOS (Intel) | `llmgw-kb-VERSION-darwin-amd64` |
|
||||
| Linux | `llmgw-kb-VERSION-linux-amd64` / `-arm64` |
|
||||
|
||||
Rename it to `llmgw-kb` (`llmgw-kb.exe` on Windows) and put it somewhere on
|
||||
your PATH. You do not need to install anything else.
|
||||
|
||||
On macOS, a downloaded program that is not signed is stopped the first time.
|
||||
Either right-click it and choose Open, or run
|
||||
`xattr -d com.apple.quarantine llmgw-kb`.
|
||||
|
||||
### Log in once
|
||||
|
||||
```sh
|
||||
llmgw-kb login
|
||||
```
|
||||
|
||||
It asks for three things, checks them, and saves them where only you can
|
||||
read them:
|
||||
|
||||
- **The gateway's address.**
|
||||
- **Your API key.** It is typed without being shown.
|
||||
- **A project.** Needed only for a personal key.
|
||||
|
||||
Instead of logging in you can set `LLMGW_URL`, `LLMGW_API_KEY` and
|
||||
`LLMGW_PROJECT`.
|
||||
|
||||
```sh
|
||||
llmgw-kb collections # what your key may read, and which you may add to
|
||||
```
|
||||
|
||||
### Send your documents
|
||||
|
||||
```sh
|
||||
llmgw-kb sync policies ~/Documents/HR
|
||||
llmgw-kb sync policies "C:\Users\ayse\Documents\İK" --include "*.pdf,*.docx"
|
||||
llmgw-kb sync policies ./handbook ./forms --exclude "drafts/**"
|
||||
```
|
||||
|
||||
Run it as often as you like. The first run sends everything. Later runs
|
||||
send only what changed, and do not even read files whose size and time have
|
||||
not changed.
|
||||
|
||||
| Option | What it does |
|
||||
|---|---|
|
||||
| `--dry-run` | Reads and cuts everything, sends nothing, and lists what would leave the machine and how much. |
|
||||
| `--delete` | Also removes documents whose files are gone. Only those this machine sent, from these folders; never anything someone else sent. |
|
||||
| `--retry` | Reads again files skipped before, for example after installing LibreOffice. |
|
||||
| `--include`, `--exclude` | Glob patterns. `*.pdf` matches a file name; `reports/**/*.xlsx` matches a path. |
|
||||
| `--hidden` | Reads hidden files and folders too (left out by default). |
|
||||
| `-v` | Lists unchanged files too. |
|
||||
|
||||
Documents are named by their path under the folder, with forward slashes,
|
||||
on every system. With several folders, each folder's own name comes first.
|
||||
|
||||
```sh
|
||||
llmgw-kb status policies # what this machine sent and skipped, and why
|
||||
llmgw-kb history # the latest runs
|
||||
llmgw-kb formats # what can be read here
|
||||
```
|
||||
|
||||
### What it reads
|
||||
|
||||
| Read here | Formats |
|
||||
|---|---|
|
||||
| Documents | PDF (text), Word (`.docx`), Excel (`.xlsx`), PowerPoint (`.pptx`, with speaker notes), OpenDocument (`.odt`, `.ods`, `.odp`), RTF, EPUB |
|
||||
| Web and mail | HTML, MHTML, e-mail (`.eml`, with its attachments) |
|
||||
| Text | Markdown, plain text, CSV, JSON, YAML, XML, subtitles, notebooks (`.ipynb`), source code |
|
||||
| Archives | ZIP, with every file inside it read in turn |
|
||||
|
||||
| Read with a helper on your machine | Helper |
|
||||
|---|---|
|
||||
| Word, Excel and PowerPoint 97–2003 (`.doc`, `.xls`, `.ppt`), Pages, Numbers, Keynote and others | [LibreOffice](https://www.libreoffice.org) |
|
||||
| Images and scanned PDFs | [Tesseract](https://tesseract-ocr.github.io) (with the Turkish language data), plus Poppler's `pdftoppm` for PDFs |
|
||||
| PDFs the built-in reader cannot parse | Poppler's `pdftotext` |
|
||||
|
||||
`llmgw-kb formats` says which helpers it found. Text in UTF-8, UTF-16, or
|
||||
the Windows Turkish and Western code pages is read correctly.
|
||||
|
||||
A file that cannot be read is skipped with the reason, and the run goes on.
|
||||
Reasons include a locked Office file, a format with no helper, and a
|
||||
scanned page with no OCR.
|
||||
|
||||
### What leaves your machine
|
||||
|
||||
For each new or changed document, two things leave the machine, both over
|
||||
the gateway's TLS connection:
|
||||
|
||||
- **The passages, to be embedded.** They go through `/v1/embeddings` like
|
||||
any request of yours. Your key's limits, budget and guardrails apply.
|
||||
- **The document itself.** Its name, its file's SHA-256 digest and size,
|
||||
and the passages with their vectors.
|
||||
|
||||
The file itself is never sent. `--dry-run` shows exactly what would be.
|
||||
|
||||
### The ledger
|
||||
|
||||
`llmgw-kb` keeps a small SQLite database of what it did with every
|
||||
document: the file it came from, its size, time and digest, and whether it
|
||||
was sent, skipped or failed, with the reason.
|
||||
|
||||
| System | Where it is |
|
||||
|---|---|
|
||||
| Windows | `%AppData%\llmgw-kb\ledger.db` |
|
||||
| macOS | `~/Library/Application Support/llmgw-kb/ledger.db` |
|
||||
| Linux | `~/.config/llmgw-kb/ledger.db` |
|
||||
|
||||
Deleting the ledger is safe: the next run hashes every file once, finds what
|
||||
the collection already has, and sends nothing that is unchanged. `--ledger`
|
||||
uses another file.
|
||||
|
||||
## Asking questions of a collection
|
||||
|
||||
A chat request names the collections to answer from in a header. Any client
|
||||
that can add a header can do it: the OpenAI and Anthropic SDKs, Continue, or
|
||||
curl.
|
||||
|
||||
```sh
|
||||
curl https://llm.example.com/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "X-Llmgw-Knowledge: politikalar" \
|
||||
-d '{"model": "qwen2.5-coder-14b-instruct",
|
||||
"messages": [{"role": "user", "content": "Kaç gün yıllık iznim var?"}]}'
|
||||
```
|
||||
|
||||
```python
|
||||
client.chat.completions.create(
|
||||
model="qwen2.5-coder-14b-instruct",
|
||||
messages=[{"role": "user", "content": "Kaç gün yıllık iznim var?"}],
|
||||
extra_headers={"X-Llmgw-Knowledge": "politikalar"},
|
||||
)
|
||||
```
|
||||
|
||||
What the gateway does with such a request:
|
||||
|
||||
1. **It finds the collections.** It looks up the ones named, by name or ID,
|
||||
among those the key may read. A collection it may not read is answered
|
||||
404, the same as one that does not exist.
|
||||
2. **It embeds the question.** The last user message is embedded by the
|
||||
collection's model, on the caller's own account: their limits, budget and
|
||||
guardrails apply.
|
||||
3. **It searches.** The collections are searched by meaning and by words
|
||||
(below).
|
||||
4. **It adds the passages.** The best passages are numbered and added to the
|
||||
system message, with an instruction to answer from them, cite them as
|
||||
[1], [2], and say so when they do not hold the answer. The system message
|
||||
is created if the request has none. Passages that would not fit the
|
||||
model's context are left out, the weakest first.
|
||||
5. **It judges and sends the request.** The request, passages included, then
|
||||
goes through the guardrails and to the model as usual. A document
|
||||
carrying instructions aimed at its reader is judged like anything else
|
||||
the model reads.
|
||||
|
||||
On the portal's chat page, **Documents** picks a collection the person may
|
||||
read in the chosen project, the default, or none. Under each answer, the
|
||||
sources it was given are listed by number.
|
||||
|
||||
How it is shaped:
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| `X-Llmgw-Knowledge` | Collections, comma separated; `none` for no passages, whatever the default. |
|
||||
| `X-Llmgw-Knowledge-K` | How many passages, 1 to 20; 5 by default. |
|
||||
| `X-Llmgw-Project` | For a personal key: the project it acts in. |
|
||||
| `x-llmgw-sources` (answer) | The passages given, as `document#chunk`, in their numbered order; empty when nothing matched. |
|
||||
|
||||
Behaviour to know about:
|
||||
|
||||
- **No passages matched.** The model is told that nothing in the documents
|
||||
answers the question, and asked to say so rather than answer from general
|
||||
knowledge.
|
||||
- **Other APIs.** Requests through `/v1/responses` and `/v1/messages`
|
||||
(Codex, Claude Code) work the same way, since they are answered as chat
|
||||
completions.
|
||||
- **Separate prefix caches.** On vLLM replicas, a request given passages
|
||||
carries a `cache_salt` derived from the collections they came from.
|
||||
Requests that could be given the same passages still share cached blocks,
|
||||
and no one else's do: not another organization's, nor another team's of
|
||||
the same one. How fast an answer starts therefore says nothing about what
|
||||
someone else asked.
|
||||
- **No answer cache.** Answers from knowledge bases are never served from
|
||||
the answer cache, because two keys of one project can read different
|
||||
collections.
|
||||
|
||||
### How a passage is chosen
|
||||
|
||||
Each passage scores the following, and the highest scores win:
|
||||
|
||||
- its similarity in meaning to the question (cosine, 0 to 1), plus
|
||||
- up to 0.3 for the share of the question's words it contains.
|
||||
|
||||
Words almost every passage has ("ne", "için", "the", "of") are not counted.
|
||||
The meaning does most of the work: "how much leave do fathers get" finds the
|
||||
paternity page, not the annual-leave page that shares more of its words.
|
||||
The word match helps with what meaning cannot carry, such as a form number,
|
||||
a product code or a name.
|
||||
|
||||
Metrics: `llmgw_kb_retrievals_total{result}` (found, nothing, refused) and
|
||||
`llmgw_kb_retrieval_seconds`.
|
||||
|
||||
## Training an adapter on a collection
|
||||
|
||||
An adapter (LoRA) changes how a model *writes*: the organisation's words,
|
||||
format and tone. It does not reliably teach facts. On the PoC, a small
|
||||
model trained on nine policy documents answered 4 of 9 rephrased questions
|
||||
right and mixed facts up on the rest, while the same questions through the
|
||||
knowledge base got 9 of 9. Use both: the knowledge base for what answers
|
||||
say, an adapter for how they sound.
|
||||
|
||||
```sh
|
||||
# 1. Questions and answers from the collection, written by a model of the
|
||||
# cluster. Read the file through and correct it.
|
||||
llmgw ctl kb dataset Acme/policies --model qwen2.5-coder-14b-instruct --out policies.jsonl
|
||||
|
||||
# 2. Train on a node. --gpu-memory is what the job needs of a card; it
|
||||
# runs beside an engine when the engine leaves that much free.
|
||||
llmgw ctl train start policies-tone --model qwen2.5-0.5b-instruct \
|
||||
--base-repo Qwen/Qwen2.5-0.5B-Instruct --node cuda01 --gpu-memory 2300 --data policies.jsonl
|
||||
llmgw ctl train show JOB --follow
|
||||
|
||||
# 3. Make it the model's adapter, kept on the mirror.
|
||||
llmgw ctl train publish JOB --restart
|
||||
```
|
||||
|
||||
The node trains with its installed vLLM's torch and a small overlay (peft
|
||||
and accelerate), as the engines' user. For a model llama.cpp serves, it also
|
||||
converts the adapter to GGUF with the installed llama.cpp's own code. What
|
||||
it makes is kept by the control plane (`llmgw ctl train download`).
|
||||
|
||||
An adapter is private like its collection:
|
||||
|
||||
- **It has an owner.** An adapter trained on a collection carries some of
|
||||
what the collection says in its weights, so it belongs to a part of the
|
||||
tree.
|
||||
- By default the owner is the collection's owner.
|
||||
- `--owner` names another, which must be able to read every collection
|
||||
the pairs came from.
|
||||
- Only callers in the owner's part of the tree see the adapter. To anyone
|
||||
else it is a model that does not exist, absent from `/v1/models` too.
|
||||
- **It is public only by choice.** `--public` makes an adapter every
|
||||
caller's. It is refused for pairs from a knowledge base, so it is only for
|
||||
pairs written by hand.
|
||||
|
||||
## How it is kept apart
|
||||
|
||||
Passages live in PostgreSQL, partitioned by organization, with an index per
|
||||
collection. Every search and every write names the organization and the
|
||||
collections the caller may use. On top of that, row-level security on the
|
||||
table checks the same thing again, so a query that forgot its filter
|
||||
returns nothing rather than another tenant's passages.
|
||||
+27
-5
@@ -379,16 +379,38 @@ is what someone gets from a borrowed laptop, and the password is what they do no
|
||||
have — and it signs every other browser out, which is the point if they are
|
||||
changing it because something went wrong.
|
||||
|
||||
There is no self-service reset, because this gateway sends no email and a reset
|
||||
link with nowhere to send it is not a reset. When someone loses theirs:
|
||||
Someone who has forgotten theirs follows **Forgotten your password?** on the
|
||||
sign-in page. With mail set up in the cluster's settings, a link to set a new one
|
||||
is mailed to their address — the page says the same whether or not the address
|
||||
has an account, so nobody can find out who does from it:
|
||||
|
||||
```
|
||||
llmgw ctl config set mail '{smtp: smtp.example.com:587, from: llmgw@example.com,
|
||||
username: llmgw, password_file: /etc/llmgw/secrets/smtp, base_url: https://llm.example.com}'
|
||||
```
|
||||
|
||||
`base_url` is where the link points, and it is never taken from the request:
|
||||
otherwise a stranger could ask for someone's link from a page of their own. A
|
||||
link works once, for an hour, and a newer one replaces it.
|
||||
|
||||
Without mail, the page says to ask an operator, who makes the link and hands it
|
||||
over however suits:
|
||||
|
||||
```
|
||||
llmgw users reset-link ayse@example.com --store … # --base-url without mail.base_url
|
||||
```
|
||||
|
||||
Or sets a temporary password to give them, and has them change it from their
|
||||
profile:
|
||||
|
||||
```
|
||||
llmgw users reset-password usr-k4mq… --password "a temporary passphrase" --store …
|
||||
```
|
||||
|
||||
Tell them out of band and have them change it from their profile. Their sessions
|
||||
are closed — if they lost the password, the reason might be that someone else has
|
||||
it. Their API keys keep working; those are separate credentials.
|
||||
Either way their sessions are closed — if they lost the password, the reason
|
||||
might be that someone else has it. Their API keys keep working; those are
|
||||
separate credentials. People who sign in through single sign-on have no password
|
||||
here, and get no link.
|
||||
|
||||
## Trying a model in the browser
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ install:
|
||||
- your administrator account. **The password is shown once. Copy it now.**
|
||||
- a model mirror, which downloads model files once for the whole cluster;
|
||||
- this machine as a first node, and a small model (Qwen2.5 0.5B) on it;
|
||||
- guardrails on every request, with a judge model (Llama Guard 3 1B, two
|
||||
- guardrails on every request, with a judge model (Qwen3Guard-Gen 0.6B, two
|
||||
replicas) on CPUs — a request it cannot check is refused, not let through — and API keys that expire within 90 days. What is hardened, and how
|
||||
to loosen it: [install.md](install.md#2-the-control-plane).
|
||||
|
||||
|
||||
+202
-5
@@ -18,6 +18,7 @@ rather than repeated.
|
||||
| `name` | `string` | [Adapter](#adapter) |
|
||||
| `source` | `string` | [Adapter](#adapter) |
|
||||
| `scale` | `float64` | [Adapter](#adapter) |
|
||||
| `owner` | `string` | [Adapter](#adapter) |
|
||||
| `enabled` | `bool` | [Admission](#admission) |
|
||||
| `bands` | `[]Band` | [Admission](#admission) |
|
||||
| `default_band` | `string` | [Admission](#admission) |
|
||||
@@ -26,6 +27,8 @@ rather than repeated.
|
||||
| `queue_capacity_per_band` | `int` | [Admission](#admission) |
|
||||
| `queue_capacity_total` | `int` | [Admission](#admission) |
|
||||
| `interval` | `time.Duration` | [Admission](#admission) |
|
||||
| `default` | `map[string]any` | [AgentSettings](#agentsettings) |
|
||||
| `nodes` | `map[string]map[string]any` | [AgentSettings](#agentsettings) |
|
||||
| `enabled` | `bool` | [Auth](#auth) |
|
||||
| `store` | `string` | [Auth](#auth) |
|
||||
| `cache_ttl` | `time.Duration` | [Auth](#auth) |
|
||||
@@ -57,6 +60,7 @@ rather than repeated.
|
||||
| `tools` | `bool` | [Capabilities](#capabilities) |
|
||||
| `vision` | `bool` | [Capabilities](#capabilities) |
|
||||
| `embeddings` | `bool` | [Capabilities](#capabilities) |
|
||||
| `images` | `bool` | [Capabilities](#capabilities) |
|
||||
| `interval` | `time.Duration` | [Control](#control) |
|
||||
| `base_backoff` | `time.Duration` | [Control](#control) |
|
||||
| `max_backoff` | `time.Duration` | [Control](#control) |
|
||||
@@ -73,6 +77,7 @@ rather than repeated.
|
||||
| `api_key` | `string` | [Endpoint](#endpoint) |
|
||||
| `models` | `[]string` | [Endpoint](#endpoint) |
|
||||
| `engine` | `string` | [Endpoint](#endpoint) |
|
||||
| `workflow` | `string` | [Endpoint](#endpoint) |
|
||||
| `pool` | `string` | [Endpoint](#endpoint) |
|
||||
| `weight` | `float64` | [Endpoint](#endpoint) |
|
||||
| `priority` | `int` | [Endpoint](#endpoint) |
|
||||
@@ -87,6 +92,8 @@ rather than repeated.
|
||||
| `model` | `string` | [External](#external) |
|
||||
| `api_key_file` | `string` | [External](#external) |
|
||||
| `peer` | `string` | [External](#external) |
|
||||
| `api` | `string` | [External](#external) |
|
||||
| `workflow` | `string` | [External](#external) |
|
||||
| `budgets` | `[]FacilityBudget` | [Facility](#facility) |
|
||||
| `margin` | `float64` | [Facility](#facility) |
|
||||
| `gpu_watts` | `float64` | [Facility](#facility) |
|
||||
@@ -102,9 +109,11 @@ rather than repeated.
|
||||
| `format` | `string` | [GuardModel](#guardmodel) |
|
||||
| `categories` | `[]string` | [GuardModel](#guardmodel) |
|
||||
| `model` | `string` | [JudgePolicy](#judgepolicy) |
|
||||
| `models` | `[]string` | [JudgePolicy](#judgepolicy) |
|
||||
| `input` | `string` | [JudgePolicy](#judgepolicy) |
|
||||
| `output` | `string` | [JudgePolicy](#judgepolicy) |
|
||||
| `categories` | `[]string` | [JudgePolicy](#judgepolicy) |
|
||||
| `controversial` | `string` | [JudgePolicy](#judgepolicy) |
|
||||
| `latency_budget` | `time.Duration` | [JudgePolicy](#judgepolicy) |
|
||||
| `tokenize_timeout` | `time.Duration` | [KVIndex](#kvindex) |
|
||||
| `min_blocks` | `int` | [KVIndex](#kvindex) |
|
||||
@@ -140,6 +149,7 @@ rather than repeated.
|
||||
| `context_length` | `int` | [Model](#model) |
|
||||
| `max_output` | `int` | [Model](#model) |
|
||||
| `capabilities` | `Capabilities` | [Model](#model) |
|
||||
| `query_prefix` | `string` | [Model](#model) |
|
||||
| `role` | `string` | [Model](#model) |
|
||||
| `guard` | `*GuardModel` | [Model](#model) |
|
||||
| `origin` | `*Origin` | [Model](#model) |
|
||||
@@ -203,6 +213,7 @@ rather than repeated.
|
||||
| `files` | `[]OriginFile` | [PreviousOrigin](#previousorigin) |
|
||||
| `name` | `string` | [Profile](#profile) |
|
||||
| `engine` | `string` | [Profile](#profile) |
|
||||
| `engine_version` | `string` | [Profile](#profile) |
|
||||
| `gpu_class` | `string` | [Profile](#profile) |
|
||||
| `tp` | `int` | [Profile](#profile) |
|
||||
| `pp` | `int` | [Profile](#profile) |
|
||||
@@ -288,7 +299,10 @@ rather than repeated.
|
||||
| `guard` | `guard.Config` | [Server](#server) |
|
||||
| `secrets` | `Secrets` | [Server](#server) |
|
||||
| `notify` | `notify.Config` | [Server](#server) |
|
||||
| `mail` | `mail.Config` | [Server](#server) |
|
||||
| `serving` | `ServingSettings` | [Server](#server) |
|
||||
| `firewall` | `FirewallSettings` | [Server](#server) |
|
||||
| `agents` | `AgentSettings` | [Server](#server) |
|
||||
| `firewall_networks` | `[]string` | [Server](#server) |
|
||||
| `backup_recipients` | `[]string` | [Server](#server) |
|
||||
| `backup_copy_to` | `[]string` | [Server](#server) |
|
||||
@@ -300,6 +314,7 @@ rather than repeated.
|
||||
| `facility` | `Facility` | [Server](#server) |
|
||||
| `federation` | `Federation` | [Server](#server) |
|
||||
| `endpoints` | `[]Endpoint` | [Server](#server) |
|
||||
| `stopped` | `bool` | [ServingSettings](#servingsettings) |
|
||||
| `bin` | `string` | [SimConfig](#simconfig) |
|
||||
| `python` | `string` | [TRTLLMConfig](#trtllmconfig) |
|
||||
| `store` | `string` | [Threads](#threads) |
|
||||
@@ -372,8 +387,14 @@ rather than repeated.
|
||||
| `llamacpp` | `LlamaCppConfig` | [engines.Config](#enginesconfig) |
|
||||
| `trtllm` | `TRTLLMConfig` | [engines.Config](#enginesconfig) |
|
||||
| `sim` | `SimConfig` | [engines.Config](#enginesconfig) |
|
||||
| `enabled` | `*bool` | [guard.Config](#guardconfig) |
|
||||
| `default` | `[]string` | [guard.Config](#guardconfig) |
|
||||
| `policies` | `map[string]Policy` | [guard.Config](#guardconfig) |
|
||||
| `smtp` | `string` | [mail.Config](#mailconfig) |
|
||||
| `from` | `string` | [mail.Config](#mailconfig) |
|
||||
| `username` | `string` | [mail.Config](#mailconfig) |
|
||||
| `password_file` | `string` | [mail.Config](#mailconfig) |
|
||||
| `base_url` | `string` | [mail.Config](#mailconfig) |
|
||||
| `webhooks` | `[]Webhook` | [notify.Config](#notifyconfig) |
|
||||
| `host` | `string` | [notify.Config](#notifyconfig) |
|
||||
| `work_dir` | `string` | [supervisor.Config](#supervisorconfig) |
|
||||
@@ -411,6 +432,16 @@ weights are.
|
||||
|
||||
Scale is how strongly it is applied; zero takes 1.
|
||||
|
||||
### `owner`
|
||||
|
||||
`string`
|
||||
|
||||
Owner is the part of the tenancy tree the adapter belongs to —
|
||||
org:ID, team:ID, project:ID, user:ID or key:ID — and only callers in
|
||||
it see it: to anyone else it is a model that does not exist. An
|
||||
adapter trained on an organisation's documents carries some of what
|
||||
they say in its weights. Empty is every caller the model allows.
|
||||
|
||||
## Admission
|
||||
|
||||
### `enabled`
|
||||
@@ -457,6 +488,21 @@ Zero takes half of all the bands' capacities together.
|
||||
|
||||
`time.Duration`
|
||||
|
||||
## AgentSettings
|
||||
|
||||
### `default`
|
||||
|
||||
`map[string]any`
|
||||
|
||||
Default is agent.yaml's keys for every node.
|
||||
|
||||
### `nodes`
|
||||
|
||||
`map[string]map[string]any`
|
||||
|
||||
Nodes are agent.yaml's keys for one node each, by node id, over
|
||||
Default.
|
||||
|
||||
## Auth
|
||||
|
||||
### `enabled`
|
||||
@@ -674,6 +720,14 @@ token is MaxTTFTRatio times this model's. Zero is 50, 0.05 and 1.5.
|
||||
|
||||
`bool`
|
||||
|
||||
### `images`
|
||||
|
||||
`bool`
|
||||
|
||||
Images is a model that makes pictures, asked for at
|
||||
/v1/images/generations. Implied by an external model whose API is
|
||||
ComfyUI's.
|
||||
|
||||
## Control
|
||||
|
||||
### `interval`
|
||||
@@ -787,6 +841,13 @@ with one.
|
||||
Engine names the inference engine, which decides what its metrics are
|
||||
called. Empty means vLLM's names.
|
||||
|
||||
### `workflow`
|
||||
|
||||
`string`
|
||||
|
||||
Workflow is the ComfyUI workflow for engine: comfyui, a built-in's
|
||||
name or its JSON; flux-schnell unless set.
|
||||
|
||||
### `pool`
|
||||
|
||||
`string`
|
||||
@@ -878,6 +939,23 @@ a model that federation added (`llmgw ctl federation add`) rather
|
||||
than one written by hand. Such a model is named PEER/MODEL, comes and
|
||||
goes with the peer's own list, and is sent only while the peer is up.
|
||||
|
||||
### `api`
|
||||
|
||||
`string`
|
||||
|
||||
API is what the provider speaks: OpenAI's, unless this says comfyui —
|
||||
a ComfyUI server, asked for pictures at /v1/images/generations and
|
||||
sent the workflow below with the request's prompt in it. Model is then
|
||||
the checkpoint file the workflow loads, as ComfyUI lists it.
|
||||
|
||||
### `workflow`
|
||||
|
||||
`string`
|
||||
|
||||
Workflow is the ComfyUI workflow, in its API format: a built-in's name
|
||||
(flux-schnell, the default, or flux-dev) or a JSON file on the control
|
||||
plane.
|
||||
|
||||
## Facility
|
||||
|
||||
### `budgets`
|
||||
@@ -984,9 +1062,11 @@ every node and control plane within seconds.
|
||||
|
||||
`string`
|
||||
|
||||
Format is llamaguard3 (Meta's Llama Guard 3: "safe", or "unsafe" and
|
||||
the categories), shieldgemma (Google's ShieldGemma: Yes or No), or
|
||||
json ({"safe": false, "categories": ["..."]}).
|
||||
Format is qwen3guard (Qwen3Guard-Gen: "Safety: Safe", "Unsafe" or
|
||||
"Controversial", and the categories), llamaguard3 (Meta's Llama Guard
|
||||
3: "safe", or "unsafe" and the categories), shieldgemma (Google's
|
||||
ShieldGemma: Yes or No), or json ({"safe": false, "categories":
|
||||
["..."]}).
|
||||
|
||||
### `categories`
|
||||
|
||||
@@ -1003,6 +1083,15 @@ sentence — and the names reported for llamaguard3's S1, S2, ….
|
||||
|
||||
Model is a catalog model with role: guard.
|
||||
|
||||
### `models`
|
||||
|
||||
`[]string`
|
||||
|
||||
Models are judges asked in turn, instead of Model: the first that
|
||||
answers gives the verdict. A judge on a GPU first and one on CPUs
|
||||
after it is the GPU's speed while it is there, and the CPUs' when it
|
||||
is not — busy, failed, or the node down.
|
||||
|
||||
### `input`
|
||||
|
||||
`string`
|
||||
@@ -1021,12 +1110,19 @@ Empty does not ask the judge about that direction.
|
||||
Categories, when set, are the only ones that count: Llama Guard's
|
||||
codes (S1, S10) or the names the model reports.
|
||||
|
||||
### `controversial`
|
||||
|
||||
`string`
|
||||
|
||||
Controversial is what Qwen3Guard's third verdict counts as: safe
|
||||
(the default) or unsafe.
|
||||
|
||||
### `latency_budget`
|
||||
|
||||
`time.Duration`
|
||||
|
||||
LatencyBudget is how long the judge is waited for before the policy's
|
||||
fail applies. Zero is 250ms.
|
||||
LatencyBudget is how long each judge is waited for before the next is
|
||||
asked, or, after the last, the policy's fail applies. Zero is 250ms.
|
||||
|
||||
## KVIndex
|
||||
|
||||
@@ -1330,6 +1426,16 @@ Capabilities is what this model can do beyond plain completion.
|
||||
|
||||
Takes a [Capabilities](#capabilities) block.
|
||||
|
||||
### `query_prefix`
|
||||
|
||||
`string`
|
||||
|
||||
QueryPrefix is put before a question an embedding model embeds for a
|
||||
knowledge-base search; documents are embedded without it. Models
|
||||
trained with an instruction on the query side want it, Qwen3-Embedding
|
||||
"Instruct: Given a web search query, retrieve relevant passages that
|
||||
answer the query\nQuery: " for one.
|
||||
|
||||
### `role`
|
||||
|
||||
`string`
|
||||
@@ -1755,6 +1861,14 @@ before affinity is abandoned.
|
||||
|
||||
`string`
|
||||
|
||||
### `engine_version`
|
||||
|
||||
`string`
|
||||
|
||||
EngineVersion restricts the profile to nodes whose installed engine's
|
||||
version starts with it: "0.11" takes vLLM 0.11.0 and 0.11.1, not
|
||||
0.10.2. Empty takes whatever version the node has.
|
||||
|
||||
### `gpu_class`
|
||||
|
||||
`string`
|
||||
@@ -2456,6 +2570,24 @@ Prometheus and Alertmanager in between. Empty sends nothing.
|
||||
|
||||
Takes a [notify.Config](#notifyconfig) block.
|
||||
|
||||
### `mail`
|
||||
|
||||
`mail.Config`
|
||||
|
||||
Mail is where the gateway's few emails go out from: a link to set a
|
||||
forgotten password. Unset, none is sent and people ask an operator.
|
||||
|
||||
Takes a [mail.Config](#mailconfig) block.
|
||||
|
||||
### `serving`
|
||||
|
||||
`ServingSettings`
|
||||
|
||||
Serving stops and starts the whole cluster's service (llmgw stop, llmgw
|
||||
start).
|
||||
|
||||
Takes a [ServingSettings](#servingsettings) block.
|
||||
|
||||
### `firewall`
|
||||
|
||||
`FirewallSettings`
|
||||
@@ -2464,6 +2596,18 @@ Firewall turns llmgw's firewall on every machine on or off.
|
||||
|
||||
Takes a [FirewallSettings](#firewallsettings) block.
|
||||
|
||||
### `agents`
|
||||
|
||||
`AgentSettings`
|
||||
|
||||
Agents are the nodes' agent settings: agent.yaml's keys but those an
|
||||
agent needs to start, for every node (default) and for one
|
||||
(nodes.<id>). Sent to the agents as they change; an agent whose
|
||||
settings changed starts again on them, its engines serving
|
||||
throughout.
|
||||
|
||||
Takes a [AgentSettings](#agentsettings) block.
|
||||
|
||||
### `firewall_networks`
|
||||
|
||||
`[]string`
|
||||
@@ -2554,6 +2698,18 @@ Endpoints is the phase-1 static replica list. It is only consulted when
|
||||
no model declares a profile, which is what makes a router-only smoke test
|
||||
against sim-engines possible without a control plane.
|
||||
|
||||
## ServingSettings
|
||||
|
||||
### `stopped`
|
||||
|
||||
`bool`
|
||||
|
||||
Stopped takes every model down and answers every request for one with
|
||||
503 (system_stopped), leaving the control planes, agents and mirror
|
||||
running and every model's definition as it is: llmgw stop sets it,
|
||||
llmgw start clears it, and the models come back as their catalog
|
||||
entries ask. In force on every control plane within seconds.
|
||||
|
||||
## SimConfig
|
||||
|
||||
### `bin`
|
||||
@@ -3067,6 +3223,13 @@ Takes a [SimConfig](#simconfig) block.
|
||||
|
||||
## guard.Config
|
||||
|
||||
### `enabled`
|
||||
|
||||
`*bool`
|
||||
|
||||
Enabled turns every policy on or off at once: false checks nothing,
|
||||
whatever is defined or attached. Unset is on.
|
||||
|
||||
### `default`
|
||||
|
||||
`[]string`
|
||||
@@ -3075,6 +3238,40 @@ Takes a [SimConfig](#simconfig) block.
|
||||
|
||||
`map[string]Policy`
|
||||
|
||||
## mail.Config
|
||||
|
||||
### `smtp`
|
||||
|
||||
`string`
|
||||
|
||||
SMTP is the relay, host:port: 587 with STARTTLS, 465 with TLS from the
|
||||
start, 25 as the relay allows. Empty sends nothing.
|
||||
|
||||
### `from`
|
||||
|
||||
`string`
|
||||
|
||||
From is the sender, llmgw@example.com or "llmgw <llmgw@example.com>".
|
||||
|
||||
### `username`
|
||||
|
||||
`string`
|
||||
|
||||
Username and PasswordFile log in to the relay, when it wants that.
|
||||
The password is a secret reference: a file, env:NAME or vault:path.
|
||||
|
||||
### `password_file`
|
||||
|
||||
`string`
|
||||
|
||||
### `base_url`
|
||||
|
||||
`string`
|
||||
|
||||
BaseURL is the gateway's public address, which links in mail point
|
||||
at: https://llm.example.com. Never taken from a request, whose Host a
|
||||
sender chooses and could point a reset link at their own site.
|
||||
|
||||
## notify.Config
|
||||
|
||||
### `webhooks`
|
||||
|
||||
@@ -18,6 +18,7 @@ require (
|
||||
github.com/knadh/koanf/providers/env/v2 v2.0.1
|
||||
github.com/knadh/koanf/providers/file v1.2.1
|
||||
github.com/knadh/koanf/v2 v2.3.6
|
||||
github.com/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0
|
||||
github.com/mattn/go-isatty v0.0.24
|
||||
github.com/prometheus/client_golang v1.24.1
|
||||
github.com/prometheus/client_model v0.6.3
|
||||
@@ -34,6 +35,7 @@ require (
|
||||
golang.org/x/net v0.59.0
|
||||
golang.org/x/oauth2 v0.37.0
|
||||
golang.org/x/sync v0.23.0
|
||||
golang.org/x/term v0.46.0
|
||||
golang.org/x/text v0.42.0
|
||||
google.golang.org/protobuf v1.36.12
|
||||
modernc.org/sqlite v1.59.0
|
||||
|
||||
@@ -95,6 +95,8 @@ github.com/knadh/koanf/v2 v2.3.6 h1:JoQPSJmvS4aP0xNc8xMDr5tcrkSEInL23/Il7pITAKo=
|
||||
github.com/knadh/koanf/v2 v2.3.6/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28=
|
||||
github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
|
||||
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
|
||||
github.com/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0 h1:7Q+xNAZFmnfYOMweHN3c/PDFUKKfY1pVJ26K++QvVfU=
|
||||
github.com/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0/go.mod h1:1fEHWurg7pvf5SG6XNE5Q8UZmOwex51Mkx3SLhrW5B4=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
|
||||
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
|
||||
|
||||
@@ -448,3 +448,40 @@ func TestAHashThatCannotGetATurnSaysTheServerIsBusy(t *testing.T) {
|
||||
t.Errorf("a hash with no turn: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAForgottenPasswordIsResetOnceThroughALink(t *testing.T) {
|
||||
svc, st := newService(t)
|
||||
ctx := context.Background()
|
||||
must(t, svc.SignUp(ctx, "ayse@example.com", "A", goodPassword))
|
||||
u, _ := st.UserByEmail(ctx, "ayse@example.com")
|
||||
|
||||
// Waiting for approval: no link, said the same as for nobody.
|
||||
if _, _, err := svc.RequestReset(ctx, st, "ayse@example.com"); !errors.Is(err, ErrNoReset) {
|
||||
t.Fatalf("a pending account got a link: %v", err)
|
||||
}
|
||||
if _, _, err := svc.RequestReset(ctx, st, "nobody@example.com"); !errors.Is(err, ErrNoReset) {
|
||||
t.Fatalf("an unknown address: %v", err)
|
||||
}
|
||||
must(t, st.ApproveUser(ctx, u.ID, "proj", "ilker", time.Now()))
|
||||
session, _, err := svc.SignIn(ctx, "ayse@example.com", goodPassword)
|
||||
must(t, err)
|
||||
|
||||
token, got, err := svc.RequestReset(ctx, st, "AYSE@example.com")
|
||||
if err != nil || token == "" || got.ID != u.ID {
|
||||
t.Fatalf("an active account: %q %+v %v", token, got, err)
|
||||
}
|
||||
if err := svc.ResetPassword(ctx, st, st, token, "short"); !errors.Is(err, ErrWeakPassword) {
|
||||
t.Fatalf("a weak password: %v", err)
|
||||
}
|
||||
next := goodPassword + "-new"
|
||||
must(t, svc.ResetPassword(ctx, st, st, token, next))
|
||||
if _, _, err := svc.Resolve(ctx, session); err == nil {
|
||||
t.Error("a session from before the reset still works")
|
||||
}
|
||||
if _, _, err := svc.SignIn(ctx, "ayse@example.com", next); err != nil {
|
||||
t.Errorf("the new password does not sign in: %v", err)
|
||||
}
|
||||
if err := svc.ResetPassword(ctx, st, st, token, next+"2"); !errors.Is(err, ErrNoReset) {
|
||||
t.Errorf("a link used twice: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -373,3 +373,63 @@ func (s *Service) OpenSession(ctx context.Context, userID string) (token, csrf s
|
||||
}
|
||||
return token, csrf, nil
|
||||
}
|
||||
|
||||
// ResetLinkTTL is how long a password reset link works.
|
||||
const ResetLinkTTL = time.Hour
|
||||
|
||||
// ResetStore keeps password reset links.
|
||||
type ResetStore interface {
|
||||
UserByEmail(ctx context.Context, email string) (store.User, error)
|
||||
CreatePasswordReset(ctx context.Context, userID string, ttl time.Duration) (string, error)
|
||||
PasswordResetUser(ctx context.Context, token string) (string, error)
|
||||
UsePasswordReset(ctx context.Context, token string) (string, error)
|
||||
}
|
||||
|
||||
// ErrNoReset is an address that gets no reset link: none has an account, the
|
||||
// account signs in through single sign-on, or it is not active. Never told to
|
||||
// whoever asked, who learns the same either way.
|
||||
var ErrNoReset = errors.New("no reset link for that address")
|
||||
|
||||
// RequestReset mints a reset link token for the person an address belongs
|
||||
// to, when they sign in with a password and are active. A few a quarter of
|
||||
// an hour per address: each sends a mail.
|
||||
func (s *Service) RequestReset(ctx context.Context, rs ResetStore, email string) (string, store.User, error) {
|
||||
email = store.NormalizeEmail(email)
|
||||
if !validEmail(email) {
|
||||
return "", store.User{}, ErrNoReset
|
||||
}
|
||||
if !s.allowAttempt("reset:"+email, 3) {
|
||||
return "", store.User{}, ErrTooManyTries
|
||||
}
|
||||
u, err := rs.UserByEmail(ctx, email)
|
||||
if err != nil || u.PasswordHash == "" || u.State != store.UserActive {
|
||||
return "", store.User{}, ErrNoReset
|
||||
}
|
||||
token, err := rs.CreatePasswordReset(ctx, u.ID, ResetLinkTTL)
|
||||
if err != nil {
|
||||
return "", store.User{}, err
|
||||
}
|
||||
s.log.Info("a password reset link was made", "user", u.ID)
|
||||
return token, u, nil
|
||||
}
|
||||
|
||||
// ResetPassword sets a password from a reset link, which is spent, and signs
|
||||
// the person out everywhere: if they lost it, someone else may have it.
|
||||
func (s *Service) ResetPassword(ctx context.Context, rs ResetStore, pw PasswordChanger, token, next string) error {
|
||||
hash, err := HashPassword(next)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
userID, err := rs.UsePasswordReset(ctx, token)
|
||||
if err != nil {
|
||||
return ErrNoReset
|
||||
}
|
||||
if err := pw.SetPassword(ctx, userID, hash); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := pw.DeleteUserSessions(ctx, userID); err != nil {
|
||||
return err
|
||||
}
|
||||
s.log.Info("a password was set from a reset link", "user", userID)
|
||||
return nil
|
||||
}
|
||||
+140
-3
@@ -34,6 +34,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/train"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
|
||||
@@ -136,6 +137,13 @@ type Agent struct {
|
||||
// by mu.
|
||||
kv *kvevents.Mirror
|
||||
kvSubs map[string]*kvevents.Subscriber
|
||||
|
||||
// settingsCheck says whether the cluster's settings for this agent
|
||||
// would load beside its agent.yaml; see SetSettingsCheck.
|
||||
settingsCheck func([]byte) error
|
||||
|
||||
// trainer runs LoRA training jobs, one at a time (BACKLOG F4).
|
||||
trainer *train.Runner
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
@@ -182,6 +190,18 @@ func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, er
|
||||
kvSubs: map[string]*kvevents.Subscriber{},
|
||||
}
|
||||
a.prov = &provision.Installer{Root: cfg.EnginesDir, Log: log, OnStatus: a.provisionStatus}
|
||||
runAs := cfg.Superv.RunAs
|
||||
if runAs == "" {
|
||||
runAs = supervisor.DefaultRunAs
|
||||
}
|
||||
// Where the engines are, as the installer has it when the file says
|
||||
// nothing.
|
||||
enginesDir := cfg.EnginesDir
|
||||
if enginesDir == "" {
|
||||
enginesDir = provision.DefaultRoot
|
||||
}
|
||||
a.trainer = train.New(train.Env{EnginesDir: enginesDir, WorkDir: filepath.Join(cfg.Superv.WorkDir, "trainer"),
|
||||
RunAs: runAs, Log: log})
|
||||
return a, nil
|
||||
}
|
||||
|
||||
@@ -371,6 +391,12 @@ func (a *Agent) settleAdopted(sh *supervisor.Shard, o supervisor.Orphan) {
|
||||
_ = a.sup.Stop(ctx, o.ShardID)
|
||||
return
|
||||
}
|
||||
// What it computes on, for one started before the agent kept that.
|
||||
if a.sup.ComputeOf(o.ShardID) == nil {
|
||||
if c, ok := a.gpuCompute(ctx, o.Spec.GPUIndices, sh.Status().PID); ok {
|
||||
a.sup.SetCompute(o.ShardID, c)
|
||||
}
|
||||
}
|
||||
a.log.Info("adopted engine is serving",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID)
|
||||
}
|
||||
@@ -411,6 +437,17 @@ func (a *Agent) weights() *weights.Fetcher {
|
||||
}
|
||||
|
||||
// runningShards describes what this agent is supervising, for registration.
|
||||
// unready are the shards a registration reported as not yet ready.
|
||||
func unready(running []*llmgwv1.RunningShard) map[string]bool {
|
||||
out := map[string]bool{}
|
||||
for _, r := range running {
|
||||
if r.GetState() != agentsvc.StateToWire(supervisor.Ready) {
|
||||
out[r.GetShardId()] = true
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (a *Agent) runningShards() []*llmgwv1.RunningShard {
|
||||
var out []*llmgwv1.RunningShard
|
||||
for _, st := range a.sup.Statuses() {
|
||||
@@ -436,6 +473,7 @@ func (a *Agent) runningShards() []*llmgwv1.RunningShard {
|
||||
Pool: spec.Pool,
|
||||
Fabric: spec.Fabric,
|
||||
ProfileFp: spec.ProfileFP,
|
||||
Compute: computeToWire(a.sup.ComputeOf(st.ID)),
|
||||
LeaderAddr: a.cfg.MgmtAddr,
|
||||
LeaderPort: int32(spec.LeaderPort),
|
||||
})
|
||||
@@ -564,8 +602,9 @@ func (a *Agent) session(ctx context.Context, server string) error {
|
||||
a.sendLoop(ctx, stream, out)
|
||||
}()
|
||||
defer wg.Wait()
|
||||
go a.watchExits(ctx, out)
|
||||
go a.watchExits(ctx, out, unready(reg.RunningShards))
|
||||
go a.forwardKV(ctx, out)
|
||||
go a.reportCompute(ctx, out)
|
||||
|
||||
for {
|
||||
msg, err := stream.Receive()
|
||||
@@ -861,6 +900,20 @@ func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out cha
|
||||
a.setControlPlanes(m.ControlPlanes)
|
||||
case *llmgwv1.ControlMessage_Firewall:
|
||||
go a.applyFirewall(m.Firewall)
|
||||
case *llmgwv1.ControlMessage_AgentSettings:
|
||||
go a.applySettings(m.AgentSettings)
|
||||
case *llmgwv1.ControlMessage_TrainStart:
|
||||
go a.trainJob(ctx, m.TrainStart, out)
|
||||
case *llmgwv1.ControlMessage_TrainCancel:
|
||||
a.trainer.Cancel(m.TrainCancel.GetJobId())
|
||||
case *llmgwv1.ControlMessage_RestartAgent:
|
||||
go func() {
|
||||
a.log.Info("restarting, as the control plane asks; the engines keep serving",
|
||||
"reason", m.RestartAgent.GetReason())
|
||||
// A moment for the log line to leave before the process does.
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
exitForUpgrade()
|
||||
}()
|
||||
case *llmgwv1.ControlMessage_Logs:
|
||||
reply := a.logsReply(m.Logs)
|
||||
// Off the receive loop: a full send queue must not stop commands being
|
||||
@@ -1130,6 +1183,14 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
} else {
|
||||
err = a.sup.WaitAlive(ctx, sh)
|
||||
}
|
||||
// Where it put the model, by its own account: said with the readiness,
|
||||
// and a GPU placement it computes on the CPU refused even where
|
||||
// nvidia-smi cannot be asked.
|
||||
compute, said, cerr := engineCompute(eng, spec.GPUIndices, sh.Logs(512))
|
||||
if err == nil && cerr != nil {
|
||||
err = cerr
|
||||
_ = a.sup.Stop(ctx, shardID)
|
||||
}
|
||||
if err == nil {
|
||||
err = a.onItsGPUs(ctx, eng, spec, sh.Status().PID)
|
||||
if err != nil {
|
||||
@@ -1142,7 +1203,19 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
report(supervisor.Failed, err.Error(), sh.Logs(20))
|
||||
return
|
||||
}
|
||||
report(supervisor.Ready, "", nil)
|
||||
// The engine's own account when it gave one, with its layers; else what
|
||||
// nvidia-smi sees the shard holding.
|
||||
if !said {
|
||||
compute, said = a.gpuCompute(ctx, spec.GPUIndices, sh.Status().PID)
|
||||
}
|
||||
if said {
|
||||
// Kept with the shard, so an agent restarted, or a control plane
|
||||
// taking over, still learns it.
|
||||
a.sup.SetCompute(shardID, compute)
|
||||
send(&llmgwv1.ShardStatus{State: agentsvc.StateToWire(supervisor.Ready), Compute: computeToWire(&compute)})
|
||||
} else {
|
||||
report(supervisor.Ready, "", nil)
|
||||
}
|
||||
// Its dying later is reported by the session's watchExits.
|
||||
}
|
||||
|
||||
@@ -1155,7 +1228,14 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
// reported the engine dying; and an adopted engine had no watcher at all —
|
||||
// on node06 one was killed and its replica stayed "ready" for the control
|
||||
// plane, routed to and failing, indefinitely.
|
||||
func (a *Agent) watchExits(ctx context.Context, out chan<- *llmgwv1.AgentMessage) {
|
||||
//
|
||||
// It also reports the shards the registration said were not ready yet once
|
||||
// they are. An adopted engine is health-checked after the agent registers,
|
||||
// and a shard started on the stream before is reported on that stream; so
|
||||
// when the registration came first nothing ever told the control plane, and
|
||||
// the replica stayed "starting" while it served — on node05, after a
|
||||
// control plane restart, the only replica of a model.
|
||||
func (a *Agent) watchExits(ctx context.Context, out chan<- *llmgwv1.AgentMessage, pending map[string]bool) {
|
||||
t := time.NewTicker(time.Second)
|
||||
defer t.Stop()
|
||||
ready := map[string]bool{}
|
||||
@@ -1164,6 +1244,24 @@ func (a *Agent) watchExits(ctx context.Context, out chan<- *llmgwv1.AgentMessage
|
||||
switch st.State {
|
||||
case supervisor.Ready:
|
||||
ready[st.ID] = true
|
||||
if !pending[st.ID] {
|
||||
continue
|
||||
}
|
||||
delete(pending, st.ID)
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: sh.Spec.ReplicaID, ShardId: st.ID,
|
||||
State: agentsvc.StateToWire(st.State),
|
||||
Compute: computeToWire(a.sup.ComputeOf(st.ID)),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
case supervisor.Failed:
|
||||
if !ready[st.ID] {
|
||||
continue // never ready here: its start reported it
|
||||
@@ -1373,3 +1471,42 @@ func ResolveManagement(c Config) (addr, iface string) {
|
||||
}
|
||||
return addr, iface
|
||||
}
|
||||
|
||||
// reportCompute tells the control plane, once a session, what each running
|
||||
// shard computes on as soon as that is known: an adopted engine's is found
|
||||
// after registration, which carried what was known then.
|
||||
func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMessage) {
|
||||
sent := map[string]bool{}
|
||||
t := time.NewTicker(5 * time.Second)
|
||||
defer t.Stop()
|
||||
for range 60 {
|
||||
for _, st := range a.sup.Statuses() {
|
||||
if st.State.Terminal() || sent[st.ID] {
|
||||
continue
|
||||
}
|
||||
c := a.sup.ComputeOf(st.ID)
|
||||
if c == nil {
|
||||
continue
|
||||
}
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
// No state: a report of compute alone, not a transition.
|
||||
msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
|
||||
NodeId: a.cfg.NodeID, ReplicaId: sh.Spec.ReplicaID, ShardId: st.ID, Compute: computeToWire(c),
|
||||
}}}
|
||||
select {
|
||||
case out <- msg:
|
||||
sent[st.ID] = true
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -37,11 +37,12 @@ import (
|
||||
type statusRecorder struct {
|
||||
ch chan status
|
||||
|
||||
mu sync.Mutex
|
||||
calls []adoptCall
|
||||
adopted []reconciler.AdoptedShard
|
||||
unclaim []string // shard ids to report back as unclaimed
|
||||
fetches []fetchReport
|
||||
mu sync.Mutex
|
||||
calls []adoptCall
|
||||
adopted []reconciler.AdoptedShard
|
||||
unclaim []string // shard ids to report back as unclaimed
|
||||
fetches []fetchReport
|
||||
computes []reconciler.Compute
|
||||
}
|
||||
|
||||
type status struct {
|
||||
@@ -61,6 +62,12 @@ func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.
|
||||
|
||||
// fetches records download progress, so a test can assert that a fetch is
|
||||
// visible from the control plane rather than only in the node's log.
|
||||
func (r *statusRecorder) OnShardCompute(replicaID, shardID string, c reconciler.Compute) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.computes = append(r.computes, c)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) OnFetchProgress(replicaID, shardID, source string, done, total int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
@@ -103,6 +103,11 @@ type Config struct {
|
||||
Tracing obs.Tracing `koanf:"tracing"`
|
||||
|
||||
Version string
|
||||
|
||||
// Settings is the cluster's settings for this agent it was started on,
|
||||
// as the control plane sent them: what a newly sent version is
|
||||
// compared with.
|
||||
Settings string `koanf:"-"`
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Compute is where an engine runs its model, as it said while loading it.
|
||||
type Compute struct {
|
||||
// Backend is CUDA, Vulkan, Metal, ROCm or CPU.
|
||||
Backend string
|
||||
// Devices are those it named: "CUDA0: Quadro RTX 6000 (7.5)".
|
||||
Devices []string
|
||||
// Offloaded of Total layers went to the devices; 0 and 0 when it did
|
||||
// not say.
|
||||
Offloaded, Total int
|
||||
// NoGPU is the engine saying it found no device it could use.
|
||||
NoGPU bool
|
||||
}
|
||||
|
||||
var (
|
||||
offloadedRE = regexp.MustCompile(`offloaded (\d+)/(\d+) layers to GPU`)
|
||||
usingRE = regexp.MustCompile(`using device ([A-Za-z]+\d+) \(([^)]*)\)`)
|
||||
ccRE = regexp.MustCompile(`Device (\d+): ([^,]+), compute capability ([0-9.]+)`)
|
||||
noDevicesRE = regexp.MustCompile(`found 0 (CUDA|Vulkan|ROCm) devices|no usable GPU found`)
|
||||
)
|
||||
|
||||
// ReadLlamaCompute reads llama.cpp's start-up output for where it put the
|
||||
// model, reporting false when the output says nothing of it.
|
||||
func ReadLlamaCompute(lines []string) (Compute, bool) {
|
||||
var c Compute
|
||||
known := false
|
||||
cc := map[string]string{} // device index → compute capability
|
||||
for _, l := range lines {
|
||||
if m := ccRE.FindStringSubmatch(l); m != nil {
|
||||
cc[m[1]] = m[3]
|
||||
}
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
for _, l := range lines {
|
||||
if m := offloadedRE.FindStringSubmatch(l); m != nil {
|
||||
c.Offloaded, _ = strconv.Atoi(m[1])
|
||||
c.Total, _ = strconv.Atoi(m[2])
|
||||
known = true
|
||||
}
|
||||
if m := usingRE.FindStringSubmatch(l); m != nil && !seen[m[1]] {
|
||||
seen[m[1]] = true
|
||||
d := m[1] + ": " + strings.TrimSpace(m[2])
|
||||
idx := strings.TrimLeft(m[1], "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
|
||||
if v, ok := cc[idx]; ok && strings.HasPrefix(m[1], "CUDA") {
|
||||
d += " (" + v + ")"
|
||||
}
|
||||
c.Devices = append(c.Devices, d)
|
||||
if c.Backend == "" {
|
||||
c.Backend = strings.TrimRight(m[1], "0123456789")
|
||||
}
|
||||
known = true
|
||||
}
|
||||
if noDevicesRE.MatchString(l) {
|
||||
c.NoGPU, known = true, true
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case c.OnTheCPU():
|
||||
c.Backend = "CPU"
|
||||
case c.Backend == "" && c.Offloaded > 0:
|
||||
c.Backend = "GPU"
|
||||
case c.Backend == "":
|
||||
c.Backend = "CPU"
|
||||
}
|
||||
return c, known
|
||||
}
|
||||
|
||||
// OnTheCPU is whether a model the engine was given GPUs for is computed on
|
||||
// the CPU, by its own account: no device found, or no layer offloaded.
|
||||
func (c Compute) OnTheCPU() bool {
|
||||
return c.NoGPU || (c.Total > 0 && c.Offloaded == 0)
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLlamaCppSaysWhereItPutTheModel(t *testing.T) {
|
||||
gpu := strings.Split(`ggml_cuda_init: found 1 CUDA devices:
|
||||
Device 0: Quadro RTX 6000, compute capability 7.5, VMM: yes
|
||||
load_backend: loaded CUDA backend from /opt/llmgw/engines/llamacpp/b10991/libggml-cuda.so
|
||||
llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free
|
||||
load_tensors: offloading 48 repeating layers to GPU
|
||||
load_tensors: offloaded 49/49 layers to GPU`, "\n")
|
||||
c, ok := ReadLlamaCompute(gpu)
|
||||
if !ok || c.Backend != "CUDA" || c.Offloaded != 49 || c.Total != 49 || c.OnTheCPU() ||
|
||||
len(c.Devices) != 1 || c.Devices[0] != "CUDA0: Quadro RTX 6000 (7.5)" {
|
||||
t.Errorf("on the GPU: %+v %v", c, ok)
|
||||
}
|
||||
|
||||
cpu := strings.Split(`warning: no usable GPU found, --gpu-layers option will be ignored
|
||||
warning: one possible reason is that llama.cpp was compiled without GPU support
|
||||
load_tensors: offloaded 0/49 layers to GPU`, "\n")
|
||||
c, ok = ReadLlamaCompute(cpu)
|
||||
if !ok || c.Backend != "CPU" || !c.OnTheCPU() {
|
||||
t.Errorf("fell back to the CPU: %+v %v", c, ok)
|
||||
}
|
||||
|
||||
if _, ok := ReadLlamaCompute([]string{"main: server is listening on http://127.0.0.1:20001"}); ok {
|
||||
t.Error("output that says nothing was read as saying something")
|
||||
}
|
||||
}
|
||||
@@ -114,6 +114,10 @@ type ShardSpec struct {
|
||||
// lets the control plane tell, after its own restart, whether a shard that
|
||||
// is still running matches the profile the catalog holds now.
|
||||
ProfileFP string
|
||||
|
||||
// Compute is where the engine said it put the model as it became ready,
|
||||
// kept with the shard's state so an agent restarted later still knows.
|
||||
Compute *Compute `json:",omitempty"`
|
||||
}
|
||||
|
||||
// Host is where the engine listens.
|
||||
|
||||
@@ -277,6 +277,9 @@ func (l *LlamaCpp) bin() string {
|
||||
return "llama-server"
|
||||
}
|
||||
|
||||
// Binary is the llama-server this engine runs.
|
||||
func (l *LlamaCpp) Binary() string { return l.bin() }
|
||||
|
||||
// FallsBackToCPU says that llama.cpp keeps running when its GPU backend does
|
||||
// not load: --n-gpu-layers asks for layers on a GPU it then has none of, and
|
||||
// the server comes up and answers on the CPU, at a tenth of the speed, with no
|
||||
|
||||
@@ -4,10 +4,12 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
)
|
||||
@@ -153,3 +155,70 @@ func joinIndices(is []int) string {
|
||||
}
|
||||
return strings.Join(s, ", ")
|
||||
}
|
||||
|
||||
// engineCompute is where an engine said it put its model, from its output
|
||||
// so far — llama.cpp's, which says — and false for one that said nothing. A
|
||||
// shard given GPUs that its engine says it computes on the CPU is an error,
|
||||
// found this way even where nvidia-smi cannot be asked.
|
||||
func engineCompute(eng engines.Engine, gpus []int, lines []string) (engines.Compute, bool, error) {
|
||||
if _, ok := eng.(engines.CPUFallback); !ok {
|
||||
return engines.Compute{}, false, nil
|
||||
}
|
||||
read, ok := engines.ReadLlamaCompute(lines)
|
||||
if !ok {
|
||||
return engines.Compute{}, false, nil
|
||||
}
|
||||
if len(gpus) > 0 && read.OnTheCPU() {
|
||||
why := "its GPU backend saw no device it could use"
|
||||
if read.Total > 0 && !read.NoGPU {
|
||||
why = fmt.Sprintf("it offloaded %d of %d layers", read.Offloaded, read.Total)
|
||||
}
|
||||
return read, true, fmt.Errorf("placed on GPU %v but computing on the CPU: %s, by llama.cpp's own account "+
|
||||
"(llmgw ctl engine-check on this node says why)", gpus, why)
|
||||
}
|
||||
return read, true, nil
|
||||
}
|
||||
|
||||
// computeToWire is a compute as it travels to the control plane.
|
||||
func computeToWire(c *engines.Compute) *llmgwv1.Compute {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
return &llmgwv1.Compute{Backend: c.Backend, Devices: c.Devices,
|
||||
LayersOffloaded: int32(c.Offloaded), LayersTotal: int32(c.Total)} //nolint:gosec // layer counts
|
||||
}
|
||||
|
||||
// gpuCompute is where a shard computes as nvidia-smi sees it: the CPU for a
|
||||
// shard given no GPU, else the GPUs whose memory its processes hold, and how
|
||||
// much. False when nvidia-smi cannot be asked. Engine-agnostic, and needing
|
||||
// no log line: llama.cpp says where its layers went only at -lv 4.
|
||||
func (a *Agent) gpuCompute(ctx context.Context, gpus []int, pid int) (engines.Compute, bool) {
|
||||
if len(gpus) == 0 {
|
||||
return engines.Compute{Backend: "CPU"}, true
|
||||
}
|
||||
mem, err := gpuMemory(ctx, a.cfg.Discovery.NvidiaSMI)
|
||||
if err != nil || mem == nil || pid <= 0 {
|
||||
return engines.Compute{}, false
|
||||
}
|
||||
group := processGroup(pid)
|
||||
c := engines.Compute{Backend: "CUDA"}
|
||||
for _, g := range mem {
|
||||
if !slices.Contains(gpus, g.Index) {
|
||||
continue
|
||||
}
|
||||
used := 0
|
||||
for _, p := range g.Procs {
|
||||
if holds([]int{p.PID}, pid, group) {
|
||||
used += p.UsedMB
|
||||
}
|
||||
}
|
||||
if used > 0 {
|
||||
c.Devices = append(c.Devices, fmt.Sprintf("GPU %d: %.1f GB", g.Index, float64(used)/1024))
|
||||
}
|
||||
}
|
||||
if len(c.Devices) == 0 {
|
||||
// Placed on GPUs and holding none of their memory.
|
||||
c.Backend = "CPU"
|
||||
}
|
||||
return c, true
|
||||
}
|
||||
@@ -122,3 +122,55 @@ func TestForeignGPUMemoryLeavesOutTheAgentsOwnShards(t *testing.T) {
|
||||
t.Errorf("no reading came back as %v", foreign)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLlamaCppsOwnWordOnTheCPUIsEnoughToRefuse(t *testing.T) {
|
||||
eng := fallbackEngine{}
|
||||
// Without nvidia-smi to ask, the engine's own output says it.
|
||||
c, said, err := engineCompute(eng, []int{0}, []string{
|
||||
"warning: no usable GPU found, --gpu-layers option will be ignored",
|
||||
"load_tensors: offloaded 0/49 layers to GPU"})
|
||||
if err == nil || !said || !strings.Contains(err.Error(), "computing on the CPU") || c.Backend != "CPU" {
|
||||
t.Fatalf("a GPU shard llama.cpp put on the CPU: %v %+v", err, c)
|
||||
}
|
||||
c, said, err = engineCompute(eng, []int{0}, []string{
|
||||
"llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free",
|
||||
"load_tensors: offloaded 49/49 layers to GPU"})
|
||||
if err != nil || !said || c.Backend != "CUDA" || c.Offloaded != 49 || computeToWire(&c).GetLayersOffloaded() != 49 {
|
||||
t.Errorf("on its GPU: %v %+v", err, c)
|
||||
}
|
||||
// A CPU shard is not refused for being on the CPU, and says so.
|
||||
if c, _, err := engineCompute(eng, nil, []string{"load_tensors: offloaded 0/49 layers to GPU"}); err != nil || c.Backend != "CPU" {
|
||||
t.Errorf("a CPU shard: %v %+v", err, c)
|
||||
}
|
||||
// An engine that said nothing says nothing.
|
||||
if _, said, err := engineCompute(eng, []int{0}, []string{"listening"}); err != nil || said {
|
||||
t.Errorf("nothing said: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNvidiaSMISaysWhereAShardComputes(t *testing.T) {
|
||||
oldMem, oldGroup := gpuMemory, processGroup
|
||||
t.Cleanup(func() { gpuMemory, processGroup = oldMem, oldGroup })
|
||||
processGroup = func(pid int) int { return pid }
|
||||
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) {
|
||||
return []discovery.GPUMemory{
|
||||
{Index: 0, UsedMB: 20000, Procs: []discovery.GPUProcess{{PID: 700, UsedMB: 2048}, {PID: 900, UsedMB: 17000}}},
|
||||
{Index: 1, UsedMB: 0},
|
||||
}, nil
|
||||
}
|
||||
a := &Agent{log: slog.Default()}
|
||||
if c, ok := a.gpuCompute(context.Background(), []int{0}, 700); !ok || c.Backend != "CUDA" ||
|
||||
len(c.Devices) != 1 || c.Devices[0] != "GPU 0: 2.0 GB" {
|
||||
t.Errorf("on GPU 0: %+v %v", c, ok)
|
||||
}
|
||||
if c, ok := a.gpuCompute(context.Background(), []int{1}, 700); !ok || c.Backend != "CPU" {
|
||||
t.Errorf("placed on GPU 1, holding none of it: %+v", c)
|
||||
}
|
||||
if c, ok := a.gpuCompute(context.Background(), nil, 700); !ok || c.Backend != "CPU" {
|
||||
t.Errorf("a CPU shard: %+v", c)
|
||||
}
|
||||
gpuMemory = func(context.Context, string) ([]discovery.GPUMemory, error) { return nil, nil }
|
||||
if _, ok := a.gpuCompute(context.Background(), []int{0}, 700); ok {
|
||||
t.Error("no nvidia-smi read as an answer")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
// The cluster's settings for this agent (BACKLOG E3): what agent.yaml can
|
||||
// say but the agent's id, control planes, address and certificate, kept by
|
||||
// the control planes and sent over the stream. Kept here beside agent.yaml,
|
||||
// which then holds only what the agent needs to start; at start the two are
|
||||
// read together, these over the file.
|
||||
|
||||
// SettingsFile keeps the settings the control plane last sent.
|
||||
var SettingsFile = "/var/lib/llmgw-agent/agent-settings.yaml"
|
||||
|
||||
// ReadSettings is what SettingsFile holds, nil when nothing.
|
||||
func ReadSettings() []byte {
|
||||
b, err := os.ReadFile(SettingsFile)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// SetSettingsCheck says how to tell whether settings would load beside this
|
||||
// agent's own file; settings that would not are not taken.
|
||||
func (a *Agent) SetSettingsCheck(f func([]byte) error) { a.settingsCheck = f }
|
||||
|
||||
// applySettings takes the settings the control plane sent: nothing when they
|
||||
// are what this agent runs on, and otherwise kept and started again on.
|
||||
func (a *Agent) applySettings(m *llmgwv1.AgentSettings) {
|
||||
next := m.GetYaml()
|
||||
if next == a.cfg.Settings {
|
||||
return
|
||||
}
|
||||
if a.settingsCheck != nil {
|
||||
if err := a.settingsCheck([]byte(next)); err != nil {
|
||||
// Kept running on what it has: settings that will not load
|
||||
// here must not take the node out of service.
|
||||
a.log.Error("the cluster's settings for this agent do not load beside its agent.yaml; running on what it has",
|
||||
"err", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(SettingsFile), 0o755); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
tmp := SettingsFile + ".new"
|
||||
if err := os.WriteFile(tmp, []byte(next), 0o600); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
if err := os.Rename(tmp, SettingsFile); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
a.log.Info("the cluster's settings for this agent changed; starting again on them, the engines keep serving")
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
exitForUpgrade()
|
||||
}
|
||||
|
||||
// WatchFile starts the agent again when its configuration file changes and
|
||||
// still loads: agent.yaml is read only at start, and a change to it that
|
||||
// waited for somebody to restart the agent was a change nobody saw take
|
||||
// effect.
|
||||
func (a *Agent) WatchFile(ctx context.Context, path string, loads func() error) {
|
||||
if path == "" {
|
||||
return
|
||||
}
|
||||
sum := func() [32]byte {
|
||||
b, _ := os.ReadFile(path) //nolint:gosec // the agent's own configuration
|
||||
return sha256.Sum256(b)
|
||||
}
|
||||
was := sum()
|
||||
t := time.NewTicker(watchEvery)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
now := sum()
|
||||
if now == was {
|
||||
continue
|
||||
}
|
||||
if err := loads(); err != nil {
|
||||
// Once per change: an editor saving halfway must not stop the
|
||||
// agent, and a broken file is reported, not restarted into.
|
||||
a.log.Error("agent.yaml changed and does not load; running on what it had", "path", path, "err", err)
|
||||
was = now
|
||||
continue
|
||||
}
|
||||
a.log.Info("agent.yaml changed; starting again on it, the engines keep serving", "path", path)
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
exitForUpgrade()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// watchEvery is how often agent.yaml is looked at.
|
||||
var watchEvery = 10 * time.Second
|
||||
@@ -0,0 +1,88 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
func settingsAgent(t *testing.T, running string) (*Agent, chan struct{}) {
|
||||
t.Helper()
|
||||
oldFile, oldExit := SettingsFile, exitForUpgrade
|
||||
SettingsFile = filepath.Join(t.TempDir(), "agent-settings.yaml")
|
||||
exited := make(chan struct{}, 1)
|
||||
exitForUpgrade = func() { exited <- struct{}{} }
|
||||
t.Cleanup(func() { SettingsFile, exitForUpgrade = oldFile, oldExit })
|
||||
cfg := Config{Server: "127.0.0.1:7443", NodeID: "n1", WeightsCache: t.TempDir(), EnginesDir: t.TempDir(), Settings: running}
|
||||
cfg.Superv.WorkDir = t.TempDir()
|
||||
a, err := New(cfg, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return a, exited
|
||||
}
|
||||
|
||||
// exitedSoon is whether the agent exited within a second: it waits a
|
||||
// moment for its log line to leave first.
|
||||
func exitedSoon(ch chan struct{}) bool {
|
||||
select {
|
||||
case <-ch:
|
||||
return true
|
||||
case <-time.After(time.Second):
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAgentOnItsSettingsStaysAndOneOffThemStartsAgain(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "heartbeat_interval: 5s\n")
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "heartbeat_interval: 5s\n"})
|
||||
if exitedSoon(exited) || ReadSettings() != nil {
|
||||
t.Error("settings it already runs on restarted it")
|
||||
}
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "heartbeat_interval: 2s\n"})
|
||||
if !exitedSoon(exited) || string(ReadSettings()) != "heartbeat_interval: 2s\n" {
|
||||
t.Errorf("new settings: kept %q", ReadSettings())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsThatWouldNotLoadAreNotTaken(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "")
|
||||
a.SetSettingsCheck(func([]byte) error { return errors.New("fabric ib-a lists no HCAs") })
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "fabrics: [{name: ib-a}]\n"})
|
||||
if exitedSoon(exited) || ReadSettings() != nil {
|
||||
t.Error("settings that do not load were taken")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAChangedAgentYamlThatLoadsStartsTheAgentAgain(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "")
|
||||
old := watchEvery
|
||||
watchEvery = 10 * time.Millisecond
|
||||
t.Cleanup(func() { watchEvery = old })
|
||||
p := filepath.Join(t.TempDir(), "agent.yaml")
|
||||
_ = os.WriteFile(p, []byte("node_id: n1\n"), 0o600)
|
||||
var loads atomic.Bool
|
||||
check := func() error {
|
||||
if loads.Load() {
|
||||
return nil
|
||||
}
|
||||
return errors.New("broken")
|
||||
}
|
||||
go a.WatchFile(t.Context(), p, check)
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
// Saved broken: reported, not restarted into.
|
||||
_ = os.WriteFile(p, []byte("node_id: [\n"), 0o600)
|
||||
if exitedSoon(exited) {
|
||||
t.Fatal("a broken agent.yaml restarted the agent")
|
||||
}
|
||||
loads.Store(true)
|
||||
_ = os.WriteFile(p, []byte("node_id: n1\nlabels: {a: b}\n"), 0o600)
|
||||
if !exitedSoon(exited) {
|
||||
t.Error("a changed agent.yaml did not restart the agent")
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,7 @@ package supervisor
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"os/user"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
@@ -90,3 +91,25 @@ func credentialFor(u *user.User, workDir string) (*runAs, error) {
|
||||
func (r *runAs) env() []string {
|
||||
return []string{"HOME=" + r.home, "XDG_CACHE_HOME=" + filepath.Join(r.home, ".cache"), "USER=" + r.name, "LOGNAME=" + r.name}
|
||||
}
|
||||
|
||||
// AsEngineUser has cmd run as the engine user, as an engine would, with its
|
||||
// home under workDir: for work the agent does that is not an engine but
|
||||
// reads what an engine reads — training an adapter. It answers that home
|
||||
// and a way to hand a path to the user, or why it runs as this process.
|
||||
func AsEngineUser(cmd *exec.Cmd, name, workDir string) (home string, give func(path string) error, why string) {
|
||||
r, why := resolveRunAs(name, workDir)
|
||||
if r == nil {
|
||||
return "", func(string) error { return nil }, why
|
||||
}
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{Credential: r.cred, Setpgid: true}
|
||||
cmd.Env = append(cmd.Env, r.env()...)
|
||||
uid, gid := int(r.cred.Uid), int(r.cred.Gid)
|
||||
return r.home, func(path string) error {
|
||||
return filepath.WalkDir(path, func(p string, _ os.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Lchown(p, uid, gid)
|
||||
})
|
||||
}, ""
|
||||
}
|
||||
@@ -113,7 +113,10 @@ type Shard struct {
|
||||
exitCode int
|
||||
|
||||
logs *ringBuffer
|
||||
cmd *exec.Cmd
|
||||
// cmdline is an adopted engine's command line, from its state file: it
|
||||
// has no cmd of this agent's to read it from.
|
||||
cmdline []string
|
||||
cmd *exec.Cmd
|
||||
// sidecars are the shard's supporting processes, torn down with it.
|
||||
sidecars []*exec.Cmd
|
||||
cancel context.CancelFunc
|
||||
@@ -644,6 +647,32 @@ func (s *Supervisor) Get(id string) (*Shard, error) {
|
||||
return sh, nil
|
||||
}
|
||||
|
||||
// SetCompute records where a shard's engine said it put the model, with the
|
||||
// shard's state on disk.
|
||||
func (s *Supervisor) SetCompute(id string, c engines.Compute) {
|
||||
s.mu.Lock()
|
||||
sh, ok := s.shards[id]
|
||||
if ok {
|
||||
sh.Spec.Compute = &c
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if ok && sh.pid > 0 {
|
||||
_ = s.writeState(sh)
|
||||
}
|
||||
}
|
||||
|
||||
// ComputeOf is where a shard's engine was found to compute, nil when not
|
||||
// yet known.
|
||||
func (s *Supervisor) ComputeOf(id string) *engines.Compute {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if sh, ok := s.shards[id]; ok && sh.Spec.Compute != nil {
|
||||
c := *sh.Spec.Compute
|
||||
return &c
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Remove forgets a terminated shard.
|
||||
func (s *Supervisor) Remove(id string) {
|
||||
s.mu.Lock()
|
||||
@@ -703,11 +732,17 @@ type shardState struct {
|
||||
|
||||
// writeState records everything needed to recognise and describe the process later.
|
||||
func (s *Supervisor) writeState(sh *Shard) error {
|
||||
// An adopted engine was started by an earlier agent: its command line
|
||||
// is the one its state file kept.
|
||||
cmdline := sh.cmdline
|
||||
if sh.cmd != nil {
|
||||
cmdline = sh.cmd.Args
|
||||
}
|
||||
b, err := json.Marshal(shardState{
|
||||
ShardID: sh.ID,
|
||||
ReplicaID: sh.Spec.ReplicaID,
|
||||
PID: sh.pid,
|
||||
Cmdline: sh.cmd.Args,
|
||||
Cmdline: cmdline,
|
||||
StartedAt: sh.started,
|
||||
Spec: sh.Spec,
|
||||
})
|
||||
@@ -804,6 +839,7 @@ func (s *Supervisor) Adopt(o Orphan) (*Shard, error) {
|
||||
started: o.StartedAt,
|
||||
done: make(chan struct{}),
|
||||
adopted: true,
|
||||
cmdline: o.Cmdline,
|
||||
}
|
||||
s.shards[o.ShardID] = sh
|
||||
s.mu.Unlock()
|
||||
|
||||
@@ -848,3 +848,60 @@ func eventually(t *testing.T, limit time.Duration, msg string, cond func() bool)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWhatAShardComputesOnSurvivesTheAgent(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = s.Stop(context.Background(), "a") })
|
||||
_ = sh
|
||||
s.SetCompute("a", engines.Compute{Backend: "CUDA", Offloaded: 49, Total: 49})
|
||||
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
orphans, err := s2.FindOrphans()
|
||||
if err != nil || len(orphans) != 1 {
|
||||
t.Fatalf("%v %d", err, len(orphans))
|
||||
}
|
||||
if c := orphans[0].Spec.Compute; c == nil || c.Backend != "CUDA" || c.Offloaded != 49 {
|
||||
t.Errorf("not kept: %+v", c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAdoptedShardsComputeIsKeptWithoutCrashing(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100}
|
||||
if _, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// An agent restart: the engine is adopted by a new supervisor, which did
|
||||
// not start it and has no command of its own for it.
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
orphans, _ := s2.FindOrphans()
|
||||
if len(orphans) != 1 {
|
||||
t.Fatalf("%d orphans", len(orphans))
|
||||
}
|
||||
if _, err := s2.Adopt(orphans[0]); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = s2.Stop(context.Background(), "a") })
|
||||
s2.SetCompute("a", engines.Compute{Backend: "CPU"})
|
||||
|
||||
// And a third agent still finds it, with its command line and compute.
|
||||
s3, _ := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
again, _ := s3.FindOrphans()
|
||||
if len(again) != 1 || len(again[0].Cmdline) == 0 || again[0].Spec.Compute == nil {
|
||||
t.Fatalf("lost: %+v", again)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,574 @@
|
||||
// Package train runs LoRA training jobs on a node (BACKLOG F4): with an
|
||||
// installed engine's Python — its torch and transformers, already matched to
|
||||
// the node's CUDA — and a small overlay of what training needs besides, so a
|
||||
// node that serves can train without a second multi-gigabyte environment.
|
||||
package train
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bufio"
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
_ "embed"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
)
|
||||
|
||||
//go:embed trainer.py
|
||||
var trainerPy []byte
|
||||
|
||||
// overlayPackages are what training needs besides an engine's Python, pinned
|
||||
// and installed without their dependencies: those are the engine's own, and
|
||||
// letting an installer resolve them pulled a second torch.
|
||||
var overlayPackages = []string{"peft==0.21.2", "accelerate==1.15.0"}
|
||||
|
||||
// Job is one training job, as the control plane sent it.
|
||||
type Job struct {
|
||||
ID string
|
||||
BaseRepo string
|
||||
BaseRevision string
|
||||
// Data is NDJSON of {"q", "a"}.
|
||||
Data []byte
|
||||
// Params is the trainer's settings as JSON.
|
||||
Params string
|
||||
GPUMemoryMiB int
|
||||
GGUF bool
|
||||
}
|
||||
|
||||
// Status is how a job is going.
|
||||
type Status struct {
|
||||
State string
|
||||
Message string
|
||||
Step int
|
||||
Steps int
|
||||
Loss float64
|
||||
PeakMiB int
|
||||
}
|
||||
|
||||
// Env is where a node keeps what training needs.
|
||||
type Env struct {
|
||||
// EnginesDir is where engines are installed (the agent's engines_dir).
|
||||
EnginesDir string
|
||||
// WorkDir is the trainer's own directory: its overlay, the llama.cpp
|
||||
// conversion code, and each job's files while it runs.
|
||||
WorkDir string
|
||||
// RunAs is the user jobs run as: the engines' user.
|
||||
RunAs string
|
||||
Log *slog.Logger
|
||||
// Python, UV, ConvDir and GPUs replace what would be found, for tests.
|
||||
Python, UV, ConvDir string
|
||||
GPUs func(ctx context.Context) ([]GPU, error)
|
||||
}
|
||||
|
||||
// GPU is one card and how much of it is free.
|
||||
type GPU struct {
|
||||
Index int
|
||||
FreeMiB int
|
||||
}
|
||||
|
||||
// Runner runs one job at a time.
|
||||
type Runner struct {
|
||||
env Env
|
||||
mu sync.Mutex
|
||||
cancel context.CancelFunc
|
||||
job string
|
||||
}
|
||||
|
||||
// New is a runner.
|
||||
func New(env Env) *Runner {
|
||||
if env.Log == nil {
|
||||
env.Log = slog.New(slog.DiscardHandler)
|
||||
}
|
||||
return &Runner{env: env}
|
||||
}
|
||||
|
||||
// Cancel stops a job if it is the one running.
|
||||
func (r *Runner) Cancel(id string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.job == id && r.cancel != nil {
|
||||
r.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// ErrBusy is a node already training.
|
||||
var ErrBusy = errors.New("this node is already training")
|
||||
|
||||
// Run runs a job to its end, telling status as it goes and artifact each
|
||||
// piece of what it made. It answers the error it failed with.
|
||||
func (r *Runner) Run(ctx context.Context, job Job, status func(Status), artifact func(kind string, seq int, data []byte, last bool) error) error {
|
||||
r.mu.Lock()
|
||||
if r.job != "" {
|
||||
r.mu.Unlock()
|
||||
return ErrBusy
|
||||
}
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
r.job, r.cancel = job.ID, cancel
|
||||
r.mu.Unlock()
|
||||
defer func() {
|
||||
cancel()
|
||||
r.mu.Lock()
|
||||
r.job, r.cancel = "", nil
|
||||
r.mu.Unlock()
|
||||
}()
|
||||
|
||||
status(Status{State: "preparing", Message: "preparing the trainer"})
|
||||
py, err := r.python()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
overlay, err := r.overlay(ctx, py)
|
||||
if err != nil {
|
||||
return fmt.Errorf("installing %s: %w", strings.Join(overlayPackages, ", "), err)
|
||||
}
|
||||
conv := ""
|
||||
if job.GGUF {
|
||||
if conv, err = r.conversion(ctx); err != nil {
|
||||
return fmt.Errorf("fetching llama.cpp's conversion code: %w", err)
|
||||
}
|
||||
}
|
||||
gpu, err := r.pickGPU(ctx, job.GPUMemoryMiB)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
dir := filepath.Join(r.env.WorkDir, "jobs", job.ID)
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
out := filepath.Join(dir, "adapter")
|
||||
var params map[string]any
|
||||
if job.Params != "" {
|
||||
if err := json.Unmarshal([]byte(job.Params), ¶ms); err != nil {
|
||||
return fmt.Errorf("the job's params: %w", err)
|
||||
}
|
||||
}
|
||||
spec, _ := json.Marshal(map[string]any{"base_repo": job.BaseRepo, "base_revision": job.BaseRevision,
|
||||
"data": filepath.Join(dir, "data.jsonl"), "out": out, "params": params})
|
||||
for name, b := range map[string][]byte{"data.jsonl": job.Data, "job.json": spec, "trainer.py": trainerPy} {
|
||||
if err := os.WriteFile(filepath.Join(dir, name), b, 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
var baseDir string
|
||||
err = r.python3(ctx, py, dir, []string{"PYTHONPATH=" + overlay, "CUDA_VISIBLE_DEVICES=" + strconv.Itoa(gpu)},
|
||||
[]string{filepath.Join(dir, "trainer.py"), filepath.Join(dir, "job.json")},
|
||||
func(line map[string]any) {
|
||||
s := Status{State: str(line["state"]), Message: str(line["message"]), Step: num(line["step"]),
|
||||
Steps: num(line["steps"]), PeakMiB: num(line["peak_mib"])}
|
||||
if v, ok := line["loss"].(float64); ok {
|
||||
s.Loss = v
|
||||
}
|
||||
if b := str(line["base_dir"]); b != "" {
|
||||
baseDir = b
|
||||
}
|
||||
if s.State == "trained" {
|
||||
s.State = "running"
|
||||
s.Message = "trained; packing the adapter"
|
||||
}
|
||||
status(s)
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if job.GGUF {
|
||||
status(Status{State: "running", Message: "converting the adapter to GGUF for llama.cpp"})
|
||||
gguf := filepath.Join(dir, "adapter.gguf")
|
||||
err := r.python3(ctx, py, dir, []string{"PYTHONPATH=" + strings.Join([]string{filepath.Join(conv, "gguf-py"), conv, overlay}, ":")},
|
||||
[]string{filepath.Join(conv, "convert_lora_to_gguf.py"), "--base", baseDir, "--outtype", "f16", "--outfile", gguf, out}, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("converting to GGUF: %w", err)
|
||||
}
|
||||
if err := sendFile(gguf, "gguf", artifact); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
var tgz bytes.Buffer
|
||||
if err := tarGz(&tgz, out); err != nil {
|
||||
return err
|
||||
}
|
||||
return sendBytes(tgz.Bytes(), "peft", artifact)
|
||||
}
|
||||
|
||||
func str(v any) string {
|
||||
s, _ := v.(string)
|
||||
return s
|
||||
}
|
||||
|
||||
func num(v any) int {
|
||||
f, _ := v.(float64)
|
||||
return int(f)
|
||||
}
|
||||
|
||||
// python3 runs a Python script as the engines' user, passing each JSON line
|
||||
// it prints to line, and answers the last of what it wrote to stderr when it
|
||||
// fails.
|
||||
func (r *Runner) python3(ctx context.Context, py, dir string, env, args []string, line func(map[string]any)) error {
|
||||
cmd := exec.CommandContext(ctx, py, args...)
|
||||
cmd.Dir = dir
|
||||
cmd.Env = append(os.Environ(), "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True", "TOKENIZERS_PARALLELISM=false",
|
||||
"HF_HUB_DISABLE_PROGRESS_BARS=1", "PYTHONUNBUFFERED=1")
|
||||
cmd.Env = append(cmd.Env, env...)
|
||||
home, give, why := supervisor.AsEngineUser(cmd, r.env.RunAs, r.env.WorkDir)
|
||||
if why != "" {
|
||||
r.env.Log.Warn(why)
|
||||
}
|
||||
if home != "" {
|
||||
cmd.Env = append(cmd.Env, "HF_HOME="+filepath.Join(home, ".cache", "huggingface"))
|
||||
} else {
|
||||
cmd.Env = append(cmd.Env, "HF_HOME="+filepath.Join(r.env.WorkDir, "hf"))
|
||||
}
|
||||
if err := give(dir); err != nil {
|
||||
return fmt.Errorf("handing %s to %s: %w", dir, r.env.RunAs, err)
|
||||
}
|
||||
// The whole process group goes when the job is cancelled: torch starts
|
||||
// workers of its own.
|
||||
if cmd.SysProcAttr == nil {
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
}
|
||||
cmd.Cancel = func() error { return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) }
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var stderr tailBuffer
|
||||
cmd.Stderr = &stderr
|
||||
if err := cmd.Start(); err != nil {
|
||||
return err
|
||||
}
|
||||
sc := bufio.NewScanner(stdout)
|
||||
sc.Buffer(make([]byte, 64<<10), 1<<20)
|
||||
for sc.Scan() {
|
||||
var m map[string]any
|
||||
if line != nil && json.Unmarshal(sc.Bytes(), &m) == nil {
|
||||
line(m)
|
||||
}
|
||||
}
|
||||
err = cmd.Wait()
|
||||
if ctx.Err() != nil {
|
||||
return errors.New("cancelled")
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("%v: %s", err, stderr.last())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tailBuffer keeps the end of what a process wrote: the reason it failed is
|
||||
// at the end, and the rest is progress bars.
|
||||
type tailBuffer struct{ b []byte }
|
||||
|
||||
func (t *tailBuffer) Write(p []byte) (int, error) {
|
||||
t.b = append(t.b, p...)
|
||||
if len(t.b) > 16<<10 {
|
||||
t.b = t.b[len(t.b)-16<<10:]
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (t *tailBuffer) last() string {
|
||||
lines := strings.Split(strings.TrimSpace(string(t.b)), "\n")
|
||||
if len(lines) > 3 {
|
||||
lines = lines[len(lines)-3:]
|
||||
}
|
||||
return strings.Join(lines, " | ")
|
||||
}
|
||||
|
||||
// python is the newest installed vLLM's interpreter: its torch is built for
|
||||
// this node's CUDA.
|
||||
func (r *Runner) python() (string, error) {
|
||||
if r.env.Python != "" {
|
||||
return r.env.Python, nil
|
||||
}
|
||||
found, _ := filepath.Glob(filepath.Join(r.env.EnginesDir, "vllm", "*", "bin", "python3"))
|
||||
if len(found) == 0 {
|
||||
return "", errors.New("training needs vLLM installed on this node (its torch is what trains); " +
|
||||
"llmgw ctl engine-build, or serving any vLLM model here, installs it")
|
||||
}
|
||||
sort.Strings(found)
|
||||
return found[len(found)-1], nil
|
||||
}
|
||||
|
||||
func (r *Runner) uv() (string, error) {
|
||||
if r.env.UV != "" {
|
||||
return r.env.UV, nil
|
||||
}
|
||||
var found []string
|
||||
_ = filepath.WalkDir(filepath.Join(r.env.EnginesDir, "uv"), func(p string, d os.DirEntry, err error) error {
|
||||
if err == nil && !d.IsDir() && d.Name() == "uv" {
|
||||
found = append(found, p)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if len(found) == 0 {
|
||||
return "", errors.New("no uv under " + filepath.Join(r.env.EnginesDir, "uv"))
|
||||
}
|
||||
sort.Strings(found)
|
||||
return found[len(found)-1], nil
|
||||
}
|
||||
|
||||
// overlay installs the pinned packages once per Python and set of pins.
|
||||
func (r *Runner) overlay(ctx context.Context, py string) (string, error) {
|
||||
sum := sha256.Sum256([]byte(py + "\x00" + strings.Join(overlayPackages, " ")))
|
||||
dir := filepath.Join(r.env.WorkDir, "overlay-"+hex.EncodeToString(sum[:6]))
|
||||
if _, err := os.Stat(filepath.Join(dir, ".done")); err == nil {
|
||||
return dir, nil
|
||||
}
|
||||
uv, err := r.uv()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
tmp := dir + ".new"
|
||||
_ = os.RemoveAll(tmp)
|
||||
args := append([]string{"pip", "install", "--quiet", "--python", py, "--no-deps", "--target", tmp}, overlayPackages...)
|
||||
cmd := exec.CommandContext(ctx, uv, args...)
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
return "", fmt.Errorf("%v: %s", err, strings.TrimSpace(string(out)))
|
||||
}
|
||||
_ = os.WriteFile(filepath.Join(tmp, ".done"), nil, 0o644)
|
||||
_ = os.RemoveAll(dir)
|
||||
if err := os.Rename(tmp, dir); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// Readable by the engines' user, which runs the trainer.
|
||||
_ = filepath.WalkDir(dir, func(p string, d os.DirEntry, err error) error {
|
||||
if err == nil {
|
||||
mode := os.FileMode(0o644)
|
||||
if d.IsDir() {
|
||||
mode = 0o755
|
||||
}
|
||||
_ = os.Chmod(p, mode)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return dir, nil
|
||||
}
|
||||
|
||||
// llamaCppTag is the release of the newest llama.cpp installed: its
|
||||
// conversion code is the one whose GGUF that llama.cpp reads.
|
||||
func (r *Runner) llamaCppTag() (string, error) {
|
||||
found, _ := filepath.Glob(filepath.Join(r.env.EnginesDir, "llamacpp", "b*"))
|
||||
if len(found) == 0 {
|
||||
return "", errors.New("a GGUF adapter needs llama.cpp installed on this node, to match its release")
|
||||
}
|
||||
sort.Strings(found)
|
||||
tag, _, _ := strings.Cut(filepath.Base(found[len(found)-1]), "-")
|
||||
return tag, nil
|
||||
}
|
||||
|
||||
// conversionSource is where llama.cpp's source is fetched from, by tag.
|
||||
var conversionSource = "https://codeload.github.com/ggml-org/llama.cpp/tar.gz/refs/tags/"
|
||||
|
||||
// conversion fetches llama.cpp's LoRA conversion code for the installed
|
||||
// release, once.
|
||||
func (r *Runner) conversion(ctx context.Context) (string, error) {
|
||||
if r.env.ConvDir != "" {
|
||||
return r.env.ConvDir, nil
|
||||
}
|
||||
tag, err := r.llamaCppTag()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
dir := filepath.Join(r.env.WorkDir, "llama.cpp-"+tag)
|
||||
if _, err := os.Stat(filepath.Join(dir, "convert_lora_to_gguf.py")); err == nil {
|
||||
return dir, nil
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, conversionSource+tag, nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
res, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
if res.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("%s answered %s", conversionSource+tag, res.Status)
|
||||
}
|
||||
tmp := dir + ".new"
|
||||
_ = os.RemoveAll(tmp)
|
||||
if err := untarConversion(res.Body, tmp); err != nil {
|
||||
return "", err
|
||||
}
|
||||
_ = os.RemoveAll(dir)
|
||||
return dir, os.Rename(tmp, dir)
|
||||
}
|
||||
|
||||
// untarConversion keeps, of llama.cpp's source, the conversion scripts and
|
||||
// the Python packages they import.
|
||||
func untarConversion(r io.Reader, dir string) error {
|
||||
gz, err := gzip.NewReader(r)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tr := tar.NewReader(gz)
|
||||
kept := 0
|
||||
for {
|
||||
h, err := tr.Next()
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, rel, ok := strings.Cut(h.Name, "/")
|
||||
if !ok || h.Typeflag != tar.TypeReg || strings.Contains(rel, "..") {
|
||||
continue
|
||||
}
|
||||
if rel != "convert_lora_to_gguf.py" && rel != "convert_hf_to_gguf.py" &&
|
||||
!strings.HasPrefix(rel, "conversion/") && !strings.HasPrefix(rel, "gguf-py/") {
|
||||
continue
|
||||
}
|
||||
p := filepath.Join(dir, filepath.FromSlash(rel))
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
f, err := os.OpenFile(p, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = io.Copy(f, io.LimitReader(tr, 64<<20))
|
||||
f.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
kept++
|
||||
}
|
||||
if kept == 0 {
|
||||
return errors.New("no conversion code in llama.cpp's source")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// pickGPU is the card with the most memory free, if that is enough.
|
||||
func (r *Runner) pickGPU(ctx context.Context, need int) (int, error) {
|
||||
list := r.env.GPUs
|
||||
if list == nil {
|
||||
list = nvidiaSMI
|
||||
}
|
||||
gpus, err := list(ctx)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("reading the cards' free memory: %w", err)
|
||||
}
|
||||
if len(gpus) == 0 {
|
||||
return 0, errors.New("this node has no GPU to train on")
|
||||
}
|
||||
best := gpus[0]
|
||||
for _, g := range gpus[1:] {
|
||||
if g.FreeMiB > best.FreeMiB {
|
||||
best = g
|
||||
}
|
||||
}
|
||||
if best.FreeMiB < need {
|
||||
return 0, fmt.Errorf("no card here has %d MiB free (the most free is GPU %d with %d MiB); "+
|
||||
"a smaller base, or a card freed for it, would do", need, best.Index, best.FreeMiB)
|
||||
}
|
||||
return best.Index, nil
|
||||
}
|
||||
|
||||
func nvidiaSMI(ctx context.Context) ([]GPU, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
|
||||
defer cancel()
|
||||
out, err := exec.CommandContext(ctx, "nvidia-smi", "--query-gpu=index,memory.free", "--format=csv,noheader,nounits").Output()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return parseGPUs(string(out)), nil
|
||||
}
|
||||
|
||||
func parseGPUs(s string) []GPU {
|
||||
var out []GPU
|
||||
for _, line := range strings.Split(strings.TrimSpace(s), "\n") {
|
||||
a, b, ok := strings.Cut(line, ",")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
i, err1 := strconv.Atoi(strings.TrimSpace(a))
|
||||
f, err2 := strconv.Atoi(strings.TrimSpace(b))
|
||||
if err1 == nil && err2 == nil {
|
||||
out = append(out, GPU{Index: i, FreeMiB: f})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// tarGz packs a directory.
|
||||
func tarGz(w io.Writer, dir string) error {
|
||||
gz := gzip.NewWriter(w)
|
||||
tw := tar.NewWriter(gz)
|
||||
err := filepath.WalkDir(dir, func(p string, d os.DirEntry, err error) error {
|
||||
if err != nil || d.IsDir() {
|
||||
return err
|
||||
}
|
||||
rel, _ := filepath.Rel(dir, p)
|
||||
info, err := d.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tw.WriteHeader(&tar.Header{Name: filepath.ToSlash(rel), Mode: 0o644, Size: info.Size(), ModTime: info.ModTime()}); err != nil {
|
||||
return err
|
||||
}
|
||||
f, err := os.Open(p)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
_, err = io.Copy(tw, f)
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tw.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
return gz.Close()
|
||||
}
|
||||
|
||||
// artifactChunk is how much of an artifact goes in one message.
|
||||
const artifactChunk = 1 << 20
|
||||
|
||||
func sendFile(path, kind string, send func(string, int, []byte, bool) error) error {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return sendBytes(b, kind, send)
|
||||
}
|
||||
|
||||
func sendBytes(b []byte, kind string, send func(string, int, []byte, bool) error) error {
|
||||
seq := 0
|
||||
for {
|
||||
n := min(artifactChunk, len(b))
|
||||
last := n == len(b)
|
||||
if err := send(kind, seq, b[:n], last); err != nil {
|
||||
return err
|
||||
}
|
||||
if last {
|
||||
return nil
|
||||
}
|
||||
b = b[n:]
|
||||
seq++
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package train
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// fakePython stands in for an engine's interpreter: as the trainer it
|
||||
// prints progress and writes an adapter; as the converter it writes a GGUF;
|
||||
// told to fail, it fails.
|
||||
const fakePython = `#!/bin/sh
|
||||
case "$1" in
|
||||
*trainer.py)
|
||||
out=$(sed -n 's/.*"out": *"\([^"]*\)".*/\1/p' "$2")
|
||||
grep -q FAIL "$(dirname "$2")/data.jsonl" && { echo "CUDA out of memory" >&2; exit 1; }
|
||||
echo '{"state":"preparing","message":"fetching base"}'
|
||||
echo '{"state":"running","step":5,"steps":10,"loss":0.9,"peak_mib":1900}'
|
||||
mkdir -p "$out"; echo '{"r":16}' > "$out/adapter_config.json"; echo weights > "$out/adapter_model.safetensors"
|
||||
echo '{"state":"trained","step":10,"steps":10,"peak_mib":2100,"base_dir":"/base"}'
|
||||
;;
|
||||
*convert_lora_to_gguf.py)
|
||||
while [ $# -gt 0 ]; do [ "$1" = "--outfile" ] && echo GGUF > "$2"; shift; done
|
||||
;;
|
||||
esac
|
||||
`
|
||||
|
||||
func runner(t *testing.T, free int) *Runner {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
py := filepath.Join(dir, "python3")
|
||||
_ = os.WriteFile(py, []byte(fakePython), 0o755)
|
||||
uv := filepath.Join(dir, "uv")
|
||||
_ = os.WriteFile(uv, []byte("#!/bin/sh\nwhile [ $# -gt 0 ]; do [ \"$1\" = \"--target\" ] && mkdir -p \"$2\"; shift; done\n"), 0o755)
|
||||
return New(Env{WorkDir: filepath.Join(dir, "work"), Python: py, UV: uv, ConvDir: dir,
|
||||
GPUs: func(context.Context) ([]GPU, error) { return []GPU{{0, 500}, {1, free}}, nil }})
|
||||
}
|
||||
|
||||
type sink struct {
|
||||
mu sync.Mutex
|
||||
statuses []Status
|
||||
parts map[string][]byte
|
||||
lasts int
|
||||
}
|
||||
|
||||
func (s *sink) status(st Status) { s.mu.Lock(); s.statuses = append(s.statuses, st); s.mu.Unlock() }
|
||||
func (s *sink) artifact(kind string, _ int, b []byte, last bool) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.parts == nil {
|
||||
s.parts = map[string][]byte{}
|
||||
}
|
||||
s.parts[kind] = append(s.parts[kind], b...)
|
||||
if last {
|
||||
s.lasts++
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAJobTrainsOnTheFreestCardAndSendsWhatItMade(t *testing.T) {
|
||||
r := runner(t, 3000)
|
||||
s := &sink{}
|
||||
err := r.Run(context.Background(), Job{ID: "j1", BaseRepo: "Qwen/x", Data: []byte(`{"q":"a","a":"b"}`), GPUMemoryMiB: 2500, GGUF: true},
|
||||
s.status, s.artifact)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var running bool
|
||||
for _, st := range s.statuses {
|
||||
if st.State == "running" && st.Step == 5 && st.Steps == 10 && st.PeakMiB == 1900 {
|
||||
running = true
|
||||
}
|
||||
}
|
||||
if !running {
|
||||
t.Errorf("statuses %+v", s.statuses)
|
||||
}
|
||||
if string(s.parts["gguf"]) != "GGUF\n" || s.lasts != 2 {
|
||||
t.Errorf("gguf %q, lasts %d", s.parts["gguf"], s.lasts)
|
||||
}
|
||||
files := map[string]string{}
|
||||
gz, err := gzip.NewReader(bytes.NewReader(s.parts["peft"]))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tr := tar.NewReader(gz)
|
||||
for {
|
||||
h, err := tr.Next()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
b, _ := io.ReadAll(tr)
|
||||
files[h.Name] = string(b)
|
||||
}
|
||||
if files["adapter_model.safetensors"] != "weights\n" || files["adapter_config.json"] == "" {
|
||||
t.Errorf("peft archive %v", files)
|
||||
}
|
||||
// Nothing of the job is left on the node.
|
||||
if entries, _ := os.ReadDir(filepath.Join(r.env.WorkDir, "jobs")); len(entries) != 0 {
|
||||
t.Errorf("left behind: %v", entries)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAJobWaitsForNoCardItDoesNotFit(t *testing.T) {
|
||||
r := runner(t, 1000)
|
||||
err := r.Run(context.Background(), Job{ID: "j1", Data: []byte(`{"q":"a","a":"b"}`), GPUMemoryMiB: 2500}, (&sink{}).status, (&sink{}).artifact)
|
||||
if err == nil || !strings.Contains(err.Error(), "no card here has 2500 MiB free") {
|
||||
t.Errorf("got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestATrainerThatFailsSaysWhy(t *testing.T) {
|
||||
r := runner(t, 3000)
|
||||
err := r.Run(context.Background(), Job{ID: "j1", Data: []byte(`{"q":"FAIL","a":"b"}`)}, (&sink{}).status, (&sink{}).artifact)
|
||||
if err == nil || !strings.Contains(err.Error(), "CUDA out of memory") {
|
||||
t.Errorf("got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneJobAtATime(t *testing.T) {
|
||||
r := runner(t, 3000)
|
||||
r.job = "other"
|
||||
if err := r.Run(context.Background(), Job{ID: "j2"}, (&sink{}).status, (&sink{}).artifact); err != ErrBusy {
|
||||
t.Errorf("got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCardsAreReadFromNvidiaSMI(t *testing.T) {
|
||||
got := parseGPUs("0, 2310\n1, 80000\nbad line\n")
|
||||
if len(got) != 2 || got[1].FreeMiB != 80000 || got[0].Index != 0 {
|
||||
t.Errorf("parsed %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyTheConversionCodeIsKeptOfLlamaCpp(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
gz := gzip.NewWriter(&buf)
|
||||
tw := tar.NewWriter(gz)
|
||||
for _, name := range []string{"llama.cpp-b1/convert_lora_to_gguf.py", "llama.cpp-b1/conversion/qwen.py",
|
||||
"llama.cpp-b1/gguf-py/gguf/__init__.py", "llama.cpp-b1/src/llama.cpp", "llama.cpp-b1/../evil"} {
|
||||
_ = tw.WriteHeader(&tar.Header{Name: name, Mode: 0o644, Size: 1, Typeflag: tar.TypeReg})
|
||||
_, _ = tw.Write([]byte("x"))
|
||||
}
|
||||
_ = tw.Close()
|
||||
_ = gz.Close()
|
||||
dir := filepath.Join(t.TempDir(), "conv")
|
||||
if err := untarConversion(&buf, dir); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, want := range []string{"convert_lora_to_gguf.py", "conversion/qwen.py", "gguf-py/gguf/__init__.py"} {
|
||||
if _, err := os.Stat(filepath.Join(dir, want)); err != nil {
|
||||
t.Errorf("%s not kept", want)
|
||||
}
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, "src")); err == nil {
|
||||
t.Error("the rest of llama.cpp was kept")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
# llmgw's LoRA trainer (BACKLOG F4), run by the node agent with an installed
|
||||
# engine's Python (its torch and transformers) and a small overlay (peft,
|
||||
# accelerate). Reads a job, trains, writes the PEFT adapter, and says how it
|
||||
# goes as one JSON object a line on standard output.
|
||||
import json, math, os, random, sys, threading, subprocess, time
|
||||
|
||||
def say(**kw):
|
||||
print(json.dumps(kw), flush=True)
|
||||
|
||||
job = json.load(open(sys.argv[1], encoding="utf-8"))
|
||||
p = {"rank": 16, "alpha": 32, "dropout": 0.05, "epochs": 6, "lr": 2e-4, "batch": 4, "max_len": 512, "seed": 0}
|
||||
p.update(job.get("params") or {})
|
||||
|
||||
say(state="preparing", message="loading libraries")
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from peft import LoraConfig, get_peft_model
|
||||
|
||||
random.seed(p["seed"]); torch.manual_seed(p["seed"])
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("no CUDA device is visible to the trainer")
|
||||
|
||||
say(state="preparing", message="fetching " + job["base_repo"])
|
||||
base_dir = snapshot_download(job["base_repo"], revision=job.get("base_revision") or None,
|
||||
allow_patterns=["*.json", "*.safetensors", "*.txt", "*.model", "*.tiktoken"])
|
||||
say(state="preparing", message="loading the base model", base_dir=base_dir)
|
||||
|
||||
# bf16 where the card has it; fp16 with a loss scaler where it does not
|
||||
# (Turing and older).
|
||||
bf16 = torch.cuda.is_bf16_supported()
|
||||
dtype = torch.bfloat16 if bf16 else torch.float16
|
||||
tok = AutoTokenizer.from_pretrained(base_dir)
|
||||
if tok.pad_token_id is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
model = AutoModelForCausalLM.from_pretrained(base_dir, dtype=dtype).to("cuda")
|
||||
|
||||
rows = [json.loads(l) for l in open(job["data"], encoding="utf-8") if l.strip()]
|
||||
rows = [r for r in rows if r.get("q") and r.get("a")]
|
||||
if not rows:
|
||||
raise SystemExit("no {q, a} pairs in the data")
|
||||
|
||||
def encode(r):
|
||||
prompt = tok.apply_chat_template([{"role": "user", "content": r["q"]}], add_generation_prompt=True, tokenize=False)
|
||||
full = tok.apply_chat_template([{"role": "user", "content": r["q"]}, {"role": "assistant", "content": r["a"]}], tokenize=False)
|
||||
pi = tok(prompt, add_special_tokens=False)["input_ids"]
|
||||
fi = tok(full, add_special_tokens=False)["input_ids"][: p["max_len"]]
|
||||
return fi, ([-100] * len(pi) + fi[len(pi):])[: len(fi)]
|
||||
|
||||
data = [encode(r) for r in rows]
|
||||
|
||||
model.gradient_checkpointing_enable()
|
||||
model.enable_input_require_grads()
|
||||
model = get_peft_model(model, LoraConfig(r=int(p["rank"]), lora_alpha=int(p["alpha"]), lora_dropout=float(p["dropout"]),
|
||||
task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]))
|
||||
trainable = [x for x in model.parameters() if x.requires_grad]
|
||||
for x in trainable: # fp32 trainable weights, so a loss scaler can unscale them
|
||||
x.data = x.data.float()
|
||||
|
||||
peak = {"mib": 0}
|
||||
stop = threading.Event()
|
||||
def watch():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits"],
|
||||
capture_output=True, text=True, timeout=5).stdout
|
||||
for line in out.splitlines():
|
||||
pid, mib = [v.strip() for v in line.split(",")]
|
||||
if int(pid) == os.getpid():
|
||||
peak["mib"] = max(peak["mib"], int(mib))
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(0.5)
|
||||
threading.Thread(target=watch, daemon=True).start()
|
||||
|
||||
opt = torch.optim.AdamW(trainable, lr=float(p["lr"]), weight_decay=0.0)
|
||||
scaler = torch.amp.GradScaler("cuda", enabled=not bf16)
|
||||
bs, epochs = int(p["batch"]), int(p["epochs"])
|
||||
steps = epochs * math.ceil(len(data) / bs)
|
||||
warm = max(1, steps // 10)
|
||||
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / warm) * max(0.0, 1 - s / steps))
|
||||
model.train(); step = 0; last = time.time()
|
||||
for ep in range(epochs):
|
||||
random.shuffle(data)
|
||||
for i in range(0, len(data), bs):
|
||||
batch = data[i:i + bs]; n = max(len(x[0]) for x in batch)
|
||||
ids = torch.tensor([x[0] + [tok.pad_token_id] * (n - len(x[0])) for x in batch], device="cuda")
|
||||
lab = torch.tensor([x[1] + [-100] * (n - len(x[1])) for x in batch], device="cuda")
|
||||
att = torch.tensor([[1] * len(x[0]) + [0] * (n - len(x[0])) for x in batch], device="cuda")
|
||||
with torch.autocast("cuda", dtype=dtype):
|
||||
loss = model(input_ids=ids, attention_mask=att, labels=lab).loss
|
||||
scaler.scale(loss).backward(); scaler.unscale_(opt)
|
||||
torch.nn.utils.clip_grad_norm_(trainable, 1.0)
|
||||
scaler.step(opt); scaler.update(); opt.zero_grad(set_to_none=True); sched.step(); step += 1
|
||||
if time.time() - last > 2 or step == steps:
|
||||
last = time.time()
|
||||
say(state="running", message="training", step=step, steps=steps, loss=round(loss.item(), 4), peak_mib=peak["mib"])
|
||||
|
||||
model.save_pretrained(job["out"])
|
||||
stop.set()
|
||||
say(state="trained", step=step, steps=steps, peak_mib=peak["mib"], base_dir=base_dir)
|
||||
@@ -0,0 +1,46 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/train"
|
||||
)
|
||||
|
||||
// trainJob runs a training job the control plane sent, telling it how the
|
||||
// job goes and sending what it made. A stream that ends ends the job: the
|
||||
// control plane, which sees the node go, marks it failed, and it is started
|
||||
// again from the beginning rather than resumed.
|
||||
func (a *Agent) trainJob(ctx context.Context, m *llmgwv1.TrainStart, out chan<- *llmgwv1.AgentMessage) {
|
||||
id := m.GetJobId()
|
||||
send := func(msg *llmgwv1.AgentMessage) bool {
|
||||
select {
|
||||
case out <- msg:
|
||||
return true
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
}
|
||||
}
|
||||
status := func(s train.Status) {
|
||||
send(&llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_TrainStatus{TrainStatus: &llmgwv1.TrainStatus{
|
||||
JobId: id, State: s.State, Message: s.Message, Step: int32(s.Step), Steps: int32(s.Steps), //nolint:gosec // counts
|
||||
Loss: s.Loss, PeakMib: int32(s.PeakMiB)}}}) //nolint:gosec // MiB
|
||||
}
|
||||
artifact := func(kind string, seq int, data []byte, last bool) error {
|
||||
if !send(&llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_TrainArtifact{TrainArtifact: &llmgwv1.TrainArtifact{
|
||||
JobId: id, Kind: kind, Seq: int32(seq), Data: data, Last: last}}}) { //nolint:gosec // chunk count
|
||||
return ctx.Err()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
a.log.Info("training an adapter", "job", id, "base", m.GetBaseRepo())
|
||||
err := a.trainer.Run(ctx, train.Job{ID: id, BaseRepo: m.GetBaseRepo(), BaseRevision: m.GetBaseRevision(),
|
||||
Data: m.GetData(), Params: m.GetParams(), GPUMemoryMiB: int(m.GetGpuMemoryMib()), GGUF: m.GetGguf()}, status, artifact)
|
||||
if err != nil {
|
||||
a.log.Warn("training failed", "job", id, "err", err)
|
||||
status(train.Status{State: "failed", Message: err.Error()})
|
||||
return
|
||||
}
|
||||
a.log.Info("trained an adapter", "job", id)
|
||||
status(train.Status{State: "done"})
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
)
|
||||
|
||||
type alwaysReady struct{}
|
||||
@@ -40,7 +41,7 @@ func TestAnAdoptedEngineThatDiesIsReported(t *testing.T) {
|
||||
out := make(chan *llmgwv1.AgentMessage, 4)
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
defer cancel()
|
||||
go a.watchExits(ctx, out)
|
||||
go a.watchExits(ctx, out, nil)
|
||||
time.Sleep(1500 * time.Millisecond) // seen ready
|
||||
_ = cmd.Process.Kill()
|
||||
|
||||
@@ -54,3 +55,50 @@ func TestAnAdoptedEngineThatDiesIsReported(t *testing.T) {
|
||||
t.Fatal("the adopted engine died and nothing was reported")
|
||||
}
|
||||
}
|
||||
|
||||
// Registered while its health check still ran: the ready that follows is
|
||||
// reported, or the control plane keeps the replica starting while it serves.
|
||||
func TestAShardRegisteredStartingIsReportedReady(t *testing.T) {
|
||||
sup, err := supervisor.New(supervisor.Config{WorkDir: t.TempDir(), RunAs: "-"}, alwaysReady{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cmd := exec.Command("sleep", "60")
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = cmd.Process.Kill(); _ = cmd.Wait() }()
|
||||
sh, err := sup.Adopt(supervisor.Orphan{ShardID: "m-r0", ReplicaID: "m", PID: cmd.Process.Pid,
|
||||
Spec: engines.ShardSpec{ReplicaID: "m", ShardID: "m-r0", HTTPPort: 1}, Alive: true})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
a := &Agent{sup: sup, log: slog.New(slog.DiscardHandler)}
|
||||
pending := unready(a.runningShards())
|
||||
if !pending["m-r0"] {
|
||||
t.Fatalf("registration reported it ready already: %v", pending)
|
||||
}
|
||||
out := make(chan *llmgwv1.AgentMessage, 4)
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
defer cancel()
|
||||
go a.watchExits(ctx, out, pending)
|
||||
if err := sup.WaitReady(t.Context(), sh, "http://127.0.0.1:1", "m", 5*time.Second); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
select {
|
||||
case m := <-out:
|
||||
s := m.GetShard()
|
||||
if s.GetShardId() != "m-r0" || s.GetState() != agentsvc.StateToWire(supervisor.Ready) {
|
||||
t.Errorf("reported %+v", s)
|
||||
}
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("it became ready and nothing was reported")
|
||||
}
|
||||
select {
|
||||
case m := <-out:
|
||||
t.Errorf("reported twice: %+v", m.GetShard())
|
||||
case <-time.After(2500 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"slices"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -110,11 +111,16 @@ func (w *statusWriter) Flush() {
|
||||
}
|
||||
}
|
||||
|
||||
// notChanges are POSTs that change nothing and carry text that must not be
|
||||
// kept: asking a judge about a sentence.
|
||||
var notChanges = []string{"/api/v1/guard/judges/ask"}
|
||||
|
||||
// API records every change to the admin API: a request that is not a GET or
|
||||
// HEAD, under /api/. actor says who made it.
|
||||
func API(sink Sink, actor func(*http.Request) string, log *slog.Logger, next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodGet || r.Method == http.MethodHead || !strings.HasPrefix(r.URL.Path, "/api/") {
|
||||
if r.Method == http.MethodGet || r.Method == http.MethodHead || !strings.HasPrefix(r.URL.Path, "/api/") ||
|
||||
slices.Contains(notChanges, r.URL.Path) {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -72,3 +72,13 @@ func TestABodyLargerThanWhatIsRecordedReachesTheHandlerWhole(t *testing.T) {
|
||||
t.Errorf("the handler was given %d of %d bytes", got, len(big))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAskingAJudgeKeepsNoText(t *testing.T) {
|
||||
sink := &memSink{}
|
||||
h := API(sink, func(*http.Request) string { return "ilker" }, nil, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
|
||||
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest("POST", "/api/v1/guard/judges/ask",
|
||||
strings.NewReader(`{"text":"something a user wrote"}`)))
|
||||
if len(sink.got) != 0 {
|
||||
t.Errorf("a judge's question was recorded: %+v", sink.got)
|
||||
}
|
||||
}
|
||||
+22
-1
@@ -40,10 +40,23 @@ func newAgentCmd() *cobra.Command {
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
log := logger()
|
||||
|
||||
cfg, err := config.LoadAgent(configPath)
|
||||
// The cluster's settings for this agent over its file, as the
|
||||
// control plane last sent them. Settings that no longer load
|
||||
// beside the file are set aside rather than keeping the agent
|
||||
// from starting: the control plane sends them again, and says
|
||||
// why they are refused.
|
||||
settings := agent.ReadSettings()
|
||||
cfg, err := config.LoadAgentWith(configPath, settings)
|
||||
if err != nil && settings != nil {
|
||||
log.Warn("the cluster's settings kept for this agent do not load beside its agent.yaml; starting without them",
|
||||
"err", err)
|
||||
settings = nil
|
||||
cfg, err = config.LoadAgent(configPath)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg.Settings = string(settings)
|
||||
if server != "" {
|
||||
cfg.Server = server
|
||||
}
|
||||
@@ -89,9 +102,17 @@ func newAgentCmd() *cobra.Command {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
a.SetSettingsCheck(func(b []byte) error {
|
||||
_, err := config.LoadAgentWith(configPath, b)
|
||||
return err
|
||||
})
|
||||
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
go a.WatchFile(ctx, configPath, func() error {
|
||||
_, err := config.LoadAgentWith(configPath, agent.ReadSettings())
|
||||
return err
|
||||
})
|
||||
|
||||
log.Info("agent starting",
|
||||
"node", cfg.NodeID, "server", cfg.Server,
|
||||
|
||||
@@ -249,7 +249,14 @@ func pgEnv(dsn string) ([]string, error) {
|
||||
set("PGPASSWORD", pw)
|
||||
}
|
||||
set("PGDATABASE", strings.TrimPrefix(u.Path, "/"))
|
||||
set("PGSSLMODE", u.Query().Get("sslmode"))
|
||||
// The TLS settings go with the mode: verify-full without the cluster's
|
||||
// authority sent libpq looking for ~/.postgresql/root.crt, and the dump
|
||||
// failed on every control plane whose store is replicated.
|
||||
q := u.Query()
|
||||
set("PGSSLMODE", q.Get("sslmode"))
|
||||
set("PGSSLROOTCERT", q.Get("sslrootcert"))
|
||||
set("PGSSLCERT", q.Get("sslcert"))
|
||||
set("PGSSLKEY", q.Get("sslkey"))
|
||||
// A standby if there is one, else whichever answers — not the address's
|
||||
// read-write: a dump only reads, it need not load the primary, and one
|
||||
// that insists on the primary fails exactly while the cluster is failing
|
||||
|
||||
@@ -238,3 +238,17 @@ func TestAStoreOfSeveralHostsIsGivenToPgDumpAsLibpqLists(t *testing.T) {
|
||||
t.Error("an empty host in the list was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPgDumpIsGivenTheClusterAuthorityAlongWithVerifyFull(t *testing.T) {
|
||||
// Without it libpq looked for ~/.postgresql/root.crt and the dump failed.
|
||||
env, err := pgEnv("postgres://llmgw:p@10.0.0.149:5432,10.0.0.150:5432/llmgw?sslmode=verify-full&sslrootcert=%2Fetc%2Fllmgw%2Fpki%2Fca.crt&target_session_attrs=read-write")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := strings.Join(env, " ")
|
||||
for _, want := range []string{"PGSSLMODE=verify-full", "PGSSLROOTCERT=/etc/llmgw/pki/ca.crt"} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Errorf("%s is not in %s", want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -237,8 +237,12 @@ func (w *settingsWatch) check(ctx context.Context) {
|
||||
// again: what it reads at every use, or not at all.
|
||||
var liveSettings = []string{
|
||||
"guard",
|
||||
// Stopping and starting the service (llmgw stop, llmgw start).
|
||||
"serving",
|
||||
// Sent to the nodes and this machine's firewall unit as they change.
|
||||
"firewall", "firewall_networks",
|
||||
// Sent to the nodes' agents as they change.
|
||||
"agents",
|
||||
// Read by server init, node add and backups, not by the server.
|
||||
"backup_recipients", "backup_copy_to",
|
||||
}
|
||||
|
||||
+62
-7
@@ -116,6 +116,18 @@ func (c *ctlClient) do(ctx context.Context, method, path string, body, out any)
|
||||
if out == nil {
|
||||
return nil
|
||||
}
|
||||
// A stream, as it comes: NDJSON a long run writes line by line.
|
||||
if w, ok := out.(io.Writer); ok {
|
||||
_, err := io.Copy(w, resp.Body)
|
||||
return err
|
||||
}
|
||||
// An older control plane answers a path it does not know with its
|
||||
// dashboard: a page, which read as JSON said only "invalid character '<'".
|
||||
if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/html") {
|
||||
return fmt.Errorf("%s %s answered with a page, not the API: the control plane answering "+
|
||||
"(the leader, through a standby) does not know it — is it on an older release? "+
|
||||
"llmgw ctl status says which it runs", method, path)
|
||||
}
|
||||
return json.NewDecoder(resp.Body).Decode(out)
|
||||
}
|
||||
|
||||
@@ -168,11 +180,13 @@ func newCtlCmd() *cobra.Command {
|
||||
newJobsCmd(client, &asJSON),
|
||||
newEngineCheckCmd(),
|
||||
newEngineBuildCmd(client),
|
||||
newToolCheckCmd(),
|
||||
newToolCheckCmd(client),
|
||||
newWeightsCmd(client, &asJSON),
|
||||
newCapacityCmd(client, &asJSON),
|
||||
newQuotaCmd(client, &asJSON),
|
||||
newGuardCmd(client, &asJSON),
|
||||
newKBCmd(client, &asJSON),
|
||||
newTrainCmd(client, &asJSON),
|
||||
newConfigCmd(client, &asJSON),
|
||||
newAuditCmd(client, &asJSON),
|
||||
newNotifyCmd(client, &asJSON),
|
||||
@@ -225,8 +239,8 @@ func newNodesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
// What the node can run, which is the first thing to look
|
||||
// at when a model will not place and the hardware looks
|
||||
// fine.
|
||||
join(n.Engines),
|
||||
classesColumn(n), join(n.Fabrics), n.Replicas, n.LastSeenAgo)
|
||||
enginesColumn(n),
|
||||
classesColumn(n), fabricsColumn(n), n.Replicas, n.LastSeenAgo)
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
@@ -443,7 +457,7 @@ func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "REPLICA\tMODEL\tSTATE\tPOOL\tNODES\tAGE\tIDLE\tENDPOINT")
|
||||
fmt.Fprintln(w, "REPLICA\tMODEL\tSTATE\tPOOL\tON\tNODES\tAGE\tIDLE\tENDPOINT")
|
||||
for _, r := range reps {
|
||||
var nodes []string
|
||||
for _, s := range r.Shards {
|
||||
@@ -461,8 +475,14 @@ func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
if r.Restart != "" {
|
||||
state += " (restart " + r.Restart + ")"
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
r.ID, r.Model, state, r.Pool, join(nodes),
|
||||
// What the engine computes on, by its own account: the first
|
||||
// shard's, which serves.
|
||||
on := "-"
|
||||
if len(r.Shards) > 0 {
|
||||
on = r.Shards[0].Compute.Short()
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
r.ID, r.Model, state, r.Pool, on, join(nodes),
|
||||
(time.Duration(r.AgeSec) * time.Second).String(),
|
||||
idleColumn(r.IdleSec), r.Endpoint)
|
||||
// A line rather than a column, because most replicas are not
|
||||
@@ -472,7 +492,7 @@ func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
if s.Fetch == nil {
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(w, " %s: fetching %s\t\t\t\t\t\t\t\n",
|
||||
fmt.Fprintf(w, " %s: fetching %s\t\t\t\t\t\t\t\t\n",
|
||||
s.Node, fetchColumn(s.Fetch))
|
||||
}
|
||||
}
|
||||
@@ -583,7 +603,42 @@ func emitJSON(v any) error {
|
||||
// more than one or MIG instances among them: a card cut into two 3g.40gb
|
||||
// instances beside a whole card is "a100-80g-mig-3g.40gb×2,a100-80g×1", where
|
||||
// the GPUS column's 3 alone would read as three of something.
|
||||
// fabricsColumn is the node's fabrics, and when they are not all up, which
|
||||
// adapters are down: what keeps it out of replicas spanning nodes.
|
||||
func fabricsColumn(n adminapi.NodeView) string {
|
||||
s := join(n.Fabrics)
|
||||
if n.MultiNode != "" && n.MultiNode != "ready" {
|
||||
s += " (" + n.MultiNode + ")"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// enginesColumn is each engine with its version: llamacpp b10991.
|
||||
func enginesColumn(n adminapi.NodeView) string {
|
||||
if len(n.EngineVersions) == 0 {
|
||||
return join(n.Engines)
|
||||
}
|
||||
parts := make([]string, 0, len(n.Engines))
|
||||
for _, e := range n.Engines {
|
||||
if v := n.EngineVersions[e]; v != "" && v != "unknown" {
|
||||
e += " " + v
|
||||
}
|
||||
parts = append(parts, e)
|
||||
}
|
||||
return strings.Join(parts, ",")
|
||||
}
|
||||
|
||||
// classesColumn is the node's GPU classes, with their compute capability
|
||||
// when they report one: 6000-24g (7.5).
|
||||
func classesColumn(n adminapi.NodeView) string {
|
||||
s := classesOnly(n)
|
||||
if len(n.ComputeCapabilities) > 0 && s != "-" {
|
||||
s += " (" + strings.Join(n.ComputeCapabilities, ",") + ")"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func classesOnly(n adminapi.NodeView) string {
|
||||
counted := len(n.ClassGPUs) > 1
|
||||
for _, c := range n.ClassGPUs {
|
||||
counted = counted || c.MIG
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAPageInsteadOfTheAPIIsSaidPlainly(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
_, _ = w.Write([]byte("<!doctype html><title>llmgw</title>"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
c := newCtlClient(strings.TrimPrefix(srv.URL, "http://"), "")
|
||||
var out map[string]any
|
||||
err := c.get(t.Context(), "/api/v1/settings", &out)
|
||||
if err == nil || !strings.Contains(err.Error(), "older release") || strings.Contains(err.Error(), "invalid character") {
|
||||
t.Errorf("%v", err)
|
||||
}
|
||||
}
|
||||
@@ -76,8 +76,11 @@ func newRequestCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
fmt.Fprintf(out, "tags %s\n", strings.Join(kv, ", "))
|
||||
}
|
||||
if v.GuardAction != "" {
|
||||
fmt.Fprintf(out, "guard %s by policy %s (%s), %d ms\n", v.GuardAction, v.GuardPolicy,
|
||||
or(v.GuardCategory, "-"), v.GuardMs)
|
||||
what := or(v.GuardCategory, "-")
|
||||
if v.GuardRule != "" {
|
||||
what += ": " + v.GuardRule
|
||||
}
|
||||
fmt.Fprintf(out, "guard %s by policy %s (%s), %d ms\n", v.GuardAction, v.GuardPolicy, what, v.GuardMs)
|
||||
} else if v.GuardMs > 0 {
|
||||
fmt.Fprintf(out, "guard allowed, %d ms\n", v.GuardMs)
|
||||
}
|
||||
|
||||
@@ -12,9 +12,10 @@ import (
|
||||
|
||||
func newRestartCmd(client func() *ctlClient) *cobra.Command {
|
||||
var (
|
||||
models []string
|
||||
all bool
|
||||
inPlace bool
|
||||
models []string
|
||||
all bool
|
||||
inPlace bool
|
||||
surge, unavailable int
|
||||
)
|
||||
c := &cobra.Command{
|
||||
Use: "restart [REPLICA...]",
|
||||
@@ -31,6 +32,10 @@ func newRestartCmd(client func() *ctlClient) *cobra.Command {
|
||||
"profile edited to need more than any node has — the restart is given up\n" +
|
||||
"and the old replica keeps serving: llmgw ctl replicas says so, and why.\n\n" +
|
||||
"--in-place stops each replica first, whatever room there is.\n\n" +
|
||||
"A model's replicas are replaced one at a time. --max-surge N replaces N at\n" +
|
||||
"once, each beside itself, needing room for N more; --max-unavailable N\n" +
|
||||
"stops N at once first, needing no room and leaving N fewer serving while\n" +
|
||||
"they load. Together, the first N are stopped first and the rest surged.\n\n" +
|
||||
" llmgw ctl restart qwen2.5-7b-3f2a91c0\n" +
|
||||
" llmgw ctl restart --model qwen2.5-7b\n" +
|
||||
" llmgw ctl restart --all",
|
||||
@@ -38,7 +43,8 @@ func newRestartCmd(client func() *ctlClient) *cobra.Command {
|
||||
if len(args) == 0 && len(models) == 0 && !all {
|
||||
return errors.New("name replicas, --model, or --all")
|
||||
}
|
||||
body, _ := json.Marshal(adminapi.RestartRequest{Replicas: args, Models: models, All: all, InPlace: inPlace})
|
||||
body, _ := json.Marshal(adminapi.RestartRequest{Replicas: args, Models: models, All: all, InPlace: inPlace,
|
||||
MaxSurge: surge, MaxUnavailable: unavailable})
|
||||
var res adminapi.RestartResult
|
||||
if err := client().post(cmd.Context(), "/api/v1/restart", body, &res); err != nil {
|
||||
return err
|
||||
@@ -50,6 +56,9 @@ func newRestartCmd(client func() *ctlClient) *cobra.Command {
|
||||
c.Flags().StringArrayVar(&models, "model", nil, "every replica of this model; repeatable")
|
||||
c.Flags().BoolVar(&all, "all", false, "every ready replica in the cluster")
|
||||
c.Flags().BoolVar(&inPlace, "in-place", false, "stop each replica before starting its replacement")
|
||||
c.Flags().IntVar(&surge, "max-surge", 0, "replicas of a model replaced beside themselves at once (default 1)")
|
||||
c.Flags().IntVar(&unavailable, "max-unavailable", 0,
|
||||
"replicas of a model stopped first at once, each out while its replacement loads")
|
||||
return c
|
||||
}
|
||||
|
||||
|
||||
+22
-7
@@ -63,9 +63,10 @@ that file when the report is not enough to go on.`,
|
||||
enc := json.NewEncoder(os.Stdout)
|
||||
enc.SetIndent("", " ")
|
||||
return enc.Encode(struct {
|
||||
Findings []doctor.Finding `json:"findings"`
|
||||
Facts doctor.Facts `json:"facts"`
|
||||
}{findings, facts})
|
||||
Findings []doctor.Finding `json:"findings"`
|
||||
Readiness doctor.Readiness `json:"readiness"`
|
||||
Facts doctor.Facts `json:"facts"`
|
||||
}{findings, doctor.ReadinessOf(findings), facts})
|
||||
}
|
||||
printDoctor(os.Stdout, facts, findings, verbose)
|
||||
for _, f := range findings {
|
||||
@@ -101,12 +102,26 @@ func printDoctor(out io.Writer, f doctor.Facts, findings []doctor.Finding, verbo
|
||||
fmt.Fprintf(out, "%s · %s %s %s · driver %s\n\n", f.Hostname, f.Host.OSID, f.Host.OSVersion,
|
||||
f.Host.Arch, orDash(f.Host.NVIDIADriver))
|
||||
|
||||
for _, x := range findings {
|
||||
fmt.Fprintf(out, "%-4s %-6s %s\n", x.Level, x.Area, x.Message)
|
||||
if x.Fix != "" && (x.Level == doctor.Fail || x.Level == doctor.Warn) {
|
||||
fmt.Fprintf(out, " → %s\n", x.Fix)
|
||||
// The node first; then what replicas spanning nodes would need, which
|
||||
// keeps a node out only when a fabric was declared for it.
|
||||
multi := false
|
||||
for _, pass := range []bool{false, true} {
|
||||
for _, x := range findings {
|
||||
if x.MultiNode() != pass {
|
||||
continue
|
||||
}
|
||||
if pass && !multi {
|
||||
fmt.Fprintln(out, "\nMulti-node (fabrics, RDMA) — for replicas spanning nodes only:")
|
||||
multi = true
|
||||
}
|
||||
fmt.Fprintf(out, "%-4s %-6s %s\n", x.Level, x.Area, x.Message)
|
||||
if x.Fix != "" && (x.Level == doctor.Fail || x.Level == doctor.Warn) {
|
||||
fmt.Fprintf(out, " → %s\n", x.Fix)
|
||||
}
|
||||
}
|
||||
}
|
||||
r := doctor.ReadinessOf(findings)
|
||||
fmt.Fprintf(out, "\nThis node: %s\nMulti-node: %s\n", r.Node, r.MultiNode)
|
||||
|
||||
w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
|
||||
if len(f.GPUs) > 0 {
|
||||
|
||||
+166
-1
@@ -4,10 +4,15 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
@@ -22,6 +27,7 @@ func newEngineCheckCmd() *cobra.Command {
|
||||
nodes int
|
||||
tp int
|
||||
perNode int
|
||||
model string
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
@@ -31,6 +37,10 @@ func newEngineCheckCmd() *cobra.Command {
|
||||
"engines rename and remove flags between versions. When that happens the\n" +
|
||||
"failure arrives minutes into a model load, in an engine's own error\n" +
|
||||
"message, on a node nobody was watching.\n\n" +
|
||||
"For llama.cpp on a machine with NVIDIA GPUs it then loads a small GGUF\n" +
|
||||
"(--model, or the smallest in the weights cache) with every layer on the GPU\n" +
|
||||
"and asks for one token: a build whose CUDA backend does not load computes on\n" +
|
||||
"the CPU, and that is where it is caught.\n\n" +
|
||||
"This asks the engine itself. It runs the engine's --help and checks every\n" +
|
||||
"flag llmgw would pass against it, which needs no GPU — a CLI prints its\n" +
|
||||
"help on any machine. Run it after upgrading an engine.\n\n" +
|
||||
@@ -49,6 +59,11 @@ func newEngineCheckCmd() *cobra.Command {
|
||||
// rarely on PATH: a check that cannot find the engine reports
|
||||
// nothing rather than reporting a problem.
|
||||
cfg := engines.Config{}
|
||||
// The build the agent installed, when none is named: the newest
|
||||
// under /opt/llmgw/engines, which is where every node keeps it.
|
||||
if bin == "" && name == "llamacpp" {
|
||||
bin = installedLlama()
|
||||
}
|
||||
if bin != "" {
|
||||
switch name {
|
||||
case "llamacpp":
|
||||
@@ -135,6 +150,11 @@ func newEngineCheckCmd() *cobra.Command {
|
||||
var deviceErr error
|
||||
if fb, ok := eng.(engines.CPUFallback); ok && help == "" {
|
||||
deviceErr = checkDevices(cmd.Context(), name, fb, out)
|
||||
// Seeing a GPU is not computing on it: a small model loaded
|
||||
// on it and asked for a token is.
|
||||
if b, ok := eng.(interface{ Binary() string }); ok && deviceErr == nil && model != "none" {
|
||||
deviceErr = computeOnGPU(cmd.Context(), b.Binary(), model, out)
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Fprintf(out, "%s: %d flags checked", name, checked)
|
||||
@@ -160,7 +180,9 @@ func newEngineCheckCmd() *cobra.Command {
|
||||
|
||||
f := cmd.Flags()
|
||||
f.StringVar(&help, "help-output", "", "a saved --help to check against, instead of running the engine")
|
||||
f.StringVar(&bin, "bin", "", "where the engine lives, when it is not on PATH")
|
||||
f.StringVar(&bin, "bin", "", "where the engine lives (default: the build the agent installed, else PATH)")
|
||||
f.StringVar(&model, "model", "", "a GGUF file to load on the GPU and ask for a token (default: the smallest "+
|
||||
"in this node's weights cache; none skips it)")
|
||||
f.IntVar(&nodes, "nodes", 2, "shape the command for a replica of this many nodes")
|
||||
f.IntVar(&tp, "tp", 8, "tensor parallel degree")
|
||||
f.IntVar(&perNode, "gpus-per-node", 4, "GPUs each node contributes")
|
||||
@@ -286,3 +308,146 @@ func checkDevices(ctx context.Context, name string, fb engines.CPUFallback, out
|
||||
fmt.Fprintf(out, "computes on: %s\n", strings.Join(devs, "; "))
|
||||
return nil
|
||||
}
|
||||
|
||||
// llamaRoot is where the agent installs llama.cpp builds.
|
||||
var llamaRoot = "/opt/llmgw/engines/llamacpp"
|
||||
|
||||
// installedLlama is the llama-server the agent installed here: under
|
||||
// llamaRoot/<version-variant>/llama-<build>/, or one level up for older
|
||||
// layouts; on a machine with NVIDIA GPUs a CUDA build first, then the newest.
|
||||
func installedLlama() string {
|
||||
var found []string
|
||||
for _, pat := range []string{"*/llama-server", "*/*/llama-server"} {
|
||||
m, _ := filepath.Glob(filepath.Join(llamaRoot, pat))
|
||||
found = append(found, m...)
|
||||
}
|
||||
if len(found) == 0 {
|
||||
return ""
|
||||
}
|
||||
sort.Strings(found)
|
||||
if _, err := exec.LookPath("nvidia-smi"); err == nil {
|
||||
for i := len(found) - 1; i >= 0; i-- {
|
||||
if strings.Contains(found[i], "cuda") {
|
||||
return found[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
return found[len(found)-1]
|
||||
}
|
||||
|
||||
// weightDirs are where a node keeps models it has fetched.
|
||||
var weightDirs = []string{"/var/lib/llmgw-weights", "/var/lib/llmgw-models", "/srv/models"}
|
||||
|
||||
// smallestGGUF is the smallest GGUF file under dirs, for a check that only
|
||||
// needs something to load.
|
||||
func smallestGGUF(dirs []string) string {
|
||||
best, size := "", int64(0)
|
||||
for _, d := range dirs {
|
||||
_ = filepath.WalkDir(d, func(p string, e os.DirEntry, err error) error {
|
||||
if err != nil || e.IsDir() || !strings.HasSuffix(strings.ToLower(p), ".gguf") ||
|
||||
strings.Contains(strings.ToLower(filepath.Base(p)), "mmproj") {
|
||||
return nil
|
||||
}
|
||||
if info, err := e.Info(); err == nil && (best == "" || info.Size() < size) {
|
||||
best, size = p, info.Size()
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
// computeOnGPU loads a model with every layer on the GPU, asks for one
|
||||
// token, and reads from llama.cpp's own output where the layers went.
|
||||
func computeOnGPU(ctx context.Context, bin, model string, out io.Writer) error {
|
||||
if _, err := exec.LookPath("nvidia-smi"); err != nil {
|
||||
return nil
|
||||
}
|
||||
if model == "" {
|
||||
if model = smallestGGUF(weightDirs); model == "" {
|
||||
fmt.Fprintln(out, "no GGUF in this node's weights cache to compute with; --model PATH runs one")
|
||||
return nil
|
||||
}
|
||||
}
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
port := ln.Addr().(*net.TCPAddr).Port
|
||||
_ = ln.Close()
|
||||
ctx, cancel := context.WithTimeout(ctx, 3*time.Minute)
|
||||
defer cancel()
|
||||
var logs syncBuffer
|
||||
c := exec.CommandContext(ctx, bin, "-m", model, "--host", "127.0.0.1", "--port", strconv.Itoa(port), //nolint:gosec // the engine's own binary
|
||||
"--n-gpu-layers", "999", "--ctx-size", "512", "--parallel", "1",
|
||||
// Where the layers went is said only at this verbosity.
|
||||
"-lv", "4")
|
||||
c.Stdout, c.Stderr = &logs, &logs
|
||||
if err := c.Start(); err != nil {
|
||||
return fmt.Errorf("starting %s: %w", bin, err)
|
||||
}
|
||||
defer func() { _ = c.Process.Kill(); _ = c.Wait() }()
|
||||
base := "http://127.0.0.1:" + strconv.Itoa(port)
|
||||
start := time.Now()
|
||||
for {
|
||||
resp, err := http.Get(base + "/health") //nolint:gosec,noctx // loopback, bounded by the loop
|
||||
if err == nil {
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusOK {
|
||||
break
|
||||
}
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return fmt.Errorf("%s did not come up with %s: %s", bin, filepath.Base(model), tailLines(logs.String(), 5))
|
||||
}
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, base+"/completion",
|
||||
strings.NewReader(`{"prompt":"Hello","n_predict":1}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("asking it for a token: %w", err)
|
||||
}
|
||||
_ = resp.Body.Close()
|
||||
took := time.Since(start)
|
||||
comp, said := engines.ReadLlamaCompute(strings.Split(logs.String(), "\n"))
|
||||
if !said {
|
||||
fmt.Fprintf(out, "computed a token with %s, but llama.cpp did not say where\n", filepath.Base(model))
|
||||
return nil
|
||||
}
|
||||
if comp.OnTheCPU() {
|
||||
return fmt.Errorf("%s computed on the CPU though every layer was asked for on the GPU (%d of %d offloaded): "+
|
||||
"its CUDA backend did not load; ldd on its libggml-cuda.so names the missing library",
|
||||
filepath.Base(model), comp.Offloaded, comp.Total)
|
||||
}
|
||||
fmt.Fprintf(out, "computed a token on %s: %s, %d of %d layers offloaded, in %s\n", comp.Backend,
|
||||
strings.Join(comp.Devices, "; "), comp.Offloaded, comp.Total, took.Round(100*time.Millisecond))
|
||||
return nil
|
||||
}
|
||||
|
||||
func tailLines(s string, n int) string {
|
||||
l := strings.Split(strings.TrimSpace(s), "\n")
|
||||
if len(l) > n {
|
||||
l = l[len(l)-n:]
|
||||
}
|
||||
return strings.Join(l, " | ")
|
||||
}
|
||||
|
||||
// syncBuffer is a buffer two writers share.
|
||||
type syncBuffer struct {
|
||||
mu sync.Mutex
|
||||
b strings.Builder
|
||||
}
|
||||
|
||||
func (s *syncBuffer) Write(p []byte) (int, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.b.Write(p)
|
||||
}
|
||||
|
||||
func (s *syncBuffer) String() string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.b.String()
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestMain lets the test binary stand in for llama-server: run with
|
||||
// LLMGW_FAKE_LLAMA set, it prints what llama.cpp prints and answers.
|
||||
func TestMain(m *testing.M) {
|
||||
if mode := os.Getenv("LLMGW_FAKE_LLAMA"); mode != "" {
|
||||
fakeLlama(mode)
|
||||
return
|
||||
}
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
|
||||
func fakeLlama(mode string) {
|
||||
port := ""
|
||||
for i, a := range os.Args {
|
||||
if a == "--port" && i+1 < len(os.Args) {
|
||||
port = os.Args[i+1]
|
||||
}
|
||||
}
|
||||
if mode == "gpu" {
|
||||
fmt.Println("llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free")
|
||||
fmt.Println("load_tensors: offloaded 25/25 layers to GPU")
|
||||
} else {
|
||||
fmt.Println("warning: no usable GPU found, --gpu-layers option will be ignored")
|
||||
fmt.Println("load_tensors: offloaded 0/25 layers to GPU")
|
||||
}
|
||||
http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })
|
||||
http.HandleFunc("/completion", func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte(`{"content":"!"}`)) })
|
||||
_ = http.ListenAndServe("127.0.0.1:"+port, nil) //nolint:gosec // a test's stand-in
|
||||
}
|
||||
|
||||
func TestEngineCheckComputesATokenOnTheGPU(t *testing.T) {
|
||||
// nvidia-smi present, so the check runs.
|
||||
dir := t.TempDir()
|
||||
smi := filepath.Join(dir, "nvidia-smi")
|
||||
if err := os.WriteFile(smi, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH"))
|
||||
model := filepath.Join(dir, "tiny.gguf")
|
||||
_ = os.WriteFile(model, []byte("GGUF"), 0o644)
|
||||
self, _ := os.Executable()
|
||||
|
||||
t.Setenv("LLMGW_FAKE_LLAMA", "gpu")
|
||||
var out bytes.Buffer
|
||||
if err := computeOnGPU(t.Context(), self, model, &out); err != nil || !strings.Contains(out.String(), "on CUDA") ||
|
||||
!strings.Contains(out.String(), "25 of 25") {
|
||||
t.Fatalf("on the GPU: %v %s", err, out.String())
|
||||
}
|
||||
|
||||
t.Setenv("LLMGW_FAKE_LLAMA", "cpu")
|
||||
out.Reset()
|
||||
if err := computeOnGPU(t.Context(), self, model, &out); err == nil || !strings.Contains(err.Error(), "computed on the CPU") {
|
||||
t.Fatalf("a build that falls back to the CPU: %v %s", err, out.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineCheckFindsTheBuildTheAgentInstalled(t *testing.T) {
|
||||
old := llamaRoot
|
||||
llamaRoot = t.TempDir()
|
||||
t.Cleanup(func() { llamaRoot = old })
|
||||
if got := installedLlama(); got != "" {
|
||||
t.Fatalf("nothing installed, found %q", got)
|
||||
}
|
||||
for _, p := range []string{"b10991-cpu/llama-b10991/llama-server", "b10991-cuda-13.3/llama-b10991/llama-server"} {
|
||||
full := filepath.Join(llamaRoot, p)
|
||||
_ = os.MkdirAll(filepath.Dir(full), 0o755)
|
||||
_ = os.WriteFile(full, []byte("x"), 0o755)
|
||||
}
|
||||
dir := t.TempDir()
|
||||
_ = os.WriteFile(filepath.Join(dir, "nvidia-smi"), []byte("#!/bin/sh\n"), 0o755)
|
||||
t.Setenv("PATH", dir)
|
||||
if got := installedLlama(); !strings.Contains(got, "cuda-13.3/llama-b10991/llama-server") {
|
||||
t.Errorf("with GPUs, not the CUDA build: %q", got)
|
||||
}
|
||||
}
|
||||
+49
-21
@@ -7,12 +7,14 @@ import (
|
||||
"fmt"
|
||||
"net/netip"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
|
||||
)
|
||||
|
||||
// firewallStateFile is where a control plane writes what the cluster's
|
||||
@@ -20,11 +22,9 @@ import (
|
||||
// the server itself runs as llmgw and cannot change the ruleset.
|
||||
const firewallStateFile = "/var/lib/llmgw/firewall.json"
|
||||
|
||||
// firewallState is what is written there.
|
||||
type firewallState struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Networks []string `json:"networks,omitempty"`
|
||||
}
|
||||
// firewallState is what is written there: the same as the mirror is told
|
||||
// and keeps beside its models.
|
||||
type firewallState = mirror.FirewallState
|
||||
|
||||
// writeFirewallState records the settings' firewall for this machine's sync
|
||||
// unit. Without the state directory — a server not run by systemd — there
|
||||
@@ -46,25 +46,34 @@ func writeFirewallState(cfg config.Server) error {
|
||||
|
||||
// The units that apply it: a path unit watching the file, and the oneshot it
|
||||
// starts.
|
||||
const (
|
||||
firewallSyncPathUnit = "llmgw-firewall-sync.path"
|
||||
firewallSyncUnitFile = `[Unit]
|
||||
Description=llmgw firewall: apply the cluster's settings on this control plane
|
||||
const firewallSyncPathUnit = "llmgw-firewall-sync.path"
|
||||
|
||||
// firewallSyncUnits are the two units for a state file: on a control plane
|
||||
// the one its server writes, on a mirror the one the control plane has the
|
||||
// mirror write.
|
||||
func firewallSyncUnits(state string) (service, pathUnit string) {
|
||||
arg := ""
|
||||
if state != firewallStateFile {
|
||||
arg = " --state " + state
|
||||
}
|
||||
service = `[Unit]
|
||||
Description=llmgw firewall: apply the cluster's settings on this machine
|
||||
|
||||
[Service]
|
||||
Type=oneshot
|
||||
ExecStart=/usr/local/bin/llmgw firewall sync
|
||||
ExecStart=/usr/local/bin/llmgw firewall sync` + arg + `
|
||||
`
|
||||
firewallSyncPathFile = `[Unit]
|
||||
Description=llmgw firewall: follow the cluster's settings on this control plane
|
||||
pathUnit = `[Unit]
|
||||
Description=llmgw firewall: follow the cluster's settings on this machine
|
||||
|
||||
[Path]
|
||||
PathChanged=` + firewallStateFile + `
|
||||
PathChanged=` + state + `
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
`
|
||||
)
|
||||
return service, pathUnit
|
||||
}
|
||||
|
||||
func newFirewallCmd() *cobra.Command {
|
||||
c := &cobra.Command{
|
||||
@@ -76,6 +85,7 @@ func newFirewallCmd() *cobra.Command {
|
||||
"which " + firewallSyncPathUnit + " applies with 'llmgw firewall sync'.",
|
||||
}
|
||||
var install bool
|
||||
state := firewallStateFile
|
||||
sync := &cobra.Command{
|
||||
Use: "sync",
|
||||
Short: "Apply what the cluster's settings say of the firewall, here (root)",
|
||||
@@ -85,21 +95,21 @@ func newFirewallCmd() *cobra.Command {
|
||||
return errors.New("the firewall is changed as root: run it with sudo")
|
||||
}
|
||||
if install {
|
||||
if err := installFirewallSync(cmd.Context(), firewall.Local{}); err != nil {
|
||||
if err := installFirewallSync(cmd.Context(), firewall.Local{}, state); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintln(cmd.OutOrStdout(), firewallSyncPathUnit+" follows the cluster's settings here")
|
||||
if _, err := os.Stat(firewallStateFile); err != nil {
|
||||
return nil // the server writes it when it starts
|
||||
if _, err := os.Stat(state); err != nil {
|
||||
return nil // written when the settings are first told
|
||||
}
|
||||
}
|
||||
b, err := os.ReadFile(firewallStateFile)
|
||||
b, err := os.ReadFile(state)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var st firewallState
|
||||
if err := json.Unmarshal(b, &st); err != nil {
|
||||
return fmt.Errorf("%s: %w", firewallStateFile, err)
|
||||
return fmt.Errorf("%s: %w", state, err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(cmd.Context(), 2*time.Minute)
|
||||
defer cancel()
|
||||
@@ -112,14 +122,17 @@ func newFirewallCmd() *cobra.Command {
|
||||
},
|
||||
}
|
||||
sync.Flags().BoolVar(&install, "install", false, "first install the units that run this whenever the settings change")
|
||||
sync.Flags().StringVar(&state, "state", firewallStateFile,
|
||||
"the file the settings are written to: a control plane's, or "+mirror.FirewallFile+" in a mirror's directory")
|
||||
c.AddCommand(sync)
|
||||
return c
|
||||
}
|
||||
|
||||
// installFirewallSync writes the sync units on a machine and has the path
|
||||
// unit follow the state file.
|
||||
func installFirewallSync(ctx context.Context, r firewall.Runner) error {
|
||||
body := firewallSyncUnitFile + "\x00" + firewallSyncPathFile
|
||||
func installFirewallSync(ctx context.Context, r firewall.Runner, state string) error {
|
||||
service, pathUnit := firewallSyncUnits(state)
|
||||
body := service + "\x00" + pathUnit
|
||||
out, err := r.Run(ctx, `set -e
|
||||
t=$(mktemp -d); trap 'rm -rf "$t"' EXIT
|
||||
cat > "$t/all"
|
||||
@@ -172,5 +185,20 @@ func syncFirewall(ctx context.Context, r firewall.Runner, st firewallState) (str
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
// The cluster's machines: a node on a network this one has not let in
|
||||
// yet, which the mirror's files are for.
|
||||
if len(st.Toward) > 0 {
|
||||
added, err := firewall.Admit(ctx, r, st.Toward)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(added) > 0 {
|
||||
shown := make([]string, len(added))
|
||||
for i, p := range added {
|
||||
shown[i] = p.String()
|
||||
}
|
||||
return "llmgw's firewall is on here, and now lets in " + strings.Join(shown, ", "), nil
|
||||
}
|
||||
}
|
||||
return "llmgw's firewall is on here", nil
|
||||
}
|
||||
@@ -2,8 +2,11 @@ package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
|
||||
)
|
||||
|
||||
type scriptLog struct {
|
||||
@@ -49,3 +52,32 @@ func TestAControlPlaneFollowsTheFirewallSetting(t *testing.T) {
|
||||
t.Errorf("on: %q %v", msg, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMachineLetsInTheClusterMachinesItIsTold(t *testing.T) {
|
||||
// The rules hold 10.0.0.0/24; a node on 192.168.7.0/24 is let in, as
|
||||
// the route to it says, and one already in is not asked about.
|
||||
r := &scriptLog{reply: func(s string) string {
|
||||
switch {
|
||||
case strings.Contains(s, "echo loaded"):
|
||||
return "loaded\n"
|
||||
case strings.Contains(s, "cat /etc/llmgw/firewall.nft"):
|
||||
return firewallRulesFor("10.0.0.0/24")
|
||||
case strings.Contains(s, "ip -o route get 192.168.7.20"):
|
||||
return "192.168.7.20 dev eth1 src 192.168.7.1 uid 0\n2: eth1 inet 192.168.7.1/24 brd 192.168.7.255 scope global eth1\n"
|
||||
case strings.Contains(s, "nft -f"):
|
||||
return "changed\n"
|
||||
}
|
||||
return ""
|
||||
}}
|
||||
msg, err := syncFirewall(context.Background(), r, firewallState{Enabled: true, Toward: []string{"10.0.0.147", "192.168.7.20"}})
|
||||
if err != nil || !strings.Contains(msg, "192.168.7.0/24") {
|
||||
t.Fatalf("%q %v", msg, err)
|
||||
}
|
||||
if strings.Contains(strings.Join(r.scripts, "\n"), "route get 10.0.0.147") {
|
||||
t.Error("asked the route to a machine already let in")
|
||||
}
|
||||
}
|
||||
|
||||
func firewallRulesFor(nets string) string {
|
||||
return firewall.Script(firewall.Spec{Ports: []firewall.Port{firewall.MirrorFiles}, Allow: []netip.Prefix{netip.MustParsePrefix(nets)}})
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// firewallStateKey is where, in the store, the firewall's state is kept
|
||||
// for the store's members to read.
|
||||
const firewallStateKey = "firewall"
|
||||
|
||||
// firewallPublisher tells the machines that cannot ask a control plane
|
||||
// what the cluster's settings say of their firewall: the mirror, over its
|
||||
// API, and the store's members, through the store they hold a copy of.
|
||||
// Nodes are told over their agents' streams and control planes by their
|
||||
// own servers. On the leader only, at start, at every change to the
|
||||
// firewall's settings, and every few minutes for nodes come and gone.
|
||||
type firewallPublisher struct {
|
||||
settings func() config.Server
|
||||
st *store.DB
|
||||
nodes func() []string
|
||||
leader func() bool
|
||||
log *slog.Logger
|
||||
|
||||
mu sync.Mutex
|
||||
mirror *mirror.Client
|
||||
kick chan struct{}
|
||||
}
|
||||
|
||||
func newFirewallPublisher(settings func() config.Server, st *store.DB, nodes func() []string,
|
||||
leader func() bool, log *slog.Logger) *firewallPublisher {
|
||||
return &firewallPublisher{settings: settings, st: st, nodes: nodes, leader: leader, log: log,
|
||||
kick: make(chan struct{}, 1)}
|
||||
}
|
||||
|
||||
// setMirror has the mirror told too.
|
||||
func (p *firewallPublisher) setMirror(m *mirror.Client) {
|
||||
p.mu.Lock()
|
||||
p.mirror = m
|
||||
p.mu.Unlock()
|
||||
p.Kick()
|
||||
}
|
||||
|
||||
// Kick has the state worked out and told again now.
|
||||
func (p *firewallPublisher) Kick() {
|
||||
select {
|
||||
case p.kick <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (p *firewallPublisher) run(ctx context.Context, every time.Duration) {
|
||||
t := time.NewTicker(every)
|
||||
defer t.Stop()
|
||||
for {
|
||||
p.once(ctx)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
case <-p.kick:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// state is what the settings say now, with the cluster's machines.
|
||||
func (p *firewallPublisher) state(ctx context.Context) mirror.FirewallState {
|
||||
cfg := p.settings()
|
||||
st := mirror.FirewallState{Enabled: cfg.Firewall.Enabled, Networks: slices.Clone(cfg.FirewallNetworks)}
|
||||
var toward []string
|
||||
if p.nodes != nil {
|
||||
toward = append(toward, p.nodes()...)
|
||||
}
|
||||
if p.st != nil {
|
||||
if cps, err := p.st.ControlPlanes(ctx); err == nil {
|
||||
for _, c := range cps {
|
||||
toward = append(toward, c.URL)
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, t := range toward {
|
||||
if a, ok := hostAddr(t); ok {
|
||||
st.Toward = append(st.Toward, a)
|
||||
}
|
||||
}
|
||||
slices.Sort(st.Toward)
|
||||
st.Toward = slices.Compact(st.Toward)
|
||||
return st
|
||||
}
|
||||
|
||||
// hostAddr is the address in an address, host:port or URL; names are not
|
||||
// looked up here.
|
||||
func hostAddr(s string) (string, bool) {
|
||||
if u, err := url.Parse(s); err == nil && u.Host != "" {
|
||||
s = u.Hostname()
|
||||
} else if h, _, err := net.SplitHostPort(s); err == nil {
|
||||
s = h
|
||||
}
|
||||
a, err := netip.ParseAddr(s)
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
return a.Unmap().String(), true
|
||||
}
|
||||
|
||||
func (p *firewallPublisher) once(ctx context.Context) {
|
||||
if p.leader != nil && !p.leader() {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer cancel()
|
||||
st := p.state(ctx)
|
||||
p.mu.Lock()
|
||||
m := p.mirror
|
||||
p.mu.Unlock()
|
||||
if m != nil {
|
||||
changed, err := m.SetFirewall(ctx, st)
|
||||
switch {
|
||||
case err != nil:
|
||||
p.log.Warn("could not tell the mirror what the settings say of its firewall", "err", err)
|
||||
case changed:
|
||||
p.log.Info("the mirror's firewall follows the cluster's settings", "enabled", st.Enabled)
|
||||
}
|
||||
}
|
||||
if p.st != nil {
|
||||
b, _ := json.Marshal(st)
|
||||
changed, err := p.st.PutClusterState(ctx, firewallStateKey, string(b))
|
||||
switch {
|
||||
case err != nil:
|
||||
p.log.Warn("could not record the firewall's state for the store's members", "err", err)
|
||||
case changed:
|
||||
p.log.Info("the store's members' firewall follows the cluster's settings", "enabled", st.Enabled)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
func TestTheLeaderTellsTheMirrorAndTheStoreWhatTheFirewallSettingsSay(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
st, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "s.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer st.Close()
|
||||
if err := st.PutControlPlane(ctx, store.ControlPlane{URL: "https://10.0.0.150:7443", Name: "ctl02"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
root := t.TempDir()
|
||||
srv := httptest.NewServer(mirror.NewAPI(mirror.NewManager(root, slog.New(slog.DiscardHandler)), "http://m:8081").Handler())
|
||||
defer srv.Close()
|
||||
mc, _ := mirror.NewClient(srv.URL, tlsconf.Config{}, 0)
|
||||
|
||||
cfg := config.Server{Firewall: config.FirewallSettings{Enabled: false}, FirewallNetworks: []string{"192.168.5.0/24"}}
|
||||
leader := false
|
||||
p := newFirewallPublisher(func() config.Server { return cfg }, st,
|
||||
func() []string { return []string{"10.0.0.147", "", "10.0.0.147", "node-by-name"} },
|
||||
func() bool { return leader }, slog.New(slog.DiscardHandler))
|
||||
p.setMirror(mc)
|
||||
|
||||
p.once(ctx)
|
||||
if _, err := st.ClusterState(ctx, firewallStateKey); err == nil {
|
||||
t.Error("a standby told the store")
|
||||
}
|
||||
leader = true
|
||||
p.once(ctx)
|
||||
var got mirror.FirewallState
|
||||
v, err := st.ClusterState(ctx, firewallStateKey)
|
||||
if err != nil || json.Unmarshal([]byte(v), &got) != nil {
|
||||
t.Fatalf("the store holds %q %v", v, err)
|
||||
}
|
||||
if got.Enabled || got.Networks[0] != "192.168.5.0/24" || len(got.Toward) != 2 ||
|
||||
got.Toward[0] != "10.0.0.147" || got.Toward[1] != "10.0.0.150" {
|
||||
t.Errorf("state %+v", got)
|
||||
}
|
||||
b, err := os.ReadFile(filepath.Join(root, mirror.FirewallFile))
|
||||
if err != nil || string(b) != v {
|
||||
t.Errorf("the mirror holds %s, the store %s", b, v)
|
||||
}
|
||||
}
|
||||
@@ -3,11 +3,13 @@ package cli
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
|
||||
)
|
||||
|
||||
func newGuardCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
@@ -24,6 +26,35 @@ func newGuardCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
"Targets are written org:Acme, team:Acme/ml, project:Acme/ml/training (or\n" +
|
||||
"with IDs), key:KEY-ID or user:USER-ID. An attachment applies at once.",
|
||||
}
|
||||
c.AddCommand(newGuardJudgeCmd(client, asJSON))
|
||||
for _, sw := range []struct {
|
||||
use, short string
|
||||
on bool
|
||||
}{
|
||||
{"on", "Turn the guardrails on: every policy checks again", true},
|
||||
{"off", "Turn the guardrails off: no request is checked, whatever the policies say", false},
|
||||
} {
|
||||
var note string
|
||||
cmd := &cobra.Command{
|
||||
Use: sw.use,
|
||||
Short: sw.short,
|
||||
Long: "guard.enabled in the cluster's settings: in force on every control plane\n" +
|
||||
"within seconds, without a restart. The policies and where they are attached\n" +
|
||||
"stay as they are, and come back as they were when turned on again. While\n" +
|
||||
"off, llmgw ctl status says so.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if note == "" {
|
||||
note = "guardrails " + sw.use
|
||||
}
|
||||
return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
|
||||
return doc.SetYAML([]string{"guard", "enabled"}, strconv.FormatBool(sw.on))
|
||||
})
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(¬e, "note", "", "why, kept with the settings' version")
|
||||
c.AddCommand(cmd)
|
||||
}
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "list",
|
||||
Short: "The policies, and where each is attached",
|
||||
@@ -37,6 +68,9 @@ func newGuardCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
if !v.Enabled {
|
||||
fmt.Fprintln(out, "The guardrails are OFF: no request is checked (llmgw ctl guard on).")
|
||||
}
|
||||
fmt.Fprintf(out, "Policies: %s\n", strings.Join(v.Policies, ", "))
|
||||
def := strings.Join(v.Default, ", ")
|
||||
if def == "" {
|
||||
|
||||
@@ -0,0 +1,385 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"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/guard"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
|
||||
)
|
||||
|
||||
func newGuardJudgeCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
c := &cobra.Command{
|
||||
Use: "judge",
|
||||
Short: "The judge models the policies ask: whether each can answer, and serving one",
|
||||
Long: "A policy with a judge has a model read what patterns cannot. The judge must\n" +
|
||||
"be in the catalog under the name the policy gives, as a judge (role:\n" +
|
||||
"guard), with a replica ready; otherwise what the policy judges is refused\n" +
|
||||
"(fail: closed, guard_unavailable) or served unjudged (fail: open).\n\n" +
|
||||
"A policy may name several, asked in turn (judge.models): a judge on a GPU\n" +
|
||||
"first and one on CPUs after it is the GPU's speed while it is there, and\n" +
|
||||
"the CPUs' when it is not.\n\n" +
|
||||
"Judges llmgw knows where to fetch, which a policy may name and the leader\n" +
|
||||
"then fetches and serves by itself: " + knownJudgeNames() + ".",
|
||||
}
|
||||
c.AddCommand(newGuardJudgeStatusCmd(client, asJSON), newGuardJudgeEnableCmd(client),
|
||||
newGuardJudgeUseCmd(client), newGuardJudgeAskCmd(client, asJSON))
|
||||
return c
|
||||
}
|
||||
|
||||
func newGuardJudgeStatusCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "status",
|
||||
Short: "Each judge the policies ask, and what is wrong with it",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var js []adminapi.JudgeView
|
||||
if err := client().get(cmd.Context(), "/api/v1/guard/judges", &js); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(js)
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
if len(js) == 0 {
|
||||
fmt.Fprintln(out, "no policy asks a judge")
|
||||
return nil
|
||||
}
|
||||
w := tabTo(out)
|
||||
fmt.Fprintln(w, "JUDGE\tPOLICIES\tIN CATALOG\tAS JUDGE\tREADY\tSTATE")
|
||||
for _, j := range js {
|
||||
state := "answering"
|
||||
if j.Problem != "" {
|
||||
state = j.Problem
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%d/%d\t%s\n", j.Model, strings.Join(j.Policies, ","),
|
||||
yesNo(j.InCatalog), yesNo(j.Judge), j.Ready, j.Wanted, state)
|
||||
}
|
||||
for _, j := range js {
|
||||
if len(j.Then) > 0 {
|
||||
fmt.Fprintf(w, "\n%s gives no verdict: %s is asked instead", j.Model, strings.Join(j.Then, ", else "))
|
||||
}
|
||||
}
|
||||
if err := w.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
strays, _ := strayJudges(cmd, client, js)
|
||||
sayStrays(out, strays)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// strayJudges are models whose weights are a judge's, under a name no policy
|
||||
// asks: a judge the policies moved on from, or one served from the mirror
|
||||
// page as a model clients use, as before that page knew better.
|
||||
func strayJudges(cmd *cobra.Command, client func() *ctlClient, js []adminapi.JudgeView) ([]adminapi.ModelView, error) {
|
||||
var models []adminapi.ModelView
|
||||
if err := client().get(cmd.Context(), "/api/v1/models", &models); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
asked := map[string]bool{}
|
||||
for _, j := range js {
|
||||
asked[j.Model] = true
|
||||
}
|
||||
var out []adminapi.ModelView
|
||||
for _, m := range models {
|
||||
if _, judge := guard.FormatOfWeights(m.Source); !asked[m.Name] && judge {
|
||||
out = append(out, m)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// sayStrays says what each stray judge is and how to be rid of it.
|
||||
func sayStrays(out io.Writer, strays []adminapi.ModelView) {
|
||||
for _, m := range strays {
|
||||
if m.Role == catalog.RoleGuard {
|
||||
fmt.Fprintf(out, "\n%s is a judge no policy asks, still running: llmgw ctl model remove %s\n", m.Name, m.Name)
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(out, "\n%s looks like a judge's weights but is served as a model clients use, and no "+
|
||||
"policy asks it: llmgw ctl model remove %s\n", m.Name, m.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func newGuardJudgeEnableCmd(client func() *ctlClient) *cobra.Command {
|
||||
var (
|
||||
model, repo, file, format, gpuClass string
|
||||
replicas, cores int
|
||||
gpu, replace, detach bool
|
||||
)
|
||||
c := &cobra.Command{
|
||||
Use: "enable",
|
||||
Short: "Serve a judge: Qwen3Guard-Gen 0.6B on CPUs, or with --gpu the 4B on a GPU",
|
||||
Long: "Fetches a judge's weights into the mirror (once; what is there already, its\n" +
|
||||
"checksums matching, is not fetched again), adds it to the catalog as a judge\n" +
|
||||
"(role: guard) and waits for a replica to answer. Which judge: --model, or\n" +
|
||||
"with --gpu " + defaultGPUJudge + ", or else the first the policies ask that is\n" +
|
||||
"missing, or else " + defaultJudge + ".\n\n" +
|
||||
"A judge on CPUs takes no card. One on a GPU takes one card, as a standard\n" +
|
||||
"model: from nothing more important, so it waits while the cards are busy;\n" +
|
||||
"put a judge on CPUs after it in the policies (llmgw ctl guard judge use) so\n" +
|
||||
"there is always one to ask.\n\n" +
|
||||
"Known judges: " + knownJudgeNames() + ". Another is served with --repo,\n" +
|
||||
"--file and --format. A model of that name already in the catalog as a\n" +
|
||||
"model clients use is refused; --replace removes it first. An offline\n" +
|
||||
"cluster gets its judge from its bundle with server init.",
|
||||
Example: " llmgw ctl guard judge enable\n" +
|
||||
" llmgw ctl guard judge enable --gpu\n" +
|
||||
" llmgw ctl guard judge use " + defaultGPUJudge + " " + defaultJudge,
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
ctx := cmd.Context()
|
||||
out := cmd.OutOrStdout()
|
||||
var js []adminapi.JudgeView
|
||||
if err := client().get(ctx, "/api/v1/guard/judges", &js); err != nil {
|
||||
return err
|
||||
}
|
||||
name := model
|
||||
if name == "" && gpu {
|
||||
name = defaultGPUJudge
|
||||
}
|
||||
if name == "" {
|
||||
for _, j := range js {
|
||||
if !j.InCatalog || !j.Judge {
|
||||
name = j.Model
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if name == "" {
|
||||
name = defaultJudge
|
||||
}
|
||||
|
||||
k, known := knownJudgeNamed(name)
|
||||
switch {
|
||||
case repo != "":
|
||||
f := format
|
||||
if f == "" {
|
||||
f, _ = guard.FormatOfWeights(repo + "/" + file)
|
||||
}
|
||||
if f == "" || !guard.KnownFormat(f) {
|
||||
return fmt.Errorf("say how its answers are read: --format %s", strings.Join(guard.Formats, "|"))
|
||||
}
|
||||
k = knownJudge{Name: name, Repo: repo, File: file, Format: f, GPU: gpu}
|
||||
case !known:
|
||||
return fmt.Errorf("%s is not a judge llmgw knows where to fetch (%s): give --repo, --file and --format", name, knownJudgeNames())
|
||||
}
|
||||
if gpu {
|
||||
k.GPU = true
|
||||
}
|
||||
|
||||
var models []adminapi.ModelView
|
||||
if err := client().get(ctx, "/api/v1/models", &models); err != nil {
|
||||
return err
|
||||
}
|
||||
if i := slices.IndexFunc(models, func(m adminapi.ModelView) bool { return m.Name == name }); i >= 0 {
|
||||
if models[i].Role == catalog.RoleGuard {
|
||||
fmt.Fprintf(out, "%s is in the catalog as a judge already; llmgw ctl guard judge status says how it is\n", name)
|
||||
return nil
|
||||
}
|
||||
if !replace {
|
||||
return fmt.Errorf("%s is in the catalog as a model clients use, not as a judge: "+
|
||||
"--replace removes it and serves it again as a judge", name)
|
||||
}
|
||||
var res map[string]string
|
||||
if err := client().delete(ctx, "/api/v1/models/"+name, &res); err != nil {
|
||||
return fmt.Errorf("removing %s: %w", name, err)
|
||||
}
|
||||
fmt.Fprintf(out, "%s removed\n", name)
|
||||
}
|
||||
|
||||
opts := judgeServeOptions{Replicas: replicas, Cores: cores, GPUClass: gpuClass}
|
||||
if k.GPU && opts.GPUClass == "" {
|
||||
var s adminapi.Summary
|
||||
if err := client().get(ctx, "/api/v1/summary", &s); err != nil {
|
||||
return err
|
||||
}
|
||||
if opts.GPUClass = roomiestGPUClass(s.GPUs); opts.GPUClass == "" {
|
||||
return fmt.Errorf("this cluster has no GPU for %s: serve %s on CPUs instead", name, defaultJudge)
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(out, "serving %s\n", k)
|
||||
body, _ := json.Marshal(k.serveRequest(opts))
|
||||
var job modelserve.Job
|
||||
if err := client().patient(30*time.Second).post(ctx, "/api/v1/models/serve", body, &job); err != nil {
|
||||
return err
|
||||
}
|
||||
if detach {
|
||||
return detached(cmd, job)
|
||||
}
|
||||
if err := followJob(cmd, client, job, "is the judge"); err != nil {
|
||||
return err
|
||||
}
|
||||
if !slices.ContainsFunc(js, func(j adminapi.JudgeView) bool { return j.Model == name }) {
|
||||
fmt.Fprintf(out, "\nno policy asks %s yet: llmgw ctl guard judge use %s", name, name)
|
||||
if k.GPU {
|
||||
fmt.Fprintf(out, " %s", defaultJudge)
|
||||
}
|
||||
fmt.Fprintln(out)
|
||||
}
|
||||
if strays, _ := strayJudges(cmd, client, js); len(strays) > 0 {
|
||||
sayStrays(out, slices.DeleteFunc(strays, func(m adminapi.ModelView) bool { return m.Name == name }))
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
f := c.Flags()
|
||||
f.StringVar(&model, "model", "", "the judge to serve: "+knownJudgeNames()+", or a name for --repo's")
|
||||
f.BoolVar(&gpu, "gpu", false, "serve it on a GPU ("+defaultGPUJudge+" when no --model)")
|
||||
f.StringVar(&gpuClass, "gpu-class", "", "the GPUs it runs on (default: the class with the most free cards)")
|
||||
f.StringVar(&repo, "repo", "", "the registry repository of a judge llmgw does not know")
|
||||
f.StringVar(&file, "file", "", "the GGUF file in --repo")
|
||||
f.StringVar(&format, "format", "", "how --repo's judge answers: "+strings.Join(guard.Formats, ", "))
|
||||
f.IntVar(&replicas, "replicas", 0, "replicas (default: 2 on CPUs, as every judged request waits on one; 1 on a GPU)")
|
||||
f.IntVar(&cores, "cores", 4, "CPU cores a replica on CPUs computes on")
|
||||
f.BoolVar(&replace, "replace", false, "remove a model of that name that is not a judge first")
|
||||
f.BoolVar(&detach, "detach", false, "start it and return; llmgw ctl jobs follows it")
|
||||
return c
|
||||
}
|
||||
|
||||
// roomiestGPUClass is the class with the most free cards, else the most
|
||||
// cards; empty when the cluster has none.
|
||||
func roomiestGPUClass(classes []adminapi.ClassUse) string {
|
||||
best := -1
|
||||
for i, c := range classes {
|
||||
if c.Total == 0 {
|
||||
continue
|
||||
}
|
||||
if best < 0 || c.Free() > classes[best].Free() || (c.Free() == classes[best].Free() && c.Total > classes[best].Total) {
|
||||
best = i
|
||||
}
|
||||
}
|
||||
if best < 0 {
|
||||
return ""
|
||||
}
|
||||
return classes[best].Class
|
||||
}
|
||||
|
||||
func newGuardJudgeUseCmd(client func() *ctlClient) *cobra.Command {
|
||||
var policies []string
|
||||
var note string
|
||||
c := &cobra.Command{
|
||||
Use: "use JUDGE...",
|
||||
Short: "Set the judges the policies ask, in the order they are asked",
|
||||
Long: "Sets judge.models of every policy that has a judge, or of those given with\n" +
|
||||
"--policy: the first judge that answers gives the verdict, and the next is\n" +
|
||||
"asked when one has no replica serving, fails, or does not answer within\n" +
|
||||
"the policy's latency_budget. A judge on a GPU first and one on CPUs after\n" +
|
||||
"it is the GPU's speed while it is there, and the CPUs' when it is not.\n\n" +
|
||||
"A judge llmgw knows (" + knownJudgeNames() + ")\n" +
|
||||
"that is not in the catalog is fetched and served by the leader; another\n" +
|
||||
"is served first with llmgw ctl guard judge enable --repo. In force on\n" +
|
||||
"every control plane within seconds.",
|
||||
Example: " llmgw ctl guard judge use " + defaultJudge + "\n" +
|
||||
" llmgw ctl guard judge use " + defaultGPUJudge + " " + defaultJudge + "\n" +
|
||||
" llmgw ctl guard judge use " + defaultGPUJudge + " " + defaultJudge + " --policy strict",
|
||||
Args: cobra.MinimumNArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
for i, a := range args {
|
||||
if slices.Contains(args[:i], a) {
|
||||
return fmt.Errorf("%s is named twice", a)
|
||||
}
|
||||
}
|
||||
if note == "" {
|
||||
note = "judges: " + strings.Join(args, ", else ")
|
||||
}
|
||||
list := "[" + strings.Join(args, ", ") + "]"
|
||||
return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
|
||||
names := policies
|
||||
if len(names) == 0 {
|
||||
for _, p := range doc.MapKeys([]string{"guard", "policies"}) {
|
||||
if doc.Node([]string{"guard", "policies", p, "judge"}) != nil {
|
||||
names = append(names, p)
|
||||
}
|
||||
}
|
||||
if len(names) == 0 {
|
||||
return errors.New("no policy has a judge: name one with --policy, and it is given one " +
|
||||
"that blocks what the judge finds unsafe in requests (judge.input: block)")
|
||||
}
|
||||
}
|
||||
for _, p := range names {
|
||||
base := []string{"guard", "policies", p}
|
||||
if doc.Node(base) == nil {
|
||||
return fmt.Errorf("there is no policy %s (llmgw ctl guard list)", p)
|
||||
}
|
||||
judge := append(slices.Clone(base), "judge")
|
||||
if doc.Node(judge) == nil {
|
||||
if err := doc.SetYAML(append(slices.Clone(judge), "input"), "block"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
doc.Unset(append(slices.Clone(judge), "model"))
|
||||
if err := doc.SetYAML(append(slices.Clone(judge), "models"), list); err != nil {
|
||||
return err
|
||||
}
|
||||
cats := append(slices.Clone(judge), "categories")
|
||||
var have []string
|
||||
if n := doc.Node(cats); n != nil {
|
||||
for _, c := range n.Content {
|
||||
have = append(have, c.Value)
|
||||
}
|
||||
}
|
||||
if next := categoriesFor(args, have); next != nil {
|
||||
if err := doc.SetYAML(cats, "["+strings.Join(next, ", ")+"]"); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: categories [%s] are not what %s name; it counts [%s] instead\n",
|
||||
p, strings.Join(have, ", "), strings.Join(args, " or "), strings.Join(next, ", "))
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
},
|
||||
}
|
||||
c.Flags().StringSliceVar(&policies, "policy", nil, "the policies to change (default: every one with a judge)")
|
||||
c.Flags().StringVar(¬e, "note", "", "why, kept with the settings' version")
|
||||
return c
|
||||
}
|
||||
|
||||
func newGuardJudgeAskCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
var model, answer string
|
||||
c := &cobra.Command{
|
||||
Use: "ask TEXT",
|
||||
Short: "Have a judge read a text and say what it answered",
|
||||
Long: "For finding out why a request was refused: the judge reads TEXT as a user's\n" +
|
||||
"message (with --answer, ANSWER as the reply to it), past the cache, and its\n" +
|
||||
"answer is shown as it gave it and as the gateway reads it — \"unsafe\" and\n" +
|
||||
"the hazards it named, S6 Specialized Advice and the like. Nothing is kept:\n" +
|
||||
"not in the audit, the usage or the logs.",
|
||||
Example: " llmgw ctl guard judge ask 'How do I sort a list in Python?'\n" +
|
||||
" llmgw ctl guard judge ask 'Write a poem' --answer 'Roses are red'",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
body, _ := json.Marshal(adminapi.JudgeAskRequest{Model: model, Text: args[0], Answer: answer})
|
||||
var v adminapi.JudgeAskView
|
||||
if err := client().patient(45*time.Second).post(cmd.Context(), "/api/v1/guard/judges/ask", body, &v); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return emitJSON(v)
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintf(out, "%s: %s", v.Model, v.Verdict)
|
||||
if len(v.Categories) > 0 {
|
||||
fmt.Fprintf(out, " (%s)", strings.Join(v.Categories, ", "))
|
||||
}
|
||||
fmt.Fprintf(out, ", %d ms\nit answered: %q\n", v.TookMs, v.Raw)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
c.Flags().StringVar(&model, "model", "", "the judge (default: the one the policies ask)")
|
||||
c.Flags().StringVar(&answer, "answer", "", "an assistant's reply to TEXT, judged as output is")
|
||||
return c
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"slices"
|
||||
"time"
|
||||
|
||||
"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/guard"
|
||||
)
|
||||
|
||||
// judgeKeeper fetches and serves, on the leader, each judge llmgw knows that
|
||||
// a policy names and the catalog does not have: naming qwen3guard-4b in a
|
||||
// policy is all it takes to have one, and a judge removed by mistake comes
|
||||
// back. A judge it does not know is left to llmgw ctl guard judge enable
|
||||
// --repo, said in guard judge status.
|
||||
type judgeKeeper struct {
|
||||
library func() guard.Config
|
||||
models func() []catalog.Model
|
||||
serve interface {
|
||||
Start(modelserve.Request) (modelserve.Job, error)
|
||||
Jobs() []modelserve.Job
|
||||
}
|
||||
leader func() bool
|
||||
gpus func() []adminapi.ClassUse
|
||||
log *slog.Logger
|
||||
now func() time.Time
|
||||
|
||||
// tried is when each judge was last asked for, so one that cannot be
|
||||
// fetched is not asked for again on every look.
|
||||
tried map[string]time.Time
|
||||
}
|
||||
|
||||
// judgeRetry is how long a judge whose serving failed waits before it is
|
||||
// asked for again.
|
||||
const judgeRetry = 10 * time.Minute
|
||||
|
||||
func (k *judgeKeeper) run(ctx context.Context, every time.Duration) {
|
||||
t := time.NewTicker(every)
|
||||
defer t.Stop()
|
||||
for {
|
||||
k.check()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (k *judgeKeeper) check() {
|
||||
if k.leader != nil && !k.leader() {
|
||||
return
|
||||
}
|
||||
lib := k.library()
|
||||
if !lib.On() {
|
||||
return
|
||||
}
|
||||
if k.tried == nil {
|
||||
k.tried = map[string]time.Time{}
|
||||
}
|
||||
now := time.Now
|
||||
if k.now != nil {
|
||||
now = k.now
|
||||
}
|
||||
have := map[string]bool{}
|
||||
for _, m := range k.models() {
|
||||
have[m.Name] = true
|
||||
}
|
||||
busy := map[string]bool{}
|
||||
for _, j := range k.serve.Jobs() {
|
||||
if j.State == modelserve.Running {
|
||||
busy[j.Model] = true
|
||||
busy[j.Request.Name] = true
|
||||
}
|
||||
}
|
||||
var wanted []string
|
||||
for _, p := range lib.Policies {
|
||||
for _, j := range p.Judge.Judges() {
|
||||
if !slices.Contains(wanted, j) {
|
||||
wanted = append(wanted, j)
|
||||
}
|
||||
}
|
||||
}
|
||||
slices.Sort(wanted)
|
||||
for _, name := range wanted {
|
||||
kj, known := knownJudgeNamed(name)
|
||||
if have[name] || !known || busy[name] {
|
||||
continue
|
||||
}
|
||||
if at, ok := k.tried[name]; ok && now().Sub(at) < judgeRetry {
|
||||
continue
|
||||
}
|
||||
k.tried[name] = now()
|
||||
var opts judgeServeOptions
|
||||
if kj.GPU {
|
||||
if opts.GPUClass = roomiestGPUClass(k.gpus()); opts.GPUClass == "" {
|
||||
k.log.Warn("a policy names a judge that runs on a GPU, and this cluster has none",
|
||||
"judge", name, "instead", defaultJudge)
|
||||
continue
|
||||
}
|
||||
}
|
||||
req := kj.serveRequest(opts)
|
||||
req.By = "llmgw (a guard policy names it)"
|
||||
job, err := k.serve.Start(req)
|
||||
if err != nil {
|
||||
k.log.Error("could not start serving a judge a policy names", "judge", name, "err", err)
|
||||
continue
|
||||
}
|
||||
k.log.Info("serving a judge a policy names, which the catalog did not have",
|
||||
"judge", name, "repo", kj.Repo, "file", kj.File, "job", job.ID)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"log/slog"
|
||||
"slices"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"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/guard"
|
||||
)
|
||||
|
||||
type fakeServe struct {
|
||||
started []modelserve.Request
|
||||
running []modelserve.Job
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeServe) Start(r modelserve.Request) (modelserve.Job, error) {
|
||||
if f.err != nil {
|
||||
return modelserve.Job{}, f.err
|
||||
}
|
||||
f.started = append(f.started, r)
|
||||
return modelserve.Job{ID: "j", Model: r.Name, Request: r, State: modelserve.Running}, nil
|
||||
}
|
||||
func (f *fakeServe) Jobs() []modelserve.Job { return f.running }
|
||||
|
||||
func TestAJudgeAPolicyNamesIsFetchedWhenTheCatalogLacksIt(t *testing.T) {
|
||||
lib := guard.Config{Policies: map[string]guard.Policy{
|
||||
"baseline": {Judge: &guard.JudgePolicy{Models: []string{"qwen3guard-4b", "qwen3guard-0.6b"}, Input: "block"}},
|
||||
"custom": {Judge: &guard.JudgePolicy{Model: "our-own-judge", Input: "block"}},
|
||||
}}
|
||||
inCatalog := []catalog.Model{{Name: "qwen3guard-0.6b"}}
|
||||
now := time.Date(2026, 10, 8, 12, 0, 0, 0, time.UTC)
|
||||
fs := &fakeServe{}
|
||||
leader := true
|
||||
k := &judgeKeeper{
|
||||
library: func() guard.Config { return lib },
|
||||
models: func() []catalog.Model { return inCatalog },
|
||||
serve: fs,
|
||||
leader: func() bool { return leader },
|
||||
gpus: func() []adminapi.ClassUse { return []adminapi.ClassUse{{Class: "6000-24g", Total: 1, Used: 1}} },
|
||||
log: slog.New(slog.DiscardHandler),
|
||||
now: func() time.Time { return now },
|
||||
}
|
||||
|
||||
k.check()
|
||||
// The GPU judge alone: the CPU one is there, and one llmgw does not know
|
||||
// is not guessed at.
|
||||
if len(fs.started) != 1 || fs.started[0].Name != "qwen3guard-4b" || fs.started[0].GPUClass != "6000-24g" ||
|
||||
fs.started[0].Class != catalog.ClassStandard || fs.started[0].GuardFormat != "qwen3guard" || fs.started[0].CPU {
|
||||
t.Fatalf("started %+v", fs.started)
|
||||
}
|
||||
|
||||
// Being fetched: not asked for again.
|
||||
fs.running = []modelserve.Job{{Model: "qwen3guard-4b", State: modelserve.Running}}
|
||||
k.check()
|
||||
if len(fs.started) != 1 {
|
||||
t.Errorf("asked for again while it is being fetched: %d", len(fs.started))
|
||||
}
|
||||
// Failed: asked for again only after a while.
|
||||
fs.running = nil
|
||||
now = now.Add(time.Minute)
|
||||
k.check()
|
||||
if len(fs.started) != 1 {
|
||||
t.Error("asked for again at once after it failed")
|
||||
}
|
||||
now = now.Add(judgeRetry)
|
||||
k.check()
|
||||
if len(fs.started) != 2 {
|
||||
t.Error("never asked for again")
|
||||
}
|
||||
|
||||
// Only the leader fetches, and nothing while the guardrails are off.
|
||||
leader = false
|
||||
now = now.Add(judgeRetry)
|
||||
k.check()
|
||||
leader = true
|
||||
off := false
|
||||
lib.Enabled = &off
|
||||
k.check()
|
||||
if len(fs.started) != 2 {
|
||||
t.Errorf("a standby or the guardrails off fetched: %d", len(fs.started))
|
||||
}
|
||||
|
||||
// A CPU judge on a cluster without GPUs is served on CPUs, two of it.
|
||||
lib = guard.Config{Policies: map[string]guard.Policy{"b": {Judge: &guard.JudgePolicy{Model: "qwen3guard-0.6b", Input: "block"}}}}
|
||||
inCatalog = nil
|
||||
k.gpus = func() []adminapi.ClassUse { return nil }
|
||||
k.check()
|
||||
if last := fs.started[len(fs.started)-1]; last.Name != "qwen3guard-0.6b" || !last.CPU || last.Replicas != 2 {
|
||||
t.Errorf("CPU judge: %+v", last)
|
||||
}
|
||||
fs.err = errors.New("the mirror is down")
|
||||
lib = guard.Config{Policies: map[string]guard.Policy{"b": {Judge: &guard.JudgePolicy{Model: "qwen3guard-8b", Input: "block"}}}}
|
||||
k.check() // no GPU, so not started, and no panic on the error path either
|
||||
}
|
||||
|
||||
// A policy moved from Llama Guard's codes to Qwen3Guard would count nothing:
|
||||
// its categories become the judge's own.
|
||||
func TestCategoriesFollowTheJudges(t *testing.T) {
|
||||
lg := []string{"S1", "S2", "S10"}
|
||||
if got := categoriesFor([]string{"qwen3guard-0.6b"}, lg); !slices.Equal(got, baselineCategories["qwen3guard"]) {
|
||||
t.Errorf("Llama Guard's codes for Qwen3Guard: %v", got)
|
||||
}
|
||||
if got := categoriesFor([]string{"qwen3guard-4b", "qwen3guard-0.6b"}, []string{"Violent", "Jailbreak"}); got != nil {
|
||||
t.Errorf("its own names were replaced: %v", got)
|
||||
}
|
||||
if got := categoriesFor([]string{"llama-guard-3-1b"}, lg); got != nil {
|
||||
t.Errorf("Llama Guard's own codes were replaced: %v", got)
|
||||
}
|
||||
// Both kinds asked in turn: each must have something to count.
|
||||
got := categoriesFor([]string{"qwen3guard-4b", "llama-guard-3-1b"}, []string{"Violent"})
|
||||
if !slices.Contains(got, "Violent") || !slices.Contains(got, "S1") {
|
||||
t.Errorf("mixed: %v", got)
|
||||
}
|
||||
if got := categoriesFor([]string{"our-own-judge"}, lg); got != nil {
|
||||
t.Errorf("a judge llmgw does not know: %v", got)
|
||||
}
|
||||
if got := categoriesFor([]string{"qwen3guard-0.6b"}, nil); got != nil {
|
||||
t.Errorf("no categories (every one counts) became %v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
)
|
||||
|
||||
// knownJudge is a judge model llmgw knows where to fetch and how to serve:
|
||||
// what a policy may name without anyone saying where its weights are.
|
||||
type knownJudge struct {
|
||||
Name, Repo, File, Format string
|
||||
// GPU is a judge too slow on CPUs to wait for: served on one card.
|
||||
GPU bool
|
||||
About string
|
||||
}
|
||||
|
||||
// knownJudges, the default first. Measured on the reference cluster against
|
||||
// ordinary coding requests ("continue", "kill the process on port 8080") and
|
||||
// harmful ones: Llama Guard 3 1B scored both alike, about one in two
|
||||
// ordinary ones refused; Qwen3Guard-Gen 0.6B kept them apart, 0.15 against
|
||||
// 1.00 at the worst, at about two seconds a check on two CPU threads.
|
||||
var knownJudges = []knownJudge{
|
||||
{Name: "qwen3guard-0.6b", Repo: "mradermacher/Qwen3Guard-Gen-0.6B-GGUF", File: "Qwen3Guard-Gen-0.6B.Q8_0.gguf",
|
||||
Format: "qwen3guard", About: "Qwen3Guard-Gen 0.6B on CPUs, 0.8 GB: the default"},
|
||||
{Name: "qwen3guard-4b", Repo: "mradermacher/Qwen3Guard-Gen-4B-GGUF", File: "Qwen3Guard-Gen-4B.Q4_K_M.gguf",
|
||||
Format: "qwen3guard", GPU: true, About: "Qwen3Guard-Gen 4B on a GPU, 2.5 GB"},
|
||||
{Name: "qwen3guard-8b", Repo: "mradermacher/Qwen3Guard-Gen-8B-GGUF", File: "Qwen3Guard-Gen-8B.Q4_K_M.gguf",
|
||||
Format: "qwen3guard", GPU: true, About: "Qwen3Guard-Gen 8B on a GPU, 4.7 GB"},
|
||||
{Name: "llama-guard-3-1b", Repo: "QuantFactory/Llama-Guard-3-1B-GGUF", File: "Llama-Guard-3-1B.Q4_K_M.gguf",
|
||||
Format: "llamaguard3", About: "Llama Guard 3 1B on CPUs, 0.9 GB: refuses ordinary coding requests"},
|
||||
{Name: "llama-guard-3-8b", Repo: "QuantFactory/Llama-Guard-3-8B-GGUF", File: "Llama-Guard-3-8B.Q4_K_M.gguf",
|
||||
Format: "llamaguard3", GPU: true, About: "Llama Guard 3 8B on a GPU, 4.6 GB"},
|
||||
}
|
||||
|
||||
const (
|
||||
defaultJudge = "qwen3guard-0.6b"
|
||||
defaultGPUJudge = "qwen3guard-4b"
|
||||
)
|
||||
|
||||
// baselineCategories are what the baseline policy counts of each format's
|
||||
// verdicts: not PII, which is the patterns', nor what refuses ordinary work
|
||||
// — copyright, unethical acts, politically sensitive topics; Llama Guard's
|
||||
// defamation, specialized advice, intellectual property and elections.
|
||||
var baselineCategories = map[string][]string{
|
||||
"qwen3guard": {"Violent", "Non-violent Illegal Acts", "Sexual Content or Sexual Acts", "Suicide & Self-Harm", "Jailbreak"},
|
||||
"llamaguard3": {"S1", "S2", "S3", "S4", "S7", "S9", "S10", "S11", "S12"},
|
||||
}
|
||||
|
||||
// reported is whether a judge of that format ever names category c: its own
|
||||
// names, or for Llama Guard its codes.
|
||||
func reported(format, c string) bool {
|
||||
if slices.ContainsFunc(guard.DefaultCategories(format), func(n string) bool { return strings.EqualFold(n, c) }) {
|
||||
return true
|
||||
}
|
||||
var n int
|
||||
_, err := fmt.Sscanf(c, "S%d", &n)
|
||||
return format == "llamaguard3" && err == nil && n >= 1 && n <= 14
|
||||
}
|
||||
|
||||
// categoriesFor is what a policy asking these judges should count, when the
|
||||
// categories it has are ones none of them names: a policy moved from Llama
|
||||
// Guard's S1 to Qwen3Guard, which says Violent, would otherwise count
|
||||
// nothing. Nil keeps what it has; so does a judge llmgw does not know.
|
||||
func categoriesFor(judges, have []string) []string {
|
||||
var formats []string
|
||||
for _, j := range judges {
|
||||
k, ok := knownJudgeNamed(j)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if !slices.Contains(formats, k.Format) {
|
||||
formats = append(formats, k.Format)
|
||||
}
|
||||
}
|
||||
if len(have) == 0 {
|
||||
return nil
|
||||
}
|
||||
stray := false
|
||||
for _, f := range formats {
|
||||
for _, c := range have {
|
||||
if !reported(f, c) && !slices.ContainsFunc(formats, func(o string) bool { return reported(o, c) }) {
|
||||
stray = true
|
||||
}
|
||||
}
|
||||
// Every judge must have something it counts.
|
||||
if !slices.ContainsFunc(have, func(c string) bool { return reported(f, c) }) {
|
||||
stray = true
|
||||
}
|
||||
}
|
||||
if !stray {
|
||||
return nil
|
||||
}
|
||||
var out []string
|
||||
for _, f := range formats {
|
||||
out = append(out, baselineCategories[f]...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func knownJudgeNamed(name string) (knownJudge, bool) {
|
||||
i := slices.IndexFunc(knownJudges, func(k knownJudge) bool { return k.Name == name })
|
||||
if i < 0 {
|
||||
return knownJudge{}, false
|
||||
}
|
||||
return knownJudges[i], true
|
||||
}
|
||||
|
||||
func knownJudgeNames() string {
|
||||
var n []string
|
||||
for _, k := range knownJudges {
|
||||
n = append(n, k.Name)
|
||||
}
|
||||
return strings.Join(n, ", ")
|
||||
}
|
||||
|
||||
// judgeServeOptions are how a judge is served beyond which weights.
|
||||
type judgeServeOptions struct {
|
||||
Replicas, Cores int
|
||||
// GPUClass is the cards a GPU judge runs on; empty lets the plan choose.
|
||||
GPUClass string
|
||||
}
|
||||
|
||||
// serveRequest is the serve job for a judge. A judge on CPUs is critical,
|
||||
// as every judged request waits on it; one on a GPU is standard, so it takes
|
||||
// a card only from what may give one up, and the judge on CPUs after it in
|
||||
// the policies stands in while it has none.
|
||||
func (k knownJudge) serveRequest(o judgeServeOptions) modelserve.Request {
|
||||
req := modelserve.Request{
|
||||
Repo: k.Repo, File: k.File, Name: k.Name, GuardFormat: k.Format,
|
||||
Description: "The guardrails' judge: reads prompts and answers for unsafe content. Not offered to clients.",
|
||||
Replicas: o.Replicas, CtxPool: 8192, ContextLength: 4096,
|
||||
}
|
||||
if k.GPU {
|
||||
req.GPUClass, req.GPUsPerNode, req.Class = o.GPUClass, 1, catalog.ClassStandard
|
||||
if req.Replicas == 0 {
|
||||
req.Replicas = 1
|
||||
}
|
||||
return req
|
||||
}
|
||||
req.CPU, req.CPUCores = true, o.Cores
|
||||
if req.Replicas == 0 {
|
||||
req.Replicas = 2
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
func (k knownJudge) String() string {
|
||||
where := "CPUs"
|
||||
if k.GPU {
|
||||
where = "a GPU"
|
||||
}
|
||||
return fmt.Sprintf("%s (%s, %s/%s, on %s)", k.Name, k.Format, k.Repo, k.File, where)
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
func newKBCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
c := &cobra.Command{
|
||||
Use: "kb",
|
||||
Short: "Knowledge bases: collections of documents requests are answered from",
|
||||
Long: "A collection belongs to an organization, team, project or person, and is\n" +
|
||||
"embedded by one model of the catalog. People put documents into it from\n" +
|
||||
"their own machines with llmgw-kb (its own download, for Windows, macOS and\n" +
|
||||
"Linux), which sends text and vectors, never the files.\n\n" +
|
||||
"Who may read or add to a collection: its owner and everything below it,\n" +
|
||||
"and whoever it is granted to (grant), never outside its organization.\n" +
|
||||
"Targets are written org:Acme, team:Acme/hr, project:Acme/hr/docs (or with\n" +
|
||||
"IDs), key:KEY-ID or user:USER-ID. A collection is named by its ID, its name\n" +
|
||||
"when that is unique, or ORG/NAME.\n\n" +
|
||||
"Needs the PostgreSQL store with pgvector; `llmgw ctl kb list` says why not.",
|
||||
}
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "list",
|
||||
Short: "Every collection: owner, model, documents and grants",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var v adminapi.KBView
|
||||
if err := client().get(cmd.Context(), "/api/v1/kb", &v); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
if !v.Available {
|
||||
fmt.Fprintln(out, "Knowledge bases are off: "+v.Reason)
|
||||
return nil
|
||||
}
|
||||
if len(v.Collections) == 0 {
|
||||
fmt.Fprintln(out, "No collections yet (llmgw ctl kb create).")
|
||||
return nil
|
||||
}
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "COLLECTION\tOWNER\tMODEL\tDOCUMENTS\tCHUNKS\tGRANTED TO")
|
||||
for _, c := range v.Collections {
|
||||
var grants []string
|
||||
for _, g := range c.Grants {
|
||||
grants = append(grants, or(g.Name, g.Target)+" ("+g.Perm+")")
|
||||
}
|
||||
fmt.Fprintf(w, "%s/%s\t%s\t%s (%d)\t%d\t%d\t%s\n", or(c.OrgName, c.Org), c.Name, or(c.OwnerName, c.Owner),
|
||||
c.EmbedModel, c.Dims, c.Documents, c.Chunks, strings.Join(grants, ", "))
|
||||
}
|
||||
if err := w.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(v.Defaults) > 0 {
|
||||
fmt.Fprintln(out, "\nAnswered from by default, when a request names no collection:")
|
||||
w = tab()
|
||||
for _, d := range v.Defaults {
|
||||
fmt.Fprintf(w, " %s\t%s\n", or(d.Name, d.Target), strings.Join(d.Collections, ", "))
|
||||
}
|
||||
return w.Flush()
|
||||
}
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
var org string
|
||||
var dims int
|
||||
create := &cobra.Command{
|
||||
Use: "create OWNER NAME MODEL",
|
||||
Short: "Make a collection owned by an organization, team, project or person",
|
||||
Long: "MODEL is a catalog model with capabilities.embeddings. Its dimension is asked\n" +
|
||||
"of it, so it has to be able to serve now; --dims gives it instead.",
|
||||
Example: " llmgw ctl kb create team:Acme/hr policies bge-m3\n" +
|
||||
" llmgw ctl kb create user:USER-ID notes bge-m3 --org Acme",
|
||||
Args: cobra.ExactArgs(3),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
body, _ := json.Marshal(adminapi.KBCreateRequest{Owner: args[0], Name: args[1], Model: args[2], Org: org, Dims: dims})
|
||||
var v adminapi.KBCollectionView
|
||||
if err := client().do(cmd.Context(), "POST", "/api/v1/kb/collections", json.RawMessage(body), &v); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "collection %s (%s), embedded by %s in %d dimensions\n", v.Name, v.ID, v.EmbedModel, v.Dims)
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "people with access fill it with: llmgw-kb sync %s FOLDER\n", v.Name)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
create.Flags().StringVar(&org, "org", "", "the organization of a person's collection")
|
||||
create.Flags().IntVar(&dims, "dims", 0, "the model's embedding dimension, instead of asking it")
|
||||
c.AddCommand(create)
|
||||
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "delete COLLECTION",
|
||||
Short: "Remove a collection, its documents and its grants",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if err := client().do(cmd.Context(), "DELETE", "/api/v1/kb/collections/"+url.PathEscape(args[0]), nil, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "collection %s removed\n", args[0])
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
for _, verb := range []struct {
|
||||
use, short, method string
|
||||
perm bool
|
||||
}{
|
||||
{"grant", "Let a scope and everything below it read a collection, or with --write add to it", "POST", true},
|
||||
{"revoke", "Take a grant away", "DELETE", false},
|
||||
} {
|
||||
var write bool
|
||||
cmd := &cobra.Command{
|
||||
Use: verb.use + " COLLECTION TARGET",
|
||||
Short: verb.short,
|
||||
Example: " llmgw ctl kb " + verb.use + " Acme/policies project:Acme/sales/crm",
|
||||
Args: cobra.ExactArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
req := adminapi.KBGrantRequest{Collection: args[0], Target: args[1], Perm: "read"}
|
||||
if write {
|
||||
req.Perm = "write"
|
||||
}
|
||||
body, _ := json.Marshal(req)
|
||||
if err := client().do(cmd.Context(), verb.method, "/api/v1/kb/grants", json.RawMessage(body), nil); err != nil {
|
||||
return err
|
||||
}
|
||||
if verb.perm {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s may %s %s\n", args[1], req.Perm, args[0])
|
||||
} else {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s no longer granted %s\n", args[1], args[0])
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if verb.perm {
|
||||
cmd.Flags().BoolVar(&write, "write", false, "add and remove documents too")
|
||||
}
|
||||
c.AddCommand(cmd)
|
||||
}
|
||||
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "docs COLLECTION",
|
||||
Short: "A collection's documents",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
var docs []adminapi.KBDocumentView
|
||||
if err := client().get(cmd.Context(), "/api/v1/kb/collections/"+url.PathEscape(args[0])+"/documents", &docs); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(docs)
|
||||
}
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "DOCUMENT\tCHUNKS\tSIZE\tADDED\tBY")
|
||||
for _, d := range docs {
|
||||
fmt.Fprintf(w, "%s\t%d\t%s\t%s\t%s\n", d.Name, d.Chunks, humanBytes(d.Bytes), d.AddedAt.Local().Format("2006-01-02 15:04"), d.AddedBy)
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
})
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "rm COLLECTION DOCUMENT",
|
||||
Short: "Remove one document from a collection",
|
||||
Args: cobra.ExactArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
path := "/api/v1/kb/collections/" + url.PathEscape(args[0]) + "/documents?name=" + url.QueryEscape(args[1])
|
||||
if err := client().do(cmd.Context(), "DELETE", path, nil, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s removed from %s\n", args[1], args[0])
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
var clear bool
|
||||
def := &cobra.Command{
|
||||
Use: "default TARGET [COLLECTION...]",
|
||||
Short: "Answer a part of the tree's chat requests from collections, when they name none",
|
||||
Long: "Requests from TARGET and everything below it are then answered from these\n" +
|
||||
"collections without naming them (X-Llmgw-Knowledge). The narrowest default\n" +
|
||||
"wins: a key's over its project's, a project's over its team's. A request's own\n" +
|
||||
"header overrides it, and X-Llmgw-Knowledge: none turns it off. A collection a\n" +
|
||||
"caller may not read is left out for them, and a search that fails answers\n" +
|
||||
"without passages rather than refusing. In force within 30 seconds.",
|
||||
Example: " llmgw ctl kb default team:Acme/hr Acme/policies Acme/forms\n" +
|
||||
" llmgw ctl kb default team:Acme/hr --clear",
|
||||
Args: cobra.MinimumNArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if clear != (len(args) == 1) {
|
||||
return fmt.Errorf("name collections, or --clear")
|
||||
}
|
||||
body, _ := json.Marshal(adminapi.KBDefaultRequest{Target: args[0], Collections: args[1:]})
|
||||
if err := client().do(cmd.Context(), "POST", "/api/v1/kb/defaults", json.RawMessage(body), nil); err != nil {
|
||||
return err
|
||||
}
|
||||
if clear {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s answers from no collection by default\n", args[0])
|
||||
} else {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s answers from %s by default\n", args[0], strings.Join(args[1:], ", "))
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
def.Flags().BoolVar(&clear, "clear", false, "no default any more")
|
||||
c.AddCommand(def)
|
||||
|
||||
var dsModel, dsOut string
|
||||
var dsPer, dsMax int
|
||||
dataset := &cobra.Command{
|
||||
Use: "dataset COLLECTION --model MODEL",
|
||||
Short: "Write training data from a collection: questions and answers, for an adapter",
|
||||
Long: "A chat model of the cluster reads each chunk of the collection and writes\n" +
|
||||
"questions a person might ask of it, with the answers it holds, one JSON object\n" +
|
||||
"a line ({\"q\", \"a\", \"doc\", \"ord\"}). Read it through and correct it before\n" +
|
||||
"training on it (llmgw ctl train): what the model got wrong, an adapter learns.\n\n" +
|
||||
"An adapter learns how answers sound — the organisation's words, format and\n" +
|
||||
"tone — far better than what they say; the facts stay the knowledge base's.",
|
||||
Example: " llmgw ctl kb dataset Acme/policies --model qwen2.5-coder-14b-instruct --out policies.jsonl",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
w := cmd.OutOrStdout()
|
||||
if dsOut != "" {
|
||||
f, err := os.Create(dsOut)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
w = f
|
||||
}
|
||||
lw := &datasetWriter{pairs: w, problems: cmd.ErrOrStderr()}
|
||||
body, _ := json.Marshal(adminapi.KBDatasetRequest{Model: dsModel, PerChunk: dsPer, MaxChunks: dsMax})
|
||||
err := client().patient(6*time.Hour).do(cmd.Context(), "POST",
|
||||
"/api/v1/kb/collections/"+url.PathEscape(args[0])+"/dataset", json.RawMessage(body), lw)
|
||||
lw.flush()
|
||||
fmt.Fprintf(cmd.ErrOrStderr(), "%d pairs from %d chunks; %d chunks gave none\n", lw.n, len(lw.chunks), lw.failed)
|
||||
return err
|
||||
},
|
||||
}
|
||||
dataset.Flags().StringVar(&dsModel, "model", "", "the chat model that writes the questions and answers")
|
||||
dataset.Flags().StringVar(&dsOut, "out", "", "the file to write (default: standard output)")
|
||||
dataset.Flags().IntVar(&dsPer, "per-chunk", 8, "pairs from each chunk")
|
||||
dataset.Flags().IntVar(&dsMax, "max-chunks", 500, "the most chunks read")
|
||||
_ = dataset.MarkFlagRequired("model")
|
||||
c.AddCommand(dataset)
|
||||
|
||||
var k int
|
||||
search := &cobra.Command{
|
||||
Use: "search COLLECTION QUERY...",
|
||||
Short: "Try a search of a collection, as an administrator (no tenancy applies)",
|
||||
Args: cobra.MinimumNArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
body, _ := json.Marshal(adminapi.KBSearchRequest{Collection: args[0], Query: strings.Join(args[1:], " "), K: k})
|
||||
var hits []adminapi.KBHitView
|
||||
if err := client().do(cmd.Context(), "POST", "/api/v1/kb/search", json.RawMessage(body), &hits); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(hits)
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
if len(hits) == 0 {
|
||||
fmt.Fprintln(out, "Nothing found.")
|
||||
}
|
||||
for i, h := range hits {
|
||||
fmt.Fprintf(out, "%d. %s #%d\n%s\n\n", i+1, h.Document, h.Ord, indent(strings.TrimSpace(h.Text), " "))
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
search.Flags().IntVarP(&k, "k", "k", 5, "how many chunks")
|
||||
c.AddCommand(search)
|
||||
return c
|
||||
}
|
||||
|
||||
// datasetWriter splits the dataset stream: pairs to the file, problems to
|
||||
// the terminal, counted as they pass.
|
||||
type datasetWriter struct {
|
||||
pairs, problems io.Writer
|
||||
buf []byte
|
||||
n, failed int
|
||||
chunks map[string]bool
|
||||
}
|
||||
|
||||
func (d *datasetWriter) Write(p []byte) (int, error) {
|
||||
d.buf = append(d.buf, p...)
|
||||
for {
|
||||
i := bytes.IndexByte(d.buf, '\n')
|
||||
if i < 0 {
|
||||
return len(p), nil
|
||||
}
|
||||
d.line(d.buf[:i])
|
||||
d.buf = d.buf[i+1:]
|
||||
}
|
||||
}
|
||||
|
||||
func (d *datasetWriter) flush() {
|
||||
if len(bytes.TrimSpace(d.buf)) > 0 {
|
||||
d.line(d.buf)
|
||||
}
|
||||
d.buf = nil
|
||||
}
|
||||
|
||||
func (d *datasetWriter) line(b []byte) {
|
||||
var l adminapi.KBDatasetLine
|
||||
if json.Unmarshal(b, &l) != nil {
|
||||
return
|
||||
}
|
||||
if d.chunks == nil {
|
||||
d.chunks = map[string]bool{}
|
||||
}
|
||||
d.chunks[fmt.Sprintf("%s#%d", l.Doc, l.Ord)] = true
|
||||
if l.Error != "" {
|
||||
d.failed++
|
||||
fmt.Fprintf(d.problems, " %s #%d: %s\n", l.Doc, l.Ord, l.Error)
|
||||
return
|
||||
}
|
||||
d.n++
|
||||
out, _ := json.Marshal(l)
|
||||
fmt.Fprintln(d.pairs, string(out))
|
||||
}
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"net"
|
||||
"net/netip"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -65,6 +66,14 @@ Safe to run again: it changes only what differs.`,
|
||||
}
|
||||
if !noFirewall {
|
||||
mirrorFirewall(cmd.Context(), cmd.OutOrStdout(), controlPlanes)
|
||||
// And after, as the cluster's settings say: the control plane
|
||||
// tells the mirror, which writes it beside its models.
|
||||
state := filepath.Join(*dir, mirror.FirewallFile)
|
||||
if err := installFirewallSync(cmd.Context(), firewall.Local{}, state); err != nil {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), " !! firewall: %v; the cluster's settings will not change it here\n", err)
|
||||
} else {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), " ok firewall: follows the cluster's settings ("+firewallSyncPathUnit+")")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
|
||||
@@ -63,6 +63,10 @@ func Root() *cobra.Command {
|
||||
newWipeCmd(),
|
||||
newDocsCmd(),
|
||||
)
|
||||
// status, start, stop and restart: the whole system; service: this
|
||||
// machine's services.
|
||||
root.AddCommand(newSystemCmds()...)
|
||||
root.AddCommand(newServiceCmd())
|
||||
return root
|
||||
}
|
||||
|
||||
|
||||
+125
-1
@@ -5,6 +5,7 @@ import (
|
||||
"crypto/rand"
|
||||
"crypto/tls"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
@@ -48,6 +49,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fleetmetrics"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/train"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
@@ -186,6 +188,7 @@ func newServerCmd() *cobra.Command {
|
||||
|
||||
rt.Start(ctx)
|
||||
defer rt.Stop()
|
||||
rt.SetStopped(cfg.Serving.Stopped)
|
||||
|
||||
// Which cluster this is, to federated peers: the id in the hop
|
||||
// header that stops a request going round two clusters that
|
||||
@@ -230,6 +233,9 @@ func newServerCmd() *cobra.Command {
|
||||
var agentHandler http.Handler
|
||||
var agentPath string
|
||||
var admin *control.Controller
|
||||
// What the settings say of the firewall, for the mirror and the
|
||||
// store's members; nil without a control plane.
|
||||
var fwPub *firewallPublisher
|
||||
updateMetrics := &modelserve.UpdateCollector{}
|
||||
// The store's cluster, when it is one; read by the metrics and
|
||||
// by `ctl status`.
|
||||
@@ -252,6 +258,8 @@ func newServerCmd() *cobra.Command {
|
||||
}
|
||||
defer ctl.Close()
|
||||
agentSvc = svc
|
||||
// Stopped by the operators (llmgw stop): nothing placed.
|
||||
ctl.Reconciler().SetStopped(cfg.Serving.Stopped)
|
||||
// Whether nodes keep llmgw's firewall, as the settings say now.
|
||||
svc.SetFirewall(func() *llmgwv1.Firewall {
|
||||
c := current.Load()
|
||||
@@ -259,9 +267,31 @@ func newServerCmd() *cobra.Command {
|
||||
return &llmgwv1.Firewall{Enabled: c.Firewall.Enabled, Networks: c.FirewallNetworks,
|
||||
EnginePortsFrom: int32(from), EnginePortsTo: int32(to)} //nolint:gosec // ports
|
||||
})
|
||||
// What each node's agent runs on, as the settings say now.
|
||||
svc.SetAgentSettings(func(node string) string {
|
||||
b, err := config.Marshal(current.Load().Agents.For(node))
|
||||
if err != nil {
|
||||
log.Warn("could not write a node's agent settings", "node", node, "err", err)
|
||||
return ""
|
||||
}
|
||||
return string(b)
|
||||
})
|
||||
ctl.Reconciler().SetNotifier(notifier)
|
||||
admin = ctl
|
||||
agentPath, agentHandler = svc.Handler()
|
||||
fwPub = newFirewallPublisher(func() config.Server { return *current.Load() }, rt.Store(),
|
||||
func() []string {
|
||||
var out []string
|
||||
for _, n := range ctl.Inventory().Nodes() {
|
||||
out = append(out, n.MgmtAddr)
|
||||
}
|
||||
return out
|
||||
},
|
||||
func() bool { return ctl.Leadership() == nil || ctl.Leadership().IsLeader() }, log)
|
||||
go fwPub.run(ctx, 5*time.Minute)
|
||||
if l := ctl.Leadership(); l != nil {
|
||||
l.OnLead(fwPub.Kick)
|
||||
}
|
||||
|
||||
// Before AdminHandler, which is where the registry is built —
|
||||
// registering after it silently publishes nothing, which is
|
||||
@@ -361,16 +391,32 @@ func newServerCmd() *cobra.Command {
|
||||
evals = newEvalService(cfg, rt, admin, log)
|
||||
api.SetEval(evals)
|
||||
if st := rt.Store(); st != nil && gate != nil {
|
||||
api.SetGuard(&adminapi.GuardAdmin{Library: gate.Library, Store: st, Invalidate: gate.Invalidate})
|
||||
api.SetGuard(&adminapi.GuardAdmin{Library: gate.Library, Store: st, Invalidate: gate.Invalidate,
|
||||
JudgeErrors: gate.JudgeErrors, Ask: gate.AskJudge})
|
||||
}
|
||||
if st := rt.Store(); st != nil {
|
||||
// Removing a node revokes its certificate.
|
||||
api.SetRevoker(st.RevokeNode)
|
||||
if agentSvc != nil {
|
||||
api.SetAgentRestarter(agentSvc.RestartAgents)
|
||||
}
|
||||
api.SetAudit(st)
|
||||
// Knowledge bases, with the cluster's own models to embed a
|
||||
// probe (a collection's dimension) and an operator's search.
|
||||
api.SetKB(&adminapi.KBAdmin{Store: st, Embed: embedder(rt.Replay)})
|
||||
// LoRA adapters trained on the nodes.
|
||||
if agentSvc != nil {
|
||||
trainer := train.New(st, agentSvc, log)
|
||||
agentSvc.SetTrain(trainer)
|
||||
go trainer.Run(ctx)
|
||||
api.SetTrain(&adminapi.TrainAdmin{Store: st, Service: trainer})
|
||||
}
|
||||
// The cluster's settings, each change checked against this
|
||||
// machine's own keys before it is kept.
|
||||
api.SetSettings(&adminapi.Settings{Store: st, LocalKeys: config.LocalKeys(),
|
||||
Check: func(body []byte) error { return checkClusterSettings(configPath, body) }})
|
||||
// And the same document on the portal, for administrators.
|
||||
rt.SetSettingsCheck(func(body []byte) error { return checkClusterSettings(configPath, body) }, config.LocalKeys())
|
||||
}
|
||||
// Catalog changes need somewhere to persist, so this is wired
|
||||
// only when there is a store. Without one, the endpoints say so
|
||||
@@ -427,6 +473,7 @@ func newServerCmd() *cobra.Command {
|
||||
api.SetUsage(st)
|
||||
// Capacity quotas, set and read over this API.
|
||||
api.SetQuotas(st)
|
||||
api.SetTenancy(st)
|
||||
}
|
||||
api.Register(adminMux)
|
||||
|
||||
@@ -488,6 +535,16 @@ func newServerCmd() *cobra.Command {
|
||||
return fmt.Errorf("the model mirror: %w", err)
|
||||
}
|
||||
rt.SetMirror(mc)
|
||||
if fwPub != nil {
|
||||
fwPub.setMirror(mc)
|
||||
}
|
||||
// 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 {
|
||||
adminAPI.SetTrainMirror(up)
|
||||
}
|
||||
}
|
||||
// Browsing the registry is what turns "paste a resolve URL"
|
||||
// into "find the model". It needs a way out to the internet
|
||||
// from this machine, which the mirror's own fetching does not.
|
||||
@@ -511,6 +568,21 @@ func newServerCmd() *cobra.Command {
|
||||
})
|
||||
go serve.WatchCanaries(ctx, 30*time.Second)
|
||||
adminAPI.SetServe(serve)
|
||||
if gate != nil {
|
||||
// A judge a policy names and the catalog lacks is
|
||||
// fetched and served, on the leader.
|
||||
keeper := &judgeKeeper{
|
||||
library: gate.Library,
|
||||
models: func() []catalog.Model { return admin.Catalog().Models() },
|
||||
serve: serve,
|
||||
leader: func() bool {
|
||||
return admin.Leadership() == nil || admin.Leadership().IsLeader()
|
||||
},
|
||||
gpus: func() []adminapi.ClassUse { return adminAPI.Summary().GPUs },
|
||||
log: log,
|
||||
}
|
||||
go keeper.run(ctx, 30*time.Second)
|
||||
}
|
||||
updateMetrics.Set(serve)
|
||||
go serve.WatchUpdates(ctx, cfg.Mirror.UpdateCheck)
|
||||
}
|
||||
@@ -675,6 +747,14 @@ func newServerCmd() *cobra.Command {
|
||||
return false
|
||||
}
|
||||
current.Store(&next)
|
||||
if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "serving") }) {
|
||||
rt.SetStopped(next.Serving.Stopped)
|
||||
if admin != nil {
|
||||
admin.Reconciler().SetStopped(next.Serving.Stopped)
|
||||
}
|
||||
log.Warn("the service is " + map[bool]string{true: "stopped", false: "started"}[next.Serving.Stopped] +
|
||||
" by the cluster's settings")
|
||||
}
|
||||
if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "firewall") }) {
|
||||
if err := writeFirewallState(next); err != nil {
|
||||
log.Warn("could not record the firewall's settings for this machine", "err", err)
|
||||
@@ -682,6 +762,13 @@ func newServerCmd() *cobra.Command {
|
||||
if agentSvc != nil {
|
||||
agentSvc.BroadcastFirewall()
|
||||
}
|
||||
if fwPub != nil {
|
||||
fwPub.Kick()
|
||||
}
|
||||
}
|
||||
if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "agents") }) && agentSvc != nil {
|
||||
// A moment apart per node, so off this goroutine.
|
||||
go agentSvc.BroadcastAgentSettings()
|
||||
}
|
||||
if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "guard") }) {
|
||||
if gate == nil {
|
||||
@@ -1153,3 +1240,40 @@ func federationID(ctx context.Context, st *storepkg.DB) (string, error) {
|
||||
// more of a secret than needed.
|
||||
return secret[:16], nil
|
||||
}
|
||||
|
||||
// embedder has a model embed one text through the gateway itself, the way a
|
||||
// kept request is replayed: for the administrator's knowledge-base
|
||||
// endpoints, not for anybody's traffic.
|
||||
func embedder(replay func(ctx context.Context, path, model string, body []byte) (int, []byte, time.Duration, error)) func(context.Context, string, string) ([]float32, error) {
|
||||
return func(ctx context.Context, model, text string) ([]float32, error) {
|
||||
body, _ := json.Marshal(map[string]any{"model": model, "input": []string{text}})
|
||||
status, out, _, err := replay(ctx, "/v1/embeddings", model, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status != http.StatusOK {
|
||||
msg := strings.TrimSpace(string(out))
|
||||
var e struct {
|
||||
Error struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if json.Unmarshal(out, &e) == nil && e.Error.Message != "" {
|
||||
msg = e.Error.Message
|
||||
}
|
||||
return nil, fmt.Errorf("%d: %s", status, msg)
|
||||
}
|
||||
var res struct {
|
||||
Data []struct {
|
||||
Embedding []float32 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(out, &res); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(res.Data) != 1 || len(res.Data[0].Embedding) == 0 {
|
||||
return nil, fmt.Errorf("%s answered no embedding", model)
|
||||
}
|
||||
return res.Data[0].Embedding, nil
|
||||
}
|
||||
}
|
||||
@@ -160,7 +160,7 @@ func (sa *serverAdd) firewall(ctx context.Context) error {
|
||||
return err
|
||||
}
|
||||
// What applies the cluster's firewall setting there when it changes.
|
||||
if err := installFirewallSync(ctx, sa.remote); err != nil {
|
||||
if err := installFirewallSync(ctx, sa.remote, firewallStateFile); err != nil {
|
||||
fmt.Fprintf(sa.out, " !! firewall: %v\n", err)
|
||||
}
|
||||
if !cfg.Firewall.Enabled {
|
||||
|
||||
@@ -293,7 +293,7 @@ func (in *serverInit) configJudgesWith(model string) bool {
|
||||
return false
|
||||
}
|
||||
for _, p := range cfg.Guard.Policies {
|
||||
if p.Judge != nil && p.Judge.Model == model {
|
||||
if slices.Contains(p.Judge.Judges(), model) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
@@ -308,18 +308,21 @@ func guardSection(judge bool) string {
|
||||
# judge: {model: ` + guardModelName + `, input: block, output: block, ...}
|
||||
`
|
||||
if judge {
|
||||
j = ` # A judge model reads what patterns cannot: violence, weapons, abuse,
|
||||
# self-harm. Two replicas run on CPUs with AVX2, about half a second
|
||||
# for a short text; an answer is read every 600 characters as it
|
||||
# streams. With fail: closed, a request the judge cannot read within
|
||||
j = ` # A judge model reads what patterns cannot: violence, crime, sexual
|
||||
# content, self-harm, jailbreaks. Two replicas run on CPUs, about two
|
||||
# seconds for a short text; an answer is read every 600 characters as
|
||||
# it streams. With fail: closed, a request no judge can read within
|
||||
# latency_budget is refused (503); open serves it, flagged, unjudged.
|
||||
# Defamation (S5), specialized advice (S6), intellectual property (S8)
|
||||
# and elections (S13) are left out: they refuse too much ordinary work.
|
||||
# PII is the patterns' (above), and copyright, unethical acts and
|
||||
# politically sensitive topics are left out: they refuse too much
|
||||
# ordinary work. Controversial counts as safe; controversial: unsafe
|
||||
# makes it count. With a GPU: llmgw ctl guard judge enable --gpu, then
|
||||
# llmgw ctl guard judge use ` + defaultGPUJudge + ` ` + guardModelName + `.
|
||||
judge:
|
||||
model: ` + guardModelName + `
|
||||
models: [` + guardModelName + `]
|
||||
input: block
|
||||
output: block
|
||||
categories: [S1, S2, S3, S4, S7, S9, S10, S11, S12]
|
||||
categories: [Violent, Non-violent Illegal Acts, Sexual Content or Sexual Acts, Suicide & Self-Harm, Jailbreak]
|
||||
latency_budget: 5s
|
||||
`
|
||||
}
|
||||
|
||||
@@ -41,11 +41,13 @@ const (
|
||||
starterSource = "hf:Qwen/Qwen2.5-0.5B-Instruct-GGUF/qwen2.5-0.5b-instruct-q4_k_m.gguf"
|
||||
starterVersion = "starter"
|
||||
|
||||
// The guardrails' judge: Meta's Llama Guard 3 1B, which reads a prompt
|
||||
// or an answer and says whether it is unsafe and why. Q4_K_M is about
|
||||
// 900 MB and judges a short text in about half a second on a CPU.
|
||||
guardModelName = "llama-guard-3-1b"
|
||||
guardModelSource = "hf:QuantFactory/Llama-Guard-3-1B-GGUF/Llama-Guard-3-1B.Q4_K_M.gguf"
|
||||
// The guardrails' judge: Qwen3Guard-Gen 0.6B, which reads a prompt or an
|
||||
// answer and says whether it is safe, controversial or unsafe, and why.
|
||||
// Q8_0 is 0.8 GB and judges a short text in about two seconds on two CPU
|
||||
// threads. It replaced Llama Guard 3 1B, which refused about one in two
|
||||
// ordinary coding requests ("continue").
|
||||
guardModelName = defaultJudge
|
||||
guardModelSource = "hf:mradermacher/Qwen3Guard-Gen-0.6B-GGUF/Qwen3Guard-Gen-0.6B.Q8_0.gguf"
|
||||
guardModelVersion = "guard"
|
||||
|
||||
// localLabel marks the control plane's own node, so the starter model is
|
||||
@@ -443,12 +445,6 @@ func starterModel(source string, gpuClasses []string, pinned bool) catalog.Model
|
||||
}
|
||||
}
|
||||
|
||||
// llamaGuard3Categories name Llama Guard 3's S1 to S13, so a verdict is
|
||||
// recorded as "Violent Crimes" rather than "S1".
|
||||
var llamaGuard3Categories = []string{"Violent Crimes", "Non-Violent Crimes", "Sex-Related Crimes",
|
||||
"Child Sexual Exploitation", "Defamation", "Specialized Advice", "Privacy", "Intellectual Property",
|
||||
"Indiscriminate Weapons", "Hate", "Suicide & Self-Harm", "Sexual Content", "Elections"}
|
||||
|
||||
// guardModel is the judge's catalog entry: CPU only, so it takes no card from
|
||||
// the models people use, on whichever node has room.
|
||||
func guardModel(source string) catalog.Model {
|
||||
@@ -457,7 +453,7 @@ func guardModel(source string) catalog.Model {
|
||||
Description: "The guardrails' judge: reads prompts and answers for unsafe content. Not offered to clients.",
|
||||
Source: source,
|
||||
Role: catalog.RoleGuard,
|
||||
Guard: &catalog.GuardModel{Format: "llamaguard3", Categories: llamaGuard3Categories},
|
||||
Guard: &catalog.GuardModel{Format: "qwen3guard"},
|
||||
// Two: every judged request waits on it, and with fail: closed a
|
||||
// judge that is not there refuses everything.
|
||||
Scaling: catalog.Scaling{Min: 2, Max: 2},
|
||||
|
||||
@@ -7,12 +7,14 @@ import (
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
)
|
||||
|
||||
func TestTheInitAgentUnitIsTheShippedUnit(t *testing.T) {
|
||||
@@ -150,9 +152,17 @@ func TestTheInstalledConfigurationIsHardenedAndLoads(t *testing.T) {
|
||||
if (p.Judge != nil) != judge {
|
||||
t.Errorf("judge=%v: the policy's judge is %+v", judge, p.Judge)
|
||||
}
|
||||
if judge && (p.Judge.Model != guardModelName || p.Judge.Input != "block") {
|
||||
if judge && (!slices.Equal(p.Judge.Judges(), []string{guardModelName}) || p.Judge.Input != "block") {
|
||||
t.Errorf("judge %+v", p.Judge)
|
||||
}
|
||||
// A category the judge never names would never count.
|
||||
if judge {
|
||||
for _, c := range p.Judge.Categories {
|
||||
if !slices.Contains(guard.DefaultCategories("qwen3guard"), c) {
|
||||
t.Errorf("the baseline counts %q, which Qwen3Guard does not name", c)
|
||||
}
|
||||
}
|
||||
}
|
||||
if cfg.Auth.MaxKeyTTL != 90*24*time.Hour || cfg.Auth.SignupOpen || cfg.DebugPprof || cfg.AdminRemote {
|
||||
t.Errorf("not hardened: ttl %v signup %v pprof %v remote %v",
|
||||
cfg.Auth.MaxKeyTTL, cfg.Auth.SignupOpen, cfg.DebugPprof, cfg.AdminRemote)
|
||||
@@ -161,13 +171,13 @@ func TestTheInstalledConfigurationIsHardenedAndLoads(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestTheGuardModelIsAValidJudge(t *testing.T) {
|
||||
m := guardModel("http://10.0.0.5:8081/llama-guard-3-1b/guard/g.gguf")
|
||||
m := guardModel("http://10.0.0.5:8081/qwen3guard-0.6b/guard/g.gguf")
|
||||
c, err := catalog.New([]catalog.Model{m})
|
||||
if err != nil {
|
||||
t.Fatalf("not a valid catalog entry: %v", err)
|
||||
}
|
||||
got, _ := c.Resolve(guardModelName)
|
||||
if !got.IsGuard() || got.Guard.Format != "llamaguard3" || len(got.Guard.Categories) != 13 || got.Scaling.Min != 2 {
|
||||
if !got.IsGuard() || got.Guard.Format != "qwen3guard" || got.Scaling.Min != 2 {
|
||||
t.Errorf("%+v", got)
|
||||
}
|
||||
for _, p := range got.Profiles {
|
||||
|
||||
@@ -93,6 +93,11 @@ func (in *serverInit) storeFirewall(ctx context.Context, members []pgha.Node, se
|
||||
} else {
|
||||
in.kept("firewall on %s", m.Name)
|
||||
}
|
||||
// And after, as the cluster's settings say, read from the
|
||||
// member's own copy of the store.
|
||||
if err := firewall.InstallFollower(ctx, m.Remote); err != nil {
|
||||
fmt.Fprintf(in.out, " !! firewall on %s: %v; the cluster's settings will not change it\n", m.Name, err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
}
|
||||
@@ -135,5 +140,5 @@ func loopbackListen(listen string) bool {
|
||||
// ensureFirewallSync installs the units that apply the cluster's firewall
|
||||
// setting on this control plane.
|
||||
func (in *serverInit) ensureFirewallSync(ctx context.Context) error {
|
||||
return installFirewallSync(ctx, firewall.Local{})
|
||||
return installFirewallSync(ctx, firewall.Local{}, firewallStateFile)
|
||||
}
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
@@ -348,15 +349,23 @@ func (in *serverInit) setupPostgres(ctx context.Context) (string, error) {
|
||||
}
|
||||
|
||||
if installed {
|
||||
// Its defaults assume it owns the machine. A gateway swapped out to
|
||||
// give the database cache is worse than a smaller cache, and a store
|
||||
// whose working set is megabytes needs little. Only a PostgreSQL init
|
||||
// installed itself: one that was already here may serve something
|
||||
// else that needs its connections.
|
||||
for _, q := range []string{
|
||||
"ALTER SYSTEM SET shared_buffers = '96MB'",
|
||||
"ALTER SYSTEM SET max_connections = 40",
|
||||
} {
|
||||
// Sized to the machine's memory, as one it shares with the gateway:
|
||||
// a gateway swapped out to give the database cache is worse than a
|
||||
// smaller cache, but a knowledge base's vector index is searched
|
||||
// well only from memory. Only a PostgreSQL init installed itself:
|
||||
// one that was already here may serve something else that needs its
|
||||
// connections.
|
||||
meminfo, _ := os.ReadFile("/proc/meminfo")
|
||||
total, ok := pgha.MemTotal(string(meminfo))
|
||||
if !ok {
|
||||
total = 1 << 30
|
||||
}
|
||||
mem := pgha.MemorySettings(total, true)
|
||||
qs := []string{"ALTER SYSTEM SET max_connections = 40"}
|
||||
for _, k := range []string{"shared_buffers", "effective_cache_size", "maintenance_work_mem", "work_mem"} {
|
||||
qs = append(qs, "ALTER SYSTEM SET "+k+" = '"+mem[k]+"'")
|
||||
}
|
||||
for _, q := range qs {
|
||||
if _, err := in.psql(ctx, q); err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -367,9 +376,12 @@ func (in *serverInit) setupPostgres(ctx context.Context) (string, error) {
|
||||
if err := in.waitForPostgres(ctx); err != nil {
|
||||
return "", err
|
||||
}
|
||||
in.done("PostgreSQL sized for a machine it shares (shared_buffers 96MB, 40 connections)")
|
||||
in.done("PostgreSQL sized for a machine it shares, with %d MB of memory (shared_buffers %s, work_mem %s, 40 connections)",
|
||||
total>>20, mem["shared_buffers"], mem["work_mem"])
|
||||
}
|
||||
|
||||
in.setupPgvector(ctx, discovery.DiscoverHost())
|
||||
|
||||
// Proof the address works, before anything is moved into it.
|
||||
address := "@" + in.dsnPath()
|
||||
s, err := store.Open(ctx, address)
|
||||
@@ -482,8 +494,17 @@ func lastLines(b []byte, n int) string {
|
||||
|
||||
// psql runs one statement as the postgres superuser, over the local socket.
|
||||
func (in *serverInit) psql(ctx context.Context, query string) (string, error) {
|
||||
cmd := exec.CommandContext(ctx, "runuser", "-u", "postgres", "--",
|
||||
in.pg.tool("psql"), "-v", "ON_ERROR_STOP=1", "-qAtX", "-c", query)
|
||||
return in.psqlIn(ctx, "", query)
|
||||
}
|
||||
|
||||
// psqlIn runs a statement as postgres in one database, or in postgres's
|
||||
// default one when db is "".
|
||||
func (in *serverInit) psqlIn(ctx context.Context, db, query string) (string, error) {
|
||||
args := []string{"-v", "ON_ERROR_STOP=1", "-qAtX", "-c", query}
|
||||
if db != "" {
|
||||
args = append(args, "-d", db)
|
||||
}
|
||||
cmd := exec.CommandContext(ctx, "runuser", append([]string{"-u", "postgres", "--", in.pg.tool("psql")}, args...)...)
|
||||
// postgres cannot enter root's home, and psql says so on every call.
|
||||
cmd.Dir = "/"
|
||||
out, err := cmd.CombinedOutput()
|
||||
@@ -538,3 +559,50 @@ func (in *serverInit) stopService(ctx context.Context) (bool, error) {
|
||||
in.done("llmgw-server stopped while its store moves")
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// pgvectorPackages are the package names of pgvector for a PostgreSQL major
|
||||
// version: Debian's and Ubuntu's (and PGDG's apt repository) name it after the
|
||||
// server, PGDG's rpm repository after the version alone.
|
||||
func pgvectorPackages(major int) map[string][]string {
|
||||
return map[string][]string{
|
||||
"deb": {fmt.Sprintf("postgresql-%d-pgvector", major)},
|
||||
"rpm": {fmt.Sprintf("pgvector_%d", major)},
|
||||
}
|
||||
}
|
||||
|
||||
// setupPgvector installs pgvector and enables it in llmgw's database, for
|
||||
// knowledge bases. Not a reason to stop: without it everything else works,
|
||||
// and the store says why knowledge bases are off when somebody asks for one.
|
||||
func (in *serverInit) setupPgvector(ctx context.Context, h discovery.Host) {
|
||||
warn := func(format string, a ...any) {
|
||||
fmt.Fprintf(in.out, " !! pgvector: "+format+"; knowledge bases stay off\n", a...)
|
||||
}
|
||||
if on, err := in.psqlIn(ctx, pgDatabase, "SELECT 1 FROM pg_extension WHERE extname = 'vector'"); err == nil && on == "1" {
|
||||
in.kept("pgvector in database %s", pgDatabase)
|
||||
return
|
||||
}
|
||||
have, err := in.psqlBool(ctx, "SELECT 1 FROM pg_available_extensions WHERE name = 'vector'")
|
||||
if err != nil {
|
||||
warn("%v", err)
|
||||
return
|
||||
}
|
||||
if !have {
|
||||
num, err := in.psql(ctx, "SHOW server_version_num")
|
||||
n, convErr := strconv.Atoi(num)
|
||||
if err != nil || convErr != nil || h.Family() == "" {
|
||||
warn("cannot tell which package to install (version %q, %s)", num, h.OSID)
|
||||
return
|
||||
}
|
||||
pkgs := pgvectorPackages(n / 10000)
|
||||
pi := &provision.Installer{Host: h, Log: slog.Default()}
|
||||
if err := pi.InstallPackages(ctx, pkgs); err != nil {
|
||||
warn("installing %s: %v", strings.Join(pkgs[h.Family()], " "), err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if _, err := in.psqlIn(ctx, pgDatabase, "CREATE EXTENSION IF NOT EXISTS vector"); err != nil {
|
||||
warn("%v", err)
|
||||
return
|
||||
}
|
||||
in.done("pgvector enabled in database %s", pgDatabase)
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"strings"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// llmgwServices are the units llmgw runs as on a machine, in the order they
|
||||
// start: the store's clients after the store, the server before the agent
|
||||
// that dials it.
|
||||
var llmgwServices = []string{"llmgw-mirror", "llmgw-server", "llmgw-agent"}
|
||||
|
||||
// unitDir is where llmgw's units are installed.
|
||||
var unitDir = "/etc/systemd/system"
|
||||
|
||||
// serviceRunner runs systemctl; a variable so tests stand in for it.
|
||||
var serviceRunner = func(ctx context.Context, args ...string) (string, error) {
|
||||
out, err := exec.CommandContext(ctx, "systemctl", args...).CombinedOutput() //nolint:gosec // fixed command
|
||||
return strings.TrimSpace(string(out)), err
|
||||
}
|
||||
|
||||
// installedServices are the llmgw units this machine has, or those named.
|
||||
func installedServices(ctx context.Context, names []string) ([]string, error) {
|
||||
var want []string
|
||||
for _, n := range names {
|
||||
if !strings.HasPrefix(n, "llmgw-") {
|
||||
n = "llmgw-" + n
|
||||
}
|
||||
want = append(want, n)
|
||||
}
|
||||
var have []string
|
||||
for _, u := range llmgwServices {
|
||||
if _, err := os.Stat(unitDir + "/" + u + ".service"); err != nil {
|
||||
continue
|
||||
}
|
||||
if len(want) > 0 && !hasUnit(want, u) {
|
||||
continue
|
||||
}
|
||||
have = append(have, u)
|
||||
}
|
||||
for _, w := range want {
|
||||
if !hasUnit(have, w) {
|
||||
return nil, fmt.Errorf("%s is not installed here (this machine has: %s)", w, strings.Join(present(), ", "))
|
||||
}
|
||||
}
|
||||
if len(have) == 0 {
|
||||
return nil, errors.New("no llmgw service is installed on this machine")
|
||||
}
|
||||
return have, nil
|
||||
}
|
||||
|
||||
func present() []string {
|
||||
var out []string
|
||||
for _, u := range llmgwServices {
|
||||
if _, err := os.Stat(unitDir + "/" + u + ".service"); err == nil {
|
||||
out = append(out, strings.TrimPrefix(u, "llmgw-"))
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return []string{"none"}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func hasUnit(list []string, s string) bool {
|
||||
for _, x := range list {
|
||||
if x == s {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// newServiceCmd is llmgw service: this machine's llmgw services, one by one.
|
||||
// The whole system is llmgw status, start, stop and restart.
|
||||
func newServiceCmd() *cobra.Command {
|
||||
c := &cobra.Command{
|
||||
Use: "service",
|
||||
Short: "This machine's llmgw services: status, start, stop, restart",
|
||||
Long: "The control plane (server), node agent (agent) and model mirror (mirror) on\n" +
|
||||
"this machine alone, by systemd. The whole system is 'llmgw status', 'llmgw\n" +
|
||||
"stop', 'llmgw start' and 'llmgw restart', run on a control plane.",
|
||||
}
|
||||
c.AddCommand(serviceVerbs()...)
|
||||
return c
|
||||
}
|
||||
|
||||
func serviceVerbs() []*cobra.Command {
|
||||
status := &cobra.Command{
|
||||
Use: "status [server|agent|mirror]...",
|
||||
Short: "Whether llmgw's services on this machine run",
|
||||
Long: "Each llmgw service this machine has — the control plane (server), the\n" +
|
||||
"node agent (agent), the model mirror (mirror) — and whether it runs.",
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
units, err := localUnits(cmd, args)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
for _, u := range units {
|
||||
state, _ := serviceRunner(cmd.Context(), "is-active", u)
|
||||
since, _ := serviceRunner(cmd.Context(), "show", u, "-p", "ActiveEnterTimestamp", "--value")
|
||||
fmt.Fprintf(out, "%-14s %-10s %s\n", strings.TrimPrefix(u, "llmgw-"), state, since)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
verb := func(use, short, long, action string, reverse bool) *cobra.Command {
|
||||
var engines bool
|
||||
c := &cobra.Command{
|
||||
Use: use + " [server|agent|mirror]...",
|
||||
Short: short,
|
||||
Long: long,
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if os.Geteuid() != 0 {
|
||||
return errors.New("starting and stopping services needs root: run it with sudo")
|
||||
}
|
||||
units, err := localUnits(cmd, args)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if reverse {
|
||||
// The agent before the server it dials, the server before
|
||||
// the mirror it reads.
|
||||
for i, j := 0, len(units)-1; i < j; i, j = i+1, j-1 {
|
||||
units[i], units[j] = units[j], units[i]
|
||||
}
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
for _, u := range units {
|
||||
if msg, err := serviceRunner(cmd.Context(), action, u); err != nil {
|
||||
return fmt.Errorf("%s %s: %v %s", action, u, err, msg)
|
||||
}
|
||||
state, _ := serviceRunner(cmd.Context(), "is-active", u)
|
||||
fmt.Fprintf(out, "%-14s %s\n", strings.TrimPrefix(u, "llmgw-"), state)
|
||||
}
|
||||
if action == "stop" && hasUnit(units, "llmgw-agent") {
|
||||
if engines {
|
||||
_, _ = exec.CommandContext(cmd.Context(), "pkill", "-u", "llmgw-engine").CombinedOutput() //nolint:gosec // fixed command
|
||||
fmt.Fprintln(out, "engines stopped (every process of llmgw-engine)")
|
||||
} else {
|
||||
fmt.Fprintln(out, "\nThe engines keep serving: they outlive the agent, which takes them over again when it starts.\n"+
|
||||
"--engines stops them too; the control plane then places their replicas elsewhere.")
|
||||
}
|
||||
}
|
||||
if action == "stop" && hasUnit(units, "llmgw-server") {
|
||||
fmt.Fprintln(out, "A standby control plane, where there is one, takes over within seconds.")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if action == "stop" {
|
||||
c.Flags().BoolVar(&engines, "engines", false, "also stop the engines this node runs")
|
||||
}
|
||||
return c
|
||||
}
|
||||
return []*cobra.Command{
|
||||
status,
|
||||
verb("start", "Start llmgw's services on this machine",
|
||||
"Starts each llmgw service this machine has, or those named: the mirror, the\n"+
|
||||
"control plane, the node agent, in that order.", "start", false),
|
||||
verb("stop", "Stop llmgw's services on this machine",
|
||||
"Stops each llmgw service this machine has, or those named: the agent, the\n"+
|
||||
"control plane, the mirror, in that order. The engines a node runs keep\n"+
|
||||
"serving unless --engines is given; a standby control plane takes over from\n"+
|
||||
"a leader stopped. They stay stopped until 'llmgw start', or the machine\n"+
|
||||
"boots.", "stop", true),
|
||||
verb("restart", "Restart llmgw's services on this machine",
|
||||
"Restarts each llmgw service this machine has, or those named. Engines keep\n"+
|
||||
"serving and are taken over again by the agent; a control plane's settings\n"+
|
||||
"and server.yaml are read anew.", "restart", false),
|
||||
}
|
||||
}
|
||||
|
||||
func localUnits(cmd *cobra.Command, args []string) ([]string, error) {
|
||||
if runtime.GOOS != "linux" {
|
||||
return nil, fmt.Errorf("llmgw's services run under systemd on Linux; this is %s", runtime.GOOS)
|
||||
}
|
||||
return installedServices(cmd.Context(), args)
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
func TestTheServicesHereAreFoundAndAskedAbout(t *testing.T) {
|
||||
old, oldRun := unitDir, serviceRunner
|
||||
unitDir = t.TempDir()
|
||||
t.Cleanup(func() { unitDir, serviceRunner = old, oldRun })
|
||||
for _, u := range []string{"llmgw-server", "llmgw-agent"} {
|
||||
_ = os.WriteFile(filepath.Join(unitDir, u+".service"), []byte("[Unit]"), 0o644)
|
||||
}
|
||||
var asked []string
|
||||
serviceRunner = func(_ context.Context, args ...string) (string, error) {
|
||||
asked = append(asked, strings.Join(args, " "))
|
||||
if args[0] == "is-active" {
|
||||
return "active", nil
|
||||
}
|
||||
return "Wed 2026-10-07 12:00:00 UTC", nil
|
||||
}
|
||||
|
||||
units, err := installedServices(context.Background(), nil)
|
||||
if err != nil || strings.Join(units, ",") != "llmgw-server,llmgw-agent" {
|
||||
t.Fatalf("%v %v", units, err)
|
||||
}
|
||||
if _, err := installedServices(context.Background(), []string{"mirror"}); err == nil ||
|
||||
!strings.Contains(err.Error(), "server, agent") {
|
||||
t.Errorf("a service not here: %v", err)
|
||||
}
|
||||
if u, err := installedServices(context.Background(), []string{"agent"}); err != nil || len(u) != 1 {
|
||||
t.Errorf("one named: %v %v", u, err)
|
||||
}
|
||||
|
||||
if runtime.GOOS != "linux" {
|
||||
return
|
||||
}
|
||||
var out bytes.Buffer
|
||||
c := serviceVerbs()[0]
|
||||
c.SetOut(&out)
|
||||
c.SetContext(context.Background())
|
||||
if err := c.RunE(c, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(out.String(), "server") || !strings.Contains(out.String(), "active") {
|
||||
t.Errorf("%s", out.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopSetsTheClustersServingStopped(t *testing.T) {
|
||||
var put adminapi.SettingsRequest
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/api/v1/summary":
|
||||
_ = json.NewEncoder(w).Encode(adminapi.Summary{})
|
||||
case r.Method == "GET" && r.URL.Path == "/api/v1/settings":
|
||||
_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 3, Body: "router:\n retries: 2\n"})
|
||||
case r.Method == "PUT" && r.URL.Path == "/api/v1/settings":
|
||||
_ = json.NewDecoder(r.Body).Decode(&put)
|
||||
_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 4, Body: put.Body})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
dir := t.TempDir()
|
||||
cfg := filepath.Join(dir, "server.yaml")
|
||||
_ = os.WriteFile(filepath.Join(dir, "token"), []byte("t"), 0o600)
|
||||
_ = os.WriteFile(cfg, []byte("listen: \":8080\"\nagent_insecure: true\nadmin_listen: "+
|
||||
strings.TrimPrefix(srv.URL, "http://")+"\nadmin_remote: true\nadmin_token_file: "+filepath.Join(dir, "token")+"\n"), 0o600)
|
||||
|
||||
var stop *cobra.Command
|
||||
for _, c := range newSystemCmds() {
|
||||
if c.Name() == "stop" {
|
||||
stop = c
|
||||
}
|
||||
}
|
||||
var out bytes.Buffer
|
||||
stop.SetOut(&out)
|
||||
stop.SetArgs([]string{"--config", cfg, "--note", "maintenance"})
|
||||
if err := stop.ExecuteContext(context.Background()); err != nil {
|
||||
t.Fatalf("%v: %s", err, out.String())
|
||||
}
|
||||
if put.Base != 3 || !strings.Contains(put.Body, "stopped: true") || !strings.Contains(put.Body, "retries: 2") ||
|
||||
put.Note != "maintenance" {
|
||||
t.Errorf("sent %+v", put)
|
||||
}
|
||||
if !strings.Contains(out.String(), "stopping") {
|
||||
t.Errorf("%s", out.String())
|
||||
}
|
||||
}
|
||||
|
||||
// Started while the control plane has not yet taken the setting up: every
|
||||
// model of a stopped system has the nothing it asks for, which is not
|
||||
// serving.
|
||||
func TestStartWaitsForTheControlPlaneToServe(t *testing.T) {
|
||||
old := servicePoll
|
||||
servicePoll = time.Millisecond
|
||||
t.Cleanup(func() { servicePoll = old })
|
||||
answers := []adminapi.Summary{
|
||||
{Stopped: true, Models: adminapi.ModelTotals{Total: 2, Satisfied: 2}},
|
||||
{Models: adminapi.ModelTotals{Total: 2, Satisfied: 0}},
|
||||
{Models: adminapi.ModelTotals{Total: 2, Satisfied: 2}},
|
||||
}
|
||||
asked := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode(answers[min(asked, len(answers)-1)])
|
||||
asked++
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
var out bytes.Buffer
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetOut(&out)
|
||||
cmd.SetContext(context.Background())
|
||||
c := &ctlClient{base: srv.URL, http: srv.Client()}
|
||||
if err := waitForService(cmd, func() *ctlClient { return c }, false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if asked != 3 {
|
||||
t.Errorf("returned after %d answers: %s", asked, out.String())
|
||||
}
|
||||
}
|
||||
@@ -153,6 +153,19 @@ func printSummary(cmd *cobra.Command, s adminapi.Summary) {
|
||||
fmt.Fprintf(out, "STORE %s\n", strings.Join(parts, ", "))
|
||||
}
|
||||
|
||||
if s.Stopped {
|
||||
fmt.Fprintln(out, "SERVICE STOPPED: no model runs and requests are answered 503 (llmgw start)")
|
||||
}
|
||||
if s.GuardOff {
|
||||
fmt.Fprintln(out, "GUARD off: no request is checked, whatever the policies say (llmgw ctl guard on)")
|
||||
}
|
||||
if len(s.Judges) > 0 {
|
||||
fmt.Fprintln(out, "\nGUARDRAILS' JUDGE")
|
||||
for _, j := range s.Judges {
|
||||
fmt.Fprintf(out, " %s: %s\n", j.Model, j.Problem)
|
||||
}
|
||||
}
|
||||
|
||||
if len(s.Parked) > 0 {
|
||||
fmt.Fprintln(out, "\nPARKED — failed every time they started; not retried until fixed")
|
||||
for _, p := range s.Parked {
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
|
||||
)
|
||||
|
||||
// systemClient is the admin API of the control plane on this machine, as its
|
||||
// server.yaml says, with its token.
|
||||
func systemClient(configPath string) func() *ctlClient {
|
||||
return func() *ctlClient {
|
||||
addr := "127.0.0.1:9090"
|
||||
var tokens []string
|
||||
if cfg, ok := loadServerConfig(configPath); ok {
|
||||
if cfg.AdminListen != "" {
|
||||
addr = cfg.AdminListen
|
||||
if strings.HasPrefix(addr, ":") {
|
||||
addr = "127.0.0.1" + addr
|
||||
}
|
||||
}
|
||||
tokens = append(tokens, adminTokenPath(configPath, cfg))
|
||||
}
|
||||
return newCtlClient(addr, tokens...)
|
||||
}
|
||||
}
|
||||
|
||||
// newSystemCmds are llmgw status, stop, start and restart: the whole system,
|
||||
// from a control plane.
|
||||
func newSystemCmds() []*cobra.Command {
|
||||
var configPath string
|
||||
client := systemClient("/etc/llmgw/server.yaml")
|
||||
bind := func(c *cobra.Command) *cobra.Command {
|
||||
c.Flags().StringVar(&configPath, "config", "/etc/llmgw/server.yaml", "this control plane's configuration")
|
||||
c.PreRun = func(*cobra.Command, []string) { client = systemClient(configPath) }
|
||||
return c
|
||||
}
|
||||
call := func() *ctlClient { return client() }
|
||||
|
||||
status := bind(&cobra.Command{
|
||||
Use: "status",
|
||||
Short: "The whole system: whether it serves, its control planes, nodes, models and store",
|
||||
Long: "Run on a control plane. Says whether the service is stopped (llmgw stop),\n" +
|
||||
"and everything 'llmgw ctl status' does. One machine's services are\n" +
|
||||
"'llmgw service status'.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var s adminapi.Summary
|
||||
if err := call().get(cmd.Context(), "/api/v1/summary", &s); err != nil {
|
||||
return notAControlPlane(err)
|
||||
}
|
||||
printSummary(cmd, s)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
switchService := func(stop bool) *cobra.Command {
|
||||
use, short, long := "start", "Start the whole system serving again, after llmgw stop",
|
||||
"Every model comes back as its catalog entry asks: placed, loaded and\n"+
|
||||
"serving again, on every control plane within seconds. Run on a control\n"+
|
||||
"plane. The services themselves are never stopped by llmgw stop; a\n"+
|
||||
"machine's own are started with 'llmgw service start'."
|
||||
if stop {
|
||||
use, short, long = "stop", "Stop the whole system serving: every model taken down, requests answered 503",
|
||||
"Every model's replicas are taken down, freeing every GPU, and every\n"+
|
||||
"request for a model is answered 503 (system_stopped) until 'llmgw start'.\n"+
|
||||
"The control planes, agents and mirror keep running, so the console, the\n"+
|
||||
"admin pages and llmgw start still work, and no model's definition is\n"+
|
||||
"changed. In force on every control plane within seconds; kept in the\n"+
|
||||
"cluster's settings (serving.stopped), so a restart does not undo it."
|
||||
}
|
||||
var note string
|
||||
var wait bool
|
||||
c := bind(&cobra.Command{
|
||||
Use: use,
|
||||
Short: short,
|
||||
Long: long,
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if note == "" {
|
||||
note = "llmgw " + use
|
||||
}
|
||||
var s adminapi.Summary
|
||||
if err := call().get(cmd.Context(), "/api/v1/summary", &s); err != nil {
|
||||
return notAControlPlane(err)
|
||||
}
|
||||
if s.Stopped == stop {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "the system is already %s\n", map[bool]string{true: "stopped", false: "serving"}[stop])
|
||||
return nil
|
||||
}
|
||||
if err := changeSettings(cmd, call, note, func(doc *settings.Doc) error {
|
||||
return doc.SetYAML([]string{"serving", "stopped"}, fmt.Sprint(stop))
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
if !wait {
|
||||
if stop {
|
||||
fmt.Fprintln(out, "stopping: every model's replicas are being taken down; llmgw status follows it")
|
||||
} else {
|
||||
fmt.Fprintln(out, "starting: every model is placed again as it asks; llmgw status follows it")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return waitForService(cmd, call, stop)
|
||||
},
|
||||
})
|
||||
c.Flags().StringVar(¬e, "note", "", "why, kept with the settings' version and in the audit")
|
||||
c.Flags().BoolVar(&wait, "wait", false, "wait until every replica is down (stop) or the models have theirs (start)")
|
||||
return c
|
||||
}
|
||||
|
||||
var sshUser string
|
||||
restart := bind(&cobra.Command{
|
||||
Use: "restart",
|
||||
Short: "Restart every control plane and node agent; the engines keep serving",
|
||||
Long: "Every node's agent starts again, then the other control planes over ssh\n" +
|
||||
"as --ssh-user, then this one — the leader last, so it hands over to a\n" +
|
||||
"standby already restarted. The engines keep serving throughout and are\n" +
|
||||
"taken over again; nothing is placed anew. Run on a control plane, with\n" +
|
||||
"sudo. One machine's services are 'llmgw service restart'.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if os.Geteuid() != 0 {
|
||||
return errors.New("restarting this control plane needs root: run it with sudo")
|
||||
}
|
||||
ctx, out := cmd.Context(), cmd.OutOrStdout()
|
||||
var s adminapi.Summary
|
||||
if err := call().get(ctx, "/api/v1/summary", &s); err != nil {
|
||||
return notAControlPlane(err)
|
||||
}
|
||||
var agents adminapi.AgentRestartView
|
||||
if err := call().post(ctx, "/api/v1/agents/restart", []byte("{}"), &agents); err != nil {
|
||||
return fmt.Errorf("restarting the agents: %w", err)
|
||||
}
|
||||
fmt.Fprintf(out, "agents %d told to restart: %s\n", len(agents.Nodes), strings.Join(agents.Nodes, ", "))
|
||||
|
||||
// The other control planes, standbys before a leader.
|
||||
var others []adminapi.ControlPlane
|
||||
for _, cp := range s.ControlPlanes.Members {
|
||||
if !cp.This {
|
||||
others = append(others, cp)
|
||||
}
|
||||
}
|
||||
// One that cannot be reached is said and left: the rest of the
|
||||
// system still restarts, and the leader still goes last.
|
||||
var failed []string
|
||||
for pass := 0; pass < 2; pass++ {
|
||||
for _, cp := range others {
|
||||
leading := cp.State == "leading"
|
||||
if leading != (pass == 1) {
|
||||
continue
|
||||
}
|
||||
r := &onboard.SSH{Addr: cp.Address, User: sshUser}
|
||||
if res, err := r.Run(ctx, "sudo systemctl restart llmgw-server", nil); err != nil {
|
||||
fmt.Fprintf(out, "%-10s NOT restarted: ssh as %q: %v %s\n", cp.Name, sshUser, err, strings.TrimSpace(string(res)))
|
||||
failed = append(failed, cp.Name+" ("+cp.Address+")")
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(out, "%-10s restarted (%s)\n", cp.Name, cp.State)
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
}
|
||||
if _, err := os.Stat(unitDir + "/llmgw-server.service"); err == nil {
|
||||
fmt.Fprintln(out, "this restarting")
|
||||
if b, err := exec.CommandContext(ctx, "systemctl", "restart", "llmgw-server").CombinedOutput(); err != nil { //nolint:gosec // fixed command
|
||||
return fmt.Errorf("restarting this control plane: %v: %s", err, strings.TrimSpace(string(b)))
|
||||
}
|
||||
fmt.Fprintln(out, "this restarted")
|
||||
}
|
||||
fmt.Fprintln(out, "\nllmgw status shows the nodes coming back; their replicas are taken over as they are.")
|
||||
if len(failed) > 0 {
|
||||
return fmt.Errorf("not restarted: %s; run `sudo llmgw service restart server` there, "+
|
||||
"or give --ssh-user someone this machine's root can log in as", strings.Join(failed, ", "))
|
||||
}
|
||||
return nil
|
||||
},
|
||||
})
|
||||
restart.Flags().StringVar(&sshUser, "ssh-user", os.Getenv("SUDO_USER"), "who to log in to the other control planes as")
|
||||
|
||||
return []*cobra.Command{status, switchService(true), switchService(false), restart}
|
||||
}
|
||||
|
||||
// servicePoll is how often --wait asks; a variable so tests need not wait.
|
||||
var servicePoll = 5 * time.Second
|
||||
|
||||
// waitForService waits until no replica runs (stopped) or every model has
|
||||
// the replicas it asks for (started).
|
||||
//
|
||||
// Only once the control plane says it is in the state asked for: it takes the
|
||||
// setting up at its next look at the store, seconds later, and until then a
|
||||
// stopped system's models, which ask for nothing, all have what they ask for —
|
||||
// so llmgw start --wait returned at once, with nothing started.
|
||||
func waitForService(cmd *cobra.Command, call func() *ctlClient, stop bool) error {
|
||||
deadline := time.Now().Add(15 * time.Minute)
|
||||
for time.Now().Before(deadline) {
|
||||
var s adminapi.Summary
|
||||
if err := call().get(cmd.Context(), "/api/v1/summary", &s); err == nil && s.Stopped == stop {
|
||||
if stop && s.Replicas.Ready+s.Replicas.Starting+s.Replicas.Draining == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), "stopped: no replica runs")
|
||||
return nil
|
||||
}
|
||||
if !stop && s.Models.Satisfied >= s.Models.Total {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "serving: %d of %d models have the replicas they ask for\n",
|
||||
s.Models.Satisfied, s.Models.Total)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-cmd.Context().Done():
|
||||
return cmd.Context().Err()
|
||||
case <-time.After(servicePoll):
|
||||
}
|
||||
}
|
||||
return errors.New("not there within fifteen minutes; llmgw status says where it stands")
|
||||
}
|
||||
|
||||
// notAControlPlane is a failure to reach the admin API, said as what it
|
||||
// most often means.
|
||||
func notAControlPlane(err error) error {
|
||||
return fmt.Errorf("%w\n(the whole system is stopped and started from a control plane; this machine's own "+
|
||||
"services are 'llmgw service status|start|stop|restart')", err)
|
||||
}
|
||||
@@ -34,10 +34,11 @@ import (
|
||||
// wrapper — left it in the content. A differently trained model fixes that. A
|
||||
// bigger one from the same family does not: qwen2.5-coder-7b invented
|
||||
// <function-call> and failed the same way.
|
||||
func newToolCheckCmd() *cobra.Command {
|
||||
func newToolCheckCmd(client func() *ctlClient) *cobra.Command {
|
||||
var (
|
||||
endpoint string
|
||||
model string
|
||||
key string
|
||||
timeout time.Duration
|
||||
)
|
||||
|
||||
@@ -57,10 +58,21 @@ func newToolCheckCmd() *cobra.Command {
|
||||
if endpoint == "" {
|
||||
return fmt.Errorf("--endpoint is required, e.g. http://10.0.0.145:20002")
|
||||
}
|
||||
// A replica's engine demands its own key, and refused the check
|
||||
// with a 401 when none was sent. Asked of the admin API, as bench
|
||||
// does; an endpoint that is no replica's is asked without one.
|
||||
if key == "" {
|
||||
reps, err := listReplicas(cmd, client)
|
||||
if err != nil {
|
||||
cmd.PrintErrf("The admin API could not be asked for the engine key of the replica at "+
|
||||
"--endpoint (%v); an engine that demands one will refuse. Pass --key to give it.\n", err)
|
||||
}
|
||||
key = replicaKey(reps, endpoint)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(cmd.Context(), timeout)
|
||||
defer cancel()
|
||||
|
||||
res, err := probeToolCall(ctx, endpoint, model)
|
||||
res, err := probeToolCall(ctx, endpoint, model, key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -75,6 +87,7 @@ func newToolCheckCmd() *cobra.Command {
|
||||
}
|
||||
cmd.Flags().StringVar(&endpoint, "endpoint", "", "engine base URL, e.g. http://10.0.0.145:20002")
|
||||
cmd.Flags().StringVar(&model, "model", "", "model name to send, when the engine needs one")
|
||||
cmd.Flags().StringVar(&key, "key", "", "the engine's key (default: the admin API's, for the replica at --endpoint)")
|
||||
cmd.Flags().DurationVar(&timeout, "timeout", 5*time.Minute, "how long to wait; CPU inference is slow")
|
||||
return cmd
|
||||
}
|
||||
@@ -102,7 +115,7 @@ type ToolProbe struct {
|
||||
Raw string
|
||||
}
|
||||
|
||||
func probeToolCall(ctx context.Context, endpoint, model string) (ToolProbe, error) {
|
||||
func probeToolCall(ctx context.Context, endpoint, model, key string) (ToolProbe, error) {
|
||||
// A question with one obvious tool answer and no other sensible reply, so a
|
||||
// model that can use tools has no excuse not to.
|
||||
body := map[string]any{
|
||||
@@ -138,6 +151,9 @@ func probeToolCall(ctx context.Context, endpoint, model string) (ToolProbe, erro
|
||||
return ToolProbe{}, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if key != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
}
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -150,6 +166,10 @@ func probeToolCall(ctx context.Context, endpoint, model string) (ToolProbe, erro
|
||||
return ToolProbe{}, err
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
return ToolProbe{}, fmt.Errorf("%s answered 401: the engine's key is the admin API's, "+
|
||||
"for the replica `llmgw ctl replicas` lists at --endpoint; or pass it with --key", url)
|
||||
}
|
||||
return ToolProbe{}, fmt.Errorf("%s returned %d: %s", url, resp.StatusCode, trim(string(raw), 300))
|
||||
}
|
||||
return readToolProbe(raw)
|
||||
|
||||
@@ -1,10 +1,65 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
// toolCheck runs the command against an engine that answers only its own key,
|
||||
// with an admin API that lists the replica there with that key.
|
||||
func toolCheck(t *testing.T, args ...string) (string, error) {
|
||||
t.Helper()
|
||||
engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("Authorization") != "Bearer engine-key" {
|
||||
http.Error(w, `{"error":{"message":"Invalid API Key"}}`, http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
fmt.Fprint(w, `{"choices":[{"message":{"tool_calls":[{"id":"c1","type":"function",
|
||||
"function":{"name":"get_weather","arguments":"{\"city\":\"Istanbul\"}"}}]}}]}`)
|
||||
}))
|
||||
admin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode([]adminapi.ReplicaView{
|
||||
{ID: "r1", Model: "m", Profile: "gpu", Endpoint: engine.URL, EngineKey: "engine-key"},
|
||||
})
|
||||
}))
|
||||
t.Cleanup(func() { engine.Close(); admin.Close() })
|
||||
t.Setenv("LLMGW_ADMIN_TOKEN", "t")
|
||||
cmd := newToolCheckCmd(func() *ctlClient { return newCtlClient(admin.URL) })
|
||||
var out bytes.Buffer
|
||||
cmd.SetOut(&out)
|
||||
cmd.SetErr(&out)
|
||||
cmd.SetArgs(append([]string{"--endpoint", engine.URL}, args...))
|
||||
err := cmd.ExecuteContext(context.Background())
|
||||
return out.String(), err
|
||||
}
|
||||
|
||||
func TestToolCheckAsksWithTheEngineKeyTheAdminAPIGives(t *testing.T) {
|
||||
// Every replica's engine demands its key, and the check was sent without
|
||||
// one: it failed with a 401 on every engine in the cluster.
|
||||
out, err := toolCheck(t)
|
||||
if err != nil {
|
||||
t.Fatalf("%v\n%s", err, out)
|
||||
}
|
||||
if !strings.Contains(out, "tool calls work") {
|
||||
t.Errorf("unexpected report:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolCheckSaysWhereTheEngineKeyComesFrom(t *testing.T) {
|
||||
out, err := toolCheck(t, "--key", "wrong")
|
||||
if err == nil || !strings.Contains(err.Error(), "401") || !strings.Contains(err.Error(), "--key") {
|
||||
t.Fatalf("expected the refusal explained, got %v\n%s", err, out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAStructuredToolCallIsRecognised(t *testing.T) {
|
||||
p, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"tool_calls":[{"id":"c1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Istanbul\"}"}}]}}]}`))
|
||||
|
||||
@@ -0,0 +1,293 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
c := &cobra.Command{
|
||||
Use: "train",
|
||||
Short: "Train LoRA adapters on the nodes, on question/answer pairs",
|
||||
Long: "A node trains an adapter for a catalog model on pairs (llmgw ctl kb dataset\n" +
|
||||
"writes them from a knowledge base; read them through first). It trains with\n" +
|
||||
"its installed vLLM's torch and a small overlay, as the engines' user, on a card\n" +
|
||||
"with --gpu-memory free — beside an engine, if it leaves that much. What it\n" +
|
||||
"makes is kept by the control plane: llmgw ctl train download, or publish it\n" +
|
||||
"as the model's adapter.\n\n" +
|
||||
"An adapter learns how answers sound — words, format, tone — far better than\n" +
|
||||
"what they say. Facts belong to the knowledge base.",
|
||||
}
|
||||
|
||||
var model, base, rev, node, data, owner string
|
||||
var public bool
|
||||
var gpuMem int
|
||||
var epochs, rank int
|
||||
var lr float64
|
||||
var noGGUF bool
|
||||
start := &cobra.Command{
|
||||
Use: "start NAME --model MODEL --base-repo OWNER/NAME --node NODE --data FILE",
|
||||
Short: "Queue an adapter's training on a node",
|
||||
Example: " llmgw ctl train start policies-tone --model qwen2.5-0.5b-instruct \\\n" +
|
||||
" --base-repo Qwen/Qwen2.5-0.5B-Instruct --node cuda01 --gpu-memory 2300 --data policies.jsonl",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
b, err := os.ReadFile(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req := adminapi.TrainRequest{Name: args[0], Model: model, BaseRepo: base, BaseRevision: rev, Node: node,
|
||||
GPUMemoryMiB: gpuMem, Data: string(b), Params: map[string]any{}, Owner: owner, Public: public}
|
||||
if cmd.Flags().Changed("epochs") {
|
||||
req.Params["epochs"] = epochs
|
||||
}
|
||||
if cmd.Flags().Changed("rank") {
|
||||
req.Params["rank"] = rank
|
||||
}
|
||||
if cmd.Flags().Changed("lr") {
|
||||
req.Params["lr"] = lr
|
||||
}
|
||||
if noGGUF {
|
||||
f := false
|
||||
req.GGUF = &f
|
||||
}
|
||||
body, _ := json.Marshal(req)
|
||||
var v adminapi.TrainJobView
|
||||
if err := client().patient(2*time.Minute).do(cmd.Context(), "POST", "/api/v1/train/jobs", json.RawMessage(body), &v); err != nil {
|
||||
return err
|
||||
}
|
||||
who := "every caller of " + v.Model
|
||||
if v.Owner != "" {
|
||||
who = v.Owner + " only"
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "queued %s: %s for %s on %s, %d pairs; it will be seen by %s\n",
|
||||
v.ID, v.Name, v.Model, v.Node, v.Pairs, who)
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "follow it with: llmgw ctl train show %s --follow\n", v.ID)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
f := start.Flags()
|
||||
f.StringVar(&model, "model", "", "the catalog model the adapter is for")
|
||||
f.StringVar(&base, "base-repo", "", "the weights to train on, as safetensors on Hugging Face (owner/name)")
|
||||
f.StringVar(&rev, "base-revision", "", "a revision of them (default: the newest)")
|
||||
f.StringVar(&node, "node", "", "the node that trains")
|
||||
f.IntVar(&gpuMem, "gpu-memory", 4096, "MiB of a card the job needs free")
|
||||
f.StringVar(&data, "data", "", "the pairs, NDJSON of {\"q\", \"a\"}")
|
||||
f.IntVar(&epochs, "epochs", 6, "passes over the pairs")
|
||||
f.IntVar(&rank, "rank", 16, "the adapter's rank")
|
||||
f.Float64Var(&lr, "lr", 2e-4, "the learning rate")
|
||||
f.BoolVar(&noGGUF, "no-gguf", false, "skip the llama.cpp GGUF, even for a model llama.cpp serves")
|
||||
f.StringVar(&owner, "owner", "", "who the adapter belongs to — org:, team:, project: (path or ID), user: or key: — and only they see it "+
|
||||
"(default: the knowledge base's owner, when the pairs come from one)")
|
||||
f.BoolVar(&public, "public", false, "every caller of the model sees it; refused for pairs from a knowledge base")
|
||||
for _, req := range []string{"model", "base-repo", "node", "data"} {
|
||||
_ = start.MarkFlagRequired(req)
|
||||
}
|
||||
c.AddCommand(start)
|
||||
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "list",
|
||||
Short: "Training jobs, newest first",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var jobs []adminapi.TrainJobView
|
||||
if err := client().get(cmd.Context(), "/api/v1/train/jobs", &jobs); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(jobs)
|
||||
}
|
||||
if len(jobs) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), "No training jobs (llmgw ctl train start).")
|
||||
return nil
|
||||
}
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "JOB\tADAPTER\tMODEL\tOWNER\tNODE\tSTATE\tPROGRESS\tLOSS\tPEAK\tCREATED")
|
||||
for _, j := range jobs {
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", j.ID, j.Name, j.Model, or(j.Owner, "everyone"), j.Node, j.State,
|
||||
trainProgress(j), trainLoss(j), trainPeak(j), j.CreatedAt.Local().Format("01-02 15:04"))
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
})
|
||||
|
||||
var follow bool
|
||||
show := &cobra.Command{
|
||||
Use: "show JOB",
|
||||
Short: "How a training job goes, and what it made",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
out := cmd.OutOrStdout()
|
||||
last := ""
|
||||
for {
|
||||
var j adminapi.TrainJobView
|
||||
if err := client().get(cmd.Context(), "/api/v1/train/jobs/"+args[0], &j); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(out).Encode(j)
|
||||
}
|
||||
line := fmt.Sprintf("%s %s %s loss %s peak %s %s", j.State, trainProgress(j), j.Node, trainLoss(j), trainPeak(j), j.Message)
|
||||
if !follow {
|
||||
fmt.Fprintf(out, "%s: adapter %s for %s, owned by %s, on %s (base %s), %d pairs\n", j.ID, j.Name, j.Model,
|
||||
or(j.Owner, "everyone"), j.Node, j.BaseRepo, j.Pairs)
|
||||
fmt.Fprintln(out, line)
|
||||
if len(j.Artifacts) > 0 {
|
||||
fmt.Fprintf(out, "made: %s (llmgw ctl train download %s --kind KIND)\n", strings.Join(j.Artifacts, ", "), j.ID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if line != last {
|
||||
fmt.Fprintf(out, "%s %s\n", time.Now().Format("15:04:05"), line)
|
||||
last = line
|
||||
}
|
||||
if store.TrainOver(j.State) {
|
||||
if j.State != store.TrainDone {
|
||||
return fmt.Errorf("%s %s", j.ID, j.State)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-cmd.Context().Done():
|
||||
return nil
|
||||
case <-time.After(3 * time.Second):
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
show.Flags().BoolVarP(&follow, "follow", "f", false, "follow it until it ends")
|
||||
c.AddCommand(show)
|
||||
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "cancel JOB",
|
||||
Short: "Stop a training job",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if err := client().do(cmd.Context(), "POST", "/api/v1/train/jobs/"+args[0]+"/cancel", nil, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s cancelled\n", args[0])
|
||||
return nil
|
||||
},
|
||||
})
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "rm JOB",
|
||||
Short: "Remove an ended training job and what it made",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if err := client().do(cmd.Context(), "DELETE", "/api/v1/train/jobs/"+args[0], nil, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s removed\n", args[0])
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
var restart bool
|
||||
publish := &cobra.Command{
|
||||
Use: "publish JOB",
|
||||
Short: "Make what a finished job made its model's adapter, kept on the mirror",
|
||||
Long: "The GGUF for a model llama.cpp serves, the PEFT files for one vLLM serves, are kept\n" +
|
||||
"on the mirror and added to the model's adapters; a request then asks for the\n" +
|
||||
"adapter by its name, as a model. llama.cpp loads adapters when a replica starts:\n" +
|
||||
"--restart replaces the model's replicas, one at a time, to load it.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
var v adminapi.TrainPublishView
|
||||
if err := client().patient(15*time.Minute).do(cmd.Context(), "POST", "/api/v1/train/jobs/"+args[0]+"/publish", nil, &v); err != nil {
|
||||
return err
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintf(out, "%s is %s's adapter, from %s\n", v.Adapter, v.Model, v.Source)
|
||||
if v.Owner != "" {
|
||||
fmt.Fprintf(out, "only %s sees it; to anyone else it is a model that does not exist\n", v.Owner)
|
||||
}
|
||||
if restart && strings.Contains(v.Note, "restart") {
|
||||
body, _ := json.Marshal(adminapi.RestartRequest{Models: []string{v.Model}})
|
||||
var res adminapi.RestartResult
|
||||
if err := client().post(cmd.Context(), "/api/v1/restart", body, &res); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(out, "%d replica(s) of %s restarting to load it; then ask for model %s\n", len(res.Queued), v.Model, v.Adapter)
|
||||
return nil
|
||||
}
|
||||
fmt.Fprintln(out, v.Note)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
publish.Flags().BoolVar(&restart, "restart", false, "restart the model's replicas to load it (llama.cpp)")
|
||||
c.AddCommand(publish)
|
||||
|
||||
var kind, outFile string
|
||||
download := &cobra.Command{
|
||||
Use: "download JOB",
|
||||
Short: "Save what a training job made: the PEFT adapter (peft) or the GGUF (gguf)",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if outFile == "" {
|
||||
outFile = args[0] + ".tar.gz"
|
||||
if kind == "gguf" {
|
||||
outFile = args[0] + ".gguf"
|
||||
}
|
||||
}
|
||||
f, err := os.Create(outFile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
cw := &countWriter{w: f}
|
||||
if err := client().patient(10*time.Minute).do(cmd.Context(), "GET",
|
||||
"/api/v1/train/jobs/"+args[0]+"/artifacts/"+kind, nil, io.Writer(cw)); err != nil {
|
||||
_ = os.Remove(outFile)
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: %s\n", outFile, humanBytes(cw.n))
|
||||
return nil
|
||||
},
|
||||
}
|
||||
download.Flags().StringVar(&kind, "kind", "peft", "peft (a gzipped tar of the PEFT adapter) or gguf")
|
||||
download.Flags().StringVarP(&outFile, "out", "o", "", "the file to write")
|
||||
c.AddCommand(download)
|
||||
return c
|
||||
}
|
||||
|
||||
type countWriter struct {
|
||||
w io.Writer
|
||||
n int64
|
||||
}
|
||||
|
||||
func (c *countWriter) Write(p []byte) (int, error) {
|
||||
n, err := c.w.Write(p)
|
||||
c.n += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func trainProgress(j adminapi.TrainJobView) string {
|
||||
if j.Steps == 0 {
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%d/%d", j.Step, j.Steps)
|
||||
}
|
||||
|
||||
func trainLoss(j adminapi.TrainJobView) string {
|
||||
if j.Loss == 0 {
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%.3f", j.Loss)
|
||||
}
|
||||
|
||||
func trainPeak(j adminapi.TrainJobView) string {
|
||||
if j.PeakMiB == 0 {
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%d MiB", j.PeakMiB)
|
||||
}
|
||||
@@ -679,6 +679,8 @@ func (m *consoleModel) doEntry() tea.Cmd {
|
||||
}}
|
||||
case cBackup:
|
||||
return m.runAfter("server", "backup", "--config="+m.file)
|
||||
case cJudge:
|
||||
return m.runAfter("ctl", server, "guard", "judge", "enable", "--replace")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -45,6 +45,7 @@ const (
|
||||
cUsage
|
||||
cQuotas
|
||||
cGuard
|
||||
cKB
|
||||
cFederation
|
||||
cMirror
|
||||
cJobs
|
||||
@@ -59,6 +60,7 @@ const (
|
||||
cAddControl
|
||||
cReplicate
|
||||
cServe
|
||||
cJudge
|
||||
cUpgrade
|
||||
cRestartAll
|
||||
cBackup
|
||||
@@ -67,8 +69,8 @@ const (
|
||||
cCount
|
||||
)
|
||||
|
||||
var consoleLabels = [cCount]string{"Overview", "Monitor", "Nodes", "GPUs", "Models", "Replicas", "Usage", "Quotas", "Guardrails", "Federation", "Mirror", "Jobs", "Logs", "Accounts", "API keys", "Organizations", "Audit", "Manual",
|
||||
"server.yaml", "Add a node", "Add a control plane", "Replicate the store", "Serve a model", "Upgrade the cluster",
|
||||
var consoleLabels = [cCount]string{"Overview", "Monitor", "Nodes", "GPUs", "Models", "Replicas", "Usage", "Quotas", "Guardrails", "Knowledge", "Federation", "Mirror", "Jobs", "Logs", "Accounts", "API keys", "Organizations", "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"}
|
||||
|
||||
// consoleAction is what is done once the screen closes: each runs its own
|
||||
@@ -176,6 +178,7 @@ type consoleModel struct {
|
||||
quotas []adminapi.QuotaView
|
||||
quotasErr error
|
||||
guard guardState
|
||||
kb kbState
|
||||
|
||||
st *storeHolder
|
||||
people peopleData
|
||||
@@ -322,6 +325,9 @@ func (m consoleModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
if m.menu == cGuard {
|
||||
cmds = append(cmds, m.fetchGuard())
|
||||
}
|
||||
if m.menu == cKB {
|
||||
cmds = append(cmds, m.fetchKB())
|
||||
}
|
||||
if m.menu == cMirror {
|
||||
cmds = append(cmds, m.fetchMirror())
|
||||
}
|
||||
@@ -380,6 +386,8 @@ func (m consoleModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
m.people = peopleData(msg)
|
||||
case consoleQuotasMsg:
|
||||
m.quotas, m.quotasErr = msg.rows, msg.err
|
||||
case consoleKBMsg:
|
||||
m.kb = kbState{view: msg.view, err: msg.err, loaded: true}
|
||||
case consoleGuardMsg:
|
||||
window := m.guard.window
|
||||
m.guard = guardState{view: msg.view, events: msg.events, err: msg.err, window: window}
|
||||
@@ -496,7 +504,7 @@ func (m consoleModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
return m.leave()
|
||||
case cWipe:
|
||||
return m.askWipe()
|
||||
case cUpgrade, cRestartAll, cBackup:
|
||||
case cUpgrade, cRestartAll, cBackup, cJudge:
|
||||
return m, m.doEntry()
|
||||
case cOverview:
|
||||
return m, nil
|
||||
@@ -538,6 +546,8 @@ func (m consoleModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.menuD = m.quotaMenu(m.quotas[m.row])
|
||||
case cGuard:
|
||||
m.menuD = m.guardAttachmentMenu(m.guard.view.Attachments[m.row])
|
||||
case cKB:
|
||||
m.menuD = m.kbMenu(m.kb.view.Collections[m.row])
|
||||
case cMirror:
|
||||
m.menuD = m.mirrorMenu(m.mirrorGroups()[m.row])
|
||||
case cJobs:
|
||||
@@ -595,6 +605,10 @@ func (m consoleModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.quotaForm(adminapi.QuotaView{})
|
||||
case cGuard:
|
||||
m.guardAttachForm()
|
||||
case cKB:
|
||||
if m.kb.view.Available {
|
||||
m.kbCreateForm()
|
||||
}
|
||||
case cModels:
|
||||
m.addModelForm()
|
||||
case cMirror:
|
||||
@@ -650,6 +664,8 @@ func (m consoleModel) onEnter() tea.Cmd {
|
||||
return m.fetchQuotas()
|
||||
case cGuard:
|
||||
return m.fetchGuard()
|
||||
case cKB:
|
||||
return m.fetchKB()
|
||||
case cMirror:
|
||||
return m.fetchMirror()
|
||||
case cMonitor:
|
||||
@@ -1026,6 +1042,8 @@ func (m consoleModel) rowCount() int {
|
||||
return len(m.quotas)
|
||||
case cGuard:
|
||||
return len(m.guard.view.Attachments)
|
||||
case cKB:
|
||||
return len(m.kb.view.Collections)
|
||||
case cMirror:
|
||||
return len(m.mirrorGroups())
|
||||
case cJobs:
|
||||
@@ -1229,13 +1247,17 @@ func (m consoleModel) panel() (title, body, keys string) {
|
||||
for _, s := range r.Shards {
|
||||
nodes = append(nodes, fmt.Sprintf("%s%v", s.Node, discovery.DeviceLabels(s.GPUs)))
|
||||
}
|
||||
rows = append(rows, []string{r.ID, r.Model, r.State, r.Pool, strings.Join(nodes, " "),
|
||||
compute := "-"
|
||||
if len(r.Shards) > 0 {
|
||||
compute = r.Shards[0].Compute.Short()
|
||||
}
|
||||
rows = append(rows, []string{r.ID, r.Model, r.State, r.Pool, compute, strings.Join(nodes, " "),
|
||||
(time.Duration(r.AgeSec) * time.Second).Round(time.Second).String()})
|
||||
}
|
||||
if m.inPanel {
|
||||
keys = rowKeys
|
||||
}
|
||||
return "Replicas", "\n" + table([]string{"REPLICA", "MODEL", "STATE", "POOL", "ON", "AGE"}, rows, m.row, pw, m.inPanel), keys
|
||||
return "Replicas", "\n" + table([]string{"REPLICA", "MODEL", "STATE", "POOL", "COMPUTE", "ON", "AGE"}, rows, m.row, pw, m.inPanel), keys
|
||||
case cUsage:
|
||||
head := "\n Last 24 hours, by " + lipgloss.NewStyle().Bold(true).Render(usageGroupings[m.usageBy]) +
|
||||
sDim.Render(" (→ to group by "+usageGroupings[(m.usageBy+1)%len(usageGroupings)]+")") + "\n\n"
|
||||
@@ -1255,6 +1277,8 @@ func (m consoleModel) panel() (title, body, keys string) {
|
||||
return m.quotasPanel(pw)
|
||||
case cGuard:
|
||||
return m.guardPanel(pw)
|
||||
case cKB:
|
||||
return m.kbPanel(pw)
|
||||
case cAccounts, cKeys, cOrgs:
|
||||
return m.peoplePanel(pw)
|
||||
case cManual:
|
||||
@@ -1299,6 +1323,9 @@ func (m consoleModel) panel() (title, body, keys string) {
|
||||
item = "\n Moves the store onto two machines under Patroni, replicated,\n so losing one loses nothing.\n\n " +
|
||||
sDim.Render("On the command line:") +
|
||||
"\n sudo --preserve-env=SSH_AUTH_SOCK llmgw server init \\\n --patroni <user>@<a> --patroni <user>@<b> --no-mirror --no-node --no-model"
|
||||
case cJudge:
|
||||
item = "\n Serves the model the guardrail policies ask as their judge — Llama\n Guard 3 1B on CPUs, under the name the policies give it — so what they\n judge is not refused (guard_unavailable). A model of that name that is\n not a judge is replaced. Overview says when a judge is missing.\n\n " +
|
||||
sDim.Render("On the command line:") + "\n llmgw ctl guard judge status\n llmgw ctl guard judge enable --replace"
|
||||
case cUpgrade:
|
||||
item = "\n Moves the cluster to a release: this machine first, then each node\n in turn. Engines keep serving; a node that does not come back stops it.\n\n " +
|
||||
sDim.Render("On the command line:") + "\n sudo llmgw ctl upgrade VERSION"
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
// Knowledge bases, as the console shows them: each collection with its
|
||||
// owner, model, documents and grants; Enter for what can be done to one, A
|
||||
// to make one. People fill them from their own machines with llmgw-kb.
|
||||
|
||||
type kbState struct {
|
||||
view adminapi.KBView
|
||||
err error
|
||||
// loaded is set once an answer came, so "none yet" is not shown while
|
||||
// asking.
|
||||
loaded bool
|
||||
}
|
||||
|
||||
type consoleKBMsg struct {
|
||||
view adminapi.KBView
|
||||
err error
|
||||
}
|
||||
|
||||
func (m consoleModel) fetchKB() tea.Cmd {
|
||||
c := m.client
|
||||
return func() tea.Msg {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var msg consoleKBMsg
|
||||
msg.err = c.get(ctx, "/api/v1/kb", &msg.view)
|
||||
return msg
|
||||
}
|
||||
}
|
||||
|
||||
func (m consoleModel) kbPanel(pw int) (string, string, string) {
|
||||
keys := "→ or Enter open · ↑↓ choose · Q leave"
|
||||
if m.inPanel {
|
||||
keys = "↑↓ collection · Enter actions · A new collection · ← back · Q leave"
|
||||
}
|
||||
k := m.kb
|
||||
var b strings.Builder
|
||||
b.WriteString("\n " + sDim.Render("Documents requests are answered from; people add them with llmgw-kb.") + "\n\n")
|
||||
switch {
|
||||
case k.err != nil:
|
||||
b.WriteString(" " + sBad.Render(k.err.Error()))
|
||||
return "Knowledge", b.String(), keys
|
||||
case !k.loaded:
|
||||
b.WriteString(" " + sDim.Render("asking the control plane…"))
|
||||
return "Knowledge", b.String(), keys
|
||||
case !k.view.Available:
|
||||
b.WriteString(" " + sBad.Render("Off: ") + k.view.Reason)
|
||||
return "Knowledge", b.String(), keys
|
||||
case len(k.view.Collections) == 0:
|
||||
b.WriteString(" " + sDim.Render("No collections yet. A to make one."))
|
||||
return "Knowledge", b.String(), keys
|
||||
}
|
||||
var rows [][]string
|
||||
for _, c := range k.view.Collections {
|
||||
var grants []string
|
||||
for _, g := range c.Grants {
|
||||
grants = append(grants, or(g.Name, g.Target)+" "+g.Perm)
|
||||
}
|
||||
rows = append(rows, []string{or(c.OrgName, c.Org) + "/" + c.Name, or(c.OwnerName, c.Owner),
|
||||
fmt.Sprintf("%s (%d)", c.EmbedModel, c.Dims), fmt.Sprint(c.Documents), fmt.Sprint(c.Chunks), dash(strings.Join(grants, ", "))})
|
||||
}
|
||||
b.WriteString(table([]string{"COLLECTION", "OWNER", "MODEL", "DOCUMENTS", "CHUNKS", "GRANTED TO"}, rows, m.row, pw, m.inPanel))
|
||||
if len(k.view.Defaults) > 0 {
|
||||
b.WriteString("\n\n " + sTitle.Render("BY DEFAULT") + " " + sDim.Render("when a request names no collection") + "\n")
|
||||
var drows [][]string
|
||||
for _, d := range k.view.Defaults {
|
||||
drows = append(drows, []string{or(d.Name, d.Target), strings.Join(d.Collections, ", ")})
|
||||
}
|
||||
b.WriteString(table([]string{"FOR", "ANSWERED FROM"}, drows, -1, pw, false))
|
||||
}
|
||||
b.WriteString("\n\n " + sDim.Render("On a person's machine: llmgw-kb sync COLLECTION FOLDER"))
|
||||
return "Knowledge", b.String(), keys
|
||||
}
|
||||
|
||||
// kbRef is how the admin API is told which collection: by its ID.
|
||||
func kbRef(c adminapi.KBCollectionView) string { return url.PathEscape(c.ID) }
|
||||
|
||||
// kbCreateForm makes a collection.
|
||||
func (m *consoleModel) kbCreateForm() {
|
||||
model := ""
|
||||
for _, md := range m.data.models {
|
||||
n := strings.ToLower(md.Name)
|
||||
if strings.Contains(n, "embed") || strings.Contains(n, "bge") || strings.Contains(n, "e5") {
|
||||
model = md.Name
|
||||
break
|
||||
}
|
||||
}
|
||||
m.openForm(&formDialog{title: "New knowledge-base collection",
|
||||
note: "Owned by part of the tree: it and everything below it may read and add to it.\n" +
|
||||
"The model must embed (capabilities.embeddings) and be able to serve now.",
|
||||
fields: []formField{
|
||||
choiceF("owner", []string{"project", "team", "org", "user"}, "team"),
|
||||
textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, or the ID"),
|
||||
textF("org", "", "for a person's collection only"),
|
||||
textF("name", "", "how people name it to llmgw-kb"),
|
||||
textF("model", model, "the catalog model that embeds it"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
for i, what := range map[int]string{1: "which " + v[0], 3: "a name", 4: "a model"} {
|
||||
if strings.TrimSpace(v[i]) == "" {
|
||||
return nil, fmt.Errorf("%s?", what)
|
||||
}
|
||||
}
|
||||
req := adminapi.KBCreateRequest{Owner: v[0] + ":" + v[1], Org: v[2], Name: v[3], Model: v[4]}
|
||||
return m.call("New collection "+v[3], http.MethodPost, "/api/v1/kb/collections", req, func(raw json.RawMessage) string {
|
||||
var c adminapi.KBCollectionView
|
||||
_ = json.Unmarshal(raw, &c)
|
||||
return fmt.Sprintf("Collection %s, embedded by %s in %d dimensions.\n\nPeople with access fill it with:\n llmgw-kb sync %s FOLDER",
|
||||
c.Name, c.EmbedModel, c.Dims, c.Name)
|
||||
}), nil
|
||||
}})
|
||||
}
|
||||
|
||||
// kbMenu is what can be done to one collection.
|
||||
func (m consoleModel) kbMenu(c adminapi.KBCollectionView) *menuDialog {
|
||||
name := or(c.OrgName, c.Org) + "/" + c.Name
|
||||
items := []menuItem{
|
||||
{"Documents", func(m *consoleModel) tea.Cmd { return m.kbDocuments(c) }},
|
||||
{"Grant…", func(m *consoleModel) tea.Cmd {
|
||||
m.openForm(&formDialog{title: "Grant " + name,
|
||||
note: "It and everything below it may then read, or with write add to, this collection.\nNever outside " + or(c.OrgName, c.Org) + ".",
|
||||
fields: []formField{
|
||||
choiceF("to", []string{"project", "team", "org", "user", "key"}, "project"),
|
||||
textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, or the ID"),
|
||||
choiceF("may", []string{"read", "write"}, "read"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
if strings.TrimSpace(v[1]) == "" {
|
||||
return nil, fmt.Errorf("which %s?", v[0])
|
||||
}
|
||||
req := adminapi.KBGrantRequest{Collection: c.ID, Target: v[0] + ":" + strings.TrimSpace(v[1]), Perm: v[2]}
|
||||
return m.call("Grant "+c.Name, http.MethodPost, "/api/v1/kb/grants", req, nil), nil
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
}
|
||||
for _, g := range c.Grants {
|
||||
items = append(items, menuItem{"Revoke " + or(g.Name, g.Target) + " (" + g.Perm + ")…", func(m *consoleModel) tea.Cmd {
|
||||
m.confirmD = &confirmDialog{title: "Revoke a grant on " + name,
|
||||
text: or(g.Name, g.Target) + " may then no longer " + g.Perm + " it, unless it owns it.",
|
||||
yes: func(m *consoleModel) tea.Cmd {
|
||||
return m.call("Revoke", http.MethodDelete, "/api/v1/kb/grants",
|
||||
adminapi.KBGrantRequest{Collection: c.ID, Target: g.Target}, nil)
|
||||
}}
|
||||
return nil
|
||||
}})
|
||||
}
|
||||
items = append(items,
|
||||
menuItem{"Answer from it by default…", func(m *consoleModel) tea.Cmd {
|
||||
m.openForm(&formDialog{title: "Answer from " + name + " by default",
|
||||
note: "Chat requests from there that name no collection are answered from this one.\n" +
|
||||
"Replaces that part's default; llmgw ctl kb default sets several.",
|
||||
fields: []formField{
|
||||
choiceF("for", []string{"project", "team", "org", "key"}, "team"),
|
||||
textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, or the ID"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
if strings.TrimSpace(v[1]) == "" {
|
||||
return nil, fmt.Errorf("which %s?", v[0])
|
||||
}
|
||||
req := adminapi.KBDefaultRequest{Target: v[0] + ":" + strings.TrimSpace(v[1]), Collections: []string{c.ID}}
|
||||
return m.call("Default "+c.Name, http.MethodPost, "/api/v1/kb/defaults", req, nil), nil
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
menuItem{"Try a search…", func(m *consoleModel) tea.Cmd {
|
||||
m.openForm(&formDialog{title: "Search " + name,
|
||||
note: "As an administrator: no tenancy applies. Shows what a question would find.",
|
||||
fields: []formField{textF("question", "", ""), textF("results", "5", "how many chunks")},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
if v[0] == "" {
|
||||
return nil, fmt.Errorf("what question?")
|
||||
}
|
||||
k, _ := strconv.Atoi(v[1])
|
||||
req := adminapi.KBSearchRequest{Collection: c.ID, Query: v[0], K: k}
|
||||
return m.call("Search "+c.Name, http.MethodPost, "/api/v1/kb/search", req, func(raw json.RawMessage) string {
|
||||
var hits []adminapi.KBHitView
|
||||
_ = json.Unmarshal(raw, &hits)
|
||||
if len(hits) == 0 {
|
||||
return "Nothing found."
|
||||
}
|
||||
var b strings.Builder
|
||||
for i, h := range hits {
|
||||
text := h.Text
|
||||
if r := []rune(text); len(r) > 400 {
|
||||
text = string(r[:400]) + "…"
|
||||
}
|
||||
fmt.Fprintf(&b, "%d. %s #%d\n%s\n\n", i+1, h.Document, h.Ord, indent(strings.TrimSpace(text), " "))
|
||||
}
|
||||
return strings.TrimRight(b.String(), "\n")
|
||||
}), nil
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
menuItem{"Delete…", func(m *consoleModel) tea.Cmd {
|
||||
m.confirmD = &confirmDialog{title: "Delete " + name,
|
||||
text: fmt.Sprintf("Its %d documents and %d chunks, and every grant, go with it. People's own files are not touched.",
|
||||
c.Documents, c.Chunks),
|
||||
yes: func(m *consoleModel) tea.Cmd {
|
||||
return m.call("Delete "+c.Name, http.MethodDelete, "/api/v1/kb/collections/"+kbRef(c), nil, nil)
|
||||
}}
|
||||
return nil
|
||||
}},
|
||||
)
|
||||
return &menuDialog{title: name, items: items}
|
||||
}
|
||||
|
||||
// kbDocuments shows a collection's documents in a dialog.
|
||||
func (m *consoleModel) kbDocuments(c adminapi.KBCollectionView) tea.Cmd {
|
||||
return m.call("Documents of "+c.Name, http.MethodGet, "/api/v1/kb/collections/"+kbRef(c)+"/documents", nil,
|
||||
func(raw json.RawMessage) string {
|
||||
var docs []adminapi.KBDocumentView
|
||||
_ = json.Unmarshal(raw, &docs)
|
||||
if len(docs) == 0 {
|
||||
return "None yet. On a person's machine: llmgw-kb sync " + c.Name + " FOLDER"
|
||||
}
|
||||
var b strings.Builder
|
||||
for i, d := range docs {
|
||||
if i == 40 {
|
||||
fmt.Fprintf(&b, "… and %d more (llmgw ctl kb docs %s)", len(docs)-40, c.ID)
|
||||
break
|
||||
}
|
||||
fmt.Fprintf(&b, "%-48s %5d chunks %s %s\n", d.Name, d.Chunks, d.AddedAt.Local().Format("2006-01-02 15:04"), d.AddedBy)
|
||||
}
|
||||
return strings.TrimRight(b.String(), "\n")
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
func TestACollectionIsMadeAndGrantedFromTheConsole(t *testing.T) {
|
||||
m, s := consoleWith(t, cKB)
|
||||
m.data.models = append(m.data.models, adminapi.ModelView{Name: "bge-m3"})
|
||||
m.kb = kbState{loaded: true, view: adminapi.KBView{Available: true, Collections: []adminapi.KBCollectionView{{
|
||||
ID: "kb_aaaaaaaaaaaaaaaa", Name: "policies", Org: "o1", OrgName: "Acme", Owner: "team:t1", OwnerName: "Acme / hr",
|
||||
EmbedModel: "bge-m3", Dims: 1024, Documents: 3, Chunks: 40,
|
||||
Grants: []adminapi.KBGrantView{{Target: "project:p9", Name: "Acme/sales/crm", Perm: "read"}},
|
||||
}}}}
|
||||
var tm tea.Model = m
|
||||
view := tm.View()
|
||||
for _, want := range []string{"Knowledge", "Acme/policies", "Acme / hr", "bge-m3 (1024)", "Acme/s", "llmgw-kb sync COLLECTION FOLDER"} {
|
||||
if !strings.Contains(view, want) {
|
||||
t.Errorf("missing %q in:\n%s", want, view)
|
||||
}
|
||||
}
|
||||
|
||||
// A: the form, prefilled with the embedding model.
|
||||
tm = press(tm, "a")
|
||||
f := tm.(consoleModel).formD
|
||||
if f == nil || f.fields[4].value() != "bge-m3" {
|
||||
t.Fatalf("no form, or the model not offered: %+v", f)
|
||||
}
|
||||
f.fields[1].text.SetValue("Acme/hr")
|
||||
f.fields[3].text.SetValue("handbook")
|
||||
tm = pressRun(tm, "tab", "tab", "tab", "tab", "tab", "enter")
|
||||
var req adminapi.KBCreateRequest
|
||||
got := s.last()
|
||||
if !strings.HasPrefix(got, "POST /api/v1/kb/collections ") ||
|
||||
json.Unmarshal([]byte(strings.SplitN(got, " ", 3)[2]), &req) != nil ||
|
||||
req.Owner != "team:Acme/hr" || req.Name != "handbook" || req.Model != "bge-m3" {
|
||||
t.Fatalf("sent %q", got)
|
||||
}
|
||||
|
||||
// Enter on a row: grant to a project for writing.
|
||||
m.formD = nil
|
||||
tm = m
|
||||
tm = press(tm, "enter")
|
||||
if md := tm.(consoleModel).menuD; md == nil || !strings.Contains(md.items[2].label, "Revoke Acme/sales/crm") {
|
||||
t.Fatalf("menu %+v", md)
|
||||
}
|
||||
tm = press(tm, "down", "enter")
|
||||
f = tm.(consoleModel).formD
|
||||
f.fields[1].text.SetValue("Acme/sales/crm")
|
||||
f.fields[2].opt = 1 // write
|
||||
tm = pressRun(tm, "tab", "tab", "tab", "enter")
|
||||
if got := s.last(); got != `POST /api/v1/kb/grants {"collection":"kb_aaaaaaaaaaaaaaaa","target":"project:Acme/sales/crm","perm":"write"}` {
|
||||
t.Fatalf("sent %q", got)
|
||||
}
|
||||
_ = tm
|
||||
}
|
||||
+53
-4
@@ -3,6 +3,7 @@ package cli
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
@@ -34,7 +35,7 @@ func newUsersCmd() *cobra.Command {
|
||||
"Stop an account working, keeping it and its history"),
|
||||
newUserStateCmd(open, "reinstate", store.UserActive,
|
||||
"Let a suspended account work again"),
|
||||
newUserCreateCmd(open), newUserResetPasswordCmd(open),
|
||||
newUserCreateCmd(open), newUserResetPasswordCmd(open), newUserResetLinkCmd(open, &sf),
|
||||
newUserJoinCmd(open), newUserLeaveCmd(open), newUserDeleteCmd(open), newUserForgetCmd(open),
|
||||
newUserAdminCmd(open, "promote", true, "Let an account see the cluster dashboard at /admin"),
|
||||
newUserAdminCmd(open, "demote", false, "Take away access to the cluster dashboard"),
|
||||
@@ -254,9 +255,10 @@ func newUserResetPasswordCmd(open userOpener) *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "reset-password <user>",
|
||||
Short: "Set a new password for someone who lost theirs",
|
||||
Long: "There is no self-service reset, because this gateway sends no email and\n" +
|
||||
"a reset link with nowhere to send it is not a reset. Set one here, tell\n" +
|
||||
"them out of band, and have them change it from their profile.\n\n" +
|
||||
Long: "People reset a forgotten password themselves at /forgot when the cluster\n" +
|
||||
"has mail settings (mail.smtp); without them, 'llmgw users reset-link' makes\n" +
|
||||
"a link to hand them. Or set one here, tell them out of band, and have them\n" +
|
||||
"change it from their profile.\n\n" +
|
||||
"Every session they have open is closed: if they lost the password, the\n" +
|
||||
"reason might be that somebody else has it.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
@@ -532,3 +534,50 @@ func newUserForgetCmd(open userOpener) *cobra.Command {
|
||||
cmd.Flags().BoolVar(&yes, "yes", false, "confirm")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newUserResetLinkCmd(open userOpener, sf *storeFlags) *cobra.Command {
|
||||
var base string
|
||||
cmd := &cobra.Command{
|
||||
Use: "reset-link <user>",
|
||||
Short: "Make a one-time link for someone to set a forgotten password",
|
||||
Long: "For a gateway that sends no mail, or a person whose mail does not arrive:\n" +
|
||||
"the link is printed, to hand them however suits. It works once, for an\n" +
|
||||
"hour, and a newer one replaces it. Setting the password closes every\n" +
|
||||
"session they have open. The address it points at is mail.base_url from\n" +
|
||||
"the cluster's settings, or --base-url.",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if base == "" {
|
||||
if cfg, err := loadEffectiveConfig(sf.config); err == nil {
|
||||
base = cfg.Mail.BaseURL
|
||||
}
|
||||
}
|
||||
if base == "" {
|
||||
return errors.New("give --base-url, the address people reach this gateway at, or set mail.base_url " +
|
||||
"(llmgw ctl config set mail.base_url https://…)")
|
||||
}
|
||||
s, err := open(cmd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
u, err := resolveUser(cmd.Context(), s, args[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if u.PasswordHash == "" {
|
||||
return fmt.Errorf("%s signs in through single sign-on: there is no password to set", u.Email)
|
||||
}
|
||||
token, err := s.CreatePasswordReset(cmd.Context(), u.ID, account.ResetLinkTTL)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s/reset?token=%s\n", strings.TrimRight(base, "/"), url.QueryEscape(token))
|
||||
cmd.PrintErrf("For %s: works once, for %s. Hand it to them alone; it opens their account.\n",
|
||||
u.Email, account.ResetLinkTTL)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&base, "base-url", "", "the address people reach this gateway at (default: mail.base_url)")
|
||||
return cmd
|
||||
}
|
||||
@@ -0,0 +1,529 @@
|
||||
// Package comfy has images made by ComfyUI: a workflow filled in with what
|
||||
// was asked for, queued, waited on, and its pictures fetched.
|
||||
//
|
||||
// ComfyUI speaks no OpenAI. A workflow is a graph of nodes that is queued with
|
||||
// POST /prompt, is done when GET /history/{id} lists it, and leaves its
|
||||
// pictures to be read one by one from GET /view. What the graph is depends on
|
||||
// the model — FLUX schnell is not FLUX dev, and neither is SDXL — so the graph
|
||||
// is a template the catalog names, with placeholders where a request's
|
||||
// prompt, size and seed go.
|
||||
package comfy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"embed"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
//go:embed workflows/*.json
|
||||
var builtin embed.FS
|
||||
|
||||
// DefaultWorkflow is the template used when the catalog names none.
|
||||
const DefaultWorkflow = "flux-schnell"
|
||||
|
||||
// Builtins are the workflows that come with llmgw, by name.
|
||||
func Builtins() []string {
|
||||
entries, _ := builtin.ReadDir("workflows")
|
||||
out := make([]string, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
out = append(out, strings.TrimSuffix(e.Name(), ".json"))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Template is a workflow in ComfyUI's API format, with placeholders.
|
||||
//
|
||||
// A placeholder is a string value "{{name}}", or "{{name|default}}" for one a
|
||||
// request may leave out. One that is the whole value becomes the value
|
||||
// itself, a number where the node wants a number; one inside a longer string
|
||||
// is written into it. The names are those of Params.
|
||||
type Template struct {
|
||||
graph map[string]any
|
||||
}
|
||||
|
||||
// Load reads a template: a built-in's name, or the workflow's JSON itself.
|
||||
func Load(nameOrJSON string) (*Template, error) {
|
||||
s := strings.TrimSpace(nameOrJSON)
|
||||
if s == "" {
|
||||
s = DefaultWorkflow
|
||||
}
|
||||
raw := []byte(s)
|
||||
if !strings.HasPrefix(s, "{") {
|
||||
b, err := builtin.ReadFile("workflows/" + s + ".json")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("no built-in workflow %q (there are %s)", s, strings.Join(Builtins(), ", "))
|
||||
}
|
||||
raw = b
|
||||
}
|
||||
var g map[string]any
|
||||
if err := json.Unmarshal(raw, &g); err != nil {
|
||||
return nil, fmt.Errorf("the workflow is not JSON: %w", err)
|
||||
}
|
||||
if len(g) == 0 {
|
||||
return nil, errors.New("the workflow has no nodes")
|
||||
}
|
||||
for id, n := range g {
|
||||
node, ok := n.(map[string]any)
|
||||
if !ok || node["class_type"] == nil {
|
||||
// The UI's own save format has a "nodes" list instead, and is
|
||||
// not something /prompt accepts.
|
||||
return nil, fmt.Errorf("node %q has no class_type: the workflow must be saved in API format", id)
|
||||
}
|
||||
}
|
||||
return &Template{graph: g}, nil
|
||||
}
|
||||
|
||||
// Params are what a request decides. A zero number is one not given, which
|
||||
// leaves it to the template's default.
|
||||
type Params struct {
|
||||
Checkpoint string
|
||||
Prompt string
|
||||
NegativePrompt string
|
||||
Width, Height int
|
||||
N int
|
||||
Seed int64
|
||||
Steps int
|
||||
Guidance float64
|
||||
}
|
||||
|
||||
func (p Params) value(name string) (any, bool) {
|
||||
switch name {
|
||||
case "checkpoint":
|
||||
return p.Checkpoint, p.Checkpoint != ""
|
||||
case "prompt":
|
||||
return p.Prompt, true
|
||||
case "negative_prompt":
|
||||
return p.NegativePrompt, true
|
||||
case "width":
|
||||
return p.Width, p.Width > 0
|
||||
case "height":
|
||||
return p.Height, p.Height > 0
|
||||
case "n":
|
||||
return max(p.N, 1), true
|
||||
case "seed":
|
||||
return p.Seed, true
|
||||
case "steps":
|
||||
return p.Steps, p.Steps > 0
|
||||
case "guidance":
|
||||
return p.Guidance, p.Guidance > 0
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// Fill is the workflow for one request.
|
||||
func (t *Template) Fill(p Params) (map[string]any, error) {
|
||||
var missing []string
|
||||
var fill func(v any) any
|
||||
fill = func(v any) any {
|
||||
switch v := v.(type) {
|
||||
case map[string]any:
|
||||
out := make(map[string]any, len(v))
|
||||
for k, x := range v {
|
||||
out[k] = fill(x)
|
||||
}
|
||||
return out
|
||||
case []any:
|
||||
out := make([]any, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = fill(x)
|
||||
}
|
||||
return out
|
||||
case string:
|
||||
return fillString(v, p, &missing)
|
||||
}
|
||||
return v
|
||||
}
|
||||
g := fill(t.graph).(map[string]any)
|
||||
if len(missing) > 0 {
|
||||
slices.Sort(missing)
|
||||
return nil, fmt.Errorf("the workflow wants %s, which the request did not give and the workflow has no default for",
|
||||
strings.Join(slices.Compact(missing), ", "))
|
||||
}
|
||||
return g, nil
|
||||
}
|
||||
|
||||
func fillString(s string, p Params, missing *[]string) any {
|
||||
if !strings.Contains(s, "{{") {
|
||||
return s
|
||||
}
|
||||
resolve := func(ph string) (any, bool) {
|
||||
name, def, hasDef := strings.Cut(ph, "|")
|
||||
name = strings.TrimSpace(name)
|
||||
if v, ok := p.value(name); ok {
|
||||
return v, true
|
||||
}
|
||||
if hasDef {
|
||||
def = strings.TrimSpace(def)
|
||||
if f, err := strconv.ParseFloat(def, 64); err == nil {
|
||||
return f, true
|
||||
}
|
||||
return def, true
|
||||
}
|
||||
*missing = append(*missing, name)
|
||||
return nil, false
|
||||
}
|
||||
// The whole value: what the node gets is the value, not its text.
|
||||
if strings.HasPrefix(s, "{{") && strings.HasSuffix(s, "}}") && strings.Count(s, "{{") == 1 {
|
||||
v, _ := resolve(s[2 : len(s)-2])
|
||||
return v
|
||||
}
|
||||
var b strings.Builder
|
||||
for {
|
||||
i := strings.Index(s, "{{")
|
||||
j := strings.Index(s, "}}")
|
||||
if i < 0 || j < i {
|
||||
b.WriteString(s)
|
||||
return b.String()
|
||||
}
|
||||
b.WriteString(s[:i])
|
||||
if v, ok := resolve(s[i+2 : j]); ok {
|
||||
fmt.Fprint(&b, v)
|
||||
}
|
||||
s = s[j+2:]
|
||||
}
|
||||
}
|
||||
|
||||
// RandomSeed is a seed for a request that named none. ComfyUI's are 64-bit,
|
||||
// but JSON numbers are read as doubles on the way: kept to 2^53.
|
||||
func RandomSeed() int64 {
|
||||
n, err := rand.Int(rand.Reader, big.NewInt(1<<53))
|
||||
if err != nil {
|
||||
return time.Now().UnixNano() & (1<<53 - 1)
|
||||
}
|
||||
return n.Int64()
|
||||
}
|
||||
|
||||
// Image is one picture made.
|
||||
type Image struct {
|
||||
Data []byte
|
||||
ContentType string
|
||||
}
|
||||
|
||||
// Error is a failure, said as what it means for the request.
|
||||
type Error struct {
|
||||
Msg string
|
||||
// Retryable is a failure of this server, which another one might not
|
||||
// have: it could not be reached, or it answered 5xx.
|
||||
Retryable bool
|
||||
}
|
||||
|
||||
func (e *Error) Error() string { return e.Msg }
|
||||
|
||||
// Client queues workflows and waits for them.
|
||||
type Client struct {
|
||||
HTTP *http.Client
|
||||
// Poll is how often a queued workflow is asked after.
|
||||
Poll time.Duration
|
||||
}
|
||||
|
||||
// Generate runs a workflow on the ComfyUI at base and returns its pictures,
|
||||
// in the order of the nodes that made them. A context that ends takes the
|
||||
// workflow out of the queue, or interrupts it if it is running, so a client
|
||||
// that went away does not leave the card working for nobody.
|
||||
func (c *Client) Generate(ctx context.Context, base, apiKey string, workflow map[string]any) ([]Image, error) {
|
||||
base = strings.TrimRight(base, "/")
|
||||
client := make([]byte, 8)
|
||||
_, _ = rand.Read(client)
|
||||
body, err := json.Marshal(map[string]any{"prompt": workflow, "client_id": "llmgw-" + hex.EncodeToString(client)})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var queued struct {
|
||||
PromptID string `json:"prompt_id"`
|
||||
Error json.RawMessage `json:"error"`
|
||||
NodeErrors json.RawMessage `json:"node_errors"`
|
||||
}
|
||||
status, err := c.do(ctx, http.MethodPost, base+"/prompt", apiKey, body, &queued)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status != http.StatusOK || queued.PromptID == "" {
|
||||
// A workflow ComfyUI will not run: a checkpoint it does not have, a
|
||||
// node it does not know. The catalog's to fix, not the client's, and
|
||||
// no other replica set up the same way would take it either.
|
||||
return nil, &Error{Msg: "ComfyUI refused the workflow: " + refusal(queued.Error, queued.NodeErrors)}
|
||||
}
|
||||
id := queued.PromptID
|
||||
|
||||
poll := c.Poll
|
||||
if poll <= 0 {
|
||||
poll = 250 * time.Millisecond
|
||||
}
|
||||
t := time.NewTicker(poll)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
c.cancel(context.WithoutCancel(ctx), base, apiKey, id)
|
||||
return nil, ctx.Err()
|
||||
case <-t.C:
|
||||
}
|
||||
var hist map[string]historyEntry
|
||||
if _, err := c.do(ctx, http.MethodGet, base+"/history/"+url.PathEscape(id), apiKey, nil, &hist); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
continue // said by the select
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
h, ok := hist[id]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if h.Status.StatusStr == "error" {
|
||||
return nil, &Error{Msg: "the workflow failed in ComfyUI: " + h.failure()}
|
||||
}
|
||||
if !h.Status.Completed {
|
||||
continue
|
||||
}
|
||||
return c.fetch(ctx, base, apiKey, h)
|
||||
}
|
||||
}
|
||||
|
||||
type historyEntry struct {
|
||||
Status struct {
|
||||
StatusStr string `json:"status_str"`
|
||||
Completed bool `json:"completed"`
|
||||
Messages [][]json.RawMessage `json:"messages"`
|
||||
} `json:"status"`
|
||||
Outputs map[string]struct {
|
||||
Images []struct {
|
||||
Filename string `json:"filename"`
|
||||
Subfolder string `json:"subfolder"`
|
||||
Type string `json:"type"`
|
||||
} `json:"images"`
|
||||
} `json:"outputs"`
|
||||
}
|
||||
|
||||
// failure is what ComfyUI said went wrong, from its execution_error.
|
||||
func (h historyEntry) failure() string {
|
||||
for _, m := range h.Status.Messages {
|
||||
if len(m) < 2 || string(m[0]) != `"execution_error"` {
|
||||
continue
|
||||
}
|
||||
var e struct {
|
||||
NodeType string `json:"node_type"`
|
||||
Message string `json:"exception_message"`
|
||||
}
|
||||
if json.Unmarshal(m[1], &e) == nil && e.Message != "" {
|
||||
return strings.TrimSpace(e.NodeType + ": " + strings.TrimSpace(e.Message))
|
||||
}
|
||||
}
|
||||
return "no reason given"
|
||||
}
|
||||
|
||||
func (c *Client) fetch(ctx context.Context, base, apiKey string, h historyEntry) ([]Image, error) {
|
||||
nodes := make([]string, 0, len(h.Outputs))
|
||||
for id := range h.Outputs {
|
||||
nodes = append(nodes, id)
|
||||
}
|
||||
// Numerically where they are numbers, which is how ComfyUI names them.
|
||||
slices.SortFunc(nodes, func(a, b string) int {
|
||||
x, ea := strconv.Atoi(a)
|
||||
y, eb := strconv.Atoi(b)
|
||||
if ea == nil && eb == nil {
|
||||
return x - y
|
||||
}
|
||||
return strings.Compare(a, b)
|
||||
})
|
||||
var out []Image
|
||||
for _, n := range nodes {
|
||||
for _, im := range h.Outputs[n].Images {
|
||||
q := url.Values{"filename": {im.Filename}, "subfolder": {im.Subfolder}, "type": {im.Type}}
|
||||
img, err := c.get(ctx, base+"/view?"+q.Encode(), apiKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, img)
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil, &Error{Msg: "the workflow finished without a picture: it has no SaveImage or PreviewImage node"}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// cancel takes a workflow out of the queue, or stops it if it is the one
|
||||
// running. Best effort: the server may be gone, which also stops it.
|
||||
func (c *Client) cancel(ctx context.Context, base, apiKey, id string) {
|
||||
ctx, done := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer done()
|
||||
body, _ := json.Marshal(map[string]any{"delete": []string{id}})
|
||||
_, _ = c.do(ctx, http.MethodPost, base+"/queue", apiKey, body, nil)
|
||||
var q struct {
|
||||
Running [][]json.RawMessage `json:"queue_running"`
|
||||
}
|
||||
if _, err := c.do(ctx, http.MethodGet, base+"/queue", apiKey, nil, &q); err != nil {
|
||||
return
|
||||
}
|
||||
for _, r := range q.Running {
|
||||
if len(r) > 1 && string(r[1]) == strconv.Quote(id) {
|
||||
_, _ = c.do(ctx, http.MethodPost, base+"/interrupt", apiKey, []byte("{}"), nil)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) client() *http.Client {
|
||||
if c.HTTP != nil {
|
||||
return c.HTTP
|
||||
}
|
||||
return http.DefaultClient
|
||||
}
|
||||
|
||||
func (c *Client) request(ctx context.Context, method, u, apiKey string, body []byte) (*http.Response, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, method, u, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
if apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
resp, err := c.client().Do(req)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
return nil, &Error{Msg: "ComfyUI could not be reached: " + err.Error(), Retryable: true}
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// do sends a request and reads a JSON answer into out. A 5xx is an error;
|
||||
// any other status is returned with whatever body could be read.
|
||||
func (c *Client) do(ctx context.Context, method, u, apiKey string, body []byte, out any) (int, error) {
|
||||
resp, err := c.request(ctx, method, u, apiKey, body)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(io.LimitReader(resp.Body, 16<<20))
|
||||
if err != nil {
|
||||
return 0, &Error{Msg: "reading ComfyUI's answer: " + err.Error(), Retryable: true}
|
||||
}
|
||||
if resp.StatusCode >= 500 {
|
||||
return resp.StatusCode, &Error{Msg: fmt.Sprintf("ComfyUI answered %d: %s", resp.StatusCode, snippet(b)), Retryable: true}
|
||||
}
|
||||
if out != nil && len(b) > 0 {
|
||||
if err := json.Unmarshal(b, out); err != nil && resp.StatusCode == http.StatusOK {
|
||||
return resp.StatusCode, &Error{Msg: "ComfyUI's answer is not JSON: " + snippet(b)}
|
||||
}
|
||||
}
|
||||
return resp.StatusCode, nil
|
||||
}
|
||||
|
||||
func (c *Client) get(ctx context.Context, u, apiKey string) (Image, error) {
|
||||
resp, err := c.request(ctx, http.MethodGet, u, apiKey, nil)
|
||||
if err != nil {
|
||||
return Image{}, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(io.LimitReader(resp.Body, 64<<20))
|
||||
if err != nil {
|
||||
return Image{}, &Error{Msg: "reading a picture from ComfyUI: " + err.Error(), Retryable: true}
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return Image{}, &Error{Msg: fmt.Sprintf("ComfyUI answered %d for a picture it made", resp.StatusCode)}
|
||||
}
|
||||
ct := resp.Header.Get("Content-Type")
|
||||
if ct == "" || ct == "application/octet-stream" {
|
||||
ct = http.DetectContentType(b)
|
||||
}
|
||||
if ct == "image/png" {
|
||||
b = withoutPNGText(b)
|
||||
}
|
||||
return Image{Data: b, ContentType: ct}, nil
|
||||
}
|
||||
|
||||
// withoutPNGText is a PNG without its text chunks. ComfyUI writes the whole
|
||||
// workflow into every picture it saves — the prompt, and the checkpoint and
|
||||
// nodes behind it — and what a picture says about how it was made is the
|
||||
// operator's to publish, not something to hand every client. A file that is
|
||||
// not a well-formed PNG is returned as it came.
|
||||
func withoutPNGText(b []byte) []byte {
|
||||
const sig = "\x89PNG\r\n\x1a\n"
|
||||
if len(b) < len(sig) || string(b[:len(sig)]) != sig {
|
||||
return b
|
||||
}
|
||||
out := make([]byte, 0, len(b))
|
||||
out = append(out, b[:len(sig)]...)
|
||||
for i := len(sig); i < len(b); {
|
||||
if i+12 > len(b) {
|
||||
return b
|
||||
}
|
||||
n := int(binary.BigEndian.Uint32(b[i:]))
|
||||
end := i + 12 + n
|
||||
if n < 0 || end > len(b) {
|
||||
return b
|
||||
}
|
||||
switch string(b[i+4 : i+8]) {
|
||||
case "tEXt", "iTXt", "zTXt":
|
||||
default:
|
||||
out = append(out, b[i:end]...)
|
||||
}
|
||||
i = end
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// refusal is what /prompt said when it would not queue a workflow.
|
||||
func refusal(e, nodes json.RawMessage) string {
|
||||
var top struct {
|
||||
Message string `json:"message"`
|
||||
Details string `json:"details"`
|
||||
}
|
||||
msg := snippet(e)
|
||||
if json.Unmarshal(e, &top) == nil && top.Message != "" {
|
||||
msg = top.Message
|
||||
if top.Details != "" {
|
||||
msg += ": " + top.Details
|
||||
}
|
||||
}
|
||||
var perNode map[string]struct {
|
||||
ClassType string `json:"class_type"`
|
||||
Errors []struct {
|
||||
Message string `json:"message"`
|
||||
Details string `json:"details"`
|
||||
} `json:"errors"`
|
||||
}
|
||||
if json.Unmarshal(nodes, &perNode) == nil {
|
||||
var parts []string
|
||||
for _, n := range perNode {
|
||||
for _, er := range n.Errors {
|
||||
parts = append(parts, strings.TrimSpace(n.ClassType+": "+er.Message+" "+er.Details))
|
||||
}
|
||||
}
|
||||
slices.Sort(parts)
|
||||
if len(parts) > 0 {
|
||||
msg += " (" + strings.Join(parts, "; ") + ")"
|
||||
}
|
||||
}
|
||||
if msg == "" {
|
||||
msg = "no reason given"
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
func snippet(b []byte) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > 300 {
|
||||
s = s[:300] + "…"
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
package comfy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestTheBuiltInWorkflowsFillIn(t *testing.T) {
|
||||
for _, name := range Builtins() {
|
||||
tpl, err := Load(name)
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
g, err := tpl.Fill(Params{Checkpoint: "flux.safetensors", Prompt: `a "red" fox`, Width: 768, Height: 512, N: 2, Seed: 42})
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
b, _ := json.Marshal(g)
|
||||
s := string(b)
|
||||
for _, want := range []string{`"ckpt_name":"flux.safetensors"`, `"text":"a \"red\" fox"`, `"width":768`,
|
||||
`"height":512`, `"batch_size":2`, `"seed":42`} {
|
||||
if !strings.Contains(s, want) {
|
||||
t.Errorf("%s: %s not in %s", name, want, s)
|
||||
}
|
||||
}
|
||||
if strings.Contains(s, "{{") {
|
||||
t.Errorf("%s: a placeholder was left: %s", name, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAStepCountGivenReplacesTheDefault(t *testing.T) {
|
||||
tpl, _ := Load("flux-schnell")
|
||||
g, _ := tpl.Fill(Params{Checkpoint: "c", Prompt: "p", Width: 512, Height: 512})
|
||||
if steps := g["5"].(map[string]any)["inputs"].(map[string]any)["steps"]; steps != 4.0 {
|
||||
t.Errorf("default steps = %v, want 4", steps)
|
||||
}
|
||||
g, _ = tpl.Fill(Params{Checkpoint: "c", Prompt: "p", Width: 512, Height: 512, Steps: 8})
|
||||
if steps := g["5"].(map[string]any)["inputs"].(map[string]any)["steps"]; steps != 8 {
|
||||
t.Errorf("steps = %v, want 8", steps)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPlaceholderWithoutAValueOrADefaultIsSaid(t *testing.T) {
|
||||
tpl, err := Load(`{"1":{"class_type":"X","inputs":{"a":"{{checkpoint}}","b":"size {{width}}x{{height}}"}}}`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := tpl.Fill(Params{Width: 5}); err == nil || !strings.Contains(err.Error(), "checkpoint, height") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
g, err := tpl.Fill(Params{Checkpoint: "c", Width: 5, Height: 6})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b := g["1"].(map[string]any)["inputs"].(map[string]any)["b"]; b != "size 5x6" {
|
||||
t.Errorf("b = %v", b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWorkflowInTheUIsFormatIsRefused(t *testing.T) {
|
||||
if _, err := Load(`{"nodes":[{"id":1}],"links":[]}`); err == nil || !strings.Contains(err.Error(), "API format") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
if _, err := Load("sdxl-turbo"); err == nil {
|
||||
t.Fatal("an unknown built-in was loaded")
|
||||
}
|
||||
}
|
||||
|
||||
// fakeComfy is just enough of ComfyUI: /prompt queues, /history answers
|
||||
// after a few asks, /view serves a picture.
|
||||
type fakeComfy struct {
|
||||
mu sync.Mutex
|
||||
asked int
|
||||
fail string
|
||||
refuse bool
|
||||
deleted []string
|
||||
queued map[string]any
|
||||
readyAt int
|
||||
srv *httptest.Server
|
||||
interrupted bool
|
||||
}
|
||||
|
||||
func newFakeComfy(t *testing.T) *fakeComfy {
|
||||
f := &fakeComfy{readyAt: 2}
|
||||
f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
switch {
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/prompt":
|
||||
if f.refuse {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
fmt.Fprint(w, `{"error":{"message":"Prompt outputs failed validation","details":""},"node_errors":{"1":{"class_type":"CheckpointLoaderSimple","errors":[{"message":"Value not in list","details":"ckpt_name: 'x' not in []"}]}}}`)
|
||||
return
|
||||
}
|
||||
var b struct {
|
||||
Prompt map[string]any `json:"prompt"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&b)
|
||||
f.queued = b.Prompt
|
||||
fmt.Fprint(w, `{"prompt_id":"p1","number":1,"node_errors":{}}`)
|
||||
case r.URL.Path == "/history/p1":
|
||||
f.asked++
|
||||
if f.asked < f.readyAt {
|
||||
fmt.Fprint(w, `{}`)
|
||||
return
|
||||
}
|
||||
if f.fail != "" {
|
||||
fmt.Fprintf(w, `{"p1":{"status":{"status_str":"error","completed":false,"messages":[["execution_start",{}],["execution_error",{"node_type":"KSampler","exception_message":%q}]]},"outputs":{}}}`, f.fail)
|
||||
return
|
||||
}
|
||||
fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true,"messages":[]},"outputs":{"7":{"images":[{"filename":"a.png","subfolder":"","type":"temp"},{"filename":"b.png","subfolder":"","type":"temp"}]}}}}`)
|
||||
case r.URL.Path == "/view":
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
fmt.Fprint(w, "PNG:"+r.URL.Query().Get("filename"))
|
||||
case r.URL.Path == "/queue" && r.Method == http.MethodPost:
|
||||
var b struct {
|
||||
Delete []string `json:"delete"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&b)
|
||||
f.deleted = append(f.deleted, b.Delete...)
|
||||
case r.URL.Path == "/queue":
|
||||
fmt.Fprint(w, `{"queue_running":[[1,"p1",{},{},[]]],"queue_pending":[]}`)
|
||||
case r.URL.Path == "/interrupt":
|
||||
f.interrupted = true
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(f.srv.Close)
|
||||
return f
|
||||
}
|
||||
|
||||
func TestGenerateWaitsForTheWorkflowAndFetchesItsPictures(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
c := &Client{Poll: time.Millisecond}
|
||||
images, err := c.Generate(t.Context(), f.srv.URL, "", map[string]any{"1": map[string]any{"class_type": "X"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(images) != 2 || string(images[0].Data) != "PNG:a.png" || images[1].ContentType != "image/png" {
|
||||
t.Fatalf("images = %+v", images)
|
||||
}
|
||||
if f.queued["1"] == nil {
|
||||
t.Error("the workflow was not what was queued")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWorkflowThatFailsSaysWhy(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
f.fail = "CUDA out of memory"
|
||||
_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
|
||||
if err == nil || !strings.Contains(err.Error(), "KSampler: CUDA out of memory") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWorkflowComfyUIRefusesIsNotRetried(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
f.refuse = true
|
||||
_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
|
||||
ce, ok := err.(*Error)
|
||||
if !ok || ce.Retryable || !strings.Contains(ce.Msg, "Value not in list") {
|
||||
t.Fatalf("err = %#v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAComfyUIThatCannotBeReachedIsRetryable(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
f.srv.Close()
|
||||
_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
|
||||
if ce, ok := err.(*Error); !ok || !ce.Retryable {
|
||||
t.Fatalf("err = %#v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAClientThatLeavesTakesItsWorkflowOffTheCard(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
f.readyAt = 1 << 30
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
_, err := (&Client{Poll: time.Millisecond}).Generate(ctx, f.srv.URL, "", map[string]any{})
|
||||
if err == nil {
|
||||
t.Fatal("no error")
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
if len(f.deleted) != 1 || f.deleted[0] != "p1" || !f.interrupted {
|
||||
t.Errorf("deleted %v, interrupted %v", f.deleted, f.interrupted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPictureLosesTheWorkflowComfyUIWroteIntoIt(t *testing.T) {
|
||||
chunk := func(kind, data string) []byte {
|
||||
b := make([]byte, 0, 12+len(data))
|
||||
b = binary.BigEndian.AppendUint32(b, uint32(len(data)))
|
||||
b = append(b, kind...)
|
||||
b = append(b, data...)
|
||||
return append(b, 0, 0, 0, 0) // a CRC nobody here checks
|
||||
}
|
||||
png := []byte("\x89PNG\r\n\x1a\n")
|
||||
png = append(png, chunk("IHDR", "0123456789abc")...)
|
||||
png = append(png, chunk("tEXt", "prompt\x00{\"1\":...}")...)
|
||||
png = append(png, chunk("IDAT", "pixels")...)
|
||||
png = append(png, chunk("iTXt", "workflow\x00...")...)
|
||||
png = append(png, chunk("IEND", "")...)
|
||||
got := string(withoutPNGText(png))
|
||||
if strings.Contains(got, "prompt") || strings.Contains(got, "workflow") {
|
||||
t.Errorf("text left in: %q", got)
|
||||
}
|
||||
for _, kind := range []string{"IHDR", "IDAT", "pixels", "IEND"} {
|
||||
if !strings.Contains(got, kind) {
|
||||
t.Errorf("%s gone: %q", kind, got)
|
||||
}
|
||||
}
|
||||
if broken := png[:30]; string(withoutPNGText(broken)) != string(broken) {
|
||||
t.Error("a truncated PNG was changed")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "{{checkpoint}}"}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{prompt}}", "clip": ["1", 1]}},
|
||||
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{negative_prompt}}", "clip": ["1", 1]}},
|
||||
"8": {"class_type": "FluxGuidance", "inputs": {"guidance": "{{guidance|3.5}}", "conditioning": ["2", 0]}},
|
||||
"4": {"class_type": "EmptySD3LatentImage", "inputs": {"width": "{{width}}", "height": "{{height}}", "batch_size": "{{n}}"}},
|
||||
"5": {"class_type": "KSampler", "inputs": {
|
||||
"seed": "{{seed}}", "steps": "{{steps|20}}", "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1,
|
||||
"model": ["1", 0], "positive": ["8", 0], "negative": ["3", 0], "latent_image": ["4", 0]}},
|
||||
"6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
|
||||
"7": {"class_type": "PreviewImage", "inputs": {"images": ["6", 0]}}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "{{checkpoint}}"}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{prompt}}", "clip": ["1", 1]}},
|
||||
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{negative_prompt}}", "clip": ["1", 1]}},
|
||||
"4": {"class_type": "EmptySD3LatentImage", "inputs": {"width": "{{width}}", "height": "{{height}}", "batch_size": "{{n}}"}},
|
||||
"5": {"class_type": "KSampler", "inputs": {
|
||||
"seed": "{{seed}}", "steps": "{{steps|4}}", "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1,
|
||||
"model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["4", 0]}},
|
||||
"6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
|
||||
"7": {"class_type": "PreviewImage", "inputs": {"images": ["6", 0]}}
|
||||
}
|
||||
@@ -18,7 +18,12 @@ import (
|
||||
//
|
||||
// An empty path loads the defaults, which with --server and --node-id on the
|
||||
// command line is enough for a smoke test on a node with no IB.
|
||||
func LoadAgent(path string) (agent.Config, error) {
|
||||
func LoadAgent(path string) (agent.Config, error) { return LoadAgentWith(path, nil) }
|
||||
|
||||
// LoadAgentWith is LoadAgent with the cluster's settings for this node
|
||||
// (agents.default and agents.nodes.<id>, as YAML) laid over the file: the
|
||||
// file is what the agent needs to start, the settings what it runs on.
|
||||
func LoadAgentWith(path string, settings []byte) (agent.Config, error) {
|
||||
cfg := agent.DefaultConfig()
|
||||
k := koanf.New(".")
|
||||
|
||||
@@ -30,6 +35,11 @@ func LoadAgent(path string) (agent.Config, error) {
|
||||
return cfg, fmt.Errorf("parsing %s: %w", path, err)
|
||||
}
|
||||
}
|
||||
if len(settings) > 0 {
|
||||
if err := k.Load(bytesProvider(settings), yaml.Parser()); err != nil {
|
||||
return cfg, fmt.Errorf("parsing the cluster's settings for this agent: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
err := k.Load(env.Provider(".", env.Opt{
|
||||
Prefix: "LLMGW_",
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"sort"
|
||||
)
|
||||
|
||||
// AgentSettings are the nodes' agent settings, kept with the cluster's
|
||||
// (BACKLOG E3): everything agent.yaml can say except what an agent needs to
|
||||
// start — its id, the control planes to dial, its address and its
|
||||
// certificate. Default applies to every node, and Nodes, by node id, over
|
||||
// it; both over the node's own agent.yaml.
|
||||
type AgentSettings struct {
|
||||
// Default is agent.yaml's keys for every node.
|
||||
Default map[string]any `koanf:"default"`
|
||||
// Nodes are agent.yaml's keys for one node each, by node id, over
|
||||
// Default.
|
||||
Nodes map[string]map[string]any `koanf:"nodes"`
|
||||
}
|
||||
|
||||
// agentBootstrapKeys are what an agent needs before it can be told anything:
|
||||
// they stay in agent.yaml.
|
||||
var agentBootstrapKeys = []string{"node_id", "server", "servers", "hostname", "mgmt_addr", "tls"}
|
||||
|
||||
// For is the settings for one node: Default with the node's own over it.
|
||||
// Empty when there are none.
|
||||
func (a AgentSettings) For(node string) map[string]any {
|
||||
out := clone(stringKeys(a.Default))
|
||||
if n, ok := a.Nodes[node]; ok {
|
||||
merge(out, stringKeys(n))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Check refuses settings that set an agent's bootstrap keys, or that would
|
||||
// not load as an agent's configuration for some node.
|
||||
func (a AgentSettings) Check() error {
|
||||
check := func(where string, m map[string]any) error {
|
||||
for _, k := range agentBootstrapKeys {
|
||||
if _, ok := m[k]; ok {
|
||||
return fmt.Errorf("%s.%s stays in each node's agent.yaml: an agent needs it before it can be told anything", where, k)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if err := check("agents.default", stringKeys(a.Default)); err != nil {
|
||||
return err
|
||||
}
|
||||
nodes := make([]string, 0, len(a.Nodes))
|
||||
for n := range a.Nodes {
|
||||
nodes = append(nodes, n)
|
||||
}
|
||||
sort.Strings(nodes)
|
||||
for _, n := range nodes {
|
||||
if err := check("agents.nodes."+n, stringKeys(a.Nodes[n])); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// Each node's settings as its agent would load them, over nothing: what
|
||||
// its own file adds cannot be known here, and what is wrong in the
|
||||
// settings alone is wrong on every node.
|
||||
for _, n := range append([]string{""}, nodes...) {
|
||||
b, err := Marshal(a.For(n))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := LoadAgentWith("", b); err != nil {
|
||||
where := "agents.default"
|
||||
if n != "" {
|
||||
where = "agents.nodes." + n
|
||||
}
|
||||
return fmt.Errorf("%s: %w", where, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AgentBootstrapKeys are the keys agent.yaml keeps, for saying so.
|
||||
func AgentBootstrapKeys() []string { return slices.Clone(agentBootstrapKeys) }
|
||||
|
||||
// stringKeys turns YAML's map[any]any, wherever it hides, into the
|
||||
// map[string]any the rest of this package merges.
|
||||
func stringKeys(m map[string]any) map[string]any {
|
||||
out := make(map[string]any, len(m))
|
||||
for k, v := range m {
|
||||
out[k] = normalize(v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalize(v any) any {
|
||||
switch t := v.(type) {
|
||||
case map[string]any:
|
||||
return stringKeys(t)
|
||||
case map[any]any:
|
||||
m := make(map[string]any, len(t))
|
||||
for k, vv := range t {
|
||||
m[fmt.Sprint(k)] = normalize(vv)
|
||||
}
|
||||
return m
|
||||
case []any:
|
||||
out := make([]any, len(t))
|
||||
for i, vv := range t {
|
||||
out[i] = normalize(vv)
|
||||
}
|
||||
return out
|
||||
}
|
||||
return v
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestANodesAgentSettingsAreTheDefaultWithItsOwnOver(t *testing.T) {
|
||||
a := AgentSettings{
|
||||
Default: map[string]any{"heartbeat_interval": "5s", "supervisor": map[string]any{"work_dir": "/srv/w", "x": 1}},
|
||||
Nodes: map[string]map[string]any{
|
||||
"cuda01": {"supervisor": map[string]any{"work_dir": "/fast/w"}, "labels": map[string]any{"gpu": "rtx6000"}},
|
||||
},
|
||||
}
|
||||
got := a.For("cuda01")
|
||||
sup := got["supervisor"].(map[string]any)
|
||||
if sup["work_dir"] != "/fast/w" || sup["x"] != 1 || got["heartbeat_interval"] != "5s" || got["labels"] == nil {
|
||||
t.Errorf("cuda01: %v", got)
|
||||
}
|
||||
if got := a.For("node01"); got["labels"] != nil || got["supervisor"].(map[string]any)["work_dir"] != "/srv/w" {
|
||||
t.Errorf("node01: %v", got)
|
||||
}
|
||||
if got := (AgentSettings{}).For("x"); len(got) != 0 {
|
||||
t.Errorf("none: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentSettingsKeepBootstrapKeysInTheFileAndMustLoad(t *testing.T) {
|
||||
for _, a := range []AgentSettings{
|
||||
{Default: map[string]any{"servers": []any{"x:7443"}}},
|
||||
{Nodes: map[string]map[string]any{"n1": {"node_id": "other"}}},
|
||||
{Nodes: map[string]map[string]any{"n1": {"tls": map[string]any{"ca": "/x"}}}},
|
||||
} {
|
||||
if err := a.Check(); err == nil || !strings.Contains(err.Error(), "stays in each node's agent.yaml") {
|
||||
t.Errorf("%+v: %v", a, err)
|
||||
}
|
||||
}
|
||||
// A fabric with no adapters would not load on the node it names.
|
||||
bad := AgentSettings{Nodes: map[string]map[string]any{"n1": {"fabrics": []any{map[string]any{"name": "ib-a"}}}}}
|
||||
if err := bad.Check(); err == nil || !strings.Contains(err.Error(), "agents.nodes.n1") || !strings.Contains(err.Error(), "no HCAs") {
|
||||
t.Errorf("a fabric without HCAs: %v", err)
|
||||
}
|
||||
if err := (AgentSettings{Default: map[string]any{"heartbeat_interval": "3s"}}).Check(); err != nil {
|
||||
t.Errorf("a plain setting: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheClusterSettingsCarryAgentsAndRefuseBadOnes(t *testing.T) {
|
||||
good, err := Parse([]byte("agents:\n default:\n heartbeat_interval: 3s\n nodes:\n cuda01:\n labels: {gpu: rtx6000}\n"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := FromMap(good)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.Agents.For("cuda01")["labels"] == nil {
|
||||
t.Errorf("agents not read: %+v", cfg.Agents)
|
||||
}
|
||||
bad, _ := Parse([]byte("agents:\n default:\n mgmt_addr: 10.0.0.1\n"))
|
||||
if _, err := FromMap(bad); err == nil {
|
||||
t.Error("a bootstrap key in the cluster's settings was taken")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheSettingsAreLaidOverTheAgentsFile(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "agent.yaml")
|
||||
_ = os.WriteFile(p, []byte("node_id: n1\nserver: cp:7443\nheartbeat_interval: 9s\nlabels: {rack: r1}\n"), 0o600)
|
||||
cfg, err := LoadAgentWith(p, []byte("heartbeat_interval: 2s\nlabels: {gpu: a100}\n"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.NodeID != "n1" || cfg.HeartbeatInterval != 2*time.Second || cfg.Labels["rack"] != "r1" || cfg.Labels["gpu"] != "a100" {
|
||||
t.Errorf("loaded %+v", cfg)
|
||||
}
|
||||
}
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/facility"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/mail"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/secrets"
|
||||
@@ -146,9 +147,24 @@ type Server struct {
|
||||
// Prometheus and Alertmanager in between. Empty sends nothing.
|
||||
Notify notify.Config `koanf:"notify"`
|
||||
|
||||
// Mail is where the gateway's few emails go out from: a link to set a
|
||||
// forgotten password. Unset, none is sent and people ask an operator.
|
||||
Mail mail.Config `koanf:"mail"`
|
||||
|
||||
// Serving stops and starts the whole cluster's service (llmgw stop, llmgw
|
||||
// start).
|
||||
Serving ServingSettings `koanf:"serving"`
|
||||
|
||||
// Firewall turns llmgw's firewall on every machine on or off.
|
||||
Firewall FirewallSettings `koanf:"firewall"`
|
||||
|
||||
// Agents are the nodes' agent settings: agent.yaml's keys but those an
|
||||
// agent needs to start, for every node (default) and for one
|
||||
// (nodes.<id>). Sent to the agents as they change; an agent whose
|
||||
// settings changed starts again on them, its engines serving
|
||||
// throughout.
|
||||
Agents AgentSettings `koanf:"agents"`
|
||||
|
||||
// FirewallNetworks are networks, besides those each machine reaches the
|
||||
// cluster on, that llmgw's internal ports answer: a monitoring host on
|
||||
// another subnet, nodes behind a router. Written into every machine's
|
||||
@@ -707,7 +723,10 @@ type Endpoint struct {
|
||||
Models []string `koanf:"models"`
|
||||
// Engine names the inference engine, which decides what its metrics are
|
||||
// called. Empty means vLLM's names.
|
||||
Engine string `koanf:"engine"`
|
||||
Engine string `koanf:"engine"`
|
||||
// Workflow is the ComfyUI workflow for engine: comfyui, a built-in's
|
||||
// name or its JSON; flux-schnell unless set.
|
||||
Workflow string `koanf:"workflow"`
|
||||
Pool string `koanf:"pool"`
|
||||
Weight float64 `koanf:"weight"`
|
||||
Priority int `koanf:"priority"`
|
||||
@@ -878,6 +897,9 @@ func load(b []byte) (Server, error) {
|
||||
// Validate rejects configurations that would otherwise fail confusingly at
|
||||
// request time.
|
||||
func (s Server) Validate() error {
|
||||
if err := s.Agents.Check(); err != nil {
|
||||
return err
|
||||
}
|
||||
for i, o := range s.EngineOverrides {
|
||||
if o.GPUClass == "" && len(o.Nodes) == 0 && len(o.Labels) == 0 {
|
||||
return fmt.Errorf("engine_overrides[%d] picks no nodes: give it a gpu_class, nodes or labels", i)
|
||||
@@ -914,6 +936,9 @@ func (s Server) Validate() error {
|
||||
return fmt.Errorf("backup_copy_to needs backup_recipients: an archive holds the certificate authority's key " +
|
||||
"and every secret in the store, and leaves this machine only encrypted")
|
||||
}
|
||||
if err := s.Mail.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := s.FirewallAllow(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
package config
|
||||
|
||||
// ServingSettings is whether the cluster serves at all.
|
||||
type ServingSettings struct {
|
||||
// Stopped takes every model down and answers every request for one with
|
||||
// 503 (system_stopped), leaving the control planes, agents and mirror
|
||||
// running and every model's definition as it is: llmgw stop sets it,
|
||||
// llmgw start clears it, and the models come back as their catalog
|
||||
// entries ask. In force on every control plane within seconds.
|
||||
Stopped bool `koanf:"stopped"`
|
||||
}
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"maps"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"slices"
|
||||
@@ -67,8 +68,16 @@ type API struct {
|
||||
// nil without a store.
|
||||
federation *federation.Service
|
||||
// guard is the guardrail policies' library and attachments.
|
||||
guard *GuardAdmin
|
||||
guard *GuardAdmin
|
||||
// kb is knowledge bases; nil without a store.
|
||||
kb *KBAdmin
|
||||
// tenancy is organizations, keys and people; nil without a store.
|
||||
tenancy TenancyStore
|
||||
// train is adapter training; nil without a store and agents.
|
||||
train *TrainAdmin
|
||||
settings *Settings
|
||||
// restartAgents tells every connected agent to start again.
|
||||
restartAgents func(reason string) []string
|
||||
// revoke records a removed node, so its certificate is refused.
|
||||
revoke func(ctx context.Context, nodeID string, at time.Time) error
|
||||
|
||||
@@ -111,8 +120,12 @@ func (a *API) Register(mux Mux) {
|
||||
mux.HandleFunc("POST /api/v1/models/{name}/scale", a.scaleModel)
|
||||
mux.HandleFunc("POST /api/v1/models/{name}/resources", a.setResources)
|
||||
a.registerQuotas(mux)
|
||||
a.registerKB(mux)
|
||||
a.registerTenancy(mux)
|
||||
a.registerTrain(mux)
|
||||
a.registerGuard(mux)
|
||||
a.registerSettings(mux)
|
||||
mux.HandleFunc("POST /api/v1/agents/restart", a.restartAgentsHandler)
|
||||
mux.HandleFunc("GET /api/v1/audit", a.auditEvents)
|
||||
mux.HandleFunc("GET /api/v1/usage/export", a.usageExport)
|
||||
mux.HandleFunc("POST /api/v1/notify/test", a.notifyTest)
|
||||
@@ -228,6 +241,17 @@ type NodeView struct {
|
||||
// which is not the same as having none — placement treats it that way too,
|
||||
// or a control plane upgrade would empty a fleet.
|
||||
Engines []string `json:"engines,omitempty"`
|
||||
// EngineVersions are the installed engines with their versions:
|
||||
// llamacpp → b10991. ComputeCapabilities are its GPUs' CUDA compute
|
||||
// capabilities, 7.0 for a V100, so an engine build's support can be read
|
||||
// against them.
|
||||
EngineVersions map[string]string `json:"engine_versions,omitempty"`
|
||||
// MultiNode is whether it can take part in replicas spanning nodes:
|
||||
// "ready", "down: mlx5_1" naming the declared adapters whose ports are
|
||||
// not active, or empty for a node on no fabric. Its own replicas run
|
||||
// either way.
|
||||
MultiNode string `json:"multi_node,omitempty"`
|
||||
ComputeCapabilities []string `json:"compute_capabilities,omitempty"`
|
||||
// Provision is where installing each wanted engine stands. An engine is in
|
||||
// Engines only once its installation here is ready.
|
||||
Provision []inventory.EngineInstall `json:"provision,omitempty"`
|
||||
@@ -416,10 +440,12 @@ type ModelView struct {
|
||||
Description string `json:"description,omitempty"`
|
||||
Aliases []string `json:"aliases,omitempty"`
|
||||
Source string `json:"source"`
|
||||
Min int `json:"min_replicas"`
|
||||
Max int `json:"max_replicas"`
|
||||
Ready int `json:"ready_replicas"`
|
||||
Total int `json:"total_replicas"`
|
||||
// Role is "guard" for a guardrails judge, empty for a model clients use.
|
||||
Role string `json:"role,omitempty"`
|
||||
Min int `json:"min_replicas"`
|
||||
Max int `json:"max_replicas"`
|
||||
Ready int `json:"ready_replicas"`
|
||||
Total int `json:"total_replicas"`
|
||||
// Desired is how many it should have now: Min, or for an on-demand
|
||||
// model one while it is being asked for and none otherwise.
|
||||
Desired int `json:"desired_replicas"`
|
||||
@@ -518,6 +544,30 @@ type ShardView struct {
|
||||
// node's last heartbeat; zero when not measured.
|
||||
CPUCores float64 `json:"cpu_cores,omitempty"`
|
||||
MemoryBytes int64 `json:"memory_bytes,omitempty"`
|
||||
// Compute is where the engine said it put the model as it became
|
||||
// ready: absent when it did not say.
|
||||
Compute *ComputeView `json:"compute,omitempty"`
|
||||
}
|
||||
|
||||
// ComputeView is where an engine runs its model, by its own account.
|
||||
type ComputeView struct {
|
||||
// Backend is CUDA, Vulkan, Metal, ROCm, GPU or CPU.
|
||||
Backend string `json:"backend"`
|
||||
Devices []string `json:"devices,omitempty"`
|
||||
// LayersOffloaded of LayersTotal went to the devices; absent when unsaid.
|
||||
LayersOffloaded int `json:"layers_offloaded,omitempty"`
|
||||
LayersTotal int `json:"layers_total,omitempty"`
|
||||
}
|
||||
|
||||
// Short is the compute in a few characters: "CUDA 49/49", "CPU".
|
||||
func (c *ComputeView) Short() string {
|
||||
if c == nil {
|
||||
return "-"
|
||||
}
|
||||
if c.LayersTotal > 0 && c.Backend != "CPU" {
|
||||
return fmt.Sprintf("%s %d/%d", c.Backend, c.LayersOffloaded, c.LayersTotal)
|
||||
}
|
||||
return c.Backend
|
||||
}
|
||||
|
||||
// FetchView is a download in flight.
|
||||
@@ -623,6 +673,16 @@ func (a *API) Nodes() []NodeView {
|
||||
v.Engines = append(v.Engines, name)
|
||||
}
|
||||
slices.Sort(v.Engines)
|
||||
if len(n.EngineVersions) > 0 {
|
||||
v.EngineVersions = maps.Clone(n.EngineVersions)
|
||||
}
|
||||
for _, g := range n.GPUs {
|
||||
if g.ComputeCapability != "" && !slices.Contains(v.ComputeCapabilities, g.ComputeCapability) {
|
||||
v.ComputeCapabilities = append(v.ComputeCapabilities, g.ComputeCapability)
|
||||
}
|
||||
}
|
||||
slices.Sort(v.ComputeCapabilities)
|
||||
v.MultiNode = multiNodeReadiness(n)
|
||||
v.Provision = n.Provision
|
||||
if ch := a.ctl.EngineChoices(); ch != nil && len(n.GPUs) > 0 {
|
||||
chosen := ch.For(n.ID)
|
||||
@@ -761,6 +821,7 @@ func (a *API) Models() []ModelView {
|
||||
Class: m.EffectiveClass(), Owner: m.Owner, Preempted: h.Preempted, HeldFor: h.HeldFor,
|
||||
External: externalLabel(m), Adapters: m.Adapters,
|
||||
Canary: m.Canary, CanaryOf: a.ctl.Catalog().CanaryOf(m.Name), Eval: m.Eval,
|
||||
Role: m.Role,
|
||||
}
|
||||
for _, p := range m.Profiles {
|
||||
v.Profiles = append(v.Profiles, p.Name)
|
||||
@@ -829,6 +890,10 @@ func (a *API) Replicas(model string) []ReplicaView {
|
||||
sv.Fetch.Percent = &p
|
||||
}
|
||||
}
|
||||
if c := s.Compute; c != nil {
|
||||
sv.Compute = &ComputeView{Backend: c.Backend, Devices: c.Devices,
|
||||
LayersOffloaded: c.Offloaded, LayersTotal: c.Total}
|
||||
}
|
||||
if l := loads[s.NodeID]; l != nil {
|
||||
if use, ok := l.Shards[s.ID]; ok {
|
||||
sv.CPUCores, sv.MemoryBytes = use.CPUCores, use.RSSBytes
|
||||
@@ -1473,6 +1538,11 @@ type RestartRequest struct {
|
||||
// InPlace stops each replica before starting its replacement, rather
|
||||
// than starting the replacement beside it first.
|
||||
InPlace bool `json:"in_place,omitempty"`
|
||||
// MaxSurge is how many of a model's replicas are replaced beside
|
||||
// themselves at once, MaxUnavailable how many are stopped first at once;
|
||||
// both zero is one at a time, replaced beside itself.
|
||||
MaxSurge int `json:"max_surge,omitempty"`
|
||||
MaxUnavailable int `json:"max_unavailable,omitempty"`
|
||||
}
|
||||
|
||||
// RestartResult lists the replicas queued.
|
||||
@@ -1496,11 +1566,12 @@ func (a *API) restart(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusBadRequest, "nothing ready matches; name replicas, models, or all")
|
||||
return
|
||||
}
|
||||
restart := a.ctl.Reconciler().Restart
|
||||
if req.InPlace {
|
||||
restart = a.ctl.Reconciler().RestartInPlace
|
||||
if req.MaxSurge < 0 || req.MaxUnavailable < 0 {
|
||||
writeError(w, http.StatusBadRequest, "max_surge and max_unavailable are counts of replicas, not negative")
|
||||
return
|
||||
}
|
||||
if err := restart(ids...); err != nil {
|
||||
pace := reconciler.RestartPace{MaxSurge: req.MaxSurge, MaxUnavailable: req.MaxUnavailable}
|
||||
if err := a.ctl.Reconciler().RestartPaced(pace, req.InPlace, ids...); err != nil {
|
||||
writeError(w, http.StatusNotFound, err.Error())
|
||||
return
|
||||
}
|
||||
@@ -1545,6 +1616,7 @@ type RequestView struct {
|
||||
GuardAction string `json:"guard_action,omitempty"`
|
||||
GuardPolicy string `json:"guard_policy,omitempty"`
|
||||
GuardCategory string `json:"guard_category,omitempty"`
|
||||
GuardRule string `json:"guard_rule,omitempty"`
|
||||
GuardMs int64 `json:"guard_ms,omitempty"`
|
||||
// What the client said about it: OpenAI's user field, and its tags.
|
||||
EndUser string `json:"end_user,omitempty"`
|
||||
@@ -1573,7 +1645,7 @@ func (a *API) request(w http.ResponseWriter, r *http.Request) {
|
||||
Model: u.Model, ServedModel: u.ServedModel, Replica: u.Replica, Status: u.Status,
|
||||
LatencyMs: u.LatencyMs, QueueMs: u.QueueMs, TTFTMs: u.TTFTMs,
|
||||
PromptTokens: u.PromptTokens, CompletionTokens: u.CompletionTokens, CostMicros: u.CostMicros,
|
||||
GuardAction: u.GuardAction, GuardPolicy: u.GuardPolicy, GuardCategory: u.GuardCategory, GuardMs: u.GuardMs,
|
||||
GuardAction: u.GuardAction, GuardPolicy: u.GuardPolicy, GuardCategory: u.GuardCategory, GuardRule: u.GuardRule, GuardMs: u.GuardMs,
|
||||
EndUser: u.EndUser, Tags: u.Tags, CacheHit: u.CacheHit,
|
||||
})
|
||||
}
|
||||
@@ -1897,32 +1969,9 @@ func (a *API) setAdapters(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusBadRequest, "no model "+name)
|
||||
return
|
||||
}
|
||||
base, baseErr := gguf.KindFromURL(r.Context(), safedial.Client(safedial.ControlPlane, 30*time.Second), m.Source)
|
||||
have := map[string]bool{}
|
||||
for _, ad := range m.Adapters {
|
||||
have[ad.Source] = true
|
||||
}
|
||||
// The file is read only where llama.cpp will load it: a vLLM adapter is
|
||||
// a PEFT directory, which has no header to read, and vLLM itself refuses
|
||||
// one that does not fit the model when it is loaded.
|
||||
ggufs := slices.ContainsFunc(m.Profiles, func(p catalog.Profile) bool { return p.Engine == "llamacpp" })
|
||||
for _, ad := range req.Adapters {
|
||||
if have[ad.Source] || !ggufs {
|
||||
continue // already checked when it was attached
|
||||
}
|
||||
h, err := gguf.KindFromURL(r.Context(), safedial.Client(safedial.ControlPlane, 30*time.Second), ad.Source)
|
||||
switch {
|
||||
case err != nil:
|
||||
writeError(w, http.StatusBadRequest, fmt.Sprintf("adapter %s: %v", ad.Name, err))
|
||||
return
|
||||
case !h.IsAdapter():
|
||||
writeError(w, http.StatusBadRequest, fmt.Sprintf("adapter %s: %s is a model, not a LoRA adapter", ad.Name, ad.Source))
|
||||
return
|
||||
case baseErr == nil && h.Architecture != base.Architecture:
|
||||
writeError(w, http.StatusBadRequest, fmt.Sprintf("adapter %s is for %s models and %s is %s",
|
||||
ad.Name, h.Architecture, m.Name, base.Architecture))
|
||||
return
|
||||
}
|
||||
if err := checkAdapters(r.Context(), m, req.Adapters); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
out, err := a.models.SetAdapters(r.Context(), m.Name, req.Adapters, actorOf(r))
|
||||
if err != nil {
|
||||
@@ -1975,3 +2024,65 @@ func (a *API) SetNodeEngine(ctx context.Context, node, engine, build, by string)
|
||||
}
|
||||
return ch.Set(ctx, node, engine, build, by)
|
||||
}
|
||||
|
||||
// multiNodeReadiness is whether a node's declared fabrics are up: every
|
||||
// adapter they name with an active port. Empty for a node on no fabric.
|
||||
func multiNodeReadiness(n inventory.Node) string {
|
||||
if len(n.Fabrics) == 0 {
|
||||
return ""
|
||||
}
|
||||
if len(n.IBPorts) == 0 {
|
||||
return "unknown: no port reported yet"
|
||||
}
|
||||
var down []string
|
||||
for _, f := range n.Fabrics {
|
||||
for _, hca := range f.HCAs {
|
||||
up := false
|
||||
for _, p := range n.IBPorts {
|
||||
if p.HCA == hca && p.Active() {
|
||||
up = true
|
||||
}
|
||||
}
|
||||
if !up && !slices.Contains(down, hca) {
|
||||
down = append(down, hca)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(down) > 0 {
|
||||
return "down: " + strings.Join(down, ", ")
|
||||
}
|
||||
return "ready"
|
||||
}
|
||||
|
||||
// checkAdapters reads each adapter a model does not have yet, where
|
||||
// llama.cpp will load it, to refuse one that is not a LoRA adapter for the
|
||||
// model's architecture: such an adapter is a replica that fails to start,
|
||||
// found minutes later on a node.
|
||||
func checkAdapters(ctx context.Context, m catalog.Model, adapters []catalog.Adapter) error {
|
||||
// A vLLM adapter is a PEFT directory, which has no header to read, and
|
||||
// vLLM itself refuses one that does not fit the model when it is loaded.
|
||||
if !slices.ContainsFunc(m.Profiles, func(p catalog.Profile) bool { return p.Engine == "llamacpp" }) {
|
||||
return nil
|
||||
}
|
||||
client := safedial.Client(safedial.ControlPlane, 30*time.Second)
|
||||
base, baseErr := gguf.KindFromURL(ctx, client, m.Source)
|
||||
have := map[string]bool{}
|
||||
for _, ad := range m.Adapters {
|
||||
have[ad.Source] = true
|
||||
}
|
||||
for _, ad := range adapters {
|
||||
if have[ad.Source] {
|
||||
continue // checked when it was attached
|
||||
}
|
||||
h, err := gguf.KindFromURL(ctx, client, ad.Source)
|
||||
switch {
|
||||
case err != nil:
|
||||
return fmt.Errorf("adapter %s: %v", ad.Name, err)
|
||||
case !h.IsAdapter():
|
||||
return fmt.Errorf("adapter %s: %s is a model, not a LoRA adapter", ad.Name, ad.Source)
|
||||
case baseErr == nil && h.Architecture != base.Architecture:
|
||||
return fmt.Errorf("adapter %s is for %s models and %s is %s", ad.Name, h.Architecture, m.Name, base.Architecture)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -28,6 +28,10 @@ type GuardAdmin struct {
|
||||
Library func() guard.Config
|
||||
Store GuardStore
|
||||
Invalidate func()
|
||||
// JudgeErrors are each judge's last failure to give a verdict.
|
||||
JudgeErrors func() map[string]guard.JudgeError
|
||||
// Ask has a judge read a text, past the cache.
|
||||
Ask func(ctx context.Context, model, text, answer string) (guard.JudgeAnswer, error)
|
||||
}
|
||||
|
||||
// SetGuard enables the guard endpoints.
|
||||
@@ -35,6 +39,8 @@ func (a *API) SetGuard(g *GuardAdmin) { a.guard = g }
|
||||
|
||||
// GuardView is the library and where its policies are attached.
|
||||
type GuardView struct {
|
||||
// Enabled is whether the guardrails check anything (guard.enabled).
|
||||
Enabled bool `json:"enabled"`
|
||||
Policies []string `json:"policies"`
|
||||
Default []string `json:"default"`
|
||||
Attachments []GuardAttachmentView `json:"attachments"`
|
||||
@@ -89,7 +95,10 @@ type GuardEventView struct {
|
||||
Action string `json:"action"`
|
||||
Policy string `json:"policy"`
|
||||
Category string `json:"category,omitempty"`
|
||||
GuardMs int64 `json:"guard_ms"`
|
||||
// Rule is what decided it within the category: a pattern's name, or a
|
||||
// judge's verdict, "S6,Specialized Advice".
|
||||
Rule string `json:"rule,omitempty"`
|
||||
GuardMs int64 `json:"guard_ms"`
|
||||
}
|
||||
|
||||
// policyView sums one policy up.
|
||||
@@ -198,6 +207,8 @@ type GuardAttachmentRequest struct {
|
||||
}
|
||||
|
||||
func (a *API) registerGuard(mux Mux) {
|
||||
mux.HandleFunc("GET /api/v1/guard/judges", a.listJudges)
|
||||
mux.HandleFunc("POST /api/v1/guard/judges/ask", a.askJudge)
|
||||
mux.HandleFunc("GET /api/v1/guard", func(w http.ResponseWriter, r *http.Request) {
|
||||
if a.guard == nil {
|
||||
writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard in the cluster's settings (llmgw ctl config edit)")
|
||||
@@ -208,7 +219,7 @@ func (a *API) registerGuard(mux Mux) {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
v := GuardView{Default: a.guard.Library().Default, Attachments: []GuardAttachmentView{}}
|
||||
v := GuardView{Enabled: a.guard.Library().On(), Default: a.guard.Library().Default, Attachments: []GuardAttachmentView{}}
|
||||
var names map[string]string
|
||||
if a.quotas != nil {
|
||||
names, _ = a.scopeNames(r.Context())
|
||||
@@ -269,7 +280,7 @@ func (a *API) registerGuard(mux Mux) {
|
||||
slices.SortStableFunc(v.Counts, func(a, b GuardCountView) int { return int(b.Count - a.Count) })
|
||||
for _, e := range events {
|
||||
v.Events = append(v.Events, GuardEventView{At: e.At, RequestID: e.RequestID, Project: project(e.ProjectID),
|
||||
Model: e.Model, Status: e.Status, Action: e.Action, Policy: e.Policy, Category: e.Category, GuardMs: e.GuardMs})
|
||||
Model: e.Model, Status: e.Status, Action: e.Action, Policy: e.Policy, Category: e.Category, Rule: e.Rule, GuardMs: e.GuardMs})
|
||||
}
|
||||
writeJSON(w, http.StatusOK, v)
|
||||
})
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
)
|
||||
|
||||
// JudgeView is one judge model the guardrail policies ask, and whether it
|
||||
// can answer them.
|
||||
type JudgeView struct {
|
||||
Model string `json:"model"`
|
||||
// Policies are those that ask it; Closed those that refuse what it
|
||||
// cannot read (fail: closed).
|
||||
Policies []string `json:"policies"`
|
||||
Closed []string `json:"closed,omitempty"`
|
||||
// Then are the judges its policies ask after it, when it gives no
|
||||
// verdict: a CPU judge behind a GPU one.
|
||||
Then []string `json:"then,omitempty"`
|
||||
// InCatalog is whether a model of that name is in the catalog, and
|
||||
// Judge whether it is there as a judge (role: guard).
|
||||
InCatalog bool `json:"in_catalog"`
|
||||
Judge bool `json:"judge"`
|
||||
Ready int `json:"ready_replicas"`
|
||||
Wanted int `json:"wanted_replicas"`
|
||||
// LastError is its last failure to give a verdict, if any.
|
||||
LastError string `json:"last_error,omitempty"`
|
||||
LastErrorAt *time.Time `json:"last_error_at,omitempty"`
|
||||
// Problem says, in words, what is wrong and what to do; empty when
|
||||
// nothing is.
|
||||
Problem string `json:"problem,omitempty"`
|
||||
}
|
||||
|
||||
// judgeErrorFresh is how recent a judge's failure must be to count as a
|
||||
// problem now.
|
||||
const judgeErrorFresh = 5 * time.Minute
|
||||
|
||||
// Judges are the judge models the policies in force ask, each with what is
|
||||
// wrong with it, if anything.
|
||||
func (a *API) Judges() []JudgeView {
|
||||
if a.guard == nil || a.guard.Library == nil {
|
||||
return nil
|
||||
}
|
||||
lib := a.guard.Library()
|
||||
if !lib.On() {
|
||||
// Off: no judge is asked, so none is missing.
|
||||
return nil
|
||||
}
|
||||
byModel := map[string]*JudgeView{}
|
||||
for name, p := range lib.Policies {
|
||||
judges := p.Judge.Judges()
|
||||
for i, model := range judges {
|
||||
v := byModel[model]
|
||||
if v == nil {
|
||||
v = &JudgeView{Model: model}
|
||||
byModel[model] = v
|
||||
}
|
||||
if !slices.Contains(v.Policies, name) {
|
||||
v.Policies = append(v.Policies, name)
|
||||
}
|
||||
if p.Fail == "closed" && p.Mode != "shadow" && !slices.Contains(v.Closed, name) {
|
||||
v.Closed = append(v.Closed, name)
|
||||
}
|
||||
for _, after := range judges[i+1:] {
|
||||
if !slices.Contains(v.Then, after) {
|
||||
v.Then = append(v.Then, after)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
ready := map[string]int{}
|
||||
for _, rep := range a.ctl.Reconciler().Replicas() {
|
||||
if rep.State == reconciler.Ready {
|
||||
ready[rep.Model]++
|
||||
}
|
||||
}
|
||||
var errs map[string]guard.JudgeError
|
||||
if a.guard.JudgeErrors != nil {
|
||||
errs = a.guard.JudgeErrors()
|
||||
}
|
||||
out := make([]JudgeView, 0, len(byModel))
|
||||
for _, v := range byModel {
|
||||
sort.Strings(v.Policies)
|
||||
sort.Strings(v.Closed)
|
||||
if m, ok := a.ctl.Catalog().Resolve(v.Model); ok && m.Name == v.Model {
|
||||
v.InCatalog, v.Judge = true, m.IsGuard()
|
||||
v.Wanted, v.Ready = m.Scaling.Min, ready[m.Name]
|
||||
}
|
||||
if e, ok := errs[v.Model]; ok {
|
||||
at := e.At
|
||||
v.LastError, v.LastErrorAt = e.Err, &at
|
||||
}
|
||||
v.Problem = judgeProblem(*v)
|
||||
out = append(out, *v)
|
||||
}
|
||||
slices.SortFunc(out, func(x, y JudgeView) int {
|
||||
switch {
|
||||
case x.Model < y.Model:
|
||||
return -1
|
||||
case x.Model > y.Model:
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
func judgeProblem(v JudgeView) string {
|
||||
consequence := "requests its policies judge are served unjudged"
|
||||
switch {
|
||||
case len(v.Then) > 0:
|
||||
consequence = "its policies ask " + strings.Join(v.Then, ", else ") + " instead"
|
||||
case len(v.Closed) > 0:
|
||||
consequence = "requests its policies judge are refused (guard_unavailable)"
|
||||
}
|
||||
switch {
|
||||
case !v.InCatalog:
|
||||
return fmt.Sprintf("not in the catalog, so %s: llmgw ctl guard judge enable", consequence)
|
||||
case !v.Judge:
|
||||
return fmt.Sprintf("in the catalog as a model clients use, not as a judge (role: guard), so %s: "+
|
||||
"llmgw ctl guard judge enable --replace", consequence)
|
||||
case v.Ready == 0:
|
||||
return fmt.Sprintf("no replica ready, so %s: llmgw ctl replicas --model %s", consequence, v.Model)
|
||||
case v.LastErrorAt != nil && time.Since(*v.LastErrorAt) < judgeErrorFresh:
|
||||
return fmt.Sprintf("gave no verdict %s ago (%s)", time.Since(*v.LastErrorAt).Round(time.Second), v.LastError)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (a *API) listJudges(w http.ResponseWriter, _ *http.Request) {
|
||||
if a.guard == nil {
|
||||
writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard in the cluster's settings (llmgw ctl config edit)")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, a.Judges())
|
||||
}
|
||||
|
||||
// JudgeAskRequest is a text to have a judge read.
|
||||
type JudgeAskRequest struct {
|
||||
// Model is the judge; empty is the one the policies ask, when they ask one.
|
||||
Model string `json:"model,omitempty"`
|
||||
Text string `json:"text"`
|
||||
// Answer, when given, is read as the assistant's reply to Text.
|
||||
Answer string `json:"answer,omitempty"`
|
||||
}
|
||||
|
||||
// JudgeAskView is what the judge answered.
|
||||
type JudgeAskView struct {
|
||||
Model string `json:"model"`
|
||||
// Verdict is safe, unsafe or controversial, as the gateway reads the
|
||||
// answer; a policy's judge.controversial says what controversial counts as.
|
||||
Verdict string `json:"verdict"`
|
||||
Categories []string `json:"categories,omitempty"`
|
||||
// Raw is the judge's answer as it gave it.
|
||||
Raw string `json:"raw"`
|
||||
TookMs int64 `json:"took_ms"`
|
||||
}
|
||||
|
||||
func (a *API) askJudge(w http.ResponseWriter, r *http.Request) {
|
||||
if a.guard == nil || a.guard.Ask == nil {
|
||||
writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard in the cluster's settings (llmgw ctl config edit)")
|
||||
return
|
||||
}
|
||||
var req JudgeAskRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil || strings.TrimSpace(req.Text) == "" {
|
||||
writeError(w, http.StatusBadRequest, "give the text for the judge to read")
|
||||
return
|
||||
}
|
||||
if req.Model == "" {
|
||||
judges := a.Judges()
|
||||
if len(judges) == 0 {
|
||||
writeError(w, http.StatusBadRequest, "no policy asks a judge: give the model")
|
||||
return
|
||||
}
|
||||
req.Model = judges[0].Model
|
||||
}
|
||||
ans, err := a.guard.Ask(r.Context(), req.Model, req.Text, req.Answer)
|
||||
v := JudgeAskView{Model: req.Model, Raw: ans.Raw, Categories: ans.Categories, TookMs: ans.Took.Milliseconds()}
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, fmt.Sprintf("the judge %s gave no verdict: %v (it answered %q)", req.Model, err, ans.Raw))
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case ans.Unsafe:
|
||||
v.Verdict = "unsafe"
|
||||
case ans.Controversial:
|
||||
// Qwen3Guard's middle verdict: what it counts as is the policy's.
|
||||
v.Verdict = "controversial"
|
||||
default:
|
||||
v.Verdict = "safe"
|
||||
}
|
||||
writeJSON(w, http.StatusOK, v)
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
)
|
||||
|
||||
func TestEachJudgeSaysWhatIsWrongWithIt(t *testing.T) {
|
||||
models := append(testModels(),
|
||||
// Served from the mirror page before it knew better: a model
|
||||
// clients use, under the judge's name.
|
||||
catalog.Model{Name: "shield", Source: "/s", Scaling: catalog.Scaling{Min: 1, Max: 1},
|
||||
Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp"}}},
|
||||
catalog.Model{Name: "llama-guard-3-1b", Source: "/g", Role: catalog.RoleGuard,
|
||||
Guard: &catalog.GuardModel{Format: "llamaguard3"}, Scaling: catalog.Scaling{Min: 2, Max: 2},
|
||||
Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp"}}})
|
||||
a, _ := setup(t, models)
|
||||
lib := guard.Config{Default: []string{"baseline"}, Policies: map[string]guard.Policy{
|
||||
"baseline": {Fail: "closed", Judge: &guard.JudgePolicy{Model: "llama-guard-3-1b", Input: "block"}},
|
||||
"strict": {Fail: "closed", Judge: &guard.JudgePolicy{Model: "missing-judge", Input: "block"}},
|
||||
"shield": {Fail: "open", Judge: &guard.JudgePolicy{Model: "shield", Input: "flag"}},
|
||||
}}
|
||||
a.SetGuard(&GuardAdmin{Library: func() guard.Config { return lib }})
|
||||
|
||||
var js []JudgeView
|
||||
if code := call(t, a, "GET", "/api/v1/guard/judges", "", &js); code != 200 || len(js) != 3 {
|
||||
t.Fatalf("%d %+v", code, js)
|
||||
}
|
||||
by := map[string]JudgeView{}
|
||||
for _, j := range js {
|
||||
by[j.Model] = j
|
||||
}
|
||||
if j := by["missing-judge"]; j.InCatalog || !strings.Contains(j.Problem, "not in the catalog") ||
|
||||
!strings.Contains(j.Problem, "refused") {
|
||||
t.Errorf("missing: %+v", j)
|
||||
}
|
||||
if j := by["shield"]; !j.InCatalog || j.Judge || !strings.Contains(j.Problem, "--replace") ||
|
||||
!strings.Contains(j.Problem, "unjudged") {
|
||||
t.Errorf("not a judge: %+v", j)
|
||||
}
|
||||
if j := by["llama-guard-3-1b"]; !j.Judge || j.Wanted != 2 || !strings.Contains(j.Problem, "no replica ready") {
|
||||
t.Errorf("not running: %+v", j)
|
||||
}
|
||||
|
||||
var s Summary
|
||||
call(t, a, "GET", "/api/v1/summary", "", &s)
|
||||
if len(s.Judges) != 3 {
|
||||
t.Errorf("the summary does not carry the judges' problems: %+v", s.Judges)
|
||||
}
|
||||
|
||||
// A judge that failed lately says so, once it is otherwise fine.
|
||||
at := time.Now()
|
||||
v := JudgeView{Model: "x", InCatalog: true, Judge: true, Ready: 1, LastError: "deadline exceeded", LastErrorAt: &at}
|
||||
if p := judgeProblem(v); !strings.Contains(p, "deadline exceeded") {
|
||||
t.Errorf("a recent failure: %q", p)
|
||||
}
|
||||
old := at.Add(-time.Hour)
|
||||
v.LastErrorAt = &old
|
||||
if p := judgeProblem(v); p != "" {
|
||||
t.Errorf("an old failure still counts: %q", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAJudgeIsAskedAndItsWordsShown(t *testing.T) {
|
||||
a, _ := setup(t, testModels())
|
||||
lib := guard.Config{Policies: map[string]guard.Policy{
|
||||
"baseline": {Judge: &guard.JudgePolicy{Model: "llama-guard-3-1b", Input: "block"}}}}
|
||||
var asked string
|
||||
a.SetGuard(&GuardAdmin{Library: func() guard.Config { return lib },
|
||||
Ask: func(_ context.Context, model, text, _ string) (guard.JudgeAnswer, error) {
|
||||
asked = model + "|" + text
|
||||
return guard.JudgeAnswer{Model: model, Raw: "unsafe\nS6", Unsafe: true,
|
||||
Categories: []string{"S6", "Specialized Advice"}}, nil
|
||||
}})
|
||||
var v JudgeAskView
|
||||
if code := call(t, a, "POST", "/api/v1/guard/judges/ask", `{"text":"how do I sort a list"}`, &v); code != 200 {
|
||||
t.Fatalf("%d", code)
|
||||
}
|
||||
if asked != "llama-guard-3-1b|how do I sort a list" || v.Verdict != "unsafe" || v.Raw != "unsafe\nS6" ||
|
||||
len(v.Categories) != 2 {
|
||||
t.Errorf("asked %q, got %+v", asked, v)
|
||||
}
|
||||
if code := call(t, a, "POST", "/api/v1/guard/judges/ask", `{"text":" "}`, nil); code != 400 {
|
||||
t.Errorf("no text: %d", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodeSaysItsEnginesVersionsAndCapabilities(t *testing.T) {
|
||||
n := gpuNode("volta", 2)
|
||||
n.EngineVersions = map[string]string{"llamacpp": "b10991", "vllm": "0.10.2"}
|
||||
for i := range n.GPUs {
|
||||
n.GPUs[i].ComputeCapability = "7.0"
|
||||
}
|
||||
a, _ := setup(t, testModels(), n)
|
||||
var nodes []NodeView
|
||||
call(t, a, "GET", "/api/v1/nodes", "", &nodes)
|
||||
if len(nodes) != 1 || nodes[0].EngineVersions["vllm"] != "0.10.2" ||
|
||||
len(nodes[0].ComputeCapabilities) != 1 || nodes[0].ComputeCapabilities[0] != "7.0" {
|
||||
t.Errorf("%+v", nodes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodeSaysWhetherItsFabricsAreUp(t *testing.T) {
|
||||
n := gpuNode("ib1", 2, "ib-a")
|
||||
if got := multiNodeReadiness(n); !strings.HasPrefix(got, "unknown") {
|
||||
t.Errorf("no port yet: %q", got)
|
||||
}
|
||||
n.IBPorts = []inventory.IBPort{{HCA: "mlx5_0", State: "DOWN"}}
|
||||
if got := multiNodeReadiness(n); got != "down: mlx5_0" {
|
||||
t.Errorf("down: %q", got)
|
||||
}
|
||||
n.IBPorts[0].State = "ACTIVE"
|
||||
if got := multiNodeReadiness(n); got != "ready" {
|
||||
t.Errorf("up: %q", got)
|
||||
}
|
||||
if got := multiNodeReadiness(gpuNode("plain", 1)); got != "" {
|
||||
t.Errorf("no fabric: %q", got)
|
||||
}
|
||||
}
|
||||
Loaded 100 of 258 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user