This commit is contained in:
mubashir1osmani 2026-11-06 22:20:31 +00:00 • committed by GitHub
commit 65642a8eef
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
39 changed files with 3320 additions and 241 deletions

View file

@ -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<SpendCounter>,
amount: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
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<SpendCounter>, PyNumber)>,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
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::<Vec<_>>())
})
}
/// 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<SpendCounter>, PyNumber)>,
) -> bool {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
}
fn set_max<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
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<SpendCounter>,
spend: PyNumber,
increment_if_lost: bool,
) -> PyResult<Bound<'py, PyAny>> {
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<f64>,
) -> PyResult<bool> {
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<SpendCounter>,
@ -179,8 +159,52 @@ impl NativeSpendCounters {
) -> PyResult<Bound<'py, PyAny>> {
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<SpendCounter>,
spend: PyNumber,
increment_if_lost: bool,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
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<SpendCounter>,
spend: PyNumber,
memory_ttl_seconds: Option<f64>,
redis_ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
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<SpendCounter>,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
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
})
}

View file

@ -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
}

View file

@ -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}")

View file

@ -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<Self, NotASpendCounter> {
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);
}
}

View file

@ -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<Duration> {
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<Number, Error> {
pub async fn increment(
&self,
counter: &SpendCounter,
amount: Number,
ttl: Option<Duration>,
) -> Result<Number, Error> {
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<Duration>,
) -> Result<Vec<Number>, 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::<Vec<_>>()).await;
@ -180,10 +188,12 @@ impl SpendCounters {
self: &Arc<Self>,
request: &RequestRedis,
increments: &[(SpendCounter, Number)],
ttl: Option<Duration>,
) -> 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<Duration>,
) -> Result<Option<Number>, 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<f64> = 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<Duration>,
) -> Result<Number, Error> {
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<Duration>,
memory_ttl: Option<Duration>,
redis_ttl: Option<Duration>,
) -> 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<bool, Error> {
pub async fn refresh_ttl(
&self,
counter: &SpendCounter,
ttl: Option<Duration>,
) -> Result<bool, Error> {
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?)
}

View file

@ -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::<String>(&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::<String>(&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::<String>(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();

View file

@ -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<litellm_circuit_breaker::Breaker>,
events: Option<Arc<dyn RedisEvents>>,
) -> RedisClient {
build(deployment, Namespace::default(), breaker, events)
}
fn build(
deployment: &RedisDeployment,
namespace: Namespace,
breaker: Option<litellm_circuit_breaker::Breaker>,
events: Option<Arc<dyn RedisEvents>>,
) -> 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,

View file

@ -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()

View file

@ -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)

View file

@ -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 "

View file

@ -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:

View file

@ -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

View file

@ -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. "

View file

@ -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

View file

@ -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

View file

@ -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",

View file

@ -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)

View file

@ -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: ...

View file

@ -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.

View file

@ -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

View file

@ -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"]),

View file

@ -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"] == []

View file

@ -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))

View file

@ -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))

View file

@ -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))

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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"):

View file

@ -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

View file

@ -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}

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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"

View file

@ -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,

View file

@ -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)

View file

@ -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:

View file

@ -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()