mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy): show committed spend in the remaining budget gauges (#49199)
* fix(proxy): report live key and team spend to budget gauges within the key cache TTL * fix(prometheus): carry live spend for the user, org and customer budget gauges too The remaining budget gauges for users, orgs and customers read the same cached rows as the key and team gauges, so they read high for about a minute after a charge. Auth now carries the live counter for each budgeted entity in one snapshot that only the gauges read, floored at the spend auth loaded. The x-litellm-key-spend header and logging spend fields keep their current values. * fix(prometheus): show committed spend in the remaining budget gauges Reads twin counters that take only charged cost, so a worst-case budget hold never shows as spent and a request that finishes last cannot set a gauge back. The twins exist only while the Prometheus logger is registered, and every path that resets a spend counter drops them. Failed requests now set the org gauge from the carried org row and committed spend, not the cached org row. * fix(prometheus): count a charge the cached customer row already holds once * fix(prometheus): keep committed spend only for budgeted entities and drop it on spend resets Auth now reads committed spend through the admission MGET and names the counters a charge may move, so the master key and unbudgeted entities get no twin. Organization spend edits and the global spend reset drop their twins, and a long request no longer reads as charged. * fix(prometheus): read the committed org spend on failed requests without the auth object * test(prometheus): cover a failed request with no carried org * fix(prometheus): move the committed counters for charges with no auth on the row A billed background response is charged from its creator's stored attribution, after the request's auth is gone. Its charge moved the live counters but not the committed ones, so every gauge read one charge behind. It now moves the committed counters that already exist, and its own gauge adds its cost to them. * fix(prometheus): keep the gauges right for pass-through rows, overlapping charges and slow DB spend writes Pass-through rows carry the auth object, so a cold committed counter gets seeded there too. The gauges read a charge still waiting in the post-call pipeline as not counted, whatever other charges reached the counter. Committed counters outlive two DB spend writes. Cold seeds read the DB once per worker, and the global spend reset reads keys and teams in bounded pages. * fix(prometheus): count a charged failure's own cost in the org gauge A stream cut mid-response still charges for what it delivered, and that charge is usually still queued when the failure logs. The failure path passed a cost of 0, so the org gauge left that charge out until the worker's next request * fix(prometheus): skip twins for window-named ids, drop twins on delete, keep auth out of PostHog * fix(prometheus): keep the auth object off pass-through rows and drop user twins on every user delete Pass-through rows now carry only the committed spend in litellm_params, so loggers that dump the row's metadata never get the auth object. Every user delete path, SCIM and bulk delete included, drops the deleted users' committed spend twins * fix(prometheus): keep committed spend values out of logger dumps of pass-through rows CommittedBudgetSpend fields are repr=False, so Athina's payload, Sentry breadcrumbs and debug logs that str() a pass-through row's litellm_params show no spend values * fix(prometheus): keep the gauge fallback on a slow committed read and refresh committed twin TTLs on cluster The success-path committed read now gets half the per-request budget, so a slow Redis falls back to stored spend before the outer timeout cancels the gauges. Post-call charges refresh a committed twin's TTL on Redis Cluster too, so a twin never expires while its charges are still on their way to the DB. Also applies cargo fmt to the PR's Rust files * fix(proxy): key the global reset's committed-spend page queries so the Prisma span check reads them --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
zoveki
GitHub
parent
60d6797983
commit
abe3b11f7a
54 changed files with 2846 additions and 121 deletions
|
|
@ -55,7 +55,9 @@ mod _native {
|
|||
#[pymodule_export]
|
||||
use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins;
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::spend_counters::{NativeSpendCounters, PySpendRead};
|
||||
use crate::redis::store::spend_counters::{
|
||||
NativeSpendCounters, PyPostCallSettlement, PySpendRead,
|
||||
};
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::{
|
||||
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
|
||||
|
|
@ -154,6 +156,7 @@ mod tests {
|
|||
"NativeSensitiveRoutePins",
|
||||
"NativeRedisStoreError",
|
||||
"NativeSpendCounters",
|
||||
"PostCallSettlement",
|
||||
"SpendRead",
|
||||
"LocalSizes",
|
||||
"native_redis_breaker_states",
|
||||
|
|
|
|||
|
|
@ -2,12 +2,12 @@ use std::{collections::HashMap, sync::Arc};
|
|||
|
||||
use litellm_redis_store::{
|
||||
Number,
|
||||
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
|
||||
spend_counters::{LostSeed, PostCallSettlement, Reading, SpendCounter, SpendCounters},
|
||||
};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use super::{
|
||||
NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote,
|
||||
NativePending, PyData, PyNumber, duration, kind_error, pending, redis_async, remote,
|
||||
request::{NativeRequestRedis, NativeSend},
|
||||
ttl, unwrap_all,
|
||||
};
|
||||
|
|
@ -29,6 +29,18 @@ impl PySpendRead {
|
|||
}
|
||||
}
|
||||
|
||||
/// Whether every increment one `increment_many_post_call` declared has its reply.
|
||||
#[pyclass(frozen, name = "PostCallSettlement")]
|
||||
pub(crate) struct PyPostCallSettlement(PostCallSettlement);
|
||||
|
||||
#[pymethods]
|
||||
impl PyPostCallSettlement {
|
||||
#[getter]
|
||||
fn settled(&self) -> bool {
|
||||
self.0.is_settled()
|
||||
}
|
||||
}
|
||||
|
||||
fn increments(increments: Vec<(PyData<SpendCounter>, PyNumber)>) -> Vec<(SpendCounter, Number)> {
|
||||
increments
|
||||
.into_iter()
|
||||
|
|
@ -73,6 +85,12 @@ impl NativeSpendCounters {
|
|||
self.0.redis_ttl().map(|ttl| ttl.as_secs_f64())
|
||||
}
|
||||
|
||||
/// The least TTL a Redis write gives a committed spend twin.
|
||||
fn set_committed_ttl(&self, ttl_seconds: f64) -> PyResult<()> {
|
||||
self.0.set_committed_ttl(duration(ttl_seconds)?);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
|
||||
fn read<'py>(
|
||||
&self,
|
||||
|
|
@ -125,14 +143,15 @@ impl NativeSpendCounters {
|
|||
})
|
||||
}
|
||||
|
||||
/// Declares the increments on the request's post-call batch; `False` when there is no Redis.
|
||||
/// Declares the increments on the request's post-call batch; `None` when there is no Redis.
|
||||
fn increment_many_post_call(
|
||||
&self,
|
||||
request: &NativeRequestRedis,
|
||||
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
|
||||
) -> bool {
|
||||
) -> Option<PyPostCallSettlement> {
|
||||
self.0
|
||||
.increment_many_post_call(&request.0, &self::increments(increments))
|
||||
.map(PyPostCallSettlement)
|
||||
}
|
||||
|
||||
fn set_max<'py>(
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`)
|
||||
- An op declared on a `Send` returns `Pending<T>`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on
|
||||
- An op declared on a `Send` returns `Pending<T>`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on, or `None` when it returns a settlement the caller waits on, like `SpendCounterStore::increment_many_post_call`
|
||||
- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s
|
||||
- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec
|
||||
- A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc<RemoteSlot>`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient`
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
/// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum SpendCounter {
|
||||
/// The charges-only twin of a key, team, user, org or end user lifetime counter, which budget
|
||||
/// holds never touch: `spend:committed:` and the twin's key past its `spend:`.
|
||||
Committed(Box<SpendCounter>),
|
||||
Key {
|
||||
token: String,
|
||||
window: Option<String>,
|
||||
|
|
@ -33,6 +36,9 @@ impl SpendCounter {
|
|||
None => format!("spend:{scope}:{id}"),
|
||||
};
|
||||
match self {
|
||||
Self::Committed(counter) => {
|
||||
format!("spend:committed:{}", &counter.key()["spend:".len()..])
|
||||
}
|
||||
Self::Key { token, window } => windowed("key", token, window),
|
||||
Self::Team { team_id, window } => windowed("team", team_id, window),
|
||||
Self::TeamMember { user_id, team_id } => {
|
||||
|
|
@ -51,6 +57,16 @@ impl SpendCounter {
|
|||
pub fn parse(key: &str) -> Result<Self, NotASpendCounter> {
|
||||
let invalid = || NotASpendCounter(key.to_owned());
|
||||
let rest = key.strip_prefix("spend:").ok_or_else(invalid)?;
|
||||
if let Some(twin) = rest.strip_prefix("committed:") {
|
||||
return match Self::parse(&format!("spend:{twin}")).map_err(|_| invalid())? {
|
||||
counter @ (Self::Key { window: None, .. }
|
||||
| Self::Team { window: None, .. }
|
||||
| Self::User(_)
|
||||
| Self::Org(_)
|
||||
| Self::EndUser(_)) => Ok(Self::Committed(Box::new(counter))),
|
||||
_ => Err(invalid()),
|
||||
};
|
||||
}
|
||||
let (scope, id) = rest.split_once(':').ok_or_else(invalid)?;
|
||||
let windowed = || match id.rsplit_once(":window:") {
|
||||
Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())),
|
||||
|
|
|
|||
|
|
@ -1,13 +1,19 @@
|
|||
mod keys;
|
||||
mod types;
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
use std::{
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_python_compat::repr::to_str;
|
||||
use litellm_redis_client::{Declared, Expiry, RedisBatch, commands};
|
||||
|
||||
pub use self::keys::{NotASpendCounter, SpendCounter};
|
||||
pub use self::types::{LostSeed, Reading};
|
||||
pub use self::types::{LostSeed, PostCallSettlement, Reading};
|
||||
use crate::{
|
||||
Error, Number, Pending, RequestRedis, Send,
|
||||
dual_cache::{DualCache, MemorySettings},
|
||||
|
|
@ -31,6 +37,7 @@ const SET_MAX: &str = concat!(
|
|||
/// `DualCache` to its Redis and its memory.
|
||||
pub struct SpendCounters {
|
||||
cache: DualCache<Number>,
|
||||
committed_ttl_millis: AtomicU64,
|
||||
}
|
||||
|
||||
impl SpendCounters {
|
||||
|
|
@ -45,6 +52,33 @@ impl SpendCounters {
|
|||
MemorySettings::default(),
|
||||
now,
|
||||
),
|
||||
committed_ttl_millis: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// The least TTL a Redis write gives a [`SpendCounter::Committed`] twin, so it outlives the DB
|
||||
/// write that holds its charges and its next seed never reads a total still missing them.
|
||||
pub fn set_committed_ttl(&self, ttl: Duration) {
|
||||
let millis = u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX);
|
||||
self.committed_ttl_millis.store(millis, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
fn ttl(&self, counter: &SpendCounter, default_ttl: Duration) -> Duration {
|
||||
match counter {
|
||||
SpendCounter::Committed(_) => default_ttl.max(Duration::from_millis(
|
||||
self.committed_ttl_millis.load(Ordering::Relaxed),
|
||||
)),
|
||||
_ => default_ttl,
|
||||
}
|
||||
}
|
||||
|
||||
/// A committed twin's expiry moves out on every charge, on a cluster too, so it never expires
|
||||
/// while a charge it holds is still on its way to the DB.
|
||||
fn post_call_expiry(&self, counter: &SpendCounter, default_ttl: Duration) -> Expiry {
|
||||
let ttl = self.ttl(counter, default_ttl);
|
||||
match counter {
|
||||
SpendCounter::Committed(_) => Expiry::Refresh(ttl),
|
||||
_ => Expiry::KeepExisting(ttl),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -127,7 +161,7 @@ impl SpendCounters {
|
|||
let increment = commands::incr_by_float(
|
||||
&key,
|
||||
amount.into(),
|
||||
Some(Expiry::Refresh(remote.default_ttl)),
|
||||
Some(Expiry::Refresh(self.ttl(counter, remote.default_ttl))),
|
||||
);
|
||||
match remote.client.send(increment).await {
|
||||
Ok(total) => {
|
||||
|
|
@ -159,7 +193,7 @@ impl SpendCounters {
|
|||
pipeline.add(commands::incr_by_float(
|
||||
&counter.key(),
|
||||
(*amount).into(),
|
||||
Some(Expiry::Refresh(remote.default_ttl)),
|
||||
Some(Expiry::Refresh(self.ttl(counter, remote.default_ttl))),
|
||||
))
|
||||
});
|
||||
let totals = Declared::all(totals.collect::<Vec<_>>()).await;
|
||||
|
|
@ -175,15 +209,14 @@ impl SpendCounters {
|
|||
}
|
||||
|
||||
/// `increment_many` on the request's post-call batch. Each total reaches memory once the batch
|
||||
/// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis.
|
||||
/// goes out, and a counter whose increment failed is dropped everywhere. `None` without Redis,
|
||||
/// else a settlement that reports when every reply is back, whoever flushed the batch.
|
||||
pub fn increment_many_post_call(
|
||||
self: &Arc<Self>,
|
||||
request: &RequestRedis,
|
||||
increments: &[(SpendCounter, Number)],
|
||||
) -> bool {
|
||||
let Some(remote) = self.cache.redis().get() else {
|
||||
return false;
|
||||
};
|
||||
) -> Option<PostCallSettlement> {
|
||||
let remote = self.cache.redis().get()?;
|
||||
let batch = Send::PostCall(request).batch(&remote.client);
|
||||
let declared = increments
|
||||
.iter()
|
||||
|
|
@ -191,12 +224,14 @@ impl SpendCounters {
|
|||
let total = batch.add(commands::incr_by_float(
|
||||
&counter.key(),
|
||||
(*amount).into(),
|
||||
Some(Expiry::KeepExisting(remote.default_ttl)),
|
||||
Some(self.post_call_expiry(counter, remote.default_ttl)),
|
||||
));
|
||||
(counter.clone(), total)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let store = Arc::clone(self);
|
||||
let settled = Arc::new(AtomicBool::new(false));
|
||||
let settlement = PostCallSettlement(Arc::clone(&settled));
|
||||
request.spawn(async move {
|
||||
for (counter, total) in declared {
|
||||
match total.settled().await {
|
||||
|
|
@ -215,8 +250,9 @@ impl SpendCounters {
|
|||
}
|
||||
}
|
||||
}
|
||||
settled.store(true, Ordering::Release);
|
||||
});
|
||||
true
|
||||
Some(settlement)
|
||||
}
|
||||
|
||||
/// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either.
|
||||
|
|
@ -243,7 +279,7 @@ impl SpendCounters {
|
|||
};
|
||||
let args = (
|
||||
to_str(&spend.to_value()),
|
||||
remote.default_ttl.as_secs().to_string(),
|
||||
self.ttl(counter, remote.default_ttl).as_secs().to_string(),
|
||||
);
|
||||
let stored: Option<f64> = remote
|
||||
.client
|
||||
|
|
@ -276,8 +312,9 @@ impl SpendCounters {
|
|||
})
|
||||
})?);
|
||||
};
|
||||
let ttl = self.ttl(counter, *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 +327,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?)
|
||||
}
|
||||
|
|
@ -313,7 +350,7 @@ impl SpendCounters {
|
|||
let Some(remote) = self.cache.redis().get() else {
|
||||
return Ok(());
|
||||
};
|
||||
let ttl = ttl.unwrap_or(remote.default_ttl);
|
||||
let ttl = ttl.unwrap_or_else(|| self.ttl(counter, remote.default_ttl));
|
||||
remote
|
||||
.client
|
||||
.send(commands::set(&key, spend.to_json(), ttl))
|
||||
|
|
@ -334,7 +371,10 @@ impl SpendCounters {
|
|||
};
|
||||
Ok(remote
|
||||
.client
|
||||
.send(commands::expire(&counter.key(), remote.default_ttl))
|
||||
.send(commands::expire(
|
||||
&counter.key(),
|
||||
self.ttl(counter, remote.default_ttl),
|
||||
))
|
||||
.await?)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,8 @@
|
|||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
};
|
||||
|
||||
use litellm_python_compat::Value;
|
||||
|
||||
use super::keys::SpendCounter;
|
||||
|
|
@ -32,6 +37,17 @@ impl Reading {
|
|||
}
|
||||
}
|
||||
|
||||
/// One `increment_many_post_call`: settled once every increment it declared has its reply and
|
||||
/// memory mirrors it, whether `flush_post` or the post-call deadline sent the batch.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PostCallSettlement(pub(super) Arc<AtomicBool>);
|
||||
|
||||
impl PostCallSettlement {
|
||||
pub fn is_settled(&self) -> bool {
|
||||
self.0.load(Ordering::Acquire)
|
||||
}
|
||||
}
|
||||
|
||||
/// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum LostSeed {
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within};
|
|||
use tokio::runtime::Handle;
|
||||
|
||||
const TTL: Duration = Duration::from_secs(60);
|
||||
const COMMITTED_TTL: Duration = Duration::from_secs(300);
|
||||
const LONG_DEADLINE: Duration = Duration::from_secs(60);
|
||||
|
||||
const READS: [(&str, Number); 4] = [
|
||||
|
|
@ -45,6 +46,10 @@ fn team(id: &str) -> SpendCounter {
|
|||
}
|
||||
}
|
||||
|
||||
fn committed(counter: SpendCounter) -> SpendCounter {
|
||||
SpendCounter::Committed(Box::new(counter))
|
||||
}
|
||||
|
||||
fn store(client: Option<RedisClient>) -> Arc<SpendCounters> {
|
||||
Arc::new(SpendCounters::new(client.map(remote)))
|
||||
}
|
||||
|
|
@ -114,6 +119,26 @@ fn a_reading_names_where_it_was_answered(
|
|||
"spend:model_access_group:gpt-group"
|
||||
)]
|
||||
#[case(SpendCounter::ShadowEval("job-7".to_owned()), "spend:shadow_eval:job-7")]
|
||||
#[case::committed_key(
|
||||
committed(SpendCounter::Key {
|
||||
token: "88dc28d0".to_owned(),
|
||||
window: None,
|
||||
}),
|
||||
"spend:committed:key:88dc28d0"
|
||||
)]
|
||||
#[case::committed_team(committed(team("team-a")), "spend:committed:team:team-a")]
|
||||
#[case::committed_user(
|
||||
committed(SpendCounter::User("user@example.com".to_owned())),
|
||||
"spend:committed:user:user@example.com"
|
||||
)]
|
||||
#[case::committed_org(
|
||||
committed(SpendCounter::Org("org-1".to_owned())),
|
||||
"spend:committed:org:org-1"
|
||||
)]
|
||||
#[case::committed_end_user(
|
||||
committed(SpendCounter::EndUser("customer é 1".to_owned())),
|
||||
"spend:committed:end_user:customer é 1"
|
||||
)]
|
||||
fn every_counter_formats_and_parses_back_to_mains_key(
|
||||
#[case] counter: SpendCounter,
|
||||
#[case] key: &str,
|
||||
|
|
@ -128,6 +153,11 @@ fn every_counter_formats_and_parses_back_to_mains_key(
|
|||
#[case("spend:unknown:abc")]
|
||||
#[case("spend:team_member:no-team")]
|
||||
#[case("spend:")]
|
||||
#[case::committed_window("spend:committed:key:abc:window:1d")]
|
||||
#[case::committed_tag("spend:committed:tag:prod")]
|
||||
#[case::committed_team_member("spend:committed:team_member:user:team-a")]
|
||||
#[case::committed_twice("spend:committed:committed:key:abc")]
|
||||
#[case::committed_nothing("spend:committed:")]
|
||||
fn keys_outside_the_spend_counter_keyspace_are_refused(#[case] key: &str) {
|
||||
assert!(SpendCounter::parse(key).is_err(), "{key}");
|
||||
}
|
||||
|
|
@ -307,7 +337,7 @@ async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_an
|
|||
.expect("a Python thread has no Tokio context, and declaring there must not panic");
|
||||
request.flush_post().await;
|
||||
|
||||
assert!(declared);
|
||||
assert!(declared.is_some());
|
||||
let mirrored =
|
||||
litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await;
|
||||
assert_eq!(mirrored, Number::Float(3.5));
|
||||
|
|
@ -331,9 +361,12 @@ 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))]));
|
||||
let settlement = store
|
||||
.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))])
|
||||
.unwrap();
|
||||
request.flush_pre().await;
|
||||
|
||||
assert!(!settlement.is_settled());
|
||||
assert!(trips.seen().is_empty());
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
|
|
@ -345,6 +378,7 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
|
|||
request.flush_post().await;
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
|
||||
assert!(settlement.is_settled());
|
||||
let increment = match topology {
|
||||
Topology::Standalone => Operation::Pipeline { ops: 1 },
|
||||
Topology::Cluster => Operation::Command("INCRBYFLOAT"),
|
||||
|
|
@ -557,13 +591,132 @@ async fn a_failed_post_call_increment_drops_the_counter_from_memory() {
|
|||
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))]));
|
||||
let settlement = store
|
||||
.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))])
|
||||
.unwrap();
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
request.flush_post().await;
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
|
||||
assert_eq!(store.get_local(&counter).unwrap(), None);
|
||||
assert!(settlement.is_settled());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_post_call_settlement_settles_when_the_deadline_sends_the_batch(
|
||||
#[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 request = RequestRedis::new(Handle::current(), Duration::from_millis(50));
|
||||
let counter = committed(team("t"));
|
||||
|
||||
let settlement = store
|
||||
.increment_many_post_call(&request, &[(counter.clone(), Number::Float(2.5))])
|
||||
.unwrap();
|
||||
|
||||
assert!(!settlement.is_settled());
|
||||
litellm_redis_testing::eventually(async || settlement.is_settled().then_some(())).await;
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some("2.5".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::committed_twin(committed, 300)]
|
||||
#[case::enforcement_counter(std::convert::identity, 60)]
|
||||
#[tokio::test]
|
||||
async fn every_write_gives_a_committed_twin_the_committed_ttl_and_other_counters_the_redis_ttl(
|
||||
#[case] counter_of: fn(SpendCounter) -> SpendCounter,
|
||||
#[case] ttl_seconds: u64,
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
store.set_committed_ttl(COMMITTED_TTL);
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
let [
|
||||
incremented,
|
||||
pipelined,
|
||||
post_call,
|
||||
seeded,
|
||||
repaired,
|
||||
refreshed,
|
||||
] = ["a", "b", "c", "d", "e", "f"].map(|id| counter_of(team(id)));
|
||||
raw.send(commands::set(
|
||||
&refreshed.key(),
|
||||
"1.0",
|
||||
Duration::from_secs(5),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
store
|
||||
.increment(&incremented, Number::Float(1.0))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.increment_many(&[(pipelined.clone(), Number::Float(1.0))])
|
||||
.await
|
||||
.unwrap();
|
||||
let settlement = store
|
||||
.increment_many_post_call(&request, &[(post_call.clone(), Number::Float(1.0))])
|
||||
.unwrap();
|
||||
request.flush_post().await;
|
||||
store
|
||||
.seed(&seeded, Number::Float(1.0), LostSeed::Answer)
|
||||
.await
|
||||
.unwrap();
|
||||
store.set_max(&repaired, Number::Float(1.0)).await.unwrap();
|
||||
assert!(store.refresh_ttl(&refreshed).await.unwrap());
|
||||
|
||||
assert!(settlement.is_settled());
|
||||
let remaining =
|
||||
async |counter: &SpendCounter| match raw.send(commands::ttl(&counter.key())).await.unwrap()
|
||||
{
|
||||
Ttl::Expires(remaining) => remaining,
|
||||
other => panic!("{} has no TTL: {other:?}", counter.key()),
|
||||
};
|
||||
assert!(within(remaining(&incremented).await, ttl_seconds));
|
||||
assert!(within(remaining(&pipelined).await, ttl_seconds));
|
||||
assert!(within(remaining(&post_call).await, ttl_seconds));
|
||||
assert!(within(remaining(&seeded).await, ttl_seconds));
|
||||
assert!(within(remaining(&repaired).await, ttl_seconds));
|
||||
assert!(within(remaining(&refreshed).await, ttl_seconds));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_post_call_charge_moves_a_committed_twins_expiry_out_to_the_committed_ttl(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
store.set_committed_ttl(COMMITTED_TTL);
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
let counter = committed(team("t"));
|
||||
raw.send(commands::set(&counter.key(), "1.0", Duration::from_secs(5)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let settlement = store
|
||||
.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))])
|
||||
.unwrap();
|
||||
request.flush_post().await;
|
||||
|
||||
assert!(settlement.is_settled());
|
||||
match raw.send(commands::ttl(&counter.key())).await.unwrap() {
|
||||
Ttl::Expires(remaining) => assert!(within(remaining, COMMITTED_TTL.as_secs())),
|
||||
other => panic!("{} has no TTL: {other:?}", counter.key()),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -592,7 +745,11 @@ 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))])
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.read_many(RedisSend::PreCall(&request), &[counter])
|
||||
|
|
|
|||
|
|
@ -452,6 +452,13 @@ class NativePending(Protocol[_T_co]):
|
|||
def resolve(self) -> Awaitable[_T_co]: ...
|
||||
|
||||
|
||||
class PostCallSettlement(Protocol):
|
||||
"""Post-call writes declared together: settled once every reply is back, whoever sent their batch."""
|
||||
|
||||
@property
|
||||
def settled(self) -> bool: ...
|
||||
|
||||
|
||||
class NativeBatchResult(Generic[_T]):
|
||||
"""Awaitable handle for a read declared on the request's native pre-call batch; awaiting it runs the
|
||||
native flush hooks, flushes that batch, then answers from what it read."""
|
||||
|
|
@ -526,6 +533,7 @@ class RequestRedisBatches:
|
|||
"_native_misses",
|
||||
"_native_send_factory",
|
||||
"_post_call",
|
||||
"_post_call_settlements",
|
||||
"post_call_deadline",
|
||||
"prefetched",
|
||||
)
|
||||
|
|
@ -538,6 +546,7 @@ class RequestRedisBatches:
|
|||
) -> None:
|
||||
self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
|
||||
self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
|
||||
self._post_call_settlements: tuple[PostCallSettlement, ...] = ()
|
||||
self.post_call_deadline: Final = post_call_deadline
|
||||
self._native_factory: Final = native_factory
|
||||
self._native_send_factory: Final = native_send_factory
|
||||
|
|
@ -641,6 +650,14 @@ class RequestRedisBatches:
|
|||
if not any(batch.pending for batch in self._post_call.values()):
|
||||
_open_post_call.discard(self)
|
||||
|
||||
def track_post_call_settlement(self, settlement: PostCallSettlement) -> None:
|
||||
"""Post-call writes a later reader in this request must know have landed, such as its committed spend."""
|
||||
self._post_call_settlements = (*self._post_call_settlements, settlement)
|
||||
|
||||
@property
|
||||
def post_call_settlements(self) -> tuple[PostCallSettlement, ...]:
|
||||
return self._post_call_settlements
|
||||
|
||||
@property
|
||||
def batches(self) -> tuple[RedisBatch, ...]:
|
||||
return tuple(self._batches.values())
|
||||
|
|
|
|||
|
|
@ -1941,6 +1941,7 @@ SPEND_ROLLUP_LOCK_TIMEOUT_MS: Final = max(1, get_env_int("SPEND_ROLLUP_LOCK_TIME
|
|||
SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = get_env_int("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)
|
||||
SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, get_env_int("SPEND_LOG_QUEUE_MAX_BYTES", 64000000))
|
||||
SPEND_LOG_QUEUE_POLL_INTERVAL: Final = get_env_float("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)
|
||||
SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS: Final = 30.0
|
||||
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", 3))
|
||||
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = get_env_float("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.2)
|
||||
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never
|
|||
import litellm
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.caching.redis_batch import active_request_redis_batches
|
||||
from litellm.constants import (
|
||||
AUTH_OBJECTS_TARGET,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
|
|
@ -75,8 +76,10 @@ from litellm.types.integrations.prometheus import (
|
|||
sanitize_prometheus_label_value,
|
||||
validate_prometheus_deployment_and_latency_caller_identity,
|
||||
)
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth, committed_budget_spend_of_row
|
||||
from litellm.types.proxy.carried_budget_state import (
|
||||
CommittedBudgetSpend,
|
||||
EntitySpendCounterKeys,
|
||||
KeyBudgetSnapshot,
|
||||
OrgBudgetSnapshot,
|
||||
TeamBudgetSnapshot,
|
||||
|
|
@ -281,6 +284,11 @@ def _logged_fields(value: object) -> Mapping[str, object]:
|
|||
return _NO_LOGGED_FIELDS
|
||||
|
||||
|
||||
def _response_cost_of(logged: Mapping[str, object]) -> float:
|
||||
cost: Final = logged.get("response_cost")
|
||||
return float(cost) if isinstance(cost, (int, float)) else 0.0
|
||||
|
||||
|
||||
_GOOGLE_PASS_THROUGH_PROVIDERS: Final = frozenset(
|
||||
(litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value)
|
||||
)
|
||||
|
|
@ -365,6 +373,104 @@ def _get_budget_metrics_per_request_timeout() -> float:
|
|||
return parsed
|
||||
|
||||
|
||||
def _stored_spend(value: object) -> float:
|
||||
return float(value) if isinstance(value, (int, float)) else 0.0
|
||||
|
||||
|
||||
_OWN_CHARGE_MATCH_RATIO: Final = 1 - 1e-9
|
||||
|
||||
|
||||
def _committed_floor(stored: float, committed: float | None) -> float:
|
||||
return stored if committed is None else max(stored, committed)
|
||||
|
||||
|
||||
def _customer_spend_after_request(cached_spend: float, committed_spend: float | None, response_cost: float) -> float:
|
||||
if committed_spend is None:
|
||||
return cached_spend + response_cost
|
||||
return max(committed_spend + response_cost, cached_spend)
|
||||
|
||||
|
||||
def _budgeted_org_counter(org_id: str | None, org: OrgBudgetSnapshot | None) -> EntitySpendCounterKeys:
|
||||
if org is None or org.max_budget is None:
|
||||
return EntitySpendCounterKeys()
|
||||
return EntitySpendCounterKeys.of(token=None, team_id=None, user_id=None, org_id=org_id, end_user_id=None)
|
||||
|
||||
|
||||
async def _committed_spend_now(counter_keys: EntitySpendCounterKeys, timeout: float) -> CommittedBudgetSpend:
|
||||
if not counter_keys.tracked():
|
||||
return CommittedBudgetSpend()
|
||||
try:
|
||||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
from litellm.proxy.spend_tracking.committed_spend import read_committed_budget_spend
|
||||
|
||||
return await asyncio.wait_for(
|
||||
read_committed_budget_spend(spend_counter_store, counter_keys),
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend
|
||||
verbose_logger.debug("[Non-Blocking] Prometheus: committed spend read failed: %s", e)
|
||||
return CommittedBudgetSpend()
|
||||
|
||||
|
||||
OwnCharge: TypeAlias = Literal["unknown", "pending", "counted"]
|
||||
|
||||
|
||||
def _own_charge() -> OwnCharge:
|
||||
"""Whether this request's committed spend charge is in Redis yet, known only inside a request with Redis."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
except ImportError:
|
||||
return "unknown"
|
||||
batches: Final = active_request_redis_batches()
|
||||
if batches is None or not spend_counter_store.has_redis:
|
||||
return "unknown"
|
||||
settlements: Final = batches.post_call_settlements
|
||||
return "counted" if settlements and all(settlement.settled for settlement in settlements) else "pending"
|
||||
|
||||
|
||||
def _committed_spend_before_request(
|
||||
at_auth: CommittedBudgetSpend, now: CommittedBudgetSpend, response_cost: float, own_charge: OwnCharge
|
||||
) -> CommittedBudgetSpend:
|
||||
def before(at_auth_spend: float | None, now_spend: float | None) -> float | None:
|
||||
if now_spend is None:
|
||||
return at_auth_spend
|
||||
match own_charge:
|
||||
case "pending":
|
||||
return now_spend
|
||||
case "counted":
|
||||
without_own_charge: Final = now_spend - response_cost
|
||||
return without_own_charge if at_auth_spend is None else max(at_auth_spend, without_own_charge)
|
||||
case "unknown":
|
||||
if at_auth_spend is None:
|
||||
return now_spend - response_cost
|
||||
own_charge_counted: Final = now_spend - at_auth_spend >= response_cost * _OWN_CHARGE_MATCH_RATIO
|
||||
return max(at_auth_spend, now_spend - response_cost) if own_charge_counted else now_spend
|
||||
case _:
|
||||
assert_never(own_charge)
|
||||
|
||||
return CommittedBudgetSpend(
|
||||
key=before(at_auth.key, now.key),
|
||||
team=before(at_auth.team, now.team),
|
||||
user=before(at_auth.user, now.user),
|
||||
org=before(at_auth.org, now.org),
|
||||
end_user=before(at_auth.end_user, now.end_user),
|
||||
)
|
||||
|
||||
|
||||
async def _committed_spend_without_own_charge(
|
||||
counter_keys: EntitySpendCounterKeys,
|
||||
at_auth: CommittedBudgetSpend | None,
|
||||
response_cost: float,
|
||||
read_timeout: float,
|
||||
) -> CommittedBudgetSpend:
|
||||
own_charge_before_read: Final = _own_charge()
|
||||
now: Final = await _committed_spend_now(counter_keys, read_timeout)
|
||||
own_charge: Final = own_charge_before_read if _own_charge() == own_charge_before_read else "unknown"
|
||||
return _committed_spend_before_request(
|
||||
at_auth=now if at_auth is None else at_auth, now=now, response_cost=response_cost, own_charge=own_charge
|
||||
)
|
||||
|
||||
|
||||
def _get_proxy_llm_router() -> Router | None:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
|
@ -2229,6 +2335,10 @@ class PrometheusLogger(CustomLogger):
|
|||
_user_max_budget: Final = _metadata.get("user_api_key_user_max_budget", None)
|
||||
_user_email: Final = _metadata.get("user_api_key_user_email", None)
|
||||
_org_alias: Final = _metadata.get("user_api_key_org_alias", None)
|
||||
_carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata)
|
||||
_committed_at_auth: Final = committed_budget_spend_of_row(
|
||||
_carried_metadata, _REQUEST_METADATA_ADAPTER.validate_python(litellm_params)
|
||||
)
|
||||
|
||||
# Bound the per-request budget-metric emission so that slow Redis/DB
|
||||
# lookups under load cannot consume the whole LoggingWorker watchdog
|
||||
|
|
@ -2238,45 +2348,65 @@ class PrometheusLogger(CustomLogger):
|
|||
# so dropping one slow per-request emission only loses sub-cron real-time
|
||||
# detail, not correctness.
|
||||
budget_metrics_timeout: Final = _get_budget_metrics_per_request_timeout()
|
||||
gather_coro: Final = asyncio.gather(
|
||||
self._set_api_key_budget_metrics_after_api_request(
|
||||
user_api_key=user_api_key,
|
||||
user_api_key_alias=user_api_key_alias,
|
||||
committed_read_timeout: Final = budget_metrics_timeout / 2
|
||||
|
||||
async def emit_budget_metrics() -> tuple[object, ...]:
|
||||
committed: Final = await _committed_spend_without_own_charge(
|
||||
counter_keys=_committed_at_auth.counter_keys
|
||||
if _committed_at_auth is not None
|
||||
else EntitySpendCounterKeys.of(
|
||||
token=user_api_key,
|
||||
team_id=user_api_team,
|
||||
user_id=user_id,
|
||||
org_id=user_api_key_org_id,
|
||||
end_user_id=end_user_id,
|
||||
),
|
||||
at_auth=_committed_at_auth,
|
||||
response_cost=response_cost,
|
||||
key_max_budget=_api_key_max_budget,
|
||||
key_spend=_api_key_spend,
|
||||
carried=KeyBudgetSnapshot.from_metadata(_metadata),
|
||||
),
|
||||
self._set_team_budget_metrics_after_api_request(
|
||||
user_api_team=user_api_team,
|
||||
user_api_team_alias=user_api_team_alias,
|
||||
team_spend=_team_spend,
|
||||
team_max_budget=_team_max_budget,
|
||||
response_cost=response_cost,
|
||||
carried=TeamBudgetSnapshot.from_metadata(_metadata),
|
||||
),
|
||||
self._set_user_budget_metrics_after_api_request(
|
||||
user_id=user_id,
|
||||
user_spend=_user_spend,
|
||||
user_max_budget=_user_max_budget,
|
||||
response_cost=response_cost,
|
||||
carried=UserBudgetSnapshot.from_metadata(_metadata),
|
||||
user_email=_user_email if isinstance(_user_email, str) else None,
|
||||
),
|
||||
self._set_org_budget_metrics_after_api_request(
|
||||
org_id=user_api_key_org_id,
|
||||
response_cost=response_cost,
|
||||
carried=OrgBudgetSnapshot.from_metadata(_metadata),
|
||||
org_alias=_org_alias if isinstance(_org_alias, str) else None,
|
||||
),
|
||||
self._set_customer_budget_metrics_after_api_request(
|
||||
end_user_id=end_user_id,
|
||||
response_cost=response_cost,
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
read_timeout=committed_read_timeout,
|
||||
)
|
||||
return await asyncio.gather(
|
||||
self._set_api_key_budget_metrics_after_api_request(
|
||||
user_api_key=user_api_key,
|
||||
user_api_key_alias=user_api_key_alias,
|
||||
response_cost=response_cost,
|
||||
key_max_budget=_api_key_max_budget,
|
||||
key_spend=_committed_floor(_stored_spend(_api_key_spend), committed.key),
|
||||
carried=KeyBudgetSnapshot.from_metadata(_carried_metadata),
|
||||
),
|
||||
self._set_team_budget_metrics_after_api_request(
|
||||
user_api_team=user_api_team,
|
||||
user_api_team_alias=user_api_team_alias,
|
||||
team_spend=_committed_floor(_stored_spend(_team_spend), committed.team),
|
||||
team_max_budget=_team_max_budget,
|
||||
response_cost=response_cost,
|
||||
carried=TeamBudgetSnapshot.from_metadata(_carried_metadata),
|
||||
),
|
||||
self._set_user_budget_metrics_after_api_request(
|
||||
user_id=user_id,
|
||||
user_spend=_committed_floor(_stored_spend(_user_spend), committed.user),
|
||||
user_max_budget=_user_max_budget,
|
||||
response_cost=response_cost,
|
||||
carried=UserBudgetSnapshot.from_metadata(_carried_metadata),
|
||||
user_email=_user_email if isinstance(_user_email, str) else None,
|
||||
),
|
||||
self._set_org_budget_metrics_after_api_request(
|
||||
org_id=user_api_key_org_id,
|
||||
response_cost=response_cost,
|
||||
carried=OrgBudgetSnapshot.from_metadata(_carried_metadata),
|
||||
org_alias=_org_alias if isinstance(_org_alias, str) else None,
|
||||
committed_spend=committed.org,
|
||||
),
|
||||
self._set_customer_budget_metrics_after_api_request(
|
||||
end_user_id=end_user_id,
|
||||
response_cost=response_cost,
|
||||
committed_spend=committed.end_user,
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
try:
|
||||
results: Final = await asyncio.wait_for(gather_coro, timeout=budget_metrics_timeout)
|
||||
results: Final = await asyncio.wait_for(emit_budget_metrics(), timeout=budget_metrics_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
verbose_logger.debug(
|
||||
"[Non-Blocking] Prometheus: per-request budget metric emission "
|
||||
|
|
@ -2745,7 +2875,7 @@ class PrometheusLogger(CustomLogger):
|
|||
_labels,
|
||||
)
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
verbose_logger.debug(
|
||||
"prometheus Logging - Enters failure logging function (kwargs keys: %s)",
|
||||
list(kwargs.keys()) if isinstance(kwargs, dict) else type(kwargs).__name__,
|
||||
|
|
@ -2799,9 +2929,27 @@ class PrometheusLogger(CustomLogger):
|
|||
label_context=PrometheusLabelFactoryContext(enum_values),
|
||||
)
|
||||
self.set_llm_deployment_failure_metrics(kwargs)
|
||||
failure_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(
|
||||
get_litellm_metadata_from_kwargs(kwargs)
|
||||
)
|
||||
failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias")
|
||||
failure_org: Final = OrgBudgetSnapshot.from_metadata(failure_metadata)
|
||||
failure_cost: Final = _response_cost_of(_LOGGED_FIELDS.validate_python(standard_logging_payload))
|
||||
failure_committed: Final = await _committed_spend_without_own_charge(
|
||||
counter_keys=_budgeted_org_counter(user_api_key_org_id, failure_org),
|
||||
at_auth=committed_budget_spend_of_row(
|
||||
failure_metadata, _REQUEST_METADATA_ADAPTER.validate_python(litellm_params)
|
||||
)
|
||||
or CommittedBudgetSpend(),
|
||||
response_cost=failure_cost,
|
||||
read_timeout=_get_budget_metrics_per_request_timeout(),
|
||||
)
|
||||
await self._set_org_budget_metrics_after_api_request(
|
||||
org_id=user_api_key_org_id,
|
||||
response_cost=0,
|
||||
response_cost=failure_cost,
|
||||
carried=failure_org,
|
||||
org_alias=failure_org_alias if isinstance(failure_org_alias, str) else None,
|
||||
committed_spend=failure_committed.org,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception("prometheus Layer Error(): Exception occured - %s", e)
|
||||
|
|
@ -4516,6 +4664,7 @@ class PrometheusLogger(CustomLogger):
|
|||
response_cost: float,
|
||||
carried: OrgBudgetSnapshot | None = None,
|
||||
org_alias: str | None = None,
|
||||
committed_spend: float | None = None,
|
||||
):
|
||||
"""
|
||||
Set org budget metrics after an LLM API request
|
||||
|
|
@ -4533,7 +4682,7 @@ class PrometheusLogger(CustomLogger):
|
|||
self._set_org_budget_metrics(
|
||||
org_id=org_id,
|
||||
org_alias=org_alias or "",
|
||||
spend=carried.spend + response_cost,
|
||||
spend=_committed_floor(carried.spend, committed_spend) + response_cost,
|
||||
max_budget=carried.max_budget,
|
||||
budget_reset_at=None,
|
||||
)
|
||||
|
|
@ -4559,7 +4708,7 @@ class PrometheusLogger(CustomLogger):
|
|||
if org_info is None:
|
||||
return
|
||||
|
||||
_total_org_spend: Final = (org_info.spend or 0.0) + response_cost
|
||||
_total_org_spend: Final = _committed_floor(org_info.spend or 0.0, committed_spend) + response_cost
|
||||
budget_table: Final = org_info.litellm_budget_table
|
||||
max_budget: Final = budget_table.max_budget if budget_table else None
|
||||
budget_reset_at: Final = getattr(budget_table, "budget_reset_at", None) if budget_table else None
|
||||
|
|
@ -4626,6 +4775,7 @@ class PrometheusLogger(CustomLogger):
|
|||
self,
|
||||
end_user_id: str | None,
|
||||
response_cost: float,
|
||||
committed_spend: float | None = None,
|
||||
):
|
||||
if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled():
|
||||
return
|
||||
|
|
@ -4651,7 +4801,9 @@ class PrometheusLogger(CustomLogger):
|
|||
budget_table: Final = cached_customer.litellm_budget_table
|
||||
self._set_customer_budget_metrics(
|
||||
end_user_id=end_user_id,
|
||||
spend=cached_customer.spend + response_cost,
|
||||
spend=_customer_spend_after_request(
|
||||
cached_spend=cached_customer.spend, committed_spend=committed_spend, response_cost=response_cost
|
||||
),
|
||||
max_budget=budget_table.max_budget if budget_table is not None else None,
|
||||
budget_reset_at=None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import secrets
|
|||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NamedTuple, Protocol, Union, cast
|
||||
|
||||
|
|
@ -170,7 +171,15 @@ from litellm.proxy.db.db_lookup_gate import bounded_db_lookup
|
|||
from litellm.proxy.db.db_span import db_span
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state
|
||||
from litellm.proxy.spend_tracking.carried_budget_state import (
|
||||
carry_committed_budget_spend,
|
||||
carry_team_and_user_budget_state,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.committed_spend import (
|
||||
bind_committed_admission_counters,
|
||||
committed_spend_tracked,
|
||||
read_admission_committed_spend,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
bind_admission_counter_keys,
|
||||
release_spend_counter_batch,
|
||||
|
|
@ -3650,6 +3659,7 @@ async def run_centralized_common_checks(
|
|||
model_max_budget_limiter,
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
spend_counter_store,
|
||||
user_api_key_cache,
|
||||
user_custom_auth,
|
||||
)
|
||||
|
|
@ -3960,6 +3970,7 @@ async def run_centralized_common_checks(
|
|||
)
|
||||
|
||||
bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id)
|
||||
bind_committed_admission_counters(user_api_key_auth_obj, end_user_id=end_user_id)
|
||||
try:
|
||||
_ = await common_checks(
|
||||
request=request,
|
||||
|
|
@ -4035,6 +4046,18 @@ async def run_centralized_common_checks(
|
|||
user_obj=user_object,
|
||||
)
|
||||
|
||||
if committed_spend_tracked():
|
||||
await carry_committed_budget_spend(
|
||||
valid_token=user_api_key_auth_obj,
|
||||
team_object=typed_team_object,
|
||||
user_object=user_object,
|
||||
end_user_object=end_user_object,
|
||||
read_committed=(
|
||||
None
|
||||
if route_skips_budget_checks(route=route)
|
||||
else partial(read_admission_committed_spend, spend_counter_store)
|
||||
),
|
||||
)
|
||||
await _reserve_budget_after_common_checks(
|
||||
user_api_key_auth_obj=user_api_key_auth_obj,
|
||||
request=request,
|
||||
|
|
@ -4348,6 +4371,7 @@ async def _prefetch_referenced_auth_objects(
|
|||
"""Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold).
|
||||
Runs after the key's model access check so a denied request costs no more than it did before."""
|
||||
bind_admission_counter_keys(valid_token, end_user_id=end_user_id or None)
|
||||
bind_committed_admission_counters(valid_token, end_user_id=end_user_id or None)
|
||||
await prefetch_auth_objects(
|
||||
refs=AuthObjectRefs.from_token(valid_token),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
|
|||
from litellm.proxy.db.db_span import db_span, db_spanned
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
|
||||
from litellm.proxy.spend_tracking.committed_spend import committed_spend_keys
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.redis.stores.spend_counters import (
|
||||
SPEND_COUNTERS_TARGET,
|
||||
|
|
@ -621,7 +622,7 @@ class ResetBudgetJob:
|
|||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
|
||||
try:
|
||||
await spend_counter_store.delete((counter_key,))
|
||||
await spend_counter_store.delete((counter_key, *committed_spend_keys((counter_key,))))
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset spend counter %s in Redis: %s. Budget may be over-enforced until counter expires.",
|
||||
|
|
@ -680,7 +681,7 @@ class ResetBudgetJob:
|
|||
|
||||
match cache:
|
||||
case "spend counter":
|
||||
await spend_counter_store.delete(keys)
|
||||
await spend_counter_store.delete((*keys, *committed_spend_keys(keys)))
|
||||
case "user_api_key_cache":
|
||||
await user_api_key_cache.async_delete_cache_keys(keys)
|
||||
case _:
|
||||
|
|
|
|||
|
|
@ -94,6 +94,10 @@ class SpendCounterReseed:
|
|||
_locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
|
||||
_registry_lock: ClassVar[asyncio.Lock | None] = None
|
||||
|
||||
@staticmethod
|
||||
async def counter_lock(counter_key: str) -> asyncio.Lock:
|
||||
return await SpendCounterReseed._get_lock(counter_key)
|
||||
|
||||
@staticmethod
|
||||
async def _get_lock(counter_key: str) -> asyncio.Lock:
|
||||
if SpendCounterReseed._registry_lock is None:
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ from litellm.proxy.db.db_spend_update_writer import (
|
|||
from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, warm_committed_counter_keys
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope
|
||||
from litellm.proxy.spend_tracking.spend_event import (
|
||||
ObjectMapping,
|
||||
|
|
@ -75,6 +76,8 @@ from litellm.proxy.utils import (
|
|||
BilledFailedAttempt,
|
||||
ProxyUpdateSpend,
|
||||
)
|
||||
from litellm.types.proxy.auth.user_api_key_auth import committed_budget_spend_of_row
|
||||
from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys
|
||||
from litellm.types.router import Deployment
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
|
|
@ -268,6 +271,11 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter:
|
|||
return proxy_logging_obj.db_spend_update_writer
|
||||
|
||||
|
||||
def _committed_spend_scope(user_api_key_dict: UserAPIKeyAuth) -> EntitySpendCounterKeys:
|
||||
committed: Final = user_api_key_dict.committed_budget_spend
|
||||
return committed.counter_keys if committed is not None else EntitySpendCounterKeys()
|
||||
|
||||
|
||||
def _proxy_spend_counters() -> "_IncrementSpendCounters":
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
|
|
@ -639,6 +647,7 @@ class _ProxyDBLogger(CustomLogger):
|
|||
team_id=user_api_key_dict.team_id,
|
||||
org_id=user_api_key_dict.org_id,
|
||||
project_id=user_api_key_dict.project_id,
|
||||
committed_spend_scope=_committed_spend_scope(user_api_key_dict),
|
||||
kwargs=dict(row_kwargs),
|
||||
completion_response=original_exception,
|
||||
start_time=start_time,
|
||||
|
|
@ -757,6 +766,7 @@ class _ProxyDBLogger(CustomLogger):
|
|||
router=get_llm_router(),
|
||||
),
|
||||
project_id=user_api_key_dict.project_id,
|
||||
committed_spend_scope=_committed_spend_scope(user_api_key_dict),
|
||||
**_ledger_counter_options(row_kwargs.get(ACCOUNTING_LEDGER_KWARG)),
|
||||
)
|
||||
except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing
|
||||
|
|
@ -1269,10 +1279,38 @@ class _IncrementSpendCounters(Protocol):
|
|||
request_started_at: datetime | None = None,
|
||||
model_access_groups: Sequence[str] | None = None,
|
||||
project_id: str | None = None,
|
||||
committed_spend_scope: EntitySpendCounterKeys | None = None,
|
||||
windows_from_ledger: bool = False,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
async def _committed_spend_scope_of_row(
|
||||
kwargs: Mapping[str, object],
|
||||
response_cost: float,
|
||||
token: str | None,
|
||||
team_id: str | None,
|
||||
user_id: str | None,
|
||||
org_id: str | None,
|
||||
end_user_id: str | None,
|
||||
) -> EntitySpendCounterKeys:
|
||||
metadata: Final = _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=dict(kwargs)))
|
||||
carried: Final = committed_budget_spend_of_row(
|
||||
metadata, _OBJECT_MAPPING.validate_python(kwargs.get("litellm_params") or {})
|
||||
)
|
||||
if carried is not None:
|
||||
return carried.counter_keys
|
||||
if response_cost <= 0 or not committed_spend_tracked():
|
||||
return EntitySpendCounterKeys()
|
||||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
|
||||
return await warm_committed_counter_keys(
|
||||
spend_counter_store,
|
||||
EntitySpendCounterKeys.of(
|
||||
token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _update_database_and_spend_counters(
|
||||
proxy_logging_obj: "ProxyLogging",
|
||||
increment_spend_counters: _IncrementSpendCounters,
|
||||
|
|
@ -1291,6 +1329,7 @@ async def _update_database_and_spend_counters(
|
|||
model_access_groups: Sequence[str] | None = None,
|
||||
project_id: str | None = None,
|
||||
spend_writer: DBSpendUpdateWriter | None = None,
|
||||
committed_spend_scope: EntitySpendCounterKeys | None = None,
|
||||
) -> bool:
|
||||
"""The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then
|
||||
spans the database write and the counter update, so the post-call counters are read with a single MGET after the
|
||||
|
|
@ -1333,6 +1372,7 @@ async def _update_database_and_spend_counters(
|
|||
request_tags=request_tags,
|
||||
model_access_groups=model_access_groups,
|
||||
project_id=project_id,
|
||||
committed_spend_scope=committed_spend_scope,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1354,6 +1394,7 @@ async def _update_database_and_spend_counters_in_batch(
|
|||
model_access_groups: Sequence[str] | None,
|
||||
project_id: str | None,
|
||||
spend_writer: DBSpendUpdateWriter | None = None,
|
||||
committed_spend_scope: EntitySpendCounterKeys | None = None,
|
||||
) -> bool:
|
||||
try:
|
||||
writer: Final = spend_writer if spend_writer is not None else proxy_logging_obj.db_spend_update_writer
|
||||
|
|
@ -1407,6 +1448,19 @@ async def _update_database_and_spend_counters_in_batch(
|
|||
request_started_at=start_time,
|
||||
model_access_groups=model_access_groups,
|
||||
project_id=project_id,
|
||||
committed_spend_scope=(
|
||||
committed_spend_scope
|
||||
if committed_spend_scope is not None
|
||||
else await _committed_spend_scope_of_row(
|
||||
kwargs=kwargs,
|
||||
response_cost=response_cost,
|
||||
token=user_api_key,
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
org_id=org_id,
|
||||
end_user_id=end_user_id,
|
||||
)
|
||||
),
|
||||
**_ledger_counter_options(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get(ACCOUNTING_LEDGER_KWARG)),
|
||||
)
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -2156,7 +2156,14 @@ class LiteLLMProxyRequestSetup:
|
|||
user_api_key_dict.matched_model_access_groups
|
||||
)
|
||||
# UserAPIKeyAuth object for MCP server access control
|
||||
data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy(
|
||||
data[_metadata_variable_name]["user_api_key_auth"] = LiteLLMProxyRequestSetup.auth_for_request_metadata(
|
||||
user_api_key_dict
|
||||
)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def auth_for_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth:
|
||||
return user_api_key_dict.model_copy(
|
||||
update={
|
||||
"metadata": strip_callback_config(user_api_key_dict.metadata),
|
||||
"team_metadata": strip_callback_config(user_api_key_dict.team_metadata),
|
||||
|
|
@ -2164,7 +2171,6 @@ class LiteLLMProxyRequestSetup:
|
|||
"organization_metadata": strip_callback_config(user_api_key_dict.organization_metadata),
|
||||
}
|
||||
)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def add_management_endpoint_metadata_to_request_metadata(
|
||||
|
|
|
|||
|
|
@ -42,7 +42,9 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa:
|
|||
handle_update_object_permission_common,
|
||||
set_object_permission,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
|
||||
from litellm.proxy.utils import handle_exception_on_proxy
|
||||
from litellm.redis.stores.spend_counters import end_user_counter
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.chunked_in import find_many_in
|
||||
from litellm.repositories.table_repositories import EndUserRepository
|
||||
|
|
@ -812,7 +814,7 @@ async def delete_end_user(
|
|||
See below for all params
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
from litellm.proxy.proxy_server import prisma_client, spend_counter_store
|
||||
|
||||
try:
|
||||
if prisma_client is None:
|
||||
|
|
@ -842,6 +844,9 @@ async def delete_end_user(
|
|||
verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response)
|
||||
|
||||
await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids))
|
||||
await drop_committed_spend(
|
||||
spend_counter_store, tuple(end_user_counter(user_id) for user_id in data.user_ids)
|
||||
)
|
||||
|
||||
return DeleteCustomersResponse(
|
||||
deleted_customers=response,
|
||||
|
|
|
|||
|
|
@ -1483,8 +1483,11 @@ async def _invalidate_user_spend_counter_if_changed(
|
|||
"""
|
||||
if non_default_values.get("spend") is not None:
|
||||
from litellm.proxy.proxy_server import invalidate_spend_counter
|
||||
from litellm.proxy.spend_tracking.committed_spend import committed_spend_key
|
||||
|
||||
await invalidate_spend_counter(counter_key=user_counter(str(non_default_values["user_id"])))
|
||||
counter_key: Final = user_counter(str(non_default_values["user_id"]))
|
||||
await invalidate_spend_counter(counter_key=counter_key)
|
||||
await invalidate_spend_counter(counter_key=committed_spend_key(counter_key))
|
||||
|
||||
|
||||
def _clears_object_permission(user_request: UpdateUserRequest) -> bool:
|
||||
|
|
|
|||
|
|
@ -158,6 +158,7 @@ from litellm.proxy.mcp_server.outbound_credentials.sso_assertion_store import (
|
|||
)
|
||||
from litellm.proxy.search_endpoints.search_tool_registry import rotate_search_tools_master_key
|
||||
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
|
||||
from litellm.proxy.spend_tracking.committed_spend import committed_spend_key
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports
|
||||
_is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
is_master_key,
|
||||
|
|
@ -4288,7 +4289,7 @@ async def update_key_fn(
|
|||
)
|
||||
|
||||
if data.spend is not None:
|
||||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
from litellm.proxy.proxy_server import invalidate_spend_counter, spend_counter_store
|
||||
|
||||
counter_key: Final = key_counter(hash_token_if_needed(key))
|
||||
try:
|
||||
|
|
@ -4300,6 +4301,7 @@ async def update_key_fn(
|
|||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
await invalidate_spend_counter(counter_key=committed_spend_key(counter_key))
|
||||
|
||||
asyncio.create_task(
|
||||
KeyManagementEventHooks.async_key_updated_hook(
|
||||
|
|
@ -6967,6 +6969,7 @@ async def reset_key_spend_fn(
|
|||
try:
|
||||
from litellm.proxy.proxy_server import (
|
||||
hash_token,
|
||||
invalidate_spend_counter,
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
set_spend_counter_with_floor_and_broadcast,
|
||||
|
|
@ -7022,6 +7025,7 @@ async def reset_key_spend_fn(
|
|||
# pre-reset value.
|
||||
_counter_key: Final = key_counter(hashed_api_key)
|
||||
await set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to)
|
||||
await invalidate_spend_counter(counter_key=committed_spend_key(_counter_key))
|
||||
await _reset_key_budget_windows(
|
||||
prisma_client=prisma_client,
|
||||
hashed_api_key=hashed_api_key,
|
||||
|
|
|
|||
|
|
@ -66,7 +66,9 @@ from litellm.proxy.management_helpers.utils import (
|
|||
get_new_internal_user_defaults,
|
||||
management_endpoint_wrapper,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.redis.stores.spend_counters import org_counter, team_counter
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.chunked_in import find_many_in
|
||||
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
|
||||
|
|
@ -93,6 +95,7 @@ if TYPE_CHECKING:
|
|||
LiteLLM_OrganizationMembership as PrismaOrganizationMembership,
|
||||
)
|
||||
from prisma.models import LiteLLM_OrganizationTable as PrismaOrganizationTable
|
||||
from prisma.models import LiteLLM_TeamTable as PrismaTeamTable
|
||||
from prisma.models import LiteLLM_UserTable as PrismaUserTable
|
||||
from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken
|
||||
|
||||
|
|
@ -180,6 +183,8 @@ class _OrganizationMembershipTableClient(Protocol):
|
|||
|
||||
|
||||
class _TeamTableClient(Protocol):
|
||||
async def find_many(self, *, where: Mapping[str, object]) -> "Sequence[PrismaTeamTable]": ...
|
||||
|
||||
async def delete_many(self, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
|
|
@ -819,6 +824,11 @@ 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 spend_counter_store
|
||||
|
||||
await drop_committed_spend(spend_counter_store, (org_counter(data.organization_id),))
|
||||
|
||||
return response
|
||||
|
||||
|
||||
|
|
@ -1045,7 +1055,8 @@ async def delete_organization(
|
|||
teams_in_deleted_orgs: Final = await find_many_in(
|
||||
_table(TeamRepository(prisma_client)), "organization_id", data.organization_ids
|
||||
)
|
||||
raise_if_teams_hold_ptu_shares(tuple(team.team_id for team in teams_in_deleted_orgs), llm_router)
|
||||
team_ids_in_deleted_orgs: Final = tuple(team.team_id for team in teams_in_deleted_orgs)
|
||||
raise_if_teams_hold_ptu_shares(team_ids_in_deleted_orgs, llm_router)
|
||||
|
||||
deleted_orgs: Final = []
|
||||
for organization_id in data.organization_ids:
|
||||
|
|
@ -1073,6 +1084,15 @@ async def delete_organization(
|
|||
)
|
||||
deleted_orgs.append(deleted_org)
|
||||
|
||||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
|
||||
await drop_committed_spend(
|
||||
spend_counter_store,
|
||||
(
|
||||
*(org_counter(organization_id) for organization_id in data.organization_ids),
|
||||
*(team_counter(team_id) for team_id in team_ids_in_deleted_orgs),
|
||||
),
|
||||
)
|
||||
return deleted_orgs
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -226,12 +226,14 @@ from litellm.proxy.management_helpers.utils import (
|
|||
add_new_member,
|
||||
management_endpoint_wrapper,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
|
||||
from litellm.proxy.utils import (
|
||||
PrismaClient,
|
||||
ProxyLogging,
|
||||
handle_exception_on_proxy,
|
||||
query_raw_within_statement_timeout,
|
||||
)
|
||||
from litellm.redis.stores.spend_counters import team_counter
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.chunked_in import delete_many_in, find_many_in
|
||||
from litellm.repositories.daily_activity_buckets import api_key_from_bucket
|
||||
|
|
@ -4863,6 +4865,7 @@ async def delete_validated_team_rows(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
proxy_logging_obj,
|
||||
spend_counter_store,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
|
|
@ -4934,7 +4937,9 @@ async def delete_validated_team_rows(
|
|||
_emit_team_members_metric(deleted_team.model_copy(update={"members_with_roles": ()}))
|
||||
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
|
||||
|
||||
return DeletedTeams(team_ids=tuple(team.team_id for team in locked))
|
||||
deleted_team_ids: Final = tuple(team.team_id for team in locked)
|
||||
await drop_committed_spend(spend_counter_store, tuple(team_counter(team_id) for team_id in deleted_team_ids))
|
||||
return DeletedTeams(team_ids=deleted_team_ids)
|
||||
|
||||
|
||||
def _locked_roster_user_ids(locked: Sequence[LiteLLM_TeamTable]) -> tuple[str, ...]:
|
||||
|
|
|
|||
|
|
@ -46,7 +46,9 @@ from litellm.proxy.management_helpers.access_group_team_sync import (
|
|||
TEAM_ADVISORY_LOCK_SQL,
|
||||
USER_ROWS_LOCK_SQL,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.redis.stores.spend_counters import user_counter
|
||||
from litellm.repositories.table_repositories import (
|
||||
OrganizationMembershipRepository,
|
||||
TeamMembershipRepository,
|
||||
|
|
@ -585,6 +587,9 @@ async def _after_users_deleted(
|
|||
)
|
||||
await evict_and_broadcast(cache_keys=deletion.jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
|
||||
await evict_and_broadcast(cache_keys=sorted(user_ids), user_api_key_cache=user_api_key_cache)
|
||||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
|
||||
await drop_committed_spend(spend_counter_store, tuple(user_counter(user_id) for user_id in sorted(user_ids)))
|
||||
for removal in deletion.removals.values():
|
||||
_emit_team_members_metric(removal.team)
|
||||
|
||||
|
|
|
|||
|
|
@ -165,6 +165,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
WebSocketSessionUsage,
|
||||
WebsocketUpstreamFrameKind,
|
||||
)
|
||||
from litellm.types.proxy.auth.user_api_key_auth import COMMITTED_BUDGET_SPEND_PARAM
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend
|
||||
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage
|
||||
|
||||
from .llm_provider_handlers.tinyfish_passthrough_logging_handler import (
|
||||
|
|
@ -779,6 +781,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
"litellm_params": {
|
||||
**litellm_params_in_body,
|
||||
"metadata": _metadata,
|
||||
COMMITTED_BUDGET_SPEND_PARAM: user_api_key_dict.committed_budget_spend or CommittedBudgetSpend(),
|
||||
"proxy_server_request": {
|
||||
"url": str(request.url),
|
||||
"method": request.method,
|
||||
|
|
|
|||
|
|
@ -93,6 +93,7 @@ from litellm.constants import (
|
|||
LITELLM_UI_SESSION_DURATION,
|
||||
OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS,
|
||||
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
|
||||
SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS,
|
||||
VLLM_CHOICE_FIELDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.access_windows import access_windows_config_error
|
||||
|
|
@ -338,7 +339,7 @@ from litellm._logging import (
|
|||
verbose_router_logger,
|
||||
)
|
||||
from litellm.caching.caching import DualCache, RedisCache
|
||||
from litellm.caching.redis_batch import drain_post_call_redis_batches
|
||||
from litellm.caching.redis_batch import active_request_redis_batches, drain_post_call_redis_batches
|
||||
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.constants import (
|
||||
|
|
@ -995,12 +996,20 @@ from litellm.proxy.spend_tracking.budget_reservation import (
|
|||
release_unbound_budget_reservation,
|
||||
stamp_budget_reservation_actual_cost,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.committed_spend import (
|
||||
committed_spend_key,
|
||||
committed_spend_keys,
|
||||
has_committed_twin,
|
||||
is_committed_spend_key,
|
||||
seed_committed_spend,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_capture_rate import (
|
||||
run_scheduled_spend_capture_rate_check,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
PendingSpendIncrement,
|
||||
active_spend_counter_batch,
|
||||
bind_spend_counter_keys,
|
||||
forget_spend_counter,
|
||||
post_call_counter_keys,
|
||||
read_batched_spend_counter,
|
||||
|
|
@ -1105,6 +1114,7 @@ from litellm.types.llms.openai import (
|
|||
ResponseInProgressEvent,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys
|
||||
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelGroupInfoProxy,
|
||||
|
|
@ -3240,6 +3250,7 @@ proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME
|
|||
proxy_budget_rescheduler_max_time = PROXY_BUDGET_RESCHEDULER_MAX_TIME
|
||||
proxy_batch_polling_interval = PROXY_BATCH_POLLING_INTERVAL
|
||||
proxy_batch_write_at = PROXY_BATCH_WRITE_AT
|
||||
_BATCH_WRITE_JITTER_MAX_SECONDS: Final = 5
|
||||
proxy_config_reload_interval_seconds = PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
|
||||
litellm_master_key_hash = None
|
||||
disable_spend_logs = False
|
||||
|
|
@ -3638,6 +3649,7 @@ async def increment_spend_counters(
|
|||
request_started_at: datetime | None = None,
|
||||
model_access_groups: Sequence[str] | None = None,
|
||||
project_id: str | None = None,
|
||||
committed_spend_scope: EntitySpendCounterKeys | None = None,
|
||||
windows_from_ledger: bool = False,
|
||||
):
|
||||
"""
|
||||
|
|
@ -3675,6 +3687,7 @@ async def increment_spend_counters(
|
|||
request_started_at=request_started_at,
|
||||
model_access_groups=model_access_groups,
|
||||
project_id=project_id,
|
||||
committed_spend_scope=committed_spend_scope,
|
||||
windows_from_ledger=windows_from_ledger,
|
||||
)
|
||||
|
||||
|
|
@ -3790,10 +3803,17 @@ async def _increment_spend_counters_batched(
|
|||
request_started_at: datetime | None,
|
||||
model_access_groups: Sequence[str] | None,
|
||||
project_id: str | None = None,
|
||||
committed_spend_scope: EntitySpendCounterKeys | None = None,
|
||||
windows_from_ledger: bool = False,
|
||||
):
|
||||
"""Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and
|
||||
the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments."""
|
||||
committed_counter_keys: Final = (
|
||||
tuple(counter_key for counter_key in committed_spend_scope.tracked() if has_committed_twin(counter_key))
|
||||
if committed_spend_scope is not None and response_cost is not None and response_cost > 0
|
||||
else ()
|
||||
)
|
||||
bind_spend_counter_keys(frozenset(committed_spend_keys(committed_counter_keys)))
|
||||
reservation_update: Final = await _reconcile_budget_reservation_for_counter_update(
|
||||
budget_reservation=budget_reservation,
|
||||
response_cost=response_cost,
|
||||
|
|
@ -3980,6 +4000,9 @@ async def _increment_spend_counters_batched(
|
|||
)
|
||||
if project_id is not None
|
||||
else None,
|
||||
_prepare_committed_spend_increments(counter_keys=committed_counter_keys, increment=cost)
|
||||
if committed_counter_keys
|
||||
else None,
|
||||
)
|
||||
if coro is not None
|
||||
)
|
||||
|
|
@ -4082,6 +4105,23 @@ async def _reconcile_budget_reservation_for_counter_update(
|
|||
return _ReservationCounterUpdate(reserved_counter_keys=frozenset(reserved_counter_keys), pending=pending)
|
||||
|
||||
|
||||
async def _committed_spend_seeded(counter_key: str) -> bool:
|
||||
return await _is_spend_counter_cache_warm(committed_spend_key(counter_key)) or await seed_committed_spend(
|
||||
spend_counter_store=spend_counter_store, prisma_client=prisma_client, counter_key=counter_key
|
||||
)
|
||||
|
||||
|
||||
async def _prepare_committed_spend_increments(
|
||||
counter_keys: tuple[str, ...], increment: float
|
||||
) -> tuple[PendingSpendIncrement, ...]:
|
||||
seeded: Final = await asyncio.gather(*(_committed_spend_seeded(counter_key) for counter_key in counter_keys))
|
||||
return tuple(
|
||||
PendingSpendIncrement(counter_key=committed_spend_key(counter_key), increment=increment)
|
||||
for counter_key, is_seeded in zip(counter_keys, seeded)
|
||||
if is_seeded
|
||||
)
|
||||
|
||||
|
||||
async def _prepare_end_user_and_tag_spend_increments(
|
||||
end_user_id: str | None,
|
||||
tags: list[str] | None,
|
||||
|
|
@ -4487,7 +4527,15 @@ def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) ->
|
|||
"""Post-call increments ride the request's native post-call pipeline with the other counters. Each counter's
|
||||
new value lands in memory when the pipeline settles; a failed one is dropped everywhere so no reader trusts a
|
||||
counter whose increment may not have applied, as ``increment_spend_counters_pipeline`` does."""
|
||||
return spend_counter_store.increment_many_post_call(tuple((item.counter_key, item.increment) for item in pending))
|
||||
settlement: Final = spend_counter_store.increment_many_post_call(
|
||||
tuple((item.counter_key, item.increment) for item in pending)
|
||||
)
|
||||
if settlement is None:
|
||||
return False
|
||||
batches: Final = active_request_redis_batches()
|
||||
if batches is not None and any(is_committed_spend_key(item.counter_key) for item in pending):
|
||||
batches.track_post_call_settlement(settlement)
|
||||
return True
|
||||
|
||||
|
||||
async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> tuple[float | None, ...]:
|
||||
|
|
@ -11651,7 +11699,10 @@ class ProxyStartupEvent:
|
|||
)
|
||||
|
||||
# Ensure minimum interval of 30 seconds for batch writing to prevent memory issues
|
||||
batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5)
|
||||
batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, _BATCH_WRITE_JITTER_MAX_SECONDS)
|
||||
spend_counter_store.set_committed_ttl(
|
||||
2 * (SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS)
|
||||
)
|
||||
|
||||
### PROXY WORKER HEARTBEAT ###
|
||||
worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client)
|
||||
|
|
|
|||
|
|
@ -1,14 +1,17 @@
|
|||
"""Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
from litellm.models.organization import LiteLLM_OrganizationTable
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.models.user import LiteLLM_UserTable
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.proxy.carried_budget_state import (
|
||||
CommittedBudgetSpend,
|
||||
EntitySpendCounterKeys,
|
||||
OrgBudgetSnapshot,
|
||||
TeamBudgetSnapshot,
|
||||
UserBudgetSnapshot,
|
||||
|
|
@ -34,6 +37,41 @@ def carry_team_and_user_budget_state(
|
|||
)
|
||||
|
||||
|
||||
def _budgeted_org_id(valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None) -> str | None:
|
||||
org_snapshot: Final = valid_token.org_budget_snapshot
|
||||
if org_snapshot is None or org_snapshot.max_budget is None:
|
||||
return None
|
||||
return valid_token.org_id or (team_object.organization_id if team_object is not None else None)
|
||||
|
||||
|
||||
def _budgeted_end_user_id(end_user_object: LiteLLM_EndUserTable | None) -> str | None:
|
||||
budget_table: Final = end_user_object.litellm_budget_table if end_user_object is not None else None
|
||||
if end_user_object is None or budget_table is None or budget_table.max_budget is None:
|
||||
return None
|
||||
return end_user_object.user_id
|
||||
|
||||
|
||||
async def carry_committed_budget_spend(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
user_object: LiteLLM_UserTable | None,
|
||||
end_user_object: LiteLLM_EndUserTable | None,
|
||||
read_committed: Callable[[EntitySpendCounterKeys], Awaitable[CommittedBudgetSpend]] | None,
|
||||
) -> None:
|
||||
counter_keys: Final = EntitySpendCounterKeys.of(
|
||||
token=valid_token.token if valid_token.max_budget is not None else None,
|
||||
team_id=team_object.team_id if team_object is not None and team_object.max_budget is not None else None,
|
||||
user_id=user_object.user_id if user_object is not None and user_object.max_budget is not None else None,
|
||||
org_id=_budgeted_org_id(valid_token=valid_token, team_object=team_object),
|
||||
end_user_id=_budgeted_end_user_id(end_user_object=end_user_object),
|
||||
)
|
||||
valid_token.committed_budget_spend = ( # rebind-ok: request-scoped credential
|
||||
CommittedBudgetSpend(counter_keys=counter_keys)
|
||||
if read_committed is None
|
||||
else await read_committed(counter_keys)
|
||||
)
|
||||
|
||||
|
||||
def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None:
|
||||
budget_table: Final = org_table.litellm_budget_table
|
||||
valid_token.organization_alias = org_table.organization_alias
|
||||
|
|
|
|||
181
litellm/proxy/spend_tracking/committed_spend.py
Normal file
181
litellm/proxy/spend_tracking/committed_spend.py
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
"""Committed spend counters: twins of the entity spend counters that only ever take charged cost, never budget holds.
|
||||
|
||||
The remaining budget gauges read them, so a request's worst-case reservation never shows as spent. They are kept
|
||||
only for entities with a budget, while a ``PrometheusLogger`` is registered, in the same ``SpendCounterStore`` as the
|
||||
counters they twin.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.redis_batch import Batch
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
admission_counter_keys,
|
||||
bind_spend_counter_keys,
|
||||
read_batched_spend_counter,
|
||||
)
|
||||
from litellm.redis.stores.spend_counters import (
|
||||
COMMITTED_COUNTER_PREFIX,
|
||||
END_USER_COUNTER_PREFIX,
|
||||
SPEND_COUNTERS_TARGET,
|
||||
SpendCounterStore,
|
||||
committed_counter,
|
||||
)
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_BASE_COUNTER_PREFIXES: Final = ("spend:key:", "spend:team:", "spend:user:", "spend:org:", END_USER_COUNTER_PREFIX)
|
||||
|
||||
|
||||
def committed_spend_key(counter_key: str) -> str:
|
||||
return committed_counter(counter_key)
|
||||
|
||||
|
||||
def has_committed_twin(counter_key: str) -> bool:
|
||||
return counter_key.startswith(_BASE_COUNTER_PREFIXES) and ":window:" not in counter_key
|
||||
|
||||
|
||||
def committed_spend_keys(counter_keys: Iterable[str]) -> tuple[str, ...]:
|
||||
return tuple(committed_spend_key(counter_key) for counter_key in counter_keys if has_committed_twin(counter_key))
|
||||
|
||||
|
||||
def is_committed_spend_key(counter_key: str) -> bool:
|
||||
return counter_key.startswith(COMMITTED_COUNTER_PREFIX)
|
||||
|
||||
|
||||
def committed_spend_tracked() -> bool:
|
||||
return litellm.logging_callback_manager.callback_is_active(callback_type=PrometheusLogger)
|
||||
|
||||
|
||||
def bind_committed_admission_counters(token: UserAPIKeyAuth, end_user_id: str | None) -> None:
|
||||
if committed_spend_tracked():
|
||||
bind_spend_counter_keys(frozenset(committed_spend_keys(admission_counter_keys(token, end_user_id))))
|
||||
|
||||
|
||||
async def _read_committed_values(
|
||||
spend_counter_store: SpendCounterStore, twin_keys: tuple[str, ...]
|
||||
) -> Mapping[str, float]:
|
||||
batched: Final = await spend_counter_store.read_many(twin_keys, batch=Batch.NONE)
|
||||
if batched is not None:
|
||||
return {twin_key: value for twin_key, value in batched.items() if value is not None}
|
||||
reads: Final = await asyncio.gather(*(spend_counter_store.read(twin_key) for twin_key in twin_keys))
|
||||
return {twin_key: value for twin_key, (value, _) in zip(twin_keys, reads, strict=True) if value is not None}
|
||||
|
||||
|
||||
def _committed_budget_spend(
|
||||
counter_keys: EntitySpendCounterKeys, committed: Mapping[str, float]
|
||||
) -> CommittedBudgetSpend:
|
||||
def spend_of(counter_key: str | None) -> float | None:
|
||||
return committed.get(committed_spend_key(counter_key)) if counter_key is not None else None
|
||||
|
||||
return CommittedBudgetSpend(
|
||||
counter_keys=counter_keys,
|
||||
key=spend_of(counter_keys.key),
|
||||
team=spend_of(counter_keys.team),
|
||||
user=spend_of(counter_keys.user),
|
||||
org=spend_of(counter_keys.org),
|
||||
end_user=spend_of(counter_keys.end_user),
|
||||
)
|
||||
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
async def read_committed_budget_spend(
|
||||
spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys
|
||||
) -> CommittedBudgetSpend:
|
||||
twin_keys: Final = committed_spend_keys(counter_keys.tracked())
|
||||
if not twin_keys:
|
||||
return CommittedBudgetSpend(counter_keys=counter_keys)
|
||||
try:
|
||||
committed: Final = await _read_committed_values(spend_counter_store, twin_keys)
|
||||
except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend
|
||||
verbose_proxy_logger.debug("committed spend read failed for %s: %s", twin_keys, e)
|
||||
return CommittedBudgetSpend(counter_keys=counter_keys)
|
||||
return _committed_budget_spend(counter_keys=counter_keys, committed=committed)
|
||||
|
||||
|
||||
async def warm_committed_counter_keys(
|
||||
spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys
|
||||
) -> EntitySpendCounterKeys:
|
||||
"""For a charge with no auth to name the budgeted entities: only a budgeted entity's twin exists."""
|
||||
committed: Final = await read_committed_budget_spend(spend_counter_store, counter_keys)
|
||||
return EntitySpendCounterKeys(
|
||||
key=counter_keys.key if committed.key is not None else None,
|
||||
team=counter_keys.team if committed.team is not None else None,
|
||||
user=counter_keys.user if committed.user is not None else None,
|
||||
org=counter_keys.org if committed.org is not None else None,
|
||||
end_user=counter_keys.end_user if committed.end_user is not None else None,
|
||||
)
|
||||
|
||||
|
||||
async def read_admission_committed_spend(
|
||||
spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys
|
||||
) -> CommittedBudgetSpend:
|
||||
"""Rides the admission MGET: a counter the batch cannot answer stays unknown here and the emit-time read has it."""
|
||||
if not spend_counter_store.has_redis:
|
||||
return await read_committed_budget_spend(spend_counter_store, counter_keys)
|
||||
twin_keys: Final = committed_spend_keys(counter_keys.tracked())
|
||||
batched: Final = await asyncio.gather(*(read_batched_spend_counter(twin_key) for twin_key in twin_keys))
|
||||
return _committed_budget_spend(
|
||||
counter_keys=counter_keys,
|
||||
committed={
|
||||
twin_key: read[0] for twin_key, read in zip(twin_keys, batched) if read is not None and read[0] is not None
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def drop_committed_spend(spend_counter_store: SpendCounterStore, counter_keys: Sequence[str]) -> None:
|
||||
"""After a direct write to the stored spend the twin holds the old total, so the next charge seeds it again."""
|
||||
twin_keys: Final = committed_spend_keys(counter_keys)
|
||||
if not twin_keys:
|
||||
return
|
||||
try:
|
||||
await spend_counter_store.delete(twin_keys)
|
||||
except Exception as e: # noqa: BLE001 # the write already committed; a kept twin lasts until charges stop for its TTL
|
||||
verbose_proxy_logger.warning("committed spend drop failed: %s", e)
|
||||
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
async def seed_committed_spend(
|
||||
spend_counter_store: SpendCounterStore, prisma_client: PrismaClient | None, counter_key: str
|
||||
) -> bool:
|
||||
"""One DB read per worker for a cold twin: charges waiting on the lock find the twin the first one wrote."""
|
||||
twin_key: Final = committed_spend_key(counter_key)
|
||||
async with await SpendCounterReseed.counter_lock(twin_key):
|
||||
try:
|
||||
stored, _ = await spend_counter_store.read(twin_key)
|
||||
except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again
|
||||
verbose_proxy_logger.debug("committed spend seed read failed for %s: %s", twin_key, e)
|
||||
return False
|
||||
return stored is not None or await _seed_committed_spend_from_db(
|
||||
spend_counter_store, prisma_client, counter_key
|
||||
)
|
||||
|
||||
|
||||
async def _seed_committed_spend_from_db(
|
||||
spend_counter_store: SpendCounterStore, prisma_client: PrismaClient | None, counter_key: str
|
||||
) -> bool:
|
||||
db_spend: Final = await (
|
||||
SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key)
|
||||
if counter_key.startswith(END_USER_COUNTER_PREFIX)
|
||||
else SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key)
|
||||
)
|
||||
if db_spend is None:
|
||||
return False
|
||||
twin_key: Final = committed_spend_key(counter_key)
|
||||
try:
|
||||
await spend_counter_store.seed(twin_key, db_spend, increment_if_lost=False)
|
||||
except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again
|
||||
verbose_proxy_logger.debug("committed spend seed failed for %s: %s", twin_key, e)
|
||||
return False
|
||||
return True
|
||||
|
|
@ -2,7 +2,7 @@
|
|||
import collections
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from functools import partial
|
||||
|
|
@ -53,6 +53,7 @@ from litellm.proxy.auth.authorization_dependencies import (
|
|||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.openai_error_payload import error_status_code, openai_error_param, openai_error_type
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, drop_committed_spend
|
||||
from litellm.proxy.spend_tracking.session_page import (
|
||||
AGENT_CALL_TYPE_SQL,
|
||||
MCP_CALL_TYPES_SQL,
|
||||
|
|
@ -74,6 +75,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
|||
get_spend_by_team_and_customer,
|
||||
)
|
||||
from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_statement_timeout
|
||||
from litellm.redis.stores.spend_counters import key_counter, team_counter
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import SpendLogsRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
|
|
@ -3595,6 +3597,50 @@ async def view_spend_logs(
|
|||
)
|
||||
|
||||
|
||||
class _EntityIdRow(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
|
||||
|
||||
_COMMITTED_DROP_PAGE_SIZE: Final = 1000
|
||||
_CommittedDropEntity: TypeAlias = Literal["key", "team"]
|
||||
_COMMITTED_DROP_PAGE_QUERIES: Final[Mapping[_CommittedDropEntity, str]] = MappingProxyType(
|
||||
{
|
||||
"key": 'SELECT token AS id FROM "LiteLLM_VerificationToken" WHERE token > $1 ORDER BY token LIMIT $2',
|
||||
"team": 'SELECT team_id AS id FROM "LiteLLM_TeamTable" WHERE team_id > $1 ORDER BY team_id LIMIT $2',
|
||||
}
|
||||
)
|
||||
_COMMITTED_DROP_COUNTERS: Final[Mapping[_CommittedDropEntity, Callable[[str], str]]] = MappingProxyType(
|
||||
{"key": key_counter, "team": team_counter}
|
||||
)
|
||||
|
||||
|
||||
async def _drop_committed_spend_page(
|
||||
prisma_client: PrismaClient, entity: _CommittedDropEntity, after_id: str
|
||||
) -> str | None:
|
||||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
|
||||
rows: Final[Sequence[_EntityIdRow]] = await _query_raw(
|
||||
prisma_client, _COMMITTED_DROP_PAGE_QUERIES[entity], after_id, _COMMITTED_DROP_PAGE_SIZE
|
||||
)
|
||||
counter_of: Final = _COMMITTED_DROP_COUNTERS[entity]
|
||||
await drop_committed_spend(spend_counter_store, tuple(counter_of(row["id"]) for row in rows))
|
||||
return rows[-1]["id"] if len(rows) == _COMMITTED_DROP_PAGE_SIZE else None
|
||||
|
||||
|
||||
async def _drop_committed_spend_pages(prisma_client: PrismaClient, entity: _CommittedDropEntity) -> None:
|
||||
after_id: str | None = "" # rebind-ok: the keyset cursor moves one page at a time
|
||||
while after_id is not None:
|
||||
after_id = await _drop_committed_spend_page(prisma_client, entity, after_id)
|
||||
|
||||
|
||||
async def _drop_all_key_and_team_committed_spend(prisma_client: PrismaClient) -> None:
|
||||
try:
|
||||
await _drop_committed_spend_pages(prisma_client, "key")
|
||||
await _drop_committed_spend_pages(prisma_client, "team")
|
||||
except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin lasts until charges stop for its TTL
|
||||
verbose_proxy_logger.warning("committed spend drop after a global reset failed: %s", e)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/global/spend/reset",
|
||||
tags=["Budget & Spend Tracking"],
|
||||
|
|
@ -3626,6 +3672,8 @@ 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={})
|
||||
if committed_spend_tracked():
|
||||
await _drop_all_key_and_team_committed_spend(prisma_client)
|
||||
|
||||
return {
|
||||
"message": "Spend for all API Keys and Teams reset successfully",
|
||||
|
|
|
|||
|
|
@ -9263,13 +9263,14 @@ async def monitor_spend_logs_queue(
|
|||
proxy_logging_obj: Proxy logging object
|
||||
"""
|
||||
from litellm.constants import (
|
||||
SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS,
|
||||
SPEND_LOG_QUEUE_POLL_INTERVAL,
|
||||
SPEND_LOG_QUEUE_SIZE_THRESHOLD,
|
||||
)
|
||||
|
||||
threshold: Final = SPEND_LOG_QUEUE_SIZE_THRESHOLD
|
||||
base_interval: Final = SPEND_LOG_QUEUE_POLL_INTERVAL
|
||||
max_backoff: Final = 30.0 # Maximum backoff interval in seconds
|
||||
max_backoff: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS
|
||||
backoff_multiplier: Final = 1.5 # Exponential backoff multiplier
|
||||
current_interval = base_interval
|
||||
flush_requested: Final = asyncio.Event()
|
||||
|
|
|
|||
|
|
@ -6,15 +6,17 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm.caching.redis_batch import Batch
|
||||
from litellm.caching.redis_batch import Batch, PostCallSettlement
|
||||
from litellm.redis.stores._route import post_call_request, ready, route
|
||||
from litellm.redis.stores.spend_counters.handle import SpendCountersHandle
|
||||
from litellm.redis.stores.spend_counters.types import (
|
||||
COMMITTED_COUNTER_PREFIX,
|
||||
END_USER_COUNTER_PREFIX,
|
||||
SPEND_COUNTER_PREFIX,
|
||||
SPEND_COUNTERS_TARGET,
|
||||
SpendIncrement,
|
||||
SpendReading,
|
||||
committed_counter,
|
||||
end_user_counter,
|
||||
is_unavailable,
|
||||
key_counter,
|
||||
|
|
@ -34,12 +36,14 @@ if TYPE_CHECKING:
|
|||
from litellm.redis.stores.connection import RedisStores
|
||||
|
||||
__all__ = [
|
||||
"COMMITTED_COUNTER_PREFIX",
|
||||
"END_USER_COUNTER_PREFIX",
|
||||
"SPEND_COUNTERS_TARGET",
|
||||
"SPEND_COUNTER_PREFIX",
|
||||
"SpendCounterStore",
|
||||
"SpendIncrement",
|
||||
"SpendReading",
|
||||
"committed_counter",
|
||||
"end_user_counter",
|
||||
"is_unavailable",
|
||||
"key_counter",
|
||||
|
|
@ -81,6 +85,11 @@ class SpendCounterStore:
|
|||
def has_redis(self) -> bool:
|
||||
return self._handle.has_redis
|
||||
|
||||
def set_committed_ttl(self, ttl: float) -> None:
|
||||
"""The least TTL a Redis write gives a committed twin, which keeps it past the DB write that holds its
|
||||
charges."""
|
||||
self._handle.set_committed_ttl(ttl)
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
async def read(self, counter: str) -> SpendReading:
|
||||
reading: Final = await self._handle.read(counter)
|
||||
|
|
@ -104,12 +113,13 @@ class SpendCounterStore:
|
|||
return tuple(await self._handle.increment_many(increments))
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> bool:
|
||||
"""Declares the increments on the request's post-call batch, False when there is no Redis or request.
|
||||
Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere."""
|
||||
def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> PostCallSettlement | None:
|
||||
"""Declares the increments on the request's post-call batch, None when there is no Redis or request.
|
||||
Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere.
|
||||
The settlement reports when every reply is back, whoever sent the batch."""
|
||||
request: Final = post_call_request() if self.has_redis and increments else None
|
||||
if request is None:
|
||||
return False
|
||||
return None
|
||||
return self._handle.increment_many_post_call(request, increments)
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from __future__ import annotations
|
|||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from typing import Protocol
|
||||
|
||||
from litellm.caching.redis_batch import NativePending
|
||||
from litellm.caching.redis_batch import NativePending, PostCallSettlement
|
||||
from litellm.redis.stores.spend_counters.types import SpendIncrement
|
||||
from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeSend, SpendRead
|
||||
|
||||
|
|
@ -13,6 +13,7 @@ class SpendCountersHandle(Protocol):
|
|||
def has_redis(self) -> bool: ...
|
||||
@property
|
||||
def redis_ttl_seconds(self) -> float | None: ...
|
||||
def set_committed_ttl(self, ttl_seconds: float, /) -> None: ...
|
||||
def attach(self, client: NativeRedisClient, /) -> bool: ...
|
||||
def read(self, counter: str) -> Awaitable[SpendRead]: ...
|
||||
def read_many(
|
||||
|
|
@ -20,7 +21,9 @@ class SpendCountersHandle(Protocol):
|
|||
) -> 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 increment_many_post_call(
|
||||
self, request: NativeRequestRedis, increments: Sequence[SpendIncrement]
|
||||
) -> PostCallSettlement | None: ...
|
||||
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]: ...
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from litellm.rust_bridge._native import NativeRedisStoreError
|
|||
SPEND_COUNTERS_TARGET: Final = "spend_counters"
|
||||
SPEND_COUNTER_PREFIX: Final = "spend:"
|
||||
END_USER_COUNTER_PREFIX: Final = "spend:end_user:"
|
||||
COMMITTED_COUNTER_PREFIX: Final = "spend:committed:"
|
||||
|
||||
SpendIncrement: TypeAlias = tuple[str, float]
|
||||
SpendReading: TypeAlias = tuple[float | None, bool]
|
||||
|
|
@ -61,6 +62,11 @@ def shadow_eval_counter(job_id: str) -> str:
|
|||
return f"spend:shadow_eval:{job_id}"
|
||||
|
||||
|
||||
def committed_counter(counter_key: str) -> str:
|
||||
"""The charges-only twin of a key, team, user, org or end user lifetime counter."""
|
||||
return f"{COMMITTED_COUNTER_PREFIX}{counter_key.removeprefix(SPEND_COUNTER_PREFIX)}"
|
||||
|
||||
|
||||
def is_unavailable(error: BaseException) -> bool:
|
||||
"""Redis timed out or its breaker is open, the failures a post-call increment drops."""
|
||||
return isinstance(error, NativeRedisStoreError) and error.args[:1] in (("timeout",), ("breaker_open",))
|
||||
|
|
|
|||
|
|
@ -358,6 +358,11 @@ class SpendRead:
|
|||
@property
|
||||
def from_redis(self) -> bool: ...
|
||||
|
||||
@final
|
||||
class PostCallSettlement:
|
||||
@property
|
||||
def settled(self) -> bool: ...
|
||||
|
||||
@final
|
||||
class NativeSpendCounters:
|
||||
def __new__(cls, client: NativeRedisClient | None) -> NativeSpendCounters: ...
|
||||
|
|
@ -366,6 +371,7 @@ class NativeSpendCounters:
|
|||
def has_redis(self) -> bool: ...
|
||||
@property
|
||||
def redis_ttl_seconds(self) -> float | None: ...
|
||||
def set_committed_ttl(self, ttl_seconds: float) -> None: ...
|
||||
def read(self, counter: str) -> Future[SpendRead]: ...
|
||||
def read_many(
|
||||
self, counters: Sequence[str], send: NativeSend
|
||||
|
|
@ -374,7 +380,7 @@ class NativeSpendCounters:
|
|||
def increment_many(self, increments: Sequence[tuple[str, float]]) -> Future[list[float]]: ...
|
||||
def increment_many_post_call(
|
||||
self, request: NativeRequestRedis, increments: Sequence[tuple[str, float]]
|
||||
) -> bool: ...
|
||||
) -> PostCallSettlement | None: ...
|
||||
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]: ...
|
||||
|
|
@ -596,6 +602,7 @@ __all__ = [
|
|||
"NativeSend",
|
||||
"NativeSensitiveRoutePins",
|
||||
"NativeSpendCounters",
|
||||
"PostCallSettlement",
|
||||
"ProcessReservedForForking",
|
||||
"ResponsesWebSocketConnection",
|
||||
"RouterCall",
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from litellm.types.agents import AgentCaller, AgentResponse
|
|||
from litellm.types.proxy.agent_identity import ManagedAgentContext
|
||||
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
|
||||
from litellm.types.proxy.carried_budget_state import (
|
||||
CommittedBudgetSpend,
|
||||
OrgBudgetSnapshot,
|
||||
TeamBudgetSnapshot,
|
||||
UserBudgetSnapshot,
|
||||
|
|
@ -206,6 +207,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True)
|
||||
user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True)
|
||||
org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True)
|
||||
committed_budget_spend: CommittedBudgetSpend | None = Field(default=None, exclude=True, repr=False)
|
||||
matched_model_access_groups: list[str] | None = Field(default=None, exclude=True)
|
||||
budget_throttle_pct: float | None = Field(default=None, exclude=True)
|
||||
user: object | None = None # Expanded user object when expand=user is used
|
||||
|
|
@ -329,6 +331,21 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
)
|
||||
|
||||
|
||||
COMMITTED_BUDGET_SPEND_PARAM: Final = "committed_budget_spend"
|
||||
|
||||
|
||||
def committed_budget_spend_of_row(
|
||||
metadata: Mapping[str, object], litellm_params: Mapping[str, object]
|
||||
) -> CommittedBudgetSpend | None:
|
||||
auth: Final = metadata.get("user_api_key_auth")
|
||||
if isinstance(auth, UserAPIKeyAuth) and auth.committed_budget_spend is not None:
|
||||
return auth.committed_budget_spend
|
||||
if "user_api_key_auth" in metadata:
|
||||
return CommittedBudgetSpend()
|
||||
carried: Final = litellm_params.get(COMMITTED_BUDGET_SPEND_PARAM)
|
||||
return carried if isinstance(carried, CommittedBudgetSpend) else None
|
||||
|
||||
|
||||
def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
"""Return True if the caller's role grants unscoped read access to all
|
||||
tenant resources (managed files, batches, vector stores, spend rows, etc).
|
||||
|
|
|
|||
|
|
@ -9,10 +9,18 @@ reads them back with ``from_metadata``. ``None`` means this request never carrie
|
|||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ConfigDict, Field, ValidationError
|
||||
from typing_extensions import Self
|
||||
|
||||
from litellm.redis.stores.spend_counters.types import (
|
||||
end_user_counter,
|
||||
key_counter,
|
||||
org_counter,
|
||||
team_counter,
|
||||
user_counter,
|
||||
)
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
|
||||
|
|
@ -25,7 +33,7 @@ class _BudgetSnapshot(LiteLLMBaseModel):
|
|||
@classmethod
|
||||
def from_metadata(cls, metadata: Mapping[str, object]) -> Self | None:
|
||||
try:
|
||||
return cls.model_validate(metadata)
|
||||
return cls.model_validate(metadata, by_name=False)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
|
@ -64,3 +72,51 @@ class OrgBudgetSnapshot(_BudgetSnapshot):
|
|||
max_budget: float | None = Field(
|
||||
validation_alias="user_api_key_org_max_budget", serialization_alias="user_api_key_org_max_budget"
|
||||
)
|
||||
|
||||
|
||||
class EntitySpendCounterKeys(LiteLLMBaseModel):
|
||||
"""Spend counter of each entity whose remaining budget a gauge reports. ``None`` for an entity with no budget."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
key: str | None = None
|
||||
team: str | None = None
|
||||
user: str | None = None
|
||||
org: str | None = None
|
||||
end_user: str | None = None
|
||||
|
||||
@classmethod
|
||||
def of(
|
||||
cls,
|
||||
token: str | None,
|
||||
team_id: str | None,
|
||||
user_id: str | None,
|
||||
org_id: str | None,
|
||||
end_user_id: str | None,
|
||||
) -> Self:
|
||||
return cls(
|
||||
key=key_counter(token) if token else None,
|
||||
team=team_counter(team_id) if team_id else None,
|
||||
user=user_counter(user_id) if user_id else None,
|
||||
org=org_counter(org_id) if org_id else None,
|
||||
end_user=end_user_counter(end_user_id) if end_user_id else None,
|
||||
)
|
||||
|
||||
def tracked(self) -> tuple[str, ...]:
|
||||
counter_keys: Final = (self.key, self.team, self.user, self.org, self.end_user)
|
||||
return tuple(counter_key for counter_key in counter_keys if counter_key is not None)
|
||||
|
||||
|
||||
class CommittedBudgetSpend(LiteLLMBaseModel):
|
||||
"""Spend each entity was charged, budget holds excluded, read from the committed spend counters for the
|
||||
remaining budget gauges only. ``None`` when the entity has no committed counter. ``counter_keys`` names the
|
||||
budgeted entities whose committed counters this request charges and reads."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
counter_keys: EntitySpendCounterKeys = Field(default=EntitySpendCounterKeys(), repr=False)
|
||||
key: float | None = Field(default=None, repr=False)
|
||||
team: float | None = Field(default=None, repr=False)
|
||||
user: float | None = Field(default=None, repr=False)
|
||||
org: float | None = Field(default=None, repr=False)
|
||||
end_user: float | None = Field(default=None, repr=False)
|
||||
|
|
|
|||
|
|
@ -424,7 +424,7 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa
|
|||
|
||||
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double
|
||||
await proxy_server._apply_spend_counter_increments(pending)
|
||||
declared = list(spend.post_call)
|
||||
declared = [(request, increments) for request, increments, _ in spend.post_call]
|
||||
assert natives.built[0].events == []
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
|
|
@ -433,6 +433,27 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa
|
|||
assert spend.commands == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_request_sees_its_committed_spend_charge_settle_whoever_sends_the_batch(monkeypatch):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
spend = FakeSpendCounters(redis={})
|
||||
monkeypatch.setattr(proxy_server, "spend_counter_store", spend.store())
|
||||
|
||||
with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
await proxy_server._apply_spend_counter_increments([PendingSpendIncrement("spend:key:k1", 0.5)])
|
||||
without_committed_spend = request.post_call_settlements
|
||||
await proxy_server._apply_spend_counter_increments(
|
||||
[PendingSpendIncrement("spend:key:k1", 0.5), PendingSpendIncrement("spend:committed:key:k1", 0.5)]
|
||||
)
|
||||
declared = [settlement.settled for settlement in request.post_call_settlements]
|
||||
spend.settle_post_call()
|
||||
sent = [settlement.settled for settlement in request.post_call_settlements]
|
||||
|
||||
assert without_committed_spend == ()
|
||||
assert (declared, sent) == ([False], [True])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary():
|
||||
natives = FakeNatives()
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ from litellm.integrations.gcs_pubsub.pub_sub import *
|
|||
from litellm.integrations.generic_api.generic_api_callback import LOG_FORMAT_TYPES, GenericAPILogger
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
|
||||
|
|
@ -795,3 +797,27 @@ async def test_generic_api_concurrent_flushes_deliver_each_event_once(log_format
|
|||
task.cancel()
|
||||
await asyncio.gather(*owned_tasks, return_exceptions=True)
|
||||
await handler.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_payload_of_a_chat_row_never_shows_the_committed_spend_auth_read() -> None:
|
||||
counter_keys: Final = EntitySpendCounterKeys.of(
|
||||
token="hashed-token", team_id="t1", user_id=None, org_id=None, end_user_id=None
|
||||
)
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
api_key="hashed-token",
|
||||
team_id="t1",
|
||||
committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys, key=0.123456, team=0.654321),
|
||||
)
|
||||
generic_logger: Final = GenericAPILogger(endpoint="https://example.com/api/logs", flush_interval=3600)
|
||||
|
||||
payload: Final = generic_logger._get_v1_logging_payload(
|
||||
kwargs={"litellm_params": {"metadata": {"user_api_key_auth": auth}}},
|
||||
response_obj={"usage": {}},
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
assert "team_id='t1'" in payload["metadata"]
|
||||
assert "0.123456" not in payload["metadata"]
|
||||
assert str(counter_keys.team) not in payload["metadata"]
|
||||
|
|
|
|||
|
|
@ -4,13 +4,22 @@ from request metadata. get_*_object only runs when that state is missing (custom
|
|||
auth, SDK callers, failure paths)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from prometheus_client import REGISTRY
|
||||
|
||||
from litellm.caching.redis_batch import (
|
||||
active_post_call_redis_batch,
|
||||
flush_post_call_redis_batches,
|
||||
request_redis_batch_scope,
|
||||
)
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_OrganizationTable,
|
||||
|
|
@ -24,11 +33,16 @@ from litellm.proxy.spend_tracking.carried_budget_state import (
|
|||
carry_organization_budget_state,
|
||||
carry_team_and_user_budget_state,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement
|
||||
from litellm.types.proxy.carried_budget_state import (
|
||||
CommittedBudgetSpend,
|
||||
EntitySpendCounterKeys,
|
||||
KeyBudgetSnapshot,
|
||||
TeamBudgetSnapshot,
|
||||
UserBudgetSnapshot,
|
||||
)
|
||||
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send, make
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
TEAM_RESET_AT = datetime(2026, 10, 1, tzinfo=timezone.utc)
|
||||
USER_RESET_AT = datetime(2026, 11, 1, tzinfo=timezone.utc)
|
||||
|
|
@ -55,6 +69,15 @@ def prometheus_logger():
|
|||
return PrometheusLogger()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def spend_counters():
|
||||
spend: Final = FakeSpendCounters(redis=None)
|
||||
with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam
|
||||
"litellm.proxy.proxy_server.spend_counter_store", spend.store()
|
||||
):
|
||||
yield spend
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def getters():
|
||||
"""Every response-path object getter, patched where prometheus imports them from."""
|
||||
|
|
@ -137,6 +160,7 @@ def _request_metadata(token: UserAPIKeyAuth) -> dict:
|
|||
return {
|
||||
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(token),
|
||||
**carried_budget_metadata(token),
|
||||
"user_api_key_auth": token,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -198,6 +222,499 @@ async def test_authed_request_sets_every_gauge_without_any_object_getter(prometh
|
|||
prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once()
|
||||
|
||||
|
||||
_COUNTER_KEYS: Final = ("spend:key:hashed", "spend:team:t1", "spend:user:u1", "spend:org:o1")
|
||||
_TWIN_KEYS: Final = (
|
||||
"spend:committed:key:hashed",
|
||||
"spend:committed:team:t1",
|
||||
"spend:committed:user:u1",
|
||||
"spend:committed:org:o1",
|
||||
)
|
||||
|
||||
|
||||
def _store(spend: FakeSpendCounters, counter_keys: Sequence[str], values: Sequence[float]) -> None:
|
||||
tier: Final = spend.memory if spend.redis is None else spend.redis
|
||||
tier.update(zip(counter_keys, values, strict=True))
|
||||
|
||||
|
||||
def _store_counters(spend: FakeSpendCounters, key: float, team: float, user: float, org: float) -> None:
|
||||
_store(spend, _COUNTER_KEYS, (key, team, user, org))
|
||||
|
||||
|
||||
def _store_committed(spend: FakeSpendCounters, key: float, team: float, user: float, org: float) -> None:
|
||||
_store(spend, _TWIN_KEYS, (key, team, user, org))
|
||||
|
||||
|
||||
def _remaining(prometheus_logger: PrometheusLogger) -> dict[str, float]:
|
||||
return {
|
||||
"key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0],
|
||||
"team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0],
|
||||
"user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0],
|
||||
"org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0],
|
||||
}
|
||||
|
||||
|
||||
def _token_committed_at_auth() -> UserAPIKeyAuth:
|
||||
token = _authed_token()
|
||||
token.committed_budget_spend = CommittedBudgetSpend(
|
||||
counter_keys=EntitySpendCounterKeys.of(
|
||||
token="hashed", team_id="t1", user_id="u1", org_id="o1", end_user_id=None
|
||||
),
|
||||
key=3.0,
|
||||
team=30.0,
|
||||
user=8.0,
|
||||
org=60.0,
|
||||
)
|
||||
return token
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_use_the_committed_spend_at_emit_when_other_requests_finished_after_auth(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
"""This request's 2.0 plus 1.0 from another request that finished first, so the gauge is never set back."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=6.0, team=33.0, user=11.0, org=63.0)
|
||||
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(4.0),
|
||||
"team": pytest.approx(267.0),
|
||||
"user": pytest.approx(39.0),
|
||||
"org": pytest.approx(437.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_add_the_cost_to_another_requests_charge_while_the_own_charge_is_not_in_yet(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
"""1.0 from another request reached the counters during this long request, this request's 2.0 has not."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=4.0, team=31.0, user=9.0, org=61.0)
|
||||
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(4.0),
|
||||
"team": pytest.approx(267.0),
|
||||
"user": pytest.approx(39.0),
|
||||
"org": pytest.approx(437.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_read_no_committed_counter_auth_did_not_name_as_budgeted(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=9.0, team=99.0, user=19.0, org=99.0)
|
||||
|
||||
await _emit(prometheus_logger, _request_metadata(_authed_token()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(7.0),
|
||||
"team": pytest.approx(278.0),
|
||||
"user": pytest.approx(43.0),
|
||||
"org": 458.0,
|
||||
}
|
||||
|
||||
|
||||
def _billed_row_metadata() -> dict:
|
||||
return {key: value for key, value in _request_metadata(_authed_token()).items() if key != "user_api_key_auth"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_row_with_no_auth_adds_its_cost_to_the_committed_spend_of_the_entities_it_is_charged_to(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
"""Counter increments flush after the success callbacks, so the row's own charge is not in the counters yet."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=6.0, team=33.0, user=11.0, org=63.0)
|
||||
|
||||
await _emit(prometheus_logger, _billed_row_metadata())
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(2.0),
|
||||
"team": pytest.approx(265.0),
|
||||
"user": pytest.approx(37.0),
|
||||
"org": 435.0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_row_with_no_auth_and_no_committed_counters_keeps_the_stored_spend(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
_stub_gauges(prometheus_logger)
|
||||
|
||||
await _emit(prometheus_logger, _billed_row_metadata())
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(7.0),
|
||||
"team": pytest.approx(278.0),
|
||||
"user": pytest.approx(43.0),
|
||||
"org": 458.0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_auth_time_spend(
|
||||
prometheus_logger, getters, monkeypatch
|
||||
):
|
||||
async def never_answers(counter: str) -> None:
|
||||
await asyncio.sleep(30)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.spend_counter_store",
|
||||
FakeSpendCounters(redis={}, batch_reads=False, read_hook=never_answers).store(),
|
||||
)
|
||||
monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.01")
|
||||
_stub_gauges(prometheus_logger)
|
||||
|
||||
await asyncio.wait_for(
|
||||
prometheus_logger.async_log_failure_event(
|
||||
_failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(440.0))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_slow_committed_read_on_a_successful_request_still_sets_the_gauges_from_the_auth_time_spend(
|
||||
prometheus_logger, getters, monkeypatch
|
||||
):
|
||||
async def never_answers(counter: str) -> None:
|
||||
await asyncio.sleep(30)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.spend_counter_store",
|
||||
FakeSpendCounters(redis={}, batch_reads=False, read_hook=never_answers).store(),
|
||||
)
|
||||
monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.5")
|
||||
_stub_gauges(prometheus_logger)
|
||||
|
||||
await asyncio.wait_for(_emit(prometheus_logger, _request_metadata(_token_committed_at_auth())), timeout=5)
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(5.0),
|
||||
"team": pytest.approx(268.0),
|
||||
"user": pytest.approx(40.0),
|
||||
"org": pytest.approx(438.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_spend_counters():
|
||||
spend: Final = FakeSpendCounters(redis={})
|
||||
with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam
|
||||
"litellm.proxy.proxy_server.spend_counter_store", spend.store()
|
||||
):
|
||||
yield spend
|
||||
|
||||
|
||||
def _request_scope() -> request_redis_batch_scope:
|
||||
return request_redis_batch_scope(
|
||||
native_factory=FakeNatives(), # pyright: ignore[reportArgumentType] # test double
|
||||
native_send_factory=fake_send, # pyright: ignore[reportArgumentType] # test double
|
||||
)
|
||||
|
||||
|
||||
async def _declare_the_own_two_dollar_charge() -> None:
|
||||
await proxy_server._apply_spend_counter_increments(
|
||||
tuple(PendingSpendIncrement(counter_key, 2.0) for counter_key in (*_COUNTER_KEYS, *_TWIN_KEYS))
|
||||
)
|
||||
|
||||
|
||||
_BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE: Final = {
|
||||
"key": pytest.approx(3.0),
|
||||
"team": pytest.approx(266.0),
|
||||
"user": pytest.approx(38.0),
|
||||
"org": pytest.approx(436.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_add_the_own_charge_to_a_same_size_charge_of_another_request_while_the_own_waits_in_the_pipeline(
|
||||
prometheus_logger, getters, redis_spend_counters
|
||||
):
|
||||
"""Another request's 2.0 reached the counters and this request's 2.0 is queued, so both count."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0)
|
||||
|
||||
with _request_scope():
|
||||
await _declare_the_own_two_dollar_charge()
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
assert _remaining(prometheus_logger) == _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_do_not_add_the_own_charge_twice_once_the_pipeline_went_out(
|
||||
prometheus_logger, getters, redis_spend_counters
|
||||
):
|
||||
"""The 2.0 over the auth time spend is this request's own charge, settled before the gauges read it."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(redis_spend_counters, key=3.0, team=30.0, user=8.0, org=60.0)
|
||||
|
||||
with _request_scope():
|
||||
await _declare_the_own_two_dollar_charge()
|
||||
redis_spend_counters.settle_post_call()
|
||||
await flush_post_call_redis_batches()
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(5.0),
|
||||
"team": pytest.approx(268.0),
|
||||
"user": pytest.approx(40.0),
|
||||
"org": pytest.approx(438.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_unrelated_post_call_flush_does_not_make_the_gauge_drop_its_own_unsent_charge(
|
||||
prometheus_logger, getters, redis_spend_counters
|
||||
):
|
||||
"""An unrelated write went out first, then another request's 2.0 reached the counters. This request's 2.0 counts
|
||||
once whether it is undeclared, declared but unsent, or settled."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0)
|
||||
unrelated_cache, _ = make()
|
||||
metadata: Final = _request_metadata(_token_committed_at_auth())
|
||||
|
||||
with _request_scope():
|
||||
active_post_call_redis_batch(unrelated_cache).increment("unrelated-write", 1.0)
|
||||
await flush_post_call_redis_batches()
|
||||
await _emit(prometheus_logger, metadata)
|
||||
undeclared: Final = _remaining(prometheus_logger)
|
||||
await _declare_the_own_two_dollar_charge()
|
||||
await _emit(prometheus_logger, metadata)
|
||||
declared_unsent: Final = _remaining(prometheus_logger)
|
||||
redis_spend_counters.settle_post_call()
|
||||
await _emit(prometheus_logger, metadata)
|
||||
settled: Final = _remaining(prometheus_logger)
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
assert {"undeclared": undeclared, "declared_unsent": declared_unsent, "settled": settled} == {
|
||||
"undeclared": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE,
|
||||
"declared_unsent": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE,
|
||||
"settled": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_count_the_own_charge_once_when_it_settles_while_the_gauges_read_the_counters(
|
||||
prometheus_logger, getters, redis_spend_counters
|
||||
):
|
||||
"""Another request's 2.0 reached the counters, and this request's 2.0 goes out during the gauges' read."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0)
|
||||
|
||||
async def settle_during_read(_counter: str) -> None:
|
||||
redis_spend_counters.settle_post_call()
|
||||
|
||||
redis_spend_counters.batch_reads = False
|
||||
redis_spend_counters.read_hook = settle_during_read
|
||||
|
||||
with _request_scope():
|
||||
await _declare_the_own_two_dollar_charge()
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
assert _remaining(prometheus_logger) == _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_without_redis_treat_a_charge_equal_to_the_own_as_the_own(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
"""Without Redis the charge increments in memory at once, so a request scope says nothing about it."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=5.0, team=32.0, user=10.0, org=62.0)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(5.0),
|
||||
"team": pytest.approx(268.0),
|
||||
"user": pytest.approx(40.0),
|
||||
"org": pytest.approx(438.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_add_the_cost_to_the_auth_time_spend_before_the_charge_reaches_the_counters(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=60.0)
|
||||
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(5.0),
|
||||
"team": pytest.approx(268.0),
|
||||
"user": pytest.approx(40.0),
|
||||
"org": pytest.approx(438.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_keep_the_auth_time_spend_when_the_proxy_spend_counters_cannot_be_read(
|
||||
prometheus_logger, getters, monkeypatch
|
||||
):
|
||||
_stub_gauges(prometheus_logger)
|
||||
monkeypatch.delattr("litellm.proxy.proxy_server.spend_counter_store")
|
||||
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(5.0),
|
||||
"team": pytest.approx(268.0),
|
||||
"user": pytest.approx(40.0),
|
||||
"org": pytest.approx(438.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_budget_hold_on_the_enforcement_counters_never_shows_in_the_gauges(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=5.0, team=32.0, user=10.0, org=62.0)
|
||||
_store_counters(spend_counters, key=9.0, team=36.0, user=14.0, org=66.0)
|
||||
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(5.0),
|
||||
"team": pytest.approx(268.0),
|
||||
"user": pytest.approx(40.0),
|
||||
"org": pytest.approx(438.0),
|
||||
}
|
||||
|
||||
|
||||
def _failed_request(metadata: dict, response_cost: float = 0.0) -> dict:
|
||||
return {
|
||||
"model": "openai/gpt-5.4-nano",
|
||||
"litellm_params": {"metadata": metadata},
|
||||
"standard_logging_object": {
|
||||
"model": "openai/gpt-5.4-nano",
|
||||
"model_id": "dep-1",
|
||||
"model_group": "chat",
|
||||
"api_base": "https://api.openai.com",
|
||||
"custom_llm_provider": "openai",
|
||||
"status": "failure",
|
||||
"response_cost": response_cost,
|
||||
"request_tags": [],
|
||||
"hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None},
|
||||
"metadata": {
|
||||
"user_api_key_hash": "hashed",
|
||||
"user_api_key_alias": "key-alias",
|
||||
"user_api_key_team_id": "t1",
|
||||
"user_api_key_team_alias": "team-alias",
|
||||
"user_api_key_user_id": "u1",
|
||||
"user_api_key_org_id": "o1",
|
||||
"requester_metadata": None,
|
||||
"user_api_key_end_user_id": None,
|
||||
},
|
||||
},
|
||||
"exception": Exception("upstream failed"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_request_sets_the_org_gauge_from_the_carried_org_and_committed_spend(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0)
|
||||
_store_counters(spend_counters, key=9.0, team=36.0, user=14.0, org=90.0)
|
||||
|
||||
await prometheus_logger.async_log_failure_event(
|
||||
_failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT
|
||||
)
|
||||
|
||||
assert getters["get_org_object"].await_count == 0
|
||||
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_request_without_the_auth_object_still_reads_the_committed_org_spend(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0)
|
||||
metadata = _request_metadata(_token_committed_at_auth())
|
||||
del metadata["user_api_key_auth"]
|
||||
|
||||
await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT)
|
||||
|
||||
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_request_without_a_carried_org_reads_no_committed_org_spend(
|
||||
prometheus_logger, getters, spend_counters
|
||||
):
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0)
|
||||
token = _token_committed_at_auth().model_copy(update={"org_budget_snapshot": None})
|
||||
metadata = {key: value for key, value in _request_metadata(token).items() if key != "user_api_key_auth"}
|
||||
|
||||
await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT)
|
||||
|
||||
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(460.0))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_request_that_was_charged_counts_its_own_charge_while_it_waits_in_the_pipeline(
|
||||
prometheus_logger, getters, redis_spend_counters
|
||||
):
|
||||
"""A stream cut after 40 events still costs 2.0, and that charge is queued when the failure logs."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0)
|
||||
|
||||
with _request_scope():
|
||||
await _declare_the_own_two_dollar_charge()
|
||||
await prometheus_logger.async_log_failure_event(
|
||||
_failed_request(_request_metadata(_token_committed_at_auth()), response_cost=2.0),
|
||||
None,
|
||||
KEY_RESET_AT,
|
||||
KEY_RESET_AT,
|
||||
)
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(436.0))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters):
|
||||
_stub_gauges(prometheus_logger)
|
||||
metadata = {
|
||||
"user_api_key_team_spend": 20.0,
|
||||
"user_api_key_team_max_budget": 300.0,
|
||||
"spend": 1000.0,
|
||||
"max_budget": 1.0,
|
||||
"budget_reset_at": None,
|
||||
"user_alias": "not-the-user",
|
||||
}
|
||||
|
||||
await _emit(prometheus_logger, metadata)
|
||||
|
||||
assert {name: getter.await_count for name, getter in getters.items()} == {
|
||||
"get_key_object": 1,
|
||||
"get_team_object": 1,
|
||||
"get_user_object": 1,
|
||||
"get_org_object": 1,
|
||||
}
|
||||
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(458.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters):
|
||||
_stub_gauges(prometheus_logger)
|
||||
|
|
|
|||
|
|
@ -1065,6 +1065,79 @@ async def test_set_customer_budget_metrics_after_api_request_uses_cached_default
|
|||
assert _customer_sample("litellm_customer_max_budget_metric", "cust-default") == pytest.approx(3.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_customer_gauge_subtracts_the_committed_spend_instead_of_the_cached_spend(
|
||||
prometheus_logger, customer_metrics_enabled
|
||||
):
|
||||
"""Auth saw 1.5 committed, and by emit the counter holds 2.5 with this request's 0.5 and another request's."""
|
||||
import sys
|
||||
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
mock_proxy_server = MagicMock()
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(
|
||||
return_value=LiteLLM_EndUserTable(
|
||||
user_id="cust-live",
|
||||
blocked=False,
|
||||
spend=0.5,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0),
|
||||
)
|
||||
)
|
||||
mock_proxy_server.spend_counter_store = FakeSpendCounters(
|
||||
redis=None, memory={"spend:committed:end_user:cust-live": 2.5}
|
||||
).store()
|
||||
token = UserAPIKeyAuth(token="hashed")
|
||||
token.committed_budget_spend = CommittedBudgetSpend(
|
||||
counter_keys=EntitySpendCounterKeys(end_user="spend:end_user:cust-live"), end_user=1.5
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
|
||||
await prometheus_logger._increment_remaining_budget_metrics(
|
||||
user_api_team=None,
|
||||
user_api_team_alias=None,
|
||||
user_api_key=None,
|
||||
user_api_key_alias=None,
|
||||
litellm_params={"metadata": {"user_api_key_auth": token}},
|
||||
response_cost=0.5,
|
||||
end_user_id="cust-live",
|
||||
)
|
||||
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-live") == pytest.approx(0.5)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_customer_gauge_counts_a_charge_the_cached_customer_row_already_holds_once(
|
||||
prometheus_logger, customer_metrics_enabled
|
||||
):
|
||||
import sys
|
||||
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
|
||||
mock_proxy_server = MagicMock()
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(
|
||||
return_value=LiteLLM_EndUserTable(
|
||||
user_id="cust-charged",
|
||||
blocked=False,
|
||||
spend=0.5,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0),
|
||||
)
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
|
||||
await prometheus_logger._set_customer_budget_metrics_after_api_request(
|
||||
end_user_id="cust-charged",
|
||||
response_cost=0.5,
|
||||
committed_spend=0.0,
|
||||
)
|
||||
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-charged") == pytest.approx(2.5)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_customer_budget_metrics_after_api_request_without_budget_only_emits_remaining(
|
||||
prometheus_logger, customer_metrics_enabled
|
||||
|
|
|
|||
|
|
@ -5094,6 +5094,70 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("prometheus_registered", [True, False], ids=["prometheus_logger", "no_prometheus_logger"])
|
||||
async def test_centralized_common_checks_carry_the_committed_spend_beside_the_stale_cached_spend(
|
||||
monkeypatch: pytest.MonkeyPatch, prometheus_registered: bool
|
||||
):
|
||||
"""The enforcement counters also hold budget holds, the committed twins only charged spend."""
|
||||
from starlette.datastructures import URL
|
||||
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy.spend_tracking.committed_spend import committed_spend_key
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.redis.stores.connection import connect
|
||||
from litellm.redis.stores.spend_counters import SpendCounterStore
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
|
||||
|
||||
for callback_list in ("success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback"):
|
||||
monkeypatch.setattr(litellm, callback_list, [])
|
||||
monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)] if prometheus_registered else [])
|
||||
token: Final = UserAPIKeyAuth(api_key="sk-live-spend", max_budget=10.0, spend=0.0, team_id="t-live", team_spend=0.0)
|
||||
user_api_key_cache: Final = UserApiKeyCache()
|
||||
user_api_key_cache.set_cache(
|
||||
key="team_id:t-live", value=LiteLLM_TeamTableCachedObj(team_id="t-live", max_budget=20.0, spend=0.0)
|
||||
)
|
||||
spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters())
|
||||
for counter_key, held, committed in ((f"spend:key:{token.token}", 0.67, 0.37), ("spend:team:t-live", 1.55, 1.25)):
|
||||
spend_counter_store.set_local(counter_key, held)
|
||||
spend_counter_store.set_local(committed_spend_key(counter_key), committed)
|
||||
attrs: Final = {
|
||||
**_proxy_attrs_for_centralized_checks(),
|
||||
"llm_router": None,
|
||||
"prisma_client": MagicMock(),
|
||||
"user_api_key_cache": user_api_key_cache,
|
||||
"spend_counter_store": spend_counter_store,
|
||||
"proxy_logging_obj": ProxyLogging(user_api_key_cache=user_api_key_cache),
|
||||
}
|
||||
originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs}
|
||||
request: Final = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
try:
|
||||
for attribute, value in attrs.items():
|
||||
setattr(_proxy_server_mod, attribute, value)
|
||||
await run_centralized_common_checks(
|
||||
user_api_key_auth_obj=token,
|
||||
request=request,
|
||||
request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1000},
|
||||
route="/chat/completions",
|
||||
)
|
||||
finally:
|
||||
for attribute, value in originals.items():
|
||||
setattr(_proxy_server_mod, attribute, value)
|
||||
|
||||
assert token.committed_budget_spend == (
|
||||
CommittedBudgetSpend(
|
||||
counter_keys=EntitySpendCounterKeys(key=f"spend:key:{token.token}", team=f"spend:team:{token.team_id}"),
|
||||
key=0.37,
|
||||
team=1.25,
|
||||
)
|
||||
if prometheus_registered
|
||||
else None
|
||||
)
|
||||
assert (token.spend, token.team_spend) == (0.0, 0.0)
|
||||
|
||||
|
||||
def _end_user_budget_row(budget_id: str, max_budget: float) -> MagicMock:
|
||||
row = MagicMock()
|
||||
row.dict = lambda: {"budget_id": budget_id, "max_budget": max_budget}
|
||||
|
|
|
|||
|
|
@ -1438,7 +1438,47 @@ def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_
|
|||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
|
||||
|
||||
assert "spend:key:sk-abc" in counter_cache.deleted()
|
||||
assert {"spend:key:sk-abc", "spend:committed:key:sk-abc"} <= counter_cache.deleted()
|
||||
|
||||
|
||||
def _expired_key_row() -> object:
|
||||
return type(
|
||||
"Key",
|
||||
(),
|
||||
{
|
||||
"spend": 100.0,
|
||||
"budget_duration": "30d",
|
||||
"budget_reset_at": datetime.now(timezone.utc),
|
||||
"id": "key-1",
|
||||
"token": "sk-abc",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_reset_budget_for_keys_without_redis_clears_the_counter_and_its_committed_twin(
|
||||
reset_budget_job, mock_prisma_client, monkeypatch
|
||||
):
|
||||
counter_cache: Final = _make_counter_invalidation_job(monkeypatch)
|
||||
counter_cache.redis = None
|
||||
counter_cache.memory.update({"spend:key:sk-abc": 100.0, "spend:committed:key:sk-abc": 100.0})
|
||||
mock_prisma_client.data["key"] = [_expired_key_row()]
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
|
||||
|
||||
assert counter_cache.memory == {}
|
||||
|
||||
|
||||
def test_a_redis_failure_on_the_key_reset_still_drops_its_committed_twin_locally(
|
||||
reset_budget_job, mock_prisma_client, monkeypatch
|
||||
):
|
||||
counter_cache: Final = _make_counter_invalidation_job(monkeypatch)
|
||||
counter_cache.failing = frozenset({"delete"})
|
||||
counter_cache.memory.update({"spend:key:sk-abc": 100.0, "spend:committed:key:sk-abc": 100.0})
|
||||
mock_prisma_client.data["key"] = [_expired_key_row()]
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
|
||||
|
||||
assert counter_cache.memory == {}
|
||||
|
||||
|
||||
def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch):
|
||||
|
|
@ -1790,7 +1830,9 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma
|
|||
assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 6
|
||||
|
||||
invalidated: Final = counter_cache.deleted()
|
||||
assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)}
|
||||
assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} | {
|
||||
f"spend:committed:end_user:cust-{i:06d}" for i in range(population)
|
||||
}
|
||||
evicted: Final = {
|
||||
key for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list for key in call.args[0]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ from litellm.proxy.utils import (
|
|||
BilledFailedAttempt,
|
||||
ProxyUpdateSpend,
|
||||
)
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
|
||||
from litellm.types.router import Deployment, LiteLLM_Params
|
||||
from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage
|
||||
from litellm.types.videos.main import VideoObject
|
||||
|
|
@ -1129,6 +1130,95 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("metadata", "expected_scope"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"user_api_key_auth": UserAPIKeyAuth(
|
||||
committed_budget_spend=CommittedBudgetSpend(
|
||||
counter_keys=EntitySpendCounterKeys(key="spend:key:test_api_key")
|
||||
)
|
||||
)
|
||||
},
|
||||
EntitySpendCounterKeys(key="spend:key:test_api_key"),
|
||||
),
|
||||
({}, EntitySpendCounterKeys()),
|
||||
],
|
||||
)
|
||||
async def test_charge_moves_only_the_committed_counters_auth_named_as_budgeted(metadata, expected_scope):
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock()
|
||||
increment_spend_counters = AsyncMock()
|
||||
|
||||
await _update_database_and_spend_counters(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
increment_spend_counters=increment_spend_counters,
|
||||
user_api_key="test_api_key",
|
||||
user_id="test_user_id",
|
||||
end_user_id=None,
|
||||
team_id="test_team_id",
|
||||
org_id="test_org_id",
|
||||
kwargs={"litellm_params": {"metadata": metadata}},
|
||||
completion_response=None,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
response_cost=0.2,
|
||||
budget_reservation=None,
|
||||
)
|
||||
|
||||
assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == expected_scope
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("tracked", "expected_key", "expected_org"),
|
||||
[(True, pytest.approx(1.2), pytest.approx(2.2)), (False, 1.0, 2.0)],
|
||||
)
|
||||
async def test_charge_with_no_auth_on_its_row_moves_the_committed_counters_that_exist(
|
||||
tracked, expected_key, expected_org
|
||||
):
|
||||
"""A background response is billed from its creator's stored attribution, after the request's auth is gone."""
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock()
|
||||
spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters())
|
||||
spend_counter_store.set_local("spend:committed:key:test_api_key", 1.0)
|
||||
spend_counter_store.set_local("spend:committed:org:test_org_id", 2.0)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.hooks.proxy_track_cost_callback.committed_spend_tracked", return_value=tracked),
|
||||
patch("litellm.proxy.proxy_server.spend_counter_store", spend_counter_store),
|
||||
patch(
|
||||
"litellm.proxy.spend_tracking.committed_spend.SpendCounterReseed.from_db",
|
||||
AsyncMock(return_value=50.0),
|
||||
),
|
||||
):
|
||||
await _update_database_and_spend_counters(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
increment_spend_counters=increment_spend_counters,
|
||||
user_api_key="test_api_key",
|
||||
user_id="test_user_id",
|
||||
end_user_id=None,
|
||||
team_id="test_team_id",
|
||||
org_id="test_org_id",
|
||||
kwargs={"litellm_params": {"metadata": {"user_api_key_hash": "test_api_key"}}},
|
||||
completion_response=None,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
response_cost=0.2,
|
||||
budget_reservation=None,
|
||||
)
|
||||
|
||||
cached: Final = spend_counter_store.get_local
|
||||
assert cached("spend:committed:key:test_api_key") == expected_key
|
||||
assert cached("spend:committed:org:test_org_id") == expected_org
|
||||
assert cached("spend:committed:team:test_team_id") is None
|
||||
assert cached("spend:committed:user:test_user_id") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails():
|
||||
proxy_logging_obj = MagicMock()
|
||||
|
|
@ -1257,6 +1347,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
|
|||
request_started_at=start_time,
|
||||
model_access_groups=("premium",),
|
||||
project_id=None,
|
||||
committed_spend_scope=EntitySpendCounterKeys(),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2776,7 +2867,9 @@ _LIVE_COUNTERS: Final = {
|
|||
}
|
||||
|
||||
|
||||
async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None:
|
||||
async def _fail_request(
|
||||
logger: ProxyDBLogger, failure: Mapping[str, object], committed: CommittedBudgetSpend | None = None
|
||||
) -> None:
|
||||
await logger.async_post_call_failure_hook(
|
||||
request_data={
|
||||
"model": "gpt-4.1",
|
||||
|
|
@ -2793,6 +2886,7 @@ async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) ->
|
|||
end_user_id="end-user-1",
|
||||
org_id="org-1",
|
||||
project_id="project-1",
|
||||
committed_budget_spend=committed,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -2831,6 +2925,43 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("failure", "charged"),
|
||||
[
|
||||
({"combined_usage_object": Usage(prompt_tokens=30), "response_cost": 0.125}, 0.125),
|
||||
(_failure_with_billed_attempts(), 1.25),
|
||||
({"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}}, 0.125),
|
||||
],
|
||||
ids=["recovered_stream_cost", "billed_retries", "completed_cost"],
|
||||
)
|
||||
async def test_a_failed_charge_moves_only_the_committed_counters_auth_named_as_budgeted(
|
||||
live_spend_counters: SpendCounterStore, failure: Mapping[str, object], charged: float
|
||||
):
|
||||
"""The remaining budget gauges read the committed counters, so a charge missing from them shows as budget the
|
||||
request never left behind, until the next reseed from the database."""
|
||||
budgeted: Final = EntitySpendCounterKeys.of(
|
||||
token="hashed-key", team_id="team-1", user_id=None, org_id=None, end_user_id=None
|
||||
)
|
||||
starting: Final = {
|
||||
"spend:committed:key:hashed-key": 0.2,
|
||||
"spend:committed:team:team-1": 0.3,
|
||||
"spend:committed:user:user-1": 0.4,
|
||||
}
|
||||
for key, spend in starting.items():
|
||||
live_spend_counters.set_local(key, spend)
|
||||
|
||||
await _fail_request(ProxyDBLogger(), failure, CommittedBudgetSpend(counter_keys=budgeted))
|
||||
|
||||
assert {key: live_spend_counters.get_local(key) for key in starting} == pytest.approx(
|
||||
{
|
||||
"spend:committed:key:hashed-key": 0.2 + charged,
|
||||
"spend:committed:team:team-1": 0.3 + charged,
|
||||
"spend:committed:user:user-1": 0.4,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: SpendCounterStore):
|
||||
"""A failure whose cost was already finalized settles through the success path's write, which raises the
|
||||
|
|
|
|||
|
|
@ -249,6 +249,29 @@ def test_delete_customer_not_found(mock_prisma_client, mock_user_api_key_auth):
|
|||
assert response_json["error"]["code"] == "404"
|
||||
|
||||
|
||||
def test_delete_customer_drops_its_committed_spend_so_a_recreated_customer_starts_from_its_own_row(
|
||||
monkeypatch, mock_prisma_client, mock_user_api_key_auth
|
||||
):
|
||||
from types import SimpleNamespace
|
||||
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
spend: Final = FakeSpendCounters(
|
||||
redis=None,
|
||||
memory={"spend:committed:end_user:cust-gone": 7.0, "spend:committed:end_user:cust-stays": 3.0},
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
|
||||
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[SimpleNamespace(user_id="cust-gone")])
|
||||
mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1)
|
||||
|
||||
response: Final = client.post(
|
||||
"/customer/delete", json={"user_ids": ["cust-gone"]}, headers={"Authorization": "Bearer test-key"}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert spend.memory == {"spend:committed:end_user:cust-stays": 3.0}
|
||||
|
||||
|
||||
def test_error_schema_consistency(mock_prisma_client, mock_user_api_key_auth):
|
||||
"""
|
||||
Test that all customer endpoints return the same error schema format.
|
||||
|
|
|
|||
|
|
@ -4222,7 +4222,10 @@ async def test_admin_user_update_spend_invalidates_counter(mocker):
|
|||
admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller)
|
||||
mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user")
|
||||
assert [call.kwargs["counter_key"] for call in mock_invalidate.await_args_list] == [
|
||||
"spend:user:target-user",
|
||||
"spend:committed:user:target-user",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -5223,6 +5226,39 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non
|
|||
broadcast.assert_awaited_once_with(cache_key=deleted.user_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_user_drops_the_committed_spend_of_the_users_it_deleted(mocker: MockerFixture) -> None:
|
||||
from litellm.proxy._types import DeleteUserRequest, LiteLLM_UserTable
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
spend: Final = FakeSpendCounters(
|
||||
redis=None, memory={"spend:committed:user:user-gone": 7.0, "spend:committed:user:user-stays": 3.0}
|
||||
)
|
||||
mocker.patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()) # test-quality-ok: substitute the spend counter dependency
|
||||
deleted: Final = LiteLLM_UserTable(user_id="user-gone", user_email="gone@example.test", teams=[])
|
||||
prisma_client: Final = mocker.MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=deleted)
|
||||
prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_jwtkeymapping.find_many = mocker.AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0)
|
||||
prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=0)
|
||||
prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0)
|
||||
prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
|
||||
prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
|
||||
_serve_locked_user_delete(prisma_client, deleted)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
|
||||
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time
|
||||
|
||||
await delete_user(
|
||||
data=DeleteUserRequest(user_ids=[deleted.user_id]),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert spend.memory == {"spend:committed:user:user-stays": 3.0}
|
||||
|
||||
|
||||
_DB_OUTAGE_503_BODY: Final = {
|
||||
"error": {
|
||||
"message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.",
|
||||
|
|
|
|||
|
|
@ -8693,6 +8693,7 @@ async def test_reset_key_spend_success(monkeypatch):
|
|||
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
|
||||
key=f"spend_db_floor:spend:key:{hashed_key}", value=50.0, ttl=5
|
||||
)
|
||||
assert spend.deleted() == {f"spend:committed:key:{hashed_key}"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -9013,7 +9014,10 @@ 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.commands == [
|
||||
f"SET spend:key:{hashed_key} 0.0 60",
|
||||
f"DEL spend:committed:key:{hashed_key}",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -1188,6 +1188,27 @@ async def _run_legacy_update_organization(
|
|||
return mock_prisma_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("body", "twin_after"),
|
||||
[
|
||||
({"organization_id": "org-1", "spend": 0.0}, None),
|
||||
({"organization_id": "org-1", "organization_alias": "renamed"}, 7.0),
|
||||
],
|
||||
)
|
||||
async def test_legacy_update_drops_the_committed_org_spend_only_when_it_writes_the_stored_spend(
|
||||
monkeypatch, body, twin_after
|
||||
):
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
spend: Final = FakeSpendCounters(redis=None, memory={"spend:committed:org:org-1": 7.0})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
|
||||
|
||||
await _run_legacy_update_organization(monkeypatch, body=body, existing_budget_id="budget-1")
|
||||
|
||||
assert spend.memory.get("spend:committed:org:org-1") == twin_after
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
|
|
@ -1582,6 +1603,49 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke
|
|||
assert jwt_table.rows == (kept_row,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_organization_drops_the_committed_spend_of_the_org_and_its_teams(monkeypatch):
|
||||
from litellm.proxy._types import DeleteOrganizationRequest, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_endpoints.organization_endpoints import delete_organization
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
spend: Final = FakeSpendCounters(
|
||||
redis=None,
|
||||
memory={
|
||||
"spend:committed:org:org-doomed": 7.0,
|
||||
"spend:committed:team:team-in-org": 5.0,
|
||||
"spend:committed:org:org-kept": 3.0,
|
||||
"spend:committed:team:team-elsewhere": 2.0,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
|
||||
prisma_client: Final = AsyncMock()
|
||||
prisma_client.db.litellm_teamtable.find_many = AsyncMock(
|
||||
return_value=[SimpleNamespace(team_id="team-in-org", organization_id="org-doomed")]
|
||||
)
|
||||
prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_organizationtable.delete = AsyncMock(return_value=MagicMock())
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction():
|
||||
yield prisma_client.db
|
||||
|
||||
prisma_client.tx = transaction
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None)
|
||||
|
||||
await delete_organization(
|
||||
data=DeleteOrganizationRequest(organization_ids=["org-doomed"]),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert spend.memory == {"spend:committed:org:org-kept": 3.0, "spend:committed:team:team-elsewhere": 2.0}
|
||||
|
||||
|
||||
@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
|
||||
|
|
|
|||
|
|
@ -9966,6 +9966,37 @@ async def test_delete_validated_team_rows_deletes_only_the_keys_its_rostered_mem
|
|||
assert deleted_keys.kwargs["where"] == {"token": {"in": ["hashed-member"]}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_validated_team_rows_drops_the_committed_spend_of_the_teams_it_deleted(
|
||||
monkeypatch,
|
||||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
spend: Final = FakeSpendCounters(
|
||||
redis=None, memory={"spend:committed:team:team-doomed": 7.0, "spend:committed:team:team-kept": 3.0}
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
|
||||
team = LiteLLM_TeamTable(team_id="team-doomed", members_with_roles=[], metadata={}, model_max_budget={}, model_spend={})
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
|
||||
mock_access_group_tx_cm = MagicMock()
|
||||
mock_access_group_tx_cm.__aenter__ = AsyncMock(return_value=AsyncMock())
|
||||
mock_access_group_tx_cm.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_prisma_client.db.tx = MagicMock(return_value=mock_access_group_tx_cm)
|
||||
_wire_team_delete_tx(mock_prisma_client, locked_teams=(team,), team_keys=())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
|
||||
|
||||
await delete_validated_team_rows(
|
||||
teams=(team,),
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-admin"),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert spend.memory == {"spend:committed:team:team-kept": 3.0}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_validated_teams_writes_the_deleted_audit_row_only_for_a_team_still_there_under_the_lock(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
|
|||
|
|
@ -2,11 +2,12 @@ import asyncio
|
|||
import copy
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
import litellm
|
||||
|
|
@ -22,6 +23,7 @@ from litellm.proxy.management_helpers.bulk_user_deletion import (
|
|||
)
|
||||
from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkDeleteUserRequest
|
||||
from litellm.types.proxy.management_endpoints.team_endpoints import BulkTeamMemberDeleteRequest, TeamMemberRef
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin")
|
||||
INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
|
|
@ -697,6 +699,34 @@ async def test_bulk_delete_evicts_jwt_key_mappings_of_the_deleted_users_keys():
|
|||
assert cache.get_cache(key=kept) is not None
|
||||
|
||||
|
||||
async def _delete_through_scim_or_user_delete(prisma: _FakePrisma, user_ids: Sequence[str]) -> object:
|
||||
return await delete_users_from_current_teams(
|
||||
prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient
|
||||
user_ids=user_ids,
|
||||
user_api_key_dict=ADMIN,
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
proxy_logging_obj=None,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("delete", [_delete, _delete_through_scim_or_user_delete])
|
||||
async def test_user_deletes_drop_the_committed_spend_of_the_users_they_deleted(
|
||||
delete: Callable[[_FakePrisma, Sequence[str]], Awaitable[object]], mocker: MockerFixture
|
||||
) -> None:
|
||||
spend: Final = FakeSpendCounters(
|
||||
redis=None, memory={"spend:committed:user:u1": 7.0, "spend:committed:user:keep": 3.0}
|
||||
)
|
||||
mocker.patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()) # test-quality-ok: substitute the spend counter dependency
|
||||
prisma: Final = _FakePrisma(users=[_user("u1"), _user("keep")])
|
||||
|
||||
_ = await delete(prisma, ["u1"])
|
||||
|
||||
assert sorted(prisma.db.litellm_usertable.rows) == ["keep"]
|
||||
assert spend.memory == {"spend:committed:user:keep": 3.0}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_delete_rejects_non_admin_callers_before_touching_the_db():
|
||||
prisma = _FakePrisma(users=[_user("u1")])
|
||||
|
|
|
|||
|
|
@ -44,6 +44,8 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro
|
|||
from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.types.proxy.auth.user_api_key_auth import COMMITTED_BUDGET_SPEND_PARAM
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
|
||||
VertexPassthroughLoggingHandler,
|
||||
|
|
@ -9693,6 +9695,31 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock:
|
|||
return increment_spend_counters
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passthrough_charge_moves_the_committed_spend_of_the_entities_auth_named_budgeted():
|
||||
"""A forged body value stays out, and loggers that dump the row never get the auth object or the spend it read."""
|
||||
counter_keys: Final = EntitySpendCounterKeys.of(
|
||||
token="hashed-token", team_id="t1", user_id="u1", org_id=None, end_user_id=None
|
||||
)
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(
|
||||
api_key="hashed-token",
|
||||
user_id="u1",
|
||||
committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys, key=0.123456),
|
||||
metadata={"logging": [{"callback_name": "langfuse"}], "kept": "yes"},
|
||||
)
|
||||
|
||||
kwargs: Final = _passthrough_kwargs_for_reservation(
|
||||
user_api_key_dict,
|
||||
parsed_body={"litellm_metadata": {"user_api_key_auth": {"committed_budget_spend": None}}},
|
||||
)
|
||||
increment_spend_counters: Final = await _track_cost_for_passthrough_kwargs(kwargs)
|
||||
|
||||
assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == counter_keys
|
||||
assert "user_api_key_auth" not in kwargs["litellm_params"]["metadata"]
|
||||
assert "committed_budget_spend" not in repr(kwargs["litellm_params"]["metadata"])
|
||||
assert "0.123456" not in json.dumps(kwargs, default=str)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passthrough_success_reconciles_budget_reservation():
|
||||
"""
|
||||
|
|
@ -11135,6 +11162,9 @@ def _split_pass_through_body(
|
|||
GEMINI_BODY: Final = '{"contents": [{"parts": [{"text": "hi"}]}], "generationConfig": {"temperature": 0}}'
|
||||
|
||||
|
||||
PROXY_SET_PARAMS: Final = ("metadata", "proxy_server_request", COMMITTED_BUDGET_SPEND_PARAM)
|
||||
|
||||
|
||||
def _metadata_of(split: _PassThroughSplit) -> Mapping[str, object]:
|
||||
return MappingProxyType(_LITELLM_PARAMS.validate_python(split.litellm_params["metadata"]))
|
||||
|
||||
|
|
@ -11146,7 +11176,7 @@ def test_passthrough_moves_every_litellm_owned_key_from_the_forwarded_body_into_
|
|||
)
|
||||
|
||||
assert frozenset(split.litellm_params) == frozenset(
|
||||
("ttl", "num_retries", "litellm_trace_id", "metadata", "proxy_server_request")
|
||||
("ttl", "num_retries", "litellm_trace_id", *PROXY_SET_PARAMS)
|
||||
)
|
||||
assert tuple(split.litellm_params[k] for k in ("ttl", "num_retries", "litellm_trace_id")) == (30, 2, "trace-a")
|
||||
assert split.forwarded_body == json.loads(GEMINI_BODY)
|
||||
|
|
@ -11180,7 +11210,7 @@ PROXY_STAMPED_NAMES: Final = frozenset(
|
|||
def test_passthrough_keeps_each_registered_litellm_owned_name_out_of_the_forwarded_body(name: str) -> None:
|
||||
split: Final = _split_pass_through_body(json.dumps({name: "owned", **json.loads(GEMINI_BODY)}))
|
||||
|
||||
assert frozenset(split.litellm_params) == frozenset((name, "metadata", "proxy_server_request"))
|
||||
assert frozenset(split.litellm_params) == frozenset((name, *PROXY_SET_PARAMS))
|
||||
assert split.litellm_params[name] == "owned"
|
||||
assert split.forwarded_body == json.loads(GEMINI_BODY)
|
||||
|
||||
|
|
@ -11224,7 +11254,7 @@ def test_passthrough_keeps_a_client_model_info_object_for_keys_that_allow_client
|
|||
def test_passthrough_drops_a_client_supplied_proxy_stamped_name(name: str) -> None:
|
||||
split: Final = _split_pass_through_body(json.dumps({name: {"forged": "by-client"}, **json.loads(GEMINI_BODY)}))
|
||||
|
||||
assert frozenset(split.litellm_params) == frozenset(("metadata", "proxy_server_request"))
|
||||
assert frozenset(split.litellm_params) == frozenset(PROXY_SET_PARAMS)
|
||||
assert split.litellm_params["proxy_server_request"] != {"forged": "by-client"}, split.litellm_params
|
||||
assert split.forwarded_body == json.loads(GEMINI_BODY)
|
||||
|
||||
|
|
@ -11235,7 +11265,7 @@ def test_passthrough_merges_both_metadata_carriers_from_the_body_into_one_metada
|
|||
' "litellm_metadata": {"lm": "b"}, "generationConfig": {"temperature": 0}}'
|
||||
)
|
||||
|
||||
assert frozenset(split.litellm_params) == frozenset(("ttl", "metadata", "proxy_server_request"))
|
||||
assert frozenset(split.litellm_params) == frozenset(("ttl", *PROXY_SET_PARAMS))
|
||||
assert _metadata_of(split) == {**_metadata_of(_split_pass_through_body(GEMINI_BODY)), "client_tag": "a", "lm": "b"}
|
||||
assert split.forwarded_body == json.loads(GEMINI_BODY)
|
||||
|
||||
|
|
@ -11274,7 +11304,7 @@ def test_passthrough_sees_a_name_appended_to_the_public_list_after_import() -> N
|
|||
finally:
|
||||
litellm.all_litellm_params.remove("registered_later")
|
||||
|
||||
assert frozenset(split.litellm_params) == frozenset(("registered_later", "metadata", "proxy_server_request"))
|
||||
assert frozenset(split.litellm_params) == frozenset(("registered_later", *PROXY_SET_PARAMS))
|
||||
assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]}
|
||||
|
||||
|
||||
|
|
@ -11283,7 +11313,7 @@ def test_passthrough_sees_the_public_list_rebound_after_import(monkeypatch: pyte
|
|||
|
||||
split: Final = _split_pass_through_body(LATE_REGISTERED_BODY)
|
||||
|
||||
assert frozenset(split.litellm_params) == frozenset(("registered_later", "metadata", "proxy_server_request"))
|
||||
assert frozenset(split.litellm_params) == frozenset(("registered_later", *PROXY_SET_PARAMS))
|
||||
assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -38,10 +38,13 @@ from litellm.models.end_user import LiteLLM_EndUserTable
|
|||
from litellm.models.tag import LiteLLM_TagTable
|
||||
from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj
|
||||
from litellm.models.user import LiteLLM_UserTable
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key
|
||||
from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team
|
||||
from litellm.proxy.spend_tracking.committed_spend import committed_spend_key
|
||||
from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys
|
||||
|
||||
from .conftest import normalize
|
||||
from litellm.rust_bridge._native import NativeRedisStoreError
|
||||
|
|
@ -679,7 +682,7 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call(
|
|||
"spend:org:org-1",
|
||||
}
|
||||
spend = _spend_counters(monkeypatch, redis=dict.fromkeys(expected_keys, 1.0))
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache())
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
|
||||
await ps.increment_spend_counters(
|
||||
|
|
@ -697,6 +700,97 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call(
|
|||
assert spend.memory == dict.fromkeys(expected_keys, 1.5)
|
||||
|
||||
|
||||
class _VerificationTokenTable:
|
||||
def __init__(self, spend: float) -> None:
|
||||
self._spend: Final = spend
|
||||
self.lookups: list[dict[str, str]] = []
|
||||
|
||||
async def find_unique(self, where: dict[str, str]) -> MagicMock:
|
||||
self.lookups.append(where)
|
||||
return MagicMock(token=where["token"], spend=self._spend)
|
||||
|
||||
|
||||
def _prisma_with_key_spend(spend: float) -> MagicMock:
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_verificationtoken = _VerificationTokenTable(spend)
|
||||
return prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_committed_twin_takes_the_charge_but_never_a_budget_hold(monkeypatch):
|
||||
"""The key counter holds the DB spend 10.0, this request's 2.0 reservation and another request's 1.0 hold."""
|
||||
spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 13.0})
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0))
|
||||
reservation: Final = {"reserved_cost": 2.0, "entries": [{"counter_key": "spend:key:hashed"}], "finalized": False}
|
||||
scope: Final = EntitySpendCounterKeys(key="spend:key:hashed")
|
||||
|
||||
await ps.increment_spend_counters(
|
||||
token="hashed",
|
||||
team_id=None,
|
||||
user_id=None,
|
||||
response_cost=0.5,
|
||||
budget_reservation=reservation,
|
||||
committed_spend_scope=scope,
|
||||
)
|
||||
await ps.increment_spend_counters(
|
||||
token="hashed", team_id=None, user_id=None, response_cost=0.25, committed_spend_scope=scope
|
||||
)
|
||||
|
||||
assert spend.memory == pytest.approx({"spend:key:hashed": 11.75, committed_spend_key("spend:key:hashed"): 10.75})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_charge_outside_the_committed_scope_keeps_no_twin_and_reads_no_stored_spend(monkeypatch):
|
||||
spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 10.0})
|
||||
prisma: Final = _prisma_with_key_spend(10.0)
|
||||
key_cache: Final = UserApiKeyCache()
|
||||
key_cache.set_cache(key="hashed", value=UserAPIKeyAuth(token="hashed"))
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", key_cache)
|
||||
monkeypatch.setattr(ps, "prisma_client", prisma)
|
||||
|
||||
await ps.increment_spend_counters(token="hashed", team_id=None, user_id=None, response_cost=0.5)
|
||||
|
||||
assert spend.memory == pytest.approx({"spend:key:hashed": 10.5})
|
||||
assert prisma.db.litellm_verificationtoken.lookups == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_the_entities_in_the_committed_scope_take_the_charge_on_their_twin(monkeypatch):
|
||||
spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 10.0, "spend:team:t1": 10.0})
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0))
|
||||
|
||||
await ps.increment_spend_counters(
|
||||
token="hashed",
|
||||
team_id="t1",
|
||||
user_id=None,
|
||||
response_cost=0.5,
|
||||
committed_spend_scope=EntitySpendCounterKeys(key="spend:key:hashed"),
|
||||
)
|
||||
|
||||
assert spend.memory.get(committed_spend_key("spend:key:hashed")) == pytest.approx(10.5)
|
||||
assert committed_spend_key("spend:team:t1") not in spend.memory
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_budgeted_team_whose_id_names_a_window_keeps_its_charge_and_gets_no_twin(monkeypatch):
|
||||
team_counter_key: Final = "spend:team:t1:window:1d"
|
||||
spend: Final = _spend_counters(monkeypatch, redis=None, memory={team_counter_key: 1.0})
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr(ps, "prisma_client", MagicMock())
|
||||
|
||||
await ps.increment_spend_counters(
|
||||
token=None,
|
||||
team_id="t1:window:1d",
|
||||
user_id=None,
|
||||
response_cost=0.5,
|
||||
committed_spend_scope=EntitySpendCounterKeys(team=team_counter_key),
|
||||
)
|
||||
|
||||
assert spend.memory == pytest.approx({team_counter_key: 1.5})
|
||||
|
||||
|
||||
_PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = (
|
||||
"spend:key:hashed",
|
||||
"spend:team:team-1",
|
||||
|
|
@ -717,7 +811,7 @@ def _warm_counters_behind_a_failing_pipeline(monkeypatch, failure: str) -> FakeS
|
|||
fail_pipelines=True,
|
||||
failure=failure,
|
||||
)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache())
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache())
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
return spend
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,9 @@
|
|||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_OrganizationTable,
|
||||
|
|
@ -11,10 +14,13 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.spend_tracking.carried_budget_state import (
|
||||
carried_budget_metadata,
|
||||
carry_committed_budget_spend,
|
||||
carry_organization_budget_state,
|
||||
carry_team_and_user_budget_state,
|
||||
)
|
||||
from litellm.types.proxy.carried_budget_state import (
|
||||
CommittedBudgetSpend,
|
||||
EntitySpendCounterKeys,
|
||||
KeyBudgetSnapshot,
|
||||
OrgBudgetSnapshot,
|
||||
TeamBudgetSnapshot,
|
||||
|
|
@ -123,10 +129,105 @@ def test_snapshots_never_reach_the_serialized_token():
|
|||
user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"),
|
||||
)
|
||||
token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0)
|
||||
metadata_without_committed_spend = dict(carried_budget_metadata(token))
|
||||
token.committed_budget_spend = CommittedBudgetSpend(key=1.0, team=2.0, user=3.0, org=4.0, end_user=5.0)
|
||||
|
||||
dumped = token.model_dump()
|
||||
|
||||
assert "team_budget_snapshot" not in dumped
|
||||
assert "user_budget_snapshot" not in dumped
|
||||
assert "org_budget_snapshot" not in dumped
|
||||
assert "committed_budget_spend" not in dumped
|
||||
assert dict(carried_budget_metadata(token)) == metadata_without_committed_spend
|
||||
assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None
|
||||
|
||||
|
||||
def _committed_reader(committed: CommittedBudgetSpend):
|
||||
reads: list[EntitySpendCounterKeys] = []
|
||||
|
||||
async def read_committed(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend:
|
||||
reads.append(counter_keys)
|
||||
return committed
|
||||
|
||||
return read_committed, reads
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_committed_spend_is_read_once_for_every_budgeted_entity():
|
||||
token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0, max_budget=10.0)
|
||||
token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0)
|
||||
committed = CommittedBudgetSpend(key=1.3, team=20.3, user=5.3, org=40.3, end_user=0.8)
|
||||
read_committed, reads = _committed_reader(committed)
|
||||
|
||||
await carry_committed_budget_spend(
|
||||
valid_token=token,
|
||||
team_object=LiteLLM_TeamTable(team_id="t1", max_budget=300.0, spend=20.0),
|
||||
user_object=LiteLLM_UserTable(user_id="u1", max_budget=50.0, spend=5.0),
|
||||
end_user_object=LiteLLM_EndUserTable(
|
||||
user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=3.0)
|
||||
),
|
||||
read_committed=read_committed,
|
||||
)
|
||||
|
||||
assert reads == [
|
||||
EntitySpendCounterKeys(
|
||||
key="spend:key:hashed",
|
||||
team="spend:team:t1",
|
||||
user="spend:user:u1",
|
||||
org="spend:org:o1",
|
||||
end_user="spend:end_user:cust-1",
|
||||
)
|
||||
]
|
||||
assert token.committed_budget_spend == committed
|
||||
assert (token.spend, token.team_spend) == (1.0, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_route_without_a_budget_check_carries_the_counters_it_charges_and_reads_nothing():
|
||||
token = UserAPIKeyAuth(token="hashed", max_budget=10.0)
|
||||
|
||||
await carry_committed_budget_spend(
|
||||
valid_token=token, team_object=None, user_object=None, end_user_object=None, read_committed=None
|
||||
)
|
||||
|
||||
assert token.committed_budget_spend == CommittedBudgetSpend(
|
||||
counter_keys=EntitySpendCounterKeys(key="spend:key:hashed")
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_committed_spend_skips_entities_without_a_max_budget():
|
||||
token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0)
|
||||
token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=None)
|
||||
read_committed, reads = _committed_reader(CommittedBudgetSpend())
|
||||
|
||||
await carry_committed_budget_spend(
|
||||
valid_token=token,
|
||||
team_object=LiteLLM_TeamTable(team_id="t1", spend=20.0),
|
||||
user_object=LiteLLM_UserTable(user_id="u1", spend=5.0),
|
||||
end_user_object=LiteLLM_EndUserTable(
|
||||
user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=None)
|
||||
),
|
||||
read_committed=read_committed,
|
||||
)
|
||||
|
||||
assert reads == [EntitySpendCounterKeys()]
|
||||
assert token.committed_budget_spend == CommittedBudgetSpend()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_committed_spend_reads_the_team_org_when_the_key_has_none():
|
||||
token = UserAPIKeyAuth(token="hashed", team_id="t1")
|
||||
token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0)
|
||||
read_committed, reads = _committed_reader(CommittedBudgetSpend(org=41.0))
|
||||
|
||||
await carry_committed_budget_spend(
|
||||
valid_token=token,
|
||||
team_object=LiteLLM_TeamTable(team_id="t1", organization_id="team-org"),
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
read_committed=read_committed,
|
||||
)
|
||||
|
||||
assert reads == [EntitySpendCounterKeys(org="spend:org:team-org")]
|
||||
assert token.committed_budget_spend == CommittedBudgetSpend(org=41.0)
|
||||
|
|
|
|||
274
tests/unit/proxy/spend_tracking/test_committed_spend.py
Normal file
274
tests/unit/proxy/spend_tracking/test_committed_spend.py
Normal file
|
|
@ -0,0 +1,274 @@
|
|||
"""Committed spend twins hold only charged cost, so the remaining budget gauges never count a budget hold."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy.spend_tracking.committed_spend import (
|
||||
bind_committed_admission_counters,
|
||||
drop_committed_spend,
|
||||
read_admission_committed_spend,
|
||||
read_committed_budget_spend,
|
||||
seed_committed_spend,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
active_spend_counter_batch,
|
||||
read_batched_spend_counter,
|
||||
spend_counter_batch_scope,
|
||||
)
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
|
||||
def _no_side_effect() -> None:
|
||||
return None
|
||||
|
||||
|
||||
class _FindUniqueTable:
|
||||
def __init__(self, row: SimpleNamespace | None, on_read: Callable[[], None]) -> None:
|
||||
self._row: Final = row
|
||||
self._on_read: Final = on_read
|
||||
self.reads = 0
|
||||
|
||||
async def find_unique(self, where: dict[str, str]) -> SimpleNamespace | None:
|
||||
self.reads += 1
|
||||
await asyncio.sleep(0)
|
||||
self._on_read()
|
||||
return self._row
|
||||
|
||||
|
||||
def _prisma_with_spend(spend: float | None, on_read: Callable[[], None] = _no_side_effect) -> SimpleNamespace:
|
||||
row: Final = None if spend is None else SimpleNamespace(spend=spend)
|
||||
return SimpleNamespace(
|
||||
db=SimpleNamespace(
|
||||
litellm_verificationtoken=_FindUniqueTable(row, on_read),
|
||||
litellm_teamtable=_FindUniqueTable(row, on_read),
|
||||
litellm_endusertable=_FindUniqueTable(row, on_read),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(params=[False, True], ids=["in_memory", "redis"])
|
||||
def has_redis(request: pytest.FixtureRequest) -> bool:
|
||||
return request.param
|
||||
|
||||
|
||||
def _spend_counters(has_redis: bool, counters: Mapping[str, float] = MappingProxyType({})) -> FakeSpendCounters:
|
||||
if has_redis:
|
||||
return FakeSpendCounters(redis=dict(counters))
|
||||
return FakeSpendCounters(redis=None, memory=dict(counters))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_returns_each_committed_twin_and_none_for_one_never_charged(has_redis: bool):
|
||||
spend: Final = _spend_counters(
|
||||
has_redis,
|
||||
{"spend:committed:key:hashed": 1.5, "spend:committed:end_user:cust-1": 0.25, "spend:team:t1": 99.0},
|
||||
)
|
||||
counter_keys: Final = EntitySpendCounterKeys.of(
|
||||
token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id="cust-1"
|
||||
)
|
||||
|
||||
committed: Final = await read_committed_budget_spend(spend.store(), counter_keys)
|
||||
|
||||
assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5, end_user=0.25)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_twin_holding_garbage_reads_as_unknown_without_hiding_the_other_twins():
|
||||
spend: Final = FakeSpendCounters(
|
||||
redis={"spend:committed:key:hashed": "not-a-number", "spend:committed:team:t1": 4.0}
|
||||
)
|
||||
counter_keys: Final = EntitySpendCounterKeys.of(
|
||||
token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id=None
|
||||
)
|
||||
|
||||
committed: Final = await read_committed_budget_spend(spend.store(), counter_keys)
|
||||
|
||||
assert committed == CommittedBudgetSpend(counter_keys=counter_keys, team=4.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("counter_keys", "expected_spend"),
|
||||
[
|
||||
(EntitySpendCounterKeys(key="spend:key:hashed"), {"key": 12.0}),
|
||||
(EntitySpendCounterKeys(end_user="spend:end_user:cust-1"), {"end_user": 12.0}),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_writes_the_db_spend_when_the_twin_is_absent(
|
||||
has_redis: bool, counter_keys: EntitySpendCounterKeys, expected_spend: dict[str, float]
|
||||
):
|
||||
store: Final = _spend_counters(has_redis).store()
|
||||
(counter_key,) = counter_keys.tracked()
|
||||
|
||||
seeded: Final = await seed_committed_spend(store, _prisma_with_spend(12.0), counter_key)
|
||||
|
||||
assert seeded is True
|
||||
assert await read_committed_budget_spend(store, counter_keys) == CommittedBudgetSpend(
|
||||
counter_keys=counter_keys, **expected_spend
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_charges_on_a_cold_twin_read_the_db_once(has_redis: bool):
|
||||
store: Final = _spend_counters(has_redis).store()
|
||||
prisma_client: Final = _prisma_with_spend(12.0)
|
||||
counter_key: Final = f"spend:key:{uuid4()}"
|
||||
|
||||
seeded: Final = await asyncio.gather(*(seed_committed_spend(store, prisma_client, counter_key) for _ in range(5)))
|
||||
|
||||
assert seeded == [True] * 5
|
||||
assert prisma_client.db.litellm_verificationtoken.reads == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_reports_no_twin_when_the_twin_cannot_be_read():
|
||||
async def unreachable(counter: str) -> None:
|
||||
raise ConnectionError("redis unreachable")
|
||||
|
||||
spend: Final = FakeSpendCounters(read_hook=unreachable)
|
||||
|
||||
assert await seed_committed_spend(spend.store(), _prisma_with_spend(12.0), "spend:team:t1") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_never_lowers_a_twin_another_replica_already_advanced(has_redis: bool):
|
||||
store: Final = _spend_counters(has_redis, {"spend:committed:team:t1": 20.0}).store()
|
||||
|
||||
seeded: Final = await seed_committed_spend(store, _prisma_with_spend(12.0), "spend:team:t1")
|
||||
|
||||
committed: Final = await read_committed_budget_spend(store, EntitySpendCounterKeys(team="spend:team:t1"))
|
||||
assert seeded is True
|
||||
assert committed.team == 20.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_keeps_the_twin_another_worker_wrote_while_this_one_read_the_db():
|
||||
spend: Final = FakeSpendCounters(redis={})
|
||||
|
||||
def another_worker_seeds_and_charges() -> None:
|
||||
assert spend.redis is not None
|
||||
spend.redis["spend:committed:team:t1"] = 20.0
|
||||
|
||||
seeded: Final = await seed_committed_spend(
|
||||
spend.store(), _prisma_with_spend(12.0, on_read=another_worker_seeds_and_charges), "spend:team:t1"
|
||||
)
|
||||
|
||||
committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys(team="spend:team:t1"))
|
||||
assert seeded is True
|
||||
assert committed.team == 20.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_without_any_counter_key_sends_redis_nothing():
|
||||
spend: Final = FakeSpendCounters(redis={})
|
||||
|
||||
committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys())
|
||||
|
||||
assert committed == CommittedBudgetSpend()
|
||||
assert spend.commands == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_reports_no_twin_when_redis_rejects_the_write():
|
||||
spend: Final = FakeSpendCounters(redis={}, fail=True)
|
||||
|
||||
seeded: Final = await seed_committed_spend(spend.store(), _prisma_with_spend(12.0), "spend:team:t1")
|
||||
|
||||
committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys(team="spend:team:t1"))
|
||||
assert seeded is False
|
||||
assert committed.team is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_creates_no_twin_when_the_db_has_no_spend_for_the_entity(has_redis: bool):
|
||||
store: Final = _spend_counters(has_redis).store()
|
||||
|
||||
seeded: Final = await seed_committed_spend(store, _prisma_with_spend(None), "spend:key:hashed")
|
||||
|
||||
committed: Final = await read_committed_budget_spend(store, EntitySpendCounterKeys(key="spend:key:hashed"))
|
||||
assert seeded is False
|
||||
assert committed.key is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admission_read_rides_the_admission_mget_with_the_enforcement_counters(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)])
|
||||
spend: Final = FakeSpendCounters(redis={"spend:committed:key:hashed": 1.5})
|
||||
store: Final = spend.store()
|
||||
token: Final = UserAPIKeyAuth(token="hashed", team_id="t1")
|
||||
counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed", team="spend:team:t1")
|
||||
|
||||
with spend_counter_batch_scope(store, counter_keys=frozenset(("spend:key:hashed",))):
|
||||
bind_committed_admission_counters(token, end_user_id=None)
|
||||
await read_batched_spend_counter("spend:key:hashed")
|
||||
committed: Final = await read_admission_committed_spend(store, counter_keys)
|
||||
|
||||
assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5)
|
||||
assert len(spend.commands) == 1
|
||||
assert {"spend:committed:key:hashed", "spend:committed:team:t1"} <= set(spend.commands[0].split()[1:])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admission_read_outside_a_batch_leaves_redis_alone_and_the_spend_unknown():
|
||||
spend: Final = FakeSpendCounters(redis={"spend:committed:key:hashed": 1.5})
|
||||
counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed")
|
||||
|
||||
committed: Final = await read_admission_committed_spend(spend.store(), counter_keys)
|
||||
|
||||
assert committed == CommittedBudgetSpend(counter_keys=counter_keys)
|
||||
assert spend.commands == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admission_read_without_redis_reads_the_in_memory_twin():
|
||||
spend: Final = FakeSpendCounters(redis=None, memory={"spend:committed:key:hashed": 1.5})
|
||||
|
||||
committed: Final = await read_admission_committed_spend(
|
||||
spend.store(), EntitySpendCounterKeys(key="spend:key:hashed")
|
||||
)
|
||||
|
||||
assert committed.key == 1.5
|
||||
|
||||
|
||||
def test_no_committed_counter_is_bound_without_a_prometheus_logger():
|
||||
with spend_counter_batch_scope(FakeSpendCounters(redis={}).store()):
|
||||
bind_committed_admission_counters(UserAPIKeyAuth(token="hashed"), end_user_id=None)
|
||||
batch: Final = active_spend_counter_batch()
|
||||
assert batch is not None
|
||||
bound: Final = batch.counter_keys
|
||||
|
||||
assert bound == frozenset()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drop_removes_the_twin_and_leaves_the_enforcement_counter(has_redis: bool):
|
||||
store: Final = _spend_counters(has_redis, {"spend:committed:org:o1": 5.0, "spend:org:o1": 5.0}).store()
|
||||
counter_keys: Final = EntitySpendCounterKeys(org="spend:org:o1")
|
||||
|
||||
await drop_committed_spend(store, ("spend:org:o1",))
|
||||
|
||||
assert (await read_committed_budget_spend(store, counter_keys)).org is None
|
||||
assert (await store.read("spend:org:o1"))[0] == 5.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_drop_never_fails_the_write_that_already_committed():
|
||||
spend: Final = FakeSpendCounters(redis={"spend:committed:org:o1": 5.0}, failing=frozenset({"delete"}))
|
||||
|
||||
await drop_committed_spend(spend.store(), ("spend:org:o1",))
|
||||
|
||||
assert spend.deleted() == frozenset({"spend:committed:org:o1"})
|
||||
|
|
@ -1,12 +1,14 @@
|
|||
import asyncio
|
||||
import collections
|
||||
import contextlib
|
||||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import sqlite3
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from datetime import timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -17,14 +19,17 @@ from fastapi.testclient import TestClient
|
|||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
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 litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_reset
|
||||
from tests.unit.proxy.db.fake_prisma_engine import (
|
||||
ReadTargetOf,
|
||||
statement_timeout_error,
|
||||
statement_timeout_reads,
|
||||
)
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
|
||||
def _default_date_range():
|
||||
|
|
@ -8336,6 +8341,109 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader
|
|||
]
|
||||
|
||||
|
||||
class _ResetTableRows:
|
||||
def __init__(self) -> None:
|
||||
self.update_many: Final = AsyncMock()
|
||||
|
||||
|
||||
class _ResetDb:
|
||||
def __init__(self) -> None:
|
||||
self.litellm_verificationtoken: Final = _ResetTableRows()
|
||||
self.litellm_teamtable: Final = _ResetTableRows()
|
||||
|
||||
async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]:
|
||||
rows: Final = ["hashed"] if "VerificationToken" in sql_query else ["t1"]
|
||||
return [{"id": row} for row in rows if row > after_id][:limit]
|
||||
|
||||
|
||||
class _ResetTransaction:
|
||||
def __init__(self, db: _ResetDb) -> None:
|
||||
self.litellm_verificationtoken: Final = db.litellm_verificationtoken
|
||||
self.litellm_teamtable: Final = db.litellm_teamtable
|
||||
|
||||
async def query_raw(self, sql_query: str) -> list[dict[str, str]]:
|
||||
return []
|
||||
|
||||
|
||||
def _reset_prisma(db: _ResetDb) -> SimpleNamespace:
|
||||
@contextlib.asynccontextmanager
|
||||
async def tx() -> AsyncIterator[_ResetTransaction]:
|
||||
yield _ResetTransaction(db)
|
||||
|
||||
return SimpleNamespace(db=db, tx=tx)
|
||||
|
||||
|
||||
def _committed_spend(monkeypatch: pytest.MonkeyPatch, twins: Sequence[str]) -> FakeSpendCounters:
|
||||
spend: Final = FakeSpendCounters(redis=None, memory=dict.fromkeys(twins, 7.0))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
|
||||
return spend
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("registered_callbacks", "twin_after"),
|
||||
[
|
||||
([MagicMock(spec=PrometheusLogger)], None),
|
||||
([], 7.0),
|
||||
],
|
||||
)
|
||||
async def test_global_spend_reset_drops_the_committed_key_and_team_spend_the_gauges_read(
|
||||
monkeypatch, registered_callbacks, twin_after
|
||||
):
|
||||
twins: Final = ("spend:committed:key:hashed", "spend:committed:team:t1")
|
||||
spend: Final = _committed_spend(monkeypatch, twins)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(_ResetDb()))
|
||||
monkeypatch.setattr(litellm, "callbacks", registered_callbacks)
|
||||
|
||||
response: Final = await global_spend_reset()
|
||||
|
||||
assert response["status"] == "success"
|
||||
assert [spend.memory.get(twin) for twin in twins] == [twin_after, twin_after]
|
||||
|
||||
|
||||
class _PagedResetDb(_ResetDb):
|
||||
def __init__(self, token_count: int) -> None:
|
||||
super().__init__()
|
||||
self._tokens: Final = sorted(f"token-{index}" for index in range(token_count))
|
||||
self.largest_page = 0
|
||||
|
||||
async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]:
|
||||
rows: Final = self._tokens if "VerificationToken" in sql_query else []
|
||||
page: Final = [{"id": row} for row in rows if row > after_id][:limit]
|
||||
self.largest_page = max(self.largest_page, len(page))
|
||||
return page
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("token_count", [999, 1000, 1001, 2500])
|
||||
async def test_global_spend_reset_drops_every_key_twin_reading_the_keys_in_bounded_pages(monkeypatch, token_count):
|
||||
db: Final = _PagedResetDb(token_count)
|
||||
twins: Final = tuple(f"spend:committed:key:token-{index}" for index in range(token_count))
|
||||
spend: Final = _committed_spend(monkeypatch, twins)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(db))
|
||||
monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)])
|
||||
|
||||
response: Final = await global_spend_reset()
|
||||
|
||||
assert response["status"] == "success"
|
||||
assert spend.memory == {}
|
||||
assert db.largest_page <= 1000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_spend_reset_still_succeeds_when_listing_the_committed_counters_fails(monkeypatch):
|
||||
db: Final = _ResetDb()
|
||||
db.query_raw = AsyncMock(side_effect=ConnectionError("db unreachable"))
|
||||
_committed_spend(monkeypatch, ("spend:committed:team:t1",))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(db))
|
||||
monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)])
|
||||
|
||||
response: Final = await global_spend_reset()
|
||||
|
||||
assert response["status"] == "success"
|
||||
db.litellm_teamtable.update_many.assert_awaited_once()
|
||||
|
||||
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -918,6 +918,49 @@ async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_depl
|
|||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("proxy_batch_write_at", [10, 60])
|
||||
async def test_committed_spend_twins_outlive_the_slowest_spend_write_to_the_db(
|
||||
monkeypatch: pytest.MonkeyPatch, proxy_batch_write_at: int
|
||||
) -> None:
|
||||
"""A twin that expires before its charges reach the DB seeds again from a total still missing them."""
|
||||
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
from litellm.constants import SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS
|
||||
from litellm.proxy.proxy_server import _BATCH_WRITE_JITTER_MAX_SECONDS, ProxyStartupEvent
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
spend: Final = FakeSpendCounters()
|
||||
mock_proxy_logging: Final = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.slack_alerting_instance = MagicMock()
|
||||
mock_proxy_logging.db_spend_update_writer = MagicMock()
|
||||
scheduler: Final = AsyncIOScheduler()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", _mock_scheduled_proxy_config()),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
|
||||
patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()),
|
||||
):
|
||||
await ProxyStartupEvent.initialize_scheduled_background_jobs(
|
||||
general_settings={},
|
||||
prisma_client=MagicMock(),
|
||||
proxy_budget_rescheduler_min_time=1,
|
||||
proxy_budget_rescheduler_max_time=2,
|
||||
proxy_batch_write_at=proxy_batch_write_at,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
slowest_db_write: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS
|
||||
assert spend.committed_ttl is not None
|
||||
assert spend.committed_ttl >= slowest_db_write
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -35,9 +35,22 @@ def _is_number(value: object) -> bool:
|
|||
return value is None or (isinstance(value, (int, float)) and not isinstance(value, bool))
|
||||
|
||||
|
||||
_TWIN_SCOPES: Final = frozenset({"key", "team", "user", "org", "end_user"})
|
||||
_WINDOWED_SCOPES: Final = frozenset({"key", "team"})
|
||||
|
||||
|
||||
def _native_store_takes_twin(twin: str) -> bool:
|
||||
scope, _, counter_id = twin.partition(":")
|
||||
_, separator, window = counter_id.rpartition(":window:")
|
||||
return scope in _TWIN_SCOPES and not (scope in _WINDOWED_SCOPES and separator and window)
|
||||
|
||||
|
||||
def _check(counter: str) -> None:
|
||||
if not counter.startswith("spend:"):
|
||||
raise ValueError(f"{counter!r} is not a spend counter")
|
||||
twin: Final = counter.removeprefix("spend:committed:")
|
||||
if twin != counter and not _native_store_takes_twin(twin):
|
||||
raise ValueError(f"{counter!r} is not a spend counter")
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -53,6 +66,11 @@ class DeclaredRead:
|
|||
return self.owner.mget(self.counters)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeSettlement:
|
||||
settled: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeSpendCounters:
|
||||
"""``redis`` is None for a worker without Redis, and ``redis_default`` answers reads of a key it lacks. ``fail``
|
||||
|
|
@ -68,11 +86,12 @@ class FakeSpendCounters:
|
|||
failure: str = "transport"
|
||||
redis_default: float | None = None
|
||||
redis_ttl: float = 60
|
||||
committed_ttl: float | None = None
|
||||
batch_reads: bool = True
|
||||
read_hook: Callable[[str], Awaitable[None]] | None = None
|
||||
on_delete: Callable[[tuple[str, ...]], None] | None = None
|
||||
deletes: list[tuple[str, ...]] = field(default_factory=list)
|
||||
post_call: list[tuple[object, tuple[SpendIncrement, ...]]] = field(default_factory=list)
|
||||
post_call: list[tuple[object, tuple[SpendIncrement, ...], FakeSettlement]] = field(default_factory=list)
|
||||
declared: list[DeclaredRead] = field(default_factory=list)
|
||||
|
||||
def store(self) -> SpendCounterStore:
|
||||
|
|
@ -86,6 +105,9 @@ class FakeSpendCounters:
|
|||
def redis_ttl_seconds(self) -> float | None:
|
||||
return None if self.redis is None else self.redis_ttl
|
||||
|
||||
def set_committed_ttl(self, ttl_seconds: float, /) -> None:
|
||||
self.committed_ttl = ttl_seconds
|
||||
|
||||
def attach(self, client: NativeRedisClient, /) -> bool:
|
||||
if self.redis is not None:
|
||||
return False
|
||||
|
|
@ -139,17 +161,19 @@ class FakeSpendCounters:
|
|||
self.commands.append(f"PIPELINE {' '.join(counter for counter, _ in increments)}")
|
||||
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]) -> FakeSettlement | None:
|
||||
if self.redis is None:
|
||||
return False
|
||||
self.post_call.append((request, tuple(increments)))
|
||||
return True
|
||||
return None
|
||||
settlement: Final = FakeSettlement()
|
||||
self.post_call.append((request, tuple(increments), settlement))
|
||||
return settlement
|
||||
|
||||
def settle_post_call(self) -> None:
|
||||
"""What the native task does once the post-call batch goes out."""
|
||||
pending: Final = tuple(self.post_call)
|
||||
self.post_call.clear()
|
||||
for _, increments in pending:
|
||||
for _, increments, settlement in pending:
|
||||
settlement.settled = True
|
||||
if self.fail:
|
||||
for counter, _ in increments:
|
||||
self.memory.pop(counter, None)
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from litellm.redis.stores.connection import RedisStores, connect
|
|||
from litellm.redis.stores.spend_counters import (
|
||||
SPEND_COUNTERS_TARGET,
|
||||
SpendCounterStore,
|
||||
committed_counter,
|
||||
end_user_counter,
|
||||
is_unavailable,
|
||||
key_counter,
|
||||
|
|
@ -28,7 +29,7 @@ from litellm.redis.store_settings import RedisStoreSettings
|
|||
from litellm.rust_bridge._native import NativeRedisClient, NativeRedisStoreError, NativeSpendCounters, SpendRead
|
||||
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
|
||||
|
||||
from .spend_counters_double import FakeSpendCounters
|
||||
from .spend_counters_double import FakeSettlement, FakeSpendCounters
|
||||
|
||||
_KEYS: Final = (
|
||||
(
|
||||
|
|
@ -46,6 +47,11 @@ _KEYS: Final = (
|
|||
(project_counter("proj-1"), "spend:project:proj-1"),
|
||||
(model_access_group_counter("gpt-group"), "spend:model_access_group:gpt-group"),
|
||||
(shadow_eval_counter("job-7"), "spend:shadow_eval:job-7"),
|
||||
(committed_counter(key_counter("88dc28d0")), "spend:committed:key:88dc28d0"),
|
||||
(committed_counter(team_counter("team-a")), "spend:committed:team:team-a"),
|
||||
(committed_counter(user_counter("user@example.com")), "spend:committed:user:user@example.com"),
|
||||
(committed_counter(org_counter("org-1")), "spend:committed:org:org-1"),
|
||||
(committed_counter(end_user_counter("customer é 1")), "spend:committed:end_user:customer é 1"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -81,7 +87,17 @@ def test_a_read_decodes_what_redis_cache_decodes(wire: str, spend: float) -> Non
|
|||
assert type(decoded) is type(spend)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", ["spend_db_floor:spend:key:k", "virtual_key_spend:k:gpt:1d", "spend:unknown:k"])
|
||||
@pytest.mark.parametrize(
|
||||
"key",
|
||||
[
|
||||
"spend_db_floor:spend:key:k",
|
||||
"virtual_key_spend:k:gpt:1d",
|
||||
"spend:unknown:k",
|
||||
"spend:committed:key:k:window:1d",
|
||||
"spend:committed:tag:prod",
|
||||
"spend:committed:committed:key:k",
|
||||
],
|
||||
)
|
||||
def test_the_native_store_refuses_keys_outside_the_spend_counter_keyspace(key: str) -> None:
|
||||
with pytest.raises(ValueError, match="a spend counter is a spend:\* key"):
|
||||
NativeSpendCounters(None).get_local(key)
|
||||
|
|
@ -153,7 +169,7 @@ 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) -> FakeSettlement | None:
|
||||
targets.append(current_service_target())
|
||||
return super().increment_many_post_call(request, increments) # pyright: ignore[reportArgumentType] # recorded as given
|
||||
|
||||
|
|
@ -165,8 +181,8 @@ async def test_post_call_increments_go_on_the_requests_native_post_call_batch_un
|
|||
inside: Final = store.increment_many_post_call((("spend:key:k", 1.0),))
|
||||
empty: Final = store.increment_many_post_call(())
|
||||
|
||||
assert (outside, inside, empty) == (False, True, False)
|
||||
assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),))]
|
||||
assert (outside, empty) == (None, None)
|
||||
assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),), inside)]
|
||||
assert targets == [SPEND_COUNTERS_TARGET]
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue