diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index c97c1d6515b..c21352cde86 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -9,6 +9,7 @@ pub mod hooks; mod key_status; mod keys; pub mod metrics; +pub mod model_capacity; pub mod rate_limits; pub mod request_hooks; mod secrets; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index b8a35eddbcd..b242a3b88ab 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -128,6 +128,11 @@ async fn main() -> Result<(), Box> { ), )); } + if let Some(capacity) = + litellm_gateway::model_capacity::ModelCapacity::for_gateway(&config, coordination.clone()) + { + hooks.push(Arc::new(capacity)); + } hooks.extend(litellm_gateway::request_hooks::from_config(&config)?); let mut observers: Vec> = Vec::new(); let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config)); diff --git a/litellm-rust/crates/gateway/src/model_capacity.rs b/litellm-rust/crates/gateway/src/model_capacity.rs new file mode 100644 index 00000000000..5cc98e808e2 --- /dev/null +++ b/litellm-rust/crates/gateway/src/model_capacity.rs @@ -0,0 +1,654 @@ +//! The model capacity hook, the proxy's `dynamic_rate_limiter_v3`: caps each model group at the +//! sum of its deployments' `rpm` and `tpm` across every replica, and with `priority_reservation` +//! splits that capacity between priorities once the group is busy. +//! +//! Each request adds one to `{model_saturation_check:}:requests` and checks the group's +//! tokens, all or nothing, with the proxy's keys and scripts. With priorities, the request's pool +//! `{priority_model::}` (or `:default_pool`) is held to its share of capacity +//! once the group's saturation reaches `saturation_threshold`, and only counted below it. Tokens +//! are added to both counters when the response ends. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use futures_util::future::BoxFuture; +use http::{HeaderName, HeaderValue}; +use litellm_config::Config; +use litellm_gateway_hooks::{ + Facts, Hook, KeyFacts, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets, +}; +use serde_json::Value; + +use crate::{ + coordination::Coordination, + rate_limits::{Kind, Pair, RateLimitSettings, Reserved, Store, Unavailable}, +}; + +/// The names a config's `callbacks` list turns this hook on with. +const NAMES: [&str; 2] = ["dynamic_rate_limiter_v3", "model_capacity"]; + +/// A priority's share of a group's capacity, as `priority_reservation` writes it. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Share { + Percent(f64), + Rpm(f64), + Tpm(f64), +} + +impl Share { + /// The share as a fraction of the group, as the proxy's `convert_priority_to_percent`. + fn fraction(self, capacity: &Capacity) -> f64 { + match self { + Self::Percent(value) => value, + Self::Rpm(value) => match capacity.rpm { + Some(rpm) if rpm > 0 => value / rpm as f64, + _ => value, + }, + Self::Tpm(value) => match capacity.tpm { + Some(tpm) if tpm > 0 => value / tpm as f64, + _ => value, + }, + } + } +} + +/// Which tokens count: `general_settings.token_rate_limit_type`. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum TokenCount { + #[default] + Total, + Input, + Output, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CapacitySettings { + pub window: u64, + pub fail_closed: bool, + /// `litellm_settings.priority_reservation`, in config order. + pub reservation: Option>, + /// `priority_reservation_settings.default_priority`: the share of keys without a listed + /// priority, together. + pub default_priority: f64, + /// `priority_reservation_settings.saturation_threshold`. + pub saturation_threshold: f64, + /// `priority_reservation_settings.saturation_check_cache_ttl`: how long a group's + /// saturation is reused before it is read again. + pub saturation_ttl: Duration, + pub tokens: TokenCount, +} + +impl Default for CapacitySettings { + fn default() -> Self { + let rate_limits = RateLimitSettings::default(); + Self { + window: rate_limits.window, + fail_closed: rate_limits.fail_closed, + reservation: None, + default_priority: 0.25, + saturation_threshold: 0.5, + saturation_ttl: Duration::from_secs(60), + tokens: TokenCount::Total, + } + } +} + +impl CapacitySettings { + /// `None` when the config's `callbacks` do not list the hook. + pub fn from_config(config: &Config) -> Option { + let litellm = &config.litellm_settings; + let listed = litellm.callbacks.iter().flatten().any(|callback| { + callback + .as_str() + .is_some_and(|name| NAMES.contains(&name.trim())) + }); + if !listed { + return None; + } + let rate_limits = RateLimitSettings::from_config(config); + let number = |value: &litellm_config::Value| match value { + litellm_config::Value::Number(number) => number.as_f64(), + litellm_config::Value::String(text) => text.trim().parse().ok(), + _ => None, + }; + let reservation = match litellm.additional_fields.get("priority_reservation") { + Some(litellm_config::Value::Mapping(mapping)) => Some( + mapping + .iter() + .filter_map(|(priority, share)| { + let priority = match priority { + litellm_config::Value::String(text) => text.clone(), + litellm_config::Value::Number(number) => number.to_string(), + _ => return None, + }; + let share = match share { + litellm_config::Value::Mapping(share) => { + let value = share.get("value").and_then(number).unwrap_or(1.0); + match share.get("type").and_then(litellm_config::Value::as_str) { + Some("rpm") => Share::Rpm(value), + Some("tpm") => Share::Tpm(value), + _ => Share::Percent(value), + } + } + other => Share::Percent(number(other)?), + }; + Some((priority, share)) + }) + .collect(), + ), + _ => None, + }; + let settings = litellm + .additional_fields + .get("priority_reservation_settings"); + let setting = |name: &str| { + settings + .and_then(|settings| settings.get(name)) + .and_then(number) + }; + let defaults = Self::default(); + Some(Self { + window: rate_limits.window, + fail_closed: rate_limits.fail_closed, + reservation, + default_priority: setting("default_priority").unwrap_or(defaults.default_priority), + saturation_threshold: setting("saturation_threshold") + .unwrap_or(defaults.saturation_threshold), + saturation_ttl: setting("saturation_check_cache_ttl") + .filter(|seconds| *seconds >= 0.0) + .map_or(defaults.saturation_ttl, Duration::from_secs_f64), + tokens: match config + .general_settings + .additional_fields + .get("token_rate_limit_type") + .and_then(litellm_config::Value::as_str) + { + Some("input") => TokenCount::Input, + Some("output") => TokenCount::Output, + _ => TokenCount::Total, + }, + }) + } +} + +/// A model group's requests and tokens per window: its deployments' `rpm` and `tpm`, summed. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Capacity { + pub rpm: Option, + pub tpm: Option, +} + +/// Each configured group's capacity, from `rpm`/`tpm` on the deployment, its `litellm_params` +/// or its `model_info`, in that order, as the proxy's router sums them. +pub fn capacities(config: &Config) -> HashMap { + let mut groups: HashMap = HashMap::new(); + for model in &config.model_list { + let configured = |field: &str| -> Option { + let from_value = |value: &litellm_config::Value| match value { + litellm_config::Value::Number(number) => number.as_f64(), + litellm_config::Value::String(text) => text.trim().parse().ok(), + _ => None, + }; + let spelled = |value: &litellm_config::Spelled| match value { + litellm_config::Spelled::Value(value) => Some(*value), + litellm_config::Spelled::Text(text) => text.trim().parse().ok(), + }; + let params = match field { + "tpm" => model.litellm_params.tpm.as_ref(), + _ => model.litellm_params.rpm.as_ref(), + }; + model + .additional_fields + .get(field) + .and_then(from_value) + .or_else(|| params.and_then(spelled)) + .or_else(|| model.model_info.get(field).and_then(from_value)) + }; + let group = groups.entry(model.model_name.clone()).or_default(); + let add = |total: &mut Option, value: Option| { + if let Some(value) = value { + *total = Some(total.unwrap_or(0) + value as i64); + } + }; + add(&mut group.tpm, configured("tpm")); + add(&mut group.rpm, configured("rpm")); + } + groups +} + +/// The pool a request draws on and its share of the group. +#[derive(Clone, Debug, PartialEq)] +struct Pool { + name: String, + share: f64, +} + +/// The `model_capacity` hook. +pub struct ModelCapacity { + settings: CapacitySettings, + groups: HashMap, + store: Arc, + saturation: Mutex, Instant)>>, +} + +impl ModelCapacity { + pub fn new( + settings: CapacitySettings, + groups: HashMap, + redis: Option>, + ) -> Self { + Self { + store: Arc::new(Store::new(redis, settings.window, settings.fail_closed)), + settings, + groups, + saturation: Mutex::default(), + } + } + + /// The hook for a gateway whose config lists it, else `None`. + pub fn for_gateway(config: &Config, redis: Option>) -> Option { + Some(Self::new( + CapacitySettings::from_config(config)?, + capacities(config), + redis, + )) + } + + /// The pool the proxy's `_get_priority_allocation` picks, with its normalized share. + fn pool(&self, group: &str, capacity: &Capacity, priority: Option<&str>) -> Pool { + let explicit = self.settings.reservation.as_ref().and_then(|reservation| { + let priority = priority?; + reservation + .iter() + .any(|(name, _)| name == priority) + .then(|| { + let fractions: Vec<(&str, f64)> = reservation + .iter() + .map(|(name, share)| (name.as_str(), share.fraction(capacity))) + .collect(); + let total: f64 = fractions.iter().map(|(_, fraction)| fraction).sum(); + let fraction = fractions + .iter() + .find(|(name, _)| *name == priority) + .map_or(0.0, |(_, fraction)| *fraction); + ( + priority, + if total > 1.0 { + fraction / total + } else { + fraction + }, + ) + }) + }); + match explicit { + Some((priority, share)) => Pool { + name: format!("{group}:{priority}"), + share, + }, + None => Pool { + name: format!("{group}:default_pool"), + share: self.settings.default_priority, + }, + } + } + + /// A counter's value, reused for `saturation_check_cache_ttl`. + async fn cached(&self, prefix: &str, kind: Kind) -> Option { + let key = format!("{prefix}:{}", kind.name()); + if let Some((value, read_at)) = self + .saturation + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&key) + && read_at.elapsed() < self.settings.saturation_ttl + { + return *value; + } + let value = self.store.read(prefix, kind).await; + let mut cache = self.saturation.lock().unwrap_or_else(|e| e.into_inner()); + if cache.len() >= 10_000 { + cache.clear(); + } + cache.insert(key, (value, Instant::now())); + value + } + + /// How full the group is: its busiest of requests and tokens against capacity. + async fn saturation(&self, prefix: &str, capacity: &Capacity) -> f64 { + let mut saturation = 0.0_f64; + for (kind, limit) in [(Kind::Requests, capacity.rpm), (Kind::Tokens, capacity.tpm)] { + if let Some(limit) = limit.filter(|limit| *limit > 0) + && let Some(count) = self.cached(prefix, kind).await + { + saturation = saturation.max(count as f64 / limit as f64); + } + } + saturation + } + + fn unavailable(&self, request: &RequestFacts) -> Refusal { + crate::rate_limits::unavailable(request.route) + } +} + +/// `team_metadata.priority`, else the key's `metadata.priority`. +fn priority(key: Option<&KeyFacts>) -> Option { + let key = key?; + let read = |value: &Value| match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + _ => None, + }; + key.team_metadata + .get("priority") + .and_then(read) + .or_else(|| key.metadata.get("priority").and_then(read)) +} + +fn limit_text(limit: Option) -> String { + limit.map_or_else(|| "not configured".to_owned(), |limit| limit.to_string()) +} + +/// `{:.%}` as Python formats a fraction. +fn percent(fraction: f64, places: usize) -> String { + format!("{:.*}%", places, fraction * 100.0) +} + +impl Hook for ModelCapacity { + fn name(&self) -> &'static str { + "model_capacity" + } + + fn reads(&self) -> Facts { + Facts::KEY | Facts::USAGE + } + + fn writes(&self) -> Targets { + Targets::ADMISSION | Targets::RESPONSE_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + let priority = priority(request.key.as_deref()); + let priority_header = priority.as_deref().unwrap_or("default"); + if !request.stream { + if priority_header.chars().all(|char| (char as u32) < 256) + && let Ok(value) = HeaderValue::from_str(priority_header) + { + edit.add_response_header(HeaderName::from_static("x-litellm-priority"), value); + } + edit.add_response_header( + HeaderName::from_static("x-litellm-rate-limiter-version"), + HeaderValue::from_static("v3"), + ); + } + let group = &*request.requested_model; + let Some(capacity) = self.groups.get(group).copied() else { + return Ok(None); + }; + if capacity.rpm.is_none() && capacity.tpm.is_none() { + return Ok(None); + } + let model_prefix = format!("{{model_saturation_check:{group}}}"); + let saturation = self.saturation(&model_prefix, &capacity).await; + let pool = self.pool(group, &capacity, priority.as_deref()); + let pool_prefix = format!("{{priority_model:{}}}", pool.name); + + // The group counts every request against its whole capacity; a zero limit is off. + let mut pairs = Vec::new(); + let push = |pairs: &mut Vec<(&'static str, Pair)>, + key: &'static str, + prefix: &str, + limits: [Option; 2]| { + for (kind, limit, increment) in + [(Kind::Requests, limits[0], 1), (Kind::Tokens, limits[1], 0)] + { + if let Some(limit) = limit { + pairs.push(( + key, + Pair { + prefix: prefix.to_owned(), + kind, + limit, + increment, + }, + )); + } + } + }; + push( + &mut pairs, + "model_saturation_check", + &model_prefix, + [ + capacity.rpm.filter(|rpm| *rpm != 0), + capacity.tpm.filter(|tpm| *tpm != 0), + ], + ); + let share = |limit: Option| limit.map(|limit| (limit as f64 * pool.share) as i64); + let pool_limits = [share(capacity.rpm), share(capacity.tpm)]; + let reserving = + self.settings.reservation.is_some() && pool_limits.iter().any(Option::is_some); + let enforced = reserving && saturation >= self.settings.saturation_threshold; + if enforced { + push(&mut pairs, "priority_model", &pool_prefix, pool_limits); + } + let checks: Vec = pairs.iter().map(|(_, pair)| pair.clone()).collect(); + let (reserved, shared) = self + .store + .check_and_add(&checks) + .await + .map_err(|Unavailable| self.unavailable(request))?; + let counts = match reserved { + Reserved::Counts(counts) => counts, + Reserved::Over { index, current } => { + let (key, pair) = &pairs[index]; + let remaining = pair.limit - current; + let kind = pair.kind.name(); + let priority_text = priority.as_deref().unwrap_or("None"); + let (message, saturation_header) = match *key { + "model_saturation_check" => ( + format!( + "Model capacity reached for {group}. Priority: {priority_text}, Rate limit type: {kind}, Model TPM: {}, Model RPM: {}, Remaining: {remaining}", + limit_text(capacity.tpm), + limit_text(capacity.rpm), + ), + None, + ), + _ => ( + format!( + "Priority-based rate limit exceeded. Model: {group}, Priority: {priority_text}, Rate limit type: {kind}, Model TPM: {}, Model RPM: {}, Remaining: {remaining}, Model saturation: {}", + limit_text(capacity.tpm), + limit_text(capacity.rpm), + percent(saturation, 1), + ), + Some(percent(saturation, 2)), + ), + }; + let mut headers = vec![ + ("retry-after", self.settings.window.to_string()), + ("rate_limit_type", kind.to_owned()), + ("x-litellm-priority", priority_header.to_owned()), + ]; + if let Some(saturation) = saturation_header { + headers.push(("x-litellm-saturation", saturation)); + } + return Err(crate::rate_limits::throttled( + request.route, + message, + headers, + )); + } + }; + let mut statuses: Vec<(&str, Kind, i64, i64)> = pairs + .iter() + .zip(&counts) + .map(|((key, pair), count)| (*key, pair.kind, pair.limit, pair.limit - count)) + .collect(); + // Below the threshold the pool's requests are only counted, so its share is known + // once it is enforced. + if reserving + && !enforced + && let Some(rpm) = pool_limits[0] + && let Ok(counts) = self + .store + .requests(std::slice::from_ref(&pool_prefix)) + .await + { + statuses.push(("priority_model", Kind::Requests, rpm, rpm - counts[0])); + } + if !request.stream { + for (key, kind, limit, remaining) in statuses { + let kind = kind.name(); + for (name, value) in [ + (format!("x-ratelimit-{key}-remaining-{kind}"), remaining), + (format!("x-ratelimit-{key}-limit-{kind}"), limit), + ] { + if let Ok(name) = HeaderName::from_bytes(name.as_bytes()) { + edit.add_response_header(name, HeaderValue::from(value)); + } + } + } + } + Ok(Some(Box::new(Charge { + store: self.store.clone(), + prefixes: [model_prefix, pool_prefix], + tokens: self.settings.tokens, + shared, + }) as Box)) + }) + } +} + +/// Adds a successful response's tokens to the group and its pool. +struct Charge { + store: Arc, + prefixes: [String; 2], + tokens: TokenCount, + shared: bool, +} + +impl Settle for Charge { + fn settle(self: Box, record: &Record) { + if !(200..300).contains(&record.outcome.status) { + return; + } + let Some(usage) = record.outcome.usage else { + return; + }; + let tokens = match self.tokens { + TokenCount::Total => usage.total_tokens(), + TokenCount::Input => usage.prompt_tokens(), + TokenCount::Output => usage.output_tokens, + }; + let Ok(tokens) = i64::try_from(tokens) else { + return; + }; + if tokens == 0 { + return; + } + let Charge { + store, + prefixes, + shared, + .. + } = *self; + tokio::spawn(async move { + let deltas = prefixes + .into_iter() + .map(|prefix| (prefix, tokens)) + .collect(); + store.charge(deltas, shared).await; + }); + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn capacity_hook(reservation: Option>) -> ModelCapacity { + ModelCapacity::new( + CapacitySettings { + reservation, + ..CapacitySettings::default() + }, + HashMap::new(), + None, + ) + } + + #[rstest] + fn pools_follow_the_proxy_allocation() { + let capacity = Capacity { + rpm: Some(100), + tpm: Some(1_000), + }; + let hook = capacity_hook(Some(vec![ + ("prod".into(), Share::Percent(0.6)), + ("dev".into(), Share::Rpm(80.0)), + ])); + // 0.6 + 0.8 is over-allocated, so both shrink to their proportion. + let prod = hook.pool("m", &capacity, Some("prod")); + assert_eq!(prod.name, "m:prod"); + assert!((prod.share - 0.6 / 1.4).abs() < 1e-12); + assert_eq!( + hook.pool("m", &capacity, Some("other")), + Pool { + name: "m:default_pool".into(), + share: 0.25 + } + ); + assert_eq!( + capacity_hook(None).pool("m", &capacity, Some("prod")).name, + "m:default_pool" + ); + } + + #[rstest] + #[case(0.625, 1, "62.5%")] + #[case(0.5, 2, "50.00%")] + fn formats_percentages_like_python( + #[case] fraction: f64, + #[case] places: usize, + #[case] expected: &str, + ) { + assert_eq!(percent(fraction, places), expected); + } + + #[rstest] + fn capacities_sum_deployments() { + let config = Config::from_yaml( + "model_list: + - model_name: a + litellm_params: {model: openai/x, rpm: 10, tpm: 1000} + - model_name: a + rpm: 5 + litellm_params: {model: openai/y, rpm: 99} + model_info: {tpm: 500} + - model_name: b + litellm_params: {model: openai/z} +", + ) + .unwrap(); + let groups = capacities(&config); + assert_eq!( + groups["a"], + Capacity { + rpm: Some(15), + tpm: Some(1_500) + } + ); + assert_eq!(groups["b"], Capacity::default()); + } +} diff --git a/litellm-rust/crates/gateway/src/rate_limits.rs b/litellm-rust/crates/gateway/src/rate_limits.rs index 2b9ca57cbbc..2948e3bc842 100644 --- a/litellm-rust/crates/gateway/src/rate_limits.rs +++ b/litellm-rust/crates/gateway/src/rate_limits.rs @@ -225,14 +225,14 @@ impl Descriptor { } #[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum Kind { +pub(crate) enum Kind { Requests, Tokens, Parallel, } impl Kind { - fn name(self) -> &'static str { + pub(crate) fn name(self) -> &'static str { match self { Self::Requests => "requests", Self::Tokens => "tokens", @@ -250,15 +250,25 @@ struct Status { remaining: i64, } -/// The token reservation went over on the descriptor at `index`, which held `current`. +/// A counter to check and add to: `{prefix}:requests` or `{prefix}:tokens` in its window. +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct Pair { + pub(crate) prefix: String, + pub(crate) kind: Kind, + pub(crate) limit: i64, + pub(crate) increment: i64, +} + +/// Every pair's count after adding, or the pair at `index` that would go over, which held +/// `current`. #[derive(Debug, PartialEq)] -enum Reserved { +pub(crate) enum Reserved { Counts(Vec), Over { index: usize, current: i64 }, } /// Redis could not be used and the settings say to refuse rather than fall back. -struct Unavailable; +pub(crate) struct Unavailable; #[derive(Debug)] struct Window { @@ -307,27 +317,24 @@ impl Local { .collect() } - fn reserve( - &mut self, - scopes: &[(String, i64)], - increment: i64, - now: u64, - size: u64, - ) -> Reserved { - for (index, (prefix, limit)) in scopes.iter().enumerate() { - let current = match self.windows.get(prefix) { - Some(window) if now.saturating_sub(window.start) < size => window.tokens, + fn check_and_add(&mut self, pairs: &[Pair], now: u64, size: u64) -> Reserved { + for (index, pair) in pairs.iter().enumerate() { + let current = match self.windows.get(&pair.prefix) { + Some(window) if now.saturating_sub(window.start) < size => match pair.kind { + Kind::Requests => window.requests, + _ => window.tokens, + }, _ => 0, }; - if blocked(current, increment, *limit) { + if blocked(current, pair.increment, pair.limit) { return Reserved::Over { index, current }; } } Reserved::Counts( - scopes + pairs .iter() - .map(|(prefix, _)| { - let (window, expired) = self.window(prefix, now, size); + .map(|pair| { + let (window, expired) = self.window(&pair.prefix, now, size); if expired { *window = Window { start: now, @@ -335,13 +342,26 @@ impl Local { tokens: 0, }; } - window.tokens += increment; - window.tokens + let count = match pair.kind { + Kind::Requests => &mut window.requests, + _ => &mut window.tokens, + }; + *count += pair.increment; + *count }) .collect(), ) } + /// A counter's value in its current window. + fn read(&self, prefix: &str, kind: Kind, now: u64, size: u64) -> Option { + let window = self.windows.get(prefix)?; + (now.saturating_sub(window.start) < size).then_some(match kind { + Kind::Requests => window.requests, + _ => window.tokens, + }) + } + fn charge(&mut self, deltas: &[(String, i64)]) { for (prefix, delta) in deltas { if let Some(window) = self.windows.get_mut(prefix) { @@ -380,7 +400,7 @@ fn blocked(current: i64, increment: i64, limit: i64) -> bool { } /// Where the counters live. -struct Store { +pub(crate) struct Store { redis: Option>, local: Mutex, window: u64, @@ -388,6 +408,15 @@ struct Store { } impl Store { + pub(crate) fn new(redis: Option>, window: u64, fail_closed: bool) -> Self { + Self { + redis, + local: Mutex::default(), + window, + fail_closed, + } + } + fn connection(&self) -> Option<(ConnectionManager, &Coordination)> { let redis = self.redis.as_deref()?; Some((redis.connection()?, redis)) @@ -414,7 +443,7 @@ impl Store { } /// Counts the request on each prefix's window and returns the counts. - async fn requests(&self, prefixes: &[String]) -> Result, Unavailable> { + pub(crate) async fn requests(&self, prefixes: &[String]) -> Result, Unavailable> { let now = epoch_seconds(); if self.redis.is_some() { let error = match self.connection() { @@ -500,21 +529,43 @@ impl Store { &self, scopes: &[(String, i64)], increment: i64, + ) -> Result<(Reserved, bool), Unavailable> { + let pairs: Vec = scopes + .iter() + .map(|(prefix, limit)| Pair { + prefix: prefix.clone(), + kind: Kind::Tokens, + limit: *limit, + increment, + }) + .collect(); + self.check_and_add(&pairs).await + } + + /// Adds every pair's increment, or none when one would go over its limit. `true` when the + /// counters are in Redis. + pub(crate) async fn check_and_add( + &self, + pairs: &[Pair], ) -> Result<(Reserved, bool), Unavailable> { if self.redis.is_some() { let error = match self.connection() { Some((mut connection, redis)) => { let mut command = redis::cmd("EVAL"); - command.arg(RESERVE).arg(scopes.len() * 2); - for (prefix, _) in scopes { + command.arg(RESERVE).arg(pairs.len() * 2); + for pair in pairs { command - .arg(redis.namespaced(&format!("{prefix}:window"))) - .arg(redis.namespaced(&format!("{prefix}:tokens"))); + .arg(redis.namespaced(&format!("{}:window", pair.prefix))) + .arg(redis.namespaced(&format!( + "{}:{}", + pair.prefix, + pair.kind.name() + ))); } - for (_, limit) in scopes { + for pair in pairs { command - .arg(*limit) - .arg(increment) + .arg(pair.limit) + .arg(pair.increment) .arg(self.window) .arg(self.window); } @@ -528,7 +579,7 @@ impl Store { )); } Ok(reply) => match reply.as_slice() { - [1, index, current, _] if (1..=scopes.len() as i64).contains(index) => { + [1, index, current, _] if (1..=pairs.len() as i64).contains(index) => { return Ok(( Reserved::Over { index: (*index - 1) as usize, @@ -548,13 +599,34 @@ impl Store { } Ok(( self.local() - .reserve(scopes, increment, epoch_seconds(), self.window), + .check_and_add(pairs, epoch_seconds(), self.window), false, )) } + /// A counter's current value, `None` when it is not set. + pub(crate) async fn read(&self, prefix: &str, kind: Kind) -> Option { + if let Some((mut connection, redis)) = self.connection() { + let key = redis.namespaced(&format!("{prefix}:{}", kind.name())); + match redis::cmd("GET") + .arg(key) + .query_async::>(&mut connection) + .await + { + Ok(value) => { + return value + .and_then(|value| value.parse::().ok()) + .map(|value| value as i64); + } + Err(error) => tracing::warn!(%error, "reading a rate limit counter failed"), + } + } + self.local() + .read(prefix, kind, epoch_seconds(), self.window) + } + /// Adds each delta to its prefix's tokens. - async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) { + pub(crate) async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) { if shared && let Some((mut connection, redis)) = self.connection() { let mut command = redis::cmd("EVAL"); command.arg(CHARGE).arg(deltas.len()); @@ -711,12 +783,7 @@ impl RateLimiter { database: Option, ) -> Self { Self { - store: Arc::new(Store { - redis, - local: Mutex::default(), - window: settings.window, - fail_closed: settings.fail_closed, - }), + store: Arc::new(Store::new(redis, settings.window, settings.fail_closed)), settings, end_users: database.map(EndUsers::new), } @@ -801,22 +868,15 @@ impl RateLimiter { status.limit, status.remaining.max(0), ); - let mut refusal = match request.route { - Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(), - Route::ChatCompletions => { - let mut refusal = Refusal::new(429, "throttling_error", message).without_param(); - for (name, value) in [ - ("retry-after", self.settings.window.to_string()), - ("rate_limit_type", status.kind.name().to_owned()), - ("reset_at", reset_at), - ] { - if let Ok(value) = HeaderValue::from_str(&value) { - refusal = refusal.with_header(HeaderName::from_static(name), value); - } - } - refusal - } - }; + let mut refusal = throttled( + request.route, + message, + vec![ + ("retry-after", self.settings.window.to_string()), + ("rate_limit_type", status.kind.name().to_owned()), + ("reset_at", reset_at), + ], + ); for (name, value) in key_limit_headers(request.key.as_deref()) { refusal = refusal.with_header(name, value); } @@ -824,11 +884,37 @@ impl RateLimiter { } fn unavailable(&self, request: &RequestFacts) -> Refusal { - match request.route { - Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(), - Route::ChatCompletions => { - Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param() + unavailable(request.route) + } +} + +/// The proxy's 429 from one of its own limiters: on Chat Completions its error body with the +/// limiter's headers, on `/v1/messages` Anthropic's envelope, where the proxy drops them. +pub(crate) fn throttled( + route: Route, + message: String, + headers: Vec<(&'static str, String)>, +) -> Refusal { + match route { + Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(), + Route::ChatCompletions => { + let mut refusal = Refusal::new(429, "throttling_error", message).without_param(); + for (name, value) in headers { + if let Ok(value) = HeaderValue::from_str(&value) { + refusal = refusal.with_header(HeaderName::from_static(name), value); + } } + refusal + } + } +} + +/// The proxy's 503 when counters cannot be verified and enforcement fails closed. +pub(crate) fn unavailable(route: Route) -> Refusal { + match route { + Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(), + Route::ChatCompletions => { + Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param() } } } @@ -1150,35 +1236,35 @@ mod tests { assert_eq!(estimate(&request, Some(&key), "other", 100_000).tokens, 10); } + fn reserve(local: &mut Local, increment: i64, now: u64) -> Reserved { + let pair = Pair { + prefix: "{api_key:h}".into(), + kind: Kind::Tokens, + limit: 10, + increment, + }; + local.check_and_add(&[pair], now, 60) + } + #[rstest] fn local_windows_count_reserve_and_reset() { let mut local = Local::default(); let prefixes = vec!["{api_key:h}".to_owned()]; assert_eq!(local.requests(&prefixes, 100, 60), [1]); assert_eq!(local.requests(&prefixes, 159, 60), [2]); - let scopes = vec![("{api_key:h}".to_owned(), 10)]; + assert_eq!(reserve(&mut local, 6, 159), Reserved::Counts(vec![6])); assert_eq!( - local.reserve(&scopes, 6, 159, 60), - Reserved::Counts(vec![6]) - ); - assert_eq!( - local.reserve(&scopes, 6, 159, 60), + reserve(&mut local, 6, 159), Reserved::Over { index: 0, current: 6 } ); local.charge(&[("{api_key:h}".to_owned(), -4)]); - assert_eq!( - local.reserve(&scopes, 8, 159, 60), - Reserved::Counts(vec![10]) - ); + assert_eq!(reserve(&mut local, 8, 159), Reserved::Counts(vec![10])); // The window started at 100 is over at 160: requests and tokens start again. assert_eq!(local.requests(&prefixes, 160, 60), [1]); - assert_eq!( - local.reserve(&scopes, 10, 160, 60), - Reserved::Counts(vec![10]) - ); + assert_eq!(reserve(&mut local, 10, 160), Reserved::Counts(vec![10])); } #[rstest] diff --git a/litellm-rust/crates/gateway/tests/model_capacity.rs b/litellm-rust/crates/gateway/tests/model_capacity.rs new file mode 100644 index 00000000000..0dddd6ad901 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/model_capacity.rs @@ -0,0 +1,306 @@ +//! The model capacity hook on real requests: a group held to its deployments' summed capacity, +//! priorities held to their share, refusals and headers as the proxy sends them, and counters +//! shared through Redis. Set `LITELLM_TEST_REDIS_URL` to also run the shared case. + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{HeaderMap, Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway::{ + coordination::{Coordination, RedisSettings}, + model_capacity::{CapacitySettings, ModelCapacity}, +}; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +/// 3 input and 4 output tokens: 7 count against the group. +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; + +async fn upstream() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(|_: Request| async { + ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response() + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + format!("http://{address}") +} + +struct Keys(HashMap); + +impl KeyLookup for Keys { + fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> { + Box::pin(async move { Ok(self.0.get(hash).cloned()) }) + } +} + +/// Two deployments of `claude`, each with `limits` (such as `rpm: 1`). +async fn gateway(limits: &str, settings: &str, redis: Option>) -> Router { + let address = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: {{model: anthropic/claude-x, api_base: {address}, api_key: sk-a, {limits}}} + - model_name: claude + litellm_params: {{model: anthropic/claude-y, api_base: {address}, api_key: sk-b, {limits}}} +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +litellm_settings: + callbacks: [dynamic_rate_limiter_v3] +{settings}" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let hook = ModelCapacity::for_gateway(&config, redis).unwrap(); + litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(hook)], Vec::new()) + .await + .unwrap(); + let mut metadata = serde_json::Map::new(); + metadata.insert("priority".into(), json!("prod")); + let key = KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: None, + profile: Arc::new(KeyProfile { + metadata: Arc::new(metadata), + ..KeyProfile::default() + }), + }; + let keys: Arc = Arc::new(Keys(HashMap::from([( + KeyHash::from_token(&SecretValue::new("sk-prod")), + key, + )]))); + litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys)) +} + +struct Reply { + status: u16, + headers: HeaderMap, + body: Value, +} + +async fn send(app: &Router, path: &str, token: &str) -> Reply { + let body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {token}")) + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + Reply { + status, + headers, + body: serde_json::from_slice(&body).unwrap_or(Value::Null), + } +} + +fn header<'a>(reply: &'a Reply, name: &str) -> Option<&'a str> { + reply.headers.get(name).map(|value| value.to_str().unwrap()) +} + +async fn settled() { + tokio::time::sleep(Duration::from_millis(100)).await; +} + +#[rstest] +#[tokio::test] +async fn a_group_is_held_to_its_deployments_capacity() { + let app = gateway("rpm: 1", "", None).await; + let first = send(&app, "/v1/messages", "sk-master").await; + assert_eq!(first.status, 200); + assert_eq!( + header( + &first, + "x-ratelimit-model_saturation_check-remaining-requests" + ), + Some("1") + ); + assert_eq!( + header(&first, "x-ratelimit-model_saturation_check-limit-requests"), + Some("2") + ); + assert_eq!(header(&first, "x-litellm-rate-limiter-version"), Some("v3")); + assert_eq!(header(&first, "x-litellm-priority"), Some("default")); + assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200); + + let message = "Model capacity reached for claude. Priority: None, Rate limit type: requests, Model TPM: not configured, Model RPM: 2, Remaining: 0"; + let refused = send(&app, "/v1/chat/completions", "sk-master").await; + assert_eq!(refused.status, 429); + assert_eq!( + refused.body, + json!({"error": {"message": message, "type": "throttling_error", "param": null, "code": "429"}}) + ); + assert_eq!(header(&refused, "retry-after"), Some("60")); + assert_eq!(header(&refused, "rate_limit_type"), Some("requests")); + assert_eq!(header(&refused, "x-litellm-priority"), Some("default")); + let refused = send(&app, "/v1/messages", "sk-master").await; + assert_eq!( + refused.body, + json!({"type": "error", "error": {"type": "rate_limit_error", "message": message}}) + ); + assert_eq!(header(&refused, "retry-after"), None); +} + +#[rstest] +#[tokio::test] +async fn tokens_count_after_the_response() { + // 2 x 5 tokens of capacity; each response uses 7. + let app = gateway("tpm: 5", "", None).await; + assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200); + settled().await; + assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200); + settled().await; + let refused = send(&app, "/v1/messages", "sk-master").await; + assert_eq!(refused.status, 429); + assert_eq!( + refused.body["error"]["message"], + "Model capacity reached for claude. Priority: None, Rate limit type: tokens, Model TPM: 10, Model RPM: not configured, Remaining: -4" + ); +} + +#[rstest] +#[tokio::test] +async fn busy_groups_hold_priorities_to_their_share() { + // 20 requests a window; `prod` gets a fifth once the group is busy, here from the start. + let app = gateway( + "rpm: 10", + " priority_reservation: {prod: 0.1}\n priority_reservation_settings: {saturation_threshold: 0.0}\n", + None, + ) + .await; + assert_eq!( + send(&app, "/v1/chat/completions", "sk-prod").await.status, + 200 + ); + let second = send(&app, "/v1/chat/completions", "sk-prod").await; + assert_eq!(second.status, 200); + assert_eq!(header(&second, "x-litellm-priority"), Some("prod")); + assert_eq!( + header(&second, "x-ratelimit-priority_model-remaining-requests"), + Some("0") + ); + let refused = send(&app, "/v1/chat/completions", "sk-prod").await; + assert_eq!(refused.status, 429); + assert_eq!( + refused.body["error"]["message"], + "Priority-based rate limit exceeded. Model: claude, Priority: prod, Rate limit type: requests, Model TPM: not configured, Model RPM: 20, Remaining: 0, Model saturation: 0.0%" + ); + assert_eq!(header(&refused, "x-litellm-saturation"), Some("0.00%")); + // Other callers share the default pool and the rest of the group. + assert_eq!( + send(&app, "/v1/chat/completions", "sk-master").await.status, + 200 + ); +} + +#[rstest] +fn the_hook_is_on_only_when_listed() { + let config = Config::from_yaml("model_list: []").unwrap(); + assert_eq!(CapacitySettings::from_config(&config), None); + let config = + Config::from_yaml("model_list: []\nlitellm_settings:\n callbacks: model_capacity\n") + .unwrap(); + assert_eq!( + CapacitySettings::from_config(&config), + Some(CapacitySettings::default()) + ); +} + +async fn redis(url: &str, namespace: &str) -> Arc { + let coordination = + Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap(); + let router: Arc = Arc::new(NoRouter( + litellm_gateway_inference::ModelRouter::from_config( + &Config::from_yaml("model_list: []").unwrap(), + ) + .unwrap(), + )); + coordination.start(router, None, tokio_util::sync::CancellationToken::new()); + for _ in 0..100 { + if coordination.connected() { + return coordination; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("Redis did not connect"); +} + +struct NoRouter(litellm_gateway_inference::ModelRouter); + +impl litellm_gateway::coordination::RouterHandle for NoRouter { + fn router(&self) -> &litellm_gateway_inference::ModelRouter { + &self.0 + } +} + +#[rstest] +#[tokio::test] +async fn replicas_share_the_group_counters() { + let Ok(url) = std::env::var("LITELLM_TEST_REDIS_URL") else { + return; + }; + let namespace = format!("model-capacity-{}", uuid::Uuid::new_v4()); + let limits = "rpm: 2, tpm: 1000"; + let first = gateway(limits, "", Some(redis(&url, &namespace).await)).await; + let second = gateway(limits, "", Some(redis(&url, &namespace).await)).await; + for app in [&first, &second, &first, &second] { + assert_eq!(send(app, "/v1/messages", "sk-master").await.status, 200); + } + settled().await; + assert_eq!(send(&second, "/v1/messages", "sk-master").await.status, 429); + + let mut connection = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let get = async |connection: &mut redis::aio::MultiplexedConnection, key: String| { + redis::cmd("GET") + .arg(key) + .query_async::>(connection) + .await + .unwrap() + }; + let group = format!("{namespace}:{{model_saturation_check:claude}}"); + // The refused request is not counted; four responses used 7 tokens each. + assert_eq!( + get(&mut connection, format!("{group}:requests")) + .await + .as_deref(), + Some("4") + ); + assert_eq!( + get(&mut connection, format!("{group}:tokens")) + .await + .as_deref(), + Some("28") + ); + let pool = format!("{namespace}:{{priority_model:claude:default_pool}}:tokens"); + assert_eq!(get(&mut connection, pool).await.as_deref(), Some("28")); +} diff --git a/litellm-rust/docs/customer-proxy/model-capacity.svg b/litellm-rust/docs/customer-proxy/model-capacity.svg new file mode 100644 index 00000000000..32a512c05bc --- /dev/null +++ b/litellm-rust/docs/customer-proxy/model-capacity.svg @@ -0,0 +1,52 @@ + + + + + + + Model capacity hook (dynamic_rate_limiter_v3): fleet-wide group capacity and priority shares + + + At startup + group capacity = sum of deployments' + rpm / tpm (deployment, litellm_params, + model_info) + priority_reservation shares: percent, + rpm or tpm; normalized when over 1 + on when callbacks list the hook + + + Per request (one Redis script, all or nothing) + priority: team metadata, else key metadata; pool <group>:<priority> or <group>:default_pool + saturation = max(requests / rpm, tokens / tpm), read and reused for saturation_check_cache_ttl + {model_saturation_check:<group>}: +1 request, tokens checked; always enforced + {priority_model:<pool>}: share x capacity, enforced at saturation >= threshold, + only counted below it + + + + Refusal: 429, the proxy's messages + "Model capacity reached for <group>. Priority: ..., Rate limit + type: ..., Model TPM: ..., Model RPM: ..., Remaining: ..." + "Priority-based rate limit exceeded. ... Model saturation: 62.5%" + chat: retry-after, rate_limit_type, x-litellm-priority, + x-litellm-saturation; messages: Anthropic envelope only + + + + Admitted: request runs + complete responses carry x-litellm-priority, + x-litellm-rate-limiter-version: v3 and + x-ratelimit-model_saturation_check-* / + x-ratelimit-priority_model-* + + + + Settle: a successful response's tokens (token_rate_limit_type: total, input or output) + added to the group's and the admitted pool's :tokens, with the window TTL + + + + Same keys and scripts as the Python proxy, under its namespace: Rust and Python replicas share one count + No Redis, or Redis unreachable: this replica's counters (fail_closed_rate_limit_enforcement: 503 instead) +