From b51cc3cf921ca4cbf0dcbfc8568aad360141bf23 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Sat, 10 Oct 2026 16:04:11 +0200 Subject: [PATCH] A 3D model is a job: queued with an id, asked after, its result a file POST /v1/3d/generations answers 202 with a job; GET says its status, progress and result file; cancel stops it in the engine's queue too. The request is kept as sent and replayed by whichever gateway holds the lease through the path a client's request takes, so it is charged and judged like any request (tagged job=), and a model still loading is waited out for as long as the job has. Migration 0044. Nothing makes 3D models yet; the door refuses every model until the ComfyUI adapter learns meshes. Co-Authored-By: Claude Fable 5.1 --- BACKLOG.md | 9 +- CHANGELOG.md | 16 + docs/install.md | 44 +++ internal/router/batchworker.go | 14 +- internal/router/handlers.go | 6 + internal/router/images.go | 9 + internal/router/jobs.go | 204 +++++++++++++ internal/router/jobs_test.go | 285 ++++++++++++++++++ internal/router/jobworker.go | 229 ++++++++++++++ internal/router/judge.go | 12 +- internal/router/router.go | 6 +- internal/store/copy.go | 2 + internal/store/jobs.go | 208 +++++++++++++ internal/store/jobs_test.go | 97 ++++++ .../store/migrations/postgres/0044_jobs.sql | 34 +++ .../store/migrations/sqlite/0044_jobs.sql | 33 ++ 16 files changed, 1198 insertions(+), 10 deletions(-) create mode 100644 internal/router/jobs.go create mode 100644 internal/router/jobs_test.go create mode 100644 internal/router/jobworker.go create mode 100644 internal/store/jobs.go create mode 100644 internal/store/jobs_test.go create mode 100644 internal/store/migrations/postgres/0044_jobs.sql create mode 100644 internal/store/migrations/sqlite/0044_jobs.sql diff --git a/BACKLOG.md b/BACKLOG.md index e3d17ce..d419b3f 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -1131,7 +1131,14 @@ to be read before it is considered). Steam's disclosure form (changed declared on the store page: the licence notice G3 writes into the GLB is what the studio answers that form from. -### G1. Jobs: store, worker, endpoints (M) +### G1. Jobs: store, worker, endpoints (M) — done 2026-10-10 (branch `jobs-g1`) +Built as written, with two choices: a job found running after its gateway +died is run again from the start (the prompt id column is there for G2 to +follow the same work instead), and the request is validated at the door +only as far as costs nothing — key, model, allowlist, the request's shape +— with limits, guardrails and admission applied when the worker replays it, +said as the job's error. The result file's `content_type` is the job's; the +files endpoint still serves every file as an octet stream. - Store: `jobs` (id, owner key/principal, kind, model, status, input (JSON), progress (JSON: queue position, percent, step), result file id, error, created/started/finished, lease owner/until), migrations for diff --git a/CHANGELOG.md b/CHANGELOG.md index 99d6d8e..10bd23b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,22 @@ APIs change between minors without a deprecation period. ## Unreleased +### Added + +- **Jobs: 3D models asked after, not waited for.** `POST + /v1/3d/generations` answers 202 with a job id; `GET + /v1/3d/generations/{id}` says its status, progress and, when done, the + file its result is in (`/v1/files/{id}/content`); `POST .../cancel` stops + it, in the engine's queue too. The request is replayed by whichever + gateway holds the job's lease through the same path a client's request + takes — key, allowlist, limits, admission, guardrails, accounting — so + it is charged like any request, tagged `job=`, and a model still + loading is waited out for as long as the job has (24 h) rather than + refused. A job survives the gateway that took it; results are kept for + `router.batch.retain_days`. Migration 0044. No model makes 3D models + yet: every request is refused at the door until the ComfyUI adapter + learns meshes (BACKLOG G3). + ### Fixed - **An engine's output no longer breaks when its agent restarts.** Engines diff --git a/docs/install.md b/docs/install.md index 234275b..4228fcd 100644 --- a/docs/install.md +++ b/docs/install.md @@ -1672,6 +1672,7 @@ and accounting exactly as an API call does. | `POST /v1/audio/transcriptions`, `/v1/audio/translations` | the multipart form passed through as it came, to an engine that serves speech to text: vLLM with Whisper | | `POST /v1/moderations` | answered by the guardrails' judge model (`role: guard`): `flagged`, and the judge's own categories | | `/v1/files`, `/v1/batches` | batches; see [Batches](#batches) | +| `/v1/3d/generations` | 3D models, as jobs asked after; see [Jobs](#jobs-3d-models) | | `GET /v1/models` | the models this key may use | ## Watching it @@ -2030,6 +2031,49 @@ Each line of the file is `{"custom_id": …, "method": "POST", "url": or past their retention — a project under a guardrails policy with `no_store` is refused it. +## Jobs: 3D models + +A 3D model from a picture takes minutes on a card that first has to be +emptied of whatever else was on it. No client should hold a connection for +that, so it is a job: queued, answered with an id, asked after, its result +downloaded. + +``` +curl http://gateway:8080/v1/3d/generations -H "Authorization: Bearer $LLMGW_API_KEY" \ + -d '{"model": "trellis2", "image": ""}' +# → 202 {"id": "job_…", "status": "queued", …} +curl http://gateway:8080/v1/3d/generations/job_… -H "Authorization: Bearer $LLMGW_API_KEY" +# → {"status": "completed", "output_file_id": "file-…", "content_type": "model/gltf-binary", "usage": {…}} +curl http://gateway:8080/v1/files/file-…/content -H "Authorization: Bearer $LLMGW_API_KEY" -o model.glb +curl -X POST http://gateway:8080/v1/3d/generations/job_…/cancel -H "Authorization: Bearer $LLMGW_API_KEY" +``` + +- **Status:** `queued`, `running`, `completed`, `failed` (with `error`), + `cancelling`, `cancelled`, `expired` (not finished within 24 h). + `progress` carries what the engine said of the work so far, once it says + anything; `GET /v1/3d/generations` lists the caller's jobs. +- **Refused at the door:** no key, a model that is not there or not in the + key's allowlist, a request the model cannot serve. Refused when served, + and said as the job's error: a spent budget, the guardrails, a model + that will not come up. +- **How it is served:** the request is kept as it was sent and replayed by + whichever gateway holds the job's lease, through the same path a + client's request takes — limits, admission in the key's band, guardrails, + accounting. It is one usage record under the job's id, tagged + `job=`, charged like any request; `GET` shows the usage once it is + done. A model still loading, a rate limit or a full queue is waited out + for as long as the job has: nobody is holding a connection, which is the + point. +- **Surviving a restart:** a gateway that stops hands its jobs to another + within two minutes; a job it was running is run again from the start. +- **What it keeps:** the request as sent, and the result as a file owned + by the key's project, kept `router.batch.retain_days` (30) like a batch's + files; the job's record goes with it. A project under a `no_store` + policy is refused jobs, as it is batches. + +No model makes 3D models yet: until the ComfyUI adapter learns meshes, +every request is refused at the door as "does not make 3D models". + ## Conversations across models Off everywhere until a project asks. llmgw otherwise stores no message content diff --git a/internal/router/batchworker.go b/internal/router/batchworker.go index 1ba0771..381c74c 100644 --- a/internal/router/batchworker.go +++ b/internal/router/batchworker.go @@ -148,7 +148,7 @@ func (r *Router) serveLines(ctx context.Context, b *store.Batch, lines []batchLi for _, d := range done { served[d.Line] = true } - p, perr := r.batchPrincipalOf(ctx, b.KeyID) + p, perr := r.principalOfKey(ctx, b.KeyID, r.batchBand()) var mu sync.Mutex stopWith := "" @@ -221,9 +221,10 @@ func (r *Router) serveLines(ctx context.Context, b *store.Batch, lines []batchLi 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) { +// principalOfKey is the key a batch or a job was made with, as it stands now: +// a key revoked since serves no more of it. band is the admission band the +// work goes in; empty is the key's own. +func (r *Router) principalOfKey(ctx context.Context, keyID, band string) (auth.Principal, error) { k, err := r.store.KeyByID(ctx, keyID) if err != nil { return auth.Principal{}, err @@ -238,10 +239,13 @@ func (r *Router) batchPrincipalOf(ctx context.Context, keyID string) (auth.Princ if err != nil { return auth.Principal{}, err } + if band == "" { + band = eff.Band + } 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, + Band: band, RoutingProfile: eff.RoutingProfile, Limits: eff, }, nil } diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 0da3c84..80a882c 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -123,6 +123,12 @@ func (r *Router) Handler() http.Handler { mux.HandleFunc("GET /v1/batches", r.handleListBatches) mux.HandleFunc("GET /v1/batches/{id}", r.handleGetBatch) mux.HandleFunc("POST /v1/batches/{id}/cancel", r.handleCancelBatch) + // Jobs: 3D models, which take minutes and are asked after rather than + // waited for. The result is a file. + mux.HandleFunc("POST "+meshesPath, r.handleCreateJob) + mux.HandleFunc("GET "+meshesPath, r.handleListJobs) + mux.HandleFunc("GET "+meshesPath+"/{id}", r.handleGetJob) + mux.HandleFunc("POST "+meshesPath+"/{id}/cancel", r.handleCancelJob) // Knowledge bases: documents a person's llmgw-kb read and embedded. r.registerKB(mux) // Liveness: this process is up. It says nothing about whether it can diff --git a/internal/router/images.go b/internal/router/images.go index 71634c9..a649870 100644 --- a/internal/router/images.go +++ b/internal/router/images.go @@ -111,6 +111,15 @@ func parseImageRequest(body []byte) (imageRequest, error) { // endpoint may serve as it likes. func (r *Router) imageRouteRefusal(path, model string, body []byte) string { m, known := r.modelMap()[model] + if path == meshesPath { + // Until a model makes them (BACKLOG G3), every model the catalog + // knows is one that does not; one it does not know is a + // hand-configured endpoint's to answer. + if known { + return fmt.Sprintf("%s does not make 3D models", m.Name) + } + return "" + } if path == imagesPath { if known && !m.MakesImages() { return fmt.Sprintf("%s does not make pictures", m.Name) diff --git a/internal/router/jobs.go b/internal/router/jobs.go new file mode 100644 index 0000000..a568b61 --- /dev/null +++ b/internal/router/jobs.go @@ -0,0 +1,204 @@ +package router + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strconv" + "strings" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/auth" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// Jobs: a request that takes minutes, answered with an id and asked after. +// +// A 3D model from a picture takes minutes on a card that first has to be +// emptied of whatever else was on it. That is not a request to hold a socket +// open for (activate.go says why), so it is a job: queued with 202, its +// status and progress read back, its result downloaded as a file. The +// request is kept as the client sent it and replayed, by whichever gateway +// holds the job's lease, through the same path a client's request takes — +// key, allowlist, limits, admission, guardrails, accounting — so a job is +// charged and judged like any request, tagged job=. + +// meshesPath is where 3D models are asked for. The synchronous path behind +// it is served to the job worker only; clients get a job. +const meshesPath = "/v1/3d/generations" + +const ( + purposeJobOutput = "job_output" + // jobWindow is how long a job may take, queue included, before it is + // given up on. + jobWindow = 24 * time.Hour +) + +func (r *Router) jobJSON(ctx context.Context, j store.Job) map[string]any { + out := map[string]any{ + "id": j.ID, "object": "job", "endpoint": j.Endpoint, "model": j.Model, "status": j.Status, + "progress": nil, "output_file_id": strOrNil(j.OutputFileID), "content_type": strOrNil(j.ContentType), + "error": nil, "created_at": j.CreatedAt.Unix(), "started_at": unixOrNil(j.StartedAt), + "finished_at": unixOrNil(j.FinishedAt), "expires_at": unixOrNil(j.ExpiresAt), "usage": nil, + } + if j.Progress != "" && json.Valid([]byte(j.Progress)) { + out["progress"] = json.RawMessage(j.Progress) + } + if j.Error != "" { + out["error"] = map[string]string{"message": j.Error} + } + // What the work cost, from the usage record the replayed request left + // under the job's id. + if store.JobDone(j.Status) && r.store != nil { + if u, ok, err := r.store.UsageByRequest(ctx, j.ID); err == nil && ok { + out["usage"] = map[string]int64{"prompt_tokens": u.PromptTokens, "completion_tokens": u.CompletionTokens, + "images": u.Images} + } + } + return out +} + +// handleCreateJob queues a request for a 3D model. Refused now: no key, a +// model that is not there or not allowed, a request the model cannot serve. +// Refused when it is served, and said as the job's error: limits, the +// guardrails, a model that will not come up. +func (r *Router) handleCreateJob(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes)) + if err != nil { + writeError(w, http.StatusBadRequest, "invalid_request_error", "could not read the request: "+err.Error()) + return + } + var in struct { + Model string `json:"model"` + } + if err := json.Unmarshal(body, &in); err != nil { + writeError(w, http.StatusBadRequest, "invalid_request_error", "the request is not valid JSON: "+err.Error()) + return + } + if in.Model == "" { + writeError(w, http.StatusBadRequest, "invalid_request_error", "model is required") + return + } + if _, known := r.modelMap()[in.Model]; !known || r.hiddenModel(in.Model) { + writeErrorCode(w, http.StatusNotFound, "invalid_request_error", "model_not_found", + fmt.Sprintf("the model %q does not exist", in.Model)) + return + } + if !r.permits(p, in.Model) { + auth.WriteForbiddenModel(w, in.Model) + return + } + if why := r.imageRouteRefusal(meshesPath, in.Model, body); why != "" { + writeError(w, http.StatusBadRequest, "invalid_request_error", why) + return + } + now := time.Now() + j := store.Job{ID: newID("job_"), Owner: ownerOf(p), KeyID: p.KeyID, Endpoint: meshesPath, Model: in.Model, + Input: body, Status: store.JobQueued, CreatedAt: now, ExpiresAt: now.Add(jobWindow)} + if err := r.store.PutJob(req.Context(), j); err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + r.wakeJobs() + writeJSON(w, http.StatusAccepted, r.jobJSON(req.Context(), j)) +} + +// ownedJob is a job the caller owns, or a 404 said. +func (r *Router) ownedJob(w http.ResponseWriter, req *http.Request, owner string) (store.Job, bool) { + j, err := r.store.JobByID(req.Context(), req.PathValue("id")) + if errors.Is(err, store.ErrNotFound) || err == nil && j.Owner != owner { + writeErrorCode(w, http.StatusNotFound, "invalid_request_error", "job_not_found", + fmt.Sprintf("no job %q", req.PathValue("id"))) + return store.Job{}, false + } + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return store.Job{}, false + } + return j, true +} + +func (r *Router) handleGetJob(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + if j, ok := r.ownedJob(w, req, ownerOf(p)); ok { + writeJSON(w, http.StatusOK, r.jobJSON(req.Context(), j)) + } +} + +func (r *Router) handleCancelJob(w http.ResponseWriter, req *http.Request) { + p, ok := r.batchPrincipal(w, req) + if !ok { + return + } + j, ok := r.ownedJob(w, req, ownerOf(p)) + if !ok { + return + } + j, err := r.store.CancelJob(req.Context(), j.ID, time.Now()) + if err != nil { + writeError(w, http.StatusInternalServerError, "server_error", err.Error()) + return + } + r.wakeJobs() + writeJSON(w, http.StatusOK, r.jobJSON(req.Context(), j)) +} + +func (r *Router) handleListJobs(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 + } + js, err := r.store.Jobs(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(js) > limit + if more { + js = js[:limit] + } + data := make([]map[string]any, 0, len(js)) + for _, j := range js { + data = append(data, r.jobJSON(req.Context(), j)) + } + out := map[string]any{"object": "list", "data": data, "has_more": more, "first_id": nil, "last_id": nil} + if len(js) > 0 { + out["first_id"], out["last_id"] = js[0].ID, js[len(js)-1].ID + } + writeJSON(w, http.StatusOK, out) +} + +// errorMessage is what an answer said went wrong: the message of an OpenAI +// error body, or the body itself, kept short. +func errorMessage(body []byte) string { + var e struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + if json.Unmarshal(body, &e) == nil && e.Error.Message != "" { + return e.Error.Message + } + s := strings.TrimSpace(string(body)) + if len(s) > 300 { + s = s[:300] + "…" + } + if s == "" { + return "no reason given" + } + return s +} diff --git a/internal/router/jobs_test.go b/internal/router/jobs_test.go new file mode 100644 index 0000000..24572fc --- /dev/null +++ b/internal/router/jobs_test.go @@ -0,0 +1,285 @@ +package router + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/comfy" + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// jobComfy is a ComfyUI that makes one picture a workflow, slowly when +// asked: stall holds every workflow until it is cleared, and the paths the +// router took to stop one are counted. +type jobComfy struct { + srv *httptest.Server + stall atomic.Bool + prompts atomic.Int32 + cancelled atomic.Int32 +} + +func newJobComfy(t *testing.T) *jobComfy { + t.Helper() + c := &jobComfy{} + c.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/prompt": + c.prompts.Add(1) + fmt.Fprint(w, `{"prompt_id":"p1"}`) + case r.URL.Path == "/history/p1": + if c.stall.Load() { + fmt.Fprint(w, `{}`) + return + } + fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true},"outputs":{"7":{"images":[{"filename":"0.png","type":"temp"}]}}}}`) + case r.URL.Path == "/view": + w.Header().Set("Content-Type", "image/png") + fmt.Fprint(w, "png-bytes") + case r.URL.Path == "/queue" && r.Method == http.MethodPost, r.URL.Path == "/interrupt": + c.cancelled.Add(1) + fmt.Fprint(w, `{}`) + case r.URL.Path == "/queue": + fmt.Fprint(w, `{"queue_running":[[0,"p1"]],"queue_pending":[]}`) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(c.srv.Close) + return c +} + +// jobRouter is a gateway with the store, a key, a text model and a ComfyUI +// model priced by the picture. +func jobRouter(t *testing.T, c *jobComfy) (*Router, string) { + t.Helper() + cfg := baseConfig() + cfg.Models = []catalog.Model{catModel("chat", nil), {Name: "flux", Source: "external", + External: &catalog.External{URL: c.srv.URL, Model: "flux1-schnell-fp8.safetensors", API: catalog.APIComfyUI}}} + cfg.Auth.Enabled = true + cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db") + cfg.Auth.Pricing = map[string]config.Price{"flux": {PerImage: 5000}} + cfg.Auth.FlushInterval = 20 * time.Millisecond + cfg.Auth.BatchSize = 1 + rt := newTestRouter(t, cfg) + rt.comfy = &comfy.Client{Poll: time.Millisecond} + rt.SetEndpoints(t.Context(), []endpoint.Info{ + {ID: "external:flux", URL: c.srv.URL, Models: []string{"flux"}, Engine: endpoint.EngineComfyUI, + Pool: "external", Weight: 1, UpstreamModel: "flux1-schnell-fp8.safetensors"}, + }) + org, _ := rt.store.CreateOrg(t.Context(), "acme") + team, _ := rt.store.CreateTeam(t.Context(), org, "platform") + project, err := rt.store.CreateProject(t.Context(), team, "gateway") + if err != nil { + t.Fatal(err) + } + secret := store.KeyPrefix + strings.Repeat("a", 43) + sum := sha256.Sum256([]byte(secret)) + if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-1", Hash: sum[:], Prefix: "aaaa", Name: "test", + ProjectID: project, CreatedAt: time.Now()}); err != nil { + t.Fatal(err) + } + return rt, secret +} + +// queueJob puts a job in the store as the create endpoint would, for an +// endpoint that serves today: until a model makes 3D models, the worker is +// tried with pictures. +func queueJob(t *testing.T, rt *Router, id, status string, body string) store.Job { + t.Helper() + now := time.Now() + j := store.Job{ID: id, Owner: "project:" + mustProjectOf(t, rt, "key-1"), KeyID: "key-1", Endpoint: imagesPath, + Model: "flux", Input: []byte(body), Status: status, CreatedAt: now, ExpiresAt: now.Add(jobWindow)} + if err := rt.store.PutJob(t.Context(), j); err != nil { + t.Fatal(err) + } + return j +} + +func waitJob(t *testing.T, h http.Handler, secret, id string, until func(map[string]any) bool) map[string]any { + t.Helper() + deadline := time.Now().Add(10 * time.Second) + for { + _, j, raw := call(t, h, secret, http.MethodGet, meshesPath+"/"+id, nil) + if until(j) { + return j + } + if time.Now().After(deadline) { + t.Fatalf("the job did not get there: %s", raw) + } + time.Sleep(20 * time.Millisecond) + } +} + +func jobDone(j map[string]any) bool { return store.JobDone(fmt.Sprint(j["status"])) } + +func TestAJobIsServedItsResultKeptAndItsWorkCharged(t *testing.T) { + c := newJobComfy(t) + rt, secret := jobRouter(t, c) + h := rt.Handler() + j := queueJob(t, rt, "job_1", store.JobQueued, `{"model":"flux","prompt":"a cube"}`) + rt.wakeJobs() + + got := waitJob(t, h, secret, j.ID, jobDone) + if got["status"] != "completed" || got["content_type"] != "application/json" || got["error"] != nil { + t.Fatalf("%v", got) + } + if got["started_at"] == nil || got["finished_at"] == nil { + t.Errorf("times: %v", got) + } + code, _, out := call(t, h, secret, http.MethodGet, "/v1/files/"+got["output_file_id"].(string)+"/content", nil) + if code != 200 || !strings.Contains(out, `"b64_json":"cG5nLWJ5dGVz"`) { + t.Errorf("result %d: %s", code, out) + } + _, f, _ := call(t, h, secret, http.MethodGet, "/v1/files/"+got["output_file_id"].(string), nil) + if f["purpose"] != purposeJobOutput || f["filename"] != "job_1.json" || f["expires_at"] == nil { + t.Errorf("file %v", f) + } + + // Charged like any request, by the picture, and findable by the job. + deadline := time.Now().Add(5 * time.Second) + for { + u, ok, _ := rt.store.UsageByRequest(t.Context(), j.ID) + if ok { + if u.Tags["job"] != j.ID || u.Images != 1 || u.CostMicros != 5000 { + t.Errorf("usage %+v", u) + } + break + } + if time.Now().After(deadline) { + t.Fatal("the job left no usage") + } + time.Sleep(20 * time.Millisecond) + } + got = waitJob(t, h, secret, j.ID, func(j map[string]any) bool { return j["usage"] != nil }) + if u := got["usage"].(map[string]any); u["images"] != 1.0 { + t.Errorf("usage shown %v", u) + } + _, list, _ := call(t, h, secret, http.MethodGet, meshesPath, nil) + if d := list["data"].([]any); len(d) != 1 { + t.Errorf("listed %v", list) + } +} + +func TestAJobIsRefusedBeforeItIsQueuedWhenItCannotBeServed(t *testing.T) { + c := newJobComfy(t) + rt, secret := jobRouter(t, c) + h := rt.Handler() + code, _, raw := call(t, h, secret, http.MethodPost, meshesPath, map[string]any{"model": "chat", "prompt": "a cube"}) + if code != 400 || !strings.Contains(raw, "does not make 3D models") { + t.Errorf("text model: %d %s", code, raw) + } + code, _, raw = call(t, h, secret, http.MethodPost, meshesPath, map[string]any{"model": "nope", "prompt": "a cube"}) + if code != 404 || !strings.Contains(raw, "model_not_found") { + t.Errorf("unknown model: %d %s", code, raw) + } + code, _, raw = call(t, h, secret, http.MethodPost, meshesPath, map[string]any{"prompt": "a cube"}) + if code != 400 || !strings.Contains(raw, "model is required") { + t.Errorf("no model: %d %s", code, raw) + } + if code, _, _ := call(t, h, "", http.MethodPost, meshesPath, map[string]any{"model": "chat"}); code != 401 { + t.Errorf("no key: %d", code) + } + if _, list, _ := call(t, h, secret, http.MethodGet, meshesPath, nil); len(list["data"].([]any)) != 0 { + t.Errorf("a refused job was queued: %v", list) + } +} + +func TestARunningJobIsCancelledAndTheEngineToldToStop(t *testing.T) { + c := newJobComfy(t) + c.stall.Store(true) + rt, secret := jobRouter(t, c) + h := rt.Handler() + j := queueJob(t, rt, "job_slow", store.JobQueued, `{"model":"flux","prompt":"a cube"}`) + rt.wakeJobs() + waitJob(t, h, secret, j.ID, func(j map[string]any) bool { return j["status"] == "running" }) + deadline := time.Now().Add(5 * time.Second) + for c.prompts.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + + code, got, raw := call(t, h, secret, http.MethodPost, meshesPath+"/"+j.ID+"/cancel", nil) + if code != 200 || got["status"] != "cancelling" { + t.Fatalf("cancel: %d %s", code, raw) + } + got = waitJob(t, h, secret, j.ID, jobDone) + if got["status"] != "cancelled" || got["output_file_id"] != nil { + t.Errorf("%v", got) + } + if c.cancelled.Load() == 0 { + t.Error("ComfyUI was not told to stop") + } + + // One not yet taken ends at once, and is never sent. + before := c.prompts.Load() + j2 := queueJob(t, rt, "job_never", store.JobQueued, `{"model":"flux","prompt":"a cube"}`) + _, got, _ = call(t, h, secret, http.MethodPost, meshesPath+"/"+j2.ID+"/cancel", nil) + if got["status"] != "cancelled" { + t.Errorf("%v", got) + } + rt.wakeJobs() + time.Sleep(100 * time.Millisecond) + if c.prompts.Load() != before { + t.Error("a cancelled job was sent") + } +} + +func TestAJobWhoseGatewayDiedIsTakenByAnother(t *testing.T) { + c := newJobComfy(t) + rt, secret := jobRouter(t, c) + h := rt.Handler() + // As a gateway that died mid-job left it: running, its lease run out. + j := queueJob(t, rt, "job_orphan", store.JobRunning, `{"model":"flux","prompt":"a cube"}`) + j.StartedAt = time.Now().Add(-time.Minute) + _ = rt.store.UpdateJob(t.Context(), j) + rt.wakeJobs() + if got := waitJob(t, h, secret, j.ID, jobDone); got["status"] != "completed" { + t.Errorf("%v", got) + } + + // A refusal of the request is the job's end, said. + bad := queueJob(t, rt, "job_bad", store.JobQueued, `{"model":"flux","prompt":"a cube","n":99}`) + rt.wakeJobs() + got := waitJob(t, h, secret, bad.ID, jobDone) + if got["status"] != "failed" || !strings.Contains(fmt.Sprint(got["error"]), "n is between 1 and 4") { + t.Errorf("%v", got) + } +} + +func TestJobsAreTheirOwners(t *testing.T) { + c := newJobComfy(t) + rt, secret := jobRouter(t, c) + h := rt.Handler() + j := queueJob(t, rt, "job_mine", store.JobQueued, `{"model":"flux","prompt":"a cube"}`) + rt.wakeJobs() + waitJob(t, h, secret, j.ID, jobDone) + + 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, meshesPath+"/"+j.ID, nil); code != 404 { + t.Errorf("another project read a job: %d", code) + } + if code, _, _ := call(t, h, other, http.MethodPost, meshesPath+"/"+j.ID+"/cancel", nil); code != 404 { + t.Errorf("another project cancelled a job: %d", code) + } + _, mine, _ := call(t, h, secret, http.MethodGet, meshesPath+"/"+j.ID, nil) + if code, _, _ := call(t, h, other, http.MethodGet, "/v1/files/"+mine["output_file_id"].(string)+"/content", nil); code != 404 { + t.Errorf("another project read a job's result: %d", code) + } +} diff --git a/internal/router/jobworker.go b/internal/router/jobworker.go new file mode 100644 index 0000000..040df3a --- /dev/null +++ b/internal/router/jobworker.go @@ -0,0 +1,229 @@ +package router + +import ( + "bytes" + "context" + "fmt" + "net/http" + "os" + "strconv" + "strings" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/auth" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// The worker that serves jobs. Any gateway on the store may run one: a job +// is worked on by whichever holds its lease, and one whose gateway died is +// taken by another when the lease runs out. + +const ( + jobLease = 2 * time.Minute + jobPoll = 5 * time.Second + // jobCancelCheck is how often a running job looks for a cancel. + jobCancelCheck = 2 * time.Second + // jobMaxResult bounds what one job may answer: a mesh, not a film. + jobMaxResult = 256 << 20 +) + +// wakeJobs has the worker look for work now rather than at its next poll. +func (r *Router) wakeJobs() { + select { + case r.jobWake <- struct{}{}: + default: + } +} + +func (r *Router) runJobs(ctx context.Context) { + host, _ := os.Hostname() + owner := fmt.Sprintf("%s:%d:%s", host, os.Getpid(), newID("")) + t := time.NewTicker(jobPoll) + defer t.Stop() + for { + for { + now := time.Now() + j, ok, err := r.store.ClaimJob(ctx, owner, now, now.Add(jobLease)) + if err != nil { + if ctx.Err() == nil { + r.log.Warn("could not look for jobs to serve", "err", err) + } + break + } + if !ok { + break + } + r.workJob(ctx, owner, j) + } + // A job's record outlives its file by nothing: both go at the + // retention. + _, _, _, days := r.batchCfg() + if n, err := r.store.DeleteJobsBefore(ctx, time.Now().AddDate(0, 0, -days)); err == nil && n > 0 { + r.log.Info("jobs past their retention forgotten", "jobs", n) + } + select { + case <-ctx.Done(): + return + case <-t.C: + case <-r.jobWake: + } + } +} + +// workJob takes one job as far as it goes, holding its lease the while. +func (r *Router) workJob(ctx context.Context, owner string, j store.Job) { + jctx, stop := context.WithCancel(ctx) + defer stop() + defer func() { _ = r.store.ReleaseJob(context.WithoutCancel(ctx), j.ID, owner) }() + // The lease kept while working; lost, the work stops and another + // gateway has it. + go func() { + t := time.NewTicker(jobLease / 4) + defer t.Stop() + for { + select { + case <-jctx.Done(): + return + case <-t.C: + if held, err := r.store.RenewJobLease(jctx, j.ID, owner, time.Now().Add(jobLease)); err == nil && !held { + stop() + return + } + } + } + }() + if err := r.serveJob(jctx, &j); err != nil && jctx.Err() == nil { + r.log.Warn("a job stopped with an error; it is picked up again", "job", j.ID, "err", err) + } +} + +// serveJob replays the job's request and keeps what came back. A failure +// of the gateway's own (the store unreachable) is returned, and the job is +// picked up again; a refusal of the request is the job's end. +func (r *Router) serveJob(ctx context.Context, j *store.Job) error { + now := time.Now() + switch { + case j.Status == store.JobCancelling: + j.Status, j.FinishedAt = store.JobCancelled, now + return r.store.UpdateJob(ctx, *j) + case now.After(j.ExpiresAt): + j.Status, j.FinishedAt = store.JobExpired, now + j.Error = fmt.Sprintf("not finished within %s", jobWindow) + return r.store.UpdateJob(ctx, *j) + } + if j.Status == store.JobQueued { + j.Status, j.StartedAt = store.JobRunning, now + if err := r.store.UpdateJob(ctx, *j); err != nil { + return err + } + } + // A job found running was another gateway's, which died with it. It is + // run again from the start: the engine's queue does not outlive the + // request that was waiting on it. + fail := func(msg string) error { + j.Status, j.Error, j.FinishedAt = store.JobFailed, msg, time.Now() + return r.store.UpdateJob(ctx, *j) + } + p, err := r.principalOfKey(ctx, j.KeyID, "") + if err != nil { + return fail("the key this job was made with can no longer be used: " + err.Error()) + } + + // A cancel asked for while it runs ends the request, which takes the + // work out of the engine's queue. + rctx, cancel := context.WithCancel(ctx) + defer cancel() + go func() { + t := time.NewTicker(jobCancelCheck) + defer t.Stop() + for { + select { + case <-rctx.Done(): + return + case <-t.C: + if cur, err := r.store.JobByID(rctx, j.ID); err == nil && cur.Status == store.JobCancelling { + cancel() + return + } + } + } + }() + + handler := r.handleInference(j.Endpoint) + for attempt := 1; ; attempt++ { + hreq, err := http.NewRequestWithContext(auth.WithPrincipal(rctx, p), http.MethodPost, j.Endpoint, bytes.NewReader(j.Input)) + if err != nil { + return fail(err.Error()) + } + hreq.Header.Set("Content-Type", "application/json") + hreq.Header.Set("X-Request-Id", j.ID) + hreq.Header.Set(tagsHeader, "job="+j.ID) + w := &bufferWriter{header: http.Header{}, max: jobMaxResult} + handler(w, hreq) + if rctx.Err() != nil { + if ctx.Err() != nil { + // The lease was lost or the gateway is stopping: another + // gateway's to pick up. + return ctx.Err() + } + j.Status, j.FinishedAt = store.JobCancelled, time.Now() + return r.store.UpdateJob(ctx, *j) + } + status := w.status + if status == 0 { + status = http.StatusOK + } + // Waited out, for as long as the job has: a model loading, a rate + // limit, a queue full. Nobody is holding a connection, which is the + // point of a job. Not a spent budget, which says no Retry-After + // because waiting does not help. + if wait := w.header.Get("Retry-After"); wait != "" && (status == http.StatusTooManyRequests || status >= 500) && + time.Now().Before(j.ExpiresAt) { + secs, _ := strconv.Atoi(wait) + select { + case <-rctx.Done(): + continue // said above + case <-time.After(time.Duration(max(secs, 1)) * time.Second): + } + continue + } + if status >= 300 { + return fail(errorMessage(w.buf.Bytes())) + } + ct := w.header.Get("Content-Type") + if ct == "" { + ct = http.DetectContentType(w.buf.Bytes()) + } + _, _, _, days := r.batchCfg() + done := time.Now() + f := store.File{ID: newID("file-"), Owner: j.Owner, Purpose: purposeJobOutput, + Filename: j.ID + extensionFor(ct), Bytes: int64(w.buf.Len()), CreatedAt: done, ExpiresAt: done.AddDate(0, 0, days)} + if err := r.store.PutFile(ctx, f, w.buf.Bytes()); err != nil { + return err + } + j.Status, j.OutputFileID, j.ContentType, j.FinishedAt = store.JobCompleted, f.ID, ct, done + return r.store.UpdateJob(ctx, *j) + } +} + +// extensionFor names a result's file by what it is. +func extensionFor(contentType string) string { + ct, _, _ := strings.Cut(contentType, ";") + switch strings.TrimSpace(ct) { + case "application/json": + return ".json" + case "model/gltf-binary": + return ".glb" + case "model/gltf+json": + return ".gltf" + case "model/obj": + return ".obj" + case "image/png": + return ".png" + case "image/jpeg": + return ".jpg" + case "image/webp": + return ".webp" + } + return ".bin" +} diff --git a/internal/router/judge.go b/internal/router/judge.go index c16376f..eb4bda3 100644 --- a/internal/router/judge.go +++ b/internal/router/judge.go @@ -3,7 +3,6 @@ package router import ( "bytes" "context" - "errors" "fmt" "net/http" @@ -60,6 +59,9 @@ type bufferWriter struct { header http.Header status int buf bytes.Buffer + // max bounds what is kept; 1 MiB when unset, which is a judge's answer + // many times over. A job's result is a file, and sets more. + max int } func (b *bufferWriter) Header() http.Header { return b.header } @@ -72,8 +74,12 @@ func (b *bufferWriter) Write(p []byte) (int, error) { if b.status == 0 { b.status = http.StatusOK } - if b.buf.Len()+len(p) > 1<<20 { - return 0, errors.New("a judge's answer over 1 MiB") + limit := b.max + if limit <= 0 { + limit = 1 << 20 + } + if b.buf.Len()+len(p) > limit { + return 0, fmt.Errorf("an answer over %d bytes", limit) } return b.buf.Write(p) } diff --git a/internal/router/router.go b/internal/router/router.go index 0c11c1c..436aec3 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -77,8 +77,10 @@ type Router struct { canaries canaryStats // slo keeps each model's recent token timings, for admission. slo sloTracker - // batchWake has the batch worker look for work at once. + // batchWake has the batch worker look for work at once; jobWake the + // job worker. batchWake chan struct{} + jobWake chan struct{} // reqLogs is each project's keeping of requests, as last read. reqLogs requestLogs // answers are the answers projects reuse; nil reuses none. @@ -304,6 +306,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { r := &Router{ batchWake: make(chan struct{}, 1), + jobWake: make(chan struct{}, 1), cfg: cfg.Router, landing: cfg.Landing, log: log, @@ -528,6 +531,7 @@ func (r *Router) Start(ctx context.Context) { } if r.store != nil { go r.runBatches(ctx) + go r.runJobs(ctx) go r.pruneRequestBodies(ctx) go r.refreshTenants(ctx) go r.purgeSessions(ctx) diff --git a/internal/store/copy.go b/internal/store/copy.go index 8795858..8a488c1 100644 --- a/internal/store/copy.go +++ b/internal/store/copy.go @@ -68,6 +68,8 @@ var copyOrder = []string{ "files", "batches", "batch_results", + // After files, which a finished job's result is one of. + "jobs", "request_bodies", // Runs before their cases, which point at them. "eval_runs", diff --git a/internal/store/jobs.go b/internal/store/jobs.go new file mode 100644 index 0000000..4a820ac --- /dev/null +++ b/internal/store/jobs.go @@ -0,0 +1,208 @@ +package store + +import ( + "context" + "database/sql" + "errors" + "time" +) + +// Job is one request that is answered with an id and asked after: work that +// takes minutes, which no client should hold a connection for. The request +// is kept as it was sent and replayed by whichever gateway holds the lease. +type Job struct { + ID string + Owner string + KeyID string + Endpoint string + Model string + // Input is the request's body, as the client sent it. + Input []byte + Status string + // Progress is what the engine said of the work so far, as JSON; empty + // until it said anything. + Progress string + // PromptID is the engine's own id for the work, so a gateway that takes + // over a running job asks after the same work rather than queueing it + // again. + PromptID string + OutputFileID string + ContentType string + Error string + CreatedAt time.Time + StartedAt time.Time + FinishedAt time.Time + // ExpiresAt is when a job not yet finished is given up on. + ExpiresAt time.Time +} + +// Job states. +const ( + JobQueued = "queued" + JobRunning = "running" + JobCompleted = "completed" + JobFailed = "failed" + JobCancelling = "cancelling" + JobCancelled = "cancelled" + JobExpired = "expired" +) + +// JobDone is whether a state is an end. +func JobDone(status string) bool { + switch status { + case JobCompleted, JobFailed, JobCancelled, JobExpired: + return true + } + return false +} + +const jobCols = `id, owner, key_id, endpoint, model, input, status, progress, prompt_id, output_file_id, + content_type, error, created_at, started_at, finished_at, expires_at` + +func scanJob(sc interface{ Scan(...any) error }) (Job, error) { + var j Job + var created, started, finished, expires int64 + err := sc.Scan(&j.ID, &j.Owner, &j.KeyID, &j.Endpoint, &j.Model, &j.Input, &j.Status, &j.Progress, &j.PromptID, + &j.OutputFileID, &j.ContentType, &j.Error, &created, &started, &finished, &expires) + if err != nil { + return Job{}, err + } + at := func(v int64) time.Time { + if v == 0 { + return time.Time{} + } + return time.UnixMilli(v) + } + j.CreatedAt, j.StartedAt, j.FinishedAt, j.ExpiresAt = at(created), at(started), at(finished), at(expires) + return j, nil +} + +// PutJob records a new job. +func (s *DB) PutJob(ctx context.Context, j Job) error { + if j.Input == nil { + j.Input = []byte{} + } + _, err := s.db.ExecContext(ctx, `INSERT INTO jobs (id, owner, key_id, endpoint, model, input, status, created_at, expires_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`, + j.ID, j.Owner, j.KeyID, j.Endpoint, j.Model, j.Input, j.Status, millis(j.CreatedAt), millis(j.ExpiresAt)) + return err +} + +// UpdateJob writes how a job stands: its state, result, error and times. +// Its progress and prompt id are written by SetJobProgress, which the work +// calls often and which must not undo an end written here. +func (s *DB) UpdateJob(ctx context.Context, j Job) error { + _, err := s.db.ExecContext(ctx, `UPDATE jobs SET status = ?, output_file_id = ?, content_type = ?, error = ?, + started_at = ?, finished_at = ? WHERE id = ?`, + j.Status, j.OutputFileID, j.ContentType, j.Error, optMillis(j.StartedAt), optMillis(j.FinishedAt), j.ID) + return err +} + +// SetJobProgress writes what the engine said of a running job, and its id +// for the work when it is known. A job that has ended keeps its end. +func (s *DB) SetJobProgress(ctx context.Context, id, progress, promptID string) error { + _, err := s.db.ExecContext(ctx, `UPDATE jobs SET progress = ?, + prompt_id = CASE WHEN ? = '' THEN prompt_id ELSE ? END + WHERE id = ? AND status IN (?, ?)`, progress, promptID, promptID, id, JobRunning, JobCancelling) + return err +} + +// CancelJob asks for a job to stop: cancelling, unless it has ended already. +// It says the job as it stands. +func (s *DB) CancelJob(ctx context.Context, id string, now time.Time) (Job, error) { + // One not yet taken ends at once: there is nothing to stop. + if _, err := s.db.ExecContext(ctx, `UPDATE jobs SET status = ?, finished_at = ? + WHERE id = ? AND status = ? AND lease_until < ?`, JobCancelled, millis(now), id, JobQueued, millis(now)); err != nil { + return Job{}, err + } + if _, err := s.db.ExecContext(ctx, `UPDATE jobs SET status = ? WHERE id = ? AND status IN (?, ?)`, + JobCancelling, id, JobQueued, JobRunning); err != nil { + return Job{}, err + } + return s.JobByID(ctx, id) +} + +// JobByID finds a job, ErrNotFound when there is none. +func (s *DB) JobByID(ctx context.Context, id string) (Job, error) { + j, err := scanJob(s.db.QueryRowContext(ctx, `SELECT `+jobCols+` FROM jobs WHERE id = ?`, id)) + if errors.Is(err, sql.ErrNoRows) { + return Job{}, ErrNotFound + } + return j, err +} + +// Jobs lists an owner's jobs, the newest first, those made before the job +// after when it is given. +func (s *DB) Jobs(ctx context.Context, owner, after string, limit int) ([]Job, error) { + q := `SELECT ` + jobCols + ` FROM jobs WHERE owner = ?` + args := []any{owner} + if after != "" { + q += ` AND created_at < (SELECT created_at FROM jobs 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 []Job + for rows.Next() { + j, err := scanJob(rows) + if err != nil { + return nil, err + } + out = append(out, j) + } + return out, rows.Err() +} + +// ClaimJob takes the lease on the oldest job with work left that no gateway +// holds, for owner until the time given. ok is false when there is none. +func (s *DB) ClaimJob(ctx context.Context, owner string, now, until time.Time) (Job, bool, error) { + var id string + err := s.db.QueryRowContext(ctx, `UPDATE jobs SET lease_owner = ?, lease_until = ? + WHERE id = (SELECT id FROM jobs WHERE status IN (?, ?, ?) AND lease_until < ? + ORDER BY created_at LIMIT 1) + AND lease_until < ? + RETURNING id`, owner, millis(until), JobQueued, JobRunning, JobCancelling, millis(now), millis(now)).Scan(&id) + if errors.Is(err, sql.ErrNoRows) { + return Job{}, false, nil + } + if err != nil { + return Job{}, false, err + } + j, err := s.JobByID(ctx, id) + return j, err == nil, err +} + +// RenewJobLease keeps a lease held; false when it is held no longer. +func (s *DB) RenewJobLease(ctx context.Context, id, owner string, until time.Time) (bool, error) { + res, err := s.db.ExecContext(ctx, `UPDATE jobs 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 +} + +// ReleaseJob gives a lease up. +func (s *DB) ReleaseJob(ctx context.Context, id, owner string) error { + _, err := s.db.ExecContext(ctx, `UPDATE jobs SET lease_until = 0, lease_owner = '' WHERE id = ? AND lease_owner = ?`, id, owner) + return err +} + +// DeleteJobsBefore forgets the jobs that ended before the time given, whose +// result files are gone by then too. +func (s *DB) DeleteJobsBefore(ctx context.Context, before time.Time) (int64, error) { + res, err := s.db.ExecContext(ctx, `DELETE FROM jobs WHERE finished_at > 0 AND finished_at < ?`, millis(before)) + if err != nil { + return 0, err + } + return res.RowsAffected() +} diff --git a/internal/store/jobs_test.go b/internal/store/jobs_test.go new file mode 100644 index 0000000..09acb99 --- /dev/null +++ b/internal/store/jobs_test.go @@ -0,0 +1,97 @@ +package store + +import ( + "testing" + "time" +) + +func TestJobs(t *testing.T) { + checkJobs(t, capacityStore(t)) +} + +func TestPostgresJobs(t *testing.T) { + checkJobs(t, openPG(t)) +} + +func checkJobs(t *testing.T, db *DB) { + ctx := t.Context() + now := time.Now() + j := Job{ID: "job-1", Owner: "p1", KeyID: "k", Endpoint: "/v1/3d/generations", Model: "trellis2", + Input: []byte(`{"model":"trellis2"}`), Status: JobQueued, CreatedAt: now, ExpiresAt: now.Add(24 * time.Hour)} + if err := db.PutJob(ctx, j); err != nil { + t.Fatal(err) + } + got, ok, err := db.ClaimJob(ctx, "gw-a", now, now.Add(time.Minute)) + if err != nil || !ok || got.ID != "job-1" || string(got.Input) != `{"model":"trellis2"}` { + t.Fatalf("claim: %+v %v %v", got, ok, err) + } + if _, ok, _ := db.ClaimJob(ctx, "gw-b", now, now.Add(time.Minute)); ok { + t.Error("a leased job was claimed twice") + } + if held, _ := db.RenewJobLease(ctx, "job-1", "gw-b", now.Add(time.Hour)); held { + t.Error("another gateway renewed the lease") + } + if held, _ := db.RenewJobLease(ctx, "job-1", "gw-a", now.Add(time.Hour)); !held { + t.Error("the holder could not renew") + } + + // Progress is written while it runs, and never over an end. + got.Status, got.StartedAt = JobRunning, now + if err := db.UpdateJob(ctx, got); err != nil { + t.Fatal(err) + } + if err := db.SetJobProgress(ctx, "job-1", `{"percent":40}`, "p-7"); err != nil { + t.Fatal(err) + } + if err := db.SetJobProgress(ctx, "job-1", `{"percent":60}`, ""); err != nil { + t.Fatal(err) + } + cur, _ := db.JobByID(ctx, "job-1") + if cur.Progress != `{"percent":60}` || cur.PromptID != "p-7" || cur.Status != JobRunning { + t.Errorf("%+v", cur) + } + + // A running job asked to stop is cancelling, for the worker to stop it. + cur, err = db.CancelJob(ctx, "job-1", now) + if err != nil || cur.Status != JobCancelling { + t.Fatalf("%+v %v", cur, err) + } + cur.Status, cur.FinishedAt = JobCancelled, now + _ = db.UpdateJob(ctx, cur) + if err := db.SetJobProgress(ctx, "job-1", `{"percent":99}`, ""); err != nil { + t.Fatal(err) + } + if cur, _ = db.JobByID(ctx, "job-1"); cur.Progress != `{"percent":60}` || cur.Status != JobCancelled { + t.Errorf("progress written over an end: %+v", cur) + } + _ = db.ReleaseJob(ctx, "job-1", "gw-a") + if _, ok, _ := db.ClaimJob(ctx, "gw-b", now.Add(2*time.Hour), now.Add(3*time.Hour)); ok { + t.Error("an ended job was claimed") + } + + // One not yet taken ends at once when cancelled. + j2 := j + j2.ID, j2.CreatedAt = "job-2", now.Add(time.Second) + _ = db.PutJob(ctx, j2) + if cur, _ = db.CancelJob(ctx, "job-2", now.Add(2*time.Second)); cur.Status != JobCancelled || cur.FinishedAt.IsZero() { + t.Errorf("%+v", cur) + } + + // Listed newest first, another owner's not at all. + j3 := j + j3.ID, j3.Owner, j3.CreatedAt = "job-3", "p2", now.Add(3*time.Second) + _ = db.PutJob(ctx, j3) + js, _ := db.Jobs(ctx, "p1", "", 10) + if len(js) != 2 || js[0].ID != "job-2" || js[1].ID != "job-1" { + t.Errorf("listed %+v", js) + } + if js, _ = db.Jobs(ctx, "p1", "job-2", 10); len(js) != 1 || js[0].ID != "job-1" { + t.Errorf("after: %+v", js) + } + if _, err := db.JobByID(ctx, "nope"); err != ErrNotFound { + t.Errorf("err %v", err) + } + if n, _ := db.DeleteJobsBefore(ctx, now.Add(time.Hour)); n != 2 { + t.Errorf("deleted %d", n) + } +} diff --git a/internal/store/migrations/postgres/0044_jobs.sql b/internal/store/migrations/postgres/0044_jobs.sql new file mode 100644 index 0000000..a8be981 --- /dev/null +++ b/internal/store/migrations/postgres/0044_jobs.sql @@ -0,0 +1,34 @@ +-- Ported from the SQLite migration of the same number. +-- Jobs (BACKLOG G1): a request that takes minutes — a 3D model from a +-- picture — answered with an id and asked after, rather than held open. +-- The request is kept as it was sent, so whichever gateway holds the lease +-- replays it through the same path a client's request takes; the result is +-- a file in the files table, kept for the batch files' retention. +CREATE TABLE jobs ( + id TEXT PRIMARY KEY, + owner TEXT NOT NULL, + key_id TEXT NOT NULL, + endpoint TEXT NOT NULL, + model TEXT NOT NULL, + input BYTEA NOT NULL, + -- status is queued, running, completed, failed, cancelling, cancelled or + -- expired. + status TEXT NOT NULL, + -- progress is what the engine said of the work so far, as JSON: queue + -- position, percent, step. Empty until it says anything. + progress TEXT NOT NULL DEFAULT '', + -- prompt_id is the engine's own id for the work (ComfyUI's prompt_id), + -- so a gateway that takes over a running job asks after the same work. + prompt_id TEXT NOT NULL DEFAULT '', + output_file_id TEXT NOT NULL DEFAULT '', + content_type TEXT NOT NULL DEFAULT '', + error TEXT NOT NULL DEFAULT '', + created_at BIGINT NOT NULL, + started_at BIGINT NOT NULL DEFAULT 0, + finished_at BIGINT NOT NULL DEFAULT 0, + expires_at BIGINT NOT NULL, + lease_owner TEXT NOT NULL DEFAULT '', + lease_until BIGINT NOT NULL DEFAULT 0 +); +CREATE INDEX jobs_owner ON jobs (owner, created_at); +CREATE INDEX jobs_status ON jobs (status, lease_until); diff --git a/internal/store/migrations/sqlite/0044_jobs.sql b/internal/store/migrations/sqlite/0044_jobs.sql new file mode 100644 index 0000000..31caceb --- /dev/null +++ b/internal/store/migrations/sqlite/0044_jobs.sql @@ -0,0 +1,33 @@ +-- Jobs (BACKLOG G1): a request that takes minutes — a 3D model from a +-- picture — answered with an id and asked after, rather than held open. +-- The request is kept as it was sent, so whichever gateway holds the lease +-- replays it through the same path a client's request takes; the result is +-- a file in the files table, kept for the batch files' retention. +CREATE TABLE jobs ( + id TEXT PRIMARY KEY, + owner TEXT NOT NULL, + key_id TEXT NOT NULL, + endpoint TEXT NOT NULL, + model TEXT NOT NULL, + input BLOB NOT NULL, + -- status is queued, running, completed, failed, cancelling, cancelled or + -- expired. + status TEXT NOT NULL, + -- progress is what the engine said of the work so far, as JSON: queue + -- position, percent, step. Empty until it says anything. + progress TEXT NOT NULL DEFAULT '', + -- prompt_id is the engine's own id for the work (ComfyUI's prompt_id), + -- so a gateway that takes over a running job asks after the same work. + prompt_id TEXT NOT NULL DEFAULT '', + output_file_id TEXT NOT NULL DEFAULT '', + content_type TEXT NOT NULL DEFAULT '', + error TEXT NOT NULL DEFAULT '', + created_at INTEGER NOT NULL, + started_at INTEGER NOT NULL DEFAULT 0, + finished_at INTEGER NOT NULL DEFAULT 0, + expires_at INTEGER NOT NULL, + lease_owner TEXT NOT NULL DEFAULT '', + lease_until INTEGER NOT NULL DEFAULT 0 +); +CREATE INDEX jobs_owner ON jobs (owner, created_at); +CREATE INDEX jobs_status ON jobs (status, lease_until);