From 84373ee1880ca86fa078406c428d6f7ab7662406 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 2 Nov 2026 00:17:51 +0000 Subject: [PATCH] feat(redis): move router cooldown and usage stores into Rust --- litellm-rust/Cargo.lock | 27 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/cache-memory/src/cache.rs | 17 + litellm-rust/crates/cache/Cargo.toml | 2 + litellm-rust/crates/cache/src/dual.rs | 744 ++++++---------- litellm-rust/crates/cache/src/lib.rs | 4 +- litellm-rust/crates/cache/tests/dual.rs | 839 ------------------ litellm-rust/crates/config/AGENTS.md | 18 +- litellm-rust/crates/config/src/redis.rs | 28 +- litellm-rust/crates/config/src/redis/error.rs | 2 - litellm-rust/crates/config/src/redis/keys.rs | 75 +- litellm-rust/crates/config/src/redis/tls.rs | 22 +- litellm-rust/crates/config/tests/redis.rs | 231 +++-- litellm-rust/crates/gateway/src/redis.rs | 69 +- litellm-rust/crates/python-bridge/Cargo.toml | 6 + litellm-rust/crates/python-bridge/src/lib.rs | 14 + .../python-bridge/src/python_settings.rs | 2 + .../crates/python-bridge/src/redis/client.rs | 145 +++ .../crates/python-bridge/src/redis/events.rs | 251 ++++++ .../crates/python-bridge/src/redis/mod.rs | 61 ++ .../crates/python-bridge/src/redis/policy.rs | 130 +++ .../crates/python-bridge/src/redis/request.rs | 40 + .../src/redis/router_cooldowns.rs | 244 +++++ .../python-bridge/src/redis/router_usage.rs | 160 ++++ litellm-rust/crates/redis-client/AGENTS.md | 1 + litellm-rust/crates/redis-client/Cargo.toml | 1 + litellm-rust/crates/redis-client/src/batch.rs | 4 + .../crates/redis-client/src/client.rs | 43 +- .../crates/redis-client/src/config.rs | 77 +- .../crates/redis-client/src/connection.rs | 26 + .../crates/redis-client/tests/batch.rs | 48 + .../crates/redis-client/tests/client.rs | 97 +- .../crates/redis-client/tests/config.rs | 23 + .../crates/redis-client/tests/connection.rs | 3 + .../crates/redis-client/tests/support/mod.rs | 2 + litellm-rust/crates/redis-connect/Cargo.toml | 1 + .../crates/redis-connect/src/connect.rs | 79 +- litellm-rust/crates/redis-connect/src/lib.rs | 2 +- .../crates/redis-connect/src/refresh.rs | 90 +- litellm-rust/crates/redis-store/AGENTS.md | 7 + litellm-rust/crates/redis-store/Cargo.toml | 24 + litellm-rust/crates/redis-store/src/codec.rs | 20 + litellm-rust/crates/redis-store/src/lib.rs | 9 + litellm-rust/crates/redis-store/src/number.rs | 68 ++ .../crates/redis-store/src/request.rs | 99 +++ .../redis-store/src/router_cooldowns.rs | 282 ++++++ .../crates/redis-store/src/router_usage.rs | 261 ++++++ .../crates/redis-store/tests/request.rs | 114 +++ .../redis-store/tests/router_cooldowns.rs | 394 ++++++++ .../crates/redis-store/tests/router_usage.rs | 240 +++++ .../crates/redis-store/tests/support/mod.rs | 188 ++++ .../crates/redis-testing/src/cluster.rs | 9 + .../crates/redis-testing/src/sentinel.rs | 18 +- litellm-rust/crates/redis-types/src/lib.rs | 7 + litellm/caching/redis_batch.py | 95 +- litellm/caching/redis_cache.py | 10 +- litellm/redis/stores/_connection.py | 14 + litellm/redis/stores/_routing_prefetch.py | 149 ---- litellm/redis/stores/router_cooldowns.py | 188 ++-- litellm/redis/stores/router_usage.py | 112 +-- litellm/router_utils/routing_prefetch.py | 165 ++++ litellm/router_utils/routing_read_batch.py | 36 +- litellm/rust_bridge/_native.pyi | 63 +- litellm/rust_bridge/host/AGENTS.md | 2 +- litellm/rust_bridge/host/redis_events.py | 46 + litellm/rust_bridge/host/settings.py | 20 + tests/unit/caching/test_redis_batch.py | 42 +- .../test_request_redis_batch_post_call.py | 30 +- .../test_request_redis_batch_pre_call.py | 660 +++++--------- tests/unit/conftest.py | 15 + .../redis/stores/test_router_cooldowns.py | 329 +++++-- tests/unit/redis/stores/test_router_usage.py | 158 ++-- .../test_router_tag_routing.py | 30 +- .../test_session_id_affinity.py | 38 +- .../unit/router_utils/test_cooldown_cache.py | 20 +- .../router_utils/test_cooldown_handlers.py | 79 +- .../router_utils/test_routing_read_batch.py | 147 ++- .../rust_bridge/host/test_redis_events.py | 81 ++ tests/unit/rust_bridge/host/test_settings.py | 25 +- tests/unit/test_internal_context.py | 2 +- tests/unit/test_redis.py | 160 +++- tests/unit/test_router/test_router.py | 77 +- .../test_router/test_router_helper_utils.py | 50 +- .../test_router_provider_endpoints.py | 29 +- 84 files changed, 5474 insertions(+), 2767 deletions(-) delete mode 100644 litellm-rust/crates/cache/tests/dual.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/client.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/events.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/policy.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/request.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/router_usage.rs create mode 100644 litellm-rust/crates/redis-store/AGENTS.md create mode 100644 litellm-rust/crates/redis-store/Cargo.toml create mode 100644 litellm-rust/crates/redis-store/src/codec.rs create mode 100644 litellm-rust/crates/redis-store/src/lib.rs create mode 100644 litellm-rust/crates/redis-store/src/number.rs create mode 100644 litellm-rust/crates/redis-store/src/request.rs create mode 100644 litellm-rust/crates/redis-store/src/router_cooldowns.rs create mode 100644 litellm-rust/crates/redis-store/src/router_usage.rs create mode 100644 litellm-rust/crates/redis-store/tests/request.rs create mode 100644 litellm-rust/crates/redis-store/tests/router_cooldowns.rs create mode 100644 litellm-rust/crates/redis-store/tests/router_usage.rs create mode 100644 litellm-rust/crates/redis-store/tests/support/mod.rs create mode 100644 litellm/redis/stores/_connection.py delete mode 100644 litellm/redis/stores/_routing_prefetch.py create mode 100644 litellm/router_utils/routing_prefetch.py create mode 100644 litellm/rust_bridge/host/redis_events.py create mode 100644 tests/unit/rust_bridge/host/test_redis_events.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e303ab4edb7..150037820e9 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3474,12 +3474,14 @@ name = "litellm-cache" version = "0.1.0" dependencies = [ "litellm-error", + "litellm-redis-client", "rstest", "serde", "serde_json", "strum", "thiserror 2.0.19", "tokio", + "tracing", ] [[package]] @@ -4358,6 +4360,8 @@ dependencies = [ "litellm-cache-redis", "litellm-cache-response", "litellm-callbacks-legacy-python", + "litellm-circuit-breaker", + "litellm-config", "litellm-core-utils", "litellm-error", "litellm-host", @@ -4372,6 +4376,9 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-redis-client", + "litellm-redis-connect", + "litellm-redis-store", "litellm-router-types", "litellm-secret-value", "litellm-secrets", @@ -4433,6 +4440,7 @@ dependencies = [ "socket2 0.6.5", "thiserror 2.0.19", "tokio", + "tracing", ] [[package]] @@ -4465,6 +4473,25 @@ dependencies = [ "veil", ] +[[package]] +name = "litellm-redis-store" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-cache", + "litellm-cache-memory", + "litellm-circuit-breaker", + "litellm-python-compat", + "litellm-redis-client", + "litellm-redis-testing", + "litellm-secret-value", + "rstest", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tracing", +] + [[package]] name = "litellm-redis-testing" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 76a588b6d9c..40f13e0f152 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -74,6 +74,7 @@ litellm-python-compat = { path = "crates/python-compat" } litellm-circuit-breaker = { path = "crates/circuit-breaker" } litellm-redis-client = { path = "crates/redis-client" } litellm-redis-connect = { path = "crates/redis-connect" } +litellm-redis-store = { path = "crates/redis-store" } litellm-redis-testing = { path = "crates/redis-testing" } litellm-redis-types = { path = "crates/redis-types" } litellm-cost = { path = "crates/cost" } diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 73fbf14b07a..b8706b5e762 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -200,6 +200,23 @@ impl InMemoryCache { .collect()) } + /// Every live entry, for callers that inspect the whole tier. + pub fn entries(&self) -> Result, Error> { + let now = (self.now)(); + let state = self.lock()?; + Ok(state + .values + .iter() + .filter(|(key, _)| { + state + .expirations + .get(*key) + .is_none_or(|expiration| *expiration >= now) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect()) + } + pub fn delete_cache(&self, key: &str) -> Result<(), Error> { let mut state = self.lock()?; Self::remove(&mut state, key); diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index 282c34e6b02..f7c054a5d52 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -7,10 +7,12 @@ repository.workspace = true [dependencies] litellm-error.workspace = true +litellm-redis-client.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } strum.workspace = true thiserror.workspace = true +tracing.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs index 3deee2560f4..9a44298f108 100644 --- a/litellm-rust/crates/cache/src/dual.rs +++ b/litellm-rust/crates/cache/src/dual.rs @@ -1,519 +1,319 @@ -use std::{sync::Arc, time::Duration}; - -use crate::{ - BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheContext, ClaimCache, CounterCache, - DeleteCache, Error, FlushCache, IncrementOperation, SetCache, ShortBatch, TtlCache, +use std::{ + collections::{HashMap, VecDeque}, + sync::{Arc, Mutex, MutexGuard}, + time::{Duration, SystemTime, UNIX_EPOCH}, }; -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub enum ReadPolicy { - #[default] - LocalThenRemote, - LocalOnly, +use litellm_redis_client::{Declared, Error as RedisError, RedisBatch, RedisClient}; + +use crate::{BaseCache, DeleteCache, Error, ExactCacheContext, FlushCache}; + +pub type Clock = Arc Duration + Send + Sync>; + +/// How a read-through family uses its memory tier, declared once by the family. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TierPolicy { + /// Python's `redis_batch_cache_expiry`: a key memory misses is read from Redis at most this often. + pub refresh_after: Duration, + /// How long a value read from Redis stays in memory. + pub backfill_ttl: Duration, + pub on_breaker_open: OnBreakerOpen, + /// Python's `default_max_redis_batch_cache_size`: how many keys the refresh throttle tracks. + pub tracked_reads: usize, } -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub enum WritePolicy { - #[default] - Both, - LocalOnly, -} - -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub enum RemoteFailurePolicy { - #[default] - Propagate, - UseLocal, +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OnBreakerOpen { + /// Serve what memory holds and let the next read try Redis again. + ServeMemory, } +/// The memory tier in front of Redis that a read-through family builds for itself, as Python's +/// `DualCache` with a `RedisCache` does. pub struct DualCache { - l1: Arc, - l2: Arc, - read_policy: ReadPolicy, - write_policy: WritePolicy, - remote_failure_policy: RemoteFailurePolicy, - promotion_ttl: Option, - delete_batch_size: usize, + memory: L1, + remote: Option, + policy: TierPolicy, + reads: Mutex, + updates: Mutex<()>, + now: Clock, } -/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`. -pub const DEFAULT_DELETE_BATCH_SIZE: usize = 1000; - -impl DualCache { - pub fn new(l1: Arc, l2: Arc) -> Self { - Self { - l1, - l2, - read_policy: ReadPolicy::default(), - write_policy: WritePolicy::default(), - remote_failure_policy: RemoteFailurePolicy::default(), - promotion_ttl: None, - delete_batch_size: DEFAULT_DELETE_BATCH_SIZE, - } - } - - pub fn with_delete_batch_size(self, delete_batch_size: usize) -> Self { - Self { - delete_batch_size, - ..self - } - } - - pub fn with_read_policy(self, read_policy: ReadPolicy) -> Self { - Self { - read_policy, - ..self - } - } - - pub fn with_write_policy(self, write_policy: WritePolicy) -> Self { - Self { - write_policy, - ..self - } - } - - pub fn with_remote_failure_policy(self, remote_failure_policy: RemoteFailurePolicy) -> Self { - Self { - remote_failure_policy, - ..self - } - } - - pub fn with_promotion_ttl(self, promotion_ttl: Duration) -> Self { - Self { - promotion_ttl: Some(promotion_ttl), - ..self - } - } - - fn reads_remote(&self) -> bool { - self.read_policy == ReadPolicy::LocalThenRemote - } - - fn writes_remote(&self) -> bool { - self.write_policy == WritePolicy::Both - } - - fn remote(&self, result: Result) -> Result, Error> { - match result { - Ok(value) => Ok(Some(value)), - Err(Error::Unavailable(_)) - if self.remote_failure_policy == RemoteFailurePolicy::UseLocal => - { - Ok(None) - } - Err(error) => Err(error), - } - } - - fn promotion_context(&self, context: &C) -> C { - context.with_ttl(self.promotion_ttl.or(context.ttl())) - } +pub fn system_clock() -> Clock { + Arc::new(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + }) } -impl DualCache +impl DualCache where V: Clone + Send + Sync + 'static, - C: CacheContext, - L1: BaseCache, - L2: BaseCache, + L1: BaseCache + DeleteCache + FlushCache, { - fn missing(entries: &[BatchEntry]) -> Vec { - entries + pub fn read_through( + memory: L1, + remote: Option, + policy: TierPolicy, + now: Clock, + ) -> Self { + Self { + memory, + remote, + reads: Mutex::new(RecentReads::new(policy.tracked_reads)), + policy, + updates: Mutex::new(()), + now, + } + } + + pub fn memory(&self) -> &L1 { + &self.memory + } + + pub fn remote(&self) -> Option<&RedisClient> { + self.remote.as_ref() + } + + pub fn now(&self) -> Duration { + (self.now)() + } + + pub fn set_memory(&self, key: &str, value: V, ttl: Duration) -> Result<(), Error> { + self.memory + .set_cache(key, value, &ExactCacheContext { ttl: Some(ttl) }) + } + + /// Read-modify-write of one memory entry, atomic against every other update of this tier. + pub fn update_memory( + &self, + key: &str, + ttl: Duration, + update: impl FnOnce(Option) -> V, + ) -> Result { + let _guard = lock(&self.updates)?; + let value = update(self.memory.get_cache(key, &ExactCacheContext::default())?); + self.set_memory(key, value.clone(), ttl)?; + Ok(value) + } + + /// Python's `DualCache.async_batch_get_cache`: memory first, then one round trip for the misses + /// whose refresh is due, written back to memory. A failed Redis read serves memory, as + /// `RedisCache` swallows it, and only a breaker refusal releases the refresh it reserved. + pub async fn get_many( + &self, + keys: &[String], + decode: impl Fn(&[u8]) -> Result, + ) -> Result>, Error> { + let declared = self.declare_many(RedisClient::batch, keys)?; + self.resolve_many(declared, decode).await + } + + /// The first half of `get_many`: reserves the due refreshes and declares their GETs on the + /// batch `batch_for` picks for this tier's client, without sending anything. + pub fn declare_many( + &self, + batch_for: impl FnOnce(&RedisClient) -> RedisBatch, + keys: &[String], + ) -> Result { + let fetch = match &self.remote { + None => None, + Some(remote) => { + let reservation = self.reserve(keys, &self.read_memory(keys)?)?; + (!reservation.keys.is_empty()).then(|| { + let batch = batch_for(remote); + let replies = reservation + .keys + .iter() + .map(|key| batch.get::>(key)) + .collect(); + Fetch { + reservation, + replies, + } + }) + } + }; + Ok(DeclaredRead { + keys: keys.to_vec(), + fetch, + }) + } + + /// The second half of `get_many`. Memory is read again here, so a value written to memory + /// since the declaration wins over what Redis answered. + pub async fn resolve_many( + &self, + declared: DeclaredRead, + decode: impl Fn(&[u8]) -> Result, + ) -> Result>, Error> { + let DeclaredRead { keys, fetch } = declared; + let Some(Fetch { + reservation, + replies, + }) = fetch + else { + return self.read_memory(&keys); + }; + let fetched = settle(replies).await; + let local = self.read_memory(&keys)?; + let fetched = match fetched { + Ok(fetched) => fetched, + Err(RedisError::BreakerOpen) => { + self.release(reservation)?; + return Ok(local); + } + Err(error) => { + tracing::warn!(%error, "Redis read failed, serving the memory tier"); + return Ok(local); + } + }; + let decoded = match fetched .iter() - .enumerate() - .filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index)) + .map(|raw| raw.as_deref().map(&decode).transpose()) + .collect::, _>>() + { + Ok(decoded) => decoded, + Err(error) => { + tracing::warn!(%error, "a Redis value did not decode, serving the memory tier"); + return Ok(local); + } + }; + let read = reservation + .keys + .into_iter() + .zip(decoded) + .collect::>(); + keys.iter() + .zip(local) + .map(|(key, local)| match local { + Some(value) => Ok(Some(value)), + None => self.backfill(key, read.get(key).cloned().flatten()), + }) .collect() } - fn merge_batch( - &self, - keys: &[String], - context: &C, - mut entries: Vec>, - missing: Vec, - remote: Vec>, - ) -> Result>, Error> { - if missing.len() != remote.len() { - return Err(Error::unavailable(ShortBatch { - requested: missing.len(), - returned: remote.len(), - })); - } - for (index, entry) in missing.into_iter().zip(remote) { - if let BatchEntry::Hit(value) = &entry { - let promotion_context = self.promotion_context(context); - self.l1 - .set_cache(&keys[index], value.clone(), &promotion_context)?; - } - entries[index] = entry; - } - Ok(entries) - } -} - -impl BaseCache for DualCache -where - V: Clone + Send + Sync + 'static, - C: CacheContext, - L1: BaseCache, - L2: BaseCache, -{ - type Value = V; - type Context = C; - - fn get_ttl(&self, context: &Self::Context) -> Option { - self.l2.get_ttl(context) + pub fn delete_memory(&self, key: &str) -> Result<(), Error> { + self.memory.delete_cache(key) } - fn set_cache(&self, key: &str, value: V, context: &C) -> Result<(), Error> { - if self.writes_remote() { - self.remote(self.l2.set_cache(key, value.clone(), context))?; - } - self.l1.set_cache(key, value, context) + pub fn flush_memory(&self) -> Result<(), Error> { + self.memory.flush_cache() } - fn get_cache(&self, key: &str, context: &C) -> Result, Error> { - if let Some(value) = self.l1.get_cache(key, context)? { - return Ok(Some(value)); - } - if !self.reads_remote() { - return Ok(None); - } - let value = self.remote(self.l2.get_cache(key, context))?.flatten(); + fn read_memory(&self, keys: &[String]) -> Result>, Error> { + let context = ExactCacheContext::default(); + keys.iter() + .map(|key| self.memory.get_cache(key, &context)) + .collect() + } + + fn backfill(&self, key: &str, value: Option) -> Result, Error> { if let Some(value) = &value { - let promotion_context = self.promotion_context(context); - self.l1.set_cache(key, value.clone(), &promotion_context)?; + self.set_memory(key, value.clone(), self.policy.backfill_ttl)?; } Ok(value) } - async fn async_set_cache(&self, key: &str, value: V, context: C) -> Result<(), Error> { - if self.writes_remote() { - self.remote( - self.l2 - .async_set_cache(key, value.clone(), context.clone()) - .await, - )?; - } - self.l1.async_set_cache(key, value, context).await - } - - async fn async_get_cache(&self, key: &str, context: &C) -> Result, Error> { - if let Some(value) = self.l1.async_get_cache(key, context).await? { - return Ok(Some(value)); - } - if !self.reads_remote() { - return Ok(None); - } - let value = self - .remote(self.l2.async_get_cache(key, context).await)? - .flatten(); - if let Some(value) = &value { - self.l1 - .async_set_cache(key, value.clone(), self.promotion_context(context)) - .await?; - } - Ok(value) - } - - async fn async_set_cache_pipeline( - &self, - entries: Vec<(String, V)>, - context: C, - ) -> Result<(), Error> { - if self.writes_remote() { - self.remote( - self.l2 - .async_set_cache_pipeline(entries.clone(), context.clone()) - .await, - )?; - } - self.l1.async_set_cache_pipeline(entries, context).await - } -} - -impl BatchCache for DualCache -where - V: Clone + Send + Sync + 'static, - C: CacheContext, - L1: BatchCache, - L2: BatchCache, -{ - fn batch_get_cache(&self, keys: &[String], context: &C) -> Result>, Error> { - let entries = self.l1.batch_get_cache(keys, context)?; - let missing = Self::missing(&entries); - if missing.is_empty() || !self.reads_remote() { - return Ok(entries); - } - let remote_keys = missing + fn reserve(&self, keys: &[String], local: &[Option]) -> Result { + let now = self.now(); + let mut reads = lock(&self.reads)?; + let due = keys .iter() - .map(|index| keys[*index].clone()) - .collect::>(); - match self.remote(self.l2.batch_get_cache(&remote_keys, context))? { - Some(remote) => self.merge_batch(keys, context, entries, missing, remote), - None => Ok(entries), - } - } - - async fn async_batch_get_cache( - &self, - keys: Vec, - context: C, - ) -> Result>, Error> { - let entries = self - .l1 - .async_batch_get_cache(keys.clone(), context.clone()) - .await?; - let missing = Self::missing(&entries); - if missing.is_empty() || !self.reads_remote() { - return Ok(entries); - } - let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect(); - match self.remote( - self.l2 - .async_batch_get_cache(remote_keys, context.clone()) - .await, - )? { - Some(remote) => self.merge_batch(&keys, &context, entries, missing, remote), - None => Ok(entries), - } - } -} - -impl DeleteCache for DualCache -where - V: Clone + Send + Sync + 'static, - C: CacheContext, - L1: DeleteCache, - L2: DeleteCache, -{ - fn delete_cache(&self, key: &str) -> Result<(), Error> { - if self.writes_remote() { - self.remote(self.l2.delete_cache(key))?; - } - self.l1.delete_cache(key) - } - - async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - if self.writes_remote() { - self.remote(self.l2.async_delete_cache(key).await)?; - } - self.l1.async_delete_cache(key).await - } -} - -impl FlushCache for DualCache -where - V: Clone + Send + Sync + 'static, - C: CacheContext, - L1: FlushCache, - L2: FlushCache, -{ - fn flush_cache(&self) -> Result<(), Error> { - if self.writes_remote() { - self.remote(self.l2.flush_cache())?; - } - self.l1.flush_cache() - } - - async fn async_flush_cache(&self) -> Result<(), Error> { - if self.writes_remote() { - self.remote(self.l2.async_flush_cache().await)?; - } - self.l1.async_flush_cache().await - } -} - -impl DualCache -where - C: CacheContext, - L1: BaseCache, -{ - /// Python's `local_only=True` increment: the local tier alone, read then written back. - fn increment_local(&self, key: &str, amount: f64, context: &C) -> Result { - let value = self.l1.get_cache(key, context)?.unwrap_or(0.0) + amount; - self.l1.set_cache(key, value, context)?; - Ok(value) - } -} - -impl CounterCache for DualCache -where - C: CacheContext, - L1: BaseCache, - L2: CounterCache, -{ - fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result { - if !self.writes_remote() { - return self.increment_local(key, amount, &context); - } - let value = self.l2.increment_cache(key, amount, context.clone())?; - self.l1.set_cache(key, value, &context)?; - Ok(value) - } - - async fn async_increment( - &self, - key: &str, - amount: f64, - context: C, - refresh_ttl: bool, - ) -> Result { - if !self.writes_remote() { - return self.increment_local(key, amount, &context); - } - let value = self - .l2 - .async_increment(key, amount, context.clone(), refresh_ttl) - .await?; - self.l1.async_set_cache(key, value, context).await?; - Ok(value) - } - - /// `async_increment_cache_pipeline`, L2-first like single increments: the local tier takes - /// each remote result. - async fn async_increment_pipeline( - &self, - operations: Vec, - ) -> Result, Error> - where - C: Default, - { - if !self.writes_remote() { - return operations - .iter() - .map(|operation| { - let context = C::default().with_ttl(operation.ttl); - self.increment_local(&operation.key, operation.amount, &context) + .zip(local) + .filter(|(_, value)| value.is_none()) + .filter_map(|(key, _)| { + let previous = reads.get(key); + let due = previous + .is_none_or(|read| now.saturating_sub(read) >= self.policy.refresh_after); + due.then(|| { + reads.insert(key.clone(), now); + (key.clone(), previous) }) - .collect(); - } - let values = self.l2.async_increment_pipeline(operations.clone()).await?; - if values.len() != operations.len() { - return Err(Error::unavailable(ShortBatch { - requested: operations.len(), - returned: values.len(), - })); - } - for (operation, value) in operations.iter().zip(&values) { - self.l1 - .async_set_cache(&operation.key, *value, C::default().with_ttl(operation.ttl)) - .await?; - } - Ok(values) + }) + .collect::>(); + let (keys, previous) = due.into_iter().unzip(); + Ok(Reservation { keys, previous }) } -} -/// `async_set_cache_sadd`: local set first, then the remote one unless writes stay local. -impl SetCache for DualCache -where - V: Clone + Send + Sync + 'static, - C: CacheContext, - S: Clone + Send + Sync + 'static, - L1: SetCache, - L2: SetCache, -{ - type SetValue = S; - type SetResult = (); - - async fn async_set_cache_sadd( - &self, - key: &str, - values: Vec, - ttl: Option, - ) -> Result<(), Error> { - self.l1 - .async_set_cache_sadd(key, values.clone(), ttl) - .await?; - if self.writes_remote() { - self.l2.async_set_cache_sadd(key, values, ttl).await?; + fn release(&self, reservation: Reservation) -> Result<(), Error> { + let mut reads = lock(&self.reads)?; + for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) { + match previous { + Some(read) => reads.insert(key, read), + None => reads.remove(&key), + } } Ok(()) } } -/// `async_delete_cache_keys`: every key leaves the local tier, then the remote tier in chunks -/// of `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, since Redis takes a chunk as one command. -impl BulkDeleteCache for DualCache -where - V: Clone + Send + Sync + 'static, - C: CacheContext, - L1: DeleteCache, - L2: BulkDeleteCache, -{ - async fn delete_cache_keys(&self, keys: Vec) -> Result { - for key in &keys { - self.l1.delete_cache(key)?; - } - let mut deleted = 0; - for chunk in keys.chunks(self.delete_batch_size.max(1)) { - deleted += self.l2.delete_cache_keys(chunk.to_vec()).await?; - } - Ok(deleted) - } +fn lock(mutex: &Mutex) -> Result, Error> { + Ok(mutex.lock()?) } -/// `async_get_ttl`: the local TTL, or the remote one when the local tier has none. -impl TtlCache for DualCache -where - V: Clone + Send + Sync + 'static, - C: CacheContext, - L1: TtlCache, - L2: TtlCache, -{ - async fn async_get_ttl(&self, key: &str) -> Result, Error> { - match self.l1.async_get_ttl(key).await? { - Some(ttl) => Ok(Some(ttl)), - None => self.l2.async_get_ttl(key).await, - } +async fn settle( + replies: Vec>>>, +) -> Result>>, RedisError> { + let mut fetched = Vec::with_capacity(replies.len()); + for reply in replies { + fetched.push(reply.resolve().await?); } + Ok(fetched) } -impl ClaimCache for DualCache -where - V: Clone + PartialEq + Send + Sync + 'static, - C: CacheContext, - L1: ClaimCache, - L2: ClaimCache, -{ - fn claim_cache(&self, key: &str, candidate: V, eligible: &[V], context: C) -> Result { - match self.remote( - self.l2 - .claim_cache(key, candidate.clone(), eligible, context.clone()), - )? { - Some(winner) => { - self.l1.set_cache(key, winner.clone(), &context)?; - Ok(winner) - } - None => self.l1.claim_cache(key, candidate, eligible, context), +/// The GETs `declare_many` put on a batch, with the refreshes they reserved. +pub struct DeclaredRead { + keys: Vec, + fetch: Option, +} + +struct Fetch { + reservation: Reservation, + replies: Vec>>>, +} + +struct Reservation { + keys: Vec, + previous: Vec>, +} + +/// Python's `LimitedSizeOrderedDict`: insertion-ordered, and any write at capacity drops the +/// oldest entry first, even when the written key is already present. +struct RecentReads { + capacity: usize, + order: VecDeque, + times: HashMap, +} + +impl RecentReads { + fn new(capacity: usize) -> Self { + Self { + capacity, + order: VecDeque::new(), + times: HashMap::new(), } } - async fn async_claim_cache( - &self, - key: &str, - candidate: V, - eligible: Vec, - context: C, - ) -> Result { - match self.remote( - self.l2 - .async_claim_cache(key, candidate.clone(), eligible.clone(), context.clone()) - .await, - )? { - Some(winner) => { - self.l1 - .async_set_cache(key, winner.clone(), context) - .await?; - Ok(winner) - } - None => { - self.l1 - .async_claim_cache(key, candidate, eligible, context) - .await - } + fn get(&self, key: &str) -> Option { + self.times.get(key).copied() + } + + fn insert(&mut self, key: String, time: Duration) { + if self.times.len() >= self.capacity + && let Some(oldest) = self.order.pop_front() + { + self.times.remove(&oldest); + } + if self.times.insert(key.clone(), time).is_none() { + self.order.push_back(key); + } + } + + fn remove(&mut self, key: &str) { + if self.times.remove(key).is_some() { + self.order.retain(|tracked| tracked != key); } } } diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index 4e07f3b4e8c..12ed6040d6b 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -20,7 +20,5 @@ pub use capabilities::{ RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache, TtlPipelineCache, }; pub use codec::{CacheCodec, JsonCodec}; -pub use dual::{ - DEFAULT_DELETE_BATCH_SIZE, DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy, -}; +pub use dual::{Clock, DeclaredRead, DualCache, OnBreakerOpen, TierPolicy, system_clock}; pub use error::{Error, ShortBatch}; diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs deleted file mode 100644 index 86af5e80abb..00000000000 --- a/litellm-rust/crates/cache/tests/dual.rs +++ /dev/null @@ -1,839 +0,0 @@ -use std::{ - collections::HashMap, - sync::{Arc, Mutex}, - time::Duration, -}; - -use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, BulkDeleteCache, ClaimCache, CounterCache, DeleteCache, - DualCache, Error, ExactCacheContext, FlushCache, IncrementOperation, ReadPolicy, - RemoteFailurePolicy, SetCache, TtlCache, WritePolicy, -}; -use rstest::{fixture, rstest}; - -struct TestCache { - value: Mutex>, - fail: bool, -} - -impl TestCache { - fn new(value: Option, fail: bool) -> Self { - Self { - value: Mutex::new(value), - fail, - } - } -} - -impl BaseCache for TestCache -where - V: Clone + Send + Sync + 'static, -{ - type Value = V; - type Context = ExactCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl.or(Some(Duration::from_secs(60))) - } - - fn set_cache(&self, _: &str, value: V, _: &ExactCacheContext) -> Result<(), Error> { - *self.value.lock().unwrap() = Some(value); - Ok(()) - } - - fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { - Ok(self.value.lock().unwrap().clone()) - } -} - -impl BatchCache for TestCache where V: Clone + Send + Sync + 'static {} - -impl DeleteCache for TestCache -where - V: Clone + Send + Sync + 'static, -{ - fn delete_cache(&self, _: &str) -> Result<(), Error> { - *self.value.lock().unwrap() = None; - Ok(()) - } -} - -impl FlushCache for TestCache -where - V: Clone + Send + Sync + 'static, -{ - fn flush_cache(&self) -> Result<(), Error> { - *self.value.lock().unwrap() = None; - Ok(()) - } -} - -impl CounterCache for TestCache { - fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result { - if self.fail { - return Err(Error::Unavailable(None)); - } - let mut value = self.value.lock().unwrap(); - let incremented = value.unwrap_or_default() + amount; - *value = Some(incremented); - Ok(incremented) - } -} - -impl ClaimCache for TestCache -where - V: Clone + PartialEq + Send + Sync + 'static, -{ - fn claim_cache( - &self, - _: &str, - candidate: V, - eligible: &[V], - _: ExactCacheContext, - ) -> Result { - if self.fail { - return Err(Error::Unavailable(None)); - } - let mut value = self.value.lock().unwrap(); - let winner = match value.as_ref() { - Some(existing) if eligible.is_empty() || eligible.contains(existing) => { - existing.clone() - } - _ => candidate, - }; - *value = Some(winner.clone()); - Ok(winner) - } -} - -#[rstest] -fn failed_l2_increment_leaves_l1_unchanged() { - let l1 = Arc::new(TestCache::new(Some(10.0), false)); - let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true))); - - assert_eq!( - cache.increment_cache("counter", 2.0, ExactCacheContext::default()), - Err(Error::Unavailable(None)) - ); - assert_eq!( - l1.get_cache("counter", &ExactCacheContext::default()) - .unwrap(), - Some(10.0) - ); -} - -#[rstest] -fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { - let l1 = Arc::new(TestCache::new(Some("first".to_string()), false)); - let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true))) - .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); - - assert_eq!( - cache - .claim_cache( - "affinity", - "second".into(), - &["first".into(), "second".into()], - ExactCacheContext { - ttl: Some(Duration::from_secs(60)), - }, - ) - .unwrap(), - "first" - ); -} - -struct SyncPanics(TestCache); - -impl BaseCache for SyncPanics { - type Value = String; - type Context = ExactCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - self.0.get_ttl(context) - } - - fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> { - panic!("sync L2 write on an async path") - } - - fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { - panic!("sync L2 read on an async path") - } - - async fn async_set_cache( - &self, - key: &str, - value: String, - context: ExactCacheContext, - ) -> Result<(), Error> { - self.0.set_cache(key, value, &context) - } - - async fn async_get_cache( - &self, - key: &str, - context: &ExactCacheContext, - ) -> Result, Error> { - self.0.get_cache(key, context) - } - - async fn async_set_cache_pipeline( - &self, - cache_list: Vec<(String, String)>, - context: ExactCacheContext, - ) -> Result<(), Error> { - for (key, value) in cache_list { - self.0.set_cache(&key, value, &context)?; - } - Ok(()) - } -} - -impl BatchCache for SyncPanics { - async fn async_batch_get_cache( - &self, - keys: Vec, - context: ExactCacheContext, - ) -> Result>, Error> { - assert_eq!(keys, ["missing"]); - Ok(vec![match self.0.get_cache("missing", &context)? { - Some(value) => BatchEntry::Hit(value), - None => BatchEntry::Miss, - }]) - } -} - -impl DeleteCache for SyncPanics { - fn delete_cache(&self, _: &str) -> Result<(), Error> { - panic!("sync L2 delete on an async path") - } - - async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - self.0.delete_cache(key) - } -} - -impl FlushCache for SyncPanics { - fn flush_cache(&self) -> Result<(), Error> { - panic!("sync L2 flush on an async path") - } -} - -#[rstest] -#[tokio::test] -async fn async_operations_use_the_async_l2_methods() { - let l1 = Arc::new(TestCache::new(None, false)); - let cache = DualCache::new( - l1.clone(), - Arc::new(SyncPanics(TestCache::new( - Some("remote".to_string()), - false, - ))), - ); - let context = ExactCacheContext::default(); - - assert_eq!( - cache.async_get_cache("missing", &context).await.unwrap(), - Some("remote".into()) - ); - assert_eq!( - l1.get_cache("missing", &context).unwrap(), - Some("remote".into()) - ); - - l1.delete_cache("missing").unwrap(); - assert_eq!( - cache - .async_batch_get_cache(vec!["missing".into()], context.clone()) - .await - .unwrap(), - [BatchEntry::Hit("remote".to_string())] - ); - cache - .async_set_cache("missing", "written".into(), context.clone()) - .await - .unwrap(); - cache - .async_set_cache_pipeline(vec![("missing".into(), "piped".into())], context.clone()) - .await - .unwrap(); - cache.async_delete_cache("missing").await.unwrap(); - assert_eq!( - cache.async_get_cache("missing", &context).await.unwrap(), - None - ); -} - -struct Unavailable; - -impl BaseCache for Unavailable { - type Value = String; - type Context = ExactCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl - } - - fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> { - Err(Error::Unavailable(None)) - } - - fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { - Err(Error::Unavailable(None)) - } -} - -impl BatchCache for Unavailable {} - -impl DeleteCache for Unavailable { - fn delete_cache(&self, _: &str) -> Result<(), Error> { - Err(Error::Unavailable(None)) - } -} - -impl FlushCache for Unavailable { - fn flush_cache(&self) -> Result<(), Error> { - Err(Error::Unavailable(None)) - } -} - -impl ClaimCache for Unavailable { - fn claim_cache( - &self, - _: &str, - _: String, - _: &[String], - _: ExactCacheContext, - ) -> Result { - Err(Error::InvalidEntry(None)) - } -} - -#[rstest] -fn remote_failure_policy_selects_propagation_or_the_local_tier() { - let context = ExactCacheContext::default(); - let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable)); - assert_eq!( - strict.set_cache("key", "value".into(), &context), - Err(Error::Unavailable(None)) - ); - assert_eq!( - strict.get_cache("key", &context), - Err(Error::Unavailable(None)) - ); - - let l1 = Arc::new(TestCache::new(None, false)); - let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable)) - .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); - assert_eq!(degraded.get_cache("key", &context), Ok(None)); - degraded.set_cache("key", "value".into(), &context).unwrap(); - assert_eq!( - degraded.get_cache("key", &context), - Ok(Some("value".into())) - ); - degraded.delete_cache("key").unwrap(); - assert_eq!(l1.get_cache("key", &context), Ok(None)); -} - -#[rstest] -fn claim_fallback_does_not_hide_non_availability_errors() { - let cache = DualCache::new( - Arc::new(TestCache::new(Some("first".to_string()), false)), - Arc::new(Unavailable), - ) - .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); - assert_eq!( - cache.claim_cache( - "affinity", - "second".into(), - &[], - ExactCacheContext::default() - ), - Err(Error::InvalidEntry(None)) - ); -} - -#[rstest] -fn local_only_policies_never_touch_l2() { - let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false)); - let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone()) - .with_read_policy(ReadPolicy::LocalOnly) - .with_write_policy(WritePolicy::LocalOnly); - let context = ExactCacheContext::default(); - - assert_eq!(cache.get_cache("key", &context), Ok(None)); - cache.set_cache("key", "local".into(), &context).unwrap(); - assert_eq!(l2.get_cache("key", &context), Ok(Some("remote".into()))); -} - -type Log = Arc>>; -type StoredSet = (Vec, Option); - -/// A keyed tier that logs every call, so tests can assert which tier ran and in what order. -struct Tier { - name: &'static str, - log: Log, - fail: bool, - counters: Mutex>, - sets: Mutex>, - ttls: HashMap, -} - -impl Tier { - fn new(name: &'static str, log: &Log) -> Self { - Self { - name, - log: log.clone(), - fail: false, - counters: Mutex::default(), - sets: Mutex::default(), - ttls: HashMap::new(), - } - } - - fn failing(self) -> Self { - Self { fail: true, ..self } - } - - fn with_counter(self, key: &str, value: f64) -> Self { - self.counters - .lock() - .unwrap() - .insert(key.into(), (value, ExactCacheContext::default())); - self - } - - fn with_ttl(mut self, key: &str, seconds: u64) -> Self { - self.ttls.insert(key.into(), Duration::from_secs(seconds)); - self - } - - fn record(&self, event: String) { - self.log - .lock() - .unwrap() - .push(format!("{} {event}", self.name)); - } - - fn counter(&self, key: &str) -> Option<(f64, ExactCacheContext)> { - self.counters.lock().unwrap().get(key).cloned() - } - - fn check(&self) -> Result<(), Error> { - if self.fail { - return Err(Error::Unavailable(None)); - } - Ok(()) - } -} - -impl BaseCache for Tier { - type Value = f64; - type Context = ExactCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl - } - - fn set_cache(&self, key: &str, value: f64, context: &ExactCacheContext) -> Result<(), Error> { - self.check()?; - self.record(format!("set {key}={value}")); - self.counters - .lock() - .unwrap() - .insert(key.into(), (value, context.clone())); - Ok(()) - } - - fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { - Ok(self.counter(key).map(|(value, _)| value)) - } -} - -impl CounterCache for Tier { - fn increment_cache( - &self, - key: &str, - amount: f64, - context: ExactCacheContext, - ) -> Result { - self.check()?; - let mut counters = self.counters.lock().unwrap(); - let value = counters.get(key).map_or(0.0, |(value, _)| *value) + amount; - counters.insert(key.into(), (value, context)); - Ok(value) - } - - async fn async_increment( - &self, - key: &str, - amount: f64, - context: ExactCacheContext, - refresh_ttl: bool, - ) -> Result { - self.record(format!( - "increment {key}+{amount} refresh_ttl={refresh_ttl}" - )); - self.increment_cache(key, amount, context) - } - - async fn async_increment_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - let keys = operations - .iter() - .map(|operation| operation.key.as_str()) - .collect::>(); - self.record(format!("pipeline {}", keys.join(","))); - operations - .into_iter() - .map(|operation| { - self.increment_cache( - &operation.key, - operation.amount, - ExactCacheContext { ttl: operation.ttl }, - ) - }) - .collect() - } -} - -impl DeleteCache for Tier { - fn delete_cache(&self, key: &str) -> Result<(), Error> { - self.check()?; - self.record(format!("delete {key}")); - self.counters.lock().unwrap().remove(key); - Ok(()) - } -} - -impl BulkDeleteCache for Tier { - async fn delete_cache_keys(&self, keys: Vec) -> Result { - self.check()?; - self.record(format!("delete_keys {}", keys.join(","))); - let mut counters = self.counters.lock().unwrap(); - Ok(keys - .iter() - .filter(|key| counters.remove(key.as_str()).is_some()) - .count()) - } -} - -impl SetCache for Tier { - type SetValue = String; - type SetResult = (); - - async fn async_set_cache_sadd( - &self, - key: &str, - values: Vec, - ttl: Option, - ) -> Result<(), Error> { - self.check()?; - self.record(format!("sadd {key} {}", values.join(","))); - self.sets.lock().unwrap().insert(key.into(), (values, ttl)); - Ok(()) - } -} - -impl TtlCache for Tier { - async fn async_get_ttl(&self, key: &str) -> Result, Error> { - self.check()?; - self.record(format!("ttl {key}")); - Ok(self.ttls.get(key).copied()) - } -} - -#[fixture] -fn log() -> Log { - Log::default() -} - -fn events(log: &Log) -> Vec { - log.lock().unwrap().clone() -} - -fn seconds(ttl: u64) -> ExactCacheContext { - ExactCacheContext { - ttl: Some(Duration::from_secs(ttl)), - } -} - -#[rstest] -#[case::window_semantics(false)] -#[case::refresh_on_every_write(true)] -#[tokio::test] -async fn async_increment_passes_refresh_ttl_to_l2_and_stores_its_result_locally( - log: Log, - #[case] refresh_ttl: bool, -) { - let l1 = Arc::new(Tier::new("l1", &log).with_counter("counter", 1.0)); - let l2 = Arc::new(Tier::new("l2", &log).with_counter("counter", 10.0)); - let cache = DualCache::new(l1.clone(), l2.clone()); - - assert_eq!( - cache - .async_increment("counter", 2.0, seconds(30), refresh_ttl) - .await, - Ok(12.0) - ); - assert_eq!( - events(&log), - [ - format!("l2 increment counter+2 refresh_ttl={refresh_ttl}"), - "l1 set counter=12".into(), - ] - ); - assert_eq!(l1.counter("counter"), Some((12.0, seconds(30)))); - assert_eq!(l2.counter("counter"), Some((12.0, seconds(30)))); -} - -#[rstest] -#[tokio::test] -async fn failed_async_l2_increment_leaves_l1_unchanged(log: Log) { - let l1 = Arc::new(Tier::new("l1", &log).with_counter("counter", 1.0)); - let cache = DualCache::new(l1.clone(), Arc::new(Tier::new("l2", &log).failing())) - .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); - - assert_eq!( - cache - .async_increment("counter", 2.0, seconds(30), true) - .await, - Err(Error::Unavailable(None)) - ); - assert_eq!( - l1.counter("counter"), - Some((1.0, ExactCacheContext::default())) - ); -} - -#[rstest] -#[case::empty(Vec::new(), Vec::new())] -#[case::one_key(vec![("a", 1.0, Some(10))], vec![6.0])] -#[case::repeated_and_mixed_ttls( - vec![("a", 1.0, Some(10)), ("b", 2.0, None), ("a", 3.0, Some(20))], - vec![6.0, 2.0, 9.0], -)] -#[tokio::test] -async fn async_increment_pipeline_runs_l2_first_and_l1_takes_each_remote_result( - log: Log, - #[case] operations: Vec<(&str, f64, Option)>, - #[case] expected: Vec, -) { - let operations = operations - .into_iter() - .map(|(key, amount, ttl)| IncrementOperation { - key: key.into(), - amount, - ttl: ttl.map(Duration::from_secs), - }) - .collect::>(); - let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 100.0)); - let l2 = Arc::new(Tier::new("l2", &log).with_counter("a", 5.0)); - let cache = DualCache::new(l1.clone(), l2); - - assert_eq!( - cache.async_increment_pipeline(operations.clone()).await, - Ok(expected.clone()) - ); - let keys = operations - .iter() - .map(|operation| operation.key.as_str()) - .collect::>(); - let mut expected_events = vec![format!("l2 pipeline {}", keys.join(","))]; - expected_events.extend( - operations - .iter() - .zip(&expected) - .map(|(operation, value)| format!("l1 set {}={value}", operation.key)), - ); - assert_eq!(events(&log), expected_events); - if let Some((operation, value)) = operations.iter().zip(&expected).next_back() { - assert_eq!( - l1.counter(&operation.key), - Some((*value, ExactCacheContext { ttl: operation.ttl })) - ); - } -} - -#[rstest] -#[tokio::test] -async fn failed_l2_increment_pipeline_leaves_l1_unchanged(log: Log) { - let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 1.0)); - let cache = DualCache::new(l1.clone(), Arc::new(Tier::new("l2", &log).failing())); - - assert_eq!( - cache - .async_increment_pipeline(vec![IncrementOperation { - key: "a".into(), - amount: 1.0, - ttl: None, - }]) - .await, - Err(Error::Unavailable(None)) - ); - assert_eq!(events(&log), ["l2 pipeline a"]); - assert_eq!(l1.counter("a"), Some((1.0, ExactCacheContext::default()))); -} - -#[rstest] -#[case::both_tiers(WritePolicy::Both, &["l1 sadd members a,b", "l2 sadd members a,b"])] -#[case::local_only(WritePolicy::LocalOnly, &["l1 sadd members a,b"])] -#[tokio::test] -async fn set_add_writes_locally_then_remotely_unless_local_only( - log: Log, - #[case] write_policy: WritePolicy, - #[case] expected: &[&str], -) { - let l1 = Arc::new(Tier::new("l1", &log)); - let l2 = Arc::new(Tier::new("l2", &log)); - let cache = DualCache::new(l1.clone(), l2.clone()).with_write_policy(write_policy); - let ttl = Some(Duration::from_secs(45)); - - cache - .async_set_cache_sadd("members", vec!["a".into(), "b".into()], ttl) - .await - .unwrap(); - assert_eq!(events(&log), expected); - let stored = Some((vec!["a".to_string(), "b".to_string()], ttl)); - assert_eq!(l1.sets.lock().unwrap().get("members").cloned(), stored); - assert_eq!( - l2.sets.lock().unwrap().get("members").cloned(), - stored.filter(|_| write_policy == WritePolicy::Both) - ); -} - -#[rstest] -#[tokio::test] -async fn failed_local_set_add_never_reaches_l2(log: Log) { - let cache = DualCache::new( - Arc::new(Tier::new("l1", &log).failing()), - Arc::new(Tier::new("l2", &log)), - ); - assert_eq!( - cache - .async_set_cache_sadd("members", vec!["a".into()], None) - .await, - Err(Error::Unavailable(None)) - ); - assert!(events(&log).is_empty()); -} - -#[rstest] -#[case::empty(None, &[], &[])] -#[case::default_batch_size(None, &["a", "b", "c"], &["a,b,c"])] -#[case::chunked(Some(2), &["a", "b", "c", "d", "e"], &["a,b", "c,d", "e"])] -#[case::exact_chunks(Some(2), &["a", "b", "c", "d"], &["a,b", "c,d"])] -#[case::zero_means_one_per_chunk(Some(0), &["a", "b"], &["a", "b"])] -#[tokio::test] -async fn bulk_delete_removes_every_key_locally_then_remotely_in_chunks( - log: Log, - #[case] batch_size: Option, - #[case] keys: &[&str], - #[case] chunks: &[&str], -) { - let l1 = Tier::new("l1", &log); - let l2 = Tier::new("l2", &log); - for key in keys.iter().step_by(2) { - l2.counters - .lock() - .unwrap() - .insert((*key).into(), (1.0, ExactCacheContext::default())); - } - let l1 = Arc::new(l1); - let mut cache = DualCache::new(l1, Arc::new(l2)); - if let Some(batch_size) = batch_size { - cache = cache.with_delete_batch_size(batch_size); - } - - assert_eq!( - cache - .delete_cache_keys(keys.iter().map(|key| (*key).into()).collect()) - .await, - Ok(keys.len().div_ceil(2)) - ); - let expected = keys - .iter() - .map(|key| format!("l1 delete {key}")) - .chain(chunks.iter().map(|chunk| format!("l2 delete_keys {chunk}"))) - .collect::>(); - assert_eq!(events(&log), expected); -} - -#[rstest] -#[tokio::test] -async fn bulk_delete_stops_before_l2_when_the_local_delete_fails(log: Log) { - let cache = DualCache::new( - Arc::new(Tier::new("l1", &log).failing()), - Arc::new(Tier::new("l2", &log)), - ); - assert_eq!( - cache.delete_cache_keys(vec!["a".into()]).await, - Err(Error::Unavailable(None)) - ); - assert!(events(&log).is_empty()); -} - -#[rstest] -#[case::local_hit("both", Some(10), &["l1 ttl both"])] -#[case::remote_fallback("remote", Some(20), &["l1 ttl remote", "l2 ttl remote"])] -#[case::missing_everywhere("missing", None, &["l1 ttl missing", "l2 ttl missing"])] -#[tokio::test] -async fn ttl_reads_local_then_remote( - log: Log, - #[case] key: &str, - #[case] expected: Option, - #[case] expected_events: &[&str], -) { - let cache = DualCache::new( - Arc::new(Tier::new("l1", &log).with_ttl("both", 10)), - Arc::new( - Tier::new("l2", &log) - .with_ttl("both", 99) - .with_ttl("remote", 20), - ), - ); - assert_eq!( - cache.async_get_ttl(key).await, - Ok(expected.map(Duration::from_secs)) - ); - assert_eq!(events(&log), expected_events); -} - -/// Python `local_only=True`: the increment and the pipeline stay on the local tier. -#[rstest] -#[tokio::test] -async fn local_only_writes_increment_the_local_tier_alone(log: Log) { - let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 1.0)); - let l2 = Arc::new(Tier::new("l2", &log).with_counter("a", 10.0)); - let cache = DualCache::new(l1.clone(), l2.clone()).with_write_policy(WritePolicy::LocalOnly); - - assert_eq!(cache.increment_cache("a", 2.0, seconds(30)), Ok(3.0)); - assert_eq!( - cache.async_increment("a", 1.0, seconds(30), true).await, - Ok(4.0) - ); - let operations = vec![ - IncrementOperation { - key: "a".into(), - amount: 1.0, - ttl: Some(Duration::from_secs(10)), - }, - IncrementOperation { - key: "b".into(), - amount: 2.0, - ttl: None, - }, - ]; - assert_eq!( - cache.async_increment_pipeline(operations).await, - Ok(vec![5.0, 2.0]) - ); - assert_eq!( - events(&log), - ["l1 set a=3", "l1 set a=4", "l1 set a=5", "l1 set b=2"] - ); - assert_eq!(l2.counter("a"), Some((10.0, ExactCacheContext::default()))); -} diff --git a/litellm-rust/crates/config/AGENTS.md b/litellm-rust/crates/config/AGENTS.md index e1f77705b68..0e1251b2558 100644 --- a/litellm-rust/crates/config/AGENTS.md +++ b/litellm-rust/crates/config/AGENTS.md @@ -1,9 +1,9 @@ ## Redis settings - `RedisSettings::resolve` follows `litellm/_redis.py` `_get_redis_client_logic` plus `_credential_provider_auth_kwargs` and the per-client allowlists of `get_redis_async_client` (redis-py 5.3.1). Python and Rust pods share one Redis during the switch, so a change there is a change to a wire contract -- That contract is `resolves_like_python_redis` in `tests/redis.rs` and its mirror `test_redis_settings_resolve_like_the_rust_resolver_expects` in `tests/unit/test_redis.py`. Both hold the same named cases with the expected view written next to the input, the Python side reading the effective redis-py 5.3.1 connection objects without connecting. A new resolution rule is a new case in both files, never a generated fixture. A case where the two sides expect different outcomes is a divergence and is listed below +- That contract is `resolves_like_python_redis` in `tests/redis.rs` and its mirror `test_redis_settings_resolve_like_the_rust_resolver_expects` in `tests/unit/test_redis.py`. Both hold the same named cases with the expected view written next to the input, the Python side reading the effective redis-py 5.3.1 connection objects without connecting. A new resolution rule is a new case in both files, never a generated fixture - `resolve` takes a `SecretSnapshot` and an environment `Lookup`. The snapshot holds every name `RedisSettings::secret_names` lists, read the way Python's `get_secret` reads it (secret manager, then environment). The host fills it once before resolving, python-bridge through `get_secret` and rust-proxy from its environment, so the resolver never calls back into the host. The environment lookup is only for what Python reads straight from `os.environ` (`REDIS_USERNAME` for Azure). Reading a name the snapshot does not list fails a debug assertion -- Every Redis setting Python reads is one row in `SETTINGS` in `src/redis/keys.rs`: how Rust handles it (read, ignore, warn, reject), whether `_get_redis_env_kwarg_mapping` reads it as `REDIS_`, and the names `_get_redis_client_logic` falls back to. The env layer, the rejections, the warnings and `secret_names` all come from that table, so a new setting is one row. It is the shape the `#[derive(Layered)]` attributes in `litellm-layered` would generate, and moves onto them once that crate lands +- Every Redis setting Python reads is one row in `SETTINGS` in `src/redis/keys.rs`: how Rust handles it (read, or warn and drop), whether `_get_redis_env_kwarg_mapping` reads it as `REDIS_`, and the names `_get_redis_client_logic` falls back to. The env layer, the warnings and `secret_names` all come from that table, so a new setting is one row. It is the shape the `#[derive(Layered)]` attributes in `litellm-layered` would generate, and moves onto them once that crate lands - Lookups go through `litellm_core_utils::settings::Lookup`, never a closure type of our own - A url carries its own TLS. redis-py's url allowlist is the plain `Connection` signature, so every `ssl_*` setting outside the url is dropped, and a credential provider strips the url's userinfo. A `rediss://` url's `ssl_cert_reqs`, `ssl_ca_certs`, `ssl_certfile`, `ssl_keyfile` and `ssl_check_hostname` query parameters are read like `parse_url` reads them - TLS options and `db` are only read for the endpoint that uses them, so a bad `ssl_cert_reqs` never fails a url target. A cluster with a non-zero `db`, and Sentinel with a `host` or `port`, are errors because redis-py raises building those clients @@ -11,9 +11,11 @@ - Precedence matches Python. The target is cluster, then URL, then Sentinel (`sentinel_nodes` plus `service_name`), then host. Auth is an explicit `credential_provider`, then GCP, then Azure, then AWS, and the first one wins with a warning - Numbers are read with `litellm_python_compat::number`, so whitespace and digit underscores parse like Python's `int()` and `float()`. A value they cannot parse (`port: abc`, `socket_timeout: slow`) is dropped with a warning and the default is used, like `_coerce_redis_kwargs_types`. A number that parses but is out of range (negative seconds, a port above 65535) fails with `RedisError::Invalid` - `validate_coordination` mirrors the proxy's `CoordinationRedisParams` pydantic model, which rejects a malformed `port` or `ssl` in `general_settings.coordination_redis` before redis-py coerces anything. Hosts call it for that block only -- Known divergences, kept on purpose. Each is a case in both contract tests with a different expected outcome on each side - - A url query parameter other than the TLS ones, or any query parameter on a plain `redis://` url, is `RedisError::Unsupported` naming `url ?`. redis-py turns each one into a client kwarg - - Settings that would change the server, identity or TLS trust if ignored are `RedisError::Unsupported` naming the setting, including `ssl_ciphers`, `ssl_min_version`, `ssl_password`, `ssl_ca_data`, OCSP options, `unix_socket_path`, `connection_pool`, `event_dispatcher`, `redis_connect_func`, `retry` - - `ssl_ca_path` is rejected on both sides. redis-py 5.3.1's Redis constructor raises TypeError and Rust returns Unsupported naming the setting - - `retry_on_timeout`, `retry_on_error`, `protocol`, `cache`, `cache_config`, `health_check_interval`, `socket_keepalive` and `max_connections` are accepted with a warning and ignored. Cluster clients still turn TCP keepalive on, like Python's async cluster default, but redis-rs has no periodic health check, so Python's 25s cluster default has no equivalent - - A boolean `REDIS_RETRY_ON_ERROR` reaches redis-py's connection constructor as a non-iterable and raises TypeError. Rust ignores it with a warning and otherwise resolves the default host connection, verified against the host-default fixture +- The resolver never fails on a setting it cannot honor. A `Warn` row is dropped with a warning that names its source (`ssl_ciphers` or `REDIS_SSL_CIPHERS`) and the consequence from `keys::consequence`, because a Python deployment that worked must keep working when its stores move to Rust. Only a value that is malformed (`RedisError::Invalid`) or a target that is missing fails. Promoting a `Warn` row to `Read` is then a pure improvement with no behavior cliff + - `ssl_ciphers`, `ssl_min_version`, `ssl_password`, the OCSP options, `connection_pool`, `event_dispatcher`, `redis_connect_func` and `retry` are warned and dropped. Ignoring OCSP connects without revocation checks, and the warning says so + - `ssl_ca_data` and `ssl_ca_path` are read into `RedisTls` and trusted next to `ssl_ca_certs`. `unix_socket_path` is read into `RedisHost` and, like redis-py's `UnixDomainSocketConnection`, replaces host, port and TLS. Python still requires `host` next to it, so Rust does too + - A url query parameter other than the TLS ones, or any query parameter on a plain `redis://` url, is warned and dropped naming `url ?`. redis-py turns each one into a client kwarg + - `retry_on_timeout`, `retry_on_error`, `protocol`, `cache`, `cache_config`, `socket_keepalive` and `max_connections` are warned and dropped. Cluster clients still turn TCP keepalive on, like Python's async cluster default + - `health_check_interval` is read as an `int` of seconds like `_coerce_redis_kwargs_types`, zero meaning off. `redis-connect` gives cluster clients Python's 25s default and `redis-client` pings every node before the next command once the client sat idle that long, outside the breaker and the call events + - A `credential_provider` object resolves to `RedisAuth::Host`, the way Python's `{"kind": "host"}` does. `redis-connect` has no host credential source yet, so python-bridge leaves that family on memory with a warning, and rust-proxy never sees one because yaml cannot carry an object +- Cases where Rust and Python disagree on purpose are cases in both contract tests with a different expected outcome on each side. `ssl_ca_path` is one, redis-py 5.3.1's async `Redis` has no such parameter and raises `TypeError` building the client, while Rust trusts the directory like `load_verify_locations(capath=...)` diff --git a/litellm-rust/crates/config/src/redis.rs b/litellm-rust/crates/config/src/redis.rs index 16fba8f0087..38165f5f8c1 100644 --- a/litellm-rust/crates/config/src/redis.rs +++ b/litellm-rust/crates/config/src/redis.rs @@ -73,6 +73,8 @@ pub struct ResolvedRedis { pub namespace: Option, pub timeout: Option, pub connect_timeout: Option, + /// redis-py's `health_check_interval`, which `_coerce_redis_kwargs_types` reads as an `int` + pub health_check_interval: Option, } #[derive(Clone)] @@ -132,7 +134,6 @@ impl RedisSettings { .filter(|(key, _)| !(key == "url" && drop_environment_url)) .chain(block) .collect(); - reject_unsupported(&merged)?; warn_ignored(&merged); let startup_nodes = match configured(&merged, "startup_nodes") { @@ -197,6 +198,7 @@ impl RedisSettings { port: integer(&merged, "port")?.unwrap_or(6379), tls: tls()?, db: db()?, + unix_socket_path: text_value(&merged, "unix_socket_path").map(Into::into), }), }, }; @@ -213,6 +215,9 @@ impl RedisSettings { namespace: text_value(&merged, "namespace"), timeout: seconds(&merged, "socket_timeout")?, connect_timeout: seconds(&merged, "socket_connect_timeout")?, + health_check_interval: integer::(&merged, "health_check_interval")? + .filter(|seconds| *seconds > 0) + .map(|seconds| Duration::from_secs(seconds.unsigned_abs())), }) } @@ -279,15 +284,6 @@ pub fn environment_has_target(lookup: &dyn Lookup) -> bool { .any(|name| lookup.get(name).is_some()) } -fn reject_unsupported(merged: &Kwargs) -> Result<(), RedisError> { - match merged.iter().find(|(key, _)| is_unsupported(key)) { - Some((_, entry)) => Err(RedisError::Unsupported { - setting: entry.source.clone(), - }), - None => Ok(()), - } -} - fn invalid(setting: &str, expected: Expected) -> RedisError { RedisError::Invalid { setting: setting.to_owned(), @@ -339,17 +335,17 @@ fn coerced( } fn warn_ignored(merged: &Kwargs) { - for key in merged.keys() { + for (key, entry) in merged { if keys::setting(key).is_some_and(|setting| setting.handling == Handling::Warn) { - tracing::warn!(setting = %key, "Redis setting is accepted but ignored"); + tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored: {}", + entry.source, + keys::consequence(key) + ); } } } -fn is_unsupported(key: &str) -> bool { - keys::setting(key).is_some_and(|setting| setting.handling == Handling::Reject) -} - fn environment_kwargs(lookup: &dyn Lookup) -> Kwargs { keys::SETTINGS .iter() diff --git a/litellm-rust/crates/config/src/redis/error.rs b/litellm-rust/crates/config/src/redis/error.rs index b770f1ed33b..67e6dc58431 100644 --- a/litellm-rust/crates/config/src/redis/error.rs +++ b/litellm-rust/crates/config/src/redis/error.rs @@ -1,7 +1,5 @@ #[derive(Debug, thiserror::Error)] pub enum RedisError { - #[error("Redis host does not yet support configured setting {setting}")] - Unsupported { setting: String }, #[error("invalid Redis setting {setting}: expected {expected}")] Invalid { setting: String, expected: Expected }, #[error("either host or url must be specified for redis")] diff --git a/litellm-rust/crates/config/src/redis/keys.rs b/litellm-rust/crates/config/src/redis/keys.rs index 38bbac26458..f5a0acf07e5 100644 --- a/litellm-rust/crates/config/src/redis/keys.rs +++ b/litellm-rust/crates/config/src/redis/keys.rs @@ -1,14 +1,13 @@ use litellm_core_utils::settings::Lookup; // Every Redis setting litellm/_redis.py reads, with how Rust handles it and where else Python looks for it. +// Warn rows are read by Python and dropped here with a warning naming the consequence, never an error. // environment marks the kwargs _get_redis_env_kwarg_mapping reads as REDIS_, and fallback lists the names // _get_redis_client_logic consults with get_secret when the merged kwargs lack the setting #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(super) enum Handling { Read, - Ignore, Warn, - Reject, } pub(super) struct Setting { @@ -76,16 +75,16 @@ pub(super) const SETTINGS: &[Setting] = &[ ), Setting::new("cache", Handling::Warn, true, &[]), Setting::new("cache_config", Handling::Warn, true, &[]), - Setting::new("charset", Handling::Ignore, true, &[]), - Setting::new("client_name", Handling::Ignore, true, &[]), - Setting::new("connection_pool", Handling::Reject, false, &[]), + Setting::new("charset", Handling::Warn, true, &[]), + Setting::new("client_name", Handling::Warn, true, &[]), + Setting::new("connection_pool", Handling::Warn, false, &[]), Setting::new("credential_provider", Handling::Read, false, &[]), Setting::new("db", Handling::Read, true, &[]), - Setting::new("decode_responses", Handling::Ignore, true, &[]), - Setting::new("encoding", Handling::Ignore, true, &[]), - Setting::new("encoding_errors", Handling::Ignore, true, &[]), - Setting::new("errors", Handling::Ignore, true, &[]), - Setting::new("event_dispatcher", Handling::Reject, true, &[]), + Setting::new("decode_responses", Handling::Warn, true, &[]), + Setting::new("encoding", Handling::Warn, true, &[]), + Setting::new("encoding_errors", Handling::Warn, true, &[]), + Setting::new("errors", Handling::Warn, true, &[]), + Setting::new("event_dispatcher", Handling::Warn, true, &[]), Setting::new( "gcp_service_account", Handling::Read, @@ -98,17 +97,17 @@ pub(super) const SETTINGS: &[Setting] = &[ true, &["REDIS_GCP_SSL_CA_CERTS"], ), - Setting::new("health_check_interval", Handling::Warn, true, &[]), + Setting::new("health_check_interval", Handling::Read, true, &[]), Setting::new("host", Handling::Read, true, &[]), - Setting::new("lib_name", Handling::Ignore, true, &[]), - Setting::new("lib_version", Handling::Ignore, true, &[]), + Setting::new("lib_name", Handling::Warn, true, &[]), + Setting::new("lib_version", Handling::Warn, true, &[]), Setting::new("max_connections", Handling::Warn, true, &[]), Setting::new("namespace", Handling::Read, false, &[]), Setting::new("password", Handling::Read, true, &[]), Setting::new("port", Handling::Read, true, &[]), Setting::new("protocol", Handling::Warn, true, &[]), - Setting::new("redis_connect_func", Handling::Reject, true, &[]), - Setting::new("retry", Handling::Reject, false, &[]), + Setting::new("redis_connect_func", Handling::Warn, true, &[]), + Setting::new("retry", Handling::Warn, false, &[]), Setting::new("retry_on_error", Handling::Warn, true, &[]), Setting::new("retry_on_timeout", Handling::Warn, true, &[]), Setting::new( @@ -129,37 +128,61 @@ pub(super) const SETTINGS: &[Setting] = &[ false, &["REDIS_SERVICE_NAME"], ), - Setting::new("single_connection_client", Handling::Ignore, true, &[]), + Setting::new("single_connection_client", Handling::Warn, true, &[]), Setting::new("socket_connect_timeout", Handling::Read, true, &[]), Setting::new("socket_keepalive", Handling::Warn, true, &[]), Setting::new("socket_keepalive_options", Handling::Warn, true, &[]), Setting::new("socket_timeout", Handling::Read, true, &[]), Setting::new("ssl", Handling::Read, true, &[]), Setting::new("ssl_ca_certs", Handling::Read, true, &[]), - Setting::new("ssl_ca_data", Handling::Reject, true, &[]), - Setting::new("ssl_ca_path", Handling::Reject, true, &[]), + Setting::new("ssl_ca_data", Handling::Read, true, &[]), + Setting::new("ssl_ca_path", Handling::Read, true, &[]), Setting::new("ssl_cert_reqs", Handling::Read, true, &[]), Setting::new("ssl_certfile", Handling::Read, true, &[]), Setting::new("ssl_check_hostname", Handling::Read, true, &[]), - Setting::new("ssl_ciphers", Handling::Reject, true, &[]), + Setting::new("ssl_ciphers", Handling::Warn, true, &[]), Setting::new("ssl_keyfile", Handling::Read, true, &[]), - Setting::new("ssl_min_version", Handling::Reject, true, &[]), - Setting::new("ssl_ocsp_context", Handling::Reject, true, &[]), - Setting::new("ssl_ocsp_expected_cert", Handling::Reject, true, &[]), - Setting::new("ssl_password", Handling::Reject, true, &[]), - Setting::new("ssl_validate_ocsp", Handling::Reject, true, &[]), - Setting::new("ssl_validate_ocsp_stapled", Handling::Reject, true, &[]), + Setting::new("ssl_min_version", Handling::Warn, true, &[]), + Setting::new("ssl_ocsp_context", Handling::Warn, true, &[]), + Setting::new("ssl_ocsp_expected_cert", Handling::Warn, true, &[]), + Setting::new("ssl_password", Handling::Warn, true, &[]), + Setting::new("ssl_validate_ocsp", Handling::Warn, true, &[]), + Setting::new("ssl_validate_ocsp_stapled", Handling::Warn, true, &[]), Setting::new( "startup_nodes", Handling::Read, false, &["REDIS_CLUSTER_NODES"], ), - Setting::new("unix_socket_path", Handling::Reject, true, &[]), + Setting::new("unix_socket_path", Handling::Read, true, &[]), Setting::new("url", Handling::Read, true, &[]), Setting::new("username", Handling::Read, true, &[]), ]; +/// What ignoring a warned setting means for the connection, so the operator reads the consequence +/// rather than only the key name. +pub(super) fn consequence(kwarg: &str) -> &'static str { + match kwarg { + "ssl_validate_ocsp" + | "ssl_validate_ocsp_stapled" + | "ssl_ocsp_context" + | "ssl_ocsp_expected_cert" => { + "rustls has no client-side OCSP, so the connection is made without revocation checks" + } + "ssl_ciphers" | "ssl_min_version" => { + "rustls negotiates TLS 1.2+ with its own ECDHE AEAD suites" + } + "ssl_password" => "an encrypted client key cannot be read, so the TLS handshake fails", + "credential_provider" | "redis_connect_func" => { + "the Python object is not callable from Rust, so static username and password are used" + } + "connection_pool" | "retry" | "event_dispatcher" => { + "the Rust client uses its own pool, retries and events" + } + _ => "the Rust client has no equivalent", + } +} + pub(super) fn setting(kwarg: &str) -> Option<&'static Setting> { SETTINGS.iter().find(|setting| setting.kwarg == kwarg) } diff --git a/litellm-rust/crates/config/src/redis/tls.rs b/litellm-rust/crates/config/src/redis/tls.rs index 9339d89d3fb..68a863daa68 100644 --- a/litellm-rust/crates/config/src/redis/tls.rs +++ b/litellm-rust/crates/config/src/redis/tls.rs @@ -4,10 +4,12 @@ use super::{Entry, Expected, Kwargs, RedisError, bool_value, parsed, text_value} use crate::Value; // redis-py's parse_url turns every query parameter into a client kwarg. Only the TLS ones have a redis-rs -// equivalent, so any other parameter, and any parameter on a plain redis:// url, is rejected +// equivalent, so any other parameter, and any parameter on a plain redis:// url, is warned and dropped const URL_TLS_SETTINGS: &[&str] = &[ "ssl_cert_reqs", "ssl_ca_certs", + "ssl_ca_data", + "ssl_ca_path", "ssl_certfile", "ssl_keyfile", "ssl_check_hostname", @@ -35,6 +37,8 @@ pub(super) fn tls_options( .map(Into::into), certfile: text_value(merged, "ssl_certfile").map(Into::into), keyfile: text_value(merged, "ssl_keyfile").map(Into::into), + ca_data: text_value(merged, "ssl_ca_data"), + ca_path: text_value(merged, "ssl_ca_path").map(Into::into), verify_certificate, check_hostname: bool_value(merged, "ssl_check_hostname"), }) @@ -57,14 +61,16 @@ pub(super) fn url_tls(url: &str) -> Result { ) }) .collect(); - if let Some((_, entry)) = query - .iter() - .find(|(name, _)| !tls_scheme || !URL_TLS_SETTINGS.contains(&name.as_str())) - { - return Err(RedisError::Unsupported { - setting: entry.source.clone(), - }); + let (honored, ignored): (Kwargs, Kwargs) = query + .into_iter() + .partition(|(name, _)| tls_scheme && URL_TLS_SETTINGS.contains(&name.as_str())); + for entry in ignored.values() { + tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored", + entry.source + ); } + let query = honored; Ok(TlsOptions { check_hostname: text_value(&query, "ssl_check_hostname") .is_some_and(|text| url_bool(&text)), diff --git a/litellm-rust/crates/config/tests/redis.rs b/litellm-rust/crates/config/tests/redis.rs index 47410696212..5b2e1313095 100644 --- a/litellm-rust/crates/config/tests/redis.rs +++ b/litellm-rust/crates/config/tests/redis.rs @@ -2,6 +2,7 @@ use std::{ collections::HashMap, fmt::Write, sync::{Arc, Mutex}, + time::Duration, }; use litellm_config::{ @@ -136,14 +137,31 @@ fn resolves_redis_through_the_parsed_config( &[], "connection_pool" )] -fn rejected_settings_name_their_source( +fn unsupported_settings_resolve_with_a_warning_naming_their_source( #[case] yaml: &str, #[case] environment: &[(&str, &str)], #[case] expected: &str, ) { - let error = resolve(yaml, environment).unwrap_err(); + let events = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone())); + let resolved = tracing::subscriber::with_default(subscriber, || resolve(yaml, environment)); + let warnings = events.lock().unwrap().join(" "); - assert!(matches!(error, RedisError::Unsupported { setting } if setting == expected)); + assert_eq!( + resolved.unwrap().endpoint, + Endpoint::Host(litellm_redis_types::HostEndpoint { + host: "redis".into(), + port: 6379, + db: 0, + tls: None, + unix_socket_path: None, + }) + ); + assert!( + warnings.contains(&format!("Redis setting {expected} is not supported")), + "{warnings}" + ); + assert!(warnings.contains("is ignored: "), "{warnings}"); } fn static_auth(username: Option<&str>, password: Option<&str>) -> RedisAuth { @@ -258,22 +276,22 @@ fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) { #[case::defaults_verify_the_certificate_but_not_the_hostname( " ssl: true", &[], - Some(TlsOptions { ca_certs: None, certfile: None, keyfile: None, verify_certificate: true, check_hostname: false }) + Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, verify_certificate: true, check_hostname: false }) )] #[case::every_option( " ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_cert_reqs: none\n ssl_check_hostname: true", &[], - Some(TlsOptions { ca_certs: Some("/ca.pem".into()), certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), verify_certificate: false, check_hostname: true }) + Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), verify_certificate: false, check_hostname: true }) )] #[case::options_from_the_environment( " ssl: true", &[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")], - Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), certfile: None, keyfile: None, verify_certificate: true, check_hostname: false }) + Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, verify_certificate: true, check_hostname: false }) )] #[case::gcp_ca_certs_replace_the_configured_ones( " ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem", &[], - Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), certfile: None, keyfile: None, verify_certificate: true, check_hostname: false }) + Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, verify_certificate: true, check_hostname: false }) )] #[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)] fn tls_options_resolve_like_redis_py( @@ -522,19 +540,23 @@ fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() { } #[rstest] -fn rejected_environment_settings_never_print_the_value() { - let error = resolve( - "model_list: []", - &[ - ("REDIS_HOST", "redis"), - ("REDIS_SSL_CIPHERS", "private-option-value"), - ], - ) - .unwrap_err(); - let message = error.to_string(); +fn unsupported_environment_settings_never_print_the_value() { + let events = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone())); + tracing::subscriber::with_default(subscriber, || { + resolve( + "model_list: []", + &[ + ("REDIS_HOST", "redis"), + ("REDIS_SSL_CIPHERS", "private-option-value"), + ], + ) + .unwrap(); + }); + let warnings = events.lock().unwrap().join(" "); - assert!(message.contains("REDIS_SSL_CIPHERS")); - assert!(!message.contains("private-option-value")); + assert!(warnings.contains("REDIS_SSL_CIPHERS"), "{warnings}"); + assert!(!warnings.contains("private-option-value"), "{warnings}"); } #[rstest] @@ -649,6 +671,8 @@ fn tls_view(tls: Option<&TlsOptions>) -> serde_json::Value { tls.map_or(serde_json::Value::Null, |tls| { serde_json::json!({ "ca_certs": tls.ca_certs, + "ca_data": tls.ca_data, + "ca_path": tls.ca_path, "certfile": tls.certfile, "keyfile": tls.keyfile, "verify_certificate": tls.verify_certificate, @@ -668,24 +692,11 @@ fn expose(value: Option<&litellm_secret_value::SecretValue>) -> Option<&str> { value.map(litellm_secret_value::SecretValue::expose) } -fn expected_view(fields: serde_json::Value) -> serde_json::Value { - let serde_json::Value::Object(fields) = fields else { - unreachable!("expected views are objects") - }; - let defaults = json!({ - "auth": {"kind": "static", "username": null, "password": null}, - "tls": null, - "socket_timeout": 5.0, - "socket_connect_timeout": 5.0, - }); - let serde_json::Value::Object(defaults) = defaults else { - unreachable!("the defaults are an object") - }; - serde_json::Value::Object(defaults.into_iter().chain(fields).collect()) -} - fn view(resolved: &ResolvedRedis) -> serde_json::Value { let target = match &resolved.endpoint { + Endpoint::Host(host) if host.unix_socket_path.is_some() => serde_json::json!({ + "mode": "unix", "path": host.unix_socket_path, "tls": null, + }), Endpoint::Host(host) => serde_json::json!({ "mode": "host", "host": host.host, @@ -740,6 +751,10 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { "auth": auth, "socket_timeout": resolved.timeout.map_or(5.0, |timeout| timeout.as_secs_f64()), "socket_connect_timeout": resolved.connect_timeout.or(resolved.timeout).map_or(5.0, |timeout| timeout.as_secs_f64()), + "health_check_interval": resolved.health_check_interval + .or_else(|| matches!(resolved.endpoint, Endpoint::Cluster(_)).then_some(Duration::from_secs(25))) + .filter(|interval| !interval.is_zero()) + .map(|interval| interval.as_secs_f64()), }); let serde_json::Value::Object(rest) = rest else { unreachable!("the view is an object") @@ -747,51 +762,68 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { serde_json::Value::Object(target.into_iter().chain(rest).collect()) } +fn expected_view(fields: serde_json::Value) -> serde_json::Value { + let serde_json::Value::Object(fields) = fields else { + unreachable!("expected views are objects") + }; + let defaults = json!({ + "auth": {"kind": "static", "username": null, "password": null}, + "tls": null, + "socket_timeout": 5.0, + "socket_connect_timeout": 5.0, + "health_check_interval": null, + }); + let serde_json::Value::Object(defaults) = defaults else { + unreachable!("the defaults are an object") + }; + serde_json::Value::Object(defaults.into_iter().chain(fields).collect()) +} + #[rstest] #[case::host_defaults( json!({"host": "redis"}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::empty_cluster_nodes_use_the_environment( json!({"startup_nodes": []}), &[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)], - Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))), + Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), )] #[case::empty_sentinel_nodes_use_the_environment( json!({"sentinel_nodes": [], "service_name": "primary"}), &[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)], - Ok(expected_view(json!({"db": 0, "mode": "sentinel", "nodes": ["sentinel-a:26379"], "sentinel_password": null, "service_name": "primary"}))), + Ok(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))), )] #[case::empty_service_name_uses_the_environment( json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}), &[("REDIS_SERVICE_NAME", "primary")], - Ok(expected_view(json!({"db": 0, "mode": "sentinel", "nodes": ["sentinel-a:26379"], "sentinel_password": null, "service_name": "primary"}))), + Ok(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))), )] #[case::plaintext_ignores_tls_requirements( json!({"host": "redis", "ssl_cert_reqs": "strict"}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::ssl_ca_path_is_unsupported( json!({"host": "redis", "ssl": true, "ssl_ca_path": "/ca"}), &[], - Err("Redis host does not yet support configured setting ssl_ca_path"), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/ca", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), )] #[case::host_with_every_basic_setting( json!({"db": 2, "host": "redis", "password": "pass", "port": 6380, "socket_timeout": 1.5, "ssl": true, "username": "user"}), &[], - Ok(expected_view(json!({"auth": {"kind": "static", "password": "pass", "username": "user"}, "db": 2, "host": "redis", "mode": "host", "port": 6380, "socket_connect_timeout": 1.5, "socket_timeout": 1.5, "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 2, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": "pass"}, "socket_timeout": 1.5, "socket_connect_timeout": 1.5}))), )] #[case::environment_host_port_and_password( json!({}), &[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")], - Ok(expected_view(json!({"auth": {"kind": "static", "password": "environment-secret", "username": null}, "db": 0, "host": "environment", "mode": "host", "port": 6390, "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))), + Ok(expected_view(json!({"mode": "host", "host": "environment", "port": 6390, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": null, "password": "environment-secret"}}))), )] #[case::string_numbers_are_coerced( json!({"db": "3", "host": "redis", "port": "6380", "socket_timeout": "2.5"}), &[], - Ok(expected_view(json!({"db": 3, "host": "redis", "mode": "host", "port": 6380, "socket_connect_timeout": 2.5, "socket_timeout": 2.5}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 3, "socket_timeout": 2.5, "socket_connect_timeout": 2.5}))), )] #[case::padded_and_underscored_numbers_parse_like_python_int_and_float( json!({"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}), @@ -806,7 +838,7 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { #[case::socket_timeout_from_the_environment( json!({"host": "redis"}), &[("REDIS_SOCKET_TIMEOUT", "0.75")], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "socket_connect_timeout": 0.75, "socket_timeout": 0.75}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 0.75, "socket_connect_timeout": 0.75}))), )] #[case::environment_url( json!({}), @@ -821,42 +853,42 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { #[case::an_explicit_host_drops_the_environment_url( json!({"host": "configured"}), &[("REDIS_URL", "redis://environment:6379")], - Ok(expected_view(json!({"db": 0, "host": "configured", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "configured", "port": 6379, "db": 0}))), )] #[case::a_url_drops_host_db_and_password_and_keeps_the_username( json!({"db": 4, "host": "dropped", "password": "dropped", "url": "rediss://cache:6380", "username": "user"}), &[], - Ok(expected_view(json!({"auth": {"kind": "static", "password": null, "username": "user"}, "mode": "url", "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}, "url": "rediss://cache:6380"}))), + Ok(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": null}}))), )] #[case::a_url_ignores_tls_settings_outside_the_url( json!({"ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true, "url": "rediss://cache:6380"}), &[], - Ok(expected_view(json!({"mode": "url", "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}, "url": "rediss://cache:6380"}))), + Ok(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), )] #[case::startup_nodes_beat_a_url( json!({"password": "cluster-secret", "startup_nodes": [{"host": "node-a", "port": 7000}], "url": "redis://single:6379"}), &[], - Ok(expected_view(json!({"auth": {"kind": "static", "password": "cluster-secret", "username": null}, "mode": "cluster", "nodes": ["node-a:7000"]}))), + Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "auth": {"kind": "static", "username": null, "password": "cluster-secret"}, "health_check_interval": 25.0}))), )] #[case::environment_cluster_nodes( json!({}), &[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)], - Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"]}))), + Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"], "health_check_interval": 25.0}))), )] #[case::cluster_tls_options( json!({"ssl": "true", "ssl_ca_certs": "/ca.pem", "ssl_check_hostname": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}), &[], - Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "tls": {"ca_certs": "/ca.pem", "certfile": null, "check_hostname": true, "keyfile": null, "verify_certificate": true}}))), + Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": true}, "health_check_interval": 25.0}))), )] #[case::sentinel_from_settings( json!({"password": "data-secret", "sentinel_nodes": [["sentinel-a", 26379], ["sentinel-b", 26380]], "sentinel_password": "sentinel-secret", "service_name": "primary"}), &[], - Ok(expected_view(json!({"auth": {"kind": "static", "password": "data-secret", "username": null}, "db": 0, "mode": "sentinel", "nodes": ["sentinel-a:26379", "sentinel-b:26380"], "sentinel_password": "sentinel-secret", "service_name": "primary"}))), + Ok(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379", "sentinel-b:26380"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0, "auth": {"kind": "static", "username": null, "password": "data-secret"}}))), )] #[case::sentinel_from_the_environment( json!({}), &[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")], - Ok(expected_view(json!({"db": 0, "mode": "sentinel", "nodes": ["sentinel-a:26379"], "sentinel_password": "sentinel-secret", "service_name": "primary"}))), + Ok(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0}))), )] #[case::a_host_breaks_sentinel( json!({"host": "ignored", "sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), @@ -871,47 +903,47 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { #[case::an_explicit_credential_provider_beats_cloud_auth( json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), &[("REDIS_AZURE_AD_TOKEN", "true")], - Ok(expected_view(json!({"auth": {"kind": "host"}, "db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "host"}}))), )] #[case::gcp_beats_azure_and_aws( json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), &[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")], - Ok(expected_view(json!({"auth": {"kind": "gcp_iam", "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}, "db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "gcp_iam", "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}}))), )] #[case::gcp_ca_certs_replace_the_configured_ones_under_ssl( json!({"gcp_service_account": "account", "gcp_ssl_ca_certs": "/gcp-ca.pem", "host": "redis", "ssl": true, "ssl_ca_certs": "/configured-ca.pem"}), &[], - Ok(expected_view(json!({"auth": {"kind": "gcp_iam", "service_account": "account"}, "db": 0, "host": "redis", "mode": "host", "port": 6379, "tls": {"ca_certs": "/gcp-ca.pem", "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/gcp-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "gcp_iam", "service_account": "account"}}))), )] #[case::azure_beats_aws_and_takes_the_environment_username( json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), &[("REDIS_USERNAME", "environment-user")], - Ok(expected_view(json!({"auth": {"kind": "azure_ad", "username": "environment-user"}, "db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": "environment-user"}}))), )] #[case::azure_without_a_username_sends_only_the_token( json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}), &[], - Ok(expected_view(json!({"auth": {"kind": "azure_ad", "username": null}, "db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))), )] #[case::azure_needs_the_literal_true( json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": "1", "azure_tenant_id": "tenant", "host": "redis"}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::aws_iam_over_tls_with_the_default_region( json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}), &[("AWS_DEFAULT_REGION", "us-east-1")], - Ok(expected_view(json!({"auth": {"cache_name": "cache-1", "kind": "aws_iam", "region": "us-east-1", "serverless": false, "user_name": "iam-user"}, "db": 0, "host": "cache-1", "mode": "host", "port": 6379, "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))), + Ok(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "us-east-1", "serverless": false}}))), )] #[case::aws_region_beats_the_default_region( json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}), &[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")], - Ok(expected_view(json!({"auth": {"cache_name": "cache-1", "kind": "aws_iam", "region": "eu-central-1", "serverless": false, "user_name": "iam-user"}, "db": 0, "host": "cache-1", "mode": "host", "port": 6379, "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))), + Ok(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-central-1", "serverless": false}}))), )] #[case::aws_iam_serverless_over_a_rediss_url_drops_its_userinfo( json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "eu-west-1", "aws_iam_serverless": true, "aws_iam_user_name": "iam-user", "url": "rediss://user:pass@cache-1:6379"}), &[], - Ok(expected_view(json!({"auth": {"cache_name": "cache-1", "kind": "aws_iam", "region": "eu-west-1", "serverless": true, "user_name": "iam-user"}, "mode": "url", "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}, "url": "rediss://cache-1:6379"}))), + Ok(expected_view(json!({"mode": "url", "url": "rediss://cache-1:6379", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-west-1", "serverless": true}}))), )] #[case::aws_iam_without_tls_fails( json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "us-east-1", "aws_iam_user_name": "iam-user", "host": "cache-1"}), @@ -926,22 +958,22 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { #[case::tls_options_from_settings( json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "tls": {"ca_certs": "/ca.pem", "certfile": null, "check_hostname": true, "keyfile": null, "verify_certificate": false}}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))), )] #[case::tls_options_from_the_environment( json!({"host": "redis"}), &[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "tls": {"ca_certs": "/env-ca.pem", "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/env-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), )] #[case::connect_timeout_is_separate_from_the_socket_timeout( json!({"host": "redis", "socket_connect_timeout": 0.25, "socket_timeout": 2.5}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "socket_connect_timeout": 0.25, "socket_timeout": 2.5}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 2.5, "socket_connect_timeout": 0.25}))), )] #[case::environment_references_resolve( json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}), &[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")], - Ok(expected_view(json!({"auth": {"kind": "static", "password": "referenced-secret", "username": null}, "db": 0, "host": "referenced", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "referenced", "port": 6379, "db": 0, "auth": {"kind": "static", "username": null, "password": "referenced-secret"}}))), )] #[case::no_target_fails( json!({"port": 6379}), @@ -951,22 +983,22 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { #[case::client_certificates( json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379, "tls": {"ca_certs": null, "certfile": "/client.pem", "check_hostname": false, "keyfile": "/client.key", "verify_certificate": true}}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": "/client.key", "verify_certificate": true, "check_hostname": false}}))), )] #[case::a_rediss_url_takes_tls_options_from_its_query( json!({"url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}), &[], - Ok(expected_view(json!({"mode": "url", "tls": {"ca_certs": "/ca.pem", "certfile": null, "check_hostname": true, "keyfile": null, "verify_certificate": false}, "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}))), + Ok(expected_view(json!({"mode": "url", "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes", "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))), )] #[case::a_url_query_setting_without_a_redis_rs_equivalent( json!({"url": "redis://cache:6379/0?socket_timeout=2"}), &[], - Err("Redis host does not yet support configured setting url ?socket_timeout"), + Ok(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2"}))), )] #[case::tls_settings_outside_a_url_never_fail_it( json!({"db": "99999999999999999999", "ssl_cert_reqs": "strict", "url": "rediss://cache:6380"}), &[], - Ok(expected_view(json!({"mode": "url", "tls": {"ca_certs": null, "certfile": null, "check_hostname": false, "keyfile": null, "verify_certificate": true}, "url": "rediss://cache:6380"}))), + Ok(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), )] #[case::db_is_not_read_for_a_cluster( json!({"db": 99_999_999_999_i64, "startup_nodes": [{"host": "node-a", "port": 7000}]}), @@ -981,7 +1013,7 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { #[case::cluster_db_false( json!({"db": false, "startup_nodes": [{"host": "node-a", "port": 7000}]}), &[], - Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))), + Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), )] #[case::cluster_db_fraction( json!({"db": 0.5, "startup_nodes": [{"host": "node-a", "port": 7000}]}), @@ -991,7 +1023,7 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { #[case::cluster_db_float_zero( json!({"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}), &[], - Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))), + Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), )] #[case::cluster_db_nonempty_list( json!({"db": [1], "startup_nodes": [{"host": "node-a", "port": 7000}]}), @@ -1001,52 +1033,87 @@ fn view(resolved: &ResolvedRedis) -> serde_json::Value { #[case::cluster_db_empty_list( json!({"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}), &[], - Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))), + Ok(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), )] #[case::ssl_ciphers( json!({"host": "redis", "ssl": true, "ssl_ciphers": "HIGH"}), &[], - Err("Redis host does not yet support configured setting ssl_ciphers"), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), )] #[case::retry_on_timeout( json!({"host": "redis", "retry_on_timeout": true}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::malformed_port( json!({"host": "redis", "port": "abc"}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::malformed_db( json!({"db": "first", "host": "redis"}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::socket_timeout_not_a_number( json!({"host": "redis", "socket_timeout": "slow"}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::malformed_port_from_the_environment( json!({}), &[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::protocol_3( json!({"host": "redis", "protocol": 3}), &[], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::retry_on_error_from_the_environment( json!({"host": "redis"}), &[("REDIS_RETRY_ON_ERROR", "true")], - Ok(expected_view(json!({"db": 0, "host": "redis", "mode": "host", "port": 6379}))), + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), )] #[case::unix_socket_path( json!({"host": "redis", "unix_socket_path": "/tmp/redis.sock"}), &[], - Err("Redis host does not yet support configured setting unix_socket_path"), + Ok(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))), +)] +#[case::unix_socket_path_keeps_tls_settings_out_of_the_target( + json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}), + &[], + Ok(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))), +)] +#[case::ssl_ca_data( + json!({"host": "redis", "ssl": true, "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----"}), + &[], + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::ssl_ca_path_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")], + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/etc/ssl/private-cas", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::health_check_interval_is_whole_seconds( + json!({"health_check_interval": "7", "host": "redis"}), + &[], + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "health_check_interval": 7.0}))), +)] +#[case::health_check_interval_zero_is_off( + json!({"health_check_interval": 0, "host": "redis"}), + &[], + Ok(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::a_cluster_rejects_a_nonzero_db( + json!({"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + Err("db must be 0 in cluster mode"), +)] +#[case::unix_socket_path_without_a_host_fails_like_python( + json!({"unix_socket_path": "/tmp/redis.sock"}), + &[], + Err("either host or url must be specified for redis"), )] fn resolves_like_python_redis( #[case] kwargs: serde_json::Value, diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis.rs index c15dacfc9cc..1992b9faa2d 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis.rs @@ -1,9 +1,8 @@ use std::{sync::Arc, time::Duration}; -use litellm_circuit_breaker::Config as BreakerConfig; use litellm_config::{ - CacheSettings, Config, Expected, RedisError, RedisSettings, ResolvedRedis, SecretSnapshot, - Value, environment_has_target, + CacheSettings, Config, RedisSettings, ResolvedRedis, SecretSnapshot, Value, + environment_has_target, }; use litellm_core_utils::settings::Lookup; use litellm_redis_client::{RedisClient, RedisEvents}; @@ -72,9 +71,11 @@ pub(crate) async fn connect( resources: &Resources, ) -> Result { let plan = plan(config, lookup)?; - let breaker = breaker_config(lookup)?; + let breaker = litellm_redis_connect::breaker_config(lookup)?; let events: Arc = Arc::new(TracingEvents::new(lookup)); - let client = |params: Params, breaker: Option| -> Result { + let client = |params: Params, + breaker: Option| + -> Result { Ok(litellm_redis_connect::connect( params, HostServices { @@ -178,62 +179,6 @@ async fn reachable(client: RedisClient) -> Option { } } -fn parse_seconds(text: &str) -> Option { - text.parse::() - .ok() - .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) -} - -fn breaker_config(lookup: &dyn Lookup) -> Result, Error> { - let enabled = lookup - .get("REDIS_CIRCUIT_BREAKER_ENABLED") - .is_none_or(|value| value.eq_ignore_ascii_case("true")); - if !enabled { - return Ok(None); - } - let defaults = BreakerConfig::default(); - Ok(Some(BreakerConfig { - failure_threshold: environment_value( - lookup, - "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", - Expected::Integer, - |text| text.parse().ok(), - )? - .unwrap_or(defaults.failure_threshold), - recovery_timeout: environment_value( - lookup, - "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", - Expected::Seconds, - parse_seconds, - )? - .unwrap_or(defaults.recovery_timeout), - timeout_min_duration: environment_value( - lookup, - "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", - Expected::Seconds, - parse_seconds, - )? - .unwrap_or(defaults.timeout_min_duration), - })) -} - -fn environment_value( - lookup: &dyn Lookup, - name: &str, - expected: Expected, - parse: impl Fn(&str) -> Option, -) -> Result, RedisError> { - lookup - .get(name) - .map(|value| { - parse(&value).ok_or_else(|| RedisError::Invalid { - setting: name.to_owned(), - expected, - }) - }) - .transpose() -} - #[cfg(test)] mod tests { use litellm_core_utils::settings::Lookup; @@ -264,6 +209,7 @@ mod tests { port, tls: None, db: 0, + unix_socket_path: None, }), auth: RedisAuth::Static { username: None, @@ -272,6 +218,7 @@ mod tests { namespace: None, timeout: None, connect_timeout: None, + health_check_interval: None, } } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e3ef45f7876..a66cbba0717 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -31,6 +31,12 @@ litellm-cache.workspace = true litellm-error.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-client.workspace = true +litellm-redis-connect.workspace = true +litellm-redis-store.workspace = true +litellm-config.workspace = true +tracing.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index fdcb969a423..dc99c2c8f30 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -12,6 +12,7 @@ mod marshal; mod model_registry; mod preflight; mod python_settings; +mod redis; mod routes; mod secrets; mod tokenizer; @@ -36,6 +37,14 @@ mod _native { #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::redis::native_redis_breaker_states; + #[pymodule_export] + use crate::redis::request::NativeRequestRedis; + #[pymodule_export] + use crate::redis::router_cooldowns::{NativeDeclaredCooldowns, NativeRouterCooldowns}; + #[pymodule_export] + use crate::redis::router_usage::NativeRouterUsage; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -113,6 +122,11 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", + "NativeRequestRedis", + "NativeDeclaredCooldowns", + "NativeRouterCooldowns", + "NativeRouterUsage", + "native_redis_breaker_states", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index eacec6859b5..0152fe8a5b9 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -13,6 +13,8 @@ pub(crate) enum PythonSettings { UrlPolicy, #[strum(serialize = "provider_defaults")] ProviderDefaults, + #[strum(serialize = "router_tiers")] + RouterTiers, #[strum(serialize = "secret_manager")] SecretManager, #[strum(serialize = "secret_manager_binding")] diff --git a/litellm-rust/crates/python-bridge/src/redis/client.rs b/litellm-rust/crates/python-bridge/src/redis/client.rs new file mode 100644 index 00000000000..34355613c56 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/client.rs @@ -0,0 +1,145 @@ +use std::{ + collections::BTreeSet, + sync::{Arc, LazyLock, Mutex, PoisonError}, +}; + +use litellm_config::{RedisSettings, ResolvedRedis, SecretSnapshot}; +use litellm_core_utils::settings::ProcessEnvironment; +use litellm_redis_client::{RedisClient, RedisEvents}; +use litellm_redis_connect::{HostServices, SharedEnvironment}; +use pyo3::{ + prelude::*, + types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple}, +}; + +use super::events::{ServiceLoggingEvents, state_label}; + +/// The `redis_kwargs` dict `RedisCache` hands redis-py, read the way `litellm/_redis.py` reads it. +/// A value that is not plain data, such as a `credential_provider` object, reaches the resolver by +/// its type name, so the resolver or the connector names the setting it cannot honor. +pub(crate) struct PythonRedisSettings(RedisSettings); + +impl<'a, 'py> FromPyObject<'a, 'py> for PythonRedisSettings { + type Error = PyErr; + + fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult { + let py = kwargs.py(); + let plain = PyDict::new(py); + for (name, value) in kwargs.cast::()?.iter() { + if value.is_none() { + continue; + } + if is_plain(&value) { + plain.set_item(&name, value)?; + } else { + plain.set_item(&name, value.get_type().name()?)?; + } + } + Ok(Self(litellm_host_python::from_py_argument(plain.as_any())?)) + } +} + +fn is_plain(value: &Bound<'_, PyAny>) -> bool { + value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() +} + +struct Shared { + resolved: ResolvedRedis, + client: RedisClient, +} + +static CLIENTS: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); + +/// The one client for a resolved Redis config, so every family in this process shares its pool +/// and its breaker. Like redis-py, which never connects at construction, a config the resolver or +/// the connector cannot honor leaves the family on memory with a warning rather than failing the +/// `Router`. +pub(crate) fn shared_client( + py: Python<'_>, + kwargs: PythonRedisSettings, + namespace: Option, +) -> PyResult> { + let environment: SharedEnvironment = Arc::new(ProcessEnvironment); + let secrets = match secret_snapshot(py, kwargs.0.secret_names())? { + Ok(secrets) => secrets, + Err(error) => return Ok(memory_only(py, &error)), + }; + let logger = crate::logger::capture(py); + let resolved = match logger.scope(|| kwargs.0.resolve(&secrets, &*environment)) { + Ok(resolved) => ResolvedRedis { + namespace, + ..resolved + }, + Err(error) => return Ok(memory_only(py, &error)), + }; + let mut clients = CLIENTS.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(shared) = clients.iter().find(|shared| shared.resolved == resolved) { + return Ok(Some(shared.client.clone())); + } + let events: Arc = Arc::new(ServiceLoggingEvents::new(py)?); + let breaker = match logger.scope(|| litellm_redis_connect::breaker_config(&*environment)) { + Ok(breaker) => breaker, + Err(error) => return Ok(memory_only(py, &error)), + }; + let host = HostServices { + events: Some(events), + breaker, + http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?, + environment, + }; + let target = resolved.clone(); + let client = match crate::execution::run_sync_value(py, async move { + Ok(litellm_redis_connect::connect(target, host)) + })? { + Ok(client) => client, + Err(error) => return Ok(memory_only(py, &error)), + }; + clients.push(Shared { + resolved, + client: client.clone(), + }); + Ok(Some(client)) +} + +/// The breaker state of every native Redis client this process built, so the Prometheus state +/// gauge counts them next to the Python `RedisCache` breakers. +#[pyfunction] +pub(crate) fn native_redis_breaker_states() -> Vec<&'static str> { + let clients = CLIENTS.lock().unwrap_or_else(PoisonError::into_inner); + clients + .iter() + .filter_map(|shared| shared.client.breaker_state()) + .map(|state| state_label(&state)) + .collect() +} + +/// Every name the resolver may read, fetched once the way Python's `get_secret` fetches it: the +/// configured secret manager first, then the process environment. +fn secret_snapshot( + py: Python<'_>, + names: BTreeSet, +) -> PyResult> { + let source = crate::secrets::source(py)?; + let owned: Vec = names.iter().cloned().collect(); + let lookup = crate::execution::run_sync_value(py, async move { + let borrowed: Vec<&str> = owned.iter().map(String::as_str).collect(); + Ok(source.resolve(&borrowed).await) + })?; + Ok(lookup.map(|lookup| SecretSnapshot::read(names, &*lookup))) +} + +fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option { + let report = litellm_error::Report::new(error).to_string(); + crate::logger::capture(py).scope(|| { + litellm_tracing::warn!( + "Redis is not usable from the Rust router stores, cooldowns and usage stay in memory: {report}" + ) + }); + None +} diff --git a/litellm-rust/crates/python-bridge/src/redis/events.rs b/litellm-rust/crates/python-bridge/src/redis/events.rs new file mode 100644 index 00000000000..8ceeaec4808 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/events.rs @@ -0,0 +1,251 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; +use pyo3::{prelude::*, types::PyModule}; + +tokio::task_local! { + static CALL: RedisCall; +} + +/// The Python caller of one store operation. Its context carries `with_service_target`, and its +/// `call_type` and `caller` are the labels `RedisCache` reports, so dashboards keep their series. +#[derive(Clone)] +pub(crate) struct RedisCall { + pub(crate) context: PythonContext, + pub(crate) call_type: &'static str, + pub(crate) caller: Option, +} + +impl RedisCall { + pub(crate) async fn scope(self, operation: F) -> F::Output { + CALL.scope(self, operation).await + } +} + +/// Reports every native Redis call and breaker transition to `ServiceLogging`. +pub(crate) struct ServiceLoggingEvents { + record_call: Arc>, + record_breaker: Arc>, + context: PythonContext, +} + +impl ServiceLoggingEvents { + pub(crate) fn new(py: Python<'_>) -> PyResult { + const MODULE: &str = "litellm.rust_bridge.host.redis_events"; + Self::from_module(py, &py.import(MODULE)?) + } + + fn from_module(py: Python<'_>, module: &Bound<'_, PyModule>) -> PyResult { + Ok(Self { + record_call: Arc::new(module.getattr("record_redis_call")?.unbind()), + record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()), + context: PythonContext::capture(py)?, + }) + } + + fn report(&self, context: PythonContext, record: Arc>, args: ReportArgs) { + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + runtime.spawn(async move { + let reported = attach_blocking(context, move |py| match args { + ReportArgs::Call(call_type, duration, failure, caller) => record + .bind(py) + .call1((call_type, duration.as_secs_f64(), failure, caller)) + .map(drop), + ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop), + }) + .await + .and_then(|reported| reported); + if let Err(error) = reported { + tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed"); + } + }); + } +} + +enum ReportArgs { + Call(&'static str, Duration, Option<&'static str>, Option), + Breaker(&'static str), +} + +fn operation_call_type(operation: Operation) -> &'static str { + match operation { + Operation::Command(name) => name, + Operation::Script => "eval", + Operation::Pipeline { .. } => "pipeline", + } +} + +fn failure_label(outcome: Outcome) -> Option<&'static str> { + match outcome { + Outcome::Succeeded => None, + Outcome::Failed(FailureKind::Timeout) => Some("timeout"), + Outcome::Failed(FailureKind::Transport) => Some("transport"), + Outcome::Failed(FailureKind::Server) => Some("server"), + Outcome::Failed(FailureKind::Decode) => Some("decode"), + Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"), + } +} + +pub(crate) fn state_label(state: &State) -> &'static str { + match state { + State::Closed { .. } => "closed", + State::Open { .. } => "open", + State::HalfOpen => "half_open", + } +} + +impl RedisEvents for ServiceLoggingEvents { + fn call(&self, event: &CallEvent) { + let caller = CALL.try_with(Clone::clone).ok(); + let (context, call_type, caller) = match caller { + Some(call) => (call.context, call.call_type, call.caller), + None => ( + self.context.clone(), + operation_call_type(event.operation), + None, + ), + }; + self.report( + context, + self.record_call.clone(), + ReportArgs::Call( + call_type, + event.duration, + failure_label(event.outcome), + caller, + ), + ); + } + + fn breaker(&self, transition: &Transition) { + self.report( + self.context.clone(), + self.record_breaker.clone(), + ReportArgs::Breaker(state_label(&transition.to)), + ); + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use litellm_host_python::PythonContext; + use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; + use pyo3::{prelude::*, types::PyModule}; + use rstest::rstest; + + use super::{RedisCall, ServiceLoggingEvents}; + + const RECORDER: &str = " +import contextvars +target = contextvars.ContextVar('target', default=None) +seen = [] +def record_redis_call(call_type, duration, failure, caller): + seen.append((call_type, round(duration, 3), failure, caller, target.get())) +def record_redis_breaker(state): + seen.append(('breaker', state)) +"; + + fn recorder(name: &str) -> (ServiceLoggingEvents, Py) { + Python::initialize(); + Python::attach(|py| { + let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap(); + let module = + PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name) + .unwrap(); + ( + ServiceLoggingEvents::from_module(py, &module).unwrap(), + module.unbind(), + ) + }) + } + + async fn seen(module: &Py, count: usize) -> String { + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + loop { + let seen: (usize, String) = Python::attach(|py| { + let seen = module.bind(py).getattr("seen").unwrap(); + (seen.len().unwrap(), seen.repr().unwrap().to_string()) + }); + if seen.0 >= count || tokio::time::Instant::now() > deadline { + return seen.1; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn a_scoped_call_reports_the_store_labels_in_the_callers_context() { + let (events, namespace) = recorder("scoped"); + let call = Python::attach(|py| { + namespace + .bind(py) + .getattr("target") + .unwrap() + .call_method1("set", ("spend_counter",)) + .unwrap(); + RedisCall { + context: PythonContext::capture(py).unwrap(), + call_type: "async_increment", + caller: Some("router.py:42".to_owned()), + } + }); + let event = CallEvent { + operation: Operation::Command("INCRBYFLOAT"), + duration: Duration::from_millis(250), + outcome: Outcome::Failed(FailureKind::Timeout), + }; + + call.scope(async { events.call(&event) }).await; + + assert_eq!( + seen(&namespace, 1).await, + "[('async_increment', 0.25, 'timeout', 'router.py:42', 'spend_counter')]" + ); + } + + #[rstest] + #[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")] + #[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")] + #[tokio::test(flavor = "multi_thread")] + async fn breaker_transitions_report_python_state_names( + #[case] to: litellm_circuit_breaker::State, + #[case] expected: &str, + ) { + let (events, namespace) = recorder(expected); + + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to, + generation: 1, + }); + + assert_eq!( + seen(&namespace, 1).await, + format!("[('breaker', '{expected}')]") + ); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn an_unscoped_call_falls_back_to_the_command_name() { + let (events, namespace) = recorder("unscoped"); + + events.call(&CallEvent { + operation: Operation::Pipeline { ops: 3 }, + duration: Duration::from_millis(5), + outcome: Outcome::Succeeded, + }); + + assert_eq!( + seen(&namespace, 1).await, + "[('pipeline', 0.005, None, None, None)]" + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/mod.rs b/litellm-rust/crates/python-bridge/src/redis/mod.rs new file mode 100644 index 00000000000..257a37efe31 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/mod.rs @@ -0,0 +1,61 @@ +mod client; +mod events; +pub(crate) use client::native_redis_breaker_states; +mod policy; +pub(crate) mod request; +pub(crate) mod router_cooldowns; +pub(crate) mod router_usage; + +use litellm_host_python::PythonContext; +use litellm_redis_store::Number; +use pyo3::{ + exceptions::{PyRuntimeError, PyTypeError}, + prelude::*, + types::{PyBool, PyFloat, PyInt}, +}; + +use self::events::RedisCall; + +/// The labels `RedisCache` reports for this operation, in the caller's context so its +/// `service_target` reaches the service event. +fn call(py: Python<'_>, call_type: &'static str) -> PyResult { + Ok(RedisCall { + context: PythonContext::capture(py)?, + call_type, + caller: None, + }) +} + +fn memory_error(error: litellm_cache::Error) -> PyErr { + PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string()) +} + +/// A Python `int` or `float` across the bridge, keeping which one it is. +pub(crate) struct PyNumber(Number); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + if value.is_instance_of::() && !value.is_instance_of::() { + return Ok(Self(Number::Int(value.extract()?))); + } + if value.is_instance_of::() { + return Ok(Self(Number::Float(value.extract()?))); + } + Err(PyTypeError::new_err("expected an int or a float")) + } +} + +impl<'py> IntoPyObject<'py> for PyNumber { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Ok(match self.0 { + Number::Int(value) => value.into_pyobject(py)?.into_any(), + Number::Float(value) => value.into_pyobject(py)?.into_any(), + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/policy.rs b/litellm-rust/crates/python-bridge/src/redis/policy.rs new file mode 100644 index 00000000000..c21b8c39f47 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/policy.rs @@ -0,0 +1,130 @@ +use std::time::Duration; + +use litellm_cache::TierPolicy; +use litellm_redis_store::{router_cooldowns, router_usage}; +use pyo3::prelude::*; + +use crate::{ + coercion::{Field, FieldSpec, ProjectionError}, + python_settings::{PythonSettings, Snapshot}, +}; + +const COOLDOWN_REFRESH: FieldSpec = + FieldSpec::new("cooldown_redis_read_interval", seconds); +const USAGE_REFRESH: FieldSpec = FieldSpec::new("usage_redis_read_interval", seconds); +const TRACKED_READS: FieldSpec = FieldSpec::new("max_redis_batch_cache_size", count); + +/// The tier values this Python process runs with. `POLICY` in each family holds `main`'s +/// defaults, the host reports what its environment and `litellm` globals changed. +#[derive(Debug)] +pub(crate) struct RouterTierPolicies { + pub(crate) cooldowns: TierPolicy, + pub(crate) usage: TierPolicy, +} + +pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult { + Ok(project(&PythonSettings::RouterTiers.read(py)?)?) +} + +fn project(snapshot: &Snapshot<'_>) -> Result { + let tracked_reads = snapshot.read(&TRACKED_READS)?; + Ok(RouterTierPolicies { + cooldowns: TierPolicy { + refresh_after: snapshot.read(&COOLDOWN_REFRESH)?, + tracked_reads, + ..router_cooldowns::POLICY + }, + usage: TierPolicy { + refresh_after: snapshot.read(&USAGE_REFRESH)?, + tracked_reads, + ..router_usage::POLICY + }, + }) +} + +fn seconds(field: &Field<'_>) -> Result { + match field.value().extract::() { + Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => { + Ok(Duration::from_secs_f64(seconds)) + } + _ => Err(field.invalid("a non-negative number of seconds")), + } +} + +fn count(field: &Field<'_>) -> Result { + field + .value() + .extract::() + .map_err(|_| field.invalid("a non-negative integer")) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pyo3::{prelude::*, types::PyDict}; + + use super::project; + use crate::python_settings::PythonSettings; + + fn projected(py: Python<'_>, fields: &str) -> Result { + let locals = PyDict::new(py); + let source = std::ffi::CString::new(format!( + "import types\nsettings = types.SimpleNamespace({fields})" + )) + .unwrap(); + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + let snapshot = + PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap()); + project(&snapshot).map_err(|error| PyErr::from(error).to_string()) + } + + #[test] + fn the_host_values_replace_only_the_refresh_and_tracked_reads_of_mains_defaults() { + Python::initialize(); + Python::attach(|py| { + let policies = projected( + py, + "cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42", + ) + .unwrap(); + assert_eq!( + policies.cooldowns.refresh_after, + Duration::from_millis(2500) + ); + assert_eq!(policies.usage.refresh_after, Duration::from_secs(30)); + assert_eq!(policies.cooldowns.tracked_reads, 42); + assert_eq!(policies.usage.tracked_reads, 42); + assert_eq!( + policies.cooldowns.backfill_ttl, + litellm_redis_store::router_cooldowns::POLICY.backfill_ttl + ); + assert_eq!( + policies.usage.backfill_ttl, + litellm_redis_store::router_usage::POLICY.backfill_ttl + ); + }); + } + + #[test] + fn a_value_python_would_not_run_with_is_a_configuration_error() { + Python::initialize(); + Python::attach(|py| { + let negative = projected( + py, + "cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000", + ) + .unwrap_err(); + let text = projected( + py, + "cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many'", + ) + .unwrap_err(); + assert!( + negative.contains("cooldown_redis_read_interval"), + "{negative}" + ); + assert!(text.contains("max_redis_batch_cache_size"), "{text}"); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/request.rs b/litellm-rust/crates/python-bridge/src/redis/request.rs new file mode 100644 index 00000000000..7c6e42a360f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/request.rs @@ -0,0 +1,40 @@ +use std::time::Duration; + +use litellm_redis_store::RequestRedis; +use pyo3::prelude::*; + +use super::call; +use crate::execution::run_async_value; + +/// One request's native Redis batches. Python builds one per request scope and passes it to each +/// request-aware store call. +#[pyclass(frozen, name = "NativeRequestRedis")] +pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis); + +#[pymethods] +impl NativeRequestRedis { + #[new] + fn new(post_call_deadline_seconds: f64) -> Self { + Self(RequestRedis::new(Duration::from_secs_f64( + post_call_deadline_seconds.max(0.0), + ))) + } + + fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "request_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_pre()).await; + Ok(()) + }) + } + + fn flush_post<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "post_call_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_post()).await; + Ok(()) + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs b/litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs new file mode 100644 index 00000000000..25c8be6f8f0 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/router_cooldowns.rs @@ -0,0 +1,244 @@ +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::Duration, +}; + +use litellm_redis_store::router_cooldowns::{Cooldown, DeclaredCooldowns, RouterCooldowns}; +use pyo3::{ + exceptions::PyRuntimeError, + prelude::*, + types::{PyDict, PyTuple}, +}; + +use super::{ + PyNumber, call, + client::{PythonRedisSettings, shared_client}, + memory_error, + policy::router_tier_policies, + request::NativeRequestRedis, +}; +use crate::execution::{run_async_value, run_sync_value}; + +/// One router's cooldown store. The router's sync completion path calls `put`, `active`, +/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py. +#[pyclass(frozen, name = "NativeRouterCooldowns")] +pub(crate) struct NativeRouterCooldowns(Arc); + +pub(crate) struct PyCooldown(Cooldown); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyCooldown { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let fields = value.cast::()?; + let item = |name: &str| fields.as_any().get_item(name); + Ok(Self(Cooldown { + exception_received: item("exception_received")?.extract()?, + status_code: item("status_code")?.extract()?, + timestamp: item("timestamp")?.extract::()?.0, + cooldown_time: item("cooldown_time")?.extract::()?.0, + })) + } +} + +impl<'py> IntoPyObject<'py> for PyCooldown { + type Target = PyDict; + type Output = Bound<'py, PyDict>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + let fields = PyDict::new(py); + fields.set_item("exception_received", self.0.exception_received)?; + fields.set_item("status_code", self.0.status_code)?; + fields.set_item("timestamp", PyNumber(self.0.timestamp))?; + fields.set_item("cooldown_time", PyNumber(self.0.cooldown_time))?; + Ok(fields) + } +} + +fn active_to_py( + py: Python<'_>, + active: Vec<(String, Cooldown)>, +) -> PyResult>> { + active + .into_iter() + .map(|(model_id, cooldown)| { + PyTuple::new( + py, + [ + model_id.into_pyobject(py)?.into_any(), + PyCooldown(cooldown).into_pyobject(py)?.into_any(), + ], + ) + }) + .collect() +} + +#[pymethods] +impl NativeRouterCooldowns { + #[new] + #[pyo3(signature = (redis_kwargs, *, namespace = None))] + fn new( + py: Python<'_>, + redis_kwargs: Option, + namespace: Option, + ) -> PyResult { + let client = redis_kwargs + .map(|kwargs| shared_client(py, kwargs, namespace)) + .transpose()? + .flatten(); + let policy = router_tier_policies(py)?.cooldowns; + Ok(Self(Arc::new(RouterCooldowns::new(client, policy)))) + } + + fn put( + &self, + py: Python<'_>, + model_id: String, + cooldown: PyCooldown, + ttl_seconds: f64, + ) -> PyResult<()> { + let store = self.0.clone(); + let scope = call(py, "set_cache")?; + run_sync_value(py, async move { + scope + .scope(store.put( + &model_id, + &cooldown.0, + Duration::from_secs_f64(ttl_seconds.max(0.0)), + )) + .await + .map_err(memory_error) + }) + } + + fn increment_allowed_fails( + &self, + py: Python<'_>, + model_id: String, + suffix: Option, + ttl_seconds: f64, + ) -> PyResult { + let store = self.0.clone(); + let scope = call(py, "increment_cache")?; + run_sync_value(py, async move { + scope + .scope(store.increment_allowed_fails( + &model_id, + suffix.as_deref(), + Duration::from_secs_f64(ttl_seconds.max(0.0)), + )) + .await + .map_err(memory_error) + }) + } + + fn active<'py>( + &self, + py: Python<'py>, + model_ids: Vec, + ) -> PyResult>> { + let store = self.0.clone(); + let scope = call(py, "batch_get_cache")?; + let active = run_sync_value(py, async move { + scope + .scope(store.active(&model_ids)) + .await + .map_err(memory_error) + })?; + active_to_py(py, active) + } + + fn aactive<'py>(&self, py: Python<'py>, model_ids: Vec) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_batch_get_cache")?; + run_async_value(py, async move { + let active = scope + .scope(store.active(&model_ids)) + .await + .map_err(memory_error)?; + Python::attach(|py| { + Ok(active_to_py(py, active)? + .into_iter() + .map(Bound::unbind) + .collect::>()) + }) + }) + } + + fn declare_active( + &self, + request: &NativeRequestRedis, + model_ids: Vec, + ) -> PyResult { + let declared = self + .0 + .declare_active(&request.0, &model_ids) + .map_err(memory_error)?; + Ok(NativeDeclaredCooldowns { + store: self.0.clone(), + declared: Mutex::new(Some(declared)), + }) + } + + fn minimum(&self, py: Python<'_>, model_ids: Vec) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "batch_get_cache")?; + let minimum = run_sync_value(py, async move { + scope + .scope(store.minimum(&model_ids)) + .await + .map_err(memory_error) + })?; + Ok(minimum.map(PyNumber)) + } + + fn evict_local(&self, model_id: &str) -> PyResult<()> { + self.0.evict_local(model_id).map_err(memory_error) + } + + fn flush_local_counters(&self) -> PyResult<()> { + self.0.flush_local_counters().map_err(memory_error) + } + + fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { + let snapshot = PyDict::new(py); + for (key, cooldown) in self.0.snapshot().map_err(memory_error)? { + snapshot.set_item(key, PyCooldown(cooldown))?; + } + Ok(snapshot) + } +} + +/// An `active` read declared on a request's pre-call batch. `resolve` answers it once. +#[pyclass(frozen, name = "NativeDeclaredCooldowns")] +pub(crate) struct NativeDeclaredCooldowns { + store: Arc, + declared: Mutex>, +} + +#[pymethods] +impl NativeDeclaredCooldowns { + fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { + let declared = self + .declared + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .ok_or_else(|| PyRuntimeError::new_err("this cooldown read was already resolved"))?; + let store = self.store.clone(); + let scope = call(py, "async_batch_get_cache")?; + run_async_value(py, async move { + let active = scope + .scope(store.resolve_active(declared)) + .await + .map_err(memory_error)?; + Python::attach(|py| { + Ok(active_to_py(py, active)? + .into_iter() + .map(Bound::unbind) + .collect::>()) + }) + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/router_usage.rs new file mode 100644 index 00000000000..c1901efd906 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/router_usage.rs @@ -0,0 +1,160 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage}; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; + +use super::{ + PyNumber, call, + client::{PythonRedisSettings, shared_client}, + memory_error, + policy::router_tier_policies, + request::NativeRequestRedis, +}; +use crate::execution::run_async_value; + +/// One router's global TPM/RPM counters. Every caller on `main` is async. +#[pyclass(frozen, name = "NativeRouterUsage")] +pub(crate) struct NativeRouterUsage(Arc); + +/// `(deployment_id, model, minute, "tpm" | "rpm")`. +pub(crate) struct PyCounter(Counter); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let (deployment_id, model, minute, kind): (String, Option, String, String) = + value.extract()?; + let kind = match kind.as_str() { + "tpm" => Kind::Tpm, + "rpm" => Kind::Rpm, + _ => { + return Err(PyValueError::new_err(format!( + "unknown usage counter kind {kind:?}" + ))); + } + }; + Ok(Self(Counter { + deployment_id, + model, + minute, + kind, + })) + } +} + +#[pymethods] +impl NativeRouterUsage { + #[new] + #[pyo3(signature = (redis_kwargs, *, namespace = None))] + fn new( + py: Python<'_>, + redis_kwargs: Option, + namespace: Option, + ) -> PyResult { + let client = redis_kwargs + .map(|kwargs| shared_client(py, kwargs, namespace)) + .transpose()? + .flatten(); + let policy = router_tier_policies(py)?.usage; + Ok(Self(Arc::new(RouterUsage::new(client, policy)))) + } + + fn increment<'py>( + &self, + py: Python<'py>, + counter: PyCounter, + value: PyNumber, + ttl_seconds: u64, + ) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_increment")?; + run_async_value(py, async move { + scope + .scope(store.increment(&counter.0, value.0, Duration::from_secs(ttl_seconds))) + .await + .map(PyNumber) + .map_err(memory_error) + }) + } + + #[pyo3(signature = (deployment_id, model, minute, tokens, requests, *, request = None))] + #[expect( + clippy::too_many_arguments, + reason = "preserves the Python increment_deployment signature" + )] + fn increment_deployment<'py>( + &self, + py: Python<'py>, + deployment_id: String, + model: Option, + minute: String, + tokens: PyNumber, + requests: i64, + request: Option>, + ) -> PyResult> { + let store = self.0.clone(); + let request = request.map(|request| request.0.clone()); + let scope = call(py, "async_increment_pipeline")?; + run_async_value(py, async move { + let model = model.as_deref(); + match request { + Some(request) => store.increment_deployment_post_call( + &request, + &deployment_id, + model, + &minute, + tokens.0, + requests, + ), + None => { + scope + .scope(store.increment_deployment( + &deployment_id, + model, + &minute, + tokens.0, + requests, + )) + .await + } + } + .map_err(memory_error) + }) + } + + fn get_many<'py>( + &self, + py: Python<'py>, + counters: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let counters = counters + .into_iter() + .map(|counter| counter.0) + .collect::>(); + let scope = call(py, "async_batch_get_cache")?; + run_async_value(py, async move { + let read = scope + .scope(store.get_many(&counters)) + .await + .map_err(memory_error)?; + Ok(read + .into_iter() + .map(|count| count.map(PyNumber)) + .collect::>()) + }) + } + + fn flush_local_counters(&self) -> PyResult<()> { + self.0.flush_local_counters().map_err(memory_error) + } + + fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { + let snapshot = PyDict::new(py); + for (key, count) in self.0.snapshot().map_err(memory_error)? { + snapshot.set_item(key, PyNumber(count))?; + } + Ok(snapshot) + } +} diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index 78c9593aa17..c42ec251872 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -42,6 +42,7 @@ - `RedisBatch` declares ops now and sends them later as one `Pipeline`, so it inherits every pipeline rule above. It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local - Each op returns a `Declared` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once - An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis +- `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client - Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes - Tests that only need the public API live in `tests/.rs`, with client builders in `tests/support`. Inline `#[cfg(test)]` modules are only for private items (`error.rs`, `namespace.rs`) - Observe connections from the server side (`RedisDeployment::clients`, `cli_on_every_node`), never by reaching into `RedisClient`'s connection diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index eb0ba02576c..824685b894f 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +tracing.workspace = true litellm-circuit-breaker.workspace = true litellm-secret-value.workspace = true litellm-redis-types.workspace = true diff --git a/litellm-rust/crates/redis-client/src/batch.rs b/litellm-rust/crates/redis-client/src/batch.rs index 37955e1350c..f8d9b9fc120 100644 --- a/litellm-rust/crates/redis-client/src/batch.rs +++ b/litellm-rust/crates/redis-client/src/batch.rs @@ -135,6 +135,10 @@ impl Declared { Err(TryRecvError::Empty) => {} } self.batch.flush().await; + self.settled().await + } + + pub async fn settled(self) -> Result { let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned())); reply.and_then(self.decode) } diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 0e3b37a8c67..cf11f269095 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -1,10 +1,10 @@ use std::{ future::Future, - sync::Arc, + sync::{Arc, Mutex, PoisonError}, time::{Duration, Instant}, }; -use litellm_circuit_breaker::{Admission, Breaker, Failure}; +use litellm_circuit_breaker::{Admission, Breaker, Failure, State}; use redis::{Cmd, FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Value}; use crate::{ @@ -20,6 +20,8 @@ pub struct RedisClient { namespace: Namespace, breaker: Option, events: Option>, + health_check: Option, + last_used: Arc>, } #[derive(Debug, PartialEq, Eq)] @@ -36,9 +38,16 @@ impl RedisClient { namespace: config.namespace, breaker: config.breaker, events: config.events, + health_check: config.health_check, + last_used: Arc::new(Mutex::new(Instant::now())), }) } + /// The breaker's current state, or none when the breaker is disabled. + pub fn breaker_state(&self) -> Option { + self.breaker.as_ref().map(Breaker::state) + } + pub async fn ping(&self) -> Result<(), Error> { self.query("PING", &commands::ping()).await } @@ -141,6 +150,10 @@ impl RedisClient { Pipeline::new(self.namespace.clone()) } + pub fn same_client(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.last_used, &other.last_used) + } + pub fn batch(&self) -> RedisBatch { RedisBatch::new(self.clone()) } @@ -163,6 +176,7 @@ impl RedisClient { .collect(); } }; + self.check_health().await; let results = pipeline.run(&self.connection).await; let transport = results.iter().find_map(|result| { result @@ -220,6 +234,7 @@ impl RedisClient { let started = Instant::now(); let result = match self.admit() { Ok(admission) => { + self.check_health().await; let result = call .await .map_err(|error| Error::classify(error, self.connection.deployment())); @@ -233,6 +248,10 @@ impl RedisClient { } fn observe(&self, operation: Operation, started: Instant, outcome: Outcome) { + *self + .last_used + .lock() + .unwrap_or_else(PoisonError::into_inner) = Instant::now(); if let Some(events) = &self.events { events.call(&CallEvent { operation, @@ -242,6 +261,26 @@ impl RedisClient { } } + /// redis-py pings a connection that sat idle longer than `health_check_interval` before + /// sending the next command on it, and reconnects when the ping fails. The ping is not a + /// call of the caller's and never reaches the breaker or the events. + async fn check_health(&self) { + let Some(interval) = self.health_check else { + return; + }; + let idle = self + .last_used + .lock() + .unwrap_or_else(PoisonError::into_inner) + .elapsed(); + if idle < interval { + return; + } + if let Err(error) = self.connection.ping_every_node().await { + tracing::debug!(%error, "idle Redis connection failed its health check, the next command reconnects"); + } + } + fn admit(&self) -> Result, Error> { self.breaker .as_ref() diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index 01af9d38879..d5952bee9ba 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -24,6 +24,10 @@ pub struct RedisConfig { pub timeouts: Timeouts, pub breaker: Option, pub events: Option>, + /// Ping every node before the next command once the client sat idle this long, like + /// redis-py's `health_check_interval`, so a connection the server dropped silently is + /// replaced before a real command pays for it. + pub health_check: Option, } #[derive(Clone, Copy, Debug, PartialEq)] @@ -95,17 +99,44 @@ pub(crate) fn tls_certificates(tls: &TlsOptions) -> Result None, }; + let from_file: Option> = tls.ca_certs.as_ref().map(read).transpose()?; + let from_directory: Vec> = tls + .ca_path + .as_ref() + .map(read_directory) + .transpose()? + .unwrap_or_default(); + let configured: Vec> = from_file + .into_iter() + .chain(tls.ca_data.iter().map(|pem| pem.clone().into_bytes())) + .chain(from_directory) + .collect(); Ok(TlsCertificates { client_tls, - root_cert: tls - .ca_certs - .as_ref() - .map(read) - .transpose()? - .map(with_system_roots), + root_cert: (!configured.is_empty()).then(|| with_system_roots(configured.join(&b'\n'))), }) } +/// Every regular file in `path`, in name order, as redis-py hands `ssl_ca_path` to +/// `load_verify_locations(capath=...)` +fn read_directory(path: &PathBuf) -> Result>, ConnectError> { + let tls_file = |source| ConnectError::TlsFile { + path: path.display().to_string(), + source, + }; + let mut files: Vec = std::fs::read_dir(path) + .map_err(tls_file)? + .map(|entry| entry.map(|entry| entry.path())) + .collect::>() + .map_err(tls_file)?; + files.sort(); + files + .iter() + .filter(|file| file.is_file()) + .map(read) + .collect() +} + pub(crate) fn apply_tls( tls: &TlsOptions, info: ConnectionInfo, @@ -179,9 +210,16 @@ pub(crate) fn connection_info( let (info, tls) = match endpoint { Endpoint::Url { url, tls } => (url.expose().into_connection_info()?, Some(tls)), Endpoint::Host(host) => { - let info = addr(&host.host, host.port, host.tls.as_ref()).into_connection_info()?; + let (addr, tls) = match &host.unix_socket_path { + Some(path) => (ConnectionAddr::Unix(path.clone()), None), + None => ( + addr(&host.host, host.port, host.tls.as_ref()), + host.tls.as_ref(), + ), + }; + let info = addr.into_connection_info()?; let settings = info.redis_settings().clone().set_db(host.db); - (info.set_redis_settings(settings), host.tls.as_ref()) + (info.set_redis_settings(settings), tls) } Endpoint::Cluster(_) | Endpoint::Sentinel(_) => { unreachable!("cluster and sentinel endpoints build their own connection info") @@ -210,3 +248,26 @@ fn with_static_credentials(info: ConnectionInfo, credentials: &Credentials) -> C }; info.set_redis_settings(settings) } + +#[cfg(test)] +mod tests { + use litellm_redis_types::HostEndpoint; + + use super::*; + + #[test] + fn a_unix_socket_path_replaces_host_port_and_tls_like_redis_py() { + let endpoint = Endpoint::Host(HostEndpoint { + host: "ignored".to_owned(), + port: 1, + tls: Some(TlsOptions::verified()), + db: 3, + unix_socket_path: Some("/tmp/redis.sock".into()), + }); + + let info = connection_info(&endpoint, &Credentials::NONE).unwrap(); + + assert_eq!(*info.addr(), ConnectionAddr::Unix("/tmp/redis.sock".into())); + assert_eq!(info.redis_settings().db(), 3); + } +} diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index 4563868e47e..5f343de4b58 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -12,6 +12,7 @@ use redis::{ aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig, MultiplexedConnection}, cluster::{ClusterClientBuilder, ClusterConfig}, cluster_async::ClusterConnection, + cluster_routing::{MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo}, io::tcp::TcpSettings, sentinel::{SentinelClient, SentinelClientBuilder, SentinelServerType}, }; @@ -131,6 +132,30 @@ impl Connection { } } +impl Connection { + /// `PING` on the one multiplexed connection, or on every cluster node, so a dead socket is + /// found and replaced by redis-rs before a real command uses it. + pub(crate) async fn ping_every_node(&self) -> RedisResult<()> { + let ping = redis::cmd("PING"); + match self { + Self::Cluster(connection) => connection + .clone() + .route_command( + ping, + RoutingInfo::MultiNode(( + MultipleNodeRoutingInfo::AllNodes, + Some(ResponsePolicy::AllSucceeded), + )), + ) + .await + .map(drop), + Self::Standalone(_) | Self::Sentinel(_) => { + ping.query_async::(&mut self.clone()).await.map(drop) + } + } + } +} + impl ConnectionLike for Connection { fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { match self { @@ -277,6 +302,7 @@ impl SentinelConnection { port, tls: self.0.endpoint.tls.clone(), db: self.0.endpoint.db, + unix_socket_path: None, }); let info = config::connection_info(&endpoint, &self.0.credentials).map_err( diff --git a/litellm-rust/crates/redis-client/tests/batch.rs b/litellm-rust/crates/redis-client/tests/batch.rs index 4bf5f6fc911..b1095fc3797 100644 --- a/litellm-rust/crates/redis-client/tests/batch.rs +++ b/litellm-rust/crates/redis-client/tests/batch.rs @@ -40,6 +40,7 @@ fn observed(server: &RedisDeployment) -> (RedisClient, Arc) { timeouts: Timeouts::default(), breaker: None, events: Some(trips.clone()), + health_check: None, }) .unwrap(); (client, trips) @@ -145,3 +146,50 @@ async fn an_op_whose_flush_was_cancelled_resolves_as_transport( assert!(cancelled.is_err()); assert!(matches!(counter.resolve().await, Err(Error::Transport(_)))); } + +#[rstest] +#[tokio::test] +async fn a_settled_op_waits_for_someone_else_to_flush_and_sends_nothing_itself( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let counter = tokio::spawn(batch.incr_by("counter", 4, KEEP).settled()); + + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(!counter.is_finished()); + assert!(trips.seen().is_empty()); + + batch.flush().await; + + assert_eq!(counter.await.unwrap().unwrap(), 4); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 1 }]); +} + +#[rstest] +#[tokio::test] +async fn a_settled_op_whose_flush_was_cancelled_resolves_as_transport() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(5)); + let batch = client.batch(); + let counter = batch.incr_by("counter", 1, KEEP); + + server.freeze().await.unwrap(); + let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; + server.thaw().await.unwrap(); + + assert!(cancelled.is_err()); + assert!(matches!(counter.settled().await, Err(Error::Transport(_)))); +} + +#[rstest] +#[tokio::test] +async fn clones_are_the_same_client_and_a_second_connection_is_not() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let twin = connect(&server, "", DEFAULT_TIMEOUT); + + assert!(client.same_client(&client.clone())); + assert!(!client.same_client(&twin)); +} diff --git a/litellm-rust/crates/redis-client/tests/client.rs b/litellm-rust/crates/redis-client/tests/client.rs index ea4e59b6922..231cb70bf05 100644 --- a/litellm-rust/crates/redis-client/tests/client.rs +++ b/litellm-rust/crates/redis-client/tests/client.rs @@ -5,7 +5,7 @@ use std::{ time::{Duration, Instant}, }; -use litellm_redis_testing::{RedisDeployment, RedisServer, Topology, eventually}; +use litellm_redis_testing::{RedisCluster, RedisDeployment, RedisServer, Topology, eventually}; use rstest::rstest; use litellm_circuit_breaker::{Breaker, Config, State}; @@ -254,6 +254,7 @@ async fn outage_events_report_the_failure_kind( breaker: None, events: Some(recorded.clone()), credentials: Credentials::NONE, + health_check: None, }) .unwrap(); client.ping().await.unwrap(); @@ -310,6 +311,7 @@ async fn every_call_reports_its_operation_and_outcome( timeouts: Timeouts::default(), breaker: None, events: Some(recorded.clone()), + health_check: None, }) .unwrap(); client.rpush("list", "item", KEEP).await.unwrap(); @@ -739,6 +741,7 @@ async fn a_host_endpoint_connects_to_its_database() { timeouts: Timeouts::default(), breaker: None, events: None, + health_check: None, }) }; let url = |url: String| Endpoint::Url { @@ -750,6 +753,7 @@ async fn a_host_endpoint_connects_to_its_database() { port, tls: None, db: 3, + unix_socket_path: None, })) .unwrap(); let in_db_0 = connect(url(server.url().to_owned())).unwrap(); @@ -766,3 +770,94 @@ async fn a_host_endpoint_connects_to_its_database() { Some("3") ); } + +async fn ping_calls(cluster: &RedisCluster) -> Vec { + cluster + .cli_outputs(&["INFO", "commandstats"]) + .await + .unwrap() + .iter() + .map(|info| { + info.lines() + .find_map(|line| line.strip_prefix("cmdstat_ping:calls=")) + .and_then(|rest| rest.split(',').next()) + .and_then(|calls| calls.trim().parse().ok()) + .unwrap_or(0) + }) + .collect() +} + +fn health_checked( + server: &RedisDeployment, + interval: Duration, + timeout: Duration, + breaker: Option, + events: Option>, +) -> RedisClient { + RedisClient::new(RedisConfig { + endpoint: endpoint(server), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(timeout, None), + breaker, + events: events.map(|events| events as Arc), + health_check: Some(interval), + }) + .unwrap() +} + +// redis-py 5.3.1 redis/asyncio/connection.py check_health: a connection idle longer than +// health_check_interval gets a PING before the next command, a busy one does not +#[rstest] +#[tokio::test] +async fn an_idle_client_pings_every_node_before_the_next_command_like_redis_py() { + let server = RedisDeployment::start(Topology::Cluster).await.unwrap(); + let RedisDeployment::Cluster(cluster) = &server else { + unreachable!("started a cluster") + }; + let client = health_checked(&server, Duration::from_secs(3), DEFAULT_TIMEOUT, None, None); + client.set("k", "1", Duration::from_secs(60)).await.unwrap(); + let before = ping_calls(cluster).await; + + client.set("k", "2", Duration::from_secs(60)).await.unwrap(); + let busy = ping_calls(cluster).await; + tokio::time::sleep(Duration::from_millis(3500)).await; + client.set("k", "3", Duration::from_secs(60)).await.unwrap(); + let idle = ping_calls(cluster).await; + + assert_eq!(busy, before); + assert_eq!( + idle, + before.iter().map(|calls| calls + 1).collect::>(), + "before {before:?}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_failed_health_check_is_neither_a_call_nor_a_breaker_failure() { + let server = RedisDeployment::start(Topology::Cluster).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let breaker = breaker(2); + let client = health_checked( + &server, + Duration::from_millis(200), + Duration::from_millis(300), + Some(breaker.clone()), + Some(recorded.clone()), + ); + client.set("k", "1", Duration::from_secs(60)).await.unwrap(); + + server.freeze().await.unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + let failed = client.get::("k").await; + server.thaw().await.unwrap(); + + assert!(matches!(failed, Err(Error::Transport(_))), "{failed:?}"); + assert!( + matches!(breaker.state(), State::Closed { timeouts: 1, .. }), + "{:?}", + breaker.state() + ); + assert_eq!(recorded.0.lock().unwrap().len(), 2); +} diff --git a/litellm-rust/crates/redis-client/tests/config.rs b/litellm-rust/crates/redis-client/tests/config.rs index eb118af7195..4f9824568d5 100644 --- a/litellm-rust/crates/redis-client/tests/config.rs +++ b/litellm-rust/crates/redis-client/tests/config.rs @@ -21,9 +21,18 @@ fn debug_output_never_shows_the_secret(#[case] debug: String) { #[derive(Clone, Copy, Debug)] enum Ca { Trusted, + TrustedData, + TrustedDirectory, SystemRoots, } +fn directory_holding(ca: &std::path::Path) -> std::path::PathBuf { + let directory = std::env::temp_dir().join(format!("litellm-ca-dir-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + std::fs::copy(ca, directory.join("ca.pem")).unwrap(); + directory +} + async fn ping_over_tls( require_client_certificate: bool, host: &str, @@ -41,6 +50,8 @@ async fn ping_over_tls( let files = server.files(); let tls = TlsOptions { ca_certs: matches!(ca, Ca::Trusted).then(|| files.ca.clone()), + ca_data: matches!(ca, Ca::TrustedData).then(|| std::fs::read_to_string(&files.ca).unwrap()), + ca_path: matches!(ca, Ca::TrustedDirectory).then(|| directory_holding(&files.ca)), certfile: client_certificate.then(|| files.client_cert.clone()), keyfile: client_certificate.then(|| files.client_key.clone()), verify_certificate, @@ -52,12 +63,14 @@ async fn ping_over_tls( port: server.port(), tls: Some(tls), db: 0, + unix_socket_path: None, }), credentials: Credentials::NONE, namespace: litellm_redis_client::Namespace::default(), timeouts: litellm_redis_client::Timeouts::default(), breaker: None, events: None, + health_check: None, }) .unwrap(); client.ping().await @@ -71,6 +84,14 @@ async fn ping_over_tls( false, true )] +#[case::trusted_ca_as_pem_text_like_ssl_ca_data("127.0.0.1", Ca::TrustedData, true, false, true)] +#[case::trusted_ca_in_a_directory_like_ssl_ca_path( + "127.0.0.1", + Ca::TrustedDirectory, + true, + false, + true +)] #[case::hostname_check_with_a_matching_name("localhost", Ca::Trusted, true, true, true)] #[case::hostname_check_against_an_unlisted_address("127.0.0.1", Ca::Trusted, true, true, false)] #[case::an_untrusted_ca_is_refused("127.0.0.1", Ca::SystemRoots, true, false, false)] @@ -128,12 +149,14 @@ fn a_missing_tls_file_names_the_path() { port: 6379, tls: Some(tls), db: 0, + unix_socket_path: None, }), credentials: Credentials::NONE, namespace: litellm_redis_client::Namespace::default(), timeouts: litellm_redis_client::Timeouts::default(), breaker: None, events: None, + health_check: None, }) .err() .unwrap(); diff --git a/litellm-rust/crates/redis-client/tests/connection.rs b/litellm-rust/crates/redis-client/tests/connection.rs index 462dd630134..a84d06d09bb 100644 --- a/litellm-rust/crates/redis-client/tests/connection.rs +++ b/litellm-rust/crates/redis-client/tests/connection.rs @@ -91,12 +91,14 @@ async fn an_out_of_range_database_is_a_server_rejection() { port, tls: None, db: 999, + unix_socket_path: None, }), credentials: Credentials::NONE, namespace: Namespace::default(), timeouts: Timeouts::new(Duration::from_secs(2), None), breaker: None, events: None, + health_check: None, }) .unwrap(); let result = client.ping().await; @@ -145,6 +147,7 @@ fn sentinel_client(sentinel: &RedisSentinel) -> RedisClient { timeouts: Timeouts::new(Duration::from_secs(1), None), breaker: None, events: None, + health_check: None, }) .unwrap() } diff --git a/litellm-rust/crates/redis-client/tests/support/mod.rs b/litellm-rust/crates/redis-client/tests/support/mod.rs index dc7d95e4d26..12af6912d51 100644 --- a/litellm-rust/crates/redis-client/tests/support/mod.rs +++ b/litellm-rust/crates/redis-client/tests/support/mod.rs @@ -50,6 +50,7 @@ pub fn connect_guarded( timeouts: Timeouts::new(timeout, None), breaker, events: None, + health_check: None, }) .unwrap() } @@ -62,6 +63,7 @@ pub fn connect_as(deployment: &RedisDeployment, credentials: Credentials) -> Red timeouts: Timeouts::new(Duration::from_secs(2), None), breaker: None, events: None, + health_check: None, }) .unwrap() } diff --git a/litellm-rust/crates/redis-connect/Cargo.toml b/litellm-rust/crates/redis-connect/Cargo.toml index 7484182e440..3a943401ecb 100644 --- a/litellm-rust/crates/redis-connect/Cargo.toml +++ b/litellm-rust/crates/redis-connect/Cargo.toml @@ -17,6 +17,7 @@ litellm-circuit-breaker.workspace = true litellm-config.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true +rustls.workspace = true litellm-redis-client.workspace = true futures-util.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs index f4b9f4eb06a..b6d230e2507 100644 --- a/litellm-rust/crates/redis-connect/src/connect.rs +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -1,11 +1,13 @@ +use std::time::Duration; + use std::sync::Arc; use litellm_circuit_breaker::{Breaker, Config as BreakerConfig}; -use litellm_config::{RedisAuth, ResolvedRedis}; +use litellm_config::{Expected, RedisAuth, RedisError, ResolvedRedis}; use litellm_core_utils::settings::Lookup; use litellm_redis_client::{ - Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Namespace, RedisClient, RedisConfig, - RedisEvents, Timeouts, + Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Endpoint, Namespace, RedisClient, + RedisConfig, RedisEvents, Timeouts, }; use crate::{ @@ -22,7 +24,13 @@ pub struct HostServices { pub environment: SharedEnvironment, } +// litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only +const CLUSTER_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(25); + pub fn connect(resolved: ResolvedRedis, host: HostServices) -> Result { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } tracing::info!(endpoint = ?resolved.endpoint, namespace = ?resolved.namespace, "redis client configured"); let credentials = credentials(resolved.auth, &host)?; let events = host.events; @@ -37,6 +45,13 @@ pub fn connect(resolved: ResolvedRedis, host: HostServices) -> Result resolved + .health_check_interval + .or(Some(CLUSTER_HEALTH_CHECK_INTERVAL)), + _ => resolved.health_check_interval, + } + .filter(|interval| !interval.is_zero()); Ok(RedisClient::new(RedisConfig { endpoint: resolved.endpoint, credentials, @@ -47,6 +62,7 @@ pub fn connect(resolved: ResolvedRedis, host: HostServices) -> Result Result Result, RedisError> { + let enabled = lookup + .get("REDIS_CIRCUIT_BREAKER_ENABLED") + .is_none_or(|value| value.eq_ignore_ascii_case("true")); + if !enabled { + return Ok(None); + } + let defaults = BreakerConfig::default(); + Ok(Some(BreakerConfig { + failure_threshold: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", + Expected::Integer, + |text| text.parse().ok(), + )? + .unwrap_or(defaults.failure_threshold), + recovery_timeout: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", + Expected::Seconds, + parse_seconds, + )? + .unwrap_or(defaults.recovery_timeout), + timeout_min_duration: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", + Expected::Seconds, + parse_seconds, + )? + .unwrap_or(defaults.timeout_min_duration), + })) +} + +fn parse_seconds(text: &str) -> Option { + text.parse::() + .ok() + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) +} + +fn environment_value( + lookup: &dyn Lookup, + name: &str, + expected: Expected, + parse: impl Fn(&str) -> Option, +) -> Result, RedisError> { + lookup + .get(name) + .map(|value| { + parse(&value).ok_or_else(|| RedisError::Invalid { + setting: name.to_owned(), + expected, + }) + }) + .transpose() +} diff --git a/litellm-rust/crates/redis-connect/src/lib.rs b/litellm-rust/crates/redis-connect/src/lib.rs index 08d7af91fb7..4e82af2db3b 100644 --- a/litellm-rust/crates/redis-connect/src/lib.rs +++ b/litellm-rust/crates/redis-connect/src/lib.rs @@ -4,7 +4,7 @@ mod events; mod refresh; mod sources; -pub use connect::{HostServices, SharedEnvironment, connect}; +pub use connect::{HostServices, SharedEnvironment, breaker_config, connect}; pub use error::{CredentialStage, Error}; pub use events::TracingEvents; pub(crate) use refresh::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture}; diff --git a/litellm-rust/crates/redis-connect/src/refresh.rs b/litellm-rust/crates/redis-connect/src/refresh.rs index 1c4e733d45f..852ba412b52 100644 --- a/litellm-rust/crates/redis-connect/src/refresh.rs +++ b/litellm-rust/crates/redis-connect/src/refresh.rs @@ -247,9 +247,11 @@ mod tests { use super::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture, refresh_after}; use litellm_redis_client::{ ClusterEndpoint, ClusterNode, Credentials, CredentialsProvider, Endpoint, Namespace, - RedisClient, RedisConfig, Timeouts, TlsOptions, + RedisClient, RedisConfig, SentinelEndpoint, Timeouts, TlsOptions, + }; + use litellm_redis_testing::{ + ClientConnection, RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually, }; - use litellm_redis_testing::{ClientConnection, RedisDeployment, Topology, eventually}; #[rstest] #[case::three_quarters_of_the_lifetime( @@ -623,6 +625,7 @@ mod tests { timeouts: Timeouts::default(), breaker: None, events: None, + health_check: None, }) .unwrap(); client.set("k", "1", Duration::from_secs(60)).await.unwrap(); @@ -658,4 +661,87 @@ mod tests { assert_eq!(connection_users(&reconnected), BTreeSet::from(["bob"])); assert!(connection_ids(&reconnected).is_disjoint(&connection_ids(&after))); } + + /// Cloud credentials rotate while Sentinel fails over: the refreshed user must be what the + /// client presents to the master Sentinel promotes, not the one it connected with first. + #[rstest] + #[tokio::test] + async fn a_refreshed_user_reaches_the_master_sentinel_promotes() { + let sentinel = RedisSentinel::start().await.unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + user("alice", "first", Duration::from_secs(2)), + user("bob", "second", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let (host, port) = sentinel.sentinel(); + let client = RedisClient::new(RedisConfig { + endpoint: Endpoint::Sentinel(SentinelEndpoint { + sentinels: vec![ClusterNode { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + sentinel_password: None, + tls: None, + db: 0, + }), + credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))), + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(1), None), + breaker: None, + events: None, + health_check: None, + }) + .unwrap(); + client + .set("before", "1", Duration::from_secs(60)) + .await + .unwrap(); + let master = sentinel.master_port().await.unwrap(); + let app_users = |connections: &[ClientConnection]| -> BTreeSet { + connections + .iter() + .filter(|connection| connection.user != "default") + .map(|connection| connection.user.clone()) + .collect() + }; + eventually(async || { + let connections = sentinel.clients_on(master).await.ok()?; + (app_users(&connections) == BTreeSet::from(["bob".to_owned()])).then_some(()) + }) + .await; + sentinel + .cli_on_data_nodes(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || client.set("after", "2", Duration::from_secs(60)).await.ok()).await; + let on_promoted = sentinel.clients_on(promoted).await.unwrap(); + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!( + app_users(&on_promoted), + BTreeSet::from(["bob".to_owned()]), + "{on_promoted:?}" + ); + assert_eq!( + client.get::("after").await.unwrap().as_deref(), + Some("2") + ); + } } diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md new file mode 100644 index 00000000000..286aa6fba37 --- /dev/null +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -0,0 +1,7 @@ +- One module per Redis family (`router_cooldowns.rs`, `router_usage.rs`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys +- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec +- A read-through family builds its own `DualCache` from a `TierPolicy` the host passes in. Its `pub const POLICY` is `main`'s defaults, and the Python host overrides what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `RedisClient` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` +- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not +- `tests/fixtures/router_wire.json` is the wire contract, checked here against real Redis and in `tests/unit/redis/stores/` against Python's own codec. Change both sides together +- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. A family exposes a request-aware variant that declares on it (`declare_active`, `increment_deployment_post_call`) next to its direct method, which runs the same declaration on a batch of its own +- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml new file mode 100644 index 00000000000..b5d3c0cf88c --- /dev/null +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-redis-store" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-python-compat.workspace = true +litellm-redis-client.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tracing.workspace = true + +[dev-dependencies] +litellm-circuit-breaker.workspace = true +litellm-redis-testing.workspace = true +litellm-secret-value.workspace = true +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/redis-store/src/codec.rs b/litellm-rust/crates/redis-store/src/codec.rs new file mode 100644 index 00000000000..a8662fcbe0c --- /dev/null +++ b/litellm-rust/crates/redis-store/src/codec.rs @@ -0,0 +1,20 @@ +use litellm_python_compat::{Value, json, literal}; + +#[derive(Debug, thiserror::Error)] +pub enum DecodeError { + #[error("Redis value is not UTF-8")] + NotUtf8, + #[error("Redis value is neither JSON nor a Python literal: {0}")] + Unparsed(#[source] litellm_python_compat::Error), + #[error("Redis value has the wrong shape: {0}")] + Shape(&'static str), +} + +/// `RedisCache._get_cache_logic`: `json.loads`, falling back to `ast.literal_eval`. +pub(crate) fn python_value(raw: &[u8]) -> Result { + let text = std::str::from_utf8(raw).map_err(|_| DecodeError::NotUtf8)?; + match serde_json::from_str::(text) { + Ok(document) => Ok(json::from_json(document)), + Err(_) => literal::literal_eval(text).map_err(DecodeError::Unparsed), + } +} diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs new file mode 100644 index 00000000000..d5999206fd6 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -0,0 +1,9 @@ +mod codec; +mod number; +mod request; +pub mod router_cooldowns; +pub mod router_usage; + +pub use codec::DecodeError; +pub use number::Number; +pub use request::RequestRedis; diff --git a/litellm-rust/crates/redis-store/src/number.rs b/litellm-rust/crates/redis-store/src/number.rs new file mode 100644 index 00000000000..65864e49bba --- /dev/null +++ b/litellm-rust/crates/redis-store/src/number.rs @@ -0,0 +1,68 @@ +use std::{cmp::Ordering, ops::Add}; + +use litellm_python_compat::{BigInt, Value}; + +/// A Python `int` or `float`, kept apart because `repr` and `json.dumps` write them differently. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Number { + Int(i64), + Float(f64), +} + +impl Number { + pub fn as_f64(self) -> f64 { + match self { + Self::Int(value) => value as f64, + Self::Float(value) => value, + } + } + + /// Python's `int(value)`, which truncates toward zero. + pub fn truncate(self) -> i64 { + match self { + Self::Int(value) => value, + Self::Float(value) => value.trunc() as i64, + } + } + + pub fn to_value(self) -> Value { + match self { + Self::Int(value) => Value::Int(BigInt::from(value)), + Self::Float(value) => Value::Float(value), + } + } + + pub fn from_value(value: &Value) -> Option { + match value { + Value::Int(value) => i64::try_from(value).ok().map(Self::Int), + Value::Float(value) => Some(Self::Float(*value)), + _ => None, + } + } + + /// Python's `min` over numbers: the first of the smallest, keeping its type. + pub fn min(values: impl IntoIterator) -> Option { + values + .into_iter() + .fold(None, |smallest, value| match smallest { + Some(smallest) + if value.as_f64().partial_cmp(&smallest.as_f64()) != Some(Ordering::Less) => + { + Some(smallest) + } + _ => Some(value), + }) + } +} + +/// Python's `+`: two ints stay an int, anything with a float is a float. +impl Add for Number { + type Output = Self; + + fn add(self, other: Self) -> Self { + match (self, other) { + (Self::Int(left), Self::Int(right)) => Self::Int(left + right), + (left, right) => Self::Float(left.as_f64() + right.as_f64()), + } + } +} diff --git a/litellm-rust/crates/redis-store/src/request.rs b/litellm-rust/crates/redis-store/src/request.rs new file mode 100644 index 00000000000..e4f8014c0bb --- /dev/null +++ b/litellm-rust/crates/redis-store/src/request.rs @@ -0,0 +1,99 @@ +use std::{ + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::Duration, +}; + +use futures_util::future::join_all; +use litellm_redis_client::{RedisBatch, RedisClient}; +use tokio::sync::oneshot; + +/// One request's Redis batches: a pre-call batch for the reads it needs before answering and a +/// post-call batch for the writes nobody waits on, one of each per Redis client. The host that owns +/// the request builds one and passes it down; it knows nothing else about the request. +#[derive(Clone)] +pub struct RequestRedis(Arc); + +struct Shared { + post_call_deadline: Duration, + batches: Mutex>, + deadline: Mutex>>, +} + +#[derive(Clone)] +struct ClientBatches { + client: RedisClient, + pre: RedisBatch, + post: RedisBatch, +} + +impl RequestRedis { + pub fn new(post_call_deadline: Duration) -> Self { + Self(Arc::new(Shared { + post_call_deadline, + batches: Mutex::new(Vec::new()), + deadline: Mutex::new(None), + })) + } + + pub fn pre(&self, client: &RedisClient) -> RedisBatch { + self.batches_for(client).pre + } + + /// The post-call batch, which goes out on `flush_post` or `post_call_deadline` after the first + /// declaration, whichever comes first. Must be called inside a Tokio runtime. + pub fn post(&self, client: &RedisClient) -> RedisBatch { + self.arm_deadline(); + self.batches_for(client).post + } + + pub async fn flush_pre(&self) { + join_all(self.batches().iter().map(|batches| batches.pre.flush())).await; + } + + /// Sends every post-call op declared so far; the shutdown drain calls it too. + pub async fn flush_post(&self) { + lock(&self.0.deadline).take(); + join_all(self.batches().iter().map(|batches| batches.post.flush())).await; + } + + fn batches(&self) -> Vec { + lock(&self.0.batches).clone() + } + + fn batches_for(&self, client: &RedisClient) -> ClientBatches { + let mut batches = lock(&self.0.batches); + if let Some(existing) = batches + .iter() + .find(|batches| batches.client.same_client(client)) + { + return existing.clone(); + } + let created = ClientBatches { + client: client.clone(), + pre: client.batch(), + post: client.batch(), + }; + batches.push(created.clone()); + created + } + + fn arm_deadline(&self) { + let mut deadline = lock(&self.0.deadline); + if deadline.is_some() { + return; + } + let (disarm, disarmed) = oneshot::channel(); + *deadline = Some(disarm); + let request = self.clone(); + tokio::spawn(async move { + tokio::select! { + () = tokio::time::sleep(request.0.post_call_deadline) => request.flush_post().await, + _ = disarmed => {} + } + }); + } +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} diff --git a/litellm-rust/crates/redis-store/src/router_cooldowns.rs b/litellm-rust/crates/redis-store/src/router_cooldowns.rs new file mode 100644 index 00000000000..da1b047798b --- /dev/null +++ b/litellm-rust/crates/redis-store/src/router_cooldowns.rs @@ -0,0 +1,282 @@ +use std::time::Duration; + +use litellm_cache::{ + Clock, DeclaredRead, DualCache, Error as MemoryError, OnBreakerOpen, TierPolicy, system_clock, +}; +use litellm_cache_memory::InMemoryCache; +use litellm_python_compat::{Value, repr::repr}; +use litellm_redis_client::{Expiry, RedisBatch, RedisClient}; + +use crate::{DecodeError, Number, RequestRedis, codec::python_value}; + +/// `main`'s cooldown tier: `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` (1 s), `InMemoryCache`'s +/// 600 s default for backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`. +pub const POLICY: TierPolicy = TierPolicy { + refresh_after: Duration::from_secs(1), + backfill_ttl: Duration::from_secs(600), + on_breaker_open: OnBreakerOpen::ServeMemory, + tracked_reads: 1000, +}; +const MEMORY_ENTRIES: usize = 200; +const MAX_CORRECTED_TTL: Duration = Duration::from_secs(60); +const CORRECTION_SLACK: Duration = Duration::from_secs(5); + +#[derive(Clone, Debug, PartialEq)] +pub struct Cooldown { + pub exception_received: String, + pub status_code: String, + pub timestamp: Number, + pub cooldown_time: Number, +} + +impl Cooldown { + fn remaining(&self, now: Duration) -> f64 { + self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64() + } + + /// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`. + pub fn encode(&self) -> String { + let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value); + repr(&Value::Dict(vec![ + field( + "exception_received", + Value::Str(self.exception_received.clone()), + ), + field("status_code", Value::Str(self.status_code.clone())), + field("timestamp", self.timestamp.to_value()), + field("cooldown_time", self.cooldown_time.to_value()), + ])) + } + + /// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does. + pub fn decode(raw: &[u8]) -> Result { + let Value::Dict(fields) = python_value(raw)? else { + return Err(DecodeError::Shape("a cooldown is a dict")); + }; + let field = |name: &str| { + fields + .iter() + .find(|(key, _)| matches!(key, Value::Str(key) if key == name)) + .map(|(_, value)| value) + }; + let text = |name: &'static str| match field(name) { + Some(Value::Str(text)) => Ok(text.clone()), + _ => Err(DecodeError::Shape(name)), + }; + let number = |name: &'static str| { + field(name) + .and_then(Number::from_value) + .ok_or(DecodeError::Shape(name)) + }; + Ok(Self { + exception_received: text("exception_received")?, + status_code: text("status_code")?, + timestamp: number("timestamp")?, + cooldown_time: number("cooldown_time")?, + }) + } +} + +fn key(model_id: &str) -> String { + format!("deployment:{model_id}:cooldown") +} + +fn keys(model_ids: &[String]) -> Vec { + model_ids.iter().map(|model_id| key(model_id)).collect() +} + +fn allowed_fails_key(model_id: &str, suffix: Option<&str>) -> String { + match suffix.filter(|suffix| !suffix.is_empty()) { + Some(suffix) => format!("deployment:{model_id}:allowed_fails:{suffix}"), + None => format!("deployment:{model_id}:allowed_fails"), + } +} + +fn memory(policy: TierPolicy, now: &Clock) -> InMemoryCache { + let now = now.clone(); + InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || { + now() + }) +} + +/// `BaseCache.get_ttl` sends `int(ttl)`, so Redis sees whole seconds. +fn whole_seconds(ttl: Duration) -> Duration { + Duration::from_secs(ttl.as_secs()) +} + +/// An `active` read declared on a request batch, answered by `resolve_active`. +pub struct DeclaredCooldowns { + model_ids: Vec, + read: DeclaredRead, +} + +/// One router's cooldowns and allowed-fails counters: memory per router, Redis shared by every pod. +pub struct RouterCooldowns { + cooldowns: DualCache, RedisClient>, + fails: DualCache, RedisClient>, +} + +impl RouterCooldowns { + pub fn new(client: Option, policy: TierPolicy) -> Self { + Self::with_clock(client, policy, system_clock()) + } + + pub fn with_clock(client: Option, policy: TierPolicy, now: Clock) -> Self { + Self { + cooldowns: DualCache::read_through( + memory(policy, &now), + client.clone(), + policy, + now.clone(), + ), + fails: DualCache::read_through(memory(policy, &now), client, policy, now), + } + } + + /// Memory first, then a `SET ... EX int(ttl)` whose failure only logs, like `DualCache.set_cache`. + pub async fn put( + &self, + model_id: &str, + cooldown: &Cooldown, + ttl: Duration, + ) -> Result<(), MemoryError> { + let key = key(model_id); + self.cooldowns.set_memory(&key, cooldown.clone(), ttl)?; + if let Some(remote) = self.cooldowns.remote() + && let Err(error) = remote + .set(&key, cooldown.encode(), whole_seconds(ttl)) + .await + { + tracing::warn!(%error, "writing a cooldown to Redis failed"); + } + Ok(()) + } + + /// The fleet-wide fail count. A Redis failure falls back to this router's own count. + pub async fn increment_allowed_fails( + &self, + model_id: &str, + suffix: Option<&str>, + ttl: Duration, + ) -> Result { + let key = allowed_fails_key(model_id, suffix); + let local = self + .fails + .update_memory(&key, ttl, |count| count.unwrap_or(0) + 1)?; + let Some(remote) = self.fails.remote() else { + return Ok(local); + }; + match remote + .incr_by(&key, 1, Expiry::KeepExisting(whole_seconds(ttl))) + .await + { + Ok(count) => Ok(count), + Err(error) => { + tracing::warn!(%error, "allowed_fails counter fell back to this worker's in-memory count"); + Ok(local) + } + } + } + + /// The deployments still cooling down. An entry whose cooldown already ended is dropped + /// from memory, and one memory would keep past its end is shortened, at most to 60 s. + pub async fn active( + &self, + model_ids: &[String], + ) -> Result, MemoryError> { + let declared = self.declare(RedisClient::batch, model_ids)?; + self.resolve_active(declared).await + } + + /// `active`, with its Redis read declared on the request's pre-call batch. + pub fn declare_active( + &self, + request: &RequestRedis, + model_ids: &[String], + ) -> Result { + self.declare(|client| request.pre(client), model_ids) + } + + pub async fn resolve_active( + &self, + declared: DeclaredCooldowns, + ) -> Result, MemoryError> { + let DeclaredCooldowns { model_ids, read } = declared; + let read = self.cooldowns.resolve_many(read, Cooldown::decode).await?; + let now = self.cooldowns.now(); + let corrected = model_ids + .into_iter() + .zip(read) + .filter_map(|(model_id, cooldown)| cooldown.map(|cooldown| (model_id, cooldown))) + .map(|(model_id, cooldown)| { + Ok(self + .correct(&model_id, &cooldown, now)? + .then_some((model_id, cooldown))) + }) + .collect::, MemoryError>>()?; + Ok(corrected.into_iter().flatten().collect()) + } + + /// The shortest stored cooldown, ended ones included, as `main`'s `get_min_cooldown` reads it. + pub async fn minimum(&self, model_ids: &[String]) -> Result, MemoryError> { + let read = self.read(model_ids).await?; + Ok(Number::min( + read.into_iter() + .flatten() + .map(|cooldown| cooldown.cooldown_time), + )) + } + + pub fn evict_local(&self, model_id: &str) -> Result<(), MemoryError> { + self.cooldowns.delete_memory(&key(model_id)) + } + + pub fn flush_local_counters(&self) -> Result<(), MemoryError> { + self.fails.flush_memory() + } + + /// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads + /// in-memory router state. + pub fn snapshot(&self) -> Result, MemoryError> { + self.cooldowns.memory().entries() + } + + async fn read(&self, model_ids: &[String]) -> Result>, MemoryError> { + self.cooldowns + .get_many(&keys(model_ids), Cooldown::decode) + .await + } + + fn declare( + &self, + batch_for: impl FnOnce(&RedisClient) -> RedisBatch, + model_ids: &[String], + ) -> Result { + Ok(DeclaredCooldowns { + model_ids: model_ids.to_vec(), + read: self.cooldowns.declare_many(batch_for, &keys(model_ids))?, + }) + } + + fn correct( + &self, + model_id: &str, + cooldown: &Cooldown, + now: Duration, + ) -> Result { + let key = key(model_id); + let remaining = cooldown.remaining(now); + if remaining <= 0.0 { + self.cooldowns.delete_memory(&key)?; + return Ok(false); + } + let remaining = Duration::from_secs_f64(remaining); + let memory = self.cooldowns.memory(); + if memory.expires_at(&key)? > Some(now + remaining + CORRECTION_SLACK) { + self.cooldowns.delete_memory(&key)?; + self.cooldowns + .set_memory(&key, cooldown.clone(), remaining.min(MAX_CORRECTED_TTL))?; + } + Ok(true) + } +} diff --git a/litellm-rust/crates/redis-store/src/router_usage.rs b/litellm-rust/crates/redis-store/src/router_usage.rs new file mode 100644 index 00000000000..10e7ae90626 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/router_usage.rs @@ -0,0 +1,261 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{ + Clock, DualCache, Error as MemoryError, OnBreakerOpen, TierPolicy, system_clock, +}; +use litellm_cache_memory::InMemoryCache; +use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient}; + +use crate::{DecodeError, Number, RequestRedis, codec::python_value}; + +/// `main`'s usage tier: `DualCache`'s 10 s default refresh, `InMemoryCache`'s 600 s default for +/// backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`. +pub const POLICY: TierPolicy = TierPolicy { + refresh_after: Duration::from_secs(10), + backfill_ttl: Duration::from_secs(600), + on_breaker_open: OnBreakerOpen::ServeMemory, + tracked_reads: 1000, +}; +const MEMORY_ENTRIES: usize = 200; +/// `RoutingArgs.ttl`, the window of a per-minute counter. +pub const COUNTER_TTL: Duration = Duration::from_secs(60); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Kind { + Tpm, + Rpm, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Counter { + pub deployment_id: String, + pub model: Option, + pub minute: String, + pub kind: Kind, +} + +impl Counter { + /// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`. + fn key(&self) -> String { + let kind = match self.kind { + Kind::Tpm => "tpm", + Kind::Rpm => "rpm", + }; + let model = self.model.as_deref().unwrap_or("None"); + format!( + "global_router:{}:{model}:{kind}:{}", + self.deployment_id, self.minute + ) + } +} + +fn decode(raw: &[u8]) -> Result { + Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape("a usage counter is a number")) +} + +fn whole_seconds(ttl: Duration) -> Duration { + Duration::from_secs(ttl.as_secs()) +} + +/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod. +pub struct RouterUsage(DualCache, RedisClient>); + +impl RouterUsage { + pub fn new(client: Option, policy: TierPolicy) -> Self { + Self::with_clock(client, policy, system_clock()) + } + + pub fn with_clock(client: Option, policy: TierPolicy, now: Clock) -> Self { + let clock = now.clone(); + let memory = + InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || { + clock() + }); + Self(DualCache::read_through(memory, client, policy, now)) + } + + /// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL. + /// A Redis failure answers with the memory count. + pub async fn increment( + &self, + counter: &Counter, + value: Number, + ttl: Duration, + ) -> Result { + let key = counter.key(); + let local = self.add_local(&key, value, ttl)?; + let Some(remote) = self.0.remote() else { + return Ok(local); + }; + match remote + .incr_by_float( + &key, + value.as_f64(), + Expiry::KeepExisting(whole_seconds(ttl)), + ) + .await + { + Ok(total) => Ok(Number::Float(total)), + Err(error) => { + tracing::warn!(%error, "Redis usage increment failed, falling back to the in-memory count"); + Ok(local) + } + } + } + + /// One deployment's tokens and requests: memory, then one pipeline that resets each + /// counter's 60 s window, mirrored back into memory as `int` totals. + pub async fn increment_deployment( + &self, + deployment_id: &str, + model: Option<&str>, + minute: &str, + tokens: Number, + requests: i64, + ) -> Result<(), MemoryError> { + let increment = Increment::new(deployment_id, model, minute, tokens, requests); + let Some(declared) = self.declare(RedisClient::batch, increment)? else { + return Ok(()); + }; + let totals = totals(declared.totals, Declared::resolve).await; + self.mirror(&declared.keys, declared.local, totals) + } + + /// `increment_deployment` on the request's post-call batch. Memory counts at once; the Redis + /// totals reach memory whenever that batch goes out, from a task this spawns, so it must be + /// called inside a Tokio runtime. + pub fn increment_deployment_post_call( + self: &Arc, + request: &RequestRedis, + deployment_id: &str, + model: Option<&str>, + minute: &str, + tokens: Number, + requests: i64, + ) -> Result<(), MemoryError> { + let increment = Increment::new(deployment_id, model, minute, tokens, requests); + let Some(declared) = self.declare(|client| request.post(client), increment)? else { + return Ok(()); + }; + let store = Arc::clone(self); + tokio::spawn(async move { + let totals = totals(declared.totals, Declared::settled).await; + if let Err(error) = store.mirror(&declared.keys, declared.local, totals) { + tracing::warn!(%error, "mirroring Redis usage totals into memory failed"); + } + }); + Ok(()) + } + + pub async fn get_many(&self, counters: &[Counter]) -> Result>, MemoryError> { + let keys = counters.iter().map(Counter::key).collect::>(); + self.0.get_many(&keys, decode).await + } + + pub fn flush_local_counters(&self) -> Result<(), MemoryError> { + self.0.flush_memory() + } + + /// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads + /// in-memory router state. + pub fn snapshot(&self) -> Result, MemoryError> { + self.0.memory().entries() + } + + fn declare( + &self, + batch_for: impl FnOnce(&RedisClient) -> RedisBatch, + increment: Increment, + ) -> Result, MemoryError> { + let local = increment + .counts + .iter() + .map(|(key, value)| self.add_local(key, *value, COUNTER_TTL)) + .collect::, _>>()?; + let Some(remote) = self.0.remote() else { + return Ok(None); + }; + let batch = batch_for(remote); + let totals = increment + .counts + .iter() + .map(|(key, value)| { + batch.incr_by_float(key, value.as_f64(), Expiry::Refresh(COUNTER_TTL)) + }) + .collect(); + Ok(Some(DeclaredIncrement { + keys: increment.counts.into_iter().map(|(key, _)| key).collect(), + local, + totals, + })) + } + + fn mirror( + &self, + keys: &[String], + local: Vec, + totals: Result, litellm_redis_client::Error>, + ) -> Result<(), MemoryError> { + let totals = match totals { + Ok(totals) => totals.into_iter().map(Number::Float).collect(), + Err(error) => { + tracing::warn!(%error, "Redis usage pipeline failed, keeping the in-memory counts"); + local + } + }; + keys.iter().zip(totals).try_for_each(|(key, total)| { + self.0 + .set_memory(key, Number::Int(total.truncate()), COUNTER_TTL) + }) + } + + fn add_local(&self, key: &str, value: Number, ttl: Duration) -> Result { + self.0.update_memory(key, ttl, |current| { + current.unwrap_or(Number::Int(0)) + value + }) + } +} + +struct Increment { + counts: [(String, Number); 2], +} + +impl Increment { + fn new( + deployment_id: &str, + model: Option<&str>, + minute: &str, + tokens: Number, + requests: i64, + ) -> Self { + let counter = |kind| Counter { + deployment_id: deployment_id.to_owned(), + model: model.map(str::to_owned), + minute: minute.to_owned(), + kind, + }; + Self { + counts: [ + (counter(Kind::Tpm).key(), tokens), + (counter(Kind::Rpm).key(), Number::Int(requests)), + ], + } + } +} + +struct DeclaredIncrement { + keys: Vec, + local: Vec, + totals: Vec>, +} + +async fn totals>>( + declared: Vec>, + wait: impl Fn(Declared) -> F, +) -> Result, litellm_redis_client::Error> { + let mut totals = Vec::with_capacity(declared.len()); + for total in declared { + totals.push(wait(total).await?); + } + Ok(totals) +} diff --git a/litellm-rust/crates/redis-store/tests/request.rs b/litellm-rust/crates/redis-store/tests/request.rs new file mode 100644 index 00000000000..f0314f2228c --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/request.rs @@ -0,0 +1,114 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_client::{Expiry, Operation}; +use litellm_redis_store::RequestRedis; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, connect_observed}; + +const KEEP: Expiry = Expiry::KeepExisting(Duration::from_secs(60)); +const LONG_DEADLINE: Duration = Duration::from_secs(60); + +#[rstest] +#[tokio::test] +async fn pre_declarations_on_one_client_share_one_pipeline_and_another_client_gets_its_own( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let (other, other_trips) = connect_observed(&deployment); + let request = RequestRedis::new(LONG_DEADLINE); + + let first = request.pre(&client).incr_by("a", 1, KEEP); + let second = request.pre(&client.clone()).incr_by("b", 2, KEEP); + let third = request.pre(&other).incr_by("c", 3, KEEP); + request.flush_pre().await; + + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]); + assert_eq!(other_trips.seen(), [Operation::Pipeline { ops: 1 }]); + assert_eq!( + [ + first.resolve().await.unwrap(), + second.resolve().await.unwrap(), + third.resolve().await.unwrap(), + ], + [1, 2, 3] + ); + assert_eq!(trips.seen().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_post_declaration_waits_for_flush_post_and_flush_pre_leaves_it_alone( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(LONG_DEADLINE); + + let total = tokio::spawn( + request + .post(&client) + .incr_by_float("spend", 2.5, KEEP) + .settled(), + ); + request.flush_pre().await; + tokio::time::sleep(Duration::from_millis(200)).await; + + assert!(!total.is_finished()); + assert!(trips.seen().is_empty()); + assert_eq!(raw.get::("spend").await.unwrap(), None); + + request.flush_post().await; + + assert_eq!(total.await.unwrap().unwrap(), 2.5); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 1 }]); +} + +#[rstest] +#[tokio::test] +async fn a_post_batch_nobody_flushes_goes_out_at_the_deadline( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Duration::from_millis(100)); + + let first = request.post(&client).incr_by("a", 1, KEEP); + let second = request.post(&client).incr_by("b", 2, KEEP); + let settled = tokio::time::timeout(Duration::from_secs(5), async { + (first.settled().await, second.settled().await) + }) + .await + .unwrap(); + + assert_eq!((settled.0.unwrap(), settled.1.unwrap()), (1, 2)); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]); +} + +#[rstest] +#[tokio::test] +async fn a_post_declaration_after_a_flush_arms_a_new_deadline( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Duration::from_millis(300)); + + let early = request.post(&client).incr_by("a", 1, KEEP); + request.flush_post().await; + tokio::time::sleep(Duration::from_millis(200)).await; + let late = request.post(&client).incr_by("a", 1, KEEP); + tokio::time::sleep(Duration::from_millis(200)).await; + let flushed_before_its_deadline = trips.seen().len() > 1; + let late = tokio::time::timeout(Duration::from_secs(5), late.settled()) + .await + .unwrap(); + + assert_eq!(early.settled().await.unwrap(), 1); + assert!(!flushed_before_its_deadline); + assert_eq!(late.unwrap(), 2); +} diff --git a/litellm-rust/crates/redis-store/tests/router_cooldowns.rs b/litellm-rust/crates/redis-store/tests/router_cooldowns.rs new file mode 100644 index 00000000000..17bfc566a60 --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/router_cooldowns.rs @@ -0,0 +1,394 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_client::{Expiry, Operation, Ttl}; +use litellm_redis_store::{ + Number, RequestRedis, + router_cooldowns::{POLICY, RouterCooldowns}, +}; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{ALLOWED_FAILS, connect, cooldowns, within}; + +/// A clock fixed just after the fixture timestamps, so every fixture cooldown is still active. +fn fixture_clock() -> litellm_cache::Clock { + std::sync::Arc::new(|| Duration::from_secs_f64(1_730_000_000.75)) +} + +#[rstest] +#[tokio::test] +async fn put_writes_mains_str_dict_bytes_with_whole_second_ttl( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY); + + for case in cooldowns() { + let ttl = Duration::from_secs_f64(case.ttl_seconds); + store.put(case.model_id, &case.cooldown, ttl).await.unwrap(); + let key = case.key; + + let written: Option = raw.get(key).await.unwrap(); + let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else { + panic!("{key} has no TTL"); + }; + assert_eq!(written.as_deref(), Some(case.str_dict)); + assert!(within(remaining, case.redis_ttl_seconds), "{remaining:?}"); + } +} + +#[rstest] +#[tokio::test] +async fn reads_cooldowns_python_wrote_as_str_dict_or_json( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values("str_dict", "json")] format: &str, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock()); + + for case in cooldowns() { + let bytes = match format { + "json" => case.json, + _ => case.str_dict, + }; + raw.set(case.key, bytes, Duration::from_secs(60)) + .await + .unwrap(); + + assert_eq!( + store.active(&[case.model_id.to_owned()]).await.unwrap(), + vec![(case.model_id.to_owned(), case.cooldown)] + ); + } +} + +#[rstest] +#[tokio::test] +async fn an_ended_cooldown_is_not_active_but_still_counts_toward_the_minimum( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let late = std::sync::Arc::new(|| Duration::from_secs(1_730_000_100)); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, late); + let models = cooldowns() + .iter() + .map(|case| case.model_id.to_owned()) + .collect::>(); + for case in cooldowns() { + raw.set(case.key, case.str_dict, Duration::from_secs(60)) + .await + .unwrap(); + } + + assert_eq!( + store.minimum(&models).await.unwrap(), + Some(Number::Float(0.5)) + ); + assert_eq!(store.active(&models).await.unwrap(), vec![]); +} + +#[rstest] +#[tokio::test] +async fn allowed_fails_count_across_routers_and_keep_the_first_ttl( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let first = RouterCooldowns::new(Some(connect(&deployment)), POLICY); + let second = RouterCooldowns::new(Some(connect(&deployment)), POLICY); + + for case in ALLOWED_FAILS { + let key = case.key; + raw.del(&[key]).await.unwrap(); + let suffix = case.suffix; + + let counts = [ + first + .increment_allowed_fails("dep-a", suffix, Duration::from_secs_f64(30.9)) + .await + .unwrap(), + second + .increment_allowed_fails("dep-a", suffix, Duration::from_secs(120)) + .await + .unwrap(), + ]; + let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else { + panic!("{key} has no TTL"); + }; + + assert_eq!(counts, [1, 2], "{key}"); + assert!(within(remaining, 30), "{remaining:?}"); + } +} + +#[rstest] +#[tokio::test] +async fn python_written_fail_counts_continue_in_rust( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY); + raw.incr_by( + "deployment:dep-a:allowed_fails", + 4, + Expiry::KeepExisting(Duration::from_secs(60)), + ) + .await + .unwrap(); + + assert_eq!( + store + .increment_allowed_fails("dep-a", None, Duration::from_secs(60)) + .await + .unwrap(), + 5 + ); +} + +#[rstest] +#[tokio::test] +async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() { + let store = RouterCooldowns::with_clock(None, POLICY, fixture_clock()); + let [case, _] = cooldowns(); + + store + .put(case.model_id, &case.cooldown, Duration::from_secs(90)) + .await + .unwrap(); + let fails = [ + store + .increment_allowed_fails("dep-a", None, Duration::from_secs(60)) + .await + .unwrap(), + store + .increment_allowed_fails("dep-a", None, Duration::from_secs(60)) + .await + .unwrap(), + ]; + store.flush_local_counters().unwrap(); + + assert_eq!( + store.active(&[case.model_id.to_owned()]).await.unwrap(), + vec![(case.model_id.to_owned(), case.cooldown)] + ); + assert_eq!(fails, [1, 2]); + assert_eq!( + store + .increment_allowed_fails("dep-a", None, Duration::from_secs(60)) + .await + .unwrap(), + 1 + ); +} + +#[rstest] +#[tokio::test] +async fn a_redis_outage_serves_memory_and_counts_fails_locally() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock()); + let [case, _] = cooldowns(); + store + .put(case.model_id, &case.cooldown, Duration::from_secs(90)) + .await + .unwrap(); + deployment.stop().await.unwrap(); + + let active = store + .active(&[case.model_id.to_owned(), "dep-missing".to_owned()]) + .await + .unwrap(); + let fails = store + .increment_allowed_fails("dep-a", None, Duration::from_secs(60)) + .await + .unwrap(); + + assert_eq!(active, vec![(case.model_id.to_owned(), case.cooldown)]); + assert_eq!(fails, 1); +} + +struct StepClock(std::sync::Mutex); + +impl StepClock { + fn at(seconds: f64) -> std::sync::Arc { + std::sync::Arc::new(Self(std::sync::Mutex::new(Duration::from_secs_f64( + seconds, + )))) + } + + fn advance(&self, by: Duration) { + *self.0.lock().unwrap() += by; + } + + fn clock(self: &std::sync::Arc) -> litellm_cache::Clock { + let this = self.clone(); + std::sync::Arc::new(move || *this.0.lock().unwrap()) + } +} + +#[rstest] +#[tokio::test] +async fn a_miss_is_read_from_redis_again_only_after_the_one_second_refresh( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let time = StepClock::at(1_730_000_000.75); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, time.clock()); + let [case, _] = cooldowns(); + + let before = store.active(&[case.model_id.to_owned()]).await.unwrap(); + raw.set(case.key, case.str_dict, Duration::from_secs(60)) + .await + .unwrap(); + let throttled = store.active(&[case.model_id.to_owned()]).await.unwrap(); + time.advance(Duration::from_secs(1)); + let refreshed = store.active(&[case.model_id.to_owned()]).await.unwrap(); + + assert_eq!(before, vec![]); + assert_eq!(throttled, vec![]); + assert_eq!( + refreshed, + vec![(case.model_id.to_owned(), case.cooldown.clone())] + ); +} + +#[rstest] +#[tokio::test] +async fn a_breaker_refusal_does_not_hold_back_the_next_read() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let raw = connect(&deployment); + let breaker = litellm_circuit_breaker::Breaker::new( + litellm_circuit_breaker::Config { + failure_threshold: 1, + recovery_timeout: Duration::from_millis(200), + ..litellm_circuit_breaker::Config::default() + }, + Box::new(|_| {}), + ); + let time = StepClock::at(1_730_000_000.75); + let guarded = support::connect_guarded(&deployment, breaker); + let store = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, time.clock()); + let [case, _] = cooldowns(); + deployment.stop().await.unwrap(); + store.active(&[case.model_id.to_owned()]).await.unwrap(); + time.advance(Duration::from_secs(1)); + let refused = store.active(&[case.model_id.to_owned()]).await.unwrap(); + deployment.resume().await.unwrap(); + litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await; + raw.set(case.key, case.str_dict, Duration::from_secs(60)) + .await + .unwrap(); + + let recovered = store.active(&[case.model_id.to_owned()]).await.unwrap(); + + assert_eq!(refused, vec![]); + assert_eq!( + recovered, + vec![(case.model_id.to_owned(), case.cooldown.clone())] + ); +} + +#[rstest] +#[tokio::test] +async fn a_declared_active_read_rides_the_request_batch_and_answers_what_active_does( + #[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 = RouterCooldowns::with_clock(Some(client), POLICY, fixture_clock()); + let models = cooldowns().map(|case| case.model_id.to_owned()); + for case in cooldowns() { + raw.set(case.key, case.str_dict, Duration::from_secs(60)) + .await + .unwrap(); + } + let request = RequestRedis::new(Duration::from_secs(60)); + + let declared = store.declare_active(&request, &models).unwrap(); + assert!(trips.seen().is_empty()); + request.flush_pre().await; + let active = store.resolve_active(declared).await.unwrap(); + + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]); + assert_eq!( + active, + cooldowns() + .map(|case| (case.model_id.to_owned(), case.cooldown)) + .to_vec() + ); + assert_eq!(store.active(&models).await.unwrap(), active); + assert_eq!(trips.seen().len(), 1, "the declared read backfilled memory"); +} + +#[rstest] +#[tokio::test] +async fn a_cooldown_recorded_after_the_declaration_wins_over_the_redis_miss( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock()); + let [case, _] = cooldowns(); + let request = RequestRedis::new(Duration::from_secs(60)); + + let declared = store + .declare_active(&request, &[case.model_id.to_owned()]) + .unwrap(); + request.flush_pre().await; + store + .put(case.model_id, &case.cooldown, Duration::from_secs(90)) + .await + .unwrap(); + + assert_eq!( + store.resolve_active(declared).await.unwrap(), + vec![(case.model_id.to_owned(), case.cooldown)] + ); +} + +#[rstest] +#[tokio::test] +async fn a_breaker_refusal_of_a_declared_read_does_not_hold_back_the_next_read() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let raw = connect(&deployment); + let breaker = litellm_circuit_breaker::Breaker::new( + litellm_circuit_breaker::Config { + failure_threshold: 1, + recovery_timeout: Duration::from_millis(200), + ..litellm_circuit_breaker::Config::default() + }, + Box::new(|_| {}), + ); + let guarded = support::connect_guarded(&deployment, breaker); + let store = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, fixture_clock()); + let [case, _] = cooldowns(); + let models = [case.model_id.to_owned()]; + deployment.stop().await.unwrap(); + let opened = RequestRedis::new(Duration::from_secs(60)); + let failed = store + .declare_active(&opened, &["dep-other".to_owned()]) + .unwrap(); + opened.flush_pre().await; + store.resolve_active(failed).await.unwrap(); + let refused_request = RequestRedis::new(Duration::from_secs(60)); + + let refused = store.declare_active(&refused_request, &models).unwrap(); + refused_request.flush_pre().await; + let refused = store.resolve_active(refused).await.unwrap(); + deployment.resume().await.unwrap(); + litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await; + raw.set(case.key, case.str_dict, Duration::from_secs(60)) + .await + .unwrap(); + + assert_eq!(refused, vec![]); + assert_eq!( + store.active(&models).await.unwrap(), + vec![(case.model_id.to_owned(), case.cooldown.clone())] + ); +} diff --git a/litellm-rust/crates/redis-store/tests/router_usage.rs b/litellm-rust/crates/redis-store/tests/router_usage.rs new file mode 100644 index 00000000000..2d934cdc02c --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/router_usage.rs @@ -0,0 +1,240 @@ +mod support; + +use std::{sync::Arc, time::Duration}; + +use litellm_redis_client::{Expiry, Operation, Ttl}; +use litellm_redis_store::{ + Number, RequestRedis, + router_usage::{COUNTER_TTL, Counter, Kind, POLICY, RouterUsage}, +}; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, usage, within}; + +fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] { + [Kind::Tpm, Kind::Rpm].map(|kind| Counter { + deployment_id: "dep".to_owned(), + model: model.map(str::to_owned), + minute: "12-34".to_owned(), + kind, + }) +} + +#[rstest] +#[tokio::test] +async fn reads_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(connect(&deployment)), POLICY); + let cases = usage(); + for (case, delta) in cases.iter().zip([7.75, 3.0]) { + raw.incr_by_float(case.key, delta, Expiry::Refresh(COUNTER_TTL)) + .await + .unwrap(); + } + + assert_eq!( + store + .get_many(&cases.map(|case| case.counter)) + .await + .unwrap(), + vec![Some(Number::Float(7.75)), Some(Number::Int(3))] + ); +} + +#[rstest] +#[tokio::test] +async fn deployment_increments_write_incrbyfloat_totals_with_a_refreshed_minute_window( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let writer = RouterUsage::new(Some(connect(&deployment)), POLICY); + let reader = RouterUsage::new(Some(connect(&deployment)), POLICY); + let [case, _] = usage(); + let model = case.counter.model.as_deref(); + let [tpm_key, rpm_key] = + ["tpm", "rpm"].map(|kind| case.key.replace(":tpm:", &format!(":{kind}:"))); + + writer + .increment_deployment("dep", model, "12-34", Number::Float(5.5), 1) + .await + .unwrap(); + raw.expire(&tpm_key, Duration::from_secs(5)).await.unwrap(); + writer + .increment_deployment("dep", model, "12-34", Number::Int(2), 1) + .await + .unwrap(); + + let written: [Option; 2] = [ + raw.get(&tpm_key).await.unwrap(), + raw.get(&rpm_key).await.unwrap(), + ]; + let Ttl::Expires(remaining) = raw.ttl(&tpm_key).await.unwrap() else { + panic!("{tpm_key} has no TTL"); + }; + assert_eq!(written, [Some("7.5".to_owned()), Some("2".to_owned())]); + assert!(within(remaining, 60), "{remaining:?}"); + assert_eq!( + reader.get_many(&tpm_and_rpm(model)).await.unwrap(), + vec![Some(Number::Float(7.5)), Some(Number::Int(2))] + ); + assert_eq!( + writer.get_many(&tpm_and_rpm(model)).await.unwrap(), + vec![Some(Number::Int(7)), Some(Number::Int(2))] + ); +} + +#[rstest] +#[tokio::test] +async fn a_request_increment_keeps_the_window_python_started( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(connect(&deployment)), POLICY); + let [_, case] = usage(); + let key = case.key; + raw.incr_by_float(key, 1.0, Expiry::KeepExisting(Duration::from_secs(20))) + .await + .unwrap(); + + let total = store + .increment(&case.counter, Number::Int(1), COUNTER_TTL) + .await + .unwrap(); + let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else { + panic!("{key} has no TTL"); + }; + + assert_eq!(total, Number::Float(2.0)); + assert!(within(remaining, 20), "{remaining:?}"); +} + +#[rstest] +#[tokio::test] +async fn without_redis_counters_accumulate_in_memory_until_flushed() { + let store = RouterUsage::new(None, POLICY); + let counters = tpm_and_rpm(Some("model")); + + store + .increment_deployment("dep", Some("model"), "12-34", Number::Int(2), 1) + .await + .unwrap(); + store + .increment_deployment("dep", Some("model"), "12-34", Number::Float(5.5), 1) + .await + .unwrap(); + let counted = store.get_many(&counters).await.unwrap(); + store.flush_local_counters().unwrap(); + + assert_eq!( + counted, + vec![Some(Number::Float(7.5)), Some(Number::Int(2))] + ); + assert_eq!(store.get_many(&counters).await.unwrap(), vec![None, None]); +} + +#[rstest] +#[tokio::test] +async fn a_failed_pipeline_keeps_the_memory_counts_as_integers() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let store = RouterUsage::new(Some(connect(&deployment)), POLICY); + deployment.stop().await.unwrap(); + + store + .increment_deployment("dep", Some("model"), "12-34", Number::Float(5.75), 1) + .await + .unwrap(); + + assert_eq!( + store.get_many(&tpm_and_rpm(Some("model"))).await.unwrap(), + vec![Some(Number::Int(5)), Some(Number::Int(1))] + ); +} + +#[rstest] +#[tokio::test] +async fn post_call_increments_wait_for_flush_post_then_mirror_the_redis_totals_into_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(client), POLICY)); + let counters = tpm_and_rpm(Some("model")); + let [tpm_key, rpm_key] = + ["tpm", "rpm"].map(|kind| format!("global_router:dep:model:{kind}:12-34")); + raw.incr_by_float(&tpm_key, 10.0, Expiry::Refresh(COUNTER_TTL)) + .await + .unwrap(); + let request = RequestRedis::new(Duration::from_secs(60)); + + store + .increment_deployment_post_call( + &request, + "dep", + Some("model"), + "12-34", + Number::Float(5.5), + 1, + ) + .unwrap(); + let before_flush = store.get_many(&counters).await.unwrap(); + let unsent: Option = raw.get(&rpm_key).await.unwrap(); + let trips_before_flush = trips.seen(); + request.flush_post().await; + let mirrored = litellm_redis_testing::eventually(async || { + let counted = store.get_many(&counters).await.unwrap(); + (counted == [Some(Number::Int(15)), Some(Number::Int(1))]).then_some(counted) + }) + .await; + + assert_eq!( + before_flush, + vec![Some(Number::Float(5.5)), Some(Number::Int(1))] + ); + assert_eq!(unsent, None); + assert!(trips_before_flush.is_empty()); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]); + assert_eq!(mirrored, vec![Some(Number::Int(15)), Some(Number::Int(1))]); + let written: [Option; 2] = [ + raw.get(&tpm_key).await.unwrap(), + raw.get(&rpm_key).await.unwrap(), + ]; + assert_eq!(written, [Some("15.5".to_owned()), Some("1".to_owned())]); + let Ttl::Expires(remaining) = raw.ttl(&rpm_key).await.unwrap() else { + panic!("{rpm_key} has no TTL"); + }; + assert!(within(remaining, 60), "{remaining:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_failed_post_call_pipeline_keeps_the_memory_counts_as_integers() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let store = Arc::new(RouterUsage::new(Some(connect(&deployment)), POLICY)); + let counters = tpm_and_rpm(Some("model")); + let request = RequestRedis::new(Duration::from_secs(60)); + deployment.stop().await.unwrap(); + + store + .increment_deployment_post_call( + &request, + "dep", + Some("model"), + "12-34", + Number::Float(5.75), + 1, + ) + .unwrap(); + request.flush_post().await; + + litellm_redis_testing::eventually(async || { + (store.get_many(&counters).await.unwrap() == [Some(Number::Int(5)), Some(Number::Int(1))]) + .then_some(()) + }) + .await; +} diff --git a/litellm-rust/crates/redis-store/tests/support/mod.rs b/litellm-rust/crates/redis-store/tests/support/mod.rs new file mode 100644 index 00000000000..7b8effbf3c1 --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/support/mod.rs @@ -0,0 +1,188 @@ +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_redis_client::{ + CallEvent, ClusterEndpoint, ClusterNode, Credentials, Endpoint, Namespace, Operation, + RedisClient, RedisConfig, RedisEvents, Timeouts, TlsOptions, +}; +use litellm_redis_store::{ + Number, + router_cooldowns::Cooldown, + router_usage::{Counter, Kind}, +}; +use litellm_redis_testing::RedisDeployment; + +pub fn connect(deployment: &RedisDeployment) -> RedisClient { + connect_with(deployment, None, None) +} + +/// Every round trip a client made, as the client reports it. +#[derive(Default)] +pub struct RoundTrips(Mutex>); + +impl RoundTrips { + pub fn seen(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +impl RedisEvents for RoundTrips { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(event.operation); + } + + fn breaker(&self, _transition: &Transition) {} +} + +pub fn connect_observed(deployment: &RedisDeployment) -> (RedisClient, Arc) { + let trips = Arc::new(RoundTrips::default()); + (connect_with(deployment, None, Some(trips.clone())), trips) +} + +pub fn connect_guarded( + deployment: &RedisDeployment, + breaker: litellm_circuit_breaker::Breaker, +) -> RedisClient { + connect_with(deployment, Some(breaker), None) +} + +fn connect_with( + deployment: &RedisDeployment, + breaker: Option, + events: Option>, +) -> RedisClient { + let endpoint = match deployment { + RedisDeployment::Standalone(server) => Endpoint::Url { + url: litellm_secret_value::SecretValue::new(server.url().to_owned()), + tls: TlsOptions::verified(), + }, + RedisDeployment::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { + nodes: cluster + .nodes() + .map(|(host, port)| ClusterNode { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + }; + RedisClient::new(RedisConfig { + endpoint, + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(500), None), + breaker, + events, + health_check: None, + }) + .unwrap() +} + +/// What main's `RedisCache` writes for a cooldown: `str(value)` with `EX int(ttl)`, read back with +/// `json.loads` then `ast.literal_eval` (`litellm/caching/redis_cache.py` `set_cache` and +/// `_get_cache_logic`). Keys come from main's `CooldownCache`. The same literals sit in +/// `tests/unit/redis/stores/test_router_cooldowns.py`, checked against main's codec. +pub struct CooldownCase { + pub model_id: &'static str, + pub cooldown: Cooldown, + pub ttl_seconds: f64, + pub key: &'static str, + pub redis_ttl_seconds: u64, + pub str_dict: &'static str, + pub json: &'static str, +} + +pub fn cooldowns() -> [CooldownCase; 2] { + [ + CooldownCase { + model_id: "dep-a", + cooldown: Cooldown { + exception_received: "RateLimitError: limited é 'quoted'".to_owned(), + status_code: "429".to_owned(), + timestamp: Number::Float(1_730_000_000.25), + cooldown_time: Number::Int(60), + }, + ttl_seconds: 90.5, + key: "deployment:dep-a:cooldown", + redis_ttl_seconds: 90, + str_dict: "{'exception_received': \"RateLimitError: limited é 'quoted'\", 'status_code': '429', 'timestamp': 1730000000.25, 'cooldown_time': 60}", + json: "{\"exception_received\": \"RateLimitError: limited \\u00e9 'quoted'\", \"status_code\": \"429\", \"timestamp\": 1730000000.25, \"cooldown_time\": 60}", + }, + CooldownCase { + model_id: "dep-b", + cooldown: Cooldown { + exception_received: "Timeout".to_owned(), + status_code: "408".to_owned(), + timestamp: Number::Float(1_730_000_000.5), + cooldown_time: Number::Float(0.5), + }, + ttl_seconds: 30.0, + key: "deployment:dep-b:cooldown", + redis_ttl_seconds: 30, + str_dict: "{'exception_received': 'Timeout', 'status_code': '408', 'timestamp': 1730000000.5, 'cooldown_time': 0.5}", + json: "{\"exception_received\": \"Timeout\", \"status_code\": \"408\", \"timestamp\": 1730000000.5, \"cooldown_time\": 0.5}", + }, + ] +} + +/// main's allowed-fails keys for `dep-a`: an empty suffix is the same key as none. +pub struct AllowedFailsCase { + pub suffix: Option<&'static str>, + pub key: &'static str, +} + +pub const ALLOWED_FAILS: [AllowedFailsCase; 3] = [ + AllowedFailsCase { + suffix: None, + key: "deployment:dep-a:allowed_fails", + }, + AllowedFailsCase { + suffix: Some(""), + key: "deployment:dep-a:allowed_fails", + }, + AllowedFailsCase { + suffix: Some("timeout"), + key: "deployment:dep-a:allowed_fails:timeout", + }, +]; + +/// main's `RouterCacheEnum` usage keys, including the `None` a missing model prints into the key. +pub struct UsageCase { + pub counter: Counter, + pub key: &'static str, +} + +pub fn usage() -> [UsageCase; 2] { + [ + UsageCase { + counter: Counter { + deployment_id: "dep".to_owned(), + model: Some("anthropic/claude-sonnet-5-5".to_owned()), + minute: "12-34".to_owned(), + kind: Kind::Tpm, + }, + key: "global_router:dep:anthropic/claude-sonnet-5-5:tpm:12-34", + }, + UsageCase { + counter: Counter { + deployment_id: "dep".to_owned(), + model: None, + minute: "09-05".to_owned(), + kind: Kind::Rpm, + }, + key: "global_router:dep:None:rpm:09-05", + }, + ] +} + +/// 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) + && ttl > Duration::from_secs(whole_seconds.saturating_sub(5)) +} diff --git a/litellm-rust/crates/redis-testing/src/cluster.rs b/litellm-rust/crates/redis-testing/src/cluster.rs index b1373167050..1a8fe882a64 100644 --- a/litellm-rust/crates/redis-testing/src/cluster.rs +++ b/litellm-rust/crates/redis-testing/src/cluster.rs @@ -71,6 +71,15 @@ impl RedisCluster { self.ports.iter().map(|&port| (HOST, port)) } + /// One redis-cli reply per node, in `nodes()` order. + pub async fn cli_outputs(&self, args: &[&str]) -> Result, Error> { + let mut replies = Vec::with_capacity(self.ports.len()); + for port in self.ports { + replies.push(crate::exec::run(&self.container, crate::exec::cli(port, args)).await?); + } + Ok(replies) + } + pub async fn cli_on_every_node(&self, args: &[&str]) -> Result<(), Error> { for port in self.ports { crate::exec::run(&self.container, crate::exec::cli(port, args)).await?; diff --git a/litellm-rust/crates/redis-testing/src/sentinel.rs b/litellm-rust/crates/redis-testing/src/sentinel.rs index 928a892cf61..d91b4c7461e 100644 --- a/litellm-rust/crates/redis-testing/src/sentinel.rs +++ b/litellm-rust/crates/redis-testing/src/sentinel.rs @@ -7,7 +7,8 @@ use testcontainers_modules::testcontainers::{ }; use crate::{ - Error, + ClientConnection, Error, + deployment::parse_client_list, exec::{cli, run}, ports::{HOST, start_on_free_ports, wait_for_pong}, standalone::REDIS_TAG, @@ -71,6 +72,21 @@ impl RedisSentinel { (HOST, self.ports[2]) } + /// Runs one redis-cli command on the primary and the replica, which keep their own ACLs. + pub async fn cli_on_data_nodes(&self, args: &[&str]) -> Result<(), Error> { + for port in [self.ports[0], self.ports[1]] { + run(&self.container, cli(port, args)).await?; + } + Ok(()) + } + + pub async fn clients_on(&self, port: u16) -> Result, Error> { + Ok(parse_client_list( + port, + &run(&self.container, cli(port, &["CLIENT", "LIST"])).await?, + )) + } + pub fn primary_port(&self) -> u16 { self.ports[0] } diff --git a/litellm-rust/crates/redis-types/src/lib.rs b/litellm-rust/crates/redis-types/src/lib.rs index 5656803f612..2bac635f2e1 100644 --- a/litellm-rust/crates/redis-types/src/lib.rs +++ b/litellm-rust/crates/redis-types/src/lib.rs @@ -16,6 +16,9 @@ pub struct HostEndpoint { pub port: u16, pub tls: Option, pub db: i64, + /// Connect over this unix socket instead of host, port and TLS, like redis-py's + /// `UnixDomainSocketConnection` + pub unix_socket_path: Option, } #[derive(Clone, Debug, PartialEq)] @@ -42,6 +45,10 @@ pub struct SentinelEndpoint { #[derive(Clone, Debug, Default, PartialEq)] pub struct TlsOptions { pub ca_certs: Option, + /// PEM text trusted alongside `ca_certs`, redis-py's `ssl_ca_data` + pub ca_data: Option, + /// A directory whose files are PEM certificates to trust, redis-py's `ssl_ca_path` + pub ca_path: Option, pub certfile: Option, pub keyfile: Option, pub verify_certificate: bool, diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..fdb9e0cb16c 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -20,7 +20,7 @@ from contextvars import ContextVar, Token from dataclasses import dataclass, field from datetime import timedelta from types import MappingProxyType, TracebackType -from typing import Final, Generic, Protocol, TypeVar +from typing import TYPE_CHECKING, Final, Generic, Protocol, TypeAlias, TypeVar from litellm._internal_context import ( REDIS_FAMILIES_METADATA_KEY, @@ -38,7 +38,11 @@ from litellm.caching.redis_cache import ( from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.types.services import ServiceTypes +if TYPE_CHECKING: + from litellm.rust_bridge._native import NativeRequestRedis + _T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) _ScriptArg = str | bytes | int | float SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None] POST_CALL_FLUSH_DEADLINE_SECONDS: Final = 1.0 @@ -443,6 +447,37 @@ def _backend_key(redis_cache: RedisCache) -> object: return (type(redis_cache), redis_cache.namespace, settings) +class NativePending(Protocol[_T_co]): + def resolve(self) -> Awaitable[_T_co]: ... + + +class NativeBatchResult(Generic[_T]): + """Awaitable handle for a read declared on the request's native pre-call batch; awaiting it runs the + native flush hooks, flushes that batch, then answers from what it read.""" + + __slots__ = ("_pending", "_request") + + def __init__(self, request: RequestRedisBatches, pending: NativePending[_T]) -> None: + self._request: Final = request + self._pending: Final = pending + + def __await__(self) -> Generator[object, None, _T]: + return self._wait().__await__() + + async def _wait(self) -> _T: + await self._request.flush_native_pre() + return await self._pending.resolve() + + +NativeRequestFactory: TypeAlias = Callable[[float], "NativeRequestRedis"] + + +def _native_request(post_call_deadline: float) -> NativeRequestRedis: + from litellm.rust_bridge._native import NativeRequestRedis + + return NativeRequestRedis(post_call_deadline) + + _open_post_call: Final[weakref.WeakSet[RequestRedisBatches]] = weakref.WeakSet() """Requests whose post-call batch still holds declared ops, so a shutdown can send them before Redis goes away.""" @@ -460,15 +495,25 @@ class RequestRedisBatches: "_batches", "_deadline", "_deadline_flush", + "_native", + "_native_factory", + "_native_flush_hooks", "_post_call", "post_call_deadline", "prefetched", ) - def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None: + def __init__( + self, + post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS, + native_factory: NativeRequestFactory = _native_request, + ) -> None: self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self.post_call_deadline: Final = post_call_deadline + self._native_factory: Final = native_factory + self._native: NativeRequestRedis | None = None + self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry self._deadline: asyncio.TimerHandle | None = None self._deadline_flush: asyncio.Task[None] | None = None # Reads declared early for a consumer that runs later in the request, keyed by consumer name. @@ -495,20 +540,49 @@ class RequestRedisBatches: _open_post_call.add(self) return batch + def native(self) -> NativeRequestRedis: + """The request's native batches, built on first use so requests that touch no native store pay nothing.""" + if self._native is None: + self._native = self._native_factory(self.post_call_deadline) + return self._native + + def native_post_call(self) -> NativeRequestRedis: + """``native`` for a post-call declaration, so a shutdown drain also flushes it.""" + _open_post_call.add(self) + return self.native() + + def native_result(self, pending: NativePending[_T]) -> NativeBatchResult[_T]: + return NativeBatchResult(self, pending) + + def add_native_flush_hook(self, hook: Callable[[], None]) -> None: + """Called before every native pre-call flush so late readers can declare into the same trip.""" + self._native_flush_hooks.append(hook) + + async def flush_native_pre(self) -> None: + for hook in self._native_flush_hooks: + hook() + if self._native is not None: + await self._native.flush_pre() + def _flush_on_deadline(self) -> None: self._deadline = None self._deadline_flush = asyncio.ensure_future(self.flush_post_call()) async def flush_all(self) -> None: """Send whatever is still declared (write-backs nobody awaits) before the request scope closes.""" - await asyncio.gather(*(batch.flush() for batch in self._batches.values() if batch.pending)) + await asyncio.gather( + *(batch.flush() for batch in self._batches.values() if batch.pending), self.flush_native_pre() + ) async def flush_post_call(self) -> None: """One pipeline per backend for the post-call writes; the deadline is disarmed since this is that flush.""" if self._deadline is not None: self._deadline.cancel() self._deadline = None - await asyncio.gather(*(batch.flush() for batch in self._post_call.values() if batch.pending)) + await asyncio.gather( + *(batch.flush() for batch in self._post_call.values() if batch.pending), + *(() if self._native is None else (self._native.flush_post(),)), + ) if not any(batch.pending for batch in self._post_call.values()): _open_post_call.discard(self) @@ -557,17 +631,24 @@ async def drain_post_call_redis_batches() -> None: class request_redis_batch_scope: """Redis reads declared inside share one pipeline per backend; nested scopes join the outer one.""" - __slots__ = ("_post_call_deadline", "_token") + __slots__ = ("_native_factory", "_post_call_deadline", "_token") - def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None: + def __init__( + self, + post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS, + native_factory: NativeRequestFactory = _native_request, + ) -> None: self._token: Token[RequestRedisBatches | None] | None = None self._post_call_deadline: Final = post_call_deadline + self._native_factory: Final = native_factory def __enter__(self) -> RequestRedisBatches: outer: Final = _active_request_batches.get() if outer is not None: return outer - batches: Final = RequestRedisBatches(post_call_deadline=self._post_call_deadline) + batches: Final = RequestRedisBatches( + post_call_deadline=self._post_call_deadline, native_factory=self._native_factory + ) self._token = _active_request_batches.set(batches) return batches diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..66c9a1c6771 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -472,8 +472,14 @@ _LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet() def redis_breaker_state_counts() -> Mapping[str, int]: - """How many live Redis circuit breakers sit in each state, for the Prometheus services logger.""" - return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS))) + """How many live Redis circuit breakers sit in each state, for the Prometheus services logger. + + Counts the Python breakers of every live ``RedisCache`` and the breaker of every native Redis + client the router stores built, since each opens and closes on its own.""" + from litellm.rust_bridge._native import native_redis_breaker_states + + python_states: Final = Counter(breaker.state for breaker in _LIVE_BREAKERS) + return MappingProxyType(dict(python_states + Counter(native_redis_breaker_states()))) class RedisCircuitBreakerOpenError(Exception): diff --git a/litellm/redis/stores/_connection.py b/litellm/redis/stores/_connection.py new file mode 100644 index 00000000000..85781fa5033 --- /dev/null +++ b/litellm/redis/stores/_connection.py @@ -0,0 +1,14 @@ +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +from litellm.caching.redis_cache import RedisCache + +_KWARGS: Final = TypeAdapter(dict[str, object]) + + +def redis_kwargs(connection: RedisCache) -> Mapping[str, object]: + return _KWARGS.validate_python( + connection.redis_kwargs # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped legacy dict + ) diff --git a/litellm/redis/stores/_routing_prefetch.py b/litellm/redis/stores/_routing_prefetch.py deleted file mode 100644 index 1e8bac1ad3f..00000000000 --- a/litellm/redis/stores/_routing_prefetch.py +++ /dev/null @@ -1,149 +0,0 @@ -from __future__ import annotations - -import asyncio -import itertools -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from typing import TYPE_CHECKING, Final - -if TYPE_CHECKING: - from opentelemetry.trace import Span - -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 BatchResult, active_request_redis_batches - -_PREFETCH_SLOT: Final = "routing_read" - - -@dataclass(frozen=True, slots=True) -class TierRead: - cache: DualCache - keys: tuple[str, ...] - - -def selector_read(cache: DualCache, keys: Sequence[str]) -> TierRead: - return TierRead(cache, tuple(keys)) - - -def routing_target(reads: Sequence[TierRead]) -> str: - if len(reads) < 2 or not reads[1].keys: - return "router_cooldowns" - return "router_cooldowns_usage" if reads[0].keys else "router_usage" - - -async def _backfill(read: TierRead, values: Mapping[str, object]) -> None: - pending: Final = await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier - list(read.keys), local_only=True - ) - missed: Final = { - key: values[key] - for key, local in zip(read.keys, pending.result) - if local is None and values.get(key) is not None - } - await read.cache._apply_batch_get(pending, missed) # pyright: ignore[reportPrivateUsage] # store owns tier - - -@dataclass(frozen=True, slots=True) -class ArmedRead: - keys: frozenset[str] - fetched: frozenset[str] - result: BatchResult[Mapping[str, object]] - reservations: tuple[tuple[TierRead, Mapping[str, float | None]], ...] - - def release(self) -> None: - for read, previous in self.reservations: - read.cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # store owns tier - dict(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: - for read, _ in self.reservations: - await _backfill(read, future.result()) - except Exception: - self.release() - raise - - -def arm(reads: Sequence[TierRead]) -> None: - request: Final = active_request_redis_batches() - redis: Final = reads[0].cache.redis_cache if reads else None - if request is None or redis is None or _PREFETCH_SLOT in request.prefetched: - return - reserved: Final = tuple((read, read.cache.reserve_redis_batch_reads(read.keys)) for read in reads) - due_reads: Final = tuple(TierRead(read.cache, tuple(due)) for read, (due, _) in reserved) - due: Final = tuple(itertools.chain.from_iterable(read.keys for read in due_reads)) - if not due: - return - with service_target(routing_target(due_reads)): - result: Final = request.batch(redis).mget(due) - prefetch: Final = ArmedRead( - keys=frozenset(itertools.chain.from_iterable(read.keys for read in reads)), - fetched=frozenset(due), - result=result, - reservations=tuple((read, reservation[1]) for read, (_, reservation) in zip(due_reads, reserved)), - ) - result.on_settled(prefetch.settle) - request.prefetched[_PREFETCH_SLOT] = prefetch - - -def armed() -> bool: - request: Final = active_request_redis_batches() - return request is not None and _PREFETCH_SLOT in request.prefetched - - -def _take(needed: Sequence[str]) -> ArmedRead | 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 taken.keys.issuperset(needed): - return taken - if isinstance(taken, ArmedRead): - taken.release() - return None - - -async def _read_prefetched(reads: Sequence[TierRead]) -> Sequence[Sequence[object | None] | None] | None: - prefetch: Final = _take(tuple(itertools.chain.from_iterable(read.keys for read in reads))) - if prefetch is None: - return None - try: - values: Final = await prefetch.result - except Exception as e: # noqa: BLE001 # shared read applies the tiers' existing Redis fallback - verbose_router_logger.debug("routing prefetch failed, reading again: %s", e) - return None - pending_reads: Final = tuple( - [ - await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier - list(read.keys), local_only=True - ) - for read in reads - ] - ) - if any( - any(key not in prefetch.fetched for key, value in zip(read.keys, pending.result) if value is None) - for read, pending in zip(reads, pending_reads) - ): - return None - return [ - await read.cache._apply_batch_get( # pyright: ignore[reportPrivateUsage] # store owns tier - pending, {key: values.get(key) for key, value in zip(read.keys, pending.result) if value is None} - ) - for read, pending in zip(reads, pending_reads) - ] - - -async def read(reads: Sequence[TierRead], parent_otel_span: Span | None) -> Sequence[Sequence[object | None] | None]: - with service_target(routing_target(reads)): - prefetched: Final = await _read_prefetched(reads) - if prefetched is not None: - return prefetched - return await DualCache.async_batch_get_cache_shared( - tuple((read.cache, list(read.keys)) for read in reads), parent_otel_span=parent_otel_span - ) diff --git a/litellm/redis/stores/router_cooldowns.py b/litellm/redis/stores/router_cooldowns.py index d07f018056e..a57b928153b 100644 --- a/litellm/redis/stores/router_cooldowns.py +++ b/litellm/redis/stores/router_cooldowns.py @@ -1,9 +1,6 @@ from __future__ import annotations -import functools -import time -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol, TypeAlias if TYPE_CHECKING: @@ -11,14 +8,11 @@ if TYPE_CHECKING: from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import service_target -from litellm._logging import verbose_router_logger -from litellm.caching.dual_cache import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_batch import NativeBatchResult, NativePending, RequestRedisBatches from litellm.caching.redis_cache import RedisCache -from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE +from litellm.rust_bridge._native import NativeRequestRedis, NativeRouterCooldowns -from . import _routing_prefetch -from ._routing_prefetch import TierRead +from ._connection import redis_kwargs class CooldownCacheValue(TypedDict): @@ -28,156 +22,96 @@ class CooldownCacheValue(TypedDict): cooldown_time: ReadOnly[float] -@dataclass(frozen=True, slots=True) -class CooldownPolicy: - redis_read_interval: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS - backfill_ttl: int = 600 - max_memory_entries: int = 200 - max_batch_entries: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE - max_corrected_ttl: float = 60.0 - - -POLICY: Final = CooldownPolicy() ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns" ActiveCooldowns: TypeAlias = Sequence[tuple[str, CooldownCacheValue]] -@dataclass(frozen=True, slots=True) -class RoutingRead: - cooldowns: ActiveCooldowns - selector_values: Sequence[object | None] | None - - -@functools.lru_cache(maxsize=1024) -def _cooldown_key(model_id: str) -> str: - return "deployment:" + model_id + ":cooldown" - - class RouterCooldowns(Protocol): def attach_redis(self, connection: RedisCache | None) -> None: ... def flush_local_counters(self) -> None: ... def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: ... def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int: ... - def arm_routing_read(self, model_ids: Sequence[str], selector: TierRead | None) -> None: ... - async def routing_read( - self, model_ids: Sequence[str], selector: TierRead | None, parent_otel_span: Span | None - ) -> RoutingRead: ... async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ... + def declare_active( + self, request: RequestRedisBatches, model_ids: Sequence[str] + ) -> NativeBatchResult[ActiveCooldowns] | None: ... def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ... def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: ... def evict_local(self, model_id: str) -> None: ... + def snapshot(self) -> Mapping[str, CooldownCacheValue]: ... + + +class CooldownsHandle(Protocol): + def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ... + def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ... + def active(self, model_ids: Sequence[str]) -> Sequence[tuple[str, CooldownCacheValue]]: ... + def aactive(self, model_ids: Sequence[str]) -> Awaitable[Sequence[tuple[str, CooldownCacheValue]]]: ... + def declare_active( + self, request: NativeRequestRedis, model_ids: Sequence[str] + ) -> NativePending[ActiveCooldowns]: ... + def minimum(self, model_ids: Sequence[str]) -> float | None: ... + def evict_local(self, model_id: str) -> None: ... + def flush_local_counters(self) -> None: ... + def snapshot(self) -> Mapping[str, CooldownCacheValue]: ... + + +HandleFactory: TypeAlias = Callable[[RedisCache | None], CooldownsHandle] + + +def native_handle(connection: RedisCache | None) -> CooldownsHandle: + if connection is None: + return NativeRouterCooldowns(None) + return NativeRouterCooldowns(redis_kwargs(connection), namespace=connection.namespace) class RouterCooldownStore: - def __init__(self, connection: RedisCache | None, *, clock: Callable[[], float] | None = None) -> None: - self._clock: Final = clock if clock is not None else lambda: time.time() - self._memory: Final = InMemoryCache( - max_size_in_memory=POLICY.max_memory_entries, default_ttl=POLICY.backfill_ttl, clock=self._clock - ) - self._cache: Final = DualCache( - redis_cache=connection, - in_memory_cache=self._memory, - default_redis_batch_cache_expiry=POLICY.redis_read_interval, - default_max_redis_batch_cache_size=POLICY.max_batch_entries, - clock=self._clock, - ) - self._fails: Final = DualCache(redis_cache=connection, in_memory_cache=InMemoryCache(clock=self._clock)) + def __init__(self, connection: RedisCache | None, *, handle: HandleFactory = native_handle) -> None: + self._connected = connection is not None + self._build: Final = handle + self._handle = handle(connection) def attach_redis(self, connection: RedisCache | None) -> None: - self._cache.attach_redis_cache(connection) - self._fails.attach_redis_cache(connection) + if connection is None or self._connected: + return + self._connected = True + self._handle = self._build(connection) def flush_local_counters(self) -> None: - self._fails.in_memory_cache.flush_cache() + self._handle.flush_local_counters() def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: with service_target(ROUTER_COOLDOWNS_TARGET): - self._cache.set_cache(key=_cooldown_key(model_id), value=value, ttl=ttl) + self._handle.put(model_id, value, ttl) def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int: - base: Final = f"deployment:{model_id}:allowed_fails" - key: Final = f"{base}:{suffix}" if suffix else base with service_target(ROUTER_COOLDOWNS_TARGET): - try: - return self._fails.increment_cache(key=key, value=1, ttl=ttl) - except Exception as e: # noqa: BLE001 # preserve worker-local counting during a Redis outage - verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e) - local: Final[object] = self._fails.get_cache(key=key, local_only=True) - return local if isinstance(local, int) else 0 - - def _tier_read(self, model_ids: Sequence[str]) -> TierRead: - return TierRead(self._cache, tuple(_cooldown_key(model_id) for model_id in model_ids)) - - def arm_routing_read(self, model_ids: Sequence[str], selector: TierRead | None) -> None: - _routing_prefetch.arm((self._tier_read(model_ids), *(() if selector is None else (selector,)))) - - async def routing_read( - self, model_ids: Sequence[str], selector: TierRead | None, parent_otel_span: Span | None - ) -> RoutingRead: - reads: Final = (self._tier_read(model_ids), *(() if selector is None else (selector,))) - results: Final = await _routing_prefetch.read(reads, parent_otel_span) - return RoutingRead( - cooldowns=self._active_from_results(model_ids, results[0]), - selector_values=None if selector is None else results[1], - ) - - def _corrected_active_cooldown( - self, model_id: str, result: Mapping[str, object], current_time: float - ) -> CooldownCacheValue | None: - value: Final = CooldownCacheValue( - **result # pyright: ignore[reportArgumentType] # preserve main's unvalidated dict codec - ) - remaining: Final = value["timestamp"] + value["cooldown_time"] - current_time - key: Final = _cooldown_key(model_id) - if remaining <= 0: - self._memory.delete_cache(key) - return None - expiry: Final = self._memory.ttl_dict.get(key) - if expiry is not None and expiry > current_time + remaining + 5: - self._memory.delete_cache(key) - self._memory.set_cache(key, result, ttl=min(remaining, POLICY.max_corrected_ttl)) - return value - - def _active_from_results(self, model_ids: Sequence[str], results: Sequence[object] | None) -> ActiveCooldowns: - if results is None: - return () - now: Final = self._clock() - return tuple( - (model_id, active) - for model_id, result in zip(model_ids, results) - if isinstance(result, dict) - and result - and (active := self._corrected_active_cooldown(model_id, result, now)) is not None - ) + return self._handle.increment_allowed_fails(model_id, suffix, ttl) async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: with service_target(ROUTER_COOLDOWNS_TARGET): - results: Final = await self._cache.async_batch_get_cache( - keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span - ) - return self._active_from_results(model_ids, results) + pending: Final = self._handle.aactive(model_ids) + return tuple(await pending) - def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: + def declare_active( + self, request: RequestRedisBatches, model_ids: Sequence[str] + ) -> NativeBatchResult[ActiveCooldowns] | None: + """``active`` declared on the request's native pre-call batch; None when there is no Redis to read.""" + if not self._connected: + return None with service_target(ROUTER_COOLDOWNS_TARGET): - results: Final = ( - self._cache.batch_get_cache( - keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span - ) - or () - ) - durations: Final = tuple( - CooldownCacheValue(**result)["cooldown_time"] # pyright: ignore[reportArgumentType] # preserve main's dict codec - for result in results - if isinstance(result, dict) and result - ) - return (min(durations) if durations else None) or default + pending: Final = self._handle.declare_active(request.native(), model_ids) + return request.native_result(pending) def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: with service_target(ROUTER_COOLDOWNS_TARGET): - results: Final = self._cache.batch_get_cache( - keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span - ) - return self._active_from_results(model_ids, results) + return tuple(self._handle.active(model_ids)) + + def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: + with service_target(ROUTER_COOLDOWNS_TARGET): + return self._handle.minimum(model_ids) or default def evict_local(self, model_id: str) -> None: - self._memory.delete_cache(_cooldown_key(model_id)) + self._handle.evict_local(model_id) + + def snapshot(self) -> Mapping[str, CooldownCacheValue]: + return self._handle.snapshot() diff --git a/litellm/redis/stores/router_usage.py b/litellm/redis/stores/router_usage.py index 623e600433d..3d4309053d2 100644 --- a/litellm/redis/stores/router_usage.py +++ b/litellm/redis/stores/router_usage.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias @@ -8,27 +8,12 @@ if TYPE_CHECKING: from opentelemetry.trace import Span from litellm._internal_context import with_service_target -from litellm.caching.dual_cache import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_batch import active_request_redis_batches from litellm.caching.redis_cache import RedisCache -from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE -from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.rust_bridge._native import NativeRequestRedis, NativeRouterUsage -from ._keys import ROUTER_RPM_KEY, ROUTER_TPM_KEY +from ._connection import redis_kwargs - -@dataclass(frozen=True, slots=True) -class UsagePolicy: - redis_read_interval: float | None = None - memory_ttl: float | None = None - redis_ttl: float | None = None - backfill_ttl: int = 600 - max_memory_entries: int = 200 - max_batch_entries: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE - counter_ttl: int = 60 - - -POLICY: Final = UsagePolicy() ROUTER_USAGE_TARGET: Final = "router_usage" CounterKind: TypeAlias = Literal["tpm", "rpm"] @@ -41,11 +26,6 @@ class UsageCounter: kind: CounterKind -def _key(counter: UsageCounter) -> str: - template: Final = ROUTER_TPM_KEY if counter.kind == "tpm" else ROUTER_RPM_KEY - return template.format(id=counter.deployment_id, model=counter.model, current_minute=counter.minute) - - class RouterUsage(Protocol): def attach_redis(self, connection: RedisCache | None) -> None: ... def flush_local_counters(self) -> None: ... @@ -62,34 +42,62 @@ class RouterUsage(Protocol): parent_otel_span: Span | None, ) -> None: ... async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None: ... + def snapshot(self) -> Mapping[str, float]: ... + + +NativeCounter: TypeAlias = tuple[str, str | None, str, CounterKind] + + +class UsageHandle(Protocol): + def increment(self, counter: NativeCounter, value: float, ttl_seconds: int) -> Awaitable[float]: ... + def increment_deployment( + self, + deployment_id: str, + model: str | None, + minute: str, + tokens: float, + requests: int, + *, + request: NativeRequestRedis | None = None, + ) -> Awaitable[None]: ... + def get_many(self, counters: Sequence[NativeCounter]) -> Awaitable[Sequence[float | None]]: ... + def flush_local_counters(self) -> None: ... + def snapshot(self) -> Mapping[str, float]: ... + + +HandleFactory: TypeAlias = Callable[[RedisCache | None], UsageHandle] + + +def native_handle(connection: RedisCache | None) -> UsageHandle: + if connection is None: + return NativeRouterUsage(None) + return NativeRouterUsage(redis_kwargs(connection), namespace=connection.namespace) + + +def _native(counter: UsageCounter) -> NativeCounter: + return (counter.deployment_id, counter.model, counter.minute, counter.kind) class RouterUsageStore: - def __init__(self, connection: RedisCache | None) -> None: - self._cache: Final = DualCache( - redis_cache=connection, - in_memory_cache=InMemoryCache( - max_size_in_memory=POLICY.max_memory_entries, default_ttl=POLICY.backfill_ttl - ), - default_in_memory_ttl=POLICY.memory_ttl, - default_redis_ttl=POLICY.redis_ttl, - default_redis_batch_cache_expiry=POLICY.redis_read_interval, - default_max_redis_batch_cache_size=POLICY.max_batch_entries, - ) + def __init__(self, connection: RedisCache | None, *, handle: HandleFactory = native_handle) -> None: + self._connected = connection is not None + self._build: Final = handle + self._handle = handle(connection) def attach_redis(self, connection: RedisCache | None) -> None: - self._cache.attach_redis_cache(connection) + if connection is None or self._connected: + return + self._connected = True + self._handle = self._build(connection) def flush_local_counters(self) -> None: - self._cache.in_memory_cache.flush_cache() + self._handle.flush_local_counters() @with_service_target(ROUTER_USAGE_TARGET) async def async_increment( self, counter: UsageCounter, value: float, ttl: int, parent_otel_span: Span | None = None ) -> float | None: - return await self._cache.async_increment_cache( - key=_key(counter), value=value, ttl=ttl, parent_otel_span=parent_otel_span - ) + return await self._handle.increment(_native(counter), value, ttl) @with_service_target(ROUTER_USAGE_TARGET) async def increment_deployment( @@ -101,18 +109,18 @@ class RouterUsageStore: requests: int, parent_otel_span: Span | None, ) -> None: - tpm: Final = UsageCounter(deployment_id, model, minute, "tpm") - rpm: Final = UsageCounter(deployment_id, model, minute, "rpm") - operations: Final = [ - RedisPipelineIncrementOperation(key=_key(counter), increment_value=value, ttl=POLICY.counter_ttl) - for counter, value in ((tpm, tokens), (rpm, requests)) - ] - values: Final = await self._cache.async_increment_cache_pipeline( - increment_list=operations, parent_otel_span=parent_otel_span + request: Final = active_request_redis_batches() + await self._handle.increment_deployment( + deployment_id, + model, + minute, + tokens, + requests, + request=None if request is None else request.native_post_call(), ) - if values is not None and self._cache.redis_cache is not None: - for operation, value in zip(operations, values): - await self._cache.async_set_cache(operation["key"], int(value), local_only=True, ttl=POLICY.counter_ttl) async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None: - return await self._cache.async_batch_get_cache(keys=[_key(counter) for counter in counters]) + return await self._handle.get_many(tuple(map(_native, counters))) + + def snapshot(self) -> Mapping[str, float]: + return self._handle.snapshot() diff --git a/litellm/router_utils/routing_prefetch.py b/litellm/router_utils/routing_prefetch.py new file mode 100644 index 00000000000..2b5425574c7 --- /dev/null +++ b/litellm/router_utils/routing_prefetch.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from opentelemetry.trace import Span + +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 BatchResult, active_request_redis_batches +from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, ActiveCooldowns, RouterCooldowns + +_PREFETCH_SLOT: Final = "routing_read" +_COOLDOWNS_SLOT: Final = "routing_cooldowns" +ROUTER_USAGE_TARGET: Final = "router_usage" + + +@dataclass(frozen=True, slots=True) +class TierRead: + cache: DualCache + keys: tuple[str, ...] + + +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: + 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: + 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 + + +@dataclass(frozen=True, slots=True) +class ArmedCooldowns: + model_ids: frozenset[str] + result: Awaitable[ActiveCooldowns] + + +def arm_cooldowns(store: RouterCooldowns, model_ids: Sequence[str]) -> None: + request: Final = active_request_redis_batches() + if request is None or _COOLDOWNS_SLOT in request.prefetched: + return + result: Final = store.declare_active(request, model_ids) + if result is not None: + request.prefetched[_COOLDOWNS_SLOT] = ArmedCooldowns(frozenset(model_ids), result) + + +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) + + +async def _armed_cooldowns(model_ids: Sequence[str]) -> ActiveCooldowns | None: + request: Final = active_request_redis_batches() + taken: Final = None if request is None else request.prefetched.pop(_COOLDOWNS_SLOT, None) + if not isinstance(taken, ArmedCooldowns): + return None + try: + active: 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("cooldown prefetch failed, reading again: %s", e) + return None + if not taken.model_ids.issuperset(model_ids): + return None + wanted: Final = frozenset(model_ids) + return tuple((model_id, cooldown) for model_id, cooldown in active if model_id in wanted) + + +async def read_cooldowns( + store: RouterCooldowns, model_ids: Sequence[str], parent_otel_span: Span | None +) -> ActiveCooldowns: + """The armed cooldown read when it covers ``model_ids``, taken once; otherwise the store reads on its own.""" + with service_target(ROUTER_COOLDOWNS_TARGET): + prefetched: Final = await _armed_cooldowns(model_ids) + if prefetched is not None: + return prefetched + return await store.active(model_ids, parent_otel_span) + + +def _take(read: TierRead) -> ArmedRead | 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: + 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 + 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): + 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} + ) + + +async def read(read: TierRead, parent_otel_span: Span | None) -> Sequence[object | None] | None: + 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] diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index c0bbaa0c859..7622991a78d 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -1,21 +1,21 @@ -from collections.abc import Iterator +import asyncio +from collections.abc import Iterator, Sequence from contextlib import contextmanager from contextvars import ContextVar from typing import TYPE_CHECKING, Final from litellm._logging import verbose_router_logger -from litellm.redis.stores import _routing_prefetch -from litellm.redis.stores._routing_prefetch import TierRead, selector_read 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 ROUTER_USAGE_TARGET, TierRead, selector_read if TYPE_CHECKING: from opentelemetry.trace import Span from litellm.router_api import RoutingContext - -ROUTER_USAGE_TARGET: Final = "router_usage" +__all__ = ("ROUTER_USAGE_TARGET", "RoutingPrefetch", "RoutingReadBatch") def _selector_keys(selector: LowestTPMLoggingHandler_v2, deployments: list) -> list[str]: @@ -27,17 +27,22 @@ def _selector_read(selector: LowestTPMLoggingHandler_v2 | None, deployments: lis return None if selector is None else selector_read(selector.router_cache, _selector_keys(selector, deployments)) +async def _nothing() -> Sequence[object | None] | None: + return None + + class RoutingPrefetch: @staticmethod @renamed_router_param("litellm_router_instance") def arm(ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: list) -> None: - ctx.health().cooldowns.cooldown_store.arm_routing_read( - ctx.registry().names(kind="ids"), _selector_read(usage_selector, deployments) - ) + routing_prefetch.arm_cooldowns(ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids")) + read: Final = _selector_read(usage_selector, deployments) + if read is not None: + routing_prefetch.arm(read) @staticmethod def armed() -> bool: - return _routing_prefetch.armed() + return routing_prefetch.armed() _active_routing_read_batch: Final[ContextVar["RoutingReadBatch | None"]] = ContextVar( @@ -74,13 +79,16 @@ class RoutingReadBatch: self, ctx: "RoutingContext", healthy_deployments: list, parent_otel_span: "Span | None" ) -> list[str]: usage_read: Final = _selector_read(self.usage_selector, healthy_deployments) - read: Final = await ctx.health().cooldowns.cooldown_store.routing_read( - ctx.registry().names(kind="ids"), usage_read, parent_otel_span + cooldowns, selector_values = await asyncio.gather( + routing_prefetch.read_cooldowns( + ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids"), parent_otel_span + ), + _nothing() if usage_read is None else routing_prefetch.read(usage_read, parent_otel_span), ) if usage_read is not None: self.prefetched_usage = PrefetchedUsage( keys=frozenset(usage_read.keys), - values=None if read.selector_values is None else dict(zip(usage_read.keys, read.selector_values)), + values=None if selector_values is None else dict(zip(usage_read.keys, selector_values)), ) - verbose_router_logger.debug("retrieve cooldown models: %s", read.cooldowns) - return [model_id for model_id, _ in read.cooldowns] + verbose_router_logger.debug("retrieve cooldown models: %s", cooldowns) + return [model_id for model_id, _ in cooldowns] diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8177e4a5cc9..c26596bc019 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,6 +1,6 @@ from asyncio import Future from collections.abc import Mapping, Sequence -from typing import Never, final +from typing import Literal, Never, TypedDict, final import httpx from pydantic import JsonValue @@ -15,6 +15,7 @@ class ProcessReservedForForking(RuntimeError): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... +def native_redis_breaker_states() -> list[str]: ... @final class NativeClickHouseSpendConfig: @@ -26,6 +27,61 @@ class NativeClickHouseSpendStorage: def ensure_schema(self) -> Future[None]: ... def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... +class _NativeCooldown(TypedDict): + exception_received: str + status_code: str + timestamp: float + cooldown_time: float + +@final +class NativeRequestRedis: + def __new__(cls, post_call_deadline_seconds: float) -> NativeRequestRedis: ... + def flush_pre(self) -> Future[None]: ... + def flush_post(self) -> Future[None]: ... + +@final +class NativeDeclaredCooldowns: + def resolve(self) -> Future[list[tuple[str, _NativeCooldown]]]: ... + +@final +class NativeRouterCooldowns: + def __new__( + cls, redis_kwargs: Mapping[str, object] | None, *, namespace: str | None = None + ) -> NativeRouterCooldowns: ... + def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ... + def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ... + def active(self, model_ids: Sequence[str]) -> list[tuple[str, _NativeCooldown]]: ... + def aactive(self, model_ids: Sequence[str]) -> Future[list[tuple[str, _NativeCooldown]]]: ... + def declare_active(self, request: NativeRequestRedis, model_ids: Sequence[str]) -> NativeDeclaredCooldowns: ... + def minimum(self, model_ids: Sequence[str]) -> float | None: ... + def evict_local(self, model_id: str) -> None: ... + def flush_local_counters(self) -> None: ... + def snapshot(self) -> dict[str, _NativeCooldown]: ... + +@final +class NativeRouterUsage: + def __new__( + cls, redis_kwargs: Mapping[str, object] | None, *, namespace: str | None = None + ) -> NativeRouterUsage: ... + def increment( + self, counter: tuple[str, str | None, str, Literal["tpm", "rpm"]], value: float, ttl_seconds: int + ) -> Future[float]: ... + def increment_deployment( + self, + deployment_id: str, + model: str | None, + minute: str, + tokens: float, + requests: int, + *, + request: NativeRequestRedis | None = None, + ) -> Future[None]: ... + def get_many( + self, counters: Sequence[tuple[str, str | None, str, Literal["tpm", "rpm"]]] + ) -> Future[list[float | None]]: ... + def flush_local_counters(self) -> None: ... + def snapshot(self) -> dict[str, float]: ... + @final class NativeDatabase: @staticmethod @@ -215,7 +271,11 @@ __all__ = [ "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", + "NativeDeclaredCooldowns", "NativeDiagnosticProcessor", + "NativeRequestRedis", + "NativeRouterCooldowns", + "NativeRouterUsage", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", @@ -233,6 +293,7 @@ __all__ = [ "embedding", "gil_stats", "messages", + "native_redis_breaker_states", "ocr", "ocr_health_check_document", "ocr_passthrough_response", diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 9b751a1ada2..d4cb25cb4fd 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -2,7 +2,7 @@ Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name -- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds) +- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`) - A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here - Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers - The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green diff --git a/litellm/rust_bridge/host/redis_events.py b/litellm/rust_bridge/host/redis_events.py new file mode 100644 index 00000000000..202d43cd4fd --- /dev/null +++ b/litellm/rust_bridge/host/redis_events.py @@ -0,0 +1,46 @@ +"""Redis call and circuit breaker events from the native Redis client, reported the way ``RedisCache`` reports them.""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Final, Literal, TypeAlias + +from litellm._service_logger import ServiceLogging +from litellm.types.services import ServiceTypes + +Failure: TypeAlias = Literal["timeout", "transport", "server", "decode", "breaker_open"] +BreakerState: TypeAlias = Literal["closed", "open", "half_open"] + +_HEALTH_FAILURE_CLASSES: Final = {"timeout": "timeout", "transport": "connectivity"} + + +class RedisCallError(Exception): + """The failure class of a native Redis call, the only detail the Rust client hands over.""" + + +@lru_cache(maxsize=1) +def _service_logging() -> ServiceLogging: + return ServiceLogging() + + +def record_redis_call(call_type: str, duration: float, failure: Failure | None, caller: str | None) -> None: + if failure == "breaker_open": + return + logger: Final = _service_logging() + if failure is None: + logger.service_success_hook(service=ServiceTypes.REDIS, duration=duration, call_type=call_type, caller=caller) + return + logger.service_failure_hook( + service=ServiceTypes.REDIS, + duration=duration, + error=RedisCallError(failure), + call_type=call_type, + caller=caller, + ) + failure_class: Final = _HEALTH_FAILURE_CLASSES.get(failure) + if failure_class is not None: + logger.redis_breaker_hook(transitions_to=None, failure_class=failure_class) + + +def record_redis_breaker(state: BreakerState) -> None: + _service_logging().redis_breaker_hook(transitions_to=state, failure_class=None) diff --git a/litellm/rust_bridge/host/settings.py b/litellm/rust_bridge/host/settings.py index 60a58f1aaba..4932a788221 100644 --- a/litellm/rust_bridge/host/settings.py +++ b/litellm/rust_bridge/host/settings.py @@ -31,6 +31,15 @@ class ProviderDefaults: enable_azure_ad_token_refresh: object +@dataclass(frozen=True, slots=True) +class RouterTierSettings: + """The memory-over-Redis tier values `main`'s router reads from the environment and `litellm` globals.""" + + cooldown_redis_read_interval: float + usage_redis_read_interval: float + max_redis_batch_cache_size: int + + @dataclass(frozen=True, slots=True) class SecretManager: readable: bool @@ -108,6 +117,17 @@ def secret_manager_binding() -> SecretManagerBinding: ) +def router_tiers() -> RouterTierSettings: + import litellm + from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE + + return RouterTierSettings( + cooldown_redis_read_interval=DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, + usage_redis_read_interval=litellm.default_redis_batch_cache_expiry or 10, + max_redis_batch_cache_size=DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, + ) + + def provider_defaults() -> ProviderDefaults: import litellm diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index cd270035416..7745bafe1b5 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -6,8 +6,9 @@ import asyncio import hashlib import json from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass, field from datetime import timedelta -from typing import Any +from typing import Any, Generic, TypeVar import pytest from redis.exceptions import NoScriptError @@ -28,6 +29,7 @@ from litellm.caching.redis_cache import ( from litellm.caching.redis_cluster_cache import RedisClusterCache SCRIPT = "return redis.call('GET', KEYS[1])" +_Answer = TypeVar("_Answer") SHA = hashlib.sha1(SCRIPT.encode()).hexdigest() # noqa: S324 @@ -93,6 +95,7 @@ class FakeRedisCache(RedisCache): def __init__(self, client: FakeClient, namespace: str | None = None) -> None: # super().__init__ needs a server self.client = client self.namespace = namespace + self.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2, "db": id(self) % 1_000_000} self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) self.service_logger_obj = ServiceLogging() self.default_ttl = None @@ -147,6 +150,43 @@ def replies(command: tuple[Any, ...]) -> Any: raise AssertionError(command) +@dataclass +class FakeNativeRequest: + """Stands in for ``NativeRequestRedis``: records each flush in ``events`` next to what the test's pendings log.""" + + post_call_deadline: float + events: list[str] = field(default_factory=list) + + async def flush_pre(self) -> None: + self.events.append("flush_pre") + + async def flush_post(self) -> None: + self.events.append("flush_post") + + +@dataclass +class FakeNatives: + built: list[FakeNativeRequest] = field(default_factory=list) + + def __call__(self, post_call_deadline: float) -> FakeNativeRequest: + self.built.append(FakeNativeRequest(post_call_deadline)) + return self.built[-1] + + +@dataclass +class FakeNativePending(Generic[_Answer]): + """A read declared on a ``FakeNativeRequest``; it only answers once that request's pre-call batch was flushed.""" + + request: FakeNativeRequest + answer: Callable[[], _Answer] + + async def resolve(self) -> _Answer: + if "flush_pre" not in self.request.events: + raise ConnectionError("resolved before its batch was flushed") + self.request.events.append("resolve") + return self.answer() + + def make(fail: Exception | None = None, namespace: str | None = None) -> tuple[FakeRedisCache, FakeClient]: client = FakeClient(replies, fail) return FakeRedisCache(client, namespace), client 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 4b73bb6dd34..b3932285220 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -38,7 +38,7 @@ from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ModelResponse -from .test_redis_batch import FakeClient, FakeRedisCache +from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache async def _script_outside_the_pipeline(keys: Sequence[str], args: Sequence[object]) -> object: @@ -514,3 +514,31 @@ async def test_a_cancelled_post_call_flush_keeps_the_shared_spend_counter_and_co assert redis_cache.alone == [], "a cancel says nothing about the shared counter, so Redis keeps it" assert spend_cache.in_memory_cache.get_cache("spend:key:k1") == 3.5, "the local copy counts the cancelled spend" assert spend_cache.in_memory_cache.get_cache("spend:team:t1") is None, "an absent local copy is not seeded" + + +@pytest.mark.asyncio +async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary(): + natives = FakeNatives() + + with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double + request.native_post_call() + await request.flush_all() + at_the_boundary = list(natives.built[0].events) + await flush_post_call_redis_batches() + await drain_post_call_redis_batches() + + assert at_the_boundary == ["flush_pre"] + assert natives.built[0].events == ["flush_pre", "flush_post"] + + +@pytest.mark.asyncio +async def test_native_post_call_ops_still_waiting_on_their_callbacks_are_drained_at_shutdown(): + natives = FakeNatives() + + with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double + request.native_post_call() + await request.flush_all() + await drain_post_call_redis_batches() + await drain_post_call_redis_batches() + + assert natives.built[0].events == ["flush_pre", "flush_post"] 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 5ad521fad72..31878658049 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -1,13 +1,13 @@ """One Redis pipeline per backend for the pre-call reads a request makes: rate limiter Lua groups, the -router's cooldown and usage read, auth identity and spend counters all join the request batch.""" +router's usage counter read, auth identity and spend counters all join the request batch.""" from __future__ import annotations import asyncio import hashlib import json +from collections.abc import Callable from datetime import datetime, timezone -from itertools import chain from typing import Any, Final from unittest.mock import AsyncMock, MagicMock @@ -24,21 +24,19 @@ from litellm.proxy.auth.auth_object_prefetch import _CacheEntry, _write_back, pr from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.parallel_request_limiter_v3 import ( CHECK_AND_INCREMENT_BY_N_SCRIPT, + PROXY_MaxParallelRequestsHandler_v3, RateLimitDescriptor, RateLimitUnverifiableError, - PROXY_MaxParallelRequestsHandler_v3, ) from litellm.proxy.utils import InternalUsageCache from litellm.router_api import RoutingContext -from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET -from litellm.redis.stores.router_cooldowns import POLICY -from litellm.redis.stores._routing_prefetch import ArmedRead, TierRead, routing_target +from litellm.router_utils.routing_prefetch import ArmedRead, TierRead from litellm.router_utils.routing_read_batch import ( ROUTER_USAGE_TARGET, RoutingPrefetch, ) -from .test_redis_batch import FakeClient, FakeRedisCache, replies +from .test_redis_batch import FakeClient, FakeNativePending, FakeNatives, FakeRedisCache, replies def _cooldown_key(model_id: str) -> str: @@ -294,12 +292,76 @@ def _router(redis_cache: FakeRedisCache, routing_strategy: str = "usage-based-ro return router -@pytest.mark.asyncio -async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issues_no_read_of_its_own(monkeypatch): - monkeypatch.setattr( - "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", - lambda: datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc), +_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 + + +def _python_cooldown_reads(client: FakeClient, redis_cache: FakeRedisCache) -> tuple[str, ...]: + pipelined: Final = tuple( + key + for pipeline in client.pipelines + for command in pipeline.commands + if command[0] == "MGET" + for key in command[1:] + ) + alone: Final = tuple(key for command, keys in redis_cache.alone if command == "MGET" for key in keys) + 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={} + ) + return deployment["model_info"]["id"] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("pinned_clock") +async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issues_no_read_of_its_own(): client = FakeClient(_lua_ok_replies) redis_cache = FakeRedisCache(client) router = _router(redis_cache) @@ -311,55 +373,40 @@ async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issue descriptors=[_descriptor("api_key", "k1", 10), _descriptor("team", "t1", 20)], increments=[{"requests": 1}, {"requests": 1}], ) - deployment = await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + deployment = await _route(router) - assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} + 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", "EVALSHA", "EVALSHA"] - mget_keys = set(commands[0][1:]) - assert {_cooldown_key("dep-a"), _cooldown_key("dep-b")} <= mget_keys - assert any(":tpm:" in key for key in mget_keys) and any(":rpm:" in key for key in mget_keys) + assert frozenset(commands[0][1:]) == _counter_keys(router) assert redis_cache.alone == [] @pytest.mark.asyncio -async def test_a_cooldown_recorded_locally_after_the_prefetch_left_still_excludes_its_deployment(): - expired = {"exception_received": "429", "status_code": "429", "timestamp": 0.0, "cooldown_time": 60} - - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": # Redis holds a stale cooldown for dep-b and nothing for dep-a - return [json.dumps(expired) if key == _cooldown_key("dep-b") else None for key in command[1:]] - return _lua_ok_replies(command) - - client = FakeClient(replies) +@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing-v2"]) +async def test_a_cooldown_recorded_through_the_store_excludes_its_deployment_without_a_python_cooldown_read( + routing_strategy: str, +): + client = FakeClient(_lua_ok_replies) redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - cooldown_store = router.cooldown_cache.cooldown_store + router = _router(redis_cache, routing_strategy=routing_strategy) + router.cooldown_cache.cooldown_store.put( + "dep-a", + {"exception_received": "429", "status_code": "429", "timestamp": _FAR_FUTURE, "cooldown_time": 60}, + 60, + ) with request_redis_batch_scope(): router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - cooldown_store.put( - "dep-a", - {"exception_received": "429", "status_code": "429", "timestamp": _FAR_FUTURE, "cooldown_time": 60}, - 60, - ) - picks = { - ( - await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) - )["model_info"]["id"] - for _ in range(5) - } + picks = {await _route(router) for _ in range(5)} assert picks == {"dep-b"} - assert len(client.pipelines) == 1 + 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) @@ -370,79 +417,56 @@ async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_ro armed = request.prefetched["routing_read"] assert isinstance(armed, ArmedRead) request.prefetched["routing_read"] = ArmedRead( - keys=frozenset({"other"}), + read=TierRead(armed.read.cache, ("other",)), fetched=armed.fetched, result=armed.result, - reservations=armed.reservations, - ) - deployment = await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + previous=armed.previous, ) + deployment = await _route(router) assert request.prefetched == {} - assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} - assert len(redis_cache.alone) == 1 # the shared cooldown+usage read, one round trip as in P1 + 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 -async def test_a_prefetch_with_incomplete_usage_keys_releases_cooldown_reservations(): +@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) - cooldown_keys: Final = frozenset( - { - _cooldown_key("dep-a"), - _cooldown_key("dep-b"), - } - ) with request_redis_batch_scope(): RoutingPrefetch.arm(RoutingContext(router), router.lowesttpm_logger_v2, router.model_list[:1]) - deployment: Final = await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + deployment: Final = await _route(router) - fallback_cooldown_mgets: Final = tuple( - keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_keys.issubset(keys) - ) - - assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} - assert len(fallback_cooldown_mgets) == 1 + 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 -async def test_a_failed_prefetch_falls_back_to_the_shared_read(): +@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 router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + deployment = await _route(router) - assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} - assert len(redis_cache.alone) == 1 - cooldown_keys: Final = frozenset( - { - _cooldown_key("dep-a"), - _cooldown_key("dep-b"), - } - ) - fallback_cooldown_mgets: Final = tuple( - keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_keys.issubset(keys) - ) - assert len(fallback_cooldown_mgets) == 1 + 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 -async def test_the_armed_routing_read_is_declared_under_the_router_cooldowns_family(): +@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_cooldowns`` instead of a bare ``redis.pipeline``.""" + 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, routing_strategy="simple-shuffle") + router = _router(redis_cache) pipeline_targets: list[str | None] = [] # mutable-ok: filled by the recording hook async def record(**kwargs: object) -> None: @@ -452,35 +476,21 @@ async def test_the_armed_routing_read_is_declared_under_the_router_cooldowns_fam with request_redis_batch_scope(): router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) - await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task())) + await _route(router) + async with asyncio.timeout(5): + while not pipeline_targets: + await asyncio.sleep(0.01) - assert pipeline_targets == [ROUTER_COOLDOWNS_TARGET] + assert pipeline_targets == [ROUTER_USAGE_TARGET] @pytest.mark.asyncio -async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkeypatch): - clock: Final = 1_000_000.0 - monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock) - cooldown_key: Final = _cooldown_key("dep-a") - active_cooldown: Final = { - "exception_received": "429", - "status_code": "429", - "timestamp": _FAR_FUTURE, - "cooldown_time": 60, - } - - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": - return [json.dumps(active_cooldown) if key == cooldown_key else None for key in command[1:]] - return _lua_ok_replies(command) - - client: Final = FakeClient(replies) +@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) - redis_cache.store[cooldown_key] = active_cooldown - router: Final = _router(redis_cache, routing_strategy="simple-shuffle") + router: Final = _router(redis_cache) + client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100}) limiter: Final = _limiter(redis_cache) with request_redis_batch_scope(): @@ -489,176 +499,66 @@ async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkey descriptors=[_descriptor("api_key", "k1", 10)], increments=[{"requests": 1}], ) + after_abandoned: Final = len(_counter_mgets(router, client, redis_cache)) - pipeline_count: Final = len(client.pipelines) - first_cooldown_mgets: Final = tuple( - command - for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[:pipeline_count]) - if command[0] == "MGET" and cooldown_key in command[1:] - ) with request_redis_batch_scope(): router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment: Final = await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) - cooldowns: Final = await router.cooldown_cache.async_get_active_cooldowns(["dep-a"], parent_otel_span=None) + deployment: Final = await _route(router) - second_cooldown_mgets: Final = tuple( - command - for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[pipeline_count:]) - if command[0] == "MGET" and cooldown_key in command[1:] - ) - - assert deployment["model_info"]["id"] == "dep-b" - assert [model_id for model_id, _ in cooldowns] == ["dep-a"] - assert len(first_cooldown_mgets) == 1 - assert second_cooldown_mgets == () + assert deployment == "dep-b" + assert after_abandoned == 1 + assert len(_counter_mgets(router, client, redis_cache)) == 1 assert redis_cache.alone == [] @pytest.mark.asyncio -async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses(monkeypatch): - clock: Final = 1_000_000.0 - monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock) - dep_a_key: Final = _cooldown_key("dep-a") - dep_b_key: Final = _cooldown_key("dep-b") - old_cooldown: Final = { - "exception_received": "429", - "status_code": "429", - "timestamp": _FAR_FUTURE, - "cooldown_time": 60, - } - newer_memory_cooldown: Final = { - "exception_received": "429", - "status_code": "429", - "timestamp": _FAR_FUTURE + 1, - "cooldown_time": 60, - } - redis_only_cooldown: Final = { - "exception_received": "429", - "status_code": "429", - "timestamp": _FAR_FUTURE + 2, - "cooldown_time": 60, - } - - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": - return [json.dumps(redis_cache.store[key]) if key in redis_cache.store else None for key in command[1:]] - return _lua_ok_replies(command) - - client: Final = FakeClient(replies) - redis_cache: Final = FakeRedisCache(client) - redis_cache.store[dep_a_key] = old_cooldown - redis_cache.store[dep_b_key] = redis_only_cooldown - router: Final = _router(redis_cache, routing_strategy="simple-shuffle") - cooldown_store: Final = router.cooldown_cache.cooldown_store - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - cooldown_store.put("dep-a", newer_memory_cooldown, 60) - redis_cache.store[dep_a_key] = old_cooldown - await request.flush_all() - - prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET") - - assert len(prefetched_mgets) == 1 - assert frozenset(prefetched_mgets[0][1:]) == frozenset({dep_a_key, dep_b_key}) - assert dict(await cooldown_store.active(("dep-a", "dep-b"), None)) == { - "dep-a": newer_memory_cooldown, - "dep-b": redis_only_cooldown, - } - - -@pytest.mark.asyncio -async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(monkeypatch): - clock: Final = 1_000_000.0 - monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock) - cooldown_key: Final = _cooldown_key("dep-a") - active_cooldown: Final = { - "exception_received": "429", - "status_code": "429", - "timestamp": _FAR_FUTURE, - "cooldown_time": 60, - } - mget_replies: Final = iter((ConnectionError("redis down"), None)) - - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": - response: Final = next(mget_replies) - if isinstance(response, Exception): - return response - return [json.dumps(active_cooldown) if key == cooldown_key else None for key in command[1:]] - return _lua_ok_replies(command) - - client: Final = FakeClient(replies) - redis_cache: Final = FakeRedisCache(client) - redis_cache.store[cooldown_key] = active_cooldown - router: Final = _router(redis_cache, routing_strategy="simple-shuffle") - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await request.flush_all() - - pipeline_count: Final = len(client.pipelines) - first_cooldown_mgets: Final = tuple( - command - for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[:pipeline_count]) - if command[0] == "MGET" and cooldown_key in command[1:] - ) - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment: Final = await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) - cooldowns: Final = await router.cooldown_cache.async_get_active_cooldowns(["dep-a"], parent_otel_span=None) - - second_cooldown_mgets: Final = tuple( - command - for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[pipeline_count:]) - if command[0] == "MGET" and cooldown_key in command[1:] - ) - - assert deployment["model_info"]["id"] == "dep-b" - assert [model_id for model_id, _ in cooldowns] == ["dep-a"] - assert len(first_cooldown_mgets) == 1 - assert len(second_cooldown_mgets) == 1 - - -@pytest.mark.asyncio -async def test_a_cooldown_that_leaves_memory_before_routing_is_read_again(monkeypatch): - clock: Final = 1_000_000.0 - monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock) - cooldown_key: Final = _cooldown_key("dep-a") - active_cooldown: Final = { - "exception_received": "429", - "status_code": "429", - "timestamp": _FAR_FUTURE, - "cooldown_time": 60, - } +@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) - redis_cache.store[cooldown_key] = active_cooldown - router: Final = _router(redis_cache, routing_strategy="simple-shuffle") - cooldown_store: Final = router.cooldown_cache.cooldown_store - cooldown_store.put("dep-a", active_cooldown, 60) + 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() - cooldown_store.evict_local("dep-a") - deployment: Final = await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + after_failure: Final = len(_counter_mgets(router, client, redis_cache)) - prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET") - fallback_cooldown_mgets: Final = tuple( - keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_key in keys - ) + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + deployment: Final = await _route(router) - assert len(prefetched_mgets) == 1 - assert prefetched_mgets[0][1:] == (_cooldown_key("dep-b"),) - assert deployment["model_info"]["id"] == "dep-b" - assert fallback_cooldown_mgets == ((cooldown_key,),) + assert after_failure == 1 + assert len(_counter_mgets(router, client, redis_cache)) == 2 + assert deployment == "dep-b" @pytest.mark.asyncio @@ -670,52 +570,35 @@ async def test_arming_outside_a_request_scope_is_a_no_op(): @pytest.mark.asyncio -async def test_simple_shuffle_prefetches_only_its_cooldown_read_into_the_admission_pipeline(): +async def test_simple_shuffle_arms_no_python_prefetch_and_reads_no_cooldown_key(): client = FakeClient(_lua_ok_replies) redis_cache = FakeRedisCache(client) router = _router(redis_cache, routing_strategy="simple-shuffle") limiter = _limiter(redis_cache) - with request_redis_batch_scope(): + with request_redis_batch_scope() as request: router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + assert set(request.prefetched) == {"routing_cooldowns"} await limiter.atomic_check_and_increment_by_n( descriptors=[_descriptor("api_key", "k1", 10)], increments=[{"requests": 1}], ) - deployment = await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + deployment = await _route(router) - assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} - assert len(client.pipelines) == 1 - commands = client.pipelines[0].commands - assert [c[0] for c in commands] == ["MGET", "EVALSHA"] - assert set(commands[0][1:]) == { - _cooldown_key("dep-a"), - _cooldown_key("dep-b"), - } + assert deployment in {"dep-a", "dep-b"} + assert [[c[0] for c in pipeline.commands] for pipeline in client.pipelines] == [["EVALSHA"]] assert redis_cache.alone == [] - shuffle = Router(model_list=[_deployment("dep-a")], routing_strategy="simple-shuffle") - shuffle.update_redis_cache(cache=redis_cache) - with request_redis_batch_scope() as request: - shuffle.arm_routing_read_prefetch(_MODEL_GROUP, {}) - armed = request.prefetched["routing_read"] - assert isinstance(armed, ArmedRead) - assert armed.keys == {_cooldown_key("dep-a")} # no usage counters for shuffle - @pytest.mark.asyncio -@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing-v2"]) +@pytest.mark.usefixtures("pinned_clock") @pytest.mark.parametrize("with_limiter", [True, False]) -async def test_requests_within_the_cooldown_read_interval_read_cooldowns_from_redis_once( - routing_strategy: str, with_limiter: bool -): +async def test_requests_within_the_read_interval_read_the_counters_from_redis_once(with_limiter: bool): client = FakeClient(_lua_ok_replies) redis_cache = FakeRedisCache(client) - router = _router(redis_cache, routing_strategy=routing_strategy) + router = _router(redis_cache) limiter = _limiter(redis_cache) - request_round_trips: list[tuple[int, int]] = [] + request_round_trips: list[tuple[int, int]] = [] # mutable-ok: one entry per simulated request for _ in range(3): pipeline_count = len(client.pipelines) @@ -727,35 +610,24 @@ async def test_requests_within_the_cooldown_read_interval_read_cooldowns_from_re descriptors=[_descriptor("api_key", "k1", 10)], increments=[{"requests": 1}], ) - await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + await _route(router) request_round_trips.append((len(client.pipelines) - pipeline_count, len(redis_cache.alone) - alone_count)) - pipeline_mgets = [command for pipeline in client.pipelines for command in pipeline.commands if command[0] == "MGET"] - alone_mgets = [keys for command, keys in redis_cache.alone if command == "MGET"] - cooldown_keys = { - _cooldown_key("dep-a"), - _cooldown_key("dep-b"), - } - cooldown_mgets = [command[1:] for command in pipeline_mgets if cooldown_keys.intersection(command[1:])] + [ - keys for keys in alone_mgets if cooldown_keys.intersection(keys) - ] - - assert len(cooldown_mgets) == 1 + assert len(_counter_mgets(router, client, redis_cache)) == 1 if not with_limiter: assert request_round_trips[1:] == [(0, 0), (0, 0)] @pytest.mark.asyncio -async def test_concurrent_requests_share_one_cooldown_read_per_interval(): +@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(): + 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(): @@ -763,126 +635,57 @@ async def test_concurrent_requests_share_one_cooldown_read_per_interval(): else: first_armed.set() await both_armed.wait() - return await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + return await _route(router) deployments: Final = await asyncio.gather(route_after_both_requests_arm(), route_after_both_requests_arm()) - cooldown_keys: Final = frozenset( - { - _cooldown_key("dep-a"), - _cooldown_key("dep-b"), - } - ) - cooldown_mgets: Final = tuple( - command - for pipeline in client.pipelines - for command in pipeline.commands - if command[0] == "MGET" and cooldown_keys.intersection(command[1:]) - ) - assert all(deployment["model_info"]["id"] in {"dep-a", "dep-b"} for deployment in deployments) - assert len(cooldown_mgets) == 1 + 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_cooldowns_again_once_the_read_interval_elapses(monkeypatch): - first_time: Final = 1_000_000.0 - monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time) - active_cooldown = { - "exception_received": "429", - "status_code": "429", - "timestamp": _FAR_FUTURE, - "cooldown_time": 60, - } - mget_results = iter((None, active_cooldown)) - - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": - result = next(mget_results) - return [ - None if result is None or key != _cooldown_key("dep-a") else json.dumps(result) for key in command[1:] - ] - return _lua_ok_replies(command) - - client = FakeClient(replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache, routing_strategy="simple-shuffle") - cooldown_store = router.cooldown_cache.cooldown_store - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) - - monkeypatch.setattr( - dual_cache_module.time, - "time", - lambda: first_time + POLICY.redis_read_interval + 1, - ) - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment = await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) - - cooldown_keys = { - _cooldown_key("dep-a"), - _cooldown_key("dep-b"), - } - cooldown_mgets = [ - command - for pipeline in client.pipelines - for command in pipeline.commands - if command[0] == "MGET" and cooldown_keys.intersection(command[1:]) - ] - assert len(cooldown_mgets) == 2 - assert deployment["model_info"]["id"] == "dep-b" - - -@pytest.mark.asyncio -async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeypatch): - first_time: Final = 1_000_000.0 - pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) - monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time) - monkeypatch.setattr( - "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", - lambda: pinned, - ) +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, routing_strategy="usage-based-routing-v2") - cooldown_store = router.cooldown_cache.cooldown_store + 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 - time_offset = POLICY.redis_read_interval + 0.5 - - assert time_offset < usage_cache.redis_batch_cache_expiry with request_redis_batch_scope(): router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + await _route(router) - monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time + time_offset) + 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 router.async_get_available_deployment( - model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} - ) + await _route(router) - cooldown_keys = frozenset( - { - _cooldown_key("dep-a"), - _cooldown_key("dep-b"), - } - ) 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:]) == cooldown_keys + assert frozenset(second_pipeline_mgets[0][1:]) == _counter_keys(router, ("dep-b",)) @pytest.mark.asyncio @@ -897,6 +700,7 @@ async def test_two_backends_flush_concurrently_one_pipeline_each(): @pytest.mark.asyncio +@pytest.mark.usefixtures("pinned_clock") async def test_a_single_lua_group_rides_the_pipeline_with_the_armed_routing_read(): client = FakeClient(_lua_ok_replies) redis_cache = FakeRedisCache(client) @@ -1081,17 +885,31 @@ async def test_identity_prefetch_is_one_mget_after_which_hits_and_misses_alike_c assert cache.in_memory_cache.get_cache("end_user_id:eu-miss") is None -@pytest.mark.parametrize( - ("cooldown_keys", "usage_keys", "expected"), - [ - (("cooldown:a",), (), ROUTER_COOLDOWNS_TARGET), - ((), ("usage:a",), ROUTER_USAGE_TARGET), - (("cooldown:a",), ("usage:a",), "router_cooldowns_usage"), - ], -) -def test_the_routing_read_family_follows_the_keys_that_are_actually_due( - cooldown_keys: tuple[str, ...], usage_keys: tuple[str, ...], expected: str -): - """A routing MGET is ``router_cooldowns`` when only cooldown keys go out, ``router_usage`` when the - cooldowns were already in memory and only usage counters go out, and the combined family otherwise.""" - assert routing_target((TierRead(DualCache(), cooldown_keys), TierRead(DualCache(), usage_keys))) == expected +@pytest.mark.asyncio +async def test_awaiting_a_native_read_runs_the_native_hooks_then_flushes_the_native_pre_batch_once_per_request(): + natives: Final = FakeNatives() + + with request_redis_batch_scope(post_call_deadline=7, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double + native: Final = request.native() + request.add_native_flush_hook(lambda: native.events.append("hook")) + result: Final = request.native_result(FakeNativePending(native, lambda: "answer")) + declared_only: Final = list(native.events) + answer: Final = await result + await request.flush_all() + + assert request.native() is native + assert [built.post_call_deadline for built in natives.built] == [7] + assert declared_only == [] + assert answer == "answer" + assert native.events == ["hook", "flush_pre", "resolve", "hook", "flush_pre"] + + +@pytest.mark.asyncio +async def test_a_request_that_no_native_store_asks_for_builds_no_native_batches(): + natives: Final = FakeNatives() + + with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double + await request.flush_all() + await request.flush_post_call() + + assert natives.built == [] diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 272da9d579a..63faf732b68 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -16,12 +16,27 @@ import pytest from pytest_socket import enable_socket, socket_allow_hosts if TYPE_CHECKING: + from unittest.mock import MagicMock + from litellm.caching.redis_cache import RedisCache RouterRedis: TypeAlias = tuple["RedisCache", asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]] RedisWriteGate: TypeAlias = tuple[asyncio.Event, asyncio.Event, asyncio.Event] +def redis_cache_double() -> "MagicMock": + """A `RedisCache` double for router tests. The native router stores connect from its settings, which name + a closed local port, so they serve their memory tier while the double answers the Python Redis calls.""" + from unittest.mock import MagicMock + + from litellm.caching.redis_cache import RedisCache + + redis: Final = MagicMock(spec=RedisCache) + redis.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2} + redis.namespace = None + return redis + + @pytest.fixture def redis_write_gate() -> RedisWriteGate: released: Final = asyncio.Event() diff --git a/tests/unit/redis/stores/test_router_cooldowns.py b/tests/unit/redis/stores/test_router_cooldowns.py index e65a0c9eeff..e469f4263da 100644 --- a/tests/unit/redis/stores/test_router_cooldowns.py +++ b/tests/unit/redis/stores/test_router_cooldowns.py @@ -1,112 +1,267 @@ -from dataclasses import dataclass -from typing import TYPE_CHECKING, Final -from unittest.mock import AsyncMock, MagicMock +import ast +import asyncio +import json +import logging +import time +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field, replace +from typing import Final import pytest +import litellm from litellm._internal_context import current_service_target -from litellm.caching.redis_batch import request_redis_batch_scope -from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_batch import RequestRedisBatches +from litellm.caching.redis_cache import RedisCache, redis_breaker_state_counts +from litellm.integrations.custom_secret_manager import CustomSecretManager +from litellm.redis.stores._connection import redis_kwargs from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore +from litellm.rust_bridge._native import NativeRouterCooldowns +from litellm.rust_bridge.host import settings +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNativeRequest, FakeNatives -if TYPE_CHECKING: - from tests.unit.conftest import RouterRedis + +@dataclass(frozen=True, slots=True) +class _WireCooldown: + """One cooldown as main's ``RedisCache`` writes it (``str(value)``) and reads it (``json.loads`` then + ``ast.literal_eval``). The same literals sit in ``litellm-rust/crates/redis-store/tests/support/mod.rs``, + where real Redis is checked to hold them after the Rust store writes.""" + + cooldown: CooldownCacheValue + str_dict: str + json: str + ttl_seconds: float + redis_ttl_seconds: int + + +_WIRE_COOLDOWNS: Final = ( + _WireCooldown( + cooldown={ + "exception_received": "RateLimitError: limited é 'quoted'", + "status_code": "429", + "timestamp": 1730000000.25, + "cooldown_time": 60, + }, + str_dict="{'exception_received': \"RateLimitError: limited é 'quoted'\", 'status_code': '429', 'timestamp': 1730000000.25, 'cooldown_time': 60}", + json='{"exception_received": "RateLimitError: limited \\u00e9 \'quoted\'", "status_code": "429", "timestamp": 1730000000.25, "cooldown_time": 60}', + ttl_seconds=90.5, + redis_ttl_seconds=90, + ), + _WireCooldown( + cooldown={ + "exception_received": "Timeout", + "status_code": "408", + "timestamp": 1730000000.5, + "cooldown_time": 0.5, + }, + str_dict="{'exception_received': 'Timeout', 'status_code': '408', 'timestamp': 1730000000.5, 'cooldown_time': 0.5}", + json='{"exception_received": "Timeout", "status_code": "408", "timestamp": 1730000000.5, "cooldown_time": 0.5}', + ttl_seconds=30.0, + redis_ttl_seconds=30, + ), +) + + +def _entry(model_id: str = "dep") -> CooldownCacheValue: + return {"exception_received": model_id, "status_code": "429", "timestamp": 1000.0, "cooldown_time": 120.0} @dataclass -class Clock: - now: float = 1000.0 - failure: RuntimeError | None = None +class FakeCooldowns: + redis_kwargs: Mapping[str, object] | None + namespace: str | None + stored: dict[str, CooldownCacheValue] = field(default_factory=dict) + fails: dict[str, int] = field(default_factory=dict) + targets: list[str | None] = field(default_factory=list) + smallest: float | None = None + declared: list[tuple[str, ...]] = field(default_factory=list) - def __call__(self) -> float: - if self.failure is not None: - raise self.failure - return self.now + def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: + self.targets.append(current_service_target()) + self.stored[model_id] = CooldownCacheValue(**cooldown) # pyright: ignore[reportArgumentType] # test double + + def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: + self.targets.append(current_service_target()) + self.fails[model_id] = self.fails.get(model_id, 0) + 1 + return self.fails[model_id] + + def active(self, model_ids: Sequence[str]) -> list[tuple[str, CooldownCacheValue]]: + self.targets.append(current_service_target()) + return [(model_id, self.stored[model_id]) for model_id in model_ids if model_id in self.stored] + + def aactive(self, model_ids: Sequence[str]) -> "asyncio.Future[list[tuple[str, CooldownCacheValue]]]": + future: Final[asyncio.Future[list[tuple[str, CooldownCacheValue]]]] = asyncio.get_running_loop().create_future() + future.set_result(self.active(model_ids)) + return future + + def declare_active( + self, request: FakeNativeRequest, model_ids: Sequence[str] + ) -> FakeNativePending[list[tuple[str, CooldownCacheValue]]]: + self.targets.append(current_service_target()) + self.declared.append(tuple(model_ids)) + return FakeNativePending(request, lambda: self.active(model_ids)) + + def minimum(self, model_ids: Sequence[str]) -> float | None: + return self.smallest + + def evict_local(self, model_id: str) -> None: + self.stored.pop(model_id, None) + + def flush_local_counters(self) -> None: + self.fails.clear() + + def snapshot(self) -> dict[str, CooldownCacheValue]: + return dict(self.stored) -def _entry(timestamp: float = 1000.0, ttl: float = 120.0) -> CooldownCacheValue: - return {"exception_received": "limited", "status_code": "429", "timestamp": timestamp, "cooldown_time": ttl} +@dataclass +class Handles: + built: list[FakeCooldowns] = field(default_factory=list) + + def __call__(self, connection: RedisCache | None) -> FakeCooldowns: + handle: Final = FakeCooldowns( + None if connection is None else redis_kwargs(connection), + None if connection is None else connection.namespace, + ) + self.built.append(handle) + return handle -@pytest.mark.parametrize("suffix", (None, "", "timeout")) -def test_allowed_fails_preserves_local_count_during_redis_failure(suffix: str | None) -> None: - redis: Final = MagicMock(spec=RedisCache) - redis.increment_cache.side_effect = ConnectionError("offline") - clock: Final = Clock() - store: Final = RouterCooldownStore(redis, clock=clock) - assert store.increment_allowed_fails("dep", suffix, 30.0) == 1 - assert store.increment_allowed_fails("dep", suffix, 30.0) == 2 - clock.now += 31.0 - assert store.increment_allowed_fails("dep", suffix, 30.0) == 1 - expected: Final = "deployment:dep:allowed_fails" + (f":{suffix}" if suffix else "") - redis.increment_cache.assert_called_with(expected, 1, ttl=30.0) +def _redis(namespace: str | None = None) -> RedisCache: + return RedisCache(host="redis.internal", port=6380, namespace=namespace) -def test_sync_wire_write_and_reads_keep_service_target_and_float_ttl() -> None: - redis: Final = MagicMock(spec=RedisCache) - seen: Final[list[str | None]] = [] - redis.set_cache.side_effect = lambda *args, **kw: seen.append(current_service_target()) - redis.batch_get_cache.side_effect = lambda **kw: {key: _entry() for key in kw["key_list"]} - store: Final = RouterCooldownStore(redis, clock=Clock()) - store.put("dep", _entry(ttl=0.5), 0.5) - redis.set_cache.assert_called_once_with("deployment:dep:cooldown", _entry(ttl=0.5), ttl=0.5) - reader: Final = RouterCooldownStore(redis, clock=Clock()) - assert reader.active_sync(["dep"], None) == (("dep", _entry()),) - assert reader.minimum(["dep"], 60.0, None) == 120.0 - assert seen == ["router_cooldowns"] +def test_the_store_hands_the_native_handle_redis_kwargs_and_namespace_not_a_client() -> None: + handles: Final = Handles() + RouterCooldownStore(_redis("tenant"), handle=handles) + assert handles.built[0].redis_kwargs == {"host": "redis.internal", "port": 6380, "socket_timeout": 5.0} + assert handles.built[0].namespace == "tenant" + + +def test_late_attach_rebuilds_once_and_later_attaches_are_ignored() -> None: + handles: Final = Handles() + store: Final = RouterCooldownStore(None, handle=handles) + store.put("dep", _entry(), 60.0) + store.attach_redis(None) + store.attach_redis(_redis()) + store.attach_redis(_redis("other")) + assert [handle.redis_kwargs is not None for handle in handles.built] == [False, True] + assert handles.built[1].namespace is None + assert store.active_sync(["dep"], None) == () + + +def test_writes_and_reads_carry_the_router_cooldowns_service_target() -> None: + handles: Final = Handles() + store: Final = RouterCooldownStore(_redis(), handle=handles) + store.put("dep", _entry(), 60.0) + assert store.increment_allowed_fails("dep", "timeout", 30.0) == 1 + assert store.active_sync(["dep", "other"], None) == (("dep", _entry()),) + assert handles.built[0].targets == ["router_cooldowns"] * 3 assert current_service_target() is None -@pytest.mark.asyncio -async def test_remote_backfill_expires_at_corrected_sixty_second_cap() -> None: - clock: Final = Clock() - redis: Final = MagicMock(spec=RedisCache) - redis.async_batch_get_cache = AsyncMock(return_value={"deployment:dep:cooldown": _entry()}) - store: Final = RouterCooldownStore(redis, clock=clock) - assert await store.active(["dep"], None) == (("dep", _entry()),) - clock.now += 61.0 - assert await store.active(["dep"], None) == (("dep", _entry()),) - assert redis.async_batch_get_cache.await_count == 2 - clock.now += 60.0 - assert await store.active(["dep"], None) == () +@pytest.mark.parametrize(("smallest", "expected"), ((None, 60.0), (0, 60.0), (5, 5))) +def test_minimum_falls_back_to_the_default_when_nothing_or_zero_is_stored( + smallest: float | None, expected: float +) -> None: + handles: Final = Handles() + store: Final = RouterCooldownStore(None, handle=handles) + handles.built[0].smallest = smallest + assert store.minimum(["dep"], 60.0, None) == expected + + +@pytest.mark.parametrize("case", _WIRE_COOLDOWNS, ids=("rate-limit", "timeout")) +def test_the_wire_cases_are_what_mains_codec_writes_and_reads(case: _WireCooldown) -> None: + assert str(case.cooldown) == case.str_dict + assert ast.literal_eval(case.str_dict) == case.cooldown + assert json.loads(case.json) == case.cooldown + assert int(case.ttl_seconds) == case.redis_ttl_seconds + + +@pytest.mark.parametrize( + "setting", + ({"ssl_ciphers": "HIGH"}, {"ssl": True, "ssl_validate_ocsp": True}, {"credential_provider": object()}), +) +def test_a_redis_setting_the_native_client_cannot_honor_never_fails_construction( + setting: dict[str, object], caplog: pytest.LogCaptureFixture +) -> None: + redis: Final = RedisCache(host="127.0.0.1", port=1, **setting) + with caplog.at_level(logging.WARNING): + store: Final = RouterCooldownStore(redis) + cooldown: Final = CooldownCacheValue( + exception_received="model", status_code="429", timestamp=time.time(), cooldown_time=120.0 + ) + store.put("model", cooldown, 60.0) + + assert store.active_sync(["model"], None) == (("model", cooldown),) + assert next(iter(setting)) in caplog.text + + +class _RecordingSecrets(CustomSecretManager): + def __init__(self, values: Mapping[str, str]) -> None: + super().__init__(secret_manager_name="router_cooldowns_test") # pyright: ignore[reportUnknownMemberType] # CustomSecretManager takes untyped **kwargs + self.values: Final = values + self.asked: tuple[str, ...] = () + + async def async_read_secret( + self, secret_name: str, optional_params: dict[str, object] | None = None, timeout: object = None + ) -> str | None: + raise AssertionError("get_secret reads custom managers synchronously") + + def sync_read_secret( + self, secret_name: str, optional_params: dict[str, object] | None = None, timeout: object = None + ) -> str | None: + self.asked = (*self.asked, secret_name) + return self.values.get(secret_name) + + +def test_the_native_store_reads_redis_secrets_through_the_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("REDIS_PASSWORD", raising=False) + manager: Final = _RecordingSecrets({"REDIS_PASSWORD": "vault-password"}) + monkeypatch.setattr(litellm, "secret_manager_client", manager) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + configured: Final = settings.secret_manager + monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=False)) + + NativeRouterCooldowns({"host": "127.0.0.1", "port": 1}, namespace="secret-manager-test") + + assert "REDIS_PASSWORD" in manager.asked + + +def test_the_native_breaker_counts_in_the_redis_breaker_state_gauge(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true") + monkeypatch.setenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", "1") + open_before: Final = redis_breaker_state_counts().get("open", 0) + store: Final = NativeRouterCooldowns({"host": "127.0.0.1", "port": 1}, namespace="breaker-gauge-test") + + store.put("model", _entry("model"), 60.0) + + assert redis_breaker_state_counts().get("open", 0) == open_before + 1 @pytest.mark.asyncio -async def test_lazy_attach_and_failure_retry_do_not_throttle_next_read() -> None: - store: Final = RouterCooldownStore(None, clock=Clock()) - assert await store.active(["dep"], None) == () - redis: Final = MagicMock(spec=RedisCache) - redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("offline")) - store.attach_redis(redis) - assert await store.active(["dep"], None) == () - redis.async_batch_get_cache.side_effect = None - redis.async_batch_get_cache.return_value = {"deployment:dep:cooldown": _entry()} - assert await store.active(["dep"], None) == (("dep", _entry()),) - assert redis.async_batch_get_cache.await_count == 2 +async def test_a_declared_active_read_answers_after_the_native_pre_flush_under_the_router_cooldowns_target() -> None: + handles: Final = Handles() + store: Final = RouterCooldownStore(_redis(), handle=handles) + natives: Final = FakeNatives() + request: Final = RequestRedisBatches(native_factory=natives) # pyright: ignore[reportArgumentType] # test double + store.put("dep", _entry(), 60.0) + + declared: Final = store.declare_active(request, ["dep", "other"]) + assert declared is not None + assert natives.built[0].events == [] + + assert tuple(await declared) == (("dep", _entry()),) + assert natives.built[0].events == ["flush_pre", "resolve"] + assert handles.built[0].declared == [("dep", "other")] + assert handles.built[0].targets[1] == "router_cooldowns" -def test_minimum_preserves_zero_default_and_does_not_correct_expired_entries() -> None: - store: Final = RouterCooldownStore(None, clock=Clock()) - store.put("zero", _entry(ttl=0), 600) - store.put("expired", _entry(timestamp=0, ttl=20), 600) - assert store.minimum(["zero"], 60, None) == 60 - assert store.minimum(["expired"], 60, None) == 20 - assert store.active_sync(["expired"], None) == () +def test_a_store_without_redis_declares_nothing_on_the_request() -> None: + natives: Final = FakeNatives() + request: Final = RequestRedisBatches(native_factory=natives) # pyright: ignore[reportArgumentType] # test double - -@pytest.mark.asyncio -async def test_prefetch_backfill_failure_allows_an_immediate_read_retry(router_redis: "RouterRedis") -> None: - redis, _ = router_redis - clock: Final = Clock() - value: Final = _entry() - writer: Final = RouterCooldownStore(redis, clock=clock) - reader: Final = RouterCooldownStore(redis, clock=clock) - writer.put("dep", value, 120.0) - - with request_redis_batch_scope() as request: - reader.arm_routing_read(("dep",), None) - clock.failure = RuntimeError("backfill clock unavailable") - await request.flush_all() - - clock.failure = None - assert await reader.active(("dep",), None) == (("dep", value),) + assert RouterCooldownStore(None, handle=Handles()).declare_active(request, ["dep"]) is None + assert natives.built == [] diff --git a/tests/unit/redis/stores/test_router_usage.py b/tests/unit/redis/stores/test_router_usage.py index baa78514738..9a822c0fef8 100644 --- a/tests/unit/redis/stores/test_router_usage.py +++ b/tests/unit/redis/stores/test_router_usage.py @@ -1,12 +1,96 @@ -from typing import Final -from unittest.mock import AsyncMock, MagicMock +import asyncio +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Final, TypeVar import pytest from litellm import Router 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.router_usage import RouterUsageStore, UsageCounter +from litellm.redis.stores.router_usage import NativeCounter, RouterUsageStore, UsageCounter +from tests.unit.caching.test_redis_batch import FakeNatives + +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 FakeUsage: + connected: bool + counts: dict[NativeCounter, float] = field(default_factory=dict) + calls: list[tuple[str, object, str | None]] = field(default_factory=list) + requests: list[object] = field(default_factory=list) + + def increment(self, counter: NativeCounter, value: float, ttl_seconds: int) -> "asyncio.Future[float]": + self.calls.append(("increment", (counter, value, ttl_seconds), current_service_target())) + self.counts[counter] = self.counts.get(counter, 0) + value + return _done(self.counts[counter]) + + def increment_deployment( + self, + deployment_id: str, + model: str | None, + minute: str, + tokens: float, + requests: int, + *, + request: object = None, + ) -> "asyncio.Future[None]": + self.calls.append(("deployment", (deployment_id, model, minute, tokens, requests), current_service_target())) + self.requests.append(request) + for kind, value in (("tpm", tokens), ("rpm", requests)): + counter: NativeCounter = (deployment_id, model, minute, kind) # pyright: ignore[reportAssignmentType] # literal kinds + self.counts[counter] = self.counts.get(counter, 0) + value + return _done(None) + + def get_many(self, counters: Sequence[NativeCounter]) -> "asyncio.Future[list[float | None]]": + return _done([self.counts.get(counter) for counter in counters]) + + def flush_local_counters(self) -> None: + self.counts.clear() + + def snapshot(self) -> dict[str, float]: + return {} + + +@dataclass +class Handles: + built: list[FakeUsage] = field(default_factory=list) + + def __call__(self, connection: RedisCache | None) -> FakeUsage: + self.built.append(FakeUsage(connection is not None)) + return self.built[-1] + + +@pytest.mark.asyncio +async def test_counters_cross_the_bridge_as_domain_tuples_under_the_router_usage_target() -> None: + handles: Final = Handles() + store: Final = RouterUsageStore(None, handle=handles) + counter: Final = UsageCounter("dep", None, "12-34", "rpm") + + await store.async_increment(counter, 1, 60) + await store.increment_deployment("dep", "model", "12-34", 7.5, 1, None) + + assert handles.built[0].calls == [ + ("increment", (("dep", None, "12-34", "rpm"), 1, 60), "router_usage"), + ("deployment", ("dep", "model", "12-34", 7.5, 1), "router_usage"), + ] + assert await store.async_batch_get((counter, UsageCounter("dep", "model", "12-34", "tpm"))) == [1, 7.5] + + +def test_late_attach_rebuilds_the_handle_once() -> None: + handles: Final = Handles() + store: Final = RouterUsageStore(None, handle=handles) + store.attach_redis(RedisCache(host="redis.internal")) + store.attach_redis(RedisCache(host="other")) + assert [handle.connected for handle in handles.built] == [False, True] @pytest.mark.asyncio @@ -25,63 +109,15 @@ async def test_router_reset_clears_usage_and_failure_counters(operation: str) -> @pytest.mark.asyncio -async def test_deployment_increment_preserves_pipeline_ttl_and_integer_mirror() -> None: - redis: Final = MagicMock(spec=RedisCache) - redis.async_increment_pipeline = AsyncMock(return_value=[7.75, 2.0]) - store: Final = RouterUsageStore(redis) - await store.increment_deployment("dep", "anthropic/model", "12-34", 7.75, 2, None) - redis.async_increment_pipeline.assert_awaited_once_with( - increment_list=[ - {"key": "global_router:dep:anthropic/model:tpm:12-34", "increment_value": 7.75, "ttl": 60}, - {"key": "global_router:dep:anthropic/model:rpm:12-34", "increment_value": 2, "ttl": 60}, - ], - parent_otel_span=None, - ) - values: Final = await store.async_batch_get( - (UsageCounter("dep", "anthropic/model", "12-34", "tpm"), UsageCounter("dep", "anthropic/model", "12-34", "rpm")) - ) - assert values == [7, 2] - assert all(type(value) is int for value in values) - redis.async_batch_get_cache.assert_not_called() - assert current_service_target() is None +async def test_inside_a_request_deployment_increments_ride_its_native_post_call_batch_and_the_shutdown_drain(): + handles: Final = Handles() + store: Final = RouterUsageStore(None, handle=handles) + natives: Final = FakeNatives() + await store.increment_deployment("dep", "model", "12-34", 7.5, 1, None) + with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double + await store.increment_deployment("dep", "model", "12-34", 2, 1, None) + await drain_post_call_redis_batches() -@pytest.mark.asyncio -async def test_memory_only_usage_increments_accumulate_without_redis() -> None: - store: Final = RouterUsageStore(None) - await store.increment_deployment("dep", "model", "12-34", 2, 1, None) - await store.increment_deployment("dep", "model", "12-34", 5, 1, None) - assert await store.async_batch_get( - (UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm")) - ) == [7, 2] - - -@pytest.mark.asyncio -async def test_pre_call_rpm_increment_uses_same_namespace_and_service_target() -> None: - redis: Final = MagicMock(spec=RedisCache) - seen: Final[list[str | None]] = [] - - async def increment(*args: object, **kwargs: object) -> float: - seen.append(current_service_target()) - return 4.0 - - redis.async_increment = AsyncMock(side_effect=increment) - store: Final = RouterUsageStore(None) - store.attach_redis(redis) - counter: Final = UsageCounter("dep", "model", "12-34", "rpm") - assert await store.async_increment(counter, value=1, ttl=60) == 4.0 - assert seen == ["router_usage"] - redis.async_increment.assert_awaited_once_with( - "global_router:dep:model:rpm:12-34", 1, parent_otel_span=None, ttl=60, refresh_ttl=False - ) - - -@pytest.mark.asyncio -async def test_failed_usage_pipeline_keeps_the_memory_increment_and_is_readable() -> None: - redis: Final = MagicMock(spec=RedisCache) - redis.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("offline")) - store: Final = RouterUsageStore(redis) - await store.increment_deployment("dep", "model", "12-34", 5, 1, None) - assert await store.async_batch_get( - (UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm")) - ) == [5, 1] + assert handles.built[0].requests == [None, request.native()] + assert natives.built[0].events == ["flush_post"] diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index 0c19cfc72f6..1ad6c0a3491 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -4,7 +4,6 @@ import logging from typing import Final, Literal -from unittest.mock import patch import pytest @@ -1846,23 +1845,20 @@ async def test_tag_fail_open_survives_deployment_narrowing( ) -> None: router: Final = _fail_open_router(route=route) requested_model: Final = "chat" if route == "model" else "chat-test" if route == "wildcard" else "public-chat" - with patch("time.time", return_value=1_700_000_000.0): - if narrowing == "cooldown": - router.cooldown_cache.add_deployment_to_cooldown( - model_id="opted-in", - original_exception=RuntimeError("unavailable"), - exception_status=500, - cooldown_time=60, - ) - candidates: Final = await router.async_get_healthy_deployments( - model=requested_model, - request_kwargs={ - "litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"}, - **( - {"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {} - ), - }, + if narrowing == "cooldown": + router.cooldown_cache.add_deployment_to_cooldown( + model_id="opted-in", + original_exception=RuntimeError("unavailable"), + exception_status=500, + cooldown_time=60, ) + candidates: Final = await router.async_get_healthy_deployments( + model=requested_model, + request_kwargs={ + "litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"}, + **({"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {}), + }, + ) assert isinstance(candidates, (list, tuple)) assert tuple(deployment["model_info"]["id"] for deployment in candidates) == ("remaining",) diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 1f3c6a8e948..de57ff6c718 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,7 +1,6 @@ import asyncio import json -from itertools import chain -from typing import TYPE_CHECKING, Final +from typing import Final from collections.abc import AsyncIterator, Mapping from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -22,9 +21,6 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( ) -if TYPE_CHECKING: - from tests.unit.conftest import RouterRedis - def _repinning_router(enabled: bool | None = True) -> litellm.Router: return litellm.Router( @@ -325,38 +321,6 @@ async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: M router.discard() -@pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"]) -async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str, router_redis: "RouterRedis") -> None: - router: Final = _repinning_router(case != "disabled") - redis, commands = router_redis - try: - router.update_redis_cache(redis) - callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) - if case == "group-disabled": - callback.model_group_affinity_config = {"repin-group": []} - await router._repin_session_after_success( - { - "model_info": {"id": "replacement"}, - "metadata": { - "deployment_model_name": "repin-group", - **({"session_id": "repin-session"} if case != "no-session" else {}), - SESSION_ID_GENERATED_METADATA_KEY: case == "generated-session", - }, - }, - {"status": "completed"}, - ) - expected: Final = ( - {f"deployment:{model_id}:cooldown" for model_id in ("original", "replacement", "other")} - if case == "active" - else set() - ) - observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize())) - reads: Final = frozenset(chain.from_iterable(args[1:] for args, _, _ in observed if args[0] == "MGET")) - assert reads == expected - finally: - router.discard() - - class MockResponse: def __init__(self, json_data, status_code): self._json_data = json_data diff --git a/tests/unit/router_utils/test_cooldown_cache.py b/tests/unit/router_utils/test_cooldown_cache.py index a679f6ff3f7..9d52b8049ad 100644 --- a/tests/unit/router_utils/test_cooldown_cache.py +++ b/tests/unit/router_utils/test_cooldown_cache.py @@ -1,3 +1,4 @@ +import time from typing import Final import pytest @@ -5,14 +6,12 @@ import pytest from litellm.redis.stores.router_cooldowns import RouterCooldownStore from litellm.router_utils.cooldown_cache import CooldownCache -_NOW: Final = 1000.0 - - @pytest.mark.parametrize("message", ("Short error", "Rate limit " + "secret" * 100)) -@pytest.mark.parametrize("ttl", (None, 0.5, 30.0)) +@pytest.mark.parametrize("ttl", (None, 5.5, 30.0)) def test_cooldown_masks_exception_and_preserves_status_timestamp_and_ttl(message: str, ttl: float | None) -> None: - store: Final = RouterCooldownStore(None, clock=lambda: _NOW) - cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: _NOW) + now: Final = time.time() + store: Final = RouterCooldownStore(None) + cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: now) cache.add_deployment_to_cooldown("dep", Exception(message), 429, ttl) expected_ttl: Final = 60.0 if ttl is None else ttl expected_message: Final = message if len(message) <= 50 else message[:50] + "*" * (len(message) - 50) @@ -22,7 +21,7 @@ def test_cooldown_masks_exception_and_preserves_status_timestamp_and_ttl(message { "exception_received": expected_message, "status_code": "429", - "timestamp": _NOW, + "timestamp": now, "cooldown_time": expected_ttl, }, ) @@ -31,11 +30,12 @@ def test_cooldown_masks_exception_and_preserves_status_timestamp_and_ttl(message @pytest.mark.asyncio async def test_sync_and_async_router_cooldown_surfaces_share_active_state() -> None: - store: Final = RouterCooldownStore(None, clock=lambda: _NOW) - cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: _NOW) + now: Final = time.time() + store: Final = RouterCooldownStore(None) + cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: now) cache.add_deployment_to_cooldown("dep", Exception("unavailable"), 503, 30.0) expected: Final = [ - ("dep", {"exception_received": "unavailable", "status_code": "503", "timestamp": _NOW, "cooldown_time": 30.0}) + ("dep", {"exception_received": "unavailable", "status_code": "503", "timestamp": now, "cooldown_time": 30.0}) ] assert cache.get_active_cooldowns(["absent", "dep"], None) == expected assert await cache.async_get_active_cooldowns(["absent", "dep"], None) == expected diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 71b53be4c34..44688207aae 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -16,12 +16,11 @@ import respx import litellm from litellm import Router -from litellm._internal_context import current_service_target from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCache -from litellm.redis.stores.router_cooldowns import RouterCooldowns, RouterCooldownStore from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials +from litellm.redis.stores.router_cooldowns import RouterCooldowns, RouterCooldownStore from litellm.router_api import RoutingContext from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback @@ -378,9 +377,8 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy: store.increment_allowed_fails.assert_called_once_with("dep-1", None, 0.0) def test_fail_counter_is_shared_across_router_instances(self): - redis: Final = MagicMock(spec=RedisCache) - redis.increment_cache.side_effect = itertools.count(1) - workers: Final = tuple(self._make_router(store=RouterCooldownStore(redis)) for _ in range(2)) + shared: Final = RouterCooldownStore(None) + workers: Final = tuple(self._make_router(store=shared) for _ in range(2)) exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4") results: Final = [ should_cooldown_based_on_allowed_fails_policy( @@ -389,30 +387,16 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy: for i in range(6) ] assert results == [False, False, False, False, False, True] - assert redis.increment_cache.call_count == 6 - def test_fleet_wide_count_from_redis_decides_cooldown(self): - redis: Final = MagicMock(spec=RedisCache) - redis.increment_cache.return_value = 6 - ctx: Final = self._make_router(store=RouterCooldownStore(redis)) + def test_fleet_wide_count_from_the_store_decides_cooldown(self): + store: Final = MagicMock(spec=RouterCooldowns) + store.increment_allowed_fails.return_value = 6 + ctx: Final = self._make_router(store=store) exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4") assert should_cooldown_based_on_allowed_fails_policy( ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5 ) - redis.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0) - - def test_redis_outage_falls_back_to_this_workers_count(self): - redis: Final = MagicMock(spec=RedisCache) - redis.increment_cache.side_effect = ConnectionError("redis down") - ctx: Final = self._make_router(store=RouterCooldownStore(redis)) - exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4") - assert [ - should_cooldown_based_on_allowed_fails_policy( - ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5 - ) - for _ in range(6) - ] == [False, False, False, False, False, True] - assert redis.increment_cache.call_count == 6 + store.increment_allowed_fails.assert_called_once_with("dep-1", None, 60.0) class TestRoutingGroupCooldownAlternatives: @@ -560,38 +544,6 @@ class TestTeamModelCooldownAlternatives: ) -class TestIncrementAllowedFailsServiceTarget: - def test_fail_counter_bump_declares_the_router_cooldowns_key_family(self): - """The allowed_fails INCR is cooldown bookkeeping, so its service span must read - ``redis.incr router_cooldowns`` rather than a bare ``redis.incr``.""" - seen: list[str | None] = [] - cache = MagicMock(spec=RedisCache) - - def _increment(*_args, **_kwargs): - seen.append(current_service_target()) - return 2 - - cache.increment_cache.side_effect = _increment - - assert RouterCooldownStore(cache).increment_allowed_fails("dep-1", None, 60.0) == 2 - assert seen == ["router_cooldowns"] - assert current_service_target() is None - - def test_in_memory_fallback_reads_under_the_same_target(self): - seen: list[str | None] = [] - cache = MagicMock(spec=RedisCache) - - def _increment(*_args, **_kwargs): - seen.append(current_service_target()) - raise ConnectionError("redis down") - - cache.increment_cache.side_effect = _increment - store = RouterCooldownStore(cache) - - assert [store.increment_allowed_fails("dep-1", None, 60.0) for _ in range(4)] == [1, 2, 3, 4] - assert seen == ["router_cooldowns"] * 4 - - @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) @@ -869,14 +821,14 @@ class TestExceptionTypeCountersTrackedIndependently: @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCooldownCacheTTLCorrection: - @pytest.mark.parametrize(("timestamp", "active"), ((880.0, False), (1000.0, True))) + @pytest.mark.parametrize(("age", "active"), ((120.0, False), (0.0, True))) @pytest.mark.parametrize("asynchronous", (False, True)) @pytest.mark.asyncio - async def test_only_unexpired_cooldowns_are_returned(self, timestamp: float, active: bool, asynchronous: bool): - store: Final = RouterCooldownStore(None, clock=lambda: 1000.0) + async def test_only_unexpired_cooldowns_are_returned(self, age: float, active: bool, asynchronous: bool): + store: Final = RouterCooldownStore(None) cc: Final = CooldownCache(store, default_cooldown_time=60.0) value: Final = CooldownCacheValue( - exception_received="Rate limit", status_code="429", timestamp=timestamp, cooldown_time=60.0 + exception_received="Rate limit", status_code="429", timestamp=time.time() - age, cooldown_time=60.0 ) store.put("deployment", value, 600) result: Final = ( @@ -2097,7 +2049,6 @@ def test_increment_deployment_successes_for_current_minute_does_not_write_to_red """ from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache - from litellm.caching.redis_cache import RedisCache from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_successes_for_current_minute, ) @@ -2595,7 +2546,6 @@ async def test_high_traffic_cooldowns_all_healthy_deployments(): await asyncio.sleep(0.0001) except litellm.InternalServerError: model_stats[model_id]["failures"] += 1 - pass except Exception as e: print("Failed test model stats=", model_stats) raise e @@ -2689,7 +2639,6 @@ async def test_high_traffic_cooldowns_one_bad_deployment(): await asyncio.sleep(0.0001) except litellm.InternalServerError: model_stats[model_id]["failures"] += 1 - pass except Exception as e: print("Failed test model stats=", model_stats) raise e @@ -2783,10 +2732,8 @@ async def test_high_traffic_cooldowns_one_rate_limited_deployment(): await asyncio.sleep(0.0001) except litellm.InternalServerError: model_stats[model_id]["failures"] += 1 - pass except litellm.RateLimitError: model_stats[bad_deployment_id]["failures"] += 1 - pass except Exception as e: print("Failed test model stats=", model_stats) raise e diff --git a/tests/unit/router_utils/test_routing_read_batch.py b/tests/unit/router_utils/test_routing_read_batch.py index fea95071b94..7a688a890cc 100644 --- a/tests/unit/router_utils/test_routing_read_batch.py +++ b/tests/unit/router_utils/test_routing_read_batch.py @@ -1,11 +1,10 @@ """ -One Redis round trip per request for the router's pre-call reads. - -Before `RoutingReadBatch`, `async_get_available_deployment` issued one MGET for the cooldown keys -(`CooldownCache`) and a second one for the tpm/rpm counters (`LowestTPMLoggingHandler_v2`). +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. """ import time +from collections.abc import Sequence from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -13,8 +12,15 @@ import pytest import litellm from litellm import Router +from litellm.caching.redis_batch import request_redis_batch_scope from litellm.caching.redis_cache import RedisCache +from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 +from litellm.router_utils import routing_prefetch +from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNativeRequest, FakeNatives +from tests.unit.conftest import redis_cache_double +from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns +from tests.unit.redis.stores.test_router_cooldowns import Handles as CooldownHandles _MODEL_GROUP = "claude" _MESSAGES = [{"role": "user", "content": "ping"}] @@ -34,7 +40,7 @@ def _redis_answering(values_by_key_prefix: dict[str, object]) -> MagicMock: 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 = MagicMock(spec=RedisCache) + redis = redis_cache_double() redis.async_batch_get_cache = AsyncMock(side_effect=_mget) return redis @@ -55,12 +61,14 @@ def _redis_key_families(redis: MagicMock) -> list[list[str]]: ] -def _cooldown(seconds: float) -> dict: - return {"exception_received": "429", "status_code": "429", "timestamp": time.time(), "cooldown_time": seconds} +def _cool_down(router: Router, model_id: str, started: float) -> None: + router.cooldown_cache.cooldown_store.put( + model_id, {"exception_received": "429", "status_code": "429", "timestamp": started, "cooldown_time": 60}, 60 + ) @pytest.mark.asyncio -async def test_usage_based_routing_reads_cooldowns_and_counters_in_one_redis_round_trip(): +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") @@ -73,10 +81,8 @@ async def test_usage_based_routing_reads_cooldowns_and_counters_in_one_redis_rou "dep-a:anthropic/claude-x:tpm", "dep-b:anthropic/claude-x:rpm", "dep-b:anthropic/claude-x:tpm", - "deployment:dep-a:cooldown", - "deployment:dep-b:cooldown", ] - ], "cooldown state and usage counters must arrive in one MGET" + ], "the usage counters must arrive in one MGET, without cooldown keys" @pytest.mark.asyncio @@ -112,21 +118,19 @@ async def test_usage_based_routing_still_batches_when_the_strategy_is_a_fixed_si "dep-a:anthropic/claude-x:tpm", "dep-b:anthropic/claude-x:rpm", "dep-b:anthropic/claude-x:tpm", - "deployment:dep-a:cooldown", - "deployment:dep-b:cooldown", ] ], "the subclassed strategy must still get the batched read, not a second MGET" router.cache.async_batch_get_cache.assert_not_awaited() @pytest.mark.asyncio -async def test_simple_shuffle_still_reads_only_cooldowns(): +async def test_simple_shuffle_reads_nothing_from_the_python_redis(): redis = _redis_answering({}) router = _router(redis, "simple-shuffle") await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) - assert _redis_key_families(redis) == [["deployment:dep-a:cooldown", "deployment:dep-b:cooldown"]] + assert _redis_key_families(redis) == [] @pytest.mark.asyncio @@ -151,14 +155,9 @@ 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, - "deployment:dep-b:cooldown": _cooldown(seconds=60), - } - ) + redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10}) router = _router(redis, "usage-based-routing-v2") + _cool_down(router, "dep-b", started=time.time()) deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) @@ -168,14 +167,9 @@ async def test_batched_read_still_excludes_a_cooled_down_deployment(): @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, - "deployment:dep-b:cooldown": _cooldown(seconds=-1), - } - ) + redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10}) router = _router(redis, "usage-based-routing-v2") + _cool_down(router, "dep-b", started=time.time() - 61) deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) @@ -184,7 +178,7 @@ async def test_batched_read_ignores_an_expired_cooldown(): @pytest.mark.asyncio async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did(): - redis = MagicMock(spec=RedisCache) + 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") @@ -198,19 +192,100 @@ async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did(): 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" - before: Final = redis.async_batch_get_cache.await_count - redis.async_batch_get_cache.side_effect = None - redis.async_batch_get_cache.return_value = {} - assert await routed.cooldown_cache.async_get_active_cooldowns(["dep-a"], None) == [] - assert redis.async_batch_get_cache.await_count == before + 1 @pytest.mark.asyncio async def test_a_failed_batched_read_leaves_simple_shuffle_routing(): - redis = MagicMock(spec=RedisCache) + redis = redis_cache_double() redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) router = _router(redis, "simple-shuffle") deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} + + +def _armed_store() -> tuple[RouterCooldownStore, CooldownHandles]: + handles: Final = CooldownHandles() + store: Final = RouterCooldownStore(RedisCache(host="redis.internal"), handle=handles) + store.put("dep-a", {"exception_received": "429", "status_code": "429", "timestamp": 1.0, "cooldown_time": 60}, 60) + return store, handles + + +@pytest.mark.asyncio +async def test_routing_takes_the_armed_cooldown_read_from_the_native_batch_once_then_reads_on_its_own(): + store, handles = _armed_store() + natives: Final = FakeNatives() + + with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm_cooldowns(store, ["dep-a", "dep-b", "dep-c"]) + armed: Final = routing_prefetch.armed() + first: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None) + reads_after_first: Final = len(handles.built[0].targets) + second: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None) + + assert armed + assert handles.built[0].declared == [("dep-a", "dep-b", "dep-c")] + assert first == second == (("dep-a", store.snapshot()["dep-a"]),) + assert natives.built[0].events == ["flush_pre", "resolve"] + assert len(handles.built[0].targets) == reads_after_first + 1, "a second read goes to the store" + + +@pytest.mark.asyncio +async def test_an_armed_cooldown_read_that_does_not_cover_the_routing_ids_is_read_again(): + store, _ = _armed_store() + natives: Final = FakeNatives() + + with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm_cooldowns(store, ["dep-b"]) + active: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None) + + assert active == (("dep-a", store.snapshot()["dep-a"]),) + assert natives.built[0].events == ["flush_pre", "resolve"] + + +def _redis_down() -> list[tuple[str, CooldownCacheValue]]: + raise ConnectionError("redis down") + + +class _UnreachableCooldowns(FakeCooldowns): + def declare_active( + self, request: FakeNativeRequest, model_ids: Sequence[str] + ) -> FakeNativePending[list[tuple[str, CooldownCacheValue]]]: + super().declare_active(request, model_ids) + return FakeNativePending(request, _redis_down) + + +@pytest.mark.asyncio +async def test_a_failed_armed_cooldown_read_falls_back_to_the_stores_own_read(): + built: Final[list[FakeCooldowns]] = [] + + def unreachable(connection: RedisCache | None) -> FakeCooldowns: + built.append(_UnreachableCooldowns(None, None)) + return built[-1] + + store: Final = RouterCooldownStore(RedisCache(host="redis.internal"), handle=unreachable) + store.put("dep-a", {"exception_received": "429", "status_code": "429", "timestamp": 1.0, "cooldown_time": 60}, 60) + + with request_redis_batch_scope(native_factory=FakeNatives()): # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm_cooldowns(store, ["dep-a"]) + active: Final = await routing_prefetch.read_cooldowns(store, ["dep-a"], None) + + assert built[0].declared == [("dep-a",)] + assert active == (("dep-a", store.snapshot()["dep-a"]),) + + +def test_arming_the_cooldown_read_outside_a_request_or_without_redis_declares_nothing(): + store, handles = _armed_store() + offline: Final = CooldownHandles() + natives: Final = FakeNatives() + + routing_prefetch.arm_cooldowns(store, ["dep-a"]) + with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm_cooldowns(RouterCooldownStore(None, handle=offline), ["dep-a"]) + armed: Final = routing_prefetch.armed() + + assert handles.built[0].declared == [] + assert offline.built[0].declared == [] + assert not armed + assert natives.built == [] diff --git a/tests/unit/rust_bridge/host/test_redis_events.py b/tests/unit/rust_bridge/host/test_redis_events.py new file mode 100644 index 00000000000..6403d0b43c4 --- /dev/null +++ b/tests/unit/rust_bridge/host/test_redis_events.py @@ -0,0 +1,81 @@ +from typing import Final + +import pytest +from prometheus_client import REGISTRY + +import litellm +from litellm.rust_bridge.host import redis_events + + +def _sample(name: str, labels: dict[str, str]) -> float: + return REGISTRY.get_sample_value(name, labels) or 0.0 + + +@pytest.fixture +def prometheus(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + redis_events._service_logging.cache_clear() + yield + redis_events._service_logging.cache_clear() + + +@pytest.mark.parametrize( + ("failure", "failure_class"), + [("timeout", "timeout"), ("transport", "connectivity")], +) +def test_a_health_failure_counts_as_a_failed_request_and_a_breaker_failure( + prometheus: None, failure: redis_events.Failure, failure_class: str +): + failed: Final = ( + "litellm_redis_failed_requests_total", + {"redis": "redis", "error_class": "RedisCallError", "function_name": "async_increment"}, + ) + breaker: Final = ("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + before: Final = (_sample(*failed), _sample(*breaker)) + + redis_events.record_redis_call("async_increment", 0.25, failure, "caller.py") + + assert (_sample(*failed), _sample(*breaker)) == (before[0] + 1, before[1] + 1) + + +@pytest.mark.parametrize("failure", ["server", "decode"]) +def test_a_rejected_command_is_a_failed_request_but_not_a_breaker_failure( + prometheus: None, failure: redis_events.Failure +): + breaker: Final = { + failure_class: _sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + for failure_class in ("timeout", "connectivity") + } + failed: Final = ( + "litellm_redis_failed_requests_total", + {"redis": "redis", "error_class": "RedisCallError", "function_name": "async_increment"}, + ) + before: Final = _sample(*failed) + + redis_events.record_redis_call("async_increment", 0.25, failure, None) + + assert _sample(*failed) == before + 1 + assert { + failure_class: _sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + for failure_class in ("timeout", "connectivity") + } == breaker + + +def test_a_success_is_counted_and_a_breaker_rejection_is_not(prometheus: None): + total: Final = ("litellm_redis_total_requests_total", {"redis": "redis"}) + before: Final = _sample(*total) + + redis_events.record_redis_call("async_get_cache", 0.01, None, None) + redis_events.record_redis_call("async_get_cache", 0.0, "breaker_open", None) + + assert _sample(*total) == before + 1 + + +@pytest.mark.parametrize("state", ["open", "half_open", "closed"]) +def test_breaker_transitions_are_counted_by_state(prometheus: None, state: redis_events.BreakerState): + labels: Final = {"state": state} + before: Final = _sample("litellm_redis_circuit_breaker_transitions_total", labels) + + redis_events.record_redis_breaker(state) + + assert _sample("litellm_redis_circuit_breaker_transitions_total", labels) == before + 1 diff --git a/tests/unit/rust_bridge/host/test_settings.py b/tests/unit/rust_bridge/host/test_settings.py index 5d3486d9a64..a73eb41bd9b 100644 --- a/tests/unit/rust_bridge/host/test_settings.py +++ b/tests/unit/rust_bridge/host/test_settings.py @@ -4,6 +4,7 @@ import httpx import pytest import litellm +from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.custom_httpx.http_handler import default_user_agent from litellm.rust_bridge.host import settings @@ -46,6 +47,26 @@ def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch ) +def test_router_tiers_read_the_env_constants_and_the_litellm_global(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.constants + + monkeypatch.setattr(litellm.constants, "DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", 2.5) + monkeypatch.setattr(litellm.constants, "DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 42) + monkeypatch.setattr(litellm, "default_redis_batch_cache_expiry", 30.0) + + assert settings.router_tiers() == settings.RouterTierSettings( + cooldown_redis_read_interval=2.5, + usage_redis_read_interval=30.0, + max_redis_batch_cache_size=42, + ) + + +def test_router_tiers_fall_back_to_dual_caches_ten_second_refresh(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "default_redis_batch_cache_expiry", None) + + assert settings.router_tiers().usage_redis_read_interval == DualCache().redis_batch_cache_expiry + + def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_USER_AGENT", "operator/1") monkeypatch.setenv("SSL_VERIFY", "false") @@ -114,9 +135,7 @@ def test_secret_manager_is_native_only_when_a_readable_system_is_configured( monkeypatch.setattr(litellm, "_key_management_system", system) monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode)) - assert settings.secret_manager() == settings.SecretManager( - readable=access_mode != "write_only", native=native - ) + assert settings.secret_manager() == settings.SecretManager(readable=access_mode != "write_only", native=native) def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 4b6780fd058..af621a29567 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -32,7 +32,7 @@ _CACHE_PARAMETER_HELPERS: Final = frozenset( "litellm/llms/anthropic/prompt_cache_prediction.py", "litellm/proxy/common_utils/cache_coordinator.py", "litellm/proxy/common_utils/user_api_key_cache.py", - "litellm/redis/stores/_routing_prefetch.py", + "litellm/router_utils/routing_prefetch.py", "litellm/utils.py", } ) diff --git a/tests/unit/test_redis.py b/tests/unit/test_redis.py index 346e129dd2d..9b1149ba71d 100644 --- a/tests/unit/test_redis.py +++ b/tests/unit/test_redis.py @@ -1972,6 +1972,8 @@ def _tls_view(connection: AbstractConnection) -> dict[str, object] | None: context: Final = connection.ssl_context return { "ca_certs": context.ca_certs, + "ca_data": context.ca_data, + "ca_path": None, "certfile": context.certfile, "keyfile": context.keyfile, "verify_certificate": context.cert_reqs != ssl.CERT_NONE, @@ -2006,6 +2008,7 @@ def _connection_view(connection: AbstractConnection, kwargs: dict[str, object]) "auth": _auth_view(kwargs), "socket_timeout": connection.socket_timeout, "socket_connect_timeout": connection.socket_connect_timeout, + "health_check_interval": float(connection.health_check_interval) or None, } @@ -2047,6 +2050,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "tls": None, "socket_timeout": 5.0, "socket_connect_timeout": 5.0, + "health_check_interval": None, **fields, } @@ -2064,31 +2068,10 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), id="host_defaults", ), - pytest.param( - {"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}, - {}, - _expected_view( - { - "db": 2, - "host": "redis", - "mode": "host", - "port": 6380, - "socket_timeout": 10.5, - "socket_connect_timeout": 10.5, - } - ), - id="padded_and_underscored_numbers_parse_like_python_int_and_float", - ), - pytest.param( - {}, - {"REDIS_HOST": "redis", "REDIS_PORT": "6_380", "REDIS_DB": "3\n"}, - _expected_view({"db": 3, "host": "redis", "mode": "host", "port": 6380}), - id="padded_environment_numbers_parse_like_python_int", - ), pytest.param( {"startup_nodes": []}, {"REDIS_CLUSTER_NODES": '[{"host": "node-a", "port": 7000}]'}, - _expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}), + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}), id="empty_cluster_nodes_use_the_environment", ), pytest.param( @@ -2158,6 +2141,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "socket_timeout": 1.5, "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2184,6 +2169,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "port": 6390, "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2208,6 +2195,27 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: ), id="string_numbers_are_coerced", ), + pytest.param( + {"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}, + {}, + _expected_view( + { + "db": 2, + "host": "redis", + "mode": "host", + "port": 6380, + "socket_timeout": 10.5, + "socket_connect_timeout": 10.5, + } + ), + id="padded_and_underscored_numbers_parse_like_python_int_and_float", + ), + pytest.param( + {}, + {"REDIS_HOST": "redis", "REDIS_PORT": "6_380", "REDIS_DB": "3\n"}, + _expected_view({"db": 3, "host": "redis", "mode": "host", "port": 6380}), + id="padded_environment_numbers_parse_like_python_int", + ), pytest.param( {"host": "redis"}, {"REDIS_SOCKET_TIMEOUT": "0.75"}, @@ -2250,6 +2258,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "mode": "url", "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2273,6 +2283,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "mode": "url", "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2293,6 +2305,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: _expected_view( { "auth": {"kind": "static", "password": "cluster-secret", "username": None}, + "health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"], } @@ -2302,7 +2315,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: pytest.param( {}, {"REDIS_CLUSTER_NODES": '[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]'}, - _expected_view({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"]}), + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"]}), id="environment_cluster_nodes", ), pytest.param( @@ -2315,10 +2328,13 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: {}, _expected_view( { + "health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"], "tls": { "ca_certs": "/ca.pem", + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": True, "keyfile": None, @@ -2430,6 +2446,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "port": 6379, "tls": { "ca_certs": "/gcp-ca.pem", + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2512,6 +2530,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "port": 6379, "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2545,6 +2565,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "port": 6379, "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2576,6 +2598,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "mode": "url", "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2632,6 +2656,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "port": 6379, "tls": { "ca_certs": "/ca.pem", + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": True, "keyfile": None, @@ -2652,6 +2678,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "port": 6379, "tls": { "ca_certs": "/env-ca.pem", + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2707,6 +2735,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "port": 6379, "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": "/client.pem", "check_hostname": False, "keyfile": "/client.key", @@ -2724,6 +2754,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "mode": "url", "tls": { "ca_certs": "/ca.pem", + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": True, "keyfile": None, @@ -2755,6 +2787,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "mode": "url", "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2790,7 +2824,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: pytest.param( {"db": False, "startup_nodes": [{"host": "node-a", "port": 7000}]}, {}, - _expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}), + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}), id="cluster_db_false", ), pytest.param( @@ -2807,7 +2841,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: pytest.param( {"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}, {}, - _expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}), + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}), id="cluster_db_float_zero", ), pytest.param( @@ -2824,7 +2858,7 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: pytest.param( {"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}, {}, - _expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}), + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}), id="cluster_db_empty_list", ), pytest.param( @@ -2838,6 +2872,8 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: "port": 6379, "tls": { "ca_certs": None, + "ca_data": None, + "ca_path": None, "certfile": None, "check_hostname": False, "keyfile": None, @@ -2895,6 +2931,78 @@ def _expected_view(fields: dict[str, object]) -> dict[str, object]: _expected_view({"mode": "unix", "path": "/tmp/redis.sock"}), id="unix_socket_path", ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}, + {}, + _expected_view({"mode": "unix", "path": "/tmp/redis.sock"}), + id="unix_socket_path_keeps_tls_settings_out_of_the_target", + ), + pytest.param( + { + "host": "redis", + "ssl": True, + "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", + }, + {}, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": None, + "ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="ssl_ca_data", + ), + pytest.param( + {"host": "redis"}, + {"REDIS_SSL": "true", "REDIS_SSL_CA_PATH": "/etc/ssl/private-cas"}, + { + "rejected": { + "message": "Redis.__init__() got an unexpected keyword argument 'ssl_ca_path'", + "type": "TypeError", + } + }, + id="ssl_ca_path_from_the_environment", + ), + pytest.param( + {"health_check_interval": "7", "host": "redis"}, + {}, + _expected_view({"db": 0, "health_check_interval": 7.0, "host": "redis", "mode": "host", "port": 6379}), + id="health_check_interval_is_whole_seconds", + ), + pytest.param( + {"health_check_interval": 0, "host": "redis"}, + {}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="health_check_interval_zero_is_off", + ), + pytest.param( + {"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + { + "rejected": { + "message": "Argument 'db' must be 0 or None in cluster mode", + "type": "RedisClusterException", + } + }, + id="a_cluster_rejects_a_nonzero_db", + ), + pytest.param( + {"unix_socket_path": "/tmp/redis.sock"}, + {}, + {"rejected": {"message": "Either 'host' or 'url' must be specified for redis.", "type": "ValueError"}}, + id="unix_socket_path_without_a_host_fails_like_python", + ), ], ) def test_redis_settings_resolve_like_the_rust_resolver_expects(monkeypatch, kwargs, environment, expected): diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 980669b057a..dffadaa6383 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -105,12 +105,12 @@ from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE from tests.large_text import text +from tests.unit.conftest import redis_cache_double import traceback import inspect from typing import List, Optional if TYPE_CHECKING: - from tests.unit.conftest import RedisWriteGate, RouterRedis from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from litellm.integrations.otel.logger import OpenTelemetryV2 @@ -15606,11 +15606,10 @@ class TestClaudeCodeSubagentSessionRouterBinding: @pytest.mark.asyncio async def test_redis_cleanup_failure_does_not_reject_a_subagent_request(self): - from litellm.caching.caching import RedisCache router = self._router() del router.complexity_routers["smart-router"] - redis_cache = MagicMock(spec=RedisCache) + redis_cache = redis_cache_double() redis_cache.async_get_cache = AsyncMock(return_value="smart-router") redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) router.update_redis_cache(cache=redis_cache) @@ -15625,7 +15624,6 @@ class TestClaudeCodeSubagentSessionRouterBinding: @pytest.mark.asyncio async def test_redis_read_failure_does_not_reject_a_subagent_request(self): - from litellm.caching.caching import RedisCache router = self._router() request_kwargs = self._request_kwargs(agent_id="agent-1234") @@ -15635,7 +15633,7 @@ class TestClaudeCodeSubagentSessionRouterBinding: cache_key, "smart-router", ) - redis_cache = MagicMock(spec=RedisCache) + redis_cache = redis_cache_double() redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open")) router.update_redis_cache(cache=redis_cache) @@ -15650,10 +15648,9 @@ class TestClaudeCodeSubagentSessionRouterBinding: @pytest.mark.asyncio async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self): - from litellm.caching.caching import RedisCache router = self._router() - redis_cache = MagicMock(spec=RedisCache) + redis_cache = redis_cache_double() redis_cache.async_get_cache = AsyncMock(return_value="smart-router") redis_cache.async_set_cache = AsyncMock(side_effect=Exception("redis unavailable")) router.update_redis_cache(cache=redis_cache) @@ -15677,10 +15674,9 @@ class TestClaudeCodeSubagentSessionRouterBinding: async def test_subagents_follow_the_main_threads_latest_router_across_workers(self): from types import SimpleNamespace - from litellm.caching.caching import RedisCache shared_binding = SimpleNamespace(value=None) - shared_redis = MagicMock(spec=RedisCache) + shared_redis = redis_cache_double() shared_redis.async_get_cache = AsyncMock(side_effect=lambda key, **_: shared_binding.value) shared_redis.async_set_cache = AsyncMock( side_effect=lambda key, value, **_: setattr(shared_binding, "value", value) @@ -15708,11 +15704,10 @@ class TestClaudeCodeSubagentSessionRouterBinding: @pytest.mark.asyncio async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self): - from litellm.caching.caching import RedisCache router = self._router() router.complexity_routers.clear() - redis_cache = MagicMock(spec=RedisCache) + redis_cache = redis_cache_double() redis_cache.async_get_cache = AsyncMock(return_value=None) redis_cache.async_set_cache = AsyncMock() redis_cache.async_delete_cache = AsyncMock() @@ -24299,11 +24294,7 @@ def test_router_keeps_wif_secret_pointers_unresolved(monkeypatch): @pytest.mark.asyncio -async def test_failure_rpm_increment_declares_the_router_usage_key_family(router_redis: "RouterRedis"): - """The RPM bump a failed call still earns is router usage bookkeeping, so its Redis span - reads ``redis.incr router_usage`` rather than a bare ``redis.incr``.""" - from litellm._internal_context import current_service_target - +async def test_a_failed_call_still_counts_one_rpm(): router = Router( model_list=[ { @@ -24313,8 +24304,6 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(router } ] ) - redis, commands = router_redis - router.update_redis_cache(redis) await router.async_deployment_callback_on_failure( kwargs={ "call_type": "acompletion", @@ -24328,11 +24317,18 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(router end_time=None, ) - observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize())) - assert tuple( - target for args, target, _ in observed if isinstance(args[0], bytes) and b"INCRBYFLOAT" in args[0] - ) == ("router_usage",) - assert current_service_target() is None + assert _usage_by_kind(router) == {"rpm": 1} + + +def _usage_by_kind(router: Router) -> dict[str, float]: + return {key.rsplit(":", 2)[-2]: value for key, value in router._core_usage_store.snapshot().items()} + + +async def _usage_settles(router: Router, kinds: int) -> dict[str, float]: + async with asyncio.timeout(5): + while len(_usage_by_kind(router)) < kinds: + await asyncio.sleep(0.01) + return _usage_by_kind(router) @pytest.fixture @@ -24352,9 +24348,7 @@ def v2_span_exporter(monkeypatch): @pytest.mark.asyncio -async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_model_group( - v2_span_exporter, router_redis: "RouterRedis" -): +async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_model_group(v2_span_exporter): """Picking a deployment opens ``route {model_group}`` (the requested group, not the deployment it picks) under the server span, and the cooldown reads it issues run inside it, so their Redis spans nest there instead of lying flat under the request.""" @@ -24374,9 +24368,6 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo }, ] ) - redis, commands = router_redis - router.update_redis_cache(redis) - with TracerProvider().get_tracer("test").start_as_current_span("POST /v1/chat/completions") as server_span: deployment = await router.async_get_available_deployment(model="gpt-group", request_kwargs={}) @@ -24385,8 +24376,6 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo assert route_span.name == "route gpt-group" assert route_span.parent is not None and route_span.parent.span_id == server_span.get_span_context().span_id assert route_span.end_time is not None - observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize())) - assert tuple(span for args, _, span in observed if args[0] == "MGET") == ("route gpt-group",) def _record_phase_events(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, dict[str, str | int]]]: @@ -25141,7 +25130,7 @@ async def test_router_get_model_group_usage_wildcard_routes(): @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio -async def test_call_router_callbacks_on_success(router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"): +async def test_call_router_callbacks_on_success(): router = Router( model_list=[ { @@ -25152,26 +25141,18 @@ async def test_call_router_callbacks_on_success(router_redis: "RouterRedis", red ] ) - redis, _ = router_redis - started, _, _ = redis_write_gate - router.update_redis_cache(redis) response: Final = await router.acompletion( model="gemini/gemini-2.5-flash", messages=[{"role": "user", "content": "Hello, how are you?"}], mock_response="Hello, I'm good.", ) - await asyncio.wait_for(started.wait(), timeout=5) - client: Final = redis.init_async_client() - keys: Final = tuple(await client.keys("global_router:*")) - values: Final = {key.rsplit(b":", 2)[-2]: float(await client.get(key)) for key in keys} - assert len(keys) == 2 - assert values == {b"tpm": response.usage.total_tokens, b"rpm": 1} + assert await _usage_settles(router, kinds=2) == {"tpm": response.usage.total_tokens, "rpm": 1} @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.serial @pytest.mark.asyncio -async def test_call_router_callbacks_on_failure(router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"): +async def test_call_router_callbacks_on_failure(): router = Router( model_list=[ { @@ -25182,9 +25163,6 @@ async def test_call_router_callbacks_on_failure(router_redis: "RouterRedis", red ] ) - redis, _ = router_redis - started, _, _ = redis_write_gate - router.update_redis_cache(redis) with pytest.raises(litellm.RateLimitError): await router.acompletion( model="gemini/gemini-2.5-flash", @@ -25192,12 +25170,7 @@ async def test_call_router_callbacks_on_failure(router_redis: "RouterRedis", red mock_response="litellm.RateLimitError", num_retries=0, ) - await asyncio.wait_for(started.wait(), timeout=5) - client: Final = redis.init_async_client() - keys: Final = tuple(await client.keys("global_router:*")) - values: Final = {key.rsplit(b":", 2)[-2]: float(await client.get(key)) for key in keys} - assert len(keys) == 1 - assert values == {b"rpm": 1} + assert await _usage_settles(router, kinds=1) == {"rpm": 1} @pytest.mark.asyncio @@ -25352,7 +25325,7 @@ def test_router_get_model_access_groups(potential_access_group, expected_result) def test_router_redis_cache(): router = Router(model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}}]) - redis_cache = MagicMock() + redis_cache = redis_cache_double() router.update_redis_cache(cache=redis_cache) diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index aa30fb93166..fdb11ef01bf 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -501,7 +501,7 @@ async def test_deployment_callback_on_success(sync_mode): @pytest.mark.asyncio -async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(router_redis: "RouterRedis"): +async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(): import time model_list = [ @@ -512,8 +512,6 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( } ] router = Router(model_list=model_list) - redis, _ = router_redis - router.update_redis_cache(redis) standard_logging_payload = create_standard_logging_payload() standard_logging_payload["total_tokens"] = 17 @@ -534,8 +532,7 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( end_time=time.time(), ) assert counter is not None - tpm_key = f"global_router:ws-deployment:None:tpm:{counter.minute}" - assert await redis.async_get_cache(key=tpm_key) == 17 + assert await router._core_usage_store.async_batch_get((counter,)) == [17] realtime_key = await router.deployment_callback_on_success( kwargs=turn_kwargs("_arealtime"), @@ -544,7 +541,7 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( end_time=time.time(), ) assert realtime_key is None - assert await redis.async_get_cache(key=tpm_key) == 17 + assert await router._core_usage_store.async_batch_get((counter,)) == [17] async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]: @@ -1762,49 +1759,28 @@ def test_deployment_tracks_usage_for_every_limit_source( assert router._deployment_tracks_usage(deployment_id, "my-unmapped-model") is tracks -def _shared_redis_stub(store: dict) -> MagicMock: - from litellm.caching.redis_cache import RedisCache - - async def increment_pipeline(increment_list, **kwargs): - for op in increment_list: - store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"] - return [store[op["key"]] for op in increment_list] - - async def batch_get(keys, **kwargs): - return {key: store.get(key) for key in keys} - - redis_stub = MagicMock(spec=RedisCache) - redis_stub.async_increment_pipeline = increment_pipeline - redis_stub.async_batch_get_cache = batch_get - return redis_stub - - @pytest.mark.asyncio @pytest.mark.usefixtures("router_minute_pinned") -async def test_headers_on_fresh_worker_reflect_shared_redis_usage(): - store: dict = {} - worker_a = _rpm_tpm_router("lit-3058-workers") - worker_b = _rpm_tpm_router("lit-3058-workers") - worker_a.update_redis_cache(_shared_redis_stub(store)) - worker_b.update_redis_cache(_shared_redis_stub(store)) +async def test_headers_reflect_usage_counted_by_earlier_requests(): + router = _rpm_tpm_router("lit-3058-workers") messages = [{"role": "user", "content": "hi"}] - tokens_on_a = 0 + earlier_tokens = 0 for _ in range(3): - response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") - tokens_on_a += response.usage.total_tokens + response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + earlier_tokens += response.usage.total_tokens - response = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") headers = _ratelimit_headers(response) assert headers["x-ratelimit-remaining-requests"] == 96 - assert headers["x-ratelimit-remaining-tokens"] == 1000 - tokens_on_a - response.usage.total_tokens + assert headers["x-ratelimit-remaining-tokens"] == 1000 - earlier_tokens - response.usage.total_tokens - counted_tokens = tokens_on_a + response.usage.total_tokens + counted_tokens = earlier_tokens + response.usage.total_tokens for _ in range(2): - response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") counted_tokens += response.usage.total_tokens - stream = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True) + stream = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True) stream_headers = _ratelimit_headers(stream) assert stream_headers["x-ratelimit-remaining-requests"] == 93 assert stream_headers["x-ratelimit-remaining-tokens"] == 1000 - counted_tokens diff --git a/tests/unit/test_router/test_router_provider_endpoints.py b/tests/unit/test_router/test_router_provider_endpoints.py index d6ad013aa6d..8728de27503 100644 --- a/tests/unit/test_router/test_router_provider_endpoints.py +++ b/tests/unit/test_router/test_router_provider_endpoints.py @@ -5,7 +5,7 @@ import json import uuid from collections.abc import Callable, Iterator from datetime import datetime, timezone -from typing import TYPE_CHECKING, Final +from typing import Final import httpx import litellm @@ -24,8 +24,8 @@ from litellm.types.utils import ImageResponse, ModelResponse, RerankResponse from litellm.types.videos.utils import encode_video_id_with_provider from openai import AsyncAzureOpenAI, AsyncOpenAI -if TYPE_CHECKING: - from tests.unit.conftest import RedisWriteGate, RouterRedis + +from tests.unit.conftest import redis_cache_double EVENT_TIMEOUT_SECONDS: Final = 5 ROUTING_SELECTIONS: Final = 40 @@ -536,15 +536,9 @@ def test_router_track_deployment_metrics_updates_observable_usage() -> None: @pytest.mark.asyncio @pytest.mark.usefixtures("router_minute_pinned") -async def test_router_success_callback_during_pre_header_increment_does_not_double_count( - router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate" -) -> None: +async def test_router_success_callback_during_pre_header_increment_does_not_double_count() -> None: capture: Final = _RouterLoggingCapture("lit-3058-race") router: Final = _rpm_tpm_router("lit-3058-race") - redis, _ = router_redis - started, released, _ = redis_write_gate - router.update_redis_cache(redis) - released.clear() request: Final = asyncio.create_task( router.acompletion( model="gpt-5-mini", @@ -554,10 +548,8 @@ async def test_router_success_callback_during_pre_header_increment_does_not_doub ) ) - await asyncio.wait_for(started.wait(), timeout=EVENT_TIMEOUT_SECONDS) await capture.next_event() assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1 - released.set() response: Final = await request assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1) @@ -565,16 +557,10 @@ async def test_router_success_callback_during_pre_header_increment_does_not_doub @pytest.mark.asyncio @pytest.mark.usefixtures("router_minute_pinned") -async def test_router_failed_redis_increment_keeps_local_usage_without_double_counting( - router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate" -) -> None: +async def test_router_failed_redis_increment_keeps_local_usage_without_double_counting() -> None: capture: Final = _RouterLoggingCapture("lit-3058-fail") router: Final = _rpm_tpm_router("lit-3058-fail") - redis, _ = router_redis - started, released, unavailable = redis_write_gate - router.update_redis_cache(redis) - released.clear() - unavailable.set() + router.update_redis_cache(redis_cache_double()) metadata: Final[dict[str, object]] = {} request: Final = asyncio.create_task( router.acompletion( @@ -586,11 +572,8 @@ async def test_router_failed_redis_increment_keeps_local_usage_without_double_co ) ) - await asyncio.wait_for(started.wait(), timeout=EVENT_TIMEOUT_SECONDS) await capture.next_event() - assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30 assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1 - released.set() response: Final = await request assert response.usage.total_tokens == 30