mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 87150af883 into ab3b12b951
This commit is contained in:
commit
65642a8eef
39 changed files with 3320 additions and 241 deletions
|
|
@ -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
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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?)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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. "
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
33
tests/integration/_support/spend_counters.py
Normal file
33
tests/integration/_support/spend_counters.py
Normal 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
|
||||
|
|
@ -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"]),
|
||||
|
|
|
|||
|
|
@ -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"] == []
|
||||
488
tests/integration/spend/test_spend_counter_redis_ttl.py
Normal file
488
tests/integration/spend/test_spend_counter_redis_ttl.py
Normal 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))
|
||||
46
tests/integration/spend/test_spend_counter_reset_fence.py
Normal file
46
tests/integration/spend/test_spend_counter_reset_fence.py
Normal 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))
|
||||
667
tests/integration/spend/test_spend_counter_reset_paths.py
Normal file
667
tests/integration/spend/test_spend_counter_reset_paths.py
Normal 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))
|
||||
292
tests/integration/spend/test_spend_counter_ttl_chaos.py
Normal file
292
tests/integration/spend/test_spend_counter_ttl_chaos.py
Normal 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
|
||||
230
tests/integration/spend/test_spend_counter_ttl_settings.py
Normal file
230
tests/integration/spend/test_spend_counter_ttl_settings.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue