diff --git a/BACKLOG.md b/BACKLOG.md index b8f7388d..628fd47d 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 f07c7992..ef8d307e 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 79d7829e..a0f40779 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 95bd2404..62253a18 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 05efa247..4290e55b 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 3d83a51e..b4cc246a 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 a7473643..b21ff259 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 dccfbb20..a51ebef3 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 821ed52d..107b1d06 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 ee4f5281..4360eca4 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 61c0f686..8ca0ceef 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 287b2661..ca60eede 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 7bd5d2b7..b777799b 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 10cb97c7..2680ab24 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 2170afd1..38377568 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 d5978b57..fd276c84 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 3e38312b..8f9ba2ae 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 d2c830bd..ef25b3b7 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 00000000..98628383 --- /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 00000000..4655ff22 --- /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 05ad9e5b..1407db67 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 00000000..5daebc14 --- /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 bb885b33..0f0654c4 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 c5e93a77..76405f9e 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 30d5e625..08614bb3 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 00594371..0b9fba17 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 00000000..8960976c --- /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 00000000..8960976c --- /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 6f0c7f51..bed35dd1 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)) +}