From 8fb0643e35da182f1340168b9dca063c3e12e787 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 14:24:02 -0700 Subject: [PATCH] perf(gateway): admit rate limits in one redis round trip and send scripts by hash (#47197) * feat(gateway): hooks and observers around inference requests, with usage and cost Adds the litellm-gateway-hooks crate: request facts built only when something reads them, on-path hooks that declare what they read and what they may change and run in config order, off-path observers on bounded queues, a usage tap for SSE and JSON responses, and per-deployment prices from the model cost map overlaid by config prices. The messages and chat handlers run through it when anything is installed; with nothing installed requests take the old path. The key store now also returns the key alias, user, team alias and key and team metadata, and router attempts carry the deployment id. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * fix(model-catalog): read token counts the cost map writes as whole floats The published cost map now has prompt_cache_min_tokens: 4096.0, which the typed u64 fields rejected, so the gateway dropped the whole map at startup and priced every request at zero. Counts accept any whole non-negative number, as Python reads any JSON number, and a test keeps the checked-in map parsing. * feat(gateway): write spend to the proxy's database from a native observer Adds a spend writer observer that records every gateway request in the proxy's own Postgres tables with the same statements the Python proxy uses: spend logs, key/user/team/team member/end user spend, and the daily user, team and end-user activity tables. Writes are batched on the proxy's timers and are additive, so replicas and Python proxies can write side by side. The database pool is now shared between virtual keys and the writer. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * feat(gateway): enforce key, team, user, member and end-user budgets on shared spend counters Adds a built-in budgets hook that runs the proxy's budget checks, in its order and with its messages, on the proxy's Redis spend:* counters. Each request reads and seeds the counters, checks them and reserves its maximum cost in one Redis script, then settles the reservation to its real cost off the request path, so replicas and Python proxies sharing the Redis enforce one budget. Without Redis the counters are kept per replica. The key store now reads user, team member and end-user budgets, and auth no longer checks budgets against stored spend. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * feat(gateway): enforce key, user, team, member and end-user rate limits on shared counters Adds a built-in rate limits hook with the proxy's v3 limiter semantics: requests per window, tokens per window and in-flight requests per key, on the proxy's own {scope:id}:window|requests|tokens|max_parallel_requests Redis keys and with its scripts, so replicas and Python proxies sharing the Redis enforce one limit. Tokens are reserved from the proxy's estimate and settled to usage when the response ends. Refusals carry the proxy's message, body and headers on each endpoint, and complete responses carry its x-ratelimit-* headers. Without Redis the counters are kept per replica. The key store now reads key, user, team and member rpm/tpm limits and the key's max_parallel_requests, and end users' budget limits are read with their budgets; keys with plain rate limits are no longer refused. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * test(gateway-auth): expect the narrowed unsupported rate limit message Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * perf(gateway): admit rate limits in one redis round trip and send scripts by hash A key with requests, in-flight and tokens limits made three sequential EVALs before its request was routed, each carrying the whole Lua script. One script now counts the request, takes the slot and reserves the tokens, stopping at the first limit it is over (giving the slot back when tokens are over), with the same keys and semantics as before. Every rate limit and budget script is sent with EVALSHA and loaded on NOSCRIPT. --------- Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/crates/gateway/Cargo.toml | 2 +- litellm-rust/crates/gateway/src/budgets.rs | 24 +- .../crates/gateway/src/rate_limits.rs | 630 +++++++++++++----- .../crates/gateway/tests/rate_limits.rs | 74 ++ 4 files changed, 535 insertions(+), 195 deletions(-) diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index b4f386f72ab..dfeeaaa8324 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -31,7 +31,7 @@ litellm-db.workspace = true litellm-http = { workspace = true, features = ["mcp"] } litellm-llms.workspace = true litellm-router.workspace = true -redis = { version = "1.7.0", default-features = false, features = ["connection-manager", "tokio-comp", "tokio-rustls-comp"] } +redis = { version = "1.7.0", default-features = false, features = ["connection-manager", "script", "tokio-comp", "tokio-rustls-comp"] } litellm-secrets.workspace = true litellm-tracing.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/gateway/src/budgets.rs b/litellm-rust/crates/gateway/src/budgets.rs index 7196bf132da..cf9b64db401 100644 --- a/litellm-rust/crates/gateway/src/budgets.rs +++ b/litellm-rust/crates/gateway/src/budgets.rs @@ -11,7 +11,7 @@ use std::{ collections::HashMap, - sync::{Arc, Mutex, OnceLock}, + sync::{Arc, LazyLock, Mutex, OnceLock}, time::{Duration, Instant}, }; @@ -117,6 +117,10 @@ end return #KEYS "#; +/// Sent by hash, loaded once per Redis. +static ADMIT_SCRIPT: LazyLock = LazyLock::new(|| redis::Script::new(ADMIT)); +static SETTLE_SCRIPT: LazyLock = LazyLock::new(|| redis::Script::new(SETTLE)); + #[derive(Clone, Debug, PartialEq)] pub struct BudgetSettings { /// `litellm_settings.budget_exceeded_status_code`, 422 by default. @@ -459,19 +463,15 @@ impl Counters { async fn settle(&self, keys: Vec, changes: Vec, shared: bool) { if shared && let Some((mut connection, redis)) = self.connection() { let keys: Vec = keys.iter().map(|key| redis.namespaced(key)).collect(); - let mut command = redis::cmd("EVAL"); - command - .arg(SETTLE) - .arg(keys.len()) - .arg(&keys) - .arg(self.ttl.as_secs()); + let mut command = SETTLE_SCRIPT.prepare_invoke(); + command.key(&keys).arg(self.ttl.as_secs()); for change in &changes { command .arg(change.floor.to_string()) .arg(change.delta.to_string()) .arg(change.fresh.to_string()); } - if let Err(error) = command.query_async::(&mut connection).await { + if let Err(error) = command.invoke_async::(&mut connection).await { // A counter left half-written would mislead every replica; drop it so the next // read reseeds it from the database, as the proxy does. tracing::warn!(%error, "updating spend counters failed; dropping them"); @@ -518,11 +518,9 @@ async fn admit_redis( estimate: f64, ttl: Duration, ) -> redis::RedisResult { - let mut command = redis::cmd("EVAL"); + let mut command = ADMIT_SCRIPT.prepare_invoke(); command - .arg(ADMIT) - .arg(keys.len()) - .arg(keys) + .key(keys) .arg(ttl.as_secs()) .arg(estimate.to_string()) .arg( @@ -553,7 +551,7 @@ async fn admit_redis( "0" }); } - let reply: Vec = command.query_async(connection).await?; + let reply: Vec = command.invoke_async(connection).await?; let text = |value: &redis::Value| -> Option { match value { redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), diff --git a/litellm-rust/crates/gateway/src/rate_limits.rs b/litellm-rust/crates/gateway/src/rate_limits.rs index cdf41a0e231..7ae2ab54e71 100644 --- a/litellm-rust/crates/gateway/src/rate_limits.rs +++ b/litellm-rust/crates/gateway/src/rate_limits.rs @@ -12,7 +12,7 @@ use std::{ collections::HashMap, - sync::{Arc, Mutex}, + sync::{Arc, LazyLock, Mutex}, time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; @@ -38,6 +38,149 @@ const MAX_OUTPUT_FLOOR: i64 = 1024; const MAX_LOCAL: usize = 100_000; const FAIL_CLOSED: &str = "Rate limit enforcement unavailable: request counters could not be verified against Redis, and fail_closed_rate_limit_enforcement is enabled, so the request was rejected to avoid exceeding the configured rate limit. Retry shortly."; +/// Admits a request in one round trip, with the proxy's semantics for each counter: counts it on +/// every `(window, requests)` pair, starting a window that is missing or over; unless one is over +/// its limit, takes an in-flight slot unless the gauge is full; then adds the token estimate to +/// every `(window, tokens)` pair, or to none (returning the slot) when one would go over. +/// +/// KEYS: each requests limit's window and counter, the in-flight set when the key caps it, then +/// each tokens limit's window and counter. ARGV: now, window size, the requests limits' count and +/// limits, the slot (`1` or `0`, limit, ttl, id), the token increment, then each tokens limit's +/// limit, ttl and window size. +/// +/// Replies `{code, request counts..., in_flight, rest...}`: code 0 admitted with each tokens +/// count, 1 over a requests limit (no in-flight), 2 the gauge was full, 3 over the tokens limit +/// at `index` holding `current`. +const ADMIT: &str = r#" +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) +local requests = tonumber(ARGV[3]) +local arg = 4 +local key = 1 +local reply = { 0 } +local over = false +for i = 1, requests do + local limit = tonumber(ARGV[arg]) + arg = arg + 1 + local window_key = KEYS[key] + local counter_key = KEYS[key + 1] + key = key + 2 + local window_start = redis.call('GET', window_key) + local counter + if not window_start or (now - tonumber(window_start)) >= window_size then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, 1) + redis.call('EXPIRE', window_key, window_size) + redis.call('EXPIRE', counter_key, window_size) + counter = 1 + else + counter = redis.call('INCR', counter_key) + if redis.call('TTL', counter_key) == -1 then + redis.call('EXPIRE', counter_key, window_size) + end + end + if counter > limit then + over = true + end + table.insert(reply, counter) +end +if over then + reply[1] = 1 + return reply +end + +local slot_key +local slot_id = ARGV[arg + 3] +local in_flight = 0 +if ARGV[arg] == '1' then + local limit = tonumber(ARGV[arg + 1]) + local slot_ttl = tonumber(ARGV[arg + 2]) + slot_key = KEYS[key] + key = key + 1 + local time = tonumber(redis.call('TIME')[1]) + redis.call('ZREMRANGEBYSCORE', slot_key, '-inf', time - slot_ttl) + in_flight = redis.call('ZCARD', slot_key) + if in_flight + 1 > limit then + reply[1] = 2 + table.insert(reply, in_flight) + return reply + end + redis.call('ZADD', slot_key, time, slot_id) + redis.call('EXPIRE', slot_key, slot_ttl) + in_flight = redis.call('ZCARD', slot_key) +end +table.insert(reply, in_flight) +arg = arg + 4 + +local increment = tonumber(ARGV[arg]) +arg = arg + 1 +local first = key +local count = (#KEYS - first + 1) / 2 +if count == 0 then + return reply +end +local time = tonumber(redis.call('TIME')[1]) +local state = {} +for i = 1, count do + local window_key = KEYS[first + (i - 1) * 2] + local counter_key = KEYS[first + (i - 1) * 2 + 1] + local limit = tonumber(ARGV[arg + (i - 1) * 3]) + local size = tonumber(ARGV[arg + (i - 1) * 3 + 2]) + local window_start = redis.call('GET', window_key) + local expired = (not window_start) or ((time - tonumber(window_start)) >= size) + local current = 0 + if not expired then + current = tonumber(redis.call('GET', counter_key) or 0) + end + local blocked + if increment > 0 then + blocked = current + increment > limit + else + blocked = current >= limit + end + if blocked then + if slot_key then + redis.call('ZREM', slot_key, slot_id) + end + reply[1] = 3 + table.insert(reply, i) + table.insert(reply, current) + return reply + end + state[i] = expired +end +local reset = {} +for i = 1, count do + local window_key = KEYS[first + (i - 1) * 2] + local counter_key = KEYS[first + (i - 1) * 2 + 1] + local ttl = tonumber(ARGV[arg + (i - 1) * 3 + 1]) + local size = tonumber(ARGV[arg + (i - 1) * 3 + 2]) + if state[i] then + if not reset[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset[window_key] = true + end + redis.call('SET', window_key, tostring(time)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(reply, increment) + else + local counter = redis.call('INCRBY', counter_key, increment) + if redis.call('TTL', counter_key) == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(reply, counter) + end +end +return reply +"#; + /// Counts a request on every `(window, requests)` pair, starting a window that is missing or /// over, and returns each window's start and count. The caller compares; the count always rises. const REQUESTS: &str = r#" @@ -134,21 +277,6 @@ end return results "#; -/// Takes an in-flight slot unless the gauge is full: `{0, in_flight}` or `{1, 1, in_flight, limit}`. -const ACQUIRE: &str = r#" -local now = tonumber(redis.call('TIME')[1]) -local limit = tonumber(ARGV[1]) -local slot_ttl = tonumber(ARGV[2]) -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now - slot_ttl) -local in_flight = redis.call('ZCARD', KEYS[1]) -if in_flight + 1 > limit then - return { 1, 1, in_flight, limit } -end -redis.call('ZADD', KEYS[1], now, ARGV[3]) -redis.call('EXPIRE', KEYS[1], slot_ttl) -return { 0, redis.call('ZCARD', KEYS[1]) } -"#; - /// Adds each token delta, giving a counter without an expiry the window's. const CHARGE: &str = r#" local results = {} @@ -163,6 +291,12 @@ end return results "#; +/// Sent by hash, loaded once per Redis. +static ADMIT_SCRIPT: LazyLock = LazyLock::new(|| redis::Script::new(ADMIT)); +static REQUESTS_SCRIPT: LazyLock = LazyLock::new(|| redis::Script::new(REQUESTS)); +static RESERVE_SCRIPT: LazyLock = LazyLock::new(|| redis::Script::new(RESERVE)); +static CHARGE_SCRIPT: LazyLock = LazyLock::new(|| redis::Script::new(CHARGE)); + #[derive(Clone, Debug, PartialEq)] pub struct RateLimitSettings { /// `LITELLM_RATE_LIMIT_WINDOW_SIZE`, in seconds; 60 by default. @@ -448,15 +582,14 @@ impl Store { if self.redis.is_some() { let error = match self.connection() { Some((mut connection, redis)) => { - let mut command = redis::cmd("EVAL"); - command.arg(REQUESTS).arg(prefixes.len() * 2); + let mut invocation = REQUESTS_SCRIPT.prepare_invoke(); for prefix in prefixes { - command - .arg(redis.namespaced(&format!("{prefix}:window"))) - .arg(redis.namespaced(&format!("{prefix}:requests"))); + invocation + .key(redis.namespaced(&format!("{prefix}:window"))) + .key(redis.namespaced(&format!("{prefix}:requests"))); } - command.arg(now).arg(self.window); - match integers(command.query_async(&mut connection).await) { + invocation.arg(now).arg(self.window); + match integers(invocation.invoke_async(&mut connection).await) { Ok(reply) if reply.len() == prefixes.len() * 2 => { return Ok(reply.iter().skip(1).step_by(2).copied().collect()); } @@ -471,32 +604,14 @@ impl Store { Ok(self.local().requests(prefixes, now, self.window)) } - /// Takes an in-flight slot: the count with it, or the full count without. `true` when the - /// slot is in Redis. - async fn acquire( - &self, - prefix: &str, - limit: i64, - slot: &str, - ) -> Result<(Result, bool), Unavailable> { + /// Counts the request, takes its in-flight slot and reserves its tokens, stopping at the first + /// limit it is over: one script with Redis, else this replica's counters. + async fn admit(&self, plan: &Plan) -> Result<(Admission, bool), Unavailable> { if self.redis.is_some() { let error = match self.connection() { Some((mut connection, redis)) => { - let reply = redis::cmd("EVAL") - .arg(ACQUIRE) - .arg(1) - .arg(redis.namespaced(&format!("{prefix}:max_parallel_requests"))) - .arg(limit) - .arg(SLOT_TTL) - .arg(slot) - .query_async(&mut connection) - .await; - match integers(reply) { - Ok(reply) => match reply.as_slice() { - [0, in_flight] => return Ok((Ok(*in_flight), true)), - [1, _, in_flight, _] => return Ok((Err(*in_flight), true)), - _ => Some(malformed()), - }, + match admit_redis(&mut connection, redis, plan, self.window).await { + Ok(admission) => return Ok((admission, true)), Err(error) => Some(error), } } @@ -504,7 +619,60 @@ impl Store { }; self.fall_back(error)?; } - Ok((self.local().acquire(prefix, limit, slot), false)) + Ok((self.admit_local(plan), false)) + } + + fn admit_local(&self, plan: &Plan) -> Admission { + let mut local = self.local(); + let now = epoch_seconds(); + let prefixes: Vec = plan + .requests + .iter() + .map(|(prefix, _)| prefix.clone()) + .collect(); + let requests = local.requests(&prefixes, now, self.window); + if plan + .requests + .iter() + .zip(&requests) + .any(|((_, limit), count)| count > limit) + { + return Admission { + requests, + outcome: Outcome::RequestsOver, + }; + } + let in_flight = match &plan.slot { + Some(slot) => match local.acquire(&slot.prefix, slot.limit, &slot.id) { + Ok(in_flight) => Some(in_flight), + Err(in_flight) => { + return Admission { + requests, + outcome: Outcome::SlotFull(in_flight), + }; + } + }, + None => None, + }; + let tokens = match plan.tokens.is_empty() { + true => Vec::new(), + false => match local.check_and_add(&plan.token_pairs(), now, self.window) { + Reserved::Counts(counts) => counts, + Reserved::Over { index, current } => { + if let Some(slot) = &plan.slot { + local.release(&slot.prefix, &slot.id); + } + return Admission { + requests, + outcome: Outcome::TokensOver { index, current }, + }; + } + }, + }; + Admission { + requests, + outcome: Outcome::Admitted { in_flight, tokens }, + } } async fn release(&self, prefix: &str, slot: &str, shared: bool) { @@ -523,25 +691,6 @@ impl Store { self.local().release(prefix, slot); } - /// Reserves `increment` tokens on every `(prefix, limit)`, or on none. `true` when the - /// reservation is in Redis. - async fn reserve( - &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( @@ -551,25 +700,24 @@ impl Store { if self.redis.is_some() { let error = match self.connection() { Some((mut connection, redis)) => { - let mut command = redis::cmd("EVAL"); - command.arg(RESERVE).arg(pairs.len() * 2); + let mut invocation = RESERVE_SCRIPT.prepare_invoke(); for pair in pairs { - command - .arg(redis.namespaced(&format!("{}:window", pair.prefix))) - .arg(redis.namespaced(&format!( + invocation + .key(redis.namespaced(&format!("{}:window", pair.prefix))) + .key(redis.namespaced(&format!( "{}:{}", pair.prefix, pair.kind.name() ))); } for pair in pairs { - command + invocation .arg(pair.limit) .arg(pair.increment) .arg(self.window) .arg(self.window); } - match integers(command.query_async(&mut connection).await) { + match integers(invocation.invoke_async(&mut connection).await) { Ok(reply) if reply.first() == Some(&0) => { return Ok(( Reserved::Counts( @@ -628,15 +776,17 @@ impl Store { /// Adds each delta to its prefix's tokens. 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()); + let mut invocation = CHARGE_SCRIPT.prepare_invoke(); for (prefix, _) in &deltas { - command.arg(redis.namespaced(&format!("{prefix}:tokens"))); + invocation.key(redis.namespaced(&format!("{prefix}:tokens"))); } for (_, delta) in &deltas { - command.arg(*delta).arg(self.window); + invocation.arg(*delta).arg(self.window); } - if let Err(error) = command.query_async::(&mut connection).await { + if let Err(error) = invocation + .invoke_async::(&mut connection) + .await + { tracing::warn!(%error, "settling reserved tokens failed"); } return; @@ -645,6 +795,123 @@ impl Store { } } +/// What one request asks of the counters, in the proxy's order. +struct Plan { + /// Each requests limit: its prefix and limit. + requests: Vec<(String, i64)>, + slot: Option, + /// Each tokens limit: its prefix and limit. + tokens: Vec<(String, i64)>, + increment: i64, +} + +impl Plan { + /// The tokens limits as the pairs to reserve the increment on. + fn token_pairs(&self) -> Vec { + self.tokens + .iter() + .map(|(prefix, limit)| Pair { + prefix: prefix.clone(), + kind: Kind::Tokens, + limit: *limit, + increment: self.increment, + }) + .collect() + } +} + +struct SlotPlan { + prefix: String, + limit: i64, + id: String, +} + +/// Each requests count, then how far the request got. +#[derive(Debug, PartialEq)] +struct Admission { + requests: Vec, + outcome: Outcome, +} + +#[derive(Debug, PartialEq)] +enum Outcome { + RequestsOver, + SlotFull(i64), + TokensOver { + index: usize, + current: i64, + }, + Admitted { + in_flight: Option, + tokens: Vec, + }, +} + +async fn admit_redis( + connection: &mut ConnectionManager, + redis: &Coordination, + plan: &Plan, + window: u64, +) -> redis::RedisResult { + let mut invocation = ADMIT_SCRIPT.prepare_invoke(); + invocation + .arg(epoch_seconds()) + .arg(window) + .arg(plan.requests.len()); + for (prefix, limit) in &plan.requests { + invocation + .key(redis.namespaced(&format!("{prefix}:window"))) + .key(redis.namespaced(&format!("{prefix}:requests"))) + .arg(*limit); + } + match &plan.slot { + Some(slot) => { + invocation + .key(redis.namespaced(&format!("{}:max_parallel_requests", slot.prefix))) + .arg(1) + .arg(slot.limit) + .arg(SLOT_TTL) + .arg(&slot.id); + } + None => { + invocation.arg(0).arg(0).arg(0).arg(""); + } + } + invocation.arg(plan.increment); + for (prefix, limit) in &plan.tokens { + invocation + .key(redis.namespaced(&format!("{prefix}:window"))) + .key(redis.namespaced(&format!("{prefix}:tokens"))) + .arg(*limit) + .arg(window) + .arg(window); + } + let reply = integers(invocation.invoke_async(connection).await)?; + let counted = plan.requests.len(); + let (code, rest) = reply.split_first().ok_or_else(malformed)?; + if rest.len() < counted { + return Err(malformed()); + } + let (requests, rest) = rest.split_at(counted); + let requests = requests.to_vec(); + let outcome = match (code, rest) { + (1, []) => Outcome::RequestsOver, + (2, [in_flight]) => Outcome::SlotFull(*in_flight), + (3, [_, index, current]) if (1..=plan.tokens.len() as i64).contains(index) => { + Outcome::TokensOver { + index: (*index - 1) as usize, + current: *current, + } + } + (0, [in_flight, tokens @ ..]) if tokens.len() == plan.tokens.len() => Outcome::Admitted { + in_flight: plan.slot.as_ref().map(|_| *in_flight), + tokens: tokens.to_vec(), + }, + _ => return Err(malformed()), + }; + Ok(Admission { requests, outcome }) +} + fn malformed() -> redis::RedisError { redis::RedisError::from(( redis::ErrorKind::UnexpectedReturnType, @@ -957,76 +1224,21 @@ impl Hook for RateLimiter { if descriptors.is_empty() { return Ok(None); } - let unavailable = |Unavailable| self.unavailable(request); - let mut statuses = Vec::new(); - - // Requests per window, counted whether or not the request goes through. + // Requests per window (counted whether or not the request goes through), then + // requests in flight for the key only, then tokens per window, reserved from an + // estimate and settled when the response ends. let windowed: Vec = (0..descriptors.len()) .filter(|index| descriptors[*index].limits.rpm.is_some()) .collect(); - if !windowed.is_empty() { - let prefixes: Vec = windowed - .iter() - .map(|index| descriptors[*index].prefix()) - .collect(); - let counts = self.store.requests(&prefixes).await.map_err(unavailable)?; - for (index, count) in windowed.iter().zip(counts) { - let limit = descriptors[*index].limits.rpm.unwrap_or_default(); - statuses.push(Status { - descriptor: *index, - kind: Kind::Requests, - limit, - remaining: limit - count, - }); - } - if let Some(over) = statuses.iter().find(|status| status.remaining < 0) { - return Err(self.refusal(request, &descriptors[over.descriptor], *over)); - } - } - - // Requests in flight, for the key only. - let mut slot = None; - if let Some(index) = descriptors + let parallel = descriptors .iter() - .position(|descriptor| descriptor.max_parallel_requests.is_some()) - { - let descriptor = &descriptors[index]; - let limit = descriptor.max_parallel_requests.unwrap_or_default(); - let prefix = descriptor.prefix(); - let id = uuid::Uuid::new_v4().simple().to_string(); - let (taken, shared) = self - .store - .acquire(&prefix, limit, &id) - .await - .map_err(unavailable)?; - let in_flight = match taken { - Ok(in_flight) => in_flight, - Err(in_flight) => { - let status = Status { - descriptor: index, - kind: Kind::Parallel, - limit, - remaining: limit - in_flight, - }; - return Err(self.refusal(request, descriptor, status)); - } - }; - statuses.push(Status { - descriptor: index, - kind: Kind::Parallel, - limit, - remaining: (limit - in_flight).max(0), - }); - slot = Some(Slot { prefix, id, shared }); - } - - // Tokens per window, reserved from an estimate and settled when the response ends. + .position(|descriptor| descriptor.max_parallel_requests.is_some()); let limited: Vec<(usize, i64)> = descriptors .iter() .enumerate() .filter_map(|(index, descriptor)| Some((index, descriptor.limits.tpm?))) .collect(); - let mut tokens = None; + let mut increment = 0; if let Some(min_tpm) = limited.iter().map(|(_, limit)| *limit).min() { let body: Map = serde_json::from_slice(&request.body()).unwrap_or_default(); @@ -1039,54 +1251,110 @@ impl Hook for RateLimiter { if let Some(max_tokens) = estimate.implicit_max_tokens { edit.set_upstream_body_field("max_tokens", Value::from(max_tokens)); } - let scopes: Vec<(String, i64)> = limited + increment = estimate.tokens; + } + let plan = Plan { + requests: windowed + .iter() + .map(|index| { + let descriptor = &descriptors[*index]; + ( + descriptor.prefix(), + descriptor.limits.rpm.unwrap_or_default(), + ) + }) + .collect(), + slot: parallel.map(|index| SlotPlan { + prefix: descriptors[index].prefix(), + limit: descriptors[index].max_parallel_requests.unwrap_or_default(), + id: uuid::Uuid::new_v4().simple().to_string(), + }), + tokens: limited .iter() .map(|(index, limit)| (descriptors[*index].prefix(), *limit)) - .collect(); - let reserved = self.store.reserve(&scopes, estimate.tokens).await; - let release = |slot: Option| async { - if let Some(slot) = slot { - self.store - .release(&slot.prefix, &slot.id, slot.shared) - .await; - } - }; - let (reserved, shared) = match reserved { - Ok(reserved) => reserved, - Err(Unavailable) => { - release(slot).await; - return Err(self.unavailable(request)); - } - }; - match reserved { - Reserved::Counts(counts) => { - for ((index, limit), count) in limited.iter().zip(counts) { - statuses.push(Status { - descriptor: *index, - kind: Kind::Tokens, - limit: *limit, - remaining: (limit - count).max(0), - }); - } - } - Reserved::Over { index, current } => { - release(slot).await; - let (descriptor, limit) = limited[index]; - let status = Status { - descriptor, - kind: Kind::Tokens, - limit, - remaining: limit - current, - }; - return Err(self.refusal(request, &descriptors[descriptor], status)); + .collect(), + increment, + }; + let (admission, shared) = self + .store + .admit(&plan) + .await + .map_err(|Unavailable| self.unavailable(request))?; + + let mut statuses: Vec = windowed + .iter() + .zip(&admission.requests) + .map(|(index, count)| { + let limit = descriptors[*index].limits.rpm.unwrap_or_default(); + Status { + descriptor: *index, + kind: Kind::Requests, + limit, + remaining: limit - count, } + }) + .collect(); + let (in_flight, counts) = match admission.outcome { + Outcome::RequestsOver => { + let over = statuses + .iter() + .find(|status| status.remaining < 0) + .copied() + .unwrap_or(statuses[0]); + return Err(self.refusal(request, &descriptors[over.descriptor], over)); } - tokens = Some(Tokens { - prefixes: scopes.into_iter().map(|(prefix, _)| prefix).collect(), - reserved: estimate.tokens, - shared, + Outcome::SlotFull(in_flight) => { + let (index, slot) = (parallel.unwrap_or_default(), &plan.slot); + let limit = slot.as_ref().map_or(0, |slot| slot.limit); + let status = Status { + descriptor: index, + kind: Kind::Parallel, + limit, + remaining: limit - in_flight, + }; + return Err(self.refusal(request, &descriptors[index], status)); + } + Outcome::TokensOver { index, current } => { + let (descriptor, limit) = limited[index]; + let status = Status { + descriptor, + kind: Kind::Tokens, + limit, + remaining: limit - current, + }; + return Err(self.refusal(request, &descriptors[descriptor], status)); + } + Outcome::Admitted { in_flight, tokens } => (in_flight, tokens), + }; + let slot = match (parallel, plan.slot, in_flight) { + (Some(index), Some(slot), Some(in_flight)) => { + statuses.push(Status { + descriptor: index, + kind: Kind::Parallel, + limit: slot.limit, + remaining: (slot.limit - in_flight).max(0), + }); + Some(Slot { + prefix: slot.prefix, + id: slot.id, + shared, + }) + } + _ => None, + }; + for ((index, limit), count) in limited.iter().zip(&counts) { + statuses.push(Status { + descriptor: *index, + kind: Kind::Tokens, + limit: *limit, + remaining: (limit - count).max(0), }); } + let tokens = (!plan.tokens.is_empty()).then(|| Tokens { + prefixes: plan.tokens.into_iter().map(|(prefix, _)| prefix).collect(), + reserved: increment, + shared, + }); // The proxy sends where each limit stands on complete responses only. if !request.stream { diff --git a/litellm-rust/crates/gateway/tests/rate_limits.rs b/litellm-rust/crates/gateway/tests/rate_limits.rs index 2e7101bd89c..3bb2b9f9efb 100644 --- a/litellm-rust/crates/gateway/tests/rate_limits.rs +++ b/litellm-rust/crates/gateway/tests/rate_limits.rs @@ -466,6 +466,80 @@ async fn replicas_share_the_proxy_counters() { assert_eq!(in_flight, 0); } +#[rstest] +#[tokio::test] +async fn shared_refusals_count_and_hold_as_the_proxy() { + let Some(url) = redis_url() else { return }; + let mut connection = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + + // Over the tokens limit: the request still counts, and its in-flight slot is given back. + let namespace = format!("rate-limits-{}", uuid::Uuid::new_v4()); + let limits = KeyLimits { + key: Limits { + rpm: Some(10), + tpm: Some(1), + }, + max_parallel_requests: Some(1), + ..KeyLimits::default() + }; + let setup = Setup { + redis: Some(redis(&url, &namespace).await), + ..Setup::default() + }; + let single = gateway(vec![("sk-key", key(None, limits))], setup).await; + let refused = messages(&single.app).await; + assert_eq!(refused.status, 429); + assert!( + refusal_message(&refused, &["error", "message"]) + .contains("Limit type: tokens. Current limit: 1, Remaining: 1.") + ); + let prefix = format!("{namespace}:{{api_key:{}}}", hash("sk-key").as_str()); + let requests: Option = redis::cmd("GET") + .arg(format!("{prefix}:requests")) + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(requests.as_deref(), Some("1")); + let in_flight: i64 = redis::cmd("ZCARD") + .arg(format!("{prefix}:max_parallel_requests")) + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(in_flight, 0); + + // A full gauge on one replica refuses on another. + let namespace = format!("rate-limits-{}", uuid::Uuid::new_v4()); + let limits = KeyLimits { + max_parallel_requests: Some(1), + ..KeyLimits::default() + }; + let mut replicas = Vec::new(); + for _ in 0..2 { + let setup = Setup { + redis: Some(redis(&url, &namespace).await), + delay: Duration::from_millis(300), + ..Setup::default() + }; + replicas.push(gateway(vec![("sk-key", key(None, limits))], setup).await); + } + let (first, second) = tokio::join!(messages(&replicas[0].app), async { + tokio::time::sleep(Duration::from_millis(100)).await; + messages(&replicas[1].app).await + }); + assert_eq!(first.status, 200); + assert_eq!(second.status, 429); + assert!( + refusal_message(&second, &["error", "message"]) + .contains("Limit type: max_parallel_requests. Current limit: 1, Remaining: 0.") + ); + settled().await; + assert_eq!(messages(&replicas[1].app).await.status, 200); +} + #[expect( clippy::disallowed_methods, reason = "test fixture rows; the lookups under test use the checked macro"