diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8b299f5455b..725d2cfef41 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2479,6 +2479,7 @@ name = "litellm-cache-redis" version = "0.1.0" dependencies = [ "litellm-cache", + "r2d2", "redis", "redis-test", "serde_json", @@ -3564,6 +3565,17 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot", + "scheduled-thread-pool", +] + [[package]] name = "rand" version = "0.8.7" @@ -3700,6 +3712,7 @@ dependencies = [ "itoa", "num-bigint 0.5.1", "percent-encoding", + "r2d2", "ryu", "sha1_smol", "socket2 0.6.5", @@ -4096,6 +4109,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot", +] + [[package]] name = "schemars" version = "0.9.0" diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 43186faf3f7..074c8dcb0fd 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -3,7 +3,10 @@ use std::collections::{BinaryHeap, HashMap}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error}; +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, + Error, +}; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; const DEFAULT_TTL: Duration = Duration::from_secs(600); @@ -168,6 +171,55 @@ impl InMemoryCache { } } +impl ClaimCache for InMemoryCache +where + V: Clone + PartialEq + Send + Sync + 'static, +{ + fn claim_cache( + &self, + key: &str, + candidate: V, + eligible: &[V], + kwargs: CacheKwargs, + ) -> Result { + let now = (self.now)(); + let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + Self::evict(&mut state, self.max_size_in_memory, now); + let winner = match state.values.get(key) { + Some(existing) if eligible.is_empty() => existing.clone(), + Some(existing) if eligible.contains(existing) => existing.clone(), + _ => candidate, + }; + let expiration = now + self.get_ttl(&kwargs); + state.values.insert(key.into(), winner.clone()); + state.expirations.insert(key.into(), expiration); + state + .expiration_heap + .push(Reverse((expiration, key.into()))); + Ok(winner) + } +} + +impl CounterCache for InMemoryCache { + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + let now = (self.now)(); + let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + Self::evict(&mut state, self.max_size_in_memory, now); + let value = state.values.get(key).copied().unwrap_or_default() + amount; + let expiration = state + .expirations + .get(key) + .copied() + .unwrap_or_else(|| now + self.get_ttl(&kwargs)); + state.values.insert(key.into(), value); + state.expirations.insert(key.into(), expiration); + state + .expiration_heap + .push(Reverse((expiration, key.into()))); + Ok(value) + } +} + impl BaseCache for InMemoryCache { type Value = V; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index 370145bebef..e5831dfd6d5 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -3,7 +3,8 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, Error, get_cache, set_cache, + BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error, + get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -183,3 +184,51 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() { None ); } + +#[test] +fn claims_are_atomic_and_refresh_eligible_winners() { + let clock = clock(); + let cache = InMemoryCache::with_clock(Some(4), Some(Duration::from_secs(60)), { + let clock = clock.clone(); + move || Duration::from_secs(clock.load(Ordering::SeqCst)) + }); + let kwargs = CacheKwargs { + ttl: Some(Duration::from_secs(10)), + ..Default::default() + }; + assert_eq!( + cache + .claim_cache("affinity", "first".to_string(), &[], kwargs.clone()) + .unwrap(), + "first" + ); + clock.store(105, Ordering::SeqCst); + assert_eq!( + cache + .claim_cache( + "affinity", + "second".to_string(), + &["first".to_string(), "second".to_string()], + kwargs, + ) + .unwrap(), + "first" + ); + assert_eq!( + cache.expires_at("affinity").unwrap(), + Some(Duration::from_secs(115)) + ); +} + +#[test] +fn counters_increment_under_one_lock() { + let cache = InMemoryCache::::default(); + assert_eq!( + CounterCache::increment_cache(&cache, "counter", 1.5, CacheKwargs::default()).unwrap(), + 1.5 + ); + assert_eq!( + CounterCache::increment_cache(&cache, "counter", 2.0, CacheKwargs::default()).unwrap(), + 3.5 + ); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index a60813b6260..1a1bb505a7c 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,7 +7,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = "1.7.0" +redis = { version = "1.7.0", features = ["r2d2"] } +r2d2 = "0.8.10" tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index d4ca0cf0522..281ffe17534 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -1,15 +1,83 @@ -use std::sync::{Arc, Mutex, MutexGuard}; +use std::sync::{Arc, Mutex}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error, + BaseCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, + ClaimCache, CounterCache, Error, }; use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); +const REDIS_TIMEOUT: Duration = Duration::from_secs(5); +const REDIS_POOL_SIZE: u32 = 16; + +enum Connections { + Pool(r2d2::Pool), + Fixed(Mutex), +} + +struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); + +impl redis::ConnectionLike for ConnectionRef<'_> { + fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { + self.0.req_packed_command(cmd) + } + + fn req_packed_commands( + &mut self, + cmd: &[u8], + offset: usize, + count: usize, + ) -> redis::RedisResult> { + self.0.req_packed_commands(cmd, offset, count) + } + + fn get_db(&self) -> i64 { + self.0.get_db() + } + + fn supports_pipelining(&self) -> bool { + self.0.supports_pipelining() + } + + fn check_connection(&mut self) -> bool { + self.0.check_connection() + } + + fn is_open(&self) -> bool { + self.0.is_open() + } +} + +impl Connections +where + C: redis::ConnectionLike + Send + 'static, +{ + fn execute( + &self, + operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, + ) -> Result { + match self { + Self::Pool(pool) => { + let mut connection = pool.get().map_err(|_| Error::Unavailable)?; + connection + .set_read_timeout(Some(REDIS_TIMEOUT)) + .map_err(|_| Error::Unavailable)?; + connection + .set_write_timeout(Some(REDIS_TIMEOUT)) + .map_err(|_| Error::Unavailable)?; + operation(&mut ConnectionRef(&mut *connection)) + } + Self::Fixed(connection) => { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut ConnectionRef(&mut *connection)) + } + } + } +} pub struct RedisCache { - connection: Arc>, + connections: Arc>, default_ttl: Duration, codec: S, namespace: Option, @@ -18,8 +86,18 @@ pub struct RedisCache { impl RedisCache { pub fn new(url: &str, default_ttl: Option, codec: S) -> Result { let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; - let connection = client.get_connection().map_err(|_| Error::Unavailable)?; - Ok(Self::with_connection(connection, default_ttl, codec)) + let pool = r2d2::Pool::builder() + .max_size(REDIS_POOL_SIZE) + .min_idle(Some(0)) + .connection_timeout(REDIS_TIMEOUT) + .build(client) + .map_err(|_| Error::Unavailable)?; + Ok(Self { + connections: Arc::new(Connections::Pool(pool)), + default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + codec, + namespace: None, + }) } } @@ -30,17 +108,13 @@ where { pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { Self { - connection: Arc::new(Mutex::new(connection)), + connections: Arc::new(Connections::Fixed(Mutex::new(connection))), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, } } - fn connection(&self) -> Result, Error> { - self.connection.lock().map_err(|_| Error::Unavailable) - } - pub fn with_namespace(self, namespace: Option) -> Self { Self { namespace: namespace.filter(|value| !value.is_empty()), @@ -72,6 +146,29 @@ where Ok(format!("{escaped}:*")) } + fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> { + let mut cursor = 0u64; + loop { + let (next_cursor, keys): (u64, Vec) = redis::cmd("SCAN") + .cursor_arg(cursor) + .arg("MATCH") + .arg(pattern) + .arg("COUNT") + .arg(1000) + .query(connection) + .map_err(|_| Error::Unavailable)?; + if !keys.is_empty() { + connection + .del::<_, usize>(keys) + .map_err(|_| Error::Unavailable)?; + } + if next_cursor == 0 { + return Ok(()); + } + cursor = next_cursor; + } + } + fn decode_response(&self, value: redis::Value) -> Result, Error> { match value { redis::Value::Nil => Ok(None), @@ -81,23 +178,29 @@ where } } + fn decode_batch_response(&self, value: redis::Value) -> Result, Error> { + match self.decode_response(value) { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + } + } + fn ttl_seconds(ttl: Duration) -> u64 { ttl.as_secs() .saturating_add(u64::from(ttl.subsec_nanos() > 0)) .max(1) } - async fn run_blocking(connection: Arc>, operation: F) -> Result + async fn run_blocking(connections: Arc>, operation: F) -> Result where T: Send + 'static, - F: FnOnce(&mut C) -> Result + Send + 'static, + F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, { - tokio::task::spawn_blocking(move || { - let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; - operation(&mut connection) - }) - .await - .map_err(|_| Error::Unavailable)? + tokio::task::spawn_blocking(move || connections.execute(operation)) + .await + .map_err(|_| Error::Unavailable)? } } @@ -115,40 +218,55 @@ where fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { let payload = self.codec.encode(&value)?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - self.connection()? - .set_ex::<_, _, ()>(self.namespaced_key(key), payload, ttl) - .map_err(|_| Error::Unavailable) + let key = self.namespaced_key(key); + self.connections.execute(|connection| { + connection + .set_ex::<_, _, ()>(key, payload, ttl) + .map_err(|_| Error::Unavailable) + }) } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { - let value = self - .connection()? - .get::<_, redis::Value>(self.namespaced_key(key)) - .map_err(|_| Error::Unavailable)?; + let key = self.namespaced_key(key); + let value = self.connections.execute(|connection| { + connection + .get::<_, redis::Value>(key) + .map_err(|_| Error::Unavailable) + })?; self.decode_response(value) } + fn get_cache_batch( + &self, + keys: &[String], + _: &CacheKwargs, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = self.connections.execute(|connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + })?; + values + .into_iter() + .map(|value| self.decode_batch_response(value)) + .collect() + } + fn delete_cache(&self, key: &str) -> Result<(), Error> { - self.connection()? - .del::<_, ()>(self.namespaced_key(key)) - .map_err(|_| Error::Unavailable) + let key = self.namespaced_key(key); + self.connections + .execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)) } fn flush_cache(&self) -> Result<(), Error> { let pattern = self.namespaced_pattern()?; - let mut connection = self.connection()?; - let keys = connection - .scan_match(pattern) - .map_err(|_| Error::Unavailable)? - .collect::>>() - .map_err(|_| Error::Unavailable)?; - if keys.is_empty() { - return Ok(()); - } - connection - .del::<_, usize>(keys) - .map(|_| ()) - .map_err(|_| Error::Unavailable) + self.connections + .execute(|connection| Self::flush_matching(connection, &pattern)) } async fn async_set_cache( @@ -160,7 +278,7 @@ where let payload = self.codec.encode(&value)?; let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable) @@ -174,7 +292,7 @@ where _: &CacheKwargs, ) -> Result, Error> { let key = self.namespaced_key(key); - let value = Self::run_blocking(Arc::clone(&self.connection), move |connection| { + let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection .get::<_, redis::Value>(key) .map_err(|_| Error::Unavailable) @@ -183,6 +301,28 @@ where self.decode_response(value) } + async fn async_get_cache_batch( + &self, + keys: Vec, + _: CacheKwargs, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + values + .into_iter() + .map(|value| self.decode_batch_response(value)) + .collect() + } + async fn async_set_cache_pipeline( &self, cache_list: Vec<(String, Self::Value)>, @@ -197,41 +337,137 @@ where }) .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); for (key, payload) in entries { - connection - .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(|_| Error::Unavailable)?; + pipeline + .cmd("SETEX") + .arg(key) + .arg(ttl) + .arg(payload) + .ignore(); } - Ok(()) + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) }) .await } async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { let key = self.namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) .await } + async fn async_flush_cache(&self) -> Result<(), Error> { + let pattern = self.namespaced_pattern()?; + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Self::flush_matching(connection, &pattern) + }) + .await + } + async fn disconnect(&self) -> Result<(), Error> { Ok(()) } async fn test_connection(&self) -> Result { - Self::run_blocking(Arc::clone(&self.connection), |connection| { - redis::cmd("PING") - .query::(connection) + match Self::run_blocking(Arc::clone(&self.connections), |connection| { + Ok(match redis::cmd("PING").query::(connection) { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + }) + .await + { + Ok(result) => Ok(result), + Err(error) => Ok(CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }), + } + } +} + +impl CounterCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + const SCRIPT: &str = concat!( + "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", + "if redis.call('TTL', KEYS[1]) == -1 then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" + ); + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + self.connections.execute(|connection| { + redis::cmd("EVAL") + .arg(SCRIPT) + .arg(1) + .arg(key) + .arg(amount) + .arg(ttl) + .query(connection) .map_err(|_| Error::Unavailable) }) - .await?; - Ok(CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: "Redis cache connection test successful".into(), - error: None, - }) + } +} + +impl ClaimCache for RedisCache +where + S: CacheCodec, + S::Value: PartialEq, + C: redis::ConnectionLike + Send + 'static, +{ + fn claim_cache( + &self, + key: &str, + candidate: S::Value, + eligible: &[S::Value], + kwargs: CacheKwargs, + ) -> Result { + const SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if current == false then redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", + "return ARGV[1]; end; if #ARGV > 2 then for index = 3, #ARGV do ", + "if current == ARGV[index] then redis.call('EXPIRE', KEYS[1], ARGV[2]); ", + "return current; end; end; redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", + "return ARGV[1]; end; if current == ARGV[1] then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return current" + ); + let key = self.namespaced_key(key); + let candidate = self.codec.encode(&candidate)?; + let eligible = eligible + .iter() + .map(|value| self.codec.encode(value)) + .collect::, _>>()?; + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let value = self.connections.execute(|connection| { + redis::cmd("EVAL") + .arg(SCRIPT) + .arg(1) + .arg(key) + .arg(candidate) + .arg(ttl) + .arg(eligible) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + self.decode_response(value)?.ok_or(Error::Unavailable) } } @@ -302,7 +538,9 @@ mod tests { redis::cmd("SCAN") .cursor_arg(0) .arg("MATCH") - .arg("litellm-cache:*"), + .arg("litellm-cache:*") + .arg("COUNT") + .arg(1000), Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])), ), MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index d5bba19a8bd..6aef9bb36bf 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,6 +1,9 @@ use std::time::Duration; -use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, Error, JsonCodec, get_cache, set_cache}; +use litellm_cache::{ + BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, Error, JsonCodec, + get_cache, set_cache, +}; use litellm_cache_redis::RedisCache; use redis_test::{MockCmd, MockRedisConnection}; @@ -174,7 +177,9 @@ fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { redis::cmd("SCAN") .cursor_arg(0) .arg("MATCH") - .arg("team\\*:*"), + .arg("team\\*:*") + .arg("COUNT") + .arg(1000), Ok(redis_test::redis_value!(["0", ["team*:key"]])), ), MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)), @@ -184,3 +189,73 @@ fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { .with_namespace(Some("team*".into())); scoped.flush_cache().unwrap(); } + +#[tokio::test] +async fn connection_failures_use_the_python_result_contract() { + let error = redis::RedisError::from((redis::ErrorKind::Io, "connection refused")); + let connection = + MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Err::(error))]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + + let result = cache.test_connection().await.unwrap(); + assert_eq!(result.status, CacheConnectionStatus::Failed); + assert!(result.message.starts_with("Redis connection failed:")); + assert!(result.error.is_some()); +} + +#[tokio::test] +async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() { + let connection = MockRedisConnection::new([MockCmd::new( + redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"), + Ok(vec![ + redis::Value::BulkString(vec![42, 7]), + redis::Value::Nil, + redis::Value::BulkString(vec![99, 7]), + ]), + )]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + + assert_eq!( + cache + .async_get_cache_batch( + vec!["hit".into(), "miss".into(), "invalid".into()], + CacheKwargs::default(), + ) + .await + .unwrap(), + vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid] + ); +} + +#[tokio::test] +async fn async_flush_deletes_each_scan_page_separately() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("team:*") + .arg("COUNT") + .arg(1000), + Ok(redis_test::redis_value!(["7", ["team:a", "team:b"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)), + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(7) + .arg("MATCH") + .arg("team:*") + .arg("COUNT") + .arg(1000), + Ok(redis_test::redis_value!(["0", ["team:c"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + cache.async_flush_cache().await.unwrap(); +} diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 309296f7773..4e6694c191f 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -28,17 +28,17 @@ cache.store(&request, json!({"answer": 7}), now)?; assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7}))); ``` -For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Each Redis constructor currently opens its own connection; shared connection pools remain follow-up work +For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers move that blocking work off the executor Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it ## Python integration boundary -The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support synchronous and asynchronous response lookup and storage +The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution -Explicit facade registration checks object identity, method overrides, and configuration changes before selecting native execution. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python +Explicit facade registration checks object identity, method overrides, effective TTL, and configuration changes before selecting native execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy @@ -50,6 +50,6 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na ## Follow-up scope -Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial batches, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths +Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths Redis cluster, disk, cloud stores, dual caching, and semantic caching remain follow-ups. Atomic counters, affinity claims, reservations, queues, and pubsub need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/caching.rs b/litellm-rust/crates/cache-response/src/caching.rs index 53b34025ccd..afae4dfe4a5 100644 --- a/litellm-rust/crates/cache-response/src/caching.rs +++ b/litellm-rust/crates/cache-response/src/caching.rs @@ -120,6 +120,7 @@ impl CacheControls { pub fn writes(self) -> bool { self.supported_call_type && self.configured + && self.caching.unwrap_or(true) && !self.no_store && (self.default_on || self.use_cache) } @@ -131,13 +132,16 @@ pub fn should_use_cache(controls: CacheControls) -> bool { #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct CacheEntry { - pub timestamp: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub timestamp: Option, pub response: Value, } impl CacheEntry { pub fn fresh(&self, now: Duration, max_age: Option) -> bool { - self.timestamp.is_finite() - && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) + self.timestamp.is_none_or(|timestamp| { + timestamp.is_finite() + && max_age.is_none_or(|age| now.as_secs_f64() - timestamp <= age.as_secs_f64()) + }) } } diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index f1d55ddefe8..eaba3d0c349 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -9,7 +9,10 @@ impl CacheCodec for ResponseCacheCodec { type Value = CacheEntry; fn encode(&self, value: &CacheEntry) -> Result, Error> { - if !value.timestamp.is_finite() { + if value + .timestamp + .is_some_and(|timestamp| !timestamp.is_finite()) + { return Err(Error::InvalidEntry); } serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) @@ -17,12 +20,21 @@ impl CacheCodec for ResponseCacheCodec { fn decode(&self, bytes: &[u8]) -> Result { let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?; - let entry: CacheEntry = - serde_json::from_value(decode_value(text)?).map_err(|_| Error::InvalidEntry)?; - if !entry.timestamp.is_finite() { + let value = decode_value(text)?; + let Some(timestamp) = value.get("timestamp") else { + return Ok(CacheEntry { + timestamp: None, + response: value, + }); + }; + let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else { return Err(Error::InvalidEntry); - } - Ok(entry) + }; + let response = value.get("response").cloned().ok_or(Error::InvalidEntry)?; + Ok(CacheEntry { + timestamp: Some(timestamp), + response, + }) } } diff --git a/litellm-rust/crates/cache-response/src/embedding.rs b/litellm-rust/crates/cache-response/src/embedding.rs new file mode 100644 index 00000000000..d1f8a2bc0a6 --- /dev/null +++ b/litellm-rust/crates/cache-response/src/embedding.rs @@ -0,0 +1,22 @@ +use serde::Serialize; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct PartialHits { + pub values: Vec>, + pub missing_indices: Vec, +} + +impl PartialHits { + pub fn new(values: Vec>) -> Self { + let missing_indices = values + .iter() + .enumerate() + .filter_map(|(index, value)| value.is_none().then_some(index)) + .collect(); + Self { + values, + missing_indices, + } + } +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index efa0b04b9f7..72a507f8ee9 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,5 +1,6 @@ mod caching; mod codec; +mod embedding; mod response; pub use caching::{ @@ -7,4 +8,5 @@ pub use caching::{ get_cache_key, should_use_cache, }; pub use codec::ResponseCacheCodec; +pub use embedding::PartialHits; pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 987a47f0554..8f0fe953df6 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,8 +1,8 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{BaseCache, CacheKwargs, Error}; +use litellm_cache::{BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, Error}; -use crate::{CacheControls, CacheEntry, CacheKeyInput, cache_key}; +use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; use serde_json::Value; #[derive(Clone)] @@ -39,6 +39,18 @@ impl> ResponseCache { Self { backend } } + pub fn default_ttl(&self) -> Duration { + self.backend.default_ttl() + } + + pub async fn async_flush(&self) -> Result<(), Error> { + self.backend.async_flush_cache().await + } + + pub async fn test_connection(&self) -> Result { + self.backend.test_connection().await + } + pub fn lookup( &self, request: &ResponseCacheRequest, @@ -47,10 +59,15 @@ impl> ResponseCache { if !request.controls.reads() { return Ok(None); } - let entry = self + let entry = match self .backend - .get_cache(&cache_key(&request.key), &request.kwargs)?; - Self::fresh_response(entry, now, request.max_age) + .get_cache(&cache_key(&request.key), &request.kwargs) + { + Ok(entry) => entry, + Err(Error::InvalidEntry) => None, + Err(error) => return Err(error), + }; + Self::fresh_or_miss(entry, now, request.max_age) } pub async fn async_lookup( @@ -61,11 +78,62 @@ impl> ResponseCache { if !request.controls.reads() { return Ok(None); } - let entry = self + let entry = match self .backend .async_get_cache(&cache_key(&request.key), &request.kwargs) - .await?; - Self::fresh_response(entry, now, request.max_age) + .await + { + Ok(entry) => entry, + Err(Error::InvalidEntry) => None, + Err(error) => return Err(error), + }; + Self::fresh_or_miss(entry, now, request.max_age) + } + + pub fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + let readable = requests + .iter() + .enumerate() + .filter(|(_, request)| request.controls.reads()) + .collect::>(); + let keys = readable + .iter() + .map(|(_, request)| cache_key(&request.key)) + .collect::>(); + let entries = if let Some((_, request)) = readable.first() { + self.backend.get_cache_batch(&keys, &request.kwargs)? + } else { + Vec::new() + }; + Self::partial_hits(requests, readable, entries, now) + } + + pub async fn async_lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + let readable = requests + .iter() + .enumerate() + .filter(|(_, request)| request.controls.reads()) + .collect::>(); + let keys = readable + .iter() + .map(|(_, request)| cache_key(&request.key)) + .collect::>(); + let entries = if let Some((_, request)) = readable.first() { + self.backend + .async_get_cache_batch(keys, request.kwargs.clone()) + .await? + } else { + Vec::new() + }; + Self::partial_hits(requests, readable, entries, now) } pub fn store( @@ -80,7 +148,7 @@ impl> ResponseCache { self.backend.set_cache( &cache_key(&request.key), CacheEntry { - timestamp: now.as_secs_f64(), + timestamp: Some(now.as_secs_f64()), response, }, request.kwargs.clone(), @@ -100,7 +168,7 @@ impl> ResponseCache { .async_set_cache( &cache_key(&request.key), CacheEntry { - timestamp: now.as_secs_f64(), + timestamp: Some(now.as_secs_f64()), response, }, request.kwargs.clone(), @@ -108,6 +176,76 @@ impl> ResponseCache { .await } + pub async fn async_store_batch( + &self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> Result<(), Error> { + let writable = entries + .into_iter() + .filter(|(request, _)| request.controls.writes()) + .map(|(request, response)| { + ( + cache_key(&request.key), + CacheEntry { + timestamp: Some(now.as_secs_f64()), + response, + }, + request.kwargs, + ) + }) + .collect::>(); + let Some((_, _, first_kwargs)) = writable.first() else { + return Ok(()); + }; + if writable.iter().all(|(_, _, kwargs)| kwargs == first_kwargs) { + let kwargs = first_kwargs.clone(); + let cache_list = writable + .into_iter() + .map(|(key, entry, _)| (key, entry)) + .collect(); + return self + .backend + .async_set_cache_pipeline(cache_list, kwargs) + .await; + } + for (key, entry, kwargs) in writable { + self.backend.async_set_cache(&key, entry, kwargs).await?; + } + Ok(()) + } + + fn partial_hits( + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, + entries: Vec>, + now: Duration, + ) -> Result { + if readable.len() != entries.len() { + return Err(Error::Unavailable); + } + let mut values = vec![None; requests.len()]; + for ((index, request), entry) in readable.into_iter().zip(entries) { + let response = match entry { + BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age)?, + BatchEntry::Miss | BatchEntry::Invalid => None, + }; + values[index] = response; + } + Ok(PartialHits::new(values)) + } + + fn fresh_or_miss( + entry: Option, + now: Duration, + max_age: Option, + ) -> Result, Error> { + match Self::fresh_response(entry, now, max_age) { + Err(Error::InvalidEntry) => Ok(None), + result => result, + } + } + fn fresh_response( entry: Option, now: Duration, @@ -115,9 +253,9 @@ impl> ResponseCache { ) -> Result, Error> { entry .filter(|entry| entry.fresh(now, max_age)) - .map(|entry| match entry.response { - Value::String(text) => crate::codec::decode_value(&text), - value => Ok(value), + .map(|entry| match (entry.timestamp, entry.response) { + (Some(_), Value::String(text)) => crate::codec::decode_value(&text), + (_, value) => Ok(value), }) .transpose() } diff --git a/litellm-rust/crates/cache-response/tests/caching.rs b/litellm-rust/crates/cache-response/tests/caching.rs index d403791e421..0e8ce9b3b1d 100644 --- a/litellm-rust/crates/cache-response/tests/caching.rs +++ b/litellm-rust/crates/cache-response/tests/caching.rs @@ -80,4 +80,11 @@ fn cache_controls_honor_default_modes_and_directives() { } .writes() ); + assert!( + !CacheControls { + caching: Some(false), + ..enabled + } + .writes() + ); } diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 88b53fdfe86..e4a04b3dec5 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -245,7 +245,7 @@ fn response_codec_accepts_python_literals_without_executing_code() { assert_eq!( ResponseCacheCodec .encode(&CacheEntry { - timestamp: f64::NAN, + timestamp: Some(f64::NAN), response: json!({}) }) .unwrap_err(), @@ -254,7 +254,7 @@ fn response_codec_accepts_python_literals_without_executing_code() { } #[tokio::test] -async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redis() { +async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() { let connection = MockRedisConnection::new([MockCmd::new( redis::cmd("GET").arg("tenant:key"), Ok(b"invalid".to_vec()), @@ -267,22 +267,19 @@ async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redi assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None); request.controls.no_cache = false; assert_eq!( - cache - .async_lookup(&request, Duration::ZERO) - .await - .unwrap_err(), - Error::InvalidEntry + cache.async_lookup(&request, Duration::ZERO).await.unwrap(), + None ); } #[test] -fn malformed_memory_entries_are_rejected_by_the_response_consumer() { +fn malformed_memory_entries_are_treated_as_misses() { let backend = Arc::new(InMemoryCache::default()); BaseCache::set_cache( backend.as_ref(), "tenant:key", CacheEntry { - timestamp: 100.0, + timestamp: Some(100.0), response: json!("not a serialized response"), }, Default::default(), @@ -290,10 +287,8 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { .unwrap(); let cache = ResponseCache::new(backend); assert_eq!( - cache - .lookup(&request(), Duration::from_secs(100)) - .unwrap_err(), - Error::InvalidEntry + cache.lookup(&request(), Duration::from_secs(100)).unwrap(), + None ); } @@ -301,10 +296,74 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { fn response_entries_preserve_the_existing_json_representation() { let codec = ResponseCacheCodec; let entry = CacheEntry { - timestamp: 123.0, + timestamp: Some(123.0), response: json!({"choices": [{"text": "cached"}]}), }; let bytes = codec.encode(&entry).unwrap(); assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); assert_eq!(codec.decode(&bytes).unwrap(), entry); } + +#[test] +fn response_codec_preserves_values_without_timestamps() { + let codec = ResponseCacheCodec; + let raw = json!({"choices": [{"text": "legacy"}]}); + let entry = codec.decode(&serde_json::to_vec(&raw).unwrap()).unwrap(); + assert_eq!(entry.timestamp, None); + assert_eq!(entry.response, raw); + + let backend = Arc::new(InMemoryCache::default()); + BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, Default::default()).unwrap(); + let cache = ResponseCache::new(backend); + assert_eq!( + cache.lookup(&request(), Duration::from_secs(100)).unwrap(), + Some(json!({"choices": [{"text": "legacy"}]})) + ); +} + +#[tokio::test] +async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() { + let cache = memory(); + let requests = ["hit", "miss", "disabled"].map(|key| { + ResponseCacheRequest::new(CacheKeyInput { + preset: Some(key.into()), + ..Default::default() + }) + }); + cache + .store(&requests[0], json!({"value": 1}), Duration::from_secs(100)) + .unwrap(); + let mut requests = requests.to_vec(); + requests[2].controls.caching = Some(false); + + let partial = cache + .async_lookup_batch(&requests, Duration::from_secs(100)) + .await + .unwrap(); + assert_eq!(partial.values, vec![Some(json!({"value": 1})), None, None]); + assert_eq!(partial.missing_indices, vec![1, 2]); + + cache + .async_store_batch( + vec![ + (requests[1].clone(), json!({"value": 2})), + (requests[2].clone(), json!({"value": 3})), + ], + Duration::from_secs(100), + ) + .await + .unwrap(); + assert_eq!( + cache + .lookup(&requests[1], Duration::from_secs(100)) + .unwrap(), + Some(json!({"value": 2})) + ); + requests[2].controls.caching = None; + assert_eq!( + cache + .lookup(&requests[2], Duration::from_secs(100)) + .unwrap(), + None + ); +} diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 1891e417cb0..5bc1ebc2945 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -6,6 +6,13 @@ use serde_json::{Map, Value}; use crate::Error; +#[derive(Clone, Debug, PartialEq)] +pub enum BatchEntry { + Hit(V), + Miss, + Invalid, +} + #[derive(Clone, Debug, Default, PartialEq)] pub struct CacheKwargs { pub ttl: Option, @@ -42,6 +49,21 @@ pub trait BaseCache: Send + Sync { fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error>; + fn get_cache_batch( + &self, + keys: &[String], + kwargs: &CacheKwargs, + ) -> Result>, Error> { + keys.iter() + .map(|key| match self.get_cache(key, kwargs) { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + }) + .collect() + } + fn async_set_cache( &self, key: &str, @@ -59,6 +81,25 @@ pub trait BaseCache: Send + Sync { async move { self.get_cache(key, kwargs) } } + fn async_get_cache_batch( + &self, + keys: Vec, + kwargs: CacheKwargs, + ) -> impl Future>, Error>> + Send { + async move { + let mut entries = Vec::with_capacity(keys.len()); + for key in keys { + entries.push(match self.async_get_cache(&key, &kwargs).await { + Ok(Some(value)) => BatchEntry::Hit(value), + Ok(None) => BatchEntry::Miss, + Err(Error::InvalidEntry) => BatchEntry::Invalid, + Err(error) => return Err(error), + }); + } + Ok(entries) + } + } + fn async_set_cache_pipeline( &self, cache_list: Vec<(String, Self::Value)>, @@ -89,6 +130,10 @@ pub trait BaseCache: Send + Sync { fn flush_cache(&self) -> Result<(), Error>; + fn async_flush_cache(&self) -> impl Future> + Send { + async move { self.flush_cache() } + } + fn disconnect(&self) -> impl Future> + Send; fn test_connection(&self) -> impl Future> + Send; diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs new file mode 100644 index 00000000000..0b9deab1f5a --- /dev/null +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -0,0 +1,39 @@ +use std::future::Future; + +use crate::{BaseCache, CacheKwargs, Error}; + +pub trait CounterCache: BaseCache { + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result; + + fn async_increment_cache( + &self, + key: &str, + amount: f64, + kwargs: CacheKwargs, + ) -> impl Future> + Send { + async move { self.increment_cache(key, amount, kwargs) } + } +} + +pub trait ClaimCache: BaseCache +where + Self::Value: PartialEq, +{ + fn claim_cache( + &self, + key: &str, + candidate: Self::Value, + eligible: &[Self::Value], + kwargs: CacheKwargs, + ) -> Result; + + fn async_claim_cache( + &self, + key: &str, + candidate: Self::Value, + eligible: Vec, + kwargs: CacheKwargs, + ) -> impl Future> + Send { + async move { self.claim_cache(key, candidate, &eligible, kwargs) } + } +} diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs new file mode 100644 index 00000000000..17a2c430ddd --- /dev/null +++ b/litellm-rust/crates/cache/src/dual.rs @@ -0,0 +1,105 @@ +use std::sync::Arc; + +use crate::{BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error}; + +pub struct DualCache { + l1: Arc, + l2: Arc, +} + +impl DualCache { + pub fn new(l1: Arc, l2: Arc) -> Self { + Self { l1, l2 } + } +} + +impl BaseCache for DualCache +where + V: Clone + Send + Sync + 'static, + L1: BaseCache, + L2: BaseCache, +{ + type Value = V; + + fn default_ttl(&self) -> std::time::Duration { + self.l2.default_ttl() + } + + fn set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { + self.l2.set_cache(key, value.clone(), kwargs.clone())?; + self.l1.set_cache(key, value, kwargs) + } + + fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error> { + if let Some(value) = self.l1.get_cache(key, kwargs)? { + return Ok(Some(value)); + } + let value = self.l2.get_cache(key, kwargs)?; + if let Some(value) = &value { + self.l1.set_cache(key, value.clone(), kwargs.clone())?; + } + Ok(value) + } + + fn delete_cache(&self, key: &str) -> Result<(), Error> { + self.l2.delete_cache(key)?; + self.l1.delete_cache(key) + } + + fn flush_cache(&self) -> Result<(), Error> { + self.l2.flush_cache()?; + self.l1.flush_cache() + } + + async fn async_flush_cache(&self) -> Result<(), Error> { + self.l2.async_flush_cache().await?; + self.l1.async_flush_cache().await + } + + async fn disconnect(&self) -> Result<(), Error> { + self.l2.disconnect().await?; + self.l1.disconnect().await + } + + async fn test_connection(&self) -> Result { + self.l2.test_connection().await + } +} + +impl CounterCache for DualCache +where + L1: BaseCache, + L2: CounterCache, +{ + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + let value = self.l2.increment_cache(key, amount, kwargs.clone())?; + self.l1.set_cache(key, value, kwargs)?; + Ok(value) + } +} + +impl ClaimCache for DualCache +where + V: Clone + PartialEq + Send + Sync + 'static, + L1: ClaimCache, + L2: ClaimCache, +{ + fn claim_cache( + &self, + key: &str, + candidate: V, + eligible: &[V], + kwargs: CacheKwargs, + ) -> Result { + match self + .l2 + .claim_cache(key, candidate.clone(), eligible, kwargs.clone()) + { + Ok(winner) => { + self.l1.set_cache(key, winner.clone(), kwargs)?; + Ok(winner) + } + Err(_) => self.l1.claim_cache(key, candidate, eligible, kwargs), + } + } +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index 4ff02319bdc..ed67eb2fe15 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -1,9 +1,14 @@ mod base_cache; mod caching; +mod capabilities; mod codec; +pub mod dual; mod error; -pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs}; +pub use base_cache::{ + BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, +}; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; +pub use capabilities::{ClaimCache, CounterCache}; pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs new file mode 100644 index 00000000000..1be1556734c --- /dev/null +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -0,0 +1,130 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, dual::DualCache, +}; + +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; + + fn set_cache(&self, _: &str, value: V, _: CacheKwargs) -> Result<(), Error> { + *self.value.lock().unwrap() = Some(value); + Ok(()) + } + + fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + Ok(self.value.lock().unwrap().clone()) + } + + fn delete_cache(&self, _: &str) -> Result<(), Error> { + *self.value.lock().unwrap() = None; + Ok(()) + } + + fn flush_cache(&self) -> Result<(), Error> { + *self.value.lock().unwrap() = None; + Ok(()) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + unreachable!() + } +} + +impl CounterCache for TestCache { + fn increment_cache(&self, _: &str, amount: f64, _: CacheKwargs) -> Result { + if self.fail { + return Err(Error::Unavailable); + } + 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], + _: CacheKwargs, + ) -> Result { + if self.fail { + return Err(Error::Unavailable); + } + 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) + } +} + +#[test] +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, CacheKwargs::default()), + Err(Error::Unavailable) + ); + assert_eq!( + l1.get_cache("counter", &CacheKwargs::default()).unwrap(), + Some(10.0) + ); +} + +#[test] +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))); + + assert_eq!( + cache + .claim_cache( + "affinity", + "second".into(), + &["first".into(), "second".into()], + CacheKwargs { + ttl: Some(Duration::from_secs(60)), + ..Default::default() + }, + ) + .unwrap(), + "first" + ); +} diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 0af55083bef..15d0d603ac7 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -22,5 +22,8 @@ ], "secret_manager": [ "readable" + ], + "cache_settings": [ + "default_redis_ttl" ] } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index eb07118e964..58550c2987d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -6,6 +6,7 @@ use pyo3::{ types::{PyDict, PyTuple, PyType}, }; use serde_json::Value; +use std::time::Duration; use super::{NativeCacheHandle, native::NativeResponseCache}; @@ -126,7 +127,12 @@ impl ObjectGuard { } impl FacadeGuard { - pub(super) fn capture(py: Python<'_>, facade: &Bound<'_, PyAny>, kind: &str) -> PyResult { + pub(super) fn capture( + py: Python<'_>, + facade: &Bound<'_, PyAny>, + kind: &str, + native_default_ttl: Duration, + ) -> PyResult { let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( @@ -146,6 +152,12 @@ impl FacadeGuard { "facade and native backend types must match", )); } + let python_default_ttl = backend.getattr("default_ttl")?.extract::()?; + if python_default_ttl != native_default_ttl.as_secs_f64() { + return Err(PyTypeError::new_err( + "facade and native backend default TTLs must match", + )); + } Ok(Self { outer: ObjectGuard::capture( py, diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 5918967009a..4bab03fe243 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -4,7 +4,7 @@ mod native; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use litellm_cache::Error; -use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; +use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest}; use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; use pyo3::{ PyTraverseError, PyVisit, @@ -15,9 +15,26 @@ use pyo3::{ use serde::Deserialize; use serde_json::Value; +use crate::python_settings::PythonSettings; use facade::FacadeGuard; use native::NativeResponseCache; +const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); + +#[derive(FromPyObject)] +struct PythonCacheSettings { + default_redis_ttl: Option, +} + +fn redis_default_ttl(py: Python<'_>) -> PyResult { + let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; + settings + .default_redis_ttl + .map(duration) + .transpose() + .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct RequestInput { @@ -29,6 +46,10 @@ struct RequestInput { fn request(value: &Bound<'_, PyAny>) -> PyResult { let input: RequestInput = from_py(value)?; + request_input(input) +} + +fn request_input(input: RequestInput) -> PyResult { let mut request = ResponseCacheRequest::new(input.key); if let Some(controls) = input.controls { request.controls = controls; @@ -38,6 +59,13 @@ fn request(value: &Bound<'_, PyAny>) -> PyResult { Ok(request) } +fn requests(value: &Bound<'_, PyAny>) -> PyResult> { + from_py::>(value)? + .into_iter() + .map(request_input) + .collect() +} + fn duration(seconds: f64) -> PyResult { Duration::try_from_secs_f64(seconds) .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) @@ -94,7 +122,10 @@ impl NativeCacheHandle { ttl_seconds: Option, namespace: Option, ) -> PyResult { - let ttl = ttl_seconds.map(duration).transpose()?; + let ttl = Some(match ttl_seconds { + Some(seconds) => duration(seconds)?, + None => redis_default_ttl(py)?, + }); let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) .map_err(cache_error)?; Ok(Self { @@ -111,7 +142,12 @@ impl NativeCacheHandle { fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { let service = self.service()?; - let guard = FacadeGuard::capture(py, facade, self.backend())?; + let guard = FacadeGuard::capture(py, facade, self.backend(), service.default_ttl())?; + let service = service.with_redis_flush_size( + facade + .getattr("redis_flush_size")? + .extract::>()?, + ); let handle = Py::new( py, Self { @@ -249,6 +285,37 @@ impl ResolvedCache { } } + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn lookup_batch( + &self, + py: Python<'_>, + requests: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + to_py(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup_batch(&requests, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(object) => object + .bind(py) + .call_method( + "batch_get_cache", + (), + Some(self::callback_kwargs(callback_kwargs)?), + ) + .map(Bound::unbind), + } + } + #[pyo3(signature = (request, *, callback_kwargs=None))] fn async_lookup<'py>( &self, @@ -292,6 +359,102 @@ impl ResolvedCache { } } + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + ready_value(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup_batch(&requests, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_batch_get_cache", + (), + Some(self::callback_kwargs(callback_kwargs)?), + ), + } + } + + #[pyo3(signature = (requests, responses, *, callback_kwargs=None))] + fn async_store_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + responses: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let responses: Vec = from_py(responses)?; + if requests.len() != responses.len() { + return Err(PyValueError::new_err( + "batch cache requests and responses must have equal lengths", + )); + } + let entries = requests.into_iter().zip(responses).collect(); + let service = service.clone(); + run_async( + py, + async move { service.async_store_batch(entries, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_set_cache_pipeline", + (responses,), + Some(self::callback_kwargs(callback_kwargs)?), + ), + } + } + + fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async(py, async move { service.async_flush().await }, cache_error) + } + CacheBinding::PythonCallback(object) => { + object.bind(py).call_method0("flush_cache")?; + ready_none(py) + } + } + } + + fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async( + py, + async move { service.test_connection().await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method0("test_connection"), + } + } + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { if let CacheBinding::PythonCallback(object) = &self.binding { visit.call(object)?; @@ -309,11 +472,18 @@ fn callback_kwargs<'a, 'py>( } fn ready_none(py: Python<'_>) -> PyResult> { + ready_value(py, &()) +} + +fn ready_value<'py, T: serde::Serialize>( + py: Python<'py>, + value: &T, +) -> PyResult> { let future = py .import("asyncio")? .call_method0("get_running_loop")? .call_method0("create_future")?; - future.call_method1("set_result", (py.None(),))?; + future.call_method1("set_result", (to_py(py, value)?,))?; Ok(future) } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 6af04bfe2b2..20891719550 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,16 +1,27 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, Error}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use serde_json::Value; +use tokio::sync::Mutex; -use litellm_cache_response::{CacheEntry, ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; +use litellm_cache_response::{ + CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, +}; #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), - Redis(Arc>>), + Redis { + cache: Arc>>, + buffer: Option>, + }, +} + +pub(super) struct RedisWriteBuffer { + flush_size: usize, + entries: Mutex>, } impl NativeResponseCache { @@ -34,7 +45,10 @@ impl NativeResponseCache { namespace: Option, ) -> Result { let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace); - Ok(Self::Redis(Arc::new(ResponseCache::new(Arc::new(backend))))) + Ok(Self::Redis { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + buffer: None, + }) } } @@ -42,7 +56,29 @@ impl NativeResponseCache { pub fn kind(&self) -> &'static str { match self { Self::Memory(_) => "memory", - Self::Redis(_) => "redis", + Self::Redis { .. } => "redis", + } + } + + pub fn default_ttl(&self) -> Duration { + match self { + Self::Memory(cache) => cache.default_ttl(), + Self::Redis { cache, .. } => cache.default_ttl(), + } + } + + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { + match self { + Self::Redis { cache, .. } => Self::Redis { + cache, + buffer: flush_size.map(|flush_size| { + Arc::new(RedisWriteBuffer { + flush_size: flush_size.max(1), + entries: Mutex::new(Vec::new()), + }) + }), + }, + memory => memory, } } @@ -53,7 +89,7 @@ impl NativeResponseCache { ) -> Result, Error> { match self { Self::Memory(cache) => cache.lookup(request, now), - Self::Redis(cache) => cache.lookup(request, now), + Self::Redis { cache, .. } => cache.lookup(request, now), } } @@ -65,7 +101,18 @@ impl NativeResponseCache { ) -> Result<(), Error> { match self { Self::Memory(cache) => cache.store(request, response, now), - Self::Redis(cache) => cache.store(request, response, now), + Self::Redis { cache, .. } => cache.store(request, response, now), + } + } + + pub fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + match self { + Self::Memory(cache) => cache.lookup_batch(requests, now), + Self::Redis { cache, .. } => cache.lookup_batch(requests, now), } } @@ -76,7 +123,7 @@ impl NativeResponseCache { ) -> Result, Error> { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, - Self::Redis(cache) => cache.async_lookup(request, now).await, + Self::Redis { cache, .. } => cache.async_lookup(request, now).await, } } @@ -88,7 +135,71 @@ impl NativeResponseCache { ) -> Result<(), Error> { match self { Self::Memory(cache) => cache.async_store(request, response, now).await, - Self::Redis(cache) => cache.async_store(request, response, now).await, + Self::Redis { + cache, + buffer: None, + } => cache.async_store(request, response, now).await, + Self::Redis { + cache, + buffer: Some(buffer), + } => { + let pending = { + let mut entries = buffer.entries.lock().await; + entries.push((request.clone(), response)); + (entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries)) + }; + let Some(pending) = pending else { + return Ok(()); + }; + if let Err(error) = cache.async_store_batch(pending.clone(), now).await { + let mut entries = buffer.entries.lock().await; + let current = std::mem::take(&mut *entries); + *entries = pending.into_iter().chain(current).collect(); + return Err(error); + } + Ok(()) + } + } + } + + pub async fn async_lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + match self { + Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, + Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + } + } + + pub async fn async_store_batch( + &self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.async_store_batch(entries, now).await, + Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + } + } + + pub async fn async_flush(&self) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.async_flush().await, + Self::Redis { cache, buffer } => { + if let Some(buffer) = buffer { + buffer.entries.lock().await.clear(); + } + cache.async_flush().await + } + } + } + + pub async fn test_connection(&self) -> Result { + match self { + Self::Memory(cache) => cache.test_connection().await, + Self::Redis { cache, .. } => cache.test_connection().await, } } } diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 7ac23a05542..90819c5b3fc 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -8,15 +8,17 @@ pub(crate) enum PythonSettings { UrlPolicy, ProviderDefaults, SecretManager, + Cache, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 4] = [ + pub(crate) const ALL: [Self; 5] = [ Self::Http, Self::UrlPolicy, Self::ProviderDefaults, Self::SecretManager, + Self::Cache, ]; pub(crate) fn name(self) -> &'static str { @@ -25,6 +27,7 @@ impl PythonSettings { Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", Self::SecretManager => "secret_manager", + Self::Cache => "cache_settings", } } diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 66be77dbb40..64a5af2c618 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -12,7 +12,7 @@ import logging import time from collections.abc import Sequence from threading import Lock -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeVar if TYPE_CHECKING: from litellm.types.caching import RedisPipelineIncrementOperation @@ -34,14 +34,20 @@ else: from collections import OrderedDict +_KeyT = TypeVar("_KeyT") +_ValueT = TypeVar("_ValueT") -class LimitedSizeOrderedDict(OrderedDict): - def __init__(self, *args, max_size=100, **kwargs): - super().__init__(*args, **kwargs) + +class LimitedSizeOrderedDict(OrderedDict[_KeyT, _ValueT]): + def __init__(self, *, max_size: int = 100) -> None: + super().__init__() self.max_size = max_size - def __setitem__(self, key, value): - # If inserting a new key exceeds max size, remove the oldest item + def __setitem__(self, key: _KeyT, value: _ValueT) -> None: + if key in self: + super().__setitem__(key, value) + self.move_to_end(key) + return if len(self) >= self.max_size: self.popitem(last=False) super().__setitem__(key, value) @@ -68,7 +74,9 @@ class DualCache(BaseCache): self.in_memory_cache = in_memory_cache or InMemoryCache() # If redis_cache is not provided, use the default RedisCache self.redis_cache = redis_cache - self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size) + self.last_redis_batch_access_time: LimitedSizeOrderedDict[str, float] = LimitedSizeOrderedDict( + max_size=default_max_redis_batch_cache_size + ) self._last_redis_batch_access_time_lock = Lock() self.redis_batch_cache_expiry = ( default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10 @@ -131,7 +139,7 @@ class DualCache(BaseCache): except Exception as e: print_verbose(e) - def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> int: + def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> float: """ Key - the key in cache @@ -140,14 +148,15 @@ class DualCache(BaseCache): Returns - int - the incremented value """ try: - result: int = value - if self.in_memory_cache is not None: - result = self.in_memory_cache.increment_cache(key, value, **kwargs) - if self.redis_cache is not None and local_only is False: - result = self.redis_cache.increment_cache(key, value, **kwargs) + result: Final = self.redis_cache.increment_cache(key, value, **kwargs) + if self.in_memory_cache is not None: + self.in_memory_cache.set_cache(key, result, **kwargs) + return result - return result + if self.in_memory_cache is not None: + return self.in_memory_cache.increment_cache(key, value, **kwargs) + return value except Exception as e: verbose_logger.error("LiteLLM Cache: Excepton async add_cache: %s", e) raise e @@ -421,29 +430,30 @@ class DualCache(BaseCache): Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ - result: float | None = None try: - if self.in_memory_cache is not None: - result = await self.in_memory_cache.async_increment(key, value, **kwargs) - if self.redis_cache is not None and local_only is False: - result = await self.redis_cache.async_increment( + result: Final = await self.redis_cache.async_increment( key, value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), refresh_ttl=refresh_ttl, ) + if self.in_memory_cache is not None: + await self.in_memory_cache.async_set_cache(key, result, **kwargs) + return result - return result + if self.in_memory_cache is not None: + return await self.in_memory_cache.async_increment(key, value, **kwargs) + return None except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache failed, falling back to in-memory result", + "Redis async_increment_cache failed; local counter unchanged", e, ) - return result + return None async def async_increment_cache_pipeline( self, @@ -452,29 +462,32 @@ class DualCache(BaseCache): parent_otel_span: Span | None = None, **kwargs, ) -> list[float] | None: - result: list[float] | None = None try: - if self.in_memory_cache is not None: - result = await self.in_memory_cache.async_increment_pipeline( - increment_list=increment_list, - parent_otel_span=parent_otel_span, - ) - if self.redis_cache is not None and local_only is False: - result = await self.redis_cache.async_increment_pipeline( + result: Final = await self.redis_cache.async_increment_pipeline( increment_list=increment_list, parent_otel_span=parent_otel_span, ) + if result is not None and self.in_memory_cache is not None: + await self.in_memory_cache.async_set_cache_pipeline( + cache_list=tuple((increment["key"], value) for increment, value in zip(increment_list, result)) + ) + return result - return result + if self.in_memory_cache is not None: + return await self.in_memory_cache.async_increment_pipeline( + increment_list=increment_list, + parent_otel_span=parent_otel_span, + ) + return None except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache_pipeline failed, falling back to in-memory result", + "Redis async_increment_cache_pipeline failed; local counters unchanged", e, ) - return result + return None async def async_set_cache_sadd(self, key, value: list, local_only: bool = False, **kwargs) -> None: """ diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 4fd2f0829a3..68b1742dc41 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -126,6 +126,12 @@ class CacheBinding: *, callback_kwargs: dict[str, object] | None = None, ) -> None: ... + def lookup_batch( + self, + requests: Sequence[Mapping[str, object]], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> object: ... def async_lookup( self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None ) -> Awaitable[object]: ... @@ -136,6 +142,21 @@ class CacheBinding: *, callback_kwargs: dict[str, object] | None = None, ) -> Awaitable[object]: ... + def async_lookup_batch( + self, + requests: Sequence[Mapping[str, object]], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> Awaitable[object]: ... + def async_store_batch( + self, + requests: Sequence[Mapping[str, object]], + responses: Sequence[object], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> Awaitable[object]: ... + def async_flush(self) -> Awaitable[None]: ... + def ping(self) -> Awaitable[dict[str, object] | None]: ... @final class TokenCounter: diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 3aa2d742862..862a116496d 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -36,6 +36,11 @@ class SecretManager: readable: bool +@dataclass(frozen=True, slots=True) +class CacheSettings: + default_redis_ttl: float | None + + def warn(message: str) -> None: from litellm._logging import verbose_logger @@ -50,6 +55,12 @@ def secret_manager() -> SecretManager: return SecretManager(readable=_should_read_secret_from_secret_manager()) +def cache_settings() -> CacheSettings: + import litellm + + return CacheSettings(default_redis_ttl=litellm.default_redis_ttl) + + def provider_defaults() -> ProviderDefaults: import litellm diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 5f59de9cca5..149c9b34bd5 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -6,18 +6,16 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE -from litellm.caching.dual_cache import DualCache +from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync +from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from litellm.types.caching import RedisPipelineIncrementOperation @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads(): - dual_cache = DualCache( - redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) keys = ["shared_a", "shared_b"] start_gate = asyncio.Event() @@ -44,9 +42,7 @@ async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_error(): - dual_cache = DualCache( - redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) keys = ["shared_a", "shared_b"] with patch.object( @@ -116,9 +112,7 @@ def test_dual_cache_batch_get_cache_only_reads_missing_keys_from_redis(): def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): mock_redis = _redis_mock_for_sync_batch({"absent_key": None}) - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) first = dual_cache.batch_get_cache(keys=["absent_key"]) second = dual_cache.batch_get_cache(keys=["absent_key"]) @@ -131,9 +125,7 @@ def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): mock_redis = MagicMock(spec=RedisCache) mock_redis.batch_get_cache.side_effect = RuntimeError("redis unavailable") - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) first_result = dual_cache.batch_get_cache(keys=["shared_a"]) second_result = dual_cache.batch_get_cache(keys=["shared_a"]) @@ -146,9 +138,7 @@ def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled(): mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"}) - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) dual_cache.last_redis_batch_access_time["throttled_key"] = time.time() result = dual_cache.batch_get_cache(keys=["throttled_key"]) @@ -257,9 +247,7 @@ async def test_dual_cache_batch_redis_backfill_injects_default_in_memory_ttl(): default_in_memory_ttl, same as the single-key path.""" in_memory_cache = InMemoryCache(default_ttl=600) mock_redis = MagicMock(spec=RedisCache) - mock_redis.async_batch_get_cache = AsyncMock( - return_value={"batch_backfill_key": "redis_value"} - ) + mock_redis.async_batch_get_cache = AsyncMock(return_value={"batch_backfill_key": "redis_value"}) dual_cache = DualCache( in_memory_cache=in_memory_cache, redis_cache=mock_redis, @@ -371,9 +359,7 @@ async def test_circuit_breaker_open_skips_redis(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker( - failure_threshold=3, recovery_timeout=60 - ) + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) self._circuit_breaker._state = "open" self._circuit_breaker._opened_at = time.time() self.call_count = 0 @@ -426,9 +412,7 @@ def test_circuit_breaker_half_open_concurrent_calls_are_fast_failed(): # All subsequent concurrent callers: HALF_OPEN → fast-fail (return True) for _ in range(10): - assert ( - cb.is_open() is True - ), "concurrent callers should be fast-failed in HALF_OPEN" + assert cb.is_open() is True, "concurrent callers should be fast-failed in HALF_OPEN" def test_circuit_breaker_disabled_never_opens(): @@ -472,9 +456,7 @@ async def test_circuit_breaker_disabled_guard_always_calls_method(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker( - failure_threshold=1, recovery_timeout=60, enabled=False - ) + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False) self.call_count = 0 @_redis_circuit_breaker_guard @@ -512,6 +494,30 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re ) +@pytest.mark.asyncio +async def test_failed_redis_increment_does_not_change_the_local_counter(): + memory = InMemoryCache() + memory.set_cache("counter", 10) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment = AsyncMock(side_effect=RuntimeError("redis down")) + cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) + + assert await cache.async_increment_cache("counter", 2) is None + assert memory.get_cache("counter") == 10 + + +@pytest.mark.asyncio +async def test_successful_redis_increment_replaces_the_local_counter_with_the_authoritative_value(): + memory = InMemoryCache() + memory.set_cache("counter", 10) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment = AsyncMock(return_value=42.0) + cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) + + assert await cache.async_increment_cache("counter", 2) == 42.0 + assert memory.get_cache("counter") == 42.0 + + def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): """ Typical lazy startup (sync): DualCache runs with in-memory only, then Redis @@ -742,7 +748,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in visible] == [ ( logging.WARNING, - "Redis async_increment_cache_pipeline failed, falling back to in-memory result:" + "Redis async_increment_cache_pipeline failed; local counters unchanged:" " Timeout reading from 127.0.0.1:6379", ) ] @@ -756,7 +762,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ ( logging.WARNING, - "Redis async_increment_cache failed, falling back to in-memory result: Timeout reading from 127.0.0.1:6379" + "Redis async_increment_cache failed; local counter unchanged: Timeout reading from 127.0.0.1:6379" " (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", ) ] @@ -791,3 +797,14 @@ async def test_async_delete_cache_keys_on_empty_list_touches_no_backend(): await dual_cache.async_delete_cache_keys([]) redis_cache.delete_cache_keys.assert_not_awaited() + + +def test_limited_ordered_dict_refreshes_recency_without_evicting_another_key(): + tracker = LimitedSizeOrderedDict(max_size=2) + tracker["hot"] = 1 + tracker["cold"] = 2 + + tracker["hot"] = 3 + tracker["new"] = 4 + + assert list(tracker.items()) == [("hot", 3), ("new", 4)] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 493baac228a..7456a1499f7 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -8,6 +8,7 @@ import weakref from collections.abc import Generator from types import SimpleNamespace from typing import Final, Protocol, cast +from urllib.parse import urlparse import fakeredis import pytest @@ -209,8 +210,12 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) + client.set("team:raw", json.dumps(response)) + client.set("team:invalid", "not a cache entry") assert binding.lookup(request("sync")) == response assert await binding.async_lookup(request("team:async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) stored: Final = client.get("team:native") assert isinstance(stored, bytes) @@ -245,3 +250,63 @@ async def test_memory_size_policy_is_applied_by_the_native_host() -> None: ).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None + + +async def test_native_batch_lookup_and_store_report_partial_hits() -> None: + binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + requests: Final = [request("hit"), request("miss"), request("disabled")] + requests[2]["controls"] = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": False, + "no_cache": False, + "no_store": False, + "use_cache": False, + } + await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) + + partial: Final = await binding.async_lookup_batch(requests) + + assert partial == { + "values": [{"value": 1}, {"value": 2}, None], + "missing_indices": [2], + } + + +async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: + client: Final = redis.Redis.from_url(redis_url) + with rebound(litellm, "default_redis_ttl", 7): + binding: Final = _native.CacheResolver( + SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url)) + ).resolve() + await binding.async_store(request("native-default"), {"value": 1}) + + assert 0 < client.ttl("native-default") <= 7 + client.close() + + +async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: + parsed: Final = urlparse(redis_url) + with rebound(litellm, "default_redis_ttl", 60): + facade: Final = Cache( + type=LiteLLMCacheType.REDIS, + host=parsed.hostname, + port=str(parsed.port), + redis_flush_size=2, + ) + with pytest.raises(TypeError, match="default TTLs must match"): + _native.NativeCacheHandle.redis(redis_url, ttl_seconds=61).bind_facade(facade) + _native.NativeCacheHandle.redis(redis_url).bind_facade(facade) + binding: Final = _native.CacheResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(redis_url) + + await binding.async_store(request("first"), {"value": 1}) + assert client.get("first") is None + await binding.async_store(request("second"), {"value": 2}) + + assert client.get("first") is not None + assert client.get("second") is not None + await facade.cache.disconnect() + client.close()