From 85025dfb030ec1fc89f2243d92d76d15d8ba7ca4 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Thu, 8 Oct 2026 21:30:50 +0200 Subject: [PATCH] Trained adapters published, private to their owners like the knowledge bases they come from An adapter now has an owner (org, team, project, user or key): only callers there see it, and to anyone else it is a model that does not exist, absent from /v1/models too. A training job's owner defaults to the one collection its pairs come from, must be able to read every collection they come from, and pairs from a knowledge base are never made public. ctl train publish keeps the GGUF or PEFT files on the mirror, which accepts uploads now, and sets the owned adapter after the usual architecture check. The knowledge bases' cache_salt follows the collections searched rather than the organization, so two teams of one organization no longer share cached passages. Co-Authored-By: Claude Opus 5.5 (1M context) --- BACKLOG.md | 21 +- CHANGELOG.md | 14 + docs/admin-api.openapi.json | 83 ++++++ docs/knowledge.md | 51 +++- docs/reference.md | 11 + internal/agent/train/trainer.py | 2 +- internal/cli/server.go | 7 + internal/cli/train.go | 57 +++- internal/control/adminapi/api.go | 62 +++-- internal/control/adminapi/kb.go | 2 + internal/control/adminapi/kbdataset.go | 17 +- internal/control/adminapi/routes.go | 3 + internal/control/adminapi/train.go | 247 +++++++++++++++++- internal/control/adminapi/train_test.go | 61 ++++- internal/control/catalog/catalog.go | 24 ++ internal/docgen/reference_gen.go | 1 + internal/mirror/api.go | 58 ++++ internal/mirror/client.go | 30 +++ internal/mirror/upload_test.go | 45 ++++ internal/pgha/memory.go | 145 ++++++++++ internal/pgha/setup.go | 3 + internal/router/adapterowner_test.go | 61 +++++ internal/router/handlers.go | 9 +- internal/router/rag.go | 20 +- internal/router/rag_test.go | 6 +- internal/router/router.go | 47 +++- .../migrations/postgres/0038_train_owner.sql | 5 + .../migrations/sqlite/0038_train_owner.sql | 5 + internal/store/train.go | 36 ++- 29 files changed, 1049 insertions(+), 84 deletions(-) create mode 100644 internal/mirror/upload_test.go create mode 100644 internal/pgha/memory.go create mode 100644 internal/router/adapterowner_test.go create mode 100644 internal/store/migrations/postgres/0038_train_owner.sql create mode 100644 internal/store/migrations/sqlite/0038_train_owner.sql diff --git a/BACKLOG.md b/BACKLOG.md index b8f7388..628fd47 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -1005,7 +1005,7 @@ 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) +### 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 @@ -1026,6 +1026,25 @@ 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). + --- ## D. Facts worth knowing before starting diff --git a/CHANGELOG.md b/CHANGELOG.md index f07c799..ef8d307 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,6 +41,20 @@ APIs change between minors without a deprecation period. 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. ## v0.14.0 — 2026-10-08 diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 79d7829..a0f4077 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -3940,6 +3940,48 @@ } } }, + "/api/v1/train/jobs/{id}/publish": { + "post": { + "operationId": "postTrainJobsIdPublish", + "summary": "Make what a finished job made the model's adapter", + "description": "Kept on the mirror — the GGUF for a model llama.cpp serves, the PEFT files for one vLLM serves — and added to the model's adapters after the same check as POST /api/v1/models/{name}/adapters.", + "tags": [ + "train" + ], + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "description": "The job's ID, as starting it answered.", + "schema": { + "type": "string" + } + }, + { + "$ref": "#/components/parameters/Actor" + } + ], + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TrainPublishView" + } + } + } + }, + "401": { + "$ref": "#/components/responses/Unauthorized" + }, + "default": { + "$ref": "#/components/responses/Error" + } + } + } + }, "/api/v1/updates": { "get": { "operationId": "getUpdates", @@ -5378,6 +5420,10 @@ "ord": { "type": "integer" }, + "collection": { + "type": "string", + "description": "Collection is where the pair came from: an adapter trained on it belongs to a part of the tree that may read it." + }, "error": { "type": "string" } @@ -7089,6 +7135,9 @@ "node": { "type": "string" }, + "owner": { + "type": "string" + }, "gpu_memory_mib": { "type": "integer" }, @@ -7143,6 +7192,28 @@ } } }, + "TrainPublishView": { + "type": "object", + "description": "TrainPublishView is an adapter published: where the nodes fetch it from, and what to do for replicas to load it.", + "properties": { + "adapter": { + "type": "string" + }, + "model": { + "type": "string" + }, + "source": { + "type": "string" + }, + "owner": { + "type": "string", + "description": "Owner is who sees it; empty is every caller of the model." + }, + "note": { + "type": "string" + } + } + }, "TrainRequest": { "type": "object", "description": "TrainRequest queues a job.", @@ -7169,6 +7240,14 @@ "gpu_memory_mib": { "type": "integer" }, + "owner": { + "type": "string", + "description": "Owner is who the adapter will belong to: org:, team:, project: (by path or ID), user: or key:. Only callers there see it. When the pairs come from one knowledge base, it defaults to that collection's owner; it must be able to read every collection they come from." + }, + "public": { + "type": "boolean", + "description": "Public makes the adapter every caller's its model allows. Refused for pairs from a knowledge base." + }, "data": { "type": "string", "description": "Data is the pairs, NDJSON of {\"q\", \"a\"} (llmgw ctl kb dataset)." @@ -7334,6 +7413,10 @@ "scale": { "type": "number", "description": "Scale is how strongly it is applied; zero takes 1." + }, + "owner": { + "type": "string", + "description": "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." } } }, diff --git a/docs/knowledge.md b/docs/knowledge.md index 95bd240..62253a1 100644 --- a/docs/knowledge.md +++ b/docs/knowledge.md @@ -317,9 +317,11 @@ Behaviour to know about: (Codex, Claude Code) work the same way, since they are answered as chat completions. - **Separate prefix caches.** On vLLM replicas, a request given passages - carries a `cache_salt` derived from its organization. One organization's - passages then never share cached blocks with another's, so how fast an - answer starts says nothing about what another tenant asked. + 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. @@ -340,6 +342,49 @@ 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 diff --git a/docs/reference.md b/docs/reference.md index 05efa24..4290e55 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -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) | @@ -426,6 +427,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` diff --git a/internal/agent/train/trainer.py b/internal/agent/train/trainer.py index 3d83a51..b4cc246 100644 --- a/internal/agent/train/trainer.py +++ b/internal/agent/train/trainer.py @@ -94,7 +94,7 @@ for ep in range(epochs): 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", step=step, steps=steps, loss=round(loss.item(), 4), peak_mib=peak["mib"]) + 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() diff --git a/internal/cli/server.go b/internal/cli/server.go index a747364..b21ff25 100644 --- a/internal/cli/server.go +++ b/internal/cli/server.go @@ -518,6 +518,13 @@ func newServerCmd() *cobra.Command { return fmt.Errorf("the model mirror: %w", err) } rt.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. diff --git a/internal/cli/train.go b/internal/cli/train.go index dccfbb2..a51ebef 100644 --- a/internal/cli/train.go +++ b/internal/cli/train.go @@ -28,7 +28,8 @@ func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { "what they say. Facts belong to the knowledge base.", } - var model, base, rev, node, data string + var model, base, rev, node, data, owner string + var public bool var gpuMem int var epochs, rank int var lr float64 @@ -45,7 +46,7 @@ func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { 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{}} + GPUMemoryMiB: gpuMem, Data: string(b), Params: map[string]any{}, Owner: owner, Public: public} if cmd.Flags().Changed("epochs") { req.Params["epochs"] = epochs } @@ -64,7 +65,12 @@ func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { if err := client().patient(2*time.Minute).do(cmd.Context(), "POST", "/api/v1/train/jobs", json.RawMessage(body), &v); err != nil { return err } - fmt.Fprintf(cmd.OutOrStdout(), "queued %s: %s for %s on %s, %d pairs\n", v.ID, v.Name, v.Model, v.Node, v.Pairs) + 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 }, @@ -80,6 +86,9 @@ func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { 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) } @@ -102,9 +111,9 @@ func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { return nil } w := tab() - fmt.Fprintln(w, "JOB\tADAPTER\tMODEL\tNODE\tSTATE\tPROGRESS\tLOSS\tPEAK\tCREATED") + 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\n", j.ID, j.Name, j.Model, j.Node, j.State, + 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() @@ -129,7 +138,8 @@ func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { } 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, on %s (base %s), %d pairs\n", j.ID, j.Name, j.Model, j.Node, j.BaseRepo, j.Pairs) + 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) @@ -182,6 +192,41 @@ func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { }, }) + 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", diff --git a/internal/control/adminapi/api.go b/internal/control/adminapi/api.go index 821ed52..107b1d0 100644 --- a/internal/control/adminapi/api.go +++ b/internal/control/adminapi/api.go @@ -1966,32 +1966,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 { @@ -2073,3 +2050,36 @@ func multiNodeReadiness(n inventory.Node) string { } 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 +} diff --git a/internal/control/adminapi/kb.go b/internal/control/adminapi/kb.go index ee4f528..4360eca 100644 --- a/internal/control/adminapi/kb.go +++ b/internal/control/adminapi/kb.go @@ -31,6 +31,8 @@ type KBStore interface { SearchChunks(ctx context.Context, q store.KBSearch) ([]store.KBHit, error) SetKBDefault(ctx context.Context, d store.KBDefault) error CollectionChunks(ctx context.Context, c store.KBCollection, limit int) ([]store.KBHit, error) + CollectionsFor(ctx context.Context, e store.Effective, write bool) ([]store.KBCollection, error) + EffectiveFor(ctx context.Context, scope store.Scope, id, org string) (store.Effective, error) KBDefaults(ctx context.Context) ([]store.KBDefault, error) } diff --git a/internal/control/adminapi/kbdataset.go b/internal/control/adminapi/kbdataset.go index 61c0f68..8ca0cee 100644 --- a/internal/control/adminapi/kbdataset.go +++ b/internal/control/adminapi/kbdataset.go @@ -26,11 +26,14 @@ type KBDatasetRequest struct { // KBDatasetLine is one line of the answer: a pair, or why a chunk gave none. type KBDatasetLine struct { - Q string `json:"q,omitempty"` - A string `json:"a,omitempty"` - Doc string `json:"doc"` - Ord int `json:"ord"` - Error string `json:"error,omitempty"` + Q string `json:"q,omitempty"` + A string `json:"a,omitempty"` + Doc string `json:"doc"` + Ord int `json:"ord"` + // Collection is where the pair came from: an adapter trained on it + // belongs to a part of the tree that may read it. + Collection string `json:"collection,omitempty"` + Error string `json:"error,omitempty"` } const kbDatasetPrompt = `From the passage of an organisation's document below, write %d different questions an employee might ask, each with a short, correct answer. @@ -93,7 +96,7 @@ func (a *API) kbDataset(w http.ResponseWriter, r *http.Request) { body, _ := json.Marshal(map[string]any{"model": m.Name, "temperature": 0.7, "max_tokens": 200 * per, "messages": []map[string]string{{"role": "user", "content": fmt.Sprintf(kbDatasetPrompt, per, ch.Document, ch.Body)}}}) status, out, _, err := a.replay(r.Context(), "/v1/chat/completions", m.Name, body) - line := KBDatasetLine{Doc: ch.Document, Ord: ch.Ord} + line := KBDatasetLine{Doc: ch.Document, Ord: ch.Ord, Collection: c.ID} if err == nil && status != http.StatusOK { err = fmt.Errorf("%d: %s", status, strings.TrimSpace(string(out))) } @@ -109,7 +112,7 @@ func (a *API) kbDataset(w http.ResponseWriter, r *http.Request) { _ = enc.Encode(line) } for _, p := range pairs { - p.Doc, p.Ord = ch.Document, ch.Ord + p.Doc, p.Ord, p.Collection = ch.Document, ch.Ord, c.ID _ = enc.Encode(p) } flush() diff --git a/internal/control/adminapi/routes.go b/internal/control/adminapi/routes.go index 287b266..ca60eed 100644 --- a/internal/control/adminapi/routes.go +++ b/internal/control/adminapi/routes.go @@ -310,6 +310,9 @@ var Routes = []Route{ Response: TrainJobView{}}, {Method: "DELETE", Path: "/api/v1/train/jobs/{id}", Summary: "Remove an ended training job and what it made", Response: map[string]any{}}, + {Method: "POST", Path: "/api/v1/train/jobs/{id}/publish", Summary: "Make what a finished job made the model's adapter", + Description: "Kept on the mirror — the GGUF for a model llama.cpp serves, the PEFT files for one vLLM serves — and added to the model's adapters after the same check as POST /api/v1/models/{name}/adapters.", + Response: TrainPublishView{}}, {Method: "GET", Path: "/api/v1/train/jobs/{id}/artifacts/{kind}", Summary: "What a training job made: peft (a gzipped tar) or gguf", Response: map[string]any{}}, diff --git a/internal/control/adminapi/train.go b/internal/control/adminapi/train.go index 7bd5d2b..b777799 100644 --- a/internal/control/adminapi/train.go +++ b/internal/control/adminapi/train.go @@ -1,19 +1,25 @@ package adminapi import ( + "archive/tar" "bufio" "bytes" + "compress/gzip" "context" "encoding/json" "errors" "fmt" + "io" "net/http" "regexp" + "slices" "strconv" "strings" "time" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" "gitea.mynodes.xyz/ilker/llmgw/internal/control/train" + "gitea.mynodes.xyz/ilker/llmgw/internal/mirror" "gitea.mynodes.xyz/ilker/llmgw/internal/store" ) @@ -36,11 +42,38 @@ type TrainStore interface { type TrainAdmin struct { Store TrainStore Service *train.Service + // Mirror keeps what a job made where the nodes fetch weights; nil + // leaves publishing off. + Mirror TrainMirror +} + +// TrainMirror is the model mirror, as publishing needs it. +type TrainMirror interface { + Upload(ctx context.Context, model, version, file string, body []byte, source string) (mirror.FetchResponse, error) +} + +// TrainPublishView is an adapter published: where the nodes fetch it from, +// and what to do for replicas to load it. +type TrainPublishView struct { + Adapter string `json:"adapter"` + Model string `json:"model"` + Source string `json:"source"` + // Owner is who sees it; empty is every caller of the model. + Owner string `json:"owner,omitempty"` + Note string `json:"note,omitempty"` } // SetTrain enables the training endpoints. func (a *API) SetTrain(t *TrainAdmin) { a.train = t } +// SetTrainMirror gives publishing the mirror, which is set up after the +// training endpoints are. +func (a *API) SetTrainMirror(m TrainMirror) { + if a.train != nil { + a.train.Mirror = m + } +} + // TrainRequest queues a job. type TrainRequest struct { // Name is the adapter's: what requests will ask for it by. @@ -55,6 +88,14 @@ type TrainRequest struct { // what an engine there leaves free. Node string `json:"node"` GPUMemoryMiB int `json:"gpu_memory_mib"` + // Owner is who the adapter will belong to: org:, team:, project: (by + // path or ID), user: or key:. Only callers there see it. When the + // pairs come from one knowledge base, it defaults to that collection's + // owner; it must be able to read every collection they come from. + Owner string `json:"owner,omitempty"` + // Public makes the adapter every caller's its model allows. Refused for + // pairs from a knowledge base. + Public bool `json:"public,omitempty"` // Data is the pairs, NDJSON of {"q", "a"} (llmgw ctl kb dataset). Data string `json:"data"` // Params tune the trainer: rank, alpha, dropout, epochs, lr, batch, @@ -73,6 +114,7 @@ type TrainJobView struct { BaseRepo string `json:"base_repo"` BaseRevision string `json:"base_revision,omitempty"` Node string `json:"node"` + Owner string `json:"owner,omitempty"` GPUMemoryMiB int `json:"gpu_memory_mib"` Params string `json:"params,omitempty"` GGUF bool `json:"gguf"` @@ -93,7 +135,7 @@ type TrainJobView struct { func trainView(j store.TrainJob) TrainJobView { return TrainJobView{ID: j.ID, Name: j.Name, Model: j.Model, BaseRepo: j.BaseRepo, BaseRevision: j.BaseRevision, - Node: j.Node, GPUMemoryMiB: j.GPUMemoryMiB, Params: j.Params, GGUF: j.GGUF, Pairs: j.Pairs, State: j.State, + Node: j.Node, Owner: j.Owner, GPUMemoryMiB: j.GPUMemoryMiB, Params: j.Params, GGUF: j.GGUF, Pairs: j.Pairs, State: j.State, Step: j.Step, Steps: j.Steps, Loss: j.Loss, Message: j.Message, PeakMiB: j.PeakMiB, CreatedAt: j.CreatedAt, StartedAt: j.StartedAt, FinishedAt: j.FinishedAt, CreatedBy: j.CreatedBy} } @@ -105,6 +147,7 @@ func (a *API) registerTrain(mux Mux) { mux.HandleFunc("POST /api/v1/train/jobs/{id}/cancel", a.trainCancel) mux.HandleFunc("DELETE /api/v1/train/jobs/{id}", a.trainDelete) mux.HandleFunc("GET /api/v1/train/jobs/{id}/artifacts/{kind}", a.trainArtifact) + mux.HandleFunc("POST /api/v1/train/jobs/{id}/publish", a.trainPublish) } func (a *API) trainReady(w http.ResponseWriter) bool { @@ -163,7 +206,12 @@ func (a *API) trainCreate(w http.ResponseWriter, r *http.Request) { bad("gpu_memory_mib is what the job needs of a card, in MiB") return } - pairs, err := countPairs(in.Data) + pairs, colls, err := countPairs(in.Data) + if err != nil { + bad("%v", err) + return + } + owner, err := a.trainOwner(r.Context(), in, colls) if err != nil { bad("%v", err) return @@ -191,7 +239,7 @@ func (a *API) trainCreate(w http.ResponseWriter, r *http.Request) { gguf = *in.GGUF } j := store.TrainJob{ID: store.NewID("train-"), Name: in.Name, Model: m.Name, BaseRepo: in.BaseRepo, - BaseRevision: in.BaseRevision, Node: in.Node, GPUMemoryMiB: in.GPUMemoryMiB, Params: params, GGUF: gguf, + BaseRevision: in.BaseRevision, Node: in.Node, Owner: owner, GPUMemoryMiB: in.GPUMemoryMiB, Params: params, GGUF: gguf, Data: []byte(in.Data), Pairs: pairs, CreatedBy: actorOf(r)} if err := a.train.Store.CreateTrainJob(r.Context(), j); err != nil { writeError(w, http.StatusInternalServerError, err.Error()) @@ -203,11 +251,67 @@ func (a *API) trainCreate(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusCreated, trainView(j)) } -// countPairs checks a job's data: NDJSON of {"q", "a"}, at least one. -func countPairs(data string) (int, error) { - if len(data) > trainMaxData { - return 0, fmt.Errorf("the data is over %d MiB", trainMaxData>>20) +// trainOwner is who a job's adapter will belong to, as scope:ID: one that +// may read every knowledge base the pairs come from. Pairs from a knowledge +// base never make an adapter everyone's. +func (a *API) trainOwner(ctx context.Context, in TrainRequest, colls []string) (string, error) { + if in.Public { + if len(colls) > 0 { + return "", errors.New("these pairs come from a knowledge base, and an adapter trained on them " + + "carries some of what it says: it belongs to a part of the tree that may read it, not to everyone") + } + if in.Owner != "" { + return "", errors.New("an adapter is public or has an owner, not both") + } + return "", nil } + if len(colls) > 0 && a.kb == nil { + return "", errors.New("the pairs name knowledge bases, which this control plane cannot read") + } + var found []store.KBCollection + for _, id := range colls { + c, err := a.kb.Store.Collection(ctx, id) + if err != nil { + return "", fmt.Errorf("the pairs come from collection %s, which is not here: %w", id, err) + } + found = append(found, c) + } + if in.Owner == "" { + if len(found) == 1 { + return string(found[0].OwnerScope) + ":" + found[0].OwnerID, nil + } + return "", errors.New("name the adapter's owner (org:, team:, project:, user: or key:), or make it public " + + "when its pairs are not a knowledge base's") + } + scope, id, err := a.kbTarget(ctx, in.Owner) + if err != nil { + return "", err + } + for _, c := range found { + e, err := a.kb.Store.EffectiveFor(ctx, scope, id, c.OrgID) + if err != nil { + return "", fmt.Errorf("%s: %w", in.Owner, err) + } + readable, err := a.kb.Store.CollectionsFor(ctx, e, false) + if err != nil { + return "", err + } + if !slices.ContainsFunc(readable, func(x store.KBCollection) bool { return x.ID == c.ID }) { + return "", fmt.Errorf("%s may not read collection %s, which the pairs come from: "+ + "an adapter trained on it would tell its callers what it says", in.Owner, c.Name) + } + } + return string(scope) + ":" + id, nil +} + +// countPairs checks a job's data: NDJSON of {"q", "a"}, at least one, and +// says which knowledge bases the pairs come from. +func countPairs(data string) (int, []string, error) { + if len(data) > trainMaxData { + return 0, nil, fmt.Errorf("the data is over %d MiB", trainMaxData>>20) + } + seen := map[string]bool{} + var colls []string n := 0 sc := bufio.NewScanner(strings.NewReader(data)) sc.Buffer(make([]byte, 64<<10), 1<<20) @@ -218,22 +322,26 @@ func countPairs(data string) (int, error) { if len(b) == 0 { continue } - var p struct{ Q, A string } + var p struct{ Q, A, Collection string } if err := json.Unmarshal(b, &p); err != nil { - return 0, fmt.Errorf("data line %d is not JSON: %v", line, err) + return 0, nil, fmt.Errorf("data line %d is not JSON: %v", line, err) } if strings.TrimSpace(p.Q) == "" || strings.TrimSpace(p.A) == "" { - return 0, fmt.Errorf("data line %d has no q or no a", line) + return 0, nil, fmt.Errorf("data line %d has no q or no a", line) + } + if p.Collection != "" && !seen[p.Collection] { + seen[p.Collection] = true + colls = append(colls, p.Collection) } n++ } if err := sc.Err(); err != nil { - return 0, err + return 0, nil, err } if n == 0 { - return 0, errors.New("the data holds no pairs: NDJSON of {\"q\", \"a\"}, as llmgw ctl kb dataset writes") + return 0, nil, errors.New("the data holds no pairs: NDJSON of {\"q\", \"a\"}, as llmgw ctl kb dataset writes") } - return n, nil + return n, colls, nil } func (a *API) trainList(w http.ResponseWriter, r *http.Request) { @@ -336,3 +444,116 @@ func (a *API) trainArtifact(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Length", strconv.Itoa(len(b))) _, _ = w.Write(b) } + +// trainPublish puts what a job made on the mirror and makes it the model's +// adapter: the GGUF for a model llama.cpp serves, the PEFT directory for +// one vLLM serves. +func (a *API) trainPublish(w http.ResponseWriter, r *http.Request) { + j, ok := a.trainJob(w, r) + if !ok { + return + } + if a.train.Mirror == nil { + writeError(w, http.StatusNotImplemented, "publishing needs the model mirror (mirror.api in the settings); "+ + "llmgw ctl train download, then llmgw mirror import there, does it by hand") + return + } + if a.models == nil { + writeError(w, http.StatusNotImplemented, "this control plane has no store for the catalog") + return + } + if j.State != store.TrainDone { + writeError(w, http.StatusConflict, "the job has not finished: it is "+j.State) + return + } + m, found := a.ctl.Catalog().Resolve(j.Model) + if !found { + writeError(w, http.StatusConflict, "the model "+j.Model+" is no longer in the catalog") + return + } + for _, ad := range m.Adapters { + if ad.Name == j.Name { + writeError(w, http.StatusConflict, j.Name+" is already "+m.Name+"'s adapter; train under another name") + return + } + } + llama := slices.ContainsFunc(m.Profiles, func(p catalog.Profile) bool { return p.Engine == "llamacpp" }) + source := "trained by llmgw job " + j.ID + " on " + j.Node + " from " + j.BaseRepo + var url string + if llama { + b, err := a.train.Store.TrainArtifact(r.Context(), j.ID, "gguf") + if err != nil { + writeError(w, http.StatusConflict, "the job made no GGUF, which llama.cpp needs; train again without --no-gguf") + return + } + res, err := a.train.Mirror.Upload(r.Context(), j.Name, j.ID, j.Name+".gguf", b, source) + if err != nil { + writeError(w, http.StatusBadGateway, "keeping it on the mirror: "+err.Error()) + return + } + url = res.URL + } else { + b, err := a.train.Store.TrainArtifact(r.Context(), j.ID, "peft") + if err != nil { + writeError(w, http.StatusConflict, "the job made no PEFT adapter") + return + } + files, err := untarAll(b) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + for name, data := range files { + res, err := a.train.Mirror.Upload(r.Context(), j.Name, j.ID, name, data, source) + if err != nil { + writeError(w, http.StatusBadGateway, "keeping "+name+" on the mirror: "+err.Error()) + return + } + url = res.URL[:strings.LastIndex(res.URL, "/")+1] + } + } + ad := catalog.Adapter{Name: j.Name, Source: url, Owner: j.Owner} + if err := checkAdapters(r.Context(), m, []catalog.Adapter{ad}); err != nil { + writeError(w, http.StatusBadRequest, err.Error()) + return + } + if _, err := a.models.SetAdapters(r.Context(), m.Name, append(slices.Clone(m.Adapters), ad), actorOf(r)); err != nil { + writeError(w, http.StatusBadRequest, err.Error()) + return + } + v := TrainPublishView{Adapter: j.Name, Model: m.Name, Source: url, Owner: j.Owner} + if llama { + v.Note = "llama.cpp loads adapters when a replica starts: restart " + m.Name + "'s replicas (llmgw ctl restart " + + m.Name + ") and ask for model " + j.Name + } else { + v.Note = "ask for model " + j.Name + "; a replica loads it on the first request" + } + writeJSON(w, http.StatusOK, v) +} + +// untarAll reads a gzipped tar of a PEFT directory's files. +func untarAll(b []byte) (map[string][]byte, error) { + gz, err := gzip.NewReader(bytes.NewReader(b)) + if err != nil { + return nil, err + } + tr := tar.NewReader(gz) + out := map[string][]byte{} + for { + h, err := tr.Next() + if errors.Is(err, io.EOF) { + return out, nil + } + if err != nil { + return nil, err + } + if h.Typeflag != tar.TypeReg || strings.ContainsAny(h.Name, "/\\") { + continue + } + data, err := io.ReadAll(io.LimitReader(tr, 4<<30)) + if err != nil { + return nil, err + } + out[h.Name] = data + } +} diff --git a/internal/control/adminapi/train_test.go b/internal/control/adminapi/train_test.go index 10cb97c..2680ab2 100644 --- a/internal/control/adminapi/train_test.go +++ b/internal/control/adminapi/train_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "os" "path/filepath" "strings" "testing" @@ -38,7 +39,8 @@ func TestATrainingJobIsCheckedBeforeItIsQueued(t *testing.T) { return rec.Code, rec.Body.String() } good := map[string]any{"name": "tone", "model": "small", "base_repo": "Qwen/Qwen2.5-0.5B-Instruct", "node": "cuda01", - "gpu_memory_mib": 2300, "data": "{\"q\":\"a\",\"a\":\"b\"}\n{\"q\":\"c\",\"a\":\"d\"}\n", "params": map[string]any{"epochs": 3}} + "gpu_memory_mib": 2300, "data": "{\"q\":\"a\",\"a\":\"b\"}\n{\"q\":\"c\",\"a\":\"d\"}\n", "params": map[string]any{"epochs": 3}, + "public": true} with := func(k string, v any) string { m := map[string]any{} for kk, vv := range good { @@ -62,6 +64,8 @@ func TestATrainingJobIsCheckedBeforeItIsQueued(t *testing.T) { {"data", "", "no pairs"}, {"params", map[string]any{"epochs": 0}, "params.epochs"}, {"params", map[string]any{"momentum": 0.9}, "params.momentum"}, + {"public", false, "name the adapter's owner"}, + {"data", "{\"q\":\"a\",\"a\":\"b\",\"collection\":\"kb_aaaaaaaaaaaaaaaa\"}\n", "from a knowledge base"}, } { if code, body := post(with(c.field, c.value)); code != http.StatusBadRequest || !strings.Contains(body, c.says) { t.Errorf("%s=%v: %d %s", c.field, c.value, code, body) @@ -84,3 +88,58 @@ func TestATrainingJobIsCheckedBeforeItIsQueued(t *testing.T) { t.Errorf("an artifact not made: %d", rec.Code) } } + +func TestAnAdapterTrainedOnAKnowledgeBaseBelongsToWhoMayReadIt(t *testing.T) { + dsn := os.Getenv("LLMGW_TEST_POSTGRES") + if dsn == "" { + t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this") + } + ctx := context.Background() + s, err := store.OpenPostgres(ctx, dsn) + if err != nil { + t.Fatal(err) + } + defer s.Close() + if s.KnowledgeAvailable() != nil { + t.Skip("no pgvector") + } + models := append(testModels(), catalog.Model{Name: "small", Source: "/s", + Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp"}}}) + a, _ := setup(t, models, gpuNode("cuda01", 1)) + a.SetQuotas(s) + a.SetKB(&KBAdmin{Store: s}) + a.SetTrain(&TrainAdmin{Store: s, Service: train.New(s, noAgents{}, nil)}) + sfx := strings.ToLower(store.NewID("")) + org, _ := s.CreateOrg(ctx, "Tr"+sfx) + hr, _ := s.CreateTeam(ctx, org, "hr") + _, _ = s.CreateProject(ctx, hr, "docs") + sales, _ := s.CreateTeam(ctx, org, "sales") + _, _ = s.CreateProject(ctx, sales, "crm") + c, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeTeam, OwnerID: hr, Name: "p", EmbedModel: "bge-m3", Dims: 3}) + if err != nil { + t.Fatal(err) + } + mux := http.NewServeMux() + a.Register(mux) + post := func(extra string) (int, TrainJobView, string) { + body := `{"name":"tone-` + store.NewID("") + `","model":"small","base_repo":"Qwen/Q","node":"cuda01","gpu_memory_mib":2000,` + + `"data":"{\"q\":\"a\",\"a\":\"b\",\"collection\":\"` + c.ID + `\"}\n"` + extra + `}` + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, httptest.NewRequest("POST", "/api/v1/train/jobs", strings.NewReader(body))) + var v TrainJobView + _ = json.Unmarshal(rec.Body.Bytes(), &v) + return rec.Code, v, rec.Body.String() + } + if code, v, body := post(""); code != http.StatusCreated || v.Owner != "team:"+hr { + t.Errorf("by default: %d %s", code, body) + } + if code, v, body := post(`,"owner":"project:Tr` + sfx + `/hr/docs"`); code != http.StatusCreated || !strings.HasPrefix(v.Owner, "project:") { + t.Errorf("a project of the team: %d %s", code, body) + } + if code, _, body := post(`,"owner":"project:Tr` + sfx + `/sales/crm"`); code != http.StatusBadRequest || !strings.Contains(body, "may not read") { + t.Errorf("a project that cannot read it: %d %s", code, body) + } + if code, _, body := post(`,"public":true`); code != http.StatusBadRequest || !strings.Contains(body, "not to everyone") { + t.Errorf("public: %d %s", code, body) + } +} diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 2170afd..3837756 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -1127,6 +1127,27 @@ type Adapter struct { Source string `json:"source" koanf:"source"` // Scale is how strongly it is applied; zero takes 1. Scale float64 `json:"scale,omitempty" koanf:"scale"` + // 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. + Owner string `json:"owner,omitempty" koanf:"owner"` +} + +// ValidAdapterOwner checks an adapter's owner is scope:ID. +func ValidAdapterOwner(owner string) error { + if owner == "" { + return nil + } + scope, id, ok := strings.Cut(owner, ":") + switch { + case !ok || id == "": + return fmt.Errorf("an adapter's owner is scope:ID, %q is not", owner) + case scope != "org" && scope != "team" && scope != "project" && scope != "user" && scope != "key": + return fmt.Errorf("an adapter's owner is an org, team, project, user or key, not %q", scope) + } + return nil } // EffectiveScale is the scale an adapter is applied at. @@ -1175,6 +1196,9 @@ func validateAdapters(m Model) error { } seen := map[string]bool{} for _, a := range m.Adapters { + if err := ValidAdapterOwner(a.Owner); err != nil { + return fmt.Errorf("model %q adapter %q: %w", m.Name, a.Name, err) + } switch { case a.Name == "" || a.Source == "": return fmt.Errorf("model %q: an adapter needs a name and a source", m.Name) diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index d5978b5..fd276c8 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -8,6 +8,7 @@ var Reference = []Block{ {Key: "name", Type: "string", Doc: "Name is what clients ask for, like a model name.", Nested: ""}, {Key: "source", Type: "string", Doc: "Source is the adapter: a GGUF file for llama.cpp; for vLLM a PEFT\nadapter directory (a mirror version ending in \"/\", a path on the node,\nor a Hugging Face repository id). Fetched by each node as the model's\nweights are.", Nested: ""}, {Key: "scale", Type: "float64", Doc: "Scale is how strongly it is applied; zero takes 1.", Nested: ""}, + {Key: "owner", Type: "string", Doc: "Owner is the part of the tenancy tree the adapter belongs to —\norg:ID, team:ID, project:ID, user:ID or key:ID — and only callers in\nit see it: to anyone else it is a model that does not exist. An\nadapter trained on an organisation's documents carries some of what\nthey say in its weights. Empty is every caller the model allows.", Nested: ""}, }}, {Package: "config", Name: "Admission", Title: "Admission", Doc: "", Fields: []Field{ {Key: "enabled", Type: "bool", Doc: "Enabled turns the layer on. Off, every request goes straight through,\nwhich is the v0.2.0 behaviour.", Nested: ""}, diff --git a/internal/mirror/api.go b/internal/mirror/api.go index 3e38312..8f9ba2a 100644 --- a/internal/mirror/api.go +++ b/internal/mirror/api.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/http" "os" "strings" @@ -37,6 +38,7 @@ func (a *API) Handler() http.Handler { mux.HandleFunc("POST /v1/models", a.fetch) mux.HandleFunc("GET /v1/models/{model}/{version}", a.status) mux.HandleFunc("DELETE /v1/models/{model}/{version}", a.remove) + mux.HandleFunc("PUT /v1/models/{model}/{version}/{file}", a.upload) mux.HandleFunc("GET /v1/health", a.health) return mux } @@ -199,3 +201,59 @@ func LastPathElement(u string) string { } return u } + +// maxUpload bounds a file sent to the mirror rather than fetched by it. +const maxUpload = 8 << 30 + +// upload keeps a file the caller sends — an adapter a node trained, which +// is on no registry to fetch from — with the same discipline as an import: +// its checksum and source beside it, and an existing version never +// replaced. X-Llmgw-Source says where it came from. +func (a *API) upload(w http.ResponseWriter, r *http.Request) { + e := Entry{Model: r.PathValue("model"), Version: r.PathValue("version"), File: r.PathValue("file")} + for _, n := range []string{e.Model, e.Version, e.File} { + if err := ValidName(n); err != nil { + fail(w, http.StatusBadRequest, err) + return + } + } + if _, err := os.Stat(e.Path(a.m.root)); err == nil { + fail(w, http.StatusConflict, fmt.Errorf("%s/%s/%s: %w", e.Model, e.Version, e.File, ErrExists)) + return + } + tmp, err := os.CreateTemp(a.m.root, ".upload-*") + if err != nil { + fail(w, http.StatusInternalServerError, err) + return + } + defer os.Remove(tmp.Name()) + n, err := io.Copy(tmp, http.MaxBytesReader(w, r.Body, maxUpload)) + if cerr := tmp.Close(); err == nil { + err = cerr + } + if err != nil { + fail(w, http.StatusBadRequest, fmt.Errorf("receiving the file: %w", err)) + return + } + if n == 0 { + fail(w, http.StatusBadRequest, errors.New("the file is empty")) + return + } + source := strings.TrimSpace(r.Header.Get("X-Llmgw-Source")) + if source == "" { + source = "uploaded" + } + added, err := Import(r.Context(), a.m.root, tmp.Name(), e, "", source, nil) + switch { + case errors.Is(err, ErrExists): + fail(w, http.StatusConflict, err) + return + case err != nil: + fail(w, http.StatusBadRequest, err) + return + } + writeJSON(w, http.StatusCreated, FetchResponse{ + Status: Status{Model: e.Model, Version: e.Version, File: e.File, State: Present, BytesDone: added.Entry.Bytes, BytesTotal: added.Entry.Bytes}, + URL: e.URL(a.base), + }) +} diff --git a/internal/mirror/client.go b/internal/mirror/client.go index d2c830b..ef25b3b 100644 --- a/internal/mirror/client.go +++ b/internal/mirror/client.go @@ -100,6 +100,36 @@ func (c *Client) Fetch(ctx context.Context, req FetchRequest) (FetchResponse, er return out, err } +// Upload sends a file for the mirror to keep at model/version/file, with +// where it came from; it answers the URL nodes fetch it from. +func (c *Client) Upload(ctx context.Context, model, version, file string, body []byte, source string) (FetchResponse, error) { + var out FetchResponse + req, err := http.NewRequestWithContext(ctx, http.MethodPut, + c.base+"/v1/models/"+url.PathEscape(model)+"/"+url.PathEscape(version)+"/"+url.PathEscape(file), bytes.NewReader(body)) + if err != nil { + return out, err + } + req.Header.Set("Content-Type", "application/octet-stream") + req.Header.Set("X-Llmgw-Source", source) + resp, err := c.hc.Do(req) + if err != nil { + return out, fmt.Errorf("reaching the mirror at %s: %w", c.base, err) + } + defer resp.Body.Close() + raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return out, err + } + if resp.StatusCode >= 300 { + var e struct { + Error string `json:"error"` + } + _ = json.Unmarshal(raw, &e) + return out, &APIError{Status: resp.StatusCode, Msg: e.Error} + } + return out, json.Unmarshal(raw, &out) +} + // Status is where one model version stands. func (c *Client) Status(ctx context.Context, model, version string) (FetchResponse, error) { var out FetchResponse diff --git a/internal/mirror/upload_test.go b/internal/mirror/upload_test.go new file mode 100644 index 0000000..9862838 --- /dev/null +++ b/internal/mirror/upload_test.go @@ -0,0 +1,45 @@ +package mirror + +import ( + "context" + "net/http/httptest" + "os" + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf" +) + +func TestAFileSentToTheMirrorIsKeptWithItsChecksumAndNeverReplaced(t *testing.T) { + root := t.TempDir() + srv := httptest.NewServer(NewAPI(NewManager(root, nil), "http://mirror:8081").Handler()) + defer srv.Close() + c, err := NewClient(srv.URL, tlsconf.Config{}, 0) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + got, err := c.Upload(ctx, "politika-tonu", "train-1", "politika-tonu.gguf", []byte("GGUF adapter"), "trained by job train-1") + if err != nil { + t.Fatal(err) + } + if got.URL != "http://mirror:8081/politika-tonu/train-1/politika-tonu.gguf" || got.Status.BytesTotal != 12 { + t.Errorf("answered %+v", got) + } + e := Entry{Model: "politika-tonu", Version: "train-1", File: "politika-tonu.gguf"} + if b, _ := os.ReadFile(e.Path(root)); string(b) != "GGUF adapter" { + t.Errorf("kept %q", b) + } + if b, _ := os.ReadFile(sourcePath(e.Path(root))); !strings.Contains(string(b), "train-1") { + t.Errorf("source %q", b) + } + if _, err := os.Stat(checksumPath(e.Path(root))); err != nil { + t.Errorf("no checksum: %v", err) + } + if _, err := c.Upload(ctx, "politika-tonu", "train-1", "politika-tonu.gguf", []byte("other"), "x"); err == nil { + t.Error("an existing version was replaced") + } + if _, err := c.Upload(ctx, "../etc", "v", "f", []byte("x"), "x"); err == nil { + t.Error("a path outside the mirror was taken") + } +} diff --git a/internal/pgha/memory.go b/internal/pgha/memory.go new file mode 100644 index 0000000..4655ff2 --- /dev/null +++ b/internal/pgha/memory.go @@ -0,0 +1,145 @@ +package pgha + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "slices" + "strconv" + "strings" + "time" +) + +// PostgreSQL's own defaults — 128MB of cache, 4MB to sort in — are for a +// machine it barely uses. A knowledge base's vector index is searched well +// only from memory: one that does not fit is read from disk on every search. +// So each member is sized to the memory it has, read when setup runs, and +// again on every later run, so a machine given more memory gets the use of it. + +// MemorySettings are PostgreSQL's memory parameters for a machine with total +// bytes of memory. shared is a machine that also runs a control plane or an +// agent: the database is given less of it, since a gateway swapped out to +// give the database cache is worse than a smaller cache. +func MemorySettings(total int64, shared bool) map[string]string { + const mb = 1 << 20 + ram := total / mb + // Shares of memory, in MB: the cache PostgreSQL keeps itself, what the + // planner may count on the kernel caching besides, what building or + // vacuuming an index may take, and what one sort may. + buffers, cache, maint, sort := ram/4, ram*3/4, ram/16, ram/256 + if shared { + buffers, cache, maint, sort = ram*3/20, ram/2, ram/32, ram/512 + } + clamp := func(v, lo, hi int64) int64 { return min(max(v, lo), hi) } + return map[string]string{ + "shared_buffers": strconv.FormatInt(clamp(buffers, 96, 32<<10), 10) + "MB", + "effective_cache_size": strconv.FormatInt(clamp(cache, 256, 1<<20), 10) + "MB", + "maintenance_work_mem": strconv.FormatInt(clamp(maint, 64, 2<<10), 10) + "MB", + "work_mem": strconv.FormatInt(clamp(sort, 4, 64), 10) + "MB", + } +} + +// MemTotal reads the machine's memory, in bytes, from /proc/meminfo. +func MemTotal(meminfo string) (int64, bool) { + for _, line := range strings.Split(meminfo, "\n") { + f := strings.Fields(line) + if len(f) >= 2 && f[0] == "MemTotal:" { + kb, err := strconv.ParseInt(f[1], 10, 64) + return kb << 10, err == nil && kb > 0 + } + } + return 0, false +} + +// memory sizes each member to its machine, in its own configuration — the +// machines need not be alike. What takes effect on reload is applied so; +// shared_buffers needs a restart, given one member at a time, the leader +// last, each waited on before the next. +func (s *Setup) memory(ctx context.Context, sec Secrets) error { + var restart []Node + for _, d := range s.DB { + out, err := d.Remote.Run(ctx, `cat /proc/meminfo +if [ -f /etc/llmgw/server.yaml ] || [ -f /etc/llmgw/agent.yaml ]; then echo "llmgw-shared"; fi`, nil) + total, ok := MemTotal(string(out)) + if err != nil || !ok { + s.emit("memory", false, "%s: could not read its memory; PostgreSQL left as it is", d.Name) + continue + } + shared := strings.Contains(string(out), "llmgw-shared") + params := MemorySettings(total, shared) + body, _ := json.Marshal(params) + out, err = d.Remote.Run(ctx, `set -e +test -f /etc/patroni/config.yml || { echo missing; exit 0; } +want=$(cat) +r=$(sudo env WANT="$want" python3 - <<'PY' +import json, os, yaml +p = "/etc/patroni/config.yml" +want = json.loads(os.environ["WANT"]) +c = yaml.safe_load(open(p)) +pg = c.setdefault("postgresql", {}) +params = pg.get("parameters") or {} +pg["parameters"] = params +if all(str(params.get(k)) == v for k, v in want.items()): + print("same") +else: + restart = str(params.get("shared_buffers")) != want["shared_buffers"] + params.update(want) + open(p + ".new", "w").write("# llmgw: the control planes' store, kept by Patroni, its leader lock in etcd.\n" + + yaml.safe_dump(c, default_flow_style=False, sort_keys=False)) + print("restart" if restart else "changed") +PY +) +if [ "$r" != same ]; then + sudo chown root:postgres /etc/patroni/config.yml.new + sudo chmod 640 /etc/patroni/config.yml.new + sudo mv -f /etc/patroni/config.yml.new /etc/patroni/config.yml + sudo systemctl reload patroni +fi +echo "$r"`, body) + if err != nil { + return fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out))) + } + what := fmt.Sprintf("%s: %d MB of memory%s; shared_buffers %s, effective_cache_size %s, maintenance_work_mem %s, work_mem %s", + d.Name, total>>20, map[bool]string{true: ", shared with llmgw"}[shared], + params["shared_buffers"], params["effective_cache_size"], params["maintenance_work_mem"], params["work_mem"]) + switch r := strings.TrimSpace(lastLine(out)); r { + case "same", "missing": + s.emit("memory", false, "%s", what) + case "restart": + s.emit("memory", true, "%s", what) + restart = append(restart, d) + default: + s.emit("memory", true, "%s", what) + } + } + if len(restart) == 0 { + return nil + } + c, err := s.Status(ctx, s.dbHosts()) + if err != nil { + return err + } + // The leader last: its restart is the only one the control planes see, + // a few seconds without writes, and by then the rest are back. + slices.SortStableFunc(restart, func(a, b Node) int { + al, bl := a.Name == c.Leader(), b.Name == c.Leader() + switch { + case al && !bl: + return 1 + case bl && !al: + return -1 + } + return 0 + }) + for _, d := range restart { + if _, err := restJSON(ctx, d.IP, http.MethodPost, "/restart", []byte(`{}`), sec.RestAPI, 2*time.Minute); err != nil { + return fmt.Errorf("restarting PostgreSQL on %s for its new shared_buffers: %w", d.Name, err) + } + if err := s.waitFor(ctx, s.dbHosts(), func(c Cluster) bool { return c.Complete(len(s.DB)) }); err != nil { + return fmt.Errorf("after restarting %s: %w", d.Name, err) + } + s.emit("memory", true, "%s: PostgreSQL restarted with the new shared_buffers", d.Name) + } + return nil +} diff --git a/internal/pgha/setup.go b/internal/pgha/setup.go index 05ad9e5..1407db6 100644 --- a/internal/pgha/setup.go +++ b/internal/pgha/setup.go @@ -128,6 +128,9 @@ func (s *Setup) Run(ctx context.Context) (string, error) { if err := s.ensurePatroniTLS(ctx, sec); err != nil { return "", fmt.Errorf("patroni api: %w", err) } + if err := s.memory(ctx, sec); err != nil { + return "", fmt.Errorf("memory: %w", err) + } if err := s.database(ctx); err != nil { return "", fmt.Errorf("database: %w", err) } diff --git a/internal/router/adapterowner_test.go b/internal/router/adapterowner_test.go new file mode 100644 index 0000000..5daebc1 --- /dev/null +++ b/internal/router/adapterowner_test.go @@ -0,0 +1,61 @@ +package router + +import ( + "crypto/sha256" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/auth" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +func TestAnOwnedAdapterIsSeenOnlyInItsPartOfTheTree(t *testing.T) { + e := newFakeEngine(t, "qwen2.5-14b-instruct") + rt, owner := pricedRouter(t, e.srv.URL, "tone") + ctx := t.Context() + key, _ := rt.store.KeyByID(ctx, "key-1") + team, _ := rt.store.CreateTeam(ctx, mustOrgOf(t, rt), "other-team") + other, _ := rt.store.CreateProject(ctx, team, "other") + otherSecret := store.KeyPrefix + strings.Repeat("b", 43) + sum := sha256.Sum256([]byte(otherSecret)) + if err := rt.store.CreateKey(ctx, store.Key{ID: "key-2", Hash: sum[:], Prefix: "bbbb", Name: "o", ProjectID: other, CreatedAt: time.Now()}); err != nil { + t.Fatal(err) + } + rt.storeModels([]catalog.Model{{Name: "qwen2.5-14b-instruct", Aliases: []string{"q14"}, + Adapters: []catalog.Adapter{{Name: "tone", Source: "/t.gguf", Owner: "project:" + key.ProjectID}}}}) + h := rt.Handler() + ask := func(secret string) int { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"tone","messages":[{"role":"user","content":"hi"}]}`)) + if secret != "" { + req.Header.Set("Authorization", "Bearer "+secret) + } + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec.Code + } + if code := ask(owner); code != http.StatusOK { + t.Errorf("its own project: %d", code) + } + if code := ask(otherSecret); code != http.StatusNotFound { + t.Errorf("another project: %d", code) + } + list := func(secret string) string { + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + req.Header.Set("Authorization", "Bearer "+secret) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec.Body.String() + } + if !strings.Contains(list(owner), `"tone"`) || strings.Contains(list(otherSecret), `"tone"`) { + t.Errorf("listed: own %s\nother %s", list(owner), list(otherSecret)) + } + // No identity, no owned adapter. + if rt.adapterHidden(auth.Principal{}, false, "tone") != true { + t.Error("an owned adapter was shown to a caller with no identity") + } +} diff --git a/internal/router/handlers.go b/internal/router/handlers.go index bb885b3..0f0654c 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -412,9 +412,9 @@ type parsedRequest struct { guardIn guard.Verdict // knowledge is set when passages from knowledge bases were added. knowledge bool - // kbOrg is the organization whose passages were added, for the prefix + // kbSalt is the collections whose passages were added, for the prefix // cache's salt. - kbOrg string + kbSalt string // noStore is a policy that keeps nothing of this request. noStore bool // inCluster is a project whose requests may not leave the cluster. @@ -484,7 +484,8 @@ func (r *Router) handleInference(path string) http.HandlerFunc { } } - if r.hiddenModel(parsed.Model) { + p0, authed := auth.FromContext(req.Context()) + if r.hiddenModel(parsed.Model) || r.adapterHidden(p0, authed, parsed.Model) { // Said as a model that is not there: it is not, for clients. writeErrorCode(w, http.StatusNotFound, "invalid_request_error", "model_not_found", fmt.Sprintf("the model %q does not exist", parsed.Model)) @@ -1041,7 +1042,7 @@ func (r *Router) handleModels(w http.ResponseWriter, req *http.Request) { principal, authed := auth.FromContext(req.Context()) names := make([]string, 0, len(seen)) for n := range seen { - if authed && !r.permits(principal, n) || r.hiddenModel(n) { + if authed && !r.permits(principal, n) || r.hiddenModel(n) || r.adapterHidden(principal, authed, n) { continue } names = append(names, n) diff --git a/internal/router/rag.go b/internal/router/rag.go index c5e93a7..76405f9 100644 --- a/internal/router/rag.go +++ b/internal/router/rag.go @@ -10,6 +10,7 @@ import ( "fmt" "net/http" "net/url" + "sort" "strconv" "strings" "time" @@ -225,7 +226,14 @@ func (r *Router) retrieve(ctx context.Context, w http.ResponseWriter, req *parse req.UpstreamBody = body req.PromptTokens = len(body) / 4 req.knowledge = true - req.kbOrg = e.OrgID + // The collections searched, which is who else could have had these + // passages added: the prefix cache is shared among exactly them. + ids := make([]string, len(colls)) + for i, c := range colls { + ids[i] = c.ID + } + sort.Strings(ids) + req.kbSalt = strings.Join(ids, ",") // Set now: the proxy writes the headers when the answer starts. w.Header().Set(sourcesHeader, strings.Join(sources, ",")) result := "found" @@ -429,14 +437,16 @@ func (r *Router) kbDefaults(ctx context.Context, p auth.Principal, project strin return names } -// withCacheSalt sets vLLM's cache_salt to one derived from an organization: -// the same for its requests, so they still share a cache among themselves. -func withCacheSalt(body []byte, org string) []byte { +// withCacheSalt sets vLLM's cache_salt to one derived from key — the +// collections a request's passages came from — so requests that can see the +// same passages still share a cache, and nobody else's do: not another +// organization's, nor another team's of the same one. +func withCacheSalt(body []byte, key string) []byte { var top map[string]json.RawMessage if json.Unmarshal(body, &top) != nil { return body } - sum := sha256.Sum256([]byte("llmgw-kb\x00" + org)) + sum := sha256.Sum256([]byte("llmgw-kb\x00" + key)) top["cache_salt"], _ = json.Marshal(hex.EncodeToString(sum[:16])) out, err := json.Marshal(top) if err != nil { diff --git a/internal/router/rag_test.go b/internal/router/rag_test.go index 30d5e62..08614bb 100644 --- a/internal/router/rag_test.go +++ b/internal/router/rag_test.go @@ -52,9 +52,9 @@ func TestPassagesJoinTheLeadingSystemMessage(t *testing.T) { } } -func TestTheCacheSaltIsAnOrganizations(t *testing.T) { +func TestTheCacheSaltFollowsTheCollections(t *testing.T) { body := []byte(`{"model":"m","messages":[],"temperature":0}`) - a1, a2, b := withCacheSalt(body, "org-a"), withCacheSalt(body, "org-a"), withCacheSalt(body, "org-b") + a1, a2, b := withCacheSalt(body, "kb_a,kb_b"), withCacheSalt(body, "kb_a,kb_b"), withCacheSalt(body, "kb_a") var x, y struct { Salt string `json:"cache_salt"` T float64 `json:"temperature"` @@ -202,7 +202,7 @@ func mustID(t *testing.T, s *store.DB, name string) string { return "" } -func TestAVLLMReplicaGetsTheOrganizationsSaltOnlyWithPassages(t *testing.T) { +func TestAVLLMReplicaGetsASaltOnlyWithPassages(t *testing.T) { dsn := os.Getenv("LLMGW_TEST_POSTGRES") if dsn == "" { t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this") diff --git a/internal/router/router.go b/internal/router/router.go index 0059437..0b9fba1 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -713,7 +713,7 @@ func (r *Router) storeModels(models []catalog.Model) { // is applied to the request on the way there. for i, a := range entry.Adapters { m[a.Name] = entry - ads[a.Name] = adapterRef{base: entry.Name, index: i, scale: a.EffectiveScale(), + ads[a.Name] = adapterRef{base: entry.Name, owner: a.Owner, index: i, scale: a.EffectiveScale(), limit: entry.MaxLoadedAdapters, path: engines.RuntimeLoRAPath(a.Name, a.Source)} } } @@ -735,7 +735,9 @@ func (r *Router) storeModels(models []catalog.Model) { // adapterRef is where an adapter runs and how it is applied. type adapterRef struct { - base string + base string + // owner is who sees it (catalog.Adapter.Owner); empty is everyone. + owner string index int scale float64 // limit is the model's max_loaded_adapters: above zero, the adapter @@ -757,6 +759,33 @@ func (r *Router) adapter(name string) (adapterRef, bool) { return adapterRef{}, false } +// adapterHidden reports whether name is an adapter its owner keeps from this +// caller: one outside the part of the tree it belongs to, or one with no +// identity at all. Such a caller is told there is no such model. +func (r *Router) adapterHidden(p auth.Principal, authed bool, name string) bool { + ad, ok := r.adapter(name) + if !ok || ad.owner == "" { + return false + } + if !authed { + return true + } + scope, id, _ := strings.Cut(ad.owner, ":") + switch scope { + case "org": + return p.OrgID != id + case "team": + return p.TeamID != id + case "project": + return p.ProjectID != id + case "user": + return p.UserID != id + case "key": + return p.KeyID != id + } + return true +} + // LastUsed reports when each endpoint was last sent a request, satisfying // control.UsageSource so the control plane can tell a busy replica from one // holding its cards out of habit. @@ -1090,13 +1119,13 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa if isAdapter && !onDemand && dec.Endpoint.Engine != "vllm" { body = withAdapter(body, model, ad.index, ad.scale) } - // Passages of one organization's documents never share cached - // blocks with another's: a prefix found warm would say what another - // tenant asked about, by how fast its answer started. vLLM keys its - // prefix cache by this salt; other engines are not sent a field - // they may refuse. - if req.kbOrg != "" && dec.Endpoint.Engine == "vllm" && req.ContentType == "" { - body = withCacheSalt(body, req.kbOrg) + // Passages share cached blocks only with requests that could have + // been given the same ones: a prefix found warm would otherwise say + // what another tenant asked about, by how fast its answer started. + // vLLM keys its prefix cache by this salt; other engines are not + // sent a field they may refuse. + if req.kbSalt != "" && dec.Endpoint.Engine == "vllm" && req.ContentType == "" { + body = withCacheSalt(body, req.kbSalt) } // What computes the answer, said before the proxy writes the headers. // A model can have replicas on cards and on CPUs at once, and the two diff --git a/internal/store/migrations/postgres/0038_train_owner.sql b/internal/store/migrations/postgres/0038_train_owner.sql new file mode 100644 index 0000000..8960976 --- /dev/null +++ b/internal/store/migrations/postgres/0038_train_owner.sql @@ -0,0 +1,5 @@ +-- Who an adapter belongs to (BACKLOG F4): scope:ID in the tenancy tree, or +-- empty for every caller its model allows. An adapter trained on a +-- knowledge base's documents carries some of what they say, and belongs to +-- a part of the tree that may read them. +ALTER TABLE train_jobs ADD COLUMN owner TEXT NOT NULL DEFAULT ''; diff --git a/internal/store/migrations/sqlite/0038_train_owner.sql b/internal/store/migrations/sqlite/0038_train_owner.sql new file mode 100644 index 0000000..8960976 --- /dev/null +++ b/internal/store/migrations/sqlite/0038_train_owner.sql @@ -0,0 +1,5 @@ +-- Who an adapter belongs to (BACKLOG F4): scope:ID in the tenancy tree, or +-- empty for every caller its model allows. An adapter trained on a +-- knowledge base's documents carries some of what they say, and belongs to +-- a part of the tree that may read them. +ALTER TABLE train_jobs ADD COLUMN owner TEXT NOT NULL DEFAULT ''; diff --git a/internal/store/train.go b/internal/store/train.go index 6f0c7f5..bed35dd 100644 --- a/internal/store/train.go +++ b/internal/store/train.go @@ -15,6 +15,8 @@ type TrainJob struct { BaseRepo string BaseRevision string Node string + // Owner is who the adapter belongs to, scope:ID; empty is everyone. + Owner string GPUMemoryMiB int Params string GGUF bool @@ -57,21 +59,21 @@ func (s *DB) CreateTrainJob(ctx context.Context, j TrainJob) error { if j.GGUF { gguf = 1 } - _, err := s.db.ExecContext(ctx, `INSERT INTO train_jobs (id, name, model, base_repo, base_revision, node, gpu_memory_mib, - params, gguf, data, pairs, state, created_at, created_by) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - j.ID, j.Name, j.Model, j.BaseRepo, j.BaseRevision, j.Node, j.GPUMemoryMiB, j.Params, gguf, j.Data, j.Pairs, + _, err := s.db.ExecContext(ctx, `INSERT INTO train_jobs (id, name, model, base_repo, base_revision, node, owner, gpu_memory_mib, + params, gguf, data, pairs, state, created_at, created_by) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + j.ID, j.Name, j.Model, j.BaseRepo, j.BaseRevision, j.Node, j.Owner, j.GPUMemoryMiB, j.Params, gguf, j.Data, j.Pairs, TrainQueued, millis(j.CreatedAt), j.CreatedBy) return err } -const trainCols = `id, name, model, base_repo, base_revision, node, gpu_memory_mib, params, gguf, pairs, state, step, steps, +const trainCols = `id, name, model, base_repo, base_revision, node, owner, gpu_memory_mib, params, gguf, pairs, state, step, steps, loss, message, peak_mib, created_at, started_at, finished_at, created_by` func scanTrainJob(sc scanner) (TrainJob, error) { var j TrainJob var gguf int var created, started, finished int64 - err := sc.Scan(&j.ID, &j.Name, &j.Model, &j.BaseRepo, &j.BaseRevision, &j.Node, &j.GPUMemoryMiB, &j.Params, &gguf, + err := sc.Scan(&j.ID, &j.Name, &j.Model, &j.BaseRepo, &j.BaseRevision, &j.Node, &j.Owner, &j.GPUMemoryMiB, &j.Params, &gguf, &j.Pairs, &j.State, &j.Step, &j.Steps, &j.Loss, &j.Message, &j.PeakMiB, &created, &started, &finished, &j.CreatedBy) j.GGUF = gguf != 0 j.CreatedAt = fromMillis(created) @@ -175,3 +177,27 @@ func (s *DB) DeleteTrainJob(ctx context.Context, id string) error { return nil }) } + +// EffectiveFor is a scope's place in the tree, as a caller in it would have +// it, for asking what it may read: an organization, a team or a project +// with what is above it, a key as it resolves, a person in an organization. +func (s *DB) EffectiveFor(ctx context.Context, scope Scope, id, org string) (Effective, error) { + switch scope { + case ScopeOrg: + return Effective{OrgID: id}, nil + case ScopeTeam: + var o string + err := s.db.QueryRowContext(ctx, `SELECT org_id FROM teams WHERE id = ?`, id).Scan(&o) + if errors.Is(err, sql.ErrNoRows) { + return Effective{}, ErrNotFound + } + return Effective{TeamID: id, OrgID: o}, err + case ScopeProject: + return s.ResolveFor(ctx, "", "", id) + case ScopeKey: + return s.Resolve(ctx, id) + case ScopeUser: + return Effective{UserID: id, OrgID: org}, nil + } + return Effective{}, errors.New("no such scope " + string(scope)) +}