diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs index 58b77222f5e..dd5f6fea90a 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -100,77 +100,57 @@ impl NativeSpendCounters { Some(pending("async_batch_get_cache", kind_error, read, numbers)) } + #[pyo3(signature = (counter, amount, ttl_seconds = None))] fn increment<'py>( &self, py: Python<'py>, counter: PyData, amount: PyNumber, + ttl_seconds: Option, ) -> PyResult> { let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; redis_async(py, "async_increment", kind_error, async move { - store.increment(&counter.0, amount.0).await.map(PyNumber) + store + .increment(&counter.0, amount.0, ttl) + .await + .map(PyNumber) }) } + #[pyo3(signature = (increments, ttl_seconds = None))] fn increment_many<'py>( &self, py: Python<'py>, increments: Vec<(PyData, PyNumber)>, + ttl_seconds: Option, ) -> PyResult> { let store = self.0.clone(); let increments = self::increments(increments); + let ttl = ttl(ttl_seconds)?; redis_async(py, "async_increment_pipeline", kind_error, async move { - let totals = store.increment_many(&increments).await?; + let totals = store.increment_many(&increments, ttl).await?; Ok(totals.into_iter().map(PyNumber).collect::>()) }) } /// Declares the increments on the request's post-call batch; `False` when there is no Redis. + #[pyo3(signature = (request, increments, ttl_seconds = None))] fn increment_many_post_call( &self, request: &NativeRequestRedis, increments: Vec<(PyData, PyNumber)>, - ) -> bool { - self.0 - .increment_many_post_call(&request.0, &self::increments(increments)) - } - - fn set_max<'py>( - &self, - py: Python<'py>, - counter: PyData, - spend: PyNumber, - ) -> PyResult> { - let store = self.0.clone(); - redis_async(py, "async_set_max", kind_error, async move { - store - .set_max(&counter.0, spend.0) - .await - .map(|stored| stored.map(PyNumber)) - }) - } - - #[pyo3(signature = (counter, spend, *, increment_if_lost))] - fn seed<'py>( - &self, - py: Python<'py>, - counter: PyData, - spend: PyNumber, - increment_if_lost: bool, - ) -> PyResult> { - let store = self.0.clone(); - let lost = if increment_if_lost { - LostSeed::Increment - } else { - LostSeed::Answer - }; - redis_async(py, "async_set_cache", kind_error, async move { - store.seed(&counter.0, spend.0, lost).await.map(PyNumber) - }) + ttl_seconds: Option, + ) -> PyResult { + Ok(self.0.increment_many_post_call( + &request.0, + &self::increments(increments), + ttl(ttl_seconds)?, + )) } #[pyo3(signature = (counter, spend, ttl_seconds = None))] - fn set<'py>( + fn set_max<'py>( &self, py: Python<'py>, counter: PyData, @@ -179,8 +159,52 @@ impl NativeSpendCounters { ) -> PyResult> { let store = self.0.clone(); let ttl = ttl(ttl_seconds)?; + redis_async(py, "async_set_max", kind_error, async move { + store + .set_max(&counter.0, spend.0, ttl) + .await + .map(|stored| stored.map(PyNumber)) + }) + } + + #[pyo3(signature = (counter, spend, *, increment_if_lost, ttl_seconds = None))] + fn seed<'py>( + &self, + py: Python<'py>, + counter: PyData, + spend: PyNumber, + increment_if_lost: bool, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; + let lost = if increment_if_lost { + LostSeed::Increment + } else { + LostSeed::Answer + }; redis_async(py, "async_set_cache", kind_error, async move { - store.set(&counter.0, spend.0, ttl).await + store + .seed(&counter.0, spend.0, lost, ttl) + .await + .map(PyNumber) + }) + } + + #[pyo3(signature = (counter, spend, memory_ttl_seconds = None, redis_ttl_seconds = None))] + fn set<'py>( + &self, + py: Python<'py>, + counter: PyData, + spend: PyNumber, + memory_ttl_seconds: Option, + redis_ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let memory_ttl = ttl(memory_ttl_seconds)?; + let redis_ttl = ttl(redis_ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.set(&counter.0, spend.0, memory_ttl, redis_ttl).await }) } @@ -196,14 +220,17 @@ impl NativeSpendCounters { }) } + #[pyo3(signature = (counter, ttl_seconds = None))] fn refresh_ttl<'py>( &self, py: Python<'py>, counter: PyData, + ttl_seconds: Option, ) -> PyResult> { let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; redis_async(py, "async_refresh_ttl", kind_error, async move { - store.refresh_ttl(&counter.0).await + store.refresh_ttl(&counter.0, ttl).await }) } diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 3cf468e0d3f..f3674a43b1c 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -164,7 +164,7 @@ impl RedisClient { Pipeline::new(self.0.namespace.clone()) } - pub(crate) fn namespace(&self) -> &Namespace { + pub fn namespace(&self) -> &Namespace { &self.0.namespace } diff --git a/litellm-rust/crates/redis-client/src/namespace.rs b/litellm-rust/crates/redis-client/src/namespace.rs index 4db02b5b7bb..07cbf305332 100644 --- a/litellm-rust/crates/redis-client/src/namespace.rs +++ b/litellm-rust/crates/redis-client/src/namespace.rs @@ -11,7 +11,7 @@ impl Namespace { self.0.as_deref() } - pub(crate) fn apply(&self, key: &str) -> String { + pub fn apply(&self, key: &str) -> String { match &self.0 { Some(namespace) if !is_prefixed(key, namespace) => { format!("{namespace}:{key}") diff --git a/litellm-rust/crates/redis-store/src/spend_counters/keys.rs b/litellm-rust/crates/redis-store/src/spend_counters/keys.rs index 061e19ed1de..7d04ead3dfb 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/keys.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/keys.rs @@ -1,3 +1,5 @@ +use litellm_redis_client::Namespace; + /// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum SpendCounter { @@ -48,6 +50,12 @@ impl SpendCounter { } } + /// `RedisCache.reset_fence_key`: hash-tagged with the namespaced key, so a cluster keeps the + /// fence and the counter on one slot. + pub(super) fn reset_fence_key(&self, namespace: &Namespace) -> String { + format!("{{{}}}:reset_fence", namespace.apply(&self.key())) + } + pub fn parse(key: &str) -> Result { let invalid = || NotASpendCounter(key.to_owned()); let rest = key.strip_prefix("spend:").ok_or_else(invalid)?; @@ -83,3 +91,27 @@ impl SpendCounter { }) } } + +#[cfg(test)] +mod tests { + use litellm_redis_client::Namespace; + use rstest::rstest; + + use super::SpendCounter; + + /// Mirrored by `tests/unit/redis/stores/test_spend_counters.py`, against `RedisCache.reset_fence_key`. + #[rstest] + #[case::no_namespace(Namespace::default(), "{spend:team:t}:reset_fence")] + #[case::namespaced(Namespace::new("litellm"), "{litellm:spend:team:t}:reset_fence")] + fn the_reset_fence_is_tagged_with_the_namespaced_counter( + #[case] namespace: Namespace, + #[case] fence: &str, + ) { + let counter = SpendCounter::Team { + team_id: "t".to_owned(), + window: None, + }; + + assert_eq!(counter.reset_fence_key(&namespace), fence); + } +} diff --git a/litellm-rust/crates/redis-store/src/spend_counters/mod.rs b/litellm-rust/crates/redis-store/src/spend_counters/mod.rs index 98f619cc501..31750e04a78 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/mod.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/mod.rs @@ -16,8 +16,9 @@ use crate::{ }; /// `RedisCache.async_set_max`, byte for byte, sent as `EVAL` like it: a counter only ever rises to -/// the authoritative spend. +/// the authoritative spend, and a reset's fence holds a repair still carrying the pre-reset total. const SET_MAX: &str = concat!( + "if KEYS[2] ~= nil and redis.call('EXISTS', KEYS[2]) == 1 then return redis.call('GET', KEYS[1]) end ", "local cur = redis.call('GET', KEYS[1]) ", "if cur == false or tonumber(cur) < tonumber(ARGV[1]) then ", "redis.call('SET', KEYS[1], ARGV[1]) ", @@ -57,7 +58,7 @@ impl SpendCounters { self.cache.redis().get().is_some() } - /// The TTL every Redis write gives a counter, `None` without Redis. + /// The TTL of a Redis write whose caller names none, `None` without Redis. pub fn redis_ttl(&self) -> Option { self.cache.redis().get().map(|remote| remote.default_ttl) } @@ -119,7 +120,12 @@ impl SpendCounters { /// `increment_spend_counter_cache`: `INCRBYFLOAT` refreshing the TTL, mirrored into memory. A /// failed increment may or may not have landed, so the counter is dropped everywhere before the /// error returns. Without Redis it is a memory increment. - pub async fn increment(&self, counter: &SpendCounter, amount: Number) -> Result { + pub async fn increment( + &self, + counter: &SpendCounter, + amount: Number, + ttl: Option, + ) -> Result { let Some(remote) = self.cache.redis().get() else { return self.increment_local(counter, amount); }; @@ -127,7 +133,7 @@ impl SpendCounters { let increment = commands::incr_by_float( &key, amount.into(), - Some(Expiry::Refresh(remote.default_ttl)), + Some(Expiry::Refresh(ttl.unwrap_or(remote.default_ttl))), ); match remote.client.send(increment).await { Ok(total) => { @@ -147,6 +153,7 @@ impl SpendCounters { pub async fn increment_many( &self, increments: &[(SpendCounter, Number)], + ttl: Option, ) -> Result, Error> { let Some(remote) = self.cache.redis().get() else { return increments @@ -154,12 +161,13 @@ impl SpendCounters { .map(|(counter, amount)| self.increment_local(counter, *amount)) .collect(); }; + let ttl = ttl.unwrap_or(remote.default_ttl); let pipeline = RedisBatch::pipelined(remote.client.clone()); let totals = increments.iter().map(|(counter, amount)| { pipeline.add(commands::incr_by_float( &counter.key(), (*amount).into(), - Some(Expiry::Refresh(remote.default_ttl)), + Some(Expiry::Refresh(ttl)), )) }); let totals = Declared::all(totals.collect::>()).await; @@ -180,10 +188,12 @@ impl SpendCounters { self: &Arc, request: &RequestRedis, increments: &[(SpendCounter, Number)], + ttl: Option, ) -> bool { let Some(remote) = self.cache.redis().get() else { return false; }; + let ttl = ttl.unwrap_or(remote.default_ttl); let batch = Send::PostCall(request).batch(&remote.client); let declared = increments .iter() @@ -191,7 +201,7 @@ impl SpendCounters { let total = batch.add(commands::incr_by_float( &counter.key(), (*amount).into(), - Some(Expiry::KeepExisting(remote.default_ttl)), + Some(Expiry::KeepExisting(ttl)), )); (counter.clone(), total) }) @@ -220,11 +230,12 @@ impl SpendCounters { } /// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either. - /// `None` without Redis, else what Redis holds afterwards. + /// `None` without Redis, else what Redis holds afterwards, a reset's fence holding the write. pub async fn set_max( &self, counter: &SpendCounter, spend: Number, + ttl: Option, ) -> Result, Error> { let key = counter.key(); { @@ -243,11 +254,12 @@ impl SpendCounters { }; let args = ( to_str(&spend.to_value()), - remote.default_ttl.as_secs().to_string(), + ttl.unwrap_or(remote.default_ttl).as_secs().to_string(), ); + let fence = counter.reset_fence_key(remote.client.namespace()); let stored: Option = remote .client - .send(commands::eval_source(SET_MAX, &[&key], args)) + .send(commands::eval_source(SET_MAX, &[&key, &fence], args)) .await?; Ok(stored.map(Number::Float)) } @@ -260,6 +272,7 @@ impl SpendCounters { counter: &SpendCounter, spend: Number, lost: LostSeed, + ttl: Option, ) -> Result { let key = counter.key(); let Some(Remote { @@ -276,8 +289,9 @@ impl SpendCounters { }) })?); }; + let ttl = ttl.unwrap_or(*default_ttl); let current = if client - .send(commands::set_nx(&key, spend.to_json(), *default_ttl)) + .send(commands::set_nx(&key, spend.to_json(), ttl)) .await? { spend @@ -290,7 +304,7 @@ impl SpendCounters { let increment = commands::incr_by_float( &key, spend.into(), - Some(Expiry::KeepExisting(*default_ttl)), + Some(Expiry::KeepExisting(ttl)), ); Number::Float(client.send(increment).await?) } @@ -301,19 +315,20 @@ impl SpendCounters { Ok(current) } - /// A spend reset: memory, then Redis, each for `ttl` or its own default. + /// A spend reset: memory, then Redis, each for its own TTL or its own default. pub async fn set( &self, counter: &SpendCounter, spend: Number, - ttl: Option, + memory_ttl: Option, + redis_ttl: Option, ) -> Result<(), Error> { let key = counter.key(); - self.cache.memory().set(&key, spend, ttl)?; + self.cache.memory().set(&key, spend, memory_ttl)?; let Some(remote) = self.cache.redis().get() else { return Ok(()); }; - let ttl = ttl.unwrap_or(remote.default_ttl); + let ttl = redis_ttl.unwrap_or(remote.default_ttl); remote .client .send(commands::set(&key, spend.to_json(), ttl)) @@ -328,13 +343,20 @@ impl SpendCounters { } /// `refresh_spend_counter_ttl`: `false` without Redis or for a counter Redis no longer holds. - pub async fn refresh_ttl(&self, counter: &SpendCounter) -> Result { + pub async fn refresh_ttl( + &self, + counter: &SpendCounter, + ttl: Option, + ) -> Result { let Some(remote) = self.cache.redis().get() else { return Ok(false); }; Ok(remote .client - .send(commands::expire(&counter.key(), remote.default_ttl)) + .send(commands::expire( + &counter.key(), + ttl.unwrap_or(remote.default_ttl), + )) .await?) } diff --git a/litellm-rust/crates/redis-store/tests/spend_counters.rs b/litellm-rust/crates/redis-store/tests/spend_counters.rs index 75b7af2c321..1519b231303 100644 --- a/litellm-rust/crates/redis-store/tests/spend_counters.rs +++ b/litellm-rust/crates/redis-store/tests/spend_counters.rs @@ -12,7 +12,7 @@ use litellm_redis_store::{ }; use litellm_redis_testing::{RedisDeployment, Topology}; use rstest::rstest; -use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within}; +use support::{COOLDOWN_POLICY, connect, connect_namespaced, connect_observed, remote, within}; use tokio::runtime::Handle; const TTL: Duration = Duration::from_secs(60); @@ -181,7 +181,7 @@ async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl( let counter = team("t"); for (spend, bytes) in SETS { - store.set(&counter, spend, None).await.unwrap(); + store.set(&counter, spend, None, None).await.unwrap(); assert_eq!( raw.send(commands::get::(&counter.key())) .await @@ -191,7 +191,7 @@ async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl( } for (spend, bytes) in SET_MAXES { raw.send(commands::del(&[&counter.key()])).await.unwrap(); - let repaired = store.set_max(&counter, spend).await.unwrap(); + let repaired = store.set_max(&counter, spend, None).await.unwrap(); assert_eq!( raw.send(commands::get::(&counter.key())) .await @@ -201,7 +201,12 @@ async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl( assert_eq!(repaired, Some(Number::Float(spend.as_f64()))); } store - .set(&counter, Number::Int(5), Some(Duration::from_secs(30))) + .set( + &counter, + Number::Int(5), + None, + Some(Duration::from_secs(30)), + ) .await .unwrap(); let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else { @@ -222,10 +227,19 @@ async fn set_max_only_ever_raises_a_counter( token: "hashed".to_owned(), window: None, }; - store.increment(&counter, Number::Float(9.5)).await.unwrap(); + store + .increment(&counter, Number::Float(9.5), None) + .await + .unwrap(); - let lower = store.set_max(&counter, Number::Float(4.0)).await.unwrap(); - let higher = store.set_max(&counter, Number::Float(12.0)).await.unwrap(); + let lower = store + .set_max(&counter, Number::Float(4.0), None) + .await + .unwrap(); + let higher = store + .set_max(&counter, Number::Float(12.0), None) + .await + .unwrap(); assert_eq!(lower, Some(Number::Float(9.5))); assert_eq!(higher, Some(Number::Float(12.0))); @@ -245,6 +259,150 @@ async fn set_max_only_ever_raises_a_counter( ); } +/// Mirrored by `tests/unit/redis/stores/test_spend_counters.py`, which checks the same fence keys +/// against `RedisCache`, the writer of the fence. +#[rstest] +#[tokio::test] +async fn a_reset_fence_holds_a_repair_until_it_is_gone( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values( + (None, "spend:team:t", "{spend:team:t}:reset_fence"), + (Some("litellm"), "litellm:spend:team:t", "litellm:{litellm:spend:team:t}:reset_fence"), + )] + wire: (Option<&str>, &str, &str), +) { + let (namespace, stored_as, fence) = wire; + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let client = match namespace { + Some(namespace) => connect_namespaced(&deployment, namespace), + None => connect(&deployment), + }; + let store = store(Some(client)); + let counter = team("t"); + store + .set(&counter, Number::Float(0.0), None, None) + .await + .unwrap(); + raw.send(commands::set(fence, "1", TTL)).await.unwrap(); + + let held = store + .set_max(&counter, Number::Float(12.0), None) + .await + .unwrap(); + + assert_eq!(held, Some(Number::Float(0.0))); + assert_eq!( + raw.send(commands::get::(stored_as)).await.unwrap(), + Some("0.0".to_owned()) + ); + + raw.send(commands::del(&[fence])).await.unwrap(); + let repaired = store + .set_max(&counter, Number::Float(12.0), None) + .await + .unwrap(); + + assert_eq!(repaired, Some(Number::Float(12.0))); +} + +#[derive(Debug, Clone, Copy)] +enum Write { + Increment, + IncrementMany, + PostCall, + SetMax, + Seed, + Set, + RefreshTtl, +} + +#[rstest] +#[case::increment(Write::Increment)] +#[case::increment_many(Write::IncrementMany)] +#[case::post_call(Write::PostCall)] +#[case::set_max(Write::SetMax)] +#[case::seed(Write::Seed)] +#[case::set(Write::Set)] +#[case::refresh_ttl(Write::RefreshTtl)] +#[tokio::test] +async fn every_write_lasts_for_the_ttl_its_caller_names( + #[case] write: Write, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + let counter = team("t"); + let long = Some(Duration::from_secs(3600)); + let one = Number::Int(1); + + match write { + Write::Increment => { + store.increment(&counter, one, long).await.unwrap(); + } + Write::IncrementMany => { + store + .increment_many(&[(counter.clone(), one)], long) + .await + .unwrap(); + } + Write::PostCall => { + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + assert!(store.increment_many_post_call(&request, &[(counter.clone(), one)], long)); + request.flush_post().await; + litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await; + } + Write::SetMax => { + store.set_max(&counter, one, long).await.unwrap(); + } + Write::Seed => { + store + .seed(&counter, one, LostSeed::Answer, long) + .await + .unwrap(); + } + Write::Set => { + store.set(&counter, one, None, long).await.unwrap(); + } + Write::RefreshTtl => { + store.increment(&counter, one, None).await.unwrap(); + assert!(store.refresh_ttl(&counter, long).await.unwrap()); + } + } + + let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else { + panic!("{write:?} left the counter without a TTL"); + }; + assert!(within(remaining, 3600), "{write:?} {remaining:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_reset_gives_memory_and_redis_their_own_ttl() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + let counter = team("t"); + + store + .set( + &counter, + Number::Int(5), + Some(Duration::from_millis(1)), + Some(Duration::from_secs(3600)), + ) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + + assert_eq!(store.get_local(&counter).unwrap(), None); + let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else { + panic!("a reset counter has no TTL"); + }; + assert!(within(remaining, 3600), "{remaining:?}"); +} + #[rstest] #[tokio::test] async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory( @@ -257,7 +415,7 @@ async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory( let store = store(Some(client)); let [first, second] = [team("a"), team("b")]; other_pod - .increment(&first, Number::Float(1.25)) + .increment(&first, Number::Float(1.25), None) .await .unwrap(); raw.send(commands::expire(&first.key(), Duration::from_secs(5))) @@ -265,10 +423,13 @@ async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory( .unwrap(); let totals = store - .increment_many(&[ - (first.clone(), Number::Float(0.5)), - (second.clone(), Number::Int(2)), - ]) + .increment_many( + &[ + (first.clone(), Number::Float(0.5)), + (second.clone(), Number::Int(2)), + ], + None, + ) .await .unwrap(); @@ -301,7 +462,7 @@ async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_an let declared = std::thread::spawn({ let (store, request, counter) = (Arc::clone(&store), request.clone(), counter.clone()); - move || store.increment_many_post_call(&request, &[(counter, Number::Float(3.5))]) + move || store.increment_many_post_call(&request, &[(counter, Number::Float(3.5))], None) }) .join() .expect("a Python thread has no Tokio context, and declaring there must not panic"); @@ -331,7 +492,11 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory( let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); let counter = team("t"); - assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))])); + assert!(store.increment_many_post_call( + &request, + &[(counter.clone(), Number::Float(3.5))], + None + )); request.flush_pre().await; assert!(trips.seen().is_empty()); @@ -417,11 +582,11 @@ async fn a_seed_lets_the_first_pod_win_and_others_answer_with_its_value( let counter = SpendCounter::Org("org".to_owned()); let won = first - .seed(&counter, Number::Float(10.0), LostSeed::Answer) + .seed(&counter, Number::Float(10.0), LostSeed::Answer, None) .await .unwrap(); let lost = second - .seed(&counter, Number::Float(99.0), LostSeed::Increment) + .seed(&counter, Number::Float(99.0), LostSeed::Increment, None) .await .unwrap(); @@ -452,13 +617,16 @@ async fn delete_drops_memory_and_redis_and_refresh_ttl_reports_whether_redis_hel let store = store(Some(connect(&deployment))); let counters = [team("a"), team("b")]; for counter in &counters { - store.increment(counter, Number::Int(1)).await.unwrap(); + store + .increment(counter, Number::Int(1), None) + .await + .unwrap(); } raw.send(commands::expire(&counters[0].key(), Duration::from_secs(5))) .await .unwrap(); - assert!(store.refresh_ttl(&counters[0]).await.unwrap()); + assert!(store.refresh_ttl(&counters[0], None).await.unwrap()); let Ttl::Expires(remaining) = raw.send(commands::ttl(&counters[0].key())).await.unwrap() else { panic!("a refreshed counter has no TTL"); }; @@ -475,7 +643,7 @@ async fn delete_drops_memory_and_redis_and_refresh_ttl_reports_whether_redis_hel ); assert_eq!(store.get_local(counter).unwrap(), None); } - assert!(!store.refresh_ttl(&counters[0]).await.unwrap()); + assert!(!store.refresh_ttl(&counters[0], None).await.unwrap()); } // main's read_spend_counter_cache_value: RedisCache.async_get_cache swallows a failed GET and @@ -487,12 +655,18 @@ async fn a_failed_read_serves_the_local_counter_and_a_failed_increment_drops_the let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); let store = store(Some(connect(&deployment))); let [kept, failed] = [team("kept"), team("failed")]; - store.increment(&kept, Number::Float(2.0)).await.unwrap(); - store.increment(&failed, Number::Float(3.0)).await.unwrap(); + store + .increment(&kept, Number::Float(2.0), None) + .await + .unwrap(); + store + .increment(&failed, Number::Float(3.0), None) + .await + .unwrap(); deployment.stop().await.unwrap(); let read = store.read(&kept).await.unwrap(); - let increment = store.increment(&failed, Number::Float(1.0)).await; + let increment = store.increment(&failed, Number::Float(1.0), None).await; assert_eq!(read, Reading::Memory(Some(Number::Float(2.0)))); assert!(increment.is_err()); @@ -521,7 +695,10 @@ async fn a_read_with_an_open_breaker_serves_the_local_counter() { ); let store = store(Some(support::connect_guarded(&deployment, breaker))); let counter = team("t"); - store.increment(&counter, Number::Float(2.0)).await.unwrap(); + store + .increment(&counter, Number::Float(2.0), None) + .await + .unwrap(); deployment.stop().await.unwrap(); let failed = store.read(&counter).await.unwrap(); @@ -556,8 +733,15 @@ async fn a_failed_post_call_increment_drops_the_counter_from_memory() { let store = store(Some(connect(&deployment))); let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); let counter = team("t"); - store.increment(&counter, Number::Float(1.0)).await.unwrap(); - assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))])); + store + .increment(&counter, Number::Float(1.0), None) + .await + .unwrap(); + assert!(store.increment_many_post_call( + &request, + &[(counter.clone(), Number::Float(1.0))], + None + )); deployment.stop().await.unwrap(); request.flush_post().await; @@ -573,17 +757,23 @@ async fn without_redis_counters_live_in_memory() { let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); let totals = store - .increment_many(&[ - (counter.clone(), Number::Int(2)), - (counter.clone(), Number::Float(0.5)), - ]) + .increment_many( + &[ + (counter.clone(), Number::Int(2)), + (counter.clone(), Number::Float(0.5)), + ], + None, + ) .await .unwrap(); let seeded = store - .seed(&counter, Number::Float(1.0), LostSeed::Answer) + .seed(&counter, Number::Float(1.0), LostSeed::Answer, None) + .await + .unwrap(); + let raised = store + .set_max(&counter, Number::Float(7.0), None) .await .unwrap(); - let raised = store.set_max(&counter, Number::Float(7.0)).await.unwrap(); assert_eq!(totals, [Number::Int(2), Number::Float(2.5)]); assert_eq!(seeded, Number::Float(2.5)); @@ -592,7 +782,7 @@ async fn without_redis_counters_live_in_memory() { store.read(&counter).await.unwrap(), Reading::Memory(Some(Number::Float(7.0))) ); - assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))])); + assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))], None)); assert!( store .read_many(RedisSend::PreCall(&request), &[counter]) @@ -616,7 +806,7 @@ async fn spend_counters_send_the_commands_main_sends( .unwrap(); store - .increment(&counters[0], Number::Float(1.5)) + .increment(&counters[0], Number::Float(1.5), None) .await .unwrap(); store @@ -626,11 +816,11 @@ async fn spend_counters_send_the_commands_main_sends( .unwrap() .unwrap(); store - .set_max(&counters[0], Number::Float(9.0)) + .set_max(&counters[0], Number::Float(9.0), None) .await .unwrap(); store - .seed(&counters[1], Number::Float(2.0), LostSeed::Answer) + .seed(&counters[1], Number::Float(2.0), LostSeed::Answer, None) .await .unwrap(); store.delete(&counters).await.unwrap(); diff --git a/litellm-rust/crates/redis-store/tests/support/mod.rs b/litellm-rust/crates/redis-store/tests/support/mod.rs index f4b05c304e3..b323569c0bb 100644 --- a/litellm-rust/crates/redis-store/tests/support/mod.rs +++ b/litellm-rust/crates/redis-store/tests/support/mod.rs @@ -35,6 +35,10 @@ pub fn connect(deployment: &RedisDeployment) -> RedisClient { connect_with(deployment, None, None) } +pub fn connect_namespaced(deployment: &RedisDeployment, namespace: &str) -> RedisClient { + build(deployment, Namespace::new(namespace), None, None) +} + /// `client` with `RedisCache`'s default TTL when `litellm.default_redis_ttl` is unset. pub fn remote(client: RedisClient) -> Remote { Remote { @@ -101,6 +105,15 @@ fn connect_with( deployment: &RedisDeployment, breaker: Option, events: Option>, +) -> RedisClient { + build(deployment, Namespace::default(), breaker, events) +} + +fn build( + deployment: &RedisDeployment, + namespace: Namespace, + breaker: Option, + events: Option>, ) -> RedisClient { let endpoint = match deployment { RedisDeployment::Standalone(server) => Target::Url(RedisUrl { @@ -124,7 +137,7 @@ fn connect_with( RedisClient::new(RedisConfig { target: endpoint, credentials: Credentials::NONE, - namespace: Namespace::default(), + namespace, timeouts: Timeouts::new(Duration::from_millis(500), None), breaker, events, diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 4d9d26fc70c..c3b79ad858e 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -20,7 +20,7 @@ import threading import time import weakref from collections import Counter -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta @@ -485,6 +485,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) +async def _first_keys(scan: AsyncIterator[object], count: int | None) -> AsyncIterator[str]: + taken: Final = itertools.count(1) + async for key in scan: + yield key.decode() if isinstance(key, bytes) else str(key) + if count is not None and next(taken) >= count: + return + + _LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet() @@ -1025,10 +1033,9 @@ class RedisCache(BaseCache): return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) @_redis_circuit_breaker_guard - async def async_scan_iter(self, pattern: str, count: int = 100) -> list: + async def async_scan_iter(self, pattern: str, count: int = 100, *, unbounded: bool = False) -> list[str]: start_time: Final = time.time() try: - keys: Final = [] _redis_client: Final = self.init_async_client() if not hasattr(_redis_client, "scan_iter"): verbose_logger.debug( @@ -1036,11 +1043,13 @@ class RedisCache(BaseCache): ) return [] - pattern = self.check_and_fix_namespace(key=pattern) - async for key in _redis_client.scan_iter(match=pattern + "*", count=count): - keys.append(key) - if len(keys) >= count: - break + namespaced_pattern: Final = self.check_and_fix_namespace(key=pattern) + keys: Final = [ + key + async for key in _first_keys( + _redis_client.scan_iter(match=namespaced_pattern + "*", count=count), None if unbounded else count + ) + ] ## LOGGING ## end_time = time.time() @@ -1343,10 +1352,15 @@ class RedisCache(BaseCache): _record_swallowed_redis_failure(self._circuit_breaker, e) @_redis_circuit_breaker_guard - async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None: - """One round trip for writes whose TTLs differ; a ``None`` TTL falls back to the default TTL.""" + async def async_set_cache_pipeline_with_ttls( + self, cache_list: Sequence[tuple[str, object, float | None]] + ) -> Exception | None: + """One round trip for writes whose TTLs differ; a ``None`` TTL falls back to the default TTL. + + Returns the Redis error when the writes were lost, so a caller that must not report + success on a lost write can tell. An open circuit breaker still raises.""" if len(cache_list) == 0: - return + return None commands: Final = tuple( (self.check_and_fix_namespace(key=cache_key), json.dumps(cache_value), self.get_ttl(ttl=ttl)) for cache_key, cache_value, ttl in cache_list @@ -1357,16 +1371,6 @@ class RedisCache(BaseCache): for cache_key, json_cache_value, ttl in commands: pipe.set(name=cache_key, value=json_cache_value, ex=None if ttl is None else timedelta(seconds=ttl)) await pipe.execute() - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.REDIS, - duration=time.time() - start_time, - call_type="async_set_cache_pipeline_with_ttls", - caller=_get_call_stack_info(), - start_time=start_time, - end_time=time.time(), - ) - ) except Exception as e: asyncio.create_task( self.service_logger_obj.async_service_failure_hook( @@ -1386,6 +1390,19 @@ class RedisCache(BaseCache): e, ) _record_swallowed_redis_failure(self._circuit_breaker, e) + return e + else: + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=time.time() - start_time, + call_type="async_set_cache_pipeline_with_ttls", + caller=_get_call_stack_info(), + start_time=start_time, + end_time=time.time(), + ) + ) + return None async def _set_cache_sadd_helper( self, @@ -1562,12 +1579,18 @@ class RedisCache(BaseCache): ) raise e + def reset_fence_key(self, key: str) -> str: + """Hash-tagged with the namespaced ``key`` so a cluster keeps the fence and the key on one slot.""" + return f"{{{self.check_and_fix_namespace(key=key)}}}:reset_fence" + @_redis_circuit_breaker_guard async def async_set_max( self, key: str, value: float, ttl: int | None = None, + *, + fence_key: str | None = None, ) -> float | None: """Atomically set ``key`` to ``value`` only when ``value`` is greater than the stored value (or the key is unset), refreshing the TTL. @@ -1577,11 +1600,17 @@ class RedisCache(BaseCache): concurrent increment that has already pushed the counter higher. The GET/compare/SET runs in a single Lua call, so it is also atomic across racing callers and pods. Returns the resulting value. + + While ``fence_key`` exists the write is held and the stored value returned + as is: a spend reset sets one (``reset_fence_key``) so a repair still + carrying the pre-reset total cannot raise the counter the reset zeroed. """ _redis_client: Final = self.init_async_client() _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) + keys: Final = (key,) if fence_key is None else (key, self.check_and_fix_namespace(key=fence_key)) lua: Final = ( + "if KEYS[2] ~= nil and redis.call('EXISTS', KEYS[2]) == 1 then return redis.call('GET', KEYS[1]) end " "local cur = redis.call('GET', KEYS[1]) " "if cur == false or tonumber(cur) < tonumber(ARGV[1]) then " "redis.call('SET', KEYS[1], ARGV[1]) " @@ -1591,7 +1620,7 @@ class RedisCache(BaseCache): ) result = cast( "str | bytes | int | float | None", - await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))), + await _redis_client.eval(lua, len(keys), *keys, str(value), str(int(_used_ttl or 0))), ) if result is None: return None @@ -1968,11 +1997,10 @@ class RedisCache(BaseCache): raise e @_redis_circuit_breaker_guard - async def delete_cache_keys(self, keys): + async def delete_cache_keys(self, keys: Iterable[str]) -> None: _redis_client: Final = self._async_commands() - keys = [self.check_and_fix_namespace(key=key) for key in keys] - # keys is a list, unpack it so it gets passed as individual elements to delete - await _redis_client.delete(*keys) + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in keys] + await _redis_client.delete(*namespaced_keys) def client_list(self) -> list: client_list: Final[list] = self.redis_client.client_list() diff --git a/litellm/constants.py b/litellm/constants.py index b43f6d9f358..fe1762c5d2b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1988,6 +1988,7 @@ _background_interaction_cost_polling_env: Final = os.getenv( BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true" BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS: Final = 300.0 PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605) +SPEND_COUNTER_REDIS_TTL_SECONDS: Final = 3600 PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b40eb451979..9746957760e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3350,7 +3350,7 @@ async def invalidate_team_member_spend_state( spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}" try: - await spend_counter_store.set(spend_counter_key, new_spend, ttl=60) + await spend_counter_store.set(spend_counter_key, new_spend, memory_ttl=60) except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up verbose_proxy_logger.warning( "Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next " diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 36650629300..f1917d9ff4a 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,15 +1,13 @@ import asyncio import json from collections.abc import Mapping, Sequence -from dataclasses import asdict, dataclass -from types import MappingProxyType +from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # legacy module exports ConfigSyncPubSub, _ConfigSyncPubSub, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -36,12 +34,14 @@ AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" _POLL_TIMEOUT_SECONDS: Final = 1.0 _MAX_PENDING_PUBLISHES: Final = 1024 _MAX_IN_FLIGHT_PUBLISHES: Final = 16 +_KEYS_PER_MESSAGE: Final = 500 +_MESSAGE_CACHE_KEYS: Final = TypeAdapter(tuple[str, ...]) +_MESSAGE_FIELDS: Final = TypeAdapter(Mapping[str, object]) _BULK_EVICTION_WAIT_SECONDS: Final = 2.0 _pending_publishes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs keep background publishes alive _in_flight_publishes: Final = asyncio.Semaphore(_MAX_IN_FLIGHT_PUBLISHES) _BACKOFF_INITIAL_SECONDS: Final = 5.0 _BACKOFF_MAX_SECONDS: Final = 60.0 -_DECODED_JSON: Final = TypeAdapter(object) def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @@ -52,16 +52,21 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @dataclass(frozen=True, slots=True) class _CacheInvalidationMessage: - cache_key: str + cache_keys: tuple[str, ...] new_value: float | None = None ttl: float | None = None -def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str: - message: Final[Mapping[str, object]] = MappingProxyType( - asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) - ) - return json.dumps({field: value for field, value in message.items() if value is not None}) +def _cache_invalidation_message_json( + cache_keys: Sequence[str], new_value: float | None = None, ttl: float | None = None +) -> str: + fields: Final = { + "cache_key": cache_keys[0] if len(cache_keys) == 1 else None, + "cache_keys": list(cache_keys) if len(cache_keys) > 1 else None, + "new_value": new_value, + "ttl": ttl, + } + return json.dumps({field: value for field, value in fields.items() if value is not None}) def _finite_number_or_none(value: object) -> float | None: @@ -70,22 +75,29 @@ def _finite_number_or_none(value: object) -> float | None: return float(value) +def _message_cache_keys(single: object, many: object) -> tuple[str, ...]: + if isinstance(single, str): + return (single,) + try: + return _MESSAGE_CACHE_KEYS.validate_python(many) + except ValidationError: + return () + + def _message_from_data(data: object) -> _CacheInvalidationMessage | None: if isinstance(data, bytes): data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str if not isinstance(data, str): return None try: - parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(data))) - except json.JSONDecodeError: + parsed: Final = _MESSAGE_FIELDS.validate_json(data) + except ValidationError: return None - if parsed is None: - return None - cache_key: Final = parsed.get("cache_key") - if not isinstance(cache_key, str): + cache_keys: Final = _message_cache_keys(parsed.get("cache_key"), parsed.get("cache_keys")) + if len(cache_keys) == 0: return None return _CacheInvalidationMessage( - cache_key=cache_key, + cache_keys=cache_keys, new_value=_finite_number_or_none(parsed.get("new_value")), ttl=_finite_number_or_none(parsed.get("ttl")), ) @@ -106,7 +118,7 @@ async def _publish_batch_to_redis(redis_cache: "RedisCache", cache_keys: tuple[s channel: Final = auth_cache_invalidation_channel(redis_cache) pipeline: Final = client.pipeline(transaction=False) for cache_key in cache_keys: - _ = pipeline.publish(channel, _cache_invalidation_message_json(cache_key)) + _ = pipeline.publish(channel, _cache_invalidation_message_json((cache_key,))) async with _in_flight_publishes: _ = await pipeline.execute() except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors @@ -126,6 +138,19 @@ def _has_room_for_publish(key_count: int) -> bool: return False +def _queue_publish( + redis_cache: "RedisCache", cache_keys: tuple[str, ...], new_value: float | None, ttl: float | None +) -> None: + if not _has_room_for_publish(len(cache_keys)): + return + task: Final = asyncio.create_task( + _publish_to_redis( + redis_cache, cache_keys[0], _cache_invalidation_message_json(cache_keys, new_value=new_value, ttl=ttl) + ) + ) + _pending_publishes.add(task) + + async def publish_auth_cache_invalidation( cache_key: str, new_value: float | None = None, ttl: float | None = None ) -> None: @@ -148,15 +173,17 @@ async def publish_auth_cache_invalidation( Redis connection at once; the rest wait in the task set, so a wedge cannot drain the shared connection pool. """ + await publish_auth_cache_invalidations(cache_keys=(cache_key,), new_value=new_value, ttl=ttl) + + +async def publish_auth_cache_invalidations( + cache_keys: Sequence[str], new_value: float | None = None, ttl: float | None = None +) -> None: redis_cache: Final = coordination_redis_cache() - if redis_cache is None or not _has_room_for_publish(1): + if redis_cache is None: return - task: Final = asyncio.create_task( - _publish_to_redis( - redis_cache, cache_key, _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl) - ) - ) - _pending_publishes.add(task) + for start in range(0, len(cache_keys), _KEYS_PER_MESSAGE): + _queue_publish(redis_cache, tuple(cache_keys[start : start + _KEYS_PER_MESSAGE]), new_value, ttl) await asyncio.sleep(0) @@ -239,18 +266,31 @@ 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") + """``counter_cache`` is the one cache a value-carrying message (a spend reset) is written into; + a delete reaches every additional cache. Writing a reset's thousands of counters into every + additional cache would evict what the others hold, the BYOK credentials among them.""" + + __slots__ = ( + "_additional_in_memory_caches", + "_counter_cache", + "_redis_cache", + "_spend_counters", + "_task", + "_user_api_key_cache", + ) def __init__( self, redis_cache: "RedisCache", user_api_key_cache: "UserApiKeyCache", additional_in_memory_caches: Sequence["InMemoryCache"] = (), + counter_cache: "InMemoryCache | None" = None, spend_counters: SpendCounterCopies | None = None, ) -> None: self._redis_cache = redis_cache self._user_api_key_cache = user_api_key_cache self._additional_in_memory_caches = tuple(additional_in_memory_caches) + self._counter_cache = counter_cache self._spend_counters: Final = spend_counters self._task: asyncio.Task[None] | None = None @@ -307,14 +347,18 @@ class AuthCacheInvalidationSubscriber: if parsed is None: return if self._spend_counters is not None: - self._spend_counters.apply_broadcast(parsed.cache_key, parsed.new_value, parsed.ttl) + for cache_key in parsed.cache_keys: + self._spend_counters.apply_broadcast(cache_key, parsed.new_value, parsed.ttl) if parsed.new_value is not None: - for additional_cache in self._additional_in_memory_caches: - additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl) + if self._counter_cache is None: + return + for cache_key in parsed.cache_keys: + self._counter_cache.set_cache(cache_key, parsed.new_value, ttl=parsed.ttl) return - _drop_from_memory(parsed.cache_key, self._user_api_key_cache) - for additional_cache in self._additional_in_memory_caches: - additional_cache.delete_cache(parsed.cache_key) + for cache_key in parsed.cache_keys: + _drop_from_memory(cache_key, self._user_api_key_cache) + for additional_cache in self._additional_in_memory_caches: + additional_cache.delete_cache(cache_key) @staticmethod async def _close_pubsub(pubsub: ConfigSyncPubSub) -> None: diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 4ed3ac30be7..c8f67b60ebb 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -132,7 +132,7 @@ class SpendCounterReseed: # Per-window key/team counters share prefixes with primary counters # but don't correspond to a DB row. Do not reject arbitrary entity IDs # or tag names that merely contain ":window:". - if SpendCounterReseed._is_key_or_team_window_counter(counter_key): + if SpendCounterReseed.is_key_or_team_window_counter(counter_key): return None try: row: Final = await bounded_db_lookup( @@ -194,7 +194,7 @@ class SpendCounterReseed: return float(row.spend or 0.0) @staticmethod - def _is_key_or_team_window_counter(counter_key: str) -> bool: + def is_key_or_team_window_counter(counter_key: str) -> bool: for prefix in ("spend:key:", "spend:team:"): if not counter_key.startswith(prefix): continue diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d4a0a35d9eb..e17a160b3a4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4270,7 +4270,7 @@ async def update_key_fn( counter_key: Final = key_counter(hash_token_if_needed(key)) try: - await spend_counter_store.set(counter_key, data.spend, ttl=60) + await spend_counter_store.set(counter_key, data.spend, memory_ttl=60) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to update spend counter %s in Redis after key spend update: %s. " diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 4e511aee24a..8071cb9dd48 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -817,6 +817,17 @@ async def update_organization( include={"members": True, "teams": True, "litellm_budget_table": True}, ) + if data.spend is not None: + from litellm.proxy.proxy_server import ( + raise_if_spend_counters_unwritten, + set_spend_counters_with_floor_and_broadcast, + ) + + lost: Final = await set_spend_counters_with_floor_and_broadcast( + counter_keys=(f"spend:org:{data.organization_id}",), value=data.spend + ) + raise_if_spend_counters_unwritten(lost, saved="The organization spend was saved") + return response diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3297a298af4..e437bce32cf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -66,6 +66,7 @@ from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, Valida from pydantic.fields import FieldInfo, PydanticUndefined from pydantic.main import IncEx from pydantic_core import PydanticSerializationError, to_jsonable_python +from redis.exceptions import RedisError from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._uuid import uuid @@ -81,6 +82,7 @@ from litellm.constants import ( BASE_MCP_ROUTE, DAILY_TAG_SPEND_BATCH_MULTIPLIER, DEFAULT_MAX_RECURSE_DEPTH, + DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, DEFAULT_SLACK_ALERTING_THRESHOLD, @@ -369,6 +371,7 @@ from litellm.constants import ( REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, SPEND_CAPTURE_RATE_CHECK_JOB_ID, + SPEND_COUNTER_REDIS_TTL_SECONDS, SPEND_LOG_SIZE_CHECK_INTERVAL, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, @@ -520,6 +523,7 @@ from litellm.proxy.common_utils.anthropic_error_payload import ( from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, publish_auth_cache_invalidation, + publish_auth_cache_invalidations, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.codex_catalog_source import codex_catalog_source @@ -3169,7 +3173,24 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) -spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters()) +_SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE: Final = 10 + + +def spend_counter_redis_ttl() -> int: + """Redis TTL for spend counters. A counter that expires before its spend is flushed to the DB reseeds from a + stale-low row and re-admits a key that is over budget, so it must outlive many batch-write intervals.""" + batch_write_at: Final = ( + proxy_batch_write_at if isinstance(proxy_batch_write_at, (int, float)) else PROXY_BATCH_WRITE_AT + ) + configured_redis_ttl: Final = int(litellm.default_redis_ttl) if litellm.default_redis_ttl is not None else 0 + return max( + SPEND_COUNTER_REDIS_TTL_SECONDS, + int(_SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE * batch_write_at), + configured_redis_ttl, + ) + + +spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters(), spend_counter_redis_ttl) cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) saml_sso_cache: Final = DualCache(default_in_memory_ttl=SAML_AUTHN_REQUEST_TTL_SECONDS) model_max_budget_limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache) @@ -3367,6 +3388,8 @@ async def _drain_spend_event_producer_on_shutdown() -> None: # Bounds authoritative DB re-reads when enforcing a budget against a # stale-low spend counter: at most one DB read per counter per window. SPEND_DB_FLOOR_CACHE_TTL_SECONDS: Final = 5 +SPEND_RESET_FENCE_TTL_SECONDS: Final = 2 * SPEND_DB_FLOOR_CACHE_TTL_SECONDS +SPEND_DB_FLOOR_MARKER_PREFIX: Final = "spend_db_floor:" def _fail_closed_budget_enforcement() -> bool: @@ -4436,6 +4459,125 @@ async def invalidate_spend_counter(counter_key: str) -> None: _invalidate_spend_counter: Final = invalidate_spend_counter +_IN_MEMORY_COUNTER_KEYS: Final = TypeAdapter(tuple[str, ...]) + + +_SPEND_COUNTER_STORE_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, ValidationError) + + +def _spend_counter_keys_in_memory(prefix: str) -> frozenset[str]: + floor_prefix: Final = f"{SPEND_DB_FLOOR_MARKER_PREFIX}{prefix}" + cached_keys: Final = _IN_MEMORY_COUNTER_KEYS.validate_python(tuple(spend_counter_cache.in_memory_cache.cache_dict)) + counters: Final = (key for key in spend_counter_store.snapshot() if key.startswith(prefix)) + floors: Final = ( + key.removeprefix(SPEND_DB_FLOOR_MARKER_PREFIX) for key in cached_keys if key.startswith(floor_prefix) + ) + return frozenset((*counters, *floors)) + + +async def _spend_counter_keys_in_redis(redis_cache: RedisCache, prefix: str) -> frozenset[str] | Exception: + namespace_prefix: Final = f"{redis_cache.namespace}:" if redis_cache.namespace else "" + try: + scanned: Final = await redis_cache.async_scan_iter(pattern=prefix, count=500, unbounded=True) + except _SPEND_COUNTER_STORE_FAILURES as e: + return e + return frozenset(key.removeprefix(namespace_prefix) for key in scanned) + + +def _keys_per_reset_pipeline(max_batch_size: int) -> int: + return max(1, max_batch_size // 2) + + +_KEYS_PER_RESET_PIPELINE: Final = _keys_per_reset_pipeline(DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE) + + +async def _write_reset_batch(redis_cache: RedisCache, batch: Sequence[str], value: float) -> Exception | None: + """The fences go first, so a repair that lands between a fence and its counter already finds the fence. The + counters go through ``spend_counter_store``, the client every request increments them with, so a reset reads as + written only once that client accepts writes again.""" + fences: Final = tuple((redis_cache.reset_fence_key(key), 1, SPEND_RESET_FENCE_TTL_SECONDS) for key in batch) + try: + fenced: Final = await redis_cache.async_set_cache_pipeline_with_ttls(fences) + except RedisCircuitBreakerOpenError as e: + return e + if fenced is not None: + return fenced + written: Final = await asyncio.gather( + *(spend_counter_store.set(key, value, memory_ttl=60) for key in batch), return_exceptions=True + ) + return next((outcome for outcome in written if isinstance(outcome, Exception)), None) + + +async def _write_spend_counters_to_redis( + redis_cache: RedisCache, keys: Sequence[str], value: float +) -> Exception | None: + for start in range(0, len(keys), _KEYS_PER_RESET_PIPELINE): + if ( + lost := await _write_reset_batch(redis_cache, keys[start : start + _KEYS_PER_RESET_PIPELINE], value) + ) is not None: + return lost + return None + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def set_spend_counters_with_floor_and_broadcast(counter_keys: Collection[str], value: float) -> Exception | None: + """A spend reset SETs each counter and its spend_db_floor marker on every pod, never + deletes them: a pod that cached the pre-reset floor within SPEND_DB_FLOOR_CACHE_TTL_SECONDS + would otherwise repair the reseeded counter right back to the old total. The floors are + broadcast before the counters so a subscriber never holds the new counter under the old + floor, and a reset fence in Redis holds any repair that still races the broadcast. + + Returns the Redis error when the counters could not be written there: this pod and the + broadcast still carry the new value, but the other pods fall back to the stale Redis + counters once their copies expire, so the caller must not report success.""" + keys: Final = tuple(sorted(counter_keys)) + if len(keys) == 0: + return None + floor_keys: Final = tuple(f"{SPEND_DB_FLOOR_MARKER_PREFIX}{key}" for key in keys) + for key, floor_key in zip(keys, floor_keys, strict=True): + spend_counter_store.set_local(key, value, 60) + spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + redis_cache: Final = spend_counter_cache.redis_cache + lost: Final = None if redis_cache is None else await _write_spend_counters_to_redis(redis_cache, keys, value) + await publish_auth_cache_invalidations(floor_keys, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + await publish_auth_cache_invalidations(keys, new_value=value, ttl=60) + return lost + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def reset_spend_counters_with_prefix(prefix: str) -> Exception | None: + """Resets every counter under ``prefix`` this pod or Redis knows; returns the Redis error + when the counters there could not be listed or written, with this pod already reset.""" + redis_cache: Final = spend_counter_cache.redis_cache + in_redis: Final[frozenset[str] | Exception] = ( + frozenset() if redis_cache is None else await _spend_counter_keys_in_redis(redis_cache, prefix) + ) + listed: Final[frozenset[str]] = frozenset() if isinstance(in_redis, Exception) else in_redis + counters: Final = frozenset( + key + for key in _spend_counter_keys_in_memory(prefix) | listed + if not SpendCounterReseed.is_key_or_team_window_counter(key) + ) + lost: Final = await set_spend_counters_with_floor_and_broadcast(counters, value=0.0) + return in_redis if isinstance(in_redis, Exception) else lost + + +def raise_if_spend_counters_unwritten(lost: Exception | None, saved: str) -> None: + """Maps a lost Redis counter write to the public error: the database write stands, so the + caller retries the same request once Redis accepts writes rather than undoing anything.""" + if lost is None: + return + raise ProxyException( + message=( + f"{saved}, but the cached spend counters could not be reset in Redis ({lost}). Budget checks on other " + "instances keep the previous spend until those counters expire, so retry this request once Redis " + "accepts writes." + ), + type=ProxyErrorTypes.internal_server_error, + param=None, + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + @with_service_target(SPEND_COUNTERS_TARGET) async def set_spend_counter_with_floor_and_broadcast(counter_key: str, value: float) -> None: @@ -4452,7 +4594,7 @@ async def set_spend_counter_with_floor_and_broadcast(counter_key: str, value: fl or floor marker keeps enforcing against it until its own TTL expires. """ try: - await spend_counter_store.set(counter_key, value, ttl=60) + await spend_counter_store.set(counter_key, value, memory_ttl=60) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to update spend counter %s in Redis: %s. Budget checks may use stale value until counter expires.", @@ -9048,6 +9190,7 @@ class ProxyConfig: redis_cache=redis_cache, user_api_key_cache=user_api_key_cache, additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache), + counter_cache=spend_counter_cache.in_memory_cache, spend_counters=spend_counter_store, ) self.auth_cache_invalidation_subscriber = subscriber diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 0348d8e1f82..585844dd7d5 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3611,7 +3611,11 @@ async def global_spend_reset(): 3. LiteLLM_TeamTable spend will be set = 0 """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + raise_if_spend_counters_unwritten, + reset_spend_counters_with_prefix, + ) if prisma_client is None: raise ProxyException( @@ -3626,6 +3630,9 @@ async def global_spend_reset(): _ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) _ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL) _ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={}) + keys_lost: Final = await reset_spend_counters_with_prefix(prefix="spend:key:") + teams_lost: Final = await reset_spend_counters_with_prefix(prefix="spend:team:") + raise_if_spend_counters_unwritten(keys_lost if keys_lost is not None else teams_lost, saved="Spend was reset") return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/litellm/redis/stores/spend_counters/__init__.py b/litellm/redis/stores/spend_counters/__init__.py index 3ad06589f53..5d6923e365a 100644 --- a/litellm/redis/stores/spend_counters/__init__.py +++ b/litellm/redis/stores/spend_counters/__init__.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -62,10 +62,16 @@ def _by_counter(counters: Sequence[str], read: Sequence[float | None] | None) -> class SpendCounterStore: """The ``spend:*`` budget counters: Redis shared by every worker, and this worker's memory copy that answers - only when Redis cannot. Counters are named by the builders in ``types``.""" + only when Redis cannot. Counters are named by the builders in ``types``. ``counter_ttl`` gives each Redis write + its TTL when it happens, since the settings it depends on load after the store is built; without it a write gets + the Redis client's default.""" - def __init__(self, handle: SpendCountersHandle) -> None: + def __init__(self, handle: SpendCountersHandle, counter_ttl: Callable[[], float] | None = None) -> None: self._handle: Final = handle + self._counter_ttl: Final = counter_ttl + + def _write_ttl(self) -> float | None: + return None if self._counter_ttl is None else self._counter_ttl() def attach_redis(self, redis: RedisStores) -> None: """Adds the Redis tier once, keeping the memory copies, as ``DualCache.attach_redis_cache`` does.""" @@ -74,7 +80,8 @@ class SpendCounterStore: @property def redis_ttl(self) -> float | None: - """The TTL every Redis write gives a counter, None when there is no Redis.""" + """The Redis client's default TTL, None when there is no Redis. A write whose store has ``counter_ttl`` lasts + longer than this.""" return self._handle.redis_ttl_seconds @property @@ -97,11 +104,11 @@ class SpendCounterStore: @with_service_target(SPEND_COUNTERS_TARGET) async def increment(self, counter: str, amount: float) -> float: """Raises when the increment may not have landed, after dropping the counter everywhere.""" - return await self._handle.increment(counter, amount) + return await self._handle.increment(counter, amount, self._write_ttl()) @with_service_target(SPEND_COUNTERS_TARGET) async def increment_many(self, increments: Sequence[SpendIncrement]) -> tuple[float, ...]: - return tuple(await self._handle.increment_many(increments)) + return tuple(await self._handle.increment_many(increments, self._write_ttl())) @with_service_target(SPEND_COUNTERS_TARGET) def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> bool: @@ -110,22 +117,25 @@ class SpendCounterStore: request: Final = post_call_request() if self.has_redis and increments else None if request is None: return False - return self._handle.increment_many_post_call(request, increments) + return self._handle.increment_many_post_call(request, increments, self._write_ttl()) @with_service_target(SPEND_COUNTERS_TARGET) async def set_max(self, counter: str, spend: float) -> float | None: - """Raise memory and Redis to ``spend``, never lowering either; what Redis holds after, None without Redis.""" - return await self._handle.set_max(counter, spend) + """Raise memory and Redis to ``spend``, never lowering either; what Redis holds after, None without Redis. A + spend reset's fence (``RedisCache.reset_fence_key``) holds the Redis write while it exists.""" + return await self._handle.set_max(counter, spend, self._write_ttl()) @with_service_target(SPEND_COUNTERS_TARGET) async def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> float: """Seed a cold counter with SET NX so one worker wins, answering with what the counter then holds.""" - return await self._handle.seed(counter, spend, increment_if_lost=increment_if_lost) + return await self._handle.seed( + counter, spend, increment_if_lost=increment_if_lost, ttl_seconds=self._write_ttl() + ) @with_service_target(SPEND_COUNTERS_TARGET) - async def set(self, counter: str, spend: float, ttl: float | None = None) -> None: - """Memory, then Redis; raises when Redis failed, with memory already written.""" - await self._handle.set(counter, spend, ttl) + async def set(self, counter: str, spend: float, memory_ttl: float | None = None) -> None: + """Memory for ``memory_ttl``, then Redis; raises when Redis failed, with memory already written.""" + await self._handle.set(counter, spend, memory_ttl, self._write_ttl()) @with_service_target(SPEND_COUNTERS_TARGET) async def delete(self, counters: Sequence[str]) -> None: @@ -134,7 +144,7 @@ class SpendCounterStore: @with_service_target(SPEND_COUNTERS_TARGET) async def refresh_ttl(self, counter: str) -> bool: - return await self._handle.refresh_ttl(counter) + return await self._handle.refresh_ttl(counter, self._write_ttl()) def get_local(self, counter: str) -> float | None: return self._handle.get_local(counter) diff --git a/litellm/redis/stores/spend_counters/handle.py b/litellm/redis/stores/spend_counters/handle.py index e4bd789c341..79eb550bcf8 100644 --- a/litellm/redis/stores/spend_counters/handle.py +++ b/litellm/redis/stores/spend_counters/handle.py @@ -18,14 +18,26 @@ class SpendCountersHandle(Protocol): def read_many( self, counters: Sequence[str], send: NativeSend ) -> NativePending[Sequence[float | None] | None] | None: ... - def increment(self, counter: str, amount: float) -> Awaitable[float]: ... - def increment_many(self, increments: Sequence[SpendIncrement]) -> Awaitable[Sequence[float]]: ... - def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[SpendIncrement]) -> bool: ... - def set_max(self, counter: str, spend: float) -> Awaitable[float | None]: ... - def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Awaitable[float]: ... - def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Awaitable[None]: ... + def increment(self, counter: str, amount: float, ttl_seconds: float | None = None) -> Awaitable[float]: ... + def increment_many( + self, increments: Sequence[SpendIncrement], ttl_seconds: float | None = None + ) -> Awaitable[Sequence[float]]: ... + def increment_many_post_call( + self, request: NativeRequestRedis, increments: Sequence[SpendIncrement], ttl_seconds: float | None = None + ) -> bool: ... + def set_max(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Awaitable[float | None]: ... + def seed( + self, counter: str, spend: float, *, increment_if_lost: bool, ttl_seconds: float | None = None + ) -> Awaitable[float]: ... + def set( + self, + counter: str, + spend: float, + memory_ttl_seconds: float | None = None, + redis_ttl_seconds: float | None = None, + ) -> Awaitable[None]: ... def delete(self, counters: Sequence[str]) -> Awaitable[None]: ... - def refresh_ttl(self, counter: str) -> Awaitable[bool]: ... + def refresh_ttl(self, counter: str, ttl_seconds: float | None = None) -> Awaitable[bool]: ... def get_local(self, counter: str) -> float | None: ... def set_local(self, counter: str, spend: float, ttl_seconds: float | None = None) -> None: ... def delete_local(self, counter: str) -> None: ... diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 773985ad427..3155f240d0f 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -117,9 +117,7 @@ class NativeRouterUsage: def attach(self, client: NativeRedisClient) -> bool: ... @property def has_redis(self) -> bool: ... - def increment( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> Future[float]: ... + def increment(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> Future[float]: ... def increment_deployment( self, deployment_id: str, @@ -128,12 +126,8 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> Future[None]: ... - def get_many( - self, counters: Sequence[_UsageCounter] - ) -> Future[list[float | None]]: ... - def increment_sync( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> float: ... + def get_many(self, counters: Sequence[_UsageCounter]) -> Future[list[float | None]]: ... + def increment_sync(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> float: ... def increment_deployment_sync( self, deployment_id: str, @@ -142,9 +136,7 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> None: ... - def get_many_sync( - self, counters: Sequence[_UsageCounter] - ) -> list[float | None]: ... + def get_many_sync(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... @@ -320,7 +312,9 @@ class NativeAuthCache: def get_sync(self, key: str) -> object: ... def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Future[None]: ... def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... - def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: NativeSend) -> NativePending[None]: ... + def declare_set( + self, key: str, value: object, ttl_seconds: float | None, send: NativeSend + ) -> NativePending[None]: ... def declare_delete(self, key: str, send: NativeSend) -> NativePending[None]: ... def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ... def write_back(self, entries: Sequence[tuple[str, object, float | None]], send: NativeSend) -> Future[None]: ... @@ -354,16 +348,26 @@ class NativeSpendCounters: def read_many( self, counters: Sequence[str], send: NativeSend ) -> NativePending[list[float | None] | None] | None: ... - def increment(self, counter: str, amount: float) -> Future[float]: ... - def increment_many(self, increments: Sequence[tuple[str, float]]) -> Future[list[float]]: ... + def increment(self, counter: str, amount: float, ttl_seconds: float | None = None) -> Future[float]: ... + def increment_many( + self, increments: Sequence[tuple[str, float]], ttl_seconds: float | None = None + ) -> Future[list[float]]: ... def increment_many_post_call( - self, request: NativeRequestRedis, increments: Sequence[tuple[str, float]] + self, request: NativeRequestRedis, increments: Sequence[tuple[str, float]], ttl_seconds: float | None = None ) -> bool: ... - def set_max(self, counter: str, spend: float) -> Future[float | None]: ... - def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Future[float]: ... - def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Future[None]: ... + def set_max(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Future[float | None]: ... + def seed( + self, counter: str, spend: float, *, increment_if_lost: bool, ttl_seconds: float | None = None + ) -> Future[float]: ... + def set( + self, + counter: str, + spend: float, + memory_ttl_seconds: float | None = None, + redis_ttl_seconds: float | None = None, + ) -> Future[None]: ... def delete(self, counters: Sequence[str]) -> Future[None]: ... - def refresh_ttl(self, counter: str) -> Future[bool]: ... + def refresh_ttl(self, counter: str, ttl_seconds: float | None = None) -> Future[bool]: ... def get_local(self, counter: str) -> float | None: ... def set_local(self, counter: str, spend: float, ttl_seconds: float | None = None) -> None: ... def delete_local(self, counter: str) -> None: ... @@ -1056,7 +1060,9 @@ class RouterCoreHandle: callbacks: Mapping[str, bool] | None = None, ) -> RouterCall[_SuccessReportSteps]: ... # Count a success against any id, driven like `report_failure`; `done` carries `key`. - def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ... + def count_success( + self, id: str, callbacks: Mapping[str, bool] | None = None + ) -> RouterCall[_SuccessReportSteps]: ... # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment. def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ... # Every deployment as a `model_list` entry, in config order. diff --git a/tests/integration/_support/spend_counters.py b/tests/integration/_support/spend_counters.py new file mode 100644 index 00000000000..a8c8b738df5 --- /dev/null +++ b/tests/integration/_support/spend_counters.py @@ -0,0 +1,33 @@ +import os +from collections.abc import Iterator +from contextlib import contextmanager +from hashlib import sha256 +from typing import Final + +from redis import Redis + +SPEND_COUNTER_TTL_SECONDS: Final = 3600 +TTL_SLACK_SECONDS: Final = 600 + + +@contextmanager +def rig_cache() -> Iterator[Redis]: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + yield cache + + +def key_counter(key: str) -> str: + return f"spend:key:{sha256(key.encode()).hexdigest()}" + + +def counter_ttl(cache: Redis, counter_key: str) -> int: + return int(str(cache.ttl(counter_key))) + + +def counter_value(cache: Redis, counter_key: str) -> float | None: + raw: Final = cache.get(counter_key) + return None if raw is None else float(str(raw, "utf-8") if isinstance(raw, bytes) else str(raw)) + + +def outlives(ttl: int, expected: int = SPEND_COUNTER_TTL_SECONDS) -> bool: + return expected - TTL_SLACK_SECONDS < ttl <= expected diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 85da60a37ad..16df7c44264 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -6,6 +6,7 @@ import base64 import json import os import struct +import time import uuid import zlib from collections import deque @@ -60,6 +61,46 @@ INTERNAL_FIELDS: Final = frozenset( ) +RESPONSES_USAGE: Final = {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40} + + +def responses_body(model: str) -> dict[str, JsonValue]: + return { + "id": f"resp_{uuid.uuid4().hex[:24]}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": model, + "output": [ + { + "id": f"msg_{uuid.uuid4().hex[:24]}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Hello.", "annotations": []}], + } + ], + "usage": dict(RESPONSES_USAGE), + } + + +async def responses_stream(completed: Mapping[str, JsonValue]) -> AsyncIterator[str]: + item: Final = cast(dict[str, JsonValue], cast(list[JsonValue], completed["output"])[0]) + events: Final = ( + ("response.created", {"response": {**completed, "status": "in_progress", "output": [], "usage": None}}), + ("response.output_item.added", {"output_index": 0, "item": {**item, "status": "in_progress", "content": []}}), + ( + "response.output_text.delta", + {"output_index": 0, "content_index": 0, "item_id": item["id"], "delta": "Hello."}, + ), + ("response.output_item.done", {"output_index": 0, "item": item}), + ("response.completed", {"response": dict(completed)}), + ) + for sequence_number, (event_type, payload) in enumerate(events): + data: Final = json.dumps({"type": event_type, "sequence_number": sequence_number, **payload}) + yield f"event: {event_type}\ndata: {data}\n\n" + + def error_type(status: int) -> str: if status == 429: return "rate_limit_error" @@ -222,24 +263,41 @@ class Provider: for message in messages ): return JSONResponse({"error": {"message": "Invalid selected message contract"}}, status_code=400) - script: Final = self.scripts.get(str(body["model"])) - if script is not None: - if not script: - return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500) - status: Final = script.popleft() - if status != 200: - return JSONResponse( - { - "error": { - "message": "Controlled provider failure", - "type": error_type(status), - "code": str(status), - } - }, - status_code=status, - ) + failure: Final = self._scripted_failure(str(body["model"])) + if failure is not None: + return failure return await chat_completions(request) + def _scripted_failure(self, model: str) -> Response | None: + script: Final = self.scripts.get(model) + if script is None: + return None + if not script: + return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500) + status: Final = script.popleft() + if status == 200: + return None + return JSONResponse( + {"error": {"message": "Controlled provider failure", "type": error_type(status), "code": str(status)}}, + status_code=status, + ) + + async def responses(self, request: Request) -> Response: + body: Final = JSON_OBJECT.validate_json(await request.body()) + self.observations.put( + Observation(request.url.path, request.headers.get("authorization", ""), body, request.method) + ) + model: Final = body.get("model") + if not isinstance(model, str) or "input" not in body: + return JSONResponse({"error": {"message": "model and input are required"}}, status_code=400) + failure: Final = self._scripted_failure(model) + if failure is not None: + return failure + completed: Final = responses_body(model) + if body.get("stream") is True: + return StreamingResponse(responses_stream(completed), media_type="text/event-stream") + return JSONResponse(completed) + async def vector_store_search(self, request: Request) -> Response: body: Final = JSON_OBJECT.validate_json(await request.body()) self.observations.put( @@ -645,6 +703,7 @@ class Provider: Route("/_cost_map", self.cost_map, methods=["GET"]), Route("/_oauth/token", self.oauth_token, methods=["POST"]), Route("/v1/chat/completions", self.chat, methods=["POST"]), + Route("/v1/responses", self.responses, methods=["POST"]), Route("/v1/completions", completions, methods=["POST"]), Route("/v1/embeddings", embeddings, methods=["POST"]), Route("/v1/moderations", moderations, methods=["POST"]), diff --git a/tests/integration/caching/test_batch_redis_requests_prefetch.py b/tests/integration/caching/test_batch_redis_requests_prefetch.py new file mode 100644 index 00000000000..ae61eb354e2 --- /dev/null +++ b/tests/integration/caching/test_batch_redis_requests_prefetch.py @@ -0,0 +1,88 @@ +"""The ``batch_redis_requests`` callback prefetches a key's response cache entries +into memory with one Redis scan, so the next identical request is served without a +Redis round trip for the entry.""" + +import os +import threading +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import OwnedRedis, owned_redis +from redis import Redis +from redis.exceptions import TimeoutError as RedisTimeoutError + +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) + + +def _config_with_batch_redis_requests(directory: Path) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + litellm_settings: Final = {**object_value(shared["litellm_settings"]), "callbacks": ["batch_redis_requests"]} + path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump({**shared, "litellm_settings": litellm_settings}, sort_keys=False)) + return path + + +@contextmanager +def _monitored(redis: OwnedRedis) -> Iterator[list[str]]: + commands: Final[list[str]] = [] + stopped: Final = threading.Event() + with Redis(host=redis.host, port=redis.port, socket_timeout=0.3) as client, client.monitor() as monitor: + + def pump() -> None: + while not stopped.is_set(): + try: + commands.append(str(monitor.next_command()["command"])) + except RedisTimeoutError: + continue + + thread: Final = threading.Thread(target=pump, daemon=True) + thread.start() + try: + yield commands + finally: + stopped.set() + thread.join(timeout=5) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_the_prefetched_response_cache_entry_is_served_from_memory_on_the_next_request( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_redis(tmp_path) as redis: + overrides: Final = {"REDIS_HOST": redis.host, "REDIS_PORT": str(redis.port)} + with ( + owned_proxy_process( + gateway, tmp_path, overrides, config=_config_with_batch_redis_requests(tmp_path) + ) as owned, + owned.gateway.scenario() as scenario, + Redis(host=redis.host, port=redis.port) as cache, + httpx.Client(base_url=gateway.upstream_url, timeout=30) as upstream, + ): + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + body: Final = {"model": model, "messages": [{"role": "user", "content": f"prefetch {uuid.uuid4().hex}"}]} + namespace: Final = f"litellm:{sha256(key.encode()).hexdigest()}:acompletion" + first: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key) + assert first.status_code == 200, first.text + entries: Final = eventually(lambda: cache.keys(f"{namespace}:*"), lambda keys: len(keys) == 1, seconds=10) + upstream.delete("/__observations").raise_for_status() + with _monitored(redis) as commands: + second: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key) + assert second.status_code == 200, second.text + assert second.json()["id"] == first.json()["id"], (first.text, second.text) + assert upstream.get("/__observations").json()["requests"] == [] + fetched: Final = [command for command in commands if command.startswith(f"GET {namespace}:")] + assert fetched == [], (entries, commands) + third: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key) + assert third.status_code == 200 and third.json()["id"] == first.json()["id"], third.text + assert upstream.get("/__observations").json()["requests"] == [] diff --git a/tests/integration/spend/test_spend_counter_redis_ttl.py b/tests/integration/spend/test_spend_counter_redis_ttl.py new file mode 100644 index 00000000000..9e763ec366f --- /dev/null +++ b/tests/integration/spend/test_spend_counter_redis_ttl.py @@ -0,0 +1,488 @@ +import asyncio +import json +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from hashlib import sha256 +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.spend_counters import ( + counter_ttl, + counter_value, + key_counter, + outlives, + rig_cache, +) +from pydantic import JsonValue +from redis import Redis + +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ONE_CALL_COST: Final = 0.06 +ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"} + + +@dataclass(frozen=True, slots=True) +class Entities: + model: str + key: str + team: str + user: str + organization: str + end_user: str + tag: str + + def counters(self) -> Mapping[str, str]: + return { + "key": key_counter(self.key), + "team": f"spend:team:{self.team}", + "team_member": f"spend:team_member:{self.user}:{self.team}", + "user": f"spend:user:{self.user}", + "end_user": f"spend:end_user:{self.end_user}", + "tag": f"spend:tag:{self.tag}", + "org": f"spend:org:{self.organization}", + } + + +def _entities(gateway: Gateway, scenario: Scenario) -> Entities: + model: Final = scenario.model(**RATES) + models: Final = [model] + organization: Final = scenario.organization(max_budget=50.0) + team: Final = scenario.team(models=models, max_budget=20.0, organization_id=organization) + user: Final = scenario.user(max_budget=20.0) + added: Final = gateway.request( + "POST", + "/team/member_add", + {"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": 20.0}, + ) + assert added.status_code == 200, added.text + tag: Final = f"integration-tag-{uuid.uuid4().hex}" + key: Final = scenario.key(team_id=team, user_id=user, models=models, max_budget=5.0, tags=[tag]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + return Entities(model, key, team, user, organization, end_user, tag) + + +def _base(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _chat_openai_sync(gateway: Gateway, entities: Entities, text: str) -> str: + with openai.OpenAI(base_url=f"{_base(gateway)}/v1", api_key=entities.key, max_retries=0, timeout=60) as client: + response: Final = client.chat.completions.create( + model=entities.model, messages=[{"role": "user", "content": text}], user=entities.end_user + ) + assert response.usage is not None and response.usage.total_tokens == 40, response + return response.id + + +def _chat_openai_async_stream(gateway: Gateway, entities: Entities, text: str) -> str: + async def run() -> str: + async with openai.AsyncOpenAI( + base_url=f"{_base(gateway)}/v1", api_key=entities.key, max_retries=0, timeout=60 + ) as client: + stream: Final = await client.chat.completions.create( + model=entities.model, + messages=[{"role": "user", "content": text}], + user=entities.end_user, + stream=True, + stream_options={"include_usage": True}, + ) + ids: Final = [chunk.id async for chunk in stream] + assert len(set(ids)) == 1, ids + return ids[0] + + return asyncio.run(run()) + + +def _chat_httpx_stream(gateway: Gateway, entities: Entities, text: str) -> str: + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": entities.model, + "messages": [{"role": "user", "content": text}], + "user": entities.end_user, + "stream": True, + "stream_options": {"include_usage": True}, + }, + headers={"Authorization": f"Bearer {entities.key}"}, + ) as response: + assert response.status_code == 200, response.read() + frames: Final = [line.removeprefix("data: ") for line in response.iter_lines() if line.startswith("data: ")] + ids: Final = {json.loads(frame)["id"] for frame in frames if frame != "[DONE]"} + assert len(ids) == 1, frames + return ids.pop() + + +def _messages_anthropic_sync(gateway: Gateway, entities: Entities, text: str) -> str: + with anthropic.Anthropic(base_url=_base(gateway), api_key=entities.key, max_retries=0, timeout=60) as client: + message: Final = client.messages.create( + model=entities.model, + max_tokens=32, + messages=[{"role": "user", "content": text}], + metadata={"user_id": entities.end_user}, + ) + assert message.usage.output_tokens > 0, message + return message.id + + +def _messages_anthropic_async_stream(gateway: Gateway, entities: Entities, text: str) -> str: + async def run() -> str: + async with anthropic.AsyncAnthropic( + base_url=_base(gateway), api_key=entities.key, max_retries=0, timeout=60 + ) as client: + async with client.messages.stream( + model=entities.model, + max_tokens=32, + messages=[{"role": "user", "content": text}], + metadata={"user_id": entities.end_user}, + ) as stream: + async for _ in stream.text_stream: + pass + final: Final = await stream.get_final_message() + assert final.usage.output_tokens > 0, final + return final.id + + return asyncio.run(run()) + + +def _messages_httpx(gateway: Gateway, entities: Entities, text: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": entities.model, + "max_tokens": 32, + "messages": [{"role": "user", "content": text}], + "metadata": {"user_id": entities.end_user}, + }, + key=entities.key, + headers=ANTHROPIC_HEADERS, + ) + assert response.status_code == 200, response.text + return str(response.json()["id"]) + + +def _responses_httpx(gateway: Gateway, entities: Entities, text: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": entities.model, "input": text, "user": entities.end_user}, + key=entities.key, + ) + assert response.status_code == 200, response.text + return str(response.json()["id"]) + + +def _responses_httpx_stream(gateway: Gateway, entities: Entities, text: str) -> str: + with gateway.client.stream( + "POST", + "/v1/responses", + json={"model": entities.model, "input": text, "user": entities.end_user, "stream": True}, + headers={"Authorization": f"Bearer {entities.key}"}, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = [line for line in response.iter_lines() if line.startswith("data: ")] + events: Final = [json.loads(line.removeprefix("data: ")) for line in lines if line != "data: [DONE]"] + completed: Final = [event for event in events if event.get("type") == "response.completed"] + assert len(completed) == 1, [event.get("type") for event in events] + return str(completed[0]["response"]["id"]) + + +@dataclass(frozen=True, slots=True) +class Caller: + call: Callable[[Gateway, Entities, str], str] + upstream_path: str + exact_cost: float | None + + +CALLERS: Final[Mapping[str, Caller]] = { + "chat-openai-sync": Caller(_chat_openai_sync, "/v1/chat/completions", ONE_CALL_COST), + "chat-openai-async-stream": Caller(_chat_openai_async_stream, "/v1/chat/completions", ONE_CALL_COST), + "chat-httpx-stream": Caller(_chat_httpx_stream, "/v1/chat/completions", ONE_CALL_COST), + "messages-anthropic-sync": Caller(_messages_anthropic_sync, "/v1/responses", ONE_CALL_COST), + "messages-anthropic-async-stream": Caller(_messages_anthropic_async_stream, "/v1/responses", ONE_CALL_COST), + "messages-httpx": Caller(_messages_httpx, "/v1/responses", ONE_CALL_COST), + "responses-httpx": Caller(_responses_httpx, "/v1/responses", ONE_CALL_COST), + "responses-httpx-stream": Caller(_responses_httpx_stream, "/v1/responses", ONE_CALL_COST), +} + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', + (sha256(key.encode()).hexdigest(),), + ) + + +def _one_spend_row(key: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 1, seconds=70) + return rows[0] + + +def _observed(cache: Redis, counters: Mapping[str, str]) -> Mapping[str, tuple[float | None, int]]: + return {name: (counter_value(cache, counter), counter_ttl(cache, counter)) for name, counter in counters.items()} + + +def _assert_every_counter_outlives( + cache: Redis, counters: Mapping[str, str], spend: float +) -> Mapping[str, tuple[float | None, int]]: + observed: Final = eventually( + lambda: _observed(cache, counters), + lambda states: all( + value is not None and value == pytest.approx(spend, abs=1e-9) for value, _ in states.values() + ), + seconds=20, + ) + for name, (_, ttl) in observed.items(): + assert outlives(ttl), (name, observed) + return observed + + +@pytest.mark.parametrize("caller", sorted(CALLERS)) +def test_every_counter_scope_outlives_the_batch_write(gateway: Gateway, caller: str) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + rig_cache() as cache, + ): + entities: Final = _entities(gateway, scenario) + text: Final = f"counter ttl {caller} {uuid.uuid4().hex}" + upstream.delete("/__observations").raise_for_status() + response_id: Final = CALLERS[caller].call(gateway, entities, text) + received: Final = upstream.get("/__observations").json()["requests"] + paths: Final = [request["path"] for request in received if text in json.dumps(request["body"])] + assert len(paths) == 1 and paths[0].endswith(CALLERS[caller].upstream_path), received + row: Final = _one_spend_row(entities.key) + spend: Final = float(str(row["spend"])) + assert spend == pytest.approx(CALLERS[caller].exact_cost or spend) and spend > 0, row + assert response_id, row + _assert_every_counter_outlives(cache, entities.counters(), spend) + + +def test_a_response_cache_hit_keeps_the_counter_and_its_ttl(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + rig_cache() as cache, + ): + entities: Final = _entities(gateway, scenario) + body: Final = { + "model": entities.model, + "messages": [{"role": "user", "content": f"cache twin {uuid.uuid4().hex}"}], + "user": entities.end_user, + } + first: Final = gateway.request("POST", "/v1/chat/completions", body, key=entities.key) + assert first.status_code == 200, first.text + row: Final = _one_spend_row(entities.key) + spend: Final = float(str(row["spend"])) + before: Final = _assert_every_counter_outlives(cache, entities.counters(), spend) + upstream.delete("/__observations").raise_for_status() + second: Final = gateway.request("POST", "/v1/chat/completions", body, key=entities.key) + assert second.status_code == 200, second.text + assert second.json()["id"] == first.json()["id"], (first.text, second.text) + assert upstream.get("/__observations").json()["requests"] == [] + rows: Final = eventually(lambda: _spend_rows(entities.key), lambda values: len(values) == 2, seconds=70) + assert str(rows[1]["cache_hit"]) == "True" and float(str(rows[1]["spend"])) == 0.0, rows + after: Final = _assert_every_counter_outlives(cache, entities.counters(), spend) + assert {name: value for name, (value, _) in after.items()} == pytest.approx( + {name: value for name, (value, _) in before.items()}, abs=1e-9 + ), (before, after) + + +def _exhausted_key(gateway: Gateway, scenario: Scenario) -> tuple[str, str]: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + served: Final = gateway.chat(model, key=key, text=f"exhaust {uuid.uuid4().hex}") + assert served["usage"]["total_tokens"] == 40, served + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + return model, key + + +def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _assert_denied(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert response.json()["error"]["type"] == "budget_exceeded", response.text + + +def test_a_stale_low_counter_is_repaired_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key( + models=[model], budget_limits=[{"max_budget": ONE_CALL_COST, "budget_duration": "1d"}] + ) + counter: Final = f"{key_counter(key)}:window:1d" + assert gateway.chat(model, key=key, text=f"exhaust window {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + assert counter_value(cache, counter) == pytest.approx(ONE_CALL_COST) + + served: Final[list[int]] = [] + + def stale_then_request() -> int: + assert cache.set(counter, "0.01") + code: Final = _chat(gateway, model, key).status_code + served.append(code) + return code + + assert eventually(stale_then_request, lambda code: code == 422, seconds=30) == 422 + assert set(served) == {200, 422}, served + repaired: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + floor: Final = repaired[0] + assert floor is not None and ONE_CALL_COST - 1e-9 <= floor <= served.count(200) * ONE_CALL_COST + 1e-9, ( + repaired, + served, + ) + assert outlives(repaired[1]), repaired + + +def test_a_lost_counter_is_reseeded_from_the_db_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key = _exhausted_key(gateway, scenario) + counter: Final = key_counter(key) + assert cache.delete(counter) == 1 + _assert_denied(_chat(gateway, model, key)) + reseeded: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and state[0] >= ONE_CALL_COST and outlives(state[1]), + seconds=15, + ) + assert reseeded[0] == pytest.approx(ONE_CALL_COST), reseeded + + +def test_key_update_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert gateway.chat(model, key=key, text=f"before update {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + updated: Final = gateway.request("POST", "/key/update", {"key": key, "spend": 1.25}) + assert updated.status_code == 200, updated.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (1.25, True), updated.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value == pytest.approx(1.25 + ONE_CALL_COST), + seconds=15, + ) + assert outlives(counter_ttl(cache, counter)) + + +def test_key_reset_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key = _exhausted_key(gateway, scenario) + counter: Final = key_counter(key) + _assert_denied(_chat(gateway, model, key)) + reset: Final = gateway.request("POST", f"/key/{key}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert reset.json()["key_hash"] == sha256(key.encode()).hexdigest(), reset.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (0.0, True), reset.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + + +def test_team_member_reset_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + team: Final = scenario.team(models=[model]) + user: Final = scenario.user() + added: Final = gateway.request( + "POST", + "/team/member_add", + {"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": ONE_CALL_COST}, + ) + assert added.status_code == 200, added.text + key: Final = scenario.key(team_id=team, user_id=user, models=[model]) + counter: Final = f"spend:team_member:{user}:{team}" + assert gateway.chat(model, key=key, text=f"member {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s', (team, user) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + _assert_denied(_chat(gateway, model, key)) + reset: Final = gateway.request("POST", f"/team/{team}/member/{user}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert reset.json()["previous_spend"] == pytest.approx(ONE_CALL_COST), reset.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (0.0, True), reset.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + + +def test_malformed_spend_writes_are_rejected_and_leave_the_counter_alone(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + bystander: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert gateway.chat(model, key=key, text=f"before bad writes {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + before: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert before[0] == pytest.approx(ONE_CALL_COST) and outlives(before[1]), before + rejected: Final = tuple( + gateway.request(*attempt) + for attempt in ( + ("POST", "/key/update", {"key": key, "spend": "a" * 5000}), + ("POST", "/key/update", {"key": key, "spend": [1.0]}), + ("POST", "/key/update", {"key": key, "spend": {"spend": 1.0}}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": "abc"}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": True}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": [0.0]}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": ONE_CALL_COST * 10}), + ) + ) + for response in rejected: + assert response.status_code in (400, 422), response.text + assert "error" in response.json() or "detail" in response.json(), response.text + assert (counter_value(cache, counter), counter_ttl(cache, counter) <= before[1]) == (before[0], True) + assert outlives(counter_ttl(cache, counter)) + served: Final = _chat(gateway, model, bystander) + assert served.status_code == 200, served.text + + +def _landed_chat(gateway: Gateway, model: str, key: str, index: int) -> str: + served: Final = gateway.chat(model, key=key, text=f"repeat {index} {uuid.uuid4().hex}") + eventually(lambda: _spend_rows(key), lambda rows: len(rows) == index + 1, seconds=70) + return str(served["id"]) + + +def test_repeated_requests_log_once_each_and_the_counter_matches_their_sum(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + ids: Final = tuple(_landed_chat(gateway, model, key, index) for index in range(3)) + assert len(set(ids)) == 3, ids + rows: Final = _spend_rows(key) + assert len(rows) == 3, rows + assert sorted(str(row["request_id"]) for row in rows) == sorted(ids), rows + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx(3 * ONE_CALL_COST), rows + eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value == pytest.approx(total), + seconds=15, + ) + assert outlives(counter_ttl(cache, counter)) diff --git a/tests/integration/spend/test_spend_counter_reset_fence.py b/tests/integration/spend/test_spend_counter_reset_fence.py new file mode 100644 index 00000000000..003737f37d9 --- /dev/null +++ b/tests/integration/spend/test_spend_counter_reset_fence.py @@ -0,0 +1,46 @@ +"""The reset fence: a spend reset SETs it in the same pipeline as the zeroed counter, and the stale-floor +repair's Lua script leaves a fenced counter alone, so a pod still holding the pre-reset floor cannot raise +the counter the reset zeroed for the whole counter TTL.""" + +import os +import uuid +from typing import Final + +from integration._support.spend_counters import rig_cache + +from litellm.caching.redis_cache import RedisCache +from litellm.redis.stores.connection import connect +from litellm.redis.stores.spend_counters import SpendCounterStore + + +async def test_set_max_holds_while_the_reset_fence_exists() -> None: + cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), namespace="fence") + counter: Final = f"spend:key:{uuid.uuid4().hex}" + fence: Final = cache.reset_fence_key(counter) + assert await cache.async_set_cache_pipeline_with_ttls(((fence, 1, 10), (counter, 0.0, 3600))) is None + + assert await cache.async_set_max(key=counter, value=12.0, ttl=3600, fence_key=fence) == 0.0 + + with rig_cache() as raw: + assert raw.delete(cache.check_and_fix_namespace(key=fence)) == 1 + assert await cache.async_set_max(key=counter, value=12.0, ttl=3600, fence_key=fence) == 12.0 + with rig_cache() as raw: + raw.delete(cache.check_and_fix_namespace(key=counter)) + + +async def test_the_native_repair_holds_while_a_fence_the_python_reset_wrote_exists() -> None: + """The reset writes through ``RedisCache`` and the repair through the native store, so they must agree on the key.""" + cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), namespace="fence") + store: Final = SpendCounterStore(connect(cache).spend_counters(), lambda: 3600) + counter: Final = f"spend:key:{uuid.uuid4().hex}" + fence: Final = cache.reset_fence_key(counter) + assert await cache.async_set_cache_pipeline_with_ttls(((fence, 1, 10), (counter, 0.0, 3600))) is None + + assert await store.set_max(counter, 12.0) == 0.0 + + with rig_cache() as raw: + assert raw.delete(cache.check_and_fix_namespace(key=fence)) == 1 + assert await store.set_max(counter, 12.0) == 12.0 + with rig_cache() as raw: + assert raw.ttl(cache.check_and_fix_namespace(key=counter)) > 60 + raw.delete(cache.check_and_fix_namespace(key=counter)) diff --git a/tests/integration/spend/test_spend_counter_reset_paths.py b/tests/integration/spend/test_spend_counter_reset_paths.py new file mode 100644 index 00000000000..1f59771c0f5 --- /dev/null +++ b/tests/integration/spend/test_spend_counter_reset_paths.py @@ -0,0 +1,667 @@ +"""The two spend reset paths that touch the Redis spend counters. + +``POST /global/spend/reset`` zeroes every key and team spend in the database. The +counters that enforce those budgets must follow it on every pod, or a key the +operator just reset stays blocked until the counter's TTL runs out. ``PATCH +/organization/update`` with ``spend`` sets one organization's counter to the given +value. Both write a reset fence next to each zeroed counter so a pod still holding +the pre-reset floor cannot repair the counter back up, and both answer 500 when the +counters could not be written to Redis, so the operator retries once it is back. +Every cell drives a real proxy: the cells that reset every key or stop Redis boot +their own proxy on a scratch database and an owned Redis. +""" + +import asyncio +import os +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import OwnedRedis, owned_redis +from integration._support.spend_counters import ( + SPEND_COUNTER_TTL_SECONDS, + counter_ttl, + counter_value, + key_counter, + outlives, + rig_cache, +) +from pydantic import BaseModel, ConfigDict +from redis import Redis + +from litellm.caching.redis_cache import RedisCache + +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ONE_CALL_COST: Final = 0.06 +INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" +REDIS_AUTH_CACHE_SETTING: Final = frozenset({"enable_redis_auth_cache"}) +REDIS_SETTINGS: Final = frozenset({"cache", "cache_params", *REDIS_AUTH_CACHE_SETTING}) +SEEDED_COUNTERS: Final = 1200 +SETTLE_SECONDS: Final = 20 +BROADCAST_SECONDS: Final = 10 +RACE_WINDOW_SECONDS: Final = 2 +DATABASE_SECONDS: Final = 70 +FENCE_SECONDS: Final = 10 +RECOVERY_SECONDS: Final = 30 +LOST_COUNTERS: Final = "but the cached spend counters could not be reset in Redis" +RETRY_HINT: Final = "retry this request once Redis accepts writes." +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) + + +class _BroadcastMessage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + cache_key: str | None = None + cache_keys: tuple[str, ...] = () + new_value: float | None = None + ttl: float | None = None + + def names(self) -> tuple[str, ...]: + return (self.cache_key,) if self.cache_key is not None else self.cache_keys + + +@dataclass(frozen=True, slots=True) +class Broadcast: + key: str + value: float | None + ttl: float | None + + +class _Broadcasts: + def __init__(self, cache: Redis) -> None: + self._pubsub: Final = cache.pubsub() + self._pubsub.subscribe(INVALIDATION_CHANNEL) + confirmation: Final = self._pubsub.get_message(timeout=5) + assert confirmation is not None and confirmation["type"] == "subscribe", confirmation + self._condition: Final = threading.Condition() + self._received: Final[list[Broadcast]] = [] + self._stopped: Final = threading.Event() + self._thread: Final = threading.Thread(target=self._pump, daemon=True) + self._thread.start() + + def _pump(self) -> None: + while not self._stopped.is_set(): + self._record(self._pubsub.get_message(timeout=0.2)) + + def _record(self, message: Mapping[str, object] | None) -> None: + if message is None or message["type"] != "message": + return + data: Final = message["data"] + assert isinstance(data, bytes | str), message + parsed: Final = _BroadcastMessage.model_validate_json(data) + with self._condition: + self._received.extend(Broadcast(key, parsed.new_value, parsed.ttl) for key in parsed.names()) + self._condition.notify_all() + + def wait_for(self, satisfied: Callable[[Sequence[Broadcast]], bool], seconds: float) -> bool: + with self._condition: + return self._condition.wait_for(lambda: satisfied(tuple(self._received)), timeout=seconds) + + def received(self) -> tuple[Broadcast, ...]: + with self._condition: + return tuple(self._received) + + def close(self) -> None: + self._stopped.set() + self._thread.join(timeout=5) + self._pubsub.close() + + +@contextmanager +def _broadcasts(cache: Redis) -> Iterator[_Broadcasts]: + broadcasts: Final = _Broadcasts(cache) + try: + yield broadcasts + finally: + broadcasts.close() + + +def _zeroed(counter: str) -> Callable[[Sequence[Broadcast]], bool]: + return lambda received: any(broadcast.key == counter and broadcast.value == 0.0 for broadcast in received) + + +def _zeroed_with_its_floor(counter: str) -> Callable[[Sequence[Broadcast]], bool]: + return lambda received: ( + sum(1 for broadcast in received if broadcast.key.endswith(counter) and broadcast.value == 0.0) >= 2 + ) + + +@dataclass(frozen=True, slots=True) +class Isolated: + gateway: Gateway + cache: Redis + redis: OwnedRedis + database_url: str + log: Path + + +def _pod_environment(database_url: str, redis: OwnedRedis) -> Mapping[str, str]: + return { + "DATABASE_URL": database_url, + "REDIS_HOST": redis.host, + "REDIS_PORT": str(redis.port), + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2", + } + + +@contextmanager +def _isolated( + gateway: Gateway, + directory: Path, + *, + workers: int = 1, + settings: tuple[tuple[str, str], ...] = (), + config: Path | None = None, +) -> Iterator[Isolated]: + with scratch_database() as database_url, owned_redis(directory) as redis: + overrides: Final[Mapping[str, str]] = {**_pod_environment(database_url, redis), **dict(settings)} + with ( + owned_proxy_process( + gateway, + directory, + overrides, + workers=workers, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned, + Redis(host=redis.host, port=redis.port) as cache, + ): + yield Isolated(owned.gateway, cache, redis, database_url, owned.log) + + +@contextmanager +def _second_pod(gateway: Gateway, directory: Path, isolated: Isolated, config: Path) -> Iterator[Gateway]: + with owned_proxy_process( + gateway, + directory, + _pod_environment(isolated.database_url, isolated.redis), + config=config, + database_setup=(), + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned: + yield owned.gateway + + +def _config_without(directory: Path, dropped: frozenset[str]) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + litellm_settings: Final = { + name: value for name, value in object_value(shared["litellm_settings"]).items() if name not in dropped + } + path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump({**shared, "litellm_settings": litellm_settings}, sort_keys=False)) + return path + + +@contextmanager +def _isolated_without_redis(gateway: Gateway, directory: Path) -> Iterator[tuple[Gateway, str]]: + with ( + scratch_database() as database_url, + owned_proxy_process( + gateway, + directory, + {"DATABASE_URL": database_url}, + config=_config_without(directory, REDIS_SETTINGS), + remove_environment=("REDIS_HOST", "REDIS_PORT", "DATABASE_URL_READ_REPLICA"), + ) as owned, + ): + yield owned.gateway, database_url + + +def _one_spend(rows: Sequence[Mapping[str, object]]) -> float: + assert len(rows) == 1, rows + return float(str(rows[0]["spend"])) + + +def _key_spend(key: str, database_url: str | None = None) -> float: + return _one_spend( + read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', + (sha256(key.encode()).hexdigest(),), + database_url=database_url, + ) + ) + + +def _team_spend(team: str, database_url: str | None = None) -> float: + return _one_spend( + read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team,), database_url=database_url) + ) + + +def _user_spend(user: str, database_url: str | None = None) -> float: + return _one_spend( + read_rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=%s', (user,), database_url=database_url) + ) + + +def _organization_spend(organization: str, database_url: str | None = None) -> float: + return _one_spend( + read_rows( + 'SELECT spend FROM "LiteLLM_OrganizationTable" WHERE organization_id=%s', + (organization,), + database_url=database_url, + ) + ) + + +def _landed(read: Callable[[], float], least: float) -> float: + return eventually(read, lambda spend: spend >= least - 1e-9, seconds=DATABASE_SECONDS) + + +def _chat(client: httpx.Client, model: str, key: str, *, end_user: str | None = None) -> httpx.Response: + return client.post( + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": f"reset path {uuid.uuid4().hex}"}], + **({} if end_user is None else {"user": end_user}), + }, + headers={"Authorization": f"Bearer {key}"}, + ) + + +def _assert_served(response: httpx.Response) -> str: + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["usage"]["total_tokens"] == 40, body + return str(body["id"]) + + +def _assert_denied(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert response.json()["error"]["type"] == "budget_exceeded", response.text + + +def _assert_counters_unwritten(response: httpx.Response, saved: str) -> None: + assert response.status_code == 500, response.text + error: Final = response.json()["error"] + assert error["type"] == "internal_server_error", response.text + message: Final = str(error["message"]) + assert message.startswith(f"{saved}, {LOST_COUNTERS}"), message + assert message.endswith(RETRY_HINT), message + + +def _accepted_once_redis_is_back(request: Callable[[], httpx.Response]) -> None: + def attempt() -> int: + response: Final = request() + assert response.status_code in {200, 500}, response.text + return response.status_code + + eventually(attempt, lambda code: code == 200, seconds=RECOVERY_SECONDS) + + +def _fence(counter: str) -> str: + return f"{{{counter}}}:reset_fence" + + +async def _repair_as_a_pod_would(redis: OwnedRedis, counter: str, value: float) -> float | None: + cache: Final = RedisCache(host=redis.host, port=redis.port) + assert cache.reset_fence_key(counter) == _fence(counter) + return await cache.async_set_max( + key=counter, value=value, ttl=SPEND_COUNTER_TTL_SECONDS, fence_key=cache.reset_fence_key(counter) + ) + + +def _state(cache: Redis, counter: str) -> tuple[float | None, int]: + return counter_value(cache, counter), counter_ttl(cache, counter) + + +def _settled(cache: Redis, counter: str, spend: float) -> int: + _, ttl = eventually( + lambda: _state(cache, counter), + lambda state: state[0] is not None and state[0] == pytest.approx(spend, abs=1e-9), + seconds=SETTLE_SECONDS, + ) + return ttl + + +def _values(cache: Redis, counters: Sequence[str]) -> tuple[float | None, ...]: + return tuple(None if raw is None else float(raw) for raw in cache.mget(counters)) + + +def _ttls(cache: Redis, counters: Sequence[str]) -> tuple[int, ...]: + with cache.pipeline() as pipeline: + for counter in counters: + pipeline.ttl(counter) + return tuple(int(ttl) for ttl in pipeline.execute()) + + +def _exhaust(gateway: Gateway, cache: Redis, database_url: str, model: str, key: str) -> str: + counter: Final = key_counter(key) + _assert_served(_chat(gateway.client, model, key)) + _settled(cache, counter, ONE_CALL_COST) + _landed(lambda: _key_spend(key, database_url), ONE_CALL_COST) + _assert_denied(_chat(gateway.client, model, key)) + return counter + + +def _global_reset(gateway: Gateway) -> None: + reset: Final = gateway.request("POST", "/global/spend/reset") + assert reset.status_code == 200, reset.text + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_admits_a_blocked_key_on_the_next_request(gateway: Gateway, tmp_path: Path) -> None: + with ( + _isolated(gateway, tmp_path, workers=2) as isolated, + isolated.gateway.scenario() as scenario, + httpx.Client(base_url=str(isolated.gateway.client.base_url), timeout=60) as other_connection, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + _global_reset(isolated.gateway) + assert outlives(_settled(isolated.cache, counter, 0.0)) + assert _key_spend(key, isolated.database_url) == 0.0 + _assert_served(_chat(other_connection, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(other_connection, model, key)) + _landed(lambda: _key_spend(key, isolated.database_url), ONE_CALL_COST) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_admits_a_key_blocked_by_its_team_budget(gateway: Gateway, tmp_path: Path) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model: Final = scenario.model(**RATES) + team: Final = scenario.team(models=[model], max_budget=ONE_CALL_COST) + key: Final = scenario.key(team_id=team, models=[model]) + counter: Final = f"spend:team:{team}" + _assert_served(_chat(isolated.gateway.client, model, key)) + _settled(isolated.cache, counter, ONE_CALL_COST) + _landed(lambda: _team_spend(team, isolated.database_url), ONE_CALL_COST) + _assert_denied(_chat(isolated.gateway.client, model, key)) + _global_reset(isolated.gateway) + assert outlives(_settled(isolated.cache, counter, 0.0)) + assert _team_spend(team, isolated.database_url) == 0.0 + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_leaves_the_other_scopes_and_a_non_master_caller_alone( + gateway: Gateway, tmp_path: Path +) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model: Final = scenario.model(**RATES) + organization: Final = scenario.organization(max_budget=50.0) + team: Final = scenario.team(models=[model], max_budget=20.0, organization_id=organization) + user: Final = scenario.user(max_budget=20.0) + added: Final = isolated.gateway.request( + "POST", + "/team/member_add", + {"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": 20.0}, + ) + assert added.status_code == 200, added.text + tag: Final = f"integration-tag-{uuid.uuid4().hex}" + key: Final = scenario.key(team_id=team, user_id=user, models=[model], max_budget=5.0, tags=[tag]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + reset_counters: Final = (key_counter(key), f"spend:team:{team}") + kept_counters: Final = ( + f"spend:team_member:{user}:{team}", + f"spend:user:{user}", + f"spend:end_user:{end_user}", + f"spend:tag:{tag}", + f"spend:org:{organization}", + ) + _assert_served(_chat(isolated.gateway.client, model, key, end_user=end_user)) + for counter in (*reset_counters, *kept_counters): + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)), counter + _landed(lambda: _user_spend(user, isolated.database_url), ONE_CALL_COST) + _landed(lambda: _organization_spend(organization, isolated.database_url), ONE_CALL_COST) + + forbidden: Final = isolated.gateway.request("POST", "/global/spend/reset", key=key) + assert forbidden.status_code == 401, forbidden.text + assert forbidden.json()["error"]["type"] == "auth_error", forbidden.text + assert _values(isolated.cache, (*reset_counters, *kept_counters)) == pytest.approx( + (ONE_CALL_COST,) * 7, abs=1e-9 + ) + + _global_reset(isolated.gateway) + for counter in reset_counters: + assert outlives(_settled(isolated.cache, counter, 0.0)), counter + assert _values(isolated.cache, kept_counters) == pytest.approx((ONE_CALL_COST,) * 5, abs=1e-9) + assert all(outlives(ttl) for ttl in _ttls(isolated.cache, kept_counters)), _ttls(isolated.cache, kept_counters) + assert _key_spend(key, isolated.database_url) == 0.0 + assert _team_spend(team, isolated.database_url) == 0.0 + assert _user_spend(user, isolated.database_url) >= ONE_CALL_COST - 1e-9 + assert _organization_spend(organization, isolated.database_url) >= ONE_CALL_COST - 1e-9 + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_zeroes_every_key_counter_past_one_scan_page_and_broadcasts_them( + gateway: Gateway, tmp_path: Path +) -> None: + with _isolated(gateway, tmp_path) as isolated: + stem: Final = f"spend:key:audit-{uuid.uuid4().hex}" + team_named_like_a_window: Final = f"spend:team:audit-{uuid.uuid4().hex}:window:prod" + counters: Final = (*(f"{stem}-{index}" for index in range(SEEDED_COUNTERS)), team_named_like_a_window) + windows: Final = tuple(f"{stem}-w{index}:window:1d" for index in range(3)) + with isolated.cache.pipeline() as pipeline: + for name in (*counters, *windows): + pipeline.set(name, "1.0", ex=3600) + pipeline.execute() + with _broadcasts(isolated.cache) as broadcasts: + _global_reset(isolated.gateway) + eventually( + lambda: _values(isolated.cache, counters), + lambda values: all(value == 0.0 for value in values), + seconds=SETTLE_SECONDS, + ) + ttls: Final = _ttls(isolated.cache, counters) + assert all(outlives(ttl) for ttl in ttls), sorted(ttls)[:5] + assert _values(isolated.cache, windows) == (1.0, 1.0, 1.0) + assert broadcasts.wait_for( + lambda received: {broadcast.key for broadcast in received if broadcast.value == 0.0} >= set(counters), + BROADCAST_SECONDS, + ), len(broadcasts.received()) + assert not any(broadcast.key in windows for broadcast in broadcasts.received()) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_with_a_redis_batch_size_of_one_resets_and_fences_every_counter( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + _isolated(gateway, tmp_path, settings=(("DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", "1"),)) as isolated, + isolated.gateway.scenario() as scenario, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + seeded: Final = tuple(f"spend:key:audit-{uuid.uuid4().hex}" for _ in range(2)) + with isolated.cache.pipeline() as pipeline: + for name in seeded: + pipeline.set(name, "1.0", ex=3600) + pipeline.execute() + counters: Final = (counter, *seeded) + _global_reset(isolated.gateway) + assert isolated.cache.exists(*(_fence(name) for name in counters)) == len(counters) + assert _values(isolated.cache, counters) == (0.0, 0.0, 0.0) + assert all(outlives(ttl) for ttl in _ttls(isolated.cache, counters)), _ttls(isolated.cache, counters) + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_pod_that_cached_the_pre_reset_floor_admits_the_next_request_after_a_reset_on_another_pod( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + _isolated(gateway, tmp_path, config=_config_without(tmp_path, REDIS_AUTH_CACHE_SETTING)) as isolated, + _second_pod(gateway, tmp_path, isolated, _config_without(tmp_path, REDIS_AUTH_CACHE_SETTING)) as other_pod, + isolated.gateway.scenario() as scenario, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + with _broadcasts(isolated.cache) as broadcasts: + assert isolated.cache.set(counter, "0.0", ex=3600) + _assert_denied(_chat(other_pod.client, model, key)) + _settled(isolated.cache, counter, ONE_CALL_COST) + _global_reset(isolated.gateway) + broadcasts.wait_for(_zeroed_with_its_floor(counter), RACE_WINDOW_SECONDS) + assert other_pod.client.get("/health/liveliness").status_code == 200 + _assert_served(_chat(other_pod.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(other_pod.client, model, key)) + assert broadcasts.wait_for(_zeroed(counter), BROADCAST_SECONDS), broadcasts.received() + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_fences_the_counter_against_a_stale_floor_repair(gateway: Gateway, tmp_path: Path) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + _global_reset(isolated.gateway) + assert outlives(_settled(isolated.cache, counter, 0.0)) + assert isolated.cache.get(_fence(counter)) == b"1" + assert 0 < counter_ttl(isolated.cache, _fence(counter)) <= FENCE_SECONDS + assert asyncio.run(_repair_as_a_pod_would(isolated.redis, counter, ONE_CALL_COST)) == 0.0 + value, ttl = _state(isolated.cache, counter) + assert value == 0.0 and outlives(ttl), (value, ttl) + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) + eventually(lambda: isolated.cache.exists(_fence(counter)), lambda found: found == 0, seconds=FENCE_SECONDS + 5) + late_total: Final = 2 * ONE_CALL_COST + assert asyncio.run(_repair_as_a_pod_would(isolated.redis, counter, late_total)) == pytest.approx(late_total) + value_after, ttl_after = _state(isolated.cache, counter) + assert value_after == pytest.approx(late_total) and outlives(ttl_after), (value_after, ttl_after) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_while_redis_is_down_answers_500_and_the_retry_resets_the_counters( + gateway: Gateway, tmp_path: Path +) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + isolated.redis.stop() + _assert_counters_unwritten(isolated.gateway.request("POST", "/global/spend/reset"), "Spend was reset") + assert _key_spend(key, isolated.database_url) == 0.0 + _assert_served(_chat(isolated.gateway.client, model, key)) + isolated.redis.start() + assert isolated.cache.ping() + _landed(lambda: _key_spend(key, isolated.database_url), ONE_CALL_COST) + _accepted_once_redis_is_back(lambda: isolated.gateway.request("POST", "/global/spend/reset")) + assert _key_spend(key, isolated.database_url) == 0.0 + assert outlives(_settled(isolated.cache, counter, 0.0)) + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_without_redis_admits_the_next_request(gateway: Gateway, tmp_path: Path) -> None: + with _isolated_without_redis(gateway, tmp_path) as (owned, database_url), owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + _assert_served(_chat(owned.client, model, key)) + _landed(lambda: _key_spend(key, database_url), ONE_CALL_COST) + _assert_denied(_chat(owned.client, model, key)) + _global_reset(owned) + assert _key_spend(key, database_url) == 0.0 + _assert_served(_chat(owned.client, model, key)) + _landed(lambda: _key_spend(key, database_url), ONE_CALL_COST) + _assert_denied(_chat(owned.client, model, key)) + + +def _organization_entities( + gateway: Gateway, scenario: Scenario, database_url: str | None = None +) -> tuple[str, str, str]: + model: Final = scenario.model(**RATES) + organization: Final = scenario.organization(max_budget=ONE_CALL_COST) + team: Final = scenario.team(models=[model], organization_id=organization) + key: Final = scenario.key(team_id=team, models=[model]) + counter: Final = f"spend:org:{organization}" + _assert_served(_chat(gateway.client, model, key)) + _landed(lambda: _organization_spend(organization, database_url), ONE_CALL_COST) + _assert_denied(_chat(gateway.client, model, key)) + return model, key, counter + + +def _organization_update(gateway: Gateway, counter: str, **fields: object) -> httpx.Response: + return gateway.request( + "PATCH", "/organization/update", {"organization_id": counter.removeprefix("spend:org:"), **fields} + ) + + +def _update_organization(gateway: Gateway, counter: str, **fields: object) -> None: + updated: Final = _organization_update(gateway, counter, **fields) + assert updated.status_code == 200, updated.text + + +def test_an_organization_spend_set_to_zero_admits_a_key_blocked_by_the_organization_budget(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key, counter = _organization_entities(gateway, scenario) + _settled(cache, counter, ONE_CALL_COST) + with _broadcasts(cache) as broadcasts: + _update_organization(gateway, counter, spend=0.0) + assert outlives(_settled(cache, counter, 0.0)) + assert _organization_spend(counter.removeprefix("spend:org:")) == 0.0 + assert broadcasts.wait_for(_zeroed(counter), BROADCAST_SECONDS), broadcasts.received() + _assert_served(_chat(gateway.client, model, key)) + assert outlives(_settled(cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(gateway.client, model, key)) + + +def test_an_organization_update_without_spend_keeps_the_counter_and_the_block(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key, counter = _organization_entities(gateway, scenario) + _settled(cache, counter, ONE_CALL_COST) + _update_organization(gateway, counter, metadata={"note": f"audit {uuid.uuid4().hex}"}) + value, ttl = _state(cache, counter) + assert value == pytest.approx(ONE_CALL_COST, abs=1e-9) and outlives(ttl), (value, ttl) + _assert_denied(_chat(gateway.client, model, key)) + + +def test_an_organization_spend_set_above_the_budget_keeps_the_block_until_set_to_zero(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key, counter = _organization_entities(gateway, scenario) + _settled(cache, counter, ONE_CALL_COST) + _update_organization(gateway, counter, spend=1.5) + assert outlives(_settled(cache, counter, 1.5)) + _assert_denied(_chat(gateway.client, model, key)) + _update_organization(gateway, counter, spend=0.0) + assert outlives(_settled(cache, counter, 0.0)) + _assert_served(_chat(gateway.client, model, key)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_an_organization_spend_set_while_redis_is_down_answers_500_and_the_retry_resets_the_counter( + gateway: Gateway, tmp_path: Path +) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model, key, counter = _organization_entities(isolated.gateway, scenario, isolated.database_url) + organization: Final = counter.removeprefix("spend:org:") + _settled(isolated.cache, counter, ONE_CALL_COST) + isolated.redis.stop() + _assert_counters_unwritten( + _organization_update(isolated.gateway, counter, spend=0.0), "The organization spend was saved" + ) + assert _organization_spend(organization, isolated.database_url) == 0.0 + _assert_served(_chat(isolated.gateway.client, model, key)) + isolated.redis.start() + assert isolated.cache.ping() + _accepted_once_redis_is_back(lambda: _organization_update(isolated.gateway, counter, spend=0.0)) + assert _organization_spend(organization, isolated.database_url) == 0.0 + assert outlives(_settled(isolated.cache, counter, 0.0)) + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) diff --git a/tests/integration/spend/test_spend_counter_ttl_chaos.py b/tests/integration/spend/test_spend_counter_ttl_chaos.py new file mode 100644 index 00000000000..1cd6d9312be --- /dev/null +++ b/tests/integration/spend/test_spend_counter_ttl_chaos.py @@ -0,0 +1,292 @@ +import uuid +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.spend_counters import ( + counter_ttl, + counter_value, + key_counter, + outlives, + rig_cache, +) +from pydantic import JsonValue +from redis import Redis + +pytestmark: Final = pytest.mark.shares_worker_state + +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"} +BURST: Final = 24 +UNBOUNDED_BUDGET: Final = 100000.0 +RESERVATION_ESTIMATE_CEILING: Final = 40.0 +SLOW_UPSTREAM_SECONDS: Final = 3.0 +LEAK_KEY_BUDGET: Final = 20.0 +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +RESPAWN_SECONDS: Final = graceful_stop_seconds() + 60 + + +def _chat(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"burst chat {index} {uuid.uuid4().hex}"}], + "stream": stream, + **({"stream_options": {"include_usage": True}} if stream else {}), + }, + key=key, + ) + + +def _messages(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 32, + "messages": [{"role": "user", "content": f"burst messages {index} {uuid.uuid4().hex}"}], + "stream": stream, + }, + key=key, + headers=ANTHROPIC_HEADERS, + ) + + +def _responses(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": f"burst responses {index} {uuid.uuid4().hex}", "stream": stream}, + key=key, + ) + + +SURFACES: Final[tuple[Callable[[Gateway, str, str, int, bool], httpx.Response], ...]] = ( + _chat, + _messages, + _responses, +) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, spend, status FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ) + + +def _successes(key: str, count: int) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: [row for row in _spend_rows(key) if str(row["status"]) == "success"], + lambda values: len(values) >= count, + seconds=90, + ) + assert len(rows) == count, rows + assert len({str(row["request_id"]) for row in rows}) == count, rows + return rows + + +def _counter_matches(cache: Redis, counter: str, total: float) -> tuple[float | None, int]: + state: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda values: values[0] is not None and values[0] == pytest.approx(total, abs=1e-9), + seconds=20, + ) + assert outlives(state[1]), state + return state + + +def test_a_burst_with_upstream_failures_lands_each_success_once_with_a_long_ttl(gateway: Gateway) -> None: + with ( + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + gateway.scenario() as scenario, + rig_cache() as cache, + ): + provider_model: Final = f"burst-{uuid.uuid4().hex}" + model: Final = scenario.model(model=f"openai/{provider_model}", **RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + statuses: Final = [500] + [200, 200, 200] * (BURST // 4 + 2) + + def remove_script() -> None: + response: Final = upstream.delete(f"/__scripts/{provider_model}") + assert response.status_code in (200, 404), response.text + + scenario.cleanups.callback(remove_script) + configured: Final = upstream.post(f"/__scripts/{provider_model}", json={"statuses": statuses}) + assert configured.status_code == 200, configured.text + + def attempt(index: int) -> httpx.Response: + stream: Final = index % 2 == 1 + return SURFACES[index % len(SURFACES)](gateway, model, key, index, stream) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + responses: Final = tuple(pool.map(attempt, range(BURST))) + codes: Final = [response.status_code for response in responses] + assert all(code in (200, 500) for code in codes), codes + succeeded: Final = codes.count(200) + assert 0 < succeeded < BURST, codes + rows: Final = _successes(key, succeeded) + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx(succeeded * 0.06), rows + _counter_matches(cache, key_counter(key), total) + + +def _workers(owned: OwnedProxy) -> tuple[int, ...]: + return tuple(sorted(owned.worker_pids())) + + +def _attempt(gateway: Gateway, model: str, key: str, index: int) -> int: + try: + return _chat(gateway, model, key, index, False).status_code + except httpx.HTTPError: + return 0 + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_killed_worker_does_not_shorten_or_lose_the_counter(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, + owned.gateway.scenario() as scenario, + rig_cache() as cache, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + before: Final = _workers(owned) + assert len(before) == 2, before + with ThreadPoolExecutor(max_workers=12) as pool: + first: Final = pool.submit(_attempt, owned.gateway, model, key, 0) + assert first.result(timeout=60) == 200 + in_flight: Final = [pool.submit(_attempt, owned.gateway, model, key, index) for index in range(1, 12)] + with owned.replacing_worker(before[0]) as kill: + kill() + survivor: Final = _attempt(owned.gateway, model, key, 99) + codes: Final = [future.result(timeout=120) for future in in_flight] + assert survivor == 200, survivor + assert all(code in (0, 200) for code in codes), codes + replaced: Final = eventually( + lambda: _workers(owned), + lambda pids: len(pids) == 2 and before[0] not in pids and before[1] in pids, + seconds=RESPAWN_SECONDS, + ) + assert replaced + after: Final = [_attempt(owned.gateway, model, key, index) for index in range(12, 18)] + assert after == [200] * 6, after + answered: Final = 1 + 1 + codes.count(200) + 6 + landed: Final = eventually( + lambda: [row for row in _spend_rows(key) if str(row["status"]) == "success"], + lambda rows: len(rows) >= 1 + 1 + 6, + seconds=90, + ) + assert len({str(row["request_id"]) for row in landed}) == len(landed) <= answered, landed + ceiling: Final = answered * 0.06 + len(in_flight) * RESERVATION_ESTIMATE_CEILING + state: Final = eventually( + lambda: (counter_value(cache, key_counter(key)), counter_ttl(cache, key_counter(key))), + lambda values: values[0] is not None and 8 * 0.06 - 1e-9 <= values[0] <= ceiling, + seconds=20, + ) + assert outlives(state[1]), state + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_reservation_leaked_by_a_killed_worker_stays_for_the_counter_ttl(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, + owned.gateway.scenario() as scenario, + rig_cache() as cache, + ): + slow: Final = scenario.model(model="openai/slow-endpoint", **RATES) + key: Final = scenario.key(models=[slow], max_budget=LEAK_KEY_BUDGET) + counter: Final = key_counter(key) + before: Final = _workers(owned) + assert len(before) == 2, before + with ThreadPoolExecutor(max_workers=1) as pool: + in_flight: Final = pool.submit(_attempt, owned.gateway, slow, key, 0) + reserved: Final = eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value >= LEAK_KEY_BUDGET - 1e-9, + seconds=SLOW_UPSTREAM_SECONDS, + ) + with owned.replacing_worker(before[0]) as first, owned.replacing_worker(before[1]) as second: + first.freeze() + second.freeze() + first() + second() + interrupted: Final = in_flight.result(timeout=60) + assert interrupted != 200, interrupted + leaked: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert leaked[0] == pytest.approx(reserved), (reserved, leaked) + respawned: Final = eventually( + lambda: _workers(owned), + lambda pids: len(pids) == 2 and not set(pids) & set(before), + seconds=RESPAWN_SECONDS, + ) + assert respawned + blocked: Final = eventually( + lambda: _attempt(owned.gateway, slow, key, 1), lambda code: code == 422, seconds=RESPAWN_SECONDS + ) + assert blocked == 422, (blocked, leaked) + held: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert held[0] == pytest.approx(reserved), (reserved, held) + assert outlives(held[1]), held + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_redis_outage_mid_burst_keeps_serving_and_rebuilds_the_counter_long(gateway: Gateway, tmp_path: Path) -> None: + with owned_redis(tmp_path) as redis: + overrides: Final[Mapping[str, str]] = { + "REDIS_HOST": redis.host, + "REDIS_PORT": str(redis.port), + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2", + } + with ( + owned_proxy_process(gateway, tmp_path, overrides) as owned, + owned.gateway.scenario() as scenario, + Redis(host=redis.host, port=redis.port) as cache, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + counter: Final = key_counter(key) + + def burst(start: int) -> list[int]: + with ThreadPoolExecutor(max_workers=6) as pool: + return list( + pool.map(lambda index: _attempt(owned.gateway, model, key, index), range(start, start + 6)) + ) + + assert burst(0) == [200] * 6 + _counter_matches(cache, counter, 6 * 0.06) + redis.stop() + during: Final = burst(6) + assert during == [200] * 6, during + redis.start() + assert cache.ping() + probes: Final[list[int]] = [] + + def probe_until_the_counter_is_back() -> float | None: + probes.append(_attempt(owned.gateway, model, key, 100 + len(probes))) + return counter_value(cache, counter) + + eventually(probe_until_the_counter_is_back, lambda value: value is not None, seconds=30) + assert set(probes) == {200}, probes + after: Final = burst(12) + assert after == [200] * 6, after + rows: Final = _successes(key, 18 + len(probes)) + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx((18 + len(probes)) * 0.06), rows + rebuilt: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and 7 * 0.06 - 1e-9 <= state[0] <= total + 1e-9, + seconds=20, + ) + assert outlives(rebuilt[1]), rebuilt diff --git a/tests/integration/spend/test_spend_counter_ttl_settings.py b/tests/integration/spend/test_spend_counter_ttl_settings.py new file mode 100644 index 00000000000..c0fe74603cc --- /dev/null +++ b/tests/integration/spend/test_spend_counter_ttl_settings.py @@ -0,0 +1,230 @@ +import json +import os +import time +import uuid +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import psycopg +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.spend_counters import counter_ttl, counter_value, key_counter, outlives, rig_cache +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ONE_CALL_COST: Final = 0.06 +SLOW_UPSTREAM_SECONDS: Final = 3.0 +SHORT_REDIS_TTL: Final = 4 +LONG_REDIS_TTL: Final = 86400 +ENV_REDIS_TTL: Final = "7200" +LONG_BATCH_WRITE_SECONDS: Final = 900 +STRICT_KEY_BUDGET: Final = 50.0 +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +FAST_RESET_TICK: Final = MappingProxyType( + {"PROXY_BUDGET_RESCHEDULER_MIN_TIME": "2", "PROXY_BUDGET_RESCHEDULER_MAX_TIME": "2"} +) + + +def _config( + directory: Path, + *, + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), + general_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + configured: Final = { + **shared, + "litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings}, + "general_settings": {**object_value(shared["general_settings"]), **general_settings}, + } + path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(configured, sort_keys=False)) + return path + + +def _chat(gateway: Gateway, model: str, key: str) -> str: + served: Final = gateway.chat(model, key=key, text=f"settings {uuid.uuid4().hex}") + assert served["usage"]["total_tokens"] == 40, served + return str(served["id"]) + + +def _counter_state(counter: str) -> tuple[float | None, int]: + with rig_cache() as cache: + return counter_value(cache, counter), counter_ttl(cache, counter) + + +def _settled(counter: str, spend: float, expected_ttl: int) -> tuple[float | None, int]: + return eventually( + lambda: _counter_state(counter), + lambda state: state[0] is not None and state[0] == pytest.approx(spend) and outlives(state[1], expected_ttl), + seconds=15, + ) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_null_default_redis_ttl_falls_back_to_the_hour_floor(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": None}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + assert _chat(owned, model, key) + assert _settled(key_counter(key), ONE_CALL_COST, 3600) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_short_default_redis_ttl_never_shortens_the_counter_and_the_lease_refresh_keeps_it_long( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": SHORT_REDIS_TTL}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(model="openai/slow-endpoint", **RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + with ThreadPoolExecutor(max_workers=1) as pool: + slow: Final = pool.submit(lambda: (_chat(owned, model, key), time.monotonic())[1]) + eventually(lambda: cache.exists(counter), lambda present: present == 1, seconds=SLOW_UPSTREAM_SECONDS) + assert cache.expire(counter, 1000) + refreshed_at: Final = eventually( + lambda: (time.monotonic(), counter_ttl(cache, counter)), + lambda state: outlives(state[1]), + seconds=SLOW_UPSTREAM_SECONDS + 10, + )[0] + finished_at: Final = slow.result(timeout=60) + assert refreshed_at < finished_at, (refreshed_at, finished_at) + assert _settled(counter, ONE_CALL_COST, 3600) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_long_default_redis_ttl_wins_on_every_write_path(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": LONG_REDIS_TTL}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + assert _settled(counter, ONE_CALL_COST, LONG_REDIS_TTL) + updated: Final = owned.request("POST", "/key/update", {"key": key, "spend": 2.5}) + assert updated.status_code == 200, updated.text + assert _settled(counter, 2.5, LONG_REDIS_TTL) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_default_redis_ttl_given_as_an_environment_string_is_read_as_seconds( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": "os.environ/SPEND_COUNTER_REDIS_TTL"}) + overrides: Final = {"SPEND_COUNTER_REDIS_TTL": ENV_REDIS_TTL} + with owned_proxy(gateway, tmp_path, overrides, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + assert _settled(counter, ONE_CALL_COST, int(ENV_REDIS_TTL)) + reset: Final = owned.request("POST", f"/key/{key}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert _settled(counter, 0.0, int(ENV_REDIS_TTL)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_long_batch_write_interval_lifts_the_counter_ttl(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, general_settings={"proxy_batch_write_at": LONG_BATCH_WRITE_SECONDS}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + assert _chat(owned, model, key) + assert _settled(key_counter(key), ONE_CALL_COST, 10 * LONG_BATCH_WRITE_SECONDS) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_strict_enforcement_repairs_a_stale_low_key_counter_with_the_long_ttl(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, general_settings={"fail_closed_budget_enforcement": True}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=STRICT_KEY_BUDGET) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + attempts: Final[list[str]] = [] + + def stale_then_request() -> tuple[float | None, int]: + assert cache.set(counter, "0.01") + attempts.append(_chat(owned, model, key)) + return eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and state[0] < STRICT_KEY_BUDGET / 2, + seconds=15, + ) + + repaired: Final = eventually( + stale_then_request, + lambda state: state[0] is not None and state[0] > 0.01 + ONE_CALL_COST + 1e-6, + seconds=30, + ) + lifted: Final = repaired[0] + assert ( + lifted is not None and 2 * ONE_CALL_COST - 1e-9 <= lifted <= (len(attempts) + 1) * ONE_CALL_COST + 1e-9 + ), ( + repaired, + attempts, + ) + assert outlives(repaired[1]), repaired + + +def _window_reset_at(key: str) -> datetime: + rows: Final = read_rows( + 'SELECT budget_limits FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ) + assert len(rows) == 1, rows + raw: Final = rows[0]["budget_limits"] + windows: Final = json.loads(raw) if isinstance(raw, str) else raw + assert isinstance(windows, list) and len(windows) == 1, raw + return datetime.fromisoformat(str(object_value(windows[0])["reset_at"]).replace("Z", "+00:00")) + + +def _make_window_due(key: str) -> None: + due: Final = (datetime.now(UTC) - timedelta(days=1)).replace(tzinfo=None).isoformat() + windows: Final = [{"max_budget": 1.0, "budget_duration": "1d", "reset_at": due}] + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute( + 'UPDATE "LiteLLM_VerificationToken" SET budget_limits = %s::jsonb WHERE token = %s', + (json.dumps(windows), sha256(key.encode()).hexdigest()), + ) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_reset_budget_window_keeps_the_long_ttl(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, FAST_RESET_TICK) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], budget_limits=[{"max_budget": 1.0, "budget_duration": "1d"}]) + window_counter: Final = f"{key_counter(key)}:window:1d" + assert _chat(owned, model, key) + assert _settled(window_counter, ONE_CALL_COST, 3600) + _make_window_due(key) + rolled: Final = eventually( + lambda: (_window_reset_at(key), counter_value(cache, window_counter), counter_ttl(cache, window_counter)), + lambda state: state[0].replace(tzinfo=None) > datetime.now(UTC).replace(tzinfo=None), + seconds=45, + ) + assert rolled[1] is not None and rolled[1] < ONE_CALL_COST, rolled + assert outlives(rolled[2]), rolled diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 0fbeb05e154..50160299b58 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -403,6 +403,44 @@ async def test_async_scan_iter_namespaces_pattern(namespace, expected_match, mon assert captured["match"] == expected_match +@pytest.mark.asyncio +async def test_async_scan_iter_decodes_bytes_keys_and_stops_at_count(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + def scan_iter(match, count): + async def gen(): + for key in (b"spend:key:a", "spend:key:b", b"spend:key:c", b"spend:key:d"): + yield key + + return gen() + + mock_redis_instance = MagicMock() + mock_redis_instance.scan_iter = scan_iter + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + keys = await redis_cache.async_scan_iter(pattern="spend:key:", count=3) + assert keys == ["spend:key:a", "spend:key:b", "spend:key:c"] + + +@pytest.mark.asyncio +async def test_async_scan_iter_unbounded_returns_every_matching_key(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + def scan_iter(match, count): + async def gen(): + for key in (b"spend:key:a", "spend:key:b", b"spend:key:c", b"spend:key:d"): + yield key + + return gen() + + mock_redis_instance = MagicMock() + mock_redis_instance.scan_iter = scan_iter + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + keys = await redis_cache.async_scan_iter(pattern="spend:key:", count=3, unbounded=True) + assert keys == ["spend:key:a", "spend:key:b", "spend:key:c", "spend:key:d"] + + @pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) def test_increment_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") @@ -1516,6 +1554,65 @@ async def test_async_set_cache_pipeline_with_ttls_keeps_each_entry_ttl(monkeypat ] +class _LostPipeline(_SetRecordingPipeline): + async def execute(self) -> list[bool]: + from redis.exceptions import ConnectionError as RedisConnectionError + + raise RedisConnectionError("redis down") + + +@pytest.mark.asyncio +async def test_async_set_cache_pipeline_with_ttls_returns_the_error_of_a_lost_write(monkeypatch, redis_no_ping): + """A spend reset must not report success on a write Redis never applied, so the lost write comes back + as a value while every other caller keeps ignoring it.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + client = MagicMock() + client.pipeline = MagicMock(side_effect=(_LostPipeline(), _SetRecordingPipeline())) + + with patch.object(redis_cache, "init_async_client", return_value=client): + lost = await redis_cache.async_set_cache_pipeline_with_ttls((("k", 1, 60),)) + kept = await redis_cache.async_set_cache_pipeline_with_ttls((("k", 1, 60),)) + + assert isinstance(lost, RedisConnectionError) + assert kept is None + + +@pytest.mark.parametrize("namespace", [None, "ns"]) +def test_reset_fence_key_shares_the_cluster_slot_of_its_counter(namespace, monkeypatch, redis_no_ping): + """The repair checks the fence and the counter in one Lua call, which a cluster only runs on one slot.""" + from redis.crc import key_slot + + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + counter = redis_cache.check_and_fix_namespace(key="spend:key:abc") + fence = redis_cache.check_and_fix_namespace(key=redis_cache.reset_fence_key("spend:key:abc")) + + assert fence != counter + assert key_slot(fence.encode()) == key_slot(counter.encode()) + + +@pytest.mark.asyncio +async def test_async_set_max_declares_the_reset_fence_as_a_second_key(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = MagicMock() + client.eval = AsyncMock(return_value=b"3.0") + + with patch.object(redis_cache, "init_async_client", return_value=client): + held = await redis_cache.async_set_max( + key="spend:key:abc", value=12.0, ttl=60, fence_key=redis_cache.reset_fence_key("spend:key:abc") + ) + unfenced = await redis_cache.async_set_max(key="spend:key:abc", value=12.0, ttl=60) + + assert (held, unfenced) == (3.0, 3.0) + fenced_call, plain_call = client.eval.await_args_list + assert fenced_call.args[1:4] == (2, "ns:spend:key:abc", "ns:{ns:spend:key:abc}:reset_fence") + assert plain_call.args[1:3] == (1, "ns:spend:key:abc") + + class _ListPipeline: def __init__(self, rows: list[str]) -> None: self.rows = rows diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 78ce9728f3b..0e761eb5fd1 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -10696,7 +10696,8 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele ) assert spend.memory == {"spend:team_member:user-1:team-1": 2.5} - assert spend.commands == ["SET spend:team_member:user-1:team-1 2.5 60"] + assert spend.memory_ttls == {"spend:team_member:user-1:team-1": 60} + assert spend.commands == [f"SET spend:team_member:user-1:team-1 2.5 {spend.counter_ttl}"] @pytest.mark.asyncio @@ -10823,6 +10824,7 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t redis_cache=_FakeRedisCache(), user_api_key_cache=remote_user_api_key_cache, additional_in_memory_caches=(remote_spend_counter_in_memory_cache,), + counter_cache=remote_spend_counter_in_memory_cache, spend_counters=remote_spend_counter_store, ) for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"): 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..72f806c249b 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 @@ -19,6 +19,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( broadcast_evictions, evict_and_broadcast, publish_auth_cache_invalidation, + publish_auth_cache_invalidations, ) from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, @@ -324,6 +325,92 @@ async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None +@pytest.mark.asyncio +async def test_publish_many_sends_one_message_per_five_hundred_keys() -> None: + client = _RecordingRedisClient() + keys = tuple(f"spend:key:{index}" for index in range(1001)) + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client), + ): + await publish_auth_cache_invalidations(cache_keys=keys, new_value=0.0, ttl=5) + await asyncio.gather(*pubsub_module._pending_publishes) + + assert [channel for channel, _ in client.published] == [AUTH_CACHE_INVALIDATION_CHANNEL] * 3 + assert [json.loads(message) for _, message in client.published] == [ + {"cache_keys": list(keys[:500]), "new_value": 0.0, "ttl": 5}, + {"cache_keys": list(keys[500:1000]), "new_value": 0.0, "ttl": 5}, + {"cache_key": keys[1000], "new_value": 0.0, "ttl": 5}, + ] + + +@pytest.mark.asyncio +async def test_subscriber_sets_every_key_of_a_multi_key_message_in_the_counter_cache_only() -> None: + """A global reset broadcasts thousands of counters; written into every additional cache they would evict + the BYOK credentials cached next to them.""" + cache = UserApiKeyCache() + spend_counter_in_memory_cache = InMemoryCache() + spend_counter_in_memory_cache.set_cache("spend:key:a", 999.0) + spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:key:a", 999.0) + byok_cache = InMemoryCache() + byok_cache.set_cache("byok:cred-1", {"api_key": "sk-cred"}) + message = { + "type": "message", + "data": json.dumps({"cache_keys": ["spend:key:a", "spend:key:b"], "new_value": 0.0, "ttl": 60}).encode(), + } + + pubsub = _QueuePubSub(initial_messages=[message]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + additional_in_memory_caches=(spend_counter_in_memory_cache, byok_cache), + counter_cache=spend_counter_in_memory_cache, + ) + subscriber.start() + try: + for _ in range(200): + if spend_counter_in_memory_cache.get_cache("spend:key:b") == 0.0: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert spend_counter_in_memory_cache.get_cache("spend:key:a") == 0.0 + assert spend_counter_in_memory_cache.get_cache("spend:key:b") == 0.0 + assert spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:key:a") == 999.0 + assert byok_cache.get_cache("byok:cred-1") == {"api_key": "sk-cred"} + assert byok_cache.get_cache("spend:key:a") is None + + +@pytest.mark.asyncio +async def test_subscriber_writes_every_key_of_a_multi_key_message_to_the_native_spend_counter_copies() -> None: + """A global reset sends thousands of counters in one message, and each counter's native copy has to take the new + value or that worker keeps answering with the old spend when Redis cannot.""" + store = SpendCounterStore(connect(None).spend_counters()) + store.set_local("spend:key:a", 999.0) + store.set_local("spend:key:b", 999.0) + pubsub = _QueuePubSub( + initial_messages=[ + { + "type": "message", + "data": json.dumps({"cache_keys": ["spend:key:a", "spend:key:b"], "new_value": 0.0, "ttl": 60}), + } + ] + ) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=UserApiKeyCache(), + spend_counters=store, + ) + subscriber.start() + try: + assert await asyncio.wait_for(pubsub.drained.wait(), timeout=5) + finally: + await subscriber.stop() + + assert store.snapshot() == {"spend:key:a": 0.0, "spend:key:b": 0.0} + + @pytest.mark.asyncio async def test_subscriber_writes_and_drops_the_native_spend_counter_copies_a_broadcast_names() -> None: """The spend counters' memory tier lives in the native store, not in a Python InMemoryCache, so a broadcast diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 4a4490afc33..92c1b7c774a 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -4812,3 +4812,29 @@ async def test_reset_budget_for_litellm_team_members_called(): assert len(team_member_writes) == 1 assert team_member_writes[0]["where"]["budget_id"]["in"] == ["budget1"] assert team_member_writes[0]["data"] == {"spend": 0} + + +@pytest.mark.asyncio +async def test_reset_expired_window_writes_the_redis_counter_with_the_spend_counter_ttl(): + """The window reset is a counter write like every other, so it must carry the spend counter TTL, not the 60s cache default.""" + from datetime import datetime + + from litellm.proxy.common_utils.reset_budget_job import BudgetResetSettings, Litellm_EntityType + + spend = FakeSpendCounters(counter_ttl=4321) + window = {"budget_duration": "1d", "max_budget": 1.0, "reset_at": "2026-01-01T00:00:00Z"} + + reset = await ResetBudgetJob._reset_expired_window( + window=window, + counter_key="spend:key:abc:window:1d", + spend_counter_store=spend.store(), + now=datetime(2026, 1, 2), + reset_settings=BudgetResetSettings(), + prisma_client=AsyncMock(), + entity_type=Litellm_EntityType.KEY, + entity_id="abc", + ) + + assert reset is True + assert spend.commands[-1] == "SET spend:key:abc:window:1d 0.0 4321" + assert spend.memory == {"spend:key:abc:window:1d": 0.0} diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 77dc6eddf51..ae4f3f98595 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -8783,7 +8783,8 @@ async def test_reset_key_spend_resets_budget_windows(monkeypatch): window_counter_key = f"spend:key:{hashed_key}:window:1d" assert spend.memory[window_counter_key] == 0.0 - assert f"SET {window_counter_key} 0.0 60" in spend.commands + assert spend.memory_ttls[window_counter_key] == 60 + assert f"SET {window_counter_key} 0.0 {spend.counter_ttl}" in spend.commands mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call( key=f"spend_db_floor:{window_counter_key}", value=0.0, ttl=5 ) @@ -8997,7 +8998,8 @@ async def test_update_key_spend_updates_counter(monkeypatch): mock_delete_cache.assert_awaited_once() assert spend.memory == {f"spend:key:{hashed_key}": 0.0} - assert spend.commands == [f"SET spend:key:{hashed_key} 0.0 60"] + assert spend.memory_ttls == {f"spend:key:{hashed_key}": 60} + assert spend.commands == [f"SET spend:key:{hashed_key} 0.0 {spend.counter_ttl}"] @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 8d2c6c0054c..8326d6d92dd 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -3,7 +3,7 @@ import json from contextlib import asynccontextmanager from types import MappingProxyType, SimpleNamespace from typing import Final, Mapping, Optional, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from fastapi import HTTPException @@ -14,6 +14,7 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( CascadingJWTMappingTable, JWTMappingRow, ) +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters @pytest.mark.asyncio @@ -648,6 +649,7 @@ async def test_organization_member_add_budget_omission_and_null_leave_budget_uns from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add user = LiteLLM_UserTable(user_id="user-1", user_role="internal_user") + async def create_membership(data): return LiteLLM_OrganizationMembershipTable( user_id="user-1", @@ -1582,6 +1584,92 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke assert jwt_table.rows == (kept_row,) +def _org_spend_counter_cache( + monkeypatch: pytest.MonkeyPatch, *, lost_write: Exception | None = None, failing: frozenset[str] = frozenset() +) -> tuple[MagicMock, FakeSpendCounters]: + from litellm.caching.in_memory_cache import InMemoryCache + + spend = FakeSpendCounters(memory={"spend:org:org-1": 4.2}, failing=failing) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache = InMemoryCache() + spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:org:org-1", value=4.2) + spend_counter_cache.redis_cache.reset_fence_key = lambda key: f"{{{key}}}:reset_fence" + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock(return_value=lost_write) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.publish_auth_cache_invalidations", AsyncMock()) + return spend_counter_cache, spend + + +@pytest.mark.asyncio +async def test_legacy_update_with_spend_sets_the_org_counter_and_its_floor_on_every_pod(monkeypatch): + """A direct spend write must SET the cross-pod counter and its floor marker: a pod that cached the old floor + would otherwise repair a deleted-and-reseeded counter back up, and the org stays at the stale value until + the counter expires.""" + from litellm.proxy import proxy_server + + spend_counter_cache, spend = _org_spend_counter_cache(monkeypatch) + + await _run_legacy_update_organization( + monkeypatch, body={"organization_id": "org-1", "spend": 1.5}, existing_budget_id="budget-1" + ) + + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.assert_awaited_once_with( + (("{spend:org:org-1}:reset_fence", 1, proxy_server.SPEND_RESET_FENCE_TTL_SECONDS),) + ) + assert spend.commands == [f"SET spend:org:org-1 1.5 {spend.counter_ttl}"] + assert proxy_server.spend_counter_store.get_local("spend:org:org-1") == 1.5 + assert spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:org:org-1") == 1.5 + assert proxy_server.publish_auth_cache_invalidations.await_args_list == [ + call(("spend_db_floor:spend:org:org-1",), new_value=1.5, ttl=proxy_server.SPEND_DB_FLOOR_CACHE_TTL_SECONDS), + call(("spend:org:org-1",), new_value=1.5, ttl=60), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lost", ["fence", "counter"]) +async def test_legacy_update_with_spend_answers_500_when_the_redis_counter_write_is_lost(monkeypatch, lost): + """The row is saved either way, but a lost counter write leaves the other pods enforcing the old spend + for up to the counter TTL, so the admin hears about it and retries instead of reading success. The fence and + the counter go through two Redis clients, and losing either one is a lost reset.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException + + _org_spend_counter_cache( + monkeypatch, + lost_write=RedisConnectionError("redis down") if lost == "fence" else None, + failing=frozenset({"set"}) if lost == "counter" else frozenset(), + ) + + with pytest.raises(ProxyException) as raised: + await _run_legacy_update_organization( + monkeypatch, body={"organization_id": "org-1", "spend": 1.5}, existing_budget_id="budget-1" + ) + + assert raised.value.code == "500" + assert raised.value.message.startswith("The organization spend was saved, but the cached spend counters") + proxy_server.prisma_client.db.litellm_organizationtable.update.assert_awaited_once() + assert proxy_server.spend_counter_store.get_local("spend:org:org-1") == 1.5 + assert proxy_server.publish_auth_cache_invalidations.await_count == 2 + + +@pytest.mark.asyncio +async def test_legacy_update_without_spend_keeps_the_org_spend_counter(monkeypatch): + spend_counter_cache, spend = _org_spend_counter_cache(monkeypatch) + + await _run_legacy_update_organization( + monkeypatch, body={"organization_id": "org-1", "organization_alias": "renamed"}, existing_budget_id="budget-1" + ) + + from litellm.proxy import proxy_server + + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.assert_not_awaited() + assert spend.commands == [] + assert proxy_server.spend_counter_store.get_local("spend:org:org-1") == 4.2 + + @pytest.mark.asyncio async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu_share(monkeypatch): """Deleting an organization deletes its teams, so a team still holding a PTU share stops it the diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 19db80b8930..179b14beae7 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -44,6 +44,7 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_u from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team from .conftest import normalize +from litellm.redis.stores.spend_counters import SpendCounterStore from litellm.rust_bridge._native import NativeRedisStoreError from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters @@ -697,6 +698,32 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert spend.memory == dict.fromkeys(expected_keys, 1.5) +def test_spend_counter_redis_ttl_never_shortens_configured_redis_ttl(monkeypatch): + """Deployments that configured a longer Redis default must not see spend counters expire sooner""" + monkeypatch.setattr(ps.litellm, "default_redis_ttl", 86400) + monkeypatch.setattr(ps, "proxy_batch_write_at", 60) + + assert ps.spend_counter_redis_ttl() == 86400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("batch_write_at", [60, 900]) +async def test_spend_counter_redis_ttl_outlives_db_batch_write_interval(monkeypatch, batch_write_at): + """A counter that expires before its buffered spend is written to the DB reseeds from a stale-low + row and re-admits an over-budget key, so every counter write must outlive the batch interval""" + spend: Final = FakeSpendCounters(redis={"spend:key:reserved": 1.0}) + monkeypatch.setattr(ps, "spend_counter_store", SpendCounterStore(spend, ps.spend_counter_redis_ttl)) + monkeypatch.setattr(ps, "proxy_batch_write_at", batch_write_at) + max_flush_delay: Final = batch_write_at + 5 + + await ps.increment_spend_counter_cache(counter_key="spend:key:direct", increment=0.5) + await ps.run_spend_counter_pipeline((ps.PendingSpendIncrement(counter_key="spend:key:batched", increment=0.5),)) + await ps.refresh_spend_counter_ttl(counter_key="spend:key:reserved") + + assert set(spend.ttls) == {"spend:key:direct", "spend:key:batched", "spend:key:reserved"} + assert all(ttl is not None and ttl > max_flush_delay for ttl in spend.ttls.values()) + + _PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( "spend:key:hashed", "spend:team:team-1", @@ -1331,6 +1358,14 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey assert result is None +@pytest.mark.parametrize(("max_batch_size", "keys_per_pipeline"), [(1, 1), (2, 1), (3, 1), (1000, 500)]) +def test_a_reset_pipeline_takes_half_the_redis_batch_size_and_never_less_than_one_key( + max_batch_size: int, keys_per_pipeline: int +) -> None: + """Each key rides with its fence, so half the batch size; a step of zero would make every reset raise.""" + assert ps._keys_per_reset_pipeline(max_batch_size) == keys_per_pipeline + + async def _drain_spend_cache_writes() -> None: for _ in range(3): await asyncio.sleep(0) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index c3e916d581d..4691df6777c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -5,10 +5,11 @@ import hashlib import json import re import sqlite3 -from collections.abc import Mapping, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence +from contextlib import asynccontextmanager from datetime import timezone from typing import Final, Literal -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import prisma import pytest @@ -20,6 +21,7 @@ import litellm.proxy.proxy_server as ps from litellm.proxy.auth.authorization import OwnedRows from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup, load_permitted_log_team_ids from litellm.proxy.auth.route_checks import RouteChecks +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters from tests.unit.proxy.db.fake_prisma_engine import ( ReadTargetOf, statement_timeout_error, @@ -8336,6 +8338,187 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader ] +def _spend_counter_cache_for_reset( + monkeypatch: pytest.MonkeyPatch, + *, + scan_results: object, + lost_write: Exception | None = None, + failing: frozenset[str] = frozenset(), +) -> tuple[MagicMock, FakeSpendCounters]: + from litellm.caching.in_memory_cache import InMemoryCache + + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.query_raw = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock() + mock_prisma_client.db.litellm_teamtable.update_many = AsyncMock() + + @asynccontextmanager + async def transaction() -> AsyncIterator[MagicMock]: + yield mock_prisma_client.db + + mock_prisma_client.tx = transaction + monkeypatch.setattr(ps, "prisma_client", mock_prisma_client) + spend: Final = FakeSpendCounters( + memory={"spend:key:aaa": 0.5, "spend:key:aaa:window:1d": 0.3, "spend:team:t1": 0.5, "spend:user:u1": 0.5}, + failing=failing, + ) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + spend_counter_cache: Final = MagicMock() + spend_counter_cache.in_memory_cache = InMemoryCache() + spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:key:aaa", value=0.5) + spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:key:floored", value=0.9) + spend_counter_cache.redis_cache.namespace = "ns" + spend_counter_cache.redis_cache.reset_fence_key = lambda key: f"{{ns:{key}}}:reset_fence" + spend_counter_cache.redis_cache.async_scan_iter = AsyncMock(side_effect=scan_results) + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock(return_value=lost_write) + monkeypatch.setattr(ps, "spend_counter_cache", spend_counter_cache) + monkeypatch.setattr(ps, "publish_auth_cache_invalidations", AsyncMock()) + return spend_counter_cache, spend + + +_RESET_KEY_COUNTERS: Final = ("spend:key:aaa", "spend:key:bbb", "spend:key:floored") +_RESET_TEAM_COUNTERS: Final = ("spend:team:sales:window:east", "spend:team:t1") + + +def _reset_fences(counters: tuple[str, ...]) -> tuple[tuple[str, int, int], ...]: + return tuple((f"{{ns:{counter}}}:reset_fence", 1, ps.SPEND_RESET_FENCE_TTL_SECONDS) for counter in counters) + + +def _reset_writes(spend: FakeSpendCounters, counters: tuple[str, ...]) -> list[str]: + return [f"SET {counter} 0.0 {spend.counter_ttl}" for counter in counters] + + +def _reset_broadcasts(counters: tuple[str, ...]) -> tuple[object, object]: + """The floors go out first: a subscriber that applied the new counter under the old floor would repair it.""" + return ( + call( + tuple(f"spend_db_floor:{counter}" for counter in counters), + new_value=0.0, + ttl=ps.SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ), + call(counters, new_value=0.0, ttl=60), + ) + + +@pytest.mark.asyncio +async def test_global_spend_reset_sets_every_key_and_team_counter_and_its_floor_to_zero( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Deleting the counters is not enough: a pod that cached the pre-reset floor within the last 5s would repair + the reseeded counter back to the old total for the whole counter TTL, so the reset SETs counter and floor + everywhere, fences each counter against that repair, and leaves budget-window counters, which the reset + does not touch in the DB, alone. A window counter is told by its trailing duration, so a team whose id + merely contains ``:window:`` is still reset.""" + spend_counter_cache, spend = _spend_counter_cache_for_reset( + monkeypatch, + scan_results=[ + ["ns:spend:key:aaa", "ns:spend:key:bbb", "ns:spend:key:bbb:window:1d"], + ["ns:spend:team:t1", "ns:spend:team:t1:window:1d", "ns:spend:team:sales:window:east"], + ], + ) + + response: Final = await spend_management_endpoints.global_spend_reset() + + assert response["status"] == "success" + assert [c.kwargs for c in spend_counter_cache.redis_cache.async_scan_iter.await_args_list] == [ + {"pattern": "spend:key:", "count": 500, "unbounded": True}, + {"pattern": "spend:team:", "count": 500, "unbounded": True}, + ] + assert [c.args for c in spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_args_list] == [ + (_reset_fences(_RESET_KEY_COUNTERS),), + (_reset_fences(_RESET_TEAM_COUNTERS),), + ] + assert spend.commands == [*_reset_writes(spend, _RESET_KEY_COUNTERS), *_reset_writes(spend, _RESET_TEAM_COUNTERS)] + in_memory: Final = spend_counter_cache.in_memory_cache + reset_counters: Final = (*_RESET_KEY_COUNTERS, *_RESET_TEAM_COUNTERS) + assert [ps.spend_counter_store.get_local(counter) for counter in reset_counters] == [0.0] * 5 + assert [in_memory.get_cache(key=f"spend_db_floor:{counter}") for counter in reset_counters] == [0.0] * 5 + assert ps.spend_counter_store.get_local("spend:key:aaa:window:1d") == 0.3 + assert ps.spend_counter_store.get_local("spend:team:t1:window:1d") is None + assert ps.spend_counter_store.get_local("spend:user:u1") == 0.5 + assert ps.publish_auth_cache_invalidations.await_args_list == [ + *_reset_broadcasts(_RESET_KEY_COUNTERS), + *_reset_broadcasts(_RESET_TEAM_COUNTERS), + ] + + +@pytest.mark.asyncio +async def test_global_spend_reset_answers_500_after_resetting_the_rows_and_this_pod_when_redis_cannot_be_listed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The counters Redis holds for the other pods stay at the old spend for the whole counter TTL when the + scan fails, so the reset says so instead of reporting success, after resetting what this pod knows.""" + spend_counter_cache, spend = _spend_counter_cache_for_reset(monkeypatch, scan_results=ConnectionError("redis down")) + + with pytest.raises(ProxyException) as raised: + await spend_management_endpoints.global_spend_reset() + + assert raised.value.code == "500" + assert "retry this request once Redis accepts writes" in raised.value.message + ps.prisma_client.db.litellm_verificationtoken.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={}) + ps.prisma_client.db.litellm_teamtable.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={}) + assert [c.args for c in spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_args_list] == [ + (_reset_fences(("spend:key:aaa", "spend:key:floored")),), + (_reset_fences(("spend:team:t1",)),), + ] + assert spend.commands == [ + *_reset_writes(spend, ("spend:key:aaa", "spend:key:floored")), + *_reset_writes(spend, ("spend:team:t1",)), + ] + assert ps.spend_counter_store.get_local("spend:key:aaa") == 0.0 + assert spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:key:floored") == 0.0 + + +@pytest.mark.asyncio +async def test_global_spend_reset_answers_500_when_the_redis_counter_write_is_lost( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A lost pipeline leaves the other pods on the old counters for up to the counter TTL, so the reset still + resets the teams, this pod, and the broadcast, then reports the loss so the admin retries.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + spend_counter_cache, spend = _spend_counter_cache_for_reset( + monkeypatch, + scan_results=[["ns:spend:key:aaa"], ["ns:spend:team:t1"]], + lost_write=RedisConnectionError("redis down"), + ) + + with pytest.raises(ProxyException) as raised: + await spend_management_endpoints.global_spend_reset() + + assert raised.value.code == "500" + assert raised.value.message.startswith("Spend was reset, but the cached spend counters could not be reset in Redis") + assert "redis down" in raised.value.message + assert spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_count == 2 + assert spend.commands == [] + assert ps.spend_counter_store.get_local("spend:team:t1") == 0.0 + assert ps.publish_auth_cache_invalidations.await_args_list == [ + *_reset_broadcasts(("spend:key:aaa", "spend:key:floored")), + *_reset_broadcasts(("spend:team:t1",)), + ] + + +@pytest.mark.asyncio +async def test_global_spend_reset_answers_500_while_the_spend_counter_store_refuses_writes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The fences and the counters go through two Redis clients with their own circuit breakers. A reset that + read as written while the counter store still refused writes would drop the next request's increment, so the + retry has to keep answering 500 until the store every request increments accepts the reset.""" + spend_counter_cache, spend = _spend_counter_cache_for_reset( + monkeypatch, scan_results=[["ns:spend:key:aaa"], ["ns:spend:team:t1"]], failing=frozenset({"set"}) + ) + + with pytest.raises(ProxyException) as raised: + await spend_management_endpoints.global_spend_reset() + + assert raised.value.code == "500" + assert raised.value.message.startswith("Spend was reset, but the cached spend counters could not be reset in Redis") + assert spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_count == 2 + assert ps.spend_counter_store.get_local("spend:key:aaa") == 0.0 + assert ps.spend_counter_store.get_local("spend:team:t1") == 0.0 + + def test_spend_calculate_rejection_body_is_openai_shaped(client): app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user", api_key="hashed-k" diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 57a07153ca1..e13036e74fd 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -40,6 +40,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.spend_tracking.budget_reservation import ( _get_model_access_group_budget_counters, + _lease_renewals, estimate_request_max_cost, get_budget_window_start, invalidate_budget_reservation_counters, @@ -2204,7 +2205,7 @@ class _ExpiringSpendCounters(FakeSpendCounters): self.refresh_count = 0 self.fail_first_refresh = fail_first_refresh - def refresh_ttl(self, counter: str) -> asyncio.Future[bool]: + def refresh_ttl(self, counter: str, ttl_seconds: float | None = None) -> asyncio.Future[bool]: self.refresh_attempts += 1 future: asyncio.Future[bool] = asyncio.get_running_loop().create_future() redis = self.redis @@ -2259,6 +2260,26 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( assert redis_cache.redis.get(counter_key) is None +@pytest.mark.asyncio +async def test_reservation_lease_renewal_stops_promptly_with_long_spend_counter_ttl( + spend_counter_state, +): + """A long spend-counter TTL must not stretch how often renewal checks for completion, or finished + requests keep their renewal tasks alive for half that TTL""" + counter_cache, key_cache = spend_counter_state + _install(_ExpiringSpendCounters(ttl=0.2)) + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-long-ttl", spend=0.0, max_budget=1.0) + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + assert len(_lease_renewals) == 1 + + await release_budget_reservation(reservation) + await asyncio.sleep(0.3) + assert len(_lease_renewals) == 0 + + @pytest.mark.asyncio async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( spend_counter_state, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 1cfb14e6bf9..776b831d0cb 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -9278,7 +9278,9 @@ async def test_window_spend_counter_seed_asks_the_store_to_increment_a_lost_seed seeds: list[tuple[str, float, bool]] = [] class RecordingSeeds(FakeSpendCounters): - def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> asyncio.Future[float]: + def seed( + self, counter: str, spend: float, *, increment_if_lost: bool, ttl_seconds: float | None = None + ) -> asyncio.Future[float]: seeds.append((counter, spend, increment_if_lost)) return super().seed(counter, spend, increment_if_lost=increment_if_lost) diff --git a/tests/unit/redis/stores/spend_counters_double.py b/tests/unit/redis/stores/spend_counters_double.py index 1f3c4b94d45..4871795eb52 100644 --- a/tests/unit/redis/stores/spend_counters_double.py +++ b/tests/unit/redis/stores/spend_counters_double.py @@ -57,7 +57,8 @@ class DeclaredRead: class FakeSpendCounters: """``redis`` is None for a worker without Redis, and ``redis_default`` answers reads of a key it lacks. ``fail`` makes every Redis call fail with ``failure``, ``fail_pipelines`` only the increment pipelines and ``failing`` only - the named writes ("set", "delete").""" + the named writes ("set", "delete"). ``counter_ttl`` is what the store gives each Redis write, ``ttls`` the last one + each counter's write carried and ``memory_ttls`` the last memory TTL of a ``set``.""" redis: dict[str, object] | None = field(default_factory=dict) memory: dict[str, float] = field(default_factory=dict) @@ -68,6 +69,9 @@ class FakeSpendCounters: failure: str = "transport" redis_default: float | None = None redis_ttl: float = 60 + counter_ttl: float = 3600 + ttls: dict[str, float | None] = field(default_factory=dict) + memory_ttls: dict[str, float | None] = field(default_factory=dict) batch_reads: bool = True read_hook: Callable[[str], Awaitable[None]] | None = None on_delete: Callable[[tuple[str, ...]], None] | None = None @@ -76,7 +80,7 @@ class FakeSpendCounters: declared: list[DeclaredRead] = field(default_factory=list) def store(self) -> SpendCounterStore: - return SpendCounterStore(self) + return SpendCounterStore(self, lambda: self.counter_ttl) @property def has_redis(self) -> bool: @@ -121,7 +125,7 @@ class FakeSpendCounters: events.append("declare spend_counters") return declared - def increment(self, counter: str, amount: float) -> asyncio.Future[float]: + def increment(self, counter: str, amount: float, ttl_seconds: float | None = None) -> asyncio.Future[float]: _check(counter) if self.redis is None: return _done(self.increment_local(counter, amount)) @@ -129,19 +133,26 @@ class FakeSpendCounters: self.memory.pop(counter, None) return _failed(unavailable(self.failure)) self.commands.append(f"INCRBYFLOAT {counter} {amount}") + self.ttls[counter] = ttl_seconds return _done(self._incr(counter, amount)) - def increment_many(self, increments: Sequence[SpendIncrement]) -> asyncio.Future[Sequence[float]]: + def increment_many( + self, increments: Sequence[SpendIncrement], ttl_seconds: float | None = None + ) -> asyncio.Future[Sequence[float]]: if self.redis is None: return _done([self.increment_local(counter, amount) for counter, amount in increments]) if self.fail or self.fail_pipelines: return _failed(unavailable(self.failure)) self.commands.append(f"PIPELINE {' '.join(counter for counter, _ in increments)}") + self.ttls.update(dict.fromkeys((counter for counter, _ in increments), ttl_seconds)) return _done([self._incr(counter, amount) for counter, amount in increments]) - def increment_many_post_call(self, request: object, increments: Sequence[SpendIncrement]) -> bool: + def increment_many_post_call( + self, request: object, increments: Sequence[SpendIncrement], ttl_seconds: float | None = None + ) -> bool: if self.redis is None: return False + self.ttls.update(dict.fromkeys((counter for counter, _ in increments), ttl_seconds)) self.post_call.append((request, tuple(increments))) return True @@ -158,7 +169,7 @@ class FakeSpendCounters: for counter, amount in increments: self._incr(counter, amount) - def set_max(self, counter: str, spend: float) -> asyncio.Future[float | None]: + def set_max(self, counter: str, spend: float, ttl_seconds: float | None = None) -> asyncio.Future[float | None]: _check(counter) if self.memory.get(counter) is None or self.memory[counter] < spend: self.memory[counter] = spend @@ -168,11 +179,14 @@ class FakeSpendCounters: return _failed(unavailable(self.failure)) self.commands.append(f"SETMAX {counter} {spend}") current: Final = self.redis.get(counter) + self.ttls[counter] = ttl_seconds if current is None or current < spend: self.redis[counter] = spend return _done(float(self.redis[counter])) - def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> asyncio.Future[float]: + def seed( + self, counter: str, spend: float, *, increment_if_lost: bool, ttl_seconds: float | None = None + ) -> asyncio.Future[float]: _check(counter) if self.redis is None: cached: Final = self.memory.get(counter) @@ -182,18 +196,27 @@ class FakeSpendCounters: if self.fail: return _failed(unavailable(self.failure)) self.commands.append(f"SETNX {counter} {spend}") + self.ttls[counter] = ttl_seconds current: Final = self.redis.setdefault(counter, spend) self.memory[counter] = float(current) return _done(float(current)) - def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> asyncio.Future[None]: + def set( + self, + counter: str, + spend: float, + memory_ttl_seconds: float | None = None, + redis_ttl_seconds: float | None = None, + ) -> asyncio.Future[None]: _check(counter) self.memory[counter] = spend + self.memory_ttls[counter] = memory_ttl_seconds if self.redis is None: return _done(None) if self.fail or "set" in self.failing: return _failed(unavailable(self.failure)) - self.commands.append(f"SET {counter} {spend} {ttl_seconds}") + self.commands.append(f"SET {counter} {spend} {redis_ttl_seconds}") + self.ttls[counter] = redis_ttl_seconds self.redis[counter] = spend return _done(None) @@ -217,13 +240,14 @@ class FakeSpendCounters: self.redis.pop(counter, None) return _done(None) - def refresh_ttl(self, counter: str) -> asyncio.Future[bool]: + def refresh_ttl(self, counter: str, ttl_seconds: float | None = None) -> asyncio.Future[bool]: _check(counter) if self.redis is None: return _done(False) if self.fail: return _failed(unavailable(self.failure)) self.commands.append(f"EXPIRE {counter}") + self.ttls[counter] = ttl_seconds return _done(counter in self.redis) def get_local(self, counter: str) -> float | None: diff --git a/tests/unit/redis/stores/test_spend_counters.py b/tests/unit/redis/stores/test_spend_counters.py index c54dc6cc81e..cfde8570af7 100644 --- a/tests/unit/redis/stores/test_spend_counters.py +++ b/tests/unit/redis/stores/test_spend_counters.py @@ -103,6 +103,67 @@ async def test_without_redis_the_native_store_counts_and_answers_from_memory() - assert not store.increment_many_post_call((("spend:key:k", 1.0),)) +@pytest.mark.asyncio +async def test_every_redis_write_takes_the_counter_ttl_as_it_stands_when_it_is_written() -> None: + """The settings the TTL depends on load after the store is built, so a TTL read once would go stale.""" + spend: Final = FakeSpendCounters(redis={"spend:key:refreshed": 1.0}) + store: Final = spend.store() + + spend.counter_ttl = 101 + await store.increment("spend:key:incremented", 1.0) + spend.counter_ttl = 102 + await store.increment_many((("spend:key:piped", 1.0),)) + spend.counter_ttl = 103 + await store.set_max("spend:key:repaired", 1.0) + spend.counter_ttl = 104 + await store.seed("spend:key:seeded", 1.0, increment_if_lost=False) + spend.counter_ttl = 105 + await store.set("spend:key:reset", 1.0, memory_ttl=60) + spend.counter_ttl = 106 + await store.refresh_ttl("spend:key:refreshed") + spend.counter_ttl = 107 + natives: Final = FakeNatives() + with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double + store.increment_many_post_call((("spend:key:post_call", 1.0),)) + + assert spend.ttls == { + "spend:key:incremented": 101, + "spend:key:piped": 102, + "spend:key:repaired": 103, + "spend:key:seeded": 104, + "spend:key:reset": 105, + "spend:key:refreshed": 106, + "spend:key:post_call": 107, + } + assert spend.memory_ttls == {"spend:key:reset": 60} + + +@pytest.mark.asyncio +async def test_a_store_without_a_counter_ttl_leaves_each_write_on_the_redis_clients_default() -> None: + spend: Final = FakeSpendCounters() + + await SpendCounterStore(spend).increment("spend:key:k", 1.0) + + assert spend.ttls == {"spend:key:k": None} + + +@pytest.mark.parametrize( + ("namespace", "fence", "sent"), + [ + (None, "{spend:team:t}:reset_fence", "{spend:team:t}:reset_fence"), + ("litellm", "{litellm:spend:team:t}:reset_fence", "litellm:{litellm:spend:team:t}:reset_fence"), + ], +) +def test_redis_cache_writes_the_reset_fence_the_native_store_checks(namespace: str | None, fence: str, sent: str) -> None: + """Mirrors ``reset_fence_key`` in the crate's ``keys.rs``, and ``a_reset_fence_holds_a_repair_until_it_is_gone`` + checks the same bytes against real Redis.""" + cache: Final = RedisCache.__new__(RedisCache) + cache.namespace = namespace + + assert cache.reset_fence_key("spend:team:t") == fence + assert cache.check_and_fix_namespace(key=fence) == sent + + def test_a_broadcast_writes_and_drops_only_spend_counter_copies() -> None: store: Final = SpendCounterStore(connect(None).spend_counters()) store.set_local("spend:key:k", 5.0) @@ -153,9 +214,11 @@ async def test_post_call_increments_go_on_the_requests_native_post_call_batch_un targets: Final[list[str | None]] = [] class Recording(FakeSpendCounters): - def increment_many_post_call(self, request: object, increments: object) -> bool: + def increment_many_post_call( + self, request: object, increments: object, ttl_seconds: float | None = None + ) -> bool: targets.append(current_service_target()) - return super().increment_many_post_call(request, increments) # pyright: ignore[reportArgumentType] # recorded as given + return super().increment_many_post_call(request, increments, ttl_seconds) # pyright: ignore[reportArgumentType] # recorded as given spend: Final = Recording() store: Final = spend.store()