From fc4682727e9d77ee585ead130c74bb277a7e23c7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 3 Nov 2026 15:55:33 +0900 Subject: [PATCH] feat(redis): move the spend counter keyspace onto the Rust store Every spend:* reader and writer (auth budget checks, reservations and lease renewal, post-call increments, reseed, resets, key/team/user/tag endpoints, shadow eval) now goes through litellm/redis/stores/spend_counters.py, backed by litellm-redis-store's spend_counters module. Keys, value bytes (json.dumps for SET, str() for the set-max script), TTLs and the set-max Lua script match what main writes, pinned by matching inline cases in the Rust and Python tests. The auth budget read rides the request's native pre-call batch through a native flush hook, so it leaves in the same pipeline as the armed cooldown read, and post-call increments ride the native post-call batch. The auth-cache invalidation subscriber forwards spend:* broadcasts to the native memory tier. Co-Authored-By: Claude Opus 5.5 --- litellm-rust/crates/python-bridge/src/lib.rs | 7 + .../python-bridge/src/redis/store/mod.rs | 1 + .../src/redis/store/spend_counters.rs | 329 +++++++ litellm-rust/crates/redis-store/src/lib.rs | 1 + .../crates/redis-store/src/spend_counters.rs | 548 ++++++++++++ .../redis-store/tests/spend_counters.rs | 484 ++++++++++ litellm/integrations/shadow_eval_logger.py | 5 +- litellm/proxy/auth/auth_checks.py | 90 +- litellm/proxy/auth/fallback_budget.py | 8 +- litellm/proxy/auth/user_api_key_auth.py | 15 +- .../auth_cache_invalidation_pubsub.py | 15 +- .../proxy/common_utils/reset_budget_job.py | 105 +-- .../proxy/common_utils/user_api_key_cache.py | 15 - litellm/proxy/db/spend_counter_reseed.py | 127 +-- .../proxy/hooks/proxy_track_cost_callback.py | 4 +- .../internal_user_endpoints.py | 5 +- .../key_management_endpoints.py | 60 +- .../tag_management_endpoints.py | 5 +- litellm/proxy/proxy_server.py | 232 ++--- .../spend_tracking/budget_reservation.py | 48 +- .../spend_tracking/spend_counter_batch.py | 115 +-- litellm/redis/stores/spend_counters.py | 225 +++++ litellm/rust_bridge/_native.pyi | 31 + .../test_request_redis_batch_post_call.py | 75 +- .../unit/llms/anthropic/messages/conftest.py | 3 + .../messages/test_caller_model_checks.py | 4 +- ...st_auth_checks_object_access_and_lookup.py | 92 +- .../auth/test_model_access_group_budgets.py | 6 +- .../test_user_api_key_auth_request_flow.py | 21 +- .../test_auth_cache_invalidation_pubsub.py | 31 + .../common_utils/test_reset_budget_job.py | 117 +-- .../proxy/db/test_spend_counter_reseed.py | 25 +- .../hooks/test_proxy_track_cost_callback.py | 30 +- .../test_key_management_endpoints.py | 39 +- .../test_team_endpoints.py | 25 +- ...t_budget_reservation_release_middleware.py | 45 +- .../test_success_handler.py | 10 +- .../proxy/proxy_server/test_spend_counters.py | 642 +++++--------- .../spend_tracking/test_budget_reservation.py | 31 +- .../test_budget_reservation_redis_failure.py | 50 +- .../spend_tracking/test_cost_finalization.py | 30 +- .../test_spend_counter_batch.py | 400 ++++----- tests/unit/proxy/test_budget_reservation.py | 611 ++++--------- ...test_proxy_server_endpoints_and_startup.py | 829 ++++++------------ .../unit/proxy/test_redis_auth_cache_flag.py | 54 +- .../redis/stores/spend_counters_double.py | 245 ++++++ .../unit/redis/stores/test_spend_counters.py | 181 ++++ 47 files changed, 3569 insertions(+), 2502 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs create mode 100644 litellm-rust/crates/redis-store/src/spend_counters.rs create mode 100644 litellm-rust/crates/redis-store/tests/spend_counters.rs create mode 100644 litellm/redis/stores/spend_counters.py create mode 100644 tests/unit/redis/stores/spend_counters_double.py create mode 100644 tests/unit/redis/stores/test_spend_counters.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 896bebcf52b..d0f2230b951 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -47,6 +47,10 @@ mod _native { #[pymodule_export] use crate::redis::store::router_usage::NativeRouterUsage; #[pymodule_export] + use crate::redis::store::spend_counters::{ + NativeDeclaredSpendCounters, NativeSpendCounterError, NativeSpendCounters, + }; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -128,6 +132,9 @@ mod tests { "NativeDeclaredCooldowns", "NativeRouterCooldowns", "NativeRouterUsage", + "NativeDeclaredSpendCounters", + "NativeSpendCounterError", + "NativeSpendCounters", "native_redis_breaker_states", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs index b33e1bd096f..cee21643c9f 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -2,6 +2,7 @@ mod policy; pub(crate) mod request; pub(crate) mod router_cooldowns; pub(crate) mod router_usage; +pub(crate) mod spend_counters; use litellm_host_python::PythonContext; use litellm_redis_store::Number; diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs new file mode 100644 index 00000000000..cf573858d6a --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -0,0 +1,329 @@ +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::Duration, +}; + +use litellm_redis_store::spend_counters::{ + DeclaredReads, Error, LostSeed, Reading, SpendCounter, SpendCounters, +}; +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::PyDict, +}; + +use super::{PyNumber, call, request::NativeRequestRedis, store_error}; +use crate::execution::run_async_value; +use crate::redis::client::NativeRedisClient; + +pyo3::create_exception!( + _native, + NativeSpendCounterError, + pyo3::exceptions::PyException, + "A spend counter operation failed. Args are (kind, message), kind one of \"breaker_open\", \"timeout\", \"transport\", \"rejected\", \"decode\" or \"memory\"." +); + +fn spend_error(error: Error) -> PyErr { + use litellm_redis_client::Error as RedisError; + let kind = match &error { + Error::Redis(RedisError::BreakerOpen) => "breaker_open", + Error::Redis(RedisError::Transport(error)) if error.is_timeout() => "timeout", + Error::Redis(RedisError::Transport(_)) => "transport", + Error::Redis(RedisError::Server(_)) => "rejected", + Error::Redis(RedisError::Decode(_)) | Error::Decode(_) => "decode", + Error::Memory(_) => "memory", + }; + NativeSpendCounterError::new_err((kind, litellm_error::Report::new(&error).to_string())) +} + +/// A `spend:*` key, refused when it names anything outside the spend counter keyspace. +pub(crate) struct PySpendCounter(SpendCounter); + +impl<'a, 'py> FromPyObject<'a, 'py> for PySpendCounter { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let key: String = value.extract()?; + SpendCounter::parse(&key) + .map(Self) + .map_err(|error| PyValueError::new_err(error.to_string())) + } +} + +fn counters(counters: Vec) -> Vec { + counters.into_iter().map(|counter| counter.0).collect() +} + +fn increments( + increments: Vec<(PySpendCounter, PyNumber)>, +) -> Vec<(SpendCounter, litellm_redis_store::Number)> { + increments + .into_iter() + .map(|(counter, amount)| (counter.0, amount.0)) + .collect() +} + +fn ttl(seconds: Option) -> Option { + seconds.map(|seconds| Duration::from_secs_f64(seconds.max(0.0))) +} + +fn numbers(values: Vec>) -> Vec> { + values + .into_iter() + .map(|value| value.map(PyNumber)) + .collect() +} + +/// The proxy's spend counters, Redis first with this worker's memory copy behind it. +#[pyclass(frozen, name = "NativeSpendCounters")] +pub(crate) struct NativeSpendCounters(Arc); + +#[pymethods] +impl NativeSpendCounters { + #[new] + #[pyo3(signature = (client, redis_ttl_seconds = 60))] + fn new(client: Option>, redis_ttl_seconds: u64) -> Self { + Self(Arc::new(SpendCounters::new( + client.map(|client| client.0.clone()), + Duration::from_secs(redis_ttl_seconds), + ))) + } + + /// `("redis" | "memory", value)`: where the answer came from, a Redis miss included. + fn read<'py>(&self, py: Python<'py>, counter: PySpendCounter) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_get_cache")?; + run_async_value(py, async move { + let reading = scope + .scope(store.read(&counter.0)) + .await + .map_err(store_error)?; + Ok(match reading { + Reading::Redis(value) => ("redis", value.map(PyNumber)), + Reading::Memory(value) => ("memory", value.map(PyNumber)), + }) + }) + } + + fn read_many<'py>( + &self, + py: Python<'py>, + counters: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let counters = self::counters(counters); + let scope = call(py, "async_batch_get_cache")?; + run_async_value(py, async move { + Ok(scope.scope(store.read_many(&counters)).await.map(numbers)) + }) + } + + fn declare_read_many( + &self, + request: &NativeRequestRedis, + counters: Vec, + ) -> Option { + self.0 + .declare_read_many_pre_call(&request.0, &self::counters(counters)) + .map(|declared| NativeDeclaredSpendCounters(Mutex::new(Some(declared)))) + } + + fn increment<'py>( + &self, + py: Python<'py>, + counter: PySpendCounter, + amount: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_increment")?; + run_async_value(py, async move { + scope + .scope(store.increment(&counter.0, amount.0)) + .await + .map(PyNumber) + .map_err(spend_error) + }) + } + + fn increment_many<'py>( + &self, + py: Python<'py>, + increments: Vec<(PySpendCounter, PyNumber)>, + ) -> PyResult> { + let store = self.0.clone(); + let increments = self::increments(increments); + let scope = call(py, "async_increment_pipeline")?; + run_async_value(py, async move { + let totals = scope + .scope(store.increment_many(&increments)) + .await + .map_err(spend_error)?; + Ok(totals.into_iter().map(PyNumber).collect::>()) + }) + } + + /// Declares the increments on the request's post-call batch; `False` when there is no Redis. + fn increment_many_post_call( + &self, + request: &NativeRequestRedis, + increments: Vec<(PySpendCounter, PyNumber)>, + ) -> bool { + self.0 + .increment_many_post_call(&request.0, &self::increments(increments)) + } + + fn set_max<'py>( + &self, + py: Python<'py>, + counter: PySpendCounter, + spend: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_set_max")?; + run_async_value(py, async move { + scope + .scope(store.set_max(&counter.0, spend.0)) + .await + .map(|stored| stored.map(PyNumber)) + .map_err(spend_error) + }) + } + + #[pyo3(signature = (counter, spend, *, increment_if_lost))] + fn seed<'py>( + &self, + py: Python<'py>, + counter: PySpendCounter, + spend: PyNumber, + increment_if_lost: bool, + ) -> PyResult> { + let store = self.0.clone(); + let lost = if increment_if_lost { + LostSeed::Increment + } else { + LostSeed::Answer + }; + let scope = call(py, "async_set_cache")?; + run_async_value(py, async move { + scope + .scope(store.seed(&counter.0, spend.0, lost)) + .await + .map(PyNumber) + .map_err(spend_error) + }) + } + + #[pyo3(signature = (counter, spend, ttl_seconds = None))] + fn set<'py>( + &self, + py: Python<'py>, + counter: PySpendCounter, + spend: PyNumber, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_set_cache")?; + run_async_value(py, async move { + scope + .scope(store.set(&counter.0, spend.0, ttl(ttl_seconds))) + .await + .map_err(spend_error) + }) + } + + fn delete<'py>( + &self, + py: Python<'py>, + counters: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let counters = self::counters(counters); + let scope = call(py, "async_delete_cache")?; + run_async_value(py, async move { + scope + .scope(store.delete(&counters)) + .await + .map_err(spend_error) + }) + } + + fn refresh_ttl<'py>( + &self, + py: Python<'py>, + counter: PySpendCounter, + ) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_refresh_ttl")?; + run_async_value(py, async move { + scope + .scope(store.refresh_ttl(&counter.0)) + .await + .map_err(spend_error) + }) + } + + fn get_local(&self, counter: PySpendCounter) -> PyResult> { + Ok(self + .0 + .get_local(&counter.0) + .map_err(store_error)? + .map(PyNumber)) + } + + #[pyo3(signature = (counter, spend, ttl_seconds = None))] + fn set_local( + &self, + counter: PySpendCounter, + spend: PyNumber, + ttl_seconds: Option, + ) -> PyResult<()> { + self.0 + .set_local(&counter.0, spend.0, ttl(ttl_seconds)) + .map_err(store_error) + } + + fn delete_local(&self, counter: PySpendCounter) -> PyResult<()> { + self.0.delete_local(&counter.0).map_err(store_error) + } + + fn increment_local(&self, counter: PySpendCounter, amount: PyNumber) -> PyResult { + self.0 + .increment_local(&counter.0, amount.0) + .map(PyNumber) + .map_err(store_error) + } + + fn flush_local(&self) -> PyResult<()> { + self.0.flush_local().map_err(store_error) + } + + fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { + let snapshot = PyDict::new(py); + for (key, spend) in self.0.snapshot().map_err(store_error)? { + snapshot.set_item(key, PyNumber(spend))?; + } + Ok(snapshot) + } +} + +/// A `read_many` declared on a request's pre-call batch. `resolve` answers it once. +#[pyclass(frozen, name = "NativeDeclaredSpendCounters")] +pub(crate) struct NativeDeclaredSpendCounters(Mutex>); + +#[pymethods] +impl NativeDeclaredSpendCounters { + fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { + let declared = self + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .ok_or_else(|| { + PyRuntimeError::new_err("this spend counter read was already resolved") + })?; + let scope = call(py, "async_batch_get_cache")?; + run_async_value(py, async move { + Ok(scope.scope(declared.resolve()).await.map(numbers)) + }) + } +} diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs index c87f54163b3..87749369287 100644 --- a/litellm-rust/crates/redis-store/src/lib.rs +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -3,6 +3,7 @@ mod number; mod request; pub mod router_cooldowns; pub mod router_usage; +pub mod spend_counters; mod tier; pub use codec::DecodeError; diff --git a/litellm-rust/crates/redis-store/src/spend_counters.rs b/litellm-rust/crates/redis-store/src/spend_counters.rs new file mode 100644 index 00000000000..261b8e8ec51 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/spend_counters.rs @@ -0,0 +1,548 @@ +use std::{ + sync::{Arc, LazyLock, Mutex, MutexGuard, PoisonError}, + time::Duration, +}; + +use futures_util::future::join_all; +use litellm_cache::Error as MemoryError; +use litellm_cache_memory::InMemoryCache; +use litellm_python_compat::{json, repr::to_str}; +use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient, Script}; + +use crate::tier::{Clock, system_clock}; +use crate::{DecodeError, Number, RequestRedis, codec::python_value}; + +/// `InMemoryCache`'s defaults, which `spend_counter_cache` runs its memory tier with. +const MEMORY_ENTRIES: usize = 200; +pub const MEMORY_TTL: Duration = Duration::from_secs(600); +/// `spend_counter_cache`'s `default_in_memory_ttl`, which only its memory-only increment applies. +pub const LOCAL_INCREMENT_TTL: Duration = Duration::from_secs(60); +/// `RedisCache`'s TTL when `litellm.default_redis_ttl` is unset. +pub const REDIS_TTL: Duration = Duration::from_secs(60); + +/// `RedisCache.async_set_max`, byte for byte: a counter only ever rises to the authoritative spend. +static SET_MAX: LazyLock