diff --git a/BACKLOG.md b/BACKLOG.md index 90ca888..95631cb 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -431,7 +431,14 @@ portal is not done — the page shows; the console and ctl change. weighted routing across two revisions/aliases (scheduler filter, alias resolution in `router/router.go:569`), a header/key pin, automatic rollback on error-rate or TTFT regression from per-replica metrics. -8. **Batch API and Files (M–L).** `/v1/files`, `/v1/batches`, JSONL +8. **Done (branch `features-c2`):** `/v1/files` and `/v1/batches` + (OpenAI's shape), worker under a store lease, lines kept as served, + lowest band, usage tagged `batch=`, output and error files with + retention, `no_store` refused. Left: the Responses API (`/v1/responses`) + as a batch endpoint; the console and `ctl` do not list batches (an + operator sees them through usage tags); a file is one row in the store, + so a store backup carries batch files with it. + **Batch API and Files (M–L).** `/v1/files`, `/v1/batches`, JSONL storage, a worker feeding the bronze band, accounting. New `internal/router/batch`. 9. **Request/response logging and replay (M).** Opt-in per project with diff --git a/CHANGELOG.md b/CHANGELOG.md index 7bda153..b8c411b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -71,6 +71,21 @@ APIs change between minors without a deprecation period. - **Other ways in:** `GET /api/v1/usage/export` (CSV, or `format=json`) and **E** on the console's Usage view. +### Batches + +- **OpenAI's Batch and Files APIs.** `POST /v1/files` (purpose batch), + `POST /v1/batches`, and reading, listing, cancelling and deleting them, so + OpenAI's clients send batches here unchanged. + - **Serving:** each line goes as a request of the key that made the batch, + in the lowest admission band, a few at a time; rate limits are waited + out. Answers and failures are written to files kept 30 days. + - **Accounting:** each line is an ordinary usage record tagged + `batch=`. + - **Restarts:** answered lines are kept as they are answered, and any + gateway on the store carries a batch on under a lease. + - **Limits:** `router.batch` bounds the file, the lines and the + concurrency. Refused to a project under a `no_store` policy. + ### New weights tried on a share of the traffic first - **Canaries.** `llmgw ctl model canary MODEL --percent 10` fetches the newer diff --git a/docs/data.md b/docs/data.md index c103c23..6d3cf02 100644 --- a/docs/data.md +++ b/docs/data.md @@ -51,10 +51,17 @@ limits attached to it. Passwords are hashed. **Conversations and agent notes** — only for a project that asked. See below. +**Batch files** — when a client uses the Batch API: the file of requests it +uploaded (prompts) and the files the batch wrote (answers), in the store, +until deleted (`DELETE /v1/files/{id}`) or, for what a batch wrote, until +`router.batch.retain_days` (30) pass. Only keys of the project that made +them can read them, and a project under a `no_store` policy cannot make +them. + ## What is not stored Prompts. Completions. Anything the model was shown or produced, unless the -project turned on conversations or notes. +project turned on conversations or notes, or sent a batch. The chat page in llmgw's own web interface keeps the conversation in a JavaScript variable in the browser. Refresh it and the conversation is gone, diff --git a/docs/install.md b/docs/install.md index fae8cd8..ae7b9f6 100644 --- a/docs/install.md +++ b/docs/install.md @@ -1130,6 +1130,42 @@ curl http://gateway:8080/v1/chat/completions \ 100 values of each apart and the rest as `other`: a tag is the client's to set, and every value is a series. +## Batches + +Work that can wait — a night's classification, embeddings for a document +store — goes as a batch: OpenAI's Batch API, so its clients work unchanged. + +``` +curl http://gateway:8080/v1/files -H "Authorization: Bearer $LLMGW_API_KEY" \ + -F purpose=batch -F file=@requests.jsonl +curl http://gateway:8080/v1/batches -H "Authorization: Bearer $LLMGW_API_KEY" \ + -d '{"input_file_id": "file-…", "endpoint": "/v1/chat/completions", "completion_window": "24h"}' +curl http://gateway:8080/v1/batches/batch_… -H "Authorization: Bearer $LLMGW_API_KEY" +curl http://gateway:8080/v1/files/file-…/content -H "Authorization: Bearer $LLMGW_API_KEY" +``` + +Each line of the file is `{"custom_id": …, "method": "POST", "url": +"/v1/chat/completions", "body": {…}}`, at most 50 000 of them +(`router.batch.max_lines`) in at most 100 MB (`router.batch.max_file_bytes`). + +- **How it is served:** each line as a request of the key that made the + batch — its limits, its budget, its guardrails — in the lowest admission + band, behind every interactive request, `router.batch.concurrency` (4) at + a time. A rate limit or a model still loading is waited out; a spent + budget fails the lines left. +- **What comes back:** answers in an output file, failures in an error + file, both OpenAI's form, kept `router.batch.retain_days` (30). Each + request is a usage record tagged `batch=`, so `llmgw ctl request` + finds it. +- **Who sees it:** files and batches belong to the key's project; every key + of the project sees them, no other does. +- **Surviving a restart:** every answered line is kept as it is answered, + and whichever gateway holds the batch's lease carries on from there; a + control plane that stops hands its batches to another within two minutes. +- **What it keeps:** a batch is prompts and answers stored until deleted + or past their retention — a project under a guardrails policy with + `no_store` is refused it. + ## Conversations across models Off everywhere until a project asks. llmgw otherwise stores no message content diff --git a/docs/reference.md b/docs/reference.md index 575d192..0360627 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -43,6 +43,10 @@ rather than repeated. | `batch_size` | `int` | [Auth](#auth) | | `buffer_size` | `int` | [Auth](#auth) | | `usage_retention` | `time.Duration` | [Auth](#auth) | +| `max_file_bytes` | `int64` | [Batch](#batch) | +| `max_lines` | `int` | [Batch](#batch) | +| `concurrency` | `int` | [Batch](#batch) | +| `retain_days` | `int` | [Batch](#batch) | | `model` | `string` | [Canary](#canary) | | `percent` | `int` | [Canary](#canary) | | `min_requests` | `int` | [Canary](#canary) | @@ -198,6 +202,7 @@ rather than repeated. | `tenant_label` | `string` | [Router](#router) | | `tag_labels` | `[]string` | [Router](#router) | | `max_body_bytes` | `int64` | [Router](#router) | +| `batch` | `Batch` | [Router](#router) | | `trusted_proxies` | `[]string` | [Router](#router) | | `timeouts` | `Timeouts` | [Router](#router) | | `scoring` | `Scoring` | [Router](#router) | @@ -517,6 +522,35 @@ five-minute sums kept beside the records, which are not pruned, so this bounds the table that grows by a row a request without changing any figure a page shows. At least a day when set. +## Batch + +### `max_file_bytes` + +`int64` + +MaxFileBytes bounds one uploaded file; 100 MB unless set. + +### `max_lines` + +`int` + +MaxLines bounds the requests in one batch; 50 000 unless set. + +### `concurrency` + +`int` + +Concurrency is how many of one batch's requests are in flight at +once; 4 unless set. They queue in the lowest admission band, behind +every interactive request. + +### `retain_days` + +`int` + +RetainDays is how long a batch's answer and error files are kept; +30 unless set. Uploaded files are kept until deleted. + ## Canary ### `model` @@ -1622,6 +1656,16 @@ a series. The first 100 values of each are counted apart, the rest as MaxBodyBytes bounds the request body the router will read. +### `batch` + +`Batch` + +Batch is the Files and Batch APIs (/v1/files, /v1/batches): a file of +requests served at the lowest priority, its answers written to +another. Needs the store. + +Takes a [Batch](#batch) block. + ### `trusted_proxies` `[]string` diff --git a/internal/config/config.go b/internal/config/config.go index d19a2e2..fcc3379 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -453,6 +453,11 @@ type Router struct { // MaxBodyBytes bounds the request body the router will read. MaxBodyBytes int64 `koanf:"max_body_bytes"` + // Batch is the Files and Batch APIs (/v1/files, /v1/batches): a file of + // requests served at the lowest priority, its answers written to + // another. Needs the store. + Batch Batch `koanf:"batch"` + // TrustedProxies are the reverse proxies in front of the client port, // as addresses or CIDR prefixes. Their X-Forwarded-For, -Proto and -Host // are believed; anyone else's are not, since they decide the address @@ -553,6 +558,21 @@ type Endpoint struct { Priority int `koanf:"priority"` } +// Batch is how the Files and Batch APIs are bounded. +type Batch struct { + // MaxFileBytes bounds one uploaded file; 100 MB unless set. + MaxFileBytes int64 `koanf:"max_file_bytes"` + // MaxLines bounds the requests in one batch; 50 000 unless set. + MaxLines int `koanf:"max_lines"` + // Concurrency is how many of one batch's requests are in flight at + // once; 4 unless set. They queue in the lowest admission band, behind + // every interactive request. + Concurrency int `koanf:"concurrency"` + // RetainDays is how long a batch's answer and error files are kept; + // 30 unless set. Uploaded files are kept until deleted. + RetainDays int `koanf:"retain_days"` +} + // Default returns the configuration used when nothing is specified. func Default() Server { return Server{ diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 9d01106..3cdf078 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -38,6 +38,12 @@ var Reference = []Block{ {Key: "buffer_size", Type: "int", Doc: "BufferSize bounds the queue. Beyond it records are dropped rather than\nblocking a request: losing accounting is bad, turning a stalled database\ninto an outage is worse.", Nested: ""}, {Key: "usage_retention", Type: "time.Duration", Doc: "UsageRetention is how long the record of each request is kept: who sent\nit, which replica answered, how long it queued — what `ctl request`\nreads. Zero keeps it for ever. Totals, charts and budgets are read from\nfive-minute sums kept beside the records, which are not pruned, so this\nbounds the table that grows by a row a request without changing any\nfigure a page shows. At least a day when set.", Nested: ""}, }}, + {Package: "config", Name: "Batch", Title: "Batch", Doc: "", Fields: []Field{ + {Key: "max_file_bytes", Type: "int64", Doc: "MaxFileBytes bounds one uploaded file; 100 MB unless set.", Nested: ""}, + {Key: "max_lines", Type: "int", Doc: "MaxLines bounds the requests in one batch; 50 000 unless set.", Nested: ""}, + {Key: "concurrency", Type: "int", Doc: "Concurrency is how many of one batch's requests are in flight at\nonce; 4 unless set. They queue in the lowest admission band, behind\nevery interactive request.", Nested: ""}, + {Key: "retain_days", Type: "int", Doc: "RetainDays is how long a batch's answer and error files are kept;\n30 unless set. Uploaded files are kept until deleted.", Nested: ""}, + }}, {Package: "catalog", Name: "Canary", Title: "Canary", Doc: "", Fields: []Field{ {Key: "model", Type: "string", Doc: "Model is the candidate: a catalog model, usually this one at newer\nweights, which clients cannot name themselves.", Nested: ""}, {Key: "percent", Type: "int", Doc: "Percent of requests go to it, 0 to 100.", Nested: ""}, @@ -240,6 +246,7 @@ var Reference = []Block{ {Key: "tenant_label", Type: "string", Doc: "TenantLabel is the level of the tree the per-tenant request series are\nlabelled by — llmgw_tenant_requests_total and its siblings: \"org\",\n\"team\" or \"project\", or \"none\" for no such series. A label per project\non a cluster of thousands is more series than Prometheus wants, and a\nteam is what an alert is usually about; \"team\" is the default.", Nested: ""}, {Key: "tag_labels", Type: "[]string", Doc: "TagLabels are the request tags (metadata, x-llmgw-tags) counted in\nllmgw_tagged_requests_total and llmgw_tagged_tokens_total, by value.\nOnly these: a tag is the client's to set, and every value of one is\na series. The first 100 values of each are counted apart, the rest as\n\"other\". Empty counts none; the usage records keep every tag anyway.", Nested: ""}, {Key: "max_body_bytes", Type: "int64", Doc: "MaxBodyBytes bounds the request body the router will read.", Nested: ""}, + {Key: "batch", Type: "Batch", Doc: "Batch is the Files and Batch APIs (/v1/files, /v1/batches): a file of\nrequests served at the lowest priority, its answers written to\nanother. Needs the store.", Nested: "config.Batch"}, {Key: "trusted_proxies", Type: "[]string", Doc: "TrustedProxies are the reverse proxies in front of the client port,\nas addresses or CIDR prefixes. Their X-Forwarded-For, -Proto and -Host\nare believed; anyone else's are not, since they decide the address\nwritten into the setup scripts and whose attempts are counted where.", Nested: ""}, {Key: "timeouts", Type: "Timeouts", Doc: "", Nested: "config.Timeouts"}, {Key: "scoring", Type: "Scoring", Doc: "", Nested: "config.Scoring"}, diff --git a/internal/router/batches.go b/internal/router/batches.go new file mode 100644 index 0000000..55373ec --- /dev/null +++ b/internal/router/batches.go @@ -0,0 +1,255 @@ +package router + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strconv" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// The Batch API: a file of requests, served one by one at the lowest +// priority, its answers written to another file. OpenAI's shape, so its +// clients work unchanged. + +// batchEndpoints are what a batch may be of. +var batchEndpoints = map[string]bool{ + "/v1/chat/completions": true, + "/v1/completions": true, + "/v1/embeddings": true, +} + +// batchWindow is the only completion window, as OpenAI has it. +const batchWindow = "24h" + +func unixOrNil(t time.Time) any { + if t.IsZero() { + return nil + } + return t.Unix() +} + +func strOrNil(s string) any { + if s == "" { + return nil + } + return s +} + +func batchJSON(b store.Batch) map[string]any { + out := map[string]any{ + "id": b.ID, "object": "batch", "endpoint": b.Endpoint, "input_file_id": b.InputFileID, + "completion_window": b.CompletionWindow, "status": b.Status, + "output_file_id": strOrNil(b.OutputFileID), "error_file_id": strOrNil(b.ErrorFileID), + "created_at": b.CreatedAt.Unix(), "in_progress_at": unixOrNil(b.InProgressAt), + "expires_at": unixOrNil(b.ExpiresAt), "finalizing_at": unixOrNil(b.FinalizingAt), + "completed_at": unixOrNil(b.CompletedAt), "failed_at": unixOrNil(b.FailedAt), + "expired_at": unixOrNil(b.ExpiredAt), "cancelling_at": unixOrNil(b.CancellingAt), + "cancelled_at": unixOrNil(b.CancelledAt), + "request_counts": map[string]int{"total": b.Total, "completed": b.Completed, "failed": b.Failed}, + "errors": nil, "metadata": nil, + } + if b.Errors != "" { + out["errors"] = json.RawMessage(b.Errors) + } + if b.Metadata != "" { + out["metadata"] = json.RawMessage(b.Metadata) + } + return out +} + +func (r *Router) handleCreateBatch(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + var in struct { + InputFileID string `json:"input_file_id"` + Endpoint string `json:"endpoint"` + CompletionWindow string `json:"completion_window"` + Metadata json.RawMessage `json:"metadata"` + } + if err := json.NewDecoder(io.LimitReader(req.Body, 1<<16)).Decode(&in); err != nil { + writeError(w, http.StatusBadRequest, "invalid_request_error", "the request is not valid JSON: "+err.Error()) + return + } + if !batchEndpoints[in.Endpoint] { + writeError(w, http.StatusBadRequest, "invalid_request_error", + fmt.Sprintf("endpoint %q: a batch is of /v1/chat/completions, /v1/completions or /v1/embeddings", in.Endpoint)) + return + } + if in.CompletionWindow != batchWindow { + writeError(w, http.StatusBadRequest, "invalid_request_error", "completion_window is 24h") + return + } + meta := "" + if len(in.Metadata) > 0 && string(in.Metadata) != "null" { + if _, err := requestTags(in.Metadata, ""); err != nil { + writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } + meta = string(in.Metadata) + } + f, err := r.store.FileByID(req.Context(), in.InputFileID) + if errors.Is(err, store.ErrNotFound) || err == nil && (f.Owner != ownerOf(p) || f.Purpose != purposeBatch) { + writeErrorCode(w, http.StatusBadRequest, "invalid_request_error", "file_not_found", + fmt.Sprintf("no batch file %q", in.InputFileID)) + return + } + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + now := time.Now() + b := store.Batch{ID: newID("batch_"), Owner: ownerOf(p), KeyID: p.KeyID, Endpoint: in.Endpoint, + InputFileID: f.ID, CompletionWindow: batchWindow, Status: store.BatchValidating, + CreatedAt: now, ExpiresAt: now.Add(24 * time.Hour), Metadata: meta} + if err := r.store.PutBatch(req.Context(), b); err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + r.wakeBatches() + writeJSON(w, http.StatusOK, batchJSON(b)) +} + +// ownedBatch is a batch the caller owns, or a 404 said. +func (r *Router) ownedBatch(w http.ResponseWriter, req *http.Request, owner string) (store.Batch, bool) { + b, err := r.store.BatchByID(req.Context(), req.PathValue("id")) + if errors.Is(err, store.ErrNotFound) || err == nil && b.Owner != owner { + writeErrorCode(w, http.StatusNotFound, "invalid_request_error", "batch_not_found", + fmt.Sprintf("no batch %q", req.PathValue("id"))) + return store.Batch{}, false + } + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return store.Batch{}, false + } + return b, true +} + +func (r *Router) handleGetBatch(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + if b, ok := r.ownedBatch(w, req, ownerOf(p)); ok { + writeJSON(w, http.StatusOK, batchJSON(b)) + } +} + +func (r *Router) handleCancelBatch(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + b, ok := r.ownedBatch(w, req, ownerOf(p)) + if !ok { + return + } + b, err := r.store.CancelBatch(req.Context(), b.ID, time.Now()) + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + r.wakeBatches() + writeJSON(w, http.StatusOK, batchJSON(b)) +} + +func (r *Router) handleListBatches(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + limit, _ := strconv.Atoi(req.URL.Query().Get("limit")) + if limit <= 0 || limit > 100 { + limit = 20 + } + bs, err := r.store.Batches(req.Context(), ownerOf(p), req.URL.Query().Get("after"), limit+1) + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + more := len(bs) > limit + if more { + bs = bs[:limit] + } + data := make([]map[string]any, 0, len(bs)) + for _, b := range bs { + data = append(data, batchJSON(b)) + } + out := map[string]any{"object": "list", "data": data, "has_more": more, "first_id": nil, "last_id": nil} + if len(bs) > 0 { + out["first_id"], out["last_id"] = bs[0].ID, bs[len(bs)-1].ID + } + writeJSON(w, http.StatusOK, out) +} + +// batchLine is one line of a batch's input file. +type batchLine struct { + CustomID string `json:"custom_id"` + Method string `json:"method"` + URL string `json:"url"` + Body json.RawMessage `json:"body"` +} + +// lineError is one reason a batch's input was refused, as OpenAI lists them. +type lineError struct { + Code string `json:"code"` + Message string `json:"message"` + Line int `json:"line,omitempty"` +} + +// parseBatchInput reads a batch's input file, or says every line that is +// wrong with it. +func parseBatchInput(content []byte, endpoint string, maxLines int) ([]batchLine, []lineError) { + var lines []batchLine + var errs []lineError + seen := map[string]bool{} + n := 0 + for raw := range bytes.SplitSeq(content, []byte("\n")) { + n++ + raw = bytes.TrimSpace(raw) + if len(raw) == 0 { + continue + } + bad := func(code, msg string) { + if len(errs) < 100 { + errs = append(errs, lineError{Code: code, Message: msg, Line: n}) + } + } + var l batchLine + if err := json.Unmarshal(raw, &l); err != nil { + bad("invalid_json_line", "not a JSON object: "+err.Error()) + continue + } + var body map[string]any + switch { + case l.CustomID == "": + bad("missing_custom_id", "custom_id is required") + case seen[l.CustomID]: + bad("duplicate_custom_id", fmt.Sprintf("custom_id %q is used twice", l.CustomID)) + case l.Method != http.MethodPost: + bad("invalid_method", "method is POST") + case l.URL != endpoint: + bad("mismatched_endpoint", fmt.Sprintf("url %q is not the batch's endpoint, %s", l.URL, endpoint)) + case json.Unmarshal(l.Body, &body) != nil || body["model"] == nil: + bad("invalid_body", "body is an object naming a model") + default: + seen[l.CustomID] = true + lines = append(lines, l) + } + } + if len(lines)+len(errs) == 0 { + errs = append(errs, lineError{Code: "empty_file", Message: "the file has no requests"}) + } + if len(lines) > maxLines { + errs = append(errs, lineError{Code: "too_many_requests", + Message: fmt.Sprintf("%d requests; a batch holds at most %d", len(lines), maxLines)}) + } + return lines, errs +} diff --git a/internal/router/batches_test.go b/internal/router/batches_test.go new file mode 100644 index 0000000..7048a4f --- /dev/null +++ b/internal/router/batches_test.go @@ -0,0 +1,241 @@ +package router + +import ( + "bytes" + "crypto/sha256" + "encoding/json" + "fmt" + "mime/multipart" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +func upload(t *testing.T, h http.Handler, secret, content string) (int, map[string]any) { + t.Helper() + var buf bytes.Buffer + mw := multipart.NewWriter(&buf) + _ = mw.WriteField("purpose", "batch") + fw, _ := mw.CreateFormFile("file", "in.jsonl") + _, _ = fw.Write([]byte(content)) + _ = mw.Close() + req := httptest.NewRequest(http.MethodPost, "/v1/files", &buf) + req.Header.Set("Content-Type", mw.FormDataContentType()) + req.Header.Set("Authorization", "Bearer "+secret) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + var out map[string]any + _ = json.Unmarshal(rec.Body.Bytes(), &out) + return rec.Code, out +} + +func call(t *testing.T, h http.Handler, secret, method, path string, body any) (int, map[string]any, string) { + t.Helper() + var rd *bytes.Reader + if body != nil { + b, _ := json.Marshal(body) + rd = bytes.NewReader(b) + } else { + rd = bytes.NewReader(nil) + } + req := httptest.NewRequest(method, path, rd) + req.Header.Set("Authorization", "Bearer "+secret) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + var out map[string]any + _ = json.Unmarshal(rec.Body.Bytes(), &out) + return rec.Code, out, rec.Body.String() +} + +func waitBatch(t *testing.T, h http.Handler, secret, id string) map[string]any { + t.Helper() + deadline := time.Now().Add(10 * time.Second) + for { + _, b, _ := call(t, h, secret, http.MethodGet, "/v1/batches/"+id, nil) + if store.BatchDone(fmt.Sprint(b["status"])) { + return b + } + if time.Now().After(deadline) { + t.Fatalf("the batch did not end: %v", b) + } + time.Sleep(50 * time.Millisecond) + } +} + +func TestABatchIsServedAndItsAnswersWritten(t *testing.T) { + e := newFakeEngine(t, "qwen2.5-14b-instruct") + rt, secret := pricedRouter(t, e.srv.URL) + h := rt.Handler() + + var in strings.Builder + for i := range 5 { + fmt.Fprintf(&in, `{"custom_id":"r%d","method":"POST","url":"/v1/chat/completions","body":{"model":"q14","messages":[{"role":"user","content":"hi"}]}}`+"\n", i) + } + // One the engine cannot serve: a model that is not there. + in.WriteString(`{"custom_id":"bad","method":"POST","url":"/v1/chat/completions","body":{"model":"nope","messages":[]}}` + "\n") + code, f := upload(t, h, secret, in.String()) + if code != 200 || f["purpose"] != "batch" { + t.Fatalf("upload: %d %v", code, f) + } + code, b, raw := call(t, h, secret, http.MethodPost, "/v1/batches", map[string]any{ + "input_file_id": f["id"], "endpoint": "/v1/chat/completions", "completion_window": "24h", + "metadata": map[string]string{"job": "nightly"}}) + if code != 200 || b["status"] != "validating" { + t.Fatalf("create: %d %s", code, raw) + } + b = waitBatch(t, h, secret, b["id"].(string)) + counts := b["request_counts"].(map[string]any) + if b["status"] != "completed" || counts["total"] != 6.0 || counts["completed"] != 5.0 || counts["failed"] != 1.0 { + t.Fatalf("ended %v", b) + } + if e.calls.Load() != 5 { + t.Errorf("the engine served %d", e.calls.Load()) + } + _, _, out := call(t, h, secret, http.MethodGet, "/v1/files/"+b["output_file_id"].(string)+"/content", nil) + lines := strings.Split(strings.TrimSpace(out), "\n") + if len(lines) != 5 || !strings.Contains(lines[0], `"custom_id":"r`) || !strings.Contains(lines[0], `"status_code":200`) { + t.Errorf("output:\n%s", out) + } + _, _, errs := call(t, h, secret, http.MethodGet, "/v1/files/"+b["error_file_id"].(string)+"/content", nil) + if !strings.Contains(errs, `"custom_id":"bad"`) { + t.Errorf("errors:\n%s", errs) + } + // Charged like any request, and findable by the batch. + var rid string + for _, l := range lines { + var o struct { + Response struct { + RequestID string `json:"request_id"` + } `json:"response"` + } + _ = json.Unmarshal([]byte(l), &o) + rid = o.Response.RequestID + } + deadline := time.Now().Add(5 * time.Second) + for { + u, ok, _ := rt.store.UsageByRequest(t.Context(), rid) + if ok { + if u.Tags["batch"] != b["id"] || u.CostMicros == 0 { + t.Errorf("usage %+v", u) + } + break + } + if time.Now().After(deadline) { + t.Fatal("a batch request left no usage") + } + time.Sleep(20 * time.Millisecond) + } + _, list, _ := call(t, h, secret, http.MethodGet, "/v1/batches", nil) + if d := list["data"].([]any); len(d) != 1 { + t.Errorf("listed %v", list) + } +} + +func TestABatchWithBadLinesFailsAndSaysWhich(t *testing.T) { + e := newFakeEngine(t, "qwen2.5-14b-instruct") + rt, secret := pricedRouter(t, e.srv.URL) + h := rt.Handler() + _, f := upload(t, h, secret, `{"custom_id":"a","method":"POST","url":"/v1/embeddings","body":{"model":"q14"}} +not json +{"custom_id":"a","method":"POST","url":"/v1/chat/completions","body":{"model":"q14"}} +`) + _, b, _ := call(t, h, secret, http.MethodPost, "/v1/batches", map[string]any{ + "input_file_id": f["id"], "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + b = waitBatch(t, h, secret, b["id"].(string)) + j, _ := json.Marshal(b["errors"]) + if b["status"] != "failed" || !strings.Contains(string(j), "mismatched_endpoint") || !strings.Contains(string(j), "invalid_json_line") { + t.Errorf("%v %s", b["status"], j) + } + if e.calls.Load() != 0 { + t.Error("a batch that failed validation was served") + } +} + +func TestBatchesAndFilesAreTheirOwners(t *testing.T) { + e := newFakeEngine(t, "qwen2.5-14b-instruct") + rt, secret := pricedRouter(t, e.srv.URL) + h := rt.Handler() + _, f := upload(t, h, secret, `{"custom_id":"a","method":"POST","url":"/v1/chat/completions","body":{"model":"q14"}}`) + + // Another project's key sees none of it. + org, _ := rt.store.CreateOrg(t.Context(), "other") + team, _ := rt.store.CreateTeam(t.Context(), org, "t") + project, _ := rt.store.CreateProject(t.Context(), team, "p") + other := store.KeyPrefix + strings.Repeat("b", 43) + sum := sha256.Sum256([]byte(other)) + if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-2", Hash: sum[:], Prefix: "bbbb", Name: "other", + ProjectID: project, CreatedAt: time.Now()}); err != nil { + t.Fatal(err) + } + if code, _, _ := call(t, h, other, http.MethodGet, "/v1/files/"+f["id"].(string), nil); code != 404 { + t.Errorf("another project read a file: %d", code) + } + if code, _, _ := call(t, h, other, http.MethodPost, "/v1/batches", map[string]any{ + "input_file_id": f["id"], "endpoint": "/v1/chat/completions", "completion_window": "24h"}); code != 400 { + t.Errorf("another project batched a file: %d", code) + } + if code, _, _ := call(t, h, secret, http.MethodDelete, "/v1/files/"+f["id"].(string), nil); code != 200 { + t.Errorf("delete: %d", code) + } +} + +func TestABatchPicksUpWhereItStoppedAndCanBeCancelled(t *testing.T) { + e := newFakeEngine(t, "qwen2.5-14b-instruct") + rt, secret := pricedRouter(t, e.srv.URL) + h := rt.Handler() + var in strings.Builder + for i := range 4 { + fmt.Fprintf(&in, `{"custom_id":"r%d","method":"POST","url":"/v1/chat/completions","body":{"model":"q14","messages":[]}}`+"\n", i) + } + _, f := upload(t, h, secret, in.String()) + + // As a gateway that died after serving the first line left it. + now := time.Now() + b := store.Batch{ID: "batch_resumed", Owner: "project:" + mustProjectOf(t, rt, "key-1"), KeyID: "key-1", + Endpoint: "/v1/chat/completions", InputFileID: f["id"].(string), CompletionWindow: "24h", + Status: store.BatchInProgress, CreatedAt: now, ExpiresAt: now.Add(24 * time.Hour)} + if err := rt.store.PutBatch(t.Context(), b); err != nil { + t.Fatal(err) + } + b.Total, b.Completed, b.InProgressAt = 4, 1, now + _ = rt.store.UpdateBatch(t.Context(), b) + _ = rt.store.PutBatchResult(t.Context(), b.ID, store.BatchResult{Line: 0, OK: true, + Data: `{"id":"batch_req_x","custom_id":"r0","response":{"status_code":200},"error":null}`}) + rt.wakeBatches() + got := waitBatch(t, h, secret, b.ID) + if c := got["request_counts"].(map[string]any); got["status"] != "completed" || c["completed"] != 4.0 { + t.Fatalf("%v", got) + } + if e.calls.Load() != 3 { + t.Errorf("served %d; the line kept already was served again", e.calls.Load()) + } + + // Cancelled before it is served: nothing is sent. + b2 := b + b2.ID, b2.Status = "batch_cancelled", store.BatchValidating + if err := rt.store.PutBatch(t.Context(), b2); err != nil { + t.Fatal(err) + } + _, _ = rt.store.CancelBatch(t.Context(), b2.ID, time.Now()) + rt.wakeBatches() + if got := waitBatch(t, h, secret, b2.ID); got["status"] != "cancelled" { + t.Errorf("%v", got) + } + if e.calls.Load() != 3 { + t.Errorf("a cancelled batch was served: %d", e.calls.Load()) + } +} + +func mustProjectOf(t *testing.T, rt *Router, keyID string) string { + t.Helper() + k, err := rt.store.KeyByID(t.Context(), keyID) + if err != nil { + t.Fatal(err) + } + return k.ProjectID +} diff --git a/internal/router/batchworker.go b/internal/router/batchworker.go new file mode 100644 index 0000000..1ba0771 --- /dev/null +++ b/internal/router/batchworker.go @@ -0,0 +1,373 @@ +package router + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "strconv" + "sync" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/auth" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// The worker that serves batches. Any gateway on the store may run one: a +// batch is worked on by whichever holds its lease, and each finished line +// is kept as it is served, so another picks up where one stopped. + +const ( + batchLease = 2 * time.Minute + batchPoll = 5 * time.Second +) + +// wakeBatches has the worker look for work now rather than at its next poll. +func (r *Router) wakeBatches() { + select { + case r.batchWake <- struct{}{}: + default: + } +} + +func (r *Router) runBatches(ctx context.Context) { + host, _ := os.Hostname() + owner := fmt.Sprintf("%s:%d:%s", host, os.Getpid(), newID("")) + t := time.NewTicker(batchPoll) + defer t.Stop() + for { + for { + now := time.Now() + b, ok, err := r.store.ClaimBatch(ctx, owner, now, now.Add(batchLease)) + if err != nil { + if ctx.Err() == nil { + r.log.Warn("could not look for batches to serve", "err", err) + } + break + } + if !ok { + break + } + r.workBatch(ctx, owner, b) + } + if _, _, _, days := r.batchCfg(); days > 0 { + if n, err := r.store.DeleteExpiredFiles(ctx, time.Now()); err == nil && n > 0 { + r.log.Info("batch files past their retention deleted", "files", n) + } + } + select { + case <-ctx.Done(): + return + case <-t.C: + case <-r.batchWake: + } + } +} + +// workBatch takes one batch as far as it goes: validated, served, written. +func (r *Router) workBatch(ctx context.Context, owner string, b store.Batch) { + bctx, stop := context.WithCancel(ctx) + defer stop() + defer func() { _ = r.store.ReleaseBatch(context.WithoutCancel(ctx), b.ID, owner) }() + // The lease kept while working; lost, the work stops. + go func() { + t := time.NewTicker(batchLease / 4) + defer t.Stop() + for { + select { + case <-bctx.Done(): + return + case <-t.C: + if held, err := r.store.RenewBatchLease(bctx, b.ID, owner, time.Now().Add(batchLease)); err == nil && !held { + stop() + return + } + } + } + }() + log := r.log.With("batch", b.ID) + if err := r.serveBatch(bctx, &b); err != nil && bctx.Err() == nil { + log.Warn("a batch stopped with an error; it is picked up again", "err", err) + } +} + +func (r *Router) serveBatch(ctx context.Context, b *store.Batch) error { + content, err := r.store.FileContent(ctx, b.InputFileID) + if err != nil && !errors.Is(err, store.ErrNotFound) { + return err + } + _, maxLines, concurrency, _ := r.batchCfg() + lines, errs := parseBatchInput(content, b.Endpoint, maxLines) + if errors.Is(err, store.ErrNotFound) { + errs = []lineError{{Code: "file_not_found", Message: "the input file was deleted"}} + } + + if b.Status == store.BatchValidating { + if len(errs) > 0 { + ej, _ := json.Marshal(map[string]any{"object": "list", "data": errs}) + b.Status, b.FailedAt, b.Errors = store.BatchFailed, time.Now(), string(ej) + return r.store.UpdateBatch(ctx, *b) + } + b.Status, b.Total, b.InProgressAt = store.BatchInProgress, len(lines), time.Now() + if err := r.store.UpdateBatch(ctx, *b); err != nil { + return err + } + } + + if b.Status == store.BatchInProgress { + end, err := r.serveLines(ctx, b, lines, concurrency) + if err != nil { + return err + } + switch end { + case store.BatchCancelling, store.BatchExpired: + b.Status = end + default: + b.Status, b.FinalizingAt = store.BatchFinalizing, time.Now() + if err := r.store.UpdateBatch(ctx, *b); err != nil { + return err + } + } + } + + // Finalizing, cancelling or expired: what was served is written out. + return r.finishBatch(ctx, b, lines) +} + +// serveLines serves the lines not yet served, a few at a time, and says why +// it stopped: "" when every line is served. +func (r *Router) serveLines(ctx context.Context, b *store.Batch, lines []batchLine, concurrency int) (string, error) { + done, err := r.store.BatchResults(ctx, b.ID) + if err != nil { + return "", err + } + served := map[int]bool{} + for _, d := range done { + served[d.Line] = true + } + p, perr := r.batchPrincipalOf(ctx, b.KeyID) + + var mu sync.Mutex + stopWith := "" + lastCheck := time.Now() + sem := make(chan struct{}, concurrency) + var wg sync.WaitGroup + for i, l := range lines { + if served[i] { + continue + } + mu.Lock() + stopped := stopWith != "" + mu.Unlock() + if stopped || ctx.Err() != nil { + break + } + // Asked to stop, or out of time: looked at now and then, not per line. + if time.Since(lastCheck) > 2*time.Second { + lastCheck = time.Now() + if cur, err := r.store.BatchByID(ctx, b.ID); err == nil && cur.Status == store.BatchCancelling { + mu.Lock() + stopWith = store.BatchCancelling + mu.Unlock() + break + } + if time.Now().After(b.ExpiresAt) { + mu.Lock() + stopWith = store.BatchExpired + mu.Unlock() + break + } + } + sem <- struct{}{} + wg.Add(1) + go func(i int, l batchLine) { + defer wg.Done() + defer func() { <-sem }() + res := r.serveLine(ctx, b, p, perr, l) + if ctx.Err() != nil { + return + } + if err := r.store.PutBatchResult(ctx, b.ID, res.withLine(i)); err != nil { + r.log.Warn("a batch's answer could not be kept", "batch", b.ID, "err", err) + return + } + mu.Lock() + if res.OK { + b.Completed++ + } else { + b.Failed++ + } + _ = r.store.UpdateBatch(ctx, *b) + mu.Unlock() + }(i, l) + } + wg.Wait() + if ctx.Err() != nil { + return "", ctx.Err() + } + if stopWith == "" { + // Counted from what is kept, which is the truth after a restart. + rs, err := r.store.BatchResults(ctx, b.ID) + if err != nil { + return "", err + } + if len(rs) < len(lines) { + return "", fmt.Errorf("%d of %d lines served", len(rs), len(lines)) + } + } + return stopWith, nil +} + +// batchPrincipalOf is the key a batch was made with, as it stands now: a key +// revoked since serves no more of it. +func (r *Router) batchPrincipalOf(ctx context.Context, keyID string) (auth.Principal, error) { + k, err := r.store.KeyByID(ctx, keyID) + if err != nil { + return auth.Principal{}, err + } + switch { + case !k.RevokedAt.IsZero(): + return auth.Principal{}, auth.ErrRevoked + case !k.ExpiresAt.IsZero() && time.Now().After(k.ExpiresAt): + return auth.Principal{}, auth.ErrExpired + } + eff, err := r.store.Resolve(ctx, k.ID) + if err != nil { + return auth.Principal{}, err + } + return auth.Principal{ + KeyID: k.ID, KeyPrefix: k.Prefix, UserID: eff.UserID, ProjectID: eff.ProjectID, + TeamID: eff.TeamID, OrgID: eff.OrgID, + Band: r.batchBand(), RoutingProfile: eff.RoutingProfile, Limits: eff, + }, nil +} + +// batchBand is the lowest admission band: a batch waits behind everyone. +func (r *Router) batchBand() string { + band, prio := "", 0 + for _, b := range r.cfg.Admission.Bands { + if band == "" || b.Priority < prio { + band, prio = b.Name, b.Priority + } + } + return band +} + +// lineResult is a served line as it is kept. +type lineResult struct { + OK bool + Data string +} + +func (l lineResult) withLine(i int) store.BatchResult { + return store.BatchResult{Line: i, OK: l.OK, Data: l.Data} +} + +// serveLine sends one request the way a client's is sent — limits, +// accounting, guardrails, admission in the lowest band — and keeps its +// answer in OpenAI's batch output form. A limit that lifts with time is +// waited out; one that does not fails the line. +func (r *Router) serveLine(ctx context.Context, b *store.Batch, p auth.Principal, perr error, l batchLine) lineResult { + id := newID("batch_req_") + out := map[string]any{"id": id, "custom_id": l.CustomID, "response": nil, "error": nil} + fail := func(code, msg string) lineResult { + out["error"] = map[string]string{"code": code, "message": msg} + j, _ := json.Marshal(out) + return lineResult{Data: string(j)} + } + if perr != nil { + return fail("key_unusable", "the key this batch was made with can no longer be used: "+perr.Error()) + } + body := withoutField(l.Body, "stream") + handler := r.handleInference(b.Endpoint) + for attempt := 1; ; attempt++ { + hreq, err := http.NewRequestWithContext(auth.WithPrincipal(ctx, p), http.MethodPost, b.Endpoint, bytes.NewReader(body)) + if err != nil { + return fail("invalid_request", err.Error()) + } + hreq.Header.Set("Content-Type", "application/json") + hreq.Header.Set("X-Request-Id", id) + hreq.Header.Set(tagsHeader, "batch="+b.ID) + w := &bufferWriter{header: http.Header{}} + handler(w, hreq) + status := w.status + if status == 0 { + status = http.StatusOK + } + // Waited out: a rate limit, a model loading, a queue full. Not a + // spent budget, which says no Retry-After because waiting does not + // help, and not past the batch's time. + if wait := w.header.Get("Retry-After"); wait != "" && (status == http.StatusTooManyRequests || status >= 500) && + time.Now().Before(b.ExpiresAt) && attempt < 1000 { + secs, _ := strconv.Atoi(wait) + select { + case <-ctx.Done(): + return fail("cancelled", "the gateway stopped") + case <-time.After(time.Duration(max(secs, 1)) * time.Second): + } + continue + } + var parsed any = json.RawMessage(bytes.TrimSpace(w.buf.Bytes())) + if !json.Valid(w.buf.Bytes()) { + parsed = w.buf.String() + } + out["response"] = map[string]any{"status_code": status, "request_id": id, "body": parsed} + j, _ := json.Marshal(out) + return lineResult{OK: status < 300, Data: string(j)} + } +} + +// finishBatch writes what was served — answers to one file, failures to +// another — and ends the batch: completed, cancelled or expired. +func (r *Router) finishBatch(ctx context.Context, b *store.Batch, lines []batchLine) error { + rs, err := r.store.BatchResults(ctx, b.ID) + if err != nil { + return err + } + var okBuf, errBuf bytes.Buffer + b.Completed, b.Failed = 0, 0 + for _, res := range rs { + if res.OK { + okBuf.WriteString(res.Data + "\n") + b.Completed++ + } else { + errBuf.WriteString(res.Data + "\n") + b.Failed++ + } + } + _, _, _, days := r.batchCfg() + now := time.Now() + keep := func(buf *bytes.Buffer, name string) (string, error) { + if buf.Len() == 0 { + return "", nil + } + f := store.File{ID: newID("file-"), Owner: b.Owner, Purpose: purposeBatchOutput, + Filename: b.ID + "_" + name + ".jsonl", CreatedAt: now, ExpiresAt: now.AddDate(0, 0, days)} + return f.ID, r.store.PutFile(ctx, f, buf.Bytes()) + } + if b.OutputFileID == "" { + if b.OutputFileID, err = keep(&okBuf, "output"); err != nil { + return err + } + } + if b.ErrorFileID == "" { + if b.ErrorFileID, err = keep(&errBuf, "error"); err != nil { + return err + } + } + switch b.Status { + case store.BatchCancelling: + b.Status, b.CancelledAt = store.BatchCancelled, now + case store.BatchExpired: + b.ExpiredAt = now + default: + b.Status, b.CompletedAt = store.BatchCompleted, now + } + if err := r.store.UpdateBatch(ctx, *b); err != nil { + return err + } + return r.store.DeleteBatchResults(ctx, b.ID) +} diff --git a/internal/router/files.go b/internal/router/files.go new file mode 100644 index 0000000..f429229 --- /dev/null +++ b/internal/router/files.go @@ -0,0 +1,207 @@ +package router + +import ( + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "io" + "net/http" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/auth" + "gitea.mynodes.xyz/ilker/llmgw/internal/guard" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// The Files API, as far as batches need it: a client uploads a file of +// requests, and reads back the files a batch writes. + +const ( + purposeBatch = "batch" + purposeBatchOutput = "batch_output" +) + +func newID(prefix string) string { + var b [12]byte + _, _ = rand.Read(b[:]) + return prefix + hex.EncodeToString(b[:]) +} + +// ownerOf is who a file or a batch belongs to: the key's project, so every +// key of it sees them, or the person for a personal key. +func ownerOf(p auth.Principal) string { + switch { + case p.ProjectID != "": + return "project:" + p.ProjectID + case p.UserID != "": + return "user:" + p.UserID + } + return "key:" + p.KeyID +} + +// batchPrincipal is the caller, or a refusal when there is no store to keep +// files and batches in. +func (r *Router) batchPrincipal(w http.ResponseWriter, req *http.Request) (auth.Principal, bool) { + p, ok := auth.FromContext(req.Context()) + if !ok || r.store == nil { + writeError(w, http.StatusNotImplemented, "invalid_request_error", + "files and batches need authentication and the store, which this gateway runs without") + return auth.Principal{}, false + } + // A batch's files are prompts and answers kept: not under a policy + // that keeps nothing. + if r.noStore(guard.Subject{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID, + TeamID: p.TeamID, OrgID: p.OrgID, Path: req.URL.Path}) { + writeError(w, http.StatusForbidden, "invalid_request_error", + "the gateway's policy for this caller keeps nothing, and a batch is prompts and answers kept") + return auth.Principal{}, false + } + return p, true +} + +func (r *Router) batchCfg() (maxBytes int64, maxLines, concurrency, retainDays int) { + b := r.cfg.Batch + maxBytes, maxLines, concurrency, retainDays = b.MaxFileBytes, b.MaxLines, b.Concurrency, b.RetainDays + if maxBytes <= 0 { + maxBytes = 100 << 20 + } + if maxLines <= 0 { + maxLines = 50000 + } + if concurrency <= 0 { + concurrency = 4 + } + if retainDays <= 0 { + retainDays = 30 + } + return +} + +func fileJSON(f store.File) map[string]any { + out := map[string]any{"id": f.ID, "object": "file", "bytes": f.Bytes, "created_at": f.CreatedAt.Unix(), + "filename": f.Filename, "purpose": f.Purpose, "status": "processed"} + if !f.ExpiresAt.IsZero() { + out["expires_at"] = f.ExpiresAt.Unix() + } + return out +} + +func (r *Router) handleUploadFile(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + maxBytes, _, _, _ := r.batchCfg() + req.Body = http.MaxBytesReader(w, req.Body, maxBytes+1<<20) + if err := req.ParseMultipartForm(32 << 20); err != nil { + writeError(w, http.StatusBadRequest, "invalid_request_error", + fmt.Sprintf("a multipart form with purpose and file, at most %d MB: %v", maxBytes>>20, err)) + return + } + if purpose := req.FormValue("purpose"); purpose != purposeBatch { + writeError(w, http.StatusBadRequest, "invalid_request_error", + fmt.Sprintf("purpose %q: this gateway keeps files for batches only (purpose batch)", purpose)) + return + } + file, header, err := req.FormFile("file") + if err != nil { + writeError(w, http.StatusBadRequest, "invalid_request_error", "no file in the form: "+err.Error()) + return + } + defer file.Close() + content, err := io.ReadAll(io.LimitReader(file, maxBytes+1)) + if err != nil { + writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } + if int64(len(content)) > maxBytes { + writeError(w, http.StatusRequestEntityTooLarge, "invalid_request_error", + fmt.Sprintf("the file is larger than %d MB", maxBytes>>20)) + return + } + f := store.File{ID: newID("file-"), Owner: ownerOf(p), Purpose: purposeBatch, + Filename: header.Filename, Bytes: int64(len(content)), CreatedAt: time.Now()} + if err := r.store.PutFile(req.Context(), f, content); err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + writeJSON(w, http.StatusOK, fileJSON(f)) +} + +// ownedFile is a file the caller owns, or a 404 said. +func (r *Router) ownedFile(w http.ResponseWriter, req *http.Request, p auth.Principal) (store.File, bool) { + f, err := r.store.FileByID(req.Context(), req.PathValue("id")) + if errors.Is(err, store.ErrNotFound) || err == nil && f.Owner != ownerOf(p) { + writeErrorCode(w, http.StatusNotFound, "invalid_request_error", "file_not_found", + fmt.Sprintf("no file %q", req.PathValue("id"))) + return store.File{}, false + } + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return store.File{}, false + } + return f, true +} + +func (r *Router) handleListFiles(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + fs, err := r.store.Files(req.Context(), ownerOf(p), req.URL.Query().Get("purpose"), 0) + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + data := make([]map[string]any, 0, len(fs)) + for _, f := range fs { + data = append(data, fileJSON(f)) + } + writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": data}) +} + +func (r *Router) handleGetFile(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + if f, ok := r.ownedFile(w, req, p); ok { + writeJSON(w, http.StatusOK, fileJSON(f)) + } +} + +func (r *Router) handleFileContent(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + f, ok := r.ownedFile(w, req, p) + if !ok { + return + } + b, err := r.store.FileContent(req.Context(), f.ID) + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + w.Header().Set("Content-Type", "application/octet-stream") + w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", f.Filename)) + _, _ = w.Write(b) +} + +func (r *Router) handleDeleteFile(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + f, ok := r.ownedFile(w, req, p) + if !ok { + return + } + if _, err := r.store.DeleteFile(req.Context(), f.ID); err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + writeJSON(w, http.StatusOK, map[string]any{"id": f.ID, "object": "file", "deleted": true}) +} diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 7cc1a4c..790a2bd 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -100,6 +100,16 @@ func (r *Router) Handler() http.Handler { // chat completions on the way in and back on the way out. mux.HandleFunc("POST /v1/responses", r.handleResponses) mux.HandleFunc("GET /v1/models", r.handleModels) + // Files and batches: a file of requests served at the lowest priority. + mux.HandleFunc("POST /v1/files", r.handleUploadFile) + mux.HandleFunc("GET /v1/files", r.handleListFiles) + mux.HandleFunc("GET /v1/files/{id}", r.handleGetFile) + mux.HandleFunc("GET /v1/files/{id}/content", r.handleFileContent) + mux.HandleFunc("DELETE /v1/files/{id}", r.handleDeleteFile) + mux.HandleFunc("POST /v1/batches", r.handleCreateBatch) + mux.HandleFunc("GET /v1/batches", r.handleListBatches) + mux.HandleFunc("GET /v1/batches/{id}", r.handleGetBatch) + mux.HandleFunc("POST /v1/batches/{id}/cancel", r.handleCancelBatch) // Liveness: this process is up. It says nothing about whether it can // serve, which is why it is not the thing to put a load balancer on. mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) { diff --git a/internal/router/router.go b/internal/router/router.go index 77fa7ac..f673985 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -68,6 +68,8 @@ type Router struct { // canaries counts how the two sides served. canaryTargets atomic.Pointer[map[string]string] canaries canaryStats + // batchWake has the batch worker look for work at once. + batchWake chan struct{} // extraMetrics are collectors other components publish on /metrics. extraMetrics []prometheus.Collector @@ -240,13 +242,14 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { } r := &Router{ - cfg: cfg.Router, - landing: cfg.Landing, - log: log, - reg: reg, - profiles: profiles, - prefix: prefixScorer, - m: metrics, + batchWake: make(chan struct{}, 1), + cfg: cfg.Router, + landing: cfg.Landing, + log: log, + reg: reg, + profiles: profiles, + prefix: prefixScorer, + m: metrics, px: proxy.New(proxy.Config{ DialTimeout: cfg.Router.Timeouts.Dial, ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader, @@ -444,6 +447,7 @@ func (r *Router) Start(ctx context.Context) { go r.pruneUsage(ctx, r.usageRetention) } if r.store != nil { + go r.runBatches(ctx) go r.refreshTenants(ctx) go r.purgeSessions(ctx) } diff --git a/internal/store/batches.go b/internal/store/batches.go new file mode 100644 index 0000000..cbe37d8 --- /dev/null +++ b/internal/store/batches.go @@ -0,0 +1,345 @@ +package store + +import ( + "context" + "database/sql" + "errors" + "time" +) + +// File is an uploaded or written file of the Files API, without its bytes. +type File struct { + ID string + Owner string + Purpose string + Filename string + Bytes int64 + CreatedAt time.Time + // ExpiresAt is when it is deleted by itself; zero is never. + ExpiresAt time.Time +} + +// PutFile keeps a file and its bytes. +func (s *DB) PutFile(ctx context.Context, f File, content []byte) error { + _, err := s.db.ExecContext(ctx, `INSERT INTO files (id, owner, purpose, filename, bytes, content, created_at, expires_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, f.ID, f.Owner, f.Purpose, f.Filename, int64(len(content)), content, + millis(f.CreatedAt), optMillis(f.ExpiresAt)) + return err +} + +const fileCols = `id, owner, purpose, filename, bytes, created_at, expires_at` + +func scanFile(sc interface{ Scan(...any) error }) (File, error) { + var f File + var created, expires int64 + if err := sc.Scan(&f.ID, &f.Owner, &f.Purpose, &f.Filename, &f.Bytes, &created, &expires); err != nil { + return File{}, err + } + f.CreatedAt = time.UnixMilli(created) + if expires > 0 { + f.ExpiresAt = time.UnixMilli(expires) + } + return f, nil +} + +// FileByID finds a file, ErrNotFound when there is none. +func (s *DB) FileByID(ctx context.Context, id string) (File, error) { + f, err := scanFile(s.db.QueryRowContext(ctx, `SELECT `+fileCols+` FROM files WHERE id = ?`, id)) + if errors.Is(err, sql.ErrNoRows) { + return File{}, ErrNotFound + } + return f, err +} + +// FileContent is a file's bytes. +func (s *DB) FileContent(ctx context.Context, id string) ([]byte, error) { + var b []byte + err := s.db.QueryRowContext(ctx, `SELECT content FROM files WHERE id = ?`, id).Scan(&b) + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrNotFound + } + return b, err +} + +// Files lists an owner's files, the newest first; purpose narrows it. +func (s *DB) Files(ctx context.Context, owner, purpose string, limit int) ([]File, error) { + q := `SELECT ` + fileCols + ` FROM files WHERE owner = ?` + args := []any{owner} + if purpose != "" { + q += ` AND purpose = ?` + args = append(args, purpose) + } + if limit <= 0 || limit > 10000 { + limit = 10000 + } + q += ` ORDER BY created_at DESC LIMIT ?` + args = append(args, limit) + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []File + for rows.Next() { + f, err := scanFile(rows) + if err != nil { + return nil, err + } + out = append(out, f) + } + return out, rows.Err() +} + +// DeleteFile removes a file; false when there was none. +func (s *DB) DeleteFile(ctx context.Context, id string) (bool, error) { + res, err := s.db.ExecContext(ctx, `DELETE FROM files WHERE id = ?`, id) + if err != nil { + return false, err + } + n, _ := res.RowsAffected() + return n > 0, nil +} + +// DeleteExpiredFiles removes the files past their expiry. +func (s *DB) DeleteExpiredFiles(ctx context.Context, now time.Time) (int64, error) { + res, err := s.db.ExecContext(ctx, `DELETE FROM files WHERE expires_at > 0 AND expires_at < ?`, millis(now)) + if err != nil { + return 0, err + } + return res.RowsAffected() +} + +// Batch is one batch of the Batch API. +type Batch struct { + ID string + Owner string + KeyID string + Endpoint string + InputFileID string + CompletionWindow string + Status string + OutputFileID string + ErrorFileID string + Total int + Completed int + Failed int + CreatedAt time.Time + InProgressAt time.Time + FinalizingAt time.Time + CompletedAt time.Time + FailedAt time.Time + ExpiredAt time.Time + CancellingAt time.Time + CancelledAt time.Time + ExpiresAt time.Time + // Metadata is the client's, as JSON; Errors why validation failed, as + // JSON. + Metadata string + Errors string +} + +// Batch states, as OpenAI names them. +const ( + BatchValidating = "validating" + BatchFailed = "failed" + BatchInProgress = "in_progress" + BatchFinalizing = "finalizing" + BatchCompleted = "completed" + BatchExpired = "expired" + BatchCancelling = "cancelling" + BatchCancelled = "cancelled" +) + +// BatchDone is whether a state is an end. +func BatchDone(status string) bool { + switch status { + case BatchFailed, BatchCompleted, BatchExpired, BatchCancelled: + return true + } + return false +} + +const batchCols = `id, owner, key_id, endpoint, input_file_id, completion_window, status, output_file_id, + error_file_id, total, completed, failed, created_at, in_progress_at, finalizing_at, completed_at, + failed_at, expired_at, cancelling_at, cancelled_at, expires_at, metadata, errors` + +func scanBatch(sc interface{ Scan(...any) error }) (Batch, error) { + var b Batch + var t [10]int64 + err := sc.Scan(&b.ID, &b.Owner, &b.KeyID, &b.Endpoint, &b.InputFileID, &b.CompletionWindow, &b.Status, + &b.OutputFileID, &b.ErrorFileID, &b.Total, &b.Completed, &b.Failed, + &t[0], &t[1], &t[2], &t[3], &t[4], &t[5], &t[6], &t[7], &t[8], &b.Metadata, &b.Errors) + if err != nil { + return Batch{}, err + } + at := func(v int64) time.Time { + if v == 0 { + return time.Time{} + } + return time.UnixMilli(v) + } + b.CreatedAt, b.InProgressAt, b.FinalizingAt, b.CompletedAt = at(t[0]), at(t[1]), at(t[2]), at(t[3]) + b.FailedAt, b.ExpiredAt, b.CancellingAt, b.CancelledAt, b.ExpiresAt = at(t[4]), at(t[5]), at(t[6]), at(t[7]), at(t[8]) + return b, nil +} + +// PutBatch records a new batch. +func (s *DB) PutBatch(ctx context.Context, b Batch) error { + _, err := s.db.ExecContext(ctx, `INSERT INTO batches (id, owner, key_id, endpoint, input_file_id, + completion_window, status, created_at, expires_at, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + b.ID, b.Owner, b.KeyID, b.Endpoint, b.InputFileID, b.CompletionWindow, b.Status, + millis(b.CreatedAt), millis(b.ExpiresAt), b.Metadata) + return err +} + +// UpdateBatch writes a batch's progress: its state, counts, files, times +// and errors. +func (s *DB) UpdateBatch(ctx context.Context, b Batch) error { + _, err := s.db.ExecContext(ctx, `UPDATE batches SET status = ?, output_file_id = ?, error_file_id = ?, + total = ?, completed = ?, failed = ?, in_progress_at = ?, finalizing_at = ?, completed_at = ?, + failed_at = ?, expired_at = ?, cancelling_at = ?, cancelled_at = ?, errors = ? WHERE id = ?`, + b.Status, b.OutputFileID, b.ErrorFileID, b.Total, b.Completed, b.Failed, + optMillis(b.InProgressAt), optMillis(b.FinalizingAt), optMillis(b.CompletedAt), optMillis(b.FailedAt), + optMillis(b.ExpiredAt), optMillis(b.CancellingAt), optMillis(b.CancelledAt), b.Errors, b.ID) + return err +} + +// CancelBatch asks for a batch to stop: cancelling, unless it has ended +// already. It says the batch as it stands. +func (s *DB) CancelBatch(ctx context.Context, id string, now time.Time) (Batch, error) { + if _, err := s.db.ExecContext(ctx, `UPDATE batches SET status = ?, cancelling_at = ? + WHERE id = ? AND status IN (?, ?, ?)`, BatchCancelling, millis(now), id, + BatchValidating, BatchInProgress, BatchFinalizing); err != nil { + return Batch{}, err + } + return s.BatchByID(ctx, id) +} + +// BatchByID finds a batch, ErrNotFound when there is none. +func (s *DB) BatchByID(ctx context.Context, id string) (Batch, error) { + b, err := scanBatch(s.db.QueryRowContext(ctx, `SELECT `+batchCols+` FROM batches WHERE id = ?`, id)) + if errors.Is(err, sql.ErrNoRows) { + return Batch{}, ErrNotFound + } + return b, err +} + +// Batches lists an owner's batches, the newest first, those made before +// the batch after when it is given. +func (s *DB) Batches(ctx context.Context, owner, after string, limit int) ([]Batch, error) { + q := `SELECT ` + batchCols + ` FROM batches WHERE owner = ?` + args := []any{owner} + if after != "" { + q += ` AND created_at < (SELECT created_at FROM batches WHERE id = ?)` + args = append(args, after) + } + if limit <= 0 || limit > 100 { + limit = 20 + } + q += ` ORDER BY created_at DESC LIMIT ?` + args = append(args, limit) + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []Batch + for rows.Next() { + b, err := scanBatch(rows) + if err != nil { + return nil, err + } + out = append(out, b) + } + return out, rows.Err() +} + +// ClaimBatch takes the lease on one batch that has work left and that no +// gateway holds, for owner until the time given. ok is false when there is +// none. +func (s *DB) ClaimBatch(ctx context.Context, owner string, now, until time.Time) (Batch, bool, error) { + var id string + err := s.db.QueryRowContext(ctx, `UPDATE batches SET lease_owner = ?, lease_until = ? + WHERE id = (SELECT id FROM batches WHERE status IN (?, ?, ?, ?) AND lease_until < ? + ORDER BY created_at LIMIT 1) + AND lease_until < ? + RETURNING id`, owner, millis(until), BatchValidating, BatchInProgress, BatchFinalizing, BatchCancelling, + millis(now), millis(now)).Scan(&id) + if errors.Is(err, sql.ErrNoRows) { + return Batch{}, false, nil + } + if err != nil { + return Batch{}, false, err + } + b, err := s.BatchByID(ctx, id) + return b, err == nil, err +} + +// RenewBatchLease keeps a lease held; false when it is held no longer. +func (s *DB) RenewBatchLease(ctx context.Context, id, owner string, until time.Time) (bool, error) { + res, err := s.db.ExecContext(ctx, `UPDATE batches SET lease_until = ? WHERE id = ? AND lease_owner = ?`, + millis(until), id, owner) + if err != nil { + return false, err + } + n, _ := res.RowsAffected() + return n > 0, nil +} + +// ReleaseBatch gives a lease up. +func (s *DB) ReleaseBatch(ctx context.Context, id, owner string) error { + _, err := s.db.ExecContext(ctx, `UPDATE batches SET lease_until = 0, lease_owner = '' WHERE id = ? AND lease_owner = ?`, id, owner) + return err +} + +// BatchResult is one line of a batch, served. +type BatchResult struct { + Line int + OK bool + Data string +} + +// PutBatchResult keeps one served line; a line kept already is left as it +// was. +func (s *DB) PutBatchResult(ctx context.Context, batchID string, r BatchResult) error { + ok := 0 + if r.OK { + ok = 1 + } + _, err := s.db.ExecContext(ctx, `INSERT INTO batch_results (batch_id, line, ok, data) VALUES (?, ?, ?, ?) + ON CONFLICT (batch_id, line) DO NOTHING`, batchID, r.Line, ok, r.Data) + return err +} + +// BatchResults are a batch's served lines, in order. +func (s *DB) BatchResults(ctx context.Context, batchID string) ([]BatchResult, error) { + rows, err := s.db.QueryContext(ctx, `SELECT line, ok, data FROM batch_results WHERE batch_id = ? ORDER BY line`, batchID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []BatchResult + for rows.Next() { + var r BatchResult + var ok int + if err := rows.Scan(&r.Line, &ok, &r.Data); err != nil { + return nil, err + } + r.OK = ok == 1 + out = append(out, r) + } + return out, rows.Err() +} + +// DeleteBatchResults forgets a batch's lines, once its files are written. +func (s *DB) DeleteBatchResults(ctx context.Context, batchID string) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM batch_results WHERE batch_id = ?`, batchID) + return err +} + +func optMillis(t time.Time) int64 { + if t.IsZero() { + return 0 + } + return t.UnixMilli() +} diff --git a/internal/store/batches_test.go b/internal/store/batches_test.go new file mode 100644 index 0000000..45dd72d --- /dev/null +++ b/internal/store/batches_test.go @@ -0,0 +1,80 @@ +package store + +import ( + "errors" + "testing" + "time" +) + +func TestFilesAndBatches(t *testing.T) { + checkFilesAndBatches(t, capacityStore(t)) +} + +func TestPostgresFilesAndBatches(t *testing.T) { + checkFilesAndBatches(t, openPG(t)) +} + +func checkFilesAndBatches(t *testing.T, db *DB) { + ctx := t.Context() + now := time.Now() + if err := db.PutFile(ctx, File{ID: "file-1", Owner: "p1", Purpose: "batch", Filename: "in.jsonl", CreatedAt: now}, + []byte("{\"a\":1}\n")); err != nil { + t.Fatal(err) + } + if err := db.PutFile(ctx, File{ID: "file-old", Owner: "p1", Purpose: "batch_output", Filename: "out.jsonl", + CreatedAt: now, ExpiresAt: now.Add(-time.Minute)}, []byte("x")); err != nil { + t.Fatal(err) + } + f, err := db.FileByID(ctx, "file-1") + if err != nil || f.Bytes != 8 || f.Owner != "p1" { + t.Fatalf("%+v %v", f, err) + } + if b, _ := db.FileContent(ctx, "file-1"); string(b) != "{\"a\":1}\n" { + t.Errorf("content %q", b) + } + if fs, _ := db.Files(ctx, "p1", "batch", 0); len(fs) != 1 { + t.Errorf("listed %v", fs) + } + if n, _ := db.DeleteExpiredFiles(ctx, now); n != 1 { + t.Errorf("expired %d", n) + } + + b := Batch{ID: "batch-1", Owner: "p1", KeyID: "k", Endpoint: "/v1/chat/completions", InputFileID: "file-1", + CompletionWindow: "24h", Status: BatchValidating, CreatedAt: now, ExpiresAt: now.Add(24 * time.Hour)} + if err := db.PutBatch(ctx, b); err != nil { + t.Fatal(err) + } + got, ok, err := db.ClaimBatch(ctx, "gw-a", now, now.Add(time.Minute)) + if err != nil || !ok || got.ID != "batch-1" { + t.Fatalf("claim: %+v %v %v", got, ok, err) + } + if _, ok, _ := db.ClaimBatch(ctx, "gw-b", now, now.Add(time.Minute)); ok { + t.Error("a leased batch was claimed twice") + } + if held, _ := db.RenewBatchLease(ctx, "batch-1", "gw-b", now.Add(time.Hour)); held { + t.Error("another gateway renewed the lease") + } + got.Status, got.Total, got.InProgressAt = BatchInProgress, 2, now + if err := db.UpdateBatch(ctx, got); err != nil { + t.Fatal(err) + } + for _, r := range []BatchResult{{Line: 1, OK: false, Data: "e"}, {Line: 0, OK: true, Data: "a"}, {Line: 0, OK: false, Data: "dup"}} { + if err := db.PutBatchResult(ctx, "batch-1", r); err != nil { + t.Fatal(err) + } + } + rs, _ := db.BatchResults(ctx, "batch-1") + if len(rs) != 2 || rs[0].Data != "a" || !rs[0].OK || rs[1].OK { + t.Errorf("results %+v", rs) + } + c, err := db.CancelBatch(ctx, "batch-1", now) + if err != nil || c.Status != BatchCancelling || c.Total != 2 || c.CancellingAt.IsZero() { + t.Errorf("cancel %+v %v", c, err) + } + if bs, _ := db.Batches(ctx, "p1", "", 0); len(bs) != 1 { + t.Errorf("listed %d", len(bs)) + } + if _, err := db.BatchByID(ctx, "nope"); !errors.Is(err, ErrNotFound) { + t.Errorf("missing: %v", err) + } +} diff --git a/internal/store/copy.go b/internal/store/copy.go index 8dbea3d..a7a630c 100644 --- a/internal/store/copy.go +++ b/internal/store/copy.go @@ -58,6 +58,9 @@ var copyOrder = []string{ // Who changed what is not lost by moving the store. "audit_events", "key_ttl", + "files", + "batches", + "batch_results", } // notCopied are tables deliberately left behind. diff --git a/internal/store/migrations/postgres/0024_batches.sql b/internal/store/migrations/postgres/0024_batches.sql new file mode 100644 index 0000000..60e7d24 --- /dev/null +++ b/internal/store/migrations/postgres/0024_batches.sql @@ -0,0 +1,56 @@ +-- Ported from the SQLite migration of the same number. +-- The Files and Batch APIs. A file is what a client uploaded (a batch's +-- requests) or what a batch wrote (its answers, its errors), owned by the +-- project of the key that made it. A batch is a file of requests served one +-- by one at the lowest priority; each finished line is kept as it is +-- served, so a batch survives the control plane restarting under it, and +-- the lease says which gateway is working on it. +CREATE TABLE files ( + id TEXT PRIMARY KEY, + owner TEXT NOT NULL, + purpose TEXT NOT NULL, + filename TEXT NOT NULL, + bytes BIGINT NOT NULL, + content BYTEA NOT NULL, + created_at BIGINT NOT NULL, + expires_at BIGINT NOT NULL DEFAULT 0 +); +CREATE INDEX files_owner ON files (owner, created_at); + +CREATE TABLE batches ( + id TEXT PRIMARY KEY, + owner TEXT NOT NULL, + key_id TEXT NOT NULL, + endpoint TEXT NOT NULL, + input_file_id TEXT NOT NULL, + completion_window TEXT NOT NULL, + status TEXT NOT NULL, + output_file_id TEXT NOT NULL DEFAULT '', + error_file_id TEXT NOT NULL DEFAULT '', + total BIGINT NOT NULL DEFAULT 0, + completed BIGINT NOT NULL DEFAULT 0, + failed BIGINT NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + in_progress_at BIGINT NOT NULL DEFAULT 0, + finalizing_at BIGINT NOT NULL DEFAULT 0, + completed_at BIGINT NOT NULL DEFAULT 0, + failed_at BIGINT NOT NULL DEFAULT 0, + expired_at BIGINT NOT NULL DEFAULT 0, + cancelling_at BIGINT NOT NULL DEFAULT 0, + cancelled_at BIGINT NOT NULL DEFAULT 0, + expires_at BIGINT NOT NULL, + metadata TEXT NOT NULL DEFAULT '', + errors TEXT NOT NULL DEFAULT '', + lease_owner TEXT NOT NULL DEFAULT '', + lease_until BIGINT NOT NULL DEFAULT 0 +); +CREATE INDEX batches_owner ON batches (owner, created_at); +CREATE INDEX batches_status ON batches (status, lease_until); + +CREATE TABLE batch_results ( + batch_id TEXT NOT NULL, + line BIGINT NOT NULL, + ok BIGINT NOT NULL, + data TEXT NOT NULL, + PRIMARY KEY (batch_id, line) +); diff --git a/internal/store/migrations/sqlite/0024_batches.sql b/internal/store/migrations/sqlite/0024_batches.sql new file mode 100644 index 0000000..1293121 --- /dev/null +++ b/internal/store/migrations/sqlite/0024_batches.sql @@ -0,0 +1,55 @@ +-- The Files and Batch APIs. A file is what a client uploaded (a batch's +-- requests) or what a batch wrote (its answers, its errors), owned by the +-- project of the key that made it. A batch is a file of requests served one +-- by one at the lowest priority; each finished line is kept as it is +-- served, so a batch survives the control plane restarting under it, and +-- the lease says which gateway is working on it. +CREATE TABLE files ( + id TEXT PRIMARY KEY, + owner TEXT NOT NULL, + purpose TEXT NOT NULL, + filename TEXT NOT NULL, + bytes INTEGER NOT NULL, + content BLOB NOT NULL, + created_at INTEGER NOT NULL, + expires_at INTEGER NOT NULL DEFAULT 0 +); +CREATE INDEX files_owner ON files (owner, created_at); + +CREATE TABLE batches ( + id TEXT PRIMARY KEY, + owner TEXT NOT NULL, + key_id TEXT NOT NULL, + endpoint TEXT NOT NULL, + input_file_id TEXT NOT NULL, + completion_window TEXT NOT NULL, + status TEXT NOT NULL, + output_file_id TEXT NOT NULL DEFAULT '', + error_file_id TEXT NOT NULL DEFAULT '', + total INTEGER NOT NULL DEFAULT 0, + completed INTEGER NOT NULL DEFAULT 0, + failed INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + in_progress_at INTEGER NOT NULL DEFAULT 0, + finalizing_at INTEGER NOT NULL DEFAULT 0, + completed_at INTEGER NOT NULL DEFAULT 0, + failed_at INTEGER NOT NULL DEFAULT 0, + expired_at INTEGER NOT NULL DEFAULT 0, + cancelling_at INTEGER NOT NULL DEFAULT 0, + cancelled_at INTEGER NOT NULL DEFAULT 0, + expires_at INTEGER NOT NULL, + metadata TEXT NOT NULL DEFAULT '', + errors TEXT NOT NULL DEFAULT '', + lease_owner TEXT NOT NULL DEFAULT '', + lease_until INTEGER NOT NULL DEFAULT 0 +); +CREATE INDEX batches_owner ON batches (owner, created_at); +CREATE INDEX batches_status ON batches (status, lease_until); + +CREATE TABLE batch_results ( + batch_id TEXT NOT NULL, + line INTEGER NOT NULL, + ok INTEGER NOT NULL, + data TEXT NOT NULL, + PRIMARY KEY (batch_id, line) +); diff --git a/internal/store/postgres_test.go b/internal/store/postgres_test.go index 47bf661..13e733c 100644 --- a/internal/store/postgres_test.go +++ b/internal/store/postgres_test.go @@ -43,7 +43,7 @@ func openPG(t *testing.T) *DB { // every subtest gets a fresh file, and failed only here. `TRUNCATE thread_messages, threads, memory_notes, usage, usage_5m, usage_1h, usage_1d, api_keys, memberships, key_claims, limits, models, sessions, users, projects, teams, - organizations, leadership, control_planes RESTART IDENTITY CASCADE`) + organizations, leadership, control_planes, files, batches, batch_results RESTART IDENTITY CASCADE`) if err != nil { t.Fatalf("clearing the test database: %v", err) }