diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 44e689f..a20a546 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -7595,6 +7595,10 @@ "type": "integer", "format": "int64" }, + "images": { + "type": "integer", + "format": "int64" + }, "cost_micros": { "type": "integer", "format": "int64" diff --git a/docs/install.md b/docs/install.md index 8a72c5c..52b3190 100644 --- a/docs/install.md +++ b/docs/install.md @@ -1308,9 +1308,17 @@ before the picture is sent. Requests for pictures pass the same checks as any other request: keys, allowlists, limits, admission, and the guardrails on the prompt. A client that -leaves takes its workflow off the ComfyUI queue. Usage is recorded, but for -now it is charged as the prompt's tokens; a per-picture price is not yet -available. +leaves takes its workflow off the ComfyUI queue. + +Each request's usage row records how many pictures it made (`llmgw ctl request +ID` shows them). A picture's price is set with `per_image`, in micro-units like +the other prices, and is charged on top of the prompt's tokens: + +```yaml +auth: + pricing: + flux-schnell: {prompt_per_1k: 0, per_image: 20000} +``` Start ComfyUI with `--disable-smart-memory` on a card that llmgw also places models on. Without it, ComfyUI keeps the model on the card after each picture. diff --git a/docs/reference.md b/docs/reference.md index e8fa916..2421227 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -347,6 +347,7 @@ rather than repeated. | `secret_file` | `string` | [Webhook](#webhook) | | `prompt_per_1k` | `int64` | [accounting.Price](#accountingprice) | | `completion_per_1k` | `int64` | [accounting.Price](#accountingprice) | +| `per_image` | `int64` | [accounting.Price](#accountingprice) | | `name` | `string` | [admission.Band](#admissionband) | | `priority` | `int` | [admission.Band](#admissionband) | | `saturation_ceiling` | `float64` | [admission.Band](#admissionband) | @@ -381,6 +382,7 @@ rather than repeated. | `saturation_ceiling` | `float64` | [config.Band](#configband) | | `prompt_per_1k` | `int64` | [config.Price](#configprice) | | `completion_per_1k` | `int64` | [config.Price](#configprice) | +| `per_image` | `int64` | [config.Price](#configprice) | | `disabled` | `[]string` | [engines.Config](#enginesconfig) | | `vllm` | `VLLMConfig` | [engines.Config](#enginesconfig) | | `sglang` | `SGLangConfig` | [engines.Config](#enginesconfig) | @@ -2936,6 +2938,14 @@ is sha256=. `int64` +### `per_image` + +`int64` + +PerImage is what one picture costs, for a model that makes them: a +picture's cost is the card's time, which its prompt's few tokens do +not measure. + ## admission.Band ### `name` @@ -3178,6 +3188,13 @@ Takes a [Tracing](#tracing) block. `int64` +### `per_image` + +`int64` + +PerImage is the cost of one picture, for a model that makes them +(/v1/images/generations), on top of its prompt's tokens. + ## engines.Config ### `disabled` diff --git a/internal/accounting/accounting.go b/internal/accounting/accounting.go index e9b2c4d..0359268 100644 --- a/internal/accounting/accounting.go +++ b/internal/accounting/accounting.go @@ -323,7 +323,7 @@ func (a *Accountant) Record(res Reservation, u store.Usage) { if u.ServedModel != "" && u.ServedModel != u.Model && a.pricing != nil { price = a.pricing.For(u.ServedModel) } - u.CostMicros = price.Cost(u.PromptTokens, u.CompletionTokens) + u.CostMicros = price.Cost(u.PromptTokens, u.CompletionTokens, u.Images) for _, b := range res.Budgets { a.budgets.Add(b.Counter, b.Period, u.CostMicros) } diff --git a/internal/accounting/accounting_test.go b/internal/accounting/accounting_test.go index 4372b16..0b27457 100644 --- a/internal/accounting/accounting_test.go +++ b/internal/accounting/accounting_test.go @@ -39,17 +39,27 @@ func TestCostRoundsUpSoSmallRequestsAreNotFree(t *testing.T) { // Rounding down would make a stream of tiny requests cost nothing, which is // exactly the traffic shape a quota exists to bound. p := Price{PromptPer1k: 1000, CompletionPer1k: 3000} - if got := p.Cost(1, 0); got != 1 { + if got := p.Cost(1, 0, 0); got != 1 { t.Errorf("one prompt token cost %d", got) } - if got := p.Cost(1000, 1000); got != 4000 { + if got := p.Cost(1000, 1000, 0); got != 4000 { t.Errorf("a thousand of each cost %d, want 4000", got) } - if got := p.Cost(0, 0); got != 0 { + if got := p.Cost(0, 0, 0); got != 0 { t.Errorf("an empty request cost %d", got) } } +func TestAPictureCostsItsPriceOnTopOfItsPrompt(t *testing.T) { + p := Price{PromptPer1k: 1000, PerImage: 50000} + if got := p.Cost(20, 0, 3); got != 20+150000 { + t.Errorf("three pictures from a 20-token prompt cost %d", got) + } + if (Price{PerImage: 1}).Zero() { + t.Error("a model priced per picture reads as free, and would pass a spent budget") + } +} + func TestZeroPricedModelsStayAvailableWhenTheBudgetIsSpent(t *testing.T) { // This is what makes budget fallback work: a key with nothing left can // still be served something that costs nothing. diff --git a/internal/accounting/cost.go b/internal/accounting/cost.go index 770eab2..a2423d5 100644 --- a/internal/accounting/cost.go +++ b/internal/accounting/cost.go @@ -18,6 +18,10 @@ import ( type Price struct { PromptPer1k int64 `koanf:"prompt_per_1k"` CompletionPer1k int64 `koanf:"completion_per_1k"` + // PerImage is what one picture costs, for a model that makes them: a + // picture's cost is the card's time, which its prompt's few tokens do + // not measure. + PerImage int64 `koanf:"per_image"` } // Zero reports whether this model spends nothing. @@ -25,14 +29,15 @@ type Price struct { // It matters beyond arithmetic: a zero-cost model stays available to a key // whose budget is spent, which is what makes the plan's budget fallback work // (ยง7.7). -func (p Price) Zero() bool { return p.PromptPer1k == 0 && p.CompletionPer1k == 0 } +func (p Price) Zero() bool { return p.PromptPer1k == 0 && p.CompletionPer1k == 0 && p.PerImage == 0 } // Cost computes what a request spent, rounding up. // // Up rather than to nearest: rounding down would make a stream of tiny requests // free, which is exactly the shape of traffic a quota exists to bound. -func (p Price) Cost(promptTokens, completionTokens int64) int64 { - return divUp(promptTokens*p.PromptPer1k, 1000) + divUp(completionTokens*p.CompletionPer1k, 1000) +func (p Price) Cost(promptTokens, completionTokens, images int64) int64 { + return divUp(promptTokens*p.PromptPer1k, 1000) + divUp(completionTokens*p.CompletionPer1k, 1000) + + max(images, 0)*p.PerImage } func divUp(n, d int64) int64 { diff --git a/internal/cli/ctl_request.go b/internal/cli/ctl_request.go index 59f2ddb..0b6c814 100644 --- a/internal/cli/ctl_request.go +++ b/internal/cli/ctl_request.go @@ -56,6 +56,9 @@ func newRequestCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { fmt.Fprintf(out, "status %d\n", v.Status) fmt.Fprintf(out, "time %d ms total: %d queued, %d to the first token\n", v.LatencyMs, v.QueueMs, v.TTFTMs) fmt.Fprintf(out, "tokens %d in, %d out\n", v.PromptTokens, v.CompletionTokens) + if v.Images > 0 { + fmt.Fprintf(out, "images %d\n", v.Images) + } if v.CacheHit { fmt.Fprintf(out, "cache answered from the project's answer cache: nothing generated, nothing charged\n") } diff --git a/internal/config/config.go b/internal/config/config.go index 582c25c..b136961 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -493,6 +493,9 @@ type Personal struct { type Price struct { PromptPer1k int64 `koanf:"prompt_per_1k"` CompletionPer1k int64 `koanf:"completion_per_1k"` + // PerImage is the cost of one picture, for a model that makes them + // (/v1/images/generations), on top of its prompt's tokens. + PerImage int64 `koanf:"per_image"` } // Control configures the control plane. diff --git a/internal/control/adminapi/api.go b/internal/control/adminapi/api.go index 6e770ed..19e0de8 100644 --- a/internal/control/adminapi/api.go +++ b/internal/control/adminapi/api.go @@ -1611,6 +1611,7 @@ type RequestView struct { TTFTMs int64 `json:"ttft_ms"` PromptTokens int64 `json:"prompt_tokens"` CompletionTokens int64 `json:"completion_tokens"` + Images int64 `json:"images,omitempty"` CostMicros int64 `json:"cost_micros"` // What the guardrails did, by name. GuardAction string `json:"guard_action,omitempty"` @@ -1644,7 +1645,7 @@ func (a *API) request(w http.ResponseWriter, r *http.Request) { RequestID: u.RequestID, At: u.At, KeyID: u.KeyID, UserID: u.UserID, ProjectID: u.ProjectID, Model: u.Model, ServedModel: u.ServedModel, Replica: u.Replica, Status: u.Status, LatencyMs: u.LatencyMs, QueueMs: u.QueueMs, TTFTMs: u.TTFTMs, - PromptTokens: u.PromptTokens, CompletionTokens: u.CompletionTokens, CostMicros: u.CostMicros, + PromptTokens: u.PromptTokens, CompletionTokens: u.CompletionTokens, Images: u.Images, CostMicros: u.CostMicros, GuardAction: u.GuardAction, GuardPolicy: u.GuardPolicy, GuardCategory: u.GuardCategory, GuardRule: u.GuardRule, GuardMs: u.GuardMs, EndUser: u.EndUser, Tags: u.Tags, CacheHit: u.CacheHit, }) diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 0afa759..73d0a96 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -449,6 +449,7 @@ var Reference = []Block{ {Package: "accounting", Name: "Price", Title: "accounting.Price", Doc: "", Fields: []Field{ {Key: "prompt_per_1k", Type: "int64", Doc: "", Nested: ""}, {Key: "completion_per_1k", Type: "int64", Doc: "", Nested: ""}, + {Key: "per_image", Type: "int64", Doc: "PerImage is what one picture costs, for a model that makes them: a\npicture's cost is the card's time, which its prompt's few tokens do\nnot measure.", Nested: ""}, }}, {Package: "admission", Name: "Band", Title: "admission.Band", Doc: "", Fields: []Field{ {Key: "name", Type: "string", Doc: "", Nested: ""}, @@ -493,6 +494,7 @@ var Reference = []Block{ {Package: "config", Name: "Price", Title: "config.Price", Doc: "", Fields: []Field{ {Key: "prompt_per_1k", Type: "int64", Doc: "", Nested: ""}, {Key: "completion_per_1k", Type: "int64", Doc: "", Nested: ""}, + {Key: "per_image", Type: "int64", Doc: "PerImage is the cost of one picture, for a model that makes them\n(/v1/images/generations), on top of its prompt's tokens.", Nested: ""}, }}, {Package: "engines", Name: "Config", Title: "engines.Config", Doc: "", Fields: []Field{ {Key: "disabled", Type: "[]string", Doc: "Disabled are engines this node never installs or advertises, whatever\nthe control plane asks: vLLM on a node of cards no build of it runs on,\nthat is meant for llama.cpp alone.", Nested: ""}, diff --git a/internal/router/handlers.go b/internal/router/handlers.go index c112574..c81200e 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -750,6 +750,7 @@ func (r *Router) account(req *parsedRequest, res dispatchResult, status int, too // the truth than charging nothing. u.PromptTokens = int64(req.PromptTokens) } + u.Images = int64(res.outcome.Images) r.m.tenant.observe(req, status, u.PromptTokens, u.CompletionTokens, res.outcome.TTFT) r.acct.Record(req.reservation, u) } diff --git a/internal/router/images.go b/internal/router/images.go index 940c751..1f05de3 100644 --- a/internal/router/images.go +++ b/internal/router/images.go @@ -190,7 +190,8 @@ func (r *Router) generateImages(ctx context.Context, w http.ResponseWriter, req w.WriteHeader(http.StatusOK) _, _ = w.Write(out) took := time.Since(start) - return proxy.Outcome{Wrote: true, Status: http.StatusOK, Bytes: int64(len(out)), TTFT: took, Decode: took}, nil + return proxy.Outcome{Wrote: true, Status: http.StatusOK, Bytes: int64(len(out)), TTFT: took, Decode: took, + Images: len(images)}, nil } // ctxCauseIsClient reports whether a context ended because the client went diff --git a/internal/router/images_test.go b/internal/router/images_test.go index 40f04b7..b6b90de 100644 --- a/internal/router/images_test.go +++ b/internal/router/images_test.go @@ -2,24 +2,33 @@ package router import ( "bytes" + "crypto/sha256" "encoding/base64" "encoding/json" "fmt" "net/http" "net/http/httptest" + "path/filepath" "strings" "sync" "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" ) // imageRouter has a ComfyUI model, flux, beside a text model, and returns // what the ComfyUI server was sent. func imageRouter(t *testing.T) (*Router, *map[string]any, *sync.Mutex) { + t.Helper() + return imageRouterWith(t, nil) +} + +func imageRouterWith(t *testing.T, edit func(*config.Server)) (*Router, *map[string]any, *sync.Mutex) { t.Helper() var mu sync.Mutex var queued map[string]any @@ -37,7 +46,18 @@ func imageRouter(t *testing.T) (*Router, *map[string]any, *sync.Mutex) { mu.Unlock() fmt.Fprint(w, `{"prompt_id":"p1"}`) case "/history/p1": - fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true},"outputs":{"7":{"images":[{"filename":"a.png","type":"temp"}]}}}}`) + mu.Lock() + n := 1 + if latent, ok := queued["4"].(map[string]any); ok { + n = int(latent["inputs"].(map[string]any)["batch_size"].(float64)) + } + mu.Unlock() + var images []string + for i := range n { + images = append(images, fmt.Sprintf(`{"filename":"%d.png","type":"temp"}`, i)) + } + fmt.Fprintf(w, `{"p1":{"status":{"status_str":"success","completed":true},"outputs":{"7":{"images":[%s]}}}}`, + strings.Join(images, ",")) case "/view": w.Header().Set("Content-Type", "image/png") fmt.Fprint(w, "png-bytes") @@ -50,6 +70,9 @@ func imageRouter(t *testing.T) (*Router, *map[string]any, *sync.Mutex) { cfg := baseConfig() cfg.Models = []catalog.Model{catModel("chat", nil), {Name: "flux", Source: "external", External: &catalog.External{URL: srv.URL, Model: "flux1-schnell-fp8.safetensors", API: catalog.APIComfyUI}}} + if edit != nil { + edit(&cfg) + } rt := newTestRouter(t, cfg) rt.comfy = &comfy.Client{Poll: time.Millisecond} rt.SetEndpoints(t.Context(), []endpoint.Info{ @@ -151,3 +174,56 @@ func TestAComfyUIModelIsCheckedInTheCatalog(t *testing.T) { } } } + +func TestPicturesAreChargedByThePictureAndRecorded(t *testing.T) { + rt, _, _ := imageRouterWith(t, func(cfg *config.Server) { + cfg.Auth.Enabled = true + cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db") + cfg.Auth.Pricing = map[string]config.Price{"flux": {PromptPer1k: 1000, PerImage: 50000}} + cfg.Auth.FlushInterval = 20 * time.Millisecond + cfg.Auth.BatchSize = 1 + }) + org, _ := rt.store.CreateOrg(t.Context(), "acme") + team, _ := rt.store.CreateTeam(t.Context(), org, "design") + project, err := rt.store.CreateProject(t.Context(), team, "posters") + if err != nil { + t.Fatal(err) + } + secret := store.KeyPrefix + strings.Repeat("b", 43) + sum := sha256.Sum256([]byte(secret)) + if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-1", Hash: sum[:], Prefix: "bbbb", Name: "t", + ProjectID: project, CreatedAt: time.Now()}); err != nil { + t.Fatal(err) + } + + b, _ := json.Marshal(map[string]any{"model": "flux", "prompt": "three posters", "n": 3}) + req := httptest.NewRequest(http.MethodPost, imagesPath, bytes.NewReader(b)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+secret) + rec := httptest.NewRecorder() + rt.Handler().ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("status %d: %s", rec.Code, rec.Body.String()) + } + id := rec.Header().Get("X-Request-Id") + deadline := time.Now().Add(5 * time.Second) + for { + u, ok, err := rt.store.UsageByRequest(t.Context(), id) + if err != nil { + t.Fatal(err) + } + if ok { + if u.Images != 3 { + t.Errorf("recorded %d pictures, want 3", u.Images) + } + if want := (u.PromptTokens*1000+999)/1000 + 3*50000; u.CostMicros != want { + t.Errorf("charged %d micros, want %d (%d prompt tokens and three pictures)", u.CostMicros, want, u.PromptTokens) + } + return + } + if time.Now().After(deadline) { + t.Fatal("no usage was recorded") + } + time.Sleep(10 * time.Millisecond) + } +} diff --git a/internal/router/proxy/proxy.go b/internal/router/proxy/proxy.go index 9afbf85..61872d8 100644 --- a/internal/router/proxy/proxy.go +++ b/internal/router/proxy/proxy.go @@ -224,6 +224,8 @@ type Outcome struct { Decode time.Duration // Usage is the engine-reported token usage, if it was present. Usage *Usage + // Images is how many pictures were made, by an endpoint that makes them. + Images int // Bytes is how many bytes were relayed to the client. Bytes int64 // Reply is the assistant's text, when CaptureReply asked for it. Truncated diff --git a/internal/router/router.go b/internal/router/router.go index 0aa4d0c..0c11c1c 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -374,7 +374,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { prices := make(map[string]accounting.Price, len(cfg.Auth.Pricing)) for name, p := range cfg.Auth.Pricing { prices[name] = accounting.Price{ - PromptPer1k: p.PromptPer1k, CompletionPer1k: p.CompletionPer1k, + PromptPer1k: p.PromptPer1k, CompletionPer1k: p.CompletionPer1k, PerImage: p.PerImage, } } r.recorder = accounting.NewRecorder(st, r.auth, accounting.RecorderConfig{ @@ -390,6 +390,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { accounting.NewPricing(prices, accounting.Price{ PromptPer1k: cfg.Auth.DefaultPrice.PromptPer1k, CompletionPer1k: cfg.Auth.DefaultPrice.CompletionPer1k, + PerImage: cfg.Auth.DefaultPrice.PerImage, }), r.recorder, log) r.threads = r.store diff --git a/internal/store/migrations/postgres/0041_usage_images.sql b/internal/store/migrations/postgres/0041_usage_images.sql new file mode 100644 index 0000000..6f33802 --- /dev/null +++ b/internal/store/migrations/postgres/0041_usage_images.sql @@ -0,0 +1,4 @@ +-- Ported from the SQLite migration of the same number. +-- How many pictures a request made, for a model that makes them: what a +-- per-picture price is charged by, and what its tokens do not say. +ALTER TABLE usage ADD COLUMN images BIGINT NOT NULL DEFAULT 0; diff --git a/internal/store/migrations/sqlite/0041_usage_images.sql b/internal/store/migrations/sqlite/0041_usage_images.sql new file mode 100644 index 0000000..1e98270 --- /dev/null +++ b/internal/store/migrations/sqlite/0041_usage_images.sql @@ -0,0 +1,3 @@ +-- How many pictures a request made, for a model that makes them: what a +-- per-picture price is charged by, and what its tokens do not say. +ALTER TABLE usage ADD COLUMN images INTEGER NOT NULL DEFAULT 0; diff --git a/internal/store/sqlite.go b/internal/store/sqlite.go index 9ebad5f..62ad830 100644 --- a/internal/store/sqlite.go +++ b/internal/store/sqlite.go @@ -442,8 +442,8 @@ func (s *DB) WriteUsage(ctx context.Context, records []Usage) error { `INSERT INTO usage (ts, key_id, project_id, model, served_model, prompt_tokens, completion_tokens, cost_micros, status, latency_ms, user_id, request_id, replica, queue_ms, ttft_ms, - guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`) + guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit, images) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`) if err != nil { return err } @@ -457,7 +457,7 @@ func (s *DB) WriteUsage(ctx context.Context, records []Usage) error { r.Model, r.ServedModel, r.PromptTokens, r.CompletionTokens, r.CostMicros, r.Status, r.LatencyMs, r.UserID, r.RequestID, r.Replica, r.QueueMs, r.TTFTMs, - r.GuardAction, r.GuardPolicy, r.GuardCategory, r.GuardRule, r.GuardMs, r.EndUser, tagsJSON(r.Tags), hit); err != nil { + r.GuardAction, r.GuardPolicy, r.GuardCategory, r.GuardRule, r.GuardMs, r.EndUser, tagsJSON(r.Tags), hit, r.Images); err != nil { return err } } @@ -871,12 +871,12 @@ func (s *DB) UsageByRequest(ctx context.Context, requestID string) (Usage, bool, err := s.db.QueryRowContext(ctx, `SELECT ts, key_id, project_id, COALESCE(user_id, ''), model, served_model, prompt_tokens, completion_tokens, cost_micros, status, latency_ms, request_id, replica, queue_ms, ttft_ms, - guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit + guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit, images FROM usage WHERE request_id = ? ORDER BY ts DESC LIMIT 1`, requestID). Scan(&ts, &u.KeyID, &u.ProjectID, &u.UserID, &u.Model, &u.ServedModel, &u.PromptTokens, &u.CompletionTokens, &u.CostMicros, &u.Status, &u.LatencyMs, &u.RequestID, &u.Replica, &u.QueueMs, &u.TTFTMs, &u.GuardAction, &u.GuardPolicy, &u.GuardCategory, &u.GuardRule, &u.GuardMs, - &u.EndUser, &tags, &hit) + &u.EndUser, &tags, &hit, &u.Images) if errors.Is(err, sql.ErrNoRows) { return Usage{}, false, nil } diff --git a/internal/store/store.go b/internal/store/store.go index ff2ed23..73d0c5f 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -241,6 +241,8 @@ type Usage struct { ServedModel string PromptTokens int64 CompletionTokens int64 + // Images is how many pictures were made, for a request for them. + Images int64 // CostMicros is what this request spent against the budget. CostMicros int64 Status int