diff --git a/.sqlx/query-478ac8aed75ee5f70690dbfbdfaa7d449ac41421177dd420a7d06318a8296249.json b/.sqlx/query-478ac8aed75ee5f70690dbfbdfaa7d449ac41421177dd420a7d06318a8296249.json new file mode 100644 index 000000000..d238fc652 --- /dev/null +++ b/.sqlx/query-478ac8aed75ee5f70690dbfbdfaa7d449ac41421177dd420a7d06318a8296249.json @@ -0,0 +1,34 @@ +{ + "db_name": "PostgreSQL", + "query": "\n SELECT p.user_id, dm.alias, p.toleration_preset AS \"toleration_preset!\"\n FROM account_model_policies p\n JOIN deployed_models dm ON dm.id = p.deployed_model_id\n WHERE p.toleration_preset IS NOT NULL\n AND dm.deleted = FALSE\n AND (p.serving_class IS NULL OR p.serving_class = 'standard')\n AND ($1::uuid IS NULL OR p.user_id = $1)\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": "toleration_preset!", + "type_info": "Text" + } + ], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [ + false, + false, + true + ] + }, + "hash": "478ac8aed75ee5f70690dbfbdfaa7d449ac41421177dd420a7d06318a8296249" +} diff --git a/.sqlx/query-2b001e26e906ff93f35846396d9c92ef37cded96d830620f737e6e2985901ee8.json b/.sqlx/query-8fbd0c37b78964f008a9ad0b0918b08d0f0dede0bf76c66a853fdaeea27ce809.json similarity index 59% rename from .sqlx/query-2b001e26e906ff93f35846396d9c92ef37cded96d830620f737e6e2985901ee8.json rename to .sqlx/query-8fbd0c37b78964f008a9ad0b0918b08d0f0dede0bf76c66a853fdaeea27ce809.json index 419b6c067..9114ea17b 100644 --- a/.sqlx/query-2b001e26e906ff93f35846396d9c92ef37cded96d830620f737e6e2985901ee8.json +++ b/.sqlx/query-8fbd0c37b78964f008a9ad0b0918b08d0f0dede0bf76c66a853fdaeea27ce809.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT ak.secret, u.zero_data_retention, u.disabled_modalities, u.toleration_preset\n FROM api_keys ak\n JOIN users u ON u.id = ak.user_id\n WHERE NOT ak.is_deleted\n ", + "query": "\n SELECT ak.secret, ak.user_id, u.zero_data_retention, u.disabled_modalities, u.toleration_preset\n FROM api_keys ak\n JOIN users u ON u.id = ak.user_id\n WHERE NOT ak.is_deleted\n ", "describe": { "columns": [ { @@ -10,16 +10,21 @@ }, { "ordinal": 1, + "name": "user_id", + "type_info": "Uuid" + }, + { + "ordinal": 2, "name": "zero_data_retention", "type_info": "Bool" }, { - "ordinal": 2, + "ordinal": 3, "name": "disabled_modalities", "type_info": "TextArray" }, { - "ordinal": 3, + "ordinal": 4, "name": "toleration_preset", "type_info": "Text" } @@ -31,8 +36,9 @@ false, false, false, + false, true ] }, - "hash": "2b001e26e906ff93f35846396d9c92ef37cded96d830620f737e6e2985901ee8" + "hash": "8fbd0c37b78964f008a9ad0b0918b08d0f0dede0bf76c66a853fdaeea27ce809" } diff --git a/docs/src/reference/configuration.md b/docs/src/reference/configuration.md index ed521742f..9372814ce 100644 --- a/docs/src/reference/configuration.md +++ b/docs/src/reference/configuration.md @@ -703,6 +703,11 @@ scheduling: | `description` | string | One or two sentences of help text shown as the dashboard tooltip. Must be non-empty. | | `tolerations` | list | The request-field list (`{key, operator?, value?, effect?}`). `operator` defaults to `Equal` and then requires `value`; `Exists` forbids `value`; `effect` may be `NoSchedule` or `PreferNoSchedule`. May be empty; validated at config load. | +An [account file](./model-provisioning.md#account-limits) can also set a +preset for one model, or one class of a model, in place of the account's own +preset for those requests. The more specific setting wins: class, then model, +then the account. + The `GET /admin/api/v1/config` endpoint exposes the presets (id, label, description) so the dashboard can render the control. An account stores only the preset's id; an id the config no longer defines is logged and the request path diff --git a/docs/src/reference/model-provisioning.md b/docs/src/reference/model-provisioning.md index 6840dbfc1..d6e84d133 100644 --- a/docs/src/reference/model-provisioning.md +++ b/docs/src/reference/model-provisioning.md @@ -384,6 +384,7 @@ priority: 100 # optional; default for every model models: example/chat-model: priority: 300 # optional; overrides the account default + toleration_preset: no-overflow # optional; overrides the account's preset 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"} @@ -398,6 +399,7 @@ models: ttft_ms: 1000 # optional latency targets itl_ms: 30 priority: 500 # optional; overrides the model's priority + toleration_preset: gpu-only # optional; overrides the model's preset ``` The older shape is still accepted and means the same as @@ -421,6 +423,13 @@ Rules: `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. +- `toleration_preset` names one of `scheduling.toleration_presets` (see + [Account Scheduling](./configuration.md#account-scheduling-tolerations-and-provider-fallback)). + It inherits class, then model, then the account's own preset (set by a + platform manager), then none. It applies to realtime, flex and background + requests, and to batch lines by each line's model. A preset id the + deployment does not configure is logged at startup and fails closed to the + empty (most restrictive) list. - 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 diff --git a/dwctl/migrations/20261021190000_account_policy_toleration_presets.sql b/dwctl/migrations/20261021190000_account_policy_toleration_presets.sql new file mode 100644 index 000000000..ebcda57a3 --- /dev/null +++ b/dwctl/migrations/20261021190000_account_policy_toleration_presets.sql @@ -0,0 +1,10 @@ +-- A per-model or per-class scheduling-toleration preset from the account +-- catalog. It overrides the account's own preset (users.toleration_preset) for +-- requests to that model or class. Written only by the account catalog +-- reconciler. +SET LOCAL lock_timeout = '5s'; + +ALTER TABLE account_model_policies ADD COLUMN toleration_preset TEXT; + +COMMENT ON COLUMN account_model_policies.toleration_preset IS + 'Id of a configured scheduling-toleration preset applied to this account''s requests to the model (serving_class NULL) or class, in place of the account''s own preset. NULL inherits. An id the config does not define fails closed to the most restrictive list.'; diff --git a/dwctl/src/account_limits.rs b/dwctl/src/account_limits.rs index 5bd3a5c24..61a19984f 100644 --- a/dwctl/src/account_limits.rs +++ b/dwctl/src/account_limits.rs @@ -13,7 +13,7 @@ //! what this catalog owns. use std::{ - collections::{BTreeMap, HashMap, HashSet}, + collections::{BTreeMap, BTreeSet, HashMap, HashSet}, fs, path::Path, }; @@ -78,6 +78,10 @@ pub struct AccountModelPolicy { /// Scheduling priority on this model, overriding the account's. #[serde(default)] pub priority: Option, + /// Scheduling-toleration preset id for this model, in place of the + /// account's own preset. + #[serde(default)] + pub toleration_preset: Option, /// The account's batch and flex prices on this model, per completion /// window. Realtime prices are set per class. #[serde(default)] @@ -110,11 +114,15 @@ pub struct AccountClassPolicy { /// Scheduling priority on this class, overriding the model's and the account's. #[serde(default)] pub priority: Option, + /// Scheduling-toleration preset id for this class, overriding the model's + /// and the account's. + #[serde(default)] + pub toleration_preset: Option, } impl AccountModelPolicy { fn is_empty(&self) -> bool { - self.priority.is_none() && self.tariffs.is_empty() && self.classes.is_empty() + self.priority.is_none() && self.toleration_preset.is_none() && self.tariffs.is_empty() && self.classes.is_empty() } } @@ -126,10 +134,11 @@ impl AccountClassPolicy { && self.ttft_ms.is_none() && self.itl_ms.is_none() && self.priority.is_none() + && self.toleration_preset.is_none() } fn has_policy(&self) -> bool { - self.priority.is_some() || self.ttft_ms.is_some() || self.itl_ms.is_some() + self.priority.is_some() || self.ttft_ms.is_some() || self.itl_ms.is_some() || self.toleration_preset.is_some() } } @@ -208,6 +217,22 @@ impl AccountLimitsCatalog { /// 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. + /// Every toleration preset id the files name, for checking against the + /// deployment's configured presets. + pub fn toleration_presets(&self) -> BTreeSet<&str> { + self.files + .iter() + .flat_map(|file| file.document.models.values()) + .flat_map(|model| { + model + .toleration_preset + .as_deref() + .into_iter() + .chain(model.classes.values().filter_map(|class| class.toleration_preset.as_deref())) + }) + .collect() + } + pub fn validate_against_overlays(&self, overlays: &OrgCatalog) -> Result<()> { let mut declared: HashMap<(String, String), &str> = HashMap::new(); for entry in &overlays.orgs { @@ -256,6 +281,7 @@ impl AccountLimitsCatalog { let context = format!("{source}: model {alias:?}"); ensure!(!model.is_empty(), "{context} configures nothing"); validate_priority(model.priority, &context)?; + validate_preset(model.toleration_preset.as_deref(), &context)?; validate_tariffs(&model.tariffs, &context)?; ensure!( model.tariffs.iter().all(|t| t.purpose == TariffPurpose::Batch), @@ -287,6 +313,7 @@ impl AccountLimitsCatalog { validate_cache_tariff(&cache.as_tariff(), &context)?; } validate_priority(policy.priority, &context)?; + validate_preset(policy.toleration_preset.as_deref(), &context)?; validate_target(policy.ttft_ms, "ttft_ms", &context)?; validate_target(policy.itl_ms, "itl_ms", &context)?; } @@ -307,6 +334,20 @@ fn validate_priority(priority: Option, context: &str) -> Result<()> { Ok(()) } +/// A preset id is the operator's slug from `scheduling.toleration_presets`. +/// Whether it is configured is checked at startup (and an unconfigured id +/// fails closed at request time), since the files are validated without the +/// deployment's config. +fn validate_preset(preset: Option<&str>, context: &str) -> Result<()> { + if let Some(preset) = preset { + ensure!( + !preset.is_empty() && preset.trim() == preset, + "{context}: toleration_preset must be a non-empty id without surrounding whitespace" + ); + } + Ok(()) +} + /// Latency targets are stored as positive `INTEGER`s. fn validate_target(target: Option, field: &str, context: &str) -> Result<()> { if let Some(target) = target { @@ -328,6 +369,7 @@ struct PolicyRow { priority: Option, ttft_ms: Option, itl_ms: Option, + toleration_preset: Option, } pub async fn apply(pool: &PgPool, catalog: &AccountLimitsCatalog) -> Result<()> { @@ -440,7 +482,7 @@ async fn apply_policies( let account = accounts[&file.document.account]; for (alias, model) in &file.document.models { let model_id = models[alias]; - if model.priority.is_some() { + if model.priority.is_some() || model.toleration_preset.is_some() { rows.push(PolicyRow { account, model: model_id, @@ -448,6 +490,7 @@ async fn apply_policies( priority: model.priority, ttft_ms: None, itl_ms: None, + toleration_preset: model.toleration_preset.clone(), }); } for (class, policy) in &model.classes { @@ -459,6 +502,7 @@ async fn apply_policies( priority: policy.priority, ttft_ms: policy.ttft_ms.map(i32::try_from).transpose()?, itl_ms: policy.itl_ms.map(i32::try_from).transpose()?, + toleration_preset: policy.toleration_preset.clone(), }); } } @@ -470,6 +514,7 @@ async fn apply_policies( 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(); + let presets_col: Vec> = rows.iter().map(|r| r.toleration_preset.clone()).collect(); sqlx::query( "DELETE FROM account_model_policies p @@ -514,19 +559,22 @@ async fn apply_policies( ); } 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) + "INSERT INTO account_model_policies + (user_id, deployed_model_id, serving_class, priority, ttft_ms, itl_ms, toleration_preset, provisioning_source) + SELECT d.user_id, d.deployed_model_id, d.serving_class, d.priority, d.ttft_ms, d.itl_ms, d.toleration_preset, $8 + FROM UNNEST($1::uuid[], $2::uuid[], $3::text[], $4::int[], $5::int[], $6::int[], $7::text[]) + AS d(user_id, deployed_model_id, serving_class, priority, ttft_ms, itl_ms, toleration_preset) 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, + toleration_preset = EXCLUDED.toleration_preset, 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)", + account_model_policies.toleration_preset, account_model_policies.provisioning_source) + IS DISTINCT FROM (EXCLUDED.priority, EXCLUDED.ttft_ms, EXCLUDED.itl_ms, EXCLUDED.toleration_preset, + EXCLUDED.provisioning_source)", ) .bind(&accounts_col) .bind(&models_col) @@ -534,6 +582,7 @@ async fn apply_policies( .bind(&priorities_col) .bind(&ttft_col) .bind(&itl_col) + .bind(&presets_col) .bind(POLICY_SOURCE) .execute(&mut *db) .await @@ -866,6 +915,14 @@ mod tests { "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 toleration_preset: ' gpu'\n".to_owned(), + "toleration_preset must be a non-empty id", + ), + ( + "account: acme\nmodels:\n org/model:\n classes:\n standard:\n toleration_preset: ''\n".to_owned(), + "toleration_preset must be a non-empty id", + ), ( "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", @@ -1158,4 +1215,55 @@ mod tests { .unwrap(); assert_eq!(untouched, (Some(7), Some("elsewhere".to_owned()))); } + + #[dwctl_test_macros::test] + async fn toleration_presets_are_written_per_model_and_class_and_cleared(pool: PgPool) { + let org = insert_account(&pool, "acme", "organization").await; + insert_virtual_model(&pool, "org/model").await; + let directory = tempdir().unwrap(); + write( + directory.path(), + "acme.yaml", + "account: acme\nmodels:\n org/model:\n toleration_preset: no-overflow\n classes:\n standard:\n toleration_preset: gpu\n", + ); + let catalog = AccountLimitsCatalog::load(directory.path()).unwrap(); + assert_eq!( + catalog.toleration_presets().into_iter().collect::>(), + vec!["gpu", "no-overflow"] + ); + apply(&pool, &catalog).await.unwrap(); + let presets = || async { + sqlx::query_as::<_, (Option, Option)>( + "SELECT serving_class, toleration_preset FROM account_model_policies WHERE user_id = $1 ORDER BY serving_class NULLS FIRST", + ) + .bind(org) + .fetch_all(&pool) + .await + .unwrap() + }; + assert_eq!( + presets().await, + vec![(None, Some("no-overflow".into())), (Some("standard".into()), Some("gpu".into()))] + ); + + // The per-key policy sees the class preset folded over the model's. + let key = format!("sk-{}", Uuid::new_v4()); + sqlx::query("INSERT INTO api_keys (name, secret, user_id, created_by, purpose) VALUES ('k', $1, $2, $2, 'realtime')") + .bind(&key) + .bind(org) + .execute(&pool) + .await + .unwrap(); + let cache = crate::sync::key_policy::KeyPolicyCache::empty(); + crate::sync::key_policy::refresh(&pool, &cache).await.unwrap(); + assert_eq!(cache.policy(&key).toleration_preset_for("org/model"), Some("gpu")); + assert_eq!(cache.policy(&key).toleration_preset_for("org/other"), None); + let mut conn = pool.acquire().await.unwrap(); + let direct = crate::sync::key_policy::load_policy(&mut conn, org).await.unwrap(); + assert_eq!(direct.toleration_preset_for("org/model"), Some("gpu")); + + write(directory.path(), "acme.yaml", "account: acme\npriority: 1\n"); + apply(&pool, &AccountLimitsCatalog::load(directory.path()).unwrap()).await.unwrap(); + assert!(presets().await.is_empty()); + } } diff --git a/dwctl/src/api/handlers/files.rs b/dwctl/src/api/handlers/files.rs index 0fbe2960c..5e1bbc699 100644 --- a/dwctl/src/api/handlers/files.rs +++ b/dwctl/src/api/handlers/files.rs @@ -29,6 +29,7 @@ use crate::db::{ use crate::errors::{Error, Result}; use crate::image_normalizer::{ImageInput, ImageNormalizer, Mode as ImageNormalizerMode, walker as image_walker}; use crate::reasoning::ModelReasoningPolicy; +use crate::sync::key_policy::AccountTolerations; use crate::types::Resource; use axum::{ Json, @@ -177,7 +178,7 @@ impl OpenAIBatchRequest { api_key: String, accessible_models: &HashMap, allowed_url_paths: &[String], - tolerations: Option<&serde_json::Value>, + tolerations: &AccountTolerations, ) -> Result { // Validate custom_id is safe for HTTP headers validate_custom_id(&self.custom_id)?; @@ -236,10 +237,11 @@ impl OpenAIBatchRequest { crate::inference::middleware::strip_scheduling_priority(&mut sanitized_body); // A preset-pinned account's batch requests carry its preset's fixed - // toleration list, exactly as its realtime and flex requests do. Applied - // after the strip above so the client can never override it. - if let Some(tolerations) = tolerations { - crate::inference::middleware::set_routing_tolerations(&mut sanitized_body, tolerations); + // toleration list for this line's model, exactly as its realtime and + // flex requests do. Applied after the strip above so the client can + // never override it. + if let Some(tolerations) = tolerations.for_model(&model) { + crate::inference::middleware::set_routing_tolerations(&mut sanitized_body, &tolerations); } // Serialize sanitized body back to string @@ -697,11 +699,10 @@ struct FileRequestContext { api_key: String, accessible_models: HashMap, allowed_url_paths: Vec, - /// The toleration list the account that owns this upload is pinned to, - /// resolved from its preset, or `None` when it is on "Default". Written - /// verbatim to `nvext.routing_constraints.tolerations` on every template - /// body. - tolerations: Option, + /// The toleration presets of the account that owns this upload, resolved + /// per line by its model. A pinned line's list is written verbatim to + /// `nvext.routing_constraints.tolerations` on its template body. + tolerations: AccountTolerations, } struct AccessibleBatchModel { @@ -946,7 +947,7 @@ fn create_file_stream( api_key.clone(), &accessible_models, &allowed_url_paths, - tolerations.as_ref(), + &tolerations, ) { Ok(mut template) => { // Normalise image URLs in the per-template body @@ -1043,7 +1044,7 @@ fn create_file_stream( api_key.clone(), &accessible_models, &allowed_url_paths, - tolerations.as_ref(), + &tolerations, ) { Ok(mut template) => { // Normalise image URLs in the trailing line as well. @@ -1267,13 +1268,13 @@ pub async fn upload_file( // account's policy straight off the users row, reusing the connection `conn` // this handler already holds — a cached account on Default is not a miss, so // an ordinary upload never takes a second connection. - let tolerations = match state.key_policy_cache.tolerations_if_known(&user_api_key) { + let tolerations = match state.key_policy_cache.account_tolerations_if_known(&user_api_key) { Some(resolved) => resolved, None => { let policy = crate::sync::key_policy::load_policy(&mut conn, target_user_id) .await .map_err(|e| Error::Database(e.into()))?; - state.key_policy_cache.presets().resolve(policy.toleration_preset.as_deref()) + state.key_policy_cache.account_tolerations(policy) } }; @@ -1319,7 +1320,7 @@ pub async fn upload_file( api_key: user_api_key, accessible_models, allowed_url_paths: config.batches.allowed_url_paths.clone(), - tolerations: tolerations.map(|t| (*t).clone()), + tolerations, }, Some(api_key_id), ); diff --git a/dwctl/src/inference/middleware.rs b/dwctl/src/inference/middleware.rs index 3cd5b2fa7..5db158934 100644 --- a/dwctl/src/inference/middleware.rs +++ b/dwctl/src/inference/middleware.rs @@ -177,19 +177,6 @@ pub async fn inference_middleware set_routing_tolerations(&mut request_value, &tolerations), - None => false, - }; - // Asynchronous requests are scheduled by their deadline, so a model // suffix (`alias:suffix`) is discarded here, at the outermost layer, and // analytics, the prompt cache and billing all key on the bare alias. A @@ -211,6 +198,24 @@ pub async fn inference_middleware set_routing_tolerations(&mut request_value, &tolerations), + None => false, + }; + // The realtime path forwards `body_bytes` verbatim and onwards forwards the // bytes it receives (COR-522), so anything a scrubber removed from // `request_value` only stays removed if the bytes are rebuilt from it. @@ -1498,8 +1503,9 @@ pub(crate) fn strip_scheduling_priority(value: &mut serde_json::Value) -> bool { pub(crate) fn pinned_tolerations( cache: &crate::sync::key_policy::KeyPolicyCache, api_key: Option<&str>, + model: &str, ) -> Option> { - api_key.and_then(|secret| cache.tolerations(secret)) + api_key.and_then(|secret| cache.tolerations(secret, model)) } /// Write `tolerations` to `nvext.routing_constraints.tolerations`, in place, @@ -1605,9 +1611,9 @@ mod tests { }, )], ); - assert!(pinned_tolerations(&cache, Some("sk-pinned")).is_some()); - assert!(pinned_tolerations(&cache, Some("sk-other")).is_none()); - assert!(pinned_tolerations(&cache, None).is_none()); + assert!(pinned_tolerations(&cache, Some("sk-pinned"), "any/model").is_some()); + assert!(pinned_tolerations(&cache, Some("sk-other"), "any/model").is_none()); + assert!(pinned_tolerations(&cache, None, "any/model").is_none()); } #[test] diff --git a/dwctl/src/lib.rs b/dwctl/src/lib.rs index 8d1d99a45..bb147b01f 100644 --- a/dwctl/src/lib.rs +++ b/dwctl/src/lib.rs @@ -1617,6 +1617,21 @@ async fn setup_database( 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)?; + // An unconfigured id fails closed at request time; say so up front. + for preset in account_limits.toleration_presets() { + if !config + .scheduling + .toleration_presets + .iter() + .any(|configured| configured.id == preset) + { + tracing::warn!( + preset, + "account files name a toleration preset this deployment does not configure; \ + its requests get the empty (most restrictive) toleration list" + ); + } + } 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/sync/key_policy.rs b/dwctl/src/sync/key_policy.rs index 8d91b3435..6ec572742 100644 --- a/dwctl/src/sync/key_policy.rs +++ b/dwctl/src/sync/key_policy.rs @@ -46,6 +46,39 @@ pub struct KeyPolicy { /// every ZDR and modality read already does) copies a refcount rather than /// allocating a fresh `String`. pub toleration_preset: Option>, + /// Per-model preset ids from the account catalog (`account_model_policies`), + /// keyed by model alias, each overriding `toleration_preset` for requests to + /// that model. A class policy has already been folded over its model's. + /// Shared by every key of the account. + pub model_toleration_presets: Option>>>, +} + +impl KeyPolicy { + /// The preset id that applies to a request for `model`: the model's (or + /// its class's) from the account catalog, else the account's own. + pub fn toleration_preset_for(&self, model: &str) -> Option<&str> { + self.model_toleration_presets + .as_ref() + .and_then(|presets| presets.get(model)) + .or(self.toleration_preset.as_ref()) + .map(|preset| &**preset) + } +} + +/// One account's toleration presets, resolvable per model: what a batch +/// upload stamps on each line, by the line's model. +#[derive(Clone, Default)] +pub struct AccountTolerations { + policy: KeyPolicy, + presets: TolerationPresets, +} + +impl AccountTolerations { + /// The toleration list for a request to `model`, or `None` when the + /// account is on "Default" for it. + pub fn for_model(&self, model: &str) -> Option> { + self.presets.resolve(self.policy.toleration_preset_for(model)) + } } /// Lock-free, cheap-to-clone handle to the shared secret-to-policy map, plus @@ -108,28 +141,30 @@ impl KeyPolicyCache { self.policy(secret).disabled_modalities } - /// The toleration list the account owning `secret` is pinned to, resolved - /// from its stored preset id, or `None` when the account is on "Default". - /// An id the config no longer defines fails closed to the empty list. Hot - /// path safe. - pub fn tolerations(&self, secret: &str) -> Option> { - self.presets.resolve(self.preset_id(secret).as_deref()) + /// The toleration list the account owning `secret` is pinned to for a + /// request to `model`: the account catalog's preset for that model, else + /// the account's own, or `None` when neither is set ("Default"). An id the + /// config no longer defines fails closed to the empty list. Hot path safe. + pub fn tolerations(&self, secret: &str, model: &str) -> Option> { + let map = self.inner.load(); + self.presets + .resolve(map.get(secret).and_then(|policy| policy.toleration_preset_for(model))) } - /// The stored preset id for `secret`, cheaply cloned (an `Arc`), so a - /// caller that only needs the id does not allocate a whole `KeyPolicy`. - fn preset_id(&self, secret: &str) -> Option> { - self.inner.load().get(secret).and_then(|p| p.toleration_preset.clone()) + /// Everything needed to resolve the tolerations of the account owning + /// `secret`, for any model, when the key is present in the map. `None` + /// means the key is absent (a key created since the last refresh), which + /// the caller may resolve from the DB with [`load_policy`]. Hot path safe. + pub fn account_tolerations_if_known(&self, secret: &str) -> Option { + self.policy_if_known(secret).map(|policy| self.account_tolerations(policy)) } - /// The tolerations for `secret` when the key is present in the map: - /// `Some(Some(list))` when pinned, `Some(None)` when the account is on - /// Default. `None` means the key is absent from the map (a key created - /// since the last refresh), which the caller may resolve from the DB. Hot - /// path safe. - pub fn tolerations_if_known(&self, secret: &str) -> Option>> { - let policy = self.policy_if_known(secret)?; - Some(self.presets.resolve(policy.toleration_preset.as_deref())) + /// Resolve tolerations for an account policy read directly from the DB. + pub fn account_tolerations(&self, policy: KeyPolicy) -> AccountTolerations { + AccountTolerations { + policy, + presets: self.presets.clone(), + } } /// The presets a stored id resolves against, for a caller that read an @@ -185,7 +220,7 @@ impl KeyPolicyCache { async fn load(pool: &PgPool) -> Result, sqlx::Error> { let rows = sqlx::query!( r#" - SELECT ak.secret, u.zero_data_retention, u.disabled_modalities, u.toleration_preset + SELECT ak.secret, ak.user_id, u.zero_data_retention, u.disabled_modalities, u.toleration_preset FROM api_keys ak JOIN users u ON u.id = ak.user_id WHERE NOT ak.is_deleted @@ -193,6 +228,7 @@ async fn load(pool: &PgPool) -> Result, sqlx::Error> ) .fetch_all(pool) .await?; + let model_presets = load_model_presets(pool, None).await?; Ok(rows .into_iter() @@ -203,12 +239,47 @@ async fn load(pool: &PgPool) -> Result, sqlx::Error> zdr: r.zero_data_retention, disabled_modalities: ModalitySet::from_db(&r.disabled_modalities), toleration_preset: r.toleration_preset.map(Arc::from), + model_toleration_presets: model_presets.get(&r.user_id).cloned(), }, ) }) .collect()) } +/// Per-model preset ids from the account catalog, for every account or for +/// one, each account's map shared by `Arc` across its keys. A class row +/// (`standard`, the only class a request resolves to today) overrides its +/// model row. +async fn load_model_presets<'c>( + conn: impl sqlx::PgExecutor<'c>, + account: Option, +) -> Result>>>, sqlx::Error> { + let rows = sqlx::query!( + r#" + SELECT p.user_id, dm.alias, p.toleration_preset AS "toleration_preset!" + FROM account_model_policies p + JOIN deployed_models dm ON dm.id = p.deployed_model_id + WHERE p.toleration_preset IS NOT NULL + AND dm.deleted = FALSE + AND (p.serving_class IS NULL OR p.serving_class = 'standard') + AND ($1::uuid IS NULL OR p.user_id = $1) + ORDER BY p.serving_class NULLS FIRST + "#, + account + ) + .fetch_all(conn) + .await?; + + let mut presets: HashMap>> = HashMap::new(); + for row in rows { + presets + .entry(row.user_id) + .or_default() + .insert(row.alias, Arc::from(row.toleration_preset)); + } + Ok(presets.into_iter().map(|(account, presets)| (account, Arc::new(presets))).collect()) +} + /// Reload `cache` in place from the DB, returning the number of keys loaded. /// Shared by [`initial_cache`], the background listener, and the test-only /// manual refresh so none of them can drift. @@ -234,14 +305,17 @@ pub async fn load_policy(conn: &mut sqlx::PgConnection, user_id: crate::types::U ) .fetch_optional(&mut *conn) .await?; - - Ok(row - .map(|r| KeyPolicy { - zdr: r.zero_data_retention, - disabled_modalities: ModalitySet::from_db(&r.disabled_modalities), - toleration_preset: r.toleration_preset.map(Arc::from), - }) - .unwrap_or_default()) + let Some(row) = row else { + return Ok(KeyPolicy::default()); + }; + + let model_toleration_presets = load_model_presets(&mut *conn, Some(user_id)).await?.remove(&user_id); + Ok(KeyPolicy { + zdr: row.zero_data_retention, + disabled_modalities: ModalitySet::from_db(&row.disabled_modalities), + toleration_preset: row.toleration_preset.map(Arc::from), + model_toleration_presets, + }) } /// Load the map once, synchronously, returning a populated cache. Call at @@ -436,26 +510,90 @@ mod tests { ], ); assert_eq!( - cache.tolerations("sk-preset").as_deref(), + cache.tolerations("sk-preset", "any/model").as_deref(), Some(&serde_json::json!([{"key": "pool", "operator": "Equal", "value": "dedicated"}])) ); - assert_eq!(cache.tolerations("sk-empty").as_deref(), Some(&serde_json::json!([]))); + assert_eq!(cache.tolerations("sk-empty", "any/model").as_deref(), Some(&serde_json::json!([]))); // The Default setting and an unknown key are both unpinned. - assert!(cache.tolerations("sk-default").is_none()); - assert!(cache.tolerations("sk-missing").is_none()); - // `tolerations_if_known` distinguishes a cached Default (`Some(None)`) - // from a key absent from the map (`None`), so the batch upload only - // pays for a DB read on a genuine miss. - assert_eq!(cache.tolerations_if_known("sk-default"), Some(None)); - assert_eq!(cache.tolerations_if_known("sk-missing"), None); + assert!(cache.tolerations("sk-default", "any/model").is_none()); + assert!(cache.tolerations("sk-missing", "any/model").is_none()); + // `account_tolerations_if_known` distinguishes a cached Default from a + // key absent from the map (`None`), so the batch upload only pays for + // a DB read on a genuine miss. + assert!( + cache + .account_tolerations_if_known("sk-default") + .unwrap() + .for_model("any/model") + .is_none() + ); + assert!(cache.account_tolerations_if_known("sk-missing").is_none()); assert_eq!( - cache.tolerations_if_known("sk-preset").map(|t| t.map(|v| v.as_ref().clone())), - Some(Some( - serde_json::json!([{"key": "pool", "operator": "Equal", "value": "dedicated"}]) - )) + cache + .account_tolerations_if_known("sk-preset") + .unwrap() + .for_model("any/model") + .as_deref(), + Some(&serde_json::json!([{"key": "pool", "operator": "Equal", "value": "dedicated"}])) ); } + #[test] + fn a_model_preset_overrides_the_account_preset_for_that_model_only() { + use crate::config::{SchedulingConfig, TolerationPreset}; + use crate::scheduling::{SchedulingOperator, SchedulingToleration}; + + let presets = TolerationPresets::new(&SchedulingConfig { + toleration_presets: vec![ + SchedulingConfig::default_preset(), + TolerationPreset { + id: "gpu".to_string(), + label: "GPU".to_string(), + description: "GPU capacity only.".to_string(), + tolerations: vec![SchedulingToleration { + key: "pool".to_string(), + operator: Some(SchedulingOperator::Equal), + value: Some("dedicated".to_string()), + effect: None, + }], + }, + ], + }); + let models = Arc::new(HashMap::from([("org/pinned".to_string(), Arc::::from("no-overflow"))])); + let cache = KeyPolicyCache::from_policies_with_presets( + presets, + [ + ( + "sk-both".to_string(), + KeyPolicy { + toleration_preset: Some("gpu".into()), + model_toleration_presets: Some(models.clone()), + ..Default::default() + }, + ), + ( + "sk-model-only".to_string(), + KeyPolicy { + model_toleration_presets: Some(models), + ..Default::default() + }, + ), + ], + ); + let gpu = serde_json::json!([{"key": "pool", "operator": "Equal", "value": "dedicated"}]); + assert_eq!(cache.tolerations("sk-both", "org/pinned").as_deref(), Some(&serde_json::json!([]))); + assert_eq!(cache.tolerations("sk-both", "org/other").as_deref(), Some(&gpu)); + assert_eq!( + cache.tolerations("sk-model-only", "org/pinned").as_deref(), + Some(&serde_json::json!([])) + ); + // "No preset" stays distinct from the empty list on other models. + assert!(cache.tolerations("sk-model-only", "org/other").is_none()); + let account = cache.account_tolerations_if_known("sk-both").unwrap(); + assert_eq!(account.for_model("org/pinned").as_deref(), Some(&serde_json::json!([]))); + assert_eq!(account.for_model("org/other").as_deref(), Some(&gpu)); + } + #[test] fn an_unconfigured_preset_fails_closed_to_the_empty_list() { let cache = KeyPolicyCache::from_policies_with_presets( @@ -470,6 +608,9 @@ mod tests { ); // A preset an operator has dropped from config resolves to the most // restrictive option rather than leaving the account unpinned. - assert_eq!(cache.tolerations("sk-retired").as_deref(), Some(&serde_json::json!([]))); + assert_eq!( + cache.tolerations("sk-retired", "any/model").as_deref(), + Some(&serde_json::json!([])) + ); } } diff --git a/dwctl/src/test/toleration_presets.rs b/dwctl/src/test/toleration_presets.rs index 2bb55d33c..1b4a4b7b5 100644 --- a/dwctl/src/test/toleration_presets.rs +++ b/dwctl/src/test/toleration_presets.rs @@ -354,6 +354,64 @@ async fn a_preset_account_stamps_its_tolerations_on_uploaded_batch_requests(pool assert_eq!(body["model"], "plain"); } +impl Fixture { + /// Write an account-catalog policy row carrying a preset for `alias` + /// (`class` None = the whole model) and refresh the per-key policy map. + async fn set_model_preset(&self, alias: &str, class: Option<&str>, preset: &str) { + sqlx::query( + "INSERT INTO account_model_policies (user_id, deployed_model_id, serving_class, toleration_preset, provisioning_source) + SELECT $1, id, $3, $4, 'account-catalog' FROM deployed_models WHERE alias = $2", + ) + .bind(self.user) + .bind(alias) + .bind(class) + .bind(preset) + .execute(&self.pool) + .await + .unwrap(); + self.services.sync_key_policy(&self.pool).await.unwrap(); + } +} + +/// An account-catalog preset on a model overrides the account's own preset for +/// that model only; a class preset overrides its model's. Realtime and batch +/// agree. +#[dwctl_test_macros::test] +async fn a_model_preset_overrides_the_account_preset_on_that_model(pool: PgPool) { + let f = Fixture::new(&pool).await; + let gpu = json!([{"key": "pool", "operator": "Equal", "value": "gpu"}]); + f.set_preset(Some("gpu")).await; + + // A preset on another model leaves this one on the account's preset. + f.set_model_preset("preset-member", None, "no-overflow").await; + assert_eq!(f.post_realtime().await["nvext"]["routing_constraints"]["tolerations"], gpu); + assert_eq!(f.upload_batch().await["nvext"]["routing_constraints"]["tolerations"], gpu); + + // A model-level preset on this model replaces the account's. + f.set_model_preset("plain", None, "no-overflow").await; + assert_eq!(f.post_realtime().await["nvext"]["routing_constraints"]["tolerations"], json!([])); + assert_eq!(f.upload_batch().await["nvext"]["routing_constraints"]["tolerations"], json!([])); + + // A class preset overrides the model's. + f.set_model_preset("plain", Some("standard"), "gpu").await; + assert_eq!(f.post_realtime().await["nvext"]["routing_constraints"]["tolerations"], gpu); + assert_eq!(f.post_responses().await["nvext"]["routing_constraints"]["tolerations"], gpu); + assert_eq!(f.upload_batch().await["nvext"]["routing_constraints"]["tolerations"], gpu); +} + +/// A model preset applies to an account with no preset of its own, and only +/// on that model: "no preset" elsewhere stays distinct from the empty list. +#[dwctl_test_macros::test] +async fn a_model_preset_pins_an_otherwise_default_account(pool: PgPool) { + let f = Fixture::new(&pool).await; + f.set_model_preset("preset-member", None, "gpu").await; + let body = f.post_realtime().await; + assert!(body["nvext"]["routing_constraints"]["tolerations"].is_null(), "{body}"); + + f.set_model_preset("plain", None, "no-overflow").await; + assert_eq!(f.post_realtime().await["nvext"]["routing_constraints"]["tolerations"], json!([])); +} + /// A platform manager sets the preset through the admin API; a non-manager /// cannot, an unknown id is refused, and the preset is only visible back to a /// platform manager.