diff --git a/.sqlx/query-12023b6843111ff492203f55fb856ebfbc9bf64de34e1649af691d5f97821459.json b/.sqlx/query-12023b6843111ff492203f55fb856ebfbc9bf64de34e1649af691d5f97821459.json new file mode 100644 index 000000000..67fb690e0 --- /dev/null +++ b/.sqlx/query-12023b6843111ff492203f55fb856ebfbc9bf64de34e1649af691d5f97821459.json @@ -0,0 +1,32 @@ +{ + "db_name": "PostgreSQL", + "query": "\n SELECT id, self_hosted_only, scheduling_priority\n FROM users\n WHERE is_deleted = FALSE\n AND (self_hosted_only OR scheduling_priority IS NOT NULL)\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": "Uuid" + }, + { + "ordinal": 1, + "name": "self_hosted_only", + "type_info": "Bool" + }, + { + "ordinal": 2, + "name": "scheduling_priority", + "type_info": "Int4" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + false, + false, + true + ] + }, + "hash": "12023b6843111ff492203f55fb856ebfbc9bf64de34e1649af691d5f97821459" +} diff --git a/.sqlx/query-95bac76d1001d7010e8921cf4d531dc0de7b833f86ee0d35929c784990c7c0eb.json b/.sqlx/query-95bac76d1001d7010e8921cf4d531dc0de7b833f86ee0d35929c784990c7c0eb.json new file mode 100644 index 000000000..fa2155e1c --- /dev/null +++ b/.sqlx/query-95bac76d1001d7010e8921cf4d531dc0de7b833f86ee0d35929c784990c7c0eb.json @@ -0,0 +1,50 @@ +{ + "db_name": "PostgreSQL", + "query": "\n SELECT p.user_id, dm.alias, p.serving_class, p.priority, p.ttft_ms, p.itl_ms\n FROM account_model_policies p\n INNER JOIN deployed_models dm ON dm.id = p.deployed_model_id\n INNER JOIN users u ON u.id = p.user_id\n WHERE dm.deleted = FALSE\n AND u.is_deleted = FALSE\n AND (p.serving_class IS NULL OR p.serving_class = 'standard')\n ORDER BY p.serving_class NULLS FIRST\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "user_id", + "type_info": "Uuid" + }, + { + "ordinal": 1, + "name": "alias", + "type_info": "Varchar" + }, + { + "ordinal": 2, + "name": "serving_class", + "type_info": "Text" + }, + { + "ordinal": 3, + "name": "priority", + "type_info": "Int4" + }, + { + "ordinal": 4, + "name": "ttft_ms", + "type_info": "Int4" + }, + { + "ordinal": 5, + "name": "itl_ms", + "type_info": "Int4" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + false, + false, + true, + true, + true, + true + ] + }, + "hash": "95bac76d1001d7010e8921cf4d531dc0de7b833f86ee0d35929c784990c7c0eb" +} diff --git a/.sqlx/query-ba34540c4fd715b7f949617bdda7d8fe20a3194a39ad1665aa7582d6b7f9a9eb.json b/.sqlx/query-ba34540c4fd715b7f949617bdda7d8fe20a3194a39ad1665aa7582d6b7f9a9eb.json deleted file mode 100644 index 9501782af..000000000 --- a/.sqlx/query-ba34540c4fd715b7f949617bdda7d8fe20a3194a39ad1665aa7582d6b7f9a9eb.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "\n SELECT id, self_hosted_only\n FROM users\n WHERE is_deleted = FALSE\n AND self_hosted_only\n ", - "describe": { - "columns": [ - { - "ordinal": 0, - "name": "id", - "type_info": "Uuid" - }, - { - "ordinal": 1, - "name": "self_hosted_only", - "type_info": "Bool" - } - ], - "parameters": { - "Left": [] - }, - "nullable": [ - false, - false - ] - }, - "hash": "ba34540c4fd715b7f949617bdda7d8fe20a3194a39ad1665aa7582d6b7f9a9eb" -} diff --git a/docs/src/reference/configuration.md b/docs/src/reference/configuration.md index 48b68c980..ed521742f 100644 --- a/docs/src/reference/configuration.md +++ b/docs/src/reference/configuration.md @@ -80,7 +80,7 @@ model_provisioning: | `enabled` | boolean | `false` | Apply the declarative model catalog during startup. | | `directory` | path | `/app/model-provisioning.d` | Directory containing one `.yaml` or `.yml` document per canonical model. | | `org_overlays_directory` | path | `/app/org-overlays.d` | Directory containing one `.yaml` or `.yml` document per organisation with per-model serving overrides (explicit `targets`, `self_hosted_only`). Applied after the model catalog when `enabled` is set. A missing directory is an empty catalog, so a deployment that does not mount one applies no overlays; an existing path that is not a directory fails startup. | -| `account_limits_directory` | path | `/app/account-limits.d` | Directory containing one `.yaml` or `.yml` document per account with its own realtime in-flight limits on virtual models. Applied after the model catalog when `enabled` is set, replacing every stored per-account limit. A missing directory changes nothing; an empty one clears every per-account limit. See [Account limits](./model-provisioning.md#account-limits). | +| `account_limits_directory` | path | `/app/account-limits.d` | Directory containing one `.yaml` or `.yml` document per account with its own realtime in-flight limits, prices, scheduling priority and latency targets on virtual models. Applied after the model catalog and organisation overlays when `enabled` is set; limits replace every stored per-account limit, and the other settings are owned row by row. A missing directory changes nothing; an empty one clears what the files own. See [Account limits](./model-provisioning.md#account-limits). | When enabled, the directory must exist. Startup fails before any provisioning writes if loading, validation, or a referenced endpoint/group lookup fails. diff --git a/docs/src/reference/model-provisioning.md b/docs/src/reference/model-provisioning.md index 847ff34c7..6840dbfc1 100644 --- a/docs/src/reference/model-provisioning.md +++ b/docs/src/reference/model-provisioning.md @@ -371,36 +371,91 @@ leaves partial model edits. Metadata-only edits are still permitted. ## Account limits -An account limit file gives one account its own realtime in-flight limit on -some virtual models, in place of each model's -`settings.realtime_inflight_limit`. The account is a user or an organisation, -named by its username, and an organisation's limit covers every key the +An account file holds one account's own settings on some virtual models: its +realtime in-flight limits, its prices, and the scheduling priority and latency +targets sent with its requests. The account is a user or an organisation, +named by its username, and an organisation's settings cover every key the organisation owns. One file per account, mounted at `model_provisioning.account_limits_directory`: +```yaml +account: acme +priority: 100 # optional; default for every model +models: + example/chat-model: + priority: 300 # optional; overrides the account default + tariffs: # batch prices only, at model level + - {name: batch-24h, purpose: batch, completion_window: 24h, + input_per_million_tokens: "0.50", output_per_million_tokens: "1.50"} + classes: + standard: + realtime_inflight: 200 + tariffs: # realtime and playground prices, per class + - {name: realtime, purpose: realtime, + input_per_million_tokens: "0.70", output_per_million_tokens: "1.40"} + cache_tariff: {read_multiplier: "0.2", write_multiplier_5m: "1", + write_multiplier_1h: "1", write_multiplier_24h: "1"} + ttft_ms: 1000 # optional latency targets + itl_ms: 30 + priority: 500 # optional; overrides the model's priority +``` + +The older shape is still accepted and means the same as +`classes.standard.realtime_inflight`: + ```yaml account: acme realtime_inflight: example/chat-model: 200 - example/fast-model: 40 ``` -Startup applies the files after the model catalog, in one transaction, and -replaces every stored per-account limit with what the files declare. Removing -a line or a file returns that account to the model's default on the next -start. An empty directory clears every per-account limit; a missing directory -changes nothing. An unknown account, or a model that is not a live virtual -model, fails startup before any per-account limit is written; the model catalog -and organisation overlays have already been applied by then. +Rules: + +- `classes` may name only `standard` for now. A class the file omits keeps the + model's defaults. +- Limits, realtime and playground prices, and the cache tariff are declared per + class; model-level `tariffs` take batch prices only. +- Priority inherits class, then model, then account, then the default. It is + sent as the scheduling priority on the account's realtime requests to + `dynamo` serving members; batch and flex keep their deadline priority. + `ttft_ms` and `itl_ms` are sent as latency targets on the same realtime + requests. Neither changes the class the request is billed as. Priority + `-2147483648` is reserved for background work and rejected. +- If a model appears in both shapes, the two limits must be equal. +- An account/model pair may be configured in its account file or in an + [organisation overlay](#organisation-catalog-validation-and-ownership), not + both. The older limits shape may sit alongside an overlay. + +Startup applies the files after the model catalog and organisation overlays, +in one transaction under one advisory lock: + +- **Limits** replace every stored per-account limit with what the files + declare. Removing a line or a file returns that account to the model's + default on the next start. +- **Priorities, latency targets and prices** are owned row by row. Declaring + an account/model pair makes the file authoritative for its current prices + and policies, adopting matching rows written by hand. Omitted prices retire, and changed + prices are versioned without altering historical amounts. Rows that another + writer owns, such as the organisation catalog's prices, are refused, never + taken over. + +An empty directory clears everything the account files own. A missing +directory changes nothing. An unknown account, or a model that is not a live +virtual model, fails startup before any account-file setting is written; the +model catalog and organisation overlays have already been applied by then. Run the same offline validation used by deployment CI: ```sh -dwctl-model-provisioning validate-account-limits ./account-limits.d --models ./model-provisioning.d +dwctl-model-provisioning validate-account-limits ./account-limits.d \ + --models ./model-provisioning.d --org-overlays ./org-overlays.d +dwctl-model-provisioning account-schema ``` -It rejects unknown fields, limits below 1, an account declared in two files and -models absent from the model catalog. Whether each account exists can only be +It rejects unknown fields and classes, limits and latency targets below 1, an +account declared in two files, entries that configure nothing, prices at the +wrong level, models absent from the model catalog and, with `--org-overlays`, +pairs configured in both catalogs. Whether each account exists can only be checked against the database, at startup. `GET /admin/api/v1/models/{id}/realtime-inflight-limits` returns a virtual diff --git a/dwctl/migrations/20261021150000_account_policies.sql b/dwctl/migrations/20261021150000_account_policies.sql new file mode 100644 index 000000000..bd8a8f3d3 --- /dev/null +++ b/dwctl/migrations/20261021150000_account_policies.sql @@ -0,0 +1,39 @@ +-- Per-account serving policy from the account catalog files: an account-wide +-- scheduling priority, and per-model or per-class priority and latency targets. +-- Written only by the account catalog reconciler. +SET LOCAL lock_timeout = '5s'; + +ALTER TABLE users ADD COLUMN scheduling_priority INTEGER; + +COMMENT ON COLUMN users.scheduling_priority IS + 'Account setting from the account catalog: the scheduling priority sent with this account''s realtime requests to self-hosted serving members, unless a model or class policy overrides it. NULL sends the default.'; + +CREATE TRIGGER users_scheduling_priority_notify + AFTER UPDATE OF scheduling_priority ON users + FOR EACH ROW + WHEN (OLD.scheduling_priority IS DISTINCT FROM NEW.scheduling_priority) + EXECUTE FUNCTION notify_config_change(); + +CREATE TABLE account_model_policies ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + deployed_model_id UUID NOT NULL REFERENCES deployed_models(id) ON DELETE CASCADE, + -- NULL applies to every class of the model; a class key narrows it. + serving_class TEXT, + priority INTEGER, + ttft_ms INTEGER CONSTRAINT chk_account_model_policies_ttft CHECK (ttft_ms > 0), + itl_ms INTEGER CONSTRAINT chk_account_model_policies_itl CHECK (itl_ms > 0), + provisioning_source TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT uq_account_model_policies_scope UNIQUE NULLS NOT DISTINCT (user_id, deployed_model_id, serving_class) +); + +COMMENT ON TABLE account_model_policies IS + 'Per-account serving policy on one model (serving_class NULL) or one of its classes, from the account catalog files.'; + +CREATE INDEX idx_account_model_policies_model ON account_model_policies (deployed_model_id); + +CREATE TRIGGER account_model_policies_notify + AFTER INSERT OR UPDATE OR DELETE ON account_model_policies + EXECUTE FUNCTION notify_config_change(); diff --git a/dwctl/src/account_limits.rs b/dwctl/src/account_limits.rs index 3bc67122a..5bd3a5c24 100644 --- a/dwctl/src/account_limits.rs +++ b/dwctl/src/account_limits.rs @@ -1,3 +1,17 @@ +//! The account catalog: one YAML file per account (an organisation or a +//! personal account) holding that account's realtime in-flight limits, its +//! scheduling priority and latency targets, and its own prices, per model and +//! per serving class. +//! +//! The older shape, a top-level `realtime_inflight: {alias: limit}`, is still +//! accepted and means the same as `models..classes.standard.realtime_inflight`. +//! +//! Ownership: limits are owned table-wide (every undeclared override is +//! removed, as before). Policies and prices are owned per row through their +//! provisioning marker, so rows written by hand or by the organisation catalog +//! are never touched. A missing directory is unmanaged; an empty one clears +//! what this catalog owns. + use std::{ collections::{BTreeMap, HashMap, HashSet}, fs, @@ -5,14 +19,27 @@ use std::{ }; use anyhow::{Context, Result, bail, ensure}; +use chrono::{DateTime, Utc}; +use schemars::JsonSchema; use serde::Deserialize; use sqlx::{PgConnection, PgPool}; use uuid::Uuid; -use crate::model_provisioning::Catalog; +use crate::db::handlers::ModelProvisioning; +use crate::db::handlers::model_provisioning::{ACCOUNT_CATALOG_TARIFF_SOURCE, retire_omitted_account_prices}; +use crate::model_provisioning::{CachePrices, Catalog, Tariff, TariffPurpose, validate_cache_tariff, validate_tariffs}; +use crate::org_overlays::OrgCatalog; const ACCOUNT_LIMITS_LOCK: i64 = 0x4457_4143_4354_4c4d; +/// Ownership marker on policy rows written by this catalog. +const POLICY_SOURCE: &str = "account-catalog"; + +/// Classes an account file may configure. Only `standard` exists until +/// models gain further serving classes. +const CONFIGURABLE_CLASSES: [&str; 1] = ["standard"]; +const STANDARD_CLASS: &str = "standard"; + #[derive(Debug, Clone)] pub struct AccountLimitsCatalog { files: Vec, @@ -22,14 +49,110 @@ pub struct AccountLimitsCatalog { #[derive(Debug, Clone)] struct AccountLimitsFile { source: String, - document: AccountLimitsDocument, + document: AccountDocument, +} + +/// One account's file. +#[derive(Debug, Clone, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct AccountDocument { + /// The account, by its username: an organisation or a personal account. + pub account: String, + /// Scheduling priority for the account's realtime requests on every model, + /// unless a model or class sets its own. Omit to send the default. + #[serde(default)] + pub priority: Option, + /// Older shape: realtime in-flight limits by model alias. Equivalent to + /// `models..classes.standard.realtime_inflight`. + #[serde(default)] + pub realtime_inflight: BTreeMap, + /// Per-model policy and prices, keyed by model alias. + #[serde(default)] + pub models: BTreeMap, } -#[derive(Debug, Clone, Deserialize)] +/// The account's policy and prices on one model. +#[derive(Debug, Clone, Default, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] -struct AccountLimitsDocument { - account: String, - realtime_inflight: BTreeMap, +pub struct AccountModelPolicy { + /// Scheduling priority on this model, overriding the account's. + #[serde(default)] + pub priority: Option, + /// The account's batch and flex prices on this model, per completion + /// window. Realtime prices are set per class. + #[serde(default)] + pub tariffs: Vec, + /// Per-class policy, keyed by class. An omitted class keeps the model's + /// defaults for that class. + #[serde(default)] + pub classes: BTreeMap, +} + +/// The account's policy and prices on one class of a model. +#[derive(Debug, Clone, Default, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct AccountClassPolicy { + /// Concurrent realtime requests the account may have in flight. + #[serde(default)] + pub realtime_inflight: Option, + /// Realtime (and playground) prices. + #[serde(default)] + pub tariffs: Vec, + /// Prompt-cache multipliers this account pays on the class. + #[serde(default)] + pub cache_tariff: Option, + /// Time-to-first-token target, milliseconds. + #[serde(default)] + pub ttft_ms: Option, + /// Inter-token-latency target, milliseconds. + #[serde(default)] + pub itl_ms: Option, + /// Scheduling priority on this class, overriding the model's and the account's. + #[serde(default)] + pub priority: Option, +} + +impl AccountModelPolicy { + fn is_empty(&self) -> bool { + self.priority.is_none() && self.tariffs.is_empty() && self.classes.is_empty() + } +} + +impl AccountClassPolicy { + fn is_empty(&self) -> bool { + self.realtime_inflight.is_none() + && self.tariffs.is_empty() + && self.cache_tariff.is_none() + && self.ttft_ms.is_none() + && self.itl_ms.is_none() + && self.priority.is_none() + } + + fn has_policy(&self) -> bool { + self.priority.is_some() || self.ttft_ms.is_some() || self.itl_ms.is_some() + } +} + +impl AccountDocument { + /// Every model alias the file names. + fn aliases(&self) -> impl Iterator { + self.realtime_inflight.keys().chain(self.models.keys()) + } + + /// The effective limits, merging both shapes (equal duplicates are one limit). + fn limits(&self) -> BTreeMap<&str, i32> { + let mut limits: BTreeMap<&str, i32> = self + .realtime_inflight + .iter() + .map(|(alias, limit)| (alias.as_str(), *limit)) + .collect(); + for (alias, model) in &self.models { + if let Some(limit) = model.classes.get(STANDARD_CLASS).and_then(|class| class.realtime_inflight) { + limits.insert(alias.as_str(), limit); + } + } + limits + } } impl AccountLimitsCatalog { @@ -63,10 +186,15 @@ impl AccountLimitsCatalog { Ok(catalog) } + /// The JSON Schema of one account document. + pub fn json_schema() -> Result { + Ok(serde_json::to_string_pretty(&schemars::schema_for!(AccountDocument))?) + } + pub fn validate_models(&self, models: &Catalog) -> Result<()> { let aliases: HashSet<_> = models.models.iter().map(|model| model.clay.alias.as_str()).collect(); for file in &self.files { - for alias in file.document.realtime_inflight.keys() { + for alias in file.document.aliases() { ensure!( aliases.contains(alias.as_str()), "{}: model {alias:?} is absent from the model catalog", @@ -77,21 +205,131 @@ impl AccountLimitsCatalog { Ok(()) } + /// An account/model pair is configured in one catalog only: either an + /// account file's `models` entry or an organisation overlay, never both. + /// A limit in the older top-level shape may sit alongside an overlay. + pub fn validate_against_overlays(&self, overlays: &OrgCatalog) -> Result<()> { + let mut declared: HashMap<(String, String), &str> = HashMap::new(); + for entry in &overlays.orgs { + for model in &entry.document.models { + declared.insert( + (entry.document.org.trim().to_lowercase(), model.alias.to_lowercase()), + entry.source.as_str(), + ); + } + } + for file in &self.files { + for alias in file.document.models.keys() { + let key = (file.document.account.trim().to_lowercase(), alias.to_lowercase()); + if let Some(overlay) = declared.get(&key) { + bail!( + "{}: model {alias:?} for account {:?} is also configured by organisation overlay {overlay}; configure it in one place", + file.source, + file.document.account + ); + } + } + } + Ok(()) + } + fn validate(&self) -> Result<()> { let mut seen: HashMap<&str, &str> = HashMap::new(); for file in &self.files { - let account = file.document.account.as_str(); - if let Some(previous) = seen.insert(account, &file.source) { - bail!("{}: account {account:?} is also declared in {previous}", file.source); + let source = file.source.as_str(); + let document = &file.document; + let account = document.account.as_str(); + ensure!(!account.trim().is_empty(), "{source}: account cannot be empty"); + if let Some(previous) = seen.insert(account, source) { + bail!("{source}: account {account:?} is also declared in {previous}"); } - for (alias, limit) in &file.document.realtime_inflight { - ensure!(*limit > 0, "{}: realtime_inflight limit on {alias:?} must be positive", file.source); + // A file that declares nothing would silently clear the account. + ensure!( + document.priority.is_some() || !document.realtime_inflight.is_empty() || !document.models.is_empty(), + "{source}: account {account:?} declares nothing; set `priority`, `models` or `realtime_inflight`" + ); + validate_priority(document.priority, &format!("{source}: account {account:?}"))?; + for (alias, limit) in &document.realtime_inflight { + ensure!(*limit > 0, "{source}: realtime_inflight limit on {alias:?} must be positive"); + } + for (alias, model) in &document.models { + let context = format!("{source}: model {alias:?}"); + ensure!(!model.is_empty(), "{context} configures nothing"); + validate_priority(model.priority, &context)?; + validate_tariffs(&model.tariffs, &context)?; + ensure!( + model.tariffs.iter().all(|t| t.purpose == TariffPurpose::Batch), + "{context}: realtime and playground prices belong under `classes..tariffs`" + ); + for (class, policy) in &model.classes { + let context = format!("{context}, class {class:?}"); + ensure!( + CONFIGURABLE_CLASSES.contains(&class.as_str()), + "{context}: unknown class; configurable classes are {}", + CONFIGURABLE_CLASSES.join(", ") + ); + ensure!(!policy.is_empty(), "{context} configures nothing"); + if let Some(limit) = policy.realtime_inflight { + ensure!(limit > 0, "{context}: realtime_inflight must be positive"); + if let Some(older) = document.realtime_inflight.get(alias) { + ensure!( + *older == limit, + "{context}: realtime_inflight {limit} conflicts with the top-level realtime_inflight {older}" + ); + } + } + validate_tariffs(&policy.tariffs, &context)?; + ensure!( + policy.tariffs.iter().all(|t| t.purpose != TariffPurpose::Batch), + "{context}: batch prices belong in the model's `tariffs`" + ); + if let Some(cache) = &policy.cache_tariff { + validate_cache_tariff(&cache.as_tariff(), &context)?; + } + validate_priority(policy.priority, &context)?; + validate_target(policy.ttft_ms, "ttft_ms", &context)?; + validate_target(policy.itl_ms, "itl_ms", &context)?; + } } } Ok(()) } } +/// `i32::MIN` is the batch daemon's background tier; realtime policy may not +/// tie with it. +fn validate_priority(priority: Option, context: &str) -> Result<()> { + ensure!( + priority != Some(i32::MIN), + "{context}: priority {} is reserved for background work", + i32::MIN + ); + Ok(()) +} + +/// Latency targets are stored as positive `INTEGER`s. +fn validate_target(target: Option, field: &str, context: &str) -> Result<()> { + if let Some(target) = target { + ensure!( + target > 0 && i32::try_from(target).is_ok(), + "{context}: {field} must be positive and at most {}", + i32::MAX + ); + } + Ok(()) +} + +/// One policy row: an account's priority and targets on a model (class +/// `None`) or on one of its classes. +struct PolicyRow { + account: Uuid, + model: Uuid, + class: Option, + priority: Option, + ttft_ms: Option, + itl_ms: Option, +} + pub async fn apply(pool: &PgPool, catalog: &AccountLimitsCatalog) -> Result<()> { if !catalog.managed { return Ok(()); @@ -104,15 +342,36 @@ pub async fn apply(pool: &PgPool, catalog: &AccountLimitsCatalog) -> Result<()> .context("acquire account limits advisory lock")?; let accounts = resolve_accounts(&mut transaction, catalog).await?; let models = resolve_models(&mut transaction, catalog).await?; + // A replica may have begun its transaction before the lock winner. Use the + // wall clock after locking so the winner's committed prices are not "future". + let effective_at: DateTime = sqlx::query_scalar("SELECT clock_timestamp()") + .fetch_one(&mut *transaction) + .await + .context("read account catalog effective timestamp")?; + + apply_limits(&mut transaction, catalog, &accounts, &models).await?; + apply_priorities(&mut transaction, catalog, &accounts).await?; + apply_policies(&mut transaction, catalog, &accounts, &models).await?; + apply_prices(&mut transaction, catalog, &accounts, &models, effective_at).await?; + transaction.commit().await.context("commit account limits")?; + Ok(()) +} + +async fn apply_limits( + db: &mut PgConnection, + catalog: &AccountLimitsCatalog, + accounts: &HashMap, + models: &HashMap, +) -> Result<()> { let mut model_ids = Vec::new(); let mut account_ids = Vec::new(); let mut limits = Vec::new(); for file in &catalog.files { - for (alias, limit) in &file.document.realtime_inflight { + for (alias, limit) in file.document.limits() { model_ids.push(models[alias]); account_ids.push(accounts[&file.document.account]); - limits.push(*limit); + limits.push(limit); } } @@ -125,7 +384,7 @@ pub async fn apply(pool: &PgPool, catalog: &AccountLimitsCatalog) -> Result<()> ) .bind(&model_ids) .bind(&account_ids) - .execute(&mut *transaction) + .execute(&mut *db) .await .context("remove undeclared account limits")?; sqlx::query( @@ -138,14 +397,224 @@ pub async fn apply(pool: &PgPool, catalog: &AccountLimitsCatalog) -> Result<()> .bind(&model_ids) .bind(&account_ids) .bind(&limits) - .execute(&mut *transaction) + .execute(&mut *db) .await .context("write account limits")?; + Ok(()) +} - transaction.commit().await.context("commit account limits")?; +/// The account-wide priority column is written only by this catalog: declared +/// accounts get their value (or none), every other account is cleared. +async fn apply_priorities(db: &mut PgConnection, catalog: &AccountLimitsCatalog, accounts: &HashMap) -> Result<()> { + let (ids, priorities): (Vec, Vec>) = catalog + .files + .iter() + .map(|file| (accounts[&file.document.account], file.document.priority)) + .unzip(); + sqlx::query( + "UPDATE users u SET scheduling_priority = d.priority + FROM UNNEST($1::uuid[], $2::int[]) AS d(id, priority) + WHERE u.id = d.id AND u.scheduling_priority IS DISTINCT FROM d.priority", + ) + .bind(&ids) + .bind(&priorities) + .execute(&mut *db) + .await + .context("write account priorities")?; + sqlx::query("UPDATE users SET scheduling_priority = NULL WHERE scheduling_priority IS NOT NULL AND NOT (id = ANY($1))") + .bind(&ids) + .execute(&mut *db) + .await + .context("clear undeclared account priorities")?; Ok(()) } +async fn apply_policies( + db: &mut PgConnection, + catalog: &AccountLimitsCatalog, + accounts: &HashMap, + models: &HashMap, +) -> Result<()> { + let mut rows = Vec::new(); + for file in &catalog.files { + let account = accounts[&file.document.account]; + for (alias, model) in &file.document.models { + let model_id = models[alias]; + if model.priority.is_some() { + rows.push(PolicyRow { + account, + model: model_id, + class: None, + priority: model.priority, + ttft_ms: None, + itl_ms: None, + }); + } + for (class, policy) in &model.classes { + if policy.has_policy() { + rows.push(PolicyRow { + account, + model: model_id, + class: Some(class.clone()), + priority: policy.priority, + ttft_ms: policy.ttft_ms.map(i32::try_from).transpose()?, + itl_ms: policy.itl_ms.map(i32::try_from).transpose()?, + }); + } + } + } + } + let accounts_col: Vec = rows.iter().map(|r| r.account).collect(); + let models_col: Vec = rows.iter().map(|r| r.model).collect(); + let classes_col: Vec> = rows.iter().map(|r| r.class.clone()).collect(); + let priorities_col: Vec> = rows.iter().map(|r| r.priority).collect(); + let ttft_col: Vec> = rows.iter().map(|r| r.ttft_ms).collect(); + let itl_col: Vec> = rows.iter().map(|r| r.itl_ms).collect(); + + sqlx::query( + "DELETE FROM account_model_policies p + WHERE p.provisioning_source = $4 + AND NOT EXISTS ( + SELECT 1 FROM UNNEST($1::uuid[], $2::uuid[], $3::text[]) AS d(user_id, deployed_model_id, serving_class) + WHERE d.user_id = p.user_id AND d.deployed_model_id = p.deployed_model_id + AND d.serving_class IS NOT DISTINCT FROM p.serving_class + )", + ) + .bind(&accounts_col) + .bind(&models_col) + .bind(&classes_col) + .bind(POLICY_SOURCE) + .execute(&mut *db) + .await + .context("remove undeclared account policies")?; + // Unowned rows on a declared scope are adopted, as manual prices are; a + // row another writer owns is never taken over. + let foreign: Vec<(String, String, Option, String)> = sqlx::query_as( + "SELECT u.username, dm.alias, p.serving_class, p.provisioning_source + FROM account_model_policies p + JOIN UNNEST($1::uuid[], $2::uuid[], $3::text[]) AS d(user_id, deployed_model_id, serving_class) + ON d.user_id = p.user_id AND d.deployed_model_id = p.deployed_model_id + AND d.serving_class IS NOT DISTINCT FROM p.serving_class + JOIN users u ON u.id = p.user_id + JOIN deployed_models dm ON dm.id = p.deployed_model_id + WHERE p.provisioning_source IS NOT NULL AND p.provisioning_source <> $4 + ORDER BY u.username, dm.alias, p.serving_class NULLS FIRST", + ) + .bind(&accounts_col) + .bind(&models_col) + .bind(&classes_col) + .bind(POLICY_SOURCE) + .fetch_all(&mut *db) + .await + .context("check account policy ownership")?; + if let Some((account, alias, class, owner)) = foreign.first() { + bail!( + "account {account:?} policy on {alias:?} (class {}) is owned by {owner:?}; remove it there first", + class.as_deref().unwrap_or("all") + ); + } + sqlx::query( + "INSERT INTO account_model_policies (user_id, deployed_model_id, serving_class, priority, ttft_ms, itl_ms, provisioning_source) + SELECT d.user_id, d.deployed_model_id, d.serving_class, d.priority, d.ttft_ms, d.itl_ms, $7 + FROM UNNEST($1::uuid[], $2::uuid[], $3::text[], $4::int[], $5::int[], $6::int[]) + AS d(user_id, deployed_model_id, serving_class, priority, ttft_ms, itl_ms) + ON CONFLICT ON CONSTRAINT uq_account_model_policies_scope DO UPDATE SET + priority = EXCLUDED.priority, + ttft_ms = EXCLUDED.ttft_ms, + itl_ms = EXCLUDED.itl_ms, + provisioning_source = EXCLUDED.provisioning_source, + updated_at = NOW() + WHERE (account_model_policies.priority, account_model_policies.ttft_ms, account_model_policies.itl_ms, + account_model_policies.provisioning_source) + IS DISTINCT FROM (EXCLUDED.priority, EXCLUDED.ttft_ms, EXCLUDED.itl_ms, EXCLUDED.provisioning_source)", + ) + .bind(&accounts_col) + .bind(&models_col) + .bind(&classes_col) + .bind(&priorities_col) + .bind(&ttft_col) + .bind(&itl_col) + .bind(POLICY_SOURCE) + .execute(&mut *db) + .await + .context("write account policies")?; + Ok(()) +} + +/// Prices for every declared account/model: a temporal ledger scoped to the +/// account, versioned like the model catalog's general prices. A declared +/// pair is wholly catalog-owned: its manual prices are adopted, and prices it +/// no longer declares are closed without deleting history. +async fn apply_prices( + db: &mut PgConnection, + catalog: &AccountLimitsCatalog, + accounts: &HashMap, + models: &HashMap, + effective_at: DateTime, +) -> Result<()> { + let mut declared = Vec::new(); + for file in &catalog.files { + let account = accounts[&file.document.account]; + let context = |alias: &str| format!("{}: account {:?}, model {alias:?}", file.source, file.document.account); + if !file.document.models.is_empty() { + ModelProvisioning::for_account_catalog(db) + .preflight_future_account_tariffs(account, effective_at) + .await + .with_context(|| format!("{}: account {:?}", file.source, file.document.account))?; + } + for (alias, model) in &file.document.models { + let model_id = models[alias]; + declared.push((account, model_id)); + let mut provisioning = ModelProvisioning::for_account_catalog(db); + provisioning + .adopt_account_model_tariffs(account, model_id, effective_at) + .await + .with_context(|| context(alias))?; + provisioning + .reconcile_tariffs(model_id, Some(account), None, &model.tariffs, effective_at) + .await + .with_context(|| context(alias))?; + provisioning + .reconcile_cache_tariff(model_id, Some(account), None, None, effective_at) + .await + .with_context(|| context(alias))?; + let standard = model.classes.get(STANDARD_CLASS); + provisioning + .reconcile_tariffs( + model_id, + Some(account), + Some(STANDARD_CLASS), + standard.map_or(&[], |class| class.tariffs.as_slice()), + effective_at, + ) + .await + .with_context(|| context(alias))?; + let cache = standard.and_then(|class| class.cache_tariff.as_ref()).map(CachePrices::as_tariff); + provisioning + .reconcile_cache_tariff(model_id, Some(account), Some(STANDARD_CLASS), cache.as_ref(), effective_at) + .await + .with_context(|| context(alias))?; + // Rows adopted under any other class scope are not declared: close them. + for table in ["model_tariffs", "model_cache_tariffs"] { + sqlx::query(&format!( + "UPDATE {table} SET valid_until = $4 + WHERE user_id = $1 AND deployed_model_id = $2 AND provisioning_source = $3 + AND serving_class IS NOT NULL AND serving_class <> 'standard' + AND valid_from <= $4 AND (valid_until IS NULL OR valid_until > $4)" + )) + .bind(account) + .bind(model_id) + .bind(ACCOUNT_CATALOG_TARIFF_SOURCE) + .bind(effective_at) + .execute(&mut *db) + .await + .with_context(|| context(alias))?; + } + } + } + retire_omitted_account_prices(db, ACCOUNT_CATALOG_TARIFF_SOURCE, &declared, effective_at).await +} + async fn resolve_accounts(db: &mut PgConnection, catalog: &AccountLimitsCatalog) -> Result> { let names: Vec = catalog.files.iter().map(|file| file.document.account.clone()).collect(); let resolved: HashMap = @@ -169,7 +638,7 @@ async fn resolve_models(db: &mut PgConnection, catalog: &AccountLimitsCatalog) - let aliases: Vec = catalog .files .iter() - .flat_map(|file| file.document.realtime_inflight.keys().cloned()) + .flat_map(|file| file.document.aliases().cloned()) .collect::>() .into_iter() .collect(); @@ -276,9 +745,9 @@ mod tests { "must be positive", ), (vec![("a.yaml", "account: acme\nrealtime:\n org/model: 10\n")], "unknown field"), - // A file with no `realtime_inflight` at all is refused rather than - // read as an empty (and therefore clearing) declaration. - (vec![("a.yaml", "account: acme\n")], "missing field `realtime_inflight`"), + // A file that declares nothing is refused rather than read as an + // empty (and therefore clearing) declaration. + (vec![("a.yaml", "account: acme\n")], "declares nothing"), ( vec![("a.yaml", "account: acme\nrealtime_inflight:\n org/model: 10\n org/model: 20\n")], "duplicate entry", @@ -364,4 +833,329 @@ mod tests { assert!(stored(&pool).await.is_empty()); } } + + const REALTIME: &str = "{name: deal, purpose: realtime, input_per_million_tokens: '1', output_per_million_tokens: '2'}"; + const BATCH: &str = + "{name: deal-24h, purpose: batch, completion_window: 24h, input_per_million_tokens: '0.5', output_per_million_tokens: '1'}"; + + #[test] + fn the_new_shape_is_validated() { + for (contents, diagnostic) in [ + ( + format!("account: acme\nmodels:\n org/model:\n tariffs: [{REALTIME}]\n"), + "realtime and playground prices belong under", + ), + ( + format!("account: acme\nmodels:\n org/model:\n classes:\n standard:\n tariffs: [{BATCH}]\n"), + "batch prices belong in the model's", + ), + ( + "account: acme\nmodels:\n org/model:\n classes:\n fast:\n realtime_inflight: 5\n".to_owned(), + "unknown class", + ), + ( + "account: acme\nrealtime_inflight:\n org/model: 10\nmodels:\n org/model:\n classes:\n standard:\n realtime_inflight: 20\n".to_owned(), + "conflicts with the top-level", + ), + ("account: acme\nmodels:\n org/model: {}\n".to_owned(), "configures nothing"), + ( + "account: acme\nmodels:\n org/model:\n classes:\n standard: {}\n".to_owned(), + "configures nothing", + ), + ( + "account: acme\nmodels:\n org/model:\n classes:\n standard:\n ttft_ms: 0\n".to_owned(), + "ttft_ms must be positive", + ), + ( + "account: acme\nmodels:\n org/model:\n classes:\n standard:\n itl_ms: 2147483648\n".to_owned(), + "itl_ms must be positive and at most", + ), + ("account: acme\npriority: -2147483648\n".to_owned(), "reserved for background work"), + ( + "account: acme\nmodels:\n org/model:\n priority: -2147483648\n".to_owned(), + "reserved for background work", + ), + ( + "account: acme\nmodels:\n org/model:\n classes:\n standard:\n realtime_inflight: 0\n".to_owned(), + "must be positive", + ), + ( + "account: acme\nmodels:\n org/model:\n classes:\n standard:\n unknown: 1\n".to_owned(), + "unknown field", + ), + ] { + let directory = tempdir().unwrap(); + write(directory.path(), "a.yaml", &contents); + let error = format!("{:#}", AccountLimitsCatalog::load(directory.path()).unwrap_err()); + assert!(error.contains(diagnostic), "{contents}\n{error}"); + } + + // Equal limits in both shapes are one limit; priority alone is a declaration. + let directory = tempdir().unwrap(); + write( + directory.path(), + "a.yaml", + "account: acme\nrealtime_inflight:\n org/model: 10\nmodels:\n org/model:\n classes:\n standard:\n realtime_inflight: 10\n", + ); + write(directory.path(), "b.yaml", "account: bob\npriority: 100\n"); + let catalog = AccountLimitsCatalog::load(directory.path()).unwrap(); + assert_eq!( + catalog.files[0].document.limits().into_iter().collect::>(), + vec![("org/model", 10)] + ); + } + + #[test] + fn an_account_model_pair_is_configured_in_one_catalog() { + let overlays = tempdir().unwrap(); + fs::write( + overlays.path().join("acme.yaml"), + "org: acme\nmodels:\n - alias: org/model\n self_hosted_only: true\n", + ) + .unwrap(); + let overlays = OrgCatalog::load(overlays.path()).unwrap(); + + let directory = tempdir().unwrap(); + // The older limits shape may sit alongside an overlay... + write( + directory.path(), + "acme.yaml", + "account: acme\nrealtime_inflight:\n org/model: 10\n", + ); + AccountLimitsCatalog::load(directory.path()) + .unwrap() + .validate_against_overlays(&overlays) + .unwrap(); + // ...a `models` entry for the same pair may not. + write( + directory.path(), + "acme.yaml", + "account: acme\nmodels:\n org/model:\n priority: 5\n", + ); + let error = AccountLimitsCatalog::load(directory.path()) + .unwrap() + .validate_against_overlays(&overlays) + .unwrap_err() + .to_string(); + assert!(error.contains("also configured by organisation overlay"), "{error}"); + } + + async fn policies(pool: &PgPool) -> Vec<(Uuid, Option, Option, Option, Option)> { + sqlx::query_as( + "SELECT user_id, serving_class, priority, ttft_ms, itl_ms FROM account_model_policies ORDER BY user_id, serving_class NULLS FIRST", + ) + .fetch_all(pool) + .await + .unwrap() + } + + async fn active_prices(pool: &PgPool, account: Uuid) -> Vec<(Option, String, Option, String)> { + sqlx::query_as( + "SELECT serving_class, api_key_purpose, completion_window, provisioning_source FROM model_tariffs + WHERE user_id = $1 AND valid_until IS NULL ORDER BY serving_class NULLS FIRST, api_key_purpose", + ) + .bind(account) + .fetch_all(pool) + .await + .unwrap() + } + + #[dwctl_test_macros::test] + async fn the_account_catalog_writes_priorities_policies_limits_and_prices(pool: PgPool) { + let org = insert_account(&pool, "acme", "organization").await; + let person = insert_account(&pool, "bob", "individual").await; + let model = insert_virtual_model(&pool, "org/model").await; + + let directory = tempdir().unwrap(); + write( + directory.path(), + "acme.yaml", + &format!( + "account: acme\npriority: 100\nmodels:\n org/model:\n priority: 300\n tariffs: [{BATCH}]\n classes:\n standard:\n realtime_inflight: 900\n tariffs: [{REALTIME}]\n cache_tariff: {{read_multiplier: '0.2', write_multiplier_5m: '1', write_multiplier_1h: '1', write_multiplier_24h: '1'}}\n ttft_ms: 1000\n itl_ms: 30\n" + ), + ); + write( + directory.path(), + "bob.yaml", + &format!("account: bob\nmodels:\n org/model:\n classes:\n standard:\n tariffs: [{REALTIME}]\n"), + ); + let catalog = AccountLimitsCatalog::load(directory.path()).unwrap(); + apply(&pool, &catalog).await.unwrap(); + // A second identical start writes nothing new. + apply(&pool, &catalog).await.unwrap(); + + let priority: Option = sqlx::query_scalar("SELECT scheduling_priority FROM users WHERE id = $1") + .bind(org) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(priority, Some(100)); + assert_eq!( + policies(&pool).await, + vec![ + (org, None, Some(300), None, None), + (org, Some("standard".into()), None, Some(1000), Some(30)) + ] + ); + assert_eq!(stored(&pool).await, vec![(model, org, 900)]); + assert_eq!( + active_prices(&pool, org).await, + vec![ + (None, "batch".into(), Some("24h".into()), ACCOUNT_CATALOG_TARIFF_SOURCE.into()), + ( + Some("standard".into()), + "realtime".into(), + None, + ACCOUNT_CATALOG_TARIFF_SOURCE.into() + ), + ] + ); + let cache: i64 = sqlx::query_scalar( + "SELECT count(*) FROM model_cache_tariffs WHERE user_id = $1 AND serving_class = 'standard' AND valid_until IS NULL", + ) + .bind(org) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(cache, 1); + assert_eq!( + active_prices(&pool, person).await, + vec![( + Some("standard".into()), + "realtime".into(), + None, + ACCOUNT_CATALOG_TARIFF_SOURCE.into() + )], + "personal accounts carry deals too" + ); + + // The gateway sees the account default and the model/class policy. + let targets = crate::sync::onwards_config::load_targets_from_db(&pool, &[], false).await.unwrap(); + assert_eq!(targets.accounts.get(&org.to_string()).unwrap().priority, Some(100)); + let overlays = targets.targets.get("org/model").unwrap().default_pool().overlays().clone(); + let overlay = &overlays[&org.to_string()]; + assert_eq!( + (overlay.priority, overlay.ttft_ms, overlay.itl_ms), + (Some(300), Some(1000), Some(30)) + ); + assert!(overlay.targets.is_none(), "policy never becomes billed custom targets"); + assert_eq!(gateway_limit(&pool, "org/model", org).await, 900); + + // Removing the files clears what this catalog owns, keeping history. + fs::remove_file(directory.path().join("acme.yaml")).unwrap(); + fs::remove_file(directory.path().join("bob.yaml")).unwrap(); + apply(&pool, &AccountLimitsCatalog::load(directory.path()).unwrap()).await.unwrap(); + let priority: Option = sqlx::query_scalar("SELECT scheduling_priority FROM users WHERE id = $1") + .bind(org) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(priority, None); + assert!(policies(&pool).await.is_empty()); + assert!(stored(&pool).await.is_empty()); + assert!(active_prices(&pool, org).await.is_empty()); + assert!(active_prices(&pool, person).await.is_empty()); + let history: i64 = sqlx::query_scalar("SELECT count(*) FROM model_tariffs WHERE user_id = $1") + .bind(org) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(history, 2, "closed price versions stay for billing history"); + } + + #[dwctl_test_macros::test] + async fn manual_prices_are_adopted_and_other_catalogs_prices_are_refused(pool: PgPool) { + let org = insert_account(&pool, "acme", "organization").await; + let model = insert_virtual_model(&pool, "org/model").await; + // A hand-written deal on the pair is adopted and replaced by the file. + sqlx::query( + "INSERT INTO model_tariffs (deployed_model_id, user_id, serving_class, name, api_key_purpose, input_price_per_token, output_price_per_token) + VALUES ($1, $2, 'standard', 'hand', 'realtime', 0.000009, 0.000009)", + ) + .bind(model) + .bind(org) + .execute(&pool) + .await + .unwrap(); + let directory = tempdir().unwrap(); + write( + directory.path(), + "acme.yaml", + &format!("account: acme\nmodels:\n org/model:\n classes:\n standard:\n tariffs: [{REALTIME}]\n"), + ); + apply(&pool, &AccountLimitsCatalog::load(directory.path()).unwrap()).await.unwrap(); + let active: Vec<(String, String)> = + sqlx::query_as("SELECT name, provisioning_source FROM model_tariffs WHERE user_id = $1 AND valid_until IS NULL") + .bind(org) + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(active, vec![("deal".into(), ACCOUNT_CATALOG_TARIFF_SOURCE.into())]); + + // A price the organisation catalog owns is never taken over. + let other = insert_account(&pool, "other", "organization").await; + sqlx::query( + "INSERT INTO model_tariffs (deployed_model_id, user_id, serving_class, name, api_key_purpose, input_price_per_token, output_price_per_token, provisioning_source) + VALUES ($1, $2, 'standard', 'overlay', 'realtime', 0.000003, 0.000003, 'org-overlays')", + ) + .bind(model) + .bind(other) + .execute(&pool) + .await + .unwrap(); + write( + directory.path(), + "other.yaml", + &format!("account: other\nmodels:\n org/model:\n classes:\n standard:\n tariffs: [{REALTIME}]\n"), + ); + let error = format!( + "{:#}", + apply(&pool, &AccountLimitsCatalog::load(directory.path()).unwrap()) + .await + .unwrap_err() + ); + assert!(error.contains("conflicts with a manually managed tariff"), "{error}"); + } + + #[dwctl_test_macros::test] + async fn unowned_policies_are_adopted_and_other_writers_policies_are_refused(pool: PgPool) { + let org = insert_account(&pool, "acme", "organization").await; + let model = insert_virtual_model(&pool, "org/model").await; + sqlx::query("INSERT INTO account_model_policies (user_id, deployed_model_id, priority) VALUES ($1, $2, 7)") + .bind(org) + .bind(model) + .execute(&pool) + .await + .unwrap(); + let directory = tempdir().unwrap(); + write( + directory.path(), + "acme.yaml", + "account: acme\nmodels:\n org/model:\n priority: 9\n", + ); + let catalog = AccountLimitsCatalog::load(directory.path()).unwrap(); + apply(&pool, &catalog).await.unwrap(); + let adopted: (Option, Option) = + sqlx::query_as("SELECT priority, provisioning_source FROM account_model_policies WHERE user_id = $1") + .bind(org) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(adopted, (Some(9), Some(POLICY_SOURCE.to_owned()))); + + sqlx::query("UPDATE account_model_policies SET priority = 7, provisioning_source = 'elsewhere' WHERE user_id = $1") + .bind(org) + .execute(&pool) + .await + .unwrap(); + let error = format!("{:#}", apply(&pool, &catalog).await.unwrap_err()); + assert!(error.contains("is owned by \"elsewhere\""), "{error}"); + let untouched: (Option, Option) = + sqlx::query_as("SELECT priority, provisioning_source FROM account_model_policies WHERE user_id = $1") + .bind(org) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(untouched, (Some(7), Some("elsewhere".to_owned()))); + } } diff --git a/dwctl/src/bin/model_provisioning.rs b/dwctl/src/bin/model_provisioning.rs index 7e6b79d65..12fa754b3 100644 --- a/dwctl/src/bin/model_provisioning.rs +++ b/dwctl/src/bin/model_provisioning.rs @@ -26,9 +26,14 @@ enum Command { directory: PathBuf, #[arg(long)] models: PathBuf, + /// Also check that no account/model pair is configured in both catalogs. + #[arg(long)] + org_overlays: Option, }, /// Print the JSON Schema for one organisation document. OrgSchema, + /// Print the JSON Schema for one account document. + AccountSchema, /// Print the JSON Schema for one model document. Schema, } @@ -46,16 +51,25 @@ fn main() -> Result<()> { overlays.validate_models(&Catalog::load(&models)?)?; println!("org overlay catalog is valid: {}", directory.display()); } - Command::ValidateAccountLimits { directory, models } => { + Command::ValidateAccountLimits { + directory, + models, + org_overlays, + } => { ensure!( directory.is_dir(), "account limits directory does not exist: {}", directory.display() ); - AccountLimitsCatalog::load(&directory)?.validate_models(&Catalog::load(&models)?)?; + let accounts = AccountLimitsCatalog::load(&directory)?; + accounts.validate_models(&Catalog::load(&models)?)?; + if let Some(org_overlays) = org_overlays { + accounts.validate_against_overlays(&OrgCatalog::load(&org_overlays)?)?; + } println!("account limits are valid: {}", directory.display()); } Command::OrgSchema => println!("{}", OrgCatalog::json_schema()?), + Command::AccountSchema => println!("{}", AccountLimitsCatalog::json_schema()?), Command::Schema => println!("{}", Catalog::json_schema()?), } Ok(()) diff --git a/dwctl/src/db/handlers/model_provisioning.rs b/dwctl/src/db/handlers/model_provisioning.rs index ce3ee75d7..1a7d7129d 100644 --- a/dwctl/src/db/handlers/model_provisioning.rs +++ b/dwctl/src/db/handlers/model_provisioning.rs @@ -19,6 +19,54 @@ use crate::model_provisioning::{ /// release it automatically and works through transaction-pooled connections. const MODEL_PROVISIONING_LOCK: i64 = 0x4457_4d4f_4445_4c50; +/// Ownership marker on prices written by the organisation catalog. +pub(crate) const ORG_CATALOG_TARIFF_SOURCE: &str = "org-overlays"; +/// Ownership marker on prices written by the account catalog. +pub(crate) const ACCOUNT_CATALOG_TARIFF_SOURCE: &str = "account-catalog"; + +/// Close the active prices a catalog owns (`source`) on every account/model +/// pair it no longer declares. Manual rows and rows another catalog owns are +/// never touched. Future-dated owned rows on a removed pair are refused: +/// startup provisioning has no scheduling semantics. +pub(crate) async fn retire_omitted_account_prices( + db: &mut PgConnection, + source: &str, + declared: &[(Uuid, Uuid)], + effective_at: DateTime, +) -> Result<()> { + let (account_ids, model_ids): (Vec, Vec) = declared.iter().copied().unzip(); + for table in ["model_tariffs", "model_cache_tariffs"] { + let omitted = "t.provisioning_source = $4 AND NOT EXISTS ( + SELECT 1 FROM UNNEST($1::uuid[], $2::uuid[]) AS d(user_id, deployed_model_id) + WHERE d.user_id = t.user_id AND d.deployed_model_id = t.deployed_model_id)"; + let future: bool = sqlx::query_scalar(&format!( + "SELECT EXISTS (SELECT 1 FROM {table} t WHERE t.valid_from > $3 + AND (t.valid_until IS NULL OR t.valid_until > t.valid_from) AND {omitted})" + )) + .bind(&account_ids) + .bind(&model_ids) + .bind(effective_at) + .bind(source) + .fetch_one(&mut *db) + .await?; + ensure!( + !future, + "removed account/model entry has future-dated prices; resolve these before removing it" + ); + sqlx::query(&format!( + "UPDATE {table} t SET valid_until = $3 WHERE t.valid_from <= $3 + AND (t.valid_until IS NULL OR t.valid_until > $3) AND {omitted}" + )) + .bind(&account_ids) + .bind(&model_ids) + .bind(effective_at) + .bind(source) + .execute(&mut *db) + .await?; + } + Ok(()) +} + pub struct ModelProvisioning<'c> { db: &'c mut PgConnection, tariff_source: Option<&'static str>, @@ -41,10 +89,24 @@ impl<'c> ModelProvisioning<'c> { pub(crate) fn for_org_catalog(db: &'c mut PgConnection) -> Self { Self { db, - tariff_source: Some("org-overlays"), + tariff_source: Some(ORG_CATALOG_TARIFF_SOURCE), } } + /// Account catalog prices carry their own ownership marker, so each catalog + /// retires only the rows it wrote. + pub(crate) fn for_account_catalog(db: &'c mut PgConnection) -> Self { + Self { + db, + tariff_source: Some(ACCOUNT_CATALOG_TARIFF_SOURCE), + } + } + + fn account_tariff_source(&self) -> Result<&'static str> { + self.tariff_source + .context("account price ownership requires an organisation or account catalog writer") + } + pub async fn apply(&mut self, catalog: &Catalog) -> Result<()> { sqlx::query("SELECT pg_advisory_xact_lock($1)") .bind(MODEL_PROVISIONING_LOCK) @@ -168,6 +230,7 @@ impl<'c> ModelProvisioning<'c> { /// YAML is authoritative for the complete declared org/model, including /// prices originally inserted manually. Historical amounts remain immutable. pub(crate) async fn adopt_account_model_tariffs(&mut self, account: Uuid, model: Uuid, effective_at: DateTime) -> Result<()> { + let source = self.account_tariff_source()?; for table in ["model_tariffs", "model_cache_tariffs"] { let future: Option<(DateTime, Option)> = sqlx::query_as(&format!( "SELECT valid_from, serving_class FROM {table} @@ -186,15 +249,18 @@ impl<'c> ModelProvisioning<'c> { class.as_deref().unwrap_or("all") ); } + // Only manual rows are adopted: a row another catalog owns stays + // with it, and reconciliation then reports the conflict. sqlx::query(&format!( - "UPDATE {table} SET provisioning_source='org-overlays' + "UPDATE {table} SET provisioning_source=$4 WHERE user_id=$1 AND deployed_model_id=$2 AND valid_from<=$3 AND (valid_until IS NULL OR valid_until>$3) - AND provisioning_source IS DISTINCT FROM 'org-overlays'" + AND provisioning_source IS NULL" )) .bind(account) .bind(model) .bind(effective_at) + .bind(source) .execute(&mut *self.db) .await?; } @@ -204,16 +270,18 @@ impl<'c> ModelProvisioning<'c> { /// Refuse to reconcile an organisation whose rows carry a future `valid_from`: /// startup provisioning has no scheduling semantics (same rule as the model catalog). pub(crate) async fn preflight_future_account_tariffs(&mut self, account: Uuid, effective_at: DateTime) -> Result<()> { + let source = self.account_tariff_source()?; for table in ["model_tariffs", "model_cache_tariffs"] { let future: Option<(DateTime, String, Option)> = sqlx::query_as(&format!( "SELECT t.valid_from, dm.alias, t.serving_class FROM {table} t JOIN deployed_models dm ON dm.id=t.deployed_model_id - WHERE t.user_id=$1 AND t.provisioning_source='org-overlays' AND t.valid_from>$2 + WHERE t.user_id=$1 AND t.provisioning_source=$3 AND t.valid_from>$2 AND (t.valid_until IS NULL OR t.valid_until>t.valid_from) ORDER BY t.valid_from, t.id LIMIT 1" )) .bind(account) .bind(effective_at) + .bind(source) .fetch_optional(&mut *self.db) .await .with_context(|| format!("check future {table} rows of an organisation"))?; diff --git a/dwctl/src/lib.rs b/dwctl/src/lib.rs index b9d085863..8d1d99a45 100644 --- a/dwctl/src/lib.rs +++ b/dwctl/src/lib.rs @@ -1616,6 +1616,7 @@ async fn setup_database( let catalog = model_provisioning::Catalog::load(&config.model_provisioning.directory)?; let overlays = org_overlays::OrgCatalog::load(&config.model_provisioning.org_overlays_directory)?; let account_limits = account_limits::AccountLimitsCatalog::load(&config.model_provisioning.account_limits_directory)?; + account_limits.validate_against_overlays(&overlays)?; model_provisioning::apply(&main.pooled.write(), &catalog).await?; org_overlays::apply(&main.pooled.write(), &overlays).await?; account_limits::apply(&main.pooled.write(), &account_limits).await?; diff --git a/dwctl/src/org_overlays.rs b/dwctl/src/org_overlays.rs index 40cfbb0fb..7042ad7b5 100644 --- a/dwctl/src/org_overlays.rs +++ b/dwctl/src/org_overlays.rs @@ -28,6 +28,7 @@ use sqlx::{PgConnection, PgPool, Row}; use uuid::Uuid; use crate::db::handlers::ModelProvisioning; +use crate::db::handlers::model_provisioning::{ORG_CATALOG_TARIFF_SOURCE, retire_omitted_account_prices}; use crate::model_provisioning::{ServingPreset, Tariff, validate_cache_tariff, validate_tariffs}; /// Prefix of the `provisioning_source` marker on rows this catalog owns. @@ -367,34 +368,11 @@ async fn apply_in(db: &mut PgConnection, catalog: &OrgCatalog) -> Result<()> { } // Rows this catalog owns but no longer declares are removed; hand rows - // (NULL source) and rows owned by anything else are left alone. - let (org_ids, model_ids): (Vec, Vec) = desired.into_iter().unzip(); - // Price ownership is independent of routing overlays. Omission only retires + // (NULL source) and rows owned by anything else are left alone. Price + // ownership is independent of routing overlays: omission only retires // prices written by this catalog; manual prices survive even on owned models. - for table in ["model_tariffs", "model_cache_tariffs"] { - let omitted = "t.provisioning_source = 'org-overlays' AND NOT EXISTS ( - SELECT 1 FROM UNNEST($1::uuid[], $2::uuid[]) AS d(user_id, deployed_model_id) - WHERE d.user_id = t.user_id AND d.deployed_model_id = t.deployed_model_id)"; - let future: bool = sqlx::query_scalar(&format!( - "SELECT EXISTS (SELECT 1 FROM {table} t WHERE t.valid_from > $3 - AND (t.valid_until IS NULL OR t.valid_until > t.valid_from) AND {omitted})" - )) - .bind(&org_ids) - .bind(&model_ids) - .bind(effective_at) - .fetch_one(&mut *db) - .await?; - ensure!(!future, "removed overlay has future-dated prices; resolve these before removing it"); - sqlx::query(&format!( - "UPDATE {table} t SET valid_until = $3 WHERE t.valid_from <= $3 - AND (t.valid_until IS NULL OR t.valid_until > $3) AND {omitted}" - )) - .bind(&org_ids) - .bind(&model_ids) - .bind(effective_at) - .execute(&mut *db) - .await?; - } + retire_omitted_account_prices(db, ORG_CATALOG_TARIFF_SOURCE, &desired, effective_at).await?; + let (org_ids, model_ids): (Vec, Vec) = desired.into_iter().unzip(); sqlx::query( r#"DELETE FROM model_overlays mo WHERE mo.provisioning_source LIKE $1 diff --git a/dwctl/src/sync/onwards_config/mod.rs b/dwctl/src/sync/onwards_config/mod.rs index f88bb14e2..844a8ff33 100644 --- a/dwctl/src/sync/onwards_config/mod.rs +++ b/dwctl/src/sync/onwards_config/mod.rs @@ -126,7 +126,8 @@ struct OnwardsApiKey { } /// Every alias's overlays, keyed by alias then by account id (as a string, -/// the value of the key's `account` label), as read from `model_overlays`. +/// the value of the key's `account` label), as read from `model_overlays` +/// and `account_model_policies`. type OverlaysByAlias = HashMap>; /// Serving account settings per account id, for accounts that have any. @@ -176,9 +177,38 @@ async fn load_overlays_from_db(db: &PgPool) -> Result Result Result { let rows = sqlx::query!( r#" - SELECT id, self_hosted_only + SELECT id, self_hosted_only, scheduling_priority FROM users WHERE is_deleted = FALSE - AND self_hosted_only + AND (self_hosted_only OR scheduling_priority IS NOT NULL) "# ) .fetch_all(db) @@ -226,6 +256,7 @@ async fn load_accounts_from_db(db: &PgPool) -> Result( // the capability strip above refused scheduling fields for it): when // neither the body nor the preset sets one, realtime's 0, so the // dynamo frontend is never left to pick a value for a missing one. + // Account and model policy hints (priority, latency targets) travel + // the same way but leave the request billed as its own class. // A priority already present is kept, and targets overwrite rather // than remove: this same crate runs again as the hop inside the // serving namespace, with no key policy of its own, and must forward @@ -1335,6 +1337,7 @@ pub async fn target_message_handler( targets.stamp(object); } if !scheduling_fields_refused { + serving_resolution.hints.stamp(object); serving::stamp_default_priority(object); } attempt_body = match serde_json::to_vec(&body) { diff --git a/onwards/src/lib.rs b/onwards/src/lib.rs index 62c507d74..d79c2a43f 100644 --- a/onwards/src/lib.rs +++ b/onwards/src/lib.rs @@ -75,7 +75,8 @@ use handlers::{models as models_handler, target_message_handler}; use models::ExtractedModel; pub use rejections::RequestIdentity; pub use serving::{ - AccountServing, ProviderKind, ServingClass, ServingClassOutcome, ServingOverlay, ServingTargets, + AccountServing, ProviderKind, ServingClass, ServingClassOutcome, ServingHints, ServingOverlay, + ServingTargets, }; /// Type alias for response transformation function diff --git a/onwards/src/serving.rs b/onwards/src/serving.rs index c128aa34b..502059bf9 100644 --- a/onwards/src/serving.rs +++ b/onwards/src/serving.rs @@ -8,9 +8,15 @@ //! [`ServingClass::Custom`]; otherwise it resolves as //! [`ServingClass::Standard`] and sends no targets. //! +//! Separately, an account may carry a scheduling **priority** (account-wide, +//! or per model on an overlay) and per-model latency targets. These are sent +//! the same way but leave the request [`ServingClass::Standard`]: they change +//! scheduling, not the price. +//! //! Daemon legs (batch and flex on the batch-purpose key, continuation resume -//! legs) never receive overlay targets: their existing deadline priorities -//! carry the ordering. +//! legs) never receive overlay targets or priorities and keep their own: +//! batch and flex carry deadline-derived priorities, background dispatch the +//! reserved `i32::MIN`, and continuation resumes their fixed priority. //! //! The targets travel in the request body as `nvext.router.{ttft_target, //! itl_target}` and, when non-zero, the priority as @@ -240,12 +246,16 @@ pub struct AccountServing { /// Never fall over to an external provider. #[serde(default)] pub self_hosted_only: bool, + /// Scheduling priority for the account's realtime requests on every + /// alias, unless an overlay sets its own. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub priority: Option, } impl AccountServing { /// True when nothing here changes a request; such accounts are not synced. pub fn is_empty(&self) -> bool { - !self.self_hosted_only + !self.self_hosted_only && self.priority.is_none() } } @@ -259,6 +269,76 @@ pub struct ServingOverlay { /// Overrides the account's `self_hosted_only` on this alias. #[serde(default)] pub self_hosted_only: Option, + /// Scheduling priority on this alias, overriding the account's. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub priority: Option, + /// Time-to-first-token target on this alias, milliseconds. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ttft_ms: Option, + /// Inter-token-latency target on this alias, milliseconds. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub itl_ms: Option, +} + +/// Scheduling hints for one request that do not change how it is billed: +/// each field is written into the body only when set. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ServingHints { + pub priority: Option, + pub ttft_ms: Option, + pub itl_ms: Option, +} + +impl ServingHints { + pub fn is_empty(&self) -> bool { + self.priority.is_none() && self.ttft_ms.is_none() && self.itl_ms.is_none() + } + + /// Write the set hints into a request body, overwriting existing values: + /// `nvext.router.{ttft_target, itl_target}` and `nvext.agent_hints.priority`. + pub fn stamp(&self, body: &mut serde_json::Map) { + if self.is_empty() { + return; + } + let nvext = body + .entry(NVEXT_FIELD) + .or_insert_with(|| serde_json::Value::Object(Default::default())); + if !nvext.is_object() { + *nvext = serde_json::Value::Object(Default::default()); + } + let nvext = nvext + .as_object_mut() + .expect("nvext was just made an object"); + if self.ttft_ms.is_some() || self.itl_ms.is_some() { + let router = nvext + .entry(NVEXT_ROUTER_FIELD) + .or_insert_with(|| serde_json::Value::Object(Default::default())); + if !router.is_object() { + *router = serde_json::Value::Object(Default::default()); + } + let router = router + .as_object_mut() + .expect("router was just made an object"); + if let Some(ttft) = self.ttft_ms { + router.insert(TTFT_TARGET_FIELD.to_string(), ttft.into()); + } + if let Some(itl) = self.itl_ms { + router.insert(ITL_TARGET_FIELD.to_string(), itl.into()); + } + } + if let Some(priority) = self.priority { + let hints = nvext + .entry(NVEXT_AGENT_HINTS_FIELD) + .or_insert_with(|| serde_json::Value::Object(Default::default())); + if !hints.is_object() { + *hints = serde_json::Value::Object(Default::default()); + } + hints + .as_object_mut() + .expect("agent_hints was just made an object") + .insert(PRIORITY_FIELD.to_string(), priority.into()); + } + } } /// Key purposes whose requests are daemon legs: the ordering already travels @@ -272,6 +352,9 @@ pub struct ServingResolution { pub resolved: ServingClass, /// What a `dynamo` member is told to serve to; `None` sends nothing. pub targets: Option, + /// Priority and latency targets that leave the class unchanged. Empty + /// whenever `targets` is set. + pub hints: ServingHints, /// Restrict the composite to its non-external members. pub self_hosted_only: bool, } @@ -304,11 +387,23 @@ pub fn resolve( Some(targets) if !daemon => ServingResolution { resolved: ServingClass::Custom, targets: Some(targets), + hints: ServingHints::default(), self_hosted_only, }, _ => ServingResolution { resolved: ServingClass::Standard, targets: None, + hints: if daemon { + ServingHints::default() + } else { + ServingHints { + priority: overlay + .and_then(|o| o.priority) + .or(account.and_then(|a| a.priority)), + ttft_ms: overlay.and_then(|o| o.ttft_ms), + itl_ms: overlay.and_then(|o| o.itl_ms), + } + }, self_hosted_only, }, } @@ -364,6 +459,7 @@ mod tests { fn self_hosted_only_comes_from_the_account_unless_the_overlay_overrides() { let a = AccountServing { self_hosted_only: true, + ..Default::default() }; assert!(resolve(Some(&a), None, Some("realtime")).self_hosted_only); // the daemon leg carries the restriction too: batch work must not spill either @@ -454,6 +550,94 @@ mod tests { assert_eq!(body["nvext"]["agent_hints"]["priority"], 0); } + #[test] + fn account_and_overlay_priorities_resolve_as_standard_hints() { + let account = AccountServing { + priority: Some(100), + ..Default::default() + }; + let r = resolve(Some(&account), None, Some("realtime")); + assert_eq!(r.resolved, ServingClass::Standard); + assert_eq!( + r.hints.priority, + Some(100), + "the account default applies everywhere" + ); + + let overlay = ServingOverlay { + priority: Some(300), + ttft_ms: Some(900), + itl_ms: Some(25), + ..Default::default() + }; + let r = resolve(Some(&account), Some(&overlay), Some("realtime")); + assert_eq!( + r.resolved, + ServingClass::Standard, + "hints never change the billed class" + ); + assert_eq!( + r.hints, + ServingHints { + priority: Some(300), + ttft_ms: Some(900), + itl_ms: Some(25) + } + ); + + // Daemon legs keep their deadline priority. + assert!( + resolve(Some(&account), Some(&overlay), Some("batch")) + .hints + .is_empty() + ); + + // Explicit overlay targets still win and carry their own priority. + let legacy = ServingOverlay { + targets: Some(DEAL), + priority: Some(1), + ..Default::default() + }; + let r = resolve(Some(&account), Some(&legacy), Some("realtime")); + assert_eq!(r.resolved, ServingClass::Custom); + assert!(r.hints.is_empty()); + } + + #[test] + fn hints_stamp_only_what_is_set() { + let mut body = + serde_json::json!({"model": "m", "nvext": {"agent_hints": {"max_batch_size": 8}}}); + ServingHints { + priority: Some(500), + ..Default::default() + } + .stamp(body.as_object_mut().unwrap()); + assert_eq!( + body["nvext"]["agent_hints"], + serde_json::json!({"max_batch_size": 8, "priority": 500}) + ); + assert!(body["nvext"].get("router").is_none()); + + let mut body = serde_json::json!({"model": "m"}); + ServingHints { + ttft_ms: Some(1000), + ..Default::default() + } + .stamp(body.as_object_mut().unwrap()); + assert_eq!( + body["nvext"], + serde_json::json!({"router": {"ttft_target": 1000}}) + ); + + let mut body = serde_json::json!({"model": "m"}); + ServingHints::default().stamp(body.as_object_mut().unwrap()); + assert_eq!( + body, + serde_json::json!({"model": "m"}), + "nothing set, nothing written" + ); + } + #[test] fn kinds_and_outcomes() { assert!(AccountServing::default().is_empty()); @@ -695,6 +879,7 @@ mod handler_tests { let mut s = with_deal(&[EXTERNAL]); s.account = Some(AccountServing { self_hosted_only: true, + ..Default::default() }); let mock = MockHttpClient::new(StatusCode::OK, OK_BODY); let request = Request::builder() @@ -731,6 +916,56 @@ mod handler_tests { } } + #[tokio::test] + async fn account_and_overlay_hints_reach_dynamo_members_as_standard() { + let mut s = setup(&[DYNAMO]); + s.account = Some(AccountServing { + priority: Some(100), + ..Default::default() + }); + let mock = MockHttpClient::new(StatusCode::OK, OK_BODY); + assert_eq!( + post(&server(&s, &mock), ALIAS, json!({})) + .await + .status_code(), + 200 + ); + assert_eq!( + body(&mock.get_requests()[0])["nvext"], + json!({"agent_hints": {"priority": 100}}) + ); + + s.overlay = Some(ServingOverlay { + priority: Some(300), + ttft_ms: Some(900), + itl_ms: Some(25), + ..Default::default() + }); + let mock = MockHttpClient::new(StatusCode::OK, OK_BODY); + assert_eq!( + post(&server(&s, &mock), ALIAS, json!({})) + .await + .status_code(), + 200 + ); + assert_eq!(sent(&mock.get_requests()[0]), Some((900, 25, Some(300)))); + + // An external member never receives them. + let mut s = setup(&[EXTERNAL]); + s.account = Some(AccountServing { + priority: Some(100), + ..Default::default() + }); + let mock = MockHttpClient::new(StatusCode::OK, OK_BODY); + assert_eq!( + post(&server(&s, &mock), ALIAS, json!({})) + .await + .status_code(), + 200 + ); + assert!(body(&mock.get_requests()[0]).get("nvext").is_none()); + } + #[tokio::test] async fn explicit_overlay_targets_are_sent_as_is() { let mut s = setup(&[DYNAMO]); @@ -890,6 +1125,7 @@ mod handler_tests { let mut s = setup(&members); s.account = Some(AccountServing { self_hosted_only: true, + ..Default::default() }); let mock = MockHttpClient::new(StatusCode::INTERNAL_SERVER_ERROR, r#"{"error":"boom"}"#); let srv = server(&s, &mock); @@ -913,6 +1149,7 @@ mod handler_tests { let mut s = setup(&[EXTERNAL]); s.account = Some(AccountServing { self_hosted_only: true, + ..Default::default() }); let mock = MockHttpClient::new(StatusCode::OK, OK_BODY); let srv = server(&s, &mock); @@ -926,6 +1163,7 @@ mod handler_tests { let mut s = setup(&[DYNAMO, HOSTED]); s.account = Some(AccountServing { self_hosted_only: true, + ..Default::default() }); let mock = MockHttpClient::new(StatusCode::INTERNAL_SERVER_ERROR, r#"{"error":"boom"}"#); let srv = server(&s, &mock); @@ -936,6 +1174,7 @@ mod handler_tests { let mut s = setup(&members); s.account = Some(AccountServing { self_hosted_only: true, + ..Default::default() }); s.overlay = Some(ServingOverlay { self_hosted_only: Some(false),