diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs index bad4085cfe6..7d8c11830dc 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -1,22 +1,28 @@ use std::{sync::Arc, time::Duration}; -use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage}; +use litellm_redis_store::router_usage::{Counter, Increment, Kind, RouterUsage, Scope}; use pyo3::{prelude::*, types::PyDict}; use super::{ - PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote, + NativePending, PyNumber, Spelled, kind_error, pending, + policy::router_tier_policies, + redis_async, redis_sync, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, }; use crate::redis::client::NativeRedisClient; /// `litellm.redis.stores.router_usage.UsageCounter`. #[derive(FromPyObject)] struct PyCounter { + #[pyo3(attribute)] + scope: Spelled, #[pyo3(attribute)] deployment_id: String, #[pyo3(attribute)] model: Option, #[pyo3(attribute)] - minute: String, + window: String, #[pyo3(attribute)] kind: Spelled, } @@ -24,9 +30,10 @@ struct PyCounter { impl From for Counter { fn from(counter: PyCounter) -> Self { Self { + scope: counter.scope.0, deployment_id: counter.deployment_id, model: counter.model, - minute: counter.minute, + window: counter.window, kind: counter.kind.0, } } @@ -36,6 +43,41 @@ fn counters(counters: Vec) -> Vec { counters.into_iter().map(Counter::from).collect() } +/// `litellm.redis.stores.router_usage.CounterIncrement`. +#[derive(FromPyObject)] +struct PyIncrement { + #[pyo3(attribute)] + counter: PyCounter, + #[pyo3(attribute)] + amount: PyNumber, + #[pyo3(attribute)] + ttl: Option, +} + +impl TryFrom for Increment { + type Error = PyErr; + + fn try_from(increment: PyIncrement) -> PyResult { + Ok(Self { + counter: increment.counter.into(), + amount: increment.amount.0, + ttl: ttl(increment.ttl)?, + }) + } +} + +fn increments(increments: Vec) -> PyResult> { + increments.into_iter().map(Increment::try_from).collect() +} + +fn numbers(read: Vec>) -> Vec> { + read.into_iter().map(|count| count.map(PyNumber)).collect() +} + +fn totals(totals: Option>) -> Option> { + totals.map(|totals| totals.into_iter().map(PyNumber).collect()) +} + /// One router's global TPM/RPM counters. The `_sync` methods block with the GIL released, for the /// Rust router's sync calls, which have no event loop to await on. #[pyclass(frozen, name = "NativeRouterUsage")] @@ -170,6 +212,108 @@ impl NativeRouterUsage { }) } + /// `get_many` declared on the batch `send` picks. + fn usage(&self, counters: Vec, send: &NativeSend) -> PyResult { + let declared = self + .0 + .declare_get_many(send.send(), &self::counters(counters)) + .map_err(kind_error)?; + Ok(pending( + "async_batch_get_cache", + kind_error, + declared, + numbers, + )) + } + + fn increment_pipeline<'py>( + &self, + py: Python<'py>, + increments: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let increments = self::increments(increments)?; + redis_async(py, "async_increment_pipeline", kind_error, async move { + Ok(totals(store.increment_pipeline(&increments).await)) + }) + } + + /// Counts in memory now and declares the increments on the request's post-call batch. + fn increment_post_call( + &self, + request: &NativeRequestRedis, + increments: Vec, + ) -> PyResult<()> { + self.0 + .increment_post_call(&request.0, &self::increments(increments)?); + Ok(()) + } + + fn increment_remote<'py>( + &self, + py: Python<'py>, + increments: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let increments = self::increments(increments)?; + redis_async(py, "async_increment_pipeline", kind_error, async move { + store.increment_remote(&increments).await.map(totals) + }) + } + + fn increment_local(&self, increment: PyIncrement) -> PyResult { + self.0 + .increment_local(&increment.try_into()?) + .map(PyNumber) + .map_err(kind_error) + } + + fn get_local(&self, counter: PyCounter) -> PyResult> { + self.0 + .get_local(&counter.into()) + .map(|count| count.map(PyNumber)) + .map_err(kind_error) + } + + fn get_local_many(&self, counters: Vec) -> PyResult>> { + self.0 + .get_local_many(&self::counters(counters)) + .map(numbers) + .map_err(kind_error) + } + + fn set_local(&self, counter: PyCounter, value: PyNumber) -> PyResult<()> { + self.0 + .set_local(&counter.into(), value.0) + .map_err(kind_error) + } + + fn current<'py>( + &self, + py: Python<'py>, + counter: PyCounter, + limit: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_async(py, "async_get_cache", kind_error, async move { + Ok(store.current(&counter, limit.0).await?.map(PyNumber)) + }) + } + + fn current_sync( + &self, + py: Python<'_>, + counter: PyCounter, + limit: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_sync(py, "get_cache", kind_error, async move { + Ok(store.current(&counter, limit.0).await?.map(PyNumber)) + }) + } + fn flush_local_counters(&self) -> PyResult<()> { self.0.flush_local_counters().map_err(kind_error) } diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index 674b85013a2..c10fa2b9c1b 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -31,7 +31,7 @@ use pyo3::{ }; use serde_json::Value; -use litellm_redis_store::router_usage::{Counter, Kind}; +use litellm_redis_store::router_usage::{Counter, Kind, Scope}; use self::host::{ CacheCall, CacheOp, CacheOutcome, DriveError, Driver, Finish, HostCall, HostCallResponse, @@ -392,9 +392,10 @@ impl RouterCoreHandle { .parse() .map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?; Ok(Counter { + scope: Scope::GlobalRouter, deployment_id: id.to_owned(), model: Some(model.to_owned()), - minute: minute.to_owned(), + window: minute.to_owned(), kind, } .key()) @@ -1341,9 +1342,10 @@ fn cache_outcome(op: CacheOp, outcome: Option>) -> PyResult String { Counter { + scope: Scope::GlobalRouter, deployment_id: window.deployment_id, model: Some(window.model), - minute: window.minute, + window: window.minute, kind: Kind::Tpm, } .key() diff --git a/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs b/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs index 71dfda49233..0fab91abdde 100644 --- a/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs +++ b/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs @@ -4,7 +4,7 @@ use std::{ time::Duration, }; -use litellm_redis_client::{Declared, Error as RedisError, commands}; +use litellm_redis_client::{Declared, Error as RedisError, RedisBatch, commands}; use super::{DualCache, MemorySettings}; use crate::{Error, FromValue, Send, codec, memory::lock, request::Misses}; @@ -66,9 +66,11 @@ impl DualCache = reservation.keys.iter().map(String::as_str).collect(); - let replies = send.batch(client).add(commands::mget::>(&keys)); + let batch = send.batch(client); + let replies = batch.add(commands::mget::>(&keys)); Fetch { reservation, + batch, replies, misses: send.misses(client), } @@ -81,24 +83,37 @@ impl DualCache Result>, Error> { + if let Some(batch) = declared.batch() { + batch.flush().await; + } + self.settle_batch_get(declared).await + } + + /// The second half of `batch_get` once the read's batch answers, without flushing it. Memory + /// is read again here, so a value written to memory since the declaration wins over what Redis + /// answered. The keys Redis lacks are noted as the request's misses. An undecodable value + /// serves memory for the whole read, where `main`'s `RedisCache` returns the raw bytes and + /// `DualCache` backfills memory with them. + pub(crate) async fn settle_batch_get( + &self, + declared: DeclaredRead, ) -> Result>, Error> { let DeclaredRead { keys, fetch } = declared; let Some(Fetch { reservation, + batch: _, replies, misses, }) = fetch else { return self.read_memory(&keys); }; - let fetched = replies.resolve().await; + let fetched = replies.settled().await; let local = self.read_memory(&keys)?; let fetched = match fetched { Ok(fetched) => fetched, @@ -191,8 +206,16 @@ pub(crate) struct DeclaredRead { fetch: Option, } +impl DeclaredRead { + /// The batch carrying the read, `None` when memory answers it alone. + pub(crate) fn batch(&self) -> Option { + self.fetch.as_ref().map(|fetch| fetch.batch.clone()) + } +} + struct Fetch { reservation: Reservation, + batch: RedisBatch, replies: Declared>>>, misses: Option, } diff --git a/litellm-rust/crates/redis-store/src/dual_cache/increment.rs b/litellm-rust/crates/redis-store/src/dual_cache/increment.rs index 705bc71b2c1..25de617a5d3 100644 --- a/litellm-rust/crates/redis-store/src/dual_cache/increment.rs +++ b/litellm-rust/crates/redis-store/src/dual_cache/increment.rs @@ -1,7 +1,7 @@ use std::time::Duration; use futures_util::future::join_all; -use litellm_redis_client::{Declared, Expiry, RedisBatch, commands}; +use litellm_redis_client::{Declared, Error as RedisError, Expiry, RedisBatch, commands}; use super::DualCache; use crate::{Error, Number, RequestRedis, Send, failure_log::redis_failure, memory::Counter}; @@ -63,10 +63,23 @@ impl DualCache { .collect::, _>>() .map_err(|error| redis_failure!(&Error::from(error), PIPELINE_FAILED)) .ok()?; - let Some(client) = self.redis.client() else { - return Some(local); - }; - let pipeline = RedisBatch::pipelined(client.clone()); + match self.increment_pipeline_remote(increments).await { + None => Some(local), + Some(Ok(totals)) => Some(totals), + Some(Err(error)) => { + redis_failure!(&error, PIPELINE_FAILED); + Some(local) + } + } + } + + /// `RedisCache.async_increment_pipeline`, Redis alone: one pipeline of `INCRBYFLOAT` with + /// `EXPIRE` per increment. `None` without Redis. + pub(crate) async fn increment_pipeline_remote( + &self, + increments: &[PipelineIncrement], + ) -> Option, RedisError>> { + let pipeline = RedisBatch::pipelined(self.redis.client()?.clone()); let totals = increments .iter() .map(|increment| { @@ -77,17 +90,13 @@ impl DualCache { )) }) .collect::>(); - match Declared::all(totals) - .await - .into_iter() - .collect::, _>>() - { - Ok(totals) => Some(totals.into_iter().map(Number::Float).collect()), - Err(error) => { - redis_failure!(&error, PIPELINE_FAILED); - Some(local) - } - } + Some( + Declared::all(totals) + .await + .into_iter() + .map(|total| total.map(Number::Float)) + .collect(), + ) } /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, stopping at the first diff --git a/litellm-rust/crates/redis-store/src/request.rs b/litellm-rust/crates/redis-store/src/request.rs index 79497a5d503..91115886cc4 100644 --- a/litellm-rust/crates/redis-store/src/request.rs +++ b/litellm-rust/crates/redis-store/src/request.rs @@ -107,6 +107,16 @@ impl RequestRedis { lock(&self.0.follow_ups).push(follow_up); } + /// Runs `task` on the request's runtime without `flush_post` waiting for it, for work that + /// settles with the pre-call batch whoever awaits it. + pub(crate) fn detach(&self, task: impl Future + std::marker::Send + 'static) { + drop( + self.0 + .runtime + .spawn(self.logged(within(current_tag(), task))), + ); + } + fn logged(&self, task: F) -> impl Future + use { task.in_current_span().with_subscriber(self.0.logs.clone()) } @@ -196,6 +206,14 @@ impl Send<'_> { } } + /// The request the op rides, `None` for one on its own. + pub(crate) fn request(&self) -> Option<&RequestRedis> { + match self { + Self::Now => None, + Self::PreCall(request) | Self::PostCall(request) => Some(request), + } + } + /// The request's misses on `client`; `None` for an op outside a request. pub(crate) fn misses(&self, client: &RedisClient) -> Option { match self { diff --git a/litellm-rust/crates/redis-store/src/router_usage/keys.rs b/litellm-rust/crates/redis-store/src/router_usage/keys.rs index 1482668ab2b..f897ed92020 100644 --- a/litellm-rust/crates/redis-store/src/router_usage/keys.rs +++ b/litellm-rust/crates/redis-store/src/router_usage/keys.rs @@ -1,5 +1,15 @@ +/// Who counts under a key. The Rust router core's counters carry a `global_router:` prefix, and +/// the counters usage-based-routing-v2 and the model rate limit check share carry none. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)] +pub enum Scope { + #[strum(serialize = "global_router")] + GlobalRouter, + #[strum(serialize = "deployment")] + Deployment, +} + /// A counter's key segment, which is also how the Python router names it. -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)] pub enum Kind { #[strum(serialize = "tpm")] Tpm, @@ -11,25 +21,36 @@ pub enum Kind { /// Output tokens per minute. #[strum(serialize = "otpm")] Otpm, + /// Tokens per UTC day, whose window is `YYYY-MM-DD`. + #[strum(serialize = "tpd")] + Tpd, + /// Requests per UTC day. + #[strum(serialize = "rpd")] + Rpd, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct Counter { + pub scope: Scope, pub deployment_id: String, pub model: Option, - pub minute: String, + /// `HH-MM` for a per-minute kind, `YYYY-MM-DD` for a per-day one. + pub window: String, pub kind: Kind, } impl Counter { - /// `global_router:{id}:{model}:{kind}:{HH-MM}`, where Python formats a missing model as `None`. - /// Public only for the hosts that hand a counter's key back to their callers. + /// `global_router:{id}:{model}:{kind}:{window}` or `{id}:{model}:{kind}:{window}`, where Python + /// formats a missing model as `None`. Public only for the hosts that hand a counter's key back + /// to their callers. pub fn key(&self) -> String { let kind: &str = self.kind.into(); let model = self.model.as_deref().unwrap_or("None"); - format!( - "global_router:{}:{model}:{kind}:{}", - self.deployment_id, self.minute - ) + let id = &self.deployment_id; + let window = &self.window; + match self.scope { + Scope::GlobalRouter => format!("global_router:{id}:{model}:{kind}:{window}"), + Scope::Deployment => format!("{id}:{model}:{kind}:{window}"), + } } } diff --git a/litellm-rust/crates/redis-store/src/router_usage/mod.rs b/litellm-rust/crates/redis-store/src/router_usage/mod.rs index 4b123a0f094..077069cb06d 100644 --- a/litellm-rust/crates/redis-store/src/router_usage/mod.rs +++ b/litellm-rust/crates/redis-store/src/router_usage/mod.rs @@ -2,10 +2,14 @@ mod keys; use std::{sync::Arc, time::Duration}; -pub use self::keys::{Counter, Kind}; +use litellm_redis_client::{Error as RedisError, commands}; +use tokio::sync::oneshot; + +pub use self::keys::{Counter, Kind, Scope}; use crate::{ - Error, Number, + Error, Number, Pending, RequestRedis, Send, codec, dual_cache::{BatchReads, DualCache, PipelineIncrement, TierPolicy}, + failure_log::redis_failure, memory::{Clock, system_clock}, remote::{Remote, RemoteSlot}, }; @@ -13,7 +17,31 @@ use crate::{ /// `RoutingArgs.ttl`, the window of a per-minute counter. pub const COUNTER_TTL: Duration = Duration::from_secs(60); -/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod. +/// A `RedisPipelineIncrementOperation` on one counter. +#[derive(Clone, Debug, PartialEq)] +pub struct Increment { + pub counter: Counter, + pub amount: Number, + pub ttl: Option, +} + +impl From<&Increment> for PipelineIncrement { + fn from(increment: &Increment) -> Self { + Self { + key: increment.counter.key(), + amount: increment.amount, + ttl: increment.ttl, + } + } +} + +fn pipeline(increments: &[Increment]) -> Vec { + increments.iter().map(PipelineIncrement::from).collect() +} + +/// One router's usage counters: memory per router, Redis shared by every pod. The router core's +/// `global_router:` counters and the per-deployment ones its strategies and pre-call checks count +/// share one `DualCache`, as they shared the router's on `main`. pub struct RouterUsage { usage: DualCache, } @@ -61,9 +89,10 @@ impl RouterUsage { ) -> Result<(), Error> { let increment = |kind, amount| PipelineIncrement { key: Counter { + scope: Scope::GlobalRouter, deployment_id: deployment_id.to_owned(), model: model.map(str::to_owned), - minute: minute.to_owned(), + window: minute.to_owned(), kind, } .key(), @@ -94,8 +123,124 @@ impl RouterUsage { } pub async fn get_many(&self, counters: &[Counter]) -> Result>, Error> { - let keys = counters.iter().map(Counter::key).collect::>(); - self.usage.batch_get(&keys).await + self.usage.batch_get(&keys(counters)).await + } + + /// `get_many` declared on the batch `send` picks, so a pick's usage read can ride the request's + /// pre-call pipeline. On a request it settles when its batch answers, whoever awaits it, so a + /// prefetch the pick never takes still backfills memory or gives its reservation back, as the + /// settled hook of `main`'s prefetch does. + pub fn declare_get_many( + self: &Arc, + send: Send<'_>, + counters: &[Counter], + ) -> Result>>, Error> { + let read = self.usage.declare_batch_get(send, &keys(counters))?; + let store = Arc::clone(self); + let (Some(request), Some(batch)) = (send.request(), read.batch()) else { + return Ok(Box::pin(async move { + store.usage.resolve_batch_get(read).await + })); + }; + let (settled, answer) = oneshot::channel(); + let settling = Arc::clone(self); + request.detach(async move { + drop(settled.send(settling.usage.settle_batch_get(read).await)); + }); + let counters = counters.to_vec(); + Ok(Box::pin(async move { + batch.flush().await; + match answer.await { + Ok(read) => read, + Err(_) => store.get_local_many(&counters), + } + })) + } + + /// `DualCache.async_increment_cache_pipeline`: memory, then one pipeline on its own round trip. + /// `None` when memory refuses an increment, and the memory totals when Redis fails. + pub async fn increment_pipeline(&self, increments: &[Increment]) -> Option> { + self.usage.increment_pipeline(&pipeline(increments)).await + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, Redis on the request's + /// post-call batch. + pub fn increment_post_call(&self, request: &RequestRedis, increments: &[Increment]) { + self.usage + .increment_pipeline_post_call(request, &pipeline(increments)); + } + + /// `RedisCache.async_increment_pipeline`, which a strategy's periodic sync sends with the + /// increments it already counted in memory. `None` without Redis, and a Redis failure is the + /// caller's to handle, as there. + pub async fn increment_remote( + &self, + increments: &[Increment], + ) -> Result>, Error> { + self.usage + .increment_pipeline_remote(&pipeline(increments)) + .await + .transpose() + .map_err(Error::from) + } + + /// `in_memory_cache.async_increment`. + pub fn increment_local(&self, increment: &Increment) -> Result { + Ok(self.usage.memory().increment( + &increment.counter.key(), + increment.amount, + increment.ttl, + )?) + } + + /// `get_cache(local_only=True)`. + pub fn get_local(&self, counter: &Counter) -> Result, Error> { + Ok(self.usage.memory().get(&counter.key())?) + } + + /// `in_memory_cache.async_batch_get_cache`. + pub fn get_local_many(&self, counters: &[Counter]) -> Result>, Error> { + counters + .iter() + .map(|counter| self.get_local(counter)) + .collect() + } + + /// `in_memory_cache.async_set_cache` without a TTL, which keeps `InMemoryCache`'s own. + pub fn set_local(&self, counter: &Counter, value: Number) -> Result<(), Error> { + Ok(self.usage.memory().set(&counter.key(), value, None)?) + } + + /// The model rate limit check's current usage: memory alone when it already reaches `limit`, + /// else `redis_cache.get_cache`. No Redis and an open breaker serve memory, and any other + /// failure reads as `None`, as `RedisCache` swallows it. + pub async fn current(&self, counter: &Counter, limit: Number) -> Result, Error> { + let local = self.get_local(counter)?; + if local.is_some_and(|local| local.as_f64() >= limit.as_f64()) { + return Ok(local); + } + let Some(client) = self.usage.redis().client() else { + return Ok(local); + }; + let key = counter.key(); + match client.send(commands::get::>(&key)).await { + Err(RedisError::BreakerOpen) => Ok(local), + Err(error) => { + redis_failure!( + &error, + "litellm.caching.caching: get() - Got exception from REDIS" + ); + Ok(None) + } + Ok(raw) => Ok(raw + .as_deref() + .map(codec::decode) + .transpose() + .unwrap_or_else(|error| { + tracing::warn!(%error, key, "a usage counter did not decode"); + None + })), + } } pub fn flush_local_counters(&self) -> Result<(), Error> { @@ -108,3 +253,7 @@ impl RouterUsage { Ok(self.usage.memory().entries()?) } } + +fn keys(counters: &[Counter]) -> Vec { + counters.iter().map(Counter::key).collect() +} diff --git a/litellm-rust/crates/redis-store/tests/router_usage.rs b/litellm-rust/crates/redis-store/tests/router_usage.rs index cc327e9aa1d..9418c03ee7b 100644 --- a/litellm-rust/crates/redis-store/tests/router_usage.rs +++ b/litellm-rust/crates/redis-store/tests/router_usage.rs @@ -11,30 +11,361 @@ use std::{ use litellm_redis_client::commands; use litellm_redis_client::{Amount, Expiry, Operation, Ttl}; use litellm_redis_store::{ - Number, TierPolicy, - router_usage::{COUNTER_TTL, Counter, Kind, RouterUsage}, + Number, RequestRedis, Send as RedisSend, TierPolicy, + router_usage::{COUNTER_TTL, Counter, Increment, Kind, RouterUsage, Scope}, }; use litellm_redis_testing::{RedisDeployment, Topology}; use rstest::rstest; -use support::{USAGE_POLICY as POLICY, connect, remote, usage, within}; +use support::{USAGE_POLICY as POLICY, connect, deployment_usage, remote, usage, within}; +use tokio::runtime::Handle; + +const DAY: Duration = Duration::from_secs(86_400); fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] { [Kind::Tpm, Kind::Rpm].map(|kind| Counter { + scope: Scope::GlobalRouter, deployment_id: "dep".to_owned(), model: model.map(str::to_owned), - minute: "12-34".to_owned(), + window: "12-34".to_owned(), kind, }) } +fn increment(counter: &Counter, amount: Number, ttl: Duration) -> Increment { + Increment { + counter: counter.clone(), + amount, + ttl: Some(ttl), + } +} + +async fn ttl_of(raw: &litellm_redis_client::RedisClient, key: &str) -> Duration { + let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else { + panic!("{key} has no TTL"); + }; + remaining +} + #[rstest] #[case::tpm(Kind::Tpm, "tpm")] #[case::rpm(Kind::Rpm, "rpm")] +#[case::itpm(Kind::Itpm, "itpm")] +#[case::otpm(Kind::Otpm, "otpm")] +#[case::tpd(Kind::Tpd, "tpd")] +#[case::rpd(Kind::Rpd, "rpd")] fn a_counter_kind_round_trips_through_its_key_segment(#[case] kind: Kind, #[case] segment: &str) { assert_eq!(<&str>::from(kind), segment); assert_eq!(segment.parse::().unwrap(), kind); } +#[rstest] +#[case::global_router(Scope::GlobalRouter, "global_router")] +#[case::deployment(Scope::Deployment, "deployment")] +fn a_counter_scope_round_trips_through_its_spelling(#[case] scope: Scope, #[case] spelling: &str) { + assert_eq!(<&str>::from(scope), spelling); + assert_eq!(spelling.parse::().unwrap(), scope); +} + +#[rstest] +#[tokio::test] +async fn reads_deployment_counters_python_incremented_under_mains_keys( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let cases = deployment_usage(); + for (case, delta) in cases.iter().zip([4.0, 1200.5]) { + raw.send(commands::incr_by_float( + case.key, + Amount::Float(delta), + Some(Expiry::Refresh(DAY)), + )) + .await + .unwrap(); + } + + assert_eq!( + store + .get_many(&cases.map(|case| case.counter)) + .await + .unwrap(), + vec![Some(Number::Int(4)), Some(Number::Float(1200.5))] + ); +} + +#[rstest] +#[tokio::test] +async fn a_declared_usage_read_rides_the_request_batch_and_backfills_memory( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = support::connect_observed(&deployment); + let store = Arc::new(RouterUsage::new(Some(remote(client)), POLICY)); + let cases = deployment_usage(); + for (case, delta) in cases.iter().zip([2.0, 30.0]) { + raw.send(commands::incr_by_float( + case.key, + Amount::Float(delta), + Some(Expiry::Refresh(DAY)), + )) + .await + .unwrap(); + } + let counters = cases.map(|case| case.counter); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + let declared = store + .declare_get_many(RedisSend::PreCall(&request), &counters) + .unwrap(); + assert!(trips.seen().is_empty()); + request.flush_pre().await; + let read = declared.await.unwrap(); + + let one_mget = match topology { + Topology::Standalone => Operation::Pipeline { ops: 1 }, + Topology::Cluster => Operation::Command("MGET"), + }; + assert_eq!(trips.seen(), [one_mget]); + assert_eq!(read, vec![Some(Number::Int(2)), Some(Number::Int(30))]); + assert_eq!( + store.get_local_many(&counters).unwrap(), + vec![Some(Number::Int(2)), Some(Number::Int(30))] + ); +} + +// main's ArmedRead settles on the batch's reply: a prefetch the pick never takes still backfills +// memory, and the throttle then keeps later reads off Redis +#[rstest] +#[tokio::test] +async fn an_abandoned_usage_prefetch_still_backfills_memory_and_holds_the_throttle( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = support::connect_observed(&deployment); + let store = Arc::new(RouterUsage::new(Some(remote(client)), POLICY)); + let [rpm, tpd] = deployment_usage(); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(6.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + let counters = [rpm.counter.clone(), tpd.counter.clone()]; + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + drop( + store + .declare_get_many(RedisSend::PreCall(&request), &counters) + .unwrap(), + ); + request.flush_pre().await; + litellm_redis_testing::eventually(async || store.get_local(&rpm.counter).unwrap()).await; + let trips_after_prefetch = trips.seen().len(); + let read = store.get_many(&counters).await.unwrap(); + + assert_eq!(read, vec![Some(Number::Int(6)), None]); + assert_eq!(trips.seen().len(), trips_after_prefetch); +} + +// The model rate limit check charges a request's rpm and rpd in one async_increment_cache_pipeline, +// whose RedisCache pipeline refreshes each window with EXPIRE +#[rstest] +#[tokio::test] +async fn a_pipeline_increment_is_one_round_trip_refreshing_each_window( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = support::connect_observed(&deployment); + let store = RouterUsage::new(Some(remote(client)), POLICY); + let [rpm, tpd] = deployment_usage(); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(3.0), + Some(Expiry::Refresh(Duration::from_secs(10))), + )) + .await + .unwrap(); + + let totals = store + .increment_pipeline(&[ + increment(&rpm.counter, Number::Int(1), COUNTER_TTL), + increment(&tpd.counter, Number::Int(1), DAY), + ]) + .await; + + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]); + assert_eq!(totals, Some(vec![Number::Float(4.0), Number::Float(1.0)])); + assert!(within(ttl_of(&raw, rpm.key).await, 60)); + assert!(within(ttl_of(&raw, tpd.key).await, 86_400)); +} + +// LowestTPMLoggingHandler_v2.async_log_success_event's async_increment_cache_post_call +#[rstest] +#[tokio::test] +async fn a_post_call_increment_counts_in_memory_now_and_in_redis_at_flush_post( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let [rpm, _] = deployment_usage(); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(5.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + store.increment_post_call( + &request, + &[increment(&rpm.counter, Number::Int(7), COUNTER_TTL)], + ); + let before_flush: Option = raw.send(commands::get(rpm.key)).await.unwrap(); + request.flush_post().await; + let after_flush: Option = raw.send(commands::get(rpm.key)).await.unwrap(); + + assert_eq!(store.get_local(&rpm.counter).unwrap(), Some(Number::Int(7))); + assert_eq!(before_flush, Some("5".to_owned())); + assert_eq!(after_flush, Some("12".to_owned())); +} + +// BaseRoutingStrategy's periodic sync counts in memory first, then sends the queued increments +// through redis_cache.async_increment_pipeline alone +#[rstest] +#[tokio::test] +async fn a_remote_increment_answers_the_redis_totals_and_leaves_memory_alone( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let [rpm, _] = deployment_usage(); + let queued = increment(&rpm.counter, Number::Int(2), COUNTER_TTL); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(10.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + + let counted = store.increment_local(&queued).unwrap(); + let totals = store + .increment_remote(std::slice::from_ref(&queued)) + .await + .unwrap(); + + assert_eq!(counted, Number::Int(2)); + assert_eq!(totals, Some(vec![Number::Float(12.0)])); + assert_eq!(store.get_local(&rpm.counter).unwrap(), Some(Number::Int(2))); +} + +#[rstest] +#[tokio::test] +async fn a_remote_increment_is_none_without_redis_and_an_error_when_redis_fails() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let detached = RouterUsage::new(None, POLICY); + let attached = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let [rpm, _] = deployment_usage(); + let queued = [increment(&rpm.counter, Number::Int(1), COUNTER_TTL)]; + deployment.stop().await.unwrap(); + + assert_eq!(detached.increment_remote(&queued).await.unwrap(), None); + assert!(attached.increment_remote(&queued).await.is_err()); +} + +#[rstest] +#[tokio::test] +async fn a_local_merge_replaces_the_memory_count_without_reaching_redis() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let (client, trips) = support::connect_observed(&deployment); + let store = RouterUsage::new(Some(remote(client)), POLICY); + let [rpm, tpd] = deployment_usage(); + + store + .increment_local(&increment(&rpm.counter, Number::Int(1), COUNTER_TTL)) + .unwrap(); + store.set_local(&rpm.counter, Number::Float(9.0)).unwrap(); + + assert_eq!( + store.get_local_many(&[rpm.counter, tpd.counter]).unwrap(), + vec![Some(Number::Float(9.0)), None] + ); + assert!(trips.seen().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn current_usage_skips_redis_once_memory_reaches_the_limit() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let (client, trips) = support::connect_observed(&deployment); + let store = RouterUsage::new(Some(remote(client)), POLICY); + let [rpm, _] = deployment_usage(); + store + .increment_local(&increment(&rpm.counter, Number::Int(5), COUNTER_TTL)) + .unwrap(); + + let at_limit = store.current(&rpm.counter, Number::Int(5)).await.unwrap(); + let skipped = trips.seen(); + let below_limit = store.current(&rpm.counter, Number::Int(6)).await.unwrap(); + + assert_eq!(at_limit, Some(Number::Int(5))); + assert!(skipped.is_empty()); + assert_eq!(below_limit, None, "Redis holds no count, and Redis wins"); + assert_eq!(trips.seen(), [Operation::Command("GET")]); +} + +#[rstest] +#[tokio::test] +async fn current_usage_reads_none_on_a_failure_and_memory_once_the_breaker_opens() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let breaker = litellm_circuit_breaker::Breaker::new( + litellm_circuit_breaker::Config { + failure_threshold: 1, + recovery_timeout: Duration::from_secs(60), + ..litellm_circuit_breaker::Config::default() + }, + Box::new(|_| {}), + ); + let store = RouterUsage::new( + Some(remote(support::connect_guarded(&deployment, breaker))), + POLICY, + ); + let [rpm, _] = deployment_usage(); + store + .increment_local(&increment(&rpm.counter, Number::Int(2), COUNTER_TTL)) + .unwrap(); + deployment.stop().await.unwrap(); + + let failed = store.current(&rpm.counter, Number::Int(10)).await.unwrap(); + let refused = store.current(&rpm.counter, Number::Int(10)).await.unwrap(); + + assert_eq!(failed, None); + assert_eq!(refused, Some(Number::Int(2))); +} + +#[rstest] +#[tokio::test] +async fn current_usage_without_redis_is_the_memory_count() { + let store = RouterUsage::new(None, POLICY); + let [rpm, _] = deployment_usage(); + store + .increment_local(&increment(&rpm.counter, Number::Int(3), COUNTER_TTL)) + .unwrap(); + + assert_eq!( + store.current(&rpm.counter, Number::Int(10)).await.unwrap(), + Some(Number::Int(3)) + ); +} + #[rstest] #[case::unknown("tokens")] #[case::wrong_case("TPM")] diff --git a/litellm-rust/crates/redis-store/tests/support/mod.rs b/litellm-rust/crates/redis-store/tests/support/mod.rs index f4b05c304e3..68966986d07 100644 --- a/litellm-rust/crates/redis-store/tests/support/mod.rs +++ b/litellm-rust/crates/redis-store/tests/support/mod.rs @@ -13,7 +13,7 @@ use litellm_redis_client::{ use litellm_redis_store::{ Number, Remote, TierPolicy, router_cooldowns::Cooldown, - router_usage::{Counter, Kind}, + router_usage::{Counter, Kind, Scope}, }; use litellm_redis_testing::RedisDeployment; @@ -212,18 +212,20 @@ pub fn usage() -> [UsageCase; 2] { [ UsageCase { counter: Counter { + scope: Scope::GlobalRouter, deployment_id: "dep".to_owned(), model: Some("anthropic/claude-sonnet-5-5".to_owned()), - minute: "12-34".to_owned(), + window: "12-34".to_owned(), kind: Kind::Tpm, }, key: "global_router:dep:anthropic/claude-sonnet-5-5:tpm:12-34", }, UsageCase { counter: Counter { + scope: Scope::GlobalRouter, deployment_id: "dep".to_owned(), model: None, - minute: "09-05".to_owned(), + window: "09-05".to_owned(), kind: Kind::Rpm, }, key: "global_router:dep:None:rpm:09-05", @@ -231,6 +233,32 @@ pub fn usage() -> [UsageCase; 2] { ] } +/// main's usage-based-routing-v2 and model rate limit check keys, which carry no prefix. +pub fn deployment_usage() -> [UsageCase; 2] { + [ + UsageCase { + counter: Counter { + scope: Scope::Deployment, + deployment_id: "dep".to_owned(), + model: Some("openai/gpt-5.4-nano".to_owned()), + window: "12-34".to_owned(), + kind: Kind::Rpm, + }, + key: "dep:openai/gpt-5.4-nano:rpm:12-34", + }, + UsageCase { + counter: Counter { + scope: Scope::Deployment, + deployment_id: "dep".to_owned(), + model: None, + window: "2026-11-10".to_owned(), + kind: Kind::Tpd, + }, + key: "dep:None:tpd:2026-11-10", + }, + ] +} + /// A TTL Redis reports for a key written `whole_seconds` ago at most a few seconds back. pub fn within(ttl: Duration, whole_seconds: u64) -> bool { ttl <= Duration::from_secs(whole_seconds) diff --git a/litellm/redis/stores/router_usage/__init__.py b/litellm/redis/stores/router_usage/__init__.py index e9d46f8a68b..b09874e3512 100644 --- a/litellm/redis/stores/router_usage/__init__.py +++ b/litellm/redis/stores/router_usage/__init__.py @@ -1,22 +1,30 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol from litellm._internal_context import with_service_target +from litellm.caching.redis_batch import Batch +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.redis.stores._route import post_call_request, ready, route from litellm.redis.stores.router_usage.handle import RouterUsageHandle from litellm.redis.stores.router_usage.types import ( ROUTER_USAGE_TARGET, + CounterIncrement, CounterKind, + CounterScope, UsageCounter, ) +from litellm.rust_bridge._native import NativeRedisStoreError if TYPE_CHECKING: from litellm.redis.stores.connection import RedisStores __all__ = [ "ROUTER_USAGE_TARGET", + "CounterIncrement", "CounterKind", + "CounterScope", "RouterUsage", "RouterUsageStore", "UsageCounter", @@ -26,7 +34,7 @@ __all__ = [ class RouterUsage(Protocol): def attach_redis(self, redis: RedisStores) -> None: ... def flush_local_counters(self) -> None: ... - async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float: ... + async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float | None: ... async def increment_deployment( self, deployment_id: str, @@ -36,7 +44,7 @@ class RouterUsage(Protocol): requests: int, ) -> None: ... async def get_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... - def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float: ... + def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float | None: ... def increment_deployment_sync( self, deployment_id: str, @@ -46,6 +54,18 @@ class RouterUsage(Protocol): requests: int, ) -> None: ... def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... + @property + def has_redis(self) -> bool: ... + def usage(self, counters: Sequence[UsageCounter], *, batch: Batch) -> Awaitable[Sequence[float | None]]: ... + async def increment_pipeline(self, increments: Sequence[CounterIncrement]) -> Sequence[float] | None: ... + async def increment_post_call(self, increments: Sequence[CounterIncrement]) -> None: ... + async def increment_remote(self, increments: Sequence[CounterIncrement]) -> Sequence[float] | None: ... + def increment_local(self, increment: CounterIncrement) -> float: ... + def get_local(self, counter: UsageCounter) -> float | None: ... + def get_local_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... + def set_local(self, counter: UsageCounter, value: float) -> None: ... + async def current(self, counter: UsageCounter, limit: float) -> float | None: ... + def current_sync(self, counter: UsageCounter, limit: float) -> float | None: ... def snapshot(self) -> Mapping[str, float]: ... @@ -66,7 +86,7 @@ class RouterUsageStore: self._handle.flush_local_counters() @with_service_target(ROUTER_USAGE_TARGET) - async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float: + async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float | None: """Without Redis every async method answers from memory without suspending, as ``DualCache``'s did.""" if not self.has_redis: return self.increment_sync(counter, value, ttl) @@ -93,7 +113,7 @@ class RouterUsageStore: return await self._handle.get_many(counters) @with_service_target(ROUTER_USAGE_TARGET) - def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float: + def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float | None: """``increment`` for a sync call, which has no event loop to await on.""" return self._handle.increment_sync(counter, value, ttl) @@ -112,5 +132,70 @@ class RouterUsageStore: def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: return self._handle.get_many_sync(counters) + @with_service_target(ROUTER_USAGE_TARGET) + def usage(self, counters: Sequence[UsageCounter], *, batch: Batch) -> Awaitable[Sequence[float | None]]: + """``get_many`` declared now on the batch ``batch`` picks, so a pick's read can share the request's + pre-call trip. Without Redis it is answered from memory without suspending.""" + if not self.has_redis: + return ready(self.get_many_sync(counters)) + target: Final = route(batch) + return target.result(self._handle.usage(counters, target.send)) + + @with_service_target(ROUTER_USAGE_TARGET) + async def increment_pipeline(self, increments: Sequence[CounterIncrement]) -> Sequence[float] | None: + """``DualCache.async_increment_cache_pipeline``, on its own trip.""" + if not self.has_redis: + return [self._handle.increment_local(increment) for increment in increments] + return await self._handle.increment_pipeline(increments) + + @with_service_target(ROUTER_USAGE_TARGET) + async def increment_post_call(self, increments: Sequence[CounterIncrement]) -> None: + """``DualCache.async_increment_cache_pipeline_post_call``: memory now, Redis on the request's post-call + pipeline, or on a trip of its own outside a request.""" + request: Final = post_call_request() + if request is None or not self.has_redis: + await self.increment_pipeline(increments) + return + self._handle.increment_post_call(request, increments) + + @with_service_target(ROUTER_USAGE_TARGET) + async def increment_remote(self, increments: Sequence[CounterIncrement]) -> Sequence[float] | None: + """``redis_cache.async_increment_pipeline``: Redis alone, None without it, raising what it raises there, + so ``log_redis_failure`` keeps an open breaker at DEBUG and throttles timeouts.""" + try: + return await self._handle.increment_remote(increments) + except NativeRedisStoreError as e: + match e.args[:1]: + case ("breaker_open",): + raise RedisCircuitBreakerOpenError(*e.args[1:]) from e + case ("timeout",): + raise TimeoutError(*e.args[1:]) from e + case _: + raise + + def increment_local(self, increment: CounterIncrement) -> float: + return self._handle.increment_local(increment) + + def get_local(self, counter: UsageCounter) -> float | None: + return self._handle.get_local(counter) + + def get_local_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: + return self._handle.get_local_many(counters) + + def set_local(self, counter: UsageCounter, value: float) -> None: + self._handle.set_local(counter, value) + + @with_service_target(ROUTER_USAGE_TARGET) + async def current(self, counter: UsageCounter, limit: float) -> float | None: + """Memory alone once it reaches ``limit``, else Redis's count. No Redis and an open breaker serve memory, + and any other Redis failure reads as None, as ``RedisCache.async_get_cache`` swallows it.""" + if not self.has_redis: + return self._handle.get_local(counter) + return await self._handle.current(counter, limit) + + @with_service_target(ROUTER_USAGE_TARGET) + def current_sync(self, counter: UsageCounter, limit: float) -> float | None: + return self._handle.current_sync(counter, limit) + def snapshot(self) -> Mapping[str, float]: return self._handle.snapshot() diff --git a/litellm/redis/stores/router_usage/handle.py b/litellm/redis/stores/router_usage/handle.py index d3e138e4cd1..cb812721d99 100644 --- a/litellm/redis/stores/router_usage/handle.py +++ b/litellm/redis/stores/router_usage/handle.py @@ -3,15 +3,16 @@ from __future__ import annotations from collections.abc import Awaitable, Mapping, Sequence from typing import Protocol -from litellm.redis.stores.router_usage.types import UsageCounter -from litellm.rust_bridge._native import NativeRedisClient +from litellm.caching.redis_batch import NativePending +from litellm.redis.stores.router_usage.types import CounterIncrement, UsageCounter +from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeSend class RouterUsageHandle(Protocol): @property def has_redis(self) -> bool: ... def attach(self, client: NativeRedisClient, /) -> bool: ... - def increment(self, counter: UsageCounter, value: float, ttl_seconds: int) -> Awaitable[float]: ... + def increment(self, counter: UsageCounter, value: float, ttl_seconds: int) -> Awaitable[float | None]: ... def increment_deployment( self, deployment_id: str, @@ -21,7 +22,7 @@ class RouterUsageHandle(Protocol): requests: int, ) -> Awaitable[None]: ... def get_many(self, counters: Sequence[UsageCounter]) -> Awaitable[Sequence[float | None]]: ... - def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: ... + def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float | None: ... def increment_deployment_sync( self, deployment_id: str, @@ -31,5 +32,15 @@ class RouterUsageHandle(Protocol): requests: int, ) -> None: ... def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... + def usage(self, counters: Sequence[UsageCounter], send: NativeSend) -> NativePending[Sequence[float | None]]: ... + def increment_pipeline(self, increments: Sequence[CounterIncrement]) -> Awaitable[Sequence[float] | None]: ... + def increment_post_call(self, request: NativeRequestRedis, increments: Sequence[CounterIncrement]) -> None: ... + def increment_remote(self, increments: Sequence[CounterIncrement]) -> Awaitable[Sequence[float] | None]: ... + def increment_local(self, increment: CounterIncrement) -> float: ... + def get_local(self, counter: UsageCounter) -> float | None: ... + def get_local_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... + def set_local(self, counter: UsageCounter, value: float) -> None: ... + def current(self, counter: UsageCounter, limit: float) -> Awaitable[float | None]: ... + def current_sync(self, counter: UsageCounter, limit: float) -> float | None: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> Mapping[str, float]: ... diff --git a/litellm/redis/stores/router_usage/types.py b/litellm/redis/stores/router_usage/types.py index 062ee0a084e..159b9779879 100644 --- a/litellm/redis/stores/router_usage/types.py +++ b/litellm/redis/stores/router_usage/types.py @@ -4,12 +4,31 @@ from dataclasses import dataclass from typing import Final, Literal, TypeAlias ROUTER_USAGE_TARGET: Final = "router_usage" -CounterKind: TypeAlias = Literal["tpm", "rpm"] +CounterKind: TypeAlias = Literal["tpm", "rpm", "itpm", "otpm", "tpd", "rpd"] +CounterScope: TypeAlias = Literal["global_router", "deployment"] @dataclass(frozen=True, slots=True) class UsageCounter: + """One deployment's count in one window: ``HH-MM`` for a per-minute kind, ``YYYY-MM-DD`` for a per-day one. + ``global_router`` counters are the router core's, ``deployment`` ones those usage-based-routing-v2 and the + model rate limit check share.""" + + scope: CounterScope deployment_id: str model: str | None - minute: str + window: str kind: CounterKind + + @staticmethod + def for_deployment(model_id: str, model: str, kind: CounterKind, window: str) -> UsageCounter: + """A ``deployment`` counter. Callers spell the id and model with ``f"{value}"``, as main's + ``f"{model_id}:{model}:..."`` keys did, so a missing one is ``"None"``.""" + return UsageCounter("deployment", model_id, model, window, kind) + + +@dataclass(frozen=True, slots=True) +class CounterIncrement: + counter: UsageCounter + amount: float + ttl: float | None diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 9d1cc35bf87..64c93a6c3e2 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -1406,7 +1406,7 @@ class RustCore( # document router-core last received for them. self._index_written: frozenset[str] = frozenset() self._index_sent: Mapping[str, object] | None = None - self._selectors: Final = StrategySelectors(host, self._cache) + self._selectors: Final = StrategySelectors(host, self._cache, self._usage) self._selectors.init_default(resolved.routing_strategy, resolved.routing_strategy_args) self._selectors.init_groups(_group_selectors(resolved.routing_groups)) self._port: Final = _EndpointPort(self) @@ -1453,7 +1453,7 @@ class RustCore( self._hosted._core_routing_plugins = list(resolved.routing_plugins) # pyright: ignore[reportPrivateUsage] # the core fills its own host registry self._rebuild() self._join_callback_holders() - self._checks: Final = OptionalChecks(self, host, self._cache, self._hosted, ctx, config.arguments) + self._checks: Final = OptionalChecks(self, host, self._cache, self._usage, self._hosted, ctx, config.arguments) self._hosted.checks = self._checks strategies.ensure_affinity = self._checks.ensure_deployment_affinity_callback strategies.hand_over_affinity(self._checks.adopt) diff --git a/litellm/router_rust/routing_stores.py b/litellm/router_rust/routing_stores.py index 1dd678d377d..e77f3b01d0c 100644 --- a/litellm/router_rust/routing_stores.py +++ b/litellm/router_rust/routing_stores.py @@ -51,7 +51,7 @@ class UsageWindow(_Step): minute: str def counter(self, kind: CounterKind) -> UsageCounter: - return UsageCounter(self.deployment_id, self.model, self.minute, kind) + return UsageCounter("global_router", self.deployment_id, self.model, self.minute, kind) class CooldownsStep(_Step): diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index 75afe38548e..b566f358ddf 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -48,6 +48,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.redis.stores.router_usage import RouterUsage from litellm.router_python_host import PythonHost from litellm.router_rust.steps.selection import StrategySelectors from litellm.router_rust.surface.registry import HostRegistry @@ -106,6 +107,7 @@ class OptionalChecks: owner: "GlobalCallbackHolder", host: "PythonHost", cache: DualCache, + usage: "RouterUsage", hosted: "HostRegistry", ctx: object, arguments: Mapping[str, object], @@ -113,6 +115,7 @@ class OptionalChecks: self._owner: Final = owner self._host: Final = host self._cache: Final = cache + self._usage: Final = usage self._hosted: Final = hosted self._ctx: Final = ctx self.optional_callbacks: list[object] | None = None # mutable-ok: the router's live list, as Python keeps it @@ -303,7 +306,7 @@ class OptionalChecks: ) self.router_budget_logger = callback elif name == "enforce_model_rate_limits": - callback = ModelRateLimitingCheck(dual_cache=self._cache) + callback = ModelRateLimitingCheck(usage=self._usage, dual_cache=self._cache) if callback is None: continue if self.optional_callbacks is not None and any( diff --git a/litellm/router_rust/steps/selection.py b/litellm/router_rust/steps/selection.py index 183350bc05f..f7a982356b6 100644 --- a/litellm/router_rust/steps/selection.py +++ b/litellm/router_rust/steps/selection.py @@ -35,6 +35,7 @@ from litellm.types.router import RoutingGroup, RoutingStrategy if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache + from litellm.redis.stores.router_usage import RouterUsage from litellm.router_python_host import PythonHost Selector: TypeAlias = ( @@ -68,6 +69,7 @@ def build_selector( strategy: str | None, args: Mapping[str, object], cache: "DualCache", + usage: "RouterUsage", reservations: "InMemoryCache | None" = None, ) -> Selector | None: """``_build_strategy_selector`` without registering: ``None`` for simple-shuffle and every @@ -80,7 +82,7 @@ def build_selector( case "usage-based-routing": return LowestTPMLoggingHandler(router_cache=cache, routing_args=routing_args) case "usage-based-routing-v2": - return LowestTPMLoggingHandler_v2(router_cache=cache, routing_args=routing_args) + return LowestTPMLoggingHandler_v2(usage=usage, routing_args=routing_args) case "latency-based-routing": return LowestLatencyLoggingHandler(router_cache=cache, routing_args=routing_args) case "cost-based-routing": @@ -98,11 +100,12 @@ class StrategySelectors: """The selectors of one router: the default group's, each routing group's, and the lazily built per-request override ones.""" - def __init__(self, host: "PythonHost", cache: "DualCache") -> None: + def __init__(self, host: "PythonHost", cache: "DualCache", usage: "RouterUsage") -> None: from litellm.caching.in_memory_cache import InMemoryCache self._host: Final = host self._cache: Final = cache + self._usage: Final = usage # `_core_least_busy_reservations`: the picks least-busy reserved, shared by its selectors. self.reservations: Final = InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS) self._strategy: str | None = None @@ -144,7 +147,7 @@ class StrategySelectors: self._strategy = normalize_strategy(strategy) if configured: self._configured = self._strategy - self._default = build_selector(self._strategy, args, self._cache, self.reservations) + self._default = build_selector(self._strategy, args, self._cache, self._usage, self.reservations) if self._default is not None: self._register(self._default) @@ -158,7 +161,7 @@ class StrategySelectors: if self._default is None or self._strategy != self._configured: return try: - rebuilt: Final = build_selector(self._strategy, args, self._cache, self.reservations) + rebuilt: Final = build_selector(self._strategy, args, self._cache, self._usage, self.reservations) except (TypeError, ValidationError): verbose_router_logger.exception( "Invalid routing_strategy_args %s for '%s'; keeping the previous ones", dict(args), self._strategy @@ -177,7 +180,7 @@ class StrategySelectors: for name, strategy, args in groups if ( selector := build_selector( - None if strategy == "priority" else strategy, args, self._cache, self.reservations + None if strategy == "priority" else strategy, args, self._cache, self._usage, self.reservations ) ) is not None @@ -239,7 +242,7 @@ class StrategySelectors: return self.default with self._lock: if strategy not in self._overrides: - self._overrides[strategy] = build_selector(strategy, {}, self._cache, self.reservations) + self._overrides[strategy] = build_selector(strategy, {}, self._cache, self._usage, self.reservations) selector: Final = self._overrides[strategy] if selector is not None and request is not None and strategy not in self._registered_strategies(): from litellm.litellm_core_utils.litellm_logging import Logging diff --git a/litellm/router_rust/surface/legacy_members.py b/litellm/router_rust/surface/legacy_members.py index bef35673aaa..fe34b99758f 100644 --- a/litellm/router_rust/surface/legacy_members.py +++ b/litellm/router_rust/surface/legacy_members.py @@ -494,7 +494,9 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): if model_group is None or deployment_id is None: return model: Final = None if deployment_name is None else str(deployment_name) - counter: Final = UsageCounter(str(deployment_id), model, get_utc_datetime().strftime("%H-%M"), "rpm") + counter: Final = UsageCounter( + "global_router", str(deployment_id), model, get_utc_datetime().strftime("%H-%M"), "rpm" + ) await self._usage.increment(counter=counter, value=1, ttl=_USAGE_TTL) # -- the plugin context ------------------------------------------------------------------------ diff --git a/litellm/router_rust/surface/model_info.py b/litellm/router_rust/surface/model_info.py index b1649617557..d16b3d19d9b 100644 --- a/litellm/router_rust/surface/model_info.py +++ b/litellm/router_rust/surface/model_info.py @@ -770,7 +770,10 @@ class HostModelInfo: ) deployments: Final = tuple(pair for pair in named if pair[0] is not None and pair[1] is not None) return tuple( - tuple(UsageCounter(id, litellm_model, current_minute, kind) for id, litellm_model in deployments) + tuple( + UsageCounter("global_router", id, litellm_model, current_minute, kind) + for id, litellm_model in deployments + ) for kind in ("tpm", "rpm") ) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 43acae1f62f..f1d0373ef25 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -9,28 +9,27 @@ from typing import Final from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger -from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation, log_redis_failure +from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DEFAULT_REDIS_SYNC_INTERVAL -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, CounterIncrement, RouterUsage, UsageCounter class BaseRoutingStrategy(ABC): def __init__( self, - dual_cache: DualCache, + usage: RouterUsage, should_batch_redis_writes: bool, default_sync_interval: float | None, ): - self.dual_cache = dual_cache - self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] + self.usage = usage + self.redis_increment_operation_queue: list[CounterIncrement] = [] self._sync_task: asyncio.Task[None] | None = None self._retired = False self._flush_after_retire: asyncio.Task[None] | None = None if should_batch_redis_writes: self.setup_sync_task(default_sync_interval) - self.in_memory_keys_to_update: set[str] = set() # Set with max size of 1000 keys + self.in_memory_keys_to_update: set[UsageCounter] = set() # Set with max size of 1000 keys def setup_sync_task(self, default_sync_interval: float | None): """Setup the sync task in a way that's compatible with FastAPI""" @@ -56,7 +55,7 @@ class BaseRoutingStrategy(ABC): def _schedule_flush_after_retire(self) -> None: if self._flush_after_retire is not None and not self._flush_after_retire.done(): return - if not self.redis_increment_operation_queue or self.dual_cache.redis_cache is None: + if not self.redis_increment_operation_queue or not self.usage.has_redis: return try: loop: Final = asyncio.get_running_loop() @@ -79,19 +78,7 @@ class BaseRoutingStrategy(ABC): except asyncio.CancelledError: pass - async def _increment_value_list_in_current_window( - self, increment_list: list[tuple[str, int]], ttl: int - ) -> list[float]: - """ - Increment a list of values in the current window - """ - results: Final = [] - for key, value in increment_list: - result = await self._increment_value_in_current_window(key=key, value=value, ttl=ttl) - results.append(result) - return results - - async def _increment_value_in_current_window(self, key: str, value: float, ttl: int): + async def _increment_value_in_current_window(self, counter: UsageCounter, value: float, ttl: int): """ Increment spend within existing budget window @@ -100,19 +87,11 @@ class BaseRoutingStrategy(ABC): - Increments the spend in memory cache (so spend instantly updated in memory) - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) """ - result: Final = await self.dual_cache.in_memory_cache.async_increment( - key=key, - value=value, - ttl=ttl, - ) - increment_op: Final = RedisPipelineIncrementOperation( - key=key, - increment_value=value, - ttl=ttl, - ) + increment_op: Final = CounterIncrement(counter=counter, amount=value, ttl=ttl) + result: Final = self.usage.increment_local(increment_op) self.redis_increment_operation_queue.append(increment_op) - self.add_to_in_memory_keys_to_update(key=key) + self.add_to_in_memory_keys_to_update(counter) if self._retired: self._schedule_flush_after_retire() return result @@ -148,35 +127,26 @@ class BaseRoutingStrategy(ABC): Only runs if Redis is initialized """ try: - if not self.dual_cache.redis_cache: + if not self.usage.has_redis: return # Redis is not initialized if len(self.redis_increment_operation_queue) > 0: - # Compress operations for the same key - compressed_ops: Final[dict[str, RedisPipelineIncrementOperation]] = {} - ops_to_remove: Final = [] - for idx, op in enumerate(self.redis_increment_operation_queue): - if op["key"] in compressed_ops: - # Add to existing increment - compressed_ops[op["key"]]["increment_value"] += op["increment_value"] - else: - compressed_ops[op["key"]] = op - - ops_to_remove.append(idx) - - # Convert back to list - compressed_queue: Final = list(compressed_ops.values()) - - increment_result: Final = await self.dual_cache.redis_cache.async_increment_pipeline( - increment_list=compressed_queue, + queued: Final = tuple(self.redis_increment_operation_queue) + compressed_queue: Final = tuple( + CounterIncrement( + counter=counter, + amount=sum(op.amount for op in queued if op.counter == counter), + ttl=next(op.ttl for op in queued if op.counter == counter), + ) + for counter in dict.fromkeys(op.counter for op in queued) ) - self.redis_increment_operation_queue = [ - op for idx, op in enumerate(self.redis_increment_operation_queue) if idx not in ops_to_remove - ] + increment_result: Final = await self.usage.increment_remote(compressed_queue) + + self.redis_increment_operation_queue = self.redis_increment_operation_queue[len(queued) :] if increment_result is not None: - return_result = {key["key"]: op for key, op in zip(compressed_queue, increment_result)} + return_result = {op.counter: total for op, total in zip(compressed_queue, increment_result)} else: return_result = {} return return_result @@ -185,19 +155,13 @@ class BaseRoutingStrategy(ABC): log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) self.redis_increment_operation_queue = [] - def add_to_in_memory_keys_to_update(self, key: str): - self.in_memory_keys_to_update.add(key) + def add_to_in_memory_keys_to_update(self, counter: UsageCounter): + self.in_memory_keys_to_update.add(counter) - def get_key_pattern_to_sync(self) -> str | None: - """ - Get the key pattern to sync - """ - return None - - def get_in_memory_keys_to_update(self) -> set[str]: + def get_in_memory_keys_to_update(self) -> set[UsageCounter]: return self.in_memory_keys_to_update - def get_and_reset_in_memory_keys_to_update(self) -> set[str]: + def get_and_reset_in_memory_keys_to_update(self) -> set[UsageCounter]: """Atomic get and reset in-memory keys to update""" keys: Final = self.in_memory_keys_to_update self.in_memory_keys_to_update = set() @@ -221,7 +185,7 @@ class BaseRoutingStrategy(ABC): try: # No need to sync if Redis cache is not initialized - if self.dual_cache.redis_cache is None: + if not self.usage.has_redis: return # 2. Fetch all current provider spend from Redis to update in-memory cache @@ -232,10 +196,9 @@ class BaseRoutingStrategy(ABC): cache_keys_list: Final = list(cache_keys) # 1. Snapshot in-memory before - in_memory_before_dict: Final = {} - in_memory_before: Final = await self.dual_cache.in_memory_cache.async_batch_get_cache(keys=cache_keys_list) - for k, v in zip(cache_keys_list, in_memory_before): - in_memory_before_dict[k] = float(v or 0) + in_memory_before_dict: Final = { + k: float(v or 0) for k, v in zip(cache_keys_list, self.usage.get_local_many(cache_keys_list)) + } # 1. Push all provider spend increments to Redis redis_values: Final = await self._push_in_memory_increments_to_redis() @@ -246,7 +209,7 @@ class BaseRoutingStrategy(ABC): for key in cache_keys_list: redis_val = float(redis_values.get(key, 0) or 0) before = float(in_memory_before_dict.get(key, 0) or 0) - after = float(await self.dual_cache.in_memory_cache.async_get_cache(key=key) or 0) + after = float(self.usage.get_local(key) or 0) delta = after - before if after <= redis_val: merged = redis_val + delta @@ -259,7 +222,7 @@ class BaseRoutingStrategy(ABC): # import os # os._exit(1) # raise Exception(f"Redis is behind in-memory cache for key: {key}. This should not happen, since we should be updating redis with in-memory cache.") - await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=merged) + self.usage.set_local(key, merged) except Exception as e: verbose_router_logger.exception("Error syncing in-memory cache with Redis: %s", e) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index d34d7edf9e4..5df915f1406 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -13,13 +13,16 @@ import httpx import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_logger, verbose_router_logger -from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models from litellm.litellm_core_utils.token_counter import embedding_aware_token_count -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ( + ROUTER_USAGE_TARGET, + CounterIncrement, + RouterUsage, + UsageCounter, +) from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose @@ -44,21 +47,21 @@ _active_prefetched_usage: Final[ContextVar["PrefetchedUsage | None"]] = ContextV @dataclass(frozen=True) class PrefetchedUsage: """ - tpm/rpm counter values another read of this request already fetched from the router cache. + tpm/rpm counter values another read of this request already fetched from the usage store. `values` is None when that read failed, which is what `async_batch_get_cache` returns on failure. """ - keys: frozenset[str] - values: Mapping[str, object] | None + counters: frozenset[UsageCounter] + values: Mapping[UsageCounter, float | None] | None - def covers(self, keys: Sequence[str]) -> bool: - return self.keys.issuperset(keys) + def covers(self, counters: Sequence[UsageCounter]) -> bool: + return self.counters.issuperset(counters) - def values_for(self, keys: Sequence[str]) -> list[object | None] | None: + def values_for(self, counters: Sequence[UsageCounter]) -> list[float | None] | None: if self.values is None: return None - return [self.values.get(key) for key in keys] + return [self.values.get(counter) for counter in counters] @staticmethod @contextmanager @@ -92,12 +95,11 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): logged_failure: int = 0 default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour - def __init__(self, router_cache: DualCache, routing_args: dict = {}): - self.router_cache = router_cache + def __init__(self, usage: RouterUsage, routing_args: dict = {}): self.routing_args = RoutingArgs(**routing_args) BaseRoutingStrategy.__init__( self, - dual_cache=router_cache, + usage=usage, should_batch_redis_writes=True, default_sync_interval=0.1, ) @@ -118,11 +120,11 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") - model_id: Final = deployment.get("model_info", {}).get("id") - deployment_name: Final = deployment.get("litellm_params", {}).get("model") - rpm_key: Final = f"{model_id}:{deployment_name}:rpm:{current_minute}" + model_id: Final[object] = deployment.get("model_info", {}).get("id") + deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model") + rpm_counter: Final = UsageCounter.for_deployment(f"{model_id}", f"{deployment_name}", "rpm", current_minute) - local_result: Final = self.router_cache.get_cache(key=rpm_key, local_only=True) # check local result first + local_result: Final = self.usage.get_local(rpm_counter) # check local result first deployment_rpm = None if deployment_rpm is None: @@ -157,7 +159,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): else: # if local result below limit, check redis ## prevent unnecessary redis checks - result: Final = self.router_cache.increment_cache(key=rpm_key, value=1, ttl=self.routing_args.ttl) + result: Final = self.usage.increment_sync(rpm_counter, 1, self.routing_args.ttl) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", @@ -197,13 +199,11 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): # ------------ dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") - model_id: Final = deployment.get("model_info", {}).get("id") - deployment_name: Final = deployment.get("litellm_params", {}).get("model") + model_id: Final[object] = deployment.get("model_info", {}).get("id") + deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model") - rpm_key: Final = f"{model_id}:{deployment_name}:rpm:{current_minute}" - local_result: Final = await self.router_cache.async_get_cache( - key=rpm_key, local_only=True - ) # check local result first + rpm_counter: Final = UsageCounter.for_deployment(f"{model_id}", f"{deployment_name}", "rpm", current_minute) + local_result: Final = self.usage.get_local(rpm_counter) # check local result first deployment_rpm = None if deployment_rpm is None: @@ -232,7 +232,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) else: # if local result below limit, check redis ## prevent unnecessary redis checks - result = await self._increment_value_in_current_window(key=rpm_key, value=1, ttl=self.routing_args.ttl) + result = await self._increment_value_in_current_window(rpm_counter, 1, self.routing_args.ttl) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", @@ -267,8 +267,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if standard_logging_object is None: raise ValueError("standard_logging_object not passed in.") model_group: Final = standard_logging_object.get("model_group") - model: Final = standard_logging_object["hidden_params"].get("litellm_model_name") - id = standard_logging_object.get("model_id") + model: Final[object] = standard_logging_object["hidden_params"].get("litellm_model_name") + id: object = standard_logging_object.get("model_id") if model_group is None or id is None or model is None: return elif isinstance(id, int): @@ -282,14 +282,12 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - tpm_key: Final = f"{id}:{model}:tpm:{current_minute}" - # ------------ - # Update usage - # ------------ - # update cache - ## TPM - self.router_cache.increment_cache(key=tpm_key, value=total_tokens, ttl=self.routing_args.ttl) + self.usage.increment_sync( + UsageCounter.for_deployment(f"{id}", f"{model}", "tpm", current_minute), + total_tokens, + self.routing_args.ttl, + ) ### TESTING ### if self.test_flag: self.logged_success += 1 @@ -310,8 +308,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if standard_logging_object is None: raise ValueError("standard_logging_object not passed in.") model_group: Final = standard_logging_object.get("model_group") - model: Final = standard_logging_object["hidden_params"]["litellm_model_name"] - id = standard_logging_object.get("model_id") + model: Final[object] = standard_logging_object["hidden_params"]["litellm_model_name"] + id: object = standard_logging_object.get("model_id") if model_group is None or id is None: return elif isinstance(id, int): @@ -323,18 +321,15 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - tpm_key: Final = f"{id}:{model}:tpm:{current_minute}" - # ------------ - # Update usage - # ------------ - # update cache - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) ## TPM - await self.router_cache.async_increment_cache_post_call( - key=tpm_key, - value=total_tokens, - ttl=self.routing_args.ttl, - parent_otel_span=parent_otel_span, + await self.usage.increment_post_call( + ( + CounterIncrement( + UsageCounter.for_deployment(f"{id}", f"{model}", "tpm", current_minute), + total_tokens, + self.routing_args.ttl, + ), + ) ) ### TESTING ### @@ -402,10 +397,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): self, model_group: str, healthy_deployments: list, - tpm_keys: list, - tpm_values: Sequence | None, - rpm_keys: list, - rpm_values: Sequence | None, + tpm_counters: Sequence[UsageCounter], + tpm_values: Sequence[float | None] | None, + rpm_counters: Sequence[UsageCounter], + rpm_values: Sequence[float | None] | None, messages: list[dict[str, str]] | None = None, input: str | list | None = None, ) -> dict | None: @@ -416,13 +411,12 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if tpm_values is None or rpm_values is None: return None - tpm_dict = {} # {model_id: 1, ..} - for idx, key in enumerate(tpm_keys): - tpm_dict[tpm_keys[idx].split(":")[0]] = tpm_values[idx] - - rpm_dict: Final = {} # {model_id: 1, ..} - for idx, key in enumerate(rpm_keys): - rpm_dict[rpm_keys[idx].split(":")[0]] = rpm_values[idx] + tpm_dict = { + counter.deployment_id.split(":")[0]: value for counter, value in zip(tpm_counters, tpm_values) + } # {model_id: 1, ..} + rpm_dict: Final = { + counter.deployment_id.split(":")[0]: value for counter, value in zip(rpm_counters, rpm_values) + } # {model_id: 1, ..} try: input_tokens = embedding_aware_token_count( @@ -461,17 +455,25 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): else: return None - def usage_counter_keys(self, healthy_deployments: list) -> tuple[list[str], list[str]]: - """The `::tpm:` and `::rpm:` counter keys selection reads.""" + def usage_counters( + self, healthy_deployments: Sequence[object] + ) -> tuple[tuple[UsageCounter, ...], tuple[UsageCounter, ...]]: + """The tpm and rpm counters selection reads, this minute's for each deployment.""" current_minute: Final = get_utc_datetime().strftime("%H-%M") - prefixes: Final = tuple( - f"{m.get('model_info', {}).get('id')}:{m.get('litellm_params', {}).get('model')}" + named: Final[tuple[tuple[object, object], ...]] = tuple( + (m.get("model_info", {}).get("id"), m.get("litellm_params", {}).get("model")) for m in healthy_deployments if isinstance(m, dict) ) return ( - [f"{prefix}:tpm:{current_minute}" for prefix in prefixes], - [f"{prefix}:rpm:{current_minute}" for prefix in prefixes], + tuple( + UsageCounter.for_deployment(f"{model_id}", f"{model}", "tpm", current_minute) + for model_id, model in named + ), + tuple( + UsageCounter.for_deployment(f"{model_id}", f"{model}", "rpm", current_minute) + for model_id, model in named + ), ) @with_service_target(ROUTER_USAGE_TARGET) @@ -496,20 +498,18 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): healthy_deployments, ) - tpm_keys, rpm_keys = self.usage_counter_keys(healthy_deployments) - combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys + tpm_counters, rpm_counters = self.usage_counters(healthy_deployments) + combined_tpm_rpm_counters: Final = tpm_counters + rpm_counters prefetched_usage: Final = PrefetchedUsage.active() - if prefetched_usage is not None and prefetched_usage.covers(combined_tpm_rpm_keys): - combined_tpm_rpm_values = prefetched_usage.values_for(combined_tpm_rpm_keys) + if prefetched_usage is not None and prefetched_usage.covers(combined_tpm_rpm_counters): + combined_tpm_rpm_values = prefetched_usage.values_for(combined_tpm_rpm_counters) else: - combined_tpm_rpm_values = await self.router_cache.async_batch_get_cache( - keys=combined_tpm_rpm_keys - ) # [1, 2, None, ..] + combined_tpm_rpm_values = await self.usage.get_many(combined_tpm_rpm_counters) # [1, 2, None, ..] if combined_tpm_rpm_values is not None: - tpm_values = combined_tpm_rpm_values[: len(tpm_keys)] - rpm_values = combined_tpm_rpm_values[len(tpm_keys) :] + tpm_values = combined_tpm_rpm_values[: len(tpm_counters)] + rpm_values = combined_tpm_rpm_values[len(tpm_counters) :] else: tpm_values = None rpm_values = None @@ -517,9 +517,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): deployment: Final = self._common_checks_available_deployment( model_group=model_group, healthy_deployments=healthy_deployments, - tpm_keys=tpm_keys, + tpm_counters=tpm_counters, tpm_values=tpm_values, - rpm_keys=rpm_keys, + rpm_counters=rpm_counters, rpm_values=rpm_values, messages=messages, input=input, @@ -602,35 +602,16 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): healthy_deployments, ) - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") - tpm_keys: Final = [] - rpm_keys: Final = [] - for m in healthy_deployments: - if isinstance(m, dict): - id = m.get("model_info", {}).get( - "id" - ) # a deployment should always have an 'id'. this is set in router.py - deployment_name = m.get("litellm_params", {}).get("model") - tpm_key = f"{id}:{deployment_name}:tpm:{current_minute}" - rpm_key = f"{id}:{deployment_name}:rpm:{current_minute}" - - tpm_keys.append(tpm_key) - rpm_keys.append(rpm_key) - - tpm_values: Final = self.router_cache.batch_get_cache( - keys=tpm_keys, parent_otel_span=parent_otel_span - ) # [1, 2, None, ..] - rpm_values: Final = self.router_cache.batch_get_cache( - keys=rpm_keys, parent_otel_span=parent_otel_span - ) # [1, 2, None, ..] + tpm_counters, rpm_counters = self.usage_counters(healthy_deployments) + tpm_values: Final = self.usage.get_many_sync(tpm_counters) # [1, 2, None, ..] + rpm_values: Final = self.usage.get_many_sync(rpm_counters) # [1, 2, None, ..] deployment: Final = self._common_checks_available_deployment( model_group=model_group, healthy_deployments=healthy_deployments, - tpm_keys=tpm_keys, + tpm_counters=tpm_counters, tpm_values=tpm_values, - rpm_keys=rpm_keys, + rpm_counters=rpm_counters, rpm_values=rpm_values, messages=messages, input=input, diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index 38ae7e0d951..a4c3098a71c 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -23,9 +23,8 @@ import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, CounterIncrement, RouterUsage, UsageCounter from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( async_io_token_pre_call_check, async_io_token_reconcile_success, @@ -36,7 +35,6 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( io_token_reconcile_success, io_token_refund_failure, ) -from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.router import RouterErrors, parse_daily_reset_time from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_utc_datetime @@ -97,12 +95,18 @@ class _DailyLimitDeployment(BaseModel): def reset_offset(self) -> timedelta: return self.litellm_params.reset_offset() - def cache_keys(self, now: datetime) -> tuple[str, str]: - prefix: Final = f"{self.model_info.id}:{self.litellm_params.model}" + def counters(self, now: datetime) -> tuple[UsageCounter, UsageCounter]: current_day: Final = daily_window_day(now, self.reset_offset()) - return f"{prefix}:tpd:{current_day}", f"{prefix}:rpd:{current_day}" + model_id: Final = self.model_info.id + model: Final = self.litellm_params.model + return ( + UsageCounter.for_deployment(f"{model_id}", f"{model}", "tpd", current_day), + UsageCounter.for_deployment(f"{model_id}", f"{model}", "rpd", current_day), + ) - def limit_error(self, kind: Literal["tpd", "rpd"], limit: int, usage: int, now: datetime) -> litellm.RateLimitError: + def limit_error( + self, kind: Literal["tpd", "rpd"], limit: int, usage: float, now: datetime + ) -> litellm.RateLimitError: return litellm.RateLimitError( message=f"Model rate limit exceeded. {kind.upper()} limit={limit}, current usage={usage}", llm_provider="", @@ -134,58 +138,47 @@ class ModelRateLimitingCheck(CustomLogger): this check actively enforces those limits across ALL routing strategies. """ - def __init__(self, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime) -> None: + def __init__( + self, usage: RouterUsage, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime + ) -> None: + self.usage = usage self.dual_cache = dual_cache self.clock = clock # model_ids already warned about itpm/otpm + tpm/rpm on the same deployment, # so the warning is logged once per deployment rather than per request. self._io_token_conflict_warned_ids: set[str] = set() - def _increment(self, key: str, value: int, ttl: int) -> int: - return self.dual_cache.increment_cache(key=key, value=value, ttl=ttl) - - async def _async_increment( - self, key: str, value: int, ttl: int, parent_otel_span: Span | None = None - ) -> int | None: - count: Final = await self.dual_cache.async_increment_cache( - key=key, value=value, ttl=ttl, parent_otel_span=parent_otel_span - ) - return None if count is None else int(count) + def _increment(self, counter: UsageCounter, value: int, ttl: int) -> float | None: + return self.usage.increment_sync(counter, value, ttl) @staticmethod def _request_count_operations( - rpm_key: str | None, rpd_key: str | None, increment_value: int - ) -> tuple[RedisPipelineIncrementOperation, ...]: + rpm: UsageCounter | None, rpd: UsageCounter | None, increment_value: int + ) -> tuple[CounterIncrement, ...]: return tuple( - RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=ttl) - for key, ttl in ((rpm_key, RoutingArgs.ttl), (rpd_key, RoutingArgs.daily_ttl)) - if key is not None + CounterIncrement(counter=counter, amount=increment_value, ttl=ttl) + for counter, ttl in ((rpm, RoutingArgs.ttl), (rpd, RoutingArgs.daily_ttl)) + if counter is not None ) - async def _async_refund_request_counts( - self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None - ) -> None: - operations: Final = self._request_count_operations(rpm_key, rpd_key, -1) + async def _async_refund_request_counts(self, rpm: UsageCounter | None, rpd: UsageCounter | None) -> None: + operations: Final = self._request_count_operations(rpm, rpd, -1) if operations: - await self.dual_cache.async_increment_cache_pipeline( - increment_list=list(operations), parent_otel_span=parent_otel_span - ) + await self.usage.increment_pipeline(operations) async def _async_increment_request_counts( - self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None + self, rpm: UsageCounter | None, rpd: UsageCounter | None ) -> tuple[int | None, int | None]: - operations: Final = self._request_count_operations(rpm_key, rpd_key, 1) + operations: Final = self._request_count_operations(rpm, rpd, 1) if not operations: return None, None - counts: Final = await self.dual_cache.async_increment_cache_pipeline( - increment_list=list(operations), parent_otel_span=parent_otel_span - ) + counts: Final = await self.usage.increment_pipeline(operations) if counts is None or len(counts) != len(operations): return None, None - count_by_key: Final = {op["key"]: int(count) for op, count in zip(operations, counts)} + count_by_counter: Final = {op.counter: int(count) for op, count in zip(operations, counts)} return ( - None if rpm_key is None else count_by_key.get(rpm_key), - None if rpd_key is None else count_by_key.get(rpd_key), + None if rpm is None else count_by_counter.get(rpm), + None if rpd is None else count_by_counter.get(rpd), ) def _warn_io_token_and_tpm_rpm_coexist_once(self, deployment: dict) -> None: @@ -255,48 +248,25 @@ class ModelRateLimitingCheck(CustomLogger): tpd_limit, rpd_limit = deployment.limits() if tpd_limit is None and rpd_limit is None: return - tpd_key, rpd_key = deployment.cache_keys(now) + tpd, rpd = deployment.counters(now) if tpd_limit is not None: - current_tpd: Final = self._get_current_tpm(tpd_key, tpd_limit) + current_tpd: Final = self.usage.current_sync(tpd, tpd_limit) if current_tpd is not None and current_tpd >= tpd_limit: raise deployment.limit_error("tpd", tpd_limit, current_tpd, now) if rpd_limit is not None: - current_rpd: Final = self._increment(rpd_key, 1, RoutingArgs.daily_ttl) - if current_rpd > rpd_limit: - self._increment(rpd_key, -1, RoutingArgs.daily_ttl) + current_rpd: Final = self._increment(rpd, 1, RoutingArgs.daily_ttl) + if current_rpd is not None and current_rpd > rpd_limit: + self._increment(rpd, -1, RoutingArgs.daily_ttl) raise deployment.limit_error("rpd", rpd_limit, current_rpd, now) - def _get_cache_keys(self, deployment: dict, current_minute: str) -> tuple[str, str]: - """Get the cache keys for TPM and RPM tracking.""" - model_id: Final = deployment.get("model_info", {}).get("id") - deployment_name: Final = deployment.get("litellm_params", {}).get("model") - - tpm_key: Final = f"{model_id}:{deployment_name}:tpm:{current_minute}" - rpm_key: Final = f"{model_id}:{deployment_name}:rpm:{current_minute}" - - return tpm_key, rpm_key - - @with_service_target(ROUTER_USAGE_TARGET) - def _get_current_tpm(self, tpm_key: str, tpm_limit: int) -> int | None: - local_tpm: Final = self.dual_cache.get_cache(key=tpm_key, local_only=True) - redis_cache: Final = self.dual_cache.redis_cache - if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): - return local_tpm - try: - return redis_cache.get_cache(key=tpm_key) - except RedisCircuitBreakerOpenError: - return local_tpm - - @with_service_target(ROUTER_USAGE_TARGET) - async def _async_get_current_tpm(self, tpm_key: str, tpm_limit: int, parent_otel_span: Span | None) -> int | None: - local_tpm: Final = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True) - redis_cache: Final = self.dual_cache.redis_cache - if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): - return local_tpm - try: - return await redis_cache.async_get_cache(key=tpm_key, parent_otel_span=parent_otel_span) - except RedisCircuitBreakerOpenError: - return local_tpm + def _counters(self, deployment: dict, current_minute: str) -> tuple[UsageCounter, UsageCounter]: + """The counters TPM and RPM are tracked in.""" + model_id: Final[object] = deployment.get("model_info", {}).get("id") + deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model") + return ( + UsageCounter.for_deployment(f"{model_id}", f"{deployment_name}", "tpm", current_minute), + UsageCounter.for_deployment(f"{model_id}", f"{deployment_name}", "rpm", current_minute), + ) @with_service_target(ROUTER_USAGE_TARGET) def pre_call_check(self, deployment: dict) -> dict | None: @@ -320,7 +290,7 @@ class ModelRateLimitingCheck(CustomLogger): tpm_limit, rpm_limit = self._get_deployment_limits(deployment) current_minute: Final = dt.strftime("%H-%M") - tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + tpm, rpm = self._counters(deployment, current_minute) model_id: Final = deployment.get("model_info", {}).get("id") model_name: Final = deployment.get("litellm_params", {}).get("model") @@ -328,7 +298,7 @@ class ModelRateLimitingCheck(CustomLogger): # Check TPM limit if tpm_limit is not None: - current_tpm: Final = self._get_current_tpm(tpm_key, tpm_limit) + current_tpm: Final = self.usage.current_sync(tpm, tpm_limit) if current_tpm is not None and current_tpm >= tpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", @@ -347,9 +317,9 @@ class ModelRateLimitingCheck(CustomLogger): # Check RPM limit (atomic increment-first to avoid race conditions) if rpm_limit is not None: - current_rpm: Final = self._increment(rpm_key, 1, RoutingArgs.ttl) + current_rpm: Final = self._increment(rpm, 1, RoutingArgs.ttl) if current_rpm is not None and current_rpm > rpm_limit: - self._increment(rpm_key, -1, RoutingArgs.ttl) + self._increment(rpm, -1, RoutingArgs.ttl) raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", llm_provider="", @@ -369,7 +339,7 @@ class ModelRateLimitingCheck(CustomLogger): self._check_daily_limits(daily, dt) except litellm.RateLimitError: if rpm_limit is not None: - self._increment(rpm_key, -1, RoutingArgs.ttl) + self._increment(rpm, -1, RoutingArgs.ttl) raise return deployment @@ -406,8 +376,8 @@ class ModelRateLimitingCheck(CustomLogger): tpm_limit, rpm_limit = self._get_deployment_limits(deployment) current_minute: Final = dt.strftime("%H-%M") - tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) - tpd_key, rpd_key = daily.cache_keys(dt) + tpm, rpm = self._counters(deployment, current_minute) + tpd, rpd = daily.counters(dt) model_id: Final = deployment.get("model_info", {}).get("id") model_name: Final = deployment.get("litellm_params", {}).get("model") @@ -415,7 +385,7 @@ class ModelRateLimitingCheck(CustomLogger): # Check TPM limit if tpm_limit is not None: - current_tpm: Final = await self._async_get_current_tpm(tpm_key, tpm_limit, parent_otel_span) + current_tpm: Final = await self.usage.current(tpm, tpm_limit) if current_tpm is not None and current_tpm >= tpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", @@ -434,17 +404,15 @@ class ModelRateLimitingCheck(CustomLogger): ) if tpd_limit is not None: - current_tpd: Final = await self._async_get_current_tpm(tpd_key, tpd_limit, parent_otel_span) + current_tpd: Final = await self.usage.current(tpd, tpd_limit) if current_tpd is not None and current_tpd >= tpd_limit: raise daily.limit_error("tpd", tpd_limit, current_tpd, dt) - charged_rpm_key: Final = rpm_key if rpm_limit is not None else None - charged_rpd_key: Final = rpd_key if rpd_limit is not None else None - current_rpm, current_rpd = await self._async_increment_request_counts( - charged_rpm_key, charged_rpd_key, parent_otel_span - ) + charged_rpm: Final = rpm if rpm_limit is not None else None + charged_rpd: Final = rpd if rpd_limit is not None else None + current_rpm, current_rpd = await self._async_increment_request_counts(charged_rpm, charged_rpd) if rpm_limit is not None and current_rpm is not None and current_rpm > rpm_limit: - await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) + await self._async_refund_request_counts(charged_rpm, charged_rpd) raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", llm_provider="", @@ -462,7 +430,7 @@ class ModelRateLimitingCheck(CustomLogger): ) if rpd_limit is not None and current_rpd is not None and current_rpd > rpd_limit: - await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) + await self._async_refund_request_counts(charged_rpm, charged_rpd) raise daily.limit_error("rpd", rpd_limit, current_rpd, dt) return deployment @@ -494,12 +462,12 @@ class ModelRateLimitingCheck(CustomLogger): if standard_logging_object is None: return - model_id: Final = standard_logging_object.get("model_id") + model_id: Final[object] = standard_logging_object.get("model_id") if model_id is None: return total_tokens: Final = standard_logging_object.get("total_tokens", 0) - model: Final = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") + model: Final[object] = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") verbose_router_logger.debug( "[TPM TRACKING] model_id=%s, total_tokens=%s, model=%s", model_id, total_tokens, model @@ -509,16 +477,17 @@ class ModelRateLimitingCheck(CustomLogger): return dt: Final = self.clock() - current_minute: Final = dt.strftime("%H-%M") - tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" + tpm: Final = UsageCounter.for_deployment(f"{model_id}", f"{model}", "tpm", dt.strftime("%H-%M")) reset_offset: Final = _reset_offset_from_kwargs(kwargs) # pyright: ignore[reportUnknownArgumentType] # CustomLogger passes an untyped kwargs dict - tpd_key: Final = f"{model_id}:{model}:tpd:{daily_window_day(dt, reset_offset)}" + tpd: Final = UsageCounter.for_deployment( + f"{model_id}", f"{model}", "tpd", daily_window_day(dt, reset_offset) + ) - verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm_key, total_tokens) + verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm, total_tokens) tokens: Final = int(total_tokens) - await self._async_increment(tpm_key, tokens, RoutingArgs.ttl) - await self._async_increment(tpd_key, tokens, RoutingArgs.daily_ttl) + await self.usage.increment(tpm, tokens, RoutingArgs.ttl) + await self.usage.increment(tpd, tokens, RoutingArgs.daily_ttl) except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.async_log_success_event: %s", e) @@ -549,25 +518,26 @@ class ModelRateLimitingCheck(CustomLogger): if standard_logging_object is None: return - model_id: Final = standard_logging_object.get("model_id") + model_id: Final[object] = standard_logging_object.get("model_id") if model_id is None: return total_tokens: Final = standard_logging_object.get("total_tokens", 0) - model: Final = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") + model: Final[object] = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") if not model or not total_tokens: return dt: Final = self.clock() - current_minute: Final = dt.strftime("%H-%M") - tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" + tpm: Final = UsageCounter.for_deployment(f"{model_id}", f"{model}", "tpm", dt.strftime("%H-%M")) reset_offset: Final = _reset_offset_from_kwargs(kwargs) # pyright: ignore[reportUnknownArgumentType] # CustomLogger passes an untyped kwargs dict - tpd_key: Final = f"{model_id}:{model}:tpd:{daily_window_day(dt, reset_offset)}" + tpd: Final = UsageCounter.for_deployment( + f"{model_id}", f"{model}", "tpd", daily_window_day(dt, reset_offset) + ) tokens: Final = int(total_tokens) - self._increment(tpm_key, tokens, RoutingArgs.ttl) - self._increment(tpd_key, tokens, RoutingArgs.daily_ttl) + self._increment(tpm, tokens, RoutingArgs.ttl) + self._increment(tpd, tokens, RoutingArgs.daily_ttl) except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.log_success_event: %s", e) diff --git a/litellm/router_utils/routing_prefetch.py b/litellm/router_utils/routing_prefetch.py index 57d76220311..51a5ec3e06a 100644 --- a/litellm/router_utils/routing_prefetch.py +++ b/litellm/router_utils/routing_prefetch.py @@ -1,7 +1,6 @@ from __future__ import annotations -import asyncio -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Awaitable, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Final @@ -10,70 +9,26 @@ if TYPE_CHECKING: from litellm._internal_context import service_target from litellm._logging import verbose_router_logger -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_batch import Batch, BatchResult, active_request_redis_batches +from litellm.caching.redis_batch import Batch, active_request_redis_batches from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, ActiveCooldowns, RouterCooldowns -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter -_PREFETCH_SLOT: Final = "routing_read" +_USAGE_SLOT: Final = "routing_read" _COOLDOWNS_SLOT: Final = "routing_cooldowns" @dataclass(frozen=True, slots=True) -class TierRead: - cache: DualCache - keys: tuple[str, ...] +class ArmedUsage: + counters: tuple[UsageCounter, ...] + result: Awaitable[Sequence[float | None]] -def selector_read(cache: DualCache, keys: Sequence[str]) -> TierRead: - return TierRead(cache, tuple(keys)) - - -@dataclass(frozen=True, slots=True) -class ArmedRead: - read: TierRead - fetched: frozenset[str] - result: BatchResult[Mapping[str, object]] - previous: Mapping[str, float | None] - - def release(self) -> None: - self.read.cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # the selector owns this tier - dict(self.previous) - ) - - async def settle(self, future: asyncio.Future[Mapping[str, object]]) -> None: - if future.cancelled() or future.exception() is not None: - self.release() - return - try: - pending: Final = await self.read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier - list(self.read.keys), local_only=True - ) - values: Final = future.result() - missed: Final = { - key: values[key] - for key, local in zip(self.read.keys, pending.result) - if local is None and values.get(key) is not None - } - await self.read.cache._apply_batch_get(pending, missed) # pyright: ignore[reportPrivateUsage] # the selector owns this tier - except Exception: - self.release() - raise - - -def arm(read: TierRead) -> None: +def arm(store: RouterUsage, counters: Sequence[UsageCounter]) -> None: + """The usage read a usage-based-routing-v2 pick takes later, declared on the request's native pre-call batch.""" request: Final = active_request_redis_batches() - redis: Final = read.cache.redis_cache - if request is None or redis is None or _PREFETCH_SLOT in request.prefetched: + if request is None or not store.has_redis or _USAGE_SLOT in request.prefetched: return - due, previous = read.cache.reserve_redis_batch_reads(read.keys) - if not due: - return - with service_target(ROUTER_USAGE_TARGET): - result: Final = request.batch(redis).mget(tuple(due)) - armed_read: Final = ArmedRead(read=read, fetched=frozenset(due), result=result, previous=previous) - result.on_settled(armed_read.settle) - request.prefetched[_PREFETCH_SLOT] = armed_read + request.prefetched[_USAGE_SLOT] = ArmedUsage(tuple(counters), store.usage(counters, batch=Batch.PRE_CALL)) @dataclass(frozen=True, slots=True) @@ -92,7 +47,7 @@ def arm_cooldowns(store: RouterCooldowns, model_ids: Sequence[str]) -> None: def armed() -> bool: request: Final = active_request_redis_batches() - return request is not None and (_PREFETCH_SLOT in request.prefetched or _COOLDOWNS_SLOT in request.prefetched) + return request is not None and (_USAGE_SLOT in request.prefetched or _COOLDOWNS_SLOT in request.prefetched) async def _armed_cooldowns(model_ids: Sequence[str]) -> ActiveCooldowns | None: @@ -122,43 +77,28 @@ async def read_cooldowns( return await store.active(model_ids, batch=Batch.NONE) -def _take(read: TierRead) -> ArmedRead | None: +async def _armed_usage(counters: Sequence[UsageCounter]) -> Sequence[float | None] | None: request: Final = active_request_redis_batches() - if request is None: - return None - taken: Final = request.prefetched.pop(_PREFETCH_SLOT, None) - if isinstance(taken, ArmedRead) and set(taken.read.keys).issuperset(read.keys): - return taken - if isinstance(taken, ArmedRead): - taken.release() - return None - - -async def _read_prefetched(read: TierRead) -> Sequence[object | None] | None: - prefetch: Final = _take(read) - if prefetch is None: + taken: Final = None if request is None else request.prefetched.pop(_USAGE_SLOT, None) + if not isinstance(taken, ArmedUsage): return None try: - values: Final = await prefetch.result - except Exception as e: # noqa: BLE001 # the read falls back to the tier's own Redis read + values: Final = await taken.result + except Exception as e: # noqa: BLE001 # the read falls back to the store's own Redis read verbose_router_logger.debug("routing prefetch failed, reading again: %s", e) return None - pending: Final = await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier - list(read.keys), local_only=True - ) - if any(key not in prefetch.fetched for key, value in zip(read.keys, pending.result) if value is None): + if not frozenset(taken.counters).issuperset(counters): return None - return await read.cache._apply_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier - pending, {key: values.get(key) for key, value in zip(read.keys, pending.result) if value is None} - ) + read: Final = dict(zip(taken.counters, values)) + return [read[counter] for counter in counters] -async def read(read: TierRead, parent_otel_span: Span | None) -> Sequence[object | None] | None: +async def read( + store: RouterUsage, counters: Sequence[UsageCounter], parent_otel_span: Span | None +) -> Sequence[float | None]: + """The armed usage read when it covers ``counters``, taken once; otherwise the store reads on its own.""" with service_target(ROUTER_USAGE_TARGET): - prefetched: Final = await _read_prefetched(read) - if prefetched is not None: - return prefetched - results: Final = await DualCache.async_batch_get_cache_shared( - ((read.cache, list(read.keys)),), parent_otel_span=parent_otel_span - ) - return results[0] + prefetched: Final = await _armed_usage(counters) + if prefetched is not None: + return prefetched + return await store.usage(counters, batch=Batch.NONE) diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index 8f54a96e305..2d5aed65883 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -5,10 +5,10 @@ from contextvars import ContextVar from typing import TYPE_CHECKING, Final from litellm._logging import verbose_router_logger +from litellm.redis.stores.router_usage import UsageCounter from litellm.router_deprecated import renamed_router_param from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage from litellm.router_utils import routing_prefetch -from litellm.router_utils.routing_prefetch import TierRead, selector_read if TYPE_CHECKING: from opentelemetry.trace import Span @@ -18,13 +18,9 @@ if TYPE_CHECKING: __all__ = ("RoutingPrefetch", "RoutingReadBatch") -def _selector_keys(selector: LowestTPMLoggingHandler_v2, deployments: Sequence[object]) -> list[str]: - tpm_keys, rpm_keys = selector.usage_counter_keys(list(deployments)) - return tpm_keys + rpm_keys - - -def _selector_read(selector: LowestTPMLoggingHandler_v2 | None, deployments: Sequence[object]) -> TierRead | None: - return None if selector is None else selector_read(selector.router_cache, _selector_keys(selector, deployments)) +def _selector_counters(selector: LowestTPMLoggingHandler_v2, deployments: Sequence[object]) -> tuple[UsageCounter, ...]: + tpm_counters, rpm_counters = selector.usage_counters(deployments) + return tpm_counters + rpm_counters class RoutingPrefetch: @@ -34,9 +30,8 @@ class RoutingPrefetch: ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: Sequence[object] ) -> None: RoutingPrefetch.arm_cooldowns(ctx) - read: Final = _selector_read(usage_selector, deployments) - if read is not None: - routing_prefetch.arm(read) + if usage_selector is not None: + routing_prefetch.arm(usage_selector.usage, _selector_counters(usage_selector, deployments)) @staticmethod def arm_cooldowns(ctx: "RoutingContext") -> None: @@ -92,11 +87,9 @@ class RoutingReadBatch: async def read_usage(self, deployments: Sequence[object], parent_otel_span: "Span | None") -> None: """The selector's tpm/rpm counters for ``deployments``, kept as ``prefetched_usage``.""" - usage_read: Final = _selector_read(self.usage_selector, deployments) - if usage_read is None: + selector: Final = self.usage_selector + if selector is None: return - values: Final = await routing_prefetch.read(usage_read, parent_otel_span) - self.prefetched_usage = PrefetchedUsage( - keys=frozenset(usage_read.keys), - values=None if values is None else dict(zip(usage_read.keys, values)), - ) + counters: Final = _selector_counters(selector, deployments) + values: Final = await routing_prefetch.read(selector.usage, counters, parent_otel_span) + self.prefetched_usage = PrefetchedUsage(counters=frozenset(counters), values=dict(zip(counters, values))) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index f4dfd69534e..0248a3a4ebf 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -113,14 +113,24 @@ class NativeSensitiveRoutePins: def set(self, tenant: str, session_id: str, model: str, ttl_seconds: float) -> Future[None]: ... class _UsageCounter(Protocol): + @property + def scope(self) -> Literal["global_router", "deployment"]: ... @property def deployment_id(self) -> str: ... @property def model(self) -> str | None: ... @property - def minute(self) -> str: ... + def window(self) -> str: ... @property - def kind(self) -> Literal["tpm", "rpm"]: ... + def kind(self) -> Literal["tpm", "rpm", "itpm", "otpm", "tpd", "rpd"]: ... + +class _CounterIncrement(Protocol): + @property + def counter(self) -> _UsageCounter: ... + @property + def amount(self) -> float: ... + @property + def ttl(self) -> float | None: ... @final class NativeRouterUsage: @@ -130,7 +140,7 @@ class NativeRouterUsage: def has_redis(self) -> bool: ... def increment( self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> Future[float]: ... + ) -> Future[float | None]: ... def increment_deployment( self, deployment_id: str, @@ -144,7 +154,7 @@ class NativeRouterUsage: ) -> Future[list[float | None]]: ... def increment_sync( self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> float: ... + ) -> float | None: ... def increment_deployment_sync( self, deployment_id: str, @@ -156,6 +166,24 @@ class NativeRouterUsage: def get_many_sync( self, counters: Sequence[_UsageCounter] ) -> list[float | None]: ... + def usage( + self, counters: Sequence[_UsageCounter], send: NativeSend + ) -> NativePending[list[float | None]]: ... + def increment_pipeline( + self, increments: Sequence[_CounterIncrement] + ) -> Future[list[float] | None]: ... + def increment_post_call( + self, request: NativeRequestRedis, increments: Sequence[_CounterIncrement] + ) -> None: ... + def increment_remote( + self, increments: Sequence[_CounterIncrement] + ) -> Future[list[float] | None]: ... + def increment_local(self, increment: _CounterIncrement) -> float: ... + def get_local(self, counter: _UsageCounter) -> float | None: ... + def get_local_many(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ... + def set_local(self, counter: _UsageCounter, value: float) -> None: ... + def current(self, counter: _UsageCounter, limit: float) -> Future[float | None]: ... + def current_sync(self, counter: _UsageCounter, limit: float) -> float | None: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 41d508a0eea..18f4834bb86 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -36,10 +36,12 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement from litellm.proxy.utils import InternalUsageCache from litellm.redis.stores.rate_limits import CounterIncrement, RateLimitScope +from litellm.redis.stores.router_usage import CounterIncrement as UsageIncrement from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ModelResponse from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache, fake_send @@ -132,10 +134,10 @@ def _response_cache(redis_cache: FakeRedisCache) -> Cache: return cache -def _tpm_router(redis_cache: FakeRedisCache) -> tuple[LowestTPMLoggingHandler_v2, DualCache]: - router_cache = DualCache() - router_cache.attach_redis_cache(redis_cache) - return LowestTPMLoggingHandler_v2(router_cache=router_cache, routing_args={"ttl": 60}), router_cache +def _tpm_selector(handle: FakeUsageHandle) -> LowestTPMLoggingHandler_v2: + selector: Final = LowestTPMLoggingHandler_v2(usage=usage_store(handle), routing_args={"ttl": 60}) + selector.cancel_sync_task() + return selector def _tpm_kwargs() -> Mapping[str, object]: @@ -151,19 +153,19 @@ def _tpm_kwargs() -> Mapping[str, object]: @pytest.mark.asyncio -async def test_post_call_owners_wait_for_the_callbacks_and_the_limiter_declares_on_the_native_post_call_batch(): +async def test_post_call_owners_declare_on_the_native_post_call_batch_until_the_callbacks_finish(): client = FakeClient(_ok_replies) redis_cache = PostCallFakeRedisCache(client) calls: list[tuple[str, list[str], list[object]]] = [] # mutable-ok: filled by the fake scripts fakes = FakeRateLimits(token_increment_script=_scripts(calls, [10.0]), parallel_release_script=_scripts(calls, [0])) limiter = _limiter(fakes) - tpm, router_cache = _tpm_router(redis_cache) + usage = FakeUsageHandle(redis={}) + tpm = _tpm_selector(usage) with request_redis_batch_scope(): await tpm.async_log_success_event(_tpm_kwargs(), None, None, None) await limiter.async_increment_tokens_with_ttl_preservation(_token_ops(_K1)) await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", _K1), None, in_logging_callback=True) - assert client.pipelines == [] # nothing goes out while the callbacks are still declaring await flush_post_call_redis_batches() assert fakes.post_call == ["token_increment", "parallel_release"] @@ -171,15 +173,22 @@ async def test_post_call_owners_wait_for_the_callbacks_and_the_limiter_declares_ ("script", ["{api_key:k1}:tokens"], [10, 60]), ("script", ["{api_key:k1}:max_parallel_requests"], ["slot-1"]), ] - assert len(client.pipelines) == 1 - assert _names(client) == ["INCRBYFLOAT", "EXPIRE"] + tpm_counter = next(iter(usage.local)) + assert usage.calls == [("post_call", (UsageIncrement(tpm_counter, 42, 60),))] + assert (tpm_counter.deployment_id, tpm_counter.model, tpm_counter.kind) == ("dep-a", "openai/gpt-4o-mini", "tpm") + assert usage.local == {tpm_counter: 42} + assert client.pipelines == [] assert redis_cache.alone == [] - assert ( - await router_cache.in_memory_cache.async_get_cache( - next(k for k in router_cache.in_memory_cache.cache_dict if ":tpm:" in k) - ) - == 42 - ) + + +@pytest.mark.asyncio +async def test_outside_a_request_the_tpm_increment_goes_out_on_its_own_pipeline(): + usage = FakeUsageHandle(redis={}) + + await _tpm_selector(usage).async_log_success_event(_tpm_kwargs(), None, None, None) + + assert [name for name, _ in usage.calls] == ["pipeline"] + assert usage.local == usage.redis == {next(iter(usage.local)): 42} @pytest.mark.asyncio diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index b4268fe7bb0..32b63341a8b 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -4,16 +4,14 @@ router's usage counter read, auth identity and spend counters all join the reque from __future__ import annotations import asyncio -import json from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime, timezone -from typing import Any, Final -from unittest.mock import AsyncMock, MagicMock +from typing import Any, Final, Literal +from unittest.mock import AsyncMock, MagicMock, patch import pytest -import litellm.caching.dual_cache as dual_cache_module from litellm import Router from litellm._internal_context import current_service_target from litellm.caching.dual_cache import DualCache @@ -30,10 +28,12 @@ from litellm.proxy.utils import InternalUsageCache from litellm.redis.stores._route import route from litellm.redis.stores.rate_limits import RateLimitScope from litellm.redis.stores.router_cooldowns import ActiveCooldowns, RouterCooldownStore -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsageStore, UsageCounter from litellm.router_api import RoutingContext -from litellm.router_utils.routing_prefetch import ArmedRead, TierRead, arm_cooldowns, read_cooldowns -from litellm.router_utils.routing_read_batch import RoutingPrefetch +from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage +from litellm.router_utils import routing_prefetch +from litellm.router_utils.routing_prefetch import arm_cooldowns, read_cooldowns +from litellm.router_utils.routing_read_batch import RoutingReadBatch from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits @@ -43,6 +43,7 @@ from .test_redis_batch import ( FakeNativeRequest, FakeNatives, FakeRedisCache, + fake_pending, fake_send, ) @@ -323,43 +324,8 @@ def _router(redis_cache: FakeRedisCache, routing_strategy: str = "usage-based-ro return router -_PINNED_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) -_START: Final = 1_000_000.0 _PING: Final = [{"role": "user", "content": "ping"}] - - -@pytest.fixture -def pinned_clock(monkeypatch: pytest.MonkeyPatch) -> list[float]: - """Pins the usage counters' minute and the router cache's read throttle clock; tests move `now[0]`.""" - now: Final = [_START] # mutable-ok: the test advances the pinned clock - monkeypatch.setattr("litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", lambda: _PINNED_MINUTE) - monkeypatch.setattr(dual_cache_module.time, "time", lambda: now[0]) - return now - - -def _counter_keys(router: Router, deployment_ids: tuple[str, ...] = ("dep-a", "dep-b")) -> frozenset[str]: - deployments: Final = [d for d in router.model_list if d["model_info"]["id"] in deployment_ids] - tpm, rpm = router.lowesttpm_logger_v2.usage_counter_keys(deployments) - return frozenset(tpm + rpm) - - -def _tpm_key(router: Router, deployment_id: str) -> str: - deployment: Final = next(d for d in router.model_list if d["model_info"]["id"] == deployment_id) - return router.lowesttpm_logger_v2.usage_counter_keys([deployment])[0][0] - - -def _counter_mgets(router: Router, client: FakeClient, redis_cache: FakeRedisCache) -> tuple[frozenset[str], ...]: - counters: Final = _counter_keys(router) - pipelined: Final = tuple( - frozenset(command[1:]) - for pipeline in client.pipelines - for command in pipeline.commands - if command[0] == "MGET" and counters.intersection(command[1:]) - ) - alone: Final = tuple( - frozenset(keys) for command, keys in redis_cache.alone if command == "MGET" and counters.intersection(keys) - ) - return pipelined + alone +_PINNED_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) def _python_cooldown_reads(client: FakeClient, redis_cache: FakeRedisCache) -> tuple[str, ...]: @@ -374,15 +340,6 @@ def _python_cooldown_reads(client: FakeClient, redis_cache: FakeRedisCache) -> t return tuple(key for key in pipelined + alone if key.endswith(":cooldown")) -def _replies_with(values: dict[str, object]) -> Callable[[tuple[Any, ...]], Any]: - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": - return [None if values.get(key) is None else json.dumps(values[key]) for key in command[1:]] - return _lua_ok_replies(command) - - return replies - - async def _route(router: Router) -> str: deployment: Final = await router.async_get_available_deployment( model=_MODEL_GROUP, messages=_PING, request_kwargs={} @@ -390,23 +347,136 @@ async def _route(router: Router) -> str: return deployment["model_info"]["id"] +def _counter(deployment_id: str, kind: Literal["tpm", "rpm"]) -> UsageCounter: + return UsageCounter.for_deployment(deployment_id, "anthropic/claude-x", kind, "12-00") + + +_DEPLOYMENT_COUNTERS: Final = tuple(_counter(d, kind) for kind in ("tpm", "rpm") for d in ("dep-a", "dep-b")) + + +@dataclass +class _FakeUsageHandle: + """A connected usage store whose ``usage`` reads land on the request's fake native batch, or alone for + ``NativeSend.Now()``, recording what each read asked for and under which service target.""" + + counts: dict[UsageCounter, float] + fail: Exception | None = None + reads: list[tuple[str, tuple[UsageCounter, ...], str | None]] = field(default_factory=list) + has_redis: bool = True + + def usage(self, counters: Sequence[UsageCounter], send: object) -> object: + batched: Final = isinstance(send, FakeNativeRequest) + self.reads.append(("batch" if batched else "alone", tuple(counters), current_service_target())) + if batched: + send.events.append("declare usage") + + fail: Final = self.fail + + def answer() -> list[float | None]: + if fail is not None: + raise fail + return [self.counts.get(counter) for counter in counters] + + return fake_pending(send, answer) + + def get_many_sync(self, counters: Sequence[UsageCounter]) -> list[float | None]: + raise AssertionError("a connected store reads through usage()") + + +def _usage_store(handle: _FakeUsageHandle) -> RouterUsageStore: + return RouterUsageStore(handle) # pyright: ignore[reportArgumentType] # test double + + @pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_armed_routing_read_is_one_pipeline_and_routing_issues_no_read_of_its_own(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) +async def test_the_armed_usage_read_and_the_armed_cooldown_read_go_out_in_one_native_pipeline(): + handle: Final = _FakeUsageHandle({_counter("dep-a", "tpm"): 40}) + usage: Final = _usage_store(handle) + cooldowns: Final = RouterCooldownStore(_FakeCooldownsHandle()) + natives: Final = FakeNatives() - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment = await _route(router) + with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + arm_cooldowns(cooldowns, ["dep-a", "dep-b"]) + routing_prefetch.arm(usage, _DEPLOYMENT_COUNTERS) + active, counts = await asyncio.gather( + read_cooldowns(cooldowns, ["dep-a", "dep-b"], None), + routing_prefetch.read(usage, _DEPLOYMENT_COUNTERS, None), + ) - assert deployment in {"dep-a", "dep-b"} - assert len(client.pipelines) == 1 - commands = client.pipelines[0].commands - assert [c[0] for c in commands] == ["MGET"] - assert frozenset(commands[0][1:]) == _counter_keys(router) - assert redis_cache.alone == [] + assert active == () + assert counts == [40, None, None, None] + assert _declared_before_the_first_flush(natives) == ["declare cooldowns", "declare usage"] + assert [kind for kind, _, _ in handle.reads] == ["batch"] + assert request.prefetched == {} + + +@pytest.mark.asyncio +async def test_the_armed_usage_read_is_declared_under_the_router_usage_family(): + handle: Final = _FakeUsageHandle({}) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm(_usage_store(handle), _DEPLOYMENT_COUNTERS) + + assert [target for _, _, target in handle.reads] == [ROUTER_USAGE_TARGET] + + +@pytest.mark.asyncio +async def test_a_usage_prefetch_that_does_not_cover_the_pick_is_dropped_and_the_store_reads_on_its_own(): + handle: Final = _FakeUsageHandle({_counter("dep-b", "rpm"): 3}) + usage: Final = _usage_store(handle) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm(usage, _DEPLOYMENT_COUNTERS[:1]) + counts: Final = await routing_prefetch.read(usage, _DEPLOYMENT_COUNTERS, None) + assert request.prefetched == {} + + assert counts == [None, None, None, 3] + assert [(kind, len(counters)) for kind, counters, _ in handle.reads] == [("batch", 1), ("alone", 4)] + + +@pytest.mark.asyncio +async def test_a_failed_usage_prefetch_falls_back_to_the_stores_own_read(): + handle: Final = _FakeUsageHandle({_counter("dep-a", "tpm"): 9}, fail=ConnectionError("redis down")) + usage: Final = _usage_store(handle) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm(usage, _DEPLOYMENT_COUNTERS) + handle.fail = None + counts: Final = await routing_prefetch.read(usage, _DEPLOYMENT_COUNTERS, None) + + assert counts == [9, None, None, None] + assert [kind for kind, _, _ in handle.reads] == ["batch", "alone"] + + +@pytest.mark.asyncio +async def test_arming_a_store_without_redis_declares_nothing(): + handle: Final = _FakeUsageHandle({}, has_redis=False) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm(_usage_store(handle), _DEPLOYMENT_COUNTERS) + + assert request.prefetched == {} + assert handle.reads == [] + + +@pytest.mark.asyncio +async def test_a_usage_based_pick_reads_the_counters_the_prefetch_armed(): + handle: Final = _FakeUsageHandle({_counter("dep-a", "tpm"): 100}) + selector: Final = LowestTPMLoggingHandler_v2(usage=_usage_store(handle)) + deployments: Final[list[object]] = [_deployment("dep-a"), _deployment("dep-b")] + read_batch: Final = RoutingReadBatch(usage_selector=selector) + + with ( + patch("litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", return_value=_PINNED_MINUTE), + request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send), # pyright: ignore[reportArgumentType] # test double + ): + routing_prefetch.arm(selector.usage, _DEPLOYMENT_COUNTERS) + await read_batch.read_usage(deployments, None) + with PrefetchedUsage.scoped(read_batch.prefetched_usage): + picked: Final = await selector.async_get_available_deployments(_MODEL_GROUP, deployments, messages=_PING) + + selector.cancel_sync_task() + assert picked["model_info"]["id"] == "dep-b" + assert [kind for kind, _, _ in handle.reads] == ["batch"] @pytest.mark.asyncio @@ -431,158 +501,6 @@ async def test_a_cooldown_recorded_through_the_store_excludes_its_deployment_wit assert _python_cooldown_reads(client, redis_cache) == () -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_routing_reads_itself(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - armed = request.prefetched["routing_read"] - assert isinstance(armed, ArmedRead) - request.prefetched["routing_read"] = ArmedRead( - read=TierRead(armed.read.cache, ("other",)), - fetched=armed.fetched, - result=armed.result, - previous=armed.previous, - ) - deployment = await _route(router) - assert request.prefetched == {} - - assert deployment in {"dep-a", "dep-b"} - assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_a_prefetch_with_incomplete_usage_keys_releases_its_reservations(): - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - - with request_redis_batch_scope(): - RoutingPrefetch.arm(RoutingContext(router), router.lowesttpm_logger_v2, router.model_list[:1]) - deployment: Final = await _route(router) - - assert deployment in {"dep-a", "dep-b"} - assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_a_failed_prefetch_falls_back_to_the_selectors_own_read(): - client = FakeClient(_lua_ok_replies, fail=ConnectionError("redis down")) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment = await _route(router) - - assert deployment in {"dep-a", "dep-b"} - assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_the_armed_routing_read_is_declared_under_the_router_usage_family(): - """The prefetch is declared before routing runs under a target of its own, so the pipeline that - carries it renders ``redis.pipeline router_usage`` instead of a bare ``redis.pipeline``.""" - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - pipeline_targets: list[str | None] = [] # mutable-ok: filled by the recording hook - - async def record(**kwargs: object) -> None: - pipeline_targets.append(current_service_target()) - - redis_cache.service_logger_obj.async_service_success_hook = record # pyright: ignore[reportAttributeAccessIssue] # fake, records the hook call - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - async with asyncio.timeout(5): - while not pipeline_targets: - await asyncio.sleep(0.01) - - assert pipeline_targets == [ROUTER_USAGE_TARGET] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_an_abandoned_prefetch_still_backfills_the_counters_it_read(): - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100}) - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await request.flush_all() - after_abandoned: Final = len(_counter_mgets(router, client, redis_cache)) - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment: Final = await _route(router) - - assert deployment == "dep-b" - assert after_abandoned == 1 - assert len(_counter_mgets(router, client, redis_cache)) == 1 - assert redis_cache.alone == [] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses(): - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - dep_a: Final = _tpm_key(router, "dep-a") - dep_b: Final = _tpm_key(router, "dep-b") - client.reply_for = _replies_with({dep_a: 5, dep_b: 7}) - usage_cache: Final = router.lowesttpm_logger_v2.router_cache - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - usage_cache.in_memory_cache.set_cache(dep_a, 50) - await request.flush_all() - - assert await usage_cache.async_batch_get_cache([dep_a, dep_b], local_only=True) == [50, 7] - - -@pytest.mark.asyncio -async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(pinned_clock: list[float]): - mget_replies: Final = iter((ConnectionError("redis down"), None)) - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": - failure: Final = next(mget_replies) - if isinstance(failure, Exception): - return failure - return ["100" if key == _tpm_key(router, "dep-a") else None for key in command[1:]] - return _lua_ok_replies(command) - - client.reply_for = replies - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await request.flush_all() - after_failure: Final = len(_counter_mgets(router, client, redis_cache)) - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment: Final = await _route(router) - - assert after_failure == 1 - assert len(_counter_mgets(router, client, redis_cache)) == 2 - assert deployment == "dep-b" - - @pytest.mark.asyncio async def test_arming_outside_a_request_scope_is_a_no_op(): redis_cache = FakeRedisCache(FakeClient(_lua_ok_replies)) @@ -607,96 +525,6 @@ async def test_simple_shuffle_cooldown_prefetch_does_not_arm_usage_reads_or_read assert redis_cache.alone == [] -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_requests_within_the_read_interval_read_the_counters_from_redis_once(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - request_round_trips: list[tuple[int, int]] = [] # mutable-ok: one entry per simulated request - - for _ in range(3): - pipeline_count = len(client.pipelines) - alone_count = len(redis_cache.alone) - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - request_round_trips.append((len(client.pipelines) - pipeline_count, len(redis_cache.alone) - alone_count)) - - assert len(_counter_mgets(router, client, redis_cache)) == 1 - assert request_round_trips[1:] == [(0, 0), (0, 0)] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_concurrent_requests_share_one_counter_read_per_interval(): - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - first_armed: Final = asyncio.Event() - both_armed: Final = asyncio.Event() - - async def route_after_both_requests_arm() -> str: - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - if first_armed.is_set(): - both_armed.set() - else: - first_armed.set() - await both_armed.wait() - return await _route(router) - - deployments: Final = await asyncio.gather(route_after_both_requests_arm(), route_after_both_requests_arm()) - - assert set(deployments) <= {"dep-a", "dep-b"} - assert len(_counter_mgets(router, client, redis_cache)) == 1 - - -@pytest.mark.asyncio -async def test_the_prefetch_reads_the_counters_again_once_the_read_interval_elapses(pinned_clock: list[float]): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - - client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100}) - pinned_clock[0] = _START + router.lowesttpm_logger_v2.router_cache.redis_batch_cache_expiry + 1 - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment = await _route(router) - - assert len(_counter_mgets(router, client, redis_cache)) == 2 - assert deployment == "dep-b" - - -@pytest.mark.asyncio -async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(pinned_clock: list[float]): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - usage_cache = router.lowesttpm_logger_v2.router_cache - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - - for key in _counter_keys(router, ("dep-a",)): - usage_cache.in_memory_cache.set_cache(key, 1) - pinned_clock[0] = _START + usage_cache.redis_batch_cache_expiry + 1 - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - - second_pipeline_mgets = tuple(command for command in client.pipelines[1].commands if command[0] == "MGET") - - assert len(client.pipelines) == 2 - assert len(second_pipeline_mgets) == 1 - assert frozenset(second_pipeline_mgets[0][1:]) == _counter_keys(router, ("dep-b",)) - - @pytest.mark.asyncio async def test_two_backends_flush_concurrently_one_pipeline_each(): a_client, b_client = FakeClient(_lua_ok_replies), FakeClient(_lua_ok_replies) diff --git a/tests/unit/redis/stores/router_usage_double.py b/tests/unit/redis/stores/router_usage_double.py new file mode 100644 index 00000000000..416db232e65 --- /dev/null +++ b/tests/unit/redis/stores/router_usage_double.py @@ -0,0 +1,133 @@ +"""A stand-in for ``NativeRouterUsage`` with the native store's semantics: memory in front of a Redis dict. +``redis`` None is no Redis, ``breaker_open`` refuses every Redis round trip and ``redis_down`` fails it, and every +call the store makes is logged in ``calls``.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Final, TypeVar + +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter +from litellm.router_python_host import PythonHost +from litellm.rust_bridge._native import NativeRedisClient +from tests.unit.caching.test_redis_batch import fake_pending + +T = TypeVar("T") + + +def _done(value: T) -> asyncio.Future[T]: + future: Final[asyncio.Future[T]] = asyncio.get_running_loop().create_future() + future.set_result(value) + return future + + +@dataclass +class FakeUsageHandle: + local: dict[UsageCounter, float] = field(default_factory=dict) + redis: dict[UsageCounter, float] | None = None + breaker_open: bool = False + redis_down: bool = False + calls: list[tuple[str, tuple[object, ...]]] = field(default_factory=list) + + @property + def has_redis(self) -> bool: + return self.redis is not None + + def attach(self, client: NativeRedisClient, /) -> bool: + attached: Final = self.redis is None + self.redis = {} if self.redis is None else self.redis + return attached + + def _reachable(self) -> dict[UsageCounter, float] | None: + return None if self.breaker_open or self.redis_down else self.redis + + def _increment(self, increment: CounterIncrement) -> float: + self.local[increment.counter] = self.local.get(increment.counter, 0) + increment.amount + redis: Final = self._reachable() + if redis is None: + return self.local[increment.counter] + redis[increment.counter] = redis.get(increment.counter, 0) + increment.amount + return redis[increment.counter] + + def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: + self.calls.append(("increment", (counter, value, ttl_seconds))) + return self._increment(CounterIncrement(counter, value, ttl_seconds)) + + def increment(self, counter: UsageCounter, value: float, ttl_seconds: int) -> asyncio.Future[float]: + return _done(self.increment_sync(counter, value, ttl_seconds)) + + def increment_pipeline(self, increments: Sequence[CounterIncrement]) -> asyncio.Future[list[float]]: + self.calls.append(("pipeline", tuple(increments))) + return _done([self._increment(increment) for increment in increments]) + + def increment_post_call(self, request: object, increments: Sequence[CounterIncrement]) -> None: + """Memory now; the Redis half is the request's, so only the declaration is logged.""" + self.calls.append(("post_call", tuple(increments))) + for increment in increments: + self.local[increment.counter] = self.local.get(increment.counter, 0) + increment.amount + + def increment_local(self, increment: CounterIncrement) -> float: + self.local[increment.counter] = self.local.get(increment.counter, 0) + increment.amount + return self.local[increment.counter] + + def get_local(self, counter: UsageCounter) -> float | None: + return self.local.get(counter) + + def get_local_many(self, counters: Sequence[UsageCounter]) -> list[float | None]: + return [self.local.get(counter) for counter in counters] + + def set_local(self, counter: UsageCounter, value: float) -> None: + self.local[counter] = value + + def increment_remote(self, increments: Sequence[CounterIncrement]) -> asyncio.Future[list[float] | None]: + self.calls.append(("remote", tuple(increments))) + redis: Final = self._reachable() + if redis is None: + return _done(None) + totals: Final = [redis.get(i.counter, 0) + i.amount for i in increments] + redis.update({i.counter: total for i, total in zip(increments, totals)}) + return _done(totals) + + def current_sync(self, counter: UsageCounter, limit: float) -> float | None: + self.calls.append(("current", (counter, limit))) + local: Final = self.local.get(counter) + if self.redis is None or self.breaker_open or (local is not None and local >= limit): + return local + return None if self.redis_down else self.redis.get(counter) + + def current(self, counter: UsageCounter, limit: float) -> asyncio.Future[float | None]: + return _done(self.current_sync(counter, limit)) + + def _read(self, counters: Sequence[UsageCounter]) -> list[float | None]: + redis: Final = self._reachable() + return [self.local.get(c, None if redis is None else redis.get(c)) for c in counters] + + def get_many_sync(self, counters: Sequence[UsageCounter]) -> list[float | None]: + self.calls.append(("get_many", tuple(counters))) + return self._read(counters) + + def get_many(self, counters: Sequence[UsageCounter]) -> asyncio.Future[list[float | None]]: + return _done(self.get_many_sync(counters)) + + def usage(self, counters: Sequence[UsageCounter], send: object) -> object: + """Declared on the fake request's batch, or alone for ``NativeSend.Now()``.""" + self.calls.append(("usage", tuple(counters))) + return fake_pending(send, lambda: self._read(counters)) + + +def usage_store(handle: FakeUsageHandle) -> RouterUsageStore: + return RouterUsageStore(handle) # pyright: ignore[reportArgumentType] # implements what the tests call + + +class UsageHost(PythonHost): + """The Python host with a router's usage store counting in ``handle``.""" + + __slots__ = ("_usage",) + + def __init__(self, handle: FakeUsageHandle) -> None: + self._usage: Final = usage_store(handle) + + def router_usage(self, redis: object) -> RouterUsageStore: + return self._usage diff --git a/tests/unit/redis/stores/test_router_usage.py b/tests/unit/redis/stores/test_router_usage.py index 5b46f590c4d..54891c22a8f 100644 --- a/tests/unit/redis/stores/test_router_usage.py +++ b/tests/unit/redis/stores/test_router_usage.py @@ -10,7 +10,7 @@ from litellm._internal_context import current_service_target from litellm.caching.redis_batch import drain_post_call_redis_batches, request_redis_batch_scope from litellm.caching.redis_cache import RedisCache from litellm.redis.stores.connection import connect -from litellm.redis.stores.router_usage import RouterUsageStore, UsageCounter +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.rust_bridge._native import NativeRedisClient from litellm.types.utils import ModelResponse from tests.unit.caching.test_redis_batch import FakeNatives, fake_send @@ -52,8 +52,8 @@ class FakeUsage: requests: int, ) -> None: self.calls.append(("deployment", (deployment_id, model, minute, tokens, requests), current_service_target())) - tpm: Final = UsageCounter(deployment_id, model, minute, "tpm") - rpm: Final = UsageCounter(deployment_id, model, minute, "rpm") + tpm: Final = UsageCounter("global_router", deployment_id, model, minute, "tpm") + rpm: Final = UsageCounter("global_router", deployment_id, model, minute, "rpm") self.counts[tpm] = self.counts.get(tpm, 0) + tokens self.counts[rpm] = self.counts.get(rpm, 0) + requests @@ -84,8 +84,8 @@ class FakeUsage: @pytest.mark.asyncio async def test_counters_cross_the_bridge_as_domain_tuples_under_the_router_usage_target() -> None: handle: Final = FakeUsage() - store: Final = RouterUsageStore(handle) - counter: Final = UsageCounter("dep", None, "12-34", "rpm") + store: Final = RouterUsageStore(handle) # pyright: ignore[reportArgumentType] # implements what the test calls + counter: Final = UsageCounter("global_router", "dep", None, "12-34", "rpm") await store.increment(counter, 1, 60) await store.increment_deployment("dep", "model", "12-34", 7.5, 1) @@ -94,13 +94,35 @@ async def test_counters_cross_the_bridge_as_domain_tuples_under_the_router_usage ("increment", (counter, 1, 60), "router_usage"), ("deployment", ("dep", "model", "12-34", 7.5, 1), "router_usage"), ] - assert await store.get_many((counter, UsageCounter("dep", "model", "12-34", "tpm"))) == [1, 7.5] + assert await store.get_many((counter, UsageCounter("global_router", "dep", "model", "12-34", "tpm"))) == [1, 7.5] + + +@pytest.mark.parametrize( + ("counter", "key"), + [ + (UsageCounter("global_router", "dep", None, "09-05", "rpm"), "global_router:dep:None:rpm:09-05"), + ( + UsageCounter.for_deployment("dep", "openai/gpt-5.4-nano", "rpm", "12-34"), + "dep:openai/gpt-5.4-nano:rpm:12-34", + ), + (UsageCounter.for_deployment(f"{7}", f"{None}", "tpd", "2026-11-10"), "7:None:tpd:2026-11-10"), + ], +) +def test_counters_count_under_mains_keys(counter: UsageCounter, key: str) -> None: + store: Final = RouterUsageStore(connect(None).router_usage()) + + store.increment_local(CounterIncrement(counter, 3, 60)) + + assert store.snapshot() == {key: 3} @pytest.mark.asyncio async def test_late_attach_keeps_the_counts_this_router_already_holds() -> None: store: Final = RouterUsageStore(connect(None).router_usage()) - counters: Final = (UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm")) + counters: Final = ( + UsageCounter("global_router", "dep", "model", "12-34", "tpm"), + UsageCounter("global_router", "dep", "model", "12-34", "rpm"), + ) await store.increment_deployment("dep", "model", "12-34", 7.5, 1) store.attach_redis(connect(RedisCache(host="redis.internal"))) assert await store.get_many(counters) == [7.5, 1] @@ -132,7 +154,7 @@ async def test_router_reset_clears_usage_and_failure_counters(operation: str) -> @pytest.mark.asyncio async def test_inside_a_request_deployment_increments_go_out_on_their_own_trip_like_main(): handle: Final = FakeUsage() - store: Final = RouterUsageStore(handle) + store: Final = RouterUsageStore(handle) # pyright: ignore[reportArgumentType] # implements what the test calls natives: Final = FakeNatives() await store.increment_deployment("dep", "model", "12-34", 7.5, 1) diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 7e044b55f5a..afeea55291c 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1775,6 +1775,14 @@ def _turn_completed_reading_cache(cached_tokens: int) -> str: ) +def _rate_limit_check(dual_cache: "DualCache") -> "ModelRateLimitingCheck": + from litellm.redis.stores.connection import connect + from litellm.redis.stores.router_usage import RouterUsageStore + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + return ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) + + def _connection_logged_through(rate_limits: "ModelRateLimitingCheck", opened_with: Mapping[str, object]) -> "Logging": import time @@ -2021,7 +2029,7 @@ class TestWebSocketTurnAdmission: websocket=_ClientSocket(second, third), backend_ws=backend_ws, request_data={}, - logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + logging_obj=_connection_logged_through(_rate_limit_check(dual_cache), opened_with), first_message=first, turn_admission=_later_turns_admitted_through(dual_cache), ) @@ -2065,7 +2073,7 @@ class TestWebSocketTurnAdmission: websocket=_ClientSocket(), backend_ws=backend_ws, request_data={}, - logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + logging_obj=_connection_logged_through(_rate_limit_check(dual_cache), opened_with), first_message=first, ) @@ -2096,7 +2104,7 @@ class TestWebSocketTurnAdmission: await _run_managed_relay( monkeypatch, _TwoOutputTokenTurns(failing_turns), - ModelRateLimitingCheck(dual_cache), + _rate_limit_check(dual_cache), _turn_frames(), _later_turns_admitted_through(dual_cache), opened_with, @@ -2126,7 +2134,7 @@ class TestWebSocketTurnAdmission: await _relay_turns_to( monkeypatch, "openai/gpt-5-mini", - ModelRateLimitingCheck(dual_cache), + _rate_limit_check(dual_cache), tuple( json.dumps({"type": "response.create", "input": word, "metadata": client_metadata}) for word in ("one", "two") diff --git a/tests/unit/router_strategy/test_base_routing_strategy.py b/tests/unit/router_strategy/test_base_routing_strategy.py index 19ff11dd3f6..21e39a06031 100644 --- a/tests/unit/router_strategy/test_base_routing_strategy.py +++ b/tests/unit/router_strategy/test_base_routing_strategy.py @@ -1,172 +1,108 @@ -import json +import asyncio import logging -from typing import Any, Dict, Final, List, Optional, Set, Union +import threading +from collections.abc import Sequence +from typing import Final import pytest - - -import asyncio -import threading -from unittest.mock import MagicMock, patch - from fakeredis import TcpFakeServer - -from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy +from litellm.rust_bridge._native import NativeRedisStoreError +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store + +_A: Final = UsageCounter.for_deployment("dep-a", "openai/gpt-5.4-nano", "rpm", "12-34") +_B: Final = UsageCounter.for_deployment("dep-b", "openai/gpt-5.4-nano", "rpm", "12-34") -@pytest.fixture -async def mock_dual_cache(): - dual_cache = MagicMock(spec=DualCache) - dual_cache.in_memory_cache = MagicMock() - dual_cache.redis_cache = MagicMock() - - # Set up async method mocks to return coroutines - future1: asyncio.Future[None] = asyncio.Future() - future1.set_result(None) - dual_cache.in_memory_cache.async_increment.return_value = future1 - - future2: asyncio.Future[None] = asyncio.Future() - future2.set_result(None) - dual_cache.redis_cache.async_increment_pipeline.return_value = future2 - - future3: asyncio.Future[None] = asyncio.Future() - future3.set_result(None) - dual_cache.in_memory_cache.async_set_cache.return_value = future3 - - # Fix for async_batch_get_cache - batch_future: asyncio.Future[Dict[str, str]] = asyncio.Future() - batch_future.set_result({"key1": "10.0", "key2": "20.0"}) - dual_cache.redis_cache.async_batch_get_cache.return_value = batch_future - - return dual_cache - - -@pytest.fixture -async def base_strategy(mock_dual_cache): - return BaseRoutingStrategy( - dual_cache=mock_dual_cache, - should_batch_redis_writes=False, - default_sync_interval=1, - ) +def _strategy(handle: FakeUsageHandle) -> BaseRoutingStrategy: + return BaseRoutingStrategy(usage=usage_store(handle), should_batch_redis_writes=False, default_sync_interval=1) @pytest.mark.asyncio -async def test_increment_value_in_current_window(base_strategy, mock_dual_cache): - # Test incrementing value in current window - key = "test_key" - value = 10.0 - ttl = 3600 +async def test_an_increment_counts_in_memory_now_and_queues_the_redis_write(): + handle: Final = FakeUsageHandle(redis={}) + strategy: Final = _strategy(handle) - await base_strategy._increment_value_in_current_window(key, value, ttl) + counted: Final = await strategy._increment_value_in_current_window(_A, 1, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Verify in-memory cache was incremented - mock_dual_cache.in_memory_cache.async_increment.assert_called_once_with( - key=key, value=value, ttl=ttl - ) - - # Verify operation was queued for Redis - assert len(base_strategy.redis_increment_operation_queue) == 1 - queued_op = base_strategy.redis_increment_operation_queue[0] - assert isinstance(queued_op, dict) - assert queued_op["key"] == key - assert queued_op["increment_value"] == value - assert queued_op["ttl"] == ttl + assert counted == 1 + assert handle.local == {_A: 1} + assert handle.redis == {} + assert strategy.redis_increment_operation_queue == [CounterIncrement(_A, 1, 60)] + assert strategy.get_in_memory_keys_to_update() == {_A} @pytest.mark.asyncio -async def test_push_in_memory_increments_to_redis(base_strategy, mock_dual_cache): - # Add some operations to the queue - base_strategy.redis_increment_operation_queue = [ - RedisPipelineIncrementOperation(key="key1", increment_value=10, ttl=3600), - RedisPipelineIncrementOperation(key="key2", increment_value=20, ttl=3600), - ] +async def test_a_push_sends_each_counters_queued_increments_summed_in_one_pipeline(): + handle: Final = FakeUsageHandle(redis={_A: 10}) + strategy: Final = _strategy(handle) + for counter in (_A, _B, _A): + await strategy._increment_value_in_current_window(counter, 1, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - await base_strategy._push_in_memory_increments_to_redis() + totals: Final = await strategy._push_in_memory_increments_to_redis() # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Verify Redis pipeline was called - mock_dual_cache.redis_cache.async_increment_pipeline.assert_called_once() - # Verify queue was cleared - assert len(base_strategy.redis_increment_operation_queue) == 0 + assert handle.calls == [("remote", (CounterIncrement(_A, 2, 60), CounterIncrement(_B, 1, 60)))] + assert totals == {_A: 12, _B: 1} + assert strategy.redis_increment_operation_queue == [] @pytest.mark.asyncio -async def test_sync_in_memory_spend_with_redis(base_strategy, mock_dual_cache): - from litellm.types.caching import RedisPipelineIncrementOperation +async def test_a_sync_moves_memory_up_to_what_every_pod_counted_in_redis(): + handle: Final = FakeUsageHandle(redis={_A: 20}) + strategy: Final = _strategy(handle) + await strategy._increment_value_in_current_window(_A, 3, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Setup test data - base_strategy.in_memory_keys_to_update = {"key1"} - base_strategy.redis_increment_operation_queue = [ - RedisPipelineIncrementOperation(key="key1", increment_value=10, ttl=3600), - ] + await strategy._sync_in_memory_spend_with_redis() # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Mock the in-memory cache batch get responses for before snapshot - in_memory_before_future: asyncio.Future[List[str]] = asyncio.Future() - in_memory_before_future.set_result(["5.0"]) # Initial values - mock_dual_cache.in_memory_cache.async_batch_get_cache.return_value = ( - in_memory_before_future - ) - - # Mock Redis batch get response - redis_future: asyncio.Future[Dict[str, str]] = asyncio.Future() - redis_future.set_result([15.0]) # Redis values - mock_dual_cache.redis_cache.async_increment_pipeline.return_value = redis_future - - # Mock in-memory get for after snapshot - in_memory_after_future: asyncio.Future[Optional[str]] = asyncio.Future() - in_memory_after_future.set_result("8.0") # Value after potential updates - mock_dual_cache.in_memory_cache.async_get_cache.return_value = ( - in_memory_after_future - ) - - await base_strategy._sync_in_memory_spend_with_redis() - - # Verify the final merged values - set_cache_calls = mock_dual_cache.in_memory_cache.async_set_cache.call_args_list - print(f"set_cache_calls: {set_cache_calls}") - assert any( - call.kwargs["key"] == "key1" and float(call.kwargs["value"]) == 18.0 - for call in set_cache_calls - ) - - # Verify cache keys still exist - assert len(base_strategy.in_memory_keys_to_update) == 1 + assert handle.redis == {_A: 23} + assert handle.local == {_A: 23} @pytest.mark.asyncio -async def test_cache_keys_management(base_strategy): - # Test adding and getting cache keys - base_strategy.add_to_in_memory_keys_to_update("key1") - base_strategy.add_to_in_memory_keys_to_update("key2") - base_strategy.add_to_in_memory_keys_to_update("key1") # Duplicate should be ignored +async def test_without_redis_a_sync_leaves_the_queue_and_memory_alone(): + handle: Final = FakeUsageHandle() + strategy: Final = _strategy(handle) + await strategy._increment_value_in_current_window(_A, 3, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - cache_keys = base_strategy.get_in_memory_keys_to_update() - assert len(cache_keys) == 2 - assert "key1" in cache_keys - assert "key2" in cache_keys + await strategy._sync_in_memory_spend_with_redis() # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Test resetting cache keys - base_strategy.reset_in_memory_keys_to_update() - assert len(base_strategy.get_in_memory_keys_to_update()) == 0 + assert handle.local == {_A: 3} + assert handle.calls == [] + assert strategy.redis_increment_operation_queue == [CounterIncrement(_A, 3, 60)] @pytest.mark.asyncio -async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(base_strategy, mock_dual_cache, caplog): +async def test_cache_keys_management(): + strategy: Final = _strategy(FakeUsageHandle()) + strategy.add_to_in_memory_keys_to_update(_A) + strategy.add_to_in_memory_keys_to_update(_B) + strategy.add_to_in_memory_keys_to_update(_A) + + assert strategy.get_in_memory_keys_to_update() == {_A, _B} + strategy.reset_in_memory_keys_to_update() + assert strategy.get_in_memory_keys_to_update() == set() + + +class _RefusedHandle(FakeUsageHandle): + def increment_remote(self, increments: Sequence[CounterIncrement]) -> asyncio.Future[list[float] | None]: + raise NativeRedisStoreError("breaker_open", "the Redis circuit breaker is open") + + +@pytest.mark.asyncio +async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(caplog): """The sync loop pushes every 100 ms under usage-based routing, so an open breaker must not add an error line per cycle.""" - mock_dual_cache.redis_cache.async_increment_pipeline.side_effect = RedisCircuitBreakerOpenError( - "Redis circuit breaker is open - skipping async_increment_pipeline" - ) - base_strategy.redis_increment_operation_queue = [{"key": "k", "increment_value": 1.0, "ttl": 60}] + strategy: Final = _strategy(_RefusedHandle(redis={})) + await strategy._increment_value_in_current_window(_A, 1, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer with caplog.at_level(logging.ERROR): - await base_strategy._push_in_memory_increments_to_redis() + await strategy._push_in_memory_increments_to_redis() # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer assert caplog.records == [] - assert base_strategy.redis_increment_operation_queue == [] + assert strategy.redis_increment_operation_queue == [] @pytest.mark.asyncio @@ -175,16 +111,16 @@ async def test_an_increment_counted_after_retire_still_reaches_redis(): threading.Thread(target=server.serve_forever, daemon=True).start() redis_cache: Final = RedisCache(host="127.0.0.1", port=server.server_address[1]) strategy: Final = BaseRoutingStrategy( - dual_cache=DualCache(redis_cache=redis_cache, in_memory_cache=InMemoryCache()), + usage=RouterUsageStore(connect(redis_cache).router_usage()), should_batch_redis_writes=True, default_sync_interval=60, ) try: strategy.retire() - await strategy._increment_value_in_current_window("counted-after-retire", 3, 60) + await strategy._increment_value_in_current_window(_A, 3, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer await asyncio.sleep(0.2) - assert float(await redis_cache.async_get_cache("counted-after-retire")) == 3 + assert float(await redis_cache.async_get_cache("dep-a:openai/gpt-5.4-nano:rpm:12-34")) == 3 finally: server.shutdown() server.server_close() diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py index f1c56b9d033..bc4ea36e2ae 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm.py @@ -1,14 +1,15 @@ import asyncio from datetime import datetime, timedelta from typing import Final -from unittest.mock import AsyncMock import pytest from litellm import Router, token_counter -from litellm.caching.dual_cache import DualCache +from litellm.redis.stores.router_usage import UsageCounter from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict, RouterErrors +from litellm.utils import get_utc_datetime +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store MODEL_GROUP: Final = "lowest-tpm-router" HIGH_USAGE_DEPLOYMENT_ID: Final = "highest-usage" @@ -87,29 +88,29 @@ def test_usage_based_routing_v1_skips_a_deployment_whose_tpm_limit_is_used_up() @pytest.mark.asyncio async def test_v2_async_selection_uses_prefetched_counters_only_when_they_cover_its_keys(): - router_cache = DualCache() - router_cache.async_batch_get_cache = AsyncMock(return_value=[100, 10, None, None]) # type: ignore[method-assign] - strategy = LowestTPMLoggingHandler_v2(router_cache=router_cache) deployments = [ {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "a"}}, {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "b"}}, ] - tpm_keys, rpm_keys = strategy.usage_counter_keys(deployments) - keys = tpm_keys + rpm_keys + minute = get_utc_datetime().strftime("%H-%M") + tpm_counters = tuple(UsageCounter.for_deployment(d, "m", "tpm", minute) for d in ("a", "b")) + counters = tpm_counters + tuple(UsageCounter.for_deployment(d, "m", "rpm", minute) for d in ("a", "b")) + handle = FakeUsageHandle(redis=dict(zip(tpm_counters, [100, 10]))) + strategy = LowestTPMLoggingHandler_v2(usage=usage_store(handle)) - covering = PrefetchedUsage(keys=frozenset(keys), values=dict(zip(keys, [10, 100, None, None]))) + covering = PrefetchedUsage(counters=frozenset(counters), values=dict(zip(counters, [10, 100, None, None]))) with PrefetchedUsage.scoped(covering): chosen: Final = await strategy.async_get_available_deployments(model_group="g", healthy_deployments=deployments) assert chosen["model_info"]["id"] == "a", "the prefetched counters say a is the lowest" - router_cache.async_batch_get_cache.assert_not_awaited() + assert handle.calls == [] - stale = PrefetchedUsage(keys=frozenset(keys[:1]), values={keys[0]: 10}) + stale = PrefetchedUsage(counters=frozenset(counters[:1]), values={counters[0]: 10}) with PrefetchedUsage.scoped(stale): chosen_stale: Final = await strategy.async_get_available_deployments( model_group="g", healthy_deployments=deployments ) assert chosen_stale["model_info"]["id"] == "b", "counters that do not cover this minute's keys are read again" - router_cache.async_batch_get_cache.assert_awaited_once_with(keys=keys) + assert handle.calls == [("get_many", counters)] @pytest.mark.asyncio @@ -133,11 +134,9 @@ async def test_v2_subclass_overriding_async_get_available_deployments_with_the_o model_list=[_deployment(HIGH_USAGE_DEPLOYMENT_ID), _deployment(LOW_USAGE_DEPLOYMENT_ID)], routing_strategy="usage-based-routing-v2", ) - router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache, routing_args={}) + router.lowesttpm_logger_v2 = OldSignatureV2(usage=router.lowesttpm_logger_v2.usage, routing_args={}) - response: Final = await router.acompletion( - model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}] - ) + response: Final = await router.acompletion(model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}]) assert response.choices[0].message.content in { f"from {HIGH_USAGE_DEPLOYMENT_ID}", diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index f04ade1ce2d..7bc14b7e8c5 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -22,7 +22,8 @@ from PIL import Image import litellm from litellm import Router from litellm._logging import verbose_logger -from litellm.caching.caching import DualCache +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import RouterUsageStore, UsageCounter from litellm.router_strategy.lowest_tpm_rpm_v2 import ( LowestTPMLoggingHandler_v2 as LowestTPMLoggingHandler, ) @@ -32,6 +33,11 @@ from litellm.utils import get_utc_datetime ### UNIT TESTS FOR TPM/RPM ROUTING ### + +def _memory_usage() -> RouterUsageStore: + return RouterUsageStore(connect(None).router_usage()) + + """ - Given 2 deployments, make sure it's shuffling deployments correctly. """ @@ -48,8 +54,8 @@ def restore_verbose_logger_level() -> Iterator[None]: def test_tpm_rpm_updated(): from tests.unit.create_mock_standard_logging_payload import create_standard_logging_payload - test_cache = DualCache() - lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) + test_cache = _memory_usage() + lowest_tpm_logger = LowestTPMLoggingHandler(usage=test_cache) model_group = "gpt-3.5-turbo" deployment_id = "1234" deployment = "azure/gpt-4.1-mini" @@ -89,11 +95,11 @@ def test_tpm_rpm_updated(): ) dt = get_utc_datetime() current_minute = dt.strftime("%H-%M") - tpm_count_api_key = f"{deployment_id}:{deployment}:tpm:{current_minute}" - rpm_count_api_key = f"{deployment_id}:{deployment}:rpm:{current_minute}" + tpm_counter = UsageCounter.for_deployment(deployment_id, deployment, "tpm", current_minute) + rpm_counter = UsageCounter.for_deployment(deployment_id, deployment, "rpm", current_minute) - assert response_obj["usage"]["total_tokens"] == test_cache.get_cache(key=tpm_count_api_key) - assert 1 == test_cache.get_cache(key=rpm_count_api_key) + assert response_obj["usage"]["total_tokens"] == test_cache.get_local(tpm_counter) + assert 1 == test_cache.get_local(rpm_counter) # test_tpm_rpm_updated() @@ -103,7 +109,7 @@ def test_tpm_rpm_updated(): def test_get_available_deployments(): from tests.unit.create_mock_standard_logging_payload import create_standard_logging_payload - test_cache = DualCache() + test_cache = _memory_usage() model_list = [ { "model_name": "gpt-3.5-turbo", @@ -116,7 +122,7 @@ def test_get_available_deployments(): "model_info": {"id": "5678"}, }, ] - lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) + lowest_tpm_logger = LowestTPMLoggingHandler(usage=test_cache) model_group = "gpt-3.5-turbo" ## DEPLOYMENT 1 ## total_tokens = 50 @@ -570,8 +576,8 @@ def test_return_potential_deployments(): Assert deployment at limit is filtered out """ - test_cache = DualCache() - lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) + test_cache = _memory_usage() + lowest_tpm_logger = LowestTPMLoggingHandler(usage=test_cache) args: Dict = { "healthy_deployments": [ diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 8a0d5b9a1d6..08b045aea82 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -10,19 +10,21 @@ import datetime import uuid from collections.abc import Mapping from typing import Final, Literal -from unittest.mock import patch import pytest from pydantic import ValidationError import litellm from litellm import Router -from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger +from litellm.redis.stores.router_usage import CounterIncrement, UsageCounter +from litellm.router import router_arguments from litellm.router_api import RoutingContext +from litellm.router_config import build_router_config from litellm.router_strategy.least_busy import picks_reserve from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup from litellm.utils import Rules, function_setup +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, UsageHost def _model_list(): @@ -1094,33 +1096,21 @@ async def test_update_settings_same_routing_strategy_args_does_not_leak_sync_tas pass -class _RecordingRedisCache: - def __init__(self) -> None: - self.increment_lists: list[list[RedisPipelineIncrementOperation]] = [] - - async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]: - self.increment_lists.append(list(increment_list)) - return [float(op["increment_value"]) for op in increment_list] - - @pytest.mark.asyncio async def test_update_settings_changed_routing_strategy_args_flushes_replaced_selector_queue( monkeypatch, ) -> None: monkeypatch.setattr(litellm, "callbacks", []) monkeypatch.setattr(litellm, "input_callback", []) - router: Final = Router( - model_list=_model_list(), - routing_strategy="usage-based-routing-v2", - routing_strategy_args={"ttl": 60}, + handle: Final = FakeUsageHandle(redis={}) + queued: Final = CounterIncrement(UsageCounter.for_deployment("dep", "m", "rpm", "12-34"), 3, 60) + arguments: Final = router_arguments( + model_list=_model_list(), routing_strategy="usage-based-routing-v2", routing_strategy_args={"ttl": 60} ) + router: Final = Router.from_config(build_router_config(arguments), UsageHost(handle)) try: - redis_cache: Final = _RecordingRedisCache() replaced: Final = router.lowesttpm_logger_v2 - replaced.dual_cache.redis_cache = redis_cache - replaced.redis_increment_operation_queue.append( - RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60) - ) + replaced.redis_increment_operation_queue.append(queued) router.update_settings(routing_strategy_args={"ttl": 120}) await asyncio.sleep(0) @@ -1134,9 +1124,7 @@ async def test_update_settings_changed_routing_strategy_args_flushes_replaced_se ) assert router.lowesttpm_logger_v2 is not replaced - assert redis_cache.increment_lists == [ - [RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60)] - ] + assert [args for name, args in handle.calls if name == "remote"] == [(queued,)] assert replaced.redis_increment_operation_queue == [] finally: for t in [ diff --git a/tests/unit/router_utils/test_routing_read_batch.py b/tests/unit/router_utils/test_routing_read_batch.py index c8f4b42924c..96f46bc9e02 100644 --- a/tests/unit/router_utils/test_routing_read_batch.py +++ b/tests/unit/router_utils/test_routing_read_batch.py @@ -1,23 +1,27 @@ """ The router's pre-call reads. Cooldowns go through the native store, and the tpm/rpm selector -counters (`LowestTPMLoggingHandler_v2`) ride one MGET on the Python Redis. +counters (`LowestTPMLoggingHandler_v2`) are one read of the native usage store. """ import time from collections.abc import Sequence from typing import Final -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest -import litellm from litellm import Router from litellm.caching.redis_batch import request_redis_batch_scope from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore +from litellm.redis.stores.router_usage import UsageCounter +from litellm.router import router_arguments +from litellm.router_config import build_router_config from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.router_utils import routing_prefetch +from litellm.utils import get_utc_datetime from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNatives, ReadyPending, fake_pending, fake_send from tests.unit.conftest import redis_cache_double +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, UsageHost from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns _MODEL_GROUP = "claude" @@ -32,30 +36,31 @@ def _deployment(deployment_id: str) -> dict: } -def _redis_answering(values_by_key_prefix: dict[str, object]) -> MagicMock: - """A Redis double that answers each key from its minute-less prefix and records every MGET.""" - - def _mget(key_list, parent_otel_span=None): - return {key: values_by_key_prefix.get(key.rsplit(":", 1)[0], values_by_key_prefix.get(key)) for key in key_list} - - redis = redis_cache_double() - redis.async_batch_get_cache = AsyncMock(side_effect=_mget) - return redis - - -def _router(redis: MagicMock, routing_strategy: str) -> Router: - router = Router( - model_list=[_deployment("dep-a"), _deployment("dep-b")], - routing_strategy=routing_strategy, +def _router(routing_strategy: str, handle: FakeUsageHandle | None = None) -> Router: + arguments: Final = router_arguments( + model_list=[_deployment("dep-a"), _deployment("dep-b")], routing_strategy=routing_strategy ) - router.update_redis_cache(cache=redis) - return router + return Router.from_config(build_router_config(arguments), UsageHost(handle or FakeUsageHandle(redis={}))) -def _redis_key_families(redis: MagicMock) -> list[list[str]]: +def _counter(deployment_id: str, kind: str) -> UsageCounter: + return UsageCounter.for_deployment( + deployment_id, + "anthropic/claude-x", + kind, # pyright: ignore[reportArgumentType] # tpm or rpm + get_utc_datetime().strftime("%H-%M"), + ) + + +def _tpm(**tokens: float | None) -> dict[UsageCounter, float]: + return {_counter(deployment_id, "tpm"): count for deployment_id, count in tokens.items() if count is not None} + + +def _usage_reads(handle: FakeUsageHandle) -> list[tuple[str, ...]]: return [ - sorted(key.rsplit(":", 1)[0] if ":tpm:" in key or ":rpm:" in key else key for key in call.args[0]) - for call in redis.async_batch_get_cache.await_args_list + tuple(f"{c.deployment_id}:{c.kind}" for c in counters) # pyright: ignore[reportAttributeAccessIssue] # usage() args are counters + for name, counters in handle.calls + if name == "usage" ] @@ -66,21 +71,16 @@ def _cool_down(router: Router, model_id: str, started: float) -> None: @pytest.mark.asyncio -async def test_usage_based_routing_reads_the_counters_in_one_python_redis_round_trip(): - redis = _redis_answering({}) - router = _router(redis, "usage-based-routing-v2") +async def test_usage_based_routing_reads_the_counters_in_one_store_read(): + handle: Final = FakeUsageHandle(redis={}) + router: Final = _router("usage-based-routing-v2", handle) - deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES + ) assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} - assert _redis_key_families(redis) == [ - [ - "dep-a:anthropic/claude-x:rpm", - "dep-a:anthropic/claude-x:tpm", - "dep-b:anthropic/claude-x:rpm", - "dep-b:anthropic/claude-x:tpm", - ] - ], "the usage counters must arrive in one MGET, without cooldown keys" + assert _usage_reads(handle) == [("dep-a:tpm", "dep-b:tpm", "dep-a:rpm", "dep-b:rpm")] @pytest.mark.asyncio @@ -100,35 +100,30 @@ async def test_usage_based_routing_still_batches_when_the_strategy_is_a_fixed_si input=input, ) - redis: Final = _redis_answering({}) - router: Final = _router(redis, "usage-based-routing-v2") - router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache) - router.cache.async_batch_get_cache = AsyncMock(wraps=router.cache.async_batch_get_cache) + handle: Final = FakeUsageHandle(redis={}) + router: Final = _router("usage-based-routing-v2", handle) + router.lowesttpm_logger_v2.cancel_sync_task() + router.lowesttpm_logger_v2 = OldSignatureV2(usage=router.lowesttpm_logger_v2.usage) deployment: Final = await router.async_get_available_deployment( model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES ) assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} - assert _redis_key_families(redis) == [ - [ - "dep-a:anthropic/claude-x:rpm", - "dep-a:anthropic/claude-x:tpm", - "dep-b:anthropic/claude-x:rpm", - "dep-b:anthropic/claude-x:tpm", - ] - ], "the subclassed strategy must still get the batched read, not a second MGET" - router.cache.async_batch_get_cache.assert_not_awaited() + assert _usage_reads(handle) == [("dep-a:tpm", "dep-b:tpm", "dep-a:rpm", "dep-b:rpm")], ( + "the subclassed strategy must still get the batched read, not a second one" + ) + assert [name for name, _ in handle.calls] == ["usage"] @pytest.mark.asyncio -async def test_simple_shuffle_reads_nothing_from_the_python_redis(): - redis = _redis_answering({}) - router = _router(redis, "simple-shuffle") +async def test_simple_shuffle_reads_no_usage(): + handle: Final = FakeUsageHandle(redis={}) + router: Final = _router("simple-shuffle", handle) await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) - assert _redis_key_families(redis) == [] + assert handle.calls == [] @pytest.mark.asyncio @@ -137,14 +132,13 @@ async def test_simple_shuffle_reads_nothing_from_the_python_redis(): [(100, 10, "dep-b"), (10, 100, "dep-a"), (None, 10, "dep-a"), (10, None, "dep-b")], ) async def test_batched_counters_pick_the_deployment_the_strategy_picks_reading_alone(tpm_a, tpm_b, expected): - counters = {"dep-a:anthropic/claude-x:tpm": tpm_a, "dep-b:anthropic/claude-x:tpm": tpm_b} - routed = _router(_redis_answering(counters), "usage-based-routing-v2") - alone = _router(_redis_answering(counters), "usage-based-routing-v2") + routed: Final = _router("usage-based-routing-v2", FakeUsageHandle(redis=_tpm(**{"dep-a": tpm_a, "dep-b": tpm_b}))) + alone: Final = _router("usage-based-routing-v2", FakeUsageHandle(redis=_tpm(**{"dep-a": tpm_a, "dep-b": tpm_b}))) - routed_choice = await routed.async_get_available_deployment( + routed_choice: Final = await routed.async_get_available_deployment( model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES ) - alone_choice = await alone.lowesttpm_logger_v2.async_get_available_deployments( + alone_choice: Final = await alone.lowesttpm_logger_v2.async_get_available_deployments( model_group=_MODEL_GROUP, healthy_deployments=alone.model_list, messages=_MESSAGES ) @@ -153,50 +147,36 @@ async def test_batched_counters_pick_the_deployment_the_strategy_picks_reading_a @pytest.mark.asyncio async def test_batched_read_still_excludes_a_cooled_down_deployment(): - redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10}) - router = _router(redis, "usage-based-routing-v2") + handle: Final = FakeUsageHandle(redis=_tpm(**{"dep-a": 100, "dep-b": 10})) + router: Final = _router("usage-based-routing-v2", handle) _cool_down(router, "dep-b", started=time.time()) - deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES + ) assert deployment["model_info"]["id"] == "dep-a", "dep-b has the lowest tpm but is cooling down" - assert redis.async_batch_get_cache.await_count == 1 + assert len(_usage_reads(handle)) == 1 @pytest.mark.asyncio async def test_batched_read_ignores_an_expired_cooldown(): - redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10}) - router = _router(redis, "usage-based-routing-v2") + router: Final = _router("usage-based-routing-v2", FakeUsageHandle(redis=_tpm(**{"dep-a": 100, "dep-b": 10}))) _cool_down(router, "dep-b", started=time.time() - 61) - deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES + ) assert deployment["model_info"]["id"] == "dep-b" -@pytest.mark.asyncio -async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did(): - redis = redis_cache_double() - redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) - routed = _router(redis, "usage-based-routing-v2") - alone = _router(redis, "usage-based-routing-v2") - - with pytest.raises(litellm.RateLimitError, match="No deployments available") as routed_error: - await routed.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) - with pytest.raises(litellm.RateLimitError, match="No deployments available") as alone_error: - await alone.lowesttpm_logger_v2.async_get_available_deployments( - model_group=_MODEL_GROUP, healthy_deployments=alone.model_list, messages=_MESSAGES - ) - - assert str(routed_error.value) == str(alone_error.value) - assert len(routed.cache.last_redis_batch_access_time) == 0, "a failed read must not throttle the next one" - - @pytest.mark.asyncio async def test_a_failed_batched_read_leaves_simple_shuffle_routing(): redis = redis_cache_double() redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) - router = _router(redis, "simple-shuffle") + router = _router("simple-shuffle") + router.update_redis_cache(cache=redis) deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index ff3ffc0db0c..ee103e2df3c 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -6,20 +6,22 @@ regardless of the routing strategy being used. """ import asyncio +from collections.abc import Callable from datetime import datetime, timedelta, timezone from typing import Final -from unittest.mock import AsyncMock, MagicMock import pytest -from pydantic import TypeAdapter import litellm from litellm import Router from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import CounterIncrement, CounterKind, RouterUsage, RouterUsageStore, UsageCounter from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +from litellm.utils import get_utc_datetime +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store TPM_DEPLOYMENT = { "tpm": 1000, @@ -29,14 +31,49 @@ TPM_DEPLOYMENT = { } -def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: - dual_cache = DualCache(redis_cache=redis_cache) - check = ModelRateLimitingCheck(dual_cache=dual_cache) - now = litellm.utils.get_utc_datetime() - for minute in (now, now + timedelta(minutes=1)): - tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) - dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) - return dual_cache +def _memory_usage() -> RouterUsageStore: + return RouterUsageStore(connect(None).router_usage()) + + +def _check( + usage: RouterUsage | None = None, clock: Callable[[], datetime] = get_utc_datetime +) -> ModelRateLimitingCheck: + return ModelRateLimitingCheck(usage=usage or _memory_usage(), dual_cache=DualCache(), clock=clock) + + +def _counter(deployment: dict, kind: CounterKind, window: str) -> UsageCounter: + model_id: Final[object] = deployment["model_info"]["id"] + model: Final[object] = deployment["litellm_params"]["model"] + return UsageCounter.for_deployment(f"{model_id}", f"{model}", kind, window) + + +def _this_minute(deployment: dict, kind: CounterKind, value: float) -> dict[UsageCounter, float]: + """``value`` under this minute's counter and the next, so a test crossing a minute still finds it.""" + now: Final = get_utc_datetime() + return {_counter(deployment, kind, minute.strftime("%H-%M")): value for minute in (now, now + timedelta(minutes=1))} + + +RPM_DEPLOYMENT: Final = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", +} + + +_SUCCESS_KWARGS: Final = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } +} + + +def _tpm_increments(handle: FakeUsageHandle) -> tuple[tuple[object, object], ...]: + """The (tokens, ttl) each success charged to test-id's gpt-4 tpm counter this minute or the next.""" + tpm_counters: Final = frozenset(_this_minute(RPM_DEPLOYMENT, "tpm", 0)) + return tuple((args[1], args[2]) for name, args in handle.calls if name == "increment" and args[0] in tpm_counters) class TestModelRateLimitingCheck: @@ -44,7 +81,7 @@ class TestModelRateLimitingCheck: def test_get_deployment_limits_from_top_level(self): """Test extracting limits from top-level deployment config.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "tpm": 1000, @@ -59,7 +96,7 @@ class TestModelRateLimitingCheck: def test_get_deployment_limits_from_litellm_params(self): """Test extracting limits from litellm_params.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, @@ -72,7 +109,7 @@ class TestModelRateLimitingCheck: def test_get_deployment_limits_from_model_info(self): """Test extracting limits from model_info.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "litellm_params": {"model": "gpt-4"}, @@ -85,7 +122,7 @@ class TestModelRateLimitingCheck: def test_get_deployment_limits_none_when_not_set(self): """Test that None is returned when limits are not set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "litellm_params": {"model": "gpt-4"}, @@ -98,7 +135,7 @@ class TestModelRateLimitingCheck: def test_pre_call_check_allows_request_when_no_limits(self): """Test that requests are allowed when no limits are set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "litellm_params": {"model": "gpt-4"}, @@ -110,78 +147,50 @@ class TestModelRateLimitingCheck: def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): """Test that RateLimitError is raised when RPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 11 # Over limit after increment - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } + check = _check(usage_store(FakeUsageHandle(redis=_this_minute(RPM_DEPLOYMENT, "rpm", 10)))) with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) + check.pre_call_check(RPM_DEPLOYMENT) assert "RPM limit=10" in str(exc_info.value) assert "current usage=11" in str(exc_info.value) def test_pre_call_check_allows_request_under_limit(self): """Test that requests are allowed when under the limit.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 6 + check = _check(usage_store(FakeUsageHandle(redis=_this_minute(RPM_DEPLOYMENT, "rpm", 5)))) - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = check.pre_call_check(deployment) - assert result == deployment + assert check.pre_call_check(RPM_DEPLOYMENT) == RPM_DEPLOYMENT def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): """Test that RateLimitError is raised when TPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.get_cache.return_value = 1000 # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } + check = _check(usage_store(FakeUsageHandle(local=_this_minute(TPM_DEPLOYMENT, "tpm", 1000)))) with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) + check.pre_call_check(TPM_DEPLOYMENT) assert "TPM limit=1000" in str(exc_info.value) assert "current usage=1000" in str(exc_info.value) def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.get_cache.return_value = 1000 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + handle = FakeUsageHandle( + local=_this_minute(TPM_DEPLOYMENT, "tpm", 5), redis=_this_minute(TPM_DEPLOYMENT, "tpm", 1000) + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: check.pre_call_check(TPM_DEPLOYMENT) assert "current usage=1000" in str(exc_info.value) - @pytest.mark.parametrize( - "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] - ) - def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + @pytest.mark.parametrize("failure", ["redis_down", "breaker_open"]) + def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, failure: str): + handle = FakeUsageHandle( + local=_this_minute(TPM_DEPLOYMENT, "tpm", 1000), + redis={}, + breaker_open=failure == "breaker_open", + redis_down=failure == "redis_down", + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: check.pre_call_check(TPM_DEPLOYMENT) @@ -189,10 +198,12 @@ class TestModelRateLimitingCheck: assert "current usage=1000" in str(exc_info.value) def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): - redis_cache = MagicMock() - redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() - redis_cache.increment_cache.return_value = 2 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + handle = FakeUsageHandle( + local={**_this_minute(TPM_DEPLOYMENT, "tpm", 5), **_this_minute(TPM_DEPLOYMENT, "rpm", 1)}, + redis={}, + breaker_open=True, + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) @@ -200,7 +211,7 @@ class TestModelRateLimitingCheck: assert "RPM limit=1" in str(exc_info.value) def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + check = _check(usage_store(FakeUsageHandle(local=_this_minute(TPM_DEPLOYMENT, "tpm", 5)))) deployment = {**TPM_DEPLOYMENT, "rpm": 1} assert check.pre_call_check(deployment) == deployment @@ -211,22 +222,12 @@ class TestModelRateLimitingCheck: def test_log_success_event_increments_cache(self): """Test that log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) + handle = FakeUsageHandle() + check = _check(usage_store(handle)) - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } + check.log_success_event(_SUCCESS_KWARGS, None, None, None) - check.log_success_event(kwargs, None, None, None) - - keys = tuple(call.kwargs["key"] for call in mock_cache.increment_cache.call_args_list) - assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys - assert all(call.kwargs["value"] == 50 for call in mock_cache.increment_cache.call_args_list) + assert _tpm_increments(handle) == ((50, 60),) class TestModelRateLimitingCheckAsync: @@ -235,10 +236,7 @@ class TestModelRateLimitingCheckAsync: @pytest.mark.asyncio async def test_async_pre_call_check_allows_request_when_no_limits(self): """Test that requests are allowed when no limits are set (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) + check = _check() deployment = { "litellm_params": {"model": "gpt-4"}, @@ -251,68 +249,36 @@ class TestModelRateLimitingCheckAsync: @pytest.mark.asyncio async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): """Test that RateLimitError is raised when RPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[11.0]) # Over limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } + check = _check(usage_store(FakeUsageHandle(redis=_this_minute(RPM_DEPLOYMENT, "rpm", 10)))) with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) + await check.async_pre_call_check(RPM_DEPLOYMENT) assert "RPM limit=10" in str(exc_info.value) @pytest.mark.asyncio async def test_async_pre_call_check_allows_request_under_limit(self): """Test that requests are allowed when under the limit (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[6.0]) + check = _check(usage_store(FakeUsageHandle(redis=_this_minute(RPM_DEPLOYMENT, "rpm", 5)))) - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = await check.async_pre_call_check(deployment) - assert result == deployment + assert await check.async_pre_call_check(RPM_DEPLOYMENT) == RPM_DEPLOYMENT @pytest.mark.asyncio async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): """Test that RateLimitError is raised when TPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } + check = _check(usage_store(FakeUsageHandle(local=_this_minute(TPM_DEPLOYMENT, "tpm", 1000)))) with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) + await check.async_pre_call_check(TPM_DEPLOYMENT) assert "TPM limit=1000" in str(exc_info.value) @pytest.mark.asyncio async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=1000) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + handle = FakeUsageHandle( + local=_this_minute(TPM_DEPLOYMENT, "tpm", 5), redis=_this_minute(TPM_DEPLOYMENT, "tpm", 1000) + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: await check.async_pre_call_check(TPM_DEPLOYMENT) @@ -320,13 +286,15 @@ class TestModelRateLimitingCheckAsync: assert "current usage=1000" in str(exc_info.value) @pytest.mark.asyncio - @pytest.mark.parametrize( - "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] - ) - async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.async_get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + @pytest.mark.parametrize("failure", ["redis_down", "breaker_open"]) + async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, failure: str): + handle = FakeUsageHandle( + local=_this_minute(TPM_DEPLOYMENT, "tpm", 1000), + redis={}, + breaker_open=failure == "breaker_open", + redis_down=failure == "redis_down", + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: await check.async_pre_call_check(TPM_DEPLOYMENT) @@ -337,10 +305,12 @@ class TestModelRateLimitingCheckAsync: async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( self, ): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) - redis_cache.async_increment_pipeline = AsyncMock(return_value=[2.0]) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + handle = FakeUsageHandle( + local={**_this_minute(TPM_DEPLOYMENT, "tpm", 5), **_this_minute(TPM_DEPLOYMENT, "rpm", 1)}, + redis={}, + breaker_open=True, + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) @@ -349,7 +319,7 @@ class TestModelRateLimitingCheckAsync: @pytest.mark.asyncio async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + check = _check(usage_store(FakeUsageHandle(local=_this_minute(TPM_DEPLOYMENT, "tpm", 5)))) deployment = {**TPM_DEPLOYMENT, "rpm": 1} assert await check.async_pre_call_check(deployment) == deployment @@ -361,23 +331,12 @@ class TestModelRateLimitingCheckAsync: @pytest.mark.asyncio async def test_async_log_success_event_increments_cache(self): """Test that async_log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - mock_cache.async_increment_cache = AsyncMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) + handle = FakeUsageHandle(redis={}) + check = _check(usage_store(handle)) - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } + await check.async_log_success_event(_SUCCESS_KWARGS, None, None, None) - await check.async_log_success_event(kwargs, None, None, None) - - keys = tuple(call.kwargs["key"] for call in mock_cache.async_increment_cache.call_args_list) - assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys - assert all(call.kwargs["value"] == 50 for call in mock_cache.async_increment_cache.call_args_list) + assert _tpm_increments(handle) == ((50, 60),) class TestRouterWithEnforceModelRateLimits: @@ -441,8 +400,7 @@ class TestModelRateLimitConcurrency: the old check-then-increment pattern would let 3+ through due to a race condition on the local cache read. """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = _check() deployment = { "rpm": 2, @@ -473,6 +431,8 @@ RPD_DEPLOYMENT = { "model_name": "daily-model", } RPM_AND_RPD_DEPLOYMENT = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpm": 1, "rpd": 2}} +_DAY_ONE_RPM: Final = UsageCounter.for_deployment("daily-test-id", "gpt-4.1-mini", "rpm", "23-59") +_DAY_ONE_RPD: Final = UsageCounter.for_deployment("daily-test-id", "gpt-4.1-mini", "rpd", "2026-10-28") class Clock: @@ -497,7 +457,7 @@ async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, ob class TestDailyLimits: def test_pre_call_check_rejects_the_request_over_rpd(self): - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT @@ -509,7 +469,7 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_rejects_the_request_over_rpd_and_lets_the_router_retry(self): - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) assert await _async_pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT assert await _async_pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT @@ -522,7 +482,7 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_rejects_once_the_day_tokens_reach_tpd(self): deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100}} - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) kwargs = { "standard_logging_object": { "model_id": "daily-test-id", @@ -541,7 +501,7 @@ class TestDailyLimits: def test_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): @@ -556,7 +516,7 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): @@ -570,7 +530,7 @@ class TestDailyLimits: def test_pre_call_check_rejected_requests_hand_back_rpm_and_rpd(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): @@ -589,7 +549,7 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_rejected_requests_hand_back_rpm_and_rpd(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): @@ -607,53 +567,43 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_refunds_rpm_and_rpd_in_one_cache_trip(self): - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(side_effect=[[1.0, 3.0], [0.0, 2.0]]) - check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) + handle = FakeUsageHandle(redis={_DAY_ONE_RPD: 2}) + check = _check(usage_store(handle), clock=Clock(DAY_ONE)) with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) - assert mock_cache.async_increment_cache_pipeline.await_count == 2 - refund = mock_cache.async_increment_cache_pipeline.call_args_list[1].kwargs["increment_list"] - assert [(op["key"], op["increment_value"], op["ttl"]) for op in refund] == [ - ("daily-test-id:gpt-4.1-mini:rpm:23-59", -1, 60), - ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", -1, 86400), + assert handle.calls == [ + ("pipeline", (CounterIncrement(_DAY_ONE_RPM, 1, 60), CounterIncrement(_DAY_ONE_RPD, 1, 86400))), + ("pipeline", (CounterIncrement(_DAY_ONE_RPM, -1, 60), CounterIncrement(_DAY_ONE_RPD, -1, 86400))), ] @pytest.mark.asyncio async def test_async_pre_call_check_increments_rpm_and_rpd_in_one_cache_trip(self): - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[1.0, 1.0]) - check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) + handle = FakeUsageHandle(redis={}) + check = _check(usage_store(handle), clock=Clock(DAY_ONE)) assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT - assert mock_cache.async_increment_cache_pipeline.await_count == 1 - operations = mock_cache.async_increment_cache_pipeline.call_args.kwargs["increment_list"] - assert [(op["key"], op["ttl"]) for op in operations] == [ - ("daily-test-id:gpt-4.1-mini:rpm:23-59", 60), - ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", 86400), + assert handle.calls == [ + ("pipeline", (CounterIncrement(_DAY_ONE_RPM, 1, 60), CounterIncrement(_DAY_ONE_RPD, 1, 86400))), ] - assert not mock_cache.async_increment_cache.called def test_pre_call_check_reads_the_shared_tpd_count_from_redis(self): - redis_cache = MagicMock() - redis_cache.get_cache.return_value = 100 + tpd: Final = UsageCounter.for_deployment("daily-test-id", "gpt-4.1-mini", "tpd", "2026-10-28") + handle = FakeUsageHandle(redis={tpd: 100}) deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100}} - check = ModelRateLimitingCheck(dual_cache=DualCache(redis_cache=redis_cache), clock=Clock(DAY_ONE)) + check = _check(usage_store(handle), clock=Clock(DAY_ONE)) with pytest.raises(litellm.RateLimitError) as exc_info: _pre_call(check, deployment) assert "TPD limit=100, current usage=100" in str(exc_info.value) - assert redis_cache.get_cache.call_args.kwargs["key"] == "daily-test-id:gpt-4.1-mini:tpd:2026-10-28" + assert handle.calls[0] == ("current", (tpd, 100)) def test_rpd_count_starts_fresh_after_utc_midnight(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) _pre_call(check, RPD_DEPLOYMENT) _pre_call(check, RPD_DEPLOYMENT) with pytest.raises(litellm.RateLimitError): @@ -677,7 +627,7 @@ class TestDailyLimits: ) def test_daily_429_retry_after_is_seconds_until_utc_midnight(self, now: datetime, seconds_to_midnight: int): deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpd": 0}} - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(now)) + check = _check(clock=Clock(now)) with pytest.raises(litellm.RateLimitError) as exc_info: _pre_call(check, deployment) @@ -712,7 +662,7 @@ class TestDailyLimits: ], num_retries=1, ) - check = ModelRateLimitingCheck(dual_cache=router.cache, clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) router.optional_callbacks = [check] litellm.logging_callback_manager.add_litellm_callback(check) try: @@ -728,18 +678,13 @@ class TestDailyLimits: assert served == ["a", "b", "b", "b"] -_LOGGED_TOKENS: Final = TypeAdapter[float | None](float | None) - - -async def _tpd_tokens_once_logged(cache: DualCache, key: str) -> int: +async def _tpd_tokens_once_logged(usage: RouterUsage, counter: UsageCounter) -> int: for _ in range(200): - logged: Final = _LOGGED_TOKENS.validate_python( - cache.in_memory_cache.get_cache(key) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # InMemoryCache.get_cache is untyped - ) + logged: Final = usage.get_local(counter) if logged is not None: return int(logged) await asyncio.sleep(0.01) - pytest.fail(f"the success callback never logged tpd tokens under {key}") + pytest.fail(f"the success callback never logged tpd tokens under {counter}") RESET_AT_EIGHT = datetime(2026, 10, 28, 8, 0, 0, tzinfo=timezone.utc) @@ -752,7 +697,7 @@ DAILY_RESET_DEPLOYMENT = { class TestDailyResetTime: def test_rpd_window_starts_over_at_daily_reset_time_not_midnight(self): clock = Clock(RESET_AT_EIGHT - timedelta(hours=9)) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert _pre_call(check, DAILY_RESET_DEPLOYMENT) == DAILY_RESET_DEPLOYMENT clock.now = RESET_AT_EIGHT - timedelta(hours=7) @@ -768,7 +713,7 @@ class TestDailyResetTime: @pytest.mark.asyncio async def test_async_rpd_window_starts_over_at_daily_reset_time_not_midnight(self): clock = Clock(RESET_AT_EIGHT - timedelta(hours=9)) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert await _async_pre_call(check, DAILY_RESET_DEPLOYMENT) == DAILY_RESET_DEPLOYMENT clock.now = RESET_AT_EIGHT - timedelta(hours=7) @@ -794,7 +739,7 @@ class TestDailyResetTime: **RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpd": 0, "daily_reset_time": "08:00"}, } - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(now)) + check = _check(clock=Clock(now)) with pytest.raises(litellm.RateLimitError) as exc_info: _pre_call(check, deployment) @@ -808,7 +753,7 @@ class TestDailyResetTime: "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100, "daily_reset_time": "08:00"}, } clock = Clock(RESET_AT_EIGHT - timedelta(minutes=1)) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) await _async_log_success( check, { @@ -834,7 +779,7 @@ class TestDailyResetTime: ], ) def test_rpd_outside_litellm_params_is_not_a_limit(self, deployment: dict[str, object]): - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) assert _pre_call(check, deployment) == deployment @@ -866,19 +811,22 @@ class TestDailyResetTime: } ] ) - check = ModelRateLimitingCheck(dual_cache=router.cache, clock=Clock(RESET_AT_EIGHT - timedelta(minutes=1))) + usage: Final = _memory_usage() + check = _check(usage, clock=Clock(RESET_AT_EIGHT - timedelta(minutes=1))) router.optional_callbacks = [check] litellm.logging_callback_manager.add_litellm_callback(check) try: await router.acompletion( # pyright: ignore[reportUnknownMemberType] # Router.acompletion kwargs are untyped model="daily-group", messages=[{"role": "user", "content": "hi"}], daily_reset_time="08:00" ) - logged = await _tpd_tokens_once_logged(router.cache, "daily-test-id:openai/gpt-4.1-mini:tpd:2026-10-28") + logged = await _tpd_tokens_once_logged( + usage, UsageCounter.for_deployment("daily-test-id", "openai/gpt-4.1-mini", "tpd", "2026-10-28") + ) finally: router.discard() assert logged > 0 - day_before: Final = _LOGGED_TOKENS.validate_python( - router.cache.in_memory_cache.get_cache("daily-test-id:openai/gpt-4.1-mini:tpd:2026-10-27") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # InMemoryCache.get_cache is untyped + day_before: Final = usage.get_local( + UsageCounter.for_deployment("daily-test-id", "openai/gpt-4.1-mini", "tpd", "2026-10-27") ) assert day_before is None diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index f09ad5973e2..7960dbe7073 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -20,6 +20,8 @@ from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( ITPM_CACHE_KEY, ITPM_RESERVED_KEY, @@ -106,7 +108,7 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -158,7 +160,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_itpm_limit_raises_429(self): dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -194,7 +196,7 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) otpm_limit = 10 max_tokens = 4 deployment = { @@ -248,7 +250,7 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) itpm_limit = 5 deployment = { "litellm_params": { @@ -293,7 +295,7 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) minute = get_utc_datetime().strftime("%H-%M") itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) @@ -319,7 +321,7 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -363,7 +365,7 @@ class TestModelRateLimitingCheckIOTokens: concurrently instead of the first one claiming the entire bucket. """ dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -388,7 +390,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -435,7 +437,9 @@ class TestModelRateLimitingCheckIOTokens: self, model: str, otpm_limit: int, use_sync_check: bool ): dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache + ) deployment: Final = { "litellm_params": {"model": model, "otpm": otpm_limit}, "model_info": {"id": f"io-no-max-tokens-{otpm_limit}"}, @@ -463,7 +467,9 @@ class TestModelRateLimitingCheckIOTokens: self, model: str, otpm_limit: int ): dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache + ) deployment: Final = { "litellm_params": {"model": model, "otpm": otpm_limit}, "model_info": {"id": "io-no-max-tokens-race-id"}, @@ -496,7 +502,9 @@ class TestModelRateLimitingCheckIOTokens: self, model: str, model_max_output: int ): dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache + ) deployment: Final = { "litellm_params": {"model": model, "otpm": model_max_output * 10}, "model_info": {"id": "io-no-max-tokens-roomy-id"}, @@ -515,7 +523,9 @@ class TestModelRateLimitingCheckIOTokens: catalog_max_output: Final = litellm.get_model_info(model=backend_model)["max_output_tokens"] assert isinstance(catalog_max_output, int) otpm_limit: Final = int(max(catalog_max_output, 987) / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) + 1 - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() + ) async def _reserved(deployment_id: str, model_info: dict[str, object]) -> int: metadata: Final[dict[str, object]] = {} @@ -542,7 +552,9 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_base_model_does_not_change_the_default_otpm_reservation(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() + ) otpm_limit: Final = AZURE_BASE_MODEL_MAX_OUTPUT * 10 async def _reserved(model_info: dict[str, str]) -> int: @@ -568,7 +580,9 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_unparseable_max_tokens_leaves_itpm_uncharged(self, use_sync_check: bool): dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache + ) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "itpm": 2000, "otpm": 2000}, "model_info": {"id": "io-bad-max-tokens-id"}, @@ -596,7 +610,9 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_explicit_max_tokens_above_default_share_is_reserved_in_full(self): dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache + ) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "otpm": 100}, "model_info": {"id": "io-explicit-max-tokens-id"}, @@ -617,7 +633,9 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_deployment_max_tokens_sets_the_reservation_unless_the_request_sets_one(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() + ) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, "max_tokens": 100}, "model_info": {"id": "io-deployment-max-tokens-id"}, @@ -638,7 +656,9 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_deployment_cap_sets_the_reservation_only_through_the_param_its_api_sends(self): - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() + ) chat: Final = {"messages": [{"role": "user", "content": "hello"}]} responses: Final = {"input": "hello"} @@ -668,7 +688,9 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_deployment_cap_above_the_default_reserves_the_default(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck( + usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache() + ) otpm_limit: Final = UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT * 5 async def reserved(deployment_cap: Mapping[str, int]) -> int: @@ -694,7 +716,7 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -735,7 +757,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_reconcile_runs_via_success_event_without_model_id(self): dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) @@ -761,7 +783,8 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = RouterUsageStore(connect(None).router_usage()) + check = ModelRateLimitingCheck(usage=usage, dual_cache=dual_cache) minute = get_utc_datetime().strftime("%H-%M") itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) @@ -801,8 +824,8 @@ class TestModelRateLimitingCheckIOTokens: # The first deployment's ITPM counter is not driven negative... assert await dual_cache.async_get_cache(key=itpm_key) == 0 # ...and the non-IO deployment's TPM usage is tracked normally. - tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" - assert await dual_cache.async_get_cache(key=tpm_key) == 12 + tpm = UsageCounter.for_deployment("non-io-second", "openai/gpt-4o-mini", "tpm", minute) + assert usage.get_local(tpm) == 12 @pytest.mark.asyncio async def test_stale_reservation_refunded_before_retry_overwrites_it(self): @@ -872,7 +895,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_forged_reservation_cannot_decrement_counter(self): dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) victim_key = "global_router:victim:model:itpm:00-00" await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) @@ -962,7 +985,7 @@ class TestModelRateLimitingCheckIOTokens: async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): import logging - check = ModelRateLimitingCheck(dual_cache=DualCache()) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=DualCache()) deployment = { "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, "model_info": {}, @@ -1167,7 +1190,8 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = RouterUsageStore(connect(None).router_usage()) + check = ModelRateLimitingCheck(usage=usage, dual_cache=dual_cache) model_id = "io-mixed-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" deployment = { @@ -1181,9 +1205,9 @@ class TestModelRateLimitingCheckIOTokens: } minute = get_utc_datetime().strftime("%H-%M") - rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" + rpm = UsageCounter.for_deployment(model_id, deployment_name, "rpm", minute) itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) + usage.increment_local(CounterIncrement(rpm, 5, 60)) request_kwargs = { "messages": [{"role": "user", "content": "hi"}], @@ -1196,7 +1220,7 @@ class TestModelRateLimitingCheckIOTokens: with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) - assert await dual_cache.async_get_cache(key=rpm_key) == 5, "the rejected request hands its rpm count back" + assert usage.get_local(rpm) == 5, "the rejected request hands its rpm count back" assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] assert any("both limit types are enforced" in record.message for record in caplog.records) @@ -1211,12 +1235,13 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = RouterUsageStore(connect(None).router_usage()) + check = ModelRateLimitingCheck(usage=usage, dual_cache=dual_cache) model_id = "io-tpm-mixed-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" minute = get_utc_datetime().strftime("%H-%M") itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + tpm = UsageCounter.for_deployment(model_id, deployment_name, "tpm", minute) await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) kwargs = { @@ -1237,18 +1262,19 @@ class TestModelRateLimitingCheckIOTokens: # ITPM reconciled down from the 5-token reservation to actual usage (3). assert await dual_cache.async_get_cache(key=itpm_key) == 3 # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. - assert await dual_cache.async_get_cache(key=tpm_key) == 7 + assert usage.get_local(tpm) == 7 def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = RouterUsageStore(connect(None).router_usage()) + check = ModelRateLimitingCheck(usage=usage, dual_cache=dual_cache) model_id = "io-tpm-mixed-sync-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" minute = get_utc_datetime().strftime("%H-%M") itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + tpm = UsageCounter.for_deployment(model_id, deployment_name, "tpm", minute) dual_cache.set_cache(key=itpm_key, value=5, ttl=60) kwargs = { @@ -1267,14 +1293,14 @@ class TestModelRateLimitingCheckIOTokens: check.log_success_event(kwargs, response, None, None) assert dual_cache.get_cache(key=itpm_key) == 3 - assert dual_cache.get_cache(key=tpm_key) == 7 + assert usage.get_local(tpm) == 7 @pytest.mark.asyncio async def test_failure_refunds_itpm_reservation(self): from litellm.utils import get_utc_datetime dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) minute = get_utc_datetime().strftime("%H-%M") itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) @@ -1619,7 +1645,7 @@ async def test_success_settles_a_reservation_held_only_in_litellm_params_metadat deliveries: tuple[str, ...], ) -> None: dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + check: Final = ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache) itpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:itpm:12-34" otpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:otpm:12-34" await dual_cache.async_increment_cache(key=itpm_key, value=40, ttl=60) @@ -1760,7 +1786,7 @@ async def test_a_turn_finishing_in_the_next_minute_counts_its_actual_tokens_on_t reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) await _deliver_success( - ModelRateLimitingCheck(dual_cache=dual_cache), + ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache), delivery, _proxy_success_kwargs(reservation), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]),