diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 191bd2418e5..2d151d0c283 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4535,6 +4535,7 @@ dependencies = [ "litellm-redis-testing", "litellm-secret-value", "litellm-tracing", + "moka", "rstest", "strum", "thiserror 2.0.19", diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index e430cfa6bca..621fcc806a6 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -43,7 +43,7 @@ mod _native { #[pymodule_export] use crate::redis::client::NativeRedisClient; #[pymodule_export] - use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes}; + use crate::redis::store::auth_cache::{NativeAuthCache, NativeAuthSubscription, PyLocalSizes}; #[pymodule_export] use crate::redis::store::rate_limits::{ NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired, @@ -273,6 +273,7 @@ mod tests { "NativeRequestRedis", "NativeSend", "NativeAuthCache", + "NativeAuthSubscription", "NativeRateLimits", "WindowCount", "AtomicApplied", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs index 8d7788b136e..0d13dda3d3b 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs @@ -1,5 +1,9 @@ use std::sync::Arc; +use futures_util::{StreamExt, stream::BoxStream}; +use litellm_redis_client::SubscriptionEvent; +use tokio::sync::Mutex; + use litellm_python_compat::Value; use litellm_redis_store::auth_cache::{AuthCache, Settings}; use pyo3::prelude::*; @@ -44,6 +48,35 @@ impl NativeAuthCache { self.0.has_redis() } + fn attach(&self, client: &NativeRedisClient) -> bool { + self.0.attach(remote(client)) + } + + fn current_generation(&self) -> u64 { + self.0.current_generation() + } + + #[pyo3(signature = (key, value, generation, ttl_seconds = None))] + fn set_if_generation( + &self, + key: &str, + value: PyValue, + generation: u64, + ttl_seconds: Option, + ) -> PyResult { + Ok(self + .0 + .set_if_generation(key, value.0, generation, ttl(ttl_seconds)?)) + } + + fn subscribe(&self, client: &NativeRedisClient, channel: &str) -> NativeAuthSubscription { + self.0.subscriber_connected(false); + NativeAuthSubscription { + cache: self.0.clone(), + stream: Arc::new(Mutex::new(Some(client.client.subscribe_events(channel)))), + } + } + /// `request` answers a key its pre-call reads already found absent without another `GET`. #[pyo3(signature = (key, request = None))] fn get<'py>( @@ -245,6 +278,65 @@ impl NativeAuthCache { } } +#[pyclass(frozen)] +pub(crate) struct NativeAuthSubscription { + cache: Arc, + stream: Arc>>>, +} + +#[pymethods] +impl NativeAuthSubscription { + fn is_connected(&self) -> bool { + self.cache.local_ready() + } + + fn next<'py>(&self, py: Python<'py>) -> PyResult> { + let cache = self.cache.clone(); + let stream = self.stream.clone(); + litellm_host_python::run_async_value(py, async move { + let mut guard = stream.lock().await; + let Some(stream) = guard.as_mut() else { + return Ok(None::>); + }; + match stream.next().await { + Some(SubscriptionEvent::Connected) => cache.subscriber_connected(true), + Some(SubscriptionEvent::Message(payload)) => { + if let Ok(message) = serde_json::from_slice::(&payload) + && let Some(key) = + message.get("cache_key").and_then(serde_json::Value::as_str) + && message + .get("new_value") + .and_then(serde_json::Value::as_f64) + .is_none() + { + cache.delete_local(key).map_err(kind_error)?; + } + return Ok(Some(Python::attach(|py| { + pyo3::types::PyBytes::new(py, &payload).unbind() + }))); + } + _ => cache.subscriber_connected(false), + } + Ok(None) + }) + } + + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + self.cache.subscriber_connected(false); + let stream = self.stream.clone(); + litellm_host_python::run_async_value(py, async move { + stream.lock().await.take(); + Ok(()) + }) + } +} + +impl Drop for NativeAuthSubscription { + fn drop(&mut self) { + self.cache.subscriber_connected(false); + } +} + /// Entries the management-object and the user-key partitions hold in memory. #[pyclass(frozen, eq, get_all, name = "LocalSizes")] #[derive(PartialEq)] diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index d9a8b3af7e8..4f477c3ee3c 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -155,6 +155,10 @@ impl RedisClient { self.0.subscriber.clone().messages(channel.to_owned()) } + pub fn subscribe_events(&self, channel: &str) -> BoxStream<'static, crate::SubscriptionEvent> { + self.0.subscriber.clone().events(channel.to_owned()) + } + pub fn pipeline(&self) -> Pipeline { Pipeline::new(self.0.namespace.clone()) } diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 2754666c329..1a995334ced 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -23,6 +23,7 @@ pub use litellm_redis_types::{ }; pub use namespace::Namespace; pub use pipeline::Pipeline; +pub use pubsub::SubscriptionEvent; pub use redis::{ BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value, cluster_routing::Slot, diff --git a/litellm-rust/crates/redis-client/src/pubsub.rs b/litellm-rust/crates/redis-client/src/pubsub.rs index b6b18d3e481..a0004a239b6 100644 --- a/litellm-rust/crates/redis-client/src/pubsub.rs +++ b/litellm-rust/crates/redis-client/src/pubsub.rs @@ -1,7 +1,10 @@ use std::{sync::Arc, time::Duration}; use futures_util::{StreamExt, stream::BoxStream}; -use redis::{ConnectionInfo, RedisResult, aio::PubSubStream}; +use redis::{ + ConnectionInfo, RedisResult, + aio::{PubSubSink, PubSubStream}, +}; use crate::{ ConnectError, Credentials, Standalone, Target, Timeouts, config, @@ -11,6 +14,18 @@ use crate::{ // litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py: _BACKOFF_INITIAL_SECONDS, _BACKOFF_MAX_SECONDS const BACKOFF_INITIAL: Duration = Duration::from_secs(5); const BACKOFF_MAX: Duration = Duration::from_secs(60); +const IDLE_PROBE: Duration = Duration::from_secs(15); + +pub enum SubscriptionEvent { + Connected, + Disconnected, + Message(Vec), +} + +struct Live { + sink: PubSubSink, + stream: PubSubStream, +} /// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends /// `SUBSCRIBE`, authenticated once when it connects. @@ -66,43 +81,76 @@ impl Subscriber { } pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec> { + self.events(channel) + .filter_map(|event| async move { + match event { + SubscriptionEvent::Message(payload) => Some(payload), + _ => None, + } + }) + .boxed() + } + + pub(crate) fn events(self, channel: String) -> BoxStream<'static, SubscriptionEvent> { futures_util::stream::unfold( - (self, channel, None::, BACKOFF_INITIAL), - |(subscriber, channel, live, backoff)| async move { + (self, channel, None::, BACKOFF_INITIAL, false), + |(subscriber, channel, live, backoff, retry)| async move { let mut live = live; let mut backoff = backoff; + if retry { + tokio::time::sleep(backoff).await; + } loop { let Some(current) = live.as_mut() else { - match subscriber.subscribe(&channel).await { + let event = match subscriber.subscribe(&channel).await { Ok(subscribed) => { live = Some(subscribed); backoff = BACKOFF_INITIAL; + SubscriptionEvent::Connected } - Err(error) => { - tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying"); - tokio::time::sleep(backoff).await; - backoff = (backoff * 2).min(BACKOFF_MAX); + Err(_) => { + if retry { + backoff = (backoff * 2).min(BACKOFF_MAX); + } + SubscriptionEvent::Disconnected } - } - continue; + }; + let retry = live.is_none(); + return Some((event, (subscriber, channel, live, backoff, retry))); }; - match current.next().await { - Some(message) => { + match tokio::time::timeout(IDLE_PROBE, current.stream.next()).await { + Ok(Some(message)) => { let payload = message.get_payload_bytes().to_vec(); - return Some((payload, (subscriber, channel, live, backoff))); + return Some(( + SubscriptionEvent::Message(payload), + (subscriber, channel, live, backoff, false), + )); } - None => live = None, + Err(_) => { + if matches!( + tokio::time::timeout( + subscriber.connect_timeout, + current.sink.ping::() + ) + .await, + Ok(Ok(_)) + ) { + continue; + } + } + Ok(None) => {} } - tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing"); - tokio::time::sleep(backoff).await; - backoff = (backoff * 2).min(BACKOFF_MAX); + return Some(( + SubscriptionEvent::Disconnected, + (subscriber, channel, None, backoff, true), + )); } }, ) .boxed() } - async fn subscribe(&self, channel: &str) -> RedisResult { + async fn subscribe(&self, channel: &str) -> RedisResult { match &self.source { Source::Sentinel(sentinel) => { let master = sentinel.master().await?; @@ -122,7 +170,7 @@ impl Subscriber { } } - async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult { + async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult { let settings = config::authenticated(info.redis_settings().clone(), &self.credentials).await?; let client = redis::Client::open(info.set_redis_settings(settings))?; @@ -135,7 +183,8 @@ impl Subscriber { )) })??; subscribe_connected(&mut pubsub, channel, self.connect_timeout).await?; - Ok(pubsub.into_on_message()) + let (sink, stream) = pubsub.split(); + Ok(Live { sink, stream }) } } diff --git a/litellm-rust/crates/redis-client/tests/pubsub.rs b/litellm-rust/crates/redis-client/tests/pubsub.rs index 2ad7cb38a5a..e3fda48d765 100644 --- a/litellm-rust/crates/redis-client/tests/pubsub.rs +++ b/litellm-rust/crates/redis-client/tests/pubsub.rs @@ -7,7 +7,7 @@ use litellm_circuit_breaker::{Breaker, Config}; use litellm_redis_client::commands; use litellm_redis_client::{ Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, - Target, Timeouts, + SubscriptionEvent, Target, Timeouts, }; use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually}; use rstest::rstest; @@ -16,6 +16,85 @@ use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation"; +#[rstest] +#[tokio::test] +async fn a_silent_dead_subscription_is_detected_and_reconnected( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let proxy = support::proxied(&server).await; + let client = RedisClient::new(RedisConfig { + target: proxy.target.clone(), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(500), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + let mut events = client.subscribe_events(CHANNEL); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Connected) + )); + proxy.blackhole_open_connections(); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(20), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Disconnected) + )); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(10), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Connected) + )); +} + +#[rstest] +#[tokio::test] +async fn lifecycle_reports_disconnect_before_reconnect( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let mut events = client.subscribe_events(CHANNEL); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Connected) + )); + server.restart().await.unwrap(); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Disconnected) + )); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(15), events.next()) + .await + .unwrap(), + Some(SubscriptionEvent::Connected) + )); + eventually(async || { + client + .send(commands::publish(CHANNEL, b"after".as_slice())) + .await + .ok() + }) + .await; + assert!( + matches!(tokio::time::timeout(Duration::from_secs(5), events.next()).await.unwrap(), Some(SubscriptionEvent::Message(bytes)) if bytes == b"after") + ); +} + fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver> { let (sender, receiver) = unbounded_channel(); let mut messages = client.subscribe(channel); diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml index 1dd61d76d3c..1cd752abbef 100644 --- a/litellm-rust/crates/redis-store/Cargo.toml +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -11,6 +11,7 @@ litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-python-compat.workspace = true litellm-redis-client.workspace = true +moka = { workspace = true, features = ["sync"] } strum.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } diff --git a/litellm-rust/crates/redis-store/src/auth_cache/local.rs b/litellm-rust/crates/redis-store/src/auth_cache/local.rs new file mode 100644 index 00000000000..ed533fd2580 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/local.rs @@ -0,0 +1,108 @@ +use std::time::Duration; + +use litellm_python_compat::Value; +use moka::sync::Cache; + +use super::{Settings, is_key_object}; + +const MAX_AUTH_TTL: Duration = Duration::from_secs(60); + +#[derive(Clone)] +struct Entry { + value: Value, + expires: Duration, +} + +pub(super) struct Local { + objects: Cache, + keys: Cache, + pub generation: u64, + pub healthy: bool, +} + +impl Local { + pub fn new(settings: Settings) -> Self { + let partition = || { + Cache::builder() + .max_capacity(settings.memory_entries as u64) + .build() + }; + Self { + objects: partition(), + keys: partition(), + generation: 0, + healthy: true, + } + } + + fn partition(&self, key: &str) -> &Cache { + if is_key_object(key) { + &self.keys + } else { + &self.objects + } + } + + pub fn get(&self, key: &str, now: Duration) -> Option { + if !self.healthy { + return None; + } + let entry = self.partition(key).get(key)?; + if now >= entry.expires { + self.partition(key).invalidate(key); + return None; + } + Some(entry.value) + } + + pub fn set(&self, key: &str, value: Value, ttl: Option, now: Duration) { + if !self.healthy || key.starts_with("sk-") { + return; + } + let ttl = ttl.unwrap_or(Duration::from_secs(600)); + let ttl = if is_key_object(key) { + ttl.min(MAX_AUTH_TTL) + } else { + ttl + }; + self.partition(key).insert( + key.to_owned(), + Entry { + value, + expires: now.saturating_add(ttl), + }, + ); + } + + pub fn delete(&mut self, key: &str) { + self.generation = self.generation.wrapping_add(1); + self.partition(key).invalidate(key); + } + + pub fn flush(&mut self) { + self.generation = self.generation.wrapping_add(1); + self.keys.invalidate_all(); + self.objects.invalidate_all(); + } + + pub fn reconfigure(&mut self, settings: Settings) { + let replacement = Self::new(settings); + for (key, value) in &self.keys { + replacement.keys.insert((*key).clone(), value); + } + for (key, value) in &self.objects { + replacement.objects.insert((*key).clone(), value); + } + self.keys = replacement.keys; + self.objects = replacement.objects; + } + + pub fn len(&self) -> (usize, usize) { + self.objects.run_pending_tasks(); + self.keys.run_pending_tasks(); + ( + self.objects.entry_count() as usize, + self.keys.entry_count() as usize, + ) + } +} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/mod.rs b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs index b5a83b5aec6..4ddb68df710 100644 --- a/litellm-rust/crates/redis-store/src/auth_cache/mod.rs +++ b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs @@ -1,4 +1,5 @@ mod keys; +mod local; mod types; use std::{ @@ -12,23 +13,19 @@ use litellm_redis_client::{Declared, RedisBatch, commands}; pub use self::keys::is_key_object; pub use self::types::Settings; use crate::{ - Error, Number, Pending, RequestRedis, Send, - dual_cache::DualCache, + Error, Number, Pending, RequestRedis, Send, codec, + failure_log::redis_failure, memory::{Clock, lock, system_clock}, remote::{Remote, RemoteSlot}, - request::ready, + request::{pending, ready}, }; /// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it. const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"; -/// The proxy's auth object cache with `enable_redis_auth_cache` on: `UserApiKeyCache`'s two -/// `DualCache` partitions, so churn in management objects cannot evict user-key objects, in front of -/// one Redis shared by every pod. Keys are the logical cache keys `UserApiKeyCache` callers already -/// use; values are what `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`. pub struct AuthCache { - key_objects: DualCache, - objects: DualCache, + local: Mutex, + now: Clock, redis: Arc, settings: Mutex, } @@ -40,10 +37,9 @@ impl AuthCache { pub fn with_clock(remote: Option, settings: Settings, now: Clock) -> Self { let redis = Arc::new(RemoteSlot::new(remote)); - let partition = || DualCache::new(redis.clone(), settings.memory(), now.clone()); Self { - key_objects: partition(), - objects: partition(), + local: Mutex::new(local::Local::new(settings)), + now, redis, settings: Mutex::new(settings), } @@ -53,6 +49,40 @@ impl AuthCache { self.redis.get().is_some() } + pub fn attach(&self, remote: Remote) -> bool { + self.redis.attach(remote) + } + + pub fn current_generation(&self) -> u64 { + lock(&self.local).generation + } + + pub fn local_ready(&self) -> bool { + lock(&self.local).healthy + } + + pub fn subscriber_connected(&self, connected: bool) { + let mut local = lock(&self.local); + local.flush(); + local.healthy = connected; + } + + pub fn set_if_generation( + &self, + key: &str, + value: Value, + generation: u64, + ttl: Option, + ) -> bool { + let ttl = ttl.or(self.settings().memory_ttl); + let local = lock(&self.local); + if local.generation != generation || !local.healthy || key.starts_with("sk-") { + return false; + } + local.set(key, value, ttl, (self.now)()); + true + } + pub fn settings(&self) -> Settings { *lock(&self.settings) } @@ -61,37 +91,33 @@ impl AuthCache { /// keep what they hold. pub fn reconfigure(&self, settings: Settings) { *lock(&self.settings) = settings; - self.key_objects.reconfigure(settings.memory()); - self.objects.reconfigure(settings.memory()); + lock(&self.local).reconfigure(settings); } pub fn get_local(&self, key: &str) -> Result, Error> { - Ok(self.partition(key).memory().get(key)?) + Ok(lock(&self.local).get(key, (self.now)())) } /// A memory write with `ttl`, or `default_in_memory_ttl` when it names none. pub fn set_local(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { - let partition = self.partition(key); - Ok(partition - .memory() - .set(key, value, ttl.or(partition.default_in_memory_ttl()))?) + let ttl = ttl.or(self.settings().memory_ttl); + lock(&self.local).set(key, value, ttl, (self.now)()); + Ok(()) } pub fn delete_local(&self, key: &str) -> Result<(), Error> { - Ok(self.partition(key).memory().delete(key)?) + lock(&self.local).delete(key); + Ok(()) } pub fn flush_local(&self) -> Result<(), Error> { - self.key_objects.memory().flush()?; - Ok(self.objects.memory().flush()?) + lock(&self.local).flush(); + Ok(()) } /// Entries held by the management-object and the user-key partitions. pub fn len_local(&self) -> Result<(usize, usize), Error> { - Ok(( - self.objects.memory().len()?, - self.key_objects.memory().len()?, - )) + Ok(lock(&self.local).len()) } /// `DualCache.async_get_cache`. A key `request` already read as missing from Redis is not read @@ -101,12 +127,42 @@ impl AuthCache { request: Option<&RequestRedis>, key: &str, ) -> Result, Error> { - self.partition(key).get(request, key, None).await + let generation = self.current_generation(); + if let Some(value) = self.get_local(key)? { + return Ok(Some(value)); + } + if !lock(&self.local).healthy || key.starts_with("sk-") { + return Ok(None); + } + let Some(client) = self.redis.client() else { + return Ok(None); + }; + if request.is_some_and(|request| request.misses(client).contains(key)) { + return Ok(None); + } + let value = client + .send(commands::get::>(key)) + .await + .map_err(Error::from) + .and_then(|raw| Ok(types::decode(raw)?)) + .unwrap_or_else(|error| { + redis_failure!(&error, "auth cache read failed"); + None + }); + if let Some(value) = &value + && !self.set_if_generation(key, value.clone(), generation, None) + { + return Ok(None); + } + Ok(value) } /// `DualCache.async_set_cache`, writing `json.dumps(value)`. pub async fn set(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { - self.partition(key).set(key, value, ttl).await + if let Err(error) = self.declare_set(Send::Now, key, value, ttl)?.await { + redis_failure!(&error, "auth cache write failed"); + } + Ok(()) } /// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`. @@ -116,7 +172,14 @@ impl AuthCache { value: Value, ttl: Option, ) -> Result<(), Error> { - self.partition(key).set_str(key, value, ttl).await + let encoded = codec::str(value.clone()); + if let Err(error) = self + .declare_encoded(Send::Now, key, value, ttl, Ok(encoded))? + .await + { + redis_failure!(&error, "auth cache write failed"); + } + Ok(()) } /// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks. @@ -127,12 +190,56 @@ impl AuthCache { value: Value, ttl: Option, ) -> Result, Error> { - self.partition(key).declare_set(send, key, value, ttl) + let encoded = codec::dumps(value.clone()); + self.declare_encoded(send, key, value, ttl, encoded) + } + + fn declare_encoded( + &self, + send: Send<'_>, + key: &str, + value: Value, + ttl: Option, + encoded: Result, + ) -> Result, Error> { + let ttl = ttl.or(self.settings().memory_ttl); + self.set_local(key, value, ttl)?; + if key.starts_with("sk-") { + return Ok(ready(())); + } + let Some(remote) = self.redis.get() else { + return Ok(ready(())); + }; + let encoded = match encoded { + Ok(encoded) => encoded, + Err(error) => { + redis_failure!(&error, "auth cache value did not encode"); + return Ok(ready(())); + } + }; + if let Some(misses) = send.misses(&remote.client) { + misses.forget(key); + } + Ok(pending(send.batch(&remote.client).add(commands::set( + key, + encoded, + ttl.unwrap_or(remote.default_ttl), + )))) } /// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks. pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { - self.partition(key).declare_delete(send, key) + self.delete_local(key)?; + let Some(client) = self.redis.client() else { + return Ok(ready(())); + }; + if let Some(misses) = send.misses(client) { + misses.note([key.to_owned()]); + } + let deleted = send.batch(client).add(commands::del(&[key])); + Ok(Box::pin( + async move { Ok(deleted.resolve().await.map(drop)?) }, + )) } /// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then @@ -142,10 +249,30 @@ impl AuthCache { entries: Vec<(String, Value)>, ttl: Option, ) -> Result<(), Error> { - let (key_objects, others): (Vec<_>, Vec<_>) = + let (keys, objects): (Vec<_>, Vec<_>) = entries.into_iter().partition(|(key, _)| is_key_object(key)); - self.key_objects.set_pipeline(key_objects, ttl).await?; - self.objects.set_pipeline(others, ttl).await + self.set_partition(keys, ttl).await?; + self.set_partition(objects, ttl).await + } + + async fn set_partition( + &self, + entries: Vec<(String, Value)>, + ttl: Option, + ) -> Result<(), Error> { + let ttl = ttl.or(self.settings().memory_ttl); + let entries = entries + .into_iter() + .map(|(key, value)| (key, value, ttl)) + .collect(); + let result = match self.write_back(Send::Now, entries) { + Ok(written) => written.await, + Err(error) => Err(error), + }; + if let Err(error) = result { + redis_failure!(&error, "auth cache pipeline write failed"); + } + Ok(()) } /// The auth prefetch write-back: memory first, an entry without a TTL taking @@ -158,6 +285,10 @@ impl AuthCache { send: Send<'_>, entries: Vec<(String, Value, Option)>, ) -> Result, Error> { + let entries: Vec<_> = entries + .into_iter() + .filter(|(key, _, _)| !key.starts_with("sk-")) + .collect(); let encoded = self .redis .get() @@ -199,13 +330,30 @@ impl AuthCache { /// `UserApiKeyCache.async_delete_cache_keys`: the key-object partition and the rest as two /// concurrent `DualCache.async_delete_cache_keys`. Both run even when one fails. pub async fn delete(&self, keys: &[String]) -> Result<(), Error> { - let (key_objects, others): (Vec, Vec) = + for key in keys { + self.delete_local(key)?; + } + let (keys, objects): (Vec<_>, Vec<_>) = keys.iter().cloned().partition(|key| is_key_object(key)); - let (key_objects, others) = tokio::join!( - self.key_objects.delete_keys(&key_objects), - self.objects.delete_keys(&others) + let (keys, objects) = tokio::join!( + self.delete_partition(&keys), + self.delete_partition(&objects) ); - key_objects.and(others) + keys.and(objects) + } + + async fn delete_partition(&self, keys: &[String]) -> Result<(), Error> { + let Some(client) = self.redis.client() else { + return Ok(()); + }; + for chunk in keys.chunks(1000) { + client + .send(commands::del( + &chunk.iter().map(String::as_str).collect::>(), + )) + .await?; + } + Ok(()) } /// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on @@ -247,12 +395,67 @@ impl AuthCache { .send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount)) .await?) } +} - fn partition(&self, key: &str) -> &DualCache { - if is_key_object(key) { - &self.key_objects - } else { - &self.objects - } +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use std::sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }; + + const KEY: &str = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"; + + #[rstest] + fn invalidation_fences_an_in_flight_fill_and_allows_a_new_one() { + let cache = AuthCache::new(None, Settings::default()); + let before = cache.current_generation(); + cache.delete_local(KEY).unwrap(); + assert!(!cache.set_if_generation(KEY, Value::from("revoked"), before, None)); + assert_eq!(cache.get_local(KEY).unwrap(), None); + let after = cache.current_generation(); + assert!(cache.set_if_generation(KEY, Value::from("fresh"), after, None)); + assert_eq!(cache.get_local(KEY).unwrap(), Some(Value::from("fresh"))); + } + + #[rstest] + fn disconnect_bypasses_writes_and_reconnect_fences_old_fills() { + let cache = AuthCache::new(None, Settings::default()); + cache.set_local(KEY, Value::Bool(true), None).unwrap(); + cache.subscriber_connected(false); + assert_eq!(cache.get_local(KEY).unwrap(), None); + cache.set_local(KEY, Value::Bool(true), None).unwrap(); + let disconnected = cache.current_generation(); + assert!(!cache.set_if_generation(KEY, Value::Bool(true), disconnected, None)); + cache.subscriber_connected(true); + assert_eq!(cache.get_local(KEY).unwrap(), None); + assert!(!cache.set_if_generation(KEY, Value::Bool(true), disconnected, None)); + } + + #[rstest] + fn auth_entries_expire_even_without_a_configured_ttl() { + let seconds = Arc::new(AtomicU64::new(100)); + let clock = seconds.clone(); + let cache = AuthCache::with_clock( + None, + Settings::default(), + Arc::new(move || Duration::from_secs(clock.load(Ordering::SeqCst))), + ); + cache.set_local(KEY, Value::Bool(true), None).unwrap(); + assert_eq!(cache.get_local(KEY).unwrap(), Some(Value::Bool(true))); + seconds.store(160, Ordering::SeqCst); + assert_eq!(cache.get_local(KEY).unwrap(), None); + } + + #[rstest] + fn raw_credentials_are_not_cached() { + let cache = AuthCache::new(None, Settings::default()); + cache + .set_local("sk-do-not-cache", Value::Bool(true), None) + .unwrap(); + assert_eq!(cache.get_local("sk-do-not-cache").unwrap(), None); + assert_eq!(cache.len_local().unwrap(), (0, 0)); } } diff --git a/litellm-rust/crates/redis-store/src/auth_cache/types.rs b/litellm-rust/crates/redis-store/src/auth_cache/types.rs index e933ac19af9..d66ae858e44 100644 --- a/litellm-rust/crates/redis-store/src/auth_cache/types.rs +++ b/litellm-rust/crates/redis-store/src/auth_cache/types.rs @@ -2,7 +2,7 @@ use std::time::Duration; use litellm_python_compat::Value; -use crate::{DecodeError, codec::python_value, dual_cache::MemorySettings}; +use crate::{DecodeError, codec::python_value}; /// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -21,15 +21,6 @@ impl Default for Settings { } } -impl Settings { - pub(super) fn memory(self) -> MemorySettings { - MemorySettings { - entries: Some(self.memory_entries), - ttl: self.memory_ttl, - } - } -} - pub(super) fn decode(raw: Option>) -> Result, DecodeError> { raw.as_deref().map(python_value).transpose() } diff --git a/litellm-rust/crates/redis-store/src/dual_cache/mod.rs b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs index b2c277a118e..4158653be4b 100644 --- a/litellm-rust/crates/redis-store/src/dual_cache/mod.rs +++ b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs @@ -88,6 +88,7 @@ impl DualCache { /// `DualCache.update_cache_ttl` and `update_in_memory_max_size` together: both apply in place, /// and the entries memory already holds stay. + #[cfg(test)] pub(crate) fn reconfigure(&self, memory: MemorySettings) { *lock(&self.default_in_memory_ttl) = memory.ttl; self.memory.update_max_size(memory.entries); diff --git a/litellm-rust/crates/redis-store/src/dual_cache/write.rs b/litellm-rust/crates/redis-store/src/dual_cache/write.rs index 6fcc6f32741..79d843783df 100644 --- a/litellm-rust/crates/redis-store/src/dual_cache/write.rs +++ b/litellm-rust/crates/redis-store/src/dual_cache/write.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_redis_client::{Declared, RedisBatch, commands}; +use litellm_redis_client::commands; use super::DualCache; use crate::{ @@ -35,69 +35,6 @@ impl DualCache { .await } - /// `DualCache._set_on_batch`: memory now, the `SET` on the batch `send` picks. Memory and Redis - /// both take `ttl`, else `default_in_memory_ttl`, and Redis falls back to its own default. - pub(crate) fn declare_set( - &self, - send: Send<'_>, - key: &str, - value: V, - ttl: Option, - ) -> Result, Error> { - self.declare_encoded(send, key, value, ttl, codec::dumps) - } - - /// `DualCache.async_set_cache_pipeline`: memory, then one pipeline of `SET`s. A value - /// `json.dumps` refuses aborts the whole pipeline, and any failure only logs, as `RedisCache` - /// does. - pub(crate) async fn set_pipeline( - &self, - entries: Vec<(String, V)>, - ttl: Option, - ) -> Result<(), Error> { - let ttl = ttl.or(self.default_in_memory_ttl()); - let encoded = self.redis.get().map(|_| { - entries - .iter() - .map(|(_, value)| codec::dumps(value.clone())) - .collect::, Error>>() - }); - let keys = entries - .into_iter() - .map(|(key, value)| { - self.memory.set(&key, value, ttl)?; - Ok(key) - }) - .collect::, Error>>()?; - let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else { - return Ok(()); - }; - let failure = match encoded { - Err(error) => Some(error), - Ok(values) => { - let pipeline = RedisBatch::pipelined(remote.client.clone()); - let ttl = ttl.unwrap_or(remote.default_ttl); - let written = keys - .iter() - .zip(values) - .map(|(key, value)| pipeline.add(commands::set(key, value, ttl))) - .collect::>(); - Declared::all(written) - .await - .into_iter() - .find_map(Result::err) - .map(Error::from) - } - }; - if let Some(error) = failure { - redis_failure!( - &error, - "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS", - ); - } - Ok(()) - } - async fn set_now( &self, key: &str, @@ -147,22 +84,6 @@ impl DualCache { } impl DualCache { - /// `DualCache.async_delete_cache_pre_call`: memory now, the `DEL` on the batch `send` picks, - /// after which the request reads the key as missing. - pub(crate) fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { - self.memory.delete(key)?; - let Some(client) = self.redis.client() else { - return Ok(ready(())); - }; - if let Some(misses) = send.misses(client) { - misses.note([key.to_owned()]); - } - let deleted = send.batch(client).add(commands::del(&[key])); - Ok(Box::pin( - async move { Ok(deleted.resolve().await.map(drop)?) }, - )) - } - /// `DualCache.async_delete_cache_keys`: memory, then one multi-key `DEL` per chunk, in order. pub(crate) async fn delete_keys(&self, keys: &[String]) -> Result<(), Error> { keys.iter().try_for_each(|key| self.memory.delete(key))?; diff --git a/litellm-rust/crates/redis-store/src/memory.rs b/litellm-rust/crates/redis-store/src/memory.rs index f89d5926760..af6ac9bb0e1 100644 --- a/litellm-rust/crates/redis-store/src/memory.rs +++ b/litellm-rust/crates/redis-store/src/memory.rs @@ -50,10 +50,6 @@ impl Memory { self.cache.flush_cache() } - pub fn len(&self) -> Result { - self.cache.len() - } - pub fn expires_at(&self, key: &str) -> Result, Error> { self.cache.expires_at(key) } @@ -63,6 +59,7 @@ impl Memory { } /// `DualCache.update_in_memory_max_size`: `None` restores the 200-entry default. + #[cfg(test)] pub fn update_max_size(&self, entries: Option) { self.cache.update_max_size_in_memory(entries); } diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 36650629300..327b4c106de 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,5 +1,6 @@ import asyncio import json +import math from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass from types import MappingProxyType @@ -18,6 +19,8 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # leg pubsub_capable_client, ) from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET +from litellm.redis.stores.connection import connect +from litellm.rust_bridge._native import NativeAuthSubscription class SpendCounterCopies(Protocol): @@ -65,7 +68,7 @@ def _cache_invalidation_message_json(cache_key: str, new_value: float | None = N def _finite_number_or_none(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, (int, float)): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): return None return float(value) @@ -239,7 +242,14 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us class AuthCacheInvalidationSubscriber: - __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_spend_counters", "_task", "_user_api_key_cache") + __slots__ = ( + "_additional_in_memory_caches", + "_native", + "_redis_cache", + "_spend_counters", + "_task", + "_user_api_key_cache", + ) def __init__( self, @@ -253,12 +263,31 @@ class AuthCacheInvalidationSubscriber: self._additional_in_memory_caches = tuple(additional_in_memory_caches) self._spend_counters: Final = spend_counters self._task: asyncio.Task[None] | None = None + self._native: NativeAuthSubscription | None = None def start(self) -> None: if self._task is not None: return + shared: Final = self._user_api_key_cache.shared + client: Final = connect(self._redis_cache).client + if shared is not None and client is not None: + self._native = shared.subscribe(client, auth_cache_invalidation_channel(self._redis_cache)) + self._task = asyncio.create_task(self._consume_native(self._native)) + return self._task = asyncio.create_task(self._run()) + async def _consume_native(self, subscription: NativeAuthSubscription) -> None: + try: + while True: + if (data := await subscription.next()) is not None: + self._apply_message({"data": data}) + finally: + await subscription.close() + + @property + def is_ready(self) -> bool: + return self._native is not None and self._native.is_connected() + async def stop(self) -> None: task: Final = self._task if task is None: @@ -269,6 +298,7 @@ class AuthCacheInvalidationSubscriber: await task except asyncio.CancelledError: pass + self._native = None async def _run(self) -> None: backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: exponential backoff accumulator across reconnects diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 91529ee36fc..a7f652895b1 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -68,9 +68,8 @@ class UserApiKeyCache(DualCache): User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition, ``key_object_cache``, so churn in the other management objects cannot evict them. - Without Redis both partitions are this class's Python memory. ``attach_redis_cache`` moves the whole - keyspace, both memory partitions included, onto the native ``AuthCacheStore`` (``shared``); redis-py never - reads or writes it, and ``redis_cache`` stays None. + Both partitions live in native Moka, including without Redis. ``attach_redis_cache`` adds the shared + Redis tier without creating another local authority; ``redis_cache`` stays None. ``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous (no other methods in between) so mypy resolves ``@overload`` + implementation correctly. @@ -103,7 +102,7 @@ class UserApiKeyCache(DualCache): ) self._auth_cache: Final = auth_cache self._redis_connection: RedisCache | None = None - self._shared: AuthCacheStore | None = None + self._shared: AuthCacheStore | None = AuthCacheStore(connect(None).auth_cache(self._settings())) self.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) @property @@ -113,7 +112,7 @@ class UserApiKeyCache(DualCache): @property def has_redis(self) -> bool: - return self._shared is not None + return self._redis_connection is not None @property def redis_connection(self) -> RedisCache | None: @@ -150,12 +149,29 @@ class UserApiKeyCache(DualCache): def attach_redis_cache( self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None ) -> None: - if redis_cache is None or self._shared is not None: + if redis_cache is None or self._redis_connection is not None: return if default_redis_ttl is not None: self.update_cache_ttl(default_in_memory_ttl=None, default_redis_ttl=default_redis_ttl) self._redis_connection = redis_cache - self._shared = self._auth_cache(redis_cache, self._settings()) + if self._auth_cache is _connected_auth_cache and self._shared is not None: + client: Final = connect(redis_cache).client + if client is not None: + self._shared.attach(client) + else: + self._shared = self._auth_cache(redis_cache, self._settings()) + + def begin_fill(self, key: str) -> int: + assert self._shared is not None + return self._shared.begin_fill(key) + + def current_generation(self) -> int: + assert self._shared is not None + return self._shared.current_generation() + + def set_if_generation(self, key: str, value: CacheValue, generation: int, ttl: float | None = None) -> bool: + assert self._shared is not None + return self._shared.set_if_generation(key, CacheCodec.serialize(value, model_type=None), generation, ttl) def get_local(self, key: str) -> object: """This worker's memory copy alone.""" diff --git a/litellm/redis/stores/auth_cache/__init__.py b/litellm/redis/stores/auth_cache/__init__.py index 7891344022f..7e9f388be9c 100644 --- a/litellm/redis/stores/auth_cache/__init__.py +++ b/litellm/redis/stores/auth_cache/__init__.py @@ -10,7 +10,7 @@ from litellm.caching.redis_batch import Batch, active_request_redis_batches from litellm.redis.stores._route import route from litellm.redis.stores.auth_cache.handle import AuthCacheHandle from litellm.redis.stores.auth_cache.types import AUTH_CACHE_TARGET, AuthCacheSettings, WriteBack -from litellm.rust_bridge._native import LocalSizes +from litellm.rust_bridge._native import LocalSizes, NativeAuthSubscription, NativeRedisClient __all__ = [ "AUTH_CACHE_TARGET", @@ -36,6 +36,22 @@ class AuthCacheStore: def __init__(self, handle: AuthCacheHandle) -> None: self._handle: Final = handle + def attach(self, client: NativeRedisClient) -> bool: + return self._handle.attach(client) + + def begin_fill(self, key: str) -> int: + return self._handle.current_generation() + + def current_generation(self) -> int: + return self._handle.current_generation() + + def set_if_generation(self, key: str, value: object, generation: int, ttl: float | None = None) -> bool: + self._forget_miss(key) + return self._handle.set_if_generation(key, value, generation, _redis_ttl(ttl)) + + def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription: + return self._handle.subscribe(client, channel) + def reconfigure(self, settings: AuthCacheSettings) -> None: """Both memory partitions keep what they hold, as ``update_cache_ttl`` and ``update_in_memory_max_size`` do.""" self._handle.reconfigure(memory_entries=settings.memory_entries, memory_ttl_seconds=settings.memory_ttl) diff --git a/litellm/redis/stores/auth_cache/handle.py b/litellm/redis/stores/auth_cache/handle.py index 9a2d70079d1..fc91ae2ef35 100644 --- a/litellm/redis/stores/auth_cache/handle.py +++ b/litellm/redis/stores/auth_cache/handle.py @@ -5,10 +5,14 @@ from typing import Protocol from litellm.caching.redis_batch import NativePending from litellm.redis.stores.auth_cache.types import WriteBack -from litellm.rust_bridge._native import LocalSizes, NativeSend +from litellm.rust_bridge._native import LocalSizes, NativeAuthSubscription, NativeRedisClient, NativeSend class AuthCacheHandle(Protocol): + def attach(self, client: NativeRedisClient) -> bool: ... + def current_generation(self) -> int: ... + def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool: ... + def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription: ... def get(self, key: str) -> Awaitable[object]: ... def get_sync(self, key: str) -> object: ... def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Awaitable[None]: ... diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 5d5370c7dee..cbb8a72ae99 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -395,6 +395,10 @@ class NativeAuthCache: ) -> NativeAuthCache: ... @property def has_redis(self) -> bool: ... + def attach(self, client: NativeRedisClient) -> bool: ... + def current_generation(self) -> int: ... + def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool: ... + def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription: ... def get(self, key: str, request: NativeRequestRedis | None = None) -> Future[object]: ... def reconfigure(self, *, memory_entries: int, memory_ttl_seconds: float | None = None) -> None: ... def get_sync(self, key: str) -> object: ... @@ -416,6 +420,12 @@ class NativeAuthCache: def flush_local(self) -> None: ... def len_local(self) -> LocalSizes: ... +@final +class NativeAuthSubscription: + def is_connected(self) -> bool: ... + def next(self) -> Future[bytes | None]: ... + def close(self) -> Future[None]: ... + @final class SpendRead: def __new__(cls, value: float | None, from_redis: bool) -> SpendRead: ... diff --git a/tests/integration/caching/test_auth_cache_native.py b/tests/integration/caching/test_auth_cache_native.py new file mode 100644 index 00000000000..725a01224c5 --- /dev/null +++ b/tests/integration/caching/test_auth_cache_native.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import asyncio +import hashlib +import json +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import pytest +from redis.asyncio import Redis + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AuthCacheInvalidationSubscriber, + auth_cache_invalidation_channel, + evict_local, + evict_shared, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.redis.stores.connection import connect +from litellm.redis.stores.spend_counters import SpendCounterStore +from tests.integration._support.redis_process import owned_redis + +KEY: Final = hashlib.sha256(b"native-cache-test-key").hexdigest() + + +async def until(predicate: Callable[[], bool], seconds: float = 12) -> None: + async with asyncio.timeout(seconds): + while not predicate(): + await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("namespace", (None, "phase3")) +@pytest.mark.parametrize("shared", (False, True)) +async def test_native_subscriber_invalidates_and_recovers_without_serving_outage_fills( + tmp_path: Path, namespace: str | None, shared: bool +) -> None: + with owned_redis(tmp_path) as server: + redis_cache: Final = RedisCache( + host=server.host, + port=server.port, + namespace=namespace, + socket_timeout=0.3, + socket_connect_timeout=0.3, + ) + cache: Final = UserApiKeyCache(redis_cache=redis_cache if shared else None) + additional: Final = InMemoryCache() + spend: Final = SpendCounterStore(connect(None).spend_counters()) + subscriber: Final = AuthCacheInvalidationSubscriber(redis_cache, cache, (additional,), spend) + publisher: Final = Redis(host=server.host, port=server.port) + channel: Final = auth_cache_invalidation_channel(redis_cache) + subscriber.start() + started: Final = cache.current_generation() + try: + assert not subscriber.is_ready + await until(lambda: subscriber.is_ready) + assert cache.current_generation() > started + fill: Final = cache.begin_fill(KEY) + assert cache.set_if_generation(KEY, {"token": KEY, "models": ["allowed"]}, fill) + assert cache.get_local(KEY) == {"token": KEY, "models": ["allowed"]} + spend.set_local("spend:key:k-1", 99.0) + assert ( + await publisher.publish( + channel, json.dumps({"cache_key": "spend:key:k-1", "new_value": 0.0, "ttl": 60}) + ) + == 1 + ) + await until(lambda: spend.get_local("spend:key:k-1") == 0.0) + assert cache.get_local(KEY) == {"token": KEY, "models": ["allowed"]} + assert ( + await publisher.publish( + channel, json.dumps({"cache_key": "spend_db_floor:spend:key:k-1", "new_value": 3.0}) + ) + == 1 + ) + await until(lambda: additional.get_cache("spend_db_floor:spend:key:k-1") == 3.0) + assert await publisher.publish(channel, json.dumps({"cache_key": "spend:key:k-1"})) == 1 + await until(lambda: spend.get_local("spend:key:k-1") is None) + additional.set_cache("model-cache", "stale") + assert await publisher.publish(channel, json.dumps({"cache_key": KEY})) == 1 + await until(lambda: cache.get_local(KEY) is None) + assert not cache.set_if_generation(KEY, {"token": KEY}, fill) + assert await publisher.publish(channel, json.dumps({"cache_key": "model-cache"})) == 1 + await until(lambda: additional.get_cache("model-cache") is None) + + cache.set_local(KEY, {"token": KEY}, None) + before_stop: Final = cache.current_generation() + server.stop() + await until(lambda: cache.current_generation() > before_stop) + assert not subscriber.is_ready + assert cache.get_local(KEY) is None + outage: Final = cache.begin_fill(KEY) + assert not cache.set_if_generation(KEY, {"token": KEY}, outage) + cache.set_local(KEY, {"token": KEY}, None) + assert cache.get_local(KEY) is None + server.start() + await until(lambda: subscriber.is_ready) + assert cache.current_generation() > outage + assert not cache.set_if_generation(KEY, {"token": KEY}, outage) + assert cache.set_if_generation(KEY, {"token": KEY}, cache.begin_fill(KEY)) + assert await publisher.publish(channel, json.dumps({"cache_key": KEY})) == 1 + await until(lambda: cache.get_local(KEY) is None) + finally: + await subscriber.stop() + await publisher.aclose() + assert cache.get_local(KEY) is None + assert not cache.set_if_generation(KEY, {"token": KEY}, cache.begin_fill(KEY)) + + +@pytest.mark.asyncio +async def test_management_eviction_reaches_native_memory_and_shared_redis(tmp_path: Path) -> None: + with owned_redis(tmp_path) as server: + redis_cache: Final = RedisCache(host=server.host, port=server.port, namespace="management") + first: Final = UserApiKeyCache(redis_cache=redis_cache) + second: Final = UserApiKeyCache(redis_cache=redis_cache) + await first.async_set_cache(KEY, {"token": KEY}) + assert await second.async_get_cache(KEY) == {"token": KEY} + evict_local((KEY,), first) + assert first.get_local(KEY) is None + await evict_shared((KEY,), first) + second.delete_local(KEY) + assert await second.async_get_cache(KEY) is None diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index d41c4595a0d..1d30cbd7b1e 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -145,6 +145,8 @@ class _ScriptedPubSubRedisClient(Redis): class _FakeRedisCache: + store_settings = None + def __init__(self, client: object, namespace: str | None = None) -> None: self._client = client self.namespace = namespace @@ -207,8 +209,8 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: instead of serving the stale object until the TTL expires. """ cache = UserApiKeyCache() - cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) - assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + cache.set_local("project_id:p-1", {"models": []}, None) + assert cache.get_local("project_id:p-1") is not None pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")]) subscriber = AuthCacheInvalidationSubscriber( @@ -221,7 +223,7 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: finally: await subscriber.stop() - assert cache.in_memory_cache.get_cache("project_id:p-1") is None + assert cache.get_local("project_id:p-1") is None assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] @@ -357,7 +359,7 @@ async def test_subscriber_writes_and_drops_the_native_spend_counter_copies_a_bro @pytest.mark.asyncio async def test_subscriber_ignores_malformed_messages() -> None: cache = UserApiKeyCache() - cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + cache.set_local("project_id:p-1", {"models": []}, None) subscriber = AuthCacheInvalidationSubscriber( redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), @@ -368,7 +370,7 @@ async def test_subscriber_ignores_malformed_messages() -> None: subscriber._apply_message("raw string") subscriber._apply_message(None) - assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + assert cache.get_local("project_id:p-1") is not None @pytest.mark.asyncio diff --git a/tests/unit/proxy/common_utils/test_user_api_key_cache.py b/tests/unit/proxy/common_utils/test_user_api_key_cache.py index d3c9043e4b6..4f7342e1b4d 100644 --- a/tests/unit/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/unit/proxy/common_utils/test_user_api_key_cache.py @@ -1,6 +1,6 @@ import datetime import hashlib -from typing import Any, Final +from typing import Final from unittest.mock import MagicMock import pytest @@ -17,7 +17,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( get_management_object_ttl, is_user_key_cache_key, ) -from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum from litellm.redis.stores.auth_cache import AuthCacheStore from litellm.rust_bridge._native import NativeAuthCache from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache @@ -25,18 +24,6 @@ from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cach HASHED_TOKEN = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest() -class CapturingInMemoryCache(InMemoryCache): - """Records ``ttl`` passed into ``set_cache`` (what DualCache injects).""" - - def __init__(self) -> None: - super().__init__() - self.last_ttl: Any = None - - def set_cache(self, key, value, **kwargs): # type: ignore[override] - self.last_ttl = kwargs.get("ttl") - super().set_cache(key, value, **kwargs) - - class _ManualClock: def __init__(self) -> None: self.seconds = 1000.0 @@ -54,39 +41,55 @@ def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: class TestUserApiKeyCache: + def test_native_fill_fence_and_request_local_projection_without_redis(self): + cache: Final = UserApiKeyCache() + stale: Final = cache.begin_fill(HASHED_TOKEN) + cache.delete_local(HASHED_TOKEN) + assert not cache.set_if_generation(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), stale) + assert cache.set_if_generation( + HASHED_TOKEN, UserAPIKeyAuth(token=HASHED_TOKEN, models=["original"]), cache.begin_fill(HASHED_TOKEN) + ) + first: Final = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert first is not None + first.models.append("request-only") + second: Final = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert second is not None and second is not first + assert second.models == ["original"] + cache.flush_cache() + assert cache.get_local(HASHED_TOKEN) is None + @pytest.mark.asyncio - async def test_async_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted( + async def test_async_set_with_zero_default_ttl_expires_without_a_python_copy( self, ): - """ - If ``general_settings.user_api_key_cache_ttl`` is absent, the proxy never - calls ``update_cache_ttl``; ``user_api_key_cache`` keeps - ``default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl``. - DualCache must forward that as the in-memory ``ttl`` kwarg on each set. - """ - mem = CapturingInMemoryCache() + mem = InMemoryCache() cache = UserApiKeyCache( in_memory_cache=mem, - default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, + default_in_memory_ttl=0, ) + cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=None) await cache.async_set_cache( "k", _make_key_obj("t"), model_type=UserAPIKeyAuth, ) - expected = UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value - assert mem.last_ttl == expected + assert cache.get_local("k") is None + assert mem.get_cache("k") is None + cache.set_cache("k", {"token": "t"}, ttl=60) + assert cache.get_local("k") == {"token": "t"} - def test_sync_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted( + def test_sync_set_with_zero_default_ttl_expires_without_a_python_copy( self, ): - mem = CapturingInMemoryCache() + mem = InMemoryCache() cache = UserApiKeyCache( in_memory_cache=mem, - default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, + default_in_memory_ttl=0, ) + cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=None) cache.set_cache("sk", _make_key_obj("s"), model_type=UserAPIKeyAuth) - assert mem.last_ttl == UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value + assert cache.get_local("sk") is None + assert mem.get_cache("sk") is None @pytest.mark.asyncio async def test_the_shared_store_gets_default_in_memory_ttl_and_writes_without_a_ttl(self): @@ -173,7 +176,9 @@ class TestUserApiKeyCache: async def test_a_model_holding_a_datetime_is_cached_as_its_json_dump(self): cache = UserApiKeyCache( redis_cache=MagicMock(spec=RedisCache), - auth_cache=lambda _redis, settings: AuthCacheStore(NativeAuthCache(None, memory_entries=settings.memory_entries)), + auth_cache=lambda _redis, settings: AuthCacheStore( + NativeAuthCache(None, memory_entries=settings.memory_entries) + ), ) key_obj = UserAPIKeyAuth(token="abc", expires=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)) @@ -273,13 +278,14 @@ class TestUserKeyObjectPartition: @pytest.mark.asyncio async def test_update_cache_ttl_applies_to_key_object_partition(self): - key_partition = CapturingInMemoryCache() + key_partition = InMemoryCache() cache = UserApiKeyCache(default_in_memory_ttl=60, key_object_in_memory_cache=key_partition) - cache.update_cache_ttl(default_in_memory_ttl=7, default_redis_ttl=7) + cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=7) await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) - assert key_partition.last_ttl == 7 + assert cache.get_local(HASHED_TOKEN) is None + assert key_partition.get_cache(HASHED_TOKEN) is None @pytest.mark.asyncio async def test_attach_redis_cache_moves_both_partitions_onto_the_shared_store(self): @@ -372,8 +378,10 @@ class TestUserKeyObjectPartition: cache.set_local(HASHED_TOKEN, {"token": HASHED_TOKEN}, None) cache.set_local(end_user_cache_key("u1"), {"user_id": "u1"}, 5) - assert cache.key_object_cache.in_memory_cache.get_cache(HASHED_TOKEN) == {"token": HASHED_TOKEN} - assert cache.in_memory_cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"} + assert cache.get_local(HASHED_TOKEN) == {"token": HASHED_TOKEN} + assert cache.get_local(end_user_cache_key("u1")) == {"user_id": "u1"} + assert cache.key_object_cache.in_memory_cache.get_cache(HASHED_TOKEN) is None + assert cache.in_memory_cache.get_cache(end_user_cache_key("u1")) is None cache.delete_local(HASHED_TOKEN) assert cache.get_local(HASHED_TOKEN) is None @@ -461,11 +469,7 @@ class TestManagementObjectTTL: @pytest.mark.asyncio async def test_management_write_uses_configured_ttl_over_constant(self): - mem = CapturingInMemoryCache() - cache = UserApiKeyCache( - in_memory_cache=mem, - default_in_memory_ttl=300, - ) + cache, fake = shared_cache(default_in_memory_ttl=300) assert get_management_object_ttl(cache) != (DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL) await cache.async_set_cache( @@ -475,4 +479,4 @@ class TestManagementObjectTTL: ttl=get_management_object_ttl(cache), ) - assert mem.last_ttl == 300 + assert fake.commands == [("SET", "team_id:abc", 300)]